ADK Middleware Integration

typescript-sdk/integrations/adk-middleware/CHANGELOG.md

@@ -7,6 +7,20 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+- **NEW**: SystemMessage support for ADK agents (issue #22) - SystemMessages as first message are now appended to agent instructions
+- **NEW**: Comprehensive tests for SystemMessage functionality including edge cases
+
+### Fixed  
+- **FIXED**: Race condition in tool result processing causing "No pending tool calls found" warnings
+- **FIXED**: Tool call removal now happens after pending check to prevent race conditions
+- **IMPROVED**: Better handling of empty tool result content with graceful JSON parsing fallback
+
+### Enhanced
+- **LOGGING**: Added debug logging for tool result processing to aid in troubleshooting
+- **ARCHITECTURE**: Consolidated agent copying logic to avoid creating multiple unnecessary copies
+- **CLEANUP**: Removed unused toolset parameter from `_run_adk_in_background` method
+
 ## [0.4.1] - 2025-07-13
 
 ### Fixed

typescript-sdk/integrations/adk-middleware/examples/fastapi_server.py

@@ -7,11 +7,15 @@
 """
 
 import uvicorn
+import logging
 from fastapi import FastAPI
 from .tool_based_generative_ui.agent import haiku_generator_agent
 from .human_in_the_loop.agent import human_in_loop_agent
 from .shared_state.agent import shared_state_agent
 
+# Basic logging configuration
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+
 # These imports will work once google.adk is available
 try:
     # from src.adk_agent import ADKAgent

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

@@ -13,7 +13,7 @@
     RunStartedEvent, RunFinishedEvent, RunErrorEvent,
     TextMessageStartEvent, TextMessageContentEvent, TextMessageEndEvent,
     StateSnapshotEvent, StateDeltaEvent,
-    Context, ToolMessage, ToolCallEndEvent
+    Context, ToolMessage, ToolCallEndEvent, SystemMessage
 )
 
 from google.adk import Runner
@@ -427,6 +427,8 @@ async def _handle_tool_result_submission(
                     # Could add more specific check here for the exact tool_call_id
                     # but for now just log that we're processing a tool result while tools are pending
                     logger.debug(f"Processing tool result {tool_call_id} for thread {thread_id} with pending tools")
+                    # Remove from pending tool calls now that we're processing it
+                    await self._remove_pending_tool_call(thread_id, tool_call_id)
                 else:
                     # No pending tools - this could be a stale result or from a different session
                     logger.warning(f"No pending tool calls found for tool result {tool_call_id} in thread {thread_id}")
@@ -477,9 +479,6 @@ async def _extract_tool_results(self, input: RunAgentInput) -> List[Dict]:
             # Debug: Log the extracted tool message
             logger.info(f"Extracted most recent ToolMessage: role={most_recent_tool_message.role}, tool_call_id={most_recent_tool_message.tool_call_id}, content='{most_recent_tool_message.content}'")
 
-            # Remove from pending tool calls when response is received
-            await self._remove_pending_tool_call(input.thread_id, most_recent_tool_message.tool_call_id)
-
             return [{
                 'tool_name': tool_name,
                 'message': most_recent_tool_message
@@ -692,13 +691,47 @@ async def _start_background_execution(
         registry = AgentRegistry.get_instance()
         adk_agent = registry.get_agent(agent_id)
 
-        # Create dynamic toolset if tools provided
+        # Prepare agent modifications (SystemMessage and tools)
+        agent_updates = {}
+
+        # Handle SystemMessage if it's the first message - append to agent instructions
+        if input.messages and isinstance(input.messages[0], SystemMessage):
+            system_content = input.messages[0].content
+            if system_content:
+                # Get existing instruction (may be None or empty)
+                current_instruction = getattr(adk_agent, 'instruction', '') or ''
+
+                # Append SystemMessage content to existing instructions
+                if current_instruction:
+                    new_instruction = f"{current_instruction}\n\n{system_content}"
+                else:
+                    new_instruction = system_content
+
+                agent_updates['instruction'] = new_instruction
+                logger.debug(f"Will append SystemMessage to agent instructions: '{system_content[:100]}...'")
+
+        # Create dynamic toolset if tools provided and prepare tool updates
         toolset = None
         if input.tools:
             toolset = ClientProxyToolset(
                 ag_ui_tools=input.tools,
                 event_queue=event_queue
             )
+
+            # Get existing tools from the agent
+            existing_tools = []
+            if hasattr(adk_agent, 'tools') and adk_agent.tools:
+                existing_tools = list(adk_agent.tools) if isinstance(adk_agent.tools, (list, tuple)) else [adk_agent.tools]
+
+            # Combine existing tools with our proxy toolset
+            combined_tools = existing_tools + [toolset]
+            agent_updates['tools'] = combined_tools
+            logger.debug(f"Will combine {len(existing_tools)} existing tools with proxy toolset")
+
+        # Create a single copy of the agent with all updates if any modifications needed
+        if agent_updates:
+            adk_agent = adk_agent.model_copy(update=agent_updates)
+            logger.debug(f"Created modified agent copy with updates: {list(agent_updates.keys())}")
 
         # Create background task
         logger.debug(f"Creating background task for thread {input.thread_id}")
@@ -708,7 +741,6 @@ async def _start_background_execution(
                 adk_agent=adk_agent,
                 user_id=user_id,
                 app_name=app_name,
-                toolset=toolset,
                 event_queue=event_queue
             )
         )
@@ -726,34 +758,20 @@ async def _run_adk_in_background(
         adk_agent: ADKBaseAgent,
         user_id: str,
         app_name: str,
-        toolset: Optional[ClientProxyToolset],
         event_queue: asyncio.Queue
     ):
         """Run ADK agent in background, emitting events to queue.
 
         Args:
             input: The run input
-            adk_agent: The ADK agent to run
+            adk_agent: The ADK agent to run (already prepared with tools and SystemMessage)
             user_id: User ID
             app_name: App name
-            toolset: Optional client proxy toolset
             event_queue: Queue for emitting events
         """
         try:
-            # Handle tool combination if toolset provided - use agent cloning to avoid mutating original
-            if toolset:
-                # Get existing tools from the agent
-                existing_tools = []
-                if hasattr(adk_agent, 'tools') and adk_agent.tools:
-                    existing_tools = list(adk_agent.tools) if isinstance(adk_agent.tools, (list, tuple)) else [adk_agent.tools]
-
-                # Combine existing tools with our proxy toolset
-                combined_tools = existing_tools + [toolset]
-
-                # Create a copy of the agent with the combined tools (avoid mutating original)
-                adk_agent = adk_agent.model_copy(update={'tools': combined_tools})
-
-                logger.debug(f"Combined {len(existing_tools)} existing tools with proxy toolset via agent cloning")
+            # Agent is already prepared with tools and SystemMessage instructions (if any)
+            # from _start_background_execution, so no additional agent copying needed here
 
             # Create runner
             runner = self._create_runner(
@@ -867,9 +885,10 @@ async def _run_adk_in_background(
             )
             await event_queue.put(None)
         finally:
-            # Clean up toolset
-            if toolset:
-                await toolset.close()
+            # Background task cleanup completed
+            # Note: toolset cleanup is handled by garbage collection
+            # since toolset is now embedded in the agent's tools
+            pass
 
     async def _cleanup_stale_executions(self):
         """Clean up stale executions."""

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

@@ -301,11 +301,11 @@ def _create_state_delta_event(
             A StateDeltaEvent
         """
         # Convert to JSON Patch format (RFC 6902)
-        # For now, we'll use a simple "replace" operation for each key
+        # Use "add" operation which works for both new and existing paths
         patches = []
         for key, value in state_delta.items():
             patches.append({
-                "op": "replace",
+                "op": "add",
                 "path": f"/{key}",
                 "value": value
             })

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

@@ -10,7 +10,7 @@
 from adk_middleware import ADKAgent, AgentRegistry,SessionManager
 from ag_ui.core import (
     RunAgentInput, EventType, UserMessage, Context,
-    RunStartedEvent, RunFinishedEvent, TextMessageChunkEvent
+    RunStartedEvent, RunFinishedEvent, TextMessageChunkEvent, SystemMessage
 )
 from google.adk.agents import Agent
 
@@ -192,6 +192,128 @@ async def test_cleanup(self, adk_agent):
         mock_execution.cancel.assert_called_once()
         assert len(adk_agent._active_executions) == 0
 
+    @pytest.mark.asyncio
+    async def test_system_message_appended_to_instructions(self, registry):
+        """Test that SystemMessage as first message gets appended to agent instructions."""
+        # Create an agent with initial instructions
+        mock_agent = Agent(
+            name="test_agent",
+            instruction="You are a helpful assistant."
+        )
+        registry.set_default_agent(mock_agent)
+
+        adk_agent = ADKAgent(app_name="test_app", user_id="test_user")
+
+        # Create input with SystemMessage as first message
+        system_input = RunAgentInput(
+            thread_id="test_thread",
+            run_id="test_run",
+            messages=[
+                SystemMessage(id="sys_1", role="system", content="Be very concise in responses."),
+                UserMessage(id="msg_1", role="user", content="Hello")
+            ],
+            context=[],
+            state={},
+            tools=[],
+            forwarded_props={}
+        )
+
+        # Mock the background execution to capture the modified agent
+        captured_agent = None
+        original_run_background = adk_agent._run_adk_in_background
+
+        async def mock_run_background(input, adk_agent, user_id, app_name, event_queue):
+            nonlocal captured_agent
+            captured_agent = adk_agent
+            # Just put a completion event in the queue and return
+            await event_queue.put(None)
+
+        with patch.object(adk_agent, '_run_adk_in_background', side_effect=mock_run_background):
+            # Start execution to trigger agent modification
+            execution = await adk_agent._start_background_execution(system_input, "default")
+
+            # Wait briefly for the background task to start
+            await asyncio.sleep(0.01)
+
+        # Verify the agent's instruction was modified
+        assert captured_agent is not None
+        expected_instruction = "You are a helpful assistant.\n\nBe very concise in responses."
+        assert captured_agent.instruction == expected_instruction
+
+    @pytest.mark.asyncio
+    async def test_system_message_not_first_ignored(self, registry):
+        """Test that SystemMessage not as first message is ignored."""
+        mock_agent = Agent(
+            name="test_agent", 
+            instruction="You are a helpful assistant."
+        )
+        registry.set_default_agent(mock_agent)
+
+        adk_agent = ADKAgent(app_name="test_app", user_id="test_user")
+
+        # Create input with SystemMessage as second message
+        system_input = RunAgentInput(
+            thread_id="test_thread",
+            run_id="test_run", 
+            messages=[
+                UserMessage(id="msg_1", role="user", content="Hello"),
+                SystemMessage(id="sys_1", role="system", content="Be very concise in responses.")
+            ],
+            context=[],
+            state={},
+            tools=[],
+            forwarded_props={}
+        )
+
+        # Mock the background execution to capture the agent
+        captured_agent = None
+
+        async def mock_run_background(input, adk_agent, user_id, app_name, event_queue):
+            nonlocal captured_agent
+            captured_agent = adk_agent
+            await event_queue.put(None)
+
+        with patch.object(adk_agent, '_run_adk_in_background', side_effect=mock_run_background):
+            execution = await adk_agent._start_background_execution(system_input, "default")
+            await asyncio.sleep(0.01)
+
+        # Verify the agent's instruction was NOT modified
+        assert captured_agent.instruction == "You are a helpful assistant."
+
+    @pytest.mark.asyncio
+    async def test_system_message_with_no_existing_instruction(self, registry):
+        """Test SystemMessage handling when agent has no existing instruction."""
+        mock_agent = Agent(name="test_agent")  # No instruction
+        registry.set_default_agent(mock_agent)
+
+        adk_agent = ADKAgent(app_name="test_app", user_id="test_user")
+
+        system_input = RunAgentInput(
+            thread_id="test_thread",
+            run_id="test_run",
+            messages=[
+                SystemMessage(id="sys_1", role="system", content="You are a math tutor.")
+            ],
+            context=[],
+            state={},
+            tools=[],
+            forwarded_props={}
+        )
+
+        captured_agent = None
+
+        async def mock_run_background(input, adk_agent, user_id, app_name, event_queue):
+            nonlocal captured_agent
+            captured_agent = adk_agent
+            await event_queue.put(None)
+
+        with patch.object(adk_agent, '_run_adk_in_background', side_effect=mock_run_background):
+            execution = await adk_agent._start_background_execution(system_input, "default")
+            await asyncio.sleep(0.01)
+
+        # Verify the SystemMessage became the instruction
+        assert captured_agent.instruction == "You are a math tutor."
+
 
 @pytest.fixture(autouse=True)
 def reset_registry():