Updating session service to address various concerns

These updates are derived from @anmolg1997 PR on the community repo: google/adk-python-community#104

Author: ScottMansfield

src/google/adk/integrations/firestore/firestore_session_service.py

@@ -21,7 +21,6 @@
 from typing import Any
 from typing import Optional
 
-from google.cloud import firestore
 from pydantic import BaseModel
 
 from ...events.event import Event
@@ -55,6 +54,14 @@ def __init__(
       root_collection: The root collection name. Defaults to 'adk-session' or
         the value of ADK_FIRESTORE_ROOT_COLLECTION env var.
     """
+    try:
+      from google.cloud import firestore
+    except ImportError as e:
+      raise ImportError(
+          "FirestoreSessionService requires google-cloud-firestore. "
+          "Install it with: pip install google-cloud-firestore"
+      ) from e
+
     self.client = client or firestore.AsyncClient()
     self.root_collection = (
          root_collection
@@ -66,6 +73,22 @@ def __init__(
     self.app_state_collection = DEFAULT_APP_STATE_COLLECTION
     self.user_state_collection = DEFAULT_USER_STATE_COLLECTION
 
+  @staticmethod
+  def _merge_state(
+      app_state: dict[str, Any],
+      user_state: dict[str, Any],
+      session_state: dict[str, Any],
+  ) -> dict[str, Any]:
+    """Merge app, user, and session states into a single state dictionary."""
+    import copy
+
+    merged_state = copy.deepcopy(session_state)
+    for key, value in app_state.items():
+      merged_state["_app_" + key] = value
+    for key, value in user_state.items():
+      merged_state["_user_" + key] = value
+    return merged_state
+
   def _get_sessions_ref(
       self, user_id: str
   ) -> firestore.AsyncCollectionReference:
@@ -84,6 +107,7 @@ async def create_session(
       session_id: Optional[str] = None,
   ) -> Session:
     """Creates a new session in Firestore."""
+    from google.cloud import firestore
     if not session_id:
       from ...platform import uuid as platform_uuid
 
@@ -202,17 +226,50 @@ async def list_sessions(
       )
       docs = await query.get()
 
+    # Fetch shared state once
+    app_ref = self.client.collection(self.app_state_collection).document(
+        app_name
+    )
+    app_doc = await app_ref.get()
+    app_state = app_doc.to_dict() if app_doc.exists else {}
+
+    user_states_map = {}
+    if user_id:
+      user_ref = (
+          self.client.collection(self.user_state_collection)
+          .document(app_name)
+          .collection("users")
+          .document(user_id)
+      )
+      user_doc = await user_ref.get()
+      if user_doc.exists:
+        user_states_map[user_id] = user_doc.to_dict()
+    else:
+      users_ref = (
+          self.client.collection(self.user_state_collection)
+          .document(app_name)
+          .collection("users")
+      )
+      users_docs = await users_ref.get()
+      for u_doc in users_docs:
+        user_states_map[u_doc.id] = u_doc.to_dict()
+
     sessions = []
     for doc in docs:
       data = doc.to_dict()
       if data:
+        u_id = data["userId"]
+        s_state = data.get("state", {})
+        u_state = user_states_map.get(u_id, {})
+        merged = self._merge_state(app_state, u_state, s_state)
+
         sessions.append(
             Session(
                 id=data["id"],
                 app_name=data["appName"],
                 user_id=data["userId"],
-                state={},  # Empty state for listing
-                events=[],  # Empty events for listing
+                state=merged,
+                events=[],
                 last_update_time=0.0,
             )
         )
@@ -226,17 +283,65 @@ async def delete_session(
     session_ref = self._get_sessions_ref(user_id).document(session_id)
 
     events_ref = session_ref.collection(self.events_collection)
-    events_docs = await events_ref.get()
-
+    
     batch = self.client.batch()
-    for event_doc in events_docs:
+    count = 0
+    async for event_doc in events_ref.stream():
       batch.delete(event_doc.reference)
-    await batch.commit()
+      count += 1
+      if count >= 500:
+        await batch.commit()
+        batch = self.client.batch()
+        count = 0
+    if count > 0:
+      await batch.commit()
 
     await session_ref.delete()
 
+  async def _update_app_state_transactional(
+      self, app_name: str, delta: dict[str, Any]
+  ) -> dict[str, Any]:
+    """Atomically applies delta to app state inside a transaction."""
+    from google.cloud import firestore
+    doc_ref = self.client.collection(self.app_state_collection).document(app_name)
+
+    @firestore.async_transactional
+    async def _txn(transaction):
+      snap = await doc_ref.get(transaction=transaction)
+      current = snap.to_dict() if snap.exists else {}
+      current.update(delta)
+      transaction.set(doc_ref, current, merge=True)
+      return current
+
+    transaction = self.client.transaction()
+    return await _txn(transaction)
+
+  async def _update_user_state_transactional(
+      self, app_name: str, user_id: str, delta: dict[str, Any]
+  ) -> dict[str, Any]:
+    """Atomically applies delta to user state inside a transaction."""
+    from google.cloud import firestore
+    doc_ref = (
+        self.client.collection(self.user_state_collection)
+        .document(app_name)
+        .collection("users")
+        .document(user_id)
+    )
+
+    @firestore.async_transactional
+    async def _txn(transaction):
+      snap = await doc_ref.get(transaction=transaction)
+      current = snap.to_dict() if snap.exists else {}
+      current.update(delta)
+      transaction.set(doc_ref, current, merge=True)
+      return current
+
+    transaction = self.client.transaction()
+    return await _txn(transaction)
+
   async def append_event(self, session: Session, event: Event) -> Event:
     """Appends an event to a session in Firestore."""
+    from google.cloud import firestore
     if event.partial:
       return event
 
@@ -259,26 +364,16 @@ async def append_event(self, session: Session, event: Event) -> Event:
         else:
           session_updates[key] = value
 
-      batch = self.client.batch()
-
       if app_updates:
-        app_ref = self.client.collection(self.app_state_collection).document(
-            session.app_name
-        )
-        batch.set(app_ref, app_updates, merge=True)
+        await self._update_app_state_transactional(session.app_name, app_updates)
 
       if user_updates:
-        user_ref = (
-            self.client.collection(self.user_state_collection)
-            .document(session.app_name)
-            .collection("users")
-            .document(session.user_id)
-        )
-        batch.set(user_ref, user_updates, merge=True)
+        await self._update_user_state_transactional(session.app_name, session.user_id, user_updates)
 
       for k, v in session_updates.items():
         session.state[k] = v
 
+      batch = self.client.batch()
       batch.update(
           session_ref,
           {

tests/unittests/integrations/firestore/test_firestore_database_runner.py

@@ -71,3 +71,22 @@ def test_create_firestore_runner_missing_bucket(mock_agent, monkeypatch):
       ValueError, match="Required property 'ADK_GCS_BUCKET_NAME' is not set"
   ):
     create_firestore_runner(mock_agent)
+
+
+def test_create_firestore_runner_with_root_collection(mock_agent, monkeypatch):
+  monkeypatch.setenv("ADK_GCS_BUCKET_NAME", "test_bucket")
+
+  with (
+      mock.patch(
+          "google.adk.integrations.firestore.firestore_database_runner.FirestoreSessionService"
+      ) as mock_session,
+      mock.patch("google.adk.integrations.firestore.firestore_database_runner.FirestoreMemoryService"),
+      mock.patch("google.adk.integrations.firestore.firestore_database_runner.GcsArtifactService"),
+  ):
+    runner = create_firestore_runner(
+        mock_agent, firestore_root_collection="custom_collection"
+    )
+
+    assert runner is not None
+    mock_session.assert_called_once_with(root_collection="custom_collection")
+

tests/unittests/integrations/firestore/test_firestore_memory_service.py

@@ -93,3 +93,75 @@ async def test_search_memory_with_results(mock_firestore_client):
   mock_firestore_client.collection_group.assert_called_with("events")
   collection_ref = mock_firestore_client.collection_group.return_value
   collection_ref.where.assert_called()
+
+
+@pytest.mark.asyncio
+async def test_search_memory_deduplication(mock_firestore_client):
+  service = FirestoreMemoryService(client=mock_firestore_client)
+  app_name = "test_app"
+  user_id = "test_user"
+  query = "quick fox"
+
+  event1 = Event(
+      invocation_id="test_inv1",
+      author="user",
+      content=types.Content(parts=[types.Part(text="quick fox jumps")]),
+      timestamp=1234567890.0,
+  )
+  event2 = Event(
+      invocation_id="test_inv2",
+      author="user",
+      content=types.Content(parts=[types.Part(text="quick fox jumps")]),
+      timestamp=1234567890.0,
+  )
+
+  doc_snapshot1 = mock.MagicMock()
+  doc_snapshot1.to_dict.return_value = {
+      "event_data": event1.model_dump(exclude_none=True, mode="json")
+  }
+
+  doc_snapshot2 = mock.MagicMock()
+  doc_snapshot2.to_dict.return_value = {
+      "event_data": event2.model_dump(exclude_none=True, mode="json")
+  }
+
+  get_mock = mock.AsyncMock(side_effect=[[doc_snapshot1], [doc_snapshot2]])
+
+  mock_firestore_client.collection_group.return_value.where.return_value.where.return_value.where.return_value.get = get_mock
+
+  response = await service.search_memory(
+      app_name=app_name, user_id=user_id, query=query
+  )
+
+  assert response.memories
+  assert len(response.memories) == 1
+  assert response.memories[0].author == "user"
+
+
+@pytest.mark.asyncio
+async def test_search_memory_parsing_error(mock_firestore_client, caplog):
+  service = FirestoreMemoryService(client=mock_firestore_client)
+  app_name = "test_app"
+  user_id = "test_user"
+  query = "quick"
+
+  doc_snapshot = mock_firestore_client.collection_group.return_value.where.return_value.where.return_value.where.return_value.get.return_value[0]
+  doc_snapshot.to_dict.return_value = {"event_data": "invalid_data"}
+
+  response = await service.search_memory(
+      app_name=app_name, user_id=user_id, query=query
+  )
+
+  assert not response.memories
+  assert "Failed to parse event from Firestore" in caplog.text
+
+
+@pytest.mark.asyncio
+async def test_search_memory_only_stop_words(mock_firestore_client):
+  service = FirestoreMemoryService(client=mock_firestore_client)
+  response = await service.search_memory(
+      app_name="test_app", user_id="test_user", query="the and or"
+  )
+  assert not response.memories
+  mock_firestore_client.collection_group.assert_not_called()
+