Skip to content

Commit d5629ec

Browse files
committed
feat(contrib/mark3labs/mcp-go): support redacting tool output in LLMObs traces
Add RedactToolOutput bool to TracingConfig. When set, the tool result body sent to LLMObs is replaced with '[REDACTED]'; the result.IsError flag is still honored for span error status. This is needed when downstream access controls (e.g. org2-dac in the internal dd-source copy) forbid forwarding tool output to LLMObs traces. Generalizes the hardcoded redaction in internal dd-source PR #429591 into a configuration option so the original always-on behavior remains the default for external users.
1 parent 952cf50 commit d5629ec

3 files changed

Lines changed: 99 additions & 4 deletions

File tree

contrib/mark3labs/mcp-go/option.go

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,13 @@ type TracingConfig struct {
3232
// If both IntentCaptureEnabled and IntentCaptureEnabledFunc are set, the
3333
// predicate wins. If both are unset, intent capture is disabled.
3434
IntentCaptureEnabledFunc func(ctx context.Context) bool
35+
// RedactToolOutput replaces the tool result body sent to LLMObs with a
36+
// fixed "[REDACTED]" string. Use this when data access control policies
37+
// forbid forwarding tool output to LLMObs traces.
38+
//
39+
// The result.IsError flag is still honored for span error status — only
40+
// the output content is redacted.
41+
RedactToolOutput bool
3542
}
3643

3744
// intentCapturePredicate returns the runtime predicate for intent capture,
@@ -83,6 +90,12 @@ func WithMCPServerTracing(options *TracingConfig) server.ServerOption {
8390

8491
server.WithHooks(hooks)(s)
8592

93+
// Register the redaction middleware first (when enabled) so its ctx
94+
// flag is set before toolHandlerMiddleware annotates the LLMObs span.
95+
if options.RedactToolOutput {
96+
server.WithToolHandlerMiddleware(redactToolOutputMiddleware)(s)
97+
}
98+
8699
// Register toolHandlerMiddleware first so it runs first (creates the span)
87100
// Note: mcp-go middleware runs in registration order (first registered runs first)
88101
server.WithToolHandlerMiddleware(toolHandlerMiddleware)(s)

contrib/mark3labs/mcp-go/tracing.go

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,19 @@ func appendTracingHooks(hooks *server.Hooks) {
3131
hooks.AddOnError(tracingHooks.onError)
3232
}
3333

34+
// redactToolOutputKey is the context key set by redactToolOutputMiddleware
35+
// to flag the request's tool output as needing redaction in the LLMObs span.
36+
type redactToolOutputKey struct{}
37+
38+
// redactToolOutputMiddleware is registered when TracingConfig.RedactToolOutput
39+
// is true; it tags the request context so toolHandlerMiddleware redacts the
40+
// tool output it later annotates onto the LLMObs span.
41+
var redactToolOutputMiddleware = func(next server.ToolHandlerFunc) server.ToolHandlerFunc {
42+
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
43+
return next(context.WithValue(ctx, redactToolOutputKey{}, struct{}{}), request)
44+
}
45+
}
46+
3447
var toolHandlerMiddleware = func(next server.ToolHandlerFunc) server.ToolHandlerFunc {
3548
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
3649
startOpts := []llmobs.StartSpanOption{
@@ -51,11 +64,15 @@ var toolHandlerMiddleware = func(next server.ToolHandlerFunc) server.ToolHandler
5164
}
5265
var outputText string
5366
if result != nil {
54-
resultJSON, marshalErr := json.Marshal(result)
55-
if marshalErr != nil {
56-
instr.Logger().Warn("mcp-go: failed to marshal tool result: %v", marshalErr)
67+
if _, redact := ctx.Value(redactToolOutputKey{}).(struct{}); redact {
68+
outputText = "[REDACTED]"
69+
} else {
70+
resultJSON, marshalErr := json.Marshal(result)
71+
if marshalErr != nil {
72+
instr.Logger().Warn("mcp-go: failed to marshal tool result: %v", marshalErr)
73+
}
74+
outputText = string(resultJSON)
5775
}
58-
outputText = string(resultJSON)
5976
}
6077

6178
toolSpan.Annotate(llmobs.WithAnnotatedTags(map[string]string{

contrib/mark3labs/mcp-go/tracing_test.go

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -347,6 +347,71 @@ func TestIntegrationToolCallStructuredError(t *testing.T) {
347347
assert.Contains(t, outputStr, "Validation failed")
348348
}
349349

350+
func TestRedactToolOutput(t *testing.T) {
351+
tt := testTracer(t)
352+
defer tt.Stop()
353+
354+
srv := server.NewMCPServer("test-server", "1.0.0",
355+
WithMCPServerTracing(&TracingConfig{RedactToolOutput: true}))
356+
357+
calcTool := mcp.NewTool("calculator", mcp.WithDescription("A simple calculator"))
358+
srv.AddTool(calcTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
359+
return mcp.NewToolResultText(`{"result":42,"secret":"hunter2"}`), nil
360+
})
361+
362+
ctx := context.Background()
363+
session := &mockSession{id: "redact-session"}
364+
session.Initialize()
365+
ctx = srv.WithContext(ctx, session)
366+
367+
srv.HandleMessage(ctx, []byte(`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"calculator","arguments":{}}}`))
368+
369+
spans := tt.WaitForLLMObsSpans(t, 1)
370+
require.Len(t, spans, 1)
371+
372+
toolSpan := spans[0]
373+
outputMeta, ok := toolSpan.Meta["output"].(map[string]interface{})
374+
require.True(t, ok)
375+
assert.Equal(t, "[REDACTED]", outputMeta["value"])
376+
// Sensitive content from the actual result must not leak.
377+
outputJSON, err := json.Marshal(outputMeta)
378+
require.NoError(t, err)
379+
assert.NotContains(t, string(outputJSON), "hunter2")
380+
assert.NotContains(t, string(outputJSON), "42")
381+
}
382+
383+
func TestRedactToolOutputStillReportsIsError(t *testing.T) {
384+
tt := testTracer(t)
385+
defer tt.Stop()
386+
387+
srv := server.NewMCPServer("test-server", "1.0.0",
388+
WithMCPServerTracing(&TracingConfig{RedactToolOutput: true}))
389+
390+
failTool := mcp.NewTool("failing", mcp.WithDescription("A failing tool"))
391+
srv.AddTool(failTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
392+
return mcp.NewToolResultError("permission denied: secret-detail"), nil
393+
})
394+
395+
ctx := context.Background()
396+
session := &mockSession{id: "redact-err"}
397+
session.Initialize()
398+
ctx = srv.WithContext(ctx, session)
399+
400+
srv.HandleMessage(ctx, []byte(`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"failing","arguments":{}}}`))
401+
402+
spans := tt.WaitForLLMObsSpans(t, 1)
403+
require.Len(t, spans, 1)
404+
405+
toolSpan := spans[0]
406+
assert.Contains(t, toolSpan.Meta, "error.message")
407+
outputMeta, ok := toolSpan.Meta["output"].(map[string]interface{})
408+
require.True(t, ok)
409+
assert.Equal(t, "[REDACTED]", outputMeta["value"])
410+
outputJSON, err := json.Marshal(outputMeta)
411+
require.NoError(t, err)
412+
assert.NotContains(t, string(outputJSON), "secret-detail")
413+
}
414+
350415
func TestWithMCPServerTracingWithCustomHooks(t *testing.T) {
351416
tt := testTracer(t)
352417
defer tt.Stop()

0 commit comments

Comments
 (0)