Validate tool results with return types

We have to do a bunch of annoying type annotation magic to extract the
result type of aggregators from from type annotations.

A follow-up question here: should we use the aggregator to populate
model_input *here*, instead of in Agent? I will leave that to a second
PR.

The invocation for parsing tools with pydantic validation is:
```
ai.messages.Message.model_validate(
    data,
    context=ai.tool_validate_context([weather]),
)
```

It is also now an argument to `ai.agents.ui.ai_sdk.to_messages`.

Happy to bikeshed about names.

Author: msullivan

src/ai/__init__.py

@@ -102,6 +102,7 @@
     tool_result_part,
     user_message,
 )
+from .types.messages import tool_context
 
 __all__ = [
     "DEFAULT",
@@ -198,6 +199,7 @@
     "text_part",
     "thinking",
     "tool",
+    "tool_context",
     "tool_message",
     "tool_result",
     "tool_result_part",

src/ai/agents/agent.py

@@ -35,7 +35,7 @@
 # Use the typing_extensions backport so this works on 3.12 too.
 from typing_extensions import TypeVar
 
-from .. import models, types, util
+from .. import models, type_utils, types, util
 from ..types import builders
 from ..types import events as events_
 from ..types.messages import MessageBundle
@@ -351,6 +351,67 @@ async def render(prompt: str) -> StreamingTextTool:
 """
 
 
+def _return_type_from_callable(fn: Callable[..., Any]) -> Any:
+    try:
+        return get_type_hints(fn, include_extras=True).get("return")
+    except Exception:
+        return None
+
+
+def _stream_item_type(return_type: Any) -> Any:
+    resolved = type_utils.resolve_type_alias(return_type)
+    if typing.get_origin(resolved) is typing.Annotated:
+        resolved = typing.get_args(resolved)[0]
+
+    if typing.get_origin(resolved) is not AsyncGenerator:
+        return None
+
+    args = typing.get_args(resolved)
+    return args[0] if args else Any
+
+
+def _aggregator_result_type(
+    aggregator: Callable[[], events_.Aggregator[Any, Any, Any]] | None,
+    stream_item_type: Any,
+) -> Any:
+    agg_cls = _aggregator_cls(aggregator)
+    if agg_cls is None:
+        return None
+
+    args = type_utils.generic_base_args(agg_cls, events_.Aggregator)
+    if args is None:
+        return None
+
+    item_type, result_type, _model_input_type = args
+    bindings = type_utils.bind_typevars(item_type, stream_item_type)
+    return type_utils.replace_typevars(result_type, bindings)
+
+
+def _tool_result_type_from_return_annotation(
+    return_type: Any,
+    aggregator: Callable[[], events_.Aggregator[Any, Any, Any]] | None,
+) -> Any:
+    """Return the type used to validate ``ToolResultPart.result``.
+
+    Normal tools store the function's return value directly, so the function's
+    return annotation is the result type.
+
+    Async-generator tools store ``aggregator.snapshot()`` instead. For those,
+    read the ``Result`` parameter from ``Aggregator[Item, Result, ModelInput]``
+    on the configured aggregator class, binding generic aggregators from the
+    stream item type when possible (for example ``StreamingStatusTool[str]`` +
+    ``LastAggregator[T]`` becomes ``str | None``).
+    """
+    if return_type is None:
+        return None
+
+    item_type = _stream_item_type(return_type)
+    if item_type is not None:
+        return _aggregator_result_type(aggregator, item_type)
+
+    return return_type
+
+
 def _aggregate_from_return_type(fn: Callable[..., Any]) -> Aggregate | None:
     """Find an ``Aggregate`` marker in *fn*'s return-type metadata, if any.
 
@@ -360,11 +421,7 @@ def _aggregate_from_return_type(fn: Callable[..., Any]) -> Aggregate | None:
     * a PEP 695 alias ``type Foo = Annotated[X, Aggregate(...)]``,
     * a parameterized alias ``type Foo[T] = Annotated[X[T], Aggregate(...)]``.
     """
-    try:
-        hints = get_type_hints(fn, include_extras=True)
-    except Exception:
-        return None
-    ret = hints.get("return")
+    ret = _return_type_from_callable(fn)
     if ret is None:
         return None
 
@@ -398,6 +455,7 @@ class AgentTool:
     validator: type[pydantic.BaseModel] | None = None
     is_gen: bool = False
     aggregator: Callable[[], events_.Aggregator[Any, Any, Any]] | None = None
+    return_type: Any = None
 
     @property
     def name(self) -> str:
@@ -435,7 +493,17 @@ def tool[**P, T](fn: Callable[P, AsyncGenerator[T]], /) -> AgentTool: ...
 
 @overload
 def tool[**P](
-    *, require_approval: bool
+    *,
+    require_approval: bool,
+    return_type: Any = None,
+) -> Callable[[Callable[P, Any]], AgentTool]: ...
+
+
+@overload
+def tool[**P](
+    *,
+    return_type: Any,
+    require_approval: bool = False,
 ) -> Callable[[Callable[P, Any]], AgentTool]: ...
 
 
@@ -444,6 +512,7 @@ def tool[**P](
     *,
     aggregator: Callable[[], events_.Aggregator[Any, Any, Any]],
     require_approval: bool = False,
+    return_type: Any = None,
 ) -> Callable[[Callable[P, AsyncGenerator[Any]]], AgentTool]: ...
 
 
@@ -455,6 +524,7 @@ def tool[**P, T, R](
     *,
     aggregator: Callable[[], events_.Aggregator[Any, Any, Any]] | None = None,
     require_approval: bool = False,
+    return_type: Any = None,
 ) -> (
     Callable[[Callable[P, AsyncGenerator[Any]]], AgentTool]
     | Callable[[Callable[P, Awaitable[R]]], AgentTool]
@@ -466,7 +536,8 @@ def tool[**P, T, R](
     ``aggregator=`` keyword argument or by annotating the return type
     with an :class:`Aggregate` marker (e.g. via the :data:`SubAgentTool`
     or :data:`StreamingStatusTool` aliases).  Specifying both raises
-    ``TypeError``.
+    ``TypeError``. Pass ``return_type=`` to override the type used when
+    validating round-tripped tool results.
     """
 
     def wrap(fn: Any) -> AgentTool:
@@ -483,6 +554,7 @@ def wrap(fn: Any) -> AgentTool:
 
         validator = pydantic.create_model(f"{fn.__name__}_Args", **fields)
 
+        annotated_return_type = _return_type_from_callable(fn)
         annotated_aggregate = _aggregate_from_return_type(fn)
         if annotated_aggregate is not None and aggregator is not None:
             raise TypeError(
@@ -508,6 +580,10 @@ def wrap(fn: Any) -> AgentTool:
             validator=validator,
             is_gen=inspect.isasyncgenfunction(fn),
             aggregator=effective_aggregator,
+            return_type=return_type
+            or _tool_result_type_from_return_annotation(
+                annotated_return_type, effective_aggregator
+            ),
         )
 
     if fn is None:

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

@@ -8,7 +8,7 @@
 
 import json
 import logging
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from ....types import messages as messages_
 from ....types.messages import MessageBundle
@@ -17,6 +17,9 @@
 from .approvals import ApprovalResponse, extract_approvals
 from .tool_utils import normalize_tool_args
 
+if TYPE_CHECKING:
+    from ...agent import AgentTool
+
 logger = logging.getLogger(__name__)
 
 
@@ -64,6 +67,7 @@ def _build_result_part(
     output: Any,
     is_error: bool,
     kind_hint: str | None = None,
+    validate_context: dict[str, Any] | None = None,
 ) -> messages_.ToolResultPart:
     """Reconstruct a tool result from its wire form.
 
@@ -93,16 +97,21 @@ def _build_result_part(
             else ui_messages_.UIMessage.model_validate(m)
             for m in raw
         ]
-        result = MessageBundle(messages=tuple(_parse(ui_msgs)))
+        result = MessageBundle(
+            messages=tuple(_parse(ui_msgs, validate_context=validate_context))
+        )
         result_kind = "special"
     else:
         result = _normalize_tool_result(output)
         result_kind = "json"
-    return messages_.ToolResultPart(
-        tool_call_id=tool_call_id,
-        tool_name=tool_name,
-        result=result,
-        result_kind=result_kind,
+    return messages_.ToolResultPart.model_validate(
+        dict(  # noqa: C408
+            tool_call_id=tool_call_id,
+            tool_name=tool_name,
+            result=result,
+            result_kind=result_kind,
+        ),
+        context=validate_context,
     )
 
 
@@ -205,6 +214,8 @@ def _patch_pending_hook_aborts(
 
 def _parse(
     ui_messages: list[ui_messages_.UIMessage],
+    *,
+    validate_context: dict[str, Any] | None = None,
 ) -> list[messages_.Message]:
     result: list[messages_.Message] = []
 
@@ -276,6 +287,7 @@ def _parse(
                                     kind_hint=result_kinds.get(
                                         inv.tool_invocation_id
                                     ),
+                                    validate_context=validate_context,
                                 )
                             )
 
@@ -340,6 +352,7 @@ def _parse(
                                 output=_tool_result_output(tp),
                                 is_error=is_error,
                                 kind_hint=result_kinds.get(tp.tool_call_id),
+                                validate_context=validate_context,
                             )
                         )
                         if tp.result_provider_metadata is not None:
@@ -507,6 +520,8 @@ def _split_assistant_parts(
 
 def to_messages(
     ui_messages: list[ui_messages_.UIMessage],
+    *,
+    tools: list[AgentTool] | None = None,
 ) -> tuple[list[messages_.Message], list[ApprovalResponse]]:
     """Parse a UI request into runtime messages + extracted approvals.
 
@@ -526,11 +541,14 @@ def to_messages(
     resolutions via :func:`apply_approvals` before calling
     :meth:`Agent.run` if the run should resume from a hook.
     """
+    validate_context = (
+        messages_.tool_context(tools) if tools is not None else None
+    )
     normalized = _normalize_ui_messages(ui_messages)
     approval_responses = extract_approvals(normalized)
     messages = [
         m
-        for m in _parse(normalized)
+        for m in _parse(normalized, validate_context=validate_context)
         if not approvals.is_resolved_approval_message(m)
     ]
     _patch_pending_hook_aborts(messages, approval_responses)