Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 28 additions & 19 deletions contrib/modelcontextprotocol/go-sdk/intent_capture.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
Comment on lines +44 to 46

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Restrict predicate evaluation to tool requests

WithIntentCapturePredicate is documented as being consulted for tools/list and tools/call, but this check runs before the method switch, so it also invokes user code for initialize and any other server method where intent capture cannot apply. In servers whose predicate expects tool/list request context or performs an expensive/fragile feature-flag lookup, ordinary session initialization can now pay for or fail inside that callback; move the predicate check into the tools/list and tools/call cases.

Useful? React with 👍 / 👎.

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)
}
}

Expand Down
68 changes: 68 additions & 0 deletions contrib/modelcontextprotocol/go-sdk/intent_capture_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ package gosdk
import (
"context"
"encoding/json"
"sync/atomic"
"testing"

"github.com/google/jsonschema-go/jsonschema"
Expand All @@ -19,6 +20,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 atomic.Bool
server := mcp.NewServer(&mcp.Implementation{Name: "test-server", Version: "1.0.0"}, nil)
AddTracing(server, WithIntentCapturePredicate(func(context.Context) bool {
return enabled.Load()
}))

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.Store(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.Store(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.
Expand Down
19 changes: 17 additions & 2 deletions contrib/modelcontextprotocol/go-sdk/option.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
}
4 changes: 2 additions & 2 deletions contrib/modelcontextprotocol/go-sdk/tracing.go
Original file line number Diff line number Diff line change
Expand Up @@ -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...)
Expand Down
Loading