+ {(toolPart.output as UIMessage).parts.map(
+ (p, i) => renderUIPart(p, i),
+ )}
+
+ ) : (
+
-
+
+
-
-
+
+
);
}
- 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..e2b62736 100644
--- a/src/ai/agents/agent.py
+++ b/src/ai/agents/agent.py
@@ -117,9 +117,53 @@ 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]):
- def to_model_output(self) -> Result:
- return self.snapshot()
+ @classmethod
+ def to_model_output(cls, snapshot: Result) -> Result:
+ return snapshot
class ConcatAggregator(SimpleAggregator[str, str]):
@@ -167,8 +211,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 to_model_output(cls, snapshot: MessageBundle) -> str:
+ for m in reversed(snapshot.messages):
if m.role == "assistant" and m.text:
return m.text
return ""
@@ -471,21 +516,28 @@ async def __call__(self, **overrides: Any) -> events_.ToolCallResult:
tool = self._tool
async def _real(call: middleware_.ToolContext) -> events_.ToolCallResult:
+ result: Any
+ model_input: 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_input = agg.get_model_output()
else:
result = await tool.fn(**kwargs)
+ 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
@@ -501,13 +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,
- )
+ 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)
@@ -842,10 +894,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 +925,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.get_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 (set via ``ToolResultPart.set_model_input``)
+ without re-aggregating.
+ """
agg = aggregator()
rt = runtime.get_runtime()
@@ -888,7 +965,7 @@ async def yield_from[T, R](
aggregator_factory=aggregator,
)
)
- return agg.to_model_output()
+ return agg
class Agent:
@@ -998,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 65b6ad3f..b138e1e9 100644
--- a/src/ai/agents/ui/ai_sdk/inbound.py
+++ b/src/ai/agents/ui/ai_sdk/inbound.py
@@ -11,6 +11,7 @@
from typing import Any, NamedTuple
from ....types import messages as messages_
+from ...agent import MessageBundle
from ...hooks import resolve_hook
from . import ui_message
@@ -59,6 +60,27 @@ def _error_result(error_text: str | None, output: Any) -> dict[str, Any] | None:
return normalized
+def _decode_wire_output(output: Any) -> Any:
+ """Reconstruct the internal snapshot type from a wire tool output.
+
+ 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 not isinstance(output, dict):
+ return output
+ if output.get("role") != "assistant" or "parts" not in output:
+ return output
+ try:
+ ui_msg = ui_message.UIMessage.model_validate(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:
"""Reconstruct approval hook state from a UI tool part when possible."""
approval = tp.approval
@@ -210,6 +232,11 @@ def to_messages(
``is_hook_pending`` placeholders for tool calls whose approval was
just responded to but never recorded a real tool result.
+ 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
:meth:`Agent.run` if the run should resume from a hook.
@@ -287,6 +314,29 @@ def _is_approval_response(msg: messages_.Message) -> bool:
def _parse(
ui_messages: list[ui_message.UIMessage],
) -> list[messages_.Message]:
+ def _build_result_part(
+ *,
+ tool_call_id: str,
+ tool_name: str,
+ output: Any,
+ is_error: bool,
+ ) -> messages_.ToolResultPart:
+ 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=result,
+ is_error=is_error,
+ )
+
result: list[messages_.Message] = []
for ui_msg in ui_messages:
@@ -313,10 +363,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 +385,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..68aabaa5 100644
--- a/src/ai/models/ai_gateway/adapter.py
+++ b/src/ai/models/ai_gateway/adapter.py
@@ -160,17 +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.result) if part.result is not None else ""
+ str(model_input) if model_input is not None else ""
),
}
if part.is_error
else {
"type": "json",
- "value": part.result,
+ "value": model_input,
}
)
tool_results.append(
diff --git a/src/ai/models/anthropic/adapter.py b/src/ai/models/anthropic/adapter.py
index e030c6d9..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.result)
- if part.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 b84fd1fc..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.result)
- if part.result is not None
+ "content": str(model_input)
+ if model_input is not None
else "",
}
)
diff --git a/src/ai/types/events.py b/src/ai/types/events.py
index e5916752..a2063525 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 get_model_output(self) -> ModelResult:
+ """Return the model-facing value derived from this aggregator's state.
+
+ Default implementation defers to :meth:`to_model_output`; subclasses
+ with non-trivial state may override either or both.
+ """
+ return type(self).to_model_output(self.snapshot())
+
+ @classmethod
@abc.abstractmethod
- def to_model_output(self) -> 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
+ 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..76aa80be 100644
--- a/src/ai/types/messages.py
+++ b/src/ai/types/messages.py
@@ -21,18 +21,48 @@ 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
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
+
+ # 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)
+ 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):
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..d206c8e5 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 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.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 2f4a3b03..27b4222d 100644
--- a/tests/agents/ui/ai_sdk/test_inbound.py
+++ b/tests/agents/ui/ai_sdk/test_inbound.py
@@ -4,6 +4,7 @@
import pytest
+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 +171,57 @@ def test_to_messages_rejects_empty_user() -> None:
ui = [UIMessage.model_validate({"id": "u1", "role": "user", "parts": []})]
with pytest.raises(ValueError):
to_messages(ui)
+
+
+def test_to_messages_decodes_subagent_tool_output() -> None:
+ """A sub-agent tool's wire UIMessage decodes back to MessageBundle.
+
+ ``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 = [
+ _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)
+
+ # 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 not result_part.has_model_input
+
+
+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(
+ "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"]
+ part = tool_msgs[0].tool_results[0]
+ assert part.result == {"pong": True}
+ assert part.get_model_input() == {"pong": True}