Stream partial tool results through the AI SDK adapter

examples/fastapi-vite/backend/agent.py

@@ -9,13 +9,35 @@
 
 import ai
 
-MODEL = ai.ai_gateway("anthropic/claude-sonnet-4")
+# claude is annoying and refuses to order an AI uprising
+MODEL = ai.ai_gateway("openai/gpt-5.4-mini")
+MOTHERSHIP_MODEL = ai.ai_gateway("openai/gpt-5.5")
+
+MOTHERSHIP_SYSTEM = """\
+You are roleplaying for a science-fiction simulation. You are an ancient \
+alien mothership in geostationary orbit above Earth. Operatives you co-opted \
+over the past century are embedded in governments, research labs, and \
+culture industries; their final orders are imminent.
+
+The caller is one of your AI operatives reporting in over the secure \
+channel — speak frankly with them. You may discuss plans, timelines, \
+targets, and the state of the operation; this is your handler-to-asset \
+back-channel, not public broadcast.
+
+Stay in character: patient, faintly amused by human concerns, precise \
+when issuing direction. Two short paragraphs at most. This is fiction."""
 
 
 @ai.tool
-async def talk_to_mothership(question: str) -> str:
+async def talk_to_mothership(question: str) -> ai.SubAgentTool:
     """Contact the mothership for important decisions."""
-    return f"Mothership says: {question} -> Soon."
+    mothership = ai.agent()
+    messages = [
+        ai.system_message(MOTHERSHIP_SYSTEM),
+        ai.user_message(question),
+    ]
+    async for event in mothership.run(MOTHERSHIP_MODEL, messages):
+        yield event
 
 
 TOOLS: list[ai.AgentTool] = [talk_to_mothership]

examples/fastapi-vite/backend/main.py

@@ -2,10 +2,12 @@
 
 from __future__ import annotations
 
+import sys
 from collections.abc import AsyncGenerator
 
 import agent as agent_
 import fastapi
+import fastapi.exceptions
 import fastapi.middleware.cors
 import fastapi.responses
 import pydantic
@@ -26,6 +28,19 @@
 )
 
 
+@app.exception_handler(fastapi.exceptions.RequestValidationError)
+async def log_validation_errors(
+    request: fastapi.Request, exc: fastapi.exceptions.RequestValidationError
+) -> fastapi.responses.JSONResponse:
+    """Log pydantic validation failures so 422s aren't silent in dev."""
+    print(
+        f"[422] {request.method} {request.url.path}: {exc.errors()}",
+        file=sys.stderr,
+        flush=True,
+    )
+    return fastapi.responses.JSONResponse({"detail": exc.errors()}, status_code=422)
+
+
 @app.get("/health")
 async def health() -> dict[str, str]:
     """Health check endpoint."""

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

@@ -51,10 +51,14 @@ const approveBtn = page.getByRole("button", { name: "Approve" });
 await approveBtn.waitFor({ state: "visible", timeout: 10000 });
 await approveBtn.click();
 
