diff --git a/src/ai/agents/_middleware.py b/src/ai/agents/_middleware.py index df1f4fe9..a8cb76ff 100644 --- a/src/ai/agents/_middleware.py +++ b/src/ai/agents/_middleware.py @@ -97,6 +97,7 @@ class HookContext: label: str payload: type[pydantic.BaseModel] metadata: dict[str, Any] + tool_call_id: str | None = None def __post_init__(self) -> None: object.__setattr__(self, "metadata", dict(self.metadata)) diff --git a/src/ai/agents/agent.py b/src/ai/agents/agent.py index 5e7ed6d9..553cf501 100644 --- a/src/ai/agents/agent.py +++ b/src/ai/agents/agent.py @@ -228,7 +228,7 @@ def __init__(self) -> None: self._index_by_id: dict[str, int] = {} def feed(self, item: events_.AgentEvent) -> None: - if isinstance(item, events_.PartialToolCallResult): + if isinstance(item, events_.PartialToolCallResult | events_.RunBlocked): return msg = item.message if msg is None: @@ -750,6 +750,7 @@ async def __call__(self) -> events_.ToolCallResult: f"approve_{tc.id}", payload=types.tools.ToolApproval, metadata={"tool": tc.name, "kwargs": hook_kwargs}, + tool_call_id=tc.id, ) except hooks_.HookPendingError as e: return pending_tool_result( @@ -980,18 +981,37 @@ def __init__( ) -> None: self._gen = gen self._context = context + self._blocked = False + self._pending_hooks: dict[str, types.messages.HookPart[Any]] = {} + + def _track(self, event: events_.AgentEvent) -> events_.AgentEvent: + # Fold delivered events into the exposed run state, so that the + # properties below are always consistent with the events the + # consumer has seen so far. + if isinstance(event, events_.RunBlocked): + self._blocked = True + elif isinstance(event, events_.HookEvent): + if event.hook.status == "pending": + self._pending_hooks[event.hook.hook_id] = event.hook + else: + # A blocked run can only resume via a hook resolution + # or cancellation (see RunBlocked), so this is also the + # unblock signal. + self._pending_hooks.pop(event.hook.hook_id, None) + self._blocked = False + return event def __aiter__(self) -> Self: return self async def __anext__(self) -> events_.AgentEvent: - return await self._gen.__anext__() + return self._track(await self._gen.__anext__()) async def asend(self, value: Any) -> events_.AgentEvent: - return await self._gen.asend(value) + return self._track(await self._gen.asend(value)) async def athrow(self, *args: Any, **kwargs: Any) -> events_.AgentEvent: - return await self._gen.athrow(*args, **kwargs) + return self._track(await self._gen.athrow(*args, **kwargs)) async def aclose(self) -> None: await self._gen.aclose() @@ -1000,6 +1020,25 @@ async def aclose(self) -> None: def context(self) -> Context: return self._context + @property + def blocked(self) -> bool: + """Whether the run is blocked awaiting external hook resolution. + + Set by a delivered :class:`~ai.types.events.RunBlocked` event, + cleared by the hook resolution that follows. True means + nothing in the run can make progress until a pending hook (see + :attr:`pending_hooks`) is resolved — e.g. a tool approval. + After the run ends this keeps its final value, so serverless + drivers can distinguish "turn complete" from "turn suspended + awaiting input". + """ + return self._blocked + + @property + def pending_hooks(self) -> list[types.messages.HookPart[Any]]: + """Hooks that are pending as of the last delivered event.""" + return list(self._pending_hooks.values()) + @property def messages(self) -> list[types.messages.Message]: return self._context.messages @@ -1318,15 +1357,21 @@ async def _run( _process_interrupted_hooks(context.messages) async def _real(call: Context) -> AsyncGenerator[events_.AgentEvent]: + tracker = events_.RunStateTracker() source = self.loop(call) async for event in runtime.run(source): + # 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. + transition = tracker.feed(event) # 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 + if not (isinstance(event, events_.BaseEvent) and event.replay): + yield event + if transition is not None: + yield transition async def _stream() -> AsyncGenerator[events_.AgentEvent]: # Activate middleware for this run (and everything it calls). diff --git a/src/ai/agents/hooks.py b/src/ai/agents/hooks.py index c1e64654..c52819d2 100644 --- a/src/ai/agents/hooks.py +++ b/src/ai/agents/hooks.py @@ -86,6 +86,7 @@ async def hook[T: pydantic.BaseModel]( *, payload: type[T], metadata: dict[str, Any] | None = None, + tool_call_id: str | None = None, ) -> T: """Create a hook suspension point and await its resolution. @@ -97,12 +98,18 @@ async def hook[T: pydantic.BaseModel]( metadata: Arbitrary metadata surfaced in the pending signal message and checkpoint. Useful for UI rendering (e.g. which tool needs approval, what arguments it received). + tool_call_id: The tool call this hook suspends, if any. Stamped + onto the emitted :class:`~ai.messages.HookPart` so consumers + (run-blocked tracking, UIs) can attribute the suspension to + its tool call. Approval gating sets this automatically; + custom gating wrappers should pass it too. """ call = middleware_.HookContext( label=_label(hook), payload=payload, metadata=metadata or {}, + tool_call_id=tool_call_id, ) chain = middleware_._build_hook_chain(_hook_impl) @@ -136,6 +143,7 @@ async def _hook_impl(call: middleware_.HookContext) -> pydantic.BaseModel: hook_type=payload.__name__, status="pending", metadata=hook_metadata, + tool_call_id=call.tool_call_id, ) await rt.put_hook(hook_part) @@ -154,6 +162,7 @@ async def _hook_impl(call: middleware_.HookContext) -> pydantic.BaseModel: status="resolved", metadata=hook_metadata, resolution=resolution, + tool_call_id=call.tool_call_id, ) ) diff --git a/src/ai/agents/ui/ai_sdk/approvals.py b/src/ai/agents/ui/ai_sdk/approvals.py index 9c931eeb..4384f479 100644 --- a/src/ai/agents/ui/ai_sdk/approvals.py +++ b/src/ai/agents/ui/ai_sdk/approvals.py @@ -8,9 +8,6 @@ from ...hooks import TOOL_APPROVAL_HOOK_TYPE, resolve_hook from . import ui_messages -_PREFIX = "approve_" - - ToolPart = ui_messages.UIToolPart | ui_messages.UIDynamicToolPart @@ -24,12 +21,10 @@ class ApprovalResponse(NamedTuple): def tool_call_id_for(hook_part: messages_.HookPart[Any]) -> str | None: - """Return the tool_call_id encoded in a ToolApproval hook id, or None.""" + """Return the tool call a ToolApproval hook suspends, or None.""" if hook_part.hook_type != TOOL_APPROVAL_HOOK_TYPE: return None - if hook_part.hook_id.startswith(_PREFIX): - return hook_part.hook_id[len(_PREFIX) :] - return None + return hook_part.tool_call_id def hook_part_from_tool_part(tp: ToolPart) -> messages_.HookPart[Any] | None: @@ -52,6 +47,7 @@ def hook_part_from_tool_part(tp: ToolPart) -> messages_.HookPart[Any] | None: hook_type=TOOL_APPROVAL_HOOK_TYPE, status="pending", metadata=metadata, + tool_call_id=tp.tool_call_id, ) if tp.state == "approval-responded" and approval.approved is not None: @@ -64,6 +60,7 @@ def hook_part_from_tool_part(tp: ToolPart) -> messages_.HookPart[Any] | None: "granted": approval.approved, "reason": approval.reason, }, + tool_call_id=tp.tool_call_id, ) if tp.state == "output-denied": @@ -76,6 +73,7 @@ def hook_part_from_tool_part(tp: ToolPart) -> messages_.HookPart[Any] | None: "granted": False, "reason": approval.reason, }, + tool_call_id=tp.tool_call_id, ) return None diff --git a/src/ai/agents/ui/ai_sdk/outbound_stream.py b/src/ai/agents/ui/ai_sdk/outbound_stream.py index e2183681..191261a1 100644 --- a/src/ai/agents/ui/ai_sdk/outbound_stream.py +++ b/src/ai/agents/ui/ai_sdk/outbound_stream.py @@ -607,6 +607,10 @@ async def to_stream( case events_.HookEvent(): for ui_event in state.on_hook(event): yield ui_event + case events_.RunBlocked(): + # No AI SDK UI equivalent; hook parts already carry the + # pending-approval state the UI renders. + pass case _: for ui_event in state.on_event(event): yield ui_event diff --git a/src/ai/types/events.py b/src/ai/types/events.py index ea6d1be0..9231cd49 100644 --- a/src/ai/types/events.py +++ b/src/ai/types/events.py @@ -332,8 +332,120 @@ class HookEvent(pydantic.BaseModel): kind: Literal["hook"] = "hook" +class RunBlocked(pydantic.BaseModel): + """The run is blocked on hooks. + + Emitted when the run stops being able to make progress without + external input: at least one hook is pending, no model stream is + producing events, and every in-flight tool call is suspended + awaiting a hook. Streaming consumers can use this to surface + "waiting for approval" state without reconstructing it from + tool/hook events. + + ``hooks`` is a snapshot of the pending hooks the run is blocked on. + + There is no mirror "unblocked" event because it would be redundant: + a blocked run can only resume via a hook resolution (or + cancellation), so the next ``HookEvent`` with a non-``pending`` + status *is* the unblock signal. Note the converse does not hold — + a ``ToolCallResult`` carrying an ``is_hook_pending`` placeholder + (serverless abort) arrives while the run stays blocked, and the run + then ends still blocked. + """ + + hooks: tuple[messages.HookPart[Any], ...] = () + + kind: Literal["run_blocked"] = "run_blocked" + + AgentMessageEvent = Event | ToolCallResult | HookEvent -AgentEvent = Event | ToolCallResult | HookEvent | PartialToolCallResult +AgentEvent = ( + Event | ToolCallResult | HookEvent | PartialToolCallResult | RunBlocked +) TerminalEvent = StreamEnd | ToolCallResult | HookEvent + + +class RunStateTracker: + """Fold an agent event stream into run state (blocked-on-hooks). + + A pure function of the event stream: feed every event in order and + :meth:`feed` returns a :class:`RunBlocked` event whenever the run + becomes blocked, else None (:attr:`blocked` flips back silently — + see :class:`RunBlocked` for why no mirror event exists). Works + identically over a live run or a serialized replay of one. + + The fold reads three things: + + * hook state from :class:`HookEvent` (``pending`` adds, ``resolved`` + / ``cancelled`` removes); + * model-stream activity from :class:`StreamStart` / :class:`StreamEnd`; + * in-flight tool calls from the assistant message on + :class:`StreamEnd` (scheduled) and :class:`ToolCallResult` + (settled), matched by ``tool_call_id``. + + The run is blocked when at least one hook is pending, no stream is + producing, and every in-flight tool call is accounted for by a + pending hook's ``tool_call_id``. Consequently the signal is only + as good as the stream: loops must yield their ``StreamEnd`` (with + the assistant message) for tool calls to be counted, and custom + gating must pass ``tool_call_id=`` to ``ai.hook()`` — an + unattributed hook while tools are in flight reads as "still busy" + and suppresses the signal. + """ + + def __init__(self) -> None: + self._pending: dict[str, messages.HookPart[Any]] = {} + self._in_flight: set[str] = set() + self._streaming = 0 + self._blocked = False + + @property + def blocked(self) -> bool: + return self._blocked + + @property + def pending_hooks(self) -> list[messages.HookPart[Any]]: + return list(self._pending.values()) + + def feed(self, event: AgentEvent) -> RunBlocked | None: + match event: + case StreamStart(): + self._streaming += 1 + case StreamEnd(): + # Loops may emit a bare StreamEnd without a StreamStart + # (e.g. when the model was streamed out-of-band), so + # clamp at zero. + self._streaming = max(0, self._streaming - 1) + self._in_flight.update( + tc.tool_call_id for tc in event.message.tool_calls + ) + case ToolCallResult(): + self._in_flight.difference_update( + r.tool_call_id for r in event.results + ) + case HookEvent(): + if event.hook.status == "pending": + self._pending[event.hook.hook_id] = event.hook + else: + self._pending.pop(event.hook.hook_id, None) + case _: + return None + + attributed = { + h.tool_call_id + for h in self._pending.values() + if h.tool_call_id is not None + } + now = ( + bool(self._pending) + and not self._streaming + and self._in_flight <= attributed + ) + if now == self._blocked: + return None + self._blocked = now + if not now: + return None + return RunBlocked(hooks=tuple(self._pending.values())) diff --git a/src/ai/types/messages.py b/src/ai/types/messages.py index 5c5fe553..61a22c6f 100644 --- a/src/ai/types/messages.py +++ b/src/ai/types/messages.py @@ -419,6 +419,13 @@ class HookPart[T](pydantic.BaseModel): status: Literal["pending", "resolved", "cancelled"] metadata: dict[str, Any] = pydantic.Field(default_factory=dict) resolution: T | None = None + tool_call_id: str | None = None + """The tool call this hook suspends, if any. + + Set via ``ai.hook(..., tool_call_id=...)``; approval gating sets it + automatically. Lets consumers attribute a hook to its tool call + (e.g. run-blocked tracking, UI rendering) without parsing labels. + """ kind: Literal["hook"] = "hook" model_config = pydantic.ConfigDict(frozen=True) diff --git a/tests/agents/test_run_blocked.py b/tests/agents/test_run_blocked.py new file mode 100644 index 00000000..732748aa --- /dev/null +++ b/tests/agents/test_run_blocked.py @@ -0,0 +1,357 @@ +"""RunBlocked: the run-is-blocked-on-hooks signal. + +A run is blocked when at least one hook is pending, no model stream +is producing events, and every in-flight tool call is suspended +awaiting a hook. ``RunStateTracker`` folds the event stream and a +``RunBlocked`` event is emitted when the run blocks; there is no +mirror event — a blocked run can only resume via a hook resolution, +so the non-``pending`` ``HookEvent`` is the unblock signal. +``AgentStream`` folds both into its ``blocked`` / ``pending_hooks`` +properties. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncGenerator +from typing import Any + +import pydantic + +import ai +from ai.types import events as events_ +from ai.types import messages as messages_ + +from ..conftest import MOCK_MODEL, mock_llm, text_msg, tool_call_msg + + +class Confirmation(pydantic.BaseModel): + approved: bool + + +@ai.tool(require_approval=True) +async def gated(x: int) -> str: + """A tool that requires approval.""" + return f"gated ran with {x}" + + +def _approve(hook: messages_.HookPart[Any]) -> None: + ai.resolve_hook( + hook.hook_id, ai.tools.ToolApproval(granted=True, reason="ok") + ) + + +def _multi_call_msg(*tc: tuple[str, str]) -> messages_.Message: + """Assistant message with several (tool_call_id, tool_name) calls.""" + return messages_.Message( + id="msg-1", + role="assistant", + parts=[ + messages_.ToolCallPart( + tool_call_id=tc_id, tool_name=name, tool_args='{"x": 1}' + ) + for tc_id, name in tc + ], + ) + + +async def test_gated_tool_block_cycle() -> None: + """Park on approval -> RunBlocked; the resolved HookEvent unblocks.""" + my_agent = ai.agent(tools=[gated]) + mock_llm( + [ + [tool_call_msg(name="gated", args='{"x": 1}')], + [text_msg("done", id="msg-2")], + ] + ) + + delivered: list[events_.AgentEvent] = [] + + async with my_agent.run(MOCK_MODEL, [ai.user_message("go")]) as stream: + async for event in stream: + delivered.append(event) + if isinstance(event, events_.RunBlocked): + assert stream.blocked + assert [h.hook_id for h in stream.pending_hooks] == [ + h.hook_id for h in event.hooks + ] + assert event.hooks[0].tool_call_id == "tc-1" + _approve(event.hooks[0]) + elif ( + isinstance(event, events_.HookEvent) + and event.hook.status == "resolved" + ): + # The resolution is the unblock signal; the stream's + # fold has already applied it when the event arrives. + assert not stream.blocked + + assert not stream.blocked + assert stream.pending_hooks == [] + + # Ordering: the run blocks after the hook parks (and only once + # the model stream stops producing). + def index(pred: Any) -> int: + [idx] = [i for i, e in enumerate(delivered) if pred(e)] + return idx + + pending_idx = index( + lambda e: isinstance(e, events_.HookEvent) + and e.hook.status == "pending" + ) + resolved_idx = index( + lambda e: isinstance(e, events_.HookEvent) + and e.hook.status == "resolved" + ) + blocked_idx = index(lambda e: isinstance(e, events_.RunBlocked)) + assert pending_idx < blocked_idx < resolved_idx + + assert stream.output == "done" + + +async def test_busy_tool_defers_block_signal() -> None: + """No block signal while an unblocked tool is still running.""" + release = asyncio.Event() + + @ai.tool + async def slow(x: int) -> str: + """A slow tool.""" + await release.wait() + return "slow done" + + my_agent = ai.agent(tools=[gated, slow]) + mock_llm( + [ + [_multi_call_msg(("tc-1", "gated"), ("tc-2", "slow"))], + [text_msg("done", id="msg-2")], + ] + ) + + delivered: list[events_.AgentEvent] = [] + async with my_agent.run(MOCK_MODEL, [ai.user_message("go")]) as stream: + async for event in stream: + delivered.append(event) + if ( + isinstance(event, events_.HookEvent) + and event.hook.status == "pending" + ): + assert not stream.blocked + release.set() + if isinstance(event, events_.RunBlocked): + _approve(event.hooks[0]) + + [slow_idx] = [ + i + for i, e in enumerate(delivered) + if isinstance(e, events_.ToolCallResult) + and e.results[0].tool_call_id == "tc-2" + ] + [blocked_idx] = [ + i for i, e in enumerate(delivered) if isinstance(e, events_.RunBlocked) + ] + assert blocked_idx > slow_idx + + +async def test_parallel_gated_tools() -> None: + """Blocking needs *every* in-flight tool parked, and the signal + re-fires when the run parks again on the remaining hook.""" + my_agent = ai.agent(tools=[gated]) + mock_llm( + [ + [_multi_call_msg(("tc-1", "gated"), ("tc-2", "gated"))], + [text_msg("done", id="msg-2")], + ] + ) + + hook_counts: list[int] = [] + + async with my_agent.run(MOCK_MODEL, [ai.user_message("go")]) as stream: + async for event in stream: + if isinstance(event, events_.RunBlocked): + hook_counts.append(len(event.hooks)) + _approve(event.hooks[0]) + + # Blocks on both hooks, resumes, then re-blocks on the remaining one. + assert hook_counts == [2, 1] + assert not stream.blocked + + +async def test_loop_level_hook() -> None: + """A hook awaited directly by the loop (no tool task) blocks too.""" + + class MyAgent(ai.Agent): + async def loop( + self, context: ai.Context + ) -> AsyncGenerator[events_.AgentEvent]: + async with ai.models.stream(context=context) as stream: + async for event in stream: + yield event + await ai.hook("confirm", payload=Confirmation) + + mock_llm([[text_msg("OK")]]) + + blocked_events: list[events_.RunBlocked] = [] + async with MyAgent().run(MOCK_MODEL, [ai.user_message("go")]) as stream: + async for event in stream: + if isinstance(event, events_.RunBlocked): + blocked_events.append(event) + assert event.hooks[0].hook_id == "confirm" + ai.resolve_hook("confirm", {"approved": True}) + + assert len(blocked_events) == 1 + assert not stream.blocked + + +async def test_abort_leaves_blocked() -> None: + """Serverless abort: the run ends still blocked, hooks still pending.""" + my_agent = ai.agent(tools=[gated]) + mock_llm([[tool_call_msg(name="gated", args='{"x": 1}')]]) + + blocked_events: list[events_.RunBlocked] = [] + async with my_agent.run(MOCK_MODEL, [ai.user_message("go")]) as stream: + async for event in stream: + if isinstance(event, events_.RunBlocked): + blocked_events.append(event) + if ( + isinstance(event, events_.HookEvent) + and event.hook.status == "pending" + ): + ai.abort_pending_hook(event.hook) + + # The abort settles the tool with an is_hook_pending placeholder, + # but no hook resolution ever arrives: the run ends still blocked. + assert len(blocked_events) == 1 + assert stream.blocked + assert len(stream.pending_hooks) == 1 + assert stream.messages[-1].tool_results[0].is_hook_pending + + +async def test_no_hooks_no_block_events() -> None: + @ai.tool + async def plain(x: int) -> str: + """A plain tool.""" + return "plain done" + + my_agent = ai.agent(tools=[plain]) + mock_llm( + [ + [tool_call_msg(name="plain", args='{"x": 1}')], + [text_msg("done", id="msg-2")], + ] + ) + + async with my_agent.run(MOCK_MODEL, [ai.user_message("go")]) as stream: + async for event in stream: + assert not isinstance(event, events_.RunBlocked) + + assert not stream.blocked + assert stream.pending_hooks == [] + + +async def test_block_signal_waits_for_stream_end() -> None: + """A park while the model is still streaming is not a blocked run: + the signal is deferred until the stream stops producing.""" + msg = messages_.Message( + id="msg-1", + role="assistant", + parts=[ + messages_.ToolCallPart( + tool_call_id="tc-1", tool_name="gated", tool_args='{"x": 1}' + ), + messages_.TextPart(text="still streaming..."), + ], + ) + my_agent = ai.agent(tools=[gated]) + mock_llm([[msg], [text_msg("done", id="msg-2")]]) + + delivered: list[events_.AgentEvent] = [] + async with my_agent.run(MOCK_MODEL, [ai.user_message("go")]) as stream: + async for event in stream: + delivered.append(event) + if isinstance(event, events_.RunBlocked): + _approve(event.hooks[0]) + + blocked = [ + i for i, e in enumerate(delivered) if isinstance(e, events_.RunBlocked) + ] + stream_ends = [ + i for i, e in enumerate(delivered) if isinstance(e, events_.StreamEnd) + ] + assert len(blocked) == 1 + assert blocked[0] > stream_ends[0] + + +async def test_unattributed_hook_in_tool_fails_closed() -> None: + """A hook awaited inside a tool without ``tool_call_id=`` cannot be + attributed, so the in-flight tool reads as busy and no block signal + is emitted (fail closed, never a false "waiting for you").""" + + @ai.tool + async def self_gated(x: int) -> str: + """Parks on its own unattributed hook.""" + await ai.hook("self_gate", payload=Confirmation) + return "ran" + + my_agent = ai.agent(tools=[self_gated]) + mock_llm( + [ + [tool_call_msg(name="self_gated", args='{"x": 1}')], + [text_msg("done", id="msg-2")], + ] + ) + + async with my_agent.run(MOCK_MODEL, [ai.user_message("go")]) as stream: + async for event in stream: + assert not isinstance(event, events_.RunBlocked) + if ( + isinstance(event, events_.HookEvent) + and event.hook.status == "pending" + ): + assert event.hook.tool_call_id is None + ai.resolve_hook("self_gate", {"approved": True}) + + assert not stream.blocked + + +def _hook_event(hook: messages_.HookPart[Any]) -> events_.HookEvent: + return events_.HookEvent( + message=messages_.Message(role="internal", parts=[hook]), hook=hook + ) + + +def test_tracker_fold_sequence() -> None: + """RunStateTracker is a pure fold usable over any event stream.""" + tracker = events_.RunStateTracker() + two_calls = _multi_call_msg(("tc-1", "gated"), ("tc-2", "slow")) + hook: messages_.HookPart[Any] = messages_.HookPart( + hook_id="h1", + hook_type="ToolApproval", + status="pending", + tool_call_id="tc-1", + ) + + assert tracker.feed(events_.StreamStart()) is None + # Parks mid-stream: suppressed while the model is producing. + assert tracker.feed(_hook_event(hook)) is None + # Stream over, but tc-2 is in flight and not blocked. + assert tracker.feed(events_.StreamEnd(message=two_calls)) is None + assert not tracker.blocked + + result = ai.tool_result(tool_call_id="tc-2", tool_name="slow", result="ok") + transition = tracker.feed(result) + assert isinstance(transition, events_.RunBlocked) + assert [h.hook_id for h in transition.hooks] == ["h1"] + assert tracker.blocked + assert [h.hook_id for h in tracker.pending_hooks] == ["h1"] + + # The resolution flips the tracker back silently -- no mirror event. + resolved = tracker.feed( + _hook_event(hook.model_copy(update={"status": "resolved"})) + ) + assert resolved is None + assert not tracker.blocked + + +def test_message_aggregator_tolerates_block_events() -> None: + agg = ai.agents.MessageAggregator() + agg.feed(events_.RunBlocked()) + assert agg.snapshot().messages == () diff --git a/tests/agents/ui/ai_sdk/test_approvals.py b/tests/agents/ui/ai_sdk/test_approvals.py index 447ef9db..ec3d00f4 100644 --- a/tests/agents/ui/ai_sdk/test_approvals.py +++ b/tests/agents/ui/ai_sdk/test_approvals.py @@ -27,20 +27,33 @@ def _tool( } -def test_tool_call_id_for_strips_prefix() -> None: +def test_tool_call_id_for_returns_field() -> None: + # Custom gating can use any label as long as tool_call_id is set. hook: messages_.HookPart[Any] = messages_.HookPart( - hook_id="approve_tc_42", + hook_id="my_custom_gate", hook_type="ToolApproval", status="pending", + tool_call_id="tc_42", ) assert approvals.tool_call_id_for(hook) == "tc_42" +def test_tool_call_id_for_field_ignored_on_non_approval() -> None: + hook: messages_.HookPart[Any] = messages_.HookPart( + hook_id="confirm_something", + hook_type="Confirmation", + status="pending", + tool_call_id="tc_42", + ) + assert approvals.tool_call_id_for(hook) is None + + def test_tool_call_id_for_rejects_non_approval_type() -> None: hook: messages_.HookPart[Any] = messages_.HookPart( hook_id="approve_tc_42", hook_type="SomethingElse", status="pending", + tool_call_id="tc_42", ) assert approvals.tool_call_id_for(hook) is None @@ -97,12 +110,3 @@ def test_extract_approvals_handles_dynamic_tool_responses() -> None: assert approval_responses[0].granted is True assert approval_responses[0].reason == "ok" assert approval_responses[0].tool_call_id == "tc1" - - -def test_tool_call_id_for_rejects_bad_prefix() -> None: - hook: messages_.HookPart[Any] = messages_.HookPart( - hook_id="tc_42", - hook_type="ToolApproval", - status="pending", - ) - assert approvals.tool_call_id_for(hook) is None diff --git a/tests/agents/ui/ai_sdk/test_outbound_messages.py b/tests/agents/ui/ai_sdk/test_outbound_messages.py index ad503e58..ccc05c35 100644 --- a/tests/agents/ui/ai_sdk/test_outbound_messages.py +++ b/tests/agents/ui/ai_sdk/test_outbound_messages.py @@ -151,6 +151,7 @@ def test_merge_approval_signals_pending_then_resolved() -> None: hook_id="approve_tc1", hook_type="ToolApproval", status="pending", + tool_call_id="tc1", ) ], ) @@ -167,6 +168,7 @@ def test_merge_approval_signals_pending_then_resolved() -> None: hook_type="ToolApproval", status="resolved", resolution={"granted": True, "reason": None}, + tool_call_id="tc1", ) ], ) @@ -359,6 +361,7 @@ def test_to_ui_messages_internal_role_merges_approval() -> None: hook_id="approve_tc1", hook_type="ToolApproval", status="pending", + tool_call_id="tc1", ) ], ), @@ -600,3 +603,66 @@ def test_duplicate_tool_copies_do_not_reach_model_integrity() -> None: next_request_history, _ = ai_sdk.to_messages(reloaded_ui) integrity.prepare_messages(next_request_history) + + +def _parked_turn(*, hook_id: str, tool_call_id: str) -> list[messages_.Message]: + """History as a run parked on a gated tool records it.""" + return [ + messages_.Message( + id="assistant-1", + turn_id="turn-1", + role="assistant", + parts=[ + messages_.ToolCallPart( + id="call-1", + tool_call_id=tool_call_id, + tool_name="bash", + tool_args='{"command": "ls"}', + ) + ], + ), + messages_.Message( + id="hook-1", + turn_id="turn-1", + role="internal", + parts=[ + messages_.HookPart( + id="part-1", + hook_id=hook_id, + hook_type="ToolApproval", + status="pending", + metadata={"tool": "bash", "kwargs": {"command": "ls"}}, + tool_call_id=tool_call_id, + ) + ], + ), + ] + + +def test_pending_approval_hook_roundtrips_tool_call_id() -> None: + # Both the conventional approve_ label and a custom label must + # survive: the approval rides the UI tool part via the hook's + # tool_call_id field, and inbound reconstruction restores it. + for hook_id in ("approve_tc-1", "my_custom_gate"): + ui = ai_sdk.to_ui_messages( + _parked_turn(hook_id=hook_id, tool_call_id="tc-1") + ) + [tool_part] = [ + p + for m in ui + for p in m.parts + if isinstance(p, UIToolPart | UIDynamicToolPart) + ] + assert tool_part.state == "approval-requested" + + back, approval_responses = ai_sdk.to_messages(ui) + assert approval_responses == [] + [hook] = [ + p + for m in back + for p in m.parts + if isinstance(p, messages_.HookPart) + ] + assert hook.hook_id == hook_id + assert hook.status == "pending" + assert hook.tool_call_id == "tc-1" diff --git a/tests/agents/ui/ai_sdk/test_outbound_stream.py b/tests/agents/ui/ai_sdk/test_outbound_stream.py index 3907e006..ea40d303 100644 --- a/tests/agents/ui/ai_sdk/test_outbound_stream.py +++ b/tests/agents/ui/ai_sdk/test_outbound_stream.py @@ -356,6 +356,7 @@ async def test_approval_request_hook_emits_approval_event() -> None: hook_id="approve_tc1", hook_type="ToolApproval", status="pending", + tool_call_id="tc1", ) ], ), @@ -363,6 +364,7 @@ async def test_approval_request_hook_emits_approval_event() -> None: hook_id="approve_tc1", hook_type="ToolApproval", status="pending", + tool_call_id="tc1", ), ), ] @@ -544,6 +546,7 @@ async def test_resolved_approval_hook_emits_response_event() -> None: "callProviderMetadata": {"provider": {"approval": True}}, }, resolution={"granted": False, "reason": "not allowed"}, + tool_call_id="tc1", ) out = await _collect(