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

agentplatform/frameworks/adk.py

@@ -114,6 +114,7 @@
 
 _DEFAULT_APP_NAME = "default_app_name"
 _DEFAULT_USER_ID = "default-user-id"
+_UPSTREAM_SESSION_ID_KEY = "agent_engine:upstream_session_id"
 _TELEMETRY_API_DISABLED_WARNING = """\
 Tracing integration for Agent Engine has migrated to a new API.
 The 'telemetry.googleapis.com' has not been enabled in project %s.
@@ -984,14 +985,18 @@ async def _init_session(
         session_service: "BaseSessionService",
         artifact_service: "BaseArtifactService",
         request: _StreamRunRequest,
+        state: Optional[Dict[str, Any]] = None,
     ):
         """Initializes the session, and returns the session id."""
         from google.adk.events.event import Event
 
-        session = await session_service.create_session(
-            app_name=self._app_name(),
-            user_id=request.user_id,
-        )
+        create_session_kwargs: Dict[str, Any] = {
+            "app_name": self._app_name(),
+            "user_id": request.user_id,
+        }
+        if state is not None:
+            create_session_kwargs["state"] = state
+        session = await session_service.create_session(**create_session_kwargs)
         if not session:
             raise RuntimeError("Create session failed.")
         if request.events:
@@ -1530,6 +1535,11 @@ async def streaming_agent_run_with_events(self, request_json: str):
         you have defined for the AdkApp. This method is primarily meant for
         invocation from AgentSpace.
 
+        When `request_json` includes a `session_id` that is not found in
+        `session_service`, a new fallback session is created and the caller's
+        `session_id` is stored in `session.state` under
+        `agent_engine:upstream_session_id` (`_UPSTREAM_SESSION_ID_KEY`).
+
         Args:
             request_json (str):
                 Required. The request to stream responses for.
@@ -1594,6 +1604,7 @@ async def streaming_agent_run_with_events(self, request_json: str):
                     session_service=session_service,
                     artifact_service=artifact_service,
                     request=request,
+                    state={_UPSTREAM_SESSION_ID_KEY: request.session_id},
                 )
         else:
             # Not providing a session ID will create a new in-memory session.

tests/unit/agentplatform/frameworks/test_frameworks_adk.py

@@ -43,6 +43,7 @@
 from google.genai import types
 import pytest
 import requests
+from google.adk.plugins.base_plugin import BasePlugin
 from google.adk.sessions.base_session_service import BaseSessionService
 
 
@@ -392,6 +393,30 @@ async def run_async(self, *args, **kwargs):
         )
 
 
+class _SessionStateCapturingPlugin(BasePlugin):
+    """Captures session state on user message and short-circuits the run."""
+
+    def __init__(self):
+        super().__init__(name="session_state_capturing_plugin")
+        self.user_message_session_states = []
+
+    async def on_user_message_callback(
+        self, *, invocation_context, user_message
+    ):
+        del user_message
+        self.user_message_session_states.append(
+            dict(invocation_context.session.state)
+        )
+        return None
+
+    async def before_run_callback(self, *, invocation_context):
+        del invocation_context
+        return types.Content(
+            role="model",
+            parts=[types.Part(text="Sweden")],
+        )
+
+
 @pytest.mark.usefixtures("google_auth_mock")
 class TestAdkApp:
     def test_adk_version(self):
@@ -834,6 +859,187 @@ 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,
+    ):
+        """Seeds agent_engine:upstream_session_id in session.state on fallback."""
+        plugin = _SessionStateCapturingPlugin()
+        app = adk_template.AdkApp(agent=_TEST_AGENT, plugins=[plugin])
+        app.set_up()
+
+        real_runner = app._tmpl_attrs[  # pylint: disable=protected-access
+            "runner"
+        ]
+        run_async_spy = mock.MagicMock(wraps=real_runner.run_async)
+        real_runner.run_async = run_async_spy
+
+        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
+        run_async_spy.assert_called_once()
+        created_session_id = run_async_spy.call_args.kwargs["session_id"]
+        assert created_session_id != "ge_originating_session_123"
+        assert run_async_spy.call_args.kwargs["state_delta"] is None
+        assert plugin.user_message_session_states == [
+            {"agent_engine:upstream_session_id": "ge_originating_session_123"}
+        ]
+        session = await app.async_get_session(
+            user_id=_TEST_USER_ID,
+            session_id=created_session_id,
+        )
+        assert session.state == {
+            "agent_engine:upstream_session_id": "ge_originating_session_123"
+        }
+
+    @pytest.mark.asyncio
+    @pytest.mark.usefixtures(
+        "default_instrumentor_builder_mock",
+        "get_project_id_mock",
+    )
+    async def test_streaming_agent_run_with_events_fallback_session_with_authorizations(
+        self,
+    ):
+        """Persists upstream_session_id in session.state while keeping tokens temp."""
+        plugin = _SessionStateCapturingPlugin()
+        app = adk_template.AdkApp(agent=_TEST_AGENT, plugins=[plugin])
+        app.set_up()
+
+        session_service = app._tmpl_attrs[  # pylint: disable=protected-access
+            "session_service"
+        ]
+        create_session_spy = mock.AsyncMock(
+            wraps=session_service.create_session
+        )
+        session_service.create_session = create_session_spy
+
+        real_runner = app._tmpl_attrs[  # pylint: disable=protected-access
+            "runner"
+        ]
+        run_async_spy = mock.MagicMock(wraps=real_runner.run_async)
+        real_runner.run_async = run_async_spy
+
+        request_json = json.dumps(
+            {
+                "authorizations": {
+                    "test_user_id1": {"access_token": "test_access_token"},
+                },
+                "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
+        create_session_spy.assert_called_once()
+        assert create_session_spy.call_args.kwargs.get("state") == {
+            "agent_engine:upstream_session_id": "ge_originating_session_123"
+        }
+        run_async_spy.assert_called_once()
+        created_session_id = run_async_spy.call_args.kwargs["session_id"]
+        assert created_session_id != "ge_originating_session_123"
+        assert run_async_spy.call_args.kwargs["state_delta"] == {
+            "temp:test_user_id1": "test_access_token"
+        }
+        assert plugin.user_message_session_states == [
+            {"agent_engine:upstream_session_id": "ge_originating_session_123"}
+        ]
+        session = await app.async_get_session(
+            user_id=_TEST_USER_ID,
+            session_id=created_session_id,
+        )
+        assert session.state == {
+            "agent_engine:upstream_session_id": "ge_originating_session_123"
+        }
+        assert "test_access_token" not in json.dumps(session.state)
+
+    @pytest.mark.asyncio
+    @pytest.mark.usefixtures(
+        "default_instrumentor_builder_mock",
+        "get_project_id_mock",
+    )
+    async def test_streaming_agent_run_with_events_existing_session_omits_upstream_session_id(
+        self,
+    ):
+        """Does not add agent_engine:upstream_session_id when get_session succeeds."""
+        plugin = _SessionStateCapturingPlugin()
+        app = adk_template.AdkApp(agent=_TEST_AGENT, plugins=[plugin])
+        app.set_up()
+
+        await app.async_create_session(
+            user_id=_TEST_USER_ID,
+            session_id="test_session_id",
+        )
+
+        real_runner = app._tmpl_attrs[  # pylint: disable=protected-access
+            "runner"
+        ]
+        run_async_spy = mock.MagicMock(wraps=real_runner.run_async)
+        real_runner.run_async = run_async_spy
+
+        request_json = json.dumps(
+            {
+                "user_id": _TEST_USER_ID,
+                "session_id": "test_session_id",
+                "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
+        run_async_spy.assert_called_once()
+        assert run_async_spy.call_args.kwargs["session_id"] == "test_session_id"
+        assert run_async_spy.call_args.kwargs["state_delta"] is None
+        assert plugin.user_message_session_states == [{}]
+        session = await app.async_get_session(
+            user_id=_TEST_USER_ID,
+            session_id="test_session_id",
+        )
+        assert "agent_engine:upstream_session_id" not in session.state
+        assert not session.state
+
     @pytest.mark.asyncio
     @mock.patch.dict(
         os.environ,