Skip to content

Commit 2d2bdea

Browse files
feat(contrib/mark3labs/mcp-go): set session ID on tool spans at start time (#4943)
This PR stack ports functionality added in the clone of this contrib in dd-source back to dd-trace-go. ## Summary Pass `llmobs.WithSessionID` at `StartToolSpan` time so LLMObs groups tool spans under their MCP session natively. The session is read from `server.ClientSessionFromContext(ctx)` if present. The existing `mcp_session_id` tag is unchanged — they are complementary: the tag exposes the session ID to span search/filtering, while `WithSessionID` populates the LLMObs session field used for trace grouping. ## Provenance No single dd-source PR — the change landed as part of a series of LLMObs-session changes. The current dd-source middleware lives at [`mcp-go/tracing.go#L36-L44`](https://github.com/DataDog/dd-source/blob/d10d624c374b69aa02168983be725532a38562b1/domains/mcp_services/libs/go/mcp/pre-release-external/dd-trace-go/contrib/mark3labs/mcp-go/tracing.go#L36-L44): ```go var toolHandlerMiddleware = func(next server.ToolHandlerFunc) server.ToolHandlerFunc { return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { // Note: WithIntegration is not yet in released llmobs package and cannot be used here var startOpts []llmobs.StartSpanOption if session := server.ClientSessionFromContext(ctx); session != nil { startOpts = append(startOpts, llmobs.WithSessionID(session.SessionID())) } toolSpan, ctx := llmobs.StartToolSpan(ctx, request.Params.Name, startOpts...) ``` This PR ports that exact pattern, additionally keeping the existing `llmobs.WithIntegration(...)` call since the released llmobs package now has it. ## Test plan - [x] Extended `TestIntegrationToolCallSuccess` to assert `toolSpan.SessionID` matches the MCP client session. - [x] Existing tracing tests still pass. Stacked on #4942. Co-authored-by: jboolean <julian.boilen@datadoghq.com>
2 parents 2ec11b5 + f6787f2 commit 2d2bdea

2 files changed

Lines changed: 10 additions & 1 deletion

File tree

contrib/mark3labs/mcp-go/tracing.go

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,13 @@ func appendTracingHooks(hooks *server.Hooks) {
3333

3434
var toolHandlerMiddleware = func(next server.ToolHandlerFunc) server.ToolHandlerFunc {
3535
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
36-
toolSpan, ctx := llmobs.StartToolSpan(ctx, request.Params.Name, llmobs.WithIntegration(string(instrumentation.PackageMark3LabsMCPGo)))
36+
startOpts := []llmobs.StartSpanOption{
37+
llmobs.WithIntegration(string(instrumentation.PackageMark3LabsMCPGo)),
38+
}
39+
if session := server.ClientSessionFromContext(ctx); session != nil {
40+
startOpts = append(startOpts, llmobs.WithSessionID(session.SessionID()))
41+
}
42+
toolSpan, ctx := llmobs.StartToolSpan(ctx, request.Params.Name, startOpts...)
3743

3844
var result *mcp.CallToolResult
3945
var err error

contrib/mark3labs/mcp-go/tracing_test.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -178,6 +178,9 @@ func TestIntegrationToolCallSuccess(t *testing.T) {
178178
assert.Contains(t, initSpan.Tags, expectedTag)
179179
assert.Contains(t, toolSpan.Tags, expectedTag)
180180

181+
// Tool span carries the session ID natively so LLMObs groups it under the session.
182+
assert.Equal(t, sessionID, toolSpan.SessionID)
183+
181184
assert.Contains(t, toolSpan.Tags, "mcp_method:tools/call")
182185
assert.Contains(t, toolSpan.Tags, "mcp_tool_kind:server")
183186
assert.Contains(t, toolSpan.Tags, "mcp_tool:calculator")

0 commit comments

Comments
 (0)