Tweak streaming tool data representation a bit more

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() {
                                 </ConfirmationActions>
                               </Confirmation>
 
-                              <ToolOutput
-                                output={toolPart.output}
-                                errorText={toolPart.errorText}
-                              />
+                              {toolPart.type === "tool-talk_to_mothership" &&
+                              toolPart.state === "output-available" ? (
+                                <div className="space-y-2">
+                                  {(toolPart.output as UIMessage).parts.map(
+                                    (p, i) => renderUIPart(p, i),
+                                  )}
+                                </div>
+                              ) : (
+                                <ToolOutput
+                                  output={toolPart.output}
+                                  errorText={toolPart.errorText}
+                                />
+                              )}
                             </ToolContent>
                           </Tool>
                         );

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 (
-        <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";
+export function renderUIPart(
+  part: UIMessage["parts"][number],
+  key: number,
+): ReactNode {
+  if (isTextUIPart(part)) {
+    return (
+      <div key={key} className="p-3 text-sm leading-relaxed">
+        <MessageResponse>{part.text}</MessageResponse>
+      </div>
+    );
+  }
+  if (isToolUIPart(part)) {
+    const isComplete = part.state === "output-available";
+    if (part.type === "dynamic-tool") {
+      const dyn = part as DynamicToolUIPart;
       return (
-        <Tool key={i} defaultOpen={isComplete}>
-          <ToolHeader type={tool.type} state={tool.state} />
+        <Tool key={key} defaultOpen={isComplete}>
+          <ToolHeader
+            type="dynamic-tool"
+            state={dyn.state}
+            toolName={dyn.toolName}
+          />
           <ToolContent>
-            <ToolInput input={tool.input} />
-            <ToolOutput output={tool.output} errorText={tool.errorText} />
+            <ToolInput input={dyn.input} />
+            <ToolOutput output={dyn.output} errorText={dyn.errorText} />
           </ToolContent>
         </Tool>
       );
     }
-    return null;
-  });
+    const tool = part as ToolUIPart;
+    return (
+      <Tool key={key} 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 = ({
@@ -186,10 +195,7 @@ export const ToolOutput = ({
 
   let Output = <div>{output as ReactNode}</div>;
 
-  const message = asUIMessage(output);
-  if (message) {
-    Output = <div className="space-y-2">{renderUIParts(message.parts ?? [])}</div>;
-  } else if (typeof output === "string") {
+  if (typeof output === "string") {
     Output = (
       <div className="p-3 text-sm leading-relaxed">
         <MessageResponse>{output}</MessageResponse>

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]: