diff --git a/src/google/adk/sessions/vertex_ai_session_service.py b/src/google/adk/sessions/vertex_ai_session_service.py index f5cafee467..1837a907eb 100644 --- a/src/google/adk/sessions/vertex_ai_session_service.py +++ b/src/google/adk/sessions/vertex_ai_session_service.py @@ -215,7 +215,7 @@ async def list_sessions( config=config, ) - for api_session in sessions_iterator: + async for api_session in sessions_iterator: sessions.append( Session( app_name=app_name, diff --git a/tests/unittests/sessions/test_vertex_ai_session_service.py b/tests/unittests/sessions/test_vertex_ai_session_service.py index 12a11a93a2..dc6ef08b4a 100644 --- a/tests/unittests/sessions/test_vertex_ai_session_service.py +++ b/tests/unittests/sessions/test_vertex_ai_session_service.py @@ -264,6 +264,14 @@ async def to_async_iterator(data): yield item +class AsyncIterableList(list): + """A list that also supports async iteration, mimicking AsyncPager behavior.""" + + async def __aiter__(self): + for item in self: + yield item + + class MockAsyncClient: """Mocks the API Client.""" @@ -300,20 +308,20 @@ async def _list_sessions(self, name: str, config: dict[str, Any]): if user_id_match: user_id = user_id_match.group(1) if user_id == 'user_with_pages': - return [ + return AsyncIterableList([ _convert_to_object(MOCK_SESSION_JSON_PAGE1), _convert_to_object(MOCK_SESSION_JSON_PAGE2), - ] - return [ + ]) + return AsyncIterableList([ _convert_to_object(session) for session in self.session_dict.values() if session['user_id'] == user_id - ] + ]) # No user filter, return all sessions - return [ + return AsyncIterableList([ _convert_to_object(session) for session in self.session_dict.values() - ] + ]) async def _delete_session(self, name: str): session_id = name.split('/')[-1]