From 758d337c76d877e3174c35f06551cc9beb1def06 Mon Sep 17 00:00:00 2001 From: Anmol Jaiswal Date: Tue, 10 Feb 2026 21:18:31 -0800 Subject: [PATCH] fix(sessions): use async iteration for VertexAiSessionService.list_sessions pagination MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merge https://github.com/google/adk-python/pull/4435 ### Link to Issue or Description of Change - Closes: #4302 **Problem:** `VertexAiSessionService.list_sessions()` only returns the first ~100 sessions. The `sessions_iterator` from `api_client.agent_engines.sessions.list()` is an `AsyncPager` — it implements `__aiter__`/`__anext__` for fetching subsequent pages, but the code uses a plain `for` loop which only calls `__iter__`/`__next__`, so it never fetches beyond the first page. **Solution:** Changed `for api_session in sessions_iterator` to `async for api_session in sessions_iterator` so the `AsyncPager` actually paginates. Updated the test mock to return an `AsyncIterableList` (supports both sync and async iteration) instead of a bare list, so the tests properly simulate real `AsyncPager` behaviour. ### Testing Plan **Unit Tests:** ``` $ pytest tests/unittests/sessions/ 115 passed, 1 warning in 2.25s ``` The existing `test_list_sessions`, `test_list_sessions_with_pagination`, and `test_list_sessions_all_users` all continue to pass with the updated mock. Co-authored-by: Liang Wu COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/4435 from anmolg1997:fix/vertex-ai-session-service-pagination 14c71b607ecbf2215f4b9ba6eb4b0ff6b9eaf740 PiperOrigin-RevId: 868466166 --- .../adk/sessions/vertex_ai_session_service.py | 2 +- .../sessions/test_vertex_ai_session_service.py | 14 +++++++------- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/src/google/adk/sessions/vertex_ai_session_service.py b/src/google/adk/sessions/vertex_ai_session_service.py index f5cafee4..1837a907 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 @@ class VertexAiSessionService(BaseSessionService): 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 12a11a93..8c77f194 100644 --- a/tests/unittests/sessions/test_vertex_ai_session_service.py +++ b/tests/unittests/sessions/test_vertex_ai_session_service.py @@ -300,20 +300,20 @@ class MockAsyncClient: if user_id_match: user_id = user_id_match.group(1) if user_id == 'user_with_pages': - return [ + return to_async_iterator([ _convert_to_object(MOCK_SESSION_JSON_PAGE1), _convert_to_object(MOCK_SESSION_JSON_PAGE2), - ] - return [ + ]) + 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 [ - _convert_to_object(session) for session in self.session_dict.values() - ] + 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]