feat(telemetry): emit gen_ai.usage attrs on streaming spans

pyproject.toml

@@ -50,6 +50,7 @@ dev = [
     "bandit[toml]>=1.7.5",
     "pre-commit>=3.5.0",
     "pytest-rerunfailures>=16.1",
+    "opentelemetry-sdk>=1.30",
 ]
 
 [build-system]

src/celeste/client.py

@@ -316,8 +316,7 @@ def _stream(
             **parameters,
         )
         sse_iterator = enrich_stream_errors(sse_iterator, self._handle_error_response)
-        sse_iterator = telemetry.trace_stream(sse_iterator, span)
-        return stream_class(
+        stream = stream_class(
             sse_iterator,
             transform_output=self._transform_output,
             stream_metadata={
@@ -327,6 +326,7 @@ def _stream(
             },
             **parameters,
         )
+        return telemetry.trace_stream(stream, span)  # type: ignore[return-value]
 
     @classmethod
     @abstractmethod

src/celeste/telemetry.py

@@ -1,13 +1,17 @@
 """OpenTelemetry GenAI telemetry for Celeste."""
 
+import asyncio
 import time
-from collections.abc import AsyncIterator, Iterator
+from collections.abc import Iterator
 from contextlib import contextmanager
+from types import TracebackType
 from typing import Any
 
 from celeste.core import Modality, Protocol, Provider
+from celeste.exceptions import StreamNotExhaustedError
 from celeste.io import Output
 from celeste.models import Model
+from celeste.streaming import Stream
 
 _PROVIDER_NAME_MAP: dict[Provider, str] = {
     Provider.OPENAI: "openai",
@@ -65,14 +69,33 @@ def _noop_use_span(span: Any, end_on_exit: bool = False) -> Iterator[Any]:
     yield span
 
 
+class _NoOpStatus:
+    """Status stand-in used when ``opentelemetry-api`` is not installed."""
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        """Discard the status arguments."""
+
+
+class _NoOpStatusCode:
+    """StatusCode stand-in used when ``opentelemetry-api`` is not installed."""
+
+    ERROR = "ERROR"
+
+
 try:
     from opentelemetry import trace as _otel_trace
+    from opentelemetry.trace import Status as _OtelStatus
+    from opentelemetry.trace import StatusCode as _OtelStatusCode
 
     tracer: Any = _otel_trace.get_tracer("celeste")
     use_span: Any = _otel_trace.use_span
+    Status: Any = _OtelStatus
+    StatusCode: Any = _OtelStatusCode
 except ImportError:
     tracer = _NoOpTracer()
     use_span = _noop_use_span
+    Status = _NoOpStatus
+    StatusCode = _NoOpStatusCode
 
 
 def request_attributes(
@@ -136,31 +159,110 @@ def output_attributes(output: Output[Any]) -> dict[str, Any]:
     return attrs
 
 
-async def trace_stream(
-    sse_iterator: AsyncIterator[dict[str, Any]],
-    span: Any,
-) -> AsyncIterator[dict[str, Any]]:
-    """Wrap an SSE iterator to record TTFC and end the span at iteration's end.
-
-    OTel's ``use_span(span, end_on_exit=True)`` context manager makes the
-    span current, auto-records exceptions, sets ERROR status on exception,
-    and ends the span on exit (including ``GeneratorExit`` from consumer
-    abandonment). No try/except needed.
-    """
-    with use_span(span, end_on_exit=True):
-        started = time.monotonic()
-        seen_first = False
-        async for event in sse_iterator:
-            if not seen_first:
-                seen_first = True
-                span.set_attribute(
-                    "gen_ai.response.time_to_first_chunk",
-                    time.monotonic() - started,
-                )
-            yield event
+class _TracedStream:
+    """Stream-shaped wrapper that emits GenAI telemetry against an OTel span."""
+
+    def __init__(self, inner: Stream[Any, Any, Any], span: Any) -> None:
+        """Initialize wrapper around a constructed Stream and a detached span."""
+        self._inner = inner
+        self._span = span
+        self._started = time.monotonic()
+        self._seen_first = False
+        self._ended = False
+
+    def __aiter__(self) -> "_TracedStream":
+        """Return self as async iterator."""
+        return self
+
+    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._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(
+                "gen_ai.response.time_to_first_chunk",
+                time.monotonic() - self._started,
+            )
+        return chunk
+
+    async def __aenter__(self) -> "_TracedStream":
+        """Enter async context — return self for iteration."""
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> bool:
+        """Exit async context — close stream and finalize span."""
+        await self.aclose()
+        return False
+
+    def __iter__(self) -> Iterator[Any]:
+        """Delegate sync iteration to the inner Stream."""
+        return iter(self._inner)
+
+    def __enter__(self) -> "_TracedStream":
+        """Enter sync context — delegate to inner Stream."""
+        self._inner.__enter__()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Exit sync context — delegate to inner Stream."""
+        self._inner.__exit__(exc_type, exc_val, exc_tb)
+
+    def __repr__(self) -> str:
+        """Developer-friendly representation showing inner stream state."""
+        return repr(self._inner)
+
+    @property
+    def output(self) -> Any:
+        """Access final Output after stream exhaustion."""
+        return self._inner.output
+
+    async def aclose(self) -> None:
+        """Close the inner stream and finalize the span."""
+        await self._inner.aclose()
+        self._finalize()
+
+    def _finalize(self) -> None:
+        """End the span exactly once; emit usage attrs if Output is populated."""
+        if self._ended:
+            return
+        self._ended = True
+        try:
+            output = self._inner.output
+        except StreamNotExhaustedError:
+            output = None
+        if output is not None:
+            self._span.set_attributes(output_attributes(output))
+        self._span.end()
+
+
+def trace_stream(stream: Stream[Any, Any, Any], span: Any) -> _TracedStream:
+    """Wrap a Stream to emit GenAI telemetry against ``span``."""
+    return _TracedStream(stream, span)
 
 
 __all__ = [
+    "Status",
+    "StatusCode",
     "output_attributes",
     "request_attributes",
     "span_name",