ADK middleware: prefer LRO routing and harden translator; add tests

contextablemark · contextablemark · commit 769ff0fb5c0c · 2025-10-04T23:11:44.000-07:00
- Prefer LRO routing in ADKAgent when long‑running tool call IDs are present in event.content.parts (prevents misrouting into streaming path and tool loops; preserves HITL pause)
- Force‑close any active streaming text before emitting LRO tool events (guarantees TEXT_MESSAGE_END precedes TOOL_CALL_START)
- Harden EventTranslator.translate to filter out long‑running tool calls from the general path; only emit non‑LRO calls (avoids duplicate tool events)
- Add tests:
  * test_lro_filtering.py (translator‑level filtering + LRO‑only emission)
  * test_integration_mixed_partials.py (streaming → non‑LRO → final LRO: order, no duplicates, correct IDs)
diff --git a/typescript-sdk/integrations/adk-middleware/python/src/ag_ui_adk/adk_agent.py b/typescript-sdk/integrations/adk-middleware/python/src/ag_ui_adk/adk_agent.py
@@ -923,9 +923,24 @@ async def _run_adk_in_background(
                     (not final_response)  # Not marked as final by is_final_response()
                 )
 
-                # Process as streaming if it's a chunk OR if it has content but no finish_reason
-                # This ensures we capture all content, regardless of usage_metadata presence
-                if is_streaming_chunk or (has_content and not getattr(adk_event, 'finish_reason', None)):
+                # Prefer LRO routing when a long-running tool call is present
+                has_lro_function_call = False
+                try:
+                    lro_ids = set(getattr(adk_event, 'long_running_tool_ids', []) or [])
+                    if lro_ids and adk_event.content and getattr(adk_event.content, 'parts', None):
+                        for part in adk_event.content.parts:
+                            func = getattr(part, 'function_call', None)
+                            func_id = getattr(func, 'id', None) if func else None
+                            if func_id and func_id in lro_ids:
+                                has_lro_function_call = True
+                                break
+                except Exception:
+                    # Be conservative: if detection fails, do not block streaming path
+                    has_lro_function_call = False
+
+                # Process as streaming if it's a chunk OR if it has content but no finish_reason,
+                # but only when there is no LRO function call present (LRO takes precedence)
+                if (not has_lro_function_call) and (is_streaming_chunk or (has_content and not getattr(adk_event, 'finish_reason', None))):
                     # Regular translation path
                     async for ag_ui_event in event_translator.translate(
                         adk_event,
@@ -937,7 +952,12 @@ async def _run_adk_in_background(
                         await event_queue.put(ag_ui_event)
                         logger.debug(f"Event queued: {type(ag_ui_event).__name__} (thread {input.thread_id}, queue size after: {event_queue.qsize()})")
                 else:
-                    # LongRunning Tool events are usually emmitted in final response
+                    # LongRunning Tool events are usually emitted in final response
+                    # Ensure any active streaming text message is closed BEFORE tool calls
+                    async for end_event in event_translator.force_close_streaming_message():
+                        await event_queue.put(end_event)
+                        logger.debug(f"Event queued (forced close): {type(end_event).__name__} (thread {input.thread_id}, queue size after: {event_queue.qsize()})")
+
                     async for ag_ui_event in event_translator.translate_lro_function_calls(
                         adk_event
                     ):
@@ -1003,4 +1023,4 @@ async def close(self):
         self._session_lookup_cache.clear()
 
         # Stop session manager cleanup task
-        await self._session_manager.stop_cleanup_task()
+        await self._session_manager.stop_cleanup_task()
diff --git a/typescript-sdk/integrations/adk-middleware/python/src/ag_ui_adk/event_translator.py b/typescript-sdk/integrations/adk-middleware/python/src/ag_ui_adk/event_translator.py
@@ -87,16 +87,24 @@ async def translate(
             if hasattr(adk_event, 'get_function_calls'):               
                 function_calls = adk_event.get_function_calls()
                 if function_calls:
-                    logger.debug(f"ADK function calls detected: {len(function_calls)} calls")
-                    
-                    # CRITICAL FIX: End any active text message stream before starting tool calls
-                    # Per AG-UI protocol: TEXT_MESSAGE_END must be sent before TOOL_CALL_START
-                    async for event in self.force_close_streaming_message():
-                        yield event
-                    
-                    # NOW ACTUALLY YIELD THE EVENTS
-                    async for event in self._translate_function_calls(function_calls):
-                        yield event
+                    # Filter out long-running tool calls; those are handled by translate_lro_function_calls
+                    try:
+                        lro_ids = set(getattr(adk_event, 'long_running_tool_ids', []) or [])
+                    except Exception:
+                        lro_ids = set()
+
+                    non_lro_calls = [fc for fc in function_calls if getattr(fc, 'id', None) not in lro_ids]
+
+                    if non_lro_calls:
+                        logger.debug(f"ADK function calls detected (non-LRO): {len(non_lro_calls)} of {len(function_calls)} total")
+                        # CRITICAL FIX: End any active text message stream before starting tool calls
+                        # Per AG-UI protocol: TEXT_MESSAGE_END must be sent before TOOL_CALL_START
+                        async for event in self.force_close_streaming_message():
+                            yield event
+                        
+                        # Yield only non-LRO function call events
+                        async for event in self._translate_function_calls(non_lro_calls):
+                            yield event
                         
             # Handle function responses and yield the tool response event
             # this is essential for scenerios when user has to render function response at frontend
@@ -466,4 +474,4 @@ def reset(self):
         self._streaming_message_id = None
         self._is_streaming = False
         self.long_running_tool_ids.clear()
-        logger.debug("Reset EventTranslator state (including streaming state)")
+        logger.debug("Reset EventTranslator state (including streaming state)")
diff --git a/typescript-sdk/integrations/adk-middleware/python/tests/test_integration_mixed_partials.py b/typescript-sdk/integrations/adk-middleware/python/tests/test_integration_mixed_partials.py
@@ -0,0 +1,111 @@
+#!/usr/bin/env python
+"""Integration test: mixed partials with non-LRO calls before final LRO.
+
+Scenario:
+- Stream text in partial chunks
+- Mid-stream, a non-LRO function call appears (should close text and emit tool events)
+- Finally, an LRO function call arrives (should close any open text and emit LRO tool events)
+
+Asserts order, deduplication, and correct tool ids.
+"""
+
+import pytest
+from unittest.mock import MagicMock, AsyncMock, Mock, patch
+
+from ag_ui.core import (
+    RunAgentInput, UserMessage
+)
+from ag_ui_adk import ADKAgent
+
+
+@pytest.fixture
+def adk_agent_instance():
+    from google.adk.agents import Agent
+    mock_agent = Mock(spec=Agent)
+    mock_agent.name = "test_agent"
+    return ADKAgent(adk_agent=mock_agent, app_name="test_app", user_id="test_user")
+
+
+@pytest.mark.asyncio
+async def test_mixed_partials_non_lro_then_lro(adk_agent_instance):
+    # Helper to create partial text events
+    def mk_partial(text):
+        e = MagicMock()
+        e.author = "assistant"
+        e.content = MagicMock(); e.content.parts = [MagicMock(text=text)]
+        e.partial = True
+        e.turn_complete = False
+        e.is_final_response = lambda: False
+        # No function responses in these partials
+        e.get_function_responses = lambda: []
+        e.get_function_calls = lambda: []
+        return e
+
+    # First partial text only
+    evt1 = mk_partial("Hello")
+
+    # Second partial: text + non-LRO function call
+    normal_id = "normal-999"
+    normal_func = MagicMock(); normal_func.id = normal_id; normal_func.name = "regular_tool"; normal_func.args = {"b": 2}
+    evt2 = mk_partial(" world")
+    evt2.get_function_calls = lambda: [normal_func]
+    evt2.long_running_tool_ids = []
+
+    # Final: LRO function call
+    lro_id = "lro-777"
+    lro_func = MagicMock(); lro_func.id = lro_id; lro_func.name = "long_running_tool"; lro_func.args = {"v": 1}
+    lro_part = MagicMock(); lro_part.function_call = lro_func
+
+    evt3 = MagicMock()
+    evt3.author = "assistant"
+    evt3.content = MagicMock(); evt3.content.parts = [lro_part]
+    evt3.partial = False
+    evt3.turn_complete = True
+    evt3.is_final_response = lambda: True
+    evt3.get_function_calls = lambda: []
+    evt3.get_function_responses = lambda: []
+    evt3.long_running_tool_ids = [lro_id]
+
+    async def mock_run_async(*args, **kwargs):
+        yield evt1
+        yield evt2
+        yield evt3
+
+    mock_runner = AsyncMock(); mock_runner.run_async = mock_run_async
+
+    sample_input = RunAgentInput(
+        thread_id="thread_mixed",
+        run_id="run_mixed",
+        messages=[UserMessage(id="u1", role="user", content="go")],
+        tools=[], context=[], state={}, forwarded_props={},
+    )
+
+    with patch.object(adk_agent_instance, "_create_runner", return_value=mock_runner):
+        events = []
+        async for e in adk_agent_instance.run(sample_input):
+            events.append(e)
+
+    types = [str(ev.type).split(".")[-1] for ev in events]
+
+    # Expect at least one START and 2 CONTENTs from streaming
+    assert types.count("TEXT_MESSAGE_START") == 1
+    assert types.count("TEXT_MESSAGE_CONTENT") >= 2
+
+    # Non-LRO tool call should appear exactly once
+    normal_starts = [i for i, ev in enumerate(events) if str(ev.type).endswith("TOOL_CALL_START") and getattr(ev, "tool_call_id", None) == normal_id]
+    normal_args = [i for i, ev in enumerate(events) if str(ev.type).endswith("TOOL_CALL_ARGS") and getattr(ev, "tool_call_id", None) == normal_id]
+    normal_ends = [i for i, ev in enumerate(events) if str(ev.type).endswith("TOOL_CALL_END") and getattr(ev, "tool_call_id", None) == normal_id]
+    assert len(normal_starts) == len(normal_args) == len(normal_ends) == 1
+
+    # Ensure a TEXT_MESSAGE_END precedes the normal tool start
+    text_ends = [i for i, t in enumerate(types) if t == "TEXT_MESSAGE_END"]
+    assert len(text_ends) >= 1
+    assert text_ends[-1] < normal_starts[0], "TEXT_MESSAGE_END must precede first non-LRO TOOL_CALL_START"
+
+    # LRO tool call should appear exactly once and after the non-LRO
+    lro_starts = [i for i, ev in enumerate(events) if str(ev.type).endswith("TOOL_CALL_START") and getattr(ev, "tool_call_id", None) == lro_id]
+    lro_args = [i for i, ev in enumerate(events) if str(ev.type).endswith("TOOL_CALL_ARGS") and getattr(ev, "tool_call_id", None) == lro_id]
+    lro_ends = [i for i, ev in enumerate(events) if str(ev.type).endswith("TOOL_CALL_END") and getattr(ev, "tool_call_id", None) == lro_id]
+    assert len(lro_starts) == len(lro_args) == len(lro_ends) == 1
+    assert lro_starts[0] > normal_starts[0]
+
diff --git a/typescript-sdk/integrations/adk-middleware/python/tests/test_lro_filtering.py b/typescript-sdk/integrations/adk-middleware/python/tests/test_lro_filtering.py
@@ -0,0 +1,106 @@
+#!/usr/bin/env python
+"""Tests for LRO-aware routing and translator filtering.
+
+These tests verify that:
+- EventTranslator.translate skips long-running tool calls and only emits non-LRO calls
+- translate_lro_function_calls emits events only for long-running tool calls
+"""
+
+import asyncio
+from unittest.mock import MagicMock
+
+from ag_ui.core import EventType
+from ag_ui_adk import EventTranslator
+
+
+async def test_translate_skips_lro_function_calls():
+    """Ensure non-LRO tool calls are emitted and LRO calls are skipped in translate."""
+    translator = EventTranslator()
+
+    # Prepare mock ADK event
+    adk_event = MagicMock()
+    adk_event.author = "assistant"
+    adk_event.content = MagicMock()
+    adk_event.content.parts = []  # no text
+
+    # Two function calls, one is long-running
+    lro_id = "tool-call-lro-1"
+    normal_id = "tool-call-normal-2"
+
+    lro_call = MagicMock()
+    lro_call.id = lro_id
+    lro_call.name = "long_running_tool"
+    lro_call.args = {"x": 1}
+
+    normal_call = MagicMock()
+    normal_call.id = normal_id
+    normal_call.name = "regular_tool"
+    normal_call.args = {"y": 2}
+
+    adk_event.get_function_calls = lambda: [lro_call, normal_call]
+    # Mark the long-running call id on the event
+    adk_event.long_running_tool_ids = [lro_id]
+
+    events = []
+    async for e in translator.translate(adk_event, "thread", "run"):
+        events.append(e)
+
+    # We expect only the non-LRO tool call events to be emitted
+    # Sequence: TOOL_CALL_START(normal), TOOL_CALL_ARGS(normal), TOOL_CALL_END(normal)
+    event_types = [str(ev.type).split('.')[-1] for ev in events]
+    assert event_types.count("TOOL_CALL_START") == 1
+    assert event_types.count("TOOL_CALL_ARGS") == 1
+    assert event_types.count("TOOL_CALL_END") == 1
+
+    # Ensure the emitted tool_call_id is the normal one
+    ids = set(getattr(ev, 'tool_call_id', None) for ev in events)
+    assert normal_id in ids
+    assert lro_id not in ids
+
+
+async def test_translate_lro_function_calls_only_emits_lro():
+    """Ensure translate_lro_function_calls emits only for long-running calls."""
+    translator = EventTranslator()
+
+    # Prepare mock ADK event with content parts containing function calls
+    lro_id = "tool-call-lro-3"
+    normal_id = "tool-call-normal-4"
+
+    lro_call = MagicMock()
+    lro_call.id = lro_id
+    lro_call.name = "long_running_tool"
+    lro_call.args = {"a": 123}
+
+    normal_call = MagicMock()
+    normal_call.id = normal_id
+    normal_call.name = "regular_tool"
+    normal_call.args = {"b": 456}
+
+    # Build parts with both calls
+    lro_part = MagicMock()
+    lro_part.function_call = lro_call
+    normal_part = MagicMock()
+    normal_part.function_call = normal_call
+
+    adk_event = MagicMock()
+    adk_event.content = MagicMock()
+    adk_event.content.parts = [lro_part, normal_part]
+    adk_event.long_running_tool_ids = [lro_id]
+
+    events = []
+    async for e in translator.translate_lro_function_calls(adk_event):
+        events.append(e)
+
+    # Expect only the LRO call events
+    # Sequence: TOOL_CALL_START(lro), TOOL_CALL_ARGS(lro), TOOL_CALL_END(lro)
+    event_types = [str(ev.type).split('.')[-1] for ev in events]
+    assert event_types == ["TOOL_CALL_START", "TOOL_CALL_ARGS", "TOOL_CALL_END"]
+    for ev in events:
+        assert getattr(ev, 'tool_call_id', None) == lro_id
+
+
+if __name__ == "__main__":
+    asyncio.run(test_translate_skips_lro_function_calls())
+    asyncio.run(test_translate_lro_function_calls_only_emits_lro())
+    print("\n✅ LRO filtering tests ran to completion")
+
