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) }) }