Make middleware composable in the nested case, add stream result proto

Author: anbuzin

examples/temporal-middleware/main.py

@@ -130,66 +130,39 @@ async def llm_call_activity(params: LLMParams) -> LLMResult:
 # wrap_tool, same as without middleware.
 
 
-class _BufferedStreamResult:
-    """Wraps a single buffered Message to look like a StreamResult."""
-
-    def __init__(self, message: ai.Message) -> None:
-        self._message = message
-
-    def __aiter__(self) -> AsyncGenerator[ai.Message]:
-        return self._generate()
-
-    async def _generate(self) -> AsyncGenerator[ai.Message]:
-        yield self._message
-
-    @property
-    def tool_calls(self) -> list[ai.ToolCallPart]:
-        return self._message.tool_calls
-
-    @property
-    def text(self) -> str:
-        return self._message.text
-
-    @property
-    def usage(self) -> ai.Usage | None:
-        return self._message.usage
-
-    @property
-    def output(self) -> Any:
-        return self._message.output
-
-
 class TemporalMiddleware(ai.Middleware):
-    """Routes LLM calls and tool executions through Temporal activities.
-
-    The middleware tracks messages itself so it can serialize the full
-    conversation to the LLM activity at each step.
-    """
+    """Routes LLM calls and tool executions through Temporal activities."""
 
-    def __init__(
-        self,
-        initial_messages: list[ai.Message],
-        tool_schemas: list[dict[str, Any]],
-    ) -> None:
-        self._messages = list(initial_messages)
+    def __init__(self, tool_schemas: list[dict[str, Any]]) -> None:
         self._tool_schemas = tool_schemas
 
-    async def wrap_model(self, call: ai.middleware.ModelContext, next: Any) -> Any:
+    async def wrap_model(
+        self,
+        call: ai.middleware.ModelContext,
+        next: Any,
+    ) -> ai.StreamResultLike:
         """LLM call → Temporal activity."""
         result = await temporalio.workflow.execute_activity(
             llm_call_activity,
             LLMParams(
-                messages=[m.model_dump() for m in self._messages],
+                messages=[m.model_dump() for m in call.messages],
                 tool_schemas=self._tool_schemas,
             ),
             start_to_close_timeout=datetime.timedelta(minutes=5),
             retry_policy=temporalio.common.RetryPolicy(maximum_attempts=3),
         )
         msg = ai.Message.model_validate(result.message)
-        self._messages.append(msg)
-        return _BufferedStreamResult(msg)
 
-    async def wrap_tool(self, call: ai.middleware.ToolContext, next: Any) -> Any:
+        async def _single() -> AsyncGenerator[ai.Message]:
+            yield msg
+
+        return ai.StreamResult.from_generator(_single())
+
+    async def wrap_tool(
+        self,
+        call: ai.middleware.ToolContext,
+        next: Any,
+    ) -> ai.Message:
         """Tool execution → Temporal activity."""
         result = await temporalio.workflow.execute_activity(
             tool_dispatch_activity,
@@ -199,16 +172,14 @@ async def wrap_tool(self, call: ai.middleware.ToolContext, next: Any) -> Any:
             ),
             start_to_close_timeout=datetime.timedelta(minutes=2),
         )
-        tool_msg = ai.tool_message(
+        return ai.tool_message(
             ai.ToolResultPart(
                 tool_call_id=call.tool_call_id,
                 tool_name=call.tool_name,
                 result=result.result,
                 is_error=result.is_error,
             )
         )
-        self._messages.append(tool_msg)
-        return tool_msg
 
 
 # ── Agent (default loop — no customization) ──────────────────────
@@ -237,9 +208,9 @@ async def run(self, user_query: str) -> str:
                 "description": t.description,
                 "param_schema": t.param_schema,
             }
-            for t in weather_agent._tools
+            for t in weather_agent.tools
         ]
-        mw = TemporalMiddleware(messages, tool_schemas)
+        mw = TemporalMiddleware(tool_schemas)
 
         final_text = ""
         async for msg in weather_agent.run(model, messages, middleware=[mw]):

src/ai/__init__.py

@@ -21,6 +21,7 @@
     ImageParams,
     Model,
     ModelCost,
+    StreamResult,
     VideoParams,
     check_connection,
     generate,
@@ -36,6 +37,7 @@
     Part,
     PartState,
     ReasoningPart,
+    StreamResultLike,
     StructuredOutputPart,
     TextPart,
     ToolCallPart,
@@ -85,6 +87,8 @@
     "ImageParams",
     "VideoParams",
     "Client",
+    "StreamResult",
+    "StreamResultLike",
     "check_connection",
     "model",
     "models",

src/ai/agents/agent.py

@@ -281,6 +281,11 @@ def __init__(
         self._tools: list[Tool[..., Any]] = tools or []
         self._loop_fn: LoopFn = _default_loop
 
+    @property
+    def tools(self) -> list[Tool[..., Any]]:
+        """The agent's registered tools (read-only copy)."""
+        return list(self._tools)
+
     def loop(self, fn: LoopFn) -> LoopFn:
         """Decorator: override the default loop function."""
         self._loop_fn = fn
@@ -332,9 +337,13 @@ async def _real(
         # 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:
-            mw_token = middleware_.activate(middleware)
+            parent = middleware_.get()
+            mw_token = middleware_.activate(parent + middleware)
         try:
             chain = middleware_._build_agent_run_chain(_real)
             async for message in chain(call):

src/ai/middleware.py

@@ -24,6 +24,7 @@
 
 from .types import messages as messages_
 from .types import tools as tools_
+from .types.stream import StreamResultLike
 
 # ---------------------------------------------------------------------------
 # Call context objects — frozen dataclasses with isolated mutable fields.
@@ -115,15 +116,7 @@ def __post_init__(self) -> None:
 # Middleware base class — override the methods you care about.
 # ---------------------------------------------------------------------------
 
-# Forward-reference type aliases to avoid importing heavy modules at the
-# top level.  The actual types are checked at call sites only.
-#
-# These are intentionally broad (Any) so that the middleware module does
-# not import from models/ or agents/ — which would create circular deps.
-
-# StreamResult from models/__init__.py — async-iterable of Message snapshots.
-_StreamResult = Any
-# Message
+# Message alias for brevity in signatures.
 _Message = messages_.Message
 
 # Agent run next-function type: call -> async generator of messages.
@@ -159,13 +152,24 @@ async def wrap_agent_run(self, call, next):
     async def wrap_model(
         self,
         call: ModelContext,
-        next: Callable[[ModelContext], Awaitable[_StreamResult]],
-    ) -> _StreamResult:
+        next: Callable[[ModelContext], Awaitable[StreamResultLike]],
+    ) -> StreamResultLike:
         """Wrap a model streaming call.
 
-        ``next(call)`` returns a :class:`~ai.models.StreamResult` that is
-        async-iterable over ``Message`` snapshots.  You can do work before,
-        iterate / transform the stream, or do cleanup after.
+        ``next(call)`` returns a :class:`~ai.types.StreamResultLike` that
+        is async-iterable over ``Message`` snapshots.  You can do work
+        before, iterate / transform the stream, or do cleanup after.
+
+        To transform the stream, use
+        :meth:`~ai.models.StreamResult.from_generator`::
+
+            async def wrap_model(self, call, next):
+                stream = await next(call)
+                async def _add_suffix():
+                    async for msg in stream:
+                        yield msg
+                from ai.models import StreamResult
+                return StreamResult.from_generator(_add_suffix())
         """
         return await next(call)
 
@@ -242,8 +246,8 @@ def deactivate(token: Token) -> None:
 
 
 def _build_model_chain(
-    real: Callable[[ModelContext], Awaitable[_StreamResult]],
-) -> Callable[[ModelContext], Awaitable[_StreamResult]]:
+    real: Callable[[ModelContext], Awaitable[StreamResultLike]],
+) -> Callable[[ModelContext], Awaitable[StreamResultLike]]:
     mw = get()
     if not mw:
         return real
@@ -252,9 +256,10 @@ def _build_model_chain(
     for m in reversed(mw):
 
         def _make(
-            m: Middleware, nxt: Callable[[ModelContext], Awaitable[_StreamResult]]
-        ) -> Callable[[ModelContext], Awaitable[_StreamResult]]:
-            async def _wrapped(call: ModelContext) -> _StreamResult:
+            m: Middleware,
+            nxt: Callable[[ModelContext], Awaitable[StreamResultLike]],
+        ) -> Callable[[ModelContext], Awaitable[StreamResultLike]]:
+            async def _wrapped(call: ModelContext) -> StreamResultLike:
                 return await m.wrap_model(call, nxt)
 
             return _wrapped
@@ -357,6 +362,7 @@ async def _wrapped(call: AgentRunContext) -> AsyncGenerator[_Message]:
     "HookContext",
     "Middleware",
     "ModelContext",
+    "StreamResultLike",
     "ToolContext",
     "activate",
     "deactivate",

src/ai/models/__init__.py

@@ -39,6 +39,7 @@
 from .. import middleware as middleware_
 from ..types import messages as messages_
 from ..types import tools as tools_
+from ..types.stream import StreamResultLike
 from .ai_gateway.types import GenerateParams, ImageParams, VideoParams
 from .core.catalog import get_models, get_providers, register_catalog
 from .core.catalog import model as model
@@ -157,12 +158,32 @@ class StreamResult:
 
     Properties like ``.text`` and ``.tool_calls`` delegate to the final
     ``Message`` snapshot and are available after iteration completes.
+
+    Satisfies :class:`~ai.types.StreamResultLike`.
     """
 
     def __init__(self, gen: AsyncGenerator[messages_.Message]) -> None:
         self._gen = gen
         self._final: messages_.Message | None = None
 
+    @classmethod
+    def from_generator(cls, gen: AsyncGenerator[messages_.Message]) -> StreamResult:
+        """Create a :class:`StreamResult` from an async generator.
+
+        This is the public API for middleware that needs to transform or
+        replace the stream returned by ``wrap_model``::
+
+            async def wrap_model(self, call, next):
+                original = await next(call)
+
+                async def _transformed():
+                    async for msg in original:
+                        yield modify(msg)
+
+                return StreamResult.from_generator(_transformed())
+        """
+        return cls(gen)
+
     def __aiter__(self) -> AsyncGenerator[messages_.Message]:
         return self._iterate()
 
@@ -197,12 +218,15 @@ async def stream(
     output_type: type[pydantic.BaseModel] | None = None,
     client: Client | None = None,
     **kwargs: Any,
-) -> StreamResult:
+) -> StreamResultLike:
     """Stream an LLM response.
 
-    Returns a :class:`StreamResult` that is async-iterable and collects
-    the final ``Message``.  After iteration, access ``.text``,
+    Returns a :class:`StreamResultLike` that is async-iterable and
+    collects the final ``Message``.  After iteration, access ``.text``,
     ``.tool_calls``, ``.usage``, etc.
+
+    Without middleware the concrete type is :class:`StreamResult`; with
+    middleware it may be any :class:`~ai.StreamResultLike`.
     """
     call = middleware_.ModelContext(
         model=model,
@@ -213,7 +237,7 @@ async def stream(
         kwargs=kwargs,
     )
 
-    async def _real(call: middleware_.ModelContext) -> StreamResult:
+    async def _real(call: middleware_.ModelContext) -> StreamResultLike:
         _ensure_adapters()
         c = call.client or _auto_client(call.model)
         adapter_fn = _stream_adapters.get(call.model.adapter)
@@ -235,8 +259,7 @@ async def _real(call: middleware_.ModelContext) -> StreamResult:
         )
 
     chain = middleware_._build_model_chain(_real)
-    result: StreamResult = await chain(call)
-    return result
+    return await chain(call)
 
 
 async def generate(
@@ -337,6 +360,7 @@ async def buffer(gen: AsyncGenerator[messages_.Message]) -> messages_.Message:
     "ModelCost",
     "StreamFn",
     "StreamResult",
+    "StreamResultLike",
     "VideoParams",
     # Catalog
     "get_models",

src/ai/types/__init__.py

@@ -19,6 +19,7 @@
     Usage,
     generate_id,
 )
+from .stream import StreamResultLike
 from .tools import ToolLike, ToolSchema
 
 __all__ = [
@@ -28,6 +29,7 @@
     "Part",
     "PartState",
     "ReasoningPart",
+    "StreamResultLike",
     "StructuredOutputPart",
     "TextPart",
     "ToolCallPart",

src/ai/types/stream.py

@@ -0,0 +1,36 @@
+"""StreamResultLike — structural protocol for stream results.
+
+Middleware authors can type-check against this protocol without depending
+on the concrete ``StreamResult`` class in ``ai.models``.
+"""
+
+from __future__ import annotations
+
+from collections.abc import AsyncGenerator
+from typing import Any, Protocol, runtime_checkable
+
+from . import messages as messages_
+
+
+@runtime_checkable
+class StreamResultLike(Protocol):
+    """Structural protocol satisfied by :class:`ai.models.StreamResult`.
+
+    Middleware that transforms or replaces the stream returned by
+    ``wrap_model`` should return an object satisfying this protocol.
+    The easiest way is ``StreamResult.from_generator(gen)``.
+    """
+
+    def __aiter__(self) -> AsyncGenerator[messages_.Message]: ...
+
+    @property
+    def text(self) -> str: ...
+
+    @property
+    def tool_calls(self) -> list[messages_.ToolCallPart]: ...
+
+    @property
+    def usage(self) -> messages_.Usage | None: ...
+
+    @property
+    def output(self) -> Any: ...