From 1f9104a7ab897a3c4df32f4163e7067c4c022879 Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Wed, 6 May 2026 15:15:29 +0200 Subject: [PATCH 1/3] fix(telemetry): drop use_span from _TracedStream.__anext__; eager-prime 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__`. --- src/celeste/client.py | 25 +++++++ src/celeste/telemetry.py | 23 ++++--- tests/unit_tests/test_telemetry_streaming.py | 69 ++++++++++++++++++++ 3 files changed, 105 insertions(+), 12 deletions(-) diff --git a/src/celeste/client.py b/src/celeste/client.py index fa26508..808fd4d 100644 --- a/src/celeste/client.py +++ b/src/celeste/client.py @@ -1,5 +1,7 @@ """Base client for modality-specific AI operations.""" +import asyncio +import contextvars import warnings from abc import ABC, abstractmethod from collections.abc import AsyncIterator @@ -28,6 +30,25 @@ from celeste.types import RawUsage +async def _prime_with_context( + inner: AsyncIterator[dict[str, Any]], + ctx: contextvars.Context, +) -> AsyncIterator[dict[str, Any]]: + """Run inner's first pull under ctx; delegate the rest to the caller's context.""" + + async def _first() -> dict[str, Any]: + return await inner.__anext__() + + task = asyncio.create_task(_first(), context=ctx) + try: + first = await task + except StopAsyncIteration: + return + yield first + async for event in inner: + yield event + + class APIMixin(ABC): """Abstract base for provider API mixins. @@ -313,6 +334,10 @@ def _stream( **parameters, ) sse_iterator = enrich_stream_errors(sse_iterator, self._handle_error_response) + # Capture span-active OTel context so the first SSE pull fires HTTP under it + with telemetry.use_span(span): + ctx_with_span = contextvars.copy_context() + sse_iterator = _prime_with_context(sse_iterator, ctx_with_span) stream = stream_class( sse_iterator, transform_output=self._transform_output, diff --git a/src/celeste/telemetry.py b/src/celeste/telemetry.py index a4fb07e..57a3d6a 100644 --- a/src/celeste/telemetry.py +++ b/src/celeste/telemetry.py @@ -435,18 +435,17 @@ def __aiter__(self) -> "_TracedStream": async def __anext__(self) -> Any: """Yield next chunk; emit TTFC on first, finalize span on terminal events.""" - with use_span(self._span, end_on_exit=False): - try: - chunk = await self._inner.__anext__() - except (StopAsyncIteration, asyncio.CancelledError): - self._finalize() - raise - except Exception as exc: - self._error = exc - self._span.record_exception(exc) - self._span.set_status(Status(StatusCode.ERROR, str(exc))) - self._finalize() - raise + try: + chunk = await self._inner.__anext__() + except (StopAsyncIteration, asyncio.CancelledError): + self._finalize() + raise + except Exception as exc: + self._error = exc + self._span.record_exception(exc) + self._span.set_status(Status(StatusCode.ERROR, str(exc))) + self._finalize() + raise if not self._seen_first: self._seen_first = True self._span.set_attribute( diff --git a/tests/unit_tests/test_telemetry_streaming.py b/tests/unit_tests/test_telemetry_streaming.py index ee95a01..a625241 100644 --- a/tests/unit_tests/test_telemetry_streaming.py +++ b/tests/unit_tests/test_telemetry_streaming.py @@ -1,7 +1,9 @@ """Tests for `_TracedStream` — span lifecycle and GenAI attribute emission.""" +import contextvars from collections.abc import AsyncIterator from typing import Any +from unittest.mock import patch import pytest from opentelemetry.sdk.trace import TracerProvider @@ -11,6 +13,7 @@ from opentelemetry.trace import StatusCode from celeste import telemetry +from celeste.client import _prime_with_context from celeste.exceptions import StreamNotExhaustedError from celeste.io import Output, Usage from tests.unit_tests._telemetry_helpers import ( @@ -235,3 +238,69 @@ def test_response_model_emitted_from_metadata(self) -> None: attrs = telemetry.output_attributes(output) assert attrs["gen_ai.response.model"] == "claude-opus-4-1-20250805" + + +class TestNoSpanActivationInAnext: + """Regression: `_TracedStream.__anext__` must not activate the span across `await`.""" + + async def test_anext_does_not_call_use_span( + self, exporter: tuple[InMemorySpanExporter, TracerProvider] + ) -> None: + """Iterating a `_TracedStream` does NOT invoke `telemetry.use_span`.""" + _, provider = exporter + events = [{"delta": "a"}, {"delta": "b"}] + wrapped = telemetry.trace_stream( + TelemetryStream(async_iter(events)), start_test_span(provider) + ) + + with patch("celeste.telemetry.use_span") as mock_use_span: + async for _ in wrapped: + pass + + mock_use_span.assert_not_called() + + +class TestPrimeWithContext: + """`_prime_with_context` runs the first pull under ctx and delegates the rest.""" + + async def test_preserves_event_order(self) -> None: + """All events from inner are yielded in original order.""" + events = [{"i": 0}, {"i": 1}, {"i": 2}] + primed = _prime_with_context(async_iter(events), contextvars.copy_context()) + + collected = [event async for event in primed] + + assert collected == events + + async def test_empty_stream_yields_nothing(self) -> None: + """Inner that immediately raises StopAsyncIteration yields no events.""" + primed = _prime_with_context(async_iter([]), contextvars.copy_context()) + + collected = [event async for event in primed] + + assert collected == [] + + async def test_first_error_propagates(self) -> None: + """Exception raised by inner's first pull propagates to the consumer.""" + primed = _prime_with_context(_failing_iter(), contextvars.copy_context()) + + with pytest.raises(RuntimeError, match="boom"): + async for _ in primed: + pass + + async def test_first_pull_runs_under_captured_ctx(self) -> None: + """First inner pull executes under captured ctx; subsequent pulls under caller's ctx.""" + marker: contextvars.ContextVar[str] = contextvars.ContextVar("marker") + marker.set("captured") + captured_ctx = contextvars.copy_context() + marker.set("caller") + + async def inner() -> AsyncIterator[dict[str, str]]: + yield {"value": marker.get()} + yield {"value": marker.get()} + + primed = _prime_with_context(inner(), captured_ctx) + collected = [event async for event in primed] + + assert collected[0] == {"value": "captured"} + assert collected[1] == {"value": "caller"} From e9e12c6066c1d2eb1dc1b17c56ab8459f0909aac Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Wed, 6 May 2026 15:24:20 +0200 Subject: [PATCH 2/3] fix(telemetry): cancel prime task when consumer is cancelled before first 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. --- src/celeste/client.py | 3 +++ tests/unit_tests/test_telemetry_streaming.py | 18 ++++++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/src/celeste/client.py b/src/celeste/client.py index 808fd4d..9a2229d 100644 --- a/src/celeste/client.py +++ b/src/celeste/client.py @@ -44,6 +44,9 @@ async def _first() -> dict[str, Any]: first = await task except StopAsyncIteration: return + except BaseException: + task.cancel() + raise yield first async for event in inner: yield event diff --git a/tests/unit_tests/test_telemetry_streaming.py b/tests/unit_tests/test_telemetry_streaming.py index a625241..43b3585 100644 --- a/tests/unit_tests/test_telemetry_streaming.py +++ b/tests/unit_tests/test_telemetry_streaming.py @@ -1,5 +1,6 @@ """Tests for `_TracedStream` — span lifecycle and GenAI attribute emission.""" +import asyncio import contextvars from collections.abc import AsyncIterator from typing import Any @@ -304,3 +305,20 @@ async def inner() -> AsyncIterator[dict[str, str]]: assert collected[0] == {"value": "captured"} assert collected[1] == {"value": "caller"} + + async def test_cancel_during_first_pull_cancels_inner_task(self) -> None: + """Cancellation while awaiting the first pull propagates to the inner task — no leak.""" + started = asyncio.Event() + + async def slow_inner() -> AsyncIterator[dict[str, str]]: + started.set() + await asyncio.sleep(60) + yield {"v": "never"} + + primed = _prime_with_context(slow_inner(), contextvars.copy_context()) + consumer: asyncio.Task[dict[str, str]] = asyncio.create_task(primed.__anext__()) + await started.wait() + consumer.cancel() + with pytest.raises(asyncio.CancelledError): + await consumer + await asyncio.sleep(0) From 0fdf8bb387551dae7ae984b63da7e4a77cc606b2 Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Wed, 6 May 2026 16:24:21 +0200 Subject: [PATCH 3/3] refactor(telemetry): move bind_first_pull_to_span helper from client.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). --- src/celeste/client.py | 29 +------- src/celeste/telemetry.py | 28 ++++++- tests/unit_tests/test_telemetry_streaming.py | 78 ++++++++++++-------- 3 files changed, 77 insertions(+), 58 deletions(-) diff --git a/src/celeste/client.py b/src/celeste/client.py index 9a2229d..375d701 100644 --- a/src/celeste/client.py +++ b/src/celeste/client.py @@ -1,7 +1,5 @@ """Base client for modality-specific AI operations.""" -import asyncio -import contextvars import warnings from abc import ABC, abstractmethod from collections.abc import AsyncIterator @@ -30,28 +28,6 @@ from celeste.types import RawUsage -async def _prime_with_context( - inner: AsyncIterator[dict[str, Any]], - ctx: contextvars.Context, -) -> AsyncIterator[dict[str, Any]]: - """Run inner's first pull under ctx; delegate the rest to the caller's context.""" - - async def _first() -> dict[str, Any]: - return await inner.__anext__() - - task = asyncio.create_task(_first(), context=ctx) - try: - first = await task - except StopAsyncIteration: - return - except BaseException: - task.cancel() - raise - yield first - async for event in inner: - yield event - - class APIMixin(ABC): """Abstract base for provider API mixins. @@ -337,10 +313,7 @@ def _stream( **parameters, ) sse_iterator = enrich_stream_errors(sse_iterator, self._handle_error_response) - # Capture span-active OTel context so the first SSE pull fires HTTP under it - with telemetry.use_span(span): - ctx_with_span = contextvars.copy_context() - sse_iterator = _prime_with_context(sse_iterator, ctx_with_span) + sse_iterator = telemetry.bind_first_pull_to_span(sse_iterator, span) stream = stream_class( sse_iterator, transform_output=self._transform_output, diff --git a/src/celeste/telemetry.py b/src/celeste/telemetry.py index 57a3d6a..f2a078b 100644 --- a/src/celeste/telemetry.py +++ b/src/celeste/telemetry.py @@ -1,10 +1,11 @@ """OpenTelemetry GenAI telemetry for Celeste.""" import asyncio +import contextvars import json import os import time -from collections.abc import Iterator +from collections.abc import AsyncIterator, Iterator from contextlib import contextmanager, suppress from types import TracebackType from typing import Any @@ -540,10 +541,35 @@ def trace_stream( return _TracedStream(stream, span, metric_attributes) +async def bind_first_pull_to_span( + inner: AsyncIterator[dict[str, Any]], + span: Any, +) -> AsyncIterator[dict[str, Any]]: + """Run inner's first pull under a context where span is active; delegate the rest.""" + with use_span(span): + ctx = contextvars.copy_context() + + async def _first() -> dict[str, Any]: + return await inner.__anext__() + + task = asyncio.create_task(_first(), context=ctx) + try: + first = await task + except StopAsyncIteration: + return + except BaseException: + task.cancel() + raise + yield first + async for event in inner: + yield event + + __all__ = [ "Status", "StatusCode", "add_input_event", + "bind_first_pull_to_span", "gen_ai_span", "meter", "output_attributes", diff --git a/tests/unit_tests/test_telemetry_streaming.py b/tests/unit_tests/test_telemetry_streaming.py index 43b3585..43ad6f6 100644 --- a/tests/unit_tests/test_telemetry_streaming.py +++ b/tests/unit_tests/test_telemetry_streaming.py @@ -1,12 +1,12 @@ """Tests for `_TracedStream` — span lifecycle and GenAI attribute emission.""" import asyncio -import contextvars from collections.abc import AsyncIterator from typing import Any from unittest.mock import patch import pytest +from opentelemetry import trace as otel_trace from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( InMemorySpanExporter, @@ -14,7 +14,6 @@ from opentelemetry.trace import StatusCode from celeste import telemetry -from celeste.client import _prime_with_context from celeste.exceptions import StreamNotExhaustedError from celeste.io import Output, Usage from tests.unit_tests._telemetry_helpers import ( @@ -261,53 +260,72 @@ async def test_anext_does_not_call_use_span( mock_use_span.assert_not_called() -class TestPrimeWithContext: - """`_prime_with_context` runs the first pull under ctx and delegates the rest.""" +class TestBindFirstPullToSpan: + """`bind_first_pull_to_span` runs the first pull with span active; delegates the rest.""" - async def test_preserves_event_order(self) -> None: + async def test_preserves_event_order( + self, exporter: tuple[InMemorySpanExporter, TracerProvider] + ) -> None: """All events from inner are yielded in original order.""" + _, provider = exporter events = [{"i": 0}, {"i": 1}, {"i": 2}] - primed = _prime_with_context(async_iter(events), contextvars.copy_context()) + bound = telemetry.bind_first_pull_to_span( + async_iter(events), start_test_span(provider) + ) - collected = [event async for event in primed] + collected = [event async for event in bound] assert collected == events - async def test_empty_stream_yields_nothing(self) -> None: + async def test_empty_stream_yields_nothing( + self, exporter: tuple[InMemorySpanExporter, TracerProvider] + ) -> None: """Inner that immediately raises StopAsyncIteration yields no events.""" - primed = _prime_with_context(async_iter([]), contextvars.copy_context()) + _, provider = exporter + bound = telemetry.bind_first_pull_to_span( + async_iter([]), start_test_span(provider) + ) - collected = [event async for event in primed] + collected = [event async for event in bound] assert collected == [] - async def test_first_error_propagates(self) -> None: + async def test_first_error_propagates( + self, exporter: tuple[InMemorySpanExporter, TracerProvider] + ) -> None: """Exception raised by inner's first pull propagates to the consumer.""" - primed = _prime_with_context(_failing_iter(), contextvars.copy_context()) + _, provider = exporter + bound = telemetry.bind_first_pull_to_span( + _failing_iter(), start_test_span(provider) + ) with pytest.raises(RuntimeError, match="boom"): - async for _ in primed: + async for _ in bound: pass - async def test_first_pull_runs_under_captured_ctx(self) -> None: - """First inner pull executes under captured ctx; subsequent pulls under caller's ctx.""" - marker: contextvars.ContextVar[str] = contextvars.ContextVar("marker") - marker.set("captured") - captured_ctx = contextvars.copy_context() - marker.set("caller") + async def test_first_pull_runs_with_span_active( + self, exporter: tuple[InMemorySpanExporter, TracerProvider] + ) -> None: + """First inner pull sees span as current OTel context; subsequent pulls don't.""" + _, provider = exporter + span = start_test_span(provider) + target_id = span.get_span_context().span_id - async def inner() -> AsyncIterator[dict[str, str]]: - yield {"value": marker.get()} - yield {"value": marker.get()} + async def inner() -> AsyncIterator[dict[str, int]]: + yield {"span_id": otel_trace.get_current_span().get_span_context().span_id} + yield {"span_id": otel_trace.get_current_span().get_span_context().span_id} - primed = _prime_with_context(inner(), captured_ctx) - collected = [event async for event in primed] + bound = telemetry.bind_first_pull_to_span(inner(), span) + collected = [event async for event in bound] - assert collected[0] == {"value": "captured"} - assert collected[1] == {"value": "caller"} + assert collected[0]["span_id"] == target_id + assert collected[1]["span_id"] != target_id - async def test_cancel_during_first_pull_cancels_inner_task(self) -> None: + async def test_cancel_during_first_pull_cancels_inner_task( + self, exporter: tuple[InMemorySpanExporter, TracerProvider] + ) -> None: """Cancellation while awaiting the first pull propagates to the inner task — no leak.""" + _, provider = exporter started = asyncio.Event() async def slow_inner() -> AsyncIterator[dict[str, str]]: @@ -315,8 +333,10 @@ async def slow_inner() -> AsyncIterator[dict[str, str]]: await asyncio.sleep(60) yield {"v": "never"} - primed = _prime_with_context(slow_inner(), contextvars.copy_context()) - consumer: asyncio.Task[dict[str, str]] = asyncio.create_task(primed.__anext__()) + bound = telemetry.bind_first_pull_to_span( + slow_inner(), start_test_span(provider) + ) + consumer: asyncio.Task[dict[str, str]] = asyncio.create_task(bound.__anext__()) await started.wait() consumer.cancel() with pytest.raises(asyncio.CancelledError):