feat: propagate MCP request trace context - #869
Conversation
✅ Docs preview readyThe preview is ready to be viewed. View the preview File Changes 0 new, 2 changed, 0 removedBuild ID: 263199513faf1a3bc48ac991 URL: https://www.apollographql.com/docs/deploy-preview/263199513faf1a3bc48ac991
|
|
Navigate logical layers of code changes, visualize relationships, and explore their blast radius. No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Advanced Run ID: 📒 Files selected for processing (7)
Included review availability: This review used your included allowance. Your plan provides up to 2 included reviews per hour; 1 remain after this review. 📝 SummarySummary by CodeRabbit
WalkthroughMCP request metadata now sets trace context for handler spans and downstream GraphQL requests. The changes add context propagation across handlers, tests for HTTP and stdio behavior, and telemetry documentation. ChangesMCP trace-context propagation
Priority: ➖ Normal Estimated code review effort: 3 (Moderate) | ~25 minutes Change: Feature Sequence Diagram(s)sequenceDiagram
participant MCPClient
participant MCPHandler
participant with_request_context
participant GraphQLClient
MCPClient->>MCPHandler: Send request with _meta
MCPHandler->>with_request_context: Build handler span and request context
with_request_context-->>MCPHandler: Return span and context
MCPHandler->>GraphQLClient: Send request with active trace context
Suggested reviewers: Merge Risk: ⚪ Minimal · up to No actionable issue remains identified for this trace-context change; it is mergeable after normal checks. Security Architecture ReviewSecurity architecture risk: 🔵 Low · up to The reviewed path carries caller-provided diagnostic context, not authentication identity. Metadata validation, header sanitization and request-scoped attachment limit the apparent risk. No new access-control bypass or cross-request disclosure was substantiated, but downstream policies and failure-path coverage remain incomplete. Retained concerns Security review detailsSecurity Blast Radius
Security Findings and Attack Paths
Trust Boundaries and Controls
Resilience and Maintainability Implications
Hardening Proposals
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Full details: Docstring CoverageExplanation Docstring coverage is 41.46% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 41 functions across 5 files. (2 skipped: 2 unsupported.)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
✅ Changeset file added - thank you! |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| /// Prepare the handler span and retain its parent context for poll-scoped | ||
| /// attachment even when the span is filtered. Set the parent before entry: | ||
| /// the OTel layer starts spans on entry, after which parentage cannot change. | ||
| pub(super) fn with_request_context( |
There was a problem hiding this comment.
[Consider] Every handler now repeats the same three steps: call with_request_context(...), wrap the body in async, then .instrument(span).with_context(parent_context).await. That is nine copies, and it forced a full re-indent of listen and call_tool. If a future handler forgets either .instrument or .with_context, it silently loses the parent or the filtered-span fallback.
A helper that takes the future would keep the ordering invariant (context outside, span inside) in one place and make the call sites flat:
pub(super) fn in_request_context<F: Future>(
span: tracing::Span,
context: &RequestContext<RoleServer>,
fut: F,
) -> impl Future<Output = F::Output> {
let parent = /* current body of with_request_context */;
fut.instrument(span).with_context(parent)
}Single-expression handlers such as list_tools, list_prompts and get_prompt could then drop their async wrapper blocks. (Ch. 1.6: prefer structure over repeated patterns.)
Review SummaryOverall: This PR has every MCP request handler read W3C trace context from Update (be2a989): This commit fixes the Earlier findings:
Test coverage: Strong. Tests cover precedence, invalid, non-string and non-header-safe metadata, baggage clearing and sanitization, filtered handler and HTTP spans, and concurrent stdio isolation. Every handler is covered across both protocol versions. Recommendation: Approve. Reviewed by Claude Code Opus 5.5 |
There was a problem hiding this comment.
I reviewed this PR and didn't find any bugs. Because it reworks trace-context propagation across all MCP handler methods and parses external W3C trace-context metadata from request _meta, a human look would still be worthwhile.
What was reviewed: the with_request_context/request_parent_context extraction and fallback logic in telemetry.rs (invalid traceparent never grafts tracestate from a different trace; baggage precedence between metadata and HTTP), the mechanical running.rs refactor from #[tracing::instrument(parent = ...)] to explicit info_span + .instrument().with_context() across all nine ServerHandler methods, and the poll/cancellation context-restoration test. The flagged baggage-size concern was ruled out as pre-existing (same normalization path already used for HTTP headers).
Extended reasoning...
The change replaces instrument-macro parenting with manual OTel context propagation across every ServerHandler method (initialize, listen, call_tool, list_tools, list_resources, read_resource, list_prompts, get_prompt, set_level) in running.rs, and adds metadata-based trace-context extraction with fallback-to-HTTP logic plus baggage merging/sanitization in telemetry.rs. Security-sensitive surface: it parses externally supplied W3C trace-context fields (traceparent/tracestate/baggage) from request _meta, with explicit care taken not to graft tracestate across trace boundaries on invalid traceparent. No CODEOWNERS objections or outstanding review threads were found in the timeline, and tests are extensive (482 new lines of dedicated trace-context tests plus poll/cancellation coverage), but the subtlety of future-poll-scoped OTel context attachment and its interaction with span filtering is enough complexity that independent human review is still worthwhile rather than auto-approving.
| let baggage = meta.get_baggage().map(normalize_baggage_list); | ||
| let carrier: HashMap<String, String> = [ | ||
| ("traceparent", meta.get_traceparent()), | ||
| ("tracestate", meta.get_tracestate()), |
There was a problem hiding this comment.
I traced _meta.tracestate from request_parent_context through the W3C propagator and into the downstream GraphQL request, and I noticed traceparent gets reformatted from the parsed hex IDs, and baggage gets percent-encoded (and its metadata is checked by sanitize_baggage_for_http_injection). But TraceState::valid_value only rejects ,, = and long values, and trace_state().header() writes the value back as is. So a valid traceparent with a tracestate like vendor=k\u0001 reaches reqwest-tracing's HeaderValue::from_str(...).expect(...) and panics, and the client hangs until it times out. I wonder if we should drop tracestate here when HeaderValue::from_str rejects it, the same way we do for baggage metadata?
There was a problem hiding this comment.
Good catch, thanks! I reproduced it (panic in reqwest-tracing, then the client hangs until timeout) and fixed it in be2a989. I went slightly stricter than HeaderValue::from_str: _meta.tracestate now has to be text an inbound header could carry (visible ASCII or tab, the same set HeaderExtractor reads), because from_str still lets non-ASCII through to the downstream request. When it fails, the whole tracestate is dropped and the traceparent kept. control_char_state and non_ascii_state cover both cases.
gocamille
left a comment
gocamille
left a comment
There was a problem hiding this comment.
I just have one minor comment! Thank you @inanna-apollo !
2 participants
MCP callers can supply W3C trace context in request
_meta, but handler spans previously only inherited HTTP context. This connects every implemented request handler to valid_meta.traceparentcontext over stdio and Streamable HTTP and propagates the active GraphQL client span downstream.Valid metadata trace context directly parents the handler span; absent or invalid context falls back to the HTTP span or a new stdio trace root. Metadata baggage replaces HTTP baggage when string-valued, with existing HTTP-injection sanitization. Handler spans retain HTTP tracing ancestry for formatted logs while the selected OTel context is set before entry, when the OTel layer starts spans.
The selected context is attached during each future poll, so propagation and Rhai trace-ID correlation also work when logging filters disable handler spans. If the HTTP span is filtered too, the extracted HTTP-header context supplies the fallback. Context attachment is restored on suspension and cancellation, keeping concurrent requests isolated.
Explicit header configuration keeps its existing behavior: automatic injection regenerates traceparent/tracestate, while explicitly supplied baggage remains when the active context has no baggage. Includes documentation and a patch changeset. rmcp 3.5.0 already supplies the required metadata accessors.
Supports both protocol versions 2025-11-25 and 2026-07-28, retaining main's initialize/discovery lifecycle.
Validation
cargo test --offline --workspace— passed (1,332 tests).cargo test -p apollo-mcp-server mcp_trace_context --offline— all 50 regression tests passed across both protocol versions.cargo clippy --offline --all-targets -- --deny warnings— passed.cargo fmt --all -- --checkandgit diff --check— passed.