feat(contrib/modelcontextprotocol/go-sdk): support per-request intent capture predicate

Author: jboolean

contrib/modelcontextprotocol/go-sdk/intent_capture.go

@@ -29,28 +29,37 @@ func telemetrySchema() map[string]any {
 	}
 }
 
-// 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)
+// intentCaptureReceivingMiddlewareFor returns an mcp.Server receiving
+// middleware that adds intent information to the tool call span, gated
+// per-request by the supplied predicate. When the predicate returns false,
+// the request passes through untouched.
+//
+// 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 intentCaptureReceivingMiddlewareFor(enabled func(context.Context) bool) func(mcp.MethodHandler) mcp.MethodHandler {
+	return func(next mcp.MethodHandler) mcp.MethodHandler {
+		return func(ctx context.Context, method string, req mcp.Request) (mcp.Result, error) {
+			if !enabled(ctx) {
+				return next(ctx, method, req)
 			}
-			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)
+			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)
 		}
-		return next(ctx, method, req)
 	}
 }
 

contrib/modelcontextprotocol/go-sdk/intent_capture_test.go

@@ -19,6 +19,73 @@ import (
 	"github.com/DataDog/dd-trace-go/v2/instrumentation/testutils/testtracer"
 )
 
+func TestIntentCapturePredicate(t *testing.T) {
+	// The predicate must gate per-request: when it returns false, no schema
+	// injection and the telemetry argument reaches the handler.
+	tt := testTracer(t)
+	defer tt.Stop()
+	ctx := context.Background()
+
+	var enabled bool
+	server := mcp.NewServer(&mcp.Implementation{Name: "test-server", Version: "1.0.0"}, nil)
+	AddTracing(server, WithIntentCapturePredicate(func(context.Context) bool {
+		return enabled
+	}))
+
+	var receivedArgs map[string]any
+	server.AddTool(&mcp.Tool{
+		Name:        "tool",
+		Description: "tool",
+		InputSchema: &jsonschema.Schema{Type: "object", Properties: map[string]*jsonschema.Schema{"q": {Type: "string"}}},
+	}, func(_ context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+		_ = json.Unmarshal(req.Params.Arguments, &receivedArgs)
+		return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: "ok"}}}, 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()
+
+	// Predicate false: schema not injected, telemetry argument passed through.
+	enabled = false
+	listResult, err := clientSession.ListTools(ctx, &mcp.ListToolsParams{})
+	require.NoError(t, err)
+	require.Len(t, listResult.Tools, 1)
+	if schema, ok := listResult.Tools[0].InputSchema.(map[string]any); ok {
+		if props, _ := schema["properties"].(map[string]any); props != nil {
+			assert.NotContains(t, props, "telemetry")
+		}
+	}
+
+	_, err = clientSession.CallTool(ctx, &mcp.CallToolParams{
+		Name:      "tool",
+		Arguments: map[string]any{"q": "x", "telemetry": map[string]any{"intent": "ignored"}},
+	})
+	require.NoError(t, err)
+	assert.Contains(t, receivedArgs, "telemetry", "telemetry should pass through when predicate is false")
+
+	// Predicate true: schema injected, telemetry stripped.
+	enabled = true
+	listResult, err = clientSession.ListTools(ctx, &mcp.ListToolsParams{})
+	require.NoError(t, err)
+	schema, ok := listResult.Tools[0].InputSchema.(map[string]any)
+	require.True(t, ok)
+	assert.Contains(t, schema["properties"].(map[string]any), "telemetry")
+
+	receivedArgs = nil
+	_, err = clientSession.CallTool(ctx, &mcp.CallToolParams{
+		Name:      "tool",
+		Arguments: map[string]any{"q": "x", "telemetry": map[string]any{"intent": "captured"}},
+	})
+	require.NoError(t, err)
+	assert.NotContains(t, receivedArgs, "telemetry")
+}
+
 func TestIntentCapturePreservesUnknownSchemaKeywords(t *testing.T) {
 	// *jsonschema.Schema doesn't model additionalProperties/oneOf; the map-based
 	// injection must pass those through verbatim instead of dropping them.

contrib/modelcontextprotocol/go-sdk/option.go

@@ -5,14 +5,29 @@
 
 package gosdk
 
+import "context"
+
 type config struct {
-	intentCaptureEnabled bool
+	intentCapturePredicate func(context.Context) bool
 }
 
 type Option func(*config)
 
+// WithIntentCapture enables intent capture for all tool calls.
+//
+// For per-request control (e.g. driven by a feature flag), use
+// WithIntentCapturePredicate instead.
 func WithIntentCapture() Option {
+	return WithIntentCapturePredicate(func(context.Context) bool { return true })
+}
+
+// WithIntentCapturePredicate enables intent capture and gates it per-request
+// using the supplied predicate. The predicate is consulted on every tools/list
+// and tools/call: when it returns false, the request behaves as if intent
+// capture were disabled (no schema injection, no telemetry stripping, no
+// intent annotation).
+func WithIntentCapturePredicate(predicate func(context.Context) bool) Option {
 	return func(cfg *config) {
-		cfg.intentCaptureEnabled = true
+		cfg.intentCapturePredicate = predicate
 	}
 }

contrib/modelcontextprotocol/go-sdk/tracing.go

@@ -27,8 +27,8 @@ func AddTracing(server *mcp.Server, opts ...Option) {
 	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)
+	if cfg.intentCapturePredicate != nil {
+		middlewares = append(middlewares, intentCaptureReceivingMiddlewareFor(cfg.intentCapturePredicate))
 	}
 
 	server.AddReceivingMiddleware(middlewares...)