Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions src/ai/agents/_middleware.py
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
59 changes: 52 additions & 7 deletions src/ai/agents/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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()
Expand All @@ -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
Expand Down Expand Up @@ -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).
Expand Down
9 changes: 9 additions & 0 deletions src/ai/agents/hooks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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)
Expand Down Expand Up @@ -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)
Expand All @@ -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,
)
)

Expand Down
12 changes: 5 additions & 7 deletions src/ai/agents/ui/ai_sdk/approvals.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand All @@ -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:
Expand All @@ -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:
Expand All @@ -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":
Expand All @@ -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
Expand Down
4 changes: 4 additions & 0 deletions src/ai/agents/ui/ai_sdk/outbound_stream.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
114 changes: 113 additions & 1 deletion src/ai/types/events.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()))
7 changes: 7 additions & 0 deletions src/ai/types/messages.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Loading
Loading