From f4bd08ac41b1683d0fa83cf67dcd73fea5cc5fb9 Mon Sep 17 00:00:00 2001
From: "Michael J. Sullivan" <sully@msully.net>
Date: Wed, 7 Oct 2026 10:20:51 -0700
Subject: [PATCH 1/2] Filter out `ToolCallResult`s for tools canceled by a
 `Retry` in ui adapter

In some variants of a durable setup, it's impossible to prevent tool
results for tools that have been aborted due to a llm retry from
showing up in the workflow stream *after* the Retry event does.
Downstream users will need to deal with this.

Add a RetryFilter class for dealing with it and also update the ui
adapter to use it. MessageHydrator and other things use it
internally to make their logic work better but don't filter the
events.

Maybe we want to expose a more ergonomic adapater also, I'm not sure?
---
 .../content/docs/reference/events.mdx         | 23 +++++
 src/ai/agents/agent.py                        |  6 ++
 src/ai/types/events.py                        | 52 +++++++++++
 src/ai/ui/ai_sdk/outbound_stream.py           |  5 ++
 tests/agents/test_message_aggregator.py       | 15 +++-
 tests/agents/test_run_blocked.py              | 69 +++++++++++++-
 tests/types/test_events.py                    | 89 +++++++++++++++++++
 tests/ui/ai_sdk/test_outbound_stream.py       | 72 +++++++++++++++
 8 files changed, 327 insertions(+), 4 deletions(-)

diff --git a/docs/ai-python/content/docs/reference/events.mdx b/docs/ai-python/content/docs/reference/events.mdx
index 051cbc36..caf4849d 100644
--- a/docs/ai-python/content/docs/reference/events.mdx
+++ b/docs/ai-python/content/docs/reference/events.mdx
@@ -88,6 +88,29 @@ Agent event union:
   `kind: "hook"`; validating against `HookEvent` dispatches to the
   per-status subclass from `hook.status`.
 
+## 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 5a836b7e..10753df4 100644
--- a/src/ai/agents/agent.py
+++ b/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.
diff --git a/src/ai/types/events.py b/src/ai/types/events.py
index 9222b242..8057c73e 100644
--- a/src/ai/types/events.py
+++ b/src/ai/types/events.py
@@ -255,6 +255,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 +265,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 +746,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).
 
@@ -780,6 +829,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 +840,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 c6a3372a..07fd6d86 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 657e66a0..69584cb3 100644
--- a/tests/types/test_events.py
+++ b/tests/types/test_events.py
@@ -396,6 +396,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())
@@ -493,3 +539,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 97ac482c..8333932d 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(
         [

From ff3d16bc3e20ef05d11d433a8fe0110dc12f188e Mon Sep 17 00:00:00 2001
From: "Michael J. Sullivan" <sully@msully.net>
Date: Tue, 27 Oct 2026 15:58:36 -0700
Subject: [PATCH 2/2] ok try doing events.drop_discarded instead

---
 .../content/docs/reference/events.mdx         | 16 +++-----
 src/ai/agents/agent.py                        |  4 +-
 src/ai/types/events.py                        | 30 +++++++++++++--
 src/ai/ui/ai_sdk/outbound_stream.py           |  2 +-
 tests/types/test_events.py                    | 38 ++++++++++---------
 tests/ui/ai_sdk/test_outbound_stream.py       |  5 +--
 6 files changed, 58 insertions(+), 37 deletions(-)

diff --git a/docs/ai-python/content/docs/reference/events.mdx b/docs/ai-python/content/docs/reference/events.mdx
index caf4849d..f120717b 100644
--- a/docs/ai-python/content/docs/reference/events.mdx
+++ b/docs/ai-python/content/docs/reference/events.mdx
@@ -88,24 +88,20 @@ Agent event union:
   `kind: "hook"`; validating against `HookEvent` dispatches to the
   per-status subclass from `hook.status`.
 
-## RetryFilter
+## 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. `RetryFilter` spots that leftover output.
+step. `drop_discarded` wraps an event stream and drops that leftover output.
 
 ```python
-retry_filter = ai.events.RetryFilter()
-async for event in events:
-    if not retry_filter.feed(event):
-        continue
+async for event in ai.events.drop_discarded(events):
     ...
 ```
 
-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.
+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
diff --git a/src/ai/agents/agent.py b/src/ai/agents/agent.py
index 10753df4..01654d51 100644
--- a/src/ai/agents/agent.py
+++ b/src/ai/agents/agent.py
@@ -234,7 +234,7 @@ 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()
+        self._retry_filter = events_._RetryFilter()
 
     def feed(self, item: events_.AgentEvent) -> None:
         if not self._retry_filter.feed(item):
@@ -1602,7 +1602,7 @@ async def _run(
 
         async def _real(call: Context) -> AsyncGenerator[events_.AgentEvent]:
             tracker = events_.RunStateTracker()
-            retry_filter = events_.RetryFilter()
+            retry_filter = events_._RetryFilter()
             source = self.loop(call)
             async with contextlib.aclosing(runtime.run(source)) as events:
                 async for event in events:
diff --git a/src/ai/types/events.py b/src/ai/types/events.py
index 8057c73e..2f6a76c1 100644
--- a/src/ai/types/events.py
+++ b/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,7 +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._retry_filter = _RetryFilter()
         self._message_selected = seed_message is not None
         self._seed_checked = False
         # A stream that exhausts without StreamEnd died mid-response.
@@ -746,7 +752,7 @@ class Retry(BaseEvent):
 ]
 
 
-class RetryFilter:
+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
@@ -792,6 +798,22 @@ def feed(self, event: BaseEvent) -> bool:
         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).
 
@@ -829,7 +851,7 @@ def __init__(self) -> None:
         self._streaming = 0
         self._stream_ended = False
         self._blocked = False
-        self._retry_filter = RetryFilter()
+        self._retry_filter = _RetryFilter()
 
     @property
     def blocked(self) -> bool:
diff --git a/src/ai/ui/ai_sdk/outbound_stream.py b/src/ai/ui/ai_sdk/outbound_stream.py
index 25728a73..430accd4 100644
--- a/src/ai/ui/ai_sdk/outbound_stream.py
+++ b/src/ai/ui/ai_sdk/outbound_stream.py
@@ -640,7 +640,7 @@ async def to_stream(
     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()
+    retry_filter = events_._RetryFilter()
 
     async for event in events:
         if not retry_filter.feed(event):
diff --git a/tests/types/test_events.py b/tests/types/test_events.py
index 69584cb3..e4d3b36f 100644
--- a/tests/types/test_events.py
+++ b/tests/types/test_events.py
@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from collections.abc import AsyncGenerator
 from typing import Any
 
 import pydantic
@@ -541,7 +542,7 @@ async def test_provider_metadata_survives_replay(self) -> None:
         assert tool.provider_metadata == {"p": {"tc": 2}}
 
 
-def test_retry_filter_drops_leftover_tool_output() -> None:
+async def test_drop_discarded_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}",
@@ -549,8 +550,8 @@ def hook(tool_call_id: str) -> events.HookEvent:
             status="pending",
             tool_call_id=tool_call_id,
         )
-        return events.HookEvent(
-            message=messages.Message(role="internal", parts=[part]), hook=part
+        return events.HookEvent.for_hook(
+            messages.Message(role="internal", parts=[part]), part
         )
 
     def result(tool_call_id: str) -> events.ToolCallResult:
@@ -567,18 +568,21 @@ def partial(tool_call_id: str) -> 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"))
-
+    stream = [
+        events.StreamStart(),
+        events.ToolStart(tool_call_id="tc1", tool_name="search"),
+        events.ToolStart(tool_call_id="tc2", tool_name="search"),
+        events.Retry(),
+        events.StreamStart(),
+        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"))
+    dropped = [result("tc1"), partial("tc1"), hook("tc1")]
+    kept = [result("tc2"), partial("tc2"), hook("tc2")]
+
+    async def source() -> AsyncGenerator[events.BaseEvent]:
+        for event in [*stream, *dropped, *kept]:
+            yield event
+
+    out = [e async for e in events.drop_discarded(source())]
+    assert out == [*stream, *kept]
diff --git a/tests/ui/ai_sdk/test_outbound_stream.py b/tests/ui/ai_sdk/test_outbound_stream.py
index 8333932d..a5f7bb51 100644
--- a/tests/ui/ai_sdk/test_outbound_stream.py
+++ b/tests/ui/ai_sdk/test_outbound_stream.py
@@ -805,9 +805,8 @@ async def test_retry_drops_late_approval_for_discarded_tool_call() -> None:
             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_.HookEvent.for_hook(
+                messages_.Message(role="internal", parts=[hook]), hook
             ),
             events_.StreamEnd(),
         ]