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

[Kamilbenkirane]

Summary

Streaming spans now carry gen_ai.usage.input_tokens and gen_ai.usage.output_tokens — matching the non-streaming path. The telemetry wrap site moves from the upstream SSE iterator to the constructed Stream so that output_attributes(stream.output) can be called after _parse_output populates the typed Output.

Why

telemetry.trace_stream previously wrapped the SSE iterator (below Stream). When upstream SSE exhausted, the wrapper exited and with use_span(span, end_on_exit=True) ended the span — before Stream.__anext__ ran _parse_output to build _output.usage. Setting attributes on an already-ended span is a silent no-op, so streaming spans only carried gen_ai.response.time_to_first_chunk.

Stream._output.usage exists synchronously after natural exhaustion. The fix is positional: wrap the Stream instead of the SSE iterator. Same output_attributes(output) function _predict already uses gets called, on the same typed Output shape.

Changes

• src/celeste/telemetry.py: replace the iterator-shaped trace_stream async generator with a _TracedStream wrapper class that delegates the public Stream surface (__aiter__/__anext__/aclose/output/__aenter__/__aexit__/__iter__/__enter__/__exit__/__repr__) and owns span lifecycle. Add _NoOpStatus / _NoOpStatusCode for the existing OTel-missing fallback path.
• src/celeste/client.py: in _stream, construct the Stream from the SSE iterator and wrap the Stream (not the SSE iterator) with telemetry.trace_stream.
• tests/unit_tests/test_telemetry_streaming.py: 7 new tests covering natural exhaustion (+ async with), early aclose(), exception path, output property pre/post exhaustion, idempotent finalize.
• pyproject.toml: add opentelemetry-sdk to dev deps for InMemorySpanExporter-based unit tests.

Zero changes to src/celeste/streaming.py.

Lifecycle

• First chunk: emit gen_ai.response.time_to_first_chunk.
• Natural exhaustion: output_attributes(stream.output) → set_attributes → span.end().
• Early aclose(): end span without ERROR; emit usage only if _output was built before abandonment.
• Exception during iteration: record_exception + ERROR status + end span + re-raise.
• asyncio.CancelledError: end span without ERROR; re-raise.

use_span(span, end_on_exit=False) is held only across the inner await __anext__() so HTTP child spans (HTTPX auto-instrumentation) nest under the GenAI span.

Behavioral note

_stream returns the wrapper, not a Stream subclass. isinstance(x, Stream) checks against the return value will fail. Verified no such checks exist in this repo. Consumers using async for / .output / aclose() / async with continue to work via delegation.

Out of scope

Widening output_attributes to emit total_tokens, reasoning_tokens, modality-specific usage fields, cached-input tokens, or GenAI metrics — tracked separately in #271.

Test plan

• uv run ruff check src tests
• uv run ruff format --check src tests
• uv run mypy src/celeste and uv run mypy tests/
• uv run pytest tests/unit_tests/ -m "not integration" -q — 611 passed

[claude]

Code review

Bug: sync iteration bypasses span lifecycle (telemetry.py lines 212-228)

_TracedStream.iter returns iter(self._inner) directly, and exit delegates to self._inner.exit. This means OTel span lifecycle is skipped entirely for synchronous iteration.

Root cause: Stream.iter (streaming.py:303-322) is its own generator driving Stream.anext via an anyio portal. It never calls _TracedStream.anext, so _finalize() is never reached. The inner Stream.exit is a no-op.

Consequences:

• Span leak: the OTel span from client._stream is never ended and never exported.
• No gen_ai.usage.input_tokens / gen_ai.usage.output_tokens / TTFC emitted for sync streaming — defeating the purpose of this PR for sync callers.
• record_exception / ERROR status never set on sync-path failures.

Regression: Pre-PR, trace_stream wrapped the SSE iterator passed into the Stream constructor, so Stream.anext (used by the sync portal) activated the span lifecycle on both paths. The new approach only wires tracing on async paths.

Sync iteration is documented in the Stream docstring and exercised by 8+ integration tests (tests/integration_tests//test_stream_.py).

Ref:

celeste-python/src/celeste/telemetry.py

Lines 211 to 229 in 49b18a3

	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)

Suggested fix: iter should be a generator wrapping the inner sync iterator, recording TTFC, calling _finalize() on exhaustion/GeneratorExit/exceptions — mirroring anext. exit should also call _finalize() so the with-stream pattern ends the span.