diff --git a/src/google/adk/sessions/vertex_ai_session_service.py b/src/google/adk/sessions/vertex_ai_session_service.py index 1f1edd3d..29b133b4 100644 --- a/src/google/adk/sessions/vertex_ai_session_service.py +++ b/src/google/adk/sessions/vertex_ai_session_service.py @@ -50,6 +50,29 @@ _USAGE_METADATA_CUSTOM_METADATA_KEY = '_usage_metadata' _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( @@ -140,6 +163,9 @@ class VertexAiSessionService(BaseSessionService): 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) @@ -171,8 +197,11 @@ class VertexAiSessionService(BaseSessionService): 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}' ) @@ -271,8 +300,11 @@ class VertexAiSessionService(BaseSessionService): 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 index b8c71701..8daba2f0 100644 --- 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.events.event_actions import EventCompaction 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 @@ -764,6 +766,70 @@ async def test_session_id_path_traversal_rejected(): ) +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/test-project/locations/test-location/' + 'reasoningEngines/123/sessions/3' + ) + assert _extract_short_session_id(resource_name) == '3' + assert ( + _extract_short_session_id(resource_name, expected_engine_id='123') == '3' + ) + + +def test_extract_short_session_id_mismatch(): + resource_name = ( + 'projects/test-project/locations/test-location/' + 'reasoningEngines/wrong/sessions/3' + ) + with pytest.raises(ValueError, match='Session resource name mismatch'): + _extract_short_session_id(resource_name, expected_engine_id='123') + + +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_accepts_a_full_resource_name(): + """Agent Engine passes the full resource name, not the short id.""" + session_service = mock_vertex_ai_session_service() + resource_name = ( + 'projects/test-project/locations/test-location/' + 'reasoningEngines/123/sessions/1' + ) + + session = await session_service.get_session( + app_name='123', user_id='user', session_id=resource_name + ) + + assert session.id == '1' + + +@pytest.mark.asyncio +@pytest.mark.usefixtures('mock_get_api_client') +async def test_get_session_rejects_a_resource_name_for_another_engine(): + session_service = mock_vertex_ai_session_service() + resource_name = ( + 'projects/test-project/locations/test-location/' + 'reasoningEngines/456/sessions/1' + ) + + with pytest.raises(ValueError, match='Session resource name mismatch'): + await session_service.get_session( + app_name='123', user_id='user', session_id=resource_name + ) + + @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_get_session_with_page_token():