fix(sessions): use async iteration for VertexAiSessionService.list_sessions pagination

[anmolg1997]

Link to Issue or Description of Change

• Closes: VertexAiSessionService does not paginate when invoking list_sessions #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.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[gemini-code-assist]

Summary of Changes

Hello @anmolg1997, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses a critical pagination issue within the VertexAiSessionService.list_sessions method, which previously failed to retrieve all available sessions beyond the initial page. By transitioning from a synchronous iteration to an asynchronous one, the service now correctly processes AsyncPager objects, ensuring comprehensive session data retrieval. The accompanying test updates provide robust validation for this fix, preventing future regressions.

Highlights

• Pagination Fix: Corrected a pagination bug in VertexAiSessionService.list_sessions that caused only the first page of sessions to be retrieved, silently dropping subsequent results.
• Asynchronous Iteration: Switched from a synchronous for loop to an async for loop to properly consume AsyncPager objects, ensuring all available pages of sessions are fetched.
• Test Mock Enhancement: Enhanced test mocks with a new AsyncIterableList class to accurately simulate asynchronous pagination behavior, improving test coverage and reliability.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• src/google/adk/sessions/vertex_ai_session_service.py
  • Changed the iteration over sessions_iterator from for to async for to enable full pagination.
• tests/unittests/sessions/test_vertex_ai_session_service.py
  • Introduced AsyncIterableList class to simulate AsyncPager behavior for testing purposes.
  • Modified the mock _list_sessions method to return AsyncIterableList instances, ensuring tests accurately reflect asynchronous pagination.

Activity

• All 19 VertexAI session service tests passed.
• All 115 session tests passed, indicating no regressions.
• The test mock was updated with AsyncIterableList to correctly simulate AsyncPager behavior.
• The pagination test (test_list_sessions_with_pagination) continues to pass.
• The author completed a self-review checklist, including reading CLA and contributing guidelines, updating documentation, adding tests, ensuring all tests pass, following coding style, running linting/formatting, rebasing, and reviewing their own code.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[adk-bot]

Response from ADK Triaging Agent

Hello @anmolg1997, thank you for your contribution!

It looks like the Contributor License Agreement (CLA) check has failed. Please sign the CLA so we can proceed with reviewing this PR.

Thanks!

[gemini-code-assist]

Code Review

This pull request correctly addresses a critical pagination bug in VertexAiSessionService.list_sessions by switching from a synchronous for loop to an async for loop, which is the proper way to consume an AsyncPager and retrieve all pages of results. The test suite has been thoughtfully updated to reflect this change, with the introduction of AsyncIterableList to accurately mock the behavior of the async pager. The changes are well-contained and the fix is solid. I have one minor suggestion to improve the conciseness of the new test helper class.

[anmolg1997]

I have signed the CLA - @googlebot please verify.

[adk-bot]

Thank you @anmolg1997 for your contribution! 🎉

Your changes have been successfully imported and merged via Copybara in commit 758d337.

Closing this PR as the changes are now in the main branch.