refactor(pydantic-ai): make tool-request coalescing opt-in (preserve sync arg streaming); ref AGX1-377

PydanticAITurn.events feeds BOTH delivery channels (yield_turn for sync,
auto_send_turn for async). Applying _coalesce_tool_requests unconditionally
would deliver tool requests as a single Full with no ToolRequestDelta tokens,
losing the sync converter's documented tool-call-argument token streaming
(Task 4 routes the sync/HTTP path through emitter.yield_turn(PydanticAITurn(...))).

- Add constructor param coalesce_tool_requests: bool = False. Default OFF means
  PydanticAITurn(...).events == bare convert_pydantic_ai_to_agentex_events output
  (Start+Delta+Done for tool calls, arg streaming preserved on yield/sync).
- stream_pydantic_ai_events builds the Turn with coalesce_tool_requests=True,
  because the foundation auto_send currently DROPS tool requests delivered as
  Start+Delta+Done (AGX1-377). Comment cites AGX1-377 as a temporary workaround
  to be removed once auto_send handles the streamed tool-request shape natively.
- Tests: default-off Turn yields a ToolRequestDelta for streamed args (matches
  bare converter); coalesce-on Turn yields a single Full(tool_request) with
  fully-accumulated args and no ToolRequestDelta. Async characterization test
  still passes (goes through coalesce=True).
- Add parts-manager invariant comment to the two corrected async tests.

auto_send.py is unchanged (final_text last-segment fix stays; AGX1-377 covers
the Start+Delta+Done handling).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: declan-scale

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

@@ -51,7 +51,17 @@ async def stream_pydantic_ai_events(
     from agentex.lib.core.harness.emitter import UnifiedEmitter
     from agentex.lib.adk._modules._pydantic_ai_turn import PydanticAITurn
 
-    turn = PydanticAITurn(stream, model=None, tracing_handler=tracing_handler)
+    # coalesce_tool_requests=True is a temporary workaround (AGX1-377): the
+    # foundation auto_send currently DROPS tool requests delivered as the
+    # streamed Start+ToolRequestDelta+Done shape, so we collapse them into a
+    # single Full here. Once auto_send handles the streamed tool-request shape
+    # natively, this flag should be removed and the default (off) used.
+    turn = PydanticAITurn(
+        stream,
+        model=None,
+        tracing_handler=tracing_handler,
+        coalesce_tool_requests=True,
+    )
     emitter = UnifiedEmitter(
         task_id=task_id,
         trace_id=None,

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

@@ -154,21 +154,25 @@ class PydanticAITurn:
     canonical ``StreamTaskMessage*`` stream; ``usage()`` returns a normalized
     ``TurnUsage`` (valid only after ``events`` is exhausted).
 
-    Tool-request events emitted by ``convert_pydantic_ai_to_agentex_events``
-    as ``Start+Done`` pairs are coalesced into ``StreamTaskMessageFull``
-    events so that ``auto_send`` can deliver them as atomic messages (Option A
-    — no streaming of argument tokens in the async path).
+    By default ``events`` is identical to the bare
+    ``convert_pydantic_ai_to_agentex_events`` output (tool calls stream as
+    ``Start + ToolRequestDelta + Done``, preserving argument-token streaming on
+    the sync/yield channel). When ``coalesce_tool_requests=True``, tool-request
+    sequences are collapsed into a single ``StreamTaskMessageFull`` (Option A —
+    no streaming of argument tokens) for the async/auto_send path.
     """
 
     def __init__(
         self,
         stream: AsyncIterator[Any],
         model: str | None = None,
         tracing_handler: "AgentexPydanticAITracingHandler | None" = None,
+        coalesce_tool_requests: bool = False,
     ) -> None:
         self._stream = stream
         self._model = model
         self._tracing_handler = tracing_handler
+        self._coalesce_tool_requests = coalesce_tool_requests
         self._usage = TurnUsage(model=model)
 
     @property
@@ -195,8 +199,12 @@ def _capture(result_event: AgentRunResultEvent) -> None:
             tracing_handler=self._tracing_handler,
             on_result=_capture,
         )
-        async for ev in _coalesce_tool_requests(raw_stream):
-            yield ev
+        if self._coalesce_tool_requests:
+            async for ev in _coalesce_tool_requests(raw_stream):
+                yield ev
+        else:
+            async for ev in raw_stream:
+                yield ev
 
     def usage(self) -> TurnUsage:
         """Return the normalized usage for this turn.

tests/lib/adk/test_pydantic_ai_async.py

@@ -268,6 +268,10 @@ async def test_tool_call_emits_full_tool_request_message_on_part_end(
         This test uses a realistic pydantic-ai event sequence: args arrive as a
         PartDeltaEvent fragment (the way OpenAI/Anthropic actually stream JSON
         tool-call arguments). The new implementation accumulates them correctly.
+
+        Parts-manager invariant: PartEnd.part is the accumulated snapshot; real
+        pydantic-ai conveys args via PartStart + PartDeltaEvent, so a
+        PartStart(None)+PartEnd(json) with no delta is not realizable.
         """
         from pydantic_ai.messages import ToolCallPartDelta
 
@@ -334,6 +338,10 @@ async def test_tool_call_with_invalid_json_args_surfaces_raw(
 
         Uses a PartDeltaEvent to deliver the invalid string (the way pydantic-ai
         actually surfaces arg tokens) so the coalescer picks it up.
+
+        Parts-manager invariant: PartEnd.part is the accumulated snapshot; real
+        pydantic-ai conveys args via PartStart + PartDeltaEvent, so a
+        PartStart(None)+PartEnd(json) with no delta is not realizable.
         """
         from pydantic_ai.messages import ToolCallPartDelta
 

tests/lib/adk/test_pydantic_ai_turn.py

@@ -228,3 +228,93 @@ async def test_no_usage_event_leaves_default_usage(self):
         assert usage.model == "openai:gpt-4o"
         assert usage.input_tokens is None
         assert usage.num_llm_calls == 0
+
+
+class TestToolRequestCoalescing:
+    """The ``coalesce_tool_requests`` flag controls tool-call delivery shape.
+
+    Default (off): tool calls stream as Start + ToolRequestDelta + Done,
+    matching the bare converter and preserving argument-token streaming on the
+    sync/yield channel.
+
+    On: tool-request sequences collapse into a single StreamTaskMessageFull
+    (Option A) for the async/auto_send path (a temporary AGX1-377 workaround).
+    """
+
+    async def test_default_off_matches_bare_converter_for_streamed_tool_call(self):
+        """Default Turn (coalesce off) yields a ToolRequestDelta for a streamed-args
+        tool call — i.e. it is byte-for-byte the bare converter output, so the
+        sync/yield channel keeps its argument-token streaming."""
+        from pydantic_ai.messages import ToolCallPart, ToolCallPartDelta
+
+        from agentex.types.tool_request_delta import ToolRequestDelta
+        from agentex.types.task_message_update import StreamTaskMessageDelta
+        from agentex.lib.adk._modules._pydantic_ai_sync import convert_pydantic_ai_to_agentex_events
+
+        tool_events = [
+            PartStartEvent(index=0, part=ToolCallPart(tool_name="get_weather", args=None, tool_call_id="c1")),
+            PartDeltaEvent(index=0, delta=ToolCallPartDelta(args_delta='{"city":"Paris"}')),
+            PartEndEvent(
+                index=0,
+                part=ToolCallPart(tool_name="get_weather", args='{"city":"Paris"}', tool_call_id="c1"),
+            ),
+        ]
+
+        turn = PydanticAITurn(_aiter(tool_events), model="openai:gpt-4o")
+        turn_out = await _collect(turn.events)
+
+        bare_out = await _collect(convert_pydantic_ai_to_agentex_events(_aiter(tool_events)))
+
+        # Default Turn is identical to the bare converter.
+        assert len(turn_out) == len(bare_out)
+        for a, b in zip(turn_out, bare_out):
+            assert type(a) is type(b)
+            assert a.model_dump() == b.model_dump()
+
+        # And the arg-streaming delta is present (not coalesced away).
+        deltas = [
+            e for e in turn_out if isinstance(e, StreamTaskMessageDelta) and isinstance(e.delta, ToolRequestDelta)
+        ]
+        assert len(deltas) == 1, "streamed tool-call args must surface as a ToolRequestDelta when coalesce is off"
+        assert deltas[0].delta.arguments_delta == '{"city":"Paris"}'
+
+    async def test_coalesce_on_emits_single_full_with_accumulated_args_and_no_delta(self):
+        """coalesce_tool_requests=True yields one StreamTaskMessageFull(tool_request)
+        with fully-accumulated arguments and NO ToolRequestDelta."""
+        from pydantic_ai.messages import ToolCallPart, ToolCallPartDelta
+
+        from agentex.types.tool_request_delta import ToolRequestDelta
+        from agentex.types.task_message_update import (
+            StreamTaskMessageDone,
+            StreamTaskMessageFull,
+            StreamTaskMessageDelta,
+            StreamTaskMessageStart,
+        )
+        from agentex.types.tool_request_content import ToolRequestContent
+
+        tool_events = [
+            PartStartEvent(index=0, part=ToolCallPart(tool_name="get_weather", args=None, tool_call_id="c1")),
+            PartDeltaEvent(index=0, delta=ToolCallPartDelta(args_delta='{"city":"Paris"}')),
+            PartEndEvent(
+                index=0,
+                part=ToolCallPart(tool_name="get_weather", args='{"city":"Paris"}', tool_call_id="c1"),
+            ),
+        ]
+
+        turn = PydanticAITurn(_aiter(tool_events), model="openai:gpt-4o", coalesce_tool_requests=True)
+        turn_out = await _collect(turn.events)
+
+        # Exactly one event: a Full(tool_request). No Start/Delta/Done leak through.
+        assert len(turn_out) == 1
+        full = turn_out[0]
+        assert isinstance(full, StreamTaskMessageFull)
+        assert isinstance(full.content, ToolRequestContent)
+        assert full.content.tool_call_id == "c1"
+        assert full.content.name == "get_weather"
+        assert full.content.arguments == {"city": "Paris"}, "args must be fully accumulated"
+
+        assert not any(isinstance(e, StreamTaskMessageStart) for e in turn_out)
+        assert not any(isinstance(e, StreamTaskMessageDone) for e in turn_out)
+        assert not any(
+            isinstance(e, StreamTaskMessageDelta) and isinstance(e.delta, ToolRequestDelta) for e in turn_out
+        ), "no ToolRequestDelta when coalescing is on"