From aa8f1f34a5d0db62b924607ed2d76afe07b100c9 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Thu, 22 Oct 2026 09:08:35 -0700 Subject: [PATCH] fix: Preserve caller session_id as upstream_session_id on fallback session creation in AdkApp. When `AdkApp.streaming_agent_run_with_events` receives a caller-supplied `request.session_id` and `session_service.get_session` does not find an existing session, `AdkApp` falls back to `self._init_session(...)` and creates a new session with a server-generated `session.id`. Populate `state_delta["upstream_session_id"] = request.session_id` in that fallback branch so custom ADK agents and plugins can correlate logs and telemetry with the caller's originating session ID via `session.state["upstream_session_id"]`. PiperOrigin-RevId: 1004821692 --- agentplatform/frameworks/adk.py | 3 + .../frameworks/test_frameworks_adk.py | 48 ++++++++++++++ .../test_agent_engine_templates_adk.py | 48 ++++++++++++++ .../test_reasoning_engine_templates_adk.py | 65 +++++++++++++++++++ vertexai/agent_engines/templates/adk.py | 3 + .../reasoning_engines/templates/adk.py | 3 + 6 files changed, 170 insertions(+) diff --git a/agentplatform/frameworks/adk.py b/agentplatform/frameworks/adk.py index 4509bb6229..219e49d275 100644 --- a/agentplatform/frameworks/adk.py +++ b/agentplatform/frameworks/adk.py @@ -1595,6 +1595,9 @@ async def streaming_agent_run_with_events(self, request_json: str): artifact_service=artifact_service, request=request, ) + if state_delta is None: + state_delta = {} + state_delta["upstream_session_id"] = request.session_id else: # Not providing a session ID will create a new in-memory session. session_service = self._tmpl_attrs.get("in_memory_session_service") diff --git a/tests/unit/agentplatform/frameworks/test_frameworks_adk.py b/tests/unit/agentplatform/frameworks/test_frameworks_adk.py index d646a5123c..13da8d7b64 100644 --- a/tests/unit/agentplatform/frameworks/test_frameworks_adk.py +++ b/tests/unit/agentplatform/frameworks/test_frameworks_adk.py @@ -834,6 +834,54 @@ async def mock_run_async(*args, **kwargs): assert run_config is not None assert run_config.labels == labels + @pytest.mark.asyncio + @pytest.mark.usefixtures( + "default_instrumentor_builder_mock", + "get_project_id_mock", + ) + async def test_streaming_agent_run_with_events_fallback_session_preserves_upstream_session_id( + self, + ): + """Preserves upstream_session_id in state_delta when creating a fallback session.""" + captured = {} + + class _StateDeltaCapturingRunner(_MockRunner): + + async def run_async(self, *args, **kwargs): + captured["session_id"] = kwargs.get("session_id") + captured["state_delta"] = kwargs.get("state_delta") + async for item in super().run_async(*args, **kwargs): + yield item + + app = adk_template.AdkApp(agent=_TEST_AGENT) + app.set_up() + app._tmpl_attrs[ # pylint: disable=protected-access + "runner" + ] = _StateDeltaCapturingRunner() + + request_json = json.dumps( + { + "user_id": _TEST_USER_ID, + "session_id": "ge_originating_session_123", + "message": { + "parts": [{"text": "What is the exchange rate from USD to SEK?"}], + "role": "user", + }, + } + ) + + events = [] + async for event in app.streaming_agent_run_with_events( + request_json=request_json, + ): + events.append(event) + + assert len(events) == 1 + assert captured["session_id"] != "ge_originating_session_123" + assert captured["state_delta"] == { + "upstream_session_id": "ge_originating_session_123" + } + @pytest.mark.asyncio @mock.patch.dict( os.environ, diff --git a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py index ef5a2a29f2..df154c9bd7 100644 --- a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py +++ b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py @@ -845,6 +845,54 @@ async def mock_run_async(*args, **kwargs): run_config=None, ) + @pytest.mark.asyncio + @pytest.mark.usefixtures( + "default_instrumentor_builder_mock", + "get_project_id_mock", + ) + async def test_streaming_agent_run_with_events_fallback_session_preserves_upstream_session_id( + self, + ): + """Preserves upstream_session_id in state_delta when creating a fallback session.""" + captured = {} + + class _StateDeltaCapturingRunner(_MockRunner): + + async def run_async(self, *args, **kwargs): + captured["session_id"] = kwargs.get("session_id") + captured["state_delta"] = kwargs.get("state_delta") + async for item in super().run_async(*args, **kwargs): + yield item + + app = agent_engines.AdkApp(agent=_TEST_AGENT) + app.set_up() + app._tmpl_attrs[ # pylint: disable=protected-access + "runner" + ] = _StateDeltaCapturingRunner() + + request_json = json.dumps( + { + "user_id": _TEST_USER_ID, + "session_id": "ge_originating_session_123", + "message": { + "parts": [{"text": "What is the exchange rate from USD to SEK?"}], + "role": "user", + }, + } + ) + + events = [] + async for event in app.streaming_agent_run_with_events( + request_json=request_json, + ): + events.append(event) + + assert len(events) == 1 + assert captured["session_id"] != "ge_originating_session_123" + assert captured["state_delta"] == { + "upstream_session_id": "ge_originating_session_123" + } + @pytest.mark.asyncio async def test_streaming_agent_run_with_events_new_session_uses_temp_state( self, diff --git a/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py b/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py index 7911eb7dde..80ff62e248 100644 --- a/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py +++ b/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py @@ -23,6 +23,7 @@ from typing import Optional from google import auth +from google.adk.events.event import Event from google.api_core import exceptions as api_exceptions import vertexai from google.cloud.aiplatform import initializer @@ -807,6 +808,45 @@ def run(self, *args, **kwargs): assert captured["run_config"] is not None assert captured["run_config"].labels == labels + def test_streaming_agent_run_with_events_fallback_session_preserves_upstream_session_id( + self, + ): + """Preserves upstream_session_id in state_delta when creating a fallback session.""" + captured = {} + + class _StateDeltaCapturingRunner(_MockRunner): + + def run(self, *args, **kwargs): + captured["session_id"] = kwargs.get("session_id") + captured["state_delta"] = kwargs.get("state_delta") + yield from super().run(*args, **kwargs) + + app = reasoning_engines.AdkApp( + agent=Agent(name=_TEST_AGENT_NAME, model=_TEST_MODEL) + ) + app.set_up() + app._tmpl_attrs[ # pylint: disable=protected-access + "runner" + ] = _StateDeltaCapturingRunner() + + request_json = json.dumps( + { + "user_id": _TEST_USER_ID, + "session_id": "ge_originating_session_123", + "message": { + "parts": [{"text": "What is the exchange rate from USD to SEK?"}], + "role": "user", + }, + } + ) + events = list(app.streaming_agent_run_with_events(request_json=request_json)) + + assert len(events) == 1 + assert captured["session_id"] != "ge_originating_session_123" + assert captured["state_delta"] == { + "upstream_session_id": "ge_originating_session_123" + } + @pytest.mark.asyncio @mock.patch.dict( os.environ, @@ -969,6 +1009,31 @@ async def test_async_add_session_to_memory(self): ) assert app._tmpl_attrs.get("memory_service") is None session = app.create_session(user_id=_TEST_USER_ID) + app._tmpl_attrs[ # pylint: disable=protected-access + "runner" + ] = _MockRunner() + + session_obj = await app._tmpl_attrs[ # pylint: disable=protected-access + "session_service" + ].get_session( + app_name=app._tmpl_attrs.get( # pylint: disable=protected-access + "app_name" + ), + user_id=_TEST_USER_ID, + session_id=session["id"], + ) + await app._tmpl_attrs[ # pylint: disable=protected-access + "session_service" + ].append_event( + session=session_obj, + event=Event( + author="user", + content={ + "parts": [{"text": "My cat's name is Garfield"}], + "role": "user", + }, + ), + ) list( app.stream_query( user_id=_TEST_USER_ID, diff --git a/vertexai/agent_engines/templates/adk.py b/vertexai/agent_engines/templates/adk.py index 8d1b28b7e8..78759e9e2a 100644 --- a/vertexai/agent_engines/templates/adk.py +++ b/vertexai/agent_engines/templates/adk.py @@ -1566,6 +1566,9 @@ async def streaming_agent_run_with_events(self, request_json: str): artifact_service=artifact_service, request=request, ) + if state_delta is None: + state_delta = {} + state_delta["upstream_session_id"] = request.session_id else: # Not providing a session ID will create a new in-memory session. session_service = self._tmpl_attrs.get("in_memory_session_service") diff --git a/vertexai/preview/reasoning_engines/templates/adk.py b/vertexai/preview/reasoning_engines/templates/adk.py index 24b28caa18..1b3aca0a81 100644 --- a/vertexai/preview/reasoning_engines/templates/adk.py +++ b/vertexai/preview/reasoning_engines/templates/adk.py @@ -1366,6 +1366,9 @@ async def _invoke_agent_async(): artifact_service=artifact_service, request=request, ) + if state_delta is None: + state_delta = {} + state_delta["upstream_session_id"] = request.session_id else: # Not providing a session ID will create a new in-memory session. session_service = self._tmpl_attrs.get("in_memory_session_service")