Skip to content

fix(integrations): bound memory recall, default indexes on, keep MongoDB credentials out of Agent Engine packages - #7533

Open
theshanbhag wants to merge 18 commits into
google:mainfrom
theshanbhag:mongodb-hardening
Open

theshanbhag wants to merge 18 commits into
google:mainfrom
theshanbhag:mongodb-hardening

Conversation

@theshanbhag

Copy link
Copy Markdown
Contributor

Addresses the three remaining review notes from #7291. Stacked on the open split PRs (#7509, #7510, #7511, #7512 — it modifies files those PRs introduce); only the last commit is new.

Please ensure you have read the contribution guide before creating a pull request.

Link to Issue or Description of Change

Problem:

Three hardening gaps called out in the #7291 review:

  1. Memory search has no result limit or ranking — a common keyword could return every matching memory.
  2. Indexes are off by default (ensure_indexes=False) — users had to opt in or reads degraded into collection scans as data grew.
  3. The database password is saved in the deployment package — for Agent Engine deployments, a literal connection_string is stored on the toolset and pickled into the uploaded package in plaintext.

Solution:

  1. Bounded, ranked recall. search_memory now ranks matches by keyword overlap (recency breaks ties) and caps results at a new max_results constructor argument (default 20). Ranking and limiting run server-side in the aggregation pipeline ($match → $addFields overlap via $filter+$size → $sort → $limit), so a common keyword returns the best few memories rather than everything — and the excess never crosses the wire.
  2. Indexes on by default. Both services now default ensure_indexes=True. A runtime user without the createIndex privilege (the least-privilege setup the README recommends) triggers one warning and falls back to unindexed reads instead of failing every operation; unexpected index-creation errors still propagate, and the flag is only marked done on success or a handled privilege denial, so transient failures retry on the next operation. ensure_indexes=False remains as the explicit opt-out.
  3. Credentials out of the package. MongoDbToolset accepts connection_string_env_var: only the environment variable's name is stored on the toolset, and the value is resolved where the client is built — locally at construction, or on the Agent Engine runtime after unpickling (__setstate__ resolves it there). Pickling therefore carries no credential, and the README's new Agent Engine section pairs it with Secret Manager-backed env vars (env_vars={"MONGODB_URI": SecretRef(secret="mongodb-uri", version="latest")}) so the plaintext never touches the deployment spec either. A unit test asserts the pickled payload contains neither the URI nor the password.

Testing Plan

Unit Tests:

  • I have added or updated unit tests for my change.
  • All unit tests pass locally.

pytest tests/unittests/integrations/mongodb/ tests/unittests/features/ → 134 passed (mongomock-backed), including new coverage: recall limit, overlap ranking with recency tiebreak, max_results validation, indexes created by default, opt-out, privilege-denied fallback (warns once, reads still work), non-privilege index errors propagating, env-var resolution, missing-env-var error, multi-source constructor validation, and a pickle test asserting the credential bytes never enter the payload. All pre-commit hooks pass.

Manual End-to-End (E2E) Tests:

Against a live Atlas 8.0.34 cluster: memory ingest + recall returns ranked, capped results; session and memory operations work with the default-on indexes; an agent built with connection_string_env_var="MONGODB_URI" round-trips through pickle.dumps/loads with a working rebuilt client while the pickled bytes contain no URI.

Checklist

  • I have read the CONTRIBUTING.md document.
  • I have performed a self-review of my own code.
  • I have commented my code, particularly in hard-to-understand areas.
  • I have added tests that prove my fix is effective or that my feature works.
  • New and existing unit tests pass locally with my changes.
  • I have manually tested my changes end-to-end.
  • Any dependent changes have been merged and published in downstream modules.

Additional context

The MongoDB session and memory services still cannot be deployed to Agent Engine; the follow-up that wires them in will extend the same pickling + connection_string_env_var handling to them.

Add `use_mongodb_auto_embedding` and `mongodb_auto_embedding_model` to
MongoDbToolSettings. When enabled, the search tools send the query as
plain text (`query.text`) instead of embedding it through the genai
client first: Atlas generates the embedding itself with the Voyage AI
model configured on the index's `autoEmbed` field (Atlas Automated
Embedding, Preview), so deployments that already store data on Atlas
need no Google embedding credentials. The default mode (Google
embedding models via `queryVector`) is unchanged.

Split from google#7291.
Agent Engine packages the app object with cloudpickle, and a live
pymongo.MongoClient (sockets, locks, background threads) cannot cross
that boundary. Constructed with `connection_string`, the toolset now
drops the client at pickle time and rebuilds it on the runtime; the
query-embedding genai client is likewise rebuilt lazily from the
ambient environment on the destination. Constructing with a
caller-owned `mongo_client` raises TypeError at pickle time with
guidance, since such a client cannot be rebuilt remotely.

Pickling support for the MongoDB session and memory services is
deliberately left to the follow-up that wires those services into
Agent Engine.

Split from google#7291.
MongoDbSessionService persists sessions, events, and app:/user:-scoped
state in MongoDB, exported from google.adk.integrations.mongodb behind
the MONGODB_SESSION_SERVICE experimental feature (default on).

Concurrency model: within a process a per-session lock serializes
append_event; across runner replicas the optimistic-concurrency
revision on the session document is the only guard, and the losing
writer gets StaleSessionError. Multi-document writes run in a
transaction on replica sets / sharded clusters / Atlas; on deployments
without transaction support (standalone mongod) the verdict is cached
and writes fall back to sequential with the revision check. Optional
`ensure_indexes` creates the recommended secondary indexes on first
use. mongomock joins the test dependencies for the new unit tests.

