diff --git a/src/celeste/client.py b/src/celeste/client.py index fa26508..375d701 100644 --- a/src/celeste/client.py +++ b/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, diff --git a/src/celeste/telemetry.py b/src/celeste/telemetry.py index a4fb07e..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 @@ -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", diff --git a/tests/unit_tests/test_telemetry_streaming.py b/tests/unit_tests/test_telemetry_streaming.py index ee95a01..43ad6f6 100644 --- a/tests/unit_tests/test_telemetry_streaming.py +++ b/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)