Wire telemetry into the framework

examples/agents/telemetry.py

@@ -0,0 +1,35 @@
+"""Agent run with telemetry: console adapter + a custom span in a tool.
+
+For an out-of-process viewer, run ``python -m ai.telemetry.utils.viewer``
+and see ``ai/telemetry/utils/viewer.py`` for the otel exporter setup.
+"""
+
+import asyncio
+
+import ai
+from ai.telemetry.utils import console
+
+
+@ai.tool
+async def get_weather(city: str) -> str:
+    """Get current weather for a city."""
+    async with ai.telemetry.span("lookup", city=city) as span:
+        await asyncio.sleep(0.1)
+        span.set(source="cache")
+    return f"Sunny, 72F in {city}"
+
+
+async def main() -> None:
+    ai.telemetry.register(console.ConsoleAdapter())
+
+    model = ai.get_model("anthropic/claude-sonnet-4.6")
+    my_agent = ai.Agent(tools=[get_weather])
+    messages = [ai.user_message("What's the weather in Tokyo?")]
+
+    async with my_agent.run(model, messages) as stream:
+        async for _ in stream:
+            pass
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

src/ai/agents/_middleware.py

@@ -15,6 +15,7 @@
 
 from __future__ import annotations
 
+import contextlib
 import contextvars
 import dataclasses
 from collections.abc import AsyncGenerator, Awaitable, Callable, Sequence
@@ -141,8 +142,11 @@ async def wrap_agent_run(self, call, next):
                     yield event
                 span.end()
         """
-        async for event in next(call):
-            yield event
+        # aclosing: propagate an early close from the consumer down the
+        # chain, so inner generators unwind in the task that runs them
+        async with contextlib.aclosing(next(call)) as events:
+            async for event in events:
+                yield event
 
     async def wrap_model(
         self,
@@ -344,8 +348,10 @@ def _build_agent_run_chain(
 
         def _make(m: _Middleware, nxt: _AgentRunNext) -> _AgentRunNext:
             async def _wrapped(call: Context) -> AsyncGenerator[_Event]:
-                async for event in m.wrap_agent_run(call, nxt):
-                    yield event
+                gen = m.wrap_agent_run(call, nxt)
+                async with contextlib.aclosing(gen) as events:
+                    async for event in events:
+                        yield event
 
             return _wrapped
 

src/ai/agents/agent.py

@@ -35,7 +35,7 @@
 # Use the typing_extensions backport so this works on 3.12 too.
 from typing_extensions import TypeVar
 
-from .. import models, type_utils, types, util
+from .. import models, telemetry, type_utils, types, util
 from ..types import builders
 from ..types import events as events_
 from ..types.messages import MessageBundle
@@ -630,14 +630,24 @@ def kwargs(self) -> dict[str, Any]:
 
     async def __call__(self, **overrides: Any) -> events_.ToolCallResult:
         """Execute the tool and return a :class:`ToolCallResult`."""
+        data = telemetry.ToolExecutionSpanData(
+            tool_name=self._part.tool_name,
+            tool_call_id=self._part.tool_call_id,
+        )
+
         # Replay-from-pending-hook short-circuit: if a prior run already
         # produced a result for this call (cached on the ToolCallPart
         # by ``_process_interrupted_hooks``), return it without
         # re-executing the tool.
         cached = self._part.cached_result
         if cached is not None:
-            msg = builders.tool_message(cached)
-            return events_.ToolCallResult(message=msg, results=msg.tool_results)
+            async with telemetry.span(data, replay=True):
+                data.result = cached.result
+                data.is_error = cached.is_error
+                msg = builders.tool_message(cached)
+                return events_.ToolCallResult(
+                    message=msg, results=msg.tool_results
+                )
 
         # Best-effort parse so middleware sees usable kwargs when possible.
         # If parsing fails, middleware still gets the raw tool_call_id /
@@ -702,8 +712,18 @@ async def _real(
             part.set_model_input(model_input)
             return tool_result(part)
 
+        data.args = base_kwargs
         chain = middleware_._build_tool_chain(_real)
-        return await chain(call)
+        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)
+            # 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.
+            sp.error = res.exception
+            return res
 
 
 class GatedToolCall:
@@ -1255,9 +1275,13 @@ async def loop(
         """Stream, execute tools, repeat.
 
         Override in a subclass to customise the agent's control flow.
+        (A custom loop keeps run/model/tool/hook spans — those live in
+        the layers it calls — but loses per-step grouping unless it
+        opens its own ``telemetry.span`` blocks.)
         """
         while context.keep_running():
             async with (
+                telemetry.span(telemetry.LoopTurnSpanData()),
                 models.stream(context=context) as stream,
                 ToolRunner() as tr,
             ):
@@ -1359,38 +1383,53 @@ async def _run(
         async def _real(call: Context) -> AsyncGenerator[events_.AgentEvent]:
             tracker = events_.RunStateTracker()
             source = self.loop(call)
-            async for event in runtime.run(source):
-                # Feed the tracker before the replay filter: replayed
-                # StreamEnds carry the tool calls the dispatcher is
-                # about to re-run, which the fold must count.
-                transition = tracker.feed(event)
-                # Drop replay-flagged events: they're a control-flow
-                # signal for the loop's tool dispatcher (which already
-                # ran by the time we see the event here), not user-
-                # facing output.
-                if not (isinstance(event, events_.BaseEvent) and event.replay):
-                    yield event
-                if transition is not None:
-                    yield transition
+            async with contextlib.aclosing(runtime.run(source)) as events:
+                async for event in events:
+                    # Feed the tracker before the replay filter: replayed
+                    # StreamEnds carry the tool calls the dispatcher is
+                    # about to re-run, which the fold must count.
+                    transition = tracker.feed(event)
+                    # Drop replay-flagged events: they're a control-flow
+                    # signal for the loop's tool dispatcher (which already
+                    # ran by the time we see the event here), not user-
+                    # facing output.
+                    if not (
+                        isinstance(event, events_.BaseEvent) and event.replay
+                    ):
+                        yield event
+                    if transition is not None:
+                        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.
-            mw_token: middleware_.Token | None = None
-            if _middleware is not None:
-                parent = middleware_.get()
-                mw_token = middleware_.activate(parent + _middleware)
-            try:
-                chain = middleware_._build_agent_run_chain(_real)
-                async for event in chain(context):
-                    yield event
-            finally:
-                if mw_token is not None:
-                    middleware_.deactivate(mw_token)
-
-        yield AgentStream(_stream(), context)
+            async with telemetry.span(run_data):
+                mw_token: middleware_.Token | None = None
+                if _middleware is not None:
+                    parent = middleware_.get()
+                    mw_token = middleware_.activate(parent + _middleware)
+                try:
+                    chain = middleware_._build_agent_run_chain(_real)
+                    async with contextlib.aclosing(chain(context)) as events:
+                        async for event in events:
+                            yield event
+                finally:
+                    if mw_token is not None:
+                        middleware_.deactivate(mw_token)
+
+        # close the event generator on exit from the ``run()`` block.
+        astream = AgentStream(_stream(), context)
+        try:
+            yield astream
+        finally:
+            await astream.aclose()

src/ai/agents/hooks.py

@@ -25,7 +25,7 @@
 
 import pydantic
 
-from .. import types
+from .. import telemetry, types
 from ..types import messages as messages_
 from . import _middleware as middleware_
 from . import runtime as runtime_
@@ -124,49 +124,69 @@ async def _hook_impl(call: middleware_.HookContext) -> pydantic.BaseModel:
     payload = call.payload
     hook_metadata = call.metadata
 
+    data = telemetry.HookSpanData(
+        label=label, hook_type=payload.__name__, metadata=hook_metadata
+    )
+
     # Pre-registered resolution (serverless re-entry).
     pre_registered = _pending_resolutions.pop(label, None)
     if pre_registered is not None:
-        if isinstance(pre_registered, BaseException):
-            raise pre_registered
-        return payload(**pre_registered)
-
-    # No resolution available — suspend.
-    future: asyncio.Future[dict[str, Any]] = asyncio.Future()
-
-    _live_hooks[label] = (future, hook_metadata, rt)
-    rt.track_hook_label(label)
-
-    # Emit pending signal.
-    hook_part: messages_.HookPart[Any] = messages_.HookPart(
-        hook_id=label,
-        hook_type=payload.__name__,
-        status="pending",
-        metadata=hook_metadata,
-        tool_call_id=call.tool_call_id,
-    )
-
-    await rt.put_hook(hook_part)
-
-    # Await resolution — may be resolved externally or cancelled.
-    resolution = await future
-
-    # Clean up live registry.
-    _live_hooks.pop(label, None)
-
-    # Emit resolved signal.
-    await rt.put_hook(
-        messages_.HookPart(
+        async with telemetry.span(data, replay=True):
+            if isinstance(pre_registered, BaseException):
+                raise pre_registered
+            data.status = "resolved"
+            return payload(**pre_registered)
+
+    # No resolution available — suspend.  The span covers the whole
+    # suspension: how long the run sat waiting on external input.
+    async with telemetry.span(data) as sp:
+        future: asyncio.Future[dict[str, Any]] = asyncio.Future()
+
+        _live_hooks[label] = (future, hook_metadata, rt)
+        rt.track_hook_label(label)
+
+        # Emit pending signal.
+        hook_part: messages_.HookPart[Any] = messages_.HookPart(
             hook_id=label,
             hook_type=payload.__name__,
-            status="resolved",
+            status="pending",
             metadata=hook_metadata,
-            resolution=resolution,
             tool_call_id=call.tool_call_id,
         )
-    )
 
-    return payload(**resolution)
+        await rt.put_hook(hook_part)
+        await sp.add_event(telemetry.HOOK_PENDING)
+
+        # Await resolution — may be resolved externally or cancelled.
+        try:
+            resolution = await future
+        except asyncio.CancelledError as exc:
+            data.status = "cancelled"
+            # ``cancel_hook(reason=...)`` rides on the CancelledError.
+            attrs: dict[str, Any] = {}
+            if exc.args and exc.args[0] is not None:
+                attrs["reason"] = exc.args[0]
+            await sp.add_event(telemetry.HOOK_CANCELLED, **attrs)
+            raise
+
+        # Clean up live registry.
+        _live_hooks.pop(label, None)
+        data.status = "resolved"
+        await sp.add_event(telemetry.HOOK_RESOLVED)
+
+        # Emit resolved signal.
+        await rt.put_hook(
+            messages_.HookPart(
+                hook_id=label,
+                hook_type=payload.__name__,
+                status="resolved",
+                metadata=hook_metadata,
+                resolution=resolution,
+                tool_call_id=call.tool_call_id,
+            )
+        )
+
+        return payload(**resolution)
 
 
 def resolve_hook(

src/ai/agents/runtime.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import contextlib
 import contextvars
 from typing import TYPE_CHECKING, Any
 
@@ -12,7 +13,7 @@
 from .mcp import client as mcp_client
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator, AsyncIterable, Awaitable
+    from collections.abc import AsyncGenerator, Awaitable
 
 
 class Runtime:
@@ -65,7 +66,7 @@ async def _stop_when_done(runtime: Runtime, task: Awaitable[None]) -> None:
 
 
 async def run(
-    source: AsyncIterable[events_.AgentEvent],
+    source: AsyncGenerator[events_.AgentEvent],
 ) -> AsyncGenerator[events_.AgentEvent]:
     """Run *source* and yield events put into the Runtime queue."""
     rt = Runtime()
@@ -80,8 +81,11 @@ async def _drain() -> None:
         mcp_token = mcp_client._pool.set(mcp_pool)
 
         try:
-            async for event in source:
-                await rt.put_event(event)
+            # aclosing: if this task is cancelled while *source* sits
+            # suspended at a yield, close it here
+            async with contextlib.aclosing(source) as events:
+                async for event in events:
+                    await rt.put_event(event)
 
         finally:
             rt.cleanup_hooks()