diff --git a/constraints-3.11.txt b/constraints-3.11.txt index 98e6570fb30..8884175c9af 100644 --- a/constraints-3.11.txt +++ b/constraints-3.11.txt @@ -1,5 +1,5 @@ # This file was autogenerated by uv via the following command: -# uv pip compile pyproject.toml --all-extras --python-version 3.11 --no-emit-package google-adk --exclude-newer 2026-10-14 --index-url https://pypi.org/simple -o constraints-3.11.txt +# uv pip compile pyproject.toml --all-extras --python-version 3.11 --no-emit-package google-adk --exclude-newer 2026-10-14 -o constraints-3.11.txt.new.tmp --constraint constraints-3.11.txt.stable.tmp a2a-sdk==1.1.1 # via # -c constraints-3.11.txt.stable.tmp @@ -1127,6 +1127,10 @@ mmh3==5.2.1 # -c constraints-3.11.txt.stable.tmp # chromadb # google-cloud-spanner +mongomock==4.3.0 + # via + # -c constraints-3.11.txt.stable.tmp + # google-adk (pyproject.toml) multidict==6.7.1 # via # -c constraints-3.11.txt.stable.tmp @@ -1383,6 +1387,7 @@ packaging==26.2 # langchain-core # langsmith # marshmallow + # mongomock # onnxruntime # opentelemetry-instrumentation # pyink @@ -1766,6 +1771,7 @@ pytube==15.0.0 pytz==2026.2 # via # -c constraints-3.11.txt.stable.tmp + # mongomock # oci # pandas pyyaml==6.0.3 @@ -1886,6 +1892,10 @@ scipy==1.17.1 # via # -c constraints-3.11.txt.stable.tmp # scikit-learn +sentinels==1.1.1 + # via + # -c constraints-3.11.txt.stable.tmp + # mongomock setuptools==83.0.0 # via # -c constraints-3.11.txt.stable.tmp diff --git a/constraints-3.12.txt b/constraints-3.12.txt index a2d574591bc..8ca6552a755 100644 --- a/constraints-3.12.txt +++ b/constraints-3.12.txt @@ -1,5 +1,5 @@ # This file was autogenerated by uv via the following command: -# uv pip compile pyproject.toml --all-extras --python-version 3.12 --no-emit-package google-adk --exclude-newer 2026-10-14 --index-url https://pypi.org/simple -o constraints-3.12.txt +# uv pip compile pyproject.toml --all-extras --python-version 3.12 --no-emit-package google-adk --exclude-newer 2026-10-14 -o constraints-3.12.txt.new.tmp --constraint constraints-3.12.txt.stable.tmp a2a-sdk==1.1.1 # via # -c constraints-3.12.txt.stable.tmp @@ -987,6 +987,10 @@ mmh3==5.2.1 # via # -c constraints-3.12.txt.stable.tmp # google-cloud-spanner +mongomock==4.3.0 + # via + # -c constraints-3.12.txt.stable.tmp + # google-adk (pyproject.toml) multidict==6.7.1 # via # -c constraints-3.12.txt.stable.tmp @@ -1206,6 +1210,7 @@ packaging==26.2 # langchain-core # langsmith # marshmallow + # mongomock # opentelemetry-instrumentation # pyink # pyproject-api @@ -1519,6 +1524,7 @@ pytokens==0.4.1 pytz==2026.2 # via # -c constraints-3.12.txt.stable.tmp + # mongomock # oci # pandas pyyaml==6.0.3 @@ -1624,6 +1630,10 @@ scipy==1.18.0 # via # -c constraints-3.12.txt.stable.tmp # scikit-learn +sentinels==1.1.1 + # via + # -c constraints-3.12.txt.stable.tmp + # mongomock setuptools==83.0.0 # via # -c constraints-3.12.txt.stable.tmp diff --git a/constraints-3.13.txt b/constraints-3.13.txt index 4073661dce2..d20aa23f978 100644 --- a/constraints-3.13.txt +++ b/constraints-3.13.txt @@ -1,5 +1,5 @@ # This file was autogenerated by uv via the following command: -# uv pip compile pyproject.toml --all-extras --python-version 3.13 --no-emit-package google-adk --exclude-newer 2026-10-14 --index-url https://pypi.org/simple -o constraints-3.13.txt +# uv pip compile pyproject.toml --all-extras --python-version 3.13 --no-emit-package google-adk --exclude-newer 2026-10-14 -o constraints-3.13.txt.new.tmp --constraint constraints-3.13.txt.stable.tmp a2a-sdk==1.1.1 # via # -c constraints-3.13.txt.stable.tmp @@ -979,6 +979,10 @@ mmh3==5.2.1 # via # -c constraints-3.13.txt.stable.tmp # google-cloud-spanner +mongomock==4.3.0 + # via + # -c constraints-3.13.txt.stable.tmp + # google-adk (pyproject.toml) multidict==6.7.1 # via # -c constraints-3.13.txt.stable.tmp @@ -1198,6 +1202,7 @@ packaging==26.2 # langchain-core # langsmith # marshmallow + # mongomock # opentelemetry-instrumentation # pyink # pyproject-api @@ -1511,6 +1516,7 @@ pytokens==0.4.1 pytz==2026.2 # via # -c constraints-3.13.txt.stable.tmp + # mongomock # oci # pandas pyyaml==6.0.3 @@ -1616,6 +1622,10 @@ scipy==1.18.0 # via # -c constraints-3.13.txt.stable.tmp # scikit-learn +sentinels==1.1.1 + # via + # -c constraints-3.13.txt.stable.tmp + # mongomock setuptools==83.0.0 # via # -c constraints-3.13.txt.stable.tmp diff --git a/constraints-3.14.txt b/constraints-3.14.txt index b076e64d363..edfc2dcbeb8 100644 --- a/constraints-3.14.txt +++ b/constraints-3.14.txt @@ -1,5 +1,5 @@ # This file was autogenerated by uv via the following command: -# uv pip compile pyproject.toml --all-extras --python-version 3.14 --no-emit-package google-adk --exclude-newer 2026-10-14 --index-url https://pypi.org/simple -o constraints-3.14.txt +# uv pip compile pyproject.toml --all-extras --python-version 3.14 --no-emit-package google-adk --exclude-newer 2026-10-14 -o constraints-3.14.txt.new.tmp --constraint constraints-3.14.txt.stable.tmp a2a-sdk==1.1.1 # via # -c constraints-3.14.txt.stable.tmp @@ -979,6 +979,10 @@ mmh3==5.2.1 # via # -c constraints-3.14.txt.stable.tmp # google-cloud-spanner +mongomock==4.3.0 + # via + # -c constraints-3.14.txt.stable.tmp + # google-adk (pyproject.toml) multidict==6.7.1 # via # -c constraints-3.14.txt.stable.tmp @@ -1198,6 +1202,7 @@ packaging==26.2 # langchain-core # langsmith # marshmallow + # mongomock # opentelemetry-instrumentation # pyink # pyproject-api @@ -1511,6 +1516,7 @@ pytokens==0.4.1 pytz==2026.2 # via # -c constraints-3.14.txt.stable.tmp + # mongomock # oci # pandas pyyaml==6.0.3 @@ -1616,6 +1622,10 @@ scipy==1.18.0 # via # -c constraints-3.14.txt.stable.tmp # scikit-learn +sentinels==1.1.1 + # via + # -c constraints-3.14.txt.stable.tmp + # mongomock setuptools==83.0.0 # via # -c constraints-3.14.txt.stable.tmp diff --git a/constraints-3.15.txt b/constraints-3.15.txt index 57ade9e4a20..dd2113a8043 100644 --- a/constraints-3.15.txt +++ b/constraints-3.15.txt @@ -1,5 +1,5 @@ # This file was autogenerated by uv via the following command: -# uv pip compile pyproject.toml --all-extras --python-version 3.15 --no-emit-package google-adk --exclude-newer 2026-10-14 --index-url https://pypi.org/simple -o constraints-3.15.txt +# uv pip compile pyproject.toml --all-extras --python-version 3.15 --no-emit-package google-adk --exclude-newer 2026-10-14 -o constraints-3.15.txt.new.tmp --constraint constraints-3.15.txt.stable.tmp a2a-sdk==1.2.2 # via # -c constraints-3.15.txt.stable.tmp @@ -1002,6 +1002,10 @@ mmh3==5.3.1 # via # -c constraints-3.15.txt.stable.tmp # google-cloud-spanner +mongomock==4.3.0 + # via + # -c constraints-3.15.txt.stable.tmp + # google-adk (pyproject.toml) multidict==6.9.1 # via # -c constraints-3.15.txt.stable.tmp @@ -1225,6 +1229,7 @@ packaging==26.3 # langchain-core # langsmith # marshmallow + # mongomock # opentelemetry-instrumentation # pyink # pyproject-api @@ -1556,6 +1561,7 @@ pytokens==0.4.1 pytz==2026.5 # via # -c constraints-3.15.txt.stable.tmp + # mongomock # oci # pandas pyyaml==6.0.3 @@ -1661,6 +1667,10 @@ scipy==1.18.1 # via # -c constraints-3.15.txt.stable.tmp # scikit-learn +sentinels==1.1.1 + # via + # -c constraints-3.15.txt.stable.tmp + # mongomock setuptools==84.0.0 # via # -c constraints-3.15.txt.stable.tmp diff --git a/pyproject.toml b/pyproject.toml index eed1f7c7796..1aaabc3bb2e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -369,6 +369,7 @@ optional-dependencies.test = [ "llama-index-readers-file>=0.4", "lxml>=5.3", "mcp>=1.24,<3", + "mongomock>=4.3,<5", # In-memory MongoDB used by the mongodb integration tests. "nltk!=3.10.1", # Transitive via rouge-score and llama-index-core; 3.10.1's import hook breaks any venv living inside the working directory (reverted upstream in nltk/nltk#3732). "openai>=2.20,<3", "openpyxl>=3.1.5,<4", diff --git a/src/google/adk/features/_feature_registry.py b/src/google/adk/features/_feature_registry.py index 7ab9dec9b62..fea6b1bf6e1 100644 --- a/src/google/adk/features/_feature_registry.py +++ b/src/google/adk/features/_feature_registry.py @@ -68,6 +68,8 @@ class FeatureName(str, Enum): # Off by default since it changes wire behavior. Enable with # `ADK_ENABLE_MCP_MODERN_PROTOCOL=1`. No effect on MCP SDK 1.x. _MCP_MODERN_PROTOCOL = "MCP_MODERN_PROTOCOL" + MONGODB_MEMORY_SERVICE = "MONGODB_MEMORY_SERVICE" + MONGODB_SESSION_SERVICE = "MONGODB_SESSION_SERVICE" MONGODB_TOOLSET = "MONGODB_TOOLSET" MONGODB_TOOL_SETTINGS = "MONGODB_TOOL_SETTINGS" PROGRESSIVE_SSE_STREAMING = "PROGRESSIVE_SSE_STREAMING" @@ -218,6 +220,12 @@ class FeatureConfig: FeatureName._MCP_MODERN_PROTOCOL: FeatureConfig( FeatureStage.EXPERIMENTAL, default_on=False ), + FeatureName.MONGODB_MEMORY_SERVICE: FeatureConfig( + FeatureStage.EXPERIMENTAL, default_on=True + ), + FeatureName.MONGODB_SESSION_SERVICE: FeatureConfig( + FeatureStage.EXPERIMENTAL, default_on=True + ), FeatureName.MONGODB_TOOLSET: FeatureConfig( FeatureStage.EXPERIMENTAL, default_on=True ), diff --git a/src/google/adk/integrations/mongodb/README.md b/src/google/adk/integrations/mongodb/README.md index a00e265be8c..2adce34deaf 100644 --- a/src/google/adk/integrations/mongodb/README.md +++ b/src/google/adk/integrations/mongodb/README.md @@ -1,8 +1,9 @@ # MongoDB Integration for ADK This integration connects the Google Agent Development Kit (ADK) to MongoDB. -It provides search tools agents can call against collections on MongoDB Atlas -or self-managed MongoDB 8.2+ deployments. +It provides search tools agents can call, plus MongoDB-backed session and +memory services. It works with MongoDB Atlas and self-managed MongoDB 8.2+ +deployments. ## Features @@ -12,6 +13,15 @@ or self-managed MongoDB 8.2+ deployments. - **Hybrid Search Tool (`mongodb_hybrid_search`):** Combines full-text search and vector search with reciprocal rank fusion (`$rankFusion`), with tunable vector/text weights. +- **Session Persistence (`MongoDbSessionService`):** Stores sessions, events, + and `app:` / `user:` / session-scoped state in MongoDB. Multi-document + writes run in a transaction (replica set, sharded cluster, or Atlas + required) so concurrent writers cannot lose each other's state updates; + on deployments without transaction support (standalone mongod) the + service falls back to sequential writes with optimistic concurrency + control on the session document. +- **Long-term Memory (`MongoDbMemoryService`):** Ingests session events into a + memory collection and recalls them by keyword matching. - **Flexible Connection Options:** Pass a connection string and let the integration create/own the PyMongo client, or bring your own pre-configured `pymongo.MongoClient`. @@ -36,6 +46,8 @@ pip install "pymongo>=4.9,<5" | :--- | :--- | | `mongodb_vector_search` | A **vector search index** on the collection, and a deployment that supports `$vectorSearch` (MongoDB Atlas 6.0.11+/7.0.2+, or self-managed MongoDB 8.2+). | | `mongodb_hybrid_search` | A **vector search index** and a **full-text search index** on the collection, and a deployment that supports `$rankFusion` (MongoDB Atlas 8.0+). | +| `MongoDbSessionService` | Any MongoDB deployment reachable by PyMongo (no search indexes needed). | +| `MongoDbMemoryService` | Any MongoDB deployment reachable by PyMongo (no search indexes needed). | ## Quick Start: Search Tools @@ -114,14 +126,143 @@ MongoDB Atlas and on self-managed deployments running MongoDB 8.3+ with `mongot` 1.70.1+ and Voyage AI API keys [configured on `mongot`](https://www.mongodb.com/docs/search/self-managed/current/configuration/automated-embedding/). +## Quick Start: Session Service + +```python +from google.adk.integrations.mongodb import MongoDbSessionService +from google.adk.runners import Runner + +session_service = MongoDbSessionService( + connection_string="mongodb+srv://user:pass@cluster.mongodb.net/", + database_name="my_app", +) + +runner = Runner( + agent=agent, + app_name="my_app", + session_service=session_service, +) +``` + +### Document Layout + +`MongoDbSessionService` stores data in four collections (names configurable): + +| Collection | Document key | Contents | +| :--- | :--- | :--- | +| `sessions` | `//` | Session-scoped state (JSON-encoded), timestamps, and an optimistic-concurrency `revision`. | +| `events` | `///` | Full serialized event under `event_data`. | +| `app_states` | `` | App-scoped state (`app:` prefixed keys). | +| `user_states` | `/` | User-scoped state (`user:` prefixed keys). | + +State buckets are stored JSON-encoded so state keys containing characters +MongoDB forbids in document fields (e.g. `.`, `$`) round-trip safely. + +### Concurrency model and deployment targets + +The service targets two deployment shapes: + +- **Replica set, sharded cluster, or Atlas:** every multi-document write + (state buckets, the session revision bump, and the event insert) runs in a + multi-document transaction, so it commits or aborts atomically. +- **Standalone mongod (or any deployment without transactions):** the first + write attempts a transaction, the failure verdict is cached, and later + writes run sequentially with an optimistic-concurrency revision check on + the session document. + +Within one process, a per-session `asyncio.Lock` serializes `append_event` +calls for the same session. That lock is intentionally **in-process only**: +when several runner replicas write to the same session, the revision check +on the session document is the cross-process guard — the losing writer gets +a `StaleSessionError`, reloads the session, and retries. Do not rely on the +in-process lock for correctness across replicas. + +## Quick Start: Memory Service + +```python +from google.adk.integrations.mongodb import MongoDbMemoryService +from google.adk.runners import Runner + +memory_service = MongoDbMemoryService( + connection_string="mongodb+srv://user:pass@cluster.mongodb.net/", + database_name="my_app", +) + +runner = Runner( + agent=agent, + app_name="my_app", + memory_service=memory_service, +) +``` + +Events ingested via `add_session_to_memory` are stored as memory documents +(one per event) keyed by `///`, each +carrying the lowercase keywords extracted from its text content (a compact +English stop-word list is ignored; pass `stop_words` to customize). +`search_memory` returns entries whose keyword array intersects the query's +keywords. Ingestion is idempotent — re-adding a session overwrites its +memories rather than duplicating them. + +Retrieval is intentionally keyword-only: it needs no embedding model, no +Google credentials, and no vector index, so it works on any MongoDB +deployment (not just Atlas) at zero per-query cost. Vector recall is a +possible follow-up that can build on this integration's own search tools. + +## Secondary indexes (`ensure_indexes`) + +The services never create indexes on their own: every write and most reads +are keyed by `_id`, which MongoDB indexes implicitly. Two read patterns are +not `_id` lookups, though, and become collection scans as data grows: + +- `get_session` reads `events` by `(app_name, user_id, session_id)` sorted + by `timestamp`. +- `search_memory` reads `memories` by `(app_name, user_id)` with a `$in` + match on the `keywords` array. + +Pass `ensure_indexes=True` to have the services create the recommended +indexes on first use: + +```python +session_service = MongoDbSessionService( + connection_string="mongodb+srv://user:pass@cluster.mongodb.net/", + database_name="my_app", + ensure_indexes=True, +) +memory_service = MongoDbMemoryService( + connection_string="mongodb+srv://user:pass@cluster.mongodb.net/", + database_name="my_app", + ensure_indexes=True, +) +``` + +This creates: + +| Collection | Index | Keys | +| :--- | :--- | :--- | +| `events` | `app_user_session_ts` | `(app_name, user_id, session_id, timestamp DESC)` | +| `sessions` | `app_user` | `(app_name, user_id)` | +| `memories` | `app_user_keywords` | `(app_name, user_id, keywords)` (multikey) | + +Index creation is idempotent — re-running with the same spec is a no-op — so +the flag is safe to leave on. It requires the MongoDB user to hold the +`createIndex` privilege on the database, and it runs outside the write +transactions (`createIndexes` is not allowed inside one). At small scale the +flag changes nothing measurably; enable it before the `events` and +`memories` collections grow large. + ## Required permissions (least-privilege setup) -Do not hand the integration an admin user. The search tools only read, so the -database user needs nothing more than `find` on each collection the agent may -search: +Do not hand the integration an admin user. The components need very little, +and two separate database users keep the privileges separated: + +| Database user | Used by | Minimum privileges | +| :--- | :--- | :--- | +| **Search user** (read-only) | `MongoDbToolset` | `find` on each collection the agent may search. | +| **State user** (read-write) | `MongoDbSessionService`, `MongoDbMemoryService` | `find`, `insert`, `update`, `remove` on `sessions` and `events` (session deletion); `find`, `insert`, `update` on `app_states`, `user_states` and `memories` (upserts only, never deletes). Add `createIndex` on those collections only if you enable `ensure_indexes`. | ```javascript -// mongosh, as an admin user. +// mongosh, as an admin user: one read-only role for the search tools, one +// read-write role scoped to the state collections for the service. use admin; db.createRole({ role: "adk_search_reader", @@ -130,13 +271,24 @@ db.createRole({ ], roles: [], }); +db.createRole({ + role: "adk_state_writer", + privileges: ["sessions", "events"].map((c) => ({ + resource: { db: "my_app", collection: c }, + actions: ["find", "insert", "update", "remove"], + })).concat(["app_states", "user_states", "memories"].map((c) => ({ + resource: { db: "my_app", collection: c }, + actions: ["find", "insert", "update"], + }))), + roles: [], +}); ``` Creating and dropping **search indexes** is an administrative action, not something the tools do at run time — run setup with a separate user holding the `createSearchIndexes`, `listSearchIndexes`, and `dropSearchIndexes` actions (on Atlas, the built-in `Project Data Access Admin` role covers -this), and keep those privileges off the runtime user above. +this), and keep those privileges off the runtime users above. ### The model chooses the collection diff --git a/src/google/adk/integrations/mongodb/__init__.py b/src/google/adk/integrations/mongodb/__init__.py index da41da744d2..fa3b9601953 100644 --- a/src/google/adk/integrations/mongodb/__init__.py +++ b/src/google/adk/integrations/mongodb/__init__.py @@ -15,7 +15,8 @@ """MongoDB Integration (Experimental). MongoDB tools for vector search and hybrid search against collections on -MongoDB Atlas or MongoDB 8.0+ deployments. +MongoDB Atlas or MongoDB 8.0+ deployments, plus MongoDB-backed session and +memory services. """ from __future__ import annotations @@ -23,11 +24,15 @@ import typing if typing.TYPE_CHECKING: + from ._memory_service import MongoDbMemoryService from ._mongodb_toolset import MongoDbToolset + from ._session_service import MongoDbSessionService from ._settings import MongoDbToolSettings # Map attribute names to relative module paths _lazy_imports = { + "MongoDbMemoryService": "._memory_service", + "MongoDbSessionService": "._session_service", "MongoDbToolset": "._mongodb_toolset", "MongoDbToolSettings": "._settings", } @@ -49,6 +54,8 @@ def __dir__() -> list[str]: __all__ = [ + "MongoDbMemoryService", + "MongoDbSessionService", "MongoDbToolset", "MongoDbToolSettings", ] diff --git a/src/google/adk/integrations/mongodb/_memory_service.py b/src/google/adk/integrations/mongodb/_memory_service.py new file mode 100644 index 00000000000..210bbd9d9d1 --- /dev/null +++ b/src/google/adk/integrations/mongodb/_memory_service.py @@ -0,0 +1,303 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import asyncio +import logging +import re +import threading +from typing import Any +from typing import TYPE_CHECKING + +from typing_extensions import override + +from . import _client +from ...features import experimental +from ...features import FeatureName +from ...memory import _utils +from ...memory.base_memory_service import BaseMemoryService +from ...memory.base_memory_service import SearchMemoryResponse +from ...memory.memory_entry import MemoryEntry + +if TYPE_CHECKING: + from pymongo import MongoClient + + from ...sessions.session import Session + +logger = logging.getLogger("google_adk." + __name__) + +DEFAULT_MEMORIES_COLLECTION = "memories" + +# Recommended secondary index, created only when `ensure_indexes=True`. It +# covers the `search_memory` lookup (app/user equality plus a multikey `$in` +# on keywords); writes are keyed by `_id`, which MongoDB indexes implicitly. +_MEMORIES_RECALL_INDEX_NAME = "app_user_keywords" +_MEMORIES_RECALL_INDEX_KEYS = [("app_name", 1), ("user_id", 1), ("keywords", 1)] + +# Compact English stop-word list ignored when extracting keywords. Kept +# intentionally short; callers can pass their own set via `stop_words`. +DEFAULT_STOP_WORDS = { + "a", + "an", + "and", + "are", + "as", + "at", + "be", + "but", + "by", + "for", + "from", + "has", + "have", + "how", + "i", + "if", + "in", + "into", + "is", + "it", + "its", + "me", + "my", + "of", + "on", + "or", + "our", + "so", + "that", + "the", + "their", + "this", + "to", + "was", + "we", + "what", + "when", + "where", + "which", + "who", + "will", + "with", + "you", + "your", +} + + +@experimental(FeatureName.MONGODB_MEMORY_SERVICE) +class MongoDbMemoryService(BaseMemoryService): + """Memory service that uses MongoDB as the backend. + + Events ingested from sessions are stored as memory documents (one per + event) in a single collection, each carrying the lowercase keywords + extracted from its text content. `search_memory` matches documents whose + keyword array intersects the query's keywords. + + Example: + ```python + memory_service = MongoDbMemoryService( + connection_string="mongodb+srv://user:pass@cluster.mongodb.net/", + database_name="my_app", + ) + runner = Runner( + agent=agent, app_name="my_app", memory_service=memory_service + ) + ``` + """ + + def __init__( + self, + *, + database_name: str, + connection_string: str | None = None, + mongo_client: MongoClient | None = None, + memories_collection: str = DEFAULT_MEMORIES_COLLECTION, + stop_words: set[str] | None = None, + ensure_indexes: bool = False, + ): + """Initializes the MongoDB memory service. + + Args: + database_name: The MongoDB database used to store memories. + connection_string: The MongoDB connection string (URI) used to create a + client owned by this service. Requires the `pymongo` package + (`pip install google-adk[mongodb]`). + mongo_client: An existing PyMongo client to use instead of creating one + from `connection_string`. The caller keeps ownership of the client. + memories_collection: Collection name for memory documents. + stop_words: Words to ignore when extracting keywords. Defaults to a + standard English stop-word list. + ensure_indexes: When True, create the recommended secondary index on + first use: `app_user_keywords` on the memories collection (the + app/user/keywords lookup behind `search_memory`). Index creation is + idempotent, so this is safe to leave on, but it requires the + `createIndex` privilege. + """ + if mongo_client is not None and connection_string is not None: + raise ValueError( + "Only one of `connection_string` and `mongo_client` may be provided." + ) + if mongo_client is not None: + self._client = mongo_client + self._owns_client = False + elif connection_string is not None: + self._client = _client.get_mongo_client(connection_string) + self._owns_client = True + else: + raise ValueError( + "Either `connection_string` or `mongo_client` must be provided." + ) + self._database_name = database_name + self.memories_collection = memories_collection + self.stop_words = ( + stop_words if stop_words is not None else DEFAULT_STOP_WORDS + ) + self._ensure_indexes = ensure_indexes + self._indexes_ensured = False + # Serializes first-use index creation across the worker threads that + # `asyncio.to_thread` runs `_ensure_indexes_once` on. + self._index_lock = threading.Lock() + + def _memories(self): + return self._client[self._database_name][self.memories_collection] + + def _ensure_indexes_once(self) -> None: + """Creates the recommended secondary index on first use, if enabled. + + Index creation is idempotent on the server, and the guard flag only + skips work in this process. + """ + if not self._ensure_indexes or self._indexes_ensured: + return + with self._index_lock: + # Re-check under the lock: a concurrent first call may have done it. + if self._indexes_ensured: + return + self._memories().create_index( + _MEMORIES_RECALL_INDEX_KEYS, name=_MEMORIES_RECALL_INDEX_NAME + ) + self._indexes_ensured = True + + def _extract_keywords(self, text: str) -> set[str]: + """Extracts lowercase keywords from text, ignoring stop words.""" + words = re.findall(r"[a-z0-9]+", text.lower()) + return {word for word in words if word not in self.stop_words} + + @override + async def add_session_to_memory(self, session: Session) -> None: + """Ingests the session's text events into the memory collection. + + Ingestion is idempotent: memory documents are keyed by + `///`, so re-adding a session + overwrites its memories rather than duplicating them. + """ + await asyncio.to_thread(self._ensure_indexes_once) + + def _add() -> None: + # pymongo is imported lazily here: it stays an optional dependency + # until a MongoDB-backed service is actually constructed. + from pymongo import ReplaceOne # pylint: disable=import-outside-toplevel + + # One bulk_write round-trip, not one replace_one per event. + operations = [] + for event in session.events: + if not event.content or not event.content.parts: + continue + text = " ".join( + [part.text for part in event.content.parts if part.text] + ) + if not text: + continue + keywords = self._extract_keywords(text) + if not keywords: + continue + memory_id = ( + f"{session.app_name}/{session.user_id}/{session.id}/{event.id}" + ) + operations.append( + ReplaceOne( + {"_id": memory_id}, + { + "app_name": session.app_name, + "user_id": session.user_id, + "session_id": session.id, + "author": event.author, + "keywords": sorted(keywords), + "content": event.content.model_dump( + exclude_none=True, mode="json" + ), + "timestamp": event.timestamp, + }, + upsert=True, + ) + ) + if operations: + self._memories().bulk_write(operations, ordered=False) + + await asyncio.to_thread(_add) + + @override + async def search_memory( + self, *, app_name: str, user_id: str, query: str + ) -> SearchMemoryResponse: + """Searches memory for events matching the query's keywords.""" + await asyncio.to_thread(self._ensure_indexes_once) + + keywords = self._extract_keywords(query) + if not keywords: + return SearchMemoryResponse() + + def _search() -> list[dict[str, Any]]: + cursor = self._memories().find({ + "app_name": app_name, + "user_id": user_id, + "keywords": {"$in": sorted(keywords)}, + }) + return list(cursor) + + docs = await asyncio.to_thread(_search) + + seen = set() + memories = [] + for doc in docs: + try: + from google.genai import types + + content = types.Content.model_validate(doc["content"]) + entry = MemoryEntry( + id=doc.get("_id"), + content=content, + author=doc.get("author"), + timestamp=_utils.format_timestamp(doc.get("timestamp", 0.0)), + ) + except Exception as exc: + logger.warning(f"Failed to parse memory entry: {exc}") + continue + content_text = ( + " ".join([part.text for part in content.parts if part.text]) + if content.parts + else "" + ) + key = (entry.author, content_text, entry.timestamp) + if key not in seen: + seen.add(key) + memories.append(entry) + + return SearchMemoryResponse(memories=memories) + + async def close(self) -> None: + """Closes the MongoDB client if it was created by this service.""" + if self._owns_client: + await asyncio.to_thread(self._client.close) diff --git a/src/google/adk/integrations/mongodb/_session_service.py b/src/google/adk/integrations/mongodb/_session_service.py new file mode 100644 index 00000000000..fc4b233a921 --- /dev/null +++ b/src/google/adk/integrations/mongodb/_session_service.py @@ -0,0 +1,694 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import asyncio +from contextlib import asynccontextmanager +import copy +import json +import logging +import threading +import time +from typing import Any +from typing import AsyncGenerator +from typing import TYPE_CHECKING + +from typing_extensions import override + +from . import _client +from ...errors._stale_session_error import StaleSessionError +from ...errors.already_exists_error import AlreadyExistsError +from ...errors.session_not_found_error import SessionNotFoundError +from ...events.event import Event +from ...features import experimental +from ...features import FeatureName +from ...platform import uuid as platform_uuid +from ...sessions import _session_util +from ...sessions.base_session_service import BaseSessionService +from ...sessions.base_session_service import GetSessionConfig +from ...sessions.base_session_service import ListSessionsResponse +from ...sessions.session import Session +from ...sessions.state import State + +if TYPE_CHECKING: + from pymongo import MongoClient + +logger = logging.getLogger("google_adk." + __name__) + +_STALE_SESSION_ERROR_MESSAGE = ( + "The session has been modified in storage since it was loaded. " + "Please reload the session before appending more events." +) + +DEFAULT_SESSIONS_COLLECTION = "sessions" +DEFAULT_EVENTS_COLLECTION = "events" +DEFAULT_APP_STATE_COLLECTION = "app_states" +DEFAULT_USER_STATE_COLLECTION = "user_states" + +# Recommended secondary indexes, created only when `ensure_indexes=True`. +# Reads and writes are otherwise keyed by `_id`, which MongoDB indexes +# implicitly; these cover the two query patterns that are not `_id` lookups. +_EVENTS_HISTORY_INDEX_NAME = "app_user_session_ts" +_EVENTS_HISTORY_INDEX_KEYS = [ + ("app_name", 1), + ("user_id", 1), + ("session_id", 1), + ("timestamp", -1), +] +_SESSIONS_LIST_INDEX_NAME = "app_user" +_SESSIONS_LIST_INDEX_KEYS = [("app_name", 1), ("user_id", 1)] + +_SessionLockKey = tuple[str, str, str] + + +def _is_transaction_unsupported_error(exc: Exception) -> bool: + """Returns True when the deployment cannot run multi-document transactions. + + mongomock raises NotImplementedError for start_session, and a standalone + mongod fails the first transactional operation with IllegalOperation + (error code 20). Both conditions are permanent for the life of the client, + so callers may cache the verdict and stop attempting transactions. + """ + if isinstance(exc, NotImplementedError): + return True + if getattr(exc, "code", None) == 20: + return True + message = str(exc).lower() + return "transaction" in message and ( + "replica set" in message or "not supported" in message + ) + + +def _dumps_state(state: dict[str, Any]) -> str: + """Serializes a state bucket, coercing non-JSON values first.""" + return json.dumps(_session_util.make_json_safe_state(state)) + + +def _loads_state(raw: Any) -> dict[str, Any]: + """Parses a stored state bucket (JSON string or legacy dict).""" + if isinstance(raw, str): + return json.loads(raw) + return dict(raw or {}) + + +@experimental(FeatureName.MONGODB_SESSION_SERVICE) +class MongoDbSessionService(BaseSessionService): + """Session service that uses MongoDB as the backend. + + Document layout within the configured database: + + - ``: one document per session, keyed by + `//`, holding the session-scoped state + (JSON-encoded), timestamps, and an optimistic-concurrency `revision`. + - ``: one document per event, keyed by + `///`, holding the full + serialized event under `event_data`. + - ``: one document per app, keyed by ``. + - ``: one document per user, keyed by + `/`. + + State buckets are stored JSON-encoded so state keys containing characters + that MongoDB forbids in document fields (e.g. `.`, `$`) round-trip safely. + + Writes that span multiple documents (state buckets, the session revision + bump, and the event insert) run inside a multi-document transaction so + they commit or abort together, and so concurrent writers to shared app or + user state lose to a retry instead of silently overwriting each other. + Transactions require a replica set, sharded cluster, or Atlas deployment; + on deployments without them (standalone mongod) the service falls back to + sequential writes and logs a warning. + + Example: + ```python + session_service = MongoDbSessionService( + connection_string="mongodb+srv://user:pass@cluster.mongodb.net/", + database_name="my_app", + ) + runner = Runner( + agent=agent, app_name="my_app", session_service=session_service + ) + ``` + """ + + def __init__( + self, + *, + database_name: str, + connection_string: str | None = None, + mongo_client: MongoClient | None = None, + sessions_collection: str = DEFAULT_SESSIONS_COLLECTION, + events_collection: str = DEFAULT_EVENTS_COLLECTION, + app_state_collection: str = DEFAULT_APP_STATE_COLLECTION, + user_state_collection: str = DEFAULT_USER_STATE_COLLECTION, + ensure_indexes: bool = False, + ): + """Initializes the MongoDB session service. + + Args: + database_name: The MongoDB database used to store sessions, events and + shared state. + connection_string: The MongoDB connection string (URI) used to create a + client owned by this service. Requires the `pymongo` package + (`pip install google-adk[mongodb]`). + mongo_client: An existing PyMongo client to use instead of creating one + from `connection_string`. The caller keeps ownership of the client. + sessions_collection: Collection name for session documents. + events_collection: Collection name for event documents. + app_state_collection: Collection name for app state documents. + user_state_collection: Collection name for user state documents. + ensure_indexes: When True, create the recommended secondary indexes on + first use: `app_user_session_ts` on the events collection (the + app/user/session/timestamp lookup behind `get_session`), and + `app_user` on the sessions collection (the app/user filter behind + `list_sessions`). Index creation is idempotent, so this is safe to + leave on, but it requires the `createIndex` privilege. All other + reads and writes are keyed by `_id`, which MongoDB indexes + implicitly. + """ + if mongo_client is not None and connection_string is not None: + raise ValueError( + "Only one of `connection_string` and `mongo_client` may be provided." + ) + if mongo_client is not None: + self._client = mongo_client + self._owns_client = False + elif connection_string is not None: + self._client = _client.get_mongo_client(connection_string) + self._owns_client = True + else: + raise ValueError( + "Either `connection_string` or `mongo_client` must be provided." + ) + self._database_name = database_name + self.sessions_collection = sessions_collection + self.events_collection = events_collection + self.app_state_collection = app_state_collection + self.user_state_collection = user_state_collection + self._ensure_indexes = ensure_indexes + self._indexes_ensured = False + + # Per-session locks used to serialize append_event calls in this process. + self._session_locks: dict[_SessionLockKey, asyncio.Lock] = {} + self._session_lock_ref_count: dict[_SessionLockKey, int] = {} + self._session_locks_guard = asyncio.Lock() + + # Serializes first-use index creation across the worker threads that + # `asyncio.to_thread` runs `_ensure_indexes_once` on. + self._index_lock = threading.Lock() + + # Set to False after the deployment proves it cannot run multi-document + # transactions (standalone mongod, mongomock), so later writes skip the + # transaction attempt instead of paying for a failure each time. + self._transactions_supported = True + + def _sessions(self): + return self._client[self._database_name][self.sessions_collection] + + def _events(self): + return self._client[self._database_name][self.events_collection] + + def _app_states(self): + return self._client[self._database_name][self.app_state_collection] + + def _user_states(self): + return self._client[self._database_name][self.user_state_collection] + + def _ensure_indexes_once(self) -> None: + """Creates the recommended secondary indexes on first use, if enabled. + + Index creation is idempotent on the server, and the guard flag only + skips work in this process. It must run outside the write transactions: + `createIndexes` is not allowed inside a multi-document transaction. + """ + if not self._ensure_indexes or self._indexes_ensured: + return + with self._index_lock: + # Re-check under the lock: a concurrent first call may have done it. + if self._indexes_ensured: + return + self._events().create_index( + _EVENTS_HISTORY_INDEX_KEYS, name=_EVENTS_HISTORY_INDEX_NAME + ) + self._sessions().create_index( + _SESSIONS_LIST_INDEX_KEYS, name=_SESSIONS_LIST_INDEX_NAME + ) + self._indexes_ensured = True + + @staticmethod + def _session_key(app_name: str, user_id: str, session_id: str) -> str: + return f"{app_name}/{user_id}/{session_id}" + + @staticmethod + def _user_key(app_name: str, user_id: str) -> str: + return f"{app_name}/{user_id}" + + @asynccontextmanager + async def _with_session_lock( + self, *, app_name: str, user_id: str, session_id: str + ) -> AsyncGenerator[None]: + """Serializes event appends for the same session within this process.""" + lock_key = (app_name, user_id, session_id) + async with self._session_locks_guard: + lock = self._session_locks.get(lock_key) + if lock is None: + lock = asyncio.Lock() + self._session_locks[lock_key] = lock + self._session_lock_ref_count[lock_key] = ( + self._session_lock_ref_count.get(lock_key, 0) + 1 + ) + + try: + async with lock: + yield + finally: + async with self._session_locks_guard: + remaining = self._session_lock_ref_count.get(lock_key, 0) - 1 + if remaining <= 0 and not lock.locked(): + self._session_lock_ref_count.pop(lock_key, None) + self._session_locks.pop(lock_key, None) + else: + self._session_lock_ref_count[lock_key] = remaining + + def _run_in_transaction(self, work): + """Runs `work(mongo_session)` as one atomic unit when possible. + + Multi-document transactions make the state-bucket merges, the session + revision bump, and the event write commit or abort together, and a + concurrent conflicting writer loses to a write-conflict retry instead of + silently overwriting. Transactions require a replica set, sharded + cluster, or Atlas deployment. Where they are unavailable (a standalone + mongod, or mongomock in tests) the work runs directly without a session, + preserving the pre-transaction behavior. + + Args: + work: Callable taking the pymongo client session (or None when + transactions are unsupported) and returning the work's result. + """ + if not self._transactions_supported: + return work(None) + try: + mongo_session = self._client.start_session() + except Exception as exc: + if not _is_transaction_unsupported_error(exc): + raise + self._transactions_supported = False + logger.warning( + "MongoDB deployment does not support sessions/transactions (%s); " + "writes will not be atomic across documents.", + exc, + ) + return work(None) + with mongo_session: + try: + # with_transaction retries the callback on TransientTransactionError + # and re-commits on UnknownTransactionCommitResult; other errors + # (StaleSessionError, SessionNotFoundError, ...) abort and propagate. + return mongo_session.with_transaction(work) + except Exception as exc: + if not _is_transaction_unsupported_error(exc): + raise + self._transactions_supported = False + logger.warning( + "MongoDB deployment does not support transactions (%s); " + "writes will not be atomic across documents.", + exc, + ) + return work(None) + + @staticmethod + def _merge_state( + app_state: dict[str, Any] | None, + user_state: dict[str, Any] | None, + session_state: dict[str, Any], + ) -> dict[str, Any]: + """Merges app, user, and session states into a single state dictionary.""" + merged_state = copy.deepcopy(session_state) + for key, value in (app_state or {}).items(): + merged_state[State.APP_PREFIX + key] = value + for key, value in (user_state or {}).items(): + merged_state[State.USER_PREFIX + key] = value + return merged_state + + def _read_app_state( + self, app_name: str, mongo_session: Any = None + ) -> dict[str, Any]: + doc = self._app_states().find_one({"_id": app_name}, session=mongo_session) + return _loads_state(doc.get("state")) if doc else {} + + def _read_user_state( + self, app_name: str, user_id: str, mongo_session: Any = None + ) -> dict[str, Any]: + doc = self._user_states().find_one( + {"_id": self._user_key(app_name, user_id)}, session=mongo_session + ) + return _loads_state(doc.get("state")) if doc else {} + + def _merge_state_bucket( + self, + collection: Any, + doc_id: str, + delta: dict[str, Any], + mongo_session: Any = None, + ) -> dict[str, Any]: + """Merges delta into a stored state bucket and returns the merged state. + + The read-modify-write is only safe against lost updates when it runs + inside a transaction (`mongo_session` set); concurrent writers otherwise + race on the read and the last write wins. + """ + existing = collection.find_one({"_id": doc_id}, session=mongo_session) + merged = _loads_state(existing.get("state")) if existing else {} + merged.update(delta) + collection.update_one( + {"_id": doc_id}, + {"$set": {"state": _dumps_state(merged)}}, + upsert=True, + session=mongo_session, + ) + return merged + + def _to_session( + self, + doc: dict[str, Any], + merged_state: dict[str, Any], + events: list[Event], + ) -> Session: + session = Session( + id=doc["id"], + app_name=doc["app_name"], + user_id=doc["user_id"], + state=merged_state, + events=events, + last_update_time=doc.get("update_time", 0.0), + ) + session._storage_update_marker = str(doc.get("revision", 0)) + return session + + @override + async def create_session( + self, + *, + app_name: str, + user_id: str, + state: dict[str, Any] | None = None, + session_id: str | None = None, + ) -> Session: + """Creates a new session in MongoDB.""" + await asyncio.to_thread(self._ensure_indexes_once) + + def _create(mongo_session) -> tuple[str, dict[str, Any]]: + sid = session_id or platform_uuid.new_uuid() + state_deltas = _session_util.extract_state_delta(state or {}) + + app_state = ( + self._merge_state_bucket( + self._app_states(), + app_name, + state_deltas["app"], + mongo_session, + ) + if state_deltas["app"] + else self._read_app_state(app_name, mongo_session) + ) + user_state = ( + self._merge_state_bucket( + self._user_states(), + self._user_key(app_name, user_id), + state_deltas["user"], + mongo_session, + ) + if state_deltas["user"] + else self._read_user_state(app_name, user_id, mongo_session) + ) + + now = time.time() + doc = { + "_id": self._session_key(app_name, user_id, sid), + "id": sid, + "app_name": app_name, + "user_id": user_id, + "state": _dumps_state(state_deltas["session"]), + "create_time": now, + "update_time": now, + "revision": 0, + } + try: + self._sessions().insert_one(doc, session=mongo_session) + except Exception as exc: + if exc.__class__.__name__ == "DuplicateKeyError": + raise AlreadyExistsError(f"Session {sid} already exists.") from exc + raise + merged = self._merge_state(app_state, user_state, state_deltas["session"]) + return sid, merged + + sid, merged_state = await asyncio.to_thread( + self._run_in_transaction, _create + ) + session = Session( + id=sid, + app_name=app_name, + user_id=user_id, + state=merged_state, + events=[], + last_update_time=time.time(), + ) + session._storage_update_marker = "0" + return session + + @override + async def get_session( + self, + *, + app_name: str, + user_id: str, + session_id: str, + config: GetSessionConfig | None = None, + ) -> Session | None: + """Gets a session from MongoDB.""" + await asyncio.to_thread(self._ensure_indexes_once) + + def _get() -> Session | None: + doc = self._sessions().find_one( + {"_id": self._session_key(app_name, user_id, session_id)} + ) + if not doc: + return None + + # A requested count of zero asks for no event history at all (callers + # use it to probe whether a session exists), so skip the events query. + events: list[Event] = [] + if config is None or config.num_recent_events != 0: + query: dict[str, Any] = { + "app_name": app_name, + "user_id": user_id, + "session_id": session_id, + } + if config and config.after_timestamp is not None: + query["timestamp"] = {"$gte": config.after_timestamp} + + cursor = self._events().find(query).sort("timestamp", -1) + if config and config.num_recent_events is not None: + cursor = cursor.limit(config.num_recent_events) + event_docs = list(cursor) + event_docs.reverse() # restore chronological order + events = [ + Event.model_validate(event_doc["event_data"]) + for event_doc in event_docs + ] + + merged = self._merge_state( + self._read_app_state(app_name), + self._read_user_state(app_name, user_id), + _loads_state(doc.get("state")), + ) + return self._to_session(doc, merged, events) + + return await asyncio.to_thread(_get) + + @override + async def list_sessions( + self, *, app_name: str, user_id: str | None = None + ) -> ListSessionsResponse: + """Lists sessions from MongoDB, oldest update first.""" + await asyncio.to_thread(self._ensure_indexes_once) + + def _list() -> list[Session]: + query: dict[str, Any] = {"app_name": app_name} + if user_id: + query["user_id"] = user_id + docs = list(self._sessions().find(query)) + + app_state = self._read_app_state(app_name) + # One query for every user present, not one query per user (N+1). + user_keys = [ + self._user_key(app_name, uid) + for uid in {doc["user_id"] for doc in docs} + ] + user_states = { + doc["_id"]: _loads_state(doc.get("state")) + for doc in self._user_states().find({"_id": {"$in": user_keys}}) + } + + sessions = [ + self._to_session( + doc, + self._merge_state( + app_state, + user_states.get(self._user_key(app_name, doc["user_id"]), {}), + _loads_state(doc.get("state")), + ), + [], + ) + for doc in docs + ] + sessions.sort(key=lambda s: (s.last_update_time, s.user_id, s.id)) + return sessions + + return ListSessionsResponse(sessions=await asyncio.to_thread(_list)) + + @override + async def delete_session( + self, *, app_name: str, user_id: str, session_id: str + ) -> None: + """Deletes a session and its events from MongoDB.""" + await asyncio.to_thread(self._ensure_indexes_once) + + def _delete(mongo_session) -> None: + self._events().delete_many( + {"app_name": app_name, "user_id": user_id, "session_id": session_id}, + session=mongo_session, + ) + self._sessions().delete_one( + {"_id": self._session_key(app_name, user_id, session_id)}, + session=mongo_session, + ) + + await asyncio.to_thread(self._run_in_transaction, _delete) + + @override + async def get_user_state( + self, *, app_name: str, user_id: str + ) -> dict[str, Any]: + """Returns the user-scoped state for the given app and user.""" + return await asyncio.to_thread(self._read_user_state, app_name, user_id) + + @override + async def append_event(self, session: Session, event: Event) -> Event: + """Appends an event to a session in MongoDB.""" + if event.partial: + return event + + self._apply_temp_state(session, event) + event = self._trim_temp_delta_state(event) + + state_delta = ( + event.actions.state_delta + if event.actions and event.actions.state_delta + else {} + ) + state_deltas = _session_util.extract_state_delta(state_delta) + + await asyncio.to_thread(self._ensure_indexes_once) + async with self._with_session_lock( + app_name=session.app_name, + user_id=session.user_id, + session_id=session.id, + ): + + def _append(mongo_session) -> int: + if state_deltas["app"]: + self._merge_state_bucket( + self._app_states(), + session.app_name, + state_deltas["app"], + mongo_session, + ) + if state_deltas["user"]: + self._merge_state_bucket( + self._user_states(), + self._user_key(session.app_name, session.user_id), + state_deltas["user"], + mongo_session, + ) + + session_only_state = { + key: value + for key, value in session.state.items() + if not key.startswith(State.APP_PREFIX) + and not key.startswith(State.USER_PREFIX) + and not key.startswith(State.TEMP_PREFIX) + } + session_only_state.update(state_deltas["session"]) + + session_doc_id = self._session_key( + session.app_name, session.user_id, session.id + ) + current = self._sessions().find_one( + {"_id": session_doc_id}, session=mongo_session + ) + if not current: + raise SessionNotFoundError(f"Session {session.id} not found.") + current_revision = current.get("revision", 0) + if session._storage_update_marker is not None and ( + session._storage_update_marker != str(current_revision) + ): + raise StaleSessionError(_STALE_SESSION_ERROR_MESSAGE) + + # The revision filter makes the update a no-op when a concurrent + # writer bumped the revision between our read and write. Inside a + # transaction the conflict instead aborts with + # TransientTransactionError, and the driver's retry re-reads the + # bumped revision and lands here. + updated = self._sessions().find_one_and_update( + {"_id": session_doc_id, "revision": current_revision}, + { + "$set": { + "state": _dumps_state(session_only_state), + "update_time": event.timestamp, + }, + "$inc": {"revision": 1}, + }, + return_document=True, + session=mongo_session, + ) + if updated is None: + raise StaleSessionError(_STALE_SESSION_ERROR_MESSAGE) + + # Upsert keeps event ingestion idempotent across retries. + self._events().replace_one( + {"_id": f"{session_doc_id}/{event.id}"}, + { + "app_name": session.app_name, + "user_id": session.user_id, + "session_id": session.id, + "timestamp": event.timestamp, + "event_data": event.model_dump(exclude_none=True, mode="json"), + }, + upsert=True, + session=mongo_session, + ) + return int(updated.get("revision", current_revision + 1)) + + new_revision = await asyncio.to_thread(self._run_in_transaction, _append) + session._storage_update_marker = str(new_revision) + session.last_update_time = event.timestamp + + await super().append_event(session, event) + return event + + async def close(self) -> None: + """Closes the MongoDB client if it was created by this service.""" + if self._owns_client: + await asyncio.to_thread(self._client.close) diff --git a/tests/unittests/integrations/mongodb/test_memory_service.py b/tests/unittests/integrations/mongodb/test_memory_service.py new file mode 100644 index 00000000000..4ad21ccb847 --- /dev/null +++ b/tests/unittests/integrations/mongodb/test_memory_service.py @@ -0,0 +1,259 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for MongoDbMemoryService, backed by mongomock.""" + +from __future__ import annotations + +import asyncio +from unittest import mock + +from google.adk.events.event import Event +from google.adk.integrations.mongodb import MongoDbMemoryService +from google.adk.sessions.session import Session +from google.genai import types +import mongomock +import pytest + + +@pytest.fixture(autouse=True) +def _mongomock_bulk_write_compat(monkeypatch): + """Drops the `sort` key pymongo 4.9+ adds to bulk write operations. + + mongomock's BulkOperationBuilder does not accept `sort` (added to + ReplaceOne/UpdateOne in pymongo 4.9), so bulk_write against mongomock + raises TypeError without this shim. Real MongoDB accepts the key; this + only patches the test double. + """ + from mongomock.collection import BulkOperationBuilder + + real_add_replace = BulkOperationBuilder.add_replace + real_add_update = BulkOperationBuilder.add_update + + def add_replace(self, selector, doc, upsert, sort=None, **kwargs): + return real_add_replace(self, selector, doc, upsert, **kwargs) + + def add_update( + self, selector, doc, multi=False, upsert=False, sort=None, **kwargs + ): + return real_add_update(self, selector, doc, multi, upsert, **kwargs) + + monkeypatch.setattr(BulkOperationBuilder, "add_replace", add_replace) + monkeypatch.setattr(BulkOperationBuilder, "add_update", add_update) + + +APP = "test_app" +USER = "test_user" + + +@pytest.fixture +def service(): + return MongoDbMemoryService( + mongo_client=mongomock.MongoClient(), database_name="test_db" + ) + + +@pytest.mark.asyncio +async def test_ensure_indexes_creates_recall_index(): + """ensure_indexes=True creates the app/user/keywords index on memories.""" + client = mongomock.MongoClient() + service = MongoDbMemoryService( + mongo_client=client, database_name="test_db", ensure_indexes=True + ) + await service.add_session_to_memory(_session_with_texts("hello world")) + + indexes = client["test_db"]["memories"].index_information() + assert "app_user_keywords" in indexes + + +@pytest.mark.asyncio +async def test_index_not_created_by_default(service): + """Without ensure_indexes the implicit _id index is the only one present.""" + await service.add_session_to_memory(_session_with_texts("hello world")) + + indexes = service._memories().index_information() + assert "app_user_keywords" not in indexes + + +def _session_with_texts(*texts: str, session_id: str = "s1") -> Session: + events = [ + Event( + invocation_id="inv", + author="user" if i % 2 == 0 else "agent", + content=types.Content( + role="user" if i % 2 == 0 else "model", + parts=[types.Part(text=text)], + ), + timestamp=float(i + 1), + ) + for i, text in enumerate(texts) + ] + return Session(id=session_id, app_name=APP, user_id=USER, events=events) + + +@pytest.mark.asyncio +async def test_add_session_then_search_by_keyword(service): + await service.add_session_to_memory( + _session_with_texts( + "I love hiking in the mountains", + "My favorite trail is the Pacific Crest Trail", + ) + ) + + response = await service.search_memory( + app_name=APP, user_id=USER, query="hiking" + ) + assert len(response.memories) == 1 + assert "hiking" in response.memories[0].content.parts[0].text + assert response.memories[0].author == "user" + + # A query can match several memories. + response = await service.search_memory( + app_name=APP, user_id=USER, query="trail mountains" + ) + assert len(response.memories) == 2 + + +@pytest.mark.asyncio +async def test_search_memory_scopes_by_app_and_user(service): + await service.add_session_to_memory(_session_with_texts("remember hiking")) + + assert ( + await service.search_memory( + app_name="other_app", user_id=USER, query="hiking" + ) + ).memories == [] + assert ( + await service.search_memory( + app_name=APP, user_id="other_user", query="hiking" + ) + ).memories == [] + + +@pytest.mark.asyncio +async def test_reingesting_session_does_not_duplicate(service): + session = _session_with_texts("I love hiking") + await service.add_session_to_memory(session) + await service.add_session_to_memory(session) + + response = await service.search_memory( + app_name=APP, user_id=USER, query="hiking" + ) + assert len(response.memories) == 1 + + +@pytest.mark.asyncio +async def test_search_memory_ignores_stop_words(service): + await service.add_session_to_memory(_session_with_texts("the cat sat")) + + # "the" is a stop word, so a stop-words-only query matches nothing even + # though the ingested text contains "the". + assert ( + await service.search_memory(app_name=APP, user_id=USER, query="the") + ).memories == [] + response = await service.search_memory( + app_name=APP, user_id=USER, query="cat" + ) + assert len(response.memories) == 1 + + +@pytest.mark.asyncio +async def test_search_memory_empty_or_stop_word_query(service): + await service.add_session_to_memory(_session_with_texts("remember hiking")) + assert ( + await service.search_memory(app_name=APP, user_id=USER, query="") + ).memories == [] + assert ( + await service.search_memory(app_name=APP, user_id=USER, query="?!") + ).memories == [] + + +@pytest.mark.asyncio +async def test_add_session_skips_events_without_text(service): + session = Session(id="s1", app_name=APP, user_id=USER) + session.events.append(Event(invocation_id="inv", author="agent")) + await service.add_session_to_memory(session) + assert ( + await service.search_memory(app_name=APP, user_id=USER, query="anything") + ).memories == [] + + +def test_constructor_validates_client_args(): + with pytest.raises(ValueError): + MongoDbMemoryService(database_name="db") + with pytest.raises(ValueError): + MongoDbMemoryService( + database_name="db", + mongo_client=mongomock.MongoClient(), + connection_string="mongodb://localhost:27017", + ) + + +@pytest.mark.asyncio +async def test_ingest_uses_one_bulk_write(service): + """Ingestion issues one bulk_write, not one round-trip per event.""" + bulk_calls = [] + replace_calls = [] + collection = service._memories() + real_bulk = collection.bulk_write + + def _counting_bulk(ops, **kwargs): + bulk_calls.append(ops) + return real_bulk(ops, **kwargs) + + def _forbidden_replace(*args, **kwargs): + replace_calls.append(args) + + with ( + mock.patch.object(collection, "bulk_write", _counting_bulk), + mock.patch.object(collection, "replace_one", _forbidden_replace), + ): + await service.add_session_to_memory( + _session_with_texts("first memory", "second memory", "third memory") + ) + + assert len(bulk_calls) == 1 + assert len(bulk_calls[0]) == 3 # one ReplaceOne per event + assert replace_calls == [] + + response = await service.search_memory( + app_name=APP, user_id=USER, query="memory" + ) + assert len(response.memories) == 3 + + +@pytest.mark.asyncio +async def test_index_creation_runs_once_under_concurrency(): + """Concurrent first calls create the recall index exactly once.""" + client = mongomock.MongoClient() + service = MongoDbMemoryService( + mongo_client=client, database_name="test_db", ensure_indexes=True + ) + create_calls = [] + memories = client["test_db"]["memories"] + real_create = memories.create_index + + def _counting_create(*args, **kwargs): + create_calls.append(kwargs.get("name")) + return real_create(*args, **kwargs) + + with mock.patch.object(memories, "create_index", _counting_create): + await asyncio.gather(*( + service.add_session_to_memory( + _session_with_texts(f"memory {i}", session_id=f"s{i}") + ) + for i in range(10) + )) + + assert create_calls == ["app_user_keywords"] diff --git a/tests/unittests/integrations/mongodb/test_session_service.py b/tests/unittests/integrations/mongodb/test_session_service.py new file mode 100644 index 00000000000..825b1769da9 --- /dev/null +++ b/tests/unittests/integrations/mongodb/test_session_service.py @@ -0,0 +1,469 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for MongoDbSessionService, backed by mongomock.""" + +from __future__ import annotations + +import asyncio +from unittest import mock + +from google.adk.errors import StaleSessionError +from google.adk.errors.already_exists_error import AlreadyExistsError +from google.adk.errors.session_not_found_error import SessionNotFoundError +from google.adk.events.event import Event +from google.adk.events.event_actions import EventActions +from google.adk.integrations.mongodb import MongoDbSessionService +from google.adk.sessions.base_session_service import GetSessionConfig +from google.genai import types +import mongomock +from pymongo.errors import OperationFailure +import pytest + +APP = "test_app" +USER = "test_user" +USER_2 = "other_user" + + +@pytest.fixture +def client(): + return mongomock.MongoClient() + + +@pytest.fixture +def service(client): + return MongoDbSessionService(mongo_client=client, database_name="test_db") + + +def _text_event( + text: str, *, author: str = "user", timestamp: float = 1.0 +) -> Event: + return Event( + invocation_id="inv", + author=author, + content=types.Content( + role="user" if author == "user" else "model", + parts=[types.Part(text=text)], + ), + timestamp=timestamp, + ) + + +@pytest.mark.asyncio +async def test_create_and_get_session(service): + session = await service.create_session( + app_name=APP, + user_id=USER, + state={"session_key": "s1", "app:app_key": "a1", "user:user_key": "u1"}, + session_id="s1", + ) + + assert session.id == "s1" + assert session.state["session_key"] == "s1" + assert session.state["app:app_key"] == "a1" + assert session.state["user:user_key"] == "u1" + assert session.events == [] + + fetched = await service.get_session( + app_name=APP, user_id=USER, session_id="s1" + ) + assert fetched is not None + assert fetched.state == session.state + assert fetched._storage_update_marker == "0" + + # Missing session returns None. + assert ( + await service.get_session(app_name=APP, user_id=USER, session_id="nope") + is None + ) + + +@pytest.mark.asyncio +async def test_create_duplicate_session_raises(service): + await service.create_session(app_name=APP, user_id=USER, session_id="s1") + with pytest.raises(AlreadyExistsError): + await service.create_session(app_name=APP, user_id=USER, session_id="s1") + + +@pytest.mark.asyncio +async def test_constructor_validates_client_args(client): + with pytest.raises(ValueError): + MongoDbSessionService(database_name="db") + with pytest.raises(ValueError): + MongoDbSessionService( + database_name="db", + mongo_client=client, + connection_string="mongodb://localhost:27017", + ) + + +@pytest.mark.asyncio +async def test_append_event_persists_and_merges_state(service): + session = await service.create_session( + app_name=APP, user_id=USER, session_id="s1" + ) + event = _text_event("hello mongodb", timestamp=10.0) + event.actions = EventActions( + state_delta={"turn": 1, "user:theme": "dark", "temp:scratch": "x"} + ) + await service.append_event(session, event) + + fetched = await service.get_session( + app_name=APP, user_id=USER, session_id="s1" + ) + assert len(fetched.events) == 1 + assert fetched.events[0].content.parts[0].text == "hello mongodb" + assert fetched.state["turn"] == 1 + assert fetched.state["user:theme"] == "dark" + # temp state lives on the in-memory session but is never persisted. + assert "temp:scratch" in session.state + assert "temp:scratch" not in fetched.state + assert fetched._storage_update_marker == "1" + # The in-memory session advanced too. + assert session._storage_update_marker == "1" + assert len(session.events) == 1 + + # User state is shared across the user's sessions. + session2 = await service.create_session( + app_name=APP, user_id=USER, session_id="s2" + ) + assert session2.state["user:theme"] == "dark" + assert await service.get_user_state(app_name=APP, user_id=USER) == { + "theme": "dark" + } + + +@pytest.mark.asyncio +async def test_append_event_missing_session_raises(service): + session = await service.create_session( + app_name=APP, user_id=USER, session_id="s1" + ) + session.id = "deleted" + with pytest.raises(SessionNotFoundError): + await service.append_event(session, _text_event("hi")) + + +@pytest.mark.asyncio +async def test_append_event_stale_session_raises(service): + session = await service.create_session( + app_name=APP, user_id=USER, session_id="s1" + ) + stale_copy = await service.get_session( + app_name=APP, user_id=USER, session_id="s1" + ) + + await service.append_event(session, _text_event("first", timestamp=1.0)) + # stale_copy still holds revision marker "0" while storage is at "1". + with pytest.raises(StaleSessionError): + await service.append_event(stale_copy, _text_event("second", timestamp=2.0)) + + +@pytest.mark.asyncio +async def test_append_event_is_idempotent_on_retry(service): + session = await service.create_session( + app_name=APP, user_id=USER, session_id="s1" + ) + event = _text_event("hello", timestamp=1.0) + await service.append_event(session, event) + + # Simulate a storage-level retry of the same event write: the document is + # replaced in place instead of duplicated. + service._events().replace_one( + {"_id": f"{APP}/{USER}/s1/{event.id}"}, + { + "app_name": APP, + "user_id": USER, + "session_id": "s1", + "timestamp": event.timestamp, + "event_data": event.model_dump(exclude_none=True, mode="json"), + }, + upsert=True, + ) + fetched = await service.get_session( + app_name=APP, user_id=USER, session_id="s1" + ) + assert len(fetched.events) == 1 + + +@pytest.mark.asyncio +async def test_get_session_config_filters_events(service): + session = await service.create_session( + app_name=APP, user_id=USER, session_id="s1" + ) + for i in range(1, 5): + await service.append_event( + session, _text_event(f"event {i}", timestamp=float(i)) + ) + + # num_recent_events=0 -> no events loaded. + fetched = await service.get_session( + app_name=APP, + user_id=USER, + session_id="s1", + config=GetSessionConfig(num_recent_events=0), + ) + assert fetched.events == [] + + # num_recent_events=2 -> two most recent, in chronological order. + fetched = await service.get_session( + app_name=APP, + user_id=USER, + session_id="s1", + config=GetSessionConfig(num_recent_events=2), + ) + assert [e.content.parts[0].text for e in fetched.events] == [ + "event 3", + "event 4", + ] + + # after_timestamp -> only events at or after the cursor. + fetched = await service.get_session( + app_name=APP, + user_id=USER, + session_id="s1", + config=GetSessionConfig(after_timestamp=2.5), + ) + assert [e.content.parts[0].text for e in fetched.events] == [ + "event 3", + "event 4", + ] + + +@pytest.mark.asyncio +async def test_list_sessions_sorted_oldest_first(service): + await service.create_session(app_name=APP, user_id=USER, session_id="s1") + await service.create_session(app_name=APP, user_id=USER, session_id="s2") + await service.create_session(app_name=APP, user_id=USER_2, session_id="s3") + + response = await service.list_sessions(app_name=APP, user_id=USER) + assert [s.id for s in response.sessions] == ["s1", "s2"] + + response = await service.list_sessions(app_name=APP) + assert {s.id for s in response.sessions} == {"s1", "s2", "s3"} + + +@pytest.mark.asyncio +async def test_list_sessions_batches_user_state_reads(service, client): + """User states are fetched in one $in query, not one query per user.""" + await service.create_session( + app_name=APP, + user_id=USER, + session_id="s1", + state={"user:theme": "dark"}, + ) + await service.create_session( + app_name=APP, + user_id=USER_2, + session_id="s2", + state={"user:theme": "light"}, + ) + await service.create_session(app_name=APP, user_id="user_3", session_id="s3") + + find_calls = [] + user_states = client["test_db"]["user_states"] + real_find = user_states.find + + def _counting_find(*args, **kwargs): + find_calls.append(args) + return real_find(*args, **kwargs) + + with mock.patch.object(user_states, "find", _counting_find): + response = await service.list_sessions(app_name=APP) + + assert len(find_calls) == 1 # not one per user + assert find_calls[0][0] == { + "_id": {"$in": mock.ANY} # the three / keys + } + by_id = {s.id: s for s in response.sessions} + assert by_id["s1"].state["user:theme"] == "dark" + assert by_id["s2"].state["user:theme"] == "light" + assert "user:theme" not in by_id["s3"].state + + +@pytest.mark.asyncio +async def test_index_creation_runs_once_under_concurrency(client): + """Concurrent first calls create each index exactly once.""" + service = MongoDbSessionService( + mongo_client=client, database_name="test_db", ensure_indexes=True + ) + create_calls = [] + events = client["test_db"]["events"] + sessions = client["test_db"]["sessions"] + real_create = events.create_index + + def _counting_create(*args, **kwargs): + create_calls.append(kwargs.get("name")) + return real_create(*args, **kwargs) + + with mock.patch.object(events, "create_index", _counting_create): + await asyncio.gather(*( + service.create_session(app_name=APP, user_id=USER, session_id=f"s{i}") + for i in range(10) + )) + + assert create_calls == ["app_user_session_ts"] + assert "app_user" in sessions.index_information() + + +@pytest.mark.asyncio +async def test_delete_session_removes_session_and_events(service): + session = await service.create_session( + app_name=APP, user_id=USER, session_id="s1" + ) + await service.append_event(session, _text_event("bye", timestamp=1.0)) + + await service.delete_session(app_name=APP, user_id=USER, session_id="s1") + + assert ( + await service.get_session(app_name=APP, user_id=USER, session_id="s1") + is None + ) + assert service._events().count_documents({}) == 0 + + +@pytest.mark.asyncio +async def test_state_keys_with_dots_round_trip(service): + """MongoDB forbids dots in document keys; JSON-encoded state does not care.""" + session = await service.create_session( + app_name=APP, user_id=USER, state={"nested.key": 1}, session_id="s1" + ) + await service.append_event( + session, + _text_event("x", timestamp=1.0), + ) + fetched = await service.get_session( + app_name=APP, user_id=USER, session_id="s1" + ) + assert fetched.state["nested.key"] == 1 + + +@pytest.mark.asyncio +async def test_ensure_indexes_creates_recommended_indexes(client): + """ensure_indexes=True creates the events and sessions secondary indexes.""" + service = MongoDbSessionService( + mongo_client=client, database_name="test_db", ensure_indexes=True + ) + session = await service.create_session( + app_name=APP, user_id=USER, session_id="s1" + ) + await service.append_event(session, _text_event("hi", timestamp=1.0)) + + events_indexes = client["test_db"]["events"].index_information() + assert "app_user_session_ts" in events_indexes + sessions_indexes = client["test_db"]["sessions"].index_information() + assert "app_user" in sessions_indexes + + +@pytest.mark.asyncio +async def test_indexes_not_created_by_default(service, client): + """Without ensure_indexes the implicit _id index is the only one present.""" + await service.create_session(app_name=APP, user_id=USER, session_id="s1") + + assert "app_user_session_ts" not in ( + client["test_db"]["events"].index_information() + ) + assert "app_user" not in (client["test_db"]["sessions"].index_information()) + + +@pytest.mark.asyncio +async def test_mongomock_falls_back_and_disables_transactions(service): + """mongomock has no sessions: writes run directly and the verdict caches.""" + session = await service.create_session( + app_name=APP, user_id=USER, session_id="s1" + ) + await service.append_event(session, _text_event("hi", timestamp=1.0)) + + assert service._transactions_supported is False + fetched = await service.get_session( + app_name=APP, user_id=USER, session_id="s1" + ) + assert fetched.events[0].content.parts[0].text == "hi" + + +@pytest.mark.asyncio +async def test_writes_run_inside_transaction_when_supported(): + """A transaction-capable client wraps every write sequence in a session.""" + client = mock.MagicMock() + mongo_session = client.start_session.return_value + mongo_session.with_transaction.side_effect = lambda callback: callback( + mongo_session + ) + db = client["test_db"] + db["app_states"].find_one.return_value = None + db["user_states"].find_one.return_value = None + db["sessions"].find_one.return_value = {"revision": 0} + db["sessions"].find_one_and_update.return_value = {"revision": 1} + + service = MongoDbSessionService(mongo_client=client, database_name="test_db") + session = await service.create_session( + app_name=APP, user_id=USER, session_id="s1" + ) + await service.append_event(session, _text_event("hi", timestamp=1.0)) + + assert service._transactions_supported is True + # create_session and append_event each ran their writes in a transaction. + assert mongo_session.with_transaction.call_count == 2 + assert db["sessions"].insert_one.call_args.kwargs["session"] is mongo_session + assert ( + db["sessions"].find_one_and_update.call_args.kwargs["session"] + is mongo_session + ) + assert db["events"].replace_one.call_args.kwargs["session"] is mongo_session + + +class _UnsupportedTransactionSession: + """Fake client session whose transaction fails like a standalone mongod.""" + + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + def with_transaction(self, callback): + raise OperationFailure( + "Transaction numbers are only allowed on a replica set member or" + " mongos", + code=20, + ) + + +@pytest.mark.asyncio +async def test_falls_back_when_transactions_not_supported( + service, client, monkeypatch +): + """Standalone mongod: first write falls back, later writes skip the attempt.""" + start_session_calls = [] + monkeypatch.setattr( + client, + "start_session", + lambda: ( + start_session_calls.append(1), + _UnsupportedTransactionSession(), + )[1], + ) + + session = await service.create_session( + app_name=APP, user_id=USER, session_id="s1" + ) + await service.append_event(session, _text_event("hi", timestamp=1.0)) + + # create_session attempted a transaction, failed, and fell back; + # append_event then skipped the attempt because the verdict was cached. + assert len(start_session_calls) == 1 + assert service._transactions_supported is False + fetched = await service.get_session( + app_name=APP, user_id=USER, session_id="s1" + ) + assert fetched.events[0].content.parts[0].text == "hi"