Skip to content

Commit 1fddf90

Browse files
committed
fix(sessions): use async iteration for VertexAiSessionService.list_sessions pagination
The sessions_iterator returned by api_client.agent_engines.sessions.list() is an AsyncPager that only fetches additional pages when consumed via async iteration (__aiter__/__anext__). Using a synchronous `for` loop only retrieved the first page (~100 sessions), silently dropping all subsequent pages. Changed `for` to `async for` to enable proper async pagination, allowing all sessions to be returned regardless of count. Fixes #4302
1 parent ec660ed commit 1fddf90

2 files changed

Lines changed: 18 additions & 7 deletions

File tree

src/google/adk/sessions/vertex_ai_session_service.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -215,7 +215,7 @@ async def list_sessions(
215215
config=config,
216216
)
217217

218-
for api_session in sessions_iterator:
218+
async for api_session in sessions_iterator:
219219
sessions.append(
220220
Session(
221221
app_name=app_name,

tests/unittests/sessions/test_vertex_ai_session_service.py

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -264,6 +264,17 @@ async def to_async_iterator(data):
264264
yield item
265265

266266

267+
class AsyncIterableList(list):
268+
"""A list that also supports async iteration, mimicking AsyncPager behavior."""
269+
270+
def __aiter__(self):
271+
return self._async_iter()
272+
273+
async def _async_iter(self):
274+
for item in self:
275+
yield item
276+
277+
267278
class MockAsyncClient:
268279
"""Mocks the API Client."""
269280

@@ -300,20 +311,20 @@ async def _list_sessions(self, name: str, config: dict[str, Any]):
300311
if user_id_match:
301312
user_id = user_id_match.group(1)
302313
if user_id == 'user_with_pages':
303-
return [
314+
return AsyncIterableList([
304315
_convert_to_object(MOCK_SESSION_JSON_PAGE1),
305316
_convert_to_object(MOCK_SESSION_JSON_PAGE2),
306-
]
307-
return [
317+
])
318+
return AsyncIterableList([
308319
_convert_to_object(session)
309320
for session in self.session_dict.values()
310321
if session['user_id'] == user_id
311-
]
322+
])
312323

313324
# No user filter, return all sessions
314-
return [
325+
return AsyncIterableList([
315326
_convert_to_object(session) for session in self.session_dict.values()
316-
]
327+
])
317328

318329
async def _delete_session(self, name: str):
319330
session_id = name.split('/')[-1]

0 commit comments

Comments
 (0)