Let stream accept an agents.Context instead of model/messages/tools

README.md

@@ -107,7 +107,7 @@ Override the default loop when you need approval gates, routing, or custom orche
 @agent.loop
 async def custom(context: ai.Context):
     while True:
-        async with ai.stream(context.model, context.messages, tools=context.tools) as s:
+        async with ai.stream(context=context) as s:
             async for event in s:
                 yield event
             context.add(s.message)

examples/fastapi-vite/backend/agent.py

@@ -56,9 +56,7 @@ async def graph(context: ai.Context) -> AsyncGenerator[ai.events.AgentEvent]:
     Reject buttons and sends the decision back on the next request.
     """
     while context.keep_running():
-        async with ai.models.stream(
-            context.model, context.messages, tools=context.tools
-        ) as s:
+        async with ai.models.stream(context=context) as s:
             async for event in s:
                 yield event
         context.add(s.message)

examples/multiagent-textual/server.py

@@ -110,7 +110,7 @@ def _gated_agent(
     async def gated_loop(context: ai.Context) -> AsyncGenerator[ai.events.AgentEvent]:
         while context.keep_running():
             async with (
-                ai.stream(context.model, context.messages, tools=context.tools) as s,
+                ai.stream(context=context) as s,
                 ai.agents.ToolRunner() as tr,
             ):
                 async for event in ai.util.merge(s, tr.events()):

examples/samples/agent_custom_loop.py

@@ -27,9 +27,7 @@ async def default_loop(
         """Stream, execute tools with logging, repeat."""
         while context.keep_running():
             async with (
-                ai.models.stream(
-                    context.model, context.messages, tools=context.tools
-                ) as stream,
+                ai.models.stream(context=context) as stream,
                 ai.ToolRunner() as tr,
             ):
                 async for event in ai.util.merge(stream, tr.events()):

examples/samples/agent_hooks.py

@@ -70,7 +70,7 @@ async def with_approval(
     ) -> AsyncGenerator[ai.events.AgentEvent]:
         while context.keep_running():
             async with (
-                ai.stream(context.model, context.messages, tools=context.tools) as s,
+                ai.stream(context=context) as s,
                 ai.ToolRunner() as tr,
             ):
                 async for event in ai.util.merge(s, tr.events()):

examples/samples/agent_hooks_inline.py

@@ -46,7 +46,7 @@ async def with_approval(
     ) -> AsyncGenerator[ai.events.AgentEvent]:
         while context.keep_running():
             async with (
-                ai.stream(context.model, context.messages, tools=context.tools) as s,
+                ai.stream(context=context) as s,
                 ai.ToolRunner() as tr,
             ):
                 async for event in ai.util.merge(s, tr.events()):

examples/samples/agent_hooks_serverless.py

@@ -47,9 +47,7 @@ async def with_confirmation(
         context: ai.Context,
     ) -> AsyncGenerator[ai.events.AgentEvent]:
         while context.keep_running():
-            async with ai.models.stream(
-                context.model, context.messages, tools=context.tools
-            ) as s:
+            async with ai.models.stream(context=context) as s:
                 async for event in s:
                     yield event
 

src/ai/agents/agent.py

@@ -770,9 +770,7 @@ async def default_loop(
         """Stream, execute tools, repeat."""
         while context.keep_running():
             async with (
-                models.stream(
-                    context.model, context.messages, tools=context.tools
-                ) as stream,
+                models.stream(context=context) as stream,
                 ToolRunner() as tr,
             ):
                 async for event in util.merge(stream, tr.events()):

src/ai/models/core/api.py

@@ -1,7 +1,8 @@
 import contextlib
 import dataclasses
 from collections.abc import AsyncGenerator, AsyncIterator, Sequence
-from typing import Any, Protocol, Self, runtime_checkable
+from contextlib import AbstractAsyncContextManager
+from typing import Any, Protocol, Self, overload, runtime_checkable
 
 import pydantic
 
@@ -302,28 +303,94 @@ async def _replay_tool_calls(
         )
 
 
-@contextlib.asynccontextmanager
-async def stream(
+@runtime_checkable
+class StreamContext(Protocol):
+    """Anything that exposes ``model``/``messages``/``tools``.
+
+    Used to let callers pass an ``agents.Context`` to :func:`stream`
+    without an import-time circular dependency.
+    """
+
+    @property
+    def model(self) -> model_.Model: ...
+    @property
+    def messages(self) -> list[types.messages.Message]: ...
+    @property
+    def tools(self) -> list[types.tools.Tool]: ...
+
+
+@overload
+def stream(
+    *,
+    context: StreamContext,
+    output_type: type[pydantic.BaseModel] | None = None,
+    params: Any = None,
+    executor: StreamExecutor = _default_executor,
+) -> AbstractAsyncContextManager[Stream]: ...
+@overload
+def stream[ProviderParamsT: pydantic.BaseModel](
     model: model_.Model,
     messages: list[types.messages.Message],
     *,
     tools: Sequence[types.tools.Tool] | None = None,
     output_type: type[pydantic.BaseModel] | None = None,
     params: Any = None,
     executor: StreamExecutor = _default_executor,
-) -> AsyncIterator[Stream]:
+) -> AbstractAsyncContextManager[Stream]: ...
+def stream(
+    model: model_.Model | None = None,
+    messages: list[types.messages.Message] | None = None,
+    *,
+    context: StreamContext | None = None,
+    tools: Sequence[types.tools.Tool] | None = None,
+    output_type: type[pydantic.BaseModel] | None = None,
+    params: Any = None,
+    executor: StreamExecutor = _default_executor,
+) -> AbstractAsyncContextManager[Stream]:
     """Stream an LLM response.
 
-    Used as an async context manager whose value is the :class:`Stream`::
+    Used as an async context manager whose value is the :class:`Stream`.
+    Pass either positional ``model, messages`` (plus optional ``tools=``)
+    or ``context=`` (an ``agents.Context`` or anything matching
+    :class:`StreamContext`)::
 
-        async with ai.stream(model, messages) as s:
-            async for event in s:
-                ...
-            print(s.message)
+        async with ai.stream(model, messages) as s: ...
+        async with ai.stream(context=context) as s: ...
 
     If the last message is marked ``replay=True``, replay that turn as
     synthetic stream events instead of calling the model.
     """
+    if context is not None:
+        if model is not None or messages is not None or tools is not None:
+            raise TypeError(
+                "stream() takes either model/messages/tools or context=, not both"
+            )
+        model = context.model
+        messages = context.messages
+        tools = context.tools
+    elif model is None or messages is None:
+        raise TypeError("stream() requires either model and messages or context=")
+
+    return _stream(
+        model=model,
+        messages=messages,
+        tools=tools,
+        output_type=output_type,
+        params=params,
+        executor=executor,
+    )
+
+
+@contextlib.asynccontextmanager
+async def _stream(
+    *,
+    model: model_.Model,
+    messages: list[types.messages.Message],
+    tools: Sequence[types.tools.Tool] | None,
+    output_type: type[pydantic.BaseModel] | None,
+    params: Any,
+    executor: StreamExecutor,
+) -> AsyncIterator[Stream]:
     if messages and messages[-1].replay:
         last = messages[-1]
         s = Stream(_replay_tool_calls(last), seed_message=last.model_copy(deep=True))

tests/agents/test_hooks.py

@@ -31,7 +31,7 @@ async def test_resolve_live_future() -> None:
     @my_agent.loop
     async def custom(context: ai.Context) -> AsyncGenerator[ai.events.Event]:
         nonlocal resolved_value
-        async with ai.models.stream(context.model, context.messages) as stream:
+        async with ai.models.stream(context=context) as stream:
             async for event in stream:
                 yield event
         result = await ai.hook("confirm_1", payload=Confirmation)
@@ -63,7 +63,7 @@ async def test_cancel_live_hook() -> None:
     @my_agent.loop
     async def custom(context: ai.Context) -> AsyncGenerator[ai.events.Event]:
         nonlocal was_cancelled
-        async with ai.models.stream(context.model, context.messages) as stream:
+        async with ai.models.stream(context=context) as stream:
             async for event in stream:
                 yield event
         try:
@@ -102,7 +102,7 @@ async def test_pre_registered_resolution_consumed() -> None:
     @my_agent.loop
     async def custom(context: ai.Context) -> AsyncGenerator[ai.events.Event]:
         nonlocal resolved_value
-        async with ai.models.stream(context.model, context.messages) as stream:
+        async with ai.models.stream(context=context) as stream:
             async for event in stream:
                 yield event
         resolved_value = await ai.hook("pre_reg_1", payload=Confirmation)
@@ -142,7 +142,7 @@ async def test_resolved_hook_emits_message() -> None:
 
     @my_agent.loop
     async def custom(context: ai.Context) -> AsyncGenerator[ai.events.Event]:
-        async with ai.models.stream(context.model, context.messages) as stream:
+        async with ai.models.stream(context=context) as stream:
             async for event in stream:
                 yield event
         await ai.hook("emit_test", payload=Confirmation)
@@ -171,7 +171,7 @@ async def test_hook_metadata_in_pending() -> None:
 
     @my_agent.loop
     async def custom(context: ai.Context) -> AsyncGenerator[ai.events.Event]:
-        async with ai.models.stream(context.model, context.messages) as stream:
+        async with ai.models.stream(context=context) as stream:
             async for event in stream:
                 yield event
         await ai.hook(

tests/models/core/test_api.py

@@ -4,6 +4,7 @@
 from typing import Any
 
 import pydantic
+import pytest
 
 import ai
 from ai import models
@@ -227,6 +228,42 @@ async def _spy_stream(
     assert received_params == [params]
 
 
+async def test_stream_accepts_context() -> None:
+    """``stream(context=ctx)`` reads model/messages/tools off the context."""
+    mock = mock_llm([[text_msg("ok")]])
+    ctx = ai.Context(
+        model=MOCK_MODEL,
+        messages=[ai.user_message("Hi")],
+        tools=[],
+    )
+    async with models.stream(context=ctx) as s:
+        async for _ in s:
+            pass
+    assert mock.call_count == 1
+    assert s.text == "ok"
+
+
+async def test_stream_rejects_context_with_positional_args() -> None:
+    """Passing both positional model/messages and ``context=`` is a TypeError."""
+    ctx = ai.Context(
+        model=MOCK_MODEL,
+        messages=[ai.user_message("Hi")],
+        tools=[],
+    )
+    with pytest.raises(TypeError, match="either model/messages/tools or context="):
+        async with models.stream(  # type: ignore[call-overload]
+            MOCK_MODEL, [ai.user_message("Hi")], context=ctx
+        ):
+            pass
+
+
+async def test_stream_requires_model_messages_or_context() -> None:
+    """Passing nothing is a TypeError."""
+    with pytest.raises(TypeError, match="either model and messages or context="):
+        async with models.stream():  # type: ignore[call-overload]
+            pass
+
+
 async def test_generate_dispatches_to_registered_adapter() -> None:
     provider = MockProvider(adapter="mock-generate")
     model = models.Model(

tests/test_middleware.py

@@ -79,7 +79,7 @@ async def wrap_hook(self, call: middleware.HookContext, next: Any) -> Any:
 
     @my_agent.loop
     async def custom(context: ai.Context) -> AsyncGenerator[ai.events.Event]:
-        async with ai.models.stream(context.model, context.messages) as stream:
+        async with ai.models.stream(context=context) as stream:
             async for event in stream:
                 yield event
         await ai.hook("test_hook", payload=Confirmation)