From 9f0e5637faa47ee5832d9a70e1bcafd13dff6833 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 26 Mar 2026 20:33:23 +0000 Subject: [PATCH 1/2] fix(adk-middleware): use O(1) lookup in /agents/state when use_thread_id_as_session_id is enabled The /agents/state endpoint always used the O(n) _find_session_by_thread_id scan path on cache miss, even when use_thread_id_as_session_id=True. This adds a direct get_session(session_id=thread_id) lookup first when the flag is enabled, falling back to the scan only for legacy sessions. Fixes #1243 https://claude.ai/code/session_01A1RL4QvwUhWNHocM7H4xkH --- .../python/src/ag_ui_adk/endpoint.py | 40 +++++---- .../tests/test_use_thread_id_as_session_id.py | 83 +++++++++++++++++++ 2 files changed, 109 insertions(+), 14 deletions(-) diff --git a/integrations/adk-middleware/python/src/ag_ui_adk/endpoint.py b/integrations/adk-middleware/python/src/ag_ui_adk/endpoint.py index 3a520fce9d..1b7605aefe 100644 --- a/integrations/adk-middleware/python/src/ag_ui_adk/endpoint.py +++ b/integrations/adk-middleware/python/src/ag_ui_adk/endpoint.py @@ -237,22 +237,34 @@ async def agents_state_endpoint(request_data: AgentStateRequest): # Cache miss - search backend by thread_id if not session: - session = await agent._session_manager._find_session_by_thread_id( - app_name=app_name, - user_id=user_id, - thread_id=thread_id - ) - if session: - # Found - cache for future lookups - session_id = session.id - agent._session_lookup_cache[(thread_id, user_id)] = (session_id, app_name, user_id) - - # Reload session to populate events (list_sessions returns metadata only) - session = await agent._session_manager._session_service.get_session( - session_id=session_id, + # O(1) direct lookup when use_thread_id_as_session_id is enabled + if agent._session_manager._use_thread_id_as_session_id: + session = await agent._session_manager.get_session( + thread_id, app_name, user_id + ) + if session: + session_id = session.id + agent._session_lookup_cache[(thread_id, user_id)] = (session_id, app_name, user_id) + + # Fallback to O(n) scan (always used when flag is False, + # also used as legacy fallback when flag is True but direct lookup misses) + if not session: + session = await agent._session_manager._find_session_by_thread_id( app_name=app_name, - user_id=user_id + user_id=user_id, + thread_id=thread_id ) + if session: + # Found - cache for future lookups + session_id = session.id + agent._session_lookup_cache[(thread_id, user_id)] = (session_id, app_name, user_id) + + # Reload session to populate events (list_sessions returns metadata only) + session = await agent._session_manager._session_service.get_session( + session_id=session_id, + app_name=app_name, + user_id=user_id + ) thread_exists = session is not None diff --git a/integrations/adk-middleware/python/tests/test_use_thread_id_as_session_id.py b/integrations/adk-middleware/python/tests/test_use_thread_id_as_session_id.py index 56f079f2e0..96b8c88742 100644 --- a/integrations/adk-middleware/python/tests/test_use_thread_id_as_session_id.py +++ b/integrations/adk-middleware/python/tests/test_use_thread_id_as_session_id.py @@ -355,3 +355,86 @@ async def test_parameter_defaults_to_false(self): user_id="user", ) assert adk._session_manager._use_thread_id_as_session_id is False + + +class TestAgentsStateEndpointWithDirectLookup: + """Tests for /agents/state endpoint with use_thread_id_as_session_id=True.""" + + @pytest.fixture(autouse=True) + def reset_session_manager(self): + SessionManager.reset_instance() + yield + SessionManager.reset_instance() + + @pytest.fixture + def mock_agent(self): + agent = Mock(spec=Agent) + agent.name = "test_agent" + agent.instruction = "Test instruction" + agent.tools = [] + return agent + + @pytest.fixture + def adk_agent(self, mock_agent): + return ADKAgent( + adk_agent=mock_agent, + app_name="test_app", + user_id="test_user", + use_in_memory_services=True, + use_thread_id_as_session_id=True, + ) + + @pytest.fixture + def app(self, adk_agent): + from fastapi import FastAPI + from ag_ui_adk import add_adk_fastapi_endpoint + app = FastAPI() + add_adk_fastapi_endpoint(app, adk_agent) + return app + + @pytest.fixture + def client(self, app): + from starlette.testclient import TestClient + return TestClient(app) + + @pytest.mark.asyncio + async def test_agents_state_uses_direct_lookup(self, adk_agent, client): + """When use_thread_id_as_session_id=True, /agents/state uses O(1) lookup.""" + # Create a session first via the session manager + session, sid = await adk_agent._session_manager.get_or_create_session( + thread_id="state-thread-123", + app_name="test_app", + user_id="test_user", + ) + assert sid == "state-thread-123" + + # Ensure the cache is clear so endpoint must look up from backend + adk_agent._session_lookup_cache.clear() + + # Spy on list_sessions to verify it's NOT called + with patch.object( + adk_agent._session_manager._session_service, + "list_sessions", + wraps=adk_agent._session_manager._session_service.list_sessions, + ) as spy: + response = client.post( + "/agents/state", + json={"threadId": "state-thread-123"}, + ) + assert response.status_code == 200 + data = response.json() + assert data["threadExists"] is True + assert data["threadId"] == "state-thread-123" + # The key assertion: list_sessions should NOT be called + spy.assert_not_called() + + @pytest.mark.asyncio + async def test_agents_state_nonexistent_thread(self, adk_agent, client): + """/agents/state returns threadExists=False for unknown thread.""" + response = client.post( + "/agents/state", + json={"threadId": "nonexistent-thread"}, + ) + assert response.status_code == 200 + data = response.json() + assert data["threadExists"] is False From a52ce073c1b70d8f6b5e4cad51f7c982de5011fe Mon Sep 17 00:00:00 2001 From: Mark Fogle Date: Thu, 26 Mar 2026 23:16:19 +0000 Subject: [PATCH 2/2] fix(adk-middleware): use getattr with strict True check for _use_thread_id_as_session_id Use `getattr(..., False) is True` instead of direct attribute access so that MagicMock-based test fixtures (which return truthy mocks for any attribute) correctly fall through to the scan path. Co-Authored-By: Claude Opus 4.6 (1M context) --- integrations/adk-middleware/python/src/ag_ui_adk/endpoint.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/integrations/adk-middleware/python/src/ag_ui_adk/endpoint.py b/integrations/adk-middleware/python/src/ag_ui_adk/endpoint.py index 1b7605aefe..874126d182 100644 --- a/integrations/adk-middleware/python/src/ag_ui_adk/endpoint.py +++ b/integrations/adk-middleware/python/src/ag_ui_adk/endpoint.py @@ -238,7 +238,7 @@ async def agents_state_endpoint(request_data: AgentStateRequest): # Cache miss - search backend by thread_id if not session: # O(1) direct lookup when use_thread_id_as_session_id is enabled - if agent._session_manager._use_thread_id_as_session_id: + if getattr(agent._session_manager, '_use_thread_id_as_session_id', False) is True: session = await agent._session_manager.get_session( thread_id, app_name, user_id )