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]