add intent annotation to llmobs module

Author: jboolean

contrib/mark3labs/mcp-go/intent_capture.go

@@ -0,0 +1,113 @@
+// 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 mcpgo
+
+import (
+	"context"
+	"slices"
+
+	"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 {
+	return map[string]any{
+		"type": "object",
+		"properties": map[string]any{
+			"intent": map[string]any{
+				"type":        "string",
+				"description": intentPrompt,
+			},
+		},
+		"required":             []string{"intent"},
+		"additionalProperties": false,
+	}
+}
+
+// Injects tracing parameters into the tool list response by mutating it.
+func injectDdtraceListToolsHook(ctx context.Context, id any, message *mcp.ListToolsRequest, result *mcp.ListToolsResult) {
+	if result == nil || result.Tools == nil {
+		return
+	}
+
+	for i := range result.Tools {
+		t := &result.Tools[i]
+
+		if t.RawInputSchema != nil {
+			instr.Logger().Warn("mcp-go intent capture: raw input schema not supported")
+			continue
+		}
+
+		if t.InputSchema.Type == "" {
+			t.InputSchema.Type = "object"
+		}
+		if t.InputSchema.Properties == nil {
+			t.InputSchema.Properties = map[string]any{}
+		}
+
+		// Insert/overwrite the ddtrace property
+		t.InputSchema.Properties[ddtraceKey] = ddtraceSchema()
+
+		// Mark ddtrace as required (idempotent)
+		if !slices.Contains(t.InputSchema.Required, ddtraceKey) {
+			t.InputSchema.Required = append(t.InputSchema.Required, ddtraceKey)
+		}
+	}
+}
+
+// Removing tracing parameters from the tool call request so its not sent to the tool.
+// This must be registered before the tool handler middleware, so that the span is available.
+// This must be registered after any user-defined middleware so that it is not visible to them.
+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 ddtraceMap, ok := ddtraceVal.(map[string]any); ok {
+					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)
+				request.Params.Arguments = m
+			}
+		}
+
+		return next(ctx, request)
+	}
+}
+
+func processDdtrace(ctx context.Context, m map[string]any) {
+	if m == nil {
+		return
+	}
+
+	intentVal, exists := m["intent"]
+	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
+	}
+	// TODO: Add fields to toolSpan to annotate intent
+	_ = toolSpan
+}

contrib/mark3labs/mcp-go/intent_capture_test.go

@@ -0,0 +1,87 @@
+// 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 mcpgo
+
+import (
+	"context"
+	"encoding/json"
+	"testing"
+
+	"github.com/DataDog/dd-trace-go/v2/ddtrace/mocktracer"
+	"github.com/mark3labs/mcp-go/mcp"
+	"github.com/mark3labs/mcp-go/server"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestIntentCapture(t *testing.T) {
+	mt := mocktracer.Start()
+	defer mt.Stop()
+
+	srv := server.NewMCPServer("test-server", "1.0.0", WithTracing(&TracingConfig{IntentCaptureEnabled: true}))
+
+	var receivedArgs map[string]any
+	calcTool := mcp.NewTool("calculator",
+		mcp.WithDescription("A simple calculator"),
+		mcp.WithString("operation", mcp.Required(), mcp.Description("The operation to perform")),
+		mcp.WithNumber("x", mcp.Required(), mcp.Description("First number")),
+		mcp.WithNumber("y", mcp.Required(), mcp.Description("Second number")))
+
+	srv.AddTool(calcTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+		receivedArgs = request.Params.Arguments.(map[string]any)
+		return mcp.NewToolResultText(`{"result":8}`), nil
+	})
+
+	ctx := context.Background()
+
+	listResp := srv.HandleMessage(ctx, []byte(`{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}`))
+	var listResult map[string]interface{}
+	json.Unmarshal(json.RawMessage(mustMarshal(listResp)), &listResult)
+
+	result := listResult["result"].(map[string]interface{})
+	tools := result["tools"].([]interface{})
+	require.Len(t, tools, 1)
+
+	tool := tools[0].(map[string]interface{})
+	schema := tool["inputSchema"].(map[string]interface{})
+	props := schema["properties"].(map[string]interface{})
+
+	assert.Contains(t, props, "operation")
+	assert.Contains(t, props, "x")
+	assert.Contains(t, props, "y")
+	assert.Contains(t, props, "ddtrace")
+
+	// Ensure ddtrace is added to schema
+	ddtraceSchema := props["ddtrace"].(map[string]interface{})
+	assert.Equal(t, "object", ddtraceSchema["type"])
+	ddtraceProps := ddtraceSchema["properties"].(map[string]interface{})
+	intentSchema := ddtraceProps["intent"].(map[string]interface{})
+	assert.Equal(t, "string", intentSchema["type"])
+	assert.Equal(t, intentPrompt, intentSchema["description"])
+
+	required := schema["required"].([]interface{})
+	assert.Contains(t, required, "operation")
+	assert.Contains(t, required, "x")
+	assert.Contains(t, required, "y")
+
+	session := &mockSession{id: "test"}
+	session.Initialize()
+	ctx = srv.WithContext(ctx, session)
+
+	srv.HandleMessage(ctx, []byte(`{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"calculator","arguments":{"operation":"add","x":5,"y":3,"ddtrace":{"intent":"test"}}}}`))
+
+	// Ensure ddtrace is removed in tool call
+	require.NotNil(t, receivedArgs)
+	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")
+}
+
+func mustMarshal(v interface{}) []byte {
+	b, _ := json.Marshal(v)
+	return b
+}

