Skip to content
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
33 commits
Select commit Hold shift + click to select a range
b00480c
docs(mcp): design hosted MCP server
AchoArnold Sep 3, 2026
4c99540
docs(mcp): add implementation plan
AchoArnold Sep 3, 2026
8a5a4c8
feat(api): trust scoped MCP tokens
AchoArnold Sep 3, 2026
97d9786
feat(api): add incoming message endpoint
AchoArnold Sep 3, 2026
bf06243
feat(mcp): add service foundation
AchoArnold Sep 3, 2026
043f66f
fix(mcp): drop firebase SDK, make KeySet config one-shot
AchoArnold Sep 3, 2026
2e17e75
feat(mcp): add OAuth state and metadata
AchoArnold Sep 3, 2026
a748407
fix(mcp): harden OAuth metadata fetching
AchoArnold Sep 3, 2026
dfcf5e2
feat(mcp): add Firebase OAuth flow
AchoArnold Sep 3, 2026
e4ab70f
fix(mcp): harden OAuth authorization flow
AchoArnold Sep 3, 2026
505ecb2
feat(mcp): add httpSMS API client
AchoArnold Sep 3, 2026
5e4d7fc
fix(mcp): redact API query traces
AchoArnold Sep 3, 2026
86afa17
feat(mcp): add messaging tools
AchoArnold Sep 3, 2026
a13bc7d
feat(mcp): add API key tools
AchoArnold Sep 3, 2026
70617c2
fix(mcp): mark rotated keys sensitive
AchoArnold Sep 3, 2026
edec15f
feat(mcp): assemble hosted server
AchoArnold Sep 3, 2026
1786941
fix(mcp): harden server assembly
AchoArnold Sep 3, 2026
1740811
fix(mcp): rate limit rotation prompts
AchoArnold Sep 3, 2026
b980af0
chore(mcp): add Cloud Run deployment
AchoArnold Sep 4, 2026
f69a96f
fix(mcp): clarify deployment defaults
AchoArnold Sep 4, 2026
b95789a
test(mcp): add full integration suite
AchoArnold Sep 4, 2026
3bb6025
fix(tests): make MCP integration deterministic
AchoArnold Sep 4, 2026
ae19eb3
fix(tests): validate rate limit success path
AchoArnold Sep 4, 2026
318be45
ci(mcp): gate deploys on MCP tests
AchoArnold Sep 4, 2026
a44f222
Merge remote-tracking branch 'origin/main' into feat/mcp-server
AchoArnold Sep 4, 2026
3a2391c
fix(auth): bound token metadata caches
AchoArnold Sep 4, 2026
914c11a
refactor(mcp): reuse thread message API
AchoArnold Sep 7, 2026
c2556b3
fix(auth): enable production delegation
AchoArnold Sep 30, 2026
83be26f
fix(mcp): harden OAuth and message reads
AchoArnold Sep 30, 2026
49d24ee
test(mcp): isolate refresh token families
AchoArnold Sep 30, 2026
250fabb
feat(mcp): trace protocol operations
AchoArnold Sep 30, 2026
7cc3e30
fix(mcp): redact transport error queries
AchoArnold Oct 1, 2026
06840e2
test: wait for rate-limit scheduling
AchoArnold Oct 1, 2026
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
feat(mcp): trace protocol operations
Export correlated MCP traces and redacted structured logs to Axiom through SDK receiving and sending middleware.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot-Session: 9ff1e38a-b018-4cf7-a5e9-5044a2efd03c
  • Loading branch information
