ag-ui-protocol
diff --git a/‎typescript-sdk/integrations/adk-middleware/python/src/ag_ui_adk/__init__.py‎
Lines changed: 54 additions & 1 deletion b/‎typescript-sdk/integrations/adk-middleware/python/src/ag_ui_adk/__init__.py‎
Lines changed: 54 additions & 1 deletion
diff --git a/‎typescript-sdk/integrations/adk-middleware/python/src/ag_ui_adk/adk_agent.py‎
Lines changed: 34 additions & 5 deletions b/‎typescript-sdk/integrations/adk-middleware/python/src/ag_ui_adk/adk_agent.py‎
Lines changed: 34 additions & 5 deletions
diff --git a/‎typescript-sdk/integrations/adk-middleware/python/src/ag_ui_adk/event_translator.py‎
Lines changed: 29 additions & 16 deletions b/‎typescript-sdk/integrations/adk-middleware/python/src/ag_ui_adk/event_translator.py‎
Lines changed: 29 additions & 16 deletions
diff --git a/‎typescript-sdk/integrations/adk-middleware/python/tests/conftest.py‎
Lines changed: 20 additions & 0 deletions b/‎typescript-sdk/integrations/adk-middleware/python/tests/conftest.py‎
Lines changed: 20 additions & 0 deletions
@@ -5,11 +5,64 @@
 This middleware enables Google ADK agents to be used with the AG-UI protocol.
 """
 
+from __future__ import annotations
+
+import logging
+import os
+from typing import Dict, Iterable
+
 from .adk_agent import ADKAgent
 from .event_translator import EventTranslator
 from .session_manager import SessionManager
 from .endpoint import add_adk_fastapi_endpoint, create_adk_app
 
 __all__ = ['ADKAgent', 'add_adk_fastapi_endpoint', 'create_adk_app', 'EventTranslator', 'SessionManager']
 
-__version__ = "0.1.0"
+__version__ = "0.1.0"
+
+
+def _configure_logging_from_env() -> None:
+    """Configure component loggers based on environment variables."""
+
+    root_level = os.getenv('LOG_ROOT_LEVEL')
+    if root_level:
+        try:
+            level = getattr(logging, root_level.upper())
+        except AttributeError:
+            logging.getLogger(__name__).warning(
+                "Invalid LOG_ROOT_LEVEL value '%s'", root_level
+            )
+        else:
+            logging.basicConfig(level=level, force=True)
+
+    component_levels: Dict[str, Iterable[str]] = {
+        'LOG_ADK_AGENT': ('ag_ui_adk.adk_agent',),
+        'LOG_EVENT_TRANSLATOR': (
+            'ag_ui_adk.event_translator',
+            'event_translator',
+        ),
+        'LOG_ENDPOINT': ('ag_ui_adk.endpoint', 'endpoint'),
+        'LOG_SESSION_MANAGER': (
+            'ag_ui_adk.session_manager',
+            'session_manager',
+        ),
+    }
+
+    for env_var, logger_names in component_levels.items():
+        level_name = os.getenv(env_var)
+        if not level_name:
+            continue
+
+        try:
+            level = getattr(logging, level_name.upper())
+        except AttributeError:
+            logging.getLogger(__name__).warning(
+                "Invalid value '%s' for %s", level_name, env_var
+            )
+            continue
+
+        for logger_name in logger_names:
+            logging.getLogger(logger_name).setLevel(level)
+
+
+_configure_logging_from_env()
@@ -916,19 +916,48 @@ async def _run_adk_in_background(
                 final_response = adk_event.is_final_response()
                 has_content = adk_event.content and hasattr(adk_event.content, 'parts') and adk_event.content.parts
 
-                if not final_response or (not adk_event.usage_metadata and has_content):
-                    # Translate and emit events
+                # Check if this is a streaming chunk that needs regular processing
+                is_streaming_chunk = (
+                    getattr(adk_event, 'partial', False) or  # Explicitly marked as partial
+                    (not getattr(adk_event, 'turn_complete', True)) or  # Live streaming not complete
+                    (not final_response)  # Not marked as final by is_final_response()
+                )
+
+                # 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,
                         input.thread_id,
                         input.run_id
                     ):
-                        
+
                         logger.debug(f"Emitting event to queue: {type(ag_ui_event).__name__} (thread {input.thread_id}, queue size before: {event_queue.qsize()})")
                         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
                     ):
@@ -994,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()
@@ -189,16 +189,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
@@ -266,12 +274,17 @@ async def _translate_text_content(
         elif hasattr(adk_event, 'is_final_response'):
             is_final_response = adk_event.is_final_response
 
-        # Handle None values: if is_final_response=True, it means streaming should end
-        should_send_end = is_final_response and not is_partial
-        
+        # Handle None values: if a turn is complete or a final chunk arrives, end streaming
+        has_finish_reason = bool(getattr(adk_event, 'finish_reason', None))
+        should_send_end = (
+            (turn_complete and not is_partial)
+            or (is_final_response and not is_partial)
+            or (has_finish_reason and self._is_streaming)
+        )
+
         logger.info(f"📥 Text event - partial={is_partial}, turn_complete={turn_complete}, "
-                    f"is_final_response={is_final_response}, should_send_end={should_send_end}, "
-                    f"currently_streaming={self._is_streaming}")
+                    f"is_final_response={is_final_response}, has_finish_reason={has_finish_reason}, "
+                    f"should_send_end={should_send_end}, currently_streaming={self._is_streaming}")
 
         if is_final_response:
 
@@ -566,4 +579,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)")
@@ -0,0 +1,20 @@
+"""Shared pytest fixtures for ADK middleware tests."""
+
+from __future__ import annotations
+
+import pytest
+
+from ag_ui.core import SystemMessage as CoreSystemMessage
+
+import ag_ui_adk.adk_agent as adk_agent_module
+
+
+@pytest.fixture(autouse=True)
+def restore_system_message_class():
+    """Ensure every test starts and ends with the real SystemMessage type."""
+
+    adk_agent_module.SystemMessage = CoreSystemMessage
+    try:
+        yield
+    finally:
+        adk_agent_module.SystemMessage = CoreSystemMessage