contrib/mark3labs/mcp-go/mcpgo.go

@@ -0,0 +1,16 @@
+// 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 2016 Datadog, Inc.
+
+package mcpgo // import "github.com/DataDog/dd-trace-go/contrib/mark3labs/mcp-go/v2"
+
+import (
+	"github.com/DataDog/dd-trace-go/v2/instrumentation"
+)
+
+var instr *instrumentation.Instrumentation
+
+func init() {
+	instr = instrumentation.Load(instrumentation.PackageMark3LabsMCPGo)
+}

contrib/mark3labs/mcp-go/option.go

@@ -9,14 +9,15 @@ import (
 	"github.com/mark3labs/mcp-go/server"
 )
 
-// The file contains methods for easily adding tracing to a MCP server.
+// The file contains methods for easily adding tracing and intent capture to a MCP server.
 
 // TracingConfig holds configuration for adding tracing to an MCP server.
 type TracingConfig struct {
 	// Hooks allows you to provide custom hooks that will be merged with Datadog tracing hooks.
 	// If nil, only Datadog tracing hooks will be added and any custom hooks provided via server.WithHooks(...) will be removed.
 	// If provided, your custom hooks will be executed alongside Datadog tracing hooks.
-	Hooks *server.Hooks
+	Hooks                *server.Hooks
+	IntentCaptureEnabled bool
 }
 
 // WithTracing adds Datadog tracing to an MCP server.
@@ -54,6 +55,12 @@ func WithTracing(options *TracingConfig) server.ServerOption {
 
 		server.WithHooks(hooks)(s)
 
+		if options.IntentCaptureEnabled {
+			hooks.AddAfterListTools(injectDdtraceListToolsHook)
+			server.WithToolHandlerMiddleware(processAndRemoveDdtraceToolMiddleware)(s)
+		}
+
+		// This must be after the intent capture middleware, so that span in available. Last registered middleware is run first.
 		server.WithToolHandlerMiddleware(toolHandlerMiddleware)(s)
 	}
 }

contrib/mark3labs/mcp-go/tracing.go

@@ -17,12 +17,6 @@ import (
 	"github.com/mark3labs/mcp-go/server"
 )
 
-var instr *instrumentation.Instrumentation
-
-func init() {
-	instr = instrumentation.Load(instrumentation.PackageMark3LabsMCPGo)
-}
-
 type hooks struct {
 	spanCache *sync.Map
 }

internal/llmobs/llmobs.go

@@ -113,6 +113,9 @@ type llmobsContext struct {
 	outputMessages  []LLMMessage
 	outputText      string
 
+	// tool specific
+	intent string
+
 	// experiment specific
 	experimentInput          any
 	experimentExpectedOutput any
@@ -513,6 +516,14 @@ func (l *LLMObs) llmobsSpanEvent(span *Span) *transport.LLMObsSpanEvent {
 		meta["tool_definitions"] = toolDefinitions
 	}
 
+	if intent := span.llmCtx.intent; intent != "" {
+		if spanKind != SpanKindTool {
+			log.Warn("llmobs: dropping intent on non-tool span kind, annotating intent is only supported for tool span kinds")
+		} else {
+			meta["intent"] = intent
+		}
+	}
+
 	spanStatus := "ok"
 	var errMsg *transport.ErrorMessage
 	if span.error != nil {

internal/llmobs/llmobs_test.go

@@ -568,6 +568,17 @@ func TestSpanAnnotate(t *testing.T) {
 			},
 			wantSessionID: "config-session-456",
 		},
+		{
+			name: "tool-span-with-intent",
+			kind: llmobs.SpanKindTool,
+			annotations: llmobs.SpanAnnotations{
+				Intent: "test intent",
+			},
+			wantMeta: map[string]any{
+				"span.kind": "tool",
+				"intent":    "test intent",
+			},
+		},
 	}
 
 	for _, tc := range testCases {

internal/llmobs/span.go

@@ -191,6 +191,9 @@ type SpanAnnotations struct {
 	// ToolDefinitions are the tool definitions for LLM spans.
 	ToolDefinitions []ToolDefinition
 
+	// Intent is a description of a reason for calling an MCP tool on tool spans
+	Intent string
+
 	// AgentManifest is the agent manifest for agent spans.
 	AgentManifest string
 
@@ -368,6 +371,14 @@ func (s *Span) Annotate(a SpanAnnotations) {
 		}
 	}
 
+	if a.Intent != "" {
+		if s.spanKind != SpanKindTool {
+			log.Warn("llmobs: intent can only be annotated on tool spans, ignoring")
+		} else {
+			s.llmCtx.intent = a.Intent
+		}
+	}
+
 	s.annotateIO(a)
 }
 

llmobs/option.go

@@ -154,3 +154,11 @@ func WithAnnotatedMetrics(metrics map[string]float64) AnnotateOption {
 		}
 	}
 }
+
+// WithIntent sets the intent for the span.
+// Intent is a description of a reason for calling an MCP tool.
+func WithIntent(intent string) AnnotateOption {
+	return func(a *illmobs.SpanAnnotations) {
+		a.Intent = intent
+	}
+}