Make serverless hooks work with parallel ToolRunner

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",
                     ),

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__":

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

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",