From 108fe3e053bf159b32b4e2bd1b3398026c722b63 Mon Sep 17 00:00:00 2001 From: "Michael J. Sullivan" Date: Thu, 7 May 2026 17:56:12 -0700 Subject: [PATCH 1/2] Make serverless hooks work better * Add a notion of an event being marked as `replay`, which can be fed to the agent loop to trigger tool calls, but won't get returned to the user. * When the last message when calling `run` is an assistant reply with tool results, we treat that as a runnable state and replay the ToolEnds into the agent loop * ai-sdk inbound strips internal hook messages and returns the list of approvals granted by hooks, instead of the existing behavior of it signalling the hooks itself. For follow-up: the serverless-style loops are much nicer now, though still don't match the stock loop's ToolRunner pattern. The current issue, which I am going to think about some more, is that in the interleaved streaming flow, if we immediately bail out of the loop when a hook is requested, then that prevents a message from getting emitted at all. --- examples/fastapi-vite/backend/agent.py | 28 ++---- examples/fastapi-vite/backend/main.py | 5 +- examples/samples/agent_hooks_serverless.py | 26 ++---- src/ai/agents/agent.py | 37 ++++++-- src/ai/agents/ui/ai_sdk/inbound.py | 59 +++++------- src/ai/models/core/api.py | 60 ++++++++++++- src/ai/types/events.py | 7 ++ src/ai/types/messages.py | 8 ++ tests/agents/ui/ai_sdk/test_inbound.py | 60 +++++++++---- tests/models/core/test_api.py | 100 +++++++++++++++++++++ 10 files changed, 288 insertions(+), 102 deletions(-) diff --git a/examples/fastapi-vite/backend/agent.py b/examples/fastapi-vite/backend/agent.py index 1d4c87f6..13163780 100644 --- a/examples/fastapi-vite/backend/agent.py +++ b/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) diff --git a/examples/fastapi-vite/backend/main.py b/examples/fastapi-vite/backend/main.py index f091ac2d..6896b613 100644 --- a/examples/fastapi-vite/backend/main.py +++ b/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]: diff --git a/examples/samples/agent_hooks_serverless.py b/examples/samples/agent_hooks_serverless.py index 8239fc4e..a0281d08 100644 --- a/examples/samples/agent_hooks_serverless.py +++ b/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."), diff --git a/src/ai/agents/agent.py b/src/ai/agents/agent.py index 18456b9b..95bd05cb 100644 --- a/src/ai/agents/agent.py +++ b/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,13 @@ 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] + if last_message.role == "assistant": + return bool(last_message.tool_calls) + return last_message.role != "internal" @overload def resolve(self, tool_part: types.messages.ToolCallPart) -> ToolCall: ... @@ -558,12 +562,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): @@ -751,6 +766,12 @@ async def _real( context._agent_tools_by_name = {tool.name: tool for tool in call.tools} 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). diff --git a/src/ai/agents/ui/ai_sdk/inbound.py b/src/ai/agents/ui/ai_sdk/inbound.py index f57918ab..e0c549ca 100644 --- a/src/ai/agents/ui/ai_sdk/inbound.py +++ b/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 is responsible for + pre-registering 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( diff --git a/src/ai/models/core/api.py b/src/ai/models/core/api.py index 604666f8..29b81a80 100644 --- a/src/ai/models/core/api.py +++ b/src/ai/models/core/api.py @@ -117,9 +117,22 @@ 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], + *, + message: types.messages.Message | None = None, + ) -> None: + """Wrap an event generator. + + ``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 = message or types.messages.Message( role="assistant", parts=[] ) self._parts: dict[str, types.messages.Part] = {} @@ -180,6 +193,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 +282,29 @@ 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 an assistant turn — 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. ``Agent.run`` drops these flagged events + from the consumer-facing stream while still using them to schedule + tool execution. + """ + 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 +314,20 @@ 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 an assistant turn with tool calls, replay + that turn as synthetic stream events instead of calling the model + — useful for resume-after-approval flows where the assistant turn + is already in history and re-asking would give a different answer. + """ + # Check replay before integrity normalization so the orphaned + # assistant tool call doesn't get a synthetic tool result appended. + if messages and messages[-1].role == "assistant" and messages[-1].tool_calls: + last = messages[-1] + seeded = last.model_copy(deep=True, update={"replay": True}) + return Stream(_replay_tool_calls(last), message=seeded) + messages = integrity.prepare_messages(messages) stream_params = _normalize_stream_params(model, params) request = StreamRequest[ProviderParamsT]( diff --git a/src/ai/types/events.py b/src/ai/types/events.py index 8c7dc9c0..e3bf0f62 100644 --- a/src/ai/types/events.py +++ b/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) diff --git a/src/ai/types/messages.py b/src/ai/types/messages.py index 03bf45a2..cfb11e0a 100644 --- a/src/ai/types/messages.py +++ b/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.""" diff --git a/tests/agents/ui/ai_sdk/test_inbound.py b/tests/agents/ui/ai_sdk/test_inbound.py index da0cd7d2..2f4a3b03 100644 --- a/tests/agents/ui/ai_sdk/test_inbound.py +++ b/tests/agents/ui/ai_sdk/test_inbound.py @@ -10,6 +10,7 @@ extract_approvals, ) from ai.agents.ui.ai_sdk.ui_message import UIMessage, UIToolPart +from ai.types import messages as messages_ def _ui(role: str, *parts: dict[str, Any], id: str = "m1") -> UIMessage: @@ -35,14 +36,15 @@ def _tool( def test_to_messages_user_text() -> None: - result = to_messages([_ui("user", _text("hello"))]) - assert len(result) == 1 - assert result[0].role == "user" - assert result[0].text == "hello" + messages, approvals = to_messages([_ui("user", _text("hello"))]) + assert len(messages) == 1 + assert messages[0].role == "user" + assert messages[0].text == "hello" + assert approvals == [] def test_to_messages_splits_at_tool_boundary() -> None: - result = to_messages( + messages, _ = to_messages( [ _ui( "assistant", @@ -58,12 +60,13 @@ def test_to_messages_splits_at_tool_boundary() -> None: ) ] ) - assert [m.role for m in result] == ["assistant", "tool", "assistant"] - assert result[1].tool_results[0].tool_call_id == "tc1" + assert [m.role for m in messages] == ["assistant", "tool", "assistant"] + assert messages[1].tool_results[0].tool_call_id == "tc1" -def test_to_messages_approval_hook_emitted_as_internal() -> None: - result = to_messages( +def test_to_messages_keeps_pending_approval_tombstone() -> None: + """Pending approvals carry no response — leave the tombstone in history.""" + messages, _ = to_messages( [ _ui( "assistant", @@ -75,16 +78,39 @@ def test_to_messages_approval_hook_emitted_as_internal() -> None: ), ) ], - apply_approvals_=False, ) - assert [m.role for m in result] == ["assistant", "internal"] - hook = result[1].parts[0] - assert hook.kind == "hook" - assert hook.hook_id == "approve_tc1" + assert [m.role for m in messages] == ["assistant", "internal"] + hook_part = messages[1].parts[0] + assert isinstance(hook_part, messages_.HookPart) + assert hook_part.hook_type == "ToolApproval" + assert hook_part.status == "pending" + + +def test_to_messages_drops_resolved_approval_tombstone() -> None: + """Resolved approvals come back via the side-channel; the tombstone is dead.""" + messages, approvals = to_messages( + [ + _ui( + "assistant", + _tool( + "delete", + "tc1", + "approval-responded", + approval={ + "id": "approve_tc1", + "approved": True, + "reason": None, + }, + ), + ) + ], + ) + assert [m.role for m in messages] == ["assistant"] + assert [a.hook_id for a in approvals] == ["approve_tc1"] def test_to_messages_keeps_trailing_assistant_when_approved() -> None: - result = to_messages( + messages, approvals = to_messages( [ _ui("user", _text("delete it"), id="u1"), _ui( @@ -98,9 +124,9 @@ def test_to_messages_keeps_trailing_assistant_when_approved() -> None: id="a1", ), ], - apply_approvals_=False, ) - assert [m.role for m in result] == ["user", "assistant", "internal"] + assert [m.role for m in messages] == ["user", "assistant"] + assert [a.hook_id for a in approvals] == ["approve_tc1"] def test_extract_approvals_returns_approved_responses() -> None: diff --git a/tests/models/core/test_api.py b/tests/models/core/test_api.py index e45196cd..895c9c64 100644 --- a/tests/models/core/test_api.py +++ b/tests/models/core/test_api.py @@ -205,3 +205,103 @@ async def test_check_connection_delegates_to_model_provider() -> None: assert result is False assert provider.received_client is explicit + + +async def test_stream_replays_last_assistant_with_tool_calls() -> None: + """If the last message is an assistant turn with tool calls, no LLM call.""" + called = False + + async def _spy_stream( + client: models.Client, + model: models.Model[pydantic.BaseModel], + messages: list[messages_.Message], + *, + tools: Sequence[ai.tools.Tool] | None = None, + output_type: type[pydantic.BaseModel] | None = None, + **kwargs: Any, + ) -> AsyncGenerator[events_.Event]: + nonlocal called + called = True + yield events_.StreamStart() + yield events_.StreamEnd() + + models.register_stream("mock", _spy_stream) + + assistant_msg = messages_.Message( + role="assistant", + parts=[ + messages_.TextPart(id="t1", text="calling tools"), + messages_.ToolCallPart( + tool_call_id="tc-1", + tool_name="weather", + tool_args='{"city":"SF"}', + ), + ], + ) + + stream = models.stream( + MOCK_MODEL, + [ai.user_message("Hi"), assistant_msg], + ) + events: list[events_.Event] = [event async for event in stream] + + assert called is False, "should not have hit the LLM" + # Stream.message is seeded from the original turn — text and tool + # calls are both preserved. + assert stream.text == "calling tools" + assert len(stream.tool_calls) == 1 + assert stream.tool_calls[0].tool_call_id == "tc-1" + assert stream.tool_calls[0].tool_args == '{"city":"SF"}' + # Replay only emits ToolEnd events, flagged for agent.run to drop. + tool_ends = [e for e in events if isinstance(e, events_.ToolEnd)] + assert len(tool_ends) == 1 + assert tool_ends[0].replay is True + assert tool_ends[0].tool_call.tool_call_id == "tc-1" + # No ToolStart/Delta/text events are re-emitted. + assert not any(isinstance(e, events_.ToolStart) for e in events) + assert not any(isinstance(e, events_.TextDelta) for e in events) + + +def test_tool_end_replay_flag_excluded_from_json() -> None: + """The replay flag is internal — it must not appear in serialized output.""" + ev = events_.ToolEnd( + tool_call_id="tc-1", + tool_call=messages_.DUMMY_TOOL_CALL, + replay=True, + ) + dumped = ev.model_dump() + assert "replay" not in dumped + dumped_json = ev.model_dump(mode="json") + assert "replay" not in dumped_json + + +async def test_stream_does_not_replay_when_assistant_has_no_tool_calls() -> None: + """Bare assistant text doesn't trigger replay — fall through to LLM.""" + called = False + + async def _spy_stream( + client: models.Client, + model: models.Model[pydantic.BaseModel], + messages: list[messages_.Message], + *, + tools: Sequence[ai.tools.Tool] | None = None, + output_type: type[pydantic.BaseModel] | None = None, + **kwargs: Any, + ) -> AsyncGenerator[events_.Event]: + nonlocal called + called = True + yield events_.StreamStart() + yield events_.StreamEnd() + + models.register_stream("mock", _spy_stream) + + assistant_text_only = messages_.Message( + role="assistant", + parts=[messages_.TextPart(text="just talking")], + ) + + stream = models.stream(MOCK_MODEL, [assistant_text_only]) + async for _ in stream: + pass + + assert called is True From 689276c5dffa21fde01f690ca88db03a304c9d42 Mon Sep 17 00:00:00 2001 From: "Michael J. Sullivan" Date: Fri, 8 May 2026 17:19:37 -0700 Subject: [PATCH 2/2] Move the main logic about when to replay into agent, and drive stream level decisions based on the replay flag --- src/ai/agents/agent.py | 17 ++++++++++++--- src/ai/agents/ui/ai_sdk/inbound.py | 6 +++--- src/ai/models/core/api.py | 34 ++++++++++++------------------ tests/models/core/test_api.py | 8 +++---- 4 files changed, 35 insertions(+), 30 deletions(-) diff --git a/src/ai/agents/agent.py b/src/ai/agents/agent.py index 95bd05cb..ae8fba4f 100644 --- a/src/ai/agents/agent.py +++ b/src/ai/agents/agent.py @@ -534,9 +534,7 @@ def keep_running(self) -> bool: return False last_message = self.messages[-1] - if last_message.role == "assistant": - return bool(last_message.tool_calls) - return last_message.role != "internal" + return last_message.replay or last_message.role not in ("assistant", "internal") @overload def resolve(self, tool_part: types.messages.ToolCallPart) -> ToolCall: ... @@ -764,6 +762,19 @@ 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 diff --git a/src/ai/agents/ui/ai_sdk/inbound.py b/src/ai/agents/ui/ai_sdk/inbound.py index e0c549ca..9706fe7d 100644 --- a/src/ai/agents/ui/ai_sdk/inbound.py +++ b/src/ai/agents/ui/ai_sdk/inbound.py @@ -206,9 +206,9 @@ def to_messages( tool boundaries), and drops the internal tombstones for approval responses. - Returns ``(messages, approvals)``. The caller is responsible for - pre-registering resolutions via :func:`apply_approvals` before - calling :meth:`Agent.run` if the run should resume from a hook. + 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) diff --git a/src/ai/models/core/api.py b/src/ai/models/core/api.py index 29b81a80..feae56a2 100644 --- a/src/ai/models/core/api.py +++ b/src/ai/models/core/api.py @@ -121,18 +121,19 @@ def __init__( self, gen: AsyncGenerator[types.events.Event], *, - message: types.messages.Message | None = None, + seed_message: types.messages.Message | None = None, ) -> None: """Wrap an event generator. - ``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. + ``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 = message or types.messages.Message( + self._message: types.messages.Message = seed_message or types.messages.Message( role="assistant", parts=[] ) self._parts: dict[str, types.messages.Part] = {} @@ -288,14 +289,12 @@ async def _replay_tool_calls( """Replay an assistant turn's tool calls as ``replay``-flagged ``ToolEnd``. Used by :func:`stream` to short-circuit when the last message is - already an assistant turn — letting resume flows (e.g. post-hook + 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. ``Agent.run`` drops these flagged events - from the consumer-facing stream while still using them to schedule - tool execution. + have without replay. """ for part in msg.tool_calls: yield types.events.ToolEnd( @@ -316,17 +315,12 @@ def stream[ProviderParamsT: pydantic.BaseModel]( ) -> Stream: """Stream an LLM response. - If the last message is an assistant turn with tool calls, replay - that turn as synthetic stream events instead of calling the model - — useful for resume-after-approval flows where the assistant turn - is already in history and re-asking would give a different answer. + If the last message is marked ``replay=True``, replay that turn as + synthetic stream events instead of calling the model. """ - # Check replay before integrity normalization so the orphaned - # assistant tool call doesn't get a synthetic tool result appended. - if messages and messages[-1].role == "assistant" and messages[-1].tool_calls: + if messages and messages[-1].replay: last = messages[-1] - seeded = last.model_copy(deep=True, update={"replay": True}) - return Stream(_replay_tool_calls(last), message=seeded) + 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) diff --git a/tests/models/core/test_api.py b/tests/models/core/test_api.py index 895c9c64..3a460887 100644 --- a/tests/models/core/test_api.py +++ b/tests/models/core/test_api.py @@ -207,8 +207,8 @@ async def test_check_connection_delegates_to_model_provider() -> None: assert provider.received_client is explicit -async def test_stream_replays_last_assistant_with_tool_calls() -> None: - """If the last message is an assistant turn with tool calls, no LLM call.""" +async def test_stream_replays_marked_last_assistant_with_tool_calls() -> None: + """If the last message is marked replay, no LLM call.""" called = False async def _spy_stream( @@ -241,7 +241,7 @@ async def _spy_stream( stream = models.stream( MOCK_MODEL, - [ai.user_message("Hi"), assistant_msg], + [ai.user_message("Hi"), assistant_msg.model_copy(update={"replay": True})], ) events: list[events_.Event] = [event async for event in stream] @@ -275,7 +275,7 @@ def test_tool_end_replay_flag_excluded_from_json() -> None: assert "replay" not in dumped_json -async def test_stream_does_not_replay_when_assistant_has_no_tool_calls() -> None: +async def test_stream_does_not_replay_when_assistant_is_unmarked() -> None: """Bare assistant text doesn't trigger replay — fall through to LLM.""" called = False