fix: investigate missing auth_id for _init_session

PiperOrigin-RevId: 828119747

Author: Tongzhou-Jiang

vertexai/agent_engines/templates/adk.py

@@ -603,18 +603,23 @@ async def _init_session(
         """Initializes the session, and returns the session id."""
         from google.adk.events.event import Event
         import random
+        from google.cloud.aiplatform import base
+
+        _LOGGER = base.Logger(__name__)
 
         session_state = None
         if request.authorizations:
             session_state = {}
             for auth_id, auth in request.authorizations.items():
                 auth = _Authorization(**auth)
-                session_state[f"temp:{auth_id}"] = auth.access_token
+                session_state[auth_id] = auth.access_token
 
         if request.session_id:
             session_id = request.session_id
         else:
             session_id = f"temp_session_{random.randbytes(8).hex()}"
+        _LOGGER.info("session_state: %s", session_state)
+        print("session_state: %s", session_state)
         session = await session_service.create_session(
             app_name=self._tmpl_attrs.get("app_name"),
             user_id=request.user_id,

vertexai/preview/reasoning_engines/templates/adk.py

@@ -551,8 +551,13 @@ async def _init_session(
         """Initializes the session, and returns the session id."""
         from google.adk.events.event import Event
         import random
+        from google.cloud.aiplatform import base
+
+        _LOGGER = base.Logger(__name__)
 
         session_state = None
+        _LOGGER.info("authorizations: %s", request.authorizations)
+        print("authorizations: %s", request.authorizations)
         if request.authorizations:
             session_state = {}
             for auth_id, auth in request.authorizations.items():
@@ -563,6 +568,8 @@ async def _init_session(
             session_id = request.session_id
         else:
             session_id = f"temp_session_{random.randbytes(8).hex()}"
+        _LOGGER.info("session_state: %s", session_state)
+        print("session_state: %s", session_state)
         session = await session_service.create_session(
             app_name=self._tmpl_attrs.get("app_name"),
             user_id=request.user_id,