Filter out ToolCallResults for tools canceled by a Retry in ui adapter

docs/ai-python/content/docs/reference/events.mdx

@@ -88,6 +88,25 @@ Agent event union:
   `kind: "hook"`; validating against `HookEvent` dispatches to the
   per-status subclass from `hook.status`.
 
+## drop_discarded
+
+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. `drop_discarded` wraps an event stream and drops that leftover output.
+
+```python
+async for event in ai.events.drop_discarded(events):
+    ...
+```
+
+It drops 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

src/ai/agents/agent.py

@@ -234,8 +234,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,
@@ -1599,9 +1602,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.

src/ai/types/events.py

@@ -2,7 +2,13 @@
 
 import abc
 import contextvars
-from collections.abc import AsyncGenerator, Callable, Mapping, Sequence
+from collections.abc import (
+    AsyncGenerator,
+    AsyncIterable,
+    Callable,
+    Mapping,
+    Sequence,
+)
 from typing import Annotated, Any, Literal, cast
 
 import pydantic
@@ -255,6 +261,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.
@@ -264,6 +271,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):
@@ -743,6 +752,68 @@ 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
+
+
+async def drop_discarded[E: BaseEvent](
+    events: AsyncIterable[E],
+) -> AsyncGenerator[E]:
+    """Yield ``events`` minus late tool output for calls a ``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``.  A tool call id that the retried
+    response reuses counts as live again.
+    """
+    retry_filter = _RetryFilter()
+    async for event in events:
+        if retry_filter.feed(event):
+            yield event
+
+
 class RunStateTracker:
     """Fold an agent event stream into run state (blocked-on-hooks).
 
@@ -780,6 +851,7 @@ def __init__(self) -> None:
         self._streaming = 0
         self._stream_ended = False
         self._blocked = False
+        self._retry_filter = _RetryFilter()
 
     @property
     def blocked(self) -> bool:
@@ -790,6 +862,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

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):

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())

tests/agents/test_run_blocked.py

@@ -431,14 +431,77 @@ 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)
 
 
 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