Add Firestore backend for FastMCP OAuth proxy client storage

[anton-jackson]

Summary

Adds an optional Firestore-backed key/value store for the FastMCP OAuth proxy's dynamic client registration data, so deployments running on serverless / multi-instance infrastructure (e.g. Cloud Run) can persist registered clients across instances and restarts instead of relying on local disk.

What changed

• auth/firestore_kv_store.py — new async KV store implementing the FastMCP storage interface, backed by google-cloud-firestore.
• core/server.py — wires the Firestore store in when configured; falls back to the existing disk store otherwise.
• pyproject.toml — adds a firestore extra (google-cloud-firestore>=2.16.0).
• Dockerfile — installs the new extra. Note: google-cloud-firestore is installed via uv pip install after uv sync --frozen because uv.lock has not yet been regenerated to include it. This is a temporary workaround — once uv lock is run with PyPI access, the RUN uv pip install line should be removed and the extra folded back into uv sync --frozen --extra firestore.
• gmail/gmail_helpers.py — minor lint/format fixes from ruff.

Test plan

• Build the Docker image locally and confirm uv sync + firestore install both succeed
• Run with the disk store (default) and confirm no regression
• Run with the Firestore store configured and verify dynamic client registration round-trips through Firestore
• Restart the process and confirm registered clients persist

Summary by CodeRabbit

• New Features

  • Added Firestore-backed storage option for OAuth client credentials with encrypted-at-rest entries and honoring entry expiration.

• Chores

  • Build now includes the optional Firestore support so deployments can enable the new backend.

• Code Quality

  • Minor refactor to error-handling logic for Gmail helpers.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 743bb75c-662b-453c-9ebc-7003555344ea

📥 Commits

Reviewing files that changed from the base of the PR and between fc9ea79 and 5f90273.

📒 Files selected for processing (2)

• Dockerfile
• core/server.py

🚧 Files skipped from review as they are similar to previous changes (1)

• core/server.py

---

📝 Walkthrough

Walkthrough

Adds a Firestore-backed key/value store and integrates it as an optional OAuth client storage backend (with Fernet envelope encryption), declares the Firestore optional dependency and installs it in the Docker build, and includes a small unrelated refactor in Gmail helpers.

Changes

Firestore Storage + Integration

Layer / File(s)	Summary
Data Shape / Adapter auth/firestore_kv_store.py	New FirestoreStore persisting serialized ManagedEntry JSON and expires_at fields in a single Firestore collection; includes dependency guard for google-cloud-firestore.
Core Implementation auth/firestore_kv_store.py	Read honours expires_at (returns None if expired), write persists JSON + expires_at, delete reports success based on prior existence.
Wiring / Server Integration core/server.py	Adds storage_backend == "firestore" branch: imports FirestoreStore, reads WORKSPACE_MCP_OAUTH_PROXY_FIRESTORE_* env vars (collection default workspace_mcp_oauth_kv), constructs store, derives storage_encryption_key from JWT signing key, wraps store with FernetEncryptionWrapper, logs config; ImportError -> RuntimeError explaining missing workspace-mcp[firestore].
Dependencies / Build pyproject.toml, Dockerfile	Adds optional extra firestore = ["google-cloud-firestore>=2.16.0"] and matching dependency group; Dockerfile's uv sync now includes --extra firestore so Firestore extra is installed in the image.
Tests / Docs (none changed)	No tests or docs added in this PR for the Firestore flow.

Minor Gmail Helper Refactor

Layer / File(s)	Summary
Logic Formatting gmail/gmail_helpers.py	Refactors _is_benign_signature_http_error to return the same boolean condition as a single-line return; no functional change.

Sequence Diagram

sequenceDiagram
    participant Server as Server (core/server.py)
    participant Env as Env (WORKSPACE_MCP_OAUTH_PROXY_FIRESTORE_*)
    participant Store as FirestoreStore (auth/firestore_kv_store.py)
    participant GCF as GoogleCloudFirestore
    participant Encrypt as FernetEncryptionWrapper

    Server->>Env: Read Firestore project/database/collection
    Env-->>Server: Firestore params
    Server->>Store: Initialize FirestoreStore(project, database, collection)
    Store->>GCF: Connect / read/write documents
    Server->>Server: Derive storage_encryption_key from JWT signing key
    Server->>Encrypt: Wrap FirestoreStore with Fernet key
    Encrypt-->>Server: Encrypted client_storage ready

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Possibly related PRs

• enh: FastMCP-native CLI Mode with Cloud Support #618 — Related server storage initialization and disk-backed, Fernet-encrypted token store work.
• fix: Issues/552 - Docker [disk] storage  #554 — Related Dockerfile uv sync --extra changes for optional extras.

Suggested labels

enhancement

Poem

🐰 In cloud fields where documents play,

I tuck your tokens safe away.
Fernet wraps with gentle hop,
Firestore keeps them — none will drop.
Hop, hop, the keys are snug — hip-hop!

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The PR description covers the main changes, what changed, and includes a test plan, but is missing several required template sections and checkboxes.	Complete the description template by adding Type of Change, Testing checkboxes, Checklist section, and confirm 'Allow edits from maintainers' is enabled as required.
Docstring Coverage	⚠️ Warning	Docstring coverage is 12.50% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately summarizes the main change: adding a Firestore backend for OAuth proxy client storage, which is the primary objective of this PR.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

⚔️ Resolve merge conflicts

• Resolve merge conflict in branch feat/firestore-oauth-store

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Review rate limit: 7/8 reviews remaining, refill in 7 minutes and 30 seconds.}

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 2

🧹 Nitpick comments (1)

auth/firestore_kv_store.py (1)

67-120: ⚡ Quick win

Please add async coverage for the expiry and CRUD semantics here.

This class owns custom logic for expired reads and the collection:key document mapping, but there’s no regression net for it in the PR. A small async test around put/get/delete plus an expired-document read would catch the main failure modes before they surface in OAuth DCR flows.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@auth/firestore_kv_store.py` around lines 67 - 120, Add an async test module
that exercises the FirestoreKVStore CRUD and expiry behavior: call
_put_managed_entry to write a ManagedEntry (use the same serialization adapter
to produce JSON), then verify _get_managed_entry returns the entry; delete it
and assert _delete_managed_entry returns True and subsequent _get_managed_entry
returns None; also write a ManagedEntry with expires_at set in the past and
assert _get_managed_entry returns None (and that a future expires_at is
preserved and returned). Use the class's helper methods (_put_managed_entry,
_get_managed_entry, _delete_managed_entry) and fields (_EXPIRES_AT_FIELD,
_JSON_FIELD) to construct/inspect payloads and run the tests as async unit tests
to cover TTL/eventual deletion semantics and collection:key mapping.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@core/server.py`:
- Around line 513-518: The current except ImportError block in core/server.py
just logs a warning (logger.warning) when google-cloud-firestore import fails;
instead, if the configuration explicitly requests firestore
(WORKSPACE_MCP_OAUTH_PROXY_STORAGE_BACKEND == "firestore"), fail startup
immediately rather than falling back: replace the warning-only behavior with
logic that checks the WORKSPACE_MCP_OAUTH_PROXY_STORAGE_BACKEND env/config value
and, when it equals "firestore", log an error and re-raise the ImportError or
call sys.exit(1)/raise RuntimeError with a clear message including exc;
otherwise (when not explicitly requested) keep the existing fallback behavior or
warning.

In `@Dockerfile`:
- Around line 16-20: The Dockerfile currently bypasses the lock by running `uv
pip install "google-cloud-firestore..."`; instead update the `uv sync`
invocation to use the declared extra and remove the separate `uv pip install`
line: replace the `RUN uv sync --frozen --no-dev --extra disk`/`RUN uv pip
install ...` pattern with a single `uv sync --frozen --no-dev --extra firestore`
(or include both extras if `disk` is still needed) and delete the `uv pip
install` line; also update the comment to reflect that `uv.lock` includes the
`firestore` extra so the lockfile is honored.

---

Nitpick comments:
In `@auth/firestore_kv_store.py`:
- Around line 67-120: Add an async test module that exercises the
FirestoreKVStore CRUD and expiry behavior: call _put_managed_entry to write a
ManagedEntry (use the same serialization adapter to produce JSON), then verify
_get_managed_entry returns the entry; delete it and assert _delete_managed_entry
returns True and subsequent _get_managed_entry returns None; also write a
ManagedEntry with expires_at set in the past and assert _get_managed_entry
returns None (and that a future expires_at is preserved and returned). Use the
class's helper methods (_put_managed_entry, _get_managed_entry,
_delete_managed_entry) and fields (_EXPIRES_AT_FIELD, _JSON_FIELD) to
construct/inspect payloads and run the tests as async unit tests to cover
TTL/eventual deletion semantics and collection:key mapping.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: e95a47b8-8161-422b-b69d-6ab339179999

📥 Commits

Reviewing files that changed from the base of the PR and between f3c7dc5 and fc9ea79.

⛔ Files ignored due to path filters (1)

• uv.lock is excluded by !**/*.lock

📒 Files selected for processing (5)

• Dockerfile
• auth/firestore_kv_store.py
• core/server.py
• gmail/gmail_helpers.py
• pyproject.toml