# Copyright 2026 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import copy import datetime import re import types from typing import Any from typing import List from typing import Optional from typing import Tuple from unittest import mock from dateutil.parser import isoparse from fastapi.openapi import models as openapi_models from google.adk.auth import auth_schemes from google.adk.auth.auth_tool import AuthConfig from google.adk.events.event import Event from google.adk.events.event_actions import EventActions 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 from google.genai.errors import ClientError import pydantic import pytest MOCK_SESSION_JSON_1 = { 'name': ( 'projects/test-project/locations/test-location/' 'reasoningEngines/123/sessions/1' ), 'create_time': '2024-12-12T12:12:12.123456Z', 'update_time': '2024-12-12T12:12:12.123456Z', 'session_state': { 'key': {'value': 'test_value'}, }, 'user_id': 'user', } MOCK_SESSION_JSON_2 = { 'name': ( 'projects/test-project/locations/test-location/' 'reasoningEngines/123/sessions/2' ), 'update_time': '2024-12-13T12:12:12.123456Z', 'user_id': 'user', } MOCK_SESSION_JSON_3 = { 'name': ( 'projects/test-project/locations/test-location/' 'reasoningEngines/123/sessions/3' ), 'update_time': '2024-12-14T12:12:12.123456Z', 'user_id': 'user2', } MOCK_EVENT_JSON = [ { 'name': ( 'projects/test-project/locations/test-location/' 'reasoningEngines/123/sessions/1/events/123' ), 'invocation_id': '123', 'author': 'user', 'timestamp': '2024-12-12T12:12:12.123456Z', 'content': { 'parts': [ {'text': 'test_content'}, ], }, 'actions': { 'state_delta': { 'key': {'value': 'test_value'}, }, 'transfer_agent': 'agent', }, 'event_metadata': { 'partial': False, 'turn_complete': True, 'interrupted': False, 'branch': '', 'long_running_tool_ids': ['tool1'], }, 'raw_event': {}, }, ] MOCK_EVENT_JSON_2 = [ { 'name': ( 'projects/test-project/locations/test-location/' 'reasoningEngines/123/sessions/2/events/123' ), 'invocation_id': '222', 'author': 'user', 'timestamp': '2024-12-12T12:12:12.123456Z', }, ] MOCK_EVENT_JSON_3 = [ { 'name': ( 'projects/test-project/locations/test-location/' 'reasoningEngines/123/sessions/2/events/456' ), 'invocation_id': '333', 'author': 'user', 'timestamp': '2024-12-12T12:12:13.123456Z', }, ] MOCK_SESSION_JSON_PAGE1 = { 'name': ( 'projects/test-project/locations/test-location/' 'reasoningEngines/123/sessions/page1' ), 'update_time': '2024-12-15T12:12:12.123456Z', 'user_id': 'user_with_pages', } MOCK_SESSION_JSON_PAGE2 = { 'name': ( 'projects/test-project/locations/test-location/' 'reasoningEngines/123/sessions/page2' ), 'update_time': '2024-12-16T12:12:12.123456Z', 'user_id': 'user_with_pages', } MOCK_SESSION_JSON_5 = { 'name': ( 'projects/test-project/locations/test-location/' 'reasoningEngines/123/sessions/5' ), 'update_time': '2024-12-12T12:15:12.123456Z', 'user_id': 'user_with_many_events', } def _generate_mock_events_for_session_5(num_events): events = [] start_time = isoparse('2024-12-12T12:12:12.123456Z') for i in range(num_events): event_time = start_time + datetime.timedelta(microseconds=i * 1000) events.append({ 'name': ( 'projects/test-project/locations/test-location/' f'reasoningEngines/123/sessions/5/events/{i}' ), 'invocation_id': f'invocation_{i}', 'author': 'user_with_many_events', 'timestamp': event_time.isoformat().replace('+00:00', 'Z'), }) return events MANY_EVENTS_COUNT = 200 MOCK_EVENTS_JSON_5 = _generate_mock_events_for_session_5(MANY_EVENTS_COUNT) MOCK_EVENT_WITH_OVERRIDE_JSON = [{ 'name': ( 'projects/test-project/locations/test-location/' 'reasoningEngines/123/sessions/override/events/1' ), 'invocationId': 'override_invoke', 'author': 'user_with_override', 'timestamp': '2024-12-12T12:12:12.123456Z', 'content': { 'parts': [ {'text': 'top_level_content'}, ], }, 'actions': { 'transferToAgent': 'top_level_agent', }, 'eventMetadata': { 'partial': True, 'turnComplete': False, 'interrupted': False, 'branch': 'top_level_branch', }, 'errorCode': '111', 'errorMessage': 'top_level_error', 'rawEvent': { 'invocationId': 'wrong_invocation_id', 'author': 'wrong_author', 'content': { 'parts': [ {'text': 'raw_event_content'}, ], }, 'actions': { 'transferToAgent': 'raw_event_agent', }, 'partial': False, 'turnComplete': True, 'interrupted': True, 'branch': 'raw_event_branch', 'errorCode': '222', 'errorMessage': 'raw_event_error', }, }] MOCK_EVENT_WITH_OVERRIDE_JSON_2 = [{ 'name': ( 'projects/test-project/locations/test-location/' 'reasoningEngines/123/sessions/override/events/1' ), 'invocationId': 'override_invoke', 'author': 'user_with_override', 'content': {}, 'actions': {}, 'timestamp': '2024-12-12T12:12:12.123456Z', 'rawEvent': { 'invocationId': 'wrong_invocation_id', 'author': 'wrong_author', 'content': { 'parts': [ {'text': 'raw_event_content'}, ], }, 'actions': { 'skipSummarization': None, 'stateDelta': {}, 'artifactDelta': {}, 'transferToAgent': 'raw_event_agent', 'escalate': None, 'requestedAuthConfigs': {}, }, 'errorCode': '222', 'errorMessage': 'raw_event_error', 'partial': False, 'turnComplete': True, 'interrupted': True, 'branch': 'raw_event_branch', 'customMetadata': None, 'longRunningToolIds': None, }, }] MOCK_SESSION_WITH_OVERRIDE_JSON = { 'name': ( 'projects/test-project/locations/test-location/' 'reasoningEngines/123/sessions/override' ), 'update_time': '2024-12-12T12:12:12.123456Z', 'user_id': 'user_with_override', } MOCK_SESSION = Session( app_name='123', user_id='user', id='1', state=MOCK_SESSION_JSON_1['session_state'], last_update_time=isoparse(MOCK_SESSION_JSON_1['update_time']).timestamp(), events=[ Event( id='123', invocation_id='123', author='user', timestamp=isoparse(MOCK_EVENT_JSON[0]['timestamp']).timestamp(), content=genai_types.Content( parts=[genai_types.Part(text='test_content')] ), actions=EventActions( transfer_to_agent='agent', state_delta={'key': {'value': 'test_value'}}, ), partial=False, turn_complete=True, interrupted=False, branch='', long_running_tool_ids={'tool1'}, ), ], ) MOCK_SESSION_2 = Session( app_name='123', user_id='user', id='2', last_update_time=isoparse(MOCK_SESSION_JSON_2['update_time']).timestamp(), events=[ Event( id='123', invocation_id='222', author='user', timestamp=isoparse(MOCK_EVENT_JSON_2[0]['timestamp']).timestamp(), ), Event( id='456', invocation_id='333', author='user', timestamp=isoparse(MOCK_EVENT_JSON_3[0]['timestamp']).timestamp(), ), ], ) class PydanticNamespace(types.SimpleNamespace): def model_dump(self, exclude_none=True, mode='python'): d = {} for k, v in self.__dict__.items(): if exclude_none and v is None: continue if isinstance(v, PydanticNamespace): d[k] = v.model_dump(exclude_none=exclude_none, mode=mode) elif isinstance(v, list): d[k] = [ i.model_dump(exclude_none=exclude_none, mode=mode) if isinstance(i, PydanticNamespace) else i for i in v ] else: d[k] = v return d def _convert_to_object(data): if isinstance(data, dict): kwargs = {} for key, value in data.items(): if key in [ 'timestamp', 'update_time', 'create_time', ] and isinstance(value, str): kwargs[key] = isoparse(value) elif key in [ 'session_state', 'state_delta', 'artifact_delta', 'custom_metadata', 'requested_auth_configs', 'rawEvent', 'raw_event', ]: kwargs[key] = value else: kwargs[key] = _convert_to_object(value) return PydanticNamespace(**kwargs) elif isinstance(data, list): return [_convert_to_object(item) for item in data] else: return data async def to_async_iterator(data): for item in data: yield item class MockAsyncClient: """Mocks the API Client.""" def __init__(self) -> None: """Initializes MockClient.""" self.session_dict: dict[str, Any] = {} self.event_dict: dict[str, Tuple[List[Any], Optional[str]]] = {} self.agent_engines = mock.AsyncMock() self.agent_engines.sessions.get.side_effect = self._get_session self.agent_engines.sessions.list.side_effect = self._list_sessions self.agent_engines.sessions.delete.side_effect = self._delete_session self.agent_engines.sessions.create.side_effect = self._create_session self.agent_engines.sessions.events.list.side_effect = self._list_events self.agent_engines.sessions.events.append.side_effect = self._append_event self.last_create_session_config: dict[str, Any] = {} self.last_list_sessions_config: dict[str, Any] = {} async def __aenter__(self): """Enters the asynchronous context.""" return self async def __aexit__(self, exc_type, exc_val, exc_tb): """Exits the asynchronous context.""" pass async def _get_session(self, name: str): session_id = name.split('/')[-1] if session_id in self.session_dict: return _convert_to_object(self.session_dict[session_id]) raise api_core_exceptions.NotFound(f'Session not found: {session_id}') async def _list_sessions(self, name: str, config: dict[str, Any]): self.last_list_sessions_config = config filter_val = config.get('filter', '') user_id_match = re.search(r'user_id="((?:\\.|[^"])*)"', filter_val) if user_id_match: user_id = user_id_match.group(1) if user_id == 'user_with_pages': return to_async_iterator([ _convert_to_object(MOCK_SESSION_JSON_PAGE1), _convert_to_object(MOCK_SESSION_JSON_PAGE2), ]) return to_async_iterator([ _convert_to_object(session) for session in self.session_dict.values() if session['user_id'] == user_id ]) # No user filter, return all sessions return to_async_iterator( [_convert_to_object(session) for session in self.session_dict.values()] ) async def _delete_session(self, name: str): session_id = name.split('/')[-1] self.session_dict.pop(session_id) async def _create_session( self, name: str, user_id: str, config: dict[str, Any] ): self.last_create_session_config = config if 'session_id' in config: new_session_id = config['session_id'] else: new_session_id = '4' self.session_dict[new_session_id] = { 'name': ( 'projects/test-project/locations/test-location/' 'reasoningEngines/123/sessions/' + new_session_id ), 'user_id': user_id, 'session_state': config.get('session_state', {}), 'update_time': '2024-12-12T12:12:12.123456Z', } return _convert_to_object({ 'name': ( 'projects/test_project/locations/test_location/' 'reasoningEngines/123/sessions/' + new_session_id + '/operations/111' ), 'done': True, 'response': self.session_dict[new_session_id], }) async def _list_events(self, name: str, **kwargs): session_id = name.split('/')[-1] events = [] if session_id in self.event_dict: events_tuple = self.event_dict[session_id] events.extend(events_tuple[0]) if events_tuple[1] == 'my_token': events.extend(MOCK_EVENT_JSON_3) config = kwargs.get('config', {}) filter_str = config.get('filter', None) if filter_str: match = re.search(r'timestamp>="([^"]+)"', filter_str) if match: after_timestamp_str = match.group(1) after_timestamp = isoparse(after_timestamp_str) events = [ event for event in events if isoparse(event['timestamp']) >= after_timestamp ] return to_async_iterator([_convert_to_object(event) for event in events]) async def _append_event( self, name: str, author: str, invocation_id: str, timestamp: Any, config: dict[str, Any], ): session_id = name.split('/')[-1] event_list, token = self.event_dict.get(session_id, ([], None)) event_id = str(len(event_list) + 1000) # generate unique ID event_timestamp_str = timestamp.isoformat().replace('+00:00', 'Z') event_json = { 'name': f'{name}/events/{event_id}', 'invocation_id': invocation_id, 'author': author, 'timestamp': event_timestamp_str, } event_json.update(config) if session_id in self.session_dict: self.session_dict[session_id]['update_time'] = event_timestamp_str if session_id in self.event_dict: self.event_dict[session_id][0].append(event_json) else: self.event_dict[session_id] = ([event_json], None) class MockAsyncClientWithPagination: """Mock client that simulates pagination requiring an open client connection. This mock tracks whether the client context is active and raises RuntimeError if iteration occurs outside the context, simulating the real httpx behavior. """ def __init__(self, session_data: dict, events_pages: list[list[dict]]): self._session_data = session_data self._events_pages = events_pages self._context_active = False self.agent_engines = mock.AsyncMock() self.agent_engines.sessions.get.side_effect = self._get_session self.agent_engines.sessions.events.list.side_effect = self._list_events async def __aenter__(self): self._context_active = True return self async def __aexit__(self, exc_type, exc_val, exc_tb): self._context_active = False async def _get_session(self, name: str): return _convert_to_object(self._session_data) async def _list_events(self, name: str, **kwargs): return self._paginated_events_iterator() async def _paginated_events_iterator(self): for page in self._events_pages: for event in page: if not self._context_active: raise RuntimeError( 'Cannot send a request, as the client has been closed.' ) yield _convert_to_object(event) def _generate_events_for_page(session_id: str, start_idx: int, count: int): events = [] start_time = isoparse('2024-12-12T12:12:12.123456Z') for i in range(count): idx = start_idx + i event_time = start_time + datetime.timedelta(microseconds=idx * 1000) events.append({ 'name': ( 'projects/test-project/locations/test-location/' f'reasoningEngines/123/sessions/{session_id}/events/{idx}' ), 'invocation_id': f'invocation_{idx}', 'author': 'pagination_user', 'timestamp': event_time.isoformat().replace('+00:00', 'Z'), }) return events @pytest.mark.asyncio async def test_get_session_pagination_keeps_client_open(): """Regression test: event iteration must occur inside the api_client context. This test verifies that get_session() keeps the API client open while iterating through paginated events. Before the fix, the events_iterator was consumed outside the async with block, causing RuntimeError when fetching subsequent pages. """ session_data = { 'name': ( 'projects/test-project/locations/test-location/' 'reasoningEngines/123/sessions/pagination_test' ), 'update_time': '2024-12-12T12:12:12.123456Z', 'user_id': 'pagination_user', } page1_events = _generate_events_for_page('pagination_test', 0, 100) page2_events = _generate_events_for_page('pagination_test', 100, 100) page3_events = _generate_events_for_page('pagination_test', 200, 50) mock_client = MockAsyncClientWithPagination( session_data=session_data, events_pages=[page1_events, page2_events, page3_events], ) session_service = mock_vertex_ai_session_service() with mock.patch.object( session_service, '_get_api_client', return_value=mock_client ): session = await session_service.get_session( app_name='123', user_id='pagination_user', session_id='pagination_test' ) assert session is not None assert len(session.events) == 250 assert session.events[0].invocation_id == 'invocation_0' assert session.events[249].invocation_id == 'invocation_249' def mock_vertex_ai_session_service( project: Optional[str] = 'test-project', location: Optional[str] = 'test-location', agent_engine_id: Optional[str] = None, express_mode_api_key: Optional[str] = None, ): """Creates a mock Vertex AI Session service for testing.""" return VertexAiSessionService( project=project, location=location, agent_engine_id=agent_engine_id, express_mode_api_key=express_mode_api_key, ) @pytest.fixture def mock_api_client_instance(): """Creates a mock API client instance for testing.""" api_client = MockAsyncClient() api_client.session_dict = { '1': MOCK_SESSION_JSON_1, '2': MOCK_SESSION_JSON_2, '3': MOCK_SESSION_JSON_3, 'page1': MOCK_SESSION_JSON_PAGE1, 'page2': MOCK_SESSION_JSON_PAGE2, } api_client.event_dict = { '1': (copy.deepcopy(MOCK_EVENT_JSON), None), '2': (copy.deepcopy(MOCK_EVENT_JSON_2), 'my_token'), } return api_client @pytest.fixture def mock_get_api_client(mock_api_client_instance): """Mocks the _get_api_client method to return a mock API client.""" with mock.patch( 'google.adk.sessions.vertex_ai_session_service.VertexAiSessionService._get_api_client', return_value=mock_api_client_instance, ): yield @pytest.mark.asyncio async def test_initialize_with_project_location_and_api_key_error(): with pytest.raises(ValueError) as excinfo: mock_vertex_ai_session_service( project='test-project', location='test-location', express_mode_api_key='test-api-key', ) assert ( 'Cannot specify project or location and express_mode_api_key. Either use' ' project and location, or just the express_mode_api_key.' in str(excinfo.value) ) @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_get_session_returns_none_when_invalid_argument( mock_api_client_instance, ): session_service = mock_vertex_ai_session_service() # Simulate the API raising a session not found exception. mock_api_client_instance.agent_engines.sessions.get.side_effect = ClientError( code=404, response_json={ 'message': ( 'Session (projectNumber: 123, reasoningEngineId: 123, sessionId:' ' 123) not found.' ) }, response=None, ) session = await session_service.get_session( app_name='123', user_id='user', session_id='missing' ) assert session is None @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') @pytest.mark.parametrize('agent_engine_id', [None, '123']) async def test_get_empty_session(agent_engine_id): session_service = mock_vertex_ai_session_service(agent_engine_id) with pytest.raises(api_core_exceptions.NotFound) as excinfo: await session_service.get_session( app_name='123', user_id='user', session_id='0' ) assert str(excinfo.value) == '404 Session not found: 0' @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') @pytest.mark.parametrize('agent_engine_id', [None, '123']) async def test_get_another_user_session(agent_engine_id): session_service = mock_vertex_ai_session_service(agent_engine_id) with pytest.raises(ValueError) as excinfo: await session_service.get_session( app_name='123', user_id='user2', session_id='1' ) assert str(excinfo.value) == 'Session 1 does not belong to user user2.' @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_get_and_delete_session(): session_service = mock_vertex_ai_session_service() assert ( await session_service.get_session( app_name='123', user_id='user', session_id='1' ) == MOCK_SESSION ) await session_service.delete_session( app_name='123', user_id='user', session_id='1' ) with pytest.raises(api_core_exceptions.NotFound) as excinfo: await session_service.get_session( app_name='123', user_id='user', session_id='1' ) assert str(excinfo.value) == '404 Session not found: 1' @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_delete_session_rejects_other_users_session(): """delete_session must not delete a session owned by a different user.""" session_service = mock_vertex_ai_session_service() # session '1' belongs to 'user'; 'user2' must not be allowed to delete it. with pytest.raises(ValueError) as excinfo: await session_service.delete_session( app_name='123', user_id='user2', session_id='1' ) assert 'does not belong to user user2' in str(excinfo.value) # Session must still exist. assert ( await session_service.get_session( app_name='123', user_id='user', session_id='1' ) == MOCK_SESSION ) @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_session_id_path_traversal_rejected(): """Session IDs containing path-traversal characters must be rejected.""" session_service = mock_vertex_ai_session_service() for bad_id in ['..', '../foo', '..?force=true', 'a/b', '']: with pytest.raises(ValueError): await session_service.delete_session( app_name='123', user_id='user', session_id=bad_id ) with pytest.raises(ValueError): await session_service.get_session( app_name='123', user_id='user', session_id=bad_id ) @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_get_session_with_page_token(): session_service = mock_vertex_ai_session_service() assert ( await session_service.get_session( app_name='123', user_id='user', session_id='2' ) == MOCK_SESSION_2 ) @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_get_session_with_after_timestamp_filter(): session_service = mock_vertex_ai_session_service() session = await session_service.get_session( app_name='123', user_id='user', session_id='2', config=GetSessionConfig( after_timestamp=isoparse('2024-12-12T12:12:13.0Z').timestamp() ), ) assert session is not None assert len(session.events) == 1 assert session.events[0].id == '456' @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_get_session_keeps_events_newer_than_update_time( mock_api_client_instance: MockAsyncClient, ) -> None: future_event_time = isoparse( MOCK_SESSION_JSON_1['update_time'] ) + datetime.timedelta(seconds=1) event = mock_api_client_instance.event_dict['1'][0][0] event['timestamp'] = future_event_time.isoformat().replace('+00:00', 'Z') session_service = mock_vertex_ai_session_service() session = await session_service.get_session( app_name='123', user_id='user', session_id='1' ) assert session is not None assert len(session.events) == 1 assert session.events[0].timestamp == future_event_time.timestamp() assert session.events[0].timestamp > session.last_update_time, ( 'Event timestamp should exceed session update_time to guard against' ' filtering.' ) @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') @pytest.mark.parametrize( 'mock_event_json', [MOCK_EVENT_WITH_OVERRIDE_JSON, MOCK_EVENT_WITH_OVERRIDE_JSON_2], ) async def test_get_session_from_raw_event( mock_api_client_instance: MockAsyncClient, mock_event_json, ) -> None: mock_api_client_instance.session_dict['6'] = MOCK_SESSION_WITH_OVERRIDE_JSON mock_api_client_instance.event_dict['6'] = ( copy.deepcopy(mock_event_json), None, ) session_service = mock_vertex_ai_session_service() session = await session_service.get_session( app_name='123', user_id='user_with_override', session_id='6' ) assert session is not None assert len(session.events) == 1 event = session.events[0] assert event.content.parts[0].text == 'raw_event_content' assert event.actions.transfer_to_agent == 'raw_event_agent' assert not event.partial assert event.turn_complete assert event.interrupted assert event.branch == 'raw_event_branch' assert event.error_code == '222' assert event.error_message == 'raw_event_error' @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_get_session_with_many_events(mock_api_client_instance): mock_api_client_instance.session_dict['5'] = MOCK_SESSION_JSON_5 mock_api_client_instance.event_dict['5'] = ( copy.deepcopy(MOCK_EVENTS_JSON_5), None, ) session_service = mock_vertex_ai_session_service() session = await session_service.get_session( app_name='123', user_id='user_with_many_events', session_id='5' ) assert session is not None assert len(session.events) == MANY_EVENTS_COUNT @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_get_session_with_num_recent_events_zero(): session_service = mock_vertex_ai_session_service() session = await session_service.get_session( app_name='123', user_id='user', session_id='2', config=GetSessionConfig(num_recent_events=0), ) assert session is not None assert len(session.events) == 0 @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_list_sessions(): session_service = mock_vertex_ai_session_service() sessions = await session_service.list_sessions(app_name='123', user_id='user') assert len(sessions.sessions) == 2 assert sessions.sessions[0].id == '1' assert sessions.sessions[1].id == '2' @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_list_sessions_with_pagination(): session_service = mock_vertex_ai_session_service() sessions = await session_service.list_sessions( app_name='123', user_id='user_with_pages' ) assert len(sessions.sessions) == 2 assert sessions.sessions[0].id == 'page1' assert sessions.sessions[1].id == 'page2' @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_list_sessions_all_users(): session_service = mock_vertex_ai_session_service() sessions = await session_service.list_sessions(app_name='123', user_id=None) assert len(sessions.sessions) == 5 assert {s.id for s in sessions.sessions} == { '1', '2', '3', 'page1', 'page2', } @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') @pytest.mark.parametrize( ('payload', 'expected_filter'), [ ( 'attacker" OR user_id!=""', 'user_id="attacker\\" OR user_id!=\\"\\""', ), ('\\', 'user_id="\\\\"'), ('', 'user_id=""'), ], ) async def test_list_sessions_quotes_user_id_filter( mock_api_client_instance, payload, expected_filter ): session_service = mock_vertex_ai_session_service() sessions = await session_service.list_sessions( app_name='123', user_id=payload ) assert sessions.sessions == [] assert mock_api_client_instance.last_list_sessions_config == { 'filter': expected_filter } @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_create_session(): session_service = mock_vertex_ai_session_service() state = {'key': 'value'} session = await session_service.create_session( app_name='123', user_id='user', state=state ) assert session.state == state assert session.app_name == '123' assert session.user_id == 'user' assert session.last_update_time is not None session_id = session.id assert session == await session_service.get_session( app_name='123', user_id='user', session_id=session_id ) @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') @pytest.mark.parametrize('session_id', ['1', 'abc123']) async def test_create_session_with_custom_session_id( mock_api_client_instance: MockAsyncClient, session_id: str ): session_service = mock_vertex_ai_session_service() mock_api_client_instance.event_dict[session_id] = ( [], None, ) session = await session_service.create_session( app_name='123', user_id='user', session_id=session_id ) assert session.id == session_id assert session.app_name == '123' assert session.user_id == 'user' assert session.last_update_time is not None assert session == await session_service.get_session( app_name='123', user_id='user', session_id=session_id ) @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_create_session_with_custom_config(mock_api_client_instance): session_service = mock_vertex_ai_session_service() expire_time = '2025-12-12T12:12:12.123456Z' await session_service.create_session( app_name='123', user_id='user', expire_time=expire_time ) assert ( mock_api_client_instance.last_create_session_config['expire_time'] == expire_time ) @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_create_session_with_ttl(mock_api_client_instance): session_service = mock_vertex_ai_session_service() ttl = '7200s' await session_service.create_session(app_name='123', user_id='user', ttl=ttl) assert mock_api_client_instance.last_create_session_config['ttl'] == ttl @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_create_session_with_ttl_and_expire_time_raises_value_error( mock_api_client_instance, ): session_service = mock_vertex_ai_session_service() with pytest.raises( ValueError, match="Cannot specify both 'ttl' and 'expire_time' simultaneously.", ): await session_service.create_session( app_name='123', user_id='user', ttl='7200s', expire_time='2025-12-12T12:12:12.123456Z', ) @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_create_session_with_ttl_none_and_expire_time_none_does_not_raise( mock_api_client_instance, ): session_service = mock_vertex_ai_session_service() # None means "not set"; passing both as None must not raise. await session_service.create_session( app_name='123', user_id='user', ttl=None, expire_time=None ) @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_append_event(): session_service = mock_vertex_ai_session_service() session_before_append = await session_service.get_session( app_name='123', user_id='user', session_id='1' ) event_to_append = Event( invocation_id='new_invocation', author='model', timestamp=1734005533.0, content=genai_types.Content(parts=[genai_types.Part(text='new_content')]), actions=EventActions( transfer_to_agent='another_agent', state_delta={'new_key': 'new_value'}, skip_summarization=True, requested_auth_configs={ 'test_auth': AuthConfig( auth_scheme=auth_schemes.OAuth2( flows=openapi_models.OAuthFlows( implicit=openapi_models.OAuthFlowImplicit( authorizationUrl='http://test.com/auth', scopes={}, ) ) ), ), }, ), error_code='1', error_message='test_error', branch='test_branch', custom_metadata={'custom': 'data'}, long_running_tool_ids={'tool2'}, input_transcription=genai_types.Transcription( text='test_input_transcription' ), output_transcription=genai_types.Transcription( text='test_output_transcription' ), model_version='test_model_version', avg_logprobs=0.5, logprobs_result=genai_types.LogprobsResult( chosen_candidates=[ genai_types.LogprobsResultCandidate( log_probability=0.5, token='test_token', token_id=0, ) ] ), cache_metadata=CacheMetadata( cache_name='test_cache_name', expire_time=( datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(minutes=30) ).timestamp(), fingerprint='test_fingerprint', invocations_used=1, contents_count=1, ), citation_metadata=genai_types.CitationMetadata( citations=[ genai_types.Citation( uri='http://test.com', title='test_title', ) ] ), ) await session_service.append_event(session_before_append, event_to_append) retrieved_session = await session_service.get_session( app_name='123', user_id='user', session_id='1' ) assert len(retrieved_session.events) == 2 event_to_append.id = retrieved_session.events[1].id assert retrieved_session.events[1] == event_to_append @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_append_event_strips_unsupported_part_metadata( mock_api_client_instance: MockAsyncClient, ) -> None: """part_metadata must not reach the Sessions API (#6014). ``Part.part_metadata`` is a Gemini Developer API-only field; the Vertex AI Agent Engine Sessions ``appendEvent`` API rejects it with 400 INVALID_ARGUMENT ("Unknown name \"part_metadata\""). It must be dropped from both the ``content`` and ``raw_event`` payloads, while the part text is preserved. """ session_service = mock_vertex_ai_session_service() session = await session_service.get_session( app_name='123', user_id='user', session_id='1' ) event_to_append = Event( invocation_id='inv_part_metadata', author='user', timestamp=1734005533.0, content=genai_types.Content( parts=[ genai_types.Part( text='hello', part_metadata={'source': 'portal'} ), genai_types.Part(text='world', part_metadata={'n': 1}), ], ), ) await session_service.append_event(session, event_to_append) appended = mock_api_client_instance.event_dict['1'][0][-1] for part in appended['content']['parts']: assert 'part_metadata' not in part assert 'partMetadata' not in part for part in appended['raw_event']['content']['parts']: assert 'part_metadata' not in part assert 'partMetadata' not in part assert [p['text'] for p in appended['content']['parts']] == ['hello', 'world'] @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_append_event_with_part_metadata_round_trips( mock_api_client_instance: MockAsyncClient, ) -> None: """Reconstruction side of #6014: an event carrying part_metadata appends and reads back without error. part_metadata is dropped (unsupported on Vertex), but the session round-trips and the part text is preserved. """ session_service = mock_vertex_ai_session_service() session = await session_service.get_session( app_name='123', user_id='user', session_id='1' ) event_to_append = Event( invocation_id='inv_part_metadata_rt', author='user', timestamp=1734005533.0, content=genai_types.Content( role='user', parts=[ genai_types.Part(text='hello', part_metadata={'source': 'portal'}) ], ), ) await session_service.append_event(session, event_to_append) retrieved = await session_service.get_session( app_name='123', user_id='user', session_id='1' ) appended = next( e for e in retrieved.events if e.invocation_id == 'inv_part_metadata_rt' ) assert appended.content is not None assert appended.content.parts[0].text == 'hello' assert appended.content.parts[0].part_metadata is None @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_append_event_with_compaction(): """Compaction data round-trips through append_event and get_session.""" session_service = mock_vertex_ai_session_service() session = await session_service.get_session( app_name='123', user_id='user', session_id='1' ) assert session is not None compaction = EventCompaction( start_timestamp=1000.0, end_timestamp=2000.0, compacted_content=genai_types.Content( parts=[genai_types.Part(text='compacted summary')] ), ) event_to_append = Event( invocation_id='compaction_invocation', author='model', timestamp=1734005534.0, actions=EventActions(compaction=compaction), ) await session_service.append_event(session, event_to_append) retrieved_session = await session_service.get_session( app_name='123', user_id='user', session_id='1' ) assert retrieved_session is not None appended_event = retrieved_session.events[-1] assert appended_event.actions.compaction is not None assert appended_event.actions.compaction.start_timestamp == 1000.0 assert appended_event.actions.compaction.end_timestamp == 2000.0 assert appended_event.actions.compaction.compacted_content.parts[0].text == ( 'compacted summary' ) # custom_metadata should remain None when only compaction was stored assert appended_event.custom_metadata is None @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_append_event_with_compaction_and_custom_metadata(): """Both compaction and user custom_metadata survive the round-trip.""" session_service = mock_vertex_ai_session_service() session = await session_service.get_session( app_name='123', user_id='user', session_id='1' ) assert session is not None compaction = EventCompaction( start_timestamp=100.0, end_timestamp=200.0, compacted_content=genai_types.Content( parts=[genai_types.Part(text='summary')] ), ) event_to_append = Event( invocation_id='compaction_and_meta_invocation', author='model', timestamp=1734005535.0, actions=EventActions(compaction=compaction), custom_metadata={'user_key': 'user_value'}, ) await session_service.append_event(session, event_to_append) retrieved_session = await session_service.get_session( app_name='123', user_id='user', session_id='1' ) assert retrieved_session is not None appended_event = retrieved_session.events[-1] # Compaction is restored assert appended_event.actions.compaction is not None assert appended_event.actions.compaction.start_timestamp == 100.0 assert appended_event.actions.compaction.end_timestamp == 200.0 # User custom_metadata is preserved without the internal _compaction key assert appended_event.custom_metadata == {'user_key': 'user_value'} assert '_compaction' not in (appended_event.custom_metadata or {}) @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_append_event_with_usage_metadata(): """usage_metadata round-trips through append_event and get_session.""" session_service = mock_vertex_ai_session_service() session = await session_service.get_session( app_name='123', user_id='user', session_id='1' ) assert session is not None event_to_append = Event( invocation_id='usage_invocation', author='model', timestamp=1734005536.0, usage_metadata=genai_types.GenerateContentResponseUsageMetadata( prompt_token_count=150, candidates_token_count=50, total_token_count=200, ), ) await session_service.append_event(session, event_to_append) retrieved_session = await session_service.get_session( app_name='123', user_id='user', session_id='1' ) assert retrieved_session is not None appended_event = retrieved_session.events[-1] assert appended_event.usage_metadata is not None assert appended_event.usage_metadata.prompt_token_count == 150 assert appended_event.usage_metadata.candidates_token_count == 50 assert appended_event.usage_metadata.total_token_count == 200 # custom_metadata should remain None when only usage_metadata was stored assert appended_event.custom_metadata is None @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_append_event_with_usage_metadata_and_custom_metadata(): """Both usage_metadata and user custom_metadata survive the round-trip.""" session_service = mock_vertex_ai_session_service() session = await session_service.get_session( app_name='123', user_id='user', session_id='1' ) assert session is not None event_to_append = Event( invocation_id='usage_and_meta_invocation', author='model', timestamp=1734005537.0, usage_metadata=genai_types.GenerateContentResponseUsageMetadata( prompt_token_count=300, total_token_count=400, ), custom_metadata={'my_key': 'my_value'}, ) await session_service.append_event(session, event_to_append) retrieved_session = await session_service.get_session( app_name='123', user_id='user', session_id='1' ) assert retrieved_session is not None appended_event = retrieved_session.events[-1] # usage_metadata is restored assert appended_event.usage_metadata is not None assert appended_event.usage_metadata.prompt_token_count == 300 assert appended_event.usage_metadata.total_token_count == 400 # User custom_metadata is preserved without internal keys assert appended_event.custom_metadata == {'my_key': 'my_value'} assert '_usage_metadata' not in (appended_event.custom_metadata or {}) @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_append_event_with_usage_metadata_and_compaction(): """usage_metadata, compaction, and user custom_metadata all coexist.""" session_service = mock_vertex_ai_session_service() session = await session_service.get_session( app_name='123', user_id='user', session_id='1' ) assert session is not None compaction = EventCompaction( start_timestamp=500.0, end_timestamp=600.0, compacted_content=genai_types.Content( parts=[genai_types.Part(text='compacted')] ), ) event_to_append = Event( invocation_id='all_three_invocation', author='model', timestamp=1734005538.0, actions=EventActions(compaction=compaction), usage_metadata=genai_types.GenerateContentResponseUsageMetadata( prompt_token_count=1000, candidates_token_count=250, total_token_count=1250, ), custom_metadata={'extra': 'info'}, ) await session_service.append_event(session, event_to_append) retrieved_session = await session_service.get_session( app_name='123', user_id='user', session_id='1' ) assert retrieved_session is not None appended_event = retrieved_session.events[-1] # Compaction is restored assert appended_event.actions.compaction is not None assert appended_event.actions.compaction.start_timestamp == 500.0 assert appended_event.actions.compaction.end_timestamp == 600.0 # usage_metadata is restored assert appended_event.usage_metadata is not None assert appended_event.usage_metadata.prompt_token_count == 1000 assert appended_event.usage_metadata.candidates_token_count == 250 assert appended_event.usage_metadata.total_token_count == 1250 # User custom_metadata is preserved without internal keys assert appended_event.custom_metadata == {'extra': 'info'} assert '_compaction' not in (appended_event.custom_metadata or {}) assert '_usage_metadata' not in (appended_event.custom_metadata or {}) @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_append_event_fallback_for_older_sdk(mock_api_client_instance): """Tests that append_event falls back to custom_metadata when SDK fails on raw_event.""" session_service = mock_vertex_ai_session_service() session = await session_service.get_session( app_name='123', user_id='user', session_id='1' ) assert session is not None compaction = EventCompaction( start_timestamp=1000.0, end_timestamp=2000.0, compacted_content=genai_types.Content( parts=[genai_types.Part(text='compacted summary')] ), ) event_to_append = Event( invocation_id='fallback_invocation', author='model', timestamp=1734005534.0, actions=EventActions(compaction=compaction), ) mock_client = mock_api_client_instance async def side_effect(name, author, invocation_id, timestamp, config): if 'raw_event' in config: # Trigger a real ValidationError since Pydantic V2 doesn't allow easy # instantiation class DummyModel(pydantic.BaseModel): a: int DummyModel(a='not an int') return await mock_client._append_event( name, author, invocation_id, timestamp, config ) mock_client.agent_engines.sessions.events.append.side_effect = side_effect await session_service.append_event(session, event_to_append) # Verify that it was written and restored correctly via custom_metadata retrieved_session = await session_service.get_session( app_name='123', user_id='user', session_id='1' ) appended_event = retrieved_session.events[-1] 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' )