You've already forked adk-python
mirror of
https://github.com/encounter/adk-python.git
synced 2026-03-30 10:57:20 -07:00
fix(sessions): use async iteration for VertexAiSessionService.list_sessions pagination
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 <wuliang@google.com> COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/4435 from anmolg1997:fix/vertex-ai-session-service-pagination 14c71b607ecbf2215f4b9ba6eb4b0ff6b9eaf740 PiperOrigin-RevId: 868466166
This commit is contained in:
committed by
Copybara-Service
parent
3cf43e3842
commit
758d337c76
@@ -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,
|
||||
|
||||
@@ -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]
|
||||
|
||||
Reference in New Issue
Block a user