Add extra properties to spans to report more stuff

src/ai/agents/agent.py

@@ -1441,8 +1441,16 @@ async def _stream() -> AsyncGenerator[events_.AgentEvent]:
                     agent=type(self).__name__,
                     model=model.id,
                     messages=list(messages),
+                    provider=model.provider.name,
+                    tool_names=[t.name for t in context.tools],
+                    output_type=(
+                        output_type.__name__
+                        if output_type is not None
+                        else None
+                    ),
+                    params=params,
                 )
-            ):
+            ) as sp:
                 mw_token: middleware_.Token | None = None
                 if _middleware is not None:
                     parent = middleware_.get()
@@ -1451,10 +1459,30 @@ async def _stream() -> AsyncGenerator[events_.AgentEvent]:
                     chain = middleware_._build_agent_run_chain(_real)
                     async with contextlib.aclosing(chain(context)) as events:
                         async for event in events:
+                            # a blocked run can only resume via a hook
+                            # resolution or cancellation, so a non-pending
+                            # HookEvent is also the unblock signal.
+                            if isinstance(event, events_.RunBlocked):
+                                sp.data.blocked = True
+                            elif (
+                                isinstance(event, events_.HookEvent)
+                                and event.hook.status != "pending"
+                            ):
+                                sp.data.blocked = False
                             yield event
                 finally:
                     if mw_token is not None:
                         middleware_.deactivate(mw_token)
+                    # Record whatever got produced, even on error or
+                    # early close.
+                    sp.data.final_message = next(
+                        (
+                            msg
+                            for msg in reversed(context.messages)
+                            if msg.role == "assistant"
+                        ),
+                        None,
+                    )
 
         async with (
             hooks_.use_hook_registry(registry),

src/ai/agents/hooks.py

@@ -195,7 +195,10 @@ async def _hook_impl(call: middleware_.HookContext) -> pydantic.BaseModel:
     hook_metadata = call.metadata
 
     data = telemetry.HookSpanData(
-        label=label, hook_type=payload.__name__, metadata=hook_metadata
+        label=label,
+        hook_type=payload.__name__,
+        metadata=hook_metadata,
+        tool_call_id=call.tool_call_id,
     )
 
     # Pre-registered resolution (serverless re-entry).
@@ -205,6 +208,7 @@ async def _hook_impl(call: middleware_.HookContext) -> pydantic.BaseModel:
             if isinstance(pre_registered, BaseException):
                 raise pre_registered
             sp.data.status = "resolved"
+            sp.data.resolution = pre_registered
             return payload(**pre_registered)
 
     # No resolution available — suspend.  The span covers the whole
@@ -242,6 +246,7 @@ async def _hook_impl(call: middleware_.HookContext) -> pydantic.BaseModel:
             registry._live_hooks.pop(label, None)
 
         sp.data.status = "resolved"
+        sp.data.resolution = resolution
         await sp.add_event(telemetry.HOOK_RESOLVED)
 
         # Emit resolved signal.

src/ai/models/core/api.py

@@ -561,6 +561,13 @@ async def _stream(
     params: params_.InferenceRequestParams | None,
     executor: _StreamExecutor,
 ) -> AsyncIterator[Stream[Any]]:
+    data = telemetry.AiStreamSpanData(
+        model=model.id,
+        messages=list(messages),
+        params=params,
+        provider=model.provider.name,
+        tool_names=[t.name for t in tools] if tools is not None else None,
+    )
     if messages and messages[-1].replay:
         last = messages[-1]
         s: Stream[Any] = Stream(
@@ -572,9 +579,6 @@ async def _stream(
         # demand a finish event from the synthetic replay generator
         # (it yields nothing when the turn has no tool calls).
         s._ended = True
-        data = telemetry.AiStreamSpanData(
-            model=model.id, messages=list(messages), params=params
-        )
         replay = True
     else:
         request = _StreamRequest(
@@ -588,9 +592,6 @@ async def _stream(
             executor._do_stream(request),
             output_type=cast("type[Any] | None", output_type),
         )
-        data = telemetry.AiStreamSpanData(
-            model=model.id, messages=list(messages), params=params
-        )
         replay = False
     # Not set as current: the caller's work while the stream is open
     # (tool dispatch, user code between events) is not part of the
@@ -618,11 +619,15 @@ async def experimental_generate(
     request = _GenerateRequest(model, list(messages), params)
     async with telemetry.span(
         telemetry.AiGenerateSpanData(
-            model=model.id, messages=messages, params=params
+            model=model.id,
+            messages=messages,
+            params=params,
+            provider=model.provider.name,
         )
     ) as sp:
         message = await _default_executor._do_generate(request)
         sp.data.message = message
+        sp.data.usage = message.usage
         return message
 
 

src/ai/telemetry/span.py

@@ -153,11 +153,23 @@ class RetrievalSpanData:
 
 @dataclasses.dataclass
 class RunSpanData:
-    """One ``Agent.run``: the whole loop."""
+    """One ``Agent.run``: the whole loop.
+
+    ``blocked``/``final_message`` are set at span end: ``blocked`` is
+    True when the run ended suspended on an unresolved hook (see
+    ``AgentStream.blocked``), ``final_message`` is the last assistant
+    message produced, if any.
+    """
 
     agent: str
     model: str
     messages: list[messages_.Message]
+    provider: str | None = None
+    tool_names: list[str] | None = None
+    output_type: str | None = None
+    params: params_.InferenceRequestParams | None = None
+    blocked: bool = False
+    final_message: messages_.Message | None = None
 
     span_name: ClassVar[str] = "run"
 
@@ -180,6 +192,8 @@ class AiStreamSpanData:
     model: str
     messages: list[messages_.Message]
     params: params_.InferenceRequestParams | None = None
+    provider: str | None = None
+    tool_names: list[str] | None = None
     message: messages_.Message | None = None
     usage: usage_.Usage | None = None
 
@@ -188,37 +202,54 @@ class AiStreamSpanData:
 
 @dataclasses.dataclass
 class AiGenerateSpanData:
-    """One non-streaming generation call (images, video, ...)."""
+    """One non-streaming generation call (images, video, ...).
+
+    ``message``/``usage`` are set at span end.
+    """
 
     model: str
     messages: list[messages_.Message]
     params: params_.GenerateParams | None = None
+    provider: str | None = None
     message: messages_.Message | None = None
+    usage: usage_.Usage | None = None
 
     span_name: ClassVar[str] = "ai_generate"
 
 
 @dataclasses.dataclass
 class ToolExecutionSpanData:
-    """One tool execution, from dispatch to result."""
+    """One tool execution, from dispatch to result.
+
+    ``model_input`` is the value the LLM sees on its next turn, set
+    only when it differs from ``result`` (aggregator-backed tools).
+    """
 
     tool_name: str
     tool_call_id: str
     args: dict[str, Any] | None = None
     result: Any = None
+    model_input: Any = None
     is_error: bool = False
 
     span_name: ClassVar[str] = "tool_execution"
 
 
 @dataclasses.dataclass
 class HookSpanData:
-    """One hook suspension, from deferred until resolved or cancelled."""
+    """One hook suspension, from deferred until resolved or cancelled.
+
+    ``tool_call_id`` links the hook to the tool call it suspends, if
+    any (e.g. approval gating).  ``resolution`` carries the resolution
+    data when the hook resolves.
+    """
 
     label: str
     hook_type: str
     metadata: dict[str, Any]
+    tool_call_id: str | None = None
     status: Literal["pending", "resolved", "cancelled"] = "pending"
+    resolution: dict[str, Any] | None = None
 
     span_name: ClassVar[str] = "hook"
 

tests/agents/test_agent.py

@@ -55,13 +55,22 @@ async def lookup(x: int) -> str:
     assert isinstance(run.data, ai.telemetry.RunSpanData)
     assert run.data.agent == "Agent"
     assert run.data.model == "mock-model"
+    assert run.data.provider == "mock"
+    assert run.data.tool_names == ["lookup"]
+    assert run.data.output_type is None
+    assert not run.data.blocked
+    assert run.data.final_message is not None
+    assert run.data.final_message.text == "done"
 
     assert isinstance(tool_span.data, ai.telemetry.ToolExecutionSpanData)
     assert tool_span.data.tool_name == "lookup"
     assert tool_span.data.args == {"x": 1}
     assert tool_span.data.result == "ok"
     assert not tool_span.data.is_error
 
+    assert isinstance(calls[0].data, ai.telemetry.AiStreamSpanData)
+    assert calls[0].data.provider == "mock"
+    assert calls[0].data.tool_names == ["lookup"]
     assert isinstance(calls[1].data, ai.telemetry.AiStreamSpanData)
     assert calls[1].data.message is not None
     assert calls[1].data.message.text == "done"

tests/agents/test_hooks.py

@@ -366,6 +366,8 @@ async def test_live_hook_span(recorder: Recorder) -> None:
     (hook_span,) = _hook_spans(recorder)
     assert isinstance(hook_span.data, ai.telemetry.HookSpanData)
     assert hook_span.data.status == "resolved"
+    assert hook_span.data.resolution == {"approved": True}
+    assert hook_span.data.tool_call_id is None
     assert not hook_span.replay
     # The suspension's timeline: deferred once resolvers can see it,
     # resolved when the external input arrived.
@@ -403,6 +405,7 @@ async def loop(
     (hook_span,) = _hook_spans(recorder)
     assert isinstance(hook_span.data, ai.telemetry.HookSpanData)
     assert hook_span.data.status == "cancelled"
+    assert hook_span.data.resolution is None
     deferred, cancelled = hook_span.span_events
     assert deferred.name == ai.telemetry.HOOK_DEFERRED
     assert cancelled.name == ai.telemetry.HOOK_CANCELLED
@@ -423,5 +426,8 @@ async def test_pre_registered_hook_is_replay_span(recorder: Recorder) -> None:
     assert hook_span.replay
     assert isinstance(hook_span.data, ai.telemetry.HookSpanData)
     assert hook_span.data.status == "resolved"
+    # Pre-registration validated against the payload type, so the
+    # recorded resolution is the full model dump, defaults included.
+    assert hook_span.data.resolution == {"approved": True, "reason": ""}
     # A replayed suspension gets no synthetic timeline.
     assert hook_span.span_events == []

tests/agents/test_run_blocked.py

@@ -22,7 +22,7 @@
 from ai.types import events as events_
 from ai.types import messages as messages_
 
-from ..conftest import MOCK_MODEL, mock_llm, text_msg, tool_call_msg
+from ..conftest import MOCK_MODEL, Recorder, mock_llm, text_msg, tool_call_msg
 
 
 class Confirmation(pydantic.BaseModel):
@@ -55,7 +55,7 @@ def _multi_call_msg(*tc: tuple[str, str]) -> messages_.Message:
     )
 
 
-async def test_gated_tool_block_cycle() -> None:
+async def test_gated_tool_block_cycle(recorder: Recorder) -> None:
     """Park on approval -> RunBlocked; the resolved HookEvent unblocks."""
     my_agent = ai.Agent(tools=[gated])
     mock_llm(
@@ -107,6 +107,16 @@ def index(pred: Any) -> int:
 
     assert stream.output == "done"
 
+    # The block/unblock cycle lands on the spans: the run resolved, so
+    # it ends unblocked; the hook records its tool call and resolution.
+    (run_span,) = [s for s in recorder.ended if s.name == "run"]
+    assert isinstance(run_span.data, ai.telemetry.RunSpanData)
+    assert not run_span.data.blocked
+    (hook_span,) = [s for s in recorder.ended if s.name == "hook"]
+    assert isinstance(hook_span.data, ai.telemetry.HookSpanData)
+    assert hook_span.data.tool_call_id == "tc-1"
+    assert hook_span.data.resolution == {"granted": True, "reason": "ok"}
+
 
 async def test_busy_tool_defers_block_signal() -> None:
     """No block signal while an unblocked tool is still running."""
@@ -201,7 +211,7 @@ async def loop(
     assert not stream.blocked
 
 
-async def test_abort_leaves_blocked() -> None:
+async def test_abort_leaves_blocked(recorder: Recorder) -> None:
     """Serverless abort: the run ends still blocked, hooks still deferred."""
     my_agent = ai.Agent(tools=[gated])
     mock_llm([[tool_call_msg(name="gated", args='{"x": 1}')]])
@@ -224,6 +234,11 @@ async def test_abort_leaves_blocked() -> None:
     assert len(stream.deferred_hooks) == 1
     assert stream.messages[-1].tool_results[0].is_hook_deferred
 
+    # The run span ends still blocked, mirroring ``stream.blocked``.
+    (run_span,) = [s for s in recorder.ended if s.name == "run"]
+    assert isinstance(run_span.data, ai.telemetry.RunSpanData)
+    assert run_span.data.blocked
+
 
 async def test_no_hooks_no_block_events() -> None:
     @ai.tool

tests/models/core/test_api.py

@@ -324,7 +324,7 @@ async def _stream() -> AsyncGenerator[events_.Event]:
     assert stream.text == "override"
 
 
-async def test_generate_dispatches_to_provider() -> None:
+async def test_generate_dispatches_to_provider(recorder: Recorder) -> None:
     provider = MockProvider()
     model = models.Model(
         id="generate-model",
@@ -333,6 +333,7 @@ async def test_generate_dispatches_to_provider() -> None:
     sentinel = messages_.Message(
         role="assistant",
         parts=[messages_.FilePart(data=b"\x89PNG", media_type="image/png")],
+        usage=ai.types.usage.Usage(input_tokens=3, output_tokens=7),
     )
     called = False
 
@@ -356,6 +357,12 @@ async def _generate(
     assert called
     assert result is sentinel
 
+    (span,) = [s for s in recorder.ended if s.name == "ai_generate"]
+    assert isinstance(span.data, ai.telemetry.AiGenerateSpanData)
+    assert span.data.provider == "mock"
+    assert span.data.message is sentinel
+    assert span.data.usage is sentinel.usage
+
 
 async def test_generate_uses_model_protocol() -> None:
     sentinel = messages_.Message(
@@ -605,6 +612,9 @@ async def test_stream_span_milestones(recorder: Recorder) -> None:
 
     (call,) = [s for s in recorder.ended if s.name == "ai_stream"]
     assert stream.span is call
+    assert isinstance(call.data, ai.telemetry.AiStreamSpanData)
+    assert call.data.provider == "mock"
+    assert call.data.tool_names is None  # no tools were sent
     assert [e.name for e in call.span_events] == [
         ai.telemetry.FIRST_TOKEN,
         ai.telemetry.RESPONSE_COMPLETE,