fix(telemetry): drop use_span from _TracedStream.__anext__; eager-prime SSE under captured ctx#277
Merged
Merged
Conversation
…me SSE under captured ctx `_TracedStream.__anext__` wrapped `await self._inner.__anext__()` in `with use_span(self._span, end_on_exit=False):` to make the GenAI span the parent of HTTPX child spans. Under FastAPI StreamingResponse the attach/detach pair fired across asyncio tasks, raising `ValueError: Token was created in a different Context` ~19x per chat session. Closes #276. Capture an OTel context snapshot with the GenAI span active in `_stream` (sync `with use_span`, no await between attach and detach). Schedule the SSE iterator's first __anext__ as `asyncio.create_task(coro, context=ctx)` so the HTTP request fires under the snapshot; HTTPX auto-instrumentation captures the GenAI span as parent at request emission. Subsequent reads run in the consumer's own context and pull bytes from the open response. Drop the `with use_span` block from `_TracedStream.__anext__`.
…irst chunk Without this, a cancellation interrupting `await task` in `_prime_with_context` left the underlying create_task running in the background, holding the open SSE/HTTP connection. Add a `BaseException` catch that calls `task.cancel()` then re-raises. Add regression test covering cancel-mid-first-pull.
Code reviewNo issues found. Checked for bugs and CLAUDE.md compliance. |
…py to telemetry.py
The first pass parked an async-task + contextvars helper at the top of
client.py and added asyncio + contextvars imports there. That leaked
telemetry mechanism into the module that defines the SDK base layering.
Move the helper to telemetry.py (where it joins trace_stream / record_*
/ use_span and where contextvars belongs); rename to
bind_first_pull_to_span; revert client.py imports and drop the inline
ctx-capture block from `_stream`. `_stream` regains its single-verb
orchestration rhythm with one new line:
sse_iterator = telemetry.bind_first_pull_to_span(sse_iterator, span)
Same behavior, same tests (signatures updated to pass a span instead of
a Context; new test asserts the first pull sees the span as current
OTel context and subsequent pulls don't).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
_TracedStream.__anext__wrappedawait self._inner.__anext__()inwith use_span(self._span, end_on_exit=False):so HTTPX child spans would parent to the GenAI span. Under async iteration that resumes across asyncio tasks, theattach/detachpair fired across mismatched contexts and raisedValueError: Token was created in a different Context. Closes _TracedStream.__anext__'s use_span context manager raises ValueError on detach across asyncio context boundaries #276.with use_spanblock from__anext__entirely. Capture an OTel context snapshot with the GenAI span active inside_stream(syncwith use_span+contextvars.copy_context()— noawaitbetween attach and detach, no Task switch) and schedule the SSE iterator's first__anext__asasyncio.create_task(coro, context=ctx). The HTTP request fires under that snapshot, so HTTPX auto-instrumentation captures the GenAI span as parent at request emission. Subsequent reads run in the consumer's own context and pull bytes from the open response — no new HTTPX spans, no per-chunk context manipulation.client.py) plus the__anext__cleanup (telemetry.py). Zero changes to provider mixins,http.py,_make_stream_requestsignature,Stream, or any public API.Pattern aligns with
opentelemetry-instrumentation-openai-v2'sBaseStreamWrapper.__anext__, which has zerouse_span/attach/detachcalls.References #272 (where the workaround originated).
Test plan
make lintcleanmake format --checkcleanmake typecheckcleanmake test— 637 unit tests pass (5 new tests intests/unit_tests/test_telemetry_streaming.pycovering:__anext__does not invokeuse_span, prime preserves event order, empty stream yields nothing, first-pull error propagates, first pull runs under captured ctx with subsequent pulls under caller's ctx, cancel-mid-first-pull cancels the prime task)