fix(telemetry): drop use_span from _TracedStream.__anext__; eager-prime SSE under captured ctx (#277)

* 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__`.

* 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.

* 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).

Author: Kamilbenkirane

src/celeste/client.py

@@ -313,6 +313,7 @@ def _stream(
             **parameters,
         )
         sse_iterator = enrich_stream_errors(sse_iterator, self._handle_error_response)
+        sse_iterator = telemetry.bind_first_pull_to_span(sse_iterator, span)
         stream = stream_class(
             sse_iterator,
             transform_output=self._transform_output,

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
@@ -435,18 +436,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(
@@ -541,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",

tests/unit_tests/test_telemetry_streaming.py

@@ -1,9 +1,12 @@
 """Tests for `_TracedStream` — span lifecycle and GenAI attribute emission."""
 
+import asyncio
 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,
@@ -235,3 +238,107 @@ 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 TestBindFirstPullToSpan:
+    """`bind_first_pull_to_span` runs the first pull with span active; delegates the rest."""
+
+    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}]
+        bound = telemetry.bind_first_pull_to_span(
+            async_iter(events), start_test_span(provider)
+        )
+
+        collected = [event async for event in bound]
+
+        assert collected == events
+
+    async def test_empty_stream_yields_nothing(
+        self, exporter: tuple[InMemorySpanExporter, TracerProvider]
+    ) -> None:
+        """Inner that immediately raises StopAsyncIteration yields no events."""
+        _, provider = exporter
+        bound = telemetry.bind_first_pull_to_span(
+            async_iter([]), start_test_span(provider)
+        )
+
+        collected = [event async for event in bound]
+
+        assert collected == []
+
+    async def test_first_error_propagates(
+        self, exporter: tuple[InMemorySpanExporter, TracerProvider]
+    ) -> None:
+        """Exception raised by inner's first pull propagates to the consumer."""
+        _, provider = exporter
+        bound = telemetry.bind_first_pull_to_span(
+            _failing_iter(), start_test_span(provider)
+        )
+
+        with pytest.raises(RuntimeError, match="boom"):
+            async for _ in bound:
+                pass
+
+    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, 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}
+
+        bound = telemetry.bind_first_pull_to_span(inner(), span)
+        collected = [event async for event in bound]
+
+        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, 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]]:
+            started.set()
+            await asyncio.sleep(60)
+            yield {"v": "never"}
+
+        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):
+            await consumer
+        await asyncio.sleep(0)