Merge pull request #49 from evgeny-l/adk-middleware-before-agent-callback

fix: support of before_agent_callback in case of a direct content response

Author: contextablemark

typescript-sdk/integrations/adk-middleware/src/adk_middleware/adk_agent.py

@@ -858,8 +858,12 @@ async def _run_adk_in_background(
                 new_message=new_message,
                 run_config=run_config
             ):
-                if not adk_event.is_final_response():
-                # Translate and emit events
+
+                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
                     async for ag_ui_event in event_translator.translate(
                         adk_event,
                         input.thread_id,

typescript-sdk/integrations/adk-middleware/src/adk_middleware/event_translator.py

@@ -168,11 +168,36 @@ async def _translate_text_content(
         should_send_end = is_final_response and not is_partial
 
         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}")
-        
-        # Skip final response events to avoid duplicate content, but send END if streaming
+                    f"is_final_response={is_final_response}, should_send_end={should_send_end}, "
+                    f"currently_streaming={self._is_streaming}")
+
         if is_final_response:
+
+            # If a final text response wasn't streamed (not generated by an LLM) then deliver it in 3 events
+            if not self._is_streaming and not adk_event.usage_metadata and should_send_end:
+                logger.info(f"⏭️ Deliver non-llm response via message events "
+                            f"event_id={adk_event.id}")
+
+                combined_text = "".join(text_parts)
+                message_events = [
+                    TextMessageStartEvent(
+                        type=EventType.TEXT_MESSAGE_START,
+                        message_id=adk_event.id,
+                        role="assistant"
+                    ),
+                    TextMessageContentEvent(
+                        type=EventType.TEXT_MESSAGE_CONTENT,
+                        message_id=adk_event.id,
+                        delta=combined_text
+                    ),
+                    TextMessageEndEvent(
+                        type=EventType.TEXT_MESSAGE_END,
+                        message_id=adk_event.id
+                    )
+                ]
+                for msg in message_events:
+                    yield msg
+
             logger.info("⏭️ Skipping final response event (content already streamed)")
 
             # If we're currently streaming, this final response means we should end the stream

typescript-sdk/integrations/adk-middleware/tests/test_event_translator_comprehensive.py

@@ -50,6 +50,7 @@ def mock_adk_event_with_content(self):
         event.partial = False
         event.turn_complete = True
         event.is_final_response = False
+        event.usage_metadata = {'tokens': 22}
         return event
 
     @pytest.mark.asyncio
@@ -275,7 +276,26 @@ async def test_translate_text_content_final_response_no_streaming(self, translat
             events.append(event)
 
         assert len(events) == 0  # No events
-    
+
+    @pytest.mark.asyncio
+    async def test_translate_text_content_final_response_from_agent_callback(self, translator, mock_adk_event_with_content):
+        """Test final response when it was received from an agent callback function."""
+        mock_adk_event_with_content.is_final_response = True
+        mock_adk_event_with_content.usage_metadata = None
+
+        # Not streaming
+        translator._is_streaming = False
+
+        events = []
+        async for event in translator.translate(mock_adk_event_with_content, "thread_1", "run_1"):
+            events.append(event)
+
+        assert len(events) == 3  # START, CONTENT , END
+        assert isinstance(events[0], TextMessageStartEvent)
+        assert isinstance(events[1], TextMessageContentEvent)
+        assert events[1].delta == mock_adk_event_with_content.content.parts[0].text
+        assert isinstance(events[2], TextMessageEndEvent)
+
     @pytest.mark.asyncio
     async def test_translate_text_content_empty_text(self, translator, mock_adk_event):
         """Test text content with empty text."""

typescript-sdk/integrations/adk-middleware/tests/test_text_events.py

@@ -10,6 +10,7 @@
 from ag_ui.core import RunAgentInput, UserMessage
 from adk_middleware import ADKAgent
 from google.adk.agents import Agent
+from google.genai import types
 
 
 async def test_message_events():
@@ -87,6 +88,142 @@ async def test_message_events():
     return validate_message_event_pattern(start_count, end_count, content_count, text_message_events)
 
 
+async def test_message_events_from_before_agent_callback():
+    """Test that we get proper message events with correct START/CONTENT/END patterns,
+    even if we return the message from before_agent_callback.
+    """
+
+    if not os.getenv("GOOGLE_API_KEY"):
+        print("⚠️ GOOGLE_API_KEY not set - using mock test")
+        return await test_with_mock()
+
+    print("🧪 Testing with real Google ADK agent...")
+
+    event_message = "This message was not generated."
+    def return_predefined_message(callback_context):
+        return types.Content(
+            parts=[types.Part(text=event_message)],
+            role="model"  # Assign model role to the overriding response
+        )
+
+    # Create real agent
+    agent = Agent(
+        name="test_agent",
+        instruction="You are a helpful assistant. Keep responses brief.",
+        before_agent_callback=return_predefined_message
+    )
+
+    # Create middleware with direct agent embedding
+    adk_agent = ADKAgent(
+        adk_agent=agent,
+        app_name="test_app",
+        user_id="test_user",
+        use_in_memory_services=True,
+    )
+
+    # Test input
+    test_input = RunAgentInput(
+        thread_id="test_thread",
+        run_id="test_run",
+        messages=[
+            UserMessage(
+                id="msg_1",
+                role="user",
+                content="Say hello in exactly 3 words."
+            )
+        ],
+        state={},
+        context=[],
+        tools=[],
+        forwarded_props={}
+    )
+
+    print("🚀 Running test request...")
+
+    events = []
+    text_message_events = []
+
+    try:
+        async for event in adk_agent.run(test_input):
+            events.append(event)
+            event_type = str(event.type)
+            print(f"📧 {event_type}")
+
+            # Track text message events specifically
+            if "TEXT_MESSAGE" in event_type:
+                text_message_events.append(event_type)
+
+    except Exception as e:
+        print(f"❌ Error during test: {e}")
+        return False
+
+    print(f"\n📊 Results:")
+    print(f"   Total events: {len(events)}")
+    print(f"   Text message events: {text_message_events}")
+
+    # Analyze message event patterns
+    start_count = text_message_events.count("EventType.TEXT_MESSAGE_START")
+    end_count = text_message_events.count("EventType.TEXT_MESSAGE_END")
+    content_count = text_message_events.count("EventType.TEXT_MESSAGE_CONTENT")
+
+    print(f"   START events: {start_count}")
+    print(f"   END events: {end_count}")
+    print(f"   CONTENT events: {content_count}")
+
+    pattern_is_valid = validate_message_event_pattern(start_count, end_count, content_count, text_message_events)
+    if not pattern_is_valid:
+        return False
+
+    expected_text_events = [
+        {
+            "type": "EventType.TEXT_MESSAGE_START",
+        },
+        {
+            "type": "EventType.TEXT_MESSAGE_CONTENT",
+            "delta": event_message
+        },
+        {
+            "type": "EventType.TEXT_MESSAGE_END",
+        }
+    ]
+    return validate_message_events(events, expected_text_events)
+
+
+def validate_message_events(events, expected_events):
+    """Compare expected events by type and delta (if delta exists)."""
+    # Filter events to only those specified in expected_events
+    event_types_to_check = {expected["type"] for expected in expected_events}
+    
+    filtered_events = []
+    for event in events:
+        event_type_str = f"EventType.{event.type.value}"
+        if event_type_str in event_types_to_check:
+            filtered_events.append(event)
+    
+    if len(filtered_events) != len(expected_events):
+        print(f"❌ Event count mismatch: expected {len(expected_events)}, got {len(filtered_events)}")
+        return False
+    
+    for i, (event, expected) in enumerate(zip(filtered_events, expected_events)):
+        # Check event type
+        event_type_str = f"EventType.{event.type.value}"
+        if event_type_str != expected["type"]:
+            print(f"❌ Event {i}: type mismatch - expected {expected['type']}, got {event_type_str}")
+            return False
+        
+        # Check delta if specified
+        if "delta" in expected:
+            if not hasattr(event, 'delta'):
+                print(f"❌ Event {i}: expected delta field but event has none")
+                return False
+            if event.delta != expected["delta"]:
+                print(f"❌ Event {i}: delta mismatch - expected '{expected['delta']}', got '{event.delta}'")
+                return False
+    
+    print("✅ All expected events validated successfully")
+    return True
+
+
 def validate_message_event_pattern(start_count, end_count, content_count, text_message_events):
     """Validate that message events follow proper patterns."""
 
@@ -319,6 +456,13 @@ async def test_text_message_events():
     assert result, "Text message events test failed"
 
 
+@pytest.mark.asyncio
+async def test_text_message_events_from_before_agent_callback():
+    """Test that we get proper message events with correct START/CONTENT/END patterns."""
+    result = await test_message_events_from_before_agent_callback()
+    assert result, "Text message events for before_agent_callback test failed"
+
+
 @pytest.mark.asyncio 
 async def test_message_event_edge_cases():
     """Test edge cases for message event patterns."""