diff --git a/examples/fastapi-vite/backend/agent.py b/examples/fastapi-vite/backend/agent.py index 1ea6f5db..1d4c87f6 100644 --- a/examples/fastapi-vite/backend/agent.py +++ b/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] diff --git a/examples/fastapi-vite/backend/main.py b/examples/fastapi-vite/backend/main.py index 32b1c607..f091ac2d 100644 --- a/examples/fastapi-vite/backend/main.py +++ b/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.""" diff --git a/examples/fastapi-vite/e2e-test/e2e-test.mjs b/examples/fastapi-vite/e2e-test/e2e-test.mjs index 768036f3..b1b2fc73 100644 --- a/examples/fastapi-vite/e2e-test/e2e-test.mjs +++ b/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); diff --git a/examples/fastapi-vite/frontend/src/components/ai-elements/tool.tsx b/examples/fastapi-vite/frontend/src/components/ai-elements/tool.tsx index de70e835..219927c6 100644 --- a/examples/fastapi-vite/frontend/src/components/ai-elements/tool.tsx +++ b/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; @@ -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 ( +
+ {part.text} +
+ ); + } + if (typeof part.type === "string" && part.type.startsWith("tool-")) { + const tool = raw as ToolUIPart; + const isComplete = tool.state === "output-available"; + return ( + + + + + + + + ); + } + return null; + }); +} + export const ToolOutput = ({ className, output, @@ -145,12 +186,19 @@ export const ToolOutput = ({ let Output =
{output as ReactNode}
; - if (typeof output === "object" && !isValidElement(output)) { + const message = asUIMessage(output); + if (message) { + Output =
{renderUIParts(message.parts ?? [])}
; + } else if (typeof output === "string") { + Output = ( +
+ {output} +
+ ); + } else if (typeof output === "object" && !isValidElement(output)) { Output = ( ); - } else if (typeof output === "string") { - Output = ; } return ( diff --git a/src/ai/agents/ui/ai_sdk/_parts.py b/src/ai/agents/ui/ai_sdk/_parts.py index 64537fdf..49ebe172 100644 --- a/src/ai/agents/ui/ai_sdk/_parts.py +++ b/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( diff --git a/src/ai/agents/ui/ai_sdk/inbound.py b/src/ai/agents/ui/ai_sdk/inbound.py index 7ed3ff0c..f57918ab 100644 --- a/src/ai/agents/ui/ai_sdk/inbound.py +++ b/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: diff --git a/src/ai/agents/ui/ai_sdk/outbound/_state.py b/src/ai/agents/ui/ai_sdk/outbound/_state.py index ed912a2c..892e912d 100644 --- a/src/ai/agents/ui/ai_sdk/outbound/_state.py +++ b/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 ------------------------------------------------------- diff --git a/src/ai/agents/ui/ai_sdk/outbound/sse.py b/src/ai/agents/ui/ai_sdk/outbound/sse.py index 4821de9f..4e784aba 100644 --- a/src/ai/agents/ui/ai_sdk/outbound/sse.py +++ b/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,6 +19,18 @@ 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) @@ -23,7 +38,7 @@ def serialize_part(part: protocol.UIMessageStreamPart) -> str: 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: diff --git a/src/ai/agents/ui/ai_sdk/ui_message.py b/src/ai/agents/ui/ai_sdk/ui_message.py index 406a699f..7463a67d 100644 --- a/src/ai/agents/ui/ai_sdk/ui_message.py +++ b/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: diff --git a/tests/agents/ui/ai_sdk/outbound/test_stream.py b/tests/agents/ui/ai_sdk/outbound/test_stream.py index 78b89154..b2d6f92b 100644 --- a/tests/agents/ui/ai_sdk/outbound/test_stream.py +++ b/tests/agents/ui/ai_sdk/outbound/test_stream.py @@ -2,6 +2,7 @@ from collections.abc import AsyncGenerator +import ai from ai.agents.ui.ai_sdk import protocol, to_stream from ai.types import events as agent_events_ from ai.types import events as events_ @@ -151,6 +152,88 @@ async def test_approval_request_hook_emits_approval_part() -> None: assert approval_parts[0].approval_id == "approve_tc1" +async def test_partial_tool_results_emit_preliminary_outputs() -> None: + """Each PartialToolCallResult feeds the aggregator and yields a preliminary part.""" + out = await _collect( + [ + agent_events_.PartialToolCallResult( + tool_call_id="tc1", + tool_name="search", + value="hit 1, ", + aggregator_factory=ai.ConcatAggregator, + ), + agent_events_.PartialToolCallResult( + tool_call_id="tc1", + tool_name="search", + value="hit 2, ", + aggregator_factory=ai.ConcatAggregator, + ), + agent_events_.PartialToolCallResult( + tool_call_id="tc1", + tool_name="search", + value="hit 3", + aggregator_factory=ai.ConcatAggregator, + ), + ] + ) + + prelim = [ + p + for p in out + if isinstance(p, protocol.ToolOutputAvailablePart) and p.preliminary + ] + assert [p.output for p in prelim] == [ + "hit 1, ", + "hit 1, hit 2, ", + "hit 1, hit 2, hit 3", + ] + assert all(p.tool_call_id == "tc1" for p in prelim) + + +async def test_partial_message_bundle_becomes_ui_message() -> None: + """MessageAggregator's MessageBundle snapshot collapses to a single UIMessage.""" + from ai.agents.ui.ai_sdk.ui_message import UIMessage + + inner_msg = messages_.Message( + role="assistant", + parts=[messages_.TextPart(text="hi from sub-agent")], + ) + + out = await _collect( + [ + agent_events_.PartialToolCallResult( + tool_call_id="tc1", + tool_name="research", + value=agent_events_.ToolCallResult(message=inner_msg, results=[]), + aggregator_factory=ai.MessageAggregator, + ), + ] + ) + + [prelim] = [ + p + for p in out + if isinstance(p, protocol.ToolOutputAvailablePart) and p.preliminary + ] + assert isinstance(prelim.output, UIMessage) + assert prelim.output.role == "assistant" + assert prelim.output.parts[0].type == "text" + + +async def test_partial_tool_result_without_factory_is_skipped() -> None: + """Without an aggregator_factory there's nothing to snapshot.""" + out = await _collect( + [ + agent_events_.PartialToolCallResult( + tool_call_id="tc1", + tool_name="search", + value="ignored", + ), + ] + ) + assert not any(isinstance(p, protocol.ToolOutputAvailablePart) for p in out) + + # NOTE: agent-change boundary detection used to be driven by # Message.source_label. That field has been removed; agent-change # routing in the AI SDK adapter now needs to come from diff --git a/tests/agents/ui/ai_sdk/test_parts.py b/tests/agents/ui/ai_sdk/test_parts.py index 69ca1566..9876e9c8 100644 --- a/tests/agents/ui/ai_sdk/test_parts.py +++ b/tests/agents/ui/ai_sdk/test_parts.py @@ -17,7 +17,7 @@ def test_to_ui_parts_text_and_reasoning() -> None: ] ui_parts = _parts.to_ui_parts(parts) assert isinstance(ui_parts[0], UIReasoningPart) - assert ui_parts[0].reasoning == "thinking" + assert ui_parts[0].text == "thinking" assert isinstance(ui_parts[1], UITextPart) assert ui_parts[1].text == "hi"