Skip to content

Commit ce524bf

Browse files
committed
feat(contrib/mark3labs/mcp-go): expose captured intent to tool handlers via context
The intent-capture middleware annotates the LLMObs span with the client- supplied telemetry.intent, but tool handlers had no way to read it without re-parsing the telemetry blob (which the same middleware strips before the handler runs). Stash the intent on context.Context alongside the existing LLMObs annotation and export IntentFromContext / ContextWithIntent helpers so handlers can forward the intent downstream (e.g. to a search API for retrieval/ranking experiments). Ported from internal dd-source PR #443664.
1 parent 5e25dad commit ce524bf

2 files changed

Lines changed: 106 additions & 9 deletions

File tree

contrib/mark3labs/mcp-go/intent_capture.go

Lines changed: 51 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,36 @@ func requiredAsStringSlice(raw any) ([]string, bool) {
141141
}
142142
}
143143

144+
// intentCtxKey is the context key used to stash the captured intent so tool
145+
// handlers can forward it downstream (e.g. to a search API). Kept unexported
146+
// to force callers through IntentFromContext.
147+
type intentCtxKey struct{}
148+
149+
// IntentFromContext returns the captured intent for the current MCP tool call,
150+
// if intent capture is enabled and the client supplied a non-empty
151+
// telemetry.intent. The boolean is false when no intent is available.
152+
func IntentFromContext(ctx context.Context) (string, bool) {
153+
if ctx == nil {
154+
return "", false
155+
}
156+
v, ok := ctx.Value(intentCtxKey{}).(string)
157+
if !ok || v == "" {
158+
return "", false
159+
}
160+
return v, true
161+
}
162+
163+
// ContextWithIntent returns a copy of ctx that carries the given intent. The
164+
// middleware uses this internally; it is exported so tests (and callers that
165+
// fabricate their own contexts outside the standard middleware chain) can seed
166+
// the value that IntentFromContext will later read.
167+
func ContextWithIntent(ctx context.Context, intent string) context.Context {
168+
if intent == "" {
169+
return ctx
170+
}
171+
return context.WithValue(ctx, intentCtxKey{}, intent)
172+
}
173+
144174
// Removing tracing parameters from the tool call request so its not sent to the tool.
145175
// This must be registered after the tool handler middleware (mcp-go runs middleware in registration order).
146176
// This removes the telemetry parameter before user-defined middleware or tool handlers can see it.
@@ -149,7 +179,10 @@ var processAndRemoveTelemetryToolMiddleware = func(next server.ToolHandlerFunc)
149179
if m, ok := request.Params.Arguments.(map[string]any); ok && m != nil {
150180
if telemetryVal, has := m[instrmcp.TelemetryKey]; has {
151181
if telemetryMap, ok := telemetryVal.(map[string]any); ok {
152-
processTelemetry(ctx, telemetryMap)
182+
if intent := extractIntent(telemetryMap); intent != "" {
183+
annotateIntentOnSpan(ctx, intent)
184+
ctx = ContextWithIntent(ctx, intent)
185+
}
153186
} else if instr != nil && instr.Logger() != nil {
154187
instr.Logger().Warn("mcp-go intent capture: telemetry value is not a map")
155188
}
@@ -161,26 +194,35 @@ var processAndRemoveTelemetryToolMiddleware = func(next server.ToolHandlerFunc)
161194
}
162195
}
163196

164-
func processTelemetry(ctx context.Context, telemetryVal map[string]any) {
197+
// extractIntent pulls the intent string out of the telemetry map supplied by
198+
// the MCP client. It returns "" when the entry is missing, the wrong type, or
199+
// empty — callers should treat that as "no intent" and skip further work.
200+
func extractIntent(telemetryVal map[string]any) string {
165201
if telemetryVal == nil {
166-
return
202+
return ""
167203
}
168-
169204
intentVal, exists := telemetryVal[instrmcp.IntentKey]
170205
if !exists {
171-
return
206+
return ""
172207
}
173-
174208
intent, ok := intentVal.(string)
175-
if !ok || intent == "" {
176-
return
209+
if !ok {
210+
return ""
177211
}
212+
return intent
213+
}
178214

215+
// annotateIntentOnSpan records intent on the active LLM Obs tool span, if one
216+
// is present on ctx. It is a no-op when no span is active or the active span
217+
// is not a tool span, so it is always safe to call.
218+
func annotateIntentOnSpan(ctx context.Context, intent string) {
219+
if intent == "" {
220+
return
221+
}
179222
span, ok := llmobs.SpanFromContext(ctx)
180223
if !ok {
181224
return
182225
}
183-
184226
toolSpan, ok := span.AsTool()
185227
if !ok {
186228
return

contrib/mark3labs/mcp-go/intent_capture_test.go

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,8 @@ func TestIntentCapture(t *testing.T) {
2525
srv := server.NewMCPServer("test-server", "1.0.0", WithMCPServerTracing(&TracingConfig{IntentCaptureEnabled: true}))
2626

2727
var receivedArgs map[string]any
28+
var receivedIntent string
29+
var receivedIntentOK bool
2830
calcTool := mcp.NewTool("calculator",
2931
mcp.WithDescription("A simple calculator"),
3032
mcp.WithString("operation", mcp.Required(), mcp.Description("The operation to perform")),
@@ -33,6 +35,7 @@ func TestIntentCapture(t *testing.T) {
3335

3436
srv.AddTool(calcTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
3537
receivedArgs = request.Params.Arguments.(map[string]any)
38+
receivedIntent, receivedIntentOK = IntentFromContext(ctx)
3639
return mcp.NewToolResultText(`{"result":8}`), nil
3740
})
3841

@@ -81,6 +84,11 @@ func TestIntentCapture(t *testing.T) {
8184
assert.Equal(t, float64(3), receivedArgs["y"])
8285
assert.NotContains(t, receivedArgs, "telemetry")
8386

87+
// Verify intent was stashed in ctx so the handler can forward it downstream
88+
// (e.g. into a search API request) without re-reading the telemetry blob.
89+
assert.True(t, receivedIntentOK, "IntentFromContext should report a value")
90+
assert.Equal(t, "test intent description", receivedIntent)
91+
8492
// Verify intent was recorded on the LLMObs span
8593
spans := tt.WaitForLLMObsSpans(t, 1)
8694
require.Len(t, spans, 1)
@@ -92,6 +100,53 @@ func TestIntentCapture(t *testing.T) {
92100
assert.Equal(t, "test intent description", toolSpan.Meta["intent"])
93101
}
94102

103+
func TestIntentFromContext(t *testing.T) {
104+
ctx := context.Background()
105+
106+
_, ok := IntentFromContext(ctx)
107+
assert.False(t, ok)
108+
109+
ctx2 := ContextWithIntent(ctx, "find recent errors")
110+
got, ok := IntentFromContext(ctx2)
111+
assert.True(t, ok)
112+
assert.Equal(t, "find recent errors", got)
113+
114+
// Empty intent does not seed the context.
115+
ctx3 := ContextWithIntent(ctx, "")
116+
_, ok = IntentFromContext(ctx3)
117+
assert.False(t, ok)
118+
}
119+
120+
func TestIntentFromContext_AbsentWhenNoTelemetry(t *testing.T) {
121+
tt := testTracer(t)
122+
defer tt.Stop()
123+
124+
srv := server.NewMCPServer("test-server", "1.0.0", WithMCPServerTracing(&TracingConfig{IntentCaptureEnabled: true}))
125+
126+
var receivedIntent string
127+
var receivedIntentOK bool
128+
calcTool := mcp.NewTool("calculator",
129+
mcp.WithDescription("A simple calculator"),
130+
mcp.WithString("operation", mcp.Required(), mcp.Description("The operation to perform")),
131+
mcp.WithNumber("x", mcp.Required(), mcp.Description("First number")),
132+
mcp.WithNumber("y", mcp.Required(), mcp.Description("Second number")))
133+
134+
srv.AddTool(calcTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
135+
receivedIntent, receivedIntentOK = IntentFromContext(ctx)
136+
return mcp.NewToolResultText(`{"result":8}`), nil
137+
})
138+
139+
ctx := context.Background()
140+
session := &mockSession{id: "test"}
141+
session.Initialize()
142+
ctx = srv.WithContext(ctx, session)
143+
144+
srv.HandleMessage(ctx, []byte(`{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"calculator","arguments":{"operation":"add","x":5,"y":3}}}`))
145+
146+
assert.False(t, receivedIntentOK, "IntentFromContext should be empty when no telemetry was supplied")
147+
assert.Empty(t, receivedIntent)
148+
}
149+
95150
func TestIntentCaptureRawInputSchemaPreservesUnknownFields(t *testing.T) {
96151
// Tools defined with NewToolWithRawSchema may declare JSON Schema
97152
// keywords beyond what mcp.ToolInputSchema models (additionalProperties,

0 commit comments

Comments
 (0)