refactor(pydantic-ai): drop coalesce_tool_requests workaround — foundation auto_send delivers streamed tool requests natively (AGX1-377/378)

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

Author: declan-scale

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

@@ -51,16 +51,10 @@ async def stream_pydantic_ai_events(
     from agentex.lib.core.harness.emitter import UnifiedEmitter
     from agentex.lib.adk._modules._pydantic_ai_turn import PydanticAITurn
 
-    # 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,

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

@@ -15,7 +15,6 @@
 
 from __future__ import annotations
 
-import json
 from typing import TYPE_CHECKING, Any, AsyncIterator
 
 from pydantic_ai.run import AgentRunResultEvent
@@ -74,105 +73,29 @@ def pydantic_ai_usage_to_turn_usage(usage: Any, model: str | None) -> TurnUsage:
     )
 
 
-async def _coalesce_tool_requests(
-    source: AsyncIterator[StreamTaskMessage],
-) -> AsyncIterator[StreamTaskMessage]:
-    """Convert Start(tool_request)+deltas+Done into Full(tool_request).
-
-    ``convert_pydantic_ai_to_agentex_events`` emits ``Start+Done`` for tool
-    calls (enabling streaming of argument tokens over the sync/HTTP channel).
-    The async/auto_send delivery path does not stream tool-call arguments —
-    it uses Option A (full messages). This wrapper coalesces the Start+Done
-    sequence into a single ``StreamTaskMessageFull``, matching the shape that
-    ``auto_send`` expects and that the harness conformance tests are designed for.
-
-    Argument delta fragments (``ToolRequestDelta.arguments_delta``) are
-    accumulated as a JSON string and parsed back into a dict. If parsing
-    fails, the raw string is stored under ``"_raw"`` so no information is lost.
-    """
-    from agentex.types.tool_request_delta import ToolRequestDelta
-    from agentex.types.tool_request_content import ToolRequestContent
-
-    # pending[index] = (ToolRequestContent, accumulated_args_delta_str)
-    pending: dict[Any, tuple[Any, str]] = {}
-
-    async for event in source:
-        if isinstance(event, StreamTaskMessageStart):
-            ctype = getattr(event.content, "type", None)
-            if ctype == "tool_request":
-                # Stage; do not yield — replaced by Full on Done.
-                pending[event.index] = (event.content, "")
-                continue
-
-        elif isinstance(event, StreamTaskMessageDelta):
-            if event.index in pending and isinstance(event.delta, ToolRequestDelta):
-                # Accumulate argument delta fragments; don't yield.
-                content, accum = pending[event.index]
-                pending[event.index] = (content, accum + (event.delta.arguments_delta or ""))
-                continue
-
-        elif isinstance(event, StreamTaskMessageDone):
-            if event.index in pending:
-                content, args_delta = pending.pop(event.index)
-                # Build final arguments: merge initial dict with accumulated delta.
-                base_args: dict[str, Any] = {}
-                if isinstance(content, ToolRequestContent):
-                    base_args = dict(content.arguments) if content.arguments else {}
-
-                if args_delta:
-                    try:
-                        parsed = json.loads(args_delta)
-                        if isinstance(parsed, dict):
-                            base_args.update(parsed)
-                        else:
-                            base_args["_raw"] = args_delta
-                    except json.JSONDecodeError:
-                        base_args["_raw"] = args_delta
-
-                # Emit as Full with the complete arguments.
-                full_content = (
-                    ToolRequestContent(
-                        type="tool_request",
-                        author=content.author if isinstance(content, ToolRequestContent) else "agent",
-                        tool_call_id=content.tool_call_id if isinstance(content, ToolRequestContent) else "",
-                        name=content.name if isinstance(content, ToolRequestContent) else "",
-                        arguments=base_args,
-                    )
-                    if isinstance(content, ToolRequestContent)
-                    else content
-                )
-                yield StreamTaskMessageFull(type="full", index=event.index, content=full_content)
-                continue
-
-        yield event
-
-
 class PydanticAITurn:
     """A single harness turn backed by a pydantic-ai event stream.
 
     Satisfies the ``HarnessTurn`` protocol: ``events`` async-generates the
     canonical ``StreamTaskMessage*`` stream; ``usage()`` returns a normalized
     ``TurnUsage`` (valid only after ``events`` is exhausted).
 
-    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.
+    ``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). The foundation
+    ``auto_send`` delivers the streamed tool-request shape natively (AGX1-377),
+    so no coalescing is needed on either channel.
     """
 
     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
