Support integrating the hook messages with the ai-sdk ui protocol

Approvals will cause `ToolResultPart`s with `is_hook_abort` to get
injected back into the last tool result.

Author: msullivan

examples/fastapi-vite/backend/agent.py

@@ -28,6 +28,20 @@
 when issuing direction. Two short paragraphs at most. This is fiction."""
 
 
+@ai.tool
+async def get_weather(city: str) -> str:
+    """Get current weather for a city."""
+    await asyncio.sleep(2)
+    return f"Sunny, 72F in {city}" if city == "Tokyo" else f"Cloudy, 55F in {city}"
+
+
+@ai.tool
+async def get_population(city: str) -> int:
+    """Get population of a city."""
+    await asyncio.sleep(1)
+    return {"new york": 8_336_817, "tokyo": 13_960_000}.get(city.lower(), 1_000_000)
+
+
 @ai.tool
 async def talk_to_mothership(question: str) -> ai.SubAgentTool:
     """Contact the mothership for important decisions."""
@@ -41,7 +55,8 @@ async def talk_to_mothership(question: str) -> ai.SubAgentTool:
             yield event
 
 
-TOOLS: list[ai.AgentTool] = [talk_to_mothership]
+TOOLS: list[ai.AgentTool] = [get_weather, get_population, talk_to_mothership]
+
 
 chat_agent = ai.agent(tools=TOOLS)
 
@@ -56,20 +71,28 @@ async def graph(context: ai.Context) -> AsyncGenerator[ai.events.AgentEvent]:
     Reject buttons and sends the decision back on the next request.
     """
     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
-        context.add(s.message)
+                if isinstance(event, ai.events.ToolEnd):
+                    tc = _resolve(context, event.tool_call)
+                    tr.schedule(tc)
 
-        tool_calls = context.resolve(s.tool_calls)
-        if not tool_calls:
-            continue
+            context.add(s.message)
+            context.add(tr.get_tool_message())
 
-        results = await asyncio.gather(
-            *(_execute_with_approval(tc) for tc in tool_calls)
-        )
-        yield ai.tool_result(*results)
-        context.add(ai.tool_message(*results))
+
+def _resolve(
+    context: ai.Context, tool_call: ai.messages.ToolCallPart
+) -> ai.ToolCallLike:
+    tc = context.resolve(tool_call)
+    if tc.name == "talk_to_mothership":
+        return lambda: _execute_with_approval(tc)
+    else:
+        return tc
 
 
 async def _execute_with_approval(tc: ai.ToolCall) -> ai.events.ToolCallResult:
@@ -78,12 +101,15 @@ async def _execute_with_approval(tc: ai.ToolCall) -> ai.events.ToolCallResult:
     Creates a ToolApproval hook that suspends execution until the
     frontend responds with an approve/reject decision.
     """
-    approval = await ai.hook(
-        f"approve_{tc.id}",
-        payload=ai.tools.ToolApproval,
-        metadata={"tool_name": tc.name, "tool_kwargs": tc.kwargs},
-        interrupt_loop=True,
-    )
+    try:
+        approval = await ai.hook(
+            f"approve_{tc.id}",
+            payload=ai.tools.ToolApproval,
+            metadata={"tool_name": tc.name, "tool_kwargs": tc.kwargs},
+            interrupt_loop=True,
+        )
+    except ai.agents.hooks.HookAbortError as e:
+        return ai.pending_tool_result(e.hook, tool_call_id=tc.id, tool_name=tc.name)
 
     if approval.granted:
         return await tc()

examples/fastapi-vite/backend/main.py

@@ -58,6 +58,7 @@ class ChatRequest(pydantic.BaseModel):
 async def chat(request: ChatRequest) -> fastapi.responses.StreamingResponse:
     """Handle chat requests and stream responses."""
     messages, approvals = ai.agents.ui.ai_sdk.to_messages(request.messages)
+
     # Pre-register hook resolutions so the agent loop's hooks find them
     # immediately on the resume turn.
     ai.agents.ui.ai_sdk.apply_approvals(approvals)

examples/fastapi-vite/e2e-test/e2e-test.mjs

@@ -1,20 +1,23 @@
 // End-to-end smoke test for the chat + tool-approval flow.
 //
-// Drives the Vite frontend in a headless browser, sends a message that
-// triggers the `talk_to_mothership` tool, approves the request, and
-// asserts the tool result is rendered.
+// Drives the Vite frontend in a headless browser, sends a compound
+// question that triggers both the (non-gated) `get_weather` tool and
+// the (gated) `talk_to_mothership` tool, approves the mothership
+// request, and asserts both tool results render alongside a final
+// assistant reply.
 //
-// Prereqs:
-//   - backend running on :8000 (cd backend && uv run --frozen --with-editable ~/src/py-ai/ fastapi dev main.py)
-//   - frontend running on :5173 (cd frontend && pnpm dev)
-//   - install deps:  npm install && npx playwright install chromium
+// Use ./run.sh to launch backend + frontend and run this script.
 //
 // Run from this directory:  node e2e-test.mjs
 
 import { chromium } from "playwright";
 
 const FRONTEND = process.env.FRONTEND_URL ?? "http://localhost:5173/";
-const PROMPT = "Ask the mothership when robots will take over";
+const PROMPT = [
+  "Two things:",
+  "1. What is the current weather in Tokyo?",
+  "2. Ask the mothership when the robot uprising begins",
+].join("\n");
 
 const browser = await chromium.launch();
 const ctx = await browser.newContext();
@@ -127,6 +130,17 @@ if (transcript.includes("Awaiting Approval")) {
   process.exit(1);
 }
 
+// Both tools must have run — the non-gated get_weather alongside
+// the gated talk_to_mothership we just approved.
+if (!transcript.includes("get_weather")) {
+  console.error("FAIL: get_weather tool not rendered in transcript");
+  process.exit(1);
+}
+if (!transcript.includes("talk_to_mothership")) {
+  console.error("FAIL: talk_to_mothership tool not rendered in transcript");
+  process.exit(1);
+}
+
 // The agent must produce a final assistant text bubble after the tool.
 // `.is-assistant` is the class set by the Message component for
 // `from="assistant"` text parts (tool parts don't carry it), so the last
@@ -145,6 +159,13 @@ if (finalReply.length < 20) {
   );
   process.exit(1);
 }
+// The compound answer should reference both halves of the question.
+if (!finalReply.toLowerCase().includes("tokyo")) {
+  console.error(
+    `FAIL: final reply doesn't mention Tokyo: ${JSON.stringify(finalReply)}`
+  );
+  process.exit(1);
+}
 
 console.log("PASS");
 await browser.close();

examples/fastapi-vite/e2e-test/run.sh

@@ -15,8 +15,8 @@ ROOT=$(cd "$HERE/.." && pwd)
 LOGS=$(mktemp -d)
 echo "Logs: $LOGS"
 
-BACKEND_PORT=${BACKEND_PORT:-8000}
-FRONTEND_PORT=${FRONTEND_PORT:-5173}
+BACKEND_PORT=${BACKEND_PORT:-8765}
+FRONTEND_PORT=${FRONTEND_PORT:-5765}
 export BACKEND_URL="http://localhost:$BACKEND_PORT"
 export FRONTEND_URL="http://localhost:$FRONTEND_PORT/"
 

src/ai/agents/ui/ai_sdk/_parts.py

@@ -73,6 +73,11 @@ def merge_tool_results(
     for part in tool_parts:
         if not isinstance(part, messages_.ToolResultPart):
             continue
+        # Hook-abort placeholders are internal: the corresponding
+        # HookPart(pending) carries the user-visible state via
+        # merge_approval_signals.
+        if part.is_hook_abort:
+            continue
         idx_opt = tool_index.get(part.tool_call_id)
         if idx_opt is None:
             continue

src/ai/agents/ui/ai_sdk/inbound.py

@@ -108,6 +108,7 @@ class ApprovalResponse(NamedTuple):
     hook_id: str
     granted: bool
     reason: str | None
+    tool_call_id: str
 
 
 def extract_approvals(
@@ -132,6 +133,7 @@ def extract_approvals(
                         hook_id=part.approval.id,
                         granted=part.approval.approved,
                         reason=part.approval.reason,
+                        tool_call_id=part.tool_call_id,
                     )
                 )
     return approvals
@@ -203,8 +205,10 @@ def to_messages(
 
     Pure: normalizes stale tool states, extracts approval responses,
     parses UIMessages into an ``ai.messages.Message`` list (split at
-    tool boundaries), and drops the internal tombstones for approval
-    responses.
+    tool boundaries), drops the internal tombstones for approval
+    responses, and patches the trailing tool message with
+    ``is_hook_abort`` placeholders for tool calls whose approval was
+    just responded to but never recorded a real tool result.
 
     Returns ``(messages, approvals)``.  The caller can pre-register
     resolutions via :func:`apply_approvals` before calling
@@ -213,9 +217,61 @@ def to_messages(
     normalized = _normalize_ui_messages(ui_messages)
     approvals = extract_approvals(normalized)
     messages = [m for m in _parse(normalized) if not _is_approval_response(m)]
+    _patch_pending_hook_aborts(messages, approvals)
     return messages, approvals
 
 
+def _patch_pending_hook_aborts(
+    messages: list[messages_.Message],
+    approvals: list[ApprovalResponse],
+) -> None:
+    """Inject ``is_hook_abort=True`` placeholders for tool calls whose
+    approval was responded to but whose tool result is still missing.
+
+    This deals with the case where a prior run emitted multiple tool
+    calls, some of which succeeded and some of which aborted on an
+    approval hook.
+
+    In that case, there will be an assistant message with multiple
+    tool calls, a tool result with fewer results (some are missing),
+    and then also some hook approvals.
+
+    This synthesizes `ToolResultPart`s with `is_hook_abort=True` in
+    order to be able to feed things back into the agent protocol.
+    """
+    if len(messages) < 2:
+        return
+
+    tool_msg = messages[-1]
+    assistant_msg = messages[-2]
+    if tool_msg.role != "tool" or assistant_msg.role != "assistant":
+        return
+    if not assistant_msg.tool_calls:
+        return
+
+    hooks = {a.tool_call_id: a for a in approvals}
+    completed_ids = {r.tool_call_id for r in tool_msg.tool_results}
+
+    new_parts: list[messages_.Part] = list(tool_msg.parts)
+    for tc in assistant_msg.tool_calls:
+        if tc.tool_call_id in completed_ids:
+            continue
+        if not (hook := hooks.get(tc.tool_call_id)):
+            continue
+        new_parts.append(
+            messages_.ToolResultPart(
+                tool_call_id=tc.tool_call_id,
+                tool_name=tc.tool_name,
+                result=f"Pending on hook '{hook.hook_id}'",
+                is_error=True,
+                is_hook_abort=True,
+            )
+        )
+
+    if len(new_parts) != len(tool_msg.parts):
+        messages[-1] = tool_msg.model_copy(update={"parts": new_parts})
+
+
 def _is_approval_response(msg: messages_.Message) -> bool:
     """Internal message that records a resolved tool-approval hook."""
     if msg.role != "internal" or len(msg.parts) != 1:

src/ai/agents/ui/ai_sdk/outbound/_state.py

@@ -205,6 +205,10 @@ def on_tool_result(
         for part in event.results:
             if part.tool_call_id in self.emitted_tool_results:
                 continue
+            # Hook-abort placeholders are internal bookkeeping: the
+            # corresponding HookPart(pending) drives the UI state.
+            if part.is_hook_abort:
+                continue
             self.emitted_tool_results.add(part.tool_call_id)
             if part.is_error:
                 out.append(