fix: fix vertex_ai_session_service crashing when Agent Engine passes full resource names instead of short session IDs

google-genai-bot · copybara-github · commit 60c55ad74570 · 2026-06-16T02:14:59.000-07:00
PiperOrigin-RevId: 932964014
diff --git a/src/google/adk/sessions/vertex_ai_session_service.py b/src/google/adk/sessions/vertex_ai_session_service.py
@@ -50,6 +50,29 @@
 _SESSION_ID_PATTERN = re.compile(r'^[A-Za-z0-9_-]+$')
 
 
+def _extract_short_session_id(
+    session_id: str, expected_engine_id: str | None = None
+) -> str:
+  """Extracts the short session ID if a full resource name is provided."""
+  if isinstance(session_id, str) and '/' in session_id:
+    parts = session_id.split('/')
+    if len(parts) >= 2 and parts[-2] == 'sessions':
+      if (
+          len(parts) >= 4
+          and parts[-4] == 'reasoningEngines'
+          and expected_engine_id
+      ):
+        passed_engine_id = parts[-3]
+        if passed_engine_id != expected_engine_id:
+          raise ValueError(
+              'Session resource name mismatch: session belongs to '
+              f'reasoningEngine {passed_engine_id!r}, but service is '
+              f'configured for {expected_engine_id!r}.'
+          )
+      return parts[-1]
+  return session_id
+
+
 def _validate_session_id(session_id: str) -> None:
   """Rejects session IDs that could escape the URL path segment."""
   if not isinstance(session_id, str) or not _SESSION_ID_PATTERN.fullmatch(
@@ -146,6 +169,9 @@ async def create_session(
 
     config = {'session_state': state} if state else {}
     if session_id:
+      session_id = _extract_short_session_id(
+          session_id, expected_engine_id=reasoning_engine_id
+      )
       _validate_session_id(session_id)
       config['session_id'] = session_id
     config.update(kwargs)
@@ -177,8 +203,11 @@ async def get_session(
       session_id: str,
       config: Optional[GetSessionConfig] = None,
   ) -> Optional[Session]:
-    _validate_session_id(session_id)
     reasoning_engine_id = self._get_reasoning_engine_id(app_name)
+    session_id = _extract_short_session_id(
+        session_id, expected_engine_id=reasoning_engine_id
+    )
+    _validate_session_id(session_id)
     session_resource_name = (
         f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}'
     )
@@ -277,8 +306,11 @@ async def list_sessions(
   async def delete_session(
       self, *, app_name: str, user_id: str, session_id: str
   ) -> None:
-    _validate_session_id(session_id)
     reasoning_engine_id = self._get_reasoning_engine_id(app_name)
+    session_id = _extract_short_session_id(
+        session_id, expected_engine_id=reasoning_engine_id
+    )
+    _validate_session_id(session_id)
     session_resource_name = (
         f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}'
     )
diff --git a/tests/unittests/sessions/test_vertex_ai_session_service.py b/tests/unittests/sessions/test_vertex_ai_session_service.py
@@ -31,6 +31,8 @@
 from google.adk.models.cache_metadata import CacheMetadata
 from google.adk.sessions.base_session_service import GetSessionConfig
 from google.adk.sessions.session import Session
+from google.adk.sessions.vertex_ai_session_service import _extract_short_session_id
+from google.adk.sessions.vertex_ai_session_service import _validate_session_id
 from google.adk.sessions.vertex_ai_session_service import VertexAiSessionService
 from google.api_core import exceptions as api_core_exceptions
 from google.genai import types as genai_types
@@ -1350,3 +1352,53 @@ class DummyModel(pydantic.BaseModel):
 
   assert appended_event.actions.compaction is not None
   assert appended_event.actions.compaction.start_timestamp == 1000.0
+
+
+def test_extract_short_session_id_short_id():
+  assert _extract_short_session_id('123') == '123'
+  assert _extract_short_session_id('session-123_abc') == 'session-123_abc'
+
+
+def test_extract_short_session_id_strips_full_resource_name():
+  resource_name = 'projects/123/locations/us-east4/reasoningEngines/456/sessions/session-123'
+  assert _extract_short_session_id(resource_name) == 'session-123'
+  assert (
+      _extract_short_session_id(resource_name, expected_engine_id='456')
+      == 'session-123'
+  )
+
+
+def test_extract_short_session_id_mismatch():
+  resource_name = 'projects/123/locations/us-east4/reasoningEngines/wrong/sessions/session-123'
+  with pytest.raises(ValueError, match='Session resource name mismatch'):
+    _extract_short_session_id(resource_name, expected_engine_id='right')
+
+
+def test_validate_session_id_rejects_invalid_chars():
+  with pytest.raises(ValueError, match='Invalid session_id'):
+    _validate_session_id('invalid@id')
+  with pytest.raises(ValueError, match='Invalid session_id'):
+    _validate_session_id('invalid/id')
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures('mock_get_api_client')
+async def test_get_session_strips_full_resource_name(
+    mock_api_client_instance,
+):
+  session_service = mock_vertex_ai_session_service()
+  mock_api_client_instance.session_dict['session-123'] = {
+      'name': (
+          'projects/123/locations/us-east4/reasoningEngines/123/sessions/session-123'
+      ),
+      'update_time': '2023-01-01T00:00:00Z',
+      'user_id': 'user',
+  }
+  resource_name = 'projects/123/locations/us-east4/reasoningEngines/123/sessions/session-123'
+  session = await session_service.get_session(
+      app_name='123', user_id='user', session_id=resource_name
+  )
+  assert session.id == 'session-123'
+  mock_api_client_instance.agent_engines.sessions.get.assert_called_once_with(
+      name='reasoningEngines/123/sessions/session-123'
+  )
