fix: Preserve caller session_id as upstream_session_id on fallback session creation in AdkApp.

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")

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,

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,

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,

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")

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")