feat(pydantic-ai): optional on_result callback to expose run result for usage capture

Adds an `on_result: Callable[[AgentRunResultEvent], Any] | None = None`
parameter to `convert_pydantic_ai_to_agentex_events`. When set, the
callback is invoked (sync or async) with the terminal `AgentRunResultEvent`
that carries the run result and usage. Streaming output is unchanged.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: declan-scale

src/agentex/lib/adk/_modules/_pydantic_ai_sync.py

@@ -21,7 +21,8 @@ async def handle_message_send(params):
 from __future__ import annotations
 
 import json
-from typing import TYPE_CHECKING, Any, AsyncIterator
+import inspect
+from typing import TYPE_CHECKING, Any, Callable, AsyncIterator
 
 from pydantic_ai.run import AgentRunResultEvent
 
@@ -105,6 +106,7 @@ def _tool_return_content(result: ToolReturnPart | Any) -> Any:
 async def convert_pydantic_ai_to_agentex_events(
     stream_response: AsyncIterator[Any],
     tracing_handler: "AgentexPydanticAITracingHandler | None" = None,
+    on_result: Callable[[AgentRunResultEvent], Any] | None = None,
 ) -> AsyncIterator[StreamTaskMessageStart | StreamTaskMessageDelta | StreamTaskMessageFull | StreamTaskMessageDone]:
     """Convert a Pydantic AI agent event stream into Agentex stream events.
 
@@ -132,6 +134,12 @@ async def convert_pydantic_ai_to_agentex_events(
             tool call in the run is also recorded as an Agentex child span
             beneath the handler's configured ``parent_span_id``. Streaming
             behavior is unchanged when omitted.
+        on_result: Optional callback invoked with the terminal
+            ``AgentRunResultEvent`` when the run completes. Both sync and
+            async callables are accepted. No ``StreamTaskMessage*`` events are
+            yielded for this terminal event; the callback is the only side
+            effect. Useful for capturing run-level usage without altering the
+            streaming output.
 
     Yields:
         Agentex ``StreamTaskMessage*`` events suitable for forwarding back over
@@ -328,6 +336,10 @@ async def convert_pydantic_ai_to_agentex_events(
             # Already covered by PartStart/PartDelta/PartEnd events above, or
             # informational only (FinalResultEvent / AgentRunResultEvent signal
             # run-level state, not new content to surface).
+            if isinstance(event, AgentRunResultEvent) and on_result is not None:
+                ret = on_result(event)
+                if inspect.iscoroutine(ret):
+                    await ret
             continue
 
         else:

tests/lib/adk/test_pydantic_ai_sync.py

@@ -6,6 +6,7 @@
 from typing import Any, AsyncIterator
 
 import pytest
+from pydantic_ai.run import AgentRunResult, AgentRunResultEvent
 from pydantic_ai.messages import (
     TextPart,
     PartEndEvent,
@@ -481,3 +482,67 @@ async def test_author_is_agent(self, events: list[Any]):
             content = getattr(e, "content", None)
             if content is not None and hasattr(content, "author"):
                 assert content.author == "agent"
+
+
+class TestOnResultCallback:
+    """on_result callback: captures the terminal AgentRunResultEvent without
+    altering streaming output."""
+
+    def _make_result_event(self, output: Any = "hello") -> AgentRunResultEvent:
+        result = AgentRunResult(output=output, _output_tool_name=None)
+        return AgentRunResultEvent(result=result)
+
+    async def test_callback_invoked_once_with_result_event(self):
+        """on_result is called exactly once, with the AgentRunResultEvent."""
+        captured: list[AgentRunResultEvent] = []
+
+        def on_result(event: AgentRunResultEvent) -> None:
+            captured.append(event)
+
+        result_event = self._make_result_event("the answer")
+        events = [result_event]
+        await _collect(convert_pydantic_ai_to_agentex_events(_aiter(events), on_result=on_result))
+
+        assert len(captured) == 1
+        assert captured[0] is result_event
+        assert captured[0].result.output == "the answer"
+
+    async def test_streaming_output_unchanged_with_callback(self):
+        """Yielded StreamTaskMessage* sequence is identical whether on_result is set or not."""
+        result_event = self._make_result_event()
+        events = [
+            PartStartEvent(index=0, part=TextPart(content="")),
+            PartDeltaEvent(index=0, delta=TextPartDelta(content_delta="hi")),
+            PartEndEvent(index=0, part=TextPart(content="hi")),
+            result_event,
+        ]
+
+        captured: list[AgentRunResultEvent] = []
+        out_with = await _collect(convert_pydantic_ai_to_agentex_events(_aiter(events), on_result=captured.append))
+        out_without = await _collect(convert_pydantic_ai_to_agentex_events(_aiter(events)))
+
+        assert len(out_with) == len(out_without)
+        for a, b in zip(out_with, out_without):
+            assert type(a) is type(b)
+        assert len(captured) == 1
+
+    async def test_no_callback_no_error(self):
+        """AgentRunResultEvent is silently ignored when on_result is None."""
+        result_event = self._make_result_event()
+        events = [result_event]
+        out = await _collect(convert_pydantic_ai_to_agentex_events(_aiter(events)))
+        assert out == []
+
+    async def test_async_callback_is_awaited(self):
+        """An async on_result callable is properly awaited."""
+        captured: list[AgentRunResultEvent] = []
+
+        async def on_result_async(event: AgentRunResultEvent) -> None:
+            captured.append(event)
+
+        result_event = self._make_result_event("async_output")
+        events = [result_event]
+        await _collect(convert_pydantic_ai_to_agentex_events(_aiter(events), on_result=on_result_async))
+
+        assert len(captured) == 1
+        assert captured[0].result.output == "async_output"