@@ -199,12 +122,8 @@ def _capture(result_event: AgentRunResultEvent) -> None:
             tracing_handler=self._tracing_handler,
             on_result=_capture,
         )
-        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
+        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

@@ -82,7 +82,9 @@ class FakeStreamingModule:
     def __init__(self) -> None:
         self.contexts: list[FakeContext] = []
 
-    def streaming_task_message_context(self, *, task_id: str, initial_content: Any) -> FakeContext:
+    def streaming_task_message_context(
+        self, *, task_id: str, initial_content: Any, streaming_mode: str = "coalesced", created_at: Any = None
+    ) -> FakeContext:
         tm = TaskMessage(
             id=f"m{len(self.contexts) + 1}",
             task_id=task_id,
@@ -255,26 +257,25 @@ async def test_empty_thinking_delta_is_skipped(
 
 
 class TestToolCallEmission:
-    async def test_tool_call_emits_full_tool_request_message_on_part_end(
+    async def test_tool_call_opens_streaming_context_with_identity(
         self, fake_adk: tuple[FakeStreamingModule, FakeMessagesModule]
     ) -> None:
-        """Tool requests are posted as full messages (open+close on streaming context).
+        """Tool requests are delivered as a streaming context (Start+Delta+Done).
 
-        AGX1-373 envelope change: tool messages now arrive via
-        streaming_task_message_context (open+close pair) instead of
-        adk.messages.create. The logical content (tool_call_id, name,
-        arguments, author) is identical; only the delivery channel changed.
+        AGX1-377 fix: auto_send now delivers streamed tool-request messages
+        natively (Start+ToolRequestDelta+Done). The streaming context is opened
+        at the Start event with the initial ToolRequestContent (tool_call_id +
+        name + empty arguments), argument tokens are streamed as deltas, and the
+        context is closed on Done.
 
         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.
+        tool-call arguments).
         """
         from pydantic_ai.messages import ToolCallPartDelta
 
+        from agentex.types.tool_request_delta import ToolRequestDelta
+
         streaming, messages = fake_adk
         events = [
             PartStartEvent(
@@ -293,23 +294,27 @@ async def test_tool_call_emits_full_tool_request_message_on_part_end(
         ]
         await stream_pydantic_ai_events(_aiter(events), TASK_ID)
 
-        # AGX1-373: tool messages arrive via streaming_task_message_context,
-        # NOT via adk.messages.create.
-        assert messages.created == [], "adk.messages.create must not be called after reimplementation"
-        assert len(streaming.contexts) == 1, "tool_request opens a streaming context (open+close)"
+        # AGX1-373: tool messages arrive via streaming_task_message_context.
+        assert messages.created == [], "adk.messages.create must not be called"
+        assert len(streaming.contexts) == 1, "tool_request opens a streaming context"
         ctx = streaming.contexts[0]
         assert ctx.closed is True
         content = ctx.initial_content
         assert isinstance(content, ToolRequestContent)
         assert content.tool_call_id == "c1"
         assert content.name == "get_weather"
-        assert content.arguments == {"city": "Paris"}
         assert content.author == "agent"
-        assert ctx.updates == [], "open+close only — no deltas for tool messages"
+        # AGX1-377 streamed shape: initial_content has empty args (args come via delta)
+        assert content.arguments == {}
+        # The arg delta is delivered as a stream_update
+        assert len(ctx.updates) == 1
+        assert isinstance(ctx.updates[0].delta, ToolRequestDelta)
+        assert ctx.updates[0].delta.arguments_delta == '{"city":"Paris"}'
 
     async def test_tool_call_with_dict_args_passes_through(
         self, fake_adk: tuple[FakeStreamingModule, FakeMessagesModule]
     ) -> None:
+        """When args arrive pre-populated as a dict in PartStart, they're in initial_content."""
         streaming, messages = fake_adk
         events = [
             PartStartEvent(
@@ -326,25 +331,26 @@ async def test_tool_call_with_dict_args_passes_through(
         # AGX1-373: tool messages via streaming_task_message_context
         assert messages.created == []
         assert len(streaming.contexts) == 1
+        # Dict args present at PartStart land directly in initial_content.arguments
         assert streaming.contexts[0].initial_content.arguments == {"q": "weather"}
+        assert streaming.contexts[0].updates == [], "no delta for pre-populated dict args"
 
     async def test_tool_call_with_invalid_json_args_surfaces_raw(
         self, fake_adk: tuple[FakeStreamingModule, FakeMessagesModule]
     ) -> None:
-        """Don't drop the tool call when the model emits malformed JSON args.
+        """Malformed JSON arg delta is surfaced as a ToolRequestDelta with the raw string.
 
-        The arguments field is preserved under ``_raw`` so the failure is
-        visible to the UI rather than silently truncated.
-
-        Uses a PartDeltaEvent to deliver the invalid string (the way pydantic-ai
-        actually surfaces arg tokens) so the coalescer picks it up.
+        The argument delta is delivered as-is by auto_send; the client-side
+        accumulator or the streaming backend handles malformed JSON gracefully.
 
         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
 
+        from agentex.types.tool_request_delta import ToolRequestDelta
+
         streaming, messages = fake_adk
         events = [
             PartStartEvent(
@@ -366,7 +372,13 @@ async def test_tool_call_with_invalid_json_args_surfaces_raw(
         # AGX1-373: tool messages via streaming_task_message_context
         assert messages.created == []
         assert len(streaming.contexts) == 1
-        assert streaming.contexts[0].initial_content.arguments == {"_raw": "not-json{"}
+        ctx = streaming.contexts[0]
+        # Initial content has empty args (args come via delta)
+        assert ctx.initial_content.arguments == {}
+        # The malformed JSON is surfaced verbatim in the ToolRequestDelta
+        assert len(ctx.updates) == 1
+        assert isinstance(ctx.updates[0].delta, ToolRequestDelta)
+        assert ctx.updates[0].delta.arguments_delta == "not-json{"
 
     async def test_tool_call_with_none_args_defaults_to_empty_dict(
         self, fake_adk: tuple[FakeStreamingModule, FakeMessagesModule]
@@ -388,6 +400,7 @@ async def test_tool_call_with_none_args_defaults_to_empty_dict(
         assert messages.created == []
         assert len(streaming.contexts) == 1
         assert streaming.contexts[0].initial_content.arguments == {}
+        assert streaming.contexts[0].updates == [], "no delta when args are absent"
 
 
 class TestToolResult:

tests/lib/adk/test_pydantic_ai_turn.py

@@ -230,21 +230,17 @@ async def test_no_usage_event_leaves_default_usage(self):
         assert usage.num_llm_calls == 0
 
 
-class TestToolRequestCoalescing:
-    """The ``coalesce_tool_requests`` flag controls tool-call delivery shape.
+class TestToolRequestStreaming:
+    """PydanticAITurn.events equals the bare converter output unconditionally.
 
-    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).
+    The foundation auto_send delivers Start+ToolRequestDelta+Done natively
+    (AGX1-377), so no coalescing is needed on either channel.
     """
 
-    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."""
+    async def test_events_match_bare_converter_for_streamed_tool_call(self):
+        """PydanticAITurn yields a ToolRequestDelta for a streamed-args tool call
+        — i.e. it is byte-for-byte the bare converter output, preserving
+        argument-token streaming on the sync/yield channel."""
         from pydantic_ai.messages import ToolCallPart, ToolCallPartDelta
 
         from agentex.types.tool_request_delta import ToolRequestDelta
@@ -265,57 +261,16 @@ async def test_default_off_matches_bare_converter_for_streamed_tool_call(self):
 
         bare_out = await _collect(convert_pydantic_ai_to_agentex_events(_aiter(tool_events)))
 
-        # Default Turn is identical to the bare converter.
+        # 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).
+        # The arg-streaming delta is present.
         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 len(deltas) == 1, "streamed tool-call args must surface as a ToolRequestDelta"
         assert isinstance(deltas[0].delta, ToolRequestDelta)
         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"