feat: Support return all sessions when user id is none

PiperOrigin-RevId: 819884236

Author: DeanChensj

src/google/adk/sessions/base_session_service.py

@@ -83,9 +83,18 @@ async def get_session(
 
   @abc.abstractmethod
   async def list_sessions(
-      self, *, app_name: str, user_id: str
+      self, *, app_name: str, user_id: Optional[str] = None
   ) -> ListSessionsResponse:
-    """Lists all the sessions."""
+    """Lists all the sessions for a user.
+
+    Args:
+      app_name: The name of the app.
+      user_id: The ID of the user. If not provided, lists all sessions for all
+        users.
+
+    Returns:
+      A ListSessionsResponse containing the sessions.
+    """
 
   @abc.abstractmethod
   async def delete_session(

src/google/adk/sessions/database_session_service.py

@@ -554,30 +554,42 @@ async def get_session(
 
   @override
   async def list_sessions(
-      self, *, app_name: str, user_id: str
+      self, *, app_name: str, user_id: Optional[str] = None
   ) -> ListSessionsResponse:
     with self.database_session_factory() as sql_session:
-      results = (
-          sql_session.query(StorageSession)
-          .filter(StorageSession.app_name == app_name)
-          .filter(StorageSession.user_id == user_id)
-          .all()
+      query = sql_session.query(StorageSession).filter(
+          StorageSession.app_name == app_name
       )
+      if user_id is not None:
+        query = query.filter(StorageSession.user_id == user_id)
+      results = query.all()
 
-      # Fetch states from storage
+      # Fetch app state from storage
       storage_app_state = sql_session.get(StorageAppState, (app_name))
-      storage_user_state = sql_session.get(
-          StorageUserState, (app_name, user_id)
-      )
-
       app_state = storage_app_state.state if storage_app_state else {}
-      user_state = storage_user_state.state if storage_user_state else {}
+
+      # Fetch user state(s) from storage
+      user_states_map = {}
+      if user_id is not None:
+        storage_user_state = sql_session.get(
+            StorageUserState, (app_name, user_id)
+        )
+        if storage_user_state:
+          user_states_map[user_id] = storage_user_state.state
+      else:
+        all_user_states_for_app = (
+            sql_session.query(StorageUserState)
+            .filter(StorageUserState.app_name == app_name)
+            .all()
+        )
+        for storage_user_state in all_user_states_for_app:
+          user_states_map[storage_user_state.user_id] = storage_user_state.state
 
       sessions = []
       for storage_session in results:
         session_state = storage_session.state
+        user_state = user_states_map.get(storage_session.user_id, {})
         merged_state = _merge_state(app_state, user_state, session_state)
-
         sessions.append(storage_session.to_session(state=merged_state))
       return ListSessionsResponse(sessions=sessions)
 

src/google/adk/sessions/in_memory_session_service.py

@@ -201,31 +201,41 @@ def _merge_state(
 
   @override
   async def list_sessions(
-      self, *, app_name: str, user_id: str
+      self, *, app_name: str, user_id: Optional[str] = None
   ) -> ListSessionsResponse:
     return self._list_sessions_impl(app_name=app_name, user_id=user_id)
 
   def list_sessions_sync(
-      self, *, app_name: str, user_id: str
+      self, *, app_name: str, user_id: Optional[str] = None
   ) -> ListSessionsResponse:
     logger.warning('Deprecated. Please migrate to the async method.')
     return self._list_sessions_impl(app_name=app_name, user_id=user_id)
 
   def _list_sessions_impl(
-      self, *, app_name: str, user_id: str
+      self, *, app_name: str, user_id: Optional[str] = None
   ) -> ListSessionsResponse:
     empty_response = ListSessionsResponse()
     if app_name not in self.sessions:
       return empty_response
-    if user_id not in self.sessions[app_name]:
+    if user_id is not None and user_id not in self.sessions[app_name]:
       return empty_response
 
     sessions_without_events = []
-    for session in self.sessions[app_name][user_id].values():
-      copied_session = copy.deepcopy(session)
-      copied_session.events = []
-      copied_session = self._merge_state(app_name, user_id, copied_session)
-      sessions_without_events.append(copied_session)
+
+    if user_id is None:
+      for user_id in self.sessions[app_name]:
+        for session_id in self.sessions[app_name][user_id]:
+          session = self.sessions[app_name][user_id][session_id]
+          copied_session = copy.deepcopy(session)
+          copied_session.events = []
+          copied_session = self._merge_state(app_name, user_id, copied_session)
+          sessions_without_events.append(copied_session)
+    else:
+      for session in self.sessions[app_name][user_id].values():
+        copied_session = copy.deepcopy(session)
+        copied_session.events = []
+        copied_session = self._merge_state(app_name, user_id, copied_session)
+        sessions_without_events.append(copied_session)
     return ListSessionsResponse(sessions=sessions_without_events)
 
   @override

src/google/adk/sessions/vertex_ai_session_service.py

@@ -200,22 +200,25 @@ async def get_session(
 
   @override
   async def list_sessions(
-      self, *, app_name: str, user_id: str
+      self, *, app_name: str, user_id: Optional[str] = None
   ) -> ListSessionsResponse:
     reasoning_engine_id = self._get_reasoning_engine_id(app_name)
     api_client = self._get_api_client()
 
     sessions = []
+    config = {}
+    if user_id is not None:
+      config['filter'] = f'user_id="{user_id}"'
     sessions_iterator = api_client.agent_engines.sessions.list(
         name=f'reasoningEngines/{reasoning_engine_id}',
-        config={'filter': f'user_id="{user_id}"'},
+        config=config,
     )
 
     for api_session in sessions_iterator:
       sessions.append(
           Session(
               app_name=app_name,
-              user_id=user_id,
+              user_id=api_session.user_id,
               id=api_session.name.split('/')[-1],
               state=getattr(api_session, 'session_state', None) or {},
               last_update_time=api_session.update_time.timestamp(),

tests/unittests/sessions/test_session_service.py

@@ -116,9 +116,70 @@ async def test_create_and_list_sessions(service_type):
       app_name=app_name, user_id=user_id
   )
   sessions = list_sessions_response.sessions
-  for i in range(len(sessions)):
-    assert sessions[i].id == session_ids[i]
-    assert sessions[i].state == {'key': 'value' + session_ids[i]}
+  assert len(sessions) == len(session_ids)
+  assert {s.id for s in sessions} == set(session_ids)
+  for session in sessions:
+    assert session.state == {'key': 'value' + session.id}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE]
+)
+async def test_list_sessions_all_users(service_type):
+  session_service = get_session_service(service_type)
+  app_name = 'my_app'
+  user_id_1 = 'user1'
+  user_id_2 = 'user2'
+
+  await session_service.create_session(
+      app_name=app_name,
+      user_id=user_id_1,
+      session_id='session1a',
+      state={'key': 'value1a'},
+  )
+  await session_service.create_session(
+      app_name=app_name,
+      user_id=user_id_1,
+      session_id='session1b',
+      state={'key': 'value1b'},
+  )
+  await session_service.create_session(
+      app_name=app_name,
+      user_id=user_id_2,
+      session_id='session2a',
+      state={'key': 'value2a'},
+  )
+
+  # List sessions for user1
+  list_sessions_response_1 = await session_service.list_sessions(
+      app_name=app_name, user_id=user_id_1
+  )
+  sessions_1 = list_sessions_response_1.sessions
+  assert len(sessions_1) == 2
+  assert {s.id for s in sessions_1} == {'session1a', 'session1b'}
+  for session in sessions_1:
+    if session.id == 'session1a':
+      assert session.state == {'key': 'value1a'}
+    else:
+      assert session.state == {'key': 'value1b'}
+
+  # List sessions for user2
+  list_sessions_response_2 = await session_service.list_sessions(
+      app_name=app_name, user_id=user_id_2
+  )
+  sessions_2 = list_sessions_response_2.sessions
+  assert len(sessions_2) == 1
+  assert sessions_2[0].id == 'session2a'
+  assert sessions_2[0].state == {'key': 'value2a'}
+
+  # List sessions for all users
+  list_sessions_response_all = await session_service.list_sessions(
+      app_name=app_name, user_id=None
+  )
+  sessions_all = list_sessions_response_all.sessions
+  assert len(sessions_all) == 3
+  assert {s.id for s in sessions_all} == {'session1a', 'session1b', 'session2a'}
 
 
 @pytest.mark.asyncio

tests/unittests/sessions/test_vertex_ai_session_service.py

@@ -252,19 +252,22 @@ def _get_session(self, name: str):
   def _list_sessions(self, name: str, config: dict[str, Any]):
     filter_val = config.get('filter', '')
     user_id_match = re.search(r'user_id="([^"]+)"', filter_val)
-    if not user_id_match:
-      raise ValueError(f'Could not find user_id in filter: {filter_val}')
-    user_id = user_id_match.group(1)
-
-    if user_id == 'user_with_pages':
+    if user_id_match:
+      user_id = user_id_match.group(1)
+      if user_id == 'user_with_pages':
+        return [
+            _convert_to_object(MOCK_SESSION_JSON_PAGE1),
+            _convert_to_object(MOCK_SESSION_JSON_PAGE2),
+        ]
       return [
-          _convert_to_object(MOCK_SESSION_JSON_PAGE1),
-          _convert_to_object(MOCK_SESSION_JSON_PAGE2),
+          _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()
-        if session['user_id'] == user_id
+        _convert_to_object(session) for session in self.session_dict.values()
     ]
 
   def _delete_session(self, name: str):
@@ -475,6 +478,15 @@ async def test_list_sessions_with_pagination():
   assert sessions.sessions[1].id == 'page2'
 
 
+@pytest.mark.asyncio
+@pytest.mark.usefixtures('mock_get_api_client')
+async def test_list_sessions_all_users():
+  session_service = mock_vertex_ai_session_service()
+  sessions = await session_service.list_sessions(app_name='123', user_id=None)
+  assert len(sessions.sessions) == 5
+  assert {s.id for s in sessions.sessions} == {'1', '2', '3', 'page1', 'page2'}
+
+
 @pytest.mark.asyncio
 @pytest.mark.usefixtures('mock_get_api_client')
 async def test_create_session():