From 494004b8f71ecd48a4f62ffd42748558261f79c7 Mon Sep 17 00:00:00 2001 From: Diego Braga Date: Wed, 11 Nov 2026 20:21:53 +0100 Subject: [PATCH] feat(telemetry): trace each MCP request and tool call under the caller's trace Wrap the MCP HTTP handler in otelhttp, so every request gets a server span that is a child of the incoming W3C traceparent. The standalone SSE stream (GET) is filtered out: it stays open for the whole session. The tools/call span is now named "tools/call " and carries mcp.method.name and gen_ai.tool.name next to the existing mcp.tool.* attributes. It is parented on the server span of the request that made the call. The SDK runs handlers on the session's context, which holds the trace context of the initialize request, so the middleware writes the server span's context into a cloned request's traceparent header and the tool middleware extracts it from CallToolRequest.Extra.Header. No header, argument or result values are added to spans. Part of #102 Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_016RAxvkp7wRXwXDx4Lx8STi Signed-off-by: Diego Braga --- cmd/main.go | 13 ++- cmd/tracing_test.go | 195 +++++++++++++++++++++++++++++++ go.mod | 2 + go.sum | 4 + internal/mcp/mcp.go | 32 ++++- internal/telemetry/middleware.go | 59 ++++++++-- 6 files changed, 288 insertions(+), 17 deletions(-) create mode 100644 cmd/tracing_test.go diff --git a/cmd/main.go b/cmd/main.go index 57e90c1..ee7ebf1 100644 --- a/cmd/main.go +++ b/cmd/main.go @@ -193,7 +193,7 @@ func run(cmd *cobra.Command, args []string) { runStdioServer(ctx, mcpSrv) }() } else { - sseServer := newStreamableHTTPHandler(mcpSrv, sessionIdleTTL) + mcpHandler := newMCPHTTPHandler(mcpSrv, sessionIdleTTL) // Create a mux to handle different routes mux := http.NewServeMux() @@ -242,9 +242,7 @@ func run(cmd *cobra.Command, args []string) { metrics.KagentToolsMCPServerInfo.WithLabelValues(Name, Version, GitCommit, BuildDate, serverMode).Set(1) // Handle all other routes with the MCP server wrapped in telemetry middleware - mux.Handle("/", telemetry.HTTPMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - sseServer.ServeHTTP(w, r) - }))) + mux.Handle("/", mcpHandler) httpServer = &http.Server{ Addr: fmt.Sprintf(":%d", port), @@ -314,6 +312,13 @@ func writeResponse(w http.ResponseWriter, data []byte) error { return err } +// newMCPHTTPHandler builds the MCP HTTP handler: the streamable HTTP transport +// wrapped in the telemetry middleware, which starts a server span per request +// and hands its trace context to the tool-call spans. +func newMCPHTTPHandler(mcpSrv *sdkmcp.Server, idleTTL time.Duration) http.Handler { + return telemetry.HTTPMiddleware(newStreamableHTTPHandler(mcpSrv, idleTTL)) +} + // newStreamableHTTPHandler builds the stateful streamable HTTP transport. // // A session is registered on the first (initialize) POST and released when the diff --git a/cmd/tracing_test.go b/cmd/tracing_test.go new file mode 100644 index 0000000..9343023 --- /dev/null +++ b/cmd/tracing_test.go @@ -0,0 +1,195 @@ +package main + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + mcpserver "github.com/kagent-dev/tools/internal/mcp" + sdkmcp "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/codes" + "go.opentelemetry.io/otel/propagation" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + "go.opentelemetry.io/otel/sdk/trace/tracetest" + "go.opentelemetry.io/otel/trace" +) + +const ( + // Each traceparent stands for a different caller trace, so the test can + // tell which request a span was parented on. + initializeTraceparent = "00-11111111111111111111111111111111-1111111111111111-01" + okCallTraceparent = "00-22222222222222222222222222222222-2222222222222222-01" + toolErrTraceparent = "00-33333333333333333333333333333333-3333333333333333-01" + goErrTraceparent = "00-44444444444444444444444444444444-4444444444444444-01" + closeTraceparent = "00-55555555555555555555555555555555-5555555555555555-01" + + // Values that must never reach a span. + tracingTestBearer = "Bearer tracing-test-credential" + tracingTestArgument = "tracing-test-argument" +) + +type tracingTestInput struct { + Value string `json:"value,omitempty" jsonschema:"an arbitrary value"` +} + +// traceparentTransport sets the traceparent and Authorization headers on every +// request, the way a traced MCP client in an agent framework would. +type traceparentTransport struct { + mu sync.Mutex + traceparent string +} + +func (t *traceparentTransport) set(traceparent string) { + t.mu.Lock() + defer t.mu.Unlock() + t.traceparent = traceparent +} + +func (t *traceparentTransport) RoundTrip(r *http.Request) (*http.Response, error) { + t.mu.Lock() + traceparent := t.traceparent + t.mu.Unlock() + + r = r.Clone(r.Context()) + r.Header.Set("traceparent", traceparent) + r.Header.Set("Authorization", tracingTestBearer) + return http.DefaultTransport.RoundTrip(r) +} + +func mustSpanContext(t *testing.T, traceparent string) trace.SpanContext { + t.Helper() + ctx := propagation.TraceContext{}.Extract(context.Background(), + propagation.MapCarrier{"traceparent": traceparent}) + sc := trace.SpanContextFromContext(ctx) + require.True(t, sc.IsValid(), "invalid traceparent %q", traceparent) + return sc +} + +// TestMCPHTTPTracing drives tool calls through the MCP HTTP handler as run() +// wires it and checks the resulting spans: one server span per request, a child +// of the incoming traceparent, and one "tools/call " span under it. +func TestMCPHTTPTracing(t *testing.T) { + exporter := tracetest.NewInMemoryExporter() + provider := sdktrace.NewTracerProvider(sdktrace.WithSyncer(exporter)) + prevProvider, prevPropagator := otel.GetTracerProvider(), otel.GetTextMapPropagator() + otel.SetTracerProvider(provider) + otel.SetTextMapPropagator(propagation.NewCompositeTextMapPropagator( + propagation.TraceContext{}, propagation.Baggage{})) + t.Cleanup(func() { + otel.SetTracerProvider(prevProvider) + otel.SetTextMapPropagator(prevPropagator) + _ = provider.Shutdown(context.Background()) + }) + + srv := sdkmcp.NewServer(&sdkmcp.Implementation{Name: "tracing", Version: "test"}, nil) + srv.AddReceivingMiddleware(mcpserver.ToolMiddleware()) + mcpserver.AddTool(srv, "test", &sdkmcp.Tool{Name: "tracing_ok", Description: "succeeds"}, + func(context.Context, *sdkmcp.CallToolRequest, tracingTestInput) (*sdkmcp.CallToolResult, mcpserver.TextOutput, error) { + return mcpserver.TextResult("ok") + }) + mcpserver.AddTool(srv, "test", &sdkmcp.Tool{Name: "tracing_tool_error", Description: "returns an error result"}, + func(context.Context, *sdkmcp.CallToolRequest, tracingTestInput) (*sdkmcp.CallToolResult, mcpserver.TextOutput, error) { + return mcpserver.TextError("resource not found") + }) + mcpserver.AddTool(srv, "test", &sdkmcp.Tool{Name: "tracing_go_error", Description: "returns a Go error"}, + func(context.Context, *sdkmcp.CallToolRequest, tracingTestInput) (*sdkmcp.CallToolResult, mcpserver.TextOutput, error) { + return nil, mcpserver.TextOutput{}, errors.New("command failed") + }) + + httpServer := httptest.NewServer(newMCPHTTPHandler(srv, time.Hour)) + t.Cleanup(httpServer.Close) + + rt := &traceparentTransport{traceparent: initializeTraceparent} + client := sdkmcp.NewClient(&sdkmcp.Implementation{Name: "tracing-client", Version: "test"}, nil) + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + session, err := client.Connect(ctx, &sdkmcp.StreamableClientTransport{ + Endpoint: httpServer.URL, + HTTPClient: &http.Client{Transport: rt}, + }, nil) + require.NoError(t, err) + + calls := []struct { + tool string + traceparent string + wantStatus codes.Code + }{ + {"tracing_ok", okCallTraceparent, codes.Ok}, + {"tracing_tool_error", toolErrTraceparent, codes.Error}, + {"tracing_go_error", goErrTraceparent, codes.Error}, + } + for _, c := range calls { + rt.set(c.traceparent) + res, err := session.CallTool(ctx, &sdkmcp.CallToolParams{ + Name: c.tool, + Arguments: map[string]string{"value": tracingTestArgument}, + }) + require.NoError(t, err, c.tool) + require.Equal(t, c.wantStatus == codes.Error, res.IsError, c.tool) + } + + rt.set(closeTraceparent) + require.NoError(t, session.Close()) + // Close waits for in-flight requests, so every server span has ended. + httpServer.Close() + + spans := exporter.GetSpans() + for _, c := range calls { + t.Run(c.tool, func(t *testing.T) { + incoming := mustSpanContext(t, c.traceparent) + + var serverSpans []tracetest.SpanStub + for _, s := range spans { + if s.SpanKind == trace.SpanKindServer && s.Parent.SpanID() == incoming.SpanID() { + serverSpans = append(serverSpans, s) + } + } + require.Len(t, serverSpans, 1, "want one server span for the tools/call request") + server := serverSpans[0] + assert.Equal(t, incoming.TraceID(), server.SpanContext.TraceID()) + assert.True(t, server.Parent.IsRemote(), "server span parent should be the remote caller") + + var toolSpans []tracetest.SpanStub + for _, s := range spans { + if s.Name == "tools/call "+c.tool { + toolSpans = append(toolSpans, s) + } + } + require.Len(t, toolSpans, 1, "want one tool-call span") + tool := toolSpans[0] + assert.Equal(t, incoming.TraceID(), tool.SpanContext.TraceID(), + "tool span must join the tools/call request's trace, not the session's first one") + assert.Equal(t, server.SpanContext.SpanID(), tool.Parent.SpanID(), + "tool span must be a child of the server span") + assert.Equal(t, c.wantStatus, tool.Status.Code) + + attrs := map[string]string{} + for _, kv := range tool.Attributes { + attrs[string(kv.Key)] = kv.Value.Emit() + } + assert.Equal(t, "tools/call", attrs["mcp.method.name"]) + assert.Equal(t, c.tool, attrs["gen_ai.tool.name"]) + }) + } + + // No span may carry request headers, credentials, or tool arguments. + for _, s := range spans { + for _, kv := range s.Attributes { + key := strings.ToLower(string(kv.Key)) + for _, banned := range []string{"header", "authorization", "token"} { + assert.NotContains(t, key, banned, "span %q attribute %q", s.Name, kv.Key) + } + value := kv.Value.Emit() + assert.NotContains(t, value, "tracing-test-credential", "span %q attribute %q", s.Name, kv.Key) + assert.NotContains(t, value, tracingTestArgument, "span %q attribute %q", s.Name, kv.Key) + } + } +} diff --git a/go.mod b/go.mod index 002f29b..fda5d25 100644 --- a/go.mod +++ b/go.mod @@ -15,6 +15,7 @@ require ( github.com/spf13/cobra v1.10.2 github.com/stretchr/testify v1.12.1 github.com/tmc/langchaingo v0.1.14 + go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.72.0 go.opentelemetry.io/otel v1.47.0 go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.47.0 go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.47.0 @@ -65,6 +66,7 @@ require ( github.com/emicklei/go-restful/v3 v3.13.0 // indirect github.com/facebookincubator/nvdtools v0.1.5 // indirect github.com/fatih/color v1.19.0 // indirect + github.com/felixge/httpsnoop v1.1.0 // indirect github.com/francoispqt/gojay v1.2.13 // indirect github.com/fsnotify/fsnotify v1.9.0 // indirect github.com/fxamacker/cbor/v2 v2.9.1 // indirect diff --git a/go.sum b/go.sum index 7df5333..01b3c37 100644 --- a/go.sum +++ b/go.sum @@ -204,6 +204,8 @@ github.com/fatih/set v0.2.1 h1:nn2CaJyknWE/6txyUDGwysr3G5QC6xWB/PtVjPBbeaA= github.com/fatih/set v0.2.1/go.mod h1:+RKtMCH+favT2+3YecHGxcc0b4KyVWA1QWWJUs4E0CI= github.com/felixge/fgprof v0.9.5 h1:8+vR6yu2vvSKn08urWyEuxx75NWPEvybbkBirEpsbVY= github.com/felixge/fgprof v0.9.5/go.mod h1:yKl+ERSa++RYOs32d8K6WEXCB4uXdLls4ZaZPpayhMM= +github.com/felixge/httpsnoop v1.1.0 h1:3YtUj32ZZkqZtt3sZZsClsymw/QDuVfpNhoA31zeORc= +github.com/felixge/httpsnoop v1.1.0/go.mod h1:Zqxgdd+1Rkcz8euOqdr7lqgCRJztwr5hp9vDSi5UZCE= github.com/flynn/go-shlex v0.0.0-20150515145356-3f9db97f8568/go.mod h1:xEzjJPgXI435gkrCt3MPfRiAkVrwSbHsst4LCFVfpJc= github.com/francoispqt/gojay v1.2.13 h1:d2m3sFjloqoIUQU3TsHBgj6qg/BVGlTBeHDUmyJnXKk= github.com/francoispqt/gojay v1.2.13/go.mod h1:ehT5mTG4ua4581f1++1WLG0vPdaA9HaiDsoyrBGkyDY= @@ -769,6 +771,8 @@ go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y= go.opentelemetry.io/contrib/bridges/otelslog v0.21.0 h1:/CdgE52i1OEPHhsGUqrggYNRaZy4v4YgVenGd52dRBI= go.opentelemetry.io/contrib/bridges/otelslog v0.21.0/go.mod h1:9pBdbVbuehIDzxN4accnzzxEUPT8Nq7grbZMo+J1VN4= +go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.72.0 h1:LxwW/9ctSCv+QkE/cLR7M91ZIkXNMqJtEMi1vCw9U8s= +go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.72.0/go.mod h1:tOsftB4SslBwwErVEPaenU2RpThXWPIU8DoJHEC4dyw= go.opentelemetry.io/contrib/instrumentation/runtime v0.72.0 h1:rv8t3kKjyC2/9D99L1MZt4tJSZO4aYUmHus7L//iKKg= go.opentelemetry.io/contrib/instrumentation/runtime v0.72.0/go.mod h1:ZzG+JmzLLvKzx6tMeLrw2l7F8ik2GeRZJALoPBqnM0c= go.opentelemetry.io/contrib/processors/minsev v0.17.0 h1:8Am2/je8uZLXKqXgMgKknKk72aG2bm3aa/npmSCNvhY= diff --git a/internal/mcp/mcp.go b/internal/mcp/mcp.go index 7acb361..b69ee28 100644 --- a/internal/mcp/mcp.go +++ b/internal/mcp/mcp.go @@ -24,6 +24,8 @@ import ( "go.opentelemetry.io/otel" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/codes" + "go.opentelemetry.io/otel/propagation" + "go.opentelemetry.io/otel/trace" ) // NewToolResultText returns a successful text result. @@ -139,8 +141,34 @@ func providerOf(tool string) string { return "" } +// toolSpanParent returns the context to start a tool-call span from. +// +// Over HTTP the SDK runs handlers on the session's context, which carries the +// trace context of the request that opened the session, not of this tools/call. +// The parent is therefore taken from this request's header, where +// telemetry.HTTPMiddleware has written the HTTP server span's context. Any span +// already in ctx is dropped first, so a request without trace context starts a +// new trace instead of joining the session's first one. Without a header (stdio, +// in-memory transports) ctx is used as is. +func toolSpanParent(ctx context.Context, req sdk.Request) context.Context { + if req == nil { + return ctx + } + extra := req.GetExtra() + if extra == nil || extra.Header == nil { + return ctx + } + ctx = trace.ContextWithSpanContext(ctx, trace.SpanContext{}) + return otel.GetTextMapPropagator().Extract(ctx, propagation.HeaderCarrier(extra.Header)) +} + // ToolMiddleware instruments every tools/call with an OTel span and Prometheus // invocation counters. Register once via server.AddReceivingMiddleware. +// +// The span is named "tools/call " and carries the MCP semantic +// convention attributes mcp.method.name and gen_ai.tool.name next to the +// existing mcp.tool.* attributes. Request headers, tool arguments and tool +// results are deliberately not recorded: headers can carry credentials. func ToolMiddleware() sdk.Middleware { return func(next sdk.MethodHandler) sdk.MethodHandler { return func(ctx context.Context, method string, req sdk.Request) (sdk.Result, error) { @@ -155,9 +183,11 @@ func ToolMiddleware() sdk.Middleware { provider := providerOf(toolName) tracer := otel.Tracer("kagent-tools/mcp") - ctx, span := tracer.Start(ctx, "mcp.tool."+toolName) + ctx, span := tracer.Start(toolSpanParent(ctx, req), method+" "+toolName) defer span.End() span.SetAttributes( + attribute.String("mcp.method.name", method), + attribute.String("gen_ai.tool.name", toolName), attribute.String("mcp.tool.name", toolName), attribute.String("mcp.tool.provider", provider), ) diff --git a/internal/telemetry/middleware.go b/internal/telemetry/middleware.go index a00a727..14416af 100644 --- a/internal/telemetry/middleware.go +++ b/internal/telemetry/middleware.go @@ -4,6 +4,7 @@ import ( "context" "net/http" + "go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp" "go.opentelemetry.io/otel" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/codes" @@ -20,22 +21,50 @@ const ( SpanIDKey contextKey = "span_id" ) -// HTTPMiddleware wraps an HTTP handler to propagate OpenTelemetry trace context -// from incoming HTTP headers and to stash a few request headers plus the trace -// and span IDs in the request context. +// HTTPMiddleware wraps the MCP HTTP handler with a server span per request and +// propagates the request's trace context to the tool-call spans. // -// Note that the migrated tool handlers no longer read these context values: the -// official SDK delivers headers on the request itself (see mcp.Header, which -// reads req.Extra.Header), so ExtractHTTPHeaders and ExtractTraceInfo currently -// have no non-test callers. They are kept because the values remain useful for -// debugging and for callers that still pass the request context through. +// The server span is started by otelhttp as a child of the incoming W3C +// traceparent. The standalone SSE stream (GET) is not traced: it stays open for +// the whole session, so its span would last for hours and parent nothing. +// +// The MCP SDK does not run tool handlers on the HTTP request's context: a +// stateful session's handlers run on the context of the request that opened the +// session (initialize). What does reach the handlers, per request, is the +// request's header (CallToolRequest.Extra.Header). So the middleware writes the +// server span's context into the traceparent header of a copy of the request, +// and the tool-call span (see internal/mcp.ToolMiddleware) takes its parent from +// there. No header value is recorded on any span. +// +// It also stashes a few request headers plus the trace and span IDs in the +// request context. The migrated tool handlers no longer read these context +// values: the official SDK delivers headers on the request itself (see +// mcp.Header, which reads req.Extra.Header), so ExtractHTTPHeaders and +// ExtractTraceInfo currently have no non-test callers. They are kept because the +// values remain useful for debugging and for callers that still pass the +// request context through. func HTTPMiddleware(next http.Handler) http.Handler { + return otelhttp.NewHandler(contextMiddleware(next), "mcp", + otelhttp.WithFilter(func(r *http.Request) bool { + return r.Method != http.MethodGet + }), + ) +} + +// contextMiddleware propagates the active trace context to the request header +// and the request context. See HTTPMiddleware. +func contextMiddleware(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { ctx := r.Context() - - // Extract OpenTelemetry context from HTTP headers propagator := otel.GetTextMapPropagator() - ctx = propagator.Extract(ctx, propagation.HeaderCarrier(r.Header)) + + // otelhttp has already extracted the incoming trace context and started + // the server span. Extract here only for requests it did not trace, and + // never on top of an active span, which would replace the server span + // with its remote parent. + if !trace.SpanContextFromContext(ctx).IsValid() { + ctx = propagator.Extract(ctx, propagation.HeaderCarrier(r.Header)) + } // Store relevant HTTP headers in context for tool handlers headers := make(map[string]string) @@ -60,8 +89,14 @@ func HTTPMiddleware(next http.Handler) http.Handler { ctx = context.WithValue(ctx, SpanIDKey, span.SpanContext().SpanID().String()) } + // Clone deep-copies the header, so the caller's request is not modified. + r = r.Clone(ctx) + if span.SpanContext().IsValid() { + propagator.Inject(ctx, propagation.HeaderCarrier(r.Header)) + } + // Call next handler with enhanced context - next.ServeHTTP(w, r.WithContext(ctx)) + next.ServeHTTP(w, r) }) }