From 2858593dbdf35875ebad83cc623238b4f9c9b53e Mon Sep 17 00:00:00 2001 From: "Michael J. Sullivan" Date: Mon, 11 May 2026 23:18:25 -0700 Subject: [PATCH 1/4] Tweak streaming tool data representation a bit more For streaming tools, make the `result` be the final aggregation snapshot, and then have a separate `model_result`. This allows the `result` of a subagent to be the full `MessageBundle`, which has the advantage that the `preliminary` type when streaming over the ai-sdk ui matches the final type. UIMessage <-> MessageBundle coversions are done on the ai-sdk ui side, and we now pass tools in to enable that. --- examples/fastapi-vite/backend/main.py | 4 +- examples/fastapi-vite/frontend/src/App.tsx | 20 +++- .../src/components/ai-elements/tool.tsx | 80 ++++++++-------- src/ai/agents/agent.py | 55 +++++++++-- src/ai/agents/ui/ai_sdk/inbound.py | 93 ++++++++++++++++++- src/ai/agents/ui/ai_sdk/outbound/_state.py | 46 ++++++--- src/ai/models/ai_gateway/adapter.py | 6 +- src/ai/models/anthropic/adapter.py | 4 +- src/ai/models/openai/adapter.py | 4 +- src/ai/types/events.py | 18 +++- src/ai/types/messages.py | 13 +++ tests/agents/test_generator_tools.py | 9 +- tests/agents/ui/ai_sdk/test_inbound.py | 62 +++++++++++++ 13 files changed, 333 insertions(+), 81 deletions(-) diff --git a/examples/fastapi-vite/backend/main.py b/examples/fastapi-vite/backend/main.py index 881f6f03..9a1932ca 100644 --- a/examples/fastapi-vite/backend/main.py +++ b/examples/fastapi-vite/backend/main.py @@ -57,7 +57,9 @@ class ChatRequest(pydantic.BaseModel): @app.post("/chat") 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) + messages, approvals = ai.agents.ui.ai_sdk.to_messages( + request.messages, tools=agent_.chat_agent.tools + ) # Pre-register hook resolutions so the agent loop's hooks find them # immediately on the resume turn. diff --git a/examples/fastapi-vite/frontend/src/App.tsx b/examples/fastapi-vite/frontend/src/App.tsx index 0fedfd98..bfd42d59 100644 --- a/examples/fastapi-vite/frontend/src/App.tsx +++ b/examples/fastapi-vite/frontend/src/App.tsx @@ -3,7 +3,7 @@ import { DefaultChatTransport, lastAssistantMessageIsCompleteWithApprovalResponses, } from "ai"; -import type { ToolUIPart } from "ai"; +import type { ToolUIPart, UIMessage } from "ai"; import { CheckIcon, XIcon } from "lucide-react"; import { Fragment } from "react"; @@ -33,6 +33,7 @@ import { PromptInputSubmit, } from "@/components/ai-elements/prompt-input"; import { + renderUIPart, Tool, ToolHeader, ToolContent, @@ -140,10 +141,19 @@ export default function App() { - + {toolPart.type === "tool-talk_to_mothership" && + toolPart.state === "output-available" ? ( +
+ {(toolPart.output as UIMessage).parts.map( + (p, i) => renderUIPart(p, i), + )} +
+ ) : ( + + )} ); 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 219927c6..e3ae1d3d 100644 --- a/examples/fastapi-vite/frontend/src/components/ai-elements/tool.tsx +++ b/examples/fastapi-vite/frontend/src/components/ai-elements/tool.tsx @@ -1,6 +1,12 @@ "use client"; -import type { DynamicToolUIPart, ToolUIPart } from "ai"; +import { + isTextUIPart, + isToolUIPart, + type DynamicToolUIPart, + type ToolUIPart, + type UIMessage, +} from "ai"; import type { ComponentProps, ReactNode } from "react"; import { Badge } from "@/components/ui/badge"; @@ -134,44 +140,47 @@ 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"; +export function renderUIPart( + part: UIMessage["parts"][number], + key: number, +): ReactNode { + if (isTextUIPart(part)) { + return ( +
+ {part.text} +
+ ); + } + if (isToolUIPart(part)) { + const isComplete = part.state === "output-available"; + if (part.type === "dynamic-tool") { + const dyn = part as DynamicToolUIPart; return ( - - + + - - + + ); } - return null; - }); + const tool = part as ToolUIPart; + return ( + + + + + + + + ); + } + return null; } export const ToolOutput = ({ @@ -186,10 +195,7 @@ export const ToolOutput = ({ let Output =
{output as ReactNode}
; - const message = asUIMessage(output); - if (message) { - Output =
{renderUIParts(message.parts ?? [])}
; - } else if (typeof output === "string") { + if (typeof output === "string") { Output = (
{output} diff --git a/src/ai/agents/agent.py b/src/ai/agents/agent.py index 01b2f63d..b74ae3b1 100644 --- a/src/ai/agents/agent.py +++ b/src/ai/agents/agent.py @@ -118,8 +118,9 @@ def _process_interrupted_hooks(messages: list[types.messages.Message]) -> None: class SimpleAggregator[Item, Result](events_.Aggregator[Item, Result, Result]): - def to_model_output(self) -> Result: - return self.snapshot() + @classmethod + def from_snapshot(cls, snapshot: Result) -> Result: + return snapshot class ConcatAggregator(SimpleAggregator[str, str]): @@ -167,8 +168,9 @@ def feed(self, item: events_.AgentEvent) -> None: def snapshot(self) -> MessageBundle: return MessageBundle(messages=tuple(self._messages)) - def to_model_output(self) -> str: - for m in reversed(self._messages): + @classmethod + def from_snapshot(cls, snapshot: MessageBundle) -> str: + for m in reversed(snapshot.messages): if m.role == "assistant" and m.text: return m.text return "" @@ -471,21 +473,28 @@ async def __call__(self, **overrides: Any) -> events_.ToolCallResult: tool = self._tool async def _real(call: middleware_.ToolContext) -> events_.ToolCallResult: + result: Any + model_result: Any try: kwargs = _validate_kwargs(tool, call.kwargs) if tool.is_gen: # Generator tool (e.g. agent-as-a-tool): drain the async - # generator, forward each yielded message to the runtime for - # real-time streaming, and return the final text as the result. + # generator, forward each yielded value to the runtime for + # real-time streaming, then capture both the aggregator + # snapshot (the rich shape that flows to the UI) and the + # model-facing value (what the LLM sees on its next turn). assert tool.aggregator - result = await yield_from( + agg = await _aggregate_from( tool.fn(**kwargs), tool_call_id=call.tool_call_id, tool_name=call.tool_name, aggregator=tool.aggregator, ) + result = agg.snapshot() + model_result = agg.to_model_output() else: result = await tool.fn(**kwargs) + model_result = result except Exception as exc: # A nested runtime (e.g. a sub-agent run inside this # tool) raises errors wrapped in a singleton TaskGroup @@ -506,6 +515,7 @@ async def _real(call: middleware_.ToolContext) -> events_.ToolCallResult: tool_call_id=call.tool_call_id, tool_name=call.tool_name, result=result, + model_result=model_result, ) ) @@ -842,10 +852,10 @@ def pending_tool_result( return events_.ToolCallResult(message=msg, results=msg.tool_results) -async def yield_from[T, R]( +async def yield_from[T, S, R]( source: AsyncGenerator[T], *, - aggregator: Callable[[], events_.Aggregator[T, object, R]], + aggregator: Callable[[], events_.Aggregator[T, S, R]], # TODO: is this what we really want for labelling? tool_name: str | None = None, tool_call_id: str | None = None, @@ -873,6 +883,31 @@ async def yield_from[T, R]( Returns the final message's text (empty string if no messages). """ + agg = await _aggregate_from( + source, + aggregator=aggregator, + tool_name=tool_name, + tool_call_id=tool_call_id, + label=label, + ) + return agg.to_model_output() + + +async def _aggregate_from[T, S, R]( + source: AsyncGenerator[T], + *, + aggregator: Callable[[], events_.Aggregator[T, S, R]], + tool_name: str | None = None, + tool_call_id: str | None = None, + label: object = None, +) -> events_.Aggregator[T, S, R]: + """Drain *source* into a fresh aggregator, forwarding partial events. + + Returns the live aggregator so callers can consume both the snapshot + (the rich shape stored on ``ToolResultPart.result``) and the + model-facing value (``ToolResultPart.model_result``) without + re-aggregating. + """ agg = aggregator() rt = runtime.get_runtime() @@ -888,7 +923,7 @@ async def yield_from[T, R]( aggregator_factory=aggregator, ) ) - return agg.to_model_output() + return agg class Agent: diff --git a/src/ai/agents/ui/ai_sdk/inbound.py b/src/ai/agents/ui/ai_sdk/inbound.py index 65b6ad3f..2b8e148d 100644 --- a/src/ai/agents/ui/ai_sdk/inbound.py +++ b/src/ai/agents/ui/ai_sdk/inbound.py @@ -8,9 +8,12 @@ import json import logging +from collections.abc import Sequence from typing import Any, NamedTuple +from ....types import events as events_ from ....types import messages as messages_ +from ...agent import AgentTool, MessageAggregator, MessageBundle from ...hooks import resolve_hook from . import ui_message @@ -59,6 +62,47 @@ def _error_result(error_text: str | None, output: Any) -> dict[str, Any] | None: return normalized +def _aggregator_cls( + factory: Any, +) -> type[events_.Aggregator[Any, Any, Any]] | None: + """Resolve a tool's aggregator factory to the underlying class. + + Tools may declare an aggregator as a class directly (``LastAggregator``) + or via an ``Aggregate`` marker that wraps it (``Aggregate(LastAggregator, + delim="\\n")``). This normalizes both forms. + """ + if factory is None: + return None + if isinstance(factory, type) and issubclass(factory, events_.Aggregator): + return factory + inner = getattr(factory, "_factory", None) + if isinstance(inner, type) and issubclass(inner, events_.Aggregator): + return inner + return None + + +def _decode_wire_output( + output: Any, + agg_cls: type[events_.Aggregator[Any, Any, Any]] | None, +) -> Any: + """Reconstruct the internal snapshot type from a wire tool output. + + For aggregator-backed tools the wire shape is a ``UIMessage`` (sub-agent + transcripts) or the aggregator's snapshot type directly (passthrough + aggregators). This function decodes UIMessage shapes back into a + ``MessageBundle`` so the parent agent's message history carries the + rich snapshot, mirroring what tool execution stored locally. + """ + if agg_cls is None or output is None: + return output + + if agg_cls is MessageAggregator: + ui_msg = ui_message.UIMessage.model_validate(output) + inner = list(_parse([ui_msg])) + return MessageBundle(messages=tuple(inner)) + return output + + def _approval_hook_part(tp: ui_message.UIToolPart) -> messages_.HookPart[Any] | None: """Reconstruct approval hook state from a UI tool part when possible.""" approval = tp.approval @@ -200,6 +244,8 @@ def _normalize_ui_messages( def to_messages( ui_messages: list[ui_message.UIMessage], + *, + tools: Sequence[AgentTool] | None = None, ) -> tuple[list[messages_.Message], list[ApprovalResponse]]: """Parse a UI request into runtime messages + extracted approvals. @@ -210,13 +256,21 @@ def to_messages( ``is_hook_pending`` placeholders for tool calls whose approval was just responded to but never recorded a real tool result. + ``tools`` lets the parser decode aggregator-backed tool outputs (e.g. + sub-agent UIMessages) back into their internal snapshot type and + populate ``ToolResultPart.model_result`` correctly. When omitted, + tool outputs are kept in their wire form — fine for caller code that + never feeds the messages back to a model. + Returns ``(messages, approvals)``. The caller can pre-register resolutions via :func:`apply_approvals` before calling :meth:`Agent.run` if the run should resume from a hook. """ normalized = _normalize_ui_messages(ui_messages) approvals = extract_approvals(normalized) - messages = [m for m in _parse(normalized) if not _is_approval_response(m)] + messages = [ + m for m in _parse(normalized, tools=tools) if not _is_approval_response(m) + ] _patch_pending_hook_aborts(messages, approvals) return messages, approvals @@ -286,7 +340,36 @@ def _is_approval_response(msg: messages_.Message) -> bool: def _parse( ui_messages: list[ui_message.UIMessage], + *, + tools: Sequence[AgentTool] | None = None, ) -> list[messages_.Message]: + tools_by_name = {t.name: t for t in tools or []} + + def _build_result_part( + *, + tool_call_id: str, + tool_name: str, + output: Any, + is_error: bool, + ) -> messages_.ToolResultPart: + tool = tools_by_name.get(tool_name) + agg_cls = _aggregator_cls(tool.aggregator) if tool else None + if not is_error and agg_cls is not None: + snapshot = _decode_wire_output(output, agg_cls) + return messages_.ToolResultPart( + tool_call_id=tool_call_id, + tool_name=tool_name, + result=snapshot, + model_result=agg_cls.from_snapshot(snapshot), + is_error=False, + ) + return messages_.ToolResultPart( + tool_call_id=tool_call_id, + tool_name=tool_name, + result=output if is_error else _normalize_tool_result(output), + is_error=is_error, + ) + result: list[messages_.Message] = [] for ui_msg in ui_messages: @@ -313,10 +396,10 @@ def _parse( ) if _is_tool_completed(inv.state): tool_result_parts.append( - messages_.ToolResultPart( + _build_result_part( tool_call_id=inv.tool_invocation_id, tool_name=inv.tool_name, - result=inv.result, + output=inv.result, is_error=_is_tool_error(inv.state), ) ) @@ -335,10 +418,10 @@ def _parse( if tp.state in _TOOL_RESULT_STATES: tool_result_parts.append( - messages_.ToolResultPart( + _build_result_part( tool_call_id=tp.tool_call_id, tool_name=tp.tool_name, - result=_normalize_tool_result(tp.output), + output=tp.output, is_error=False, ) ) diff --git a/src/ai/agents/ui/ai_sdk/outbound/_state.py b/src/ai/agents/ui/ai_sdk/outbound/_state.py index 4f7081b9..3302c986 100644 --- a/src/ai/agents/ui/ai_sdk/outbound/_state.py +++ b/src/ai/agents/ui/ai_sdk/outbound/_state.py @@ -23,6 +23,23 @@ def _tool_error_text(part: messages_.ToolResultPart) -> str: return "Tool execution failed" +def _to_wire_output(snapshot: Any) -> Any: + """Convert an aggregator snapshot to its UI wire representation. + + For ``MessageBundle`` (sub-agent transcripts) this produces a single + ``UIMessage`` assistant bubble — the canonical AI SDK shape. Other + snapshot types pass through unchanged. + + Returns ``None`` if the bundle has no assistant anchor yet (e.g. a + streaming sub-agent that has produced no messages); callers should + skip emitting in that case. + """ + if isinstance(snapshot, MessageBundle): + ui_msgs = history.to_ui_messages(list(snapshot.messages)) + return ui_msgs[-1] if ui_msgs else None + return snapshot + + class _StreamState: """Single-pass state across one ``to_stream()`` call.""" @@ -218,10 +235,17 @@ def on_tool_result( ) ) else: + wire_output = _to_wire_output(part.result) + if wire_output is None: + # Aggregator produced no anchor (e.g. sub-agent + # tool that yielded nothing). Skip the final + # output emit; preliminaries already covered the + # streaming view if any. + continue out.append( protocol.ToolOutputAvailablePart( tool_call_id=part.tool_call_id, - output=part.result, + output=wire_output, ) ) @@ -254,24 +278,16 @@ def on_partial_tool_result( 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] + wire_output = _to_wire_output(agg.snapshot()) + if wire_output is None: + # Sub-agent bundle without an assistant anchor yet — wait + # for more events before emitting. + return out out.append( protocol.ToolOutputAvailablePart( tool_call_id=tcid, - output=snapshot, + output=wire_output, preliminary=True, ) ) diff --git a/src/ai/models/ai_gateway/adapter.py b/src/ai/models/ai_gateway/adapter.py index 32800b52..4c27f78b 100644 --- a/src/ai/models/ai_gateway/adapter.py +++ b/src/ai/models/ai_gateway/adapter.py @@ -164,13 +164,15 @@ async def _messages_to_prompt( { "type": "error-text", "value": ( - str(part.result) if part.result is not None else "" + str(part.model_result) + if part.model_result is not None + else "" ), } if part.is_error else { "type": "json", - "value": part.result, + "value": part.model_result, } ) tool_results.append( diff --git a/src/ai/models/anthropic/adapter.py b/src/ai/models/anthropic/adapter.py index e030c6d9..46f486b0 100644 --- a/src/ai/models/anthropic/adapter.py +++ b/src/ai/models/anthropic/adapter.py @@ -260,8 +260,8 @@ async def _messages_to_anthropic( entry: dict[str, Any] = { "type": "tool_result", "tool_use_id": part.tool_call_id, - "content": str(part.result) - if part.result is not None + "content": str(part.model_result) + if part.model_result is not None else "", } if part.is_error: diff --git a/src/ai/models/openai/adapter.py b/src/ai/models/openai/adapter.py index b84fd1fc..087e9148 100644 --- a/src/ai/models/openai/adapter.py +++ b/src/ai/models/openai/adapter.py @@ -162,8 +162,8 @@ async def _messages_to_openai( { "role": "tool", "tool_call_id": part.tool_call_id, - "content": str(part.result) - if part.result is not None + "content": str(part.model_result) + if part.model_result is not None else "", } ) diff --git a/src/ai/types/events.py b/src/ai/types/events.py index e5916752..f9bec45a 100644 --- a/src/ai/types/events.py +++ b/src/ai/types/events.py @@ -271,8 +271,24 @@ def feed(self, item: Item) -> None: ... @abc.abstractmethod def snapshot(self) -> Result: ... + def to_model_output(self) -> ModelResult: + """Return the model-facing value derived from this aggregator's state. + + Default implementation defers to :meth:`from_snapshot`; subclasses + with non-trivial state may override either or both. + """ + return type(self).from_snapshot(self.snapshot()) + + @classmethod @abc.abstractmethod - def to_model_output(self) -> ModelResult: ... + def from_snapshot(cls, snapshot: Result) -> ModelResult: + """Stateless conversion: snapshot -> model-facing value. + + Called on inbound (when a tool result round-trips back from the + wire) and anywhere else a snapshot needs to be re-derived + without a live aggregator instance. + """ + ... class PartialToolCallResult(pydantic.BaseModel): diff --git a/src/ai/types/messages.py b/src/ai/types/messages.py index 5f4271ec..1d7cfc8b 100644 --- a/src/ai/types/messages.py +++ b/src/ai/types/messages.py @@ -30,9 +30,22 @@ class ToolResultPart(pydantic.BaseModel): is_hook_pending: bool = False provider_metadata: dict[str, Any] | None = None + # The value the LLM sees on its next turn. For most tools this is + # identical to ``result``; for aggregator-backed tools (sub-agents, + # streaming-text) it's a scalar derived from the rich snapshot via + # ``Aggregator.from_snapshot``. + model_result: Any = pydantic.Field(default=None, repr=False) + kind: Literal["tool_result"] = "tool_result" model_config = pydantic.ConfigDict(frozen=True) + @pydantic.model_validator(mode="before") + @classmethod + def _default_model_result(cls, data: Any) -> Any: + if isinstance(data, dict) and "model_result" not in data: + data = {**data, "model_result": data.get("result")} + return data + class ToolCallPart(pydantic.BaseModel): id: str = pydantic.Field(default_factory=generate_id) diff --git a/tests/agents/test_generator_tools.py b/tests/agents/test_generator_tools.py index f298ce6a..4caf3809 100644 --- a/tests/agents/test_generator_tools.py +++ b/tests/agents/test_generator_tools.py @@ -9,6 +9,7 @@ import ai from ai import models +from ai.agents.agent import MessageBundle from ai.types import events as agent_events_ from ai.types import events as events_ from ai.types import messages as messages_ @@ -168,7 +169,13 @@ async def test_yield_from_nested_agent() -> None: tool_results = [ e for e in all_events if isinstance(e, agent_events_.ToolCallResult) ] - assert tool_results[0].results[0].result == "Mars has two moons." + # MessageAggregator stores the rich MessageBundle as `result` and the + # extracted assistant text as `model_result` (the value the parent + # LLM sees on its next turn). + sub_part = tool_results[0].results[0] + assert isinstance(sub_part.result, MessageBundle) + assert sub_part.result.messages[0].text == "Mars has two moons." + assert sub_part.model_result == "Mars has two moons." # The outer LLM's second call (index 2) must NOT contain any inner # agent messages. It should only see: the original user message, diff --git a/tests/agents/ui/ai_sdk/test_inbound.py b/tests/agents/ui/ai_sdk/test_inbound.py index 2f4a3b03..d7b695a7 100644 --- a/tests/agents/ui/ai_sdk/test_inbound.py +++ b/tests/agents/ui/ai_sdk/test_inbound.py @@ -4,6 +4,8 @@ import pytest +import ai +from ai.agents.agent import MessageBundle from ai.agents.ui.ai_sdk import to_messages from ai.agents.ui.ai_sdk.inbound import ( _normalize_ui_messages, @@ -170,3 +172,63 @@ def test_to_messages_rejects_empty_user() -> None: ui = [UIMessage.model_validate({"id": "u1", "role": "user", "parts": []})] with pytest.raises(ValueError): to_messages(ui) + + +@ai.tool +async def _research_tool(topic: str) -> ai.SubAgentTool: + """Sub-agent tool used by the inbound round-trip test.""" + if False: + yield # pragma: no cover + _ = topic + + +def test_to_messages_decodes_subagent_tool_output() -> None: + """A sub-agent tool's wire UIMessage decodes back to MessageBundle. + + Round-trip: ``model_result`` is recomputed via the aggregator's + ``from_snapshot``, and ``result`` carries the rich MessageBundle so + a subsequent UI render gets the same shape we sent. + """ + # Wire shape: a tool-_research_tool part with output = UIMessage{parts=[text]}. + ui = [ + _ui("user", _text("research mars"), id="u1"), + _ui( + "assistant", + _tool( + "_research_tool", + "tc1", + "output-available", + input={"topic": "mars"}, + output={ + "id": "sub-1", + "role": "assistant", + "parts": [{"type": "text", "text": "Mars has two moons."}], + }, + ), + id="a1", + ), + ] + messages, _ = to_messages(ui, tools=[_research_tool]) + + # Find the tool message with the decoded result. + tool_msgs = [m for m in messages if m.role == "tool"] + assert len(tool_msgs) == 1 + result_part = tool_msgs[0].tool_results[0] + assert isinstance(result_part.result, MessageBundle) + assert result_part.model_result == "Mars has two moons." + + +def test_to_messages_without_tools_keeps_wire_shape() -> None: + """No tools arg → tool outputs stay in their wire form (unchanged behavior).""" + ui = [ + _ui("user", _text("hi"), id="u1"), + _ui( + "assistant", + _tool("ping", "tc1", "output-available", input={}, output={"pong": True}), + id="a1", + ), + ] + messages, _ = to_messages(ui) + tool_msgs = [m for m in messages if m.role == "tool"] + assert tool_msgs[0].tool_results[0].result == {"pong": True} + assert tool_msgs[0].tool_results[0].model_result == {"pong": True} From 688f4a052d14a9afdd3f0ad5af2b6e065f6e6177 Mon Sep 17 00:00:00 2001 From: "Michael J. Sullivan" Date: Tue, 12 May 2026 11:18:55 -0700 Subject: [PATCH 2/4] doc a tiny bit --- src/ai/types/messages.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src/ai/types/messages.py b/src/ai/types/messages.py index 1d7cfc8b..88a1c3c2 100644 --- a/src/ai/types/messages.py +++ b/src/ai/types/messages.py @@ -25,15 +25,17 @@ 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 + # The "real" result of the tool call + result: Any = None + # The value the LLM sees on its next turn. For most tools this is - # identical to ``result``; for aggregator-backed tools (sub-agents, - # streaming-text) it's a scalar derived from the rich snapshot via - # ``Aggregator.from_snapshot``. + # identical to ``result``; for aggregator-backed tools + # (sub-agents, streaming-text) it's derived from the aggregator's + # `to_model_output`. model_result: Any = pydantic.Field(default=None, repr=False) kind: Literal["tool_result"] = "tool_result" From 307eae20574c08110222d88e373c9f80489c8cbe Mon Sep 17 00:00:00 2001 From: "Michael J. Sullivan" Date: Tue, 12 May 2026 11:23:59 -0700 Subject: [PATCH 3/4] renames --- src/ai/agents/agent.py | 8 ++++---- src/ai/agents/ui/ai_sdk/inbound.py | 2 +- src/ai/types/events.py | 8 ++++---- src/ai/types/messages.py | 2 +- tests/agents/ui/ai_sdk/test_inbound.py | 2 +- 5 files changed, 11 insertions(+), 11 deletions(-) diff --git a/src/ai/agents/agent.py b/src/ai/agents/agent.py index b74ae3b1..e498ffdb 100644 --- a/src/ai/agents/agent.py +++ b/src/ai/agents/agent.py @@ -119,7 +119,7 @@ def _process_interrupted_hooks(messages: list[types.messages.Message]) -> None: class SimpleAggregator[Item, Result](events_.Aggregator[Item, Result, Result]): @classmethod - def from_snapshot(cls, snapshot: Result) -> Result: + def to_model_output(cls, snapshot: Result) -> Result: return snapshot @@ -169,7 +169,7 @@ def snapshot(self) -> MessageBundle: return MessageBundle(messages=tuple(self._messages)) @classmethod - def from_snapshot(cls, snapshot: MessageBundle) -> str: + def to_model_output(cls, snapshot: MessageBundle) -> str: for m in reversed(snapshot.messages): if m.role == "assistant" and m.text: return m.text @@ -491,7 +491,7 @@ async def _real(call: middleware_.ToolContext) -> events_.ToolCallResult: aggregator=tool.aggregator, ) result = agg.snapshot() - model_result = agg.to_model_output() + model_result = agg.get_model_output() else: result = await tool.fn(**kwargs) model_result = result @@ -890,7 +890,7 @@ async def yield_from[T, S, R]( tool_call_id=tool_call_id, label=label, ) - return agg.to_model_output() + return agg.get_model_output() async def _aggregate_from[T, S, R]( diff --git a/src/ai/agents/ui/ai_sdk/inbound.py b/src/ai/agents/ui/ai_sdk/inbound.py index 2b8e148d..effad7b9 100644 --- a/src/ai/agents/ui/ai_sdk/inbound.py +++ b/src/ai/agents/ui/ai_sdk/inbound.py @@ -360,7 +360,7 @@ def _build_result_part( tool_call_id=tool_call_id, tool_name=tool_name, result=snapshot, - model_result=agg_cls.from_snapshot(snapshot), + model_result=agg_cls.to_model_output(snapshot), is_error=False, ) return messages_.ToolResultPart( diff --git a/src/ai/types/events.py b/src/ai/types/events.py index f9bec45a..a2063525 100644 --- a/src/ai/types/events.py +++ b/src/ai/types/events.py @@ -271,17 +271,17 @@ def feed(self, item: Item) -> None: ... @abc.abstractmethod def snapshot(self) -> Result: ... - def to_model_output(self) -> ModelResult: + def get_model_output(self) -> ModelResult: """Return the model-facing value derived from this aggregator's state. - Default implementation defers to :meth:`from_snapshot`; subclasses + Default implementation defers to :meth:`to_model_output`; subclasses with non-trivial state may override either or both. """ - return type(self).from_snapshot(self.snapshot()) + return type(self).to_model_output(self.snapshot()) @classmethod @abc.abstractmethod - def from_snapshot(cls, snapshot: Result) -> ModelResult: + def to_model_output(cls, snapshot: Result) -> ModelResult: """Stateless conversion: snapshot -> model-facing value. Called on inbound (when a tool result round-trips back from the diff --git a/src/ai/types/messages.py b/src/ai/types/messages.py index 88a1c3c2..d08471c2 100644 --- a/src/ai/types/messages.py +++ b/src/ai/types/messages.py @@ -35,7 +35,7 @@ class ToolResultPart(pydantic.BaseModel): # The value the LLM sees on its next turn. For most tools this is # identical to ``result``; for aggregator-backed tools # (sub-agents, streaming-text) it's derived from the aggregator's - # `to_model_output`. + # `get_model_output`. model_result: Any = pydantic.Field(default=None, repr=False) kind: Literal["tool_result"] = "tool_result" diff --git a/tests/agents/ui/ai_sdk/test_inbound.py b/tests/agents/ui/ai_sdk/test_inbound.py index d7b695a7..a11fe9a5 100644 --- a/tests/agents/ui/ai_sdk/test_inbound.py +++ b/tests/agents/ui/ai_sdk/test_inbound.py @@ -186,7 +186,7 @@ def test_to_messages_decodes_subagent_tool_output() -> None: """A sub-agent tool's wire UIMessage decodes back to MessageBundle. Round-trip: ``model_result`` is recomputed via the aggregator's - ``from_snapshot``, and ``result`` carries the rich MessageBundle so + ``to_model_output``, and ``result`` carries the rich MessageBundle so a subsequent UI render gets the same shape we sent. """ # Wire shape: a tool-_research_tool part with output = UIMessage{parts=[text]}. From d3732f8fa3afba369065ce7910bd3e581e8abde9 Mon Sep 17 00:00:00 2001 From: "Michael J. Sullivan" Date: Tue, 12 May 2026 13:25:30 -0700 Subject: [PATCH 4/4] discussed reworks --- examples/fastapi-vite/backend/main.py | 4 +- src/ai/agents/agent.py | 67 +++++++++++++++---- src/ai/agents/ui/ai_sdk/inbound.py | 91 ++++++++------------------ src/ai/models/ai_gateway/adapter.py | 7 +- src/ai/models/anthropic/adapter.py | 5 +- src/ai/models/openai/adapter.py | 5 +- src/ai/types/messages.py | 37 +++++++---- tests/agents/test_generator_tools.py | 4 +- tests/agents/ui/ai_sdk/test_inbound.py | 29 ++++---- 9 files changed, 133 insertions(+), 116 deletions(-) diff --git a/examples/fastapi-vite/backend/main.py b/examples/fastapi-vite/backend/main.py index 9a1932ca..881f6f03 100644 --- a/examples/fastapi-vite/backend/main.py +++ b/examples/fastapi-vite/backend/main.py @@ -57,9 +57,7 @@ class ChatRequest(pydantic.BaseModel): @app.post("/chat") 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, tools=agent_.chat_agent.tools - ) + 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. diff --git a/src/ai/agents/agent.py b/src/ai/agents/agent.py index e498ffdb..e2b62736 100644 --- a/src/ai/agents/agent.py +++ b/src/ai/agents/agent.py @@ -117,6 +117,49 @@ def _process_interrupted_hooks(messages: list[types.messages.Message]) -> None: messages.pop() +def _aggregator_cls( + factory: Any, +) -> type[events_.Aggregator[Any, Any, Any]] | None: + """Resolve a tool's aggregator factory to the underlying class. + + Tools may declare the aggregator as a class directly (``LastAggregator``) + or via an ``Aggregate`` marker that wraps it. This normalizes both forms. + """ + if factory is None: + return None + if isinstance(factory, type) and issubclass(factory, events_.Aggregator): + return factory + inner = getattr(factory, "_factory", None) + if isinstance(inner, type) and issubclass(inner, events_.Aggregator): + return inner + return None + + +def _populate_model_inputs( + messages: Sequence[types.messages.Message], + tools_by_name: dict[str, AgentTool], +) -> None: + """Set ``model_input`` on tool results that arrived without one. + + Tool execution sets ``model_input`` directly; this fills in the + value for tool results that were reconstructed from a wire round- + trip (e.g. the AI SDK UI inbound path) and never had it computed. + """ + for msg in messages: + if msg.role != "tool": + continue + for part in msg.tool_results: + if part.has_model_input or part.is_error or part.is_hook_pending: + continue + tool = tools_by_name.get(part.tool_name) + if tool is None: + continue + agg_cls = _aggregator_cls(tool.aggregator) + if agg_cls is None: + continue + part.set_model_input(agg_cls.to_model_output(part.result)) + + class SimpleAggregator[Item, Result](events_.Aggregator[Item, Result, Result]): @classmethod def to_model_output(cls, snapshot: Result) -> Result: @@ -474,7 +517,7 @@ async def __call__(self, **overrides: Any) -> events_.ToolCallResult: async def _real(call: middleware_.ToolContext) -> events_.ToolCallResult: result: Any - model_result: Any + model_input: Any try: kwargs = _validate_kwargs(tool, call.kwargs) if tool.is_gen: @@ -491,10 +534,10 @@ async def _real(call: middleware_.ToolContext) -> events_.ToolCallResult: aggregator=tool.aggregator, ) result = agg.snapshot() - model_result = agg.get_model_output() + model_input = agg.get_model_output() else: result = await tool.fn(**kwargs) - model_result = result + model_input = result except Exception as exc: # A nested runtime (e.g. a sub-agent run inside this # tool) raises errors wrapped in a singleton TaskGroup @@ -510,14 +553,13 @@ async def _real(call: middleware_.ToolContext) -> events_.ToolCallResult: ), exception=unwrapped, ) - return tool_result( - types.messages.ToolResultPart( - tool_call_id=call.tool_call_id, - tool_name=call.tool_name, - result=result, - model_result=model_result, - ) + part = types.messages.ToolResultPart( + tool_call_id=call.tool_call_id, + tool_name=call.tool_name, + result=result, ) + part.set_model_input(model_input) + return tool_result(part) chain = middleware_._build_tool_chain(_real) return await chain(call) @@ -905,8 +947,8 @@ async def _aggregate_from[T, S, R]( Returns the live aggregator so callers can consume both the snapshot (the rich shape stored on ``ToolResultPart.result``) and the - model-facing value (``ToolResultPart.model_result``) without - re-aggregating. + model-facing value (set via ``ToolResultPart.set_model_input``) + without re-aggregating. """ agg = aggregator() @@ -1033,6 +1075,7 @@ async def _run( output_type=output_type, ) context._agent_tools_by_name = {t.name: t for t in self._tools} + _populate_model_inputs(context.messages, context._agent_tools_by_name) _process_interrupted_hooks(context.messages) async def _real(call: Context) -> AsyncGenerator[events_.AgentEvent]: diff --git a/src/ai/agents/ui/ai_sdk/inbound.py b/src/ai/agents/ui/ai_sdk/inbound.py index effad7b9..b138e1e9 100644 --- a/src/ai/agents/ui/ai_sdk/inbound.py +++ b/src/ai/agents/ui/ai_sdk/inbound.py @@ -8,12 +8,10 @@ import json import logging -from collections.abc import Sequence from typing import Any, NamedTuple -from ....types import events as events_ from ....types import messages as messages_ -from ...agent import AgentTool, MessageAggregator, MessageBundle +from ...agent import MessageBundle from ...hooks import resolve_hook from . import ui_message @@ -62,45 +60,25 @@ def _error_result(error_text: str | None, output: Any) -> dict[str, Any] | None: return normalized -def _aggregator_cls( - factory: Any, -) -> type[events_.Aggregator[Any, Any, Any]] | None: - """Resolve a tool's aggregator factory to the underlying class. - - Tools may declare an aggregator as a class directly (``LastAggregator``) - or via an ``Aggregate`` marker that wraps it (``Aggregate(LastAggregator, - delim="\\n")``). This normalizes both forms. - """ - if factory is None: - return None - if isinstance(factory, type) and issubclass(factory, events_.Aggregator): - return factory - inner = getattr(factory, "_factory", None) - if isinstance(inner, type) and issubclass(inner, events_.Aggregator): - return inner - return None - - -def _decode_wire_output( - output: Any, - agg_cls: type[events_.Aggregator[Any, Any, Any]] | None, -) -> Any: +def _decode_wire_output(output: Any) -> Any: """Reconstruct the internal snapshot type from a wire tool output. - For aggregator-backed tools the wire shape is a ``UIMessage`` (sub-agent - transcripts) or the aggregator's snapshot type directly (passthrough - aggregators). This function decodes UIMessage shapes back into a - ``MessageBundle`` so the parent agent's message history carries the - rich snapshot, mirroring what tool execution stored locally. + Hacky special case: when the wire output looks like a ``UIMessage`` + (the wire shape we emit for sub-agent / ``MessageAggregator`` tools), + decode it back to a ``MessageBundle``. Other shapes pass through + unchanged. This avoids requiring callers to thread the tool + registry into inbound parsing. """ - if agg_cls is None or output is None: + if not isinstance(output, dict): return output - - if agg_cls is MessageAggregator: + if output.get("role") != "assistant" or "parts" not in output: + return output + try: ui_msg = ui_message.UIMessage.model_validate(output) - inner = list(_parse([ui_msg])) - return MessageBundle(messages=tuple(inner)) - return output + except Exception: + return output + inner = list(_parse([ui_msg])) + return MessageBundle(messages=tuple(inner)) def _approval_hook_part(tp: ui_message.UIToolPart) -> messages_.HookPart[Any] | None: @@ -244,8 +222,6 @@ def _normalize_ui_messages( def to_messages( ui_messages: list[ui_message.UIMessage], - *, - tools: Sequence[AgentTool] | None = None, ) -> tuple[list[messages_.Message], list[ApprovalResponse]]: """Parse a UI request into runtime messages + extracted approvals. @@ -256,11 +232,10 @@ def to_messages( ``is_hook_pending`` placeholders for tool calls whose approval was just responded to but never recorded a real tool result. - ``tools`` lets the parser decode aggregator-backed tool outputs (e.g. - sub-agent UIMessages) back into their internal snapshot type and - populate ``ToolResultPart.model_result`` correctly. When omitted, - tool outputs are kept in their wire form — fine for caller code that - never feeds the messages back to a model. + Sub-agent tool outputs (UIMessage wire shape) are decoded back to + ``MessageBundle`` so the parent agent's message history carries the + rich snapshot. Per-tool model-facing values are populated by + :meth:`Agent.run` (which has the tool registry), not here. Returns ``(messages, approvals)``. The caller can pre-register resolutions via :func:`apply_approvals` before calling @@ -268,9 +243,7 @@ def to_messages( """ normalized = _normalize_ui_messages(ui_messages) approvals = extract_approvals(normalized) - messages = [ - m for m in _parse(normalized, tools=tools) if not _is_approval_response(m) - ] + messages = [m for m in _parse(normalized) if not _is_approval_response(m)] _patch_pending_hook_aborts(messages, approvals) return messages, approvals @@ -340,11 +313,7 @@ def _is_approval_response(msg: messages_.Message) -> bool: def _parse( ui_messages: list[ui_message.UIMessage], - *, - tools: Sequence[AgentTool] | None = None, ) -> list[messages_.Message]: - tools_by_name = {t.name: t for t in tools or []} - def _build_result_part( *, tool_call_id: str, @@ -352,21 +321,19 @@ def _build_result_part( output: Any, is_error: bool, ) -> messages_.ToolResultPart: - tool = tools_by_name.get(tool_name) - agg_cls = _aggregator_cls(tool.aggregator) if tool else None - if not is_error and agg_cls is not None: - snapshot = _decode_wire_output(output, agg_cls) - return messages_.ToolResultPart( - tool_call_id=tool_call_id, - tool_name=tool_name, - result=snapshot, - model_result=agg_cls.to_model_output(snapshot), - is_error=False, + if is_error: + result: Any = output + else: + decoded = _decode_wire_output(output) + result = ( + decoded + if isinstance(decoded, MessageBundle) + else (_normalize_tool_result(decoded)) ) return messages_.ToolResultPart( tool_call_id=tool_call_id, tool_name=tool_name, - result=output if is_error else _normalize_tool_result(output), + result=result, is_error=is_error, ) diff --git a/src/ai/models/ai_gateway/adapter.py b/src/ai/models/ai_gateway/adapter.py index 4c27f78b..68aabaa5 100644 --- a/src/ai/models/ai_gateway/adapter.py +++ b/src/ai/models/ai_gateway/adapter.py @@ -160,19 +160,18 @@ async def _messages_to_prompt( tool_results: list[dict[str, Any]] = [] for part in msg.parts: if isinstance(part, types.messages.ToolResultPart): + model_input = part.get_model_input() output = ( { "type": "error-text", "value": ( - str(part.model_result) - if part.model_result is not None - else "" + str(model_input) if model_input is not None else "" ), } if part.is_error else { "type": "json", - "value": part.model_result, + "value": model_input, } ) tool_results.append( diff --git a/src/ai/models/anthropic/adapter.py b/src/ai/models/anthropic/adapter.py index 46f486b0..0129d410 100644 --- a/src/ai/models/anthropic/adapter.py +++ b/src/ai/models/anthropic/adapter.py @@ -257,11 +257,12 @@ async def _messages_to_anthropic( tool_results: list[dict[str, Any]] = [] for part in msg.parts: if isinstance(part, types.messages.ToolResultPart): + model_input = part.get_model_input() entry: dict[str, Any] = { "type": "tool_result", "tool_use_id": part.tool_call_id, - "content": str(part.model_result) - if part.model_result is not None + "content": str(model_input) + if model_input is not None else "", } if part.is_error: diff --git a/src/ai/models/openai/adapter.py b/src/ai/models/openai/adapter.py index 087e9148..baf30a7c 100644 --- a/src/ai/models/openai/adapter.py +++ b/src/ai/models/openai/adapter.py @@ -158,12 +158,13 @@ async def _messages_to_openai( case "tool": for part in msg.parts: if isinstance(part, types.messages.ToolResultPart): + model_input = part.get_model_input() result.append( { "role": "tool", "tool_call_id": part.tool_call_id, - "content": str(part.model_result) - if part.model_result is not None + "content": str(model_input) + if model_input is not None else "", } ) diff --git a/src/ai/types/messages.py b/src/ai/types/messages.py index d08471c2..76aa80be 100644 --- a/src/ai/types/messages.py +++ b/src/ai/types/messages.py @@ -21,6 +21,9 @@ class TextPart(pydantic.BaseModel): kind: Literal["text"] = "text" +_MODEL_INPUT_UNSET: Any = object() + + class ToolResultPart(pydantic.BaseModel): id: str = pydantic.Field(default_factory=generate_id) tool_call_id: str @@ -32,21 +35,33 @@ class ToolResultPart(pydantic.BaseModel): # The "real" result of the tool call result: Any = None - # The value the LLM sees on its next turn. For most tools this is - # identical to ``result``; for aggregator-backed tools - # (sub-agents, streaming-text) it's derived from the aggregator's - # `get_model_output`. - model_result: Any = pydantic.Field(default=None, repr=False) + # Value the LLM sees on its next turn. For most tools this is + # identical to ``result``; for aggregator-backed tools (sub-agents, + # streaming-text) it's derived from the aggregator's + # ``get_model_output``. Not part of the wire model: it's populated + # by tool execution and by ``Agent.run`` (which has the tool + # registry) rather than carried across serialization. ``default_factory`` + # preserves singleton identity so the unset sentinel survives pydantic's + # default-copying. + _model_input: Any = pydantic.PrivateAttr(default_factory=lambda: _MODEL_INPUT_UNSET) kind: Literal["tool_result"] = "tool_result" model_config = pydantic.ConfigDict(frozen=True) - @pydantic.model_validator(mode="before") - @classmethod - def _default_model_result(cls, data: Any) -> Any: - if isinstance(data, dict) and "model_result" not in data: - data = {**data, "model_result": data.get("result")} - return data + def get_model_input(self) -> Any: + """Return the value the LLM should see, falling back to ``result``.""" + if self._model_input is _MODEL_INPUT_UNSET: + return self.result + return self._model_input + + def set_model_input(self, value: Any) -> None: + """Set the model-facing value (overrides the ``result`` fallback).""" + self._model_input = value + + @property + def has_model_input(self) -> bool: + """Whether ``set_model_input`` has been called on this part.""" + return self._model_input is not _MODEL_INPUT_UNSET class ToolCallPart(pydantic.BaseModel): diff --git a/tests/agents/test_generator_tools.py b/tests/agents/test_generator_tools.py index 4caf3809..d206c8e5 100644 --- a/tests/agents/test_generator_tools.py +++ b/tests/agents/test_generator_tools.py @@ -170,12 +170,12 @@ async def test_yield_from_nested_agent() -> None: e for e in all_events if isinstance(e, agent_events_.ToolCallResult) ] # MessageAggregator stores the rich MessageBundle as `result` and the - # extracted assistant text as `model_result` (the value the parent + # extracted assistant text as the model input (the value the parent # LLM sees on its next turn). sub_part = tool_results[0].results[0] assert isinstance(sub_part.result, MessageBundle) assert sub_part.result.messages[0].text == "Mars has two moons." - assert sub_part.model_result == "Mars has two moons." + assert sub_part.get_model_input() == "Mars has two moons." # The outer LLM's second call (index 2) must NOT contain any inner # agent messages. It should only see: the original user message, diff --git a/tests/agents/ui/ai_sdk/test_inbound.py b/tests/agents/ui/ai_sdk/test_inbound.py index a11fe9a5..27b4222d 100644 --- a/tests/agents/ui/ai_sdk/test_inbound.py +++ b/tests/agents/ui/ai_sdk/test_inbound.py @@ -4,7 +4,6 @@ import pytest -import ai from ai.agents.agent import MessageBundle from ai.agents.ui.ai_sdk import to_messages from ai.agents.ui.ai_sdk.inbound import ( @@ -174,20 +173,13 @@ def test_to_messages_rejects_empty_user() -> None: to_messages(ui) -@ai.tool -async def _research_tool(topic: str) -> ai.SubAgentTool: - """Sub-agent tool used by the inbound round-trip test.""" - if False: - yield # pragma: no cover - _ = topic - - def test_to_messages_decodes_subagent_tool_output() -> None: """A sub-agent tool's wire UIMessage decodes back to MessageBundle. - Round-trip: ``model_result`` is recomputed via the aggregator's - ``to_model_output``, and ``result`` carries the rich MessageBundle so - a subsequent UI render gets the same shape we sent. + ``result`` carries the rich MessageBundle so a subsequent UI render + gets the same shape we sent. ``model_input`` is left unset here — + populating it requires the tool registry, which lives in + :meth:`Agent.run`. """ # Wire shape: a tool-_research_tool part with output = UIMessage{parts=[text]}. ui = [ @@ -208,18 +200,18 @@ def test_to_messages_decodes_subagent_tool_output() -> None: id="a1", ), ] - messages, _ = to_messages(ui, tools=[_research_tool]) + messages, _ = to_messages(ui) # Find the tool message with the decoded result. tool_msgs = [m for m in messages if m.role == "tool"] assert len(tool_msgs) == 1 result_part = tool_msgs[0].tool_results[0] assert isinstance(result_part.result, MessageBundle) - assert result_part.model_result == "Mars has two moons." + assert not result_part.has_model_input -def test_to_messages_without_tools_keeps_wire_shape() -> None: - """No tools arg → tool outputs stay in their wire form (unchanged behavior).""" +def test_to_messages_passthrough_keeps_wire_shape() -> None: + """Non-UIMessage tool outputs stay in their wire form.""" ui = [ _ui("user", _text("hi"), id="u1"), _ui( @@ -230,5 +222,6 @@ def test_to_messages_without_tools_keeps_wire_shape() -> None: ] messages, _ = to_messages(ui) tool_msgs = [m for m in messages if m.role == "tool"] - assert tool_msgs[0].tool_results[0].result == {"pong": True} - assert tool_msgs[0].tool_results[0].model_result == {"pong": True} + part = tool_msgs[0].tool_results[0] + assert part.result == {"pong": True} + assert part.get_model_input() == {"pong": True}