Expose context and messages on the stream run returns

examples/samples/agent_hooks_serverless.py

@@ -95,11 +95,6 @@ async def with_confirmation(
 
     async with my_agent.run(model, messages) as stream:
         async for event in stream:
-            # HACK?: When we get a complete assistant message, add it to
-            # messages so it can get replayed easily.
-            if isinstance(event, ai.events.StreamEnd):
-                messages.append(event.message)
-
             if isinstance(event, ai.events.TextDelta):
                 print(event.chunk, end="", flush=True)
             elif (
@@ -112,6 +107,9 @@ async def with_confirmation(
                     f"  Hook pending: {hook_part.hook_id} "
                     f"(metadata={hook_part.metadata})"
                 )
+        # Pick up the assistant turn that the loop appended so the
+        # next run replays from the same point.
+        messages = stream.messages
 
     print("\n  Run interrupted; approval will be pre-registered for re-entry.\n")
 

examples/samples/middleware_simple.py

@@ -26,8 +26,8 @@ class PrintMiddleware(ai.Middleware):
 
     async def wrap_agent_run(
         self,
-        call: ai.middleware.AgentRunContext,
-        next: Callable[[ai.middleware.AgentRunContext], AsyncGenerator[Any]],
+        call: ai.Context,
+        next: Callable[[ai.Context], AsyncGenerator[Any]],
     ) -> AsyncGenerator[Any]:
         print(f">>> [run] agent starting  tools={len(call.tools)}")
         t0 = time.perf_counter()

src/ai/__init__.py

@@ -27,7 +27,7 @@
     tool_result,
     yield_from,
 )
-from .agents.middleware import AgentRunContext, Middleware
+from .agents.middleware import Middleware
 from .models import (
     Client,
     Executor,
@@ -119,7 +119,6 @@
     "cancel_hook",
     "TOOL_APPROVAL_HOOK_TYPE",
     # Middleware
-    "AgentRunContext",
     "Middleware",
     "middleware",
     # Submodules

src/ai/agents/agent.py

@@ -595,6 +595,55 @@ class LoopFn(Protocol):
     def __call__(self, context: Context) -> AsyncGenerator[events_.AgentEvent]: ...
 
 
+class AgentStream:
+    """Async-iterable wrapper around an agent run's event stream.
+
+    Exposes the run's :class:`Context` via :attr:`context` so callers can
+    inspect (or use) the live messages/tools without threading them
+    through their own bookkeeping::
+
+        async with agent.run(model, messages) as stream:
+            async for event in stream:
+                ...
+            print(stream.context.messages)
+
+    Structurally satisfies the ``AsyncGenerator`` protocol by delegating
+    to the underlying generator, so it can be passed directly to
+    :func:`yield_from` and other APIs that expect an async generator.
+    """
+
+    def __init__(
+        self,
+        gen: AsyncGenerator[events_.AgentEvent],
+        context: Context,
+    ) -> None:
+        self._gen = gen
+        self._context = context
+
+    def __aiter__(self) -> Self:
+        return self
+
+    async def __anext__(self) -> events_.AgentEvent:
+        return await self._gen.__anext__()
+
+    async def asend(self, value: Any) -> events_.AgentEvent:
+        return await self._gen.asend(value)
+
+    async def athrow(self, *args: Any, **kwargs: Any) -> events_.AgentEvent:
+        return await self._gen.athrow(*args, **kwargs)
+
+    async def aclose(self) -> None:
+        await self._gen.aclose()
+
+    @property
+    def context(self) -> Context:
+        return self._context
+
+    @property
+    def messages(self) -> list[types.messages.Message]:
+        return self._context.messages
+
+
 def tool_result(
     *items: types.messages.Message
     | types.messages.ToolResultPart
@@ -746,14 +795,16 @@ async def run(
         messages: list[types.messages.Message],
         *,
         middleware: list[middleware_.Middleware] | None = None,
-    ) -> AsyncIterator[AsyncGenerator[events_.AgentEvent]]:
+    ) -> AsyncIterator[AgentStream]:
         """Run the agent loop, yielding events to the consumer.
 
-        Used as an async context manager whose value is the event stream::
+        Used as an async context manager whose value the event stream,
+        extended with the `context` and `messages` of the stream::
 
             async with agent.run(model, messages) as stream:
                 async for event in stream:
                     ...
+                print(stream.messages)
 
         Args:
             model: The model to use for LLM calls.
@@ -766,37 +817,29 @@ async def run(
         ``yield_from(..., label=...)`` — the label flows via
         ``PartialToolCallResult`` rather than on individual messages.
         """
-        call = middleware_.AgentRunContext(
+        context = Context(
             model=model,
-            messages=messages,
-            tools=self._tools,
+            messages=list(messages),
+            tools=[t.tool for t in self._tools],
         )
+        context._agent_tools_by_name = {t.name: t for t in self._tools}
+        # If the final message is an assistant call with tool
+        # calls, then probably the situation is that we bailed out
+        # earlier due to unresolved hooks, and we need to arrange
+        # to replay the message now.
+        if (
+            context.messages
+            and context.messages[-1].role == "assistant"
+            and context.messages[-1].tool_calls
+        ):
+            context.messages[-1] = context.messages[-1].model_copy(
+                update={"replay": True}
+            )
 
         loop_fn = self._loop_fn or self.default_loop
 
-        async def _real(
-            call: middleware_.AgentRunContext,
-        ) -> AsyncGenerator[events_.AgentEvent]:
-            context = Context(
-                model=call.model,
-                messages=list(call.messages),
-                tools=[tool.tool for tool in call.tools],
-            )
-            context._agent_tools_by_name = {tool.name: tool for tool in call.tools}
-            # If the final message is an assistant call with tool
-            # calls, then probably the situation is that we bailed out
-            # earlier due to unresolved hooks, and we need to arrange
-            # to replay the message now.
-            if (
-                context.messages
-                and context.messages[-1].role == "assistant"
-                and context.messages[-1].tool_calls
-            ):
-                context.messages[-1] = context.messages[-1].model_copy(
-                    update={"replay": True}
-                )
-
-            source = loop_fn(context)
+        async def _real(call: Context) -> AsyncGenerator[events_.AgentEvent]:
+            source = loop_fn(call)
             async for event in runtime.run(source):
                 # Drop replay-flagged events: they're a control-flow
                 # signal for the loop's tool dispatcher (which already
@@ -820,13 +863,13 @@ async def _stream() -> AsyncGenerator[events_.AgentEvent]:
                 mw_token = middleware_.activate(parent + middleware)
             try:
                 chain = middleware_._build_agent_run_chain(_real)
-                async for event in chain(call):
+                async for event in chain(context):
                     yield event
             finally:
                 if mw_token is not None:
                     middleware_.deactivate(mw_token)
 
-        yield _stream()
+        yield AgentStream(_stream(), context)
 
 
 def agent(

src/ai/agents/middleware.py

@@ -45,7 +45,7 @@
 if TYPE_CHECKING:
     from ..models.core.model import Model
     from ..types import events as events_
-    from .agent import AgentTool
+    from .agent import Context
 
 
 @dataclasses.dataclass(frozen=True)
@@ -102,19 +102,6 @@ def __post_init__(self) -> None:
         object.__setattr__(self, "metadata", dict(self.metadata))
 
 
-@dataclasses.dataclass(frozen=True)
-class AgentRunContext:
-    """Context for an agent run."""
-
-    model: Model[Any]
-    messages: list[messages_.Message]
-    tools: list[AgentTool]
-
-    def __post_init__(self) -> None:
-        object.__setattr__(self, "messages", list(self.messages))
-        object.__setattr__(self, "tools", list(self.tools))
-
-
 # ---------------------------------------------------------------------------
 # Middleware base class — override the methods you care about.
 # ---------------------------------------------------------------------------
@@ -127,7 +114,7 @@ def __post_init__(self) -> None:
 _Message = messages_.Message
 
 # Agent run next-function type: call -> async generator of events.
-_AgentRunNext = Callable[[AgentRunContext], AsyncGenerator[_Event]]
+_AgentRunNext = Callable[["Context"], AsyncGenerator[_Event]]
 
 
 class Middleware:
@@ -138,7 +125,7 @@ class Middleware:
 
     async def wrap_agent_run(
         self,
-        call: AgentRunContext,
+        call: Context,
         next: _AgentRunNext,
     ) -> AsyncGenerator[_Event]:
         """Wrap an agent run.
@@ -354,7 +341,7 @@ def _build_agent_run_chain(
     for m in reversed(mw):
 
         def _make(m: Middleware, nxt: _AgentRunNext) -> _AgentRunNext:
-            async def _wrapped(call: AgentRunContext) -> AsyncGenerator[_Event]:
+            async def _wrapped(call: Context) -> AsyncGenerator[_Event]:
                 async for event in m.wrap_agent_run(call, nxt):
                     yield event
 
@@ -365,7 +352,6 @@ async def _wrapped(call: AgentRunContext) -> AsyncGenerator[_Event]:
 
 
 __all__ = [
-    "AgentRunContext",
     "GenerateContext",
     "HookContext",
     "Middleware",

tests/test_middleware.py

@@ -110,7 +110,7 @@ async def test_wrap_agent_run_ordering() -> None:
 
     class Outer(ai.Middleware):
         async def wrap_agent_run(
-            self, call: middleware.AgentRunContext, next: Any
+            self, call: ai.Context, next: Any
         ) -> AsyncGenerator[ai.events.Event]:
             order.append("outer-before")
             async for event in next(call):
@@ -119,7 +119,7 @@ async def wrap_agent_run(
 
     class Inner(ai.Middleware):
         async def wrap_agent_run(
-            self, call: middleware.AgentRunContext, next: Any
+            self, call: ai.Context, next: Any
         ) -> AsyncGenerator[ai.events.Event]:
             order.append("inner-before")
             async for event in next(call):