Skip to content

fix: per-message trace context propagation for persistent sessions#1128

Open
Syed-Umer-Ali wants to merge 2 commits into
anthropics:mainfrom
Syed-Umer-Ali:fix/per-message-trace-context-propagation
Open

fix: per-message trace context propagation for persistent sessions#1128
Syed-Umer-Ali wants to merge 2 commits into
anthropics:mainfrom
Syed-Umer-Ali:fix/per-message-trace-context-propagation

Conversation

@Syed-Umer-Ali

Copy link
Copy Markdown

Inject the ambient OTel trace context into every user message sent to the CLI (instead of only at connect() time), so outbound MCP/tool spans attribute to the caller's current turn trace in long-lived sessions.

Changes:

  • Add _internal/_trace_helpers.py with shared inject_trace_into_message()
  • ClaudeSDKClient.query(): inject trace context into string/AsyncIterable messages
  • InternalClient._process_query_inner(): inject trace context into initial message
  • Query.stream_input(): inject trace context into each streamed message
  • Query.set_trace_context(): new control request to update CLI trace context

Closes #1126

Inject the ambient OTel trace context into every user message sent to the
CLI (instead of only at connect() time), so outbound MCP/tool spans
attribute to the caller's current turn trace in long-lived sessions.

Changes:
- Add _internal/_trace_helpers.py with shared inject_trace_into_message()
- ClaudeSDKClient.query(): inject trace context into string/AsyncIterable messages
- InternalClient._process_query_inner(): inject trace context into initial message
- Query.stream_input(): inject trace context into each streamed message
- Query.set_trace_context(): new control request to update CLI trace context

Closes anthropics#1126
Adds 16 tests across 5 test classes:

- TestInjectTraceIntoMessage: direct unit tests for the helper
  (traceparent/tracestate injection, no-op without OTel, error handling)
- TestClaudeSDKClientQueryTraceInjection: string + AsyncIterable prompts
  via ClaudeSDKClient.query()
- TestInternalClientProcessQueryTraceInjection: string + AsyncIterable
  prompts via query()
- TestQueryStreamInputTraceInjection: per-message injection in
  Query.stream_input()
- TestQuerySetTraceContext: control request with trace context
@Syed-Umer-Ali

Copy link
Copy Markdown
Author

Reply for knoal:
Thanks for the feedback — valid points. I've added 16 tests covering:

a) Single messages: TestInjectTraceIntoMessage verifies traceparent/tracestate injection, no-op without OTel, no-op without active span, and error handling.

b) Multi-message streams: TestClaudeSDKClientQueryTraceInjection.test_async_iterable_prompt_each_message_gets_trace_context and TestQueryStreamInputTraceInjection.test_stream_input_injects_trace_context verify every message in an AsyncIterable gets its own trace context. TestInternalClientProcessQueryTraceInjection.test_async_iterable_prompt_injects_trace_context covers the query(prompt=stream) path.

c) No propagation without OTel: Both the string and iterable paths have tests verifying traceparent is absent when opentelemetry is not installed.

d) Query.set_trace_context() is tested separately, including no-op and error handling.

All 16 tests pass. PR updated.
Summary of changes pushed to PR #1128:

  • tests/test_trace_helpers.py — 16 tests across 5 classes (624 lines)
  • Covers: helper unit tests, ClaudeSDKClient.query() (string + stream), query() function (string + stream), Query.stream_input(), Query.set_trace_context()

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

1 participant