NO_UNIT_GUIDE=The integration README documents the service, including
its document layout, concurrency model, and least-privilege setup; a
per-unit guide will follow when the experimental API stabilizes.

Split from google#7291.
MongoDbMemoryService ingests session events into a memory collection
and recalls them by keyword matching, exported from
google.adk.integrations.mongodb behind the MONGODB_MEMORY_SERVICE
experimental feature (default on). Ingestion is idempotent (memory
documents are keyed by app/user/session/event), and `ensure_indexes`
creates the app/user/keywords index on first use.

Retrieval is intentionally keyword-only: it needs no embedding model,
no Google credentials, and no vector index, so it works on any MongoDB
deployment at zero per-query cost. Vector recall is a possible
follow-up building on this integration's own search tools.

NO_UNIT_GUIDE=The integration README documents the service, including
its storage layout and least-privilege setup; a per-unit guide will
follow when the experimental API stabilizes.

Split from google#7291.
…oDB credentials out of Agent Engine packages

Three hardening fixes from the google#7291 review:

- MongoDbMemoryService.search_memory ranked nothing and returned every
  keyword match. Results are now ranked by keyword overlap (recency
  breaks ties) and capped by a new max_results constructor argument
  (default 20); ranking and limiting run server-side in the
  aggregation pipeline so a common keyword no longer returns
  everything.
- ensure_indexes defaulted to False, so users had to opt in or reads
  degraded into collection scans as data grew. It now defaults to True
  for both services; a deployment without the createIndex privilege
  logs one warning and falls back to unindexed reads instead of
  failing (unexpected index-creation errors still propagate and are
  retried on the next operation).
- A literal connection_string stored on the toolset was pickled into
  the Agent Engine deployment package in plaintext, password included.
  MongoDbToolset now accepts connection_string_env_var: only the
  variable's name is stored, and the value is resolved where the
  client is built — on the runtime after unpickling. The README's new
  Agent Engine section pairs it with Secret Manager-backed env vars
  (SecretRef) so the credential never touches the deployment spec.

Split from google#7291.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces MongoDB integration for ADK, adding MongoDbSessionService for session persistence, MongoDbMemoryService for keyword-based long-term memory, and support for MongoDB Automated Embedding (Atlas Preview) in MongoDbToolset. It also adds pickling support to allow deployment to Agent Engine. The review feedback highlights several optimization opportunities, including resolving an N+1 query pattern in list_sessions, utilizing bulk_write for batch memory ingestion, and introducing thread locks to synchronize index creation on startup in both services.

Comment thread src/google/adk/integrations/mongodb/_session_service.py Outdated
Comment thread src/google/adk/integrations/mongodb/_memory_service.py
Comment thread src/google/adk/integrations/mongodb/_memory_service.py Outdated
Comment thread src/google/adk/integrations/mongodb/_session_service.py Outdated
@theshanbhag

Copy link
Copy Markdown
Contributor Author

Note on the failing zizmor-output check: the findings are pre-existing on main, not introduced by this PR. This branch changes no files under .github/ — I verified the directory is byte-identical to upstream main (diff -r clean), and running zizmor 1.25.2 locally with the same mandatory check set produces the identical 13 findings on both main and this branch (exit 14 either way).

The findings are all in repo-owned workflows this PR doesn't touch: continuous-integration.yml (cache-poisoning), issue-maintenance.yml and release-artifact-check.yml (excessive-permissions), pr-triage.yml (dangerous-triggers), release-cut.yml/release-cherry-pick.yml/release-finalize.yml (artipacked), release-publish.yml (use-trusted-publishing). They look like they need a maintainer pass (or repo-level zizmor config) rather than a fix in this feature PR — happy to help if maintainers prefer a separate workflow-hardening PR.

(The agent-triage-pull-request failure is the triage bot crashing on a model API error — also unrelated to these changes.)

…b-hardening

# Conflicts:
#	constraints-3.11.txt
#	constraints-3.12.txt
#	constraints-3.13.txt
#	constraints-3.14.txt
…b-session-service

# Conflicts:
#	constraints-3.11.txt
#	constraints-3.12.txt
#	constraints-3.13.txt
#	constraints-3.14.txt
…n MongoDbSessionService

Review follow-up from google#7533: list_sessions fetched each user's state
with a separate find_one per user (N+1); it now issues a single $in
query against the user_states collection. Index creation on first use
is also serialized with a thread lock (double-checked flag), so
concurrent first calls no longer issue duplicate create_index calls.

Addresses gemini-code-assist review comments on google#7533.
Review follow-up from google#7533: add_session_to_memory issued one
replace_one round-trip per event; it now batches all upserts into a
single bulk_write (ordered=False). First-use index creation is
serialized with a thread lock (double-checked flag) so concurrent
startup calls no longer issue duplicate create_index calls. Tests
carry a shim for mongomock's BulkOperationBuilder, which rejects the
`sort` key pymongo 4.9+ adds to bulk operations.

Addresses gemini-code-assist review comments on google#7533.
# Conflicts:
#	src/google/adk/integrations/mongodb/_memory_service.py
#	src/google/adk/integrations/mongodb/_session_service.py
pyink reflow of the merged _ensure_indexes_once bodies, plus
constraints regenerated against upstream's 2026-10-14 cutoff (now
including the new constraints-3.15.txt) with the branch's mongomock
entries applied consistently across all five Python versions.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant