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 index b18d6fcf7..a28eb769c 100644 --- 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 @@ -834,20 +834,21 @@ async def _run_adk_in_background( app_name: App name event_queue: Queue for emitting events """ + runner: Optional[Runner] = None try: # 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( adk_agent=adk_agent, user_id=user_id, app_name=app_name ) - + # Create RunConfig run_config = self._run_config_factory(input) - + # Ensure session exists await self._ensure_session_exists( app_name, user_id, input.thread_id, input.state @@ -994,9 +995,20 @@ async def _run_adk_in_background( await event_queue.put(None) finally: # Background task cleanup completed - # Note: toolset cleanup is handled by garbage collection - # since toolset is now embedded in the agent's tools - pass + # Ensure the ADK runner releases any resources (e.g. toolsets) + if runner is not None: + close_method = getattr(runner, "close", None) + if close_method is not None: + try: + close_result = close_method() + if inspect.isawaitable(close_result): + await close_result + except Exception as close_error: + logger.warning( + "Error while closing ADK runner for thread %s: %s", + input.thread_id, + close_error, + ) async def _cleanup_stale_executions(self): """Clean up stale executions.""" diff --git a/typescript-sdk/integrations/adk-middleware/python/tests/test_adk_agent.py b/typescript-sdk/integrations/adk-middleware/python/tests/test_adk_agent.py index 92c586ea9..192b2abcb 100644 --- a/typescript-sdk/integrations/adk-middleware/python/tests/test_adk_agent.py +++ b/typescript-sdk/integrations/adk-middleware/python/tests/test_adk_agent.py @@ -113,6 +113,7 @@ async def test_run_basic_flow(self, adk_agent, sample_input, mock_agent): with patch.object(adk_agent, '_create_runner') as mock_create_runner: # Create a mock runner mock_runner = AsyncMock() + mock_runner.close = AsyncMock() mock_event = Mock() mock_event.id = "event1" mock_event.author = "test_agent" @@ -139,6 +140,31 @@ async def mock_run_async(*args, **kwargs): assert len(events) >= 2 # At least RUN_STARTED and RUN_FINISHED assert events[0].type == EventType.RUN_STARTED assert events[-1].type == EventType.RUN_FINISHED + mock_runner.close.assert_awaited_once() + + @pytest.mark.asyncio + async def test_runner_close_called_on_run_error(self, adk_agent, sample_input): + """Runner.close should still be awaited when execution errors.""" + + with patch.object(adk_agent, '_create_runner') as mock_create_runner: + mock_runner = AsyncMock() + mock_runner.close = AsyncMock() + + async def failing_run_async(*args, **kwargs): + if False: # pragma: no cover - keep async generator semantics + yield None + raise RuntimeError("boom") + + mock_runner.run_async = failing_run_async + mock_create_runner.return_value = mock_runner + + events = [] + async for event in adk_agent.run(sample_input): + events.append(event) + + # Ensure RUN_ERROR emitted and runner closed + assert any(event.type == EventType.RUN_ERROR for event in events) + mock_runner.close.assert_awaited_once() @pytest.mark.asyncio async def test_turn_complete_falls_back_to_streaming_translator(