diff --git a/contrib/mark3labs/mcp-go/README.md b/contrib/mark3labs/mcp-go/README.md index b81035c306..0eb8f2adbe 100644 --- a/contrib/mark3labs/mcp-go/README.md +++ b/contrib/mark3labs/mcp-go/README.md @@ -34,5 +34,5 @@ The integration automatically traces: - **Tool calls**: Creates LLMObs tool spans with input/output annotation for all tool invocations - **Session initialization**: Create LLMObs task spans for session initialization, including client information. -The integration can optionally capture "intent" on MCP tool calls. When enabled, this adds a parameter to the schema of each tool to request that the client include an explaination. +The integration can optionally capture "intent" on MCP tool calls. When enabled, this adds a parameter to the schema of each tool to request that the client include an explanation. This can help provide context in natural language about why tools are being used. \ No newline at end of file diff --git a/contrib/mark3labs/mcp-go/intent_capture.go b/contrib/mark3labs/mcp-go/intent_capture.go index c0da469111..c59551ea54 100644 --- a/contrib/mark3labs/mcp-go/intent_capture.go +++ b/contrib/mark3labs/mcp-go/intent_capture.go @@ -9,25 +9,22 @@ import ( "context" "slices" + instrmcp "github.com/DataDog/dd-trace-go/v2/instrumentation/mcp" "github.com/DataDog/dd-trace-go/v2/llmobs" "github.com/mark3labs/mcp-go/mcp" "github.com/mark3labs/mcp-go/server" ) -const ddtraceKey = "ddtrace" - -const intentPrompt string = "Briefly describe the wider context task, and why this tool was chosen. Omit argument values, PII/secrets. Use English." - -func ddtraceSchema() map[string]any { +func ddTraceSchema() map[string]any { return map[string]any{ "type": "object", "properties": map[string]any{ - "intent": map[string]any{ + instrmcp.IntentKey: map[string]any{ "type": "string", - "description": intentPrompt, + "description": instrmcp.IntentPrompt, }, }, - "required": []string{"intent"}, + "required": []string{instrmcp.IntentKey}, "additionalProperties": false, } } @@ -54,11 +51,11 @@ func injectDdtraceListToolsHook(ctx context.Context, id any, message *mcp.ListTo } // Insert/overwrite the ddtrace property - t.InputSchema.Properties[ddtraceKey] = ddtraceSchema() + t.InputSchema.Properties[instrmcp.DDTraceKey] = ddTraceSchema() // Mark ddtrace as required (idempotent) - if !slices.Contains(t.InputSchema.Required, ddtraceKey) { - t.InputSchema.Required = append(t.InputSchema.Required, ddtraceKey) + if !slices.Contains(t.InputSchema.Required, instrmcp.DDTraceKey) { + t.InputSchema.Required = append(t.InputSchema.Required, instrmcp.DDTraceKey) } } } @@ -66,16 +63,16 @@ func injectDdtraceListToolsHook(ctx context.Context, id any, message *mcp.ListTo // Removing tracing parameters from the tool call request so its not sent to the tool. // This must be registered after the tool handler middleware (mcp-go runs middleware in registration order). // This removes the ddtrace parameter before user-defined middleware or tool handlers can see it. -var processAndRemoveDdtraceToolMiddleware = func(next server.ToolHandlerFunc) server.ToolHandlerFunc { +var processAndRemoveDDTraceToolMiddleware = func(next server.ToolHandlerFunc) server.ToolHandlerFunc { return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { if m, ok := request.Params.Arguments.(map[string]any); ok && m != nil { - if ddtraceVal, has := m[ddtraceKey]; has { + if ddtraceVal, has := m[instrmcp.DDTraceKey]; has { if ddtraceMap, ok := ddtraceVal.(map[string]any); ok { - processDdtrace(ctx, ddtraceMap) + processDDTrace(ctx, ddtraceMap) } else if instr != nil && instr.Logger() != nil { instr.Logger().Warn("mcp-go intent capture: ddtrace value is not a map") } - delete(m, ddtraceKey) + delete(m, instrmcp.DDTraceKey) } } @@ -83,12 +80,12 @@ var processAndRemoveDdtraceToolMiddleware = func(next server.ToolHandlerFunc) se } } -func processDdtrace(ctx context.Context, m map[string]any) { - if m == nil { +func processDDTrace(ctx context.Context, ddTraceVal map[string]any) { + if ddTraceVal == nil { return } - intentVal, exists := m["intent"] + intentVal, exists := ddTraceVal[instrmcp.IntentKey] if !exists { return } diff --git a/contrib/mark3labs/mcp-go/intent_capture_test.go b/contrib/mark3labs/mcp-go/intent_capture_test.go index 6ca32b4787..680cc52609 100644 --- a/contrib/mark3labs/mcp-go/intent_capture_test.go +++ b/contrib/mark3labs/mcp-go/intent_capture_test.go @@ -10,6 +10,7 @@ import ( "encoding/json" "testing" + instrmcp "github.com/DataDog/dd-trace-go/v2/instrumentation/mcp" "github.com/mark3labs/mcp-go/mcp" "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" @@ -59,7 +60,7 @@ func TestIntentCapture(t *testing.T) { ddtraceProps := ddtraceSchema["properties"].(map[string]interface{}) intentSchema := ddtraceProps["intent"].(map[string]interface{}) assert.Equal(t, "string", intentSchema["type"]) - assert.Equal(t, intentPrompt, intentSchema["description"]) + assert.Equal(t, instrmcp.IntentPrompt, intentSchema["description"]) required := schema["required"].([]interface{}) assert.Contains(t, required, "operation") diff --git a/contrib/mark3labs/mcp-go/option.go b/contrib/mark3labs/mcp-go/option.go index 9cea6b8f6e..f5f2956fb7 100644 --- a/contrib/mark3labs/mcp-go/option.go +++ b/contrib/mark3labs/mcp-go/option.go @@ -64,7 +64,7 @@ func WithMCPServerTracing(options *TracingConfig) server.ServerOption { if options.IntentCaptureEnabled { hooks.AddAfterListTools(injectDdtraceListToolsHook) // Register intent capture middleware second so it runs second (after span is created) - server.WithToolHandlerMiddleware(processAndRemoveDdtraceToolMiddleware)(s) + server.WithToolHandlerMiddleware(processAndRemoveDDTraceToolMiddleware)(s) } } } diff --git a/contrib/mark3labs/mcp-go/testing.go b/contrib/mark3labs/mcp-go/shared_test.go similarity index 100% rename from contrib/mark3labs/mcp-go/testing.go rename to contrib/mark3labs/mcp-go/shared_test.go diff --git a/contrib/modelcontextprotocol/go-sdk/README.md b/contrib/modelcontextprotocol/go-sdk/README.md index 80e6c2a07d..5f942fff10 100644 --- a/contrib/modelcontextprotocol/go-sdk/README.md +++ b/contrib/modelcontextprotocol/go-sdk/README.md @@ -16,13 +16,22 @@ func main() { defer tracer.Stop() server := mcp.NewServer(&mcp.Implementation{Name: "my-server", Version: "1.0.0"}, nil) - gosdktrace.AddTracingMiddleware(server) + + // Add tracing middleware + gosdktrace.AddTracing(server) + + // Or with intent capture enabled: + // gosdktrace.AddTracing(server, gosdktrace.WithIntentCapture()) } ``` ## Features -The integration automatically traces: +`AddTracing` automatically traces: - **Tool calls**: Creates LLMObs tool spans with input/output annotation for all tool invocations - **Session initialization**: Creates LLMObs task spans for session initialization, including client information +`WithIntentCapture()` option enables context capture: +When enabled, this adds a parameter to the schema of each tool to request that the client include an explanation of its use. +This can help provide context in natural language about why tools are being used. + diff --git a/contrib/modelcontextprotocol/go-sdk/example_test.go b/contrib/modelcontextprotocol/go-sdk/example_test.go index 80dbcdad22..938c825c64 100644 --- a/contrib/modelcontextprotocol/go-sdk/example_test.go +++ b/contrib/modelcontextprotocol/go-sdk/example_test.go @@ -16,6 +16,6 @@ func Example() { defer tracer.Stop() server := mcp.NewServer(&mcp.Implementation{Name: "my-server", Version: "1.0.0"}, nil) - gosdktrace.AddTracingMiddleware(server) + gosdktrace.AddTracing(server, gosdktrace.WithIntentCapture()) _ = server } diff --git a/contrib/modelcontextprotocol/go-sdk/intent_capture.go b/contrib/modelcontextprotocol/go-sdk/intent_capture.go new file mode 100644 index 0000000000..e78bdbf484 --- /dev/null +++ b/contrib/modelcontextprotocol/go-sdk/intent_capture.go @@ -0,0 +1,131 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2025 Datadog, Inc. + +package gosdk + +import ( + "context" + "encoding/json" + + instrmcp "github.com/DataDog/dd-trace-go/v2/instrumentation/mcp" + "github.com/DataDog/dd-trace-go/v2/llmobs" + "github.com/google/jsonschema-go/jsonschema" + "github.com/modelcontextprotocol/go-sdk/mcp" +) + +func ddTraceSchema() *jsonschema.Schema { + return &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + instrmcp.IntentKey: { + Type: "string", + Description: instrmcp.IntentPrompt, + }, + }, + Required: []string{instrmcp.IntentKey}, + } +} + +// intentCaptureReceivingMiddleware is an mcp.Server receiving middleware +// adding intent information to the tool call span. +// Intent capture works by injecting an additional required parameter on tools that the +// client agent will fill in to explain context about the task. +// The middleware records this intent on the span, and then removes it from the arguments before the tool is called. +func intentCaptureReceivingMiddleware(next mcp.MethodHandler) mcp.MethodHandler { + return func(ctx context.Context, method string, req mcp.Request) (mcp.Result, error) { + switch method { + case "tools/list": + // The additional parameter is added to the tool arguments returned by tools/list. + res, err := next(ctx, method, req) + if toolListRes, ok := res.(*mcp.ListToolsResult); ok { + injectToolsListResponse(toolListRes) + } + return res, err + case "tools/call": + // The intent is recorded and the argument is removed. + if toolReq, ok := req.(*mcp.CallToolRequest); ok { + return processToolCallIntent(next, ctx, method, toolReq) + } + } + return next(ctx, method, req) + } +} + +func injectToolsListResponse(res *mcp.ListToolsResult) { + for i := range res.Tools { + inputSchema, ok := res.Tools[i].InputSchema.(*jsonschema.Schema) + if !ok { + instr.Logger().Warn("go-sdk intent capture: unexpected input schema type: %T", res.Tools[i].InputSchema) + continue + } + + if inputSchema.Type == "" { + inputSchema.Type = "object" + } + if inputSchema.Properties == nil { + inputSchema.Properties = map[string]*jsonschema.Schema{} + } + + inputSchema.Properties[instrmcp.DDTraceKey] = ddTraceSchema() + inputSchema.Required = append(inputSchema.Required, instrmcp.DDTraceKey) + } +} + +func processToolCallIntent(next mcp.MethodHandler, ctx context.Context, method string, req *mcp.CallToolRequest) (mcp.Result, error) { + if len(req.Params.Arguments) > 0 { + var argsMap map[string]any + if err := json.Unmarshal(req.Params.Arguments, &argsMap); err != nil { + if instr != nil && instr.Logger() != nil { + instr.Logger().Warn("go-sdk intent capture: failed to unmarshal arguments: %v", err) + } + return next(ctx, method, req) + } + + if ddtraceVal, has := argsMap[instrmcp.DDTraceKey]; has { + if ddtraceMap, ok := ddtraceVal.(map[string]any); ok { + annotateSpanWithIntent(ctx, ddtraceMap) + } else { + instr.Logger().Warn("go-sdk intent capture: ddtrace value is not a map") + } + + delete(argsMap, instrmcp.DDTraceKey) + + modifiedArgs, err := json.Marshal(argsMap) + if err != nil { + instr.Logger().Warn("go-sdk intent capture: failed to marshal modified arguments: %v", err) + } else { + req.Params.Arguments = modifiedArgs + } + } + } + return next(ctx, method, req) +} + +func annotateSpanWithIntent(ctx context.Context, ddTraceVal map[string]any) { + if ddTraceVal == nil { + return + } + + intentVal, exists := ddTraceVal[instrmcp.IntentKey] + if !exists { + return + } + + intent, ok := intentVal.(string) + if !ok || intent == "" { + return + } + + span, ok := llmobs.SpanFromContext(ctx) + if !ok { + return + } + + toolSpan, ok := span.AsTool() + if !ok { + return + } + toolSpan.Annotate(llmobs.WithIntent(intent)) +} diff --git a/contrib/modelcontextprotocol/go-sdk/intent_capture_test.go b/contrib/modelcontextprotocol/go-sdk/intent_capture_test.go new file mode 100644 index 0000000000..bf87e2ba39 --- /dev/null +++ b/contrib/modelcontextprotocol/go-sdk/intent_capture_test.go @@ -0,0 +1,146 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2025 Datadog, Inc. + +package gosdk + +import ( + "context" + "testing" + + instrmcp "github.com/DataDog/dd-trace-go/v2/instrumentation/mcp" + "github.com/DataDog/dd-trace-go/v2/instrumentation/testutils/testtracer" + "github.com/google/jsonschema-go/jsonschema" + "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestIntentCapture(t *testing.T) { + tt := testTracer(t) + defer tt.Stop() + + ctx := context.Background() + + server := mcp.NewServer(&mcp.Implementation{Name: "test-server", Version: "1.0.0"}, nil) + AddTracing(server, WithIntentCapture()) + + type CalcArgs struct { + Operation string `json:"operation"` + X float64 `json:"x"` + Y float64 `json:"y"` + } + + var receivedArgs map[string]any + var receivedRequest *mcp.CallToolRequest + mcp.AddTool(server, + &mcp.Tool{ + Name: "calculator", + Description: "A simple calculator", + InputSchema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "operation": {Type: "string", Description: "The operation to perform"}, + "x": {Type: "number", Description: "First number"}, + "y": {Type: "number", Description: "Second number"}, + }, + Required: []string{"operation", "x", "y"}, + }, + }, + func(ctx context.Context, req *mcp.CallToolRequest, args map[string]any) (*mcp.CallToolResult, any, error) { + receivedArgs = args + receivedRequest = req + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.TextContent{Text: `{"result":8}`}, + }, + }, nil, nil + }, + ) + + client := mcp.NewClient(&mcp.Implementation{Name: "test-client", Version: "1.0.0"}, nil) + + clientTransport, serverTransport := mcp.NewInMemoryTransports() + + serverSession, err := server.Connect(ctx, serverTransport, nil) + require.NoError(t, err) + defer serverSession.Close() + + clientSession, err := client.Connect(ctx, clientTransport, nil) + require.NoError(t, err) + defer clientSession.Close() + + listResult, err := clientSession.ListTools(ctx, &mcp.ListToolsParams{}) + require.NoError(t, err) + require.Len(t, listResult.Tools, 1) + + tool := listResult.Tools[0] + schemaMap, ok := tool.InputSchema.(map[string]any) + require.True(t, ok, "expected input schema to be a map[string]any, got %T", tool.InputSchema) + + // Verify the input schema has the ddtrace property added + props := schemaMap["properties"].(map[string]any) + assert.Contains(t, props, "operation") + assert.Contains(t, props, "x") + assert.Contains(t, props, "y") + assert.Contains(t, props, "ddtrace") + + ddtraceSchema := props["ddtrace"].(map[string]any) + assert.Equal(t, "object", ddtraceSchema["type"]) + ddtraceProps := ddtraceSchema["properties"].(map[string]any) + intentSchema := ddtraceProps["intent"].(map[string]any) + assert.Equal(t, "string", intentSchema["type"]) + assert.Equal(t, instrmcp.IntentPrompt, intentSchema["description"]) + + // Ensure ddtrace is required, and others are not affected + required := schemaMap["required"].([]any) + assert.Contains(t, required, "operation") + assert.Contains(t, required, "x") + assert.Contains(t, required, "y") + assert.Contains(t, required, "ddtrace") + + result, err := clientSession.CallTool(ctx, &mcp.CallToolParams{ + Name: "calculator", + Arguments: map[string]any{ + "operation": "add", + "x": float64(5), + "y": float64(3), + "ddtrace": map[string]any{ + "intent": "test intent description", + }, + }, + }) + require.NoError(t, err) + assert.NotNil(t, result) + + // Received arguments + assert.Equal(t, "add", receivedArgs["operation"]) + assert.Equal(t, float64(5), receivedArgs["x"]) + assert.Equal(t, float64(3), receivedArgs["y"]) + assert.NotContains(t, receivedArgs, "ddtrace") + + // Received request also does not contain ddtrace + assert.NotContains(t, receivedRequest.Params.Arguments, "ddtrace") + + spans := tt.WaitForLLMObsSpans(t, 2) + require.Len(t, spans, 2) + + var toolSpan *testtracer.LLMObsSpan + for i := range spans { + if spans[i].Name == "calculator" { + toolSpan = &spans[i] + } + } + + require.NotNil(t, toolSpan, "tool span not found") + + assert.Equal(t, "tool", toolSpan.Meta["span.kind"]) + assert.Equal(t, "calculator", toolSpan.Name) + assert.Contains(t, toolSpan.Meta, "intent") + // ddtrace should not be recorded in the input + assert.NotContains(t, toolSpan.Meta["input"], "ddtrace") + + // The intent *is* captured on the span + assert.Equal(t, "test intent description", toolSpan.Meta["intent"]) +} diff --git a/contrib/modelcontextprotocol/go-sdk/option.go b/contrib/modelcontextprotocol/go-sdk/option.go new file mode 100644 index 0000000000..57fc7a9d76 --- /dev/null +++ b/contrib/modelcontextprotocol/go-sdk/option.go @@ -0,0 +1,18 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2025 Datadog, Inc. + +package gosdk + +type config struct { + intentCaptureEnabled bool +} + +type Option func(*config) + +func WithIntentCapture() Option { + return func(cfg *config) { + cfg.intentCaptureEnabled = true + } +} diff --git a/contrib/modelcontextprotocol/go-sdk/tracing.go b/contrib/modelcontextprotocol/go-sdk/tracing.go index f11e8e1aaf..bb9a10a9f0 100644 --- a/contrib/modelcontextprotocol/go-sdk/tracing.go +++ b/contrib/modelcontextprotocol/go-sdk/tracing.go @@ -16,8 +16,21 @@ import ( "github.com/modelcontextprotocol/go-sdk/mcp" ) -func AddTracingMiddleware(server *mcp.Server) { - server.AddReceivingMiddleware(tracingMiddleware) +func AddTracing(server *mcp.Server, opts ...Option) { + cfg := &config{} + for _, opt := range opts { + opt(cfg) + } + + // Middleware in run in the ordering in this slice. + middlewares := []mcp.Middleware{tracingMiddleware} + + // Intent capture is added after tracing so that the intent can be annotated on the existing span. + if cfg.intentCaptureEnabled { + middlewares = append(middlewares, intentCaptureReceivingMiddleware) + } + + server.AddReceivingMiddleware(middlewares...) } func tracingMiddleware(next mcp.MethodHandler) mcp.MethodHandler { diff --git a/contrib/modelcontextprotocol/go-sdk/tracing_test.go b/contrib/modelcontextprotocol/go-sdk/tracing_test.go index aa9cd58638..0a3615fe60 100644 --- a/contrib/modelcontextprotocol/go-sdk/tracing_test.go +++ b/contrib/modelcontextprotocol/go-sdk/tracing_test.go @@ -29,7 +29,7 @@ func TestIntegrationSessionInitialize(t *testing.T) { ctx := context.Background() server := mcp.NewServer(&mcp.Implementation{Name: "test-server", Version: "1.0.0"}, nil) - AddTracingMiddleware(server) + AddTracing(server) // go-sdk only assigns session ids on streamable transports. // Using a streamable http transport in this test allows testing session id tagging behavior. @@ -91,7 +91,7 @@ func TestIntegrationToolCallSuccess(t *testing.T) { ctx := context.Background() server := mcp.NewServer(&mcp.Implementation{Name: "test-server", Version: "1.0.0"}, nil) - AddTracingMiddleware(server) + AddTracing(server, WithIntentCapture()) type CalcArgs struct { Operation string `json:"operation"` @@ -146,6 +146,8 @@ func TestIntegrationToolCallSuccess(t *testing.T) { require.NoError(t, err) defer clientSession.Close() + clientSession.ListTools(ctx, &mcp.ListToolsParams{}) + sessionID := clientSession.ID() require.NotEmpty(t, sessionID) @@ -235,7 +237,7 @@ func TestIntegrationToolCallError(t *testing.T) { ctx := context.Background() server := mcp.NewServer(&mcp.Implementation{Name: "test-server", Version: "1.0.0"}, nil) - AddTracingMiddleware(server) + AddTracing(server) mcp.AddTool(server, &mcp.Tool{ @@ -312,7 +314,7 @@ func TestIntegrationToolCallStructuredError(t *testing.T) { ctx := context.Background() server := mcp.NewServer(&mcp.Implementation{Name: "test-server", Version: "1.0.0"}, nil) - AddTracingMiddleware(server) + AddTracing(server) type ValidationArgs struct { Name string `json:"name"` @@ -415,6 +417,8 @@ func TestIntegrationToolCallStructuredError(t *testing.T) { assert.Contains(t, outputStr, "invalid input") } +// Shared helpers + // testTracer creates a testtracer with LLMObs enabled for integration tests func testTracer(t *testing.T, opts ...testtracer.Option) *testtracer.TestTracer { defaultOpts := []testtracer.Option{ diff --git a/instrumentation/mcp/mcp.go b/instrumentation/mcp/mcp.go new file mode 100644 index 0000000000..3a4c9c8880 --- /dev/null +++ b/instrumentation/mcp/mcp.go @@ -0,0 +1,12 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2025 Datadog, Inc. + +package mcp + +const DDTraceKey = "ddtrace" + +const IntentKey = "intent" + +const IntentPrompt string = "Briefly describe the wider context task, and why this tool was chosen. Omit argument values, PII/secrets. Use English."