feat(py): FirestoreSessionStore (google-cloud plugin)#5685
Conversation
github-actions
Bot
added
docs
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a Firestore-backed session store (FirestoreSessionStore) for agent snapshots in the Google Cloud plugin, along with corresponding unit tests. Feedback on the implementation highlights several critical improvements: unsubscribing from the Firestore snapshot listener upon reaching a terminal status to prevent resource leaks, ensuring only a single listener is started per snapshot ID, querying Firestore directly by sessionId instead of streaming the entire collection, and correcting the _sanitize function to properly filter out None values.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
huangjeff5
force-pushed
the
jh-py-agents-firestore-session
branch
from
9efc34b to
0feada2
Compare
huangjeff5
force-pushed
the
jh-py-agents-firestore-session
branch
from
07ad162 to
16ffefc
Compare
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces a Firestore-backed session store (FirestoreSessionStore) for agent snapshots, adding the google-cloud-firestore dependency and implementing the store across both the genkit-google-cloud package and the google-cloud plugin, along with corresponding unit tests. The review feedback correctly identifies a critical concurrency issue in the listener implementation of both files, where calling loop.create_task from a background thread is not thread-safe and a race condition exists when unsubscribing from the snapshot listener. Applying the suggested thread-safe scheduling and unsubscription fixes is highly recommended.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces a Firestore-backed session store (FirestoreSessionStore) for persisting agent snapshots in Google Cloud, along with corresponding unit tests and updated sample imports. Feedback highlights critical runtime issues in _start_listener due to unsupported async on_snapshot calls and callback signature mismatches. Additionally, it is recommended to use async transactions and server-side timestamps in _maybe_update_pointer to prevent concurrency race conditions, and to update the mock tests to reflect the correct SDK behavior.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
|
@gemini-code-assist review |
There was a problem hiding this comment.
Code Review
This pull request introduces a Firestore-backed session store (FirestoreSessionStore) for agent snapshots in both the genkit-google-cloud and google-cloud plugins, and updates various agent samples to use the correct package imports. The review feedback highlights a performance and resource leak issue in _start_listener, recommending the lazy initialization and reuse of a single synchronous firestore.Client instance to avoid leaking gRPC channels, as well as updating the deprecated asyncio.get_event_loop() call to asyncio.get_running_loop().
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
|
@gemini-code-assist review |
There was a problem hiding this comment.
Code Review
This pull request introduces a Firestore-backed session store (FirestoreSessionStore) for persisting agent snapshots in Google Cloud Firestore. It includes the implementation of the store, associated unit tests, and dependency updates in pyproject.toml and uv.lock. Additionally, it updates several agent sample files to use the correct imports for GoogleAI and middleware components. No review comments were provided, so there is no feedback to address.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
|
@gemini-code-assist review |
1 similar comment
|
@gemini-code-assist review |
There was a problem hiding this comment.
Code Review
This pull request introduces a Firestore-backed session store (FirestoreSessionStore) for agent snapshots in the genkit-google-cloud package, including comprehensive unit tests and updated import paths in several sample files. The review feedback highlights a potential issue where instantiating asyncio.Lock() directly in init can raise a RuntimeError if called outside of an active event loop, and suggests lazily initializing the lock to resolve this.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
…storeSessionStore
…d functions in firestore.py
…hout leading underscore demarcation
|
@gemini-code-assist review |
huangjeff5
force-pushed
the
jh-py-agents-firestore-session
branch
from
01779a6 to
e72c75a
Compare
There was a problem hiding this comment.
Code Review
This pull request introduces a Firestore-backed session store (FirestoreSessionStore) for persisting agent snapshots in Google Cloud, along with corresponding unit tests. It also updates several sample files to use the new standalone package imports for GoogleAI and middleware. The review feedback focuses on improving the robustness of the session store listener subscription by handling exceptions during startup, and enhancing test coverage to properly exercise the synchronous fallback path when the async document reference lacks the on_snapshot attribute.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
…nd fallback path test coverage
|
@gemini-code-assist review |
There was a problem hiding this comment.
Code Review
This pull request introduces a Firestore-backed session store (FirestoreSessionStore) for persisting agent snapshots, along with corresponding unit tests and dependency updates. It also refactors imports in several sample agent files. The code review identified a critical bug in the pointer update logic where the session pointer is not updated on subsequent turns because it compares the current pointer to the new snapshot ID instead of the parent snapshot ID. Additionally, the reviewer suggested optimizing the real-time listener to skip starting for snapshots already in a terminal status, and recommended adding a test case to verify pointer updates for existing sessions.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
There was a problem hiding this comment.
Code Review
This pull request introduces a Firestore-backed session store (FirestoreSessionStore) for persisting agent snapshots in Google Cloud, along with corresponding unit tests and import cleanups across several sample files. The review feedback highlights a critical issue in the pointer update logic during snapshot saves: using a simple boolean flag (is_new) prevents the session pointer from updating when appending to an existing session. To resolve this, the reviewer suggests implementing a proper compare-and-swap mechanism by passing the parent snapshot ID to maybe_update_pointer and adding a test case to verify appending to an existing session.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
…al listener startup
|
@gemini-code-assist review |
There was a problem hiding this comment.
Code Review
This pull request introduces a Firestore-backed session store (FirestoreSessionStore) for persisting agent snapshots in Google Cloud Firestore, along with corresponding unit tests and updates to sample imports. The feedback recommends preserving client_options when falling back to the synchronous Firestore client and running the blocking unsubscribe() call in an executor to prevent blocking the asyncio event loop.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
…h unsubscribe in executor
|
@gemini-code-assist review |
There was a problem hiding this comment.
Code Review
This pull request introduces a Firestore-backed session store (FirestoreSessionStore) for persisting agent snapshots in Google Cloud Firestore. It includes the implementation of the store, associated unit tests, and updates to the project's dependencies (pyproject.toml and uv.lock). Additionally, several sample agent files have been updated to use the correct import paths for genkit_google_genai and genkit_middleware. I have no feedback to provide as there are no review comments.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
|
| Test Case | Scenario Exercised | Verification & Assertions |
|---|---|---|
test_bug_bash_multiturn_linear_chat_pointer_advancement |
Multiturn linear chat (Turns 1, 2, 3). | Verified save_snapshot appends link parent_id sequentially. Verified pointer document currentSnapshotId advances atomicaly from Turn 1 -> Turn 2 -> Turn 3. Verified get_snapshot(session_id=...) returns Turn 3 tip. |
test_bug_bash_branching_and_forked_history |
History forking (Branch A tip vs Branch B fork from Root). | Verified appending Branch B to Root leaves pointer intact at Branch A tip (CAS isolation). Verified un-pointed fallback with reject_ambiguous_session=False returns latest leaf chronologically, while reject_ambiguous_session=True raises GenkitError(FAILED_PRECONDITION, 'branching snapshots'). |
test_bug_bash_cas_concurrent_append_race_condition |
Concurrent append race condition (Turn 2X and Turn 2Y both targeting Turn 1). | Simulated simultaneous transaction execution. Verified Turn 2X commits and updates pointer to turn-2x. Verified Turn 2Y CAS check (pointer == 'turn-1') fails and leaves pointer at turn-2x without overwriting. |
test_bug_bash_inplace_update_vs_append |
In-place snapshot update vs appending new turns. | Verified calling save_snapshot(turn1_id, fn) updates status from PENDING to COMPLETED idempotently without altering snapshot ID or creating duplicate records. Verified pointer remains stable. |
test_bug_bash_realtime_status_listener_and_cleanup |
Real-time status subscription and thread-safe listener cleanup. | Verified queue receives live transitions (PENDING -> COMPLETED -> None). Verified watch.unsubscribe() is executed in an executor and subscription is popped from store.subs. |
test_bug_bash_terminal_status_optimization |
Subscribing to an already-terminal snapshot (ABORTED). |
Verified start_listener is bypassed. Verified queue receives ABORTED followed immediately by None sentinel. |
test_bug_bash_error_handling_and_invalid_selectors |
Invalid selector combinations and missing metadata. | Verified saving without sessionId raises GenkitError(INVALID_ARGUMENT). Verified calling get_snapshot with 0 or 2 selectors raises GenkitError(INVALID_ARGUMENT, 'exactly one of'). |
4. Live Google Cloud Firestore E2E Suite (firestore_session_store_e2e_test.py)
We executed end-to-end integration tests against the live Google Cloud Firestore service on project aim-testing using tenant-isolated prefixes (test-e2e-{uuid4().hex[:8]}). All test data was verified directly via Firestore DocumentReference queries and wiped clean during teardown.
Workflow 1: E2E Agent Workflow & Tool Execution (test_firestore_session_store_e2e_agent_workflow_and_tools)
- Execution: Defined
e2eAgentwith a customcalculatortool (input_val * 2). Executed Turn 1 ("Calculate 21 * 2") and Turn 2 ("What was the answer again?") viaAgentChat. - Database Verification:
- Queried
genkit-e2e-sessions/global/snapshots/{turn1_id}directly in Firestore. Verifiedstatus: 'completed',sessionId, and exact message history structure: 4 parts durably stored (role: 'user',role: 'model'withToolRequestPart,role: 'tool'withToolResponsePartoutput42, androle: 'model'final reply). - Queried
genkit-e2e-sessions-pointers/global/pointers/{session_id}in Firestore. VerifiedcurrentSnapshotIdpointed to Turn 1 after Turn 1, and atomically advanced to Turn 2 after Turn 2. - Verified Turn 2 document in Firestore correctly referenced
parentId: {turn1_id}and contained all 6 accumulated message parts.
- Queried
Workflow 2: E2E Detach, Abort, and Real-Time gRPC Watch Stream (test_firestore_session_store_e2e_detach_abort_and_watch)
- Execution: Configured
slowAgentwith a 4-second delay. Initiated a detached background turn viatask = await chat.detach('Run slow background analysis'). Subscribed to real-time status updates viae2e_store.on_snapshot_status_change(task.snapshot_id)over a live gRPC watch stream against Google Cloud Firestore. - Database & Stream Verification:
- Verified initial status received from watch stream and direct DB query was
PENDING. - Executed client-side abort via
await task.abort(). - Verified live gRPC watch stream pushed
SnapshotStatus.ABORTEDfollowed by sentinelNone. - Queried Firestore document directly; verified durable status transitioned to
aborted. - Verified awaiting the aborted turn raised
asyncio.CancelledError. - Sent follow-up prompt on the same chat; verified conversation resumed cleanly from the last completed turn without stranding history.
- Verified initial status received from watch stream and direct DB query was
Workflow 3: E2E Branching, CAS Isolation, & Strict Ambiguous Rejection (test_firestore_session_store_e2e_branching_and_cas_isolation)
- Execution: Executed Root turn -> Branch A turn. Verified pointer in Firestore pointed to Branch A tip. Called
agent.load_chat(snapshot_id=root_id)and executed Branch B turn ("Branch B prompt"). - Database Verification:
- Queried Branch B document in Firestore; verified
parentIdmatched Root ID. - Verified pointer in Firestore remained safely at Branch A tip due to CAS isolation (
currentSnapshotId == root_idevaluated to False when pointer was at Branch A). - Deleted pointer document in Firestore to simulate a legacy/un-pointed session with multiple branch tips.
- Initialized
strict_storewithreject_ambiguous_session=True. Attempted to load session viastrict_agent.load_chat(session_id=session_id). - Verified Firestore query detected 2 leaf documents and raised
GenkitError(FAILED_PRECONDITION)matching'branching snapshots'.
- Queried Branch B document in Firestore; verified
5. Verification & Test Run Logs
uv run --no-sync pytest packages/genkit-google-cloud/tests -v --no-cov============================= test session starts ==============================
platform darwin -- Python 3.12.9, pytest-9.0.2, pluggy-1.6.0
rootdir: /Users/huangjeff/Desktop/genkit-firestore-session/py
configfile: pyproject.toml
plugins: anyio-4.12.1, mock-3.15.1, Faker-40.4.0, asyncio-1.3.0, typeguard-4.5.0, cov-7.0.0
asyncio: mode=Mode.STRICT, debug=False, asyncio_default_fixture_loop_scope=session, asyncio_default_test_loop_scope=function
collecting ... collected 83 items
packages/genkit-google-cloud/tests/firestore_session_store_bug_bash_test.py::test_bug_bash_multiturn_linear_chat_pointer_advancement PASSED [ 1%]
packages/genkit-google-cloud/tests/firestore_session_store_bug_bash_test.py::test_bug_bash_branching_and_forked_history PASSED [ 2%]
packages/genkit-google-cloud/tests/firestore_session_store_bug_bash_test.py::test_bug_bash_cas_concurrent_append_race_condition PASSED [ 3%]
packages/genkit-google-cloud/tests/firestore_session_store_bug_bash_test.py::test_bug_bash_inplace_update_vs_append PASSED [ 4%]
packages/genkit-google-cloud/tests/firestore_session_store_bug_bash_test.py::test_bug_bash_realtime_status_listener_and_cleanup PASSED [ 6%]
packages/genkit-google-cloud/tests/firestore_session_store_bug_bash_test.py::test_bug_bash_terminal_status_optimization PASSED [ 7%]
packages/genkit-google-cloud/tests/firestore_session_store_bug_bash_test.py::test_bug_bash_error_handling_and_invalid_selectors PASSED [ 8%]
packages/genkit-google-cloud/tests/firestore_session_store_e2e_test.py::test_firestore_session_store_e2e_agent_workflow_and_tools PASSED [ 9%]
packages/genkit-google-cloud/tests/firestore_session_store_e2e_test.py::test_firestore_session_store_e2e_detach_abort_and_watch PASSED [ 10%]
packages/genkit-google-cloud/tests/firestore_session_store_e2e_test.py::test_firestore_session_store_e2e_branching_and_cas_isolation PASSED [ 12%]
... [73 existing unit tests passed] ...
======================= 83 passed, 6 warnings in 13.48s ========================
All formatting (ruff format --check), linter rules (ruff check), and type checkers (pyright, pyrefly, ty) pass with 0 errors across genkit-google-cloud.
huangjeff5
force-pushed
the
jh-py-agents-firestore-session
branch
2 times, most recently
from
e1a7574 to
cdc333a
Compare
…for FirestoreSessionStore
huangjeff5
force-pushed
the
jh-py-agents-firestore-session
branch
from
cdc333a to
d464bd7
Compare
1 participant
Adds a Firestore-backed agent session store. PR 1 of 3 in the FastAPI-agents stack; PR2 (serve_agent/serve_flow) and PR3 (sample) build on this.