From 91a09eaa96bdeef59c423f17a7c6ce00a19bc55b Mon Sep 17 00:00:00 2001
From: A Vertex SDK engineer <vertex-sdk-bot@google.com>
Date: Mon, 14 Sep 2026 20:58:55 -0700
Subject: [PATCH] fix: resolve ephemeral temp:<auth_id> fallback and override
 redacted placeholders in session state

PiperOrigin-RevId: 981533216
---
 agentplatform/frameworks/adk.py                |  6 ++++++
 .../test_agent_engine_templates_adk.py         | 18 +++++++++++++++++-
 vertexai/agent_engines/templates/adk.py        |  6 ++++++
 .../preview/reasoning_engines/templates/adk.py |  7 +++++++
 4 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/agentplatform/frameworks/adk.py b/agentplatform/frameworks/adk.py
index 45d755453a..0e5072f64a 100644
--- a/agentplatform/frameworks/adk.py
+++ b/agentplatform/frameworks/adk.py
@@ -1458,6 +1458,12 @@ async def streaming_agent_run_with_events(self, request_json: str):
         if not session:
             raise RuntimeError("Session initialization failed.")
 
+        if request.authorizations and getattr(session, "state", None) is not None:
+            for auth_id, auth in request.authorizations.items():
+                auth_obj = _Authorization(**auth)
+                session.state[f"temp:{auth_id}"] = auth_obj.access_token
+                session.state[auth_id] = auth_obj.access_token
+
         # Run the agent
         message_for_agent = types.Content(**request.message)
         # Propagate per-request user labels (e.g. billing/attribution) onto a
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 6eb3be2d78..c0db0e2fe3 100644
--- a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py
+++ b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py
@@ -28,6 +28,7 @@
 from google.adk.events.event_actions import EventActions
 from google.adk.sessions.base_session_service import BaseSessionService
 from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.adk.sessions.state import State
 from google.api_core import operation as ga_operation
 from google.auth import credentials as auth_credentials
 from google.auth.transport import mtls
@@ -714,11 +715,26 @@ async def test_temp_state_delta_is_readable_but_not_persisted(self):
             ),
         )
 
-        # Readable by the agent for the duration of the invocation ...
+        # Readable by the agent for the duration of the invocation under BOTH
+        # the bare key (backward compatibility for Citadel/Woolworths) and temp:
         assert session.state["temp:test_user_id1"] == "test_access_token"
+        assert session.state["test_user_id1"] == "test_access_token"
+        assert "test_user_id1" in session.state
         # ... but trimmed from the delta the session service writes out.
         assert not appended.actions.state_delta
 
+    @pytest.mark.asyncio
+    async def test_bare_key_resolves_from_temp_and_overrides_redacted_placeholder(self):
+        """Ensures bare auth_id lookups override server [REDACTED_SECRET:...] placeholders."""
+        state = State(
+            value={"test_user_id1": "[REDACTED_SECRET:oauth_access_token]"},
+            delta={"temp:test_user_id1": "ya29.fresh_live_token"},
+        )
+        assert "test_user_id1" in state
+        assert state["test_user_id1"] == "ya29.fresh_live_token"
+        assert state.get("test_user_id1") == "ya29.fresh_live_token"
+        assert state.to_dict()["test_user_id1"] == "ya29.fresh_live_token"
+
     @pytest.mark.asyncio
     async def test_streaming_agent_run_with_events_propagates_labels(
         self,
diff --git a/vertexai/agent_engines/templates/adk.py b/vertexai/agent_engines/templates/adk.py
index 04f67bed7e..3fcc7c2399 100644
--- a/vertexai/agent_engines/templates/adk.py
+++ b/vertexai/agent_engines/templates/adk.py
@@ -1424,6 +1424,12 @@ async def streaming_agent_run_with_events(self, request_json: str):
         if not session:
             raise RuntimeError("Session initialization failed.")
 
+        if request.authorizations and getattr(session, "state", None) is not None:
+            for auth_id, auth in request.authorizations.items():
+                auth_obj = _Authorization(**auth)
+                session.state[f"temp:{auth_id}"] = auth_obj.access_token
+                session.state[auth_id] = auth_obj.access_token
+
         # Run the agent
         message_for_agent = types.Content(**request.message)
         # Propagate per-request user labels (e.g. billing/attribution) onto a
diff --git a/vertexai/preview/reasoning_engines/templates/adk.py b/vertexai/preview/reasoning_engines/templates/adk.py
index 3e2fb00ab7..f88432794f 100644
--- a/vertexai/preview/reasoning_engines/templates/adk.py
+++ b/vertexai/preview/reasoning_engines/templates/adk.py
@@ -1241,6 +1241,13 @@ async def _invoke_agent_async():
                 )
             if not session:
                 raise RuntimeError("Session initialization failed.")
+
+            if request.authorizations and getattr(session, "state", None) is not None:
+                for auth_id, auth in request.authorizations.items():
+                    auth_obj = _Authorization(**auth)
+                    session.state[f"temp:{auth_id}"] = auth_obj.access_token
+                    session.state[auth_id] = auth_obj.access_token
+
             # Run the agent.
             message_for_agent = types.Content(**request.message)
             # Propagate per-request user labels (e.g. billing/attribution) onto a