AchoArnold committed Sep 30, 2026
commit 250fabbd70a48bc14b9cd08e9b291714ef6e7e3f
31 changes: 30 additions & 1 deletion mcp/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,10 @@ absent. In `production` (`ENV=production`), every URL-valued setting must use
| `SEND_TOOLS_PER_MINUTE` | `30` | Per-user rate limit for `send_sms`. |
| `KEY_CREATES_PER_HOUR` | `10` | Per-user rate limit for `create_phone_api_key`. |
| `KEY_ROTATIONS_PER_HOUR` | `3` | Per-user rate limit for `rotate_user_api_key`. |
| `AXIOM_TOKEN` | unset | Axiom ingest token. When set with `AXIOM_DATASET_EVENTS`, structured logs and OpenTelemetry traces are sent to Axiom as well as stdout. Cloud Build maps this from the `axiom-token` Secret Manager secret. |
| `AXIOM_DATASET_EVENTS` | unset | Axiom dataset for logs and traces. Configure this directly on the MCP Cloud Run service; deploys preserve it with `--update-env-vars`. |
| `AXIOM_OTLP_ENDPOINT` | `us-east-1.aws.edge.axiom.co` | Override the Axiom OTLP/HTTP endpoint, for example for another Axiom region. |
| `OTEL_EXPORTER_OTLP_ENDPOINT` / `OTEL_EXPORTER_OTLP_TRACES_ENDPOINT` | unset | Standard OpenTelemetry trace-export destination. When either is set it takes precedence over the Axiom trace exporter; all standard OTLP/HTTP headers and TLS variables remain supported. |

None of the values above have a safe committed default for secrets: `.env`
files, `*.pem`/`*.key` files, and anything matching `mcp/.dockerignore` must
Expand Down Expand Up @@ -113,6 +117,28 @@ Cloud Run supplies the listen port through `PORT`; the container's `EXPOSE
overridden in Cloud Run, `--port` in `cloudbuild.yaml` must change to match,
or the health check will fail and the revision will never become ready.

## Observability

The server emits nested OpenTelemetry spans for:

- every inbound HTTP request;
- every inbound MCP method, including tool calls and prompt requests;
- every outbound MCP method sent from the server to the client, including
sampling and elicitation requests;
- every outbound call to the httpSMS API.

MCP spans include only the protocol method, direction, outcome, and the
tool/prompt name where applicable. Tool arguments, prompt text, sampled
messages, resource URIs, response bodies, OAuth values, API keys, and bearer
tokens are never attached to spans or logs. Structured HTTP and MCP logs carry
the matching `trace_id` and `span_id`, allowing an Axiom query to move between
an event and its trace.

The official `modelcontextprotocol/go-sdk` does not currently provide a
built-in OpenTelemetry integration. It exposes `AddReceivingMiddleware` and
`AddSendingMiddleware` as its tracing/metrics integration points; this server
instruments both directions through those APIs.

## Cloud Build invocation

`cloudbuild.yaml` mirrors `api/cloudbuild.yaml`: it builds `mcp/Dockerfile`
Expand All @@ -131,7 +157,8 @@ wired, scoped to changes under `mcp/`.

Non-sensitive configuration (`ENV`, `MCP_BASE_URL`, `HTTPSMS_API_URL`,
`FIREBASE_PROJECT_ID`, `FIREBASE_AUTH_DOMAIN`) is passed with
`--set-env-vars` from `cloudbuild.yaml` substitutions. **Secrets are never
`--update-env-vars` from `cloudbuild.yaml` substitutions so manually configured
values such as `AXIOM_DATASET_EVENTS` are preserved. **Secrets are never
placed in `cloudbuild.yaml` or Cloud Build substitutions.** They are
referenced from Google Secret Manager with `--set-secrets`:

Expand All @@ -141,6 +168,7 @@ referenced from Google Secret Manager with `--set-secrets`:
| `MCP_SIGNING_KEY_ID` | `mcp-signing-key-id` |
| `REDIS_URL` | `mcp-redis-url` |
| `FIREBASE_API_KEY` | `mcp-firebase-api-key` |
| `AXIOM_TOKEN` | `axiom-token` |

Create/update these once with, e.g.:

Expand All @@ -149,6 +177,7 @@ printf '%s' "$PRIVATE_KEY_PEM" | gcloud secrets create mcp-signing-private-key -
printf '%s' "prod-mcp-key-1" | gcloud secrets create mcp-signing-key-id --data-file=-
printf '%s' "$REDIS_URL" | gcloud secrets create mcp-redis-url --data-file=-
printf '%s' "$FIREBASE_API_KEY" | gcloud secrets create mcp-firebase-api-key --data-file=-
printf '%s' "$AXIOM_TOKEN" | gcloud secrets create axiom-token --data-file=-
```

The Cloud Run service's runtime service account needs
Expand Down
4 changes: 2 additions & 2 deletions mcp/cloudbuild.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,8 @@ steps:
--image=us.gcr.io/$PROJECT_ID/$_SERVICE_NAME:$SHORT_SHA \
--region=$_REGION --platform managed --allow-unauthenticated \
--port=8080 \
--set-env-vars="ENV=production,MCP_BASE_URL=$_MCP_BASE_URL,HTTPSMS_API_URL=$_HTTPSMS_API_URL,FIREBASE_PROJECT_ID=$_FIREBASE_PROJECT_ID,FIREBASE_AUTH_DOMAIN=$_FIREBASE_AUTH_DOMAIN" \
--set-secrets="MCP_SIGNING_PRIVATE_KEY=mcp-signing-private-key:latest,MCP_SIGNING_KEY_ID=mcp-signing-key-id:latest,REDIS_URL=mcp-redis-url:latest,FIREBASE_API_KEY=mcp-firebase-api-key:latest"
--update-env-vars="ENV=production,MCP_BASE_URL=$_MCP_BASE_URL,HTTPSMS_API_URL=$_HTTPSMS_API_URL,FIREBASE_PROJECT_ID=$_FIREBASE_PROJECT_ID,FIREBASE_AUTH_DOMAIN=$_FIREBASE_AUTH_DOMAIN" \
--set-secrets="MCP_SIGNING_PRIVATE_KEY=mcp-signing-private-key:latest,MCP_SIGNING_KEY_ID=mcp-signing-key-id:latest,REDIS_URL=mcp-redis-url:latest,FIREBASE_API_KEY=mcp-firebase-api-key:latest,AXIOM_TOKEN=axiom-token:latest"
options:
substitutionOption: ALLOW_LOOSE

Expand Down
7 changes: 6 additions & 1 deletion mcp/go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ go 1.25.0

require (
github.com/alicebob/miniredis/v2 v2.35.0
github.com/axiomhq/axiom-go v0.32.0
github.com/golang-jwt/jwt/v5 v5.3.1
github.com/google/jsonschema-go v0.4.3
github.com/google/uuid v1.6.0
Expand All @@ -15,15 +16,20 @@ require (
go.opentelemetry.io/otel v1.46.0
go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.46.0
go.opentelemetry.io/otel/sdk v1.46.0
go.opentelemetry.io/otel/trace v1.46.0
)

require (
github.com/buger/jsonparser v1.1.2 // indirect
github.com/cenkalti/backoff/v4 v4.3.0 // indirect
github.com/cenkalti/backoff/v5 v5.0.3 // indirect
github.com/cespare/xxhash/v2 v2.3.0 // indirect
github.com/felixge/httpsnoop v1.1.0 // indirect
github.com/go-logr/logr v1.4.4 // indirect
github.com/go-logr/stdr v1.2.2 // indirect
github.com/google/go-querystring v1.2.0 // indirect
github.com/grpc-ecosystem/grpc-gateway/v2 v2.30.0 // indirect
github.com/klauspost/compress v1.18.5 // indirect
github.com/mattn/go-colorable v0.1.14 // indirect
github.com/mattn/go-isatty v0.0.20 // indirect
github.com/segmentio/asm v1.1.3 // indirect
Expand All @@ -33,7 +39,6 @@ require (
go.opentelemetry.io/auto/sdk v1.2.1 // indirect
go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.46.0 // indirect
go.opentelemetry.io/otel/metric v1.46.0 // indirect
go.opentelemetry.io/otel/trace v1.46.0 // indirect
go.opentelemetry.io/proto/otlp v1.11.0 // indirect
go.uber.org/atomic v1.11.0 // indirect
go.yaml.in/yaml/v3 v3.0.5 // indirect
Expand Down
19 changes: 19 additions & 0 deletions mcp/go.sum
Original file line number Diff line number Diff line change
@@ -1,9 +1,15 @@
github.com/alicebob/miniredis/v2 v2.35.0 h1:QwLphYqCEAo1eu1TqPRN2jgVMPBweeQcR21jeqDCONI=
github.com/alicebob/miniredis/v2 v2.35.0/go.mod h1:TcL7YfarKPGDAthEtl5NBeHZfeUQj6OXMm/+iu5cLMM=
github.com/axiomhq/axiom-go v0.32.0 h1:aRpbqUAn01hY8aJXQftvWHyXfnrNB2KzN5ZquBWvFcE=
github.com/axiomhq/axiom-go v0.32.0/go.mod h1:3Gmr5M4tINm7Ti00GVfzAduO92Uhd0pghr4ZehIhFxc=
github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs=
github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c=
github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA=
github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0=
github.com/buger/jsonparser v1.1.2 h1:frqHqw7otoVbk5M8LlE/L7HTnIq2v9RX6EJ48i9AxJk=
github.com/buger/jsonparser v1.1.2/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0=
github.com/cenkalti/backoff/v4 v4.3.0 h1:MyRJ/UdXutAwSAT+s3wNd7MfTIcy71VQueUuFK343L8=
github.com/cenkalti/backoff/v4 v4.3.0/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE=
github.com/cenkalti/backoff/v5 v5.0.3 h1:ZN+IMa753KfX5hd8vVaMixjnqRZ3y8CuJKRKj1xcsSM=
github.com/cenkalti/backoff/v5 v5.0.3/go.mod h1:rkhZdG3JZukswDf7f0cwqPNk4K0sa+F97BxZthm/crw=
github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
Expand All @@ -19,14 +25,19 @@ github.com/golang-jwt/jwt/v5 v5.3.1 h1:kYf81DTWFe7t+1VvL7eS+jKFVWaUnK9cB1qbwn63Y
github.com/golang-jwt/jwt/v5 v5.3.1/go.mod h1:fxCRLWMO43lRc8nhHWY6LGqRcf+1gQWArsqaEUEa5bE=
github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek=
github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=
github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
github.com/google/go-querystring v1.2.0 h1:yhqkPbu2/OH+V9BfpCVPZkNmUXhb2gBxJArfhIxNtP0=
github.com/google/go-querystring v1.2.0/go.mod h1:8IFJqpSRITyJ8QhQ13bmbeMBDfmeEJZD5A0egEOmkqU=
github.com/google/jsonschema-go v0.4.3 h1:/DBOLZTfDow7pe2GmaJNhltueGTtDKICi8V8p+DQPd0=
github.com/google/jsonschema-go v0.4.3/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE=
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/grpc-ecosystem/grpc-gateway/v2 v2.30.0 h1:/Tnpcb2E0Pz/tN9s3bfEY2Q8ePCEX9iuS+cneUwncnw=
github.com/grpc-ecosystem/grpc-gateway/v2 v2.30.0/go.mod h1:zOBXOsUaBSjKgmH4OGzV1esUpR3oUSCPYVd2cUBjKYY=
github.com/klauspost/compress v1.18.5 h1:/h1gH5Ce+VWNLSWqPzOVn6XBO+vJbCNGvjoaGBFW2IE=
github.com/klauspost/compress v1.18.5/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ=
github.com/klauspost/cpuid/v2 v2.2.10 h1:tBs3QSyvjDyFTq3uoc/9xFpCuOsJQFNPiAhYdw2skhE=
github.com/klauspost/cpuid/v2 v2.2.10/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0=
github.com/mattn/go-colorable v0.1.14 h1:9A9LHSqF/7dyVVX6g0U9cwm9pG3kP9gSzcuIPHPsaIE=
Expand All @@ -45,6 +56,14 @@ github.com/segmentio/encoding v0.5.4 h1:OW1VRern8Nw6ITAtwSZ7Idrl3MXCFwXHPgqESYfv
github.com/segmentio/encoding v0.5.4/go.mod h1:HS1ZKa3kSN32ZHVZ7ZLPLXWvOVIiZtyJnO1gPH1sKt0=
github.com/stretchr/testify v1.12.1 h1:EuwCh5fleGS7H32xRwO3wRGT7DxrDhLAT6FF8MpWDWE=
github.com/stretchr/testify v1.12.1/go.mod h1:MDEgiDPPsNp5cuIrHPPCyornHKgEVbtFUmoNlxoYthg=
github.com/tidwall/gjson v1.18.0 h1:FIDeeyB800efLX89e5a8Y0BNH+LOngJyGrIWxG2FKQY=
github.com/tidwall/gjson v1.18.0/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=
github.com/tidwall/match v1.2.0 h1:0pt8FlkOwjN2fPt4bIl4BoNxb98gGHN2ObFEDkrfZnM=
github.com/tidwall/match v1.2.0/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM=
github.com/tidwall/pretty v1.2.1 h1:qjsOFOWWQl+N3RsoF5/ssm1pHmJJwhjlSbZ51I6wMl4=
github.com/tidwall/pretty v1.2.1/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU=
github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY=
github.com/tidwall/sjson v1.2.5/go.mod h1:Fvgq9kS/6ociJEDnK0Fk1cpYF4FIW6ZF7LAe+6jwd28=
github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4=
github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4=
github.com/yuin/gopher-lua v1.1.1 h1:kYKnWBjvbNP4XLT3+bPEwAXJx262OhaHDWDVOPjL46M=
Expand Down
115 changes: 115 additions & 0 deletions mcp/internal/observability/mcp.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,115 @@
package observability

import (
"context"
"fmt"
"time"

"github.com/modelcontextprotocol/go-sdk/mcp"
"github.com/rs/zerolog"
"go.opentelemetry.io/otel"
"go.opentelemetry.io/otel/attribute"
"go.opentelemetry.io/otel/codes"
"go.opentelemetry.io/otel/trace"
)

const mcpInstrumentationName = "github.com/NdoleStudio/httpsms/mcp"

// MCPMiddleware traces and logs one MCP protocol method in direction. It
// deliberately records only bounded protocol metadata: method, direction,
// outcome, and tool/prompt name. Arguments, prompt text, sampled messages,
// resource URIs, response content, tokens, and request metadata are never
// recorded.
func MCPMiddleware(logger zerolog.Logger, direction string, kind trace.SpanKind) mcp.Middleware {
tracer := otel.Tracer(mcpInstrumentationName)

return func(next mcp.MethodHandler) mcp.MethodHandler {
return func(ctx context.Context, method string, req mcp.Request) (result mcp.Result, err error) {
attributes := []attribute.KeyValue{
attribute.String("rpc.system", "mcp"),
attribute.String("rpc.method", method),
attribute.String("mcp.message.direction", direction),
}
attributes = append(attributes, safeRequestAttributes(req)...)

ctx, span := tracer.Start(
ctx,
"mcp "+method,
trace.WithSpanKind(kind),
trace.WithAttributes(attributes...),
)
defer span.End()

start := time.Now()
span.AddEvent("mcp.request", trace.WithAttributes(
attribute.String("mcp.message.direction", direction),
))

result, err = next(ctx, method, req)

outcome := "success"
if err != nil {
outcome = "error"
errorType := fmt.Sprintf("%T", err)
span.SetAttributes(attribute.String("error.type", errorType))
span.SetStatus(codes.Error, "MCP method failed")
} else {
span.SetStatus(codes.Ok, "")
}
span.SetAttributes(attribute.String("mcp.outcome", outcome))
span.AddEvent("mcp.response", trace.WithAttributes(
attribute.String("mcp.outcome", outcome),
))

event := logger.Info()
if err != nil {
event = logger.Error().Str("error.type", fmt.Sprintf("%T", err))
}
spanContext := span.SpanContext()
if spanContext.IsValid() {
event = event.
Str("trace_id", spanContext.TraceID().String()).
Str("span_id", spanContext.SpanID().String())
}
for _, attr := range attributes {
event = addLogAttribute(event, attr)
}
event.
Str("outcome", outcome).
Dur("duration", time.Since(start)).
Msg("mcp request")

return result, err
}
}
}

func safeRequestAttributes(req mcp.Request) []attribute.KeyValue {
if req == nil || req.GetParams() == nil {
return nil
}

switch params := req.GetParams().(type) {
case *mcp.CallToolParamsRaw:
return []attribute.KeyValue{attribute.String("mcp.tool.name", params.Name)}
case *mcp.GetPromptParams:
return []attribute.KeyValue{attribute.String("mcp.prompt.name", params.Name)}
default:
return nil
}
}

func addLogAttribute(event *zerolog.Event, attr attribute.KeyValue) *zerolog.Event {
switch attr.Value.Type() {
case attribute.STRING:
return event.Str(string(attr.Key), attr.Value.AsString())
case attribute.BOOL:
return event.Bool(string(attr.Key), attr.Value.AsBool())
case attribute.INT64:
return event.Int64(string(attr.Key), attr.Value.AsInt64())
case attribute.FLOAT64:
return event.Float64(string(attr.Key), attr.Value.AsFloat64())
default:
return event
}
}
105 changes: 105 additions & 0 deletions mcp/internal/observability/mcp_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,105 @@
package observability

import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"testing"

"github.com/modelcontextprotocol/go-sdk/mcp"
"github.com/rs/zerolog"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"go.opentelemetry.io/otel"
sdktrace "go.opentelemetry.io/otel/sdk/trace"
"go.opentelemetry.io/otel/sdk/trace/tracetest"
"go.opentelemetry.io/otel/trace"
)

func TestMCPMiddlewareTracesSafeProtocolMetadata(t *testing.T) {
const secret = "secret-message-body"

exporter := tracetest.NewInMemoryExporter()
provider := sdktrace.NewTracerProvider(sdktrace.WithSyncer(exporter))
previousProvider := otel.GetTracerProvider()
otel.SetTracerProvider(provider)
t.Cleanup(func() {
otel.SetTracerProvider(previousProvider)
require.NoError(t, provider.Shutdown(context.Background()))
})

var logs bytes.Buffer
logger := zerolog.New(&logs)
middleware := MCPMiddleware(logger, "receive", trace.SpanKindServer)
handler := middleware(func(context.Context, string, mcp.Request) (mcp.Result, error) {
return &mcp.CallToolResult{}, nil
})
request := &mcp.CallToolRequest{
Params: &mcp.CallToolParamsRaw{
Name: "send_sms",
Arguments: json.RawMessage(fmt.Sprintf(`{"content":%q}`, secret)),
},
}

_, err := handler(context.Background(), "tools/call", request)
require.NoError(t, err)

spans := exporter.GetSpans()
require.Len(t, spans, 1)
span := spans[0]
assert.Equal(t, "mcp tools/call", span.Name)
assert.Equal(t, trace.SpanKindServer, span.SpanKind)
assert.Equal(t, "mcp", spanAttribute(span, "rpc.system"))
assert.Equal(t, "tools/call", spanAttribute(span, "rpc.method"))
assert.Equal(t, "receive", spanAttribute(span, "mcp.message.direction"))
assert.Equal(t, "send_sms", spanAttribute(span, "mcp.tool.name"))
assert.Equal(t, "success", spanAttribute(span, "mcp.outcome"))
assert.NotContains(t, fmt.Sprint(span), secret)

logOutput := logs.String()
assert.Contains(t, logOutput, `"rpc.method":"tools/call"`)
assert.Contains(t, logOutput, `"mcp.tool.name":"send_sms"`)
assert.NotContains(t, logOutput, secret)
}

func TestMCPMiddlewareDoesNotRecordErrorMessages(t *testing.T) {
const secret = "sensitive-downstream-error"

exporter := tracetest.NewInMemoryExporter()
provider := sdktrace.NewTracerProvider(sdktrace.WithSyncer(exporter))
previousProvider := otel.GetTracerProvider()
otel.SetTracerProvider(provider)
t.Cleanup(func() {
otel.SetTracerProvider(previousProvider)
require.NoError(t, provider.Shutdown(context.Background()))
})

var logs bytes.Buffer
middleware := MCPMiddleware(zerolog.New(&logs), "send", trace.SpanKindClient)
handler := middleware(func(context.Context, string, mcp.Request) (mcp.Result, error) {
return nil, errors.New(secret)
})

_, err := handler(context.Background(), "sampling/createMessage", &mcp.CreateMessageRequest{
Params: &mcp.CreateMessageParams{SystemPrompt: secret},
})
require.EqualError(t, err, secret)

spans := exporter.GetSpans()
require.Len(t, spans, 1)
assert.Equal(t, "error", spanAttribute(spans[0], "mcp.outcome"))
assert.Equal(t, "*errors.errorString", spanAttribute(spans[0], "error.type"))
assert.NotContains(t, fmt.Sprint(spans[0]), secret)
assert.NotContains(t, logs.String(), secret)
}

func spanAttribute(span tracetest.SpanStub, key string) string {
for _, attr := range span.Attributes {
if string(attr.Key) == key {
return attr.Value.AsString()
}
}
return ""
}
Loading
Loading