diff --git a/docs/ai-python/content/docs/reference/events.mdx b/docs/ai-python/content/docs/reference/events.mdx index 871d3edd..391014eb 100644 --- a/docs/ai-python/content/docs/reference/events.mdx +++ b/docs/ai-python/content/docs/reference/events.mdx @@ -81,6 +81,29 @@ Agent event union: so use it as a field or `TypeAdapter` type when deserializing events and Pydantic picks the right event class from `kind`. +## RetryFilter + +A tool scheduled for a response that later gets retried can still produce +output after the `Retry`, for example when the tool runs in a separate durable +step. `RetryFilter` spots that leftover output. + +```python +retry_filter = ai.events.RetryFilter() +async for event in events: + if not retry_filter.feed(event): + continue + ... +``` + +Feed it every event in order. `feed` returns `False` for a `ToolCallResult`, +`PartialToolCallResult`, or `HookEvent` whose tool call belonged to a discarded +response. If the retried response reuses a tool call id, output for that id is +kept again. + +`Agent.run`, `MessageHydrator`, `RunStateTracker`, `MessageAggregator`, and +the AI SDK UI adapter already filter this way, so only custom consumers of raw +event streams need it. + ## Aggregator `Aggregator[Item, Result, ModelInput]` is the interface used by streaming tools diff --git a/src/ai/agents/agent.py b/src/ai/agents/agent.py index 161afff5..8745c897 100644 --- a/src/ai/agents/agent.py +++ b/src/ai/agents/agent.py @@ -233,8 +233,11 @@ def __init__(self) -> None: self._messages: list[types.messages.Message] = [] self._index_by_id: dict[str, int] = {} self._streaming_index: int | None = None + self._retry_filter = events_.RetryFilter() def feed(self, item: events_.AgentEvent) -> None: + if not self._retry_filter.feed(item): + return if isinstance(item, events_.Retry): # Drop the in-progress model response and everything after # it (its tool results); the retried stream re-adds them, @@ -1573,9 +1576,12 @@ async def _run( async def _real(call: Context) -> AsyncGenerator[events_.AgentEvent]: tracker = events_.RunStateTracker() + retry_filter = events_.RetryFilter() source = self.loop(call) async with contextlib.aclosing(runtime.run(source)) as events: async for event in events: + if not retry_filter.feed(event): + continue # Feed the tracker before the replay filter: replayed # StreamEnds carry the tool calls the dispatcher is # about to re-run, which the fold must count. diff --git a/src/ai/types/events.py b/src/ai/types/events.py index 56a36a1d..39cd9c91 100644 --- a/src/ai/types/events.py +++ b/src/ai/types/events.py @@ -253,6 +253,7 @@ def __init__(self, seed_message: messages.Message | None = None) -> None: self.messages_by_id: dict[str, messages.Message] = {} self._parts_by_message_id: dict[str, dict[str, messages.Part]] = {} self._streaming_index: int | None = None + self._retry_filter = RetryFilter() self._message_selected = seed_message is not None self._seed_checked = False # A stream that exhausts without StreamEnd died mid-response. @@ -262,6 +263,8 @@ def __init__(self, seed_message: messages.Message | None = None) -> None: self.response_model: str | None = None def feed[T: BaseEvent](self, event: T) -> T: + if not self._retry_filter.feed(event): + return event # ToolCallResult and HookEvent always carry full messages, so we just # use them. if isinstance(event, ToolCallResult | HookEvent): @@ -670,6 +673,52 @@ class Retry(BaseEvent): ] +class RetryFilter: + """Spot tool output for tool calls that a :class:`Retry` discarded. + + In a distributed setup, a tool scheduled for a response that later + gets retried can still deliver its result (or a partial result, or + a hook event) after the ``Retry``. Feed every event in order; + :meth:`feed` returns False for such leftovers. A tool call id that + the retried response reuses counts as live again. + """ + + def __init__(self) -> None: + self._response: set[str] = set() + self._discarded: set[str] = set() + + def feed(self, event: BaseEvent) -> bool: + match event: + case StreamStart(): + self._response = set() + case ( + ToolStart(tool_call_id=id) + | ToolDelta(tool_call_id=id) + | ToolEnd(tool_call_id=id) + | BuiltinToolStart(tool_call_id=id) + | BuiltinToolDelta(tool_call_id=id) + | BuiltinToolEnd(tool_call_id=id) + ): + self._response.add(id) + self._discarded.discard(id) + case StreamEnd(): + ids = {tc.tool_call_id for tc in event.message.tool_calls} + self._response |= ids + self._discarded -= ids + case Retry(): + self._discarded |= self._response + self._response = set() + case ToolCallResult(): + return not event.results or not all( + r.tool_call_id in self._discarded for r in event.results + ) + case PartialToolCallResult(): + return event.tool_call_id not in self._discarded + case HookEvent(): + return event.hook.tool_call_id not in self._discarded + return True + + class RunStateTracker: """Fold an agent event stream into run state (blocked-on-hooks). @@ -706,6 +755,7 @@ def __init__(self) -> None: self._streaming = 0 self._stream_ended = False self._blocked = False + self._retry_filter = RetryFilter() @property def blocked(self) -> bool: @@ -716,6 +766,8 @@ def deferred_hooks(self) -> list[messages.HookPart[Any]]: return list(self._deferred.values()) def feed(self, event: AgentEvent) -> RunBlocked | None: + if not self._retry_filter.feed(event): + return None match event: case StreamStart(): self._streaming += 1 diff --git a/src/ai/ui/ai_sdk/outbound_stream.py b/src/ai/ui/ai_sdk/outbound_stream.py index a29439f0..25728a73 100644 --- a/src/ai/ui/ai_sdk/outbound_stream.py +++ b/src/ai/ui/ai_sdk/outbound_stream.py @@ -638,8 +638,13 @@ async def to_stream( ) -> AsyncGenerator[ui_events.UIMessageStreamEvent]: """Walk internal events once, emitting AI SDK UI stream events.""" state = _StreamState() + # The client drops a discarded step's tool parts on reset-step, so + # any later output for those tool calls would make it error out. + retry_filter = events_.RetryFilter() async for event in events: + if not retry_filter.feed(event): + continue match event: case events_.ToolCallResult(): for ui_event in state.on_tool_result(event): diff --git a/tests/agents/test_message_aggregator.py b/tests/agents/test_message_aggregator.py index 8d7c63b2..e8fa07d7 100644 --- a/tests/agents/test_message_aggregator.py +++ b/tests/agents/test_message_aggregator.py @@ -5,7 +5,7 @@ import ai from ai.types import events as events_ -from ..conftest import text_msg, tool_result_msg +from ..conftest import text_msg, tool_call_msg, tool_result_msg def test_consecutive_snapshots_replace() -> None: @@ -69,6 +69,19 @@ def test_retry_drops_response_and_its_tool_results() -> None: assert bundle.messages[0].text == "again" +def test_retry_drops_late_result_for_discarded_tool_call() -> None: + agg = ai.agents.MessageAggregator() + agg.feed(events_.StreamEnd(message=tool_call_msg(id="msg-a", tc_id="tc-a"))) + agg.feed(events_.Retry()) + agg.feed(events_.StreamEnd(message=tool_call_msg(id="msg-b", tc_id="tc-b"))) + stale = tool_result_msg(tc_id="tc-a", result="stale") + fresh = tool_result_msg(tc_id="tc-b", result="fresh") + agg.feed(events_.ToolCallResult(message=stale, results=stale.tool_results)) + agg.feed(events_.ToolCallResult(message=fresh, results=fresh.tool_results)) + + assert [m.id for m in agg.snapshot().messages] == ["msg-b", fresh.id] + + def test_retry_without_stream_in_progress_is_noop() -> None: agg = ai.agents.MessageAggregator() agg.feed(events_.Retry()) diff --git a/tests/agents/test_run_blocked.py b/tests/agents/test_run_blocked.py index c28dbf88..c91c6fcd 100644 --- a/tests/agents/test_run_blocked.py +++ b/tests/agents/test_run_blocked.py @@ -431,10 +431,13 @@ def test_tracker_retry_after_stream_end_forgets_scheduled_calls() -> None: # tc-2 from the discarded response must not keep the run "busy". assert tracker.feed(events_.StreamStart()) is None - assert tracker.feed(_hook_event(hook)) is None - transition = tracker.feed( - events_.StreamEnd(message=_multi_call_msg(("tc-1", "gated"))) + assert ( + tracker.feed( + events_.StreamEnd(message=_multi_call_msg(("tc-1", "gated"))) + ) + is None ) + transition = tracker.feed(_hook_event(hook)) assert isinstance(transition, events_.RunBlocked) @@ -442,3 +445,63 @@ def test_message_aggregator_tolerates_block_events() -> None: agg = ai.agents.MessageAggregator() agg.feed(events_.RunBlocked()) assert agg.snapshot().messages == () + + +def test_tracker_ignores_hook_for_retried_tool_call() -> None: + """A late pending hook for a discarded call must not block the run.""" + tracker = events_.RunStateTracker() + hook: messages_.HookPart[Any] = messages_.HookPart( + hook_id="h1", + hook_type="ToolApproval", + status="pending", + tool_call_id="tc-1", + ) + + tracker.feed(events_.StreamStart()) + tracker.feed(events_.ToolStart(tool_call_id="tc-1", tool_name="gated")) + tracker.feed(events_.Retry()) + tracker.feed(events_.StreamStart()) + tracker.feed(events_.StreamEnd(message=text_msg("done"))) + + assert tracker.feed(_hook_event(hook)) is None + assert not tracker.blocked + assert tracker.deferred_hooks == [] + + +async def test_run_drops_leftover_tool_output_after_retry() -> None: + hook: messages_.HookPart[Any] = messages_.HookPart( + hook_id="h1", + hook_type="ToolApproval", + status="pending", + tool_call_id="tc-1", + ) + result = ai.tool_message( + ai.tool_result_part("tc-1", tool_name="gated", result="stale") + ) + + class MyAgent(ai.Agent): + async def loop( + self, context: ai.Context + ) -> AsyncGenerator[events_.AgentEvent]: + yield events_.StreamStart() + yield events_.ToolStart(tool_call_id="tc-1", tool_name="gated") + yield events_.Retry() + yield events_.StreamStart() + yield _hook_event(hook) + yield events_.ToolCallResult( + message=result, results=result.tool_results + ) + yield events_.StreamEnd(message=text_msg("done")) + + async with MyAgent().run(MOCK_MODEL, [ai.user_message("go")]) as stream: + kinds = [event.kind async for event in stream] + + assert kinds == [ + "stream_start", + "tool_start", + "retry", + "stream_start", + "stream_end", + ] + assert stream.deferred_hooks == [] + assert not stream.blocked diff --git a/tests/types/test_events.py b/tests/types/test_events.py index f746c8e2..040a8773 100644 --- a/tests/types/test_events.py +++ b/tests/types/test_events.py @@ -2,6 +2,8 @@ from __future__ import annotations +from typing import Any + import pydantic import pytest @@ -381,6 +383,52 @@ def test_message_hydrator_retry_discards_response_and_tool_results() -> None: assert hydrator.message.tool_calls[0].tool_args == '{"q":"full"}' +@pytest.mark.parametrize("retried_id", ["tc2", "tc1"]) +def test_message_hydrator_drops_late_result_for_retried_tool_call( + retried_id: str, +) -> None: + hydrator = events.MessageHydrator() + dropped = messages.Message(id="assistant-1", role="assistant", parts=[]) + retried = messages.Message(id="assistant-2", role="assistant", parts=[]) + + def call(message: messages.Message, id: str) -> list[events.Event]: + part = messages.ToolCallPart( + tool_call_id=id, tool_name="search", tool_args="{}" + ) + return [ + events.StreamStart(message=message), + events.ToolStart( + message=message, tool_call_id=id, tool_name="search" + ), + events.ToolEnd(message=message, tool_call_id=id, tool_call=part), + ] + + def result(id: str) -> events.ToolCallResult: + part = messages.ToolResultPart( + tool_call_id=id, tool_name="search", result=id + ) + return events.ToolCallResult( + message=messages.Message( + id=f"tool-{id}", role="tool", parts=[part] + ), + results=[part], + ) + + for event in call(dropped, "tc1"): + hydrator.feed(event) + hydrator.feed(events.Retry()) + for event in call(retried, retried_id): + hydrator.feed(event) + hydrator.feed(result("tc1")) + hydrator.feed(result(retried_id)) + + # A result for tc1 only survives if the retried response also called it. + assert [m.id for m in hydrator.messages] == [ + "assistant-2", + f"tool-{retried_id}", + ] + + def test_message_hydrator_retry_before_any_stream_is_noop() -> None: hydrator = events.MessageHydrator() retry = hydrator.feed(events.Retry()) @@ -478,3 +526,46 @@ async def test_provider_metadata_survives_replay(self) -> None: p for p in rebuilt.parts if isinstance(p, messages.ToolCallPart) ) assert tool.provider_metadata == {"p": {"tc": 2}} + + +def test_retry_filter_drops_leftover_tool_output() -> None: + def hook(tool_call_id: str) -> events.HookEvent: + part: messages.HookPart[Any] = messages.HookPart( + hook_id=f"h-{tool_call_id}", + hook_type="ToolApproval", + status="pending", + tool_call_id=tool_call_id, + ) + return events.HookEvent( + message=messages.Message(role="internal", parts=[part]), hook=part + ) + + def result(tool_call_id: str) -> events.ToolCallResult: + part = messages.ToolResultPart( + tool_call_id=tool_call_id, tool_name="search", result="r" + ) + return events.ToolCallResult( + message=messages.Message(role="tool", parts=[part]), + results=[part], + ) + + def partial(tool_call_id: str) -> events.PartialToolCallResult: + return events.PartialToolCallResult( + tool_call_id=tool_call_id, value="p" + ) + + f = events.RetryFilter() + assert f.feed(events.StreamStart()) + assert f.feed(events.ToolStart(tool_call_id="tc1", tool_name="search")) + assert f.feed(events.ToolStart(tool_call_id="tc2", tool_name="search")) + assert f.feed(events.Retry()) + assert f.feed(events.StreamStart()) + assert f.feed(events.ToolStart(tool_call_id="tc2", tool_name="search")) + + # tc1 was discarded; tc2 came back in the retried response. + assert not f.feed(result("tc1")) + assert not f.feed(partial("tc1")) + assert not f.feed(hook("tc1")) + assert f.feed(result("tc2")) + assert f.feed(partial("tc2")) + assert f.feed(hook("tc2")) diff --git a/tests/ui/ai_sdk/test_outbound_stream.py b/tests/ui/ai_sdk/test_outbound_stream.py index fdf272cf..6544fec4 100644 --- a/tests/ui/ai_sdk/test_outbound_stream.py +++ b/tests/ui/ai_sdk/test_outbound_stream.py @@ -746,6 +746,78 @@ async def test_retry_resets_step_and_reemits_retried_parts() -> None: ] +async def test_retry_drops_late_results_for_discarded_tool_calls() -> None: + """The client has no part for a discarded tool call, so a late result + for it would make the AI SDK error out.""" + + def call(id: str) -> list[agent_events_.AgentEvent]: + part = messages_.ToolCallPart( + tool_call_id=id, tool_name="search", tool_args="{}" + ) + return [ + events_.ToolStart(tool_call_id=id, tool_name="search"), + events_.ToolEnd(tool_call_id=id, tool_call=part), + ] + + def result(id: str) -> events_.ToolCallResult: + msg = ai.tool_message( + ai.tool_result_part(id, tool_name="search", result=id) + ) + return events_.ToolCallResult(message=msg, results=msg.tool_results) + + out = await _collect( + [ + events_.StreamStart(), + *call("tc1"), + events_.Retry(), + events_.StreamStart(), + *call("tc2"), + events_.StreamEnd(), + events_.PartialToolCallResult( + tool_call_id="tc1", + tool_name="search", + value="partial", + aggregator_factory=ai.agents.ConcatAggregator, + ), + result("tc1"), + result("tc2"), + ] + ) + + outputs = [ + e.tool_call_id + for e in out + if isinstance(e, ui_events.UIToolOutputAvailableEvent) + ] + assert outputs == ["tc2"] + + +async def test_retry_drops_late_approval_for_discarded_tool_call() -> None: + hook: messages_.HookPart[Any] = messages_.HookPart( + hook_id="h1", + hook_type="ToolApproval", + status="pending", + tool_call_id="tc1", + ) + out = await _collect( + [ + events_.StreamStart(), + events_.ToolStart(tool_call_id="tc1", tool_name="bash"), + events_.Retry(), + events_.StreamStart(), + events_.HookEvent( + message=messages_.Message(role="internal", parts=[hook]), + hook=hook, + ), + events_.StreamEnd(), + ] + ) + + assert not any( + isinstance(e, ui_events.UIToolApprovalRequestEvent) for e in out + ) + + async def test_retry_before_any_output_is_dropped() -> None: out = await _collect( [