Commit 2d2bdea
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 files changed
Lines changed: 10 additions & 1 deletion
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
33 | 33 | | |
34 | 34 | | |
35 | 35 | | |
36 | | - | |
| 36 | + | |
| 37 | + | |
| 38 | + | |
| 39 | + | |
| 40 | + | |
| 41 | + | |
| 42 | + | |
37 | 43 | | |
38 | 44 | | |
39 | 45 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
178 | 178 | | |
179 | 179 | | |
180 | 180 | | |
| 181 | + | |
| 182 | + | |
| 183 | + | |
181 | 184 | | |
182 | 185 | | |
183 | 186 | | |
| |||
0 commit comments