diff --git a/contrib/mark3labs/mcp-go/option.go b/contrib/mark3labs/mcp-go/option.go index 0abab682261..187373e23e5 100644 --- a/contrib/mark3labs/mcp-go/option.go +++ b/contrib/mark3labs/mcp-go/option.go @@ -32,6 +32,13 @@ type TracingConfig struct { // If both IntentCaptureEnabled and IntentCaptureEnabledFunc are set, the // predicate wins. If both are unset, intent capture is disabled. IntentCaptureEnabledFunc func(ctx context.Context) bool + // RedactToolOutput replaces the tool result body sent to LLMObs with a + // fixed "[REDACTED]" string. Use this when data access control policies + // forbid forwarding tool output to LLMObs traces. + // + // The result.IsError flag is still honored for span error status — only + // the output content is redacted. + RedactToolOutput bool } // intentCapturePredicate returns the runtime predicate for intent capture, @@ -83,6 +90,12 @@ func WithMCPServerTracing(options *TracingConfig) server.ServerOption { server.WithHooks(hooks)(s) + // Register the redaction middleware first (when enabled) so its ctx + // flag is set before toolHandlerMiddleware annotates the LLMObs span. + if options.RedactToolOutput { + server.WithToolHandlerMiddleware(redactToolOutputMiddleware)(s) + } + // Register toolHandlerMiddleware first so it runs first (creates the span) // Note: mcp-go middleware runs in registration order (first registered runs first) server.WithToolHandlerMiddleware(toolHandlerMiddleware)(s) diff --git a/contrib/mark3labs/mcp-go/tracing.go b/contrib/mark3labs/mcp-go/tracing.go index 2b9bb3a8d38..188611436fe 100644 --- a/contrib/mark3labs/mcp-go/tracing.go +++ b/contrib/mark3labs/mcp-go/tracing.go @@ -31,6 +31,19 @@ func appendTracingHooks(hooks *server.Hooks) { hooks.AddOnError(tracingHooks.onError) } +// redactToolOutputKey is the context key set by redactToolOutputMiddleware +// to flag the request's tool output as needing redaction in the LLMObs span. +type redactToolOutputKey struct{} + +// redactToolOutputMiddleware is registered when TracingConfig.RedactToolOutput +// is true; it tags the request context so toolHandlerMiddleware redacts the +// tool output it later annotates onto the LLMObs span. +var redactToolOutputMiddleware = func(next server.ToolHandlerFunc) server.ToolHandlerFunc { + return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return next(context.WithValue(ctx, redactToolOutputKey{}, struct{}{}), request) + } +} + var toolHandlerMiddleware = func(next server.ToolHandlerFunc) server.ToolHandlerFunc { return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { startOpts := []llmobs.StartSpanOption{ @@ -51,11 +64,15 @@ var toolHandlerMiddleware = func(next server.ToolHandlerFunc) server.ToolHandler } var outputText string if result != nil { - resultJSON, marshalErr := json.Marshal(result) - if marshalErr != nil { - instr.Logger().Warn("mcp-go: failed to marshal tool result: %v", marshalErr) + if _, redact := ctx.Value(redactToolOutputKey{}).(struct{}); redact { + outputText = "[REDACTED]" + } else { + resultJSON, marshalErr := json.Marshal(result) + if marshalErr != nil { + instr.Logger().Warn("mcp-go: failed to marshal tool result: %v", marshalErr) + } + outputText = string(resultJSON) } - outputText = string(resultJSON) } toolSpan.Annotate(llmobs.WithAnnotatedTags(map[string]string{ diff --git a/contrib/mark3labs/mcp-go/tracing_test.go b/contrib/mark3labs/mcp-go/tracing_test.go index bcebaa9d62a..8b097b2c324 100644 --- a/contrib/mark3labs/mcp-go/tracing_test.go +++ b/contrib/mark3labs/mcp-go/tracing_test.go @@ -347,6 +347,71 @@ func TestIntegrationToolCallStructuredError(t *testing.T) { assert.Contains(t, outputStr, "Validation failed") } +func TestRedactToolOutput(t *testing.T) { + tt := testTracer(t) + defer tt.Stop() + + srv := server.NewMCPServer("test-server", "1.0.0", + WithMCPServerTracing(&TracingConfig{RedactToolOutput: true})) + + calcTool := mcp.NewTool("calculator", mcp.WithDescription("A simple calculator")) + srv.AddTool(calcTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return mcp.NewToolResultText(`{"result":42,"secret":"hunter2"}`), nil + }) + + ctx := context.Background() + session := &mockSession{id: "redact-session"} + session.Initialize() + ctx = srv.WithContext(ctx, session) + + srv.HandleMessage(ctx, []byte(`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"calculator","arguments":{}}}`)) + + spans := tt.WaitForLLMObsSpans(t, 1) + require.Len(t, spans, 1) + + toolSpan := spans[0] + outputMeta, ok := toolSpan.Meta["output"].(map[string]interface{}) + require.True(t, ok) + assert.Equal(t, "[REDACTED]", outputMeta["value"]) + // Sensitive content from the actual result must not leak. + outputJSON, err := json.Marshal(outputMeta) + require.NoError(t, err) + assert.NotContains(t, string(outputJSON), "hunter2") + assert.NotContains(t, string(outputJSON), "42") +} + +func TestRedactToolOutputStillReportsIsError(t *testing.T) { + tt := testTracer(t) + defer tt.Stop() + + srv := server.NewMCPServer("test-server", "1.0.0", + WithMCPServerTracing(&TracingConfig{RedactToolOutput: true})) + + failTool := mcp.NewTool("failing", mcp.WithDescription("A failing tool")) + srv.AddTool(failTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return mcp.NewToolResultError("permission denied: secret-detail"), nil + }) + + ctx := context.Background() + session := &mockSession{id: "redact-err"} + session.Initialize() + ctx = srv.WithContext(ctx, session) + + srv.HandleMessage(ctx, []byte(`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"failing","arguments":{}}}`)) + + spans := tt.WaitForLLMObsSpans(t, 1) + require.Len(t, spans, 1) + + toolSpan := spans[0] + assert.Contains(t, toolSpan.Meta, "error.message") + outputMeta, ok := toolSpan.Meta["output"].(map[string]interface{}) + require.True(t, ok) + assert.Equal(t, "[REDACTED]", outputMeta["value"]) + outputJSON, err := json.Marshal(outputMeta) + require.NoError(t, err) + assert.NotContains(t, string(outputJSON), "secret-detail") +} + func TestWithMCPServerTracingWithCustomHooks(t *testing.T) { tt := testTracer(t) defer tt.Stop()