Make serverless hooks work better

examples/fastapi-vite/backend/agent.py

@@ -54,29 +54,17 @@ async def graph(context: ai.Context) -> AsyncGenerator[ai.events.AgentEvent]:
     request event on the SSE stream.  The frontend displays Approve /
     Reject buttons and sends the decision back on the next request.
     """
-    while True:
-        # If the last message is already an assistant turn, we're being
-        # replayed after a hook resolution — skip the model call and go
-        # straight to processing its tool calls (resolutions are
-        # pre-registered, so the hook returns immediately).
-        if context.keep_running():
-            s = ai.models.stream(context.model, context.messages, tools=context.tools)
+    while context.keep_running():
+        async with ai.models.stream(
+            context.model, context.messages, tools=context.tools
+        ) as s:
             async for event in s:
                 yield event
-            context.add(s.message)
-
-        # Pull tool calls off the most recent assistant message.
-        # ``to_messages`` may append "internal" hook messages after the
-        # assistant, so ``messages[-1]`` isn't always the assistant turn.
-        last_assistant = next(
-            (m for m in reversed(context.messages) if m.role == "assistant"),
-            None,
-        )
-        tool_calls = (
-            context.resolve(last_assistant.tool_calls) if last_assistant else []
-        )
+        context.add(s.message)
+
+        tool_calls = context.resolve(s.tool_calls)
         if not tool_calls:
-            return
+            continue
 
         results = await asyncio.gather(
             *(_execute_with_approval(tc) for tc in tool_calls)

examples/fastapi-vite/backend/main.py

@@ -57,7 +57,10 @@ class ChatRequest(pydantic.BaseModel):
 @app.post("/chat")
 async def chat(request: ChatRequest) -> fastapi.responses.StreamingResponse:
     """Handle chat requests and stream responses."""
-    messages = ai.agents.ui.ai_sdk.to_messages(request.messages)
+    messages, approvals = ai.agents.ui.ai_sdk.to_messages(request.messages)
+    # Pre-register hook resolutions so the agent loop's hooks find them
+    # immediately on the resume turn.
+    ai.agents.ui.ai_sdk.apply_approvals(approvals)
     result = agent_.chat_agent.run(agent_.MODEL, messages)
 
     async def stream_response() -> AsyncGenerator[str]:

examples/samples/agent_hooks_serverless.py

@@ -10,8 +10,7 @@
   2. Second run: resolve_hook() pre-registers the answer, agent.run()
      replays from the same input, and hook finds the resolution immediately.
 
-TODO: The implementation *works* but requires some hacks on the user
-side. We need to come up with cleaner API support.
+TODO: This works, but currently requires not using ToolRunner!
 """
 
 import asyncio
@@ -47,24 +46,16 @@ async def main() -> None:
     async def with_confirmation(
         context: ai.Context,
     ) -> AsyncGenerator[ai.events.AgentEvent]:
-        while True:
-            # HACK: If there isn't anything to do, it's because we hit
-            # a hook and bailed out. Skip running the stream in that
-            # case, but process its tool calls (which should have
-            # resolved now).
-            if context.keep_running():
-                s = ai.models.stream(
-                    context.model, context.messages, tools=context.tools
-                )
+        while context.keep_running():
+            async with ai.models.stream(
+                context.model, context.messages, tools=context.tools
+            ) as s:
                 async for event in s:
                     yield event
 
-                context.add(s.message)
-
-            tool_calls = context.resolve(context.messages[-1].tool_calls)
-            if not tool_calls:
-                return
+            context.add(s.message)
 
+            tool_calls = context.resolve(s.tool_calls)
             results: list[ai.events.ToolCallResult] = []
             for tc in tool_calls:
                 try:
@@ -90,7 +81,8 @@ async def with_confirmation(
                         )
                     )
 
-            context.add(ai.tool_message(*results))
+            if results:
+                context.add(ai.tool_message(*results))
 
     messages = [
         ai.system_message("Delete files when asked. Always use the delete_file tool."),

src/ai/agents/agent.py

@@ -458,7 +458,7 @@ async def __aenter__(self) -> Self:
         return self
 
     async def __aexit__(self, *args: Any) -> None:
-        await self._tg_base.__aexit__(*args)
+        return await self._tg_base.__aexit__(*args)
 
     def events(self) -> AsyncGenerator[events_.ToolCallResult]:
         return self._iterate()
@@ -530,9 +530,11 @@ def model_post_init(self, __context: Any) -> None:
 
     def keep_running(self) -> bool:
         """Call at top of an agent loop to see whether to keep running."""
-        return bool(
-            self.messages and self.messages[-1].role not in ("assistant", "internal")
-        )
+        if not self.messages:
+            return False
+
+        last_message = self.messages[-1]
+        return last_message.replay or last_message.role not in ("assistant", "internal")
 
     @overload
     def resolve(self, tool_part: types.messages.ToolCallPart) -> ToolCall: ...
@@ -558,12 +560,23 @@ def resolve(
     def add(
         self, message: types.messages.Message | Sequence[types.messages.Message] | None
     ) -> None:
+        """Append message(s) to the context, skipping any flagged ``replay``.
+
+        Replay-flagged messages come from ``models.stream`` short-
+        circuiting an existing assistant turn (resume-after-approval).
+        The default loop calls ``context.add(stream.message)`` after
+        every stream — the flag lets that be a no-op on replay rather
+        than producing a duplicate turn.
+        """
         if message is None:
             return
-        if isinstance(message, types.messages.Message):
-            self.messages.append(message)
-        else:
-            self.messages.extend(message)
+        msgs = (
+            [message] if isinstance(message, types.messages.Message) else list(message)
+        )
+        for msg in msgs:
+            if msg.replay:
+                continue
+            self.messages.append(msg)
 
 
 class LoopFn(Protocol):
@@ -749,8 +762,27 @@ async def _real(
                 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 for event in runtime.run(source):
+                # 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 isinstance(event, events_.BaseEvent) and event.replay:
+                    continue
                 yield event
 
         # Activate middleware for this run (and everything it calls).

src/ai/agents/ui/ai_sdk/inbound.py

@@ -198,47 +198,34 @@ def _normalize_ui_messages(
 
 def to_messages(
     ui_messages: list[ui_message.UIMessage],
-    *,
-    apply_approvals_: bool = True,
-) -> list[messages_.Message]:
-    """Parse a UI request into runtime messages.
+) -> tuple[list[messages_.Message], list[ApprovalResponse]]:
+    """Parse a UI request into runtime messages + extracted approvals.
 
-    Pipeline:
+    Pure: normalizes stale tool states, extracts approval responses,
+    parses UIMessages into an ``ai.messages.Message`` list (split at
+    tool boundaries), and drops the internal tombstones for approval
+    responses.
 
-    1. normalize stale tool states (``call`` → ``input-available``, etc.)
-    2. extract approval responses
-    3. parse UIMessages into ``ai.messages.Message`` list, splitting at tool boundaries
-    4. if *apply_approvals_* is True, pre-register resolutions via ``resolve_hook``
-    5. strip trailing assistant message when approval responses are present
+    Returns ``(messages, approvals)``.  The caller can pre-register
+    resolutions via :func:`apply_approvals` before calling
+    :meth:`Agent.run` if the run should resume from a hook.
     """
     normalized = _normalize_ui_messages(ui_messages)
     approvals = extract_approvals(normalized)
-    messages = _parse(normalized)
-
-    if apply_approvals_:
-        apply_approvals(approvals)
-
-    # TODO: Stripping the trailing assistant message broke the approval
-    # replay flow — the loop needs that message (with the original
-    # tool_call_ids) to match the pre-registered hook resolutions.
-    # Without it, the loop re-calls the LLM, gets new tool_call_ids,
-    # and the resolutions never match. Loops should instead check
-    # `context.keep_running()` and skip the model call when the last
-    # message is already an assistant turn.
-    #
-    # if approvals and messages:
-    #     # The assistant message that originated the approval-responded tool
-    #     # call would re-send the duplicate tool-use to the LLM on replay.
-    #     # Walk past any trailing internal (hook) messages and drop the
-    #     # assistant message beneath them.
-    #     idx = len(messages) - 1
-    #     while idx >= 0 and messages[idx].role == "internal":
-    #         idx -= 1
-    #     if idx >= 0 and messages[idx].role == "assistant":
-    #         logger.info("Stripping assistant message originating responded approvals")
-    #         messages = messages[:idx] + messages[idx + 1 :]
-
-    return messages
+    messages = [m for m in _parse(normalized) if not _is_approval_response(m)]
+    return messages, approvals
+
+
+def _is_approval_response(msg: messages_.Message) -> bool:
+    """Internal message that records a resolved tool-approval hook."""
+    if msg.role != "internal" or len(msg.parts) != 1:
+        return False
+    part = msg.parts[0]
+    return (
+        isinstance(part, messages_.HookPart)
+        and part.hook_type == "ToolApproval"
+        and part.status == "resolved"
+    )
 
 
 def _parse(

src/ai/models/core/api.py

@@ -117,9 +117,23 @@ def _normalize_stream_params[ProviderParamsT: pydantic.BaseModel](
 class Stream:
     """Async-iterable wrapper around an adapter's event stream."""
 
-    def __init__(self, gen: AsyncGenerator[types.events.Event]) -> None:
+    def __init__(
+        self,
+        gen: AsyncGenerator[types.events.Event],
+        *,
+        seed_message: types.messages.Message | None = None,
+    ) -> None:
+        """Wrap an event generator.
+
+        ``seed_message`` seeds the in-progress assistant message. Pass
+        a copy of an existing turn when replaying so
+        ``stream.message`` ends up identical to that turn instead of
+        being rebuilt from synthetic events.  When ``None`` (default),
+        an empty assistant message is created and rebuilt from the
+        incoming events.
+        """
         self._gen = gen
-        self._message: types.messages.Message = types.messages.Message(
+        self._message: types.messages.Message = seed_message or types.messages.Message(
             role="assistant", parts=[]
         )
         self._parts: dict[str, types.messages.Part] = {}
@@ -180,6 +194,11 @@ def output(self) -> Any:
     def _aggregate_event(self, event: types.events.Event) -> dict[str, Any]:
         updates: dict[str, Any] = {}
 
+        # Replay events carry no new state — the seeded message already
+        # has everything they would have produced.
+        if event.replay:
+            return updates
+
         # grab usage from any event that carries one
         if event.usage is not None:
             self._message.usage = event.usage
@@ -264,6 +283,27 @@ def _aggregate_event(self, event: types.events.Event) -> dict[str, Any]:
         return updates
 
 
+async def _replay_tool_calls(
+    msg: types.messages.Message,
+) -> AsyncGenerator[types.events.Event]:
+    """Replay an assistant turn's tool calls as ``replay``-flagged ``ToolEnd``.
+
+    Used by :func:`stream` to short-circuit when the last message is
+    already marked for replay — letting resume flows (e.g. post-hook
+    re-entry) re-dispatch the existing tool calls without hitting the
+    LLM and without re-streaming the original text/reasoning to the
+    consumer.  The wrapping :class:`Stream`'s ``message`` is seeded
+    with the original turn so callers see the same parts they would
+    have without replay.
+    """
+    for part in msg.tool_calls:
+        yield types.events.ToolEnd(
+            tool_call_id=part.tool_call_id,
+            tool_call=part,
+            replay=True,
+        )
+
+
 def stream[ProviderParamsT: pydantic.BaseModel](
     model: model_.Model[ProviderParamsT],
     messages: list[types.messages.Message],
@@ -273,7 +313,15 @@ def stream[ProviderParamsT: pydantic.BaseModel](
     params: params_.StreamParams[ProviderParamsT] | None = None,
     executor: StreamExecutor = _default_executor,
 ) -> Stream:
-    """Stream an LLM response."""
+    """Stream an LLM response.
+
+    If the last message is marked ``replay=True``, replay that turn as
+    synthetic stream events instead of calling the model.
+    """
+    if messages and messages[-1].replay:
+        last = messages[-1]
+        return Stream(_replay_tool_calls(last), seed_message=last.model_copy(deep=True))
+
     messages = integrity.prepare_messages(messages)
     stream_params = _normalize_stream_params(model, params)
     request = StreamRequest[ProviderParamsT](

src/ai/types/events.py

@@ -24,10 +24,17 @@ class BaseEvent(pydantic.BaseModel):
     streaming layer aggregates parts into it as deltas arrive and stamps
     a reference onto each yielded event. ``usage`` carries the latest
     usage value reported by the provider (latest-wins across the stream).
+
+    ``replay`` is set on synthetic events emitted when ``models.stream``
+    short-circuits an existing assistant turn (resume-after-approval
+    flows).  ``Agent.run`` drops replay-flagged events from the consumer-
+    facing stream — the loop's tool dispatcher still consumes them
+    internally.  Excluded from JSON: it's a control flag, not data.
     """
 
     message: messages.Message = _DUMMY_MESSAGE
     usage: usage_.Usage | None = None
+    replay: bool = pydantic.Field(default=False, exclude=True, repr=False)
 
     model_config = pydantic.ConfigDict(frozen=True)
 

src/ai/types/messages.py

@@ -229,6 +229,14 @@ class Message(pydantic.BaseModel):
     turn_id: str | None = None
     usage: usage_.Usage | None = None
 
+    # Set on the seeded message that ``models.stream`` returns when
+    # short-circuiting an existing assistant turn (resume-after-approval
+    # flows).  ``Context.add`` skips replay-flagged messages so the loop
+    # can call ``context.add(stream.message)`` unconditionally without
+    # producing a duplicate turn.  Excluded from JSON: control flag,
+    # not data.
+    replay: bool = pydantic.Field(default=False, exclude=True, repr=False)
+
     @property
     def text(self) -> str:
         """Concatenated text parts."""