feat(telemetry): trace each MCP request and tool call under the caller's trace

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

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 <tool>" 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)
+		}
+	}
+}

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

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=

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 <tool name>" 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),
 			)