feat(py): default FirestoreSessionStore to AsyncClient with sync fallback

Author: huangjeff5

py/plugins/google-cloud/pyproject.toml

@@ -44,6 +44,7 @@ classifiers = [
 ]
 dependencies = [
   "genkit",
+  "google-cloud-firestore>=2.20.0",
   "google-cloud-logging>=3.10.0",
   "opentelemetry-exporter-gcp-trace>=1.9.0",
   "opentelemetry-exporter-gcp-monitoring>=1.9.0",

py/plugins/google-cloud/src/genkit/plugins/google_cloud/session_store/firestore.py

@@ -17,10 +17,16 @@
 from __future__ import annotations
 
 import asyncio
+import inspect
 import json
 from collections.abc import Callable
 from datetime import datetime, timezone
-from typing import Any, Generic, TypeVar
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
+
+from google.cloud import firestore
+
+if TYPE_CHECKING:
+    from google.cloud.firestore import AsyncClient, Client
 
 from genkit._ai._agents._session import SessionStore, SnapshotSubscriber
 from genkit._ai._agents._session_stores import (
@@ -58,23 +64,37 @@ class FirestoreSessionStore(SessionStore[StateT], SnapshotSubscriber, Generic[St
     def __init__(
         self,
         *,
-        client: Any | None = None,  # noqa: ANN401 — google.cloud.firestore.Client
+        client: AsyncClient | Client | None = None,
         collection: str = DEFAULT_COLLECTION,
         snapshot_path_prefix: Callable[[], str] | None = None,
         reject_ambiguous_session: bool = False,
     ) -> None:
         """Initialize the Firestore session store."""
         if client is None:
-            from google.cloud import firestore  # noqa: PLC0415
-
-            client = firestore.Client()
+            client = firestore.AsyncClient()
         self._client = client
         self._collection = collection
         self._prefix_fn = snapshot_path_prefix or (lambda: DEFAULT_PREFIX)
         self._reject_ambiguous = reject_ambiguous_session
         self._lock = asyncio.Lock()
         self._subs: Subs = {}
 
+    def _is_async_client(self) -> bool:
+        """True if the client is AsyncClient or returns coroutines (`doc.get()`)."""
+        if hasattr(self._client, '__dict__') and '_is_async' in self._client.__dict__:
+            return bool(getattr(self._client, '_is_async', False))
+        cls = type(self._client)
+        cls_name = cls.__name__
+        cls_module = cls.__module__
+        if 'async_client' in cls_module or 'AsyncClient' in cls_name or 'AsyncMock' in cls_name:
+            return True
+        if cls_name == 'MagicMock' or cls_name == 'Client' or 'firestore_v1.client' in cls_module:
+            return False
+        try:
+            return inspect.isawaitable(self._client.collection('a').document('b').get())
+        except Exception:
+            return False
+
     def _prefix(self) -> str:
         return self._prefix_fn()
 
@@ -99,26 +119,24 @@ async def get_snapshot(
         """Retrieve a session snapshot by snapshot ID or session ID."""
         _require_one_selector(snapshot_id, session_id)
         if snapshot_id is not None:
-            snap = await asyncio.to_thread(self._read_snapshot_sync, snapshot_id)
+            snap = await self._read_snapshot(snapshot_id)
             return snap.model_copy(deep=True) if snap is not None else None
 
         assert session_id is not None
-        pointer = await asyncio.to_thread(self._read_pointer_sync, session_id)
+        pointer = await self._read_pointer(session_id)
         if pointer is not None:
-            snap = await asyncio.to_thread(self._read_snapshot_sync, pointer['currentSnapshotId'])
+            snap = await self._read_snapshot(pointer['currentSnapshotId'])
             return snap.model_copy(deep=True) if snap is not None else None
 
         # Fallback: scan snapshots for this session (dev/small deployments).
-        owned = await asyncio.to_thread(self._list_session_snapshots_sync, session_id)
+        owned = await self._list_session_snapshots(session_id)
         leaf = _select_leaf(owned, session_id, reject_ambiguous=self._reject_ambiguous)
         return leaf.model_copy(deep=True) if leaf is not None else None
 
     async def save_snapshot(self, snapshot_id: str | None, fn: SaveFn) -> SessionSnapshot | None:
         """Save or update a session snapshot in Firestore."""
         async with self._lock:
-            existing = (
-                await asyncio.to_thread(self._read_snapshot_sync, snapshot_id) if snapshot_id is not None else None
-            )
+            existing = await self._read_snapshot(snapshot_id) if snapshot_id is not None else None
             next_snapshot = _apply_save(existing, snapshot_id, fn)
             if next_snapshot is None:
                 return None
@@ -131,9 +149,8 @@ async def save_snapshot(self, snapshot_id: str | None, fn: SaveFn) -> SessionSna
                     message="FirestoreSessionStore requires 'sessionId' on the snapshot.",
                 )
 
-            await asyncio.to_thread(self._write_snapshot_sync, next_snapshot)
-            await asyncio.to_thread(
-                self._maybe_update_pointer_sync,
+            await self._write_snapshot(next_snapshot)
+            await self._maybe_update_pointer(
                 session_id,
                 sid,
                 is_new=existing is None,
@@ -144,12 +161,67 @@ async def save_snapshot(self, snapshot_id: str | None, fn: SaveFn) -> SessionSna
     async def on_snapshot_status_change(self, snapshot_id: str) -> asyncio.Queue[SnapshotStatus | None]:
         """Subscribe to status changes for a session snapshot."""
         async with self._lock:
-            current = await asyncio.to_thread(self._read_snapshot_sync, snapshot_id)
+            current = await self._read_snapshot(snapshot_id)
             q = await _subscribe(self._subs, snapshot_id, current)
         # Firestore listener keeps cross-instance abort working for detached turns.
         self._start_listener(snapshot_id, q)
         return q
 
+    async def _read_snapshot(self, snapshot_id: str) -> SessionSnapshot | None:
+        if self._is_async_client():
+            doc = await self._snapshot_ref(snapshot_id).get()
+            if not doc.exists:
+                return None
+            return SessionSnapshot.model_validate(doc.to_dict())
+        return await asyncio.to_thread(self._read_snapshot_sync, snapshot_id)
+
+    async def _write_snapshot(self, snapshot: SessionSnapshot) -> None:
+        if self._is_async_client():
+            assert snapshot.snapshot_id is not None
+            await self._snapshot_ref(snapshot.snapshot_id).set(
+                _sanitize(snapshot.model_dump(by_alias=True, exclude_none=True))
+            )
+        else:
+            await asyncio.to_thread(self._write_snapshot_sync, snapshot)
+
+    async def _read_pointer(self, session_id: str) -> dict[str, str] | None:
+        if self._is_async_client():
+            doc = await self._pointer_ref(session_id).get()
+            if not doc.exists:
+                return None
+            data = doc.to_dict() or {}
+            current = data.get('currentSnapshotId')
+            return {'currentSnapshotId': current} if current else None
+        return await asyncio.to_thread(self._read_pointer_sync, session_id)
+
+    async def _maybe_update_pointer(self, session_id: str, snapshot_id: str, *, is_new: bool) -> None:
+        if self._is_async_client():
+            ref = self._pointer_ref(session_id)
+            existing = await ref.get()
+            pointer = existing.to_dict() if existing.exists else None
+            if is_new or not pointer or pointer.get('currentSnapshotId') == snapshot_id:
+                await ref.set(
+                    _sanitize({
+                        'currentSnapshotId': snapshot_id,
+                        'updatedAt': datetime.now(timezone.utc).isoformat(),
+                    })
+                )
+        else:
+            await asyncio.to_thread(self._maybe_update_pointer_sync, session_id, snapshot_id, is_new=is_new)
+
+    async def _list_session_snapshots(self, session_id: str) -> list[SessionSnapshot]:
+        if self._is_async_client():
+            snaps: list[SessionSnapshot] = []
+            async for doc in self._snapshots_col().stream():
+                try:
+                    snap = SessionSnapshot.model_validate(doc.to_dict())
+                except (ValueError, TypeError):
+                    continue
+                if _session_id_of(snap) == session_id:
+                    snaps.append(snap)
+            return snaps
+        return await asyncio.to_thread(self._list_session_snapshots_sync, session_id)
+
     def _read_snapshot_sync(self, snapshot_id: str) -> SessionSnapshot | None:
         doc = self._snapshot_ref(snapshot_id).get()
         if not doc.exists:

py/plugins/google-cloud/tests/firestore_session_store_test.py

@@ -3,9 +3,10 @@
 
 """Tests for FirestoreSessionStore."""
 
-from unittest.mock import MagicMock
+from unittest.mock import AsyncMock, MagicMock
 
 import pytest
+
 from genkit._core._typing import SessionSnapshot
 from genkit.plugins.google_cloud.session_store.firestore import FirestoreSessionStore
 
@@ -57,3 +58,53 @@ def save_fn(existing: SessionSnapshot | None) -> SessionSnapshot:
     assert saved is not None
     assert saved.session_id == 'sess-456'
     assert isinstance(saved.snapshot_id, str)
+
+
+@pytest.mark.asyncio
+async def test_firestore_session_store_save_and_get_async() -> None:
+    """Test saving and getting a session snapshot using an AsyncMock Firestore client."""
+    mock_client = MagicMock()
+    mock_client._is_async = True
+    mock_doc_snapshot = MagicMock()
+
+    mock_doc_ref = MagicMock()
+    mock_doc_ref.get = AsyncMock(return_value=mock_doc_snapshot)
+    mock_doc_ref.set = AsyncMock()
+
+    mock_col = MagicMock()
+    mock_col.document.return_value = mock_doc_ref
+
+    mock_doc_ref.collection.return_value = mock_col
+    mock_client.collection.return_value = mock_col
+
+    store = FirestoreSessionStore(client=mock_client)
+    assert store._is_async_client() is True
+
+    snapshot_data = {
+        'snapshotId': 'snap-async-123',
+        'sessionId': 'sess-async-456',
+        'createdAt': '2026-07-03T00:00:00Z',
+    }
+    mock_doc_snapshot.exists = True
+    mock_doc_snapshot.to_dict.return_value = snapshot_data
+
+    retrieved = await store.get_snapshot(snapshot_id='snap-async-123')
+    assert retrieved is not None
+    assert retrieved.snapshot_id == 'snap-async-123'
+    assert retrieved.session_id == 'sess-async-456'
+
+    mock_doc_snapshot.exists = False
+    mock_doc_snapshot.to_dict.return_value = None
+
+    def save_fn(existing: SessionSnapshot | None) -> SessionSnapshot:
+        return SessionSnapshot(
+            snapshot_id='snap-async-789',
+            session_id='sess-async-456',
+            created_at='2026-07-03T00:00:01Z',
+        )
+
+    saved = await store.save_snapshot(None, save_fn)
+    assert saved is not None
+    assert saved.session_id == 'sess-async-456'
+    assert isinstance(saved.snapshot_id, str)
+    mock_doc_ref.set.assert_called()