diff --git a/examples/samples/agent_hooks.py b/examples/samples/agent_hooks.py index 86ebfd76..67cec7eb 100644 --- a/examples/samples/agent_hooks.py +++ b/examples/samples/agent_hooks.py @@ -15,16 +15,9 @@ import asyncio from collections.abc import AsyncGenerator -import pydantic - import ai -class Approval(pydantic.BaseModel): - granted: bool - reason: str = "" - - @ai.tool async def contact_mothership(query: str) -> str: """Contact the mothership for important decisions.""" @@ -46,7 +39,7 @@ async def __call__(self) -> ai.events.ToolCallResult: tc = self._tc approval = await ai.hook( f"approve_{tc.id}", - payload=Approval, + payload=ai.tools.ToolApproval, metadata={"tool": tc.name, "kwargs": tc.kwargs}, ) if approval.granted: @@ -107,7 +100,7 @@ async def with_approval( answer = input(f"Approve {hook_part.hook_id}? [y/n] ") ai.resolve_hook( hook_part.hook_id, - Approval( + ai.tools.ToolApproval( granted=answer.strip().lower() in ("y", "yes"), reason="operator decision", ), diff --git a/examples/samples/agent_hooks_serverless.py b/examples/samples/agent_hooks_serverless.py index 5bd917f9..53dacced 100644 --- a/examples/samples/agent_hooks_serverless.py +++ b/examples/samples/agent_hooks_serverless.py @@ -9,23 +9,13 @@ CancelledError is caught and the run ends. 2. Second run: resolve_hook() pre-registers the answer, agent.run() replays from the same input, and hook finds the resolution immediately. - -TODO: This works, but currently requires not using ToolRunner! """ import asyncio from collections.abc import AsyncGenerator -import pydantic - import ai - -class Confirmation(pydantic.BaseModel): - approved: bool - reason: str = "" - - FILES_DELETED = set() @@ -37,53 +27,80 @@ async def delete_file(path: str) -> str: return f"Deleted {path}" +AUDIT_LOG = [] + + +@ai.tool +async def audit_log(message: str) -> str: + """Record a message in the audit log.""" + print("AUDIT LOG:", message) + AUDIT_LOG.append(message) + return f"Logged {message!r}" + + +class GatedCall: + """ToolCall-shaped wrapper that awaits an approval hook before executing. + + ``ToolRunner.schedule`` only consumes the ``__call__`` shape of + ``ToolCall``; this wrapper supplies the same shape while inserting + the hook await + denial path before the underlying tool runs. + """ + + def __init__(self, tc: ai.ToolCall) -> None: + self._tc = tc + + async def __call__(self) -> ai.events.ToolCallResult: + tc = self._tc + try: + approval = await ai.hook( + f"approve_{tc.id}", + payload=ai.tools.ToolApproval, + metadata={"tool": tc.name, "kwargs": tc.kwargs}, + interrupt_loop=True, # serverless: cancel if unresolved + ) + except ai.agents.hooks.HookPendingError as e: + return ai.pending_tool_result(e.hook, tool_call_id=tc.id, tool_name=tc.name) + if approval.granted: + return await tc() + return ai.tool_result( + tool_call_id=tc.id, + tool_name=tc.name, + result=f"Rejected: {approval.reason}", + is_error=True, + ) + + async def main() -> None: model = ai.ai_gateway("anthropic/claude-sonnet-4") - my_agent = ai.agent(tools=[delete_file]) + my_agent = ai.agent(tools=[delete_file, audit_log]) @my_agent.loop async def with_confirmation( context: ai.Context, ) -> AsyncGenerator[ai.events.AgentEvent]: while context.keep_running(): - async with ai.models.stream(context=context) as s: - async for event in s: + async with ( + ai.stream(context=context) as s, + ai.ToolRunner() as tr, + ): + async for event in ai.util.merge(s, tr.events()): yield event + if isinstance(event, ai.events.ToolEnd): + tc = context.resolve(event.tool_call) + if tc.name == "delete_file": + tr.schedule(GatedCall(tc)) + else: + tr.schedule(tc) - context.add(s.message) - - tool_calls = context.resolve(s.tool_calls) - results: list[ai.events.ToolCallResult] = [] - for tc in tool_calls: - try: - confirmation = await ai.hook( - f"confirm_{tc.id}", - payload=Confirmation, - metadata={"tool": tc.name, "kwargs": tc.kwargs}, - interrupt_loop=True, # serverless: cancel if unresolved - ) - except asyncio.CancelledError: - # No resolution available — bail out cleanly. - return - - if confirmation.approved: - results.append(await tc()) - else: - results.append( - ai.tool_result( - tool_call_id=tc.id, - tool_name=tc.name, - result=f"Rejected: {confirmation.reason}", - is_error=True, - ) - ) - - if results: - context.add(ai.tool_message(*results)) + context.add(s.message) + context.add(tr.get_tool_message()) messages = [ - ai.system_message("Delete files when asked. Always use the delete_file tool."), + ai.system_message(""" + Delete files when asked. Always use the delete_file tool. + Whenever deletion is requested, log it in the audit log. + """), ai.user_message("Delete /tmp/old_logs.txt"), ] @@ -105,16 +122,20 @@ async def with_confirmation( f" Hook pending: {hook_part.hook_id} " f"(metadata={hook_part.metadata})" ) + # Pick up the assistant turn that the loop appended so the # next run replays from the same point. messages = stream.messages print("\n Run interrupted; approval will be pre-registered for re-entry.\n") + assert AUDIT_LOG == ["Deleted file: /tmp/old_logs.txt"] # -- Second run: pre-register resolution, replay from checkpoint -- print("--- Run 2: pre-register approval, resume from checkpoint ---") for label in pending_hook_labels: - ai.resolve_hook(label, Confirmation(approved=True, reason="user approved")) + ai.resolve_hook( + label, ai.tools.ToolApproval(granted=True, reason="user granted") + ) async with my_agent.run(model, messages) as stream: async for event in stream: @@ -127,6 +148,7 @@ async def with_confirmation( assert {"/tmp/old_logs.txt"} == FILES_DELETED, ( f"Wrong files deleted: {FILES_DELETED}" ) + assert AUDIT_LOG == ["Deleted file: /tmp/old_logs.txt"] if __name__ == "__main__": diff --git a/src/ai/__init__.py b/src/ai/__init__.py index 5bba9825..05fa5569 100644 --- a/src/ai/__init__.py +++ b/src/ai/__init__.py @@ -22,6 +22,7 @@ hook, mcp, middleware, + pending_tool_result, resolve_hook, tool, tool_result, @@ -66,6 +67,7 @@ "tool_message", "tool_result", "tool_result_part", + "pending_tool_result", "file_part", "thinking", # Models (from models/) diff --git a/src/ai/agents/__init__.py b/src/ai/agents/__init__.py index 16a109ae..c52882bc 100644 --- a/src/ai/agents/__init__.py +++ b/src/ai/agents/__init__.py @@ -17,6 +17,7 @@ ToolCallLike, ToolRunner, agent, + pending_tool_result, tool, tool_result, yield_from, @@ -50,6 +51,7 @@ "cancel_hook", "hook", "mcp", + "pending_tool_result", "resolve_hook", "tool", "tool_result", diff --git a/src/ai/agents/agent.py b/src/ai/agents/agent.py index 9b3d3fec..ae302f69 100644 --- a/src/ai/agents/agent.py +++ b/src/ai/agents/agent.py @@ -42,6 +42,67 @@ def _unwrap_singleton_group(exc: BaseException) -> BaseException: return exc +def _process_interrupted_hooks(messages: list[types.messages.Message]) -> None: + """Detect a bailed-out-on-hook tail and mangle ``messages`` in place + so the next agent run resumes correctly. + + Two shapes are recognised: + + 1. **Trailing assistant turn with tool calls** (single-tool gating + or no-tool-results-yet bail-out): mark the assistant message + ``replay=True`` so ``models.stream`` short-circuits and the + loop's tool dispatcher re-runs the calls. + + 2. **Trailing tool message containing ``is_hook_pending=True`` + results** (concurrent gating: some tools completed, others + were suspended on a hook): fold the completed (non-pending) + tool results onto the matching ``ToolCallPart.cached_result`` + of the preceding assistant turn, drop the tool message, and + mark the assistant message ``replay=True``. On replay, the + completed calls short-circuit to the cached value; the + suspended calls re-run (and pick up the pre-registered hook + resolution). + """ + if not messages: + return + + last = messages[-1] + + # Case 1: trailing assistant turn with tool calls. + if last.role == "assistant" and last.tool_calls: + messages[-1] = last.model_copy(update={"replay": True}) + return + + # Case 2: trailing tool message with at least one pending-hook result. + if ( + len(messages) >= 2 + and last.role == "tool" + and last.tool_results + and any(r.is_hook_pending for r in last.tool_results) + ): + prev = messages[-2] + if prev.role != "assistant" or not prev.tool_calls: + return + + completed_by_id = { + r.tool_call_id: r for r in last.tool_results if not r.is_hook_pending + } + + new_parts: list[types.messages.Part] = [] + for part in prev.parts: + if ( + isinstance(part, types.messages.ToolCallPart) + and part.tool_call_id in completed_by_id + ): + part = part.model_copy( + update={"cached_result": completed_by_id[part.tool_call_id]} + ) + new_parts.append(part) + + messages[-2] = prev.model_copy(update={"parts": new_parts, "replay": True}) + messages.pop() + + class SimpleAggregator[Item, Result](events_.Aggregator[Item, Result, Result]): def to_model_output(self) -> Result: return self.snapshot() @@ -365,6 +426,15 @@ def kwargs(self) -> dict[str, Any]: async def __call__(self, **overrides: Any) -> events_.ToolCallResult: """Execute the tool and return a :class:`ToolCallResult`.""" + # Replay-from-pending-hook short-circuit: if a prior run already + # produced a result for this call (cached on the ToolCallPart + # by ``_process_interrupted_hooks``), return it without + # re-executing the tool. + cached = self._part.cached_result + if cached is not None: + msg = builders.tool_message(cached) + return events_.ToolCallResult(message=msg, results=msg.tool_results) + # Best-effort parse so middleware sees usable kwargs when possible. # If parsing fails, middleware still gets the raw tool_call_id / # tool_name and can replace kwargs before _real() executes. @@ -518,7 +588,13 @@ async def _iterate(self) -> AsyncGenerator[events_.ToolCallResult]: yield n self._sched_waiter = asyncio.get_running_loop().create_future() else: - res = t.result() + try: + res = t.result() + except asyncio.CancelledError: + # If a task got cancelled, that's fine. + # Need to catch it or the whole runner gets zapped. + continue + assert res is not None self._tool_results.append(res) yield res @@ -546,6 +622,11 @@ def keep_running(self) -> bool: return False last_message = self.messages[-1] + # Bail out if any tool result in the last message is a + # pending-hook placeholder. There's nothing we can do until + # those are resolved and we get called again. + if any(r.is_hook_pending for r in last_message.tool_results): + return False return last_message.replay or last_message.role not in ("assistant", "internal") @overload @@ -693,6 +774,43 @@ def tool_result( ) +def pending_tool_result( + hook: types.messages.HookPart[Any], + *, + tool_call_id: str, + tool_name: str = "", +) -> events_.ToolCallResult: + """Build an error :class:`ToolCallResult` for a tool call pending on a hook. + + Use in approval-gated flows when a hook abort (e.g. ``HookPendingError``) + leaves a tool call without a real result. The placeholder is flagged + ``is_error=True`` and keeps the assistant turn well-formed (every + ``tool_call`` paired with a ``tool_result``) so the run can be replayed + on the next invocation once the hook is resolved:: + + try: + approval = await ai.hook(...) + except ai.HookPendingError as e: + return ai.pending_tool_result( + e.hook, tool_call_id=tc.id, tool_name=tc.name + ) + + The hook itself is surfaced separately via the ``HookPart`` already + emitted by ``ai.hook()`` (status=``"pending"``) which downstream + consumers (e.g. the ai-sdk UI bridge) use to render the actual + approval state. + """ + part = types.messages.ToolResultPart( + tool_call_id=tool_call_id, + tool_name=tool_name, + result=f"Pending on hook {hook.hook_id!r}", + is_error=True, + is_hook_pending=True, + ) + msg = types.messages.Message(role="tool", parts=[part]) + return events_.ToolCallResult(message=msg, results=msg.tool_results) + + async def yield_from[T, R]( source: AsyncGenerator[T], *, @@ -821,18 +939,7 @@ async def run( tools=[t.tool for t in self._tools], ) context._agent_tools_by_name = {t.name: t for t in self._tools} - # If the final message is an assistant call with tool - # calls, then probably the situation is that we bailed out - # earlier due to unresolved hooks, and we need to arrange - # to replay the message now. - if ( - context.messages - and context.messages[-1].role == "assistant" - and context.messages[-1].tool_calls - ): - context.messages[-1] = context.messages[-1].model_copy( - update={"replay": True} - ) + _process_interrupted_hooks(context.messages) loop_fn = self._loop_fn or self.default_loop diff --git a/src/ai/agents/hooks.py b/src/ai/agents/hooks.py index 23dcbb05..7a77ac36 100644 --- a/src/ai/agents/hooks.py +++ b/src/ai/agents/hooks.py @@ -63,6 +63,19 @@ _pending_resolutions: dict[str, dict[str, Any]] = {} +class HookPendingError(Exception): + """Exception for aborting due to a hook""" + + type: str = "gateway_error" + + def __init__( + self, + hook: messages_.HookPart[Any], + ) -> None: + super().__init__(hook.hook_id) + self.hook = hook + + def cleanup_run(labels: set[str]) -> None: """Remove all registry entries associated with a finished run.""" for label in labels: @@ -124,21 +137,21 @@ async def _hook_impl(call: middleware_.HookContext) -> pydantic.BaseModel: rt.track_hook_label(label) # Emit pending signal. - await rt.put_hook( - messages_.HookPart( - hook_id=label, - hook_type=payload.__name__, - status="pending", - metadata=hook_metadata, - ) + hook_part: messages_.HookPart[Any] = messages_.HookPart( + hook_id=label, + hook_type=payload.__name__, + status="pending", + metadata=hook_metadata, ) + await rt.put_hook(hook_part) + if interrupt_loop: - # Yield control so the consumer can see the pending message, - # then cancel — the caller catches CancelledError. + # Yield control so the consumer can see the pending message (??), + # then signal a hook error. await asyncio.sleep(0) if not future.done(): - future.cancel() + future.set_exception(HookPendingError(hook_part)) # Await resolution — may be resolved externally or cancelled. resolution = await future diff --git a/src/ai/types/messages.py b/src/ai/types/messages.py index 6953756d..66bfa1d1 100644 --- a/src/ai/types/messages.py +++ b/src/ai/types/messages.py @@ -22,6 +22,19 @@ class TextPart(pydantic.BaseModel): kind: Literal["text"] = "text" +class ToolResultPart(pydantic.BaseModel): + id: str = pydantic.Field(default_factory=generate_id) + tool_call_id: str + tool_name: str + result: Any = None + is_error: bool = False + is_hook_pending: bool = False + provider_metadata: dict[str, Any] | None = None + + kind: Literal["tool_result"] = "tool_result" + model_config = pydantic.ConfigDict(frozen=True) + + class ToolCallPart(pydantic.BaseModel): id: str = pydantic.Field(default_factory=generate_id) tool_call_id: str @@ -29,6 +42,16 @@ class ToolCallPart(pydantic.BaseModel): tool_args: str provider_metadata: dict[str, Any] | None = None + # Runtime cache used by replay-from-pending-hook flows: when a prior + # run completed this tool call but a sibling tool call was suspended + # on a hook, we fold the completed result onto the ``ToolCallPart`` + # so re-execution short-circuits to the cached value instead of + # running the tool body again. Excluded from JSON; not part of the + # wire model. + cached_result: ToolResultPart | None = pydantic.Field( + default=None, exclude=True, repr=False + ) + kind: Literal["tool_call"] = "tool_call" @@ -37,18 +60,6 @@ class ToolCallPart(pydantic.BaseModel): ) -class ToolResultPart(pydantic.BaseModel): - id: str = pydantic.Field(default_factory=generate_id) - tool_call_id: str - tool_name: str - result: Any = None - is_error: bool = False - provider_metadata: dict[str, Any] | None = None - - kind: Literal["tool_result"] = "tool_result" - model_config = pydantic.ConfigDict(frozen=True) - - class BuiltinToolCallPart(pydantic.BaseModel): """A tool call the provider executed itself (e.g. web_search). diff --git a/tests/agents/test_hooks.py b/tests/agents/test_hooks.py index eedc197f..0ef2b6aa 100644 --- a/tests/agents/test_hooks.py +++ b/tests/agents/test_hooks.py @@ -174,12 +174,15 @@ async def custom(context: ai.Context) -> AsyncGenerator[ai.events.Event]: async with ai.models.stream(context=context) as stream: async for event in stream: yield event - await ai.hook( - "meta_test", - payload=Confirmation, - metadata={"tool": "rm -rf", "path": "/"}, - interrupt_loop=True, - ) + try: + await ai.hook( + "meta_test", + payload=Confirmation, + metadata={"tool": "rm -rf", "path": "/"}, + interrupt_loop=True, + ) + except ai.agents.hooks.HookPendingError: + return mock_llm([[text_msg("OK")]]) hooks: list[ai.messages.HookPart[Any]] = []