Make span generic wrt span_data

Author: anbuzin

src/ai/agents/agent.py

@@ -641,9 +641,9 @@ async def __call__(self, **overrides: Any) -> events_.ToolCallResult:
         # re-executing the tool.
         cached = self._part.cached_result
         if cached is not None:
-            async with telemetry.span(data, replay=True):
-                data.result = cached.result
-                data.is_error = cached.is_error
+            async with telemetry.span(data, replay=True) as sp:
+                sp.data.result = cached.result
+                sp.data.is_error = cached.is_error
                 msg = builders.tool_message(cached)
                 return events_.ToolCallResult(
                     message=msg, results=msg.tool_results
@@ -717,8 +717,8 @@ async def _real(
         async with telemetry.span(data) as sp:
             res = await chain(call)
             if res.results:
-                data.result = res.results[0].result
-                data.is_error = any(p.is_error for p in res.results)
+                sp.data.result = res.results[0].result
+                sp.data.is_error = any(p.is_error for p in res.results)
             # A tool exception is caught and converted to an error
             # result before it reaches this block, so it never hits the
             # span's own except path — thread it through explicitly.
@@ -1401,19 +1401,20 @@ async def _real(call: Context) -> AsyncGenerator[events_.AgentEvent]:
                         yield transition
 
         async def _stream() -> AsyncGenerator[events_.AgentEvent]:
-            run_data = telemetry.RunSpanData(
-                agent=type(self).__name__,
-                model=model.id,
-                messages=list(messages),
-            )
             # Activate middleware for this run (and everything it calls).
             # When middleware is None (default), inherit the parent's
             # middleware from the context var — this lets nested agents
             # share middleware.  When middleware is explicitly provided,
             # *extend* the parent stack so that outer cross-cutting
             # concerns (tracing, durability) are preserved.  Pass
             # ``_middleware=[]`` to clear the stack entirely.
-            async with telemetry.span(run_data):
+            async with telemetry.span(
+                telemetry.RunSpanData(
+                    agent=type(self).__name__,
+                    model=model.id,
+                    messages=list(messages),
+                )
+            ):
                 mw_token: middleware_.Token | None = None
                 if _middleware is not None:
                     parent = middleware_.get()

src/ai/agents/hooks.py

@@ -131,10 +131,10 @@ async def _hook_impl(call: middleware_.HookContext) -> pydantic.BaseModel:
     # Pre-registered resolution (serverless re-entry).
     pre_registered = _pending_resolutions.pop(label, None)
     if pre_registered is not None:
-        async with telemetry.span(data, replay=True):
+        async with telemetry.span(data, replay=True) as sp:
             if isinstance(pre_registered, BaseException):
                 raise pre_registered
-            data.status = "resolved"
+            sp.data.status = "resolved"
             return payload(**pre_registered)
 
     # No resolution available — suspend.  The span covers the whole
@@ -161,7 +161,7 @@ async def _hook_impl(call: middleware_.HookContext) -> pydantic.BaseModel:
         try:
             resolution = await future
         except asyncio.CancelledError as exc:
-            data.status = "cancelled"
+            sp.data.status = "cancelled"
             # ``cancel_hook(reason=...)`` rides on the CancelledError.
             attrs: dict[str, Any] = {}
             if exc.args and exc.args[0] is not None:
@@ -171,7 +171,7 @@ async def _hook_impl(call: middleware_.HookContext) -> pydantic.BaseModel:
 
         # Clean up live registry.
         _live_hooks.pop(label, None)
-        data.status = "resolved"
+        sp.data.status = "resolved"
         await sp.add_event(telemetry.HOOK_RESOLVED)
 
         # Emit resolved signal.

src/ai/models/core/api.py

@@ -137,7 +137,7 @@ def __init__(
         # The telemetry span bracketing this stream, attached by
         # ``stream()``.  None for directly constructed streams
         # (``Stream(gen)``, ``Stream.replay_message``).
-        self._span: telemetry.Span | None = None
+        self._span: telemetry.Span[telemetry.AiStreamSpanData] | None = None
         self._first_output_seen = False
 
     @classmethod
@@ -244,7 +244,7 @@ async def __anext__(self: Self) -> types.events.Event:
         return event.model_copy(update={"message": self._message, **updates})
 
     @property
-    def span(self) -> telemetry.Span | None:
+    def span(self) -> telemetry.Span[telemetry.AiStreamSpanData] | None:
         """The telemetry span bracketing this stream.
 
         Set when the stream came from :func:`stream` (live and replay
@@ -608,8 +608,8 @@ async def _stream(
             yield s
         finally:
             # Record whatever got built, even a partial message.
-            data.message = s.message
-            data.usage = s.usage
+            sp.data.message = s.message
+            sp.data.usage = s.usage
             await s.aclose()
 
 
@@ -623,12 +623,13 @@ async def generate(
     """Generate a non-streaming response (images, video, etc.)."""
     messages = integrity.prepare_messages(messages)
     request = GenerateRequest(model, messages, params)
-    data = telemetry.AiGenerateSpanData(
-        model=model.id, messages=messages, params=params
-    )
-    async with telemetry.span(data):
+    async with telemetry.span(
+        telemetry.AiGenerateSpanData(
+            model=model.id, messages=messages, params=params
+        )
+    ) as sp:
         message = await executor._do_generate(request)
-        data.message = message
+        sp.data.message = message
         return message
 
 

src/ai/telemetry/span.py

@@ -40,7 +40,19 @@ async def on_span_event(self, span, event): ...
 import inspect
 import logging
 import time
-from typing import TYPE_CHECKING, Any, ClassVar, Literal, Protocol
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    ClassVar,
+    Generic,
+    Literal,
+    Protocol,
+    overload,
+)
+
+# ``typing.TypeVar`` lacks the ``default=`` kwarg on Python <3.13.
+# Use the typing_extensions backport so this works on 3.12 too.
+from typing_extensions import TypeVar
 
 from ..types import messages as messages_
 
@@ -81,6 +93,15 @@ class RetrievalSpanData:
     span_name: ClassVar[str]
 
 
+# Covariant: adapters take ``Span`` (= ``Span[SpanData]``) and must
+# accept any concretely-typed span.  They read ``data``, never replace
+# it, so the widened reference is safe in practice.
+DataT_co = TypeVar("DataT_co", bound=SpanData, default=SpanData, covariant=True)
+# Function-scoped variant for ``span()``: a covariant variable can't
+# appear as a parameter.
+DataT = TypeVar("DataT", bound=SpanData)
+
+
 @dataclasses.dataclass
 class RunSpanData:
     """One ``Agent.run``: the whole loop."""
@@ -181,9 +202,13 @@ class SpanEvent:
 
 
 @dataclasses.dataclass
-class Span:
+class Span(Generic[DataT_co]):
     """A record of a unit of work.
 
+    Generic in its data type: ``span(RetrievalSpanData(...))`` gives a
+    ``Span[RetrievalSpanData]``, so late assignments to ``sp.data``
+    fields are type checked.  A bare ``Span`` is ``Span[SpanData]``.
+
     ``replay=True`` marks work that is being replayed (resume,
     serverless re-entry) rather than performed live.
 
@@ -197,7 +222,7 @@ class Span:
     """
 
     name: str
-    data: SpanData
+    data: DataT_co
     id: str
     trace_id: str
     parent_id: str | None
@@ -285,25 +310,65 @@ async def _dispatch(method: str, span_: Span, *args: Any) -> None:
             )
 
 
-@contextlib.asynccontextmanager
-async def span(
+@overload
+def span(
+    name_or_data: str,
+    /,
+    *,
+    replay: bool = False,
+    set_as_current: bool = True,
+    **attributes: Any,
+) -> contextlib.AbstractAsyncContextManager[Span[CustomSpanData]]: ...
+
+
+@overload
+def span(
+    name_or_data: DataT,
+    /,
+    *,
+    replay: bool = False,
+    set_as_current: bool = True,
+) -> contextlib.AbstractAsyncContextManager[Span[DataT]]: ...
+
+
+def span(
     name_or_data: str | SpanData,
     /,
     *,
     replay: bool = False,
     set_as_current: bool = True,
     **attributes: Any,
-) -> AsyncIterator[Span]:
+) -> contextlib.AbstractAsyncContextManager[Span[Any]]:
     """Open a span; it is "current" (parents new spans) inside the block.
 
     Pass a name plus attributes for a user span, or a :class:`SpanData`
-    instance for a typed one.  Exceptions are recorded on the span and
-    re-raised.
+    instance for a typed one — the span is generic in it, so late
+    assignments to ``sp.data`` fields are type checked.  Exceptions are
+    recorded on the span and re-raised.
 
     ``set_as_current=False`` keeps the span from becoming current:
     work done while it is open parents to *its* parent instead. Used by
     ai.stream because of the context manager api.
     """
+    # The indirection exists because type checkers can't apply
+    # ``asynccontextmanager`` to an overloaded function directly.
+    return _span_impl(
+        name_or_data,
+        replay=replay,
+        set_as_current=set_as_current,
+        **attributes,
+    )
+
+
+@contextlib.asynccontextmanager
+async def _span_impl(
+    name_or_data: str | SpanData,
+    /,
+    *,
+    replay: bool,
+    set_as_current: bool,
+    **attributes: Any,
+) -> AsyncIterator[Span[Any]]:
     if isinstance(name_or_data, str):
         name = name_or_data
         data: SpanData = CustomSpanData(attributes=dict(attributes))

tests/telemetry/test_telemetry.py

@@ -68,7 +68,13 @@ async def test_set_rejects_framework_spans() -> None:
 
 async def test_data_with_attributes_rejected() -> None:
     with pytest.raises(TypeError):
-        async with ai.telemetry.span(ai.telemetry.LoopTurnSpanData(), a=1):
+        # The overloads reject this statically too; the runtime check
+        # covers untyped callers.
+        data = ai.telemetry.LoopTurnSpanData()
+        async with ai.telemetry.span(
+            data,  # ty: ignore[invalid-argument-type]
+            a=1,  # type: ignore[call-overload]
+        ):
             pass