feat(contrib/mark3labs/mcp-go): support per-request intent capture

Add IntentCaptureEnabledFunc to TracingConfig: a per-request predicate that
gates intent capture at runtime, enabling toggling per org/feature flag
without restarting the server.

When non-nil, the predicate is consulted on every tools/list and tools/call:
returning false makes that request behave as if intent capture were disabled
(no schema injection, no telemetry stripping, no intent annotation).

The static IntentCaptureEnabled bool is preserved for back-compat; if both
are set the predicate wins. The injection hook and tool-call middleware
were refactored into factory functions taking the predicate.

Generalizes the conditional intent-capture pattern that the internal
dd-source copy uses (no single PR — landed across several commits there).

Author: jboolean

contrib/mark3labs/mcp-go/intent_capture.go

@@ -32,8 +32,21 @@ func telemetrySchema() map[string]any {
 	}
 }
 
-// Injects tracing parameters into the tool list response by mutating it.
-func injectTelemetryListToolsHook(ctx context.Context, id any, message *mcp.ListToolsRequest, result *mcp.ListToolsResult) {
+// injectTelemetryListToolsHookFor returns an AfterListTools hook that injects
+// the telemetry parameter into the schemas of model-callable tools, but only
+// when the supplied predicate returns true for the request's context.
+func injectTelemetryListToolsHookFor(enabled func(context.Context) bool) func(context.Context, any, *mcp.ListToolsRequest, *mcp.ListToolsResult) {
+	return func(ctx context.Context, id any, message *mcp.ListToolsRequest, result *mcp.ListToolsResult) {
+		if !enabled(ctx) {
+			return
+		}
+		injectTelemetryListTools(result)
+	}
+}
+
+// injectTelemetryListTools mutates result to add the telemetry parameter to
+// each model-callable tool's input schema.
+func injectTelemetryListTools(result *mcp.ListToolsResult) {
 	if result == nil || result.Tools == nil {
 		return
 	}
@@ -152,26 +165,36 @@ func ContextWithIntent(ctx context.Context, intent string) context.Context {
 	return context.WithValue(ctx, intentCtxKey{}, intent)
 }
 
-// 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 telemetry parameter before user-defined middleware or tool handlers can see it.
-var processAndRemoveTelemetryToolMiddleware = 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 telemetryVal, has := m[instrmcp.TelemetryKey]; has {
-				if telemetryMap, ok := telemetryVal.(map[string]any); ok {
-					if intent := extractIntent(telemetryMap); intent != "" {
-						annotateIntentOnSpan(ctx, intent)
-						ctx = ContextWithIntent(ctx, intent)
+// processAndRemoveTelemetryToolMiddlewareFor returns a tool handler middleware
+// that strips the telemetry parameter from a tools/call request and annotates
+// the active LLMObs span with the supplied intent, but only when the supplied
+// predicate returns true for the request's context. When the predicate
+// returns false, the middleware is a transparent pass-through.
+//
+// This must be registered after the tool handler middleware (mcp-go runs
+// middleware in registration order). This removes the telemetry parameter
+// before user-defined middleware or tool handlers can see it.
+func processAndRemoveTelemetryToolMiddlewareFor(enabled func(context.Context) bool) func(server.ToolHandlerFunc) server.ToolHandlerFunc {
+	return func(next server.ToolHandlerFunc) server.ToolHandlerFunc {
+		return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+			if !enabled(ctx) {
+				return next(ctx, request)
+			}
+			if m, ok := request.Params.Arguments.(map[string]any); ok && m != nil {
+				if telemetryVal, has := m[instrmcp.TelemetryKey]; has {
+					if telemetryMap, ok := telemetryVal.(map[string]any); ok {
+						if intent := extractIntent(telemetryMap); intent != "" {
+							annotateIntentOnSpan(ctx, intent)
+							ctx = ContextWithIntent(ctx, intent)
+						}
+					} else if instr != nil && instr.Logger() != nil {
+						instr.Logger().Warn("mcp-go intent capture: telemetry value is not a map")
 					}
-				} else if instr != nil && instr.Logger() != nil {
-					instr.Logger().Warn("mcp-go intent capture: telemetry value is not a map")
+					delete(m, instrmcp.TelemetryKey)
 				}
-				delete(m, instrmcp.TelemetryKey)
 			}
+			return next(ctx, request)
 		}
-
-		return next(ctx, request)
 	}
 }
 

contrib/mark3labs/mcp-go/intent_capture_test.go

@@ -100,6 +100,80 @@ func TestIntentCapture(t *testing.T) {
 	assert.Equal(t, "test intent description", toolSpan.Meta["intent"])
 }
 
+type enabledKey struct{}
+
+func TestIntentCaptureEnabledFunc(t *testing.T) {
+	tt := testTracer(t)
+	defer tt.Stop()
+
+	srv := server.NewMCPServer("test-server", "1.0.0", WithMCPServerTracing(&TracingConfig{
+		IntentCaptureEnabledFunc: func(ctx context.Context) bool {
+			v, _ := ctx.Value(enabledKey{}).(bool)
+			return v
+		},
+	}))
+
+	calcTool := mcp.NewTool("calculator",
+		mcp.WithDescription("A simple calculator"),
+		mcp.WithString("operation", mcp.Required()))
+
+	var seenArgs map[string]any
+	srv.AddTool(calcTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+		seenArgs = request.Params.Arguments.(map[string]any)
+		return mcp.NewToolResultText(`{"ok":true}`), nil
+	})
+
+	// Disabled per ctx: schema has no telemetry; telemetry in arg passes through.
+	disabledCtx := context.WithValue(context.Background(), enabledKey{}, false)
+	listResp := srv.HandleMessage(disabledCtx, []byte(`{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}`))
+	var listResult map[string]interface{}
+	require.NoError(t, json.Unmarshal(mustMarshal(listResp), &listResult))
+	props := listResult["result"].(map[string]interface{})["tools"].([]interface{})[0].(map[string]interface{})["inputSchema"].(map[string]interface{})["properties"].(map[string]interface{})
+	assert.NotContains(t, props, "telemetry")
+
+	session := &mockSession{id: "s1"}
+	session.Initialize()
+	disabledCtx = srv.WithContext(disabledCtx, session)
+	srv.HandleMessage(disabledCtx, []byte(`{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"calculator","arguments":{"operation":"add","telemetry":{"intent":"x"}}}}`))
+	// Predicate false → middleware is a pass-through, so the telemetry argument is NOT stripped.
+	require.NotNil(t, seenArgs)
+	assert.Contains(t, seenArgs, "telemetry")
+
+	// Enabled per ctx: schema gets telemetry; tool argument is stripped before handler.
+	enabledCtx := context.WithValue(context.Background(), enabledKey{}, true)
+	listResp = srv.HandleMessage(enabledCtx, []byte(`{"jsonrpc":"2.0","id":3,"method":"tools/list","params":{}}`))
+	require.NoError(t, json.Unmarshal(mustMarshal(listResp), &listResult))
+	props = listResult["result"].(map[string]interface{})["tools"].([]interface{})[0].(map[string]interface{})["inputSchema"].(map[string]interface{})["properties"].(map[string]interface{})
+	assert.Contains(t, props, "telemetry")
+
+	session2 := &mockSession{id: "s2"}
+	session2.Initialize()
+	enabledCtx = srv.WithContext(enabledCtx, session2)
+	seenArgs = nil
+	srv.HandleMessage(enabledCtx, []byte(`{"jsonrpc":"2.0","id":4,"method":"tools/call","params":{"name":"calculator","arguments":{"operation":"add","telemetry":{"intent":"x"}}}}`))
+	require.NotNil(t, seenArgs)
+	assert.NotContains(t, seenArgs, "telemetry")
+}
+
+func TestIntentCaptureEnabledFuncOverridesBool(t *testing.T) {
+	// When both static bool and the predicate are set, the predicate wins.
+	srv := server.NewMCPServer("test-server", "1.0.0", WithMCPServerTracing(&TracingConfig{
+		IntentCaptureEnabled:     true,
+		IntentCaptureEnabledFunc: func(context.Context) bool { return false },
+	}))
+
+	tool := mcp.NewTool("t", mcp.WithDescription("t"), mcp.WithString("q", mcp.Required()))
+	srv.AddTool(tool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+		return mcp.NewToolResultText("ok"), nil
+	})
+
+	listResp := srv.HandleMessage(context.Background(), []byte(`{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}`))
+	var listResult map[string]interface{}
+	require.NoError(t, json.Unmarshal(mustMarshal(listResp), &listResult))
+	props := listResult["result"].(map[string]interface{})["tools"].([]interface{})[0].(map[string]interface{})["inputSchema"].(map[string]interface{})["properties"].(map[string]interface{})
+	assert.NotContains(t, props, "telemetry")
+}
+
 func TestIntentCaptureSkipsUIOnlyTools(t *testing.T) {
 	tt := testTracer(t)
 	defer tt.Stop()

contrib/mark3labs/mcp-go/option.go

@@ -6,6 +6,8 @@
 package mcpgo
 
 import (
+	"context"
+
 	"github.com/mark3labs/mcp-go/server"
 )
 
@@ -17,9 +19,33 @@ type TracingConfig struct {
 	// 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
-	// Enables intent capture for tool spans.
-	// This will modify the tool schemas to include a parameter for the client to provide the intent.
+	// IntentCaptureEnabled enables intent capture for tool spans unconditionally.
+	// This modifies the tool schemas to include a parameter for the client to provide the intent.
+	// For per-request control (e.g. driven by a feature flag), use IntentCaptureEnabledFunc instead.
 	IntentCaptureEnabled bool
+	// IntentCaptureEnabledFunc is a per-request predicate that gates intent capture
+	// at runtime. When non-nil, the intent-capture hook and middleware are registered
+	// and consult the predicate on every tools/list and tools/call: if it returns
+	// false, that request behaves as if intent capture were disabled (no schema
+	// injection, no telemetry stripping, no intent annotation).
+	//
+	// If both IntentCaptureEnabled and IntentCaptureEnabledFunc are set, the
+	// predicate wins. If both are unset, intent capture is disabled.
+	IntentCaptureEnabledFunc func(ctx context.Context) bool
+}
+
+// intentCapturePredicate returns the runtime predicate for intent capture,
+// or nil if intent capture is disabled. The predicate from
+// IntentCaptureEnabledFunc takes precedence over the static
+// IntentCaptureEnabled flag.
+func intentCapturePredicate(options *TracingConfig) func(context.Context) bool {
+	if options.IntentCaptureEnabledFunc != nil {
+		return options.IntentCaptureEnabledFunc
+	}
+	if options.IntentCaptureEnabled {
+		return func(context.Context) bool { return true }
+	}
+	return nil
 }
 
 // WithMCPServerTracing adds Datadog tracing to an MCP server.
@@ -61,10 +87,11 @@ func WithMCPServerTracing(options *TracingConfig) server.ServerOption {
 		// Note: mcp-go middleware runs in registration order (first registered runs first)
 		server.WithToolHandlerMiddleware(toolHandlerMiddleware)(s)
 
-		if options.IntentCaptureEnabled {
-			hooks.AddAfterListTools(injectTelemetryListToolsHook)
+		predicate := intentCapturePredicate(options)
+		if predicate != nil {
+			hooks.AddAfterListTools(injectTelemetryListToolsHookFor(predicate))
 			// Register intent capture middleware second so it runs second (after span is created)
-			server.WithToolHandlerMiddleware(processAndRemoveTelemetryToolMiddleware)(s)
+			server.WithToolHandlerMiddleware(processAndRemoveTelemetryToolMiddlewareFor(predicate))(s)
 		}
 	}
 }