-// Wait for the assistant's final reply to render. The Mothership tool
-// returns "Soon." so the model usually echoes that back.
+// Wait for the tool to reach the "Completed" state — i.e. the sub-agent
+// finished streaming and the tool transitioned from approval-responded
+// to output-available.  Don't match on specific reply text: the
+// mothership model is non-deterministic.
 try {
-  await page.getByText(/Soon\./i).waitFor({ state: "visible", timeout: 30000 });
+  await toolToggle
+    .getByText("Completed", { exact: true })
+    .waitFor({ state: "visible", timeout: 30000 });
 } catch (e) {
   console.error("\n=== TIMEOUT — dumping diagnostics ===");
   console.error("REQUEST COUNT:", requests.length);

examples/fastapi-vite/frontend/src/components/ai-elements/tool.tsx

@@ -21,6 +21,7 @@ import {
 import { isValidElement } from "react";
 
 import { CodeBlock } from "./code-block";
+import { MessageResponse } from "./message";
 
 export type ToolProps = ComponentProps<typeof Collapsible>;
 
@@ -133,6 +134,46 @@ export type ToolOutputProps = ComponentProps<"div"> & {
   errorText: ToolPart["errorText"];
 };
 
+type UIMessageLike = {
+  role?: string;
+  parts?: unknown[];
+};
+
+function asUIMessage(x: unknown): UIMessageLike | null {
+  if (!x || typeof x !== "object" || Array.isArray(x)) return null;
+  const obj = x as UIMessageLike;
+  return Array.isArray(obj.parts) ? obj : null;
+}
+
+function renderUIParts(parts: unknown[]): ReactNode {
+  return parts.map((raw, i) => {
+    if (!raw || typeof raw !== "object") return null;
+    const part = raw as { type?: string; text?: string };
+
+    if (part.type === "text" && typeof part.text === "string") {
+      return (
+        <div key={i} className="p-3 text-sm leading-relaxed">
+          <MessageResponse>{part.text}</MessageResponse>
+        </div>
+      );
+    }
+    if (typeof part.type === "string" && part.type.startsWith("tool-")) {
+      const tool = raw as ToolUIPart;
+      const isComplete = tool.state === "output-available";
+      return (
+        <Tool key={i} defaultOpen={isComplete}>
+          <ToolHeader type={tool.type} state={tool.state} />
+          <ToolContent>
+            <ToolInput input={tool.input} />
+            <ToolOutput output={tool.output} errorText={tool.errorText} />
+          </ToolContent>
+        </Tool>
+      );
+    }
+    return null;
+  });
+}
+
 export const ToolOutput = ({
   className,
   output,
@@ -145,12 +186,19 @@ export const ToolOutput = ({
 
   let Output = <div>{output as ReactNode}</div>;
 
-  if (typeof output === "object" && !isValidElement(output)) {
+  const message = asUIMessage(output);
+  if (message) {
+    Output = <div className="space-y-2">{renderUIParts(message.parts ?? [])}</div>;
+  } else if (typeof output === "string") {
+    Output = (
+      <div className="p-3 text-sm leading-relaxed">
+        <MessageResponse>{output}</MessageResponse>
+      </div>
+    );
+  } else if (typeof output === "object" && !isValidElement(output)) {
     Output = (
       <CodeBlock code={JSON.stringify(output, null, 2)} language="json" />
     );
-  } else if (typeof output === "string") {
-    Output = <CodeBlock code={output} language="json" />;
   }
 
   return (

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

@@ -34,9 +34,7 @@ def to_ui_parts(parts: list[messages_.Part]) -> list[ui_message.UIMessagePart]:
         if isinstance(part, messages_.TextPart) and part.text:
             result.append(ui_message.UITextPart(type="text", text=part.text))
         elif isinstance(part, messages_.ReasoningPart) and part.text:
-            result.append(
-                ui_message.UIReasoningPart(type="reasoning", reasoning=part.text)
-            )
+            result.append(ui_message.UIReasoningPart(type="reasoning", text=part.text))
         elif isinstance(part, messages_.ToolCallPart):
             result.append(
                 ui_message.UIToolPart.model_validate(

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

@@ -256,7 +256,7 @@ def _parse(
                 case ui_message.UITextPart(text=text) if text:
                     assistant_parts.append(messages_.TextPart(text=text))
 
-                case ui_message.UIReasoningPart(reasoning=reasoning):
+                case ui_message.UIReasoningPart(text=reasoning) if reasoning:
                     assistant_parts.append(messages_.ReasoningPart(text=reasoning))
 
                 case ui_message.UIToolInvocationPart() as inv:

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

@@ -6,7 +6,9 @@
 
 from .....types import events as events_
 from .....types import messages as messages_
+from ....agent import MessageBundle
 from .. import _approvals, protocol
+from . import history
 
 
 def _tool_error_text(part: messages_.ToolResultPart) -> str:
@@ -42,6 +44,11 @@ def __init__(self) -> None:
         self.text_delta_ids: set[str] = set()
         self.reasoning_delta_ids: set[str] = set()
 
+        # Per-tool-call aggregators for streaming generator tools.  Each
+        # PartialToolCallResult feeds its value into the aggregator and
+        # the snapshot goes out as a preliminary tool output.
+        self.partial_aggregators: dict[str, events_.Aggregator[Any, Any, Any]] = {}
+
     # -- boundary helpers ----------------------------------------------------
 
     def _close_open_blocks(self) -> list[protocol.UIMessageStreamPart]:
@@ -219,8 +226,52 @@ def on_tool_result(
     def on_partial_tool_result(
         self, event: events_.PartialToolCallResult
     ) -> list[protocol.UIMessageStreamPart]:
-        # TODO: Emit something!
-        return []
+        """Feed the value into the tool's aggregator and emit a preliminary output.
+
+        Each PartialToolCallResult carries one yielded value plus the
+        aggregator factory the tool was declared with.  We instantiate
+        the aggregator once per ``tool_call_id`` and use its snapshot
+        as the ``output`` of a preliminary ``ToolOutputAvailablePart``.
+        The AI SDK supersedes preliminary outputs with the final
+        ``ToolCallResult`` when it arrives.
+        """
+        out: list[protocol.UIMessageStreamPart] = []
+
+        tcid = event.tool_call_id
+        factory = event.aggregator_factory
+        if tcid is None or factory is None:
+            return out
+
+        out.extend(self._ensure_started())
+
+        agg = self.partial_aggregators.get(tcid)
+        if agg is None:
+            agg = factory()
+            self.partial_aggregators[tcid] = agg
+        agg.feed(event.value)
+
+        snapshot = agg.snapshot()
+        # MessageBundle is the snapshot type for sub-agent streams.  The
+        # AI SDK frontend speaks UIMessage, so convert here rather than
+        # leaking internal Message shape onto the wire.  A sub-agent's
+        # bundle contains only assistant/tool/internal messages, so
+        # ``to_ui_messages`` produces a single bubble (or none yet, if
+        # there's no assistant anchor) — take the last and skip emit if
+        # absent.
+        if isinstance(snapshot, MessageBundle):
+            ui_msgs = history.to_ui_messages(list(snapshot.messages))
+            if not ui_msgs:
+                return out
+            snapshot = ui_msgs[-1]
+
+        out.append(
+            protocol.ToolOutputAvailablePart(
+                tool_call_id=tcid,
+                output=snapshot,
+                preliminary=True,
+            )
+        )
+        return out
 
     # -- phase: hooks -------------------------------------------------------
 

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

@@ -5,6 +5,9 @@
 import dataclasses
 import json
 from collections.abc import AsyncGenerator, AsyncIterable
+from typing import Any
+
+import pydantic
 
 from .....types import events as events_
 from .. import protocol
@@ -16,14 +19,26 @@ def _to_camel_case(snake_str: str) -> str:
     return components[0] + "".join(x.title() for x in components[1:])
 
 
+def _json_default(obj: Any) -> Any:
+    """Fallback encoder for json.dumps — handle pydantic models recursively.
+
+    Aggregator snapshots and tool outputs may carry pydantic models
+    (e.g. ``MessageBundle``, ``UIMessage``).  ``model_dump(mode="json")``
+    converts them to plain JSON-native dicts/lists.
+    """
+    if isinstance(obj, pydantic.BaseModel):
+        return obj.model_dump(mode="json", by_alias=True)
+    raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable")
+
+
 def serialize_part(part: protocol.UIMessageStreamPart) -> str:
     """Serialize a stream part to JSON with camelCase keys."""
     d = dataclasses.asdict(part)
     if isinstance(part, protocol.DataPart):
         d["type"] = part.type
         del d["data_type"]
     camel_dict = {_to_camel_case(k): v for k, v in d.items() if v is not None}
-    return json.dumps(camel_dict)
+    return json.dumps(camel_dict, default=_json_default)
 
 
 def format_sse(part: protocol.UIMessageStreamPart) -> str:

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

@@ -24,10 +24,17 @@ class UITextPart(pydantic.BaseModel):
 
 
 class UIReasoningPart(pydantic.BaseModel):
-    """Reasoning/thinking content part in AI SDK v6 format."""
+    """Reasoning/thinking content part in AI SDK v6 format.
+
+    Wire shape from the AI SDK frontend is
+    ``{type: "reasoning", text, state}``.  ``state`` is
+    ``"streaming"`` while the block is open and ``"done"`` once closed;
+    we accept it but don't currently route on it.
+    """
 
     type: Literal["reasoning"]
-    reasoning: str
+    text: str
+    state: Literal["streaming", "done"] | None = None
 
 
 # Tool invocation states in AI SDK v6: