From b94537775749d00f6a50b2d0d430dd2b0da389c8 Mon Sep 17 00:00:00 2001 From: Ruslan Dautkhanov Date: Mon, 26 Oct 2026 12:25:17 -0600 Subject: [PATCH 1/2] [agentbricks] Add a LangGraph BaseStore backed by Agent Bricks managed memory memory_store() returns DatabricksMemoryStore, a LangGraph BaseStore over the managed memory API, so LangGraph and LangMem code that takes store= gets cross-conversation memory with no Postgres connection to run. The user comes from each run's config (thread_config), so one store serves every user of an app; entries are kept under the storage actor ":langgraph", and the session and memory helpers refuse actors with that suffix. Values are kept in a small JSON envelope, with optional client-side TTL and a per-user sweep. Includes unit tests over a wire-level fake of the memory API, a gated live suite, a README section, and a migration page from DatabricksStore. Co-authored-by: Isaac --- integrations/agentbricks/NOTICE | 4 + integrations/agentbricks/README.md | 86 +- .../docs/migrate-from-databricks-store.md | 159 ++++ .../src/databricks_agentkit/_actors.py | 26 + .../databricks_agentkit/langgraph/__init__.py | 5 + .../databricks_agentkit/langgraph/memory.py | 6 +- .../langgraph/session_store.py | 8 +- .../databricks_agentkit/langgraph/store.py | 782 ++++++++++++++++++ .../src/databricks_agentkit/openai/memory.py | 6 +- .../databricks_agentkit/openai/sessions.py | 7 +- .../tests/suites/test_langgraph_store.py | 166 ++++ .../tests/unit_tests/_memory_api_fake.py | 353 ++++++++ .../unit_tests/langgraph_store_edge_test.py | 315 +++++++ .../unit_tests/langgraph_store_graph_test.py | 341 ++++++++ .../langgraph_store_isolation_ttl_test.py | 325 ++++++++ .../tests/unit_tests/langgraph_store_test.py | 481 +++++++++++ 16 files changed, 3060 insertions(+), 10 deletions(-) create mode 100644 integrations/agentbricks/docs/migrate-from-databricks-store.md create mode 100644 integrations/agentbricks/src/databricks_agentkit/_actors.py create mode 100644 integrations/agentbricks/src/databricks_agentkit/langgraph/store.py create mode 100644 integrations/agentbricks/tests/suites/test_langgraph_store.py create mode 100644 integrations/agentbricks/tests/unit_tests/_memory_api_fake.py create mode 100644 integrations/agentbricks/tests/unit_tests/langgraph_store_edge_test.py create mode 100644 integrations/agentbricks/tests/unit_tests/langgraph_store_graph_test.py create mode 100644 integrations/agentbricks/tests/unit_tests/langgraph_store_isolation_ttl_test.py create mode 100644 integrations/agentbricks/tests/unit_tests/langgraph_store_test.py diff --git a/integrations/agentbricks/NOTICE b/integrations/agentbricks/NOTICE index 04b9a4020..fb15e4895 100644 --- a/integrations/agentbricks/NOTICE +++ b/integrations/agentbricks/NOTICE @@ -6,6 +6,10 @@ This Software is the property of Databricks, Inc. ("Databricks") and no license, This Software contains code from the following open source projects, licensed under the MIT license (https://opensource.org/licenses/MIT). +Package: langgraph +Project URL: https://github.com/langchain-ai/langgraph +Copyright: 2024 LangChain, Inc. + Package: PyYAML Project URL: https://github.com/yaml/pyyaml Copyright: 2017-2021 Ingy döt Net, 2006-2016 Kirill Simonov diff --git a/integrations/agentbricks/README.md b/integrations/agentbricks/README.md index d9f1778cc..aac58a1cc 100644 --- a/integrations/agentbricks/README.md +++ b/integrations/agentbricks/README.md @@ -585,8 +585,9 @@ an entry came from, for provenance). An entry is uniquely identified by its `act optional `session_id`. Write an entry when the agent learns something durable, then recall it in a later, separate -conversation with a natural-language search — results are ranked by full-text (BM25) relevance, up to -100 entries, with no pagination or vector similarity: +conversation with a natural-language search — results are ranked by relevance, full-text (BM25) +combined with semantic matching as the managed memory docs describe, up to 100 entries, with no +pagination: ```python from databricks.sdk import WorkspaceClient @@ -650,7 +651,9 @@ In an agent generated by `agentbricks init`, each invocation picks its session a partition: - **Session store:** one conversation's state, keyed by the invocation's top-level `session_id`. -- **Memory store:** facts recalled across conversations, keyed by `input.actor`. +- **Memory store:** facts recalled across conversations, keyed by `input.actor`. The + [LangGraph memory store](#langgraph-memory-store-as-a-basestore) keeps its entries under + `:langgraph` instead, so `memory_tools` doesn't list them. If a request leaves out `actor`, it defaults to the `session_id`, so long-term memory won't carry across conversations. For memory that persists, pass a stable `actor`, such as the end user's ID, from @@ -695,6 +698,83 @@ removes the corresponding `AGENT_MEMORY_STORE` or `AGENT_SESSION_STORE` entry fr Unbinding does not delete the store or revoke an existing grant to the App's service principal. If access must be revoked, coordinate that separately with the store owner or workspace admin. +### LangGraph memory store as a BaseStore + +`memory_store()` returns a LangGraph `BaseStore` backed by the managed memory store, so a graph keeps facts across conversations without a Postgres connection, and LangGraph and LangMem code that takes `store=` uses it as is. Build one store per app at startup, with no `actor=`, pass it to the graph with `store=`, and name each request's user through `thread_config`: + +```python +from databricks_agentkit.langgraph import checkpointer, memory_store, thread_config +from langchain.agents import create_agent + +store = memory_store() # once, at startup: no actor= here + +agent = create_agent(model, tools=tools, checkpointer=checkpointer(), store=store) + +# Per request. user_id is the signed-in user, from your app's trusted identity. +config = thread_config(session_id, actor=user_id) +result = await agent.ainvoke({"messages": inputs}, config=config) +``` + +Inside a node or tool, `get_store()` from `langgraph.config` returns the same store, for example `await get_store().aput(("preferences",), "style", {"value": "terse"})`. Values are JSON objects (dicts). + +**Where entries live.** Each item is one cross-session memory entry at `/langgraph//`, holding the value in a small JSON envelope (`{"lg": 1, "v": }`, plus `"x"`, the expiry in epoch milliseconds, for an item with a TTL). Entries are kept under the storage actor `:langgraph`, not the plain actor, so `memory_tools` and sessions, which use the plain actor, don't show them. Memory pipelines find the actors they process through sessions, and no session uses that id, so automatic pipeline runs leave these entries alone. That's a convention rather than a service guarantee: `thread_config`, the checkpointer, `memory_tools` and the OpenAI session and memory helpers refuse actors ending with `:langgraph`, but sessions created another way aren't checked, and a pipeline run that names its actors explicitly can still reach the entries. For the strongest separation, give the store a memory store that no pipeline writes to. Deleting a user's data includes deleting `:langgraph`. `DatabricksMemoryStore.storage_actor(user_id)` returns the id. To look at a user's entries: + +```bash +agentbricks -o json memory entries list --store --actor-id ":langgraph" --path-prefix /langgraph/ +``` + +**Who the actor is.** Take it from the identity your app receives from the trusted ingress (the bundled UI reads the `x-forwarded-email` or `x-forwarded-user` header), never from a request body field or a model-chosen value. The bundled `agent-langgraph` template still takes `input.actor` from the request body, so don't add `memory_store()` to it until its actor comes from the forwarded identity. A resumed run, such as one continuing after an approval, needs the same actor as the original run. `memory_store(actor=...)` is for calls outside a LangGraph run, such as scripts and tests; inside a run, a run for a different actor raises instead of mixing users. `thread_config(session_id)` without an actor uses the session id, so memories last for one conversation, and the store logs a warning. + +**No store configured.** The store resolves `store=`, then `AGENT_MEMORY_STORE`, on first use. Deployed Apps get the variable from the `[memory_store]` binding in agent.toml. `agentbricks dev` doesn't set it, so export it in your shell to use a dev store locally. With neither, a local run uses an in-process stand-in partitioned by actor, which keeps nothing across restarts, and a deployed App raises. + +**LangMem.** Use the namespace `("memories",)`, since the actor already partitions entries: + +```python +from langmem import create_manage_memory_tool + +manage = create_manage_memory_tool(namespace=("memories",), store=store) +await manage.ainvoke({"action": "create", "content": "Prefers email over phone."}, config=thread_config(session_id, actor=user_id)) +``` + +LangMem's default namespace, `("memories", "{langgraph_user_id}")`, needs `langgraph_user_id` in the run config. Without it LangMem raises `ConfigurationError`, a `BaseException` that `except Exception` misses; in a background `ReflectionExecutor` it kills the worker thread, so later submissions never complete. If you keep that namespace, set `langgraph_user_id` to the same user and build the store with `memory_store(strict_user_key="langgraph_user_id")`, which refuses a run where the two differ. Keep LangMem's memory manager and its manage-memory tool on separate namespaces, because the manager can't read the tool's entries. + +**Python 3.10.** LangGraph's `get_config()` and `get_store()` don't work in async code on 3.10. There, build the store per request with `memory_store(actor=user_id)`, compile the graph with it, and use that store object directly instead of `get_store()`. + +Limits: + +- **Search** ranks with the managed memory service, which the Agent Bricks docs describe as full-text (BM25) combined with semantic matching over each entry's content and description. The response doesn't say which matching ran (see the check below). Scores are the service's rank-based values: higher is better, but they aren't similarities and there's no cutoff, so a query can return entries that share no words with it. Client-side embeddings and `index` configs aren't supported. +- **Results.** A query returns at most 100 entries, with no pagination; `limit` defaults to 10, and an `offset` of 100 or more returns nothing. `filter` runs after ranking: when the service returns its full 100, a filtered page can come back short, and the store logs a warning. Expired entries count toward the 100 until they're swept. +- **Browsing.** `asearch(namespace)` with no query reads that namespace's entries, 100 per request, until the page is full; `alist_namespaces()` reads every entry the user has. Keep both out of hot paths. +- **Writes.** `aput` is last-writer-wins and not atomic, and so are batches. A call that raised may still have been applied; retrying it is safe. An existing key costs three API calls, a new one one. A put can raise `ABORTED` when another writer changes the same key at once; retry it. The SDK retries transient errors for up to five minutes before a call fails, and cancelling an async call doesn't stop a write already sent. Deletes are hard deletes, and deleting an entry that is already gone succeeds. +- **Values.** Values must be JSON objects: other types raise `TypeError`, and NaN, infinity and values JSON can't encode are refused. Tuples come back as lists and non-string keys as strings. Each put also stores the value's text, up to 500 characters, as the entry's description, which listings show and search matches. Requests are capped at about 900 KB and the actor plus path at about 2.4 KB. An entry changed with `agentbricks memory entries update --content` no longer holds the envelope, so the store skips it. + +**TTL.** `aput(..., ttl=minutes)`, or `memory_store(ttl={"default_ttl": minutes})` for puts that don't pass `ttl=`, sets an expiry. It differs from LangGraph's Postgres store: + +- A TTL counts from the last write, on the app's clock (keep hosts NTP-synced). Reads never refresh it: the store sets `refresh_on_read` to False, `refresh_on_read=True` raises, and per-call `refresh_ttl` arguments are ignored. +- Expired items aren't returned (`omit_expired` defaults to True), but they stay in the service until they're overwritten or swept, and the CLI still lists them under `:langgraph`. +- There's no background sweeper: `sweep_interval_minutes` raises, because the memory API lists one actor at a time. Call `store.sweep_ttl(actor=user_id)` (or `asweep_ttl`) to delete one user's expired items, for example after a request or from a scheduled job. A put that lands between a sweep's check and its delete can be lost. +- A put without a TTL clears an earlier one. Zero, negative, non-finite and unknown TTL settings raise. + +**Checking whether a workspace has semantic matching.** The response doesn't say, so test it with a throwaway store: + +1. Write two entries under one test actor: "Prefers to be contacted by email rather than phone." and "Allergic to peanuts." +2. Search "electronic mail correspondence", which shares no word with either entry, retrying for up to two minutes while the service indexes the entries. If the email entry comes back first, semantic matching is on. If nothing comes back, matching is full-text only. +3. Delete the test store. + +The gated live test `tests/suites/test_langgraph_store.py::test_search_mode` runs the same check. + +**The actor partitions data; it is not access control. The store is the boundary.** Calls run with the app's credentials, so any principal with store access can read every actor's entries, and any workspace user can list store names. Use a separate store per tenant or trust domain. `agentbricks deploy` grants the app's service principal access to a store declared in agent.toml. A failed grant is reported and does not fail the deploy. For a store named only in code, grant access with the app's `service_principal_client_id`: + +```python +from databricks_agentkit import workspace_client +from databricks_agentkit.client import AgentKitClient + +client = AgentKitClient(workspace_client()) +client.memory_stores.get("").grant_permission("") +``` + +**Coming from `DatabricksStore`.** `databricks_langchain.DatabricksStore` keeps memories in a Lakebase Postgres table you own and ranks with an embedding endpoint you configure. This store needs neither. [Migrate from DatabricksStore](docs/migrate-from-databricks-store.md) covers the code and data changes. + ## Model and prompt upgrades (experimental) `agentbricks experimental models` and `agentbricks experimental prompts` find the best model and diff --git a/integrations/agentbricks/docs/migrate-from-databricks-store.md b/integrations/agentbricks/docs/migrate-from-databricks-store.md new file mode 100644 index 000000000..43c3dfc0b --- /dev/null +++ b/integrations/agentbricks/docs/migrate-from-databricks-store.md @@ -0,0 +1,159 @@ +# Migrate from DatabricksStore to managed memory + +This guide covers moving an agent from the `DatabricksStore` in `databricks-langchain`, which keeps memories in a Lakebase table you own, to the managed memory store behind `memory_store()`. There is no data export or import tool. Existing memories are not copied automatically; the [data section](#data-migration) outlines a copy you script yourself. + +## Why migrate + +- **Simpler deployment**: No Postgres connection and no Lakebase project to manage. The managed memory store is reached through the AgentKit client, the same way sessions are. +- **Provisioned by deploy**: A `[memory_store]` binding in agent.toml creates the store on `agentbricks deploy`, and deploy grants the app access. + +## Before you start + +- Review the limits in the [LangGraph memory store section](../README.md#langgraph-memory-store-as-a-basestore) of the README: the 100-entry search window, last-writer-wins writes, and TTL that counts from the last write and is never refreshed by reads. +- **Recall may change.** `DatabricksStore` ranks by your embedding endpoint only when you configured one (`embedding_endpoint` or `embeddings`). Without one, its search ignores the query and returns rows newest first, with no score. The managed store ranks with the service, which the managed memory docs describe as full-text (BM25) combined with semantic matching over content and description; the README has a short check for whether semantic matching is on in your workspace. Measure recall on your own data before cutover. The managed store has no client-side embeddings and no custom index fields; if you need to control the embedding model, stay on `DatabricksStore`. +- Existing memories are not migrated. Decide before cutover whether to script a copy or accept the loss. + +## Migration steps + +### 1. Remove the Postgres connection + +Remove the code that builds the `DatabricksStore` and the Postgres or Lakebase connection it uses, such as the `psycopg` connection, SQLAlchemy engine, or `LakebasePool`. Then create the store: + +```python +from databricks_agentkit.langgraph import memory_store + +store = memory_store() # store id: store= argument, then AGENT_MEMORY_STORE, on first use +``` + +### 2. Update imports + +Replace the `databricks-langchain` store and checkpointer imports (`DatabricksStore`, `AsyncDatabricksStore`, `CheckpointSaver`, `AsyncCheckpointSaver`) with the AgentKit ones: + +```python +from databricks_agentkit.langgraph import checkpointer, memory_store, thread_config +``` + +### 3. Wire the store into the graph + +Build the store once, at startup, without `actor=`. Pass it to the agent with `store=`, the same way you pass the checkpointer. Putting the store in graph state does not work. Pass the user on every request, through `thread_config`: + +```python +from databricks_agentkit.langgraph import checkpointer, memory_store, thread_config +from langchain.agents import create_agent + +store = memory_store() # once, at startup: no actor= here + +agent = create_agent(model, tools=tools, checkpointer=checkpointer(), store=store) + +config = thread_config(session_id, actor=user_id) # per request +result = await agent.ainvoke({"messages": inputs}, config=config) +``` + +Inside a node, call `get_store()` from `langgraph.config`. Compile the store on the top-level graph; subgraphs inherit it from the run, but a store set only on a subgraph is not seen by its nodes. + +Do not build a store with `actor=` and reuse it for every user. Such a store is pinned to that actor, and a run for a different user raises. + +### 4. Update agent.toml + +Declare the memory store binding so deploy creates it and grants the app access: + +```toml +[memory_store] +name = "my-agent-memory" +``` + +If the store already exists, bind it instead: + +```sh +agentbricks memory bind existing-store-name +``` + +### 5. Remove custom Postgres and Lakebase setup + +Delete any deploy script or setup code that creates a Lakebase project, creates schemas, or manages Postgres connections for memory. The managed memory store handles provisioning. Once the binding from step 4 is in place, deploy is: + +```sh +agentbricks deploy my-agent +``` + +### 6. Take the actor from the trusted identity + +The actor is the signed-in user. Read it from the identity the Apps ingress forwards, never from a request body field. The bundled UI reads `x-forwarded-email`, then `x-forwarded-user`: + +```python +from collections.abc import Mapping + +FORWARDED_USER_HEADERS = ("x-forwarded-email", "x-forwarded-user") + + +def actor_from_headers(headers: Mapping[str, str]) -> str: + """The signed-in user, from headers set by the Apps ingress. Never the request body.""" + for name in FORWARDED_USER_HEADERS: + if value := headers.get(name): + return value + raise PermissionError("no forwarded user identity on the request") + + +user_id = actor_from_headers({"x-forwarded-email": "user@example.com"}) +``` + +The bundled agent-langgraph template still reads `actor` from the request payload (`runtime/adapter.py`). Replace that with the forwarded identity before the app serves real users. A resumed run, such as one continuing after an approval, needs the same actor as the original run. + +### 7. Test locally + +`agentbricks dev` removes the memory binding from the manifest it runs, but it passes your shell environment through, so an `AGENT_MEMORY_STORE` set in the shell reaches the app: + +```sh +AGENT_MEMORY_STORE= agentbricks dev +``` + +Treat this as a local override: the dev help still says memory is off. The variable takes the store id, not the display name. A local run with it set reads and writes that live store with your own credentials, so use a dev store. Without a store id, the store is an in-process stand-in that keeps nothing across restarts. + +### 8. Deploy and check the grant + +```sh +agentbricks deploy my-agent +``` + +Deploy creates the declared store and grants the app's service principal access. The grant is best-effort: if it fails, deploy still succeeds and prints a next step. Check the deploy output before the first memory call. If the grant failed, run the grant shown in the [README section](../README.md#langgraph-memory-store-as-a-basestore) before use. A deployed App with no store bound raises on its first memory call rather than falling back to the in-process stand-in. + +## Data migration + +Memories in the `DatabricksStore` table and in the managed store are not compatible, and nothing copies them for you. There is no export or import command. `agentbricks memory entries create` writes one entry at a time, so it is not a bulk copy route. Its `--content` is raw text; for the store to read an entry, the content must be the store's envelope JSON, the path must be under `/langgraph/`, and the actor must be the storage actor (`:langgraph`). + +If old memories matter, script a one-time copy: + +1. **Export read-only.** Connect to the legacy Lakebase with a read-only role. Read `prefix`, `key`, `value`, `created_at`, `updated_at`, `expires_at` and `ttl_minutes` from the table `store`, in the schema the legacy store was built with (`schema=`, default `public`). These names match langgraph-checkpoint-postgres 3.2.0; older 2.x releases may lack the TTL columns, so check them against your version first. Write the export to a private directory, outside the repo. +2. **Map rows to actors.** The namespace is the prefix split on dots, unchanged, so existing `aget` and `asearch` calls still hit the same tuples. The legacy table has no actor column, so use a mapping file that says which label is the user id. Rows with no matching rule go to a report, not the store. +3. **Decide per row.** Skip expired rows. For a row with a TTL, pass `ttl=` with its remaining minutes if they are above zero, or skip it. Skip values that aren't JSON objects, since the store holds JSON objects only. +4. **Write through the public API.** Write each row with a store pinned to its mapped actor, and pass the store id explicitly. A script may pin an actor outside a run. The store writes under the storage actor `:langgraph`, so a CLI check must use that id too, for example `agentbricks memory entries list --store --actor-id :langgraph --path-prefix /langgraph/`: + +```python +from databricks_agentkit.langgraph import memory_store + +STORE_ID = "" # required: without it memory_store() falls back to an in-process stand-in + + +async def import_row(actor: str, namespace: tuple[str, ...], key: str, value: dict) -> bool: + if not isinstance(value, dict): + return False # the store holds JSON objects only + store = memory_store(STORE_ID, actor=actor) # scripts pin one actor per write + if await store.aget(namespace, key) is not None: + return False # the app or an earlier run already wrote this key + await store.aput(namespace, key, value) + item = await store.aget(namespace, key) + return item is not None and item.value == value + + +await import_row("user-123", ("preferences",), "communication", {"style": "terse"}) +``` + +Pass the store id explicitly. Without one, `memory_store` returns an in-process stand-in that keeps nothing, and a read-back check against it still passes. + +5. **Make re-runs safe.** Skip a row whose key already exists in the target, as `import_row` does; that covers app writes made after cutover and rows from an earlier run. Freeze legacy writes during the window. Timestamps reset to import time, since the managed store keeps its own. +6. **Verify by reading back.** Count the source rows against written, skipped (by reason), and unmapped. Read each written row with `aget` and compare it with `item.value == value`. Do not verify with search, because search results depend on ranking. +7. **Keep a manifest.** Record a hash, path, actor, and status per row, with no values. Rollback can then delete only what the script created, and only where the content still matches. + +## Rollback + +To roll back to `DatabricksStore`, redeploy the old code branch and re-establish its Postgres or Lakebase connection. Entries written to the managed store after cutover stay there, since nothing copies them back to the legacy table. The copy in the data section is one-way. diff --git a/integrations/agentbricks/src/databricks_agentkit/_actors.py b/integrations/agentbricks/src/databricks_agentkit/_actors.py new file mode 100644 index 000000000..22b26440d --- /dev/null +++ b/integrations/agentbricks/src/databricks_agentkit/_actors.py @@ -0,0 +1,26 @@ +"""The actor rule shared by the LangGraph memory store and the session and memory helpers.""" + +from __future__ import annotations + +from typing import TypeVar + +# The LangGraph memory store keeps each user's entries under ``actor + STORE_ACTOR_SUFFIX``. Memory +# pipelines find the actors to process through their sessions, and no session should use that id, +# so automatic pipeline runs leave the store's entries alone. +STORE_ACTOR_SUFFIX = ":langgraph" + +_A = TypeVar("_A") + + +def check_actor(actor: _A) -> _A: + """Return ``actor``, or raise if it ends with the suffix the LangGraph memory store reserves. + + A session or memory-tools actor with that ending would reach another user's store entries, so + the session and memory helpers refuse it. Values other than strings pass through unchanged. + """ + if isinstance(actor, str) and actor.casefold().endswith(STORE_ACTOR_SUFFIX): + raise ValueError( + f"Actor ids can't end with {STORE_ACTOR_SUFFIX!r}: the LangGraph memory store keeps " + f"each user's entries under '{STORE_ACTOR_SUFFIX}'." + ) + return actor diff --git a/integrations/agentbricks/src/databricks_agentkit/langgraph/__init__.py b/integrations/agentbricks/src/databricks_agentkit/langgraph/__init__.py index e598958f1..c07938afd 100644 --- a/integrations/agentbricks/src/databricks_agentkit/langgraph/__init__.py +++ b/integrations/agentbricks/src/databricks_agentkit/langgraph/__init__.py @@ -37,6 +37,7 @@ from databricks_agentkit.langgraph.mcp import mcp_tools from databricks_agentkit.langgraph.memory import memory_tools from databricks_agentkit.langgraph.session_store import checkpointer, thread_config + from databricks_agentkit.langgraph.store import memory_store from databricks_agentkit.runtime import ( start_trace, workspace_client, @@ -63,6 +64,9 @@ def configure_tracing() -> None: "mcp_tools", # Long-term memory tools (opt-in via AGENT_MEMORY_STORE) — add to your tool list. "memory_tools", + # LangGraph BaseStore backed by managed memory (store= or AGENT_MEMORY_STORE; an in-process + # stand-in for local runs without one). + "memory_store", # Session persistence — pass checkpointer() to create_agent(checkpointer=...) and # thread_config(session_id) as the per-request run config. "checkpointer", @@ -84,6 +88,7 @@ def configure_tracing() -> None: "genie_tools": "databricks_agentkit.langgraph.genie", "mcp_tools": "databricks_agentkit.langgraph.mcp", "memory_tools": "databricks_agentkit.langgraph.memory", + "memory_store": "databricks_agentkit.langgraph.store", "checkpointer": "databricks_agentkit.langgraph.session_store", "thread_config": "databricks_agentkit.langgraph.session_store", "read_history": "databricks_agentkit.langgraph.history", diff --git a/integrations/agentbricks/src/databricks_agentkit/langgraph/memory.py b/integrations/agentbricks/src/databricks_agentkit/langgraph/memory.py index 880e759fe..114e83510 100644 --- a/integrations/agentbricks/src/databricks_agentkit/langgraph/memory.py +++ b/integrations/agentbricks/src/databricks_agentkit/langgraph/memory.py @@ -30,6 +30,7 @@ from langchain_core.tools import BaseTool, tool +from databricks_agentkit._actors import check_actor from databricks_agentkit.client import AgentKitClient from databricks_agentkit.runtime.workspace import workspace_client @@ -111,10 +112,13 @@ def memory_tools(actor: str, store: str | None = None) -> list[BaseTool]: The store resolves ``store`` arg → ``AGENT_MEMORY_STORE`` env → the ``[memory_store]`` binding in agent.toml (`agentbricks memory bind`) → none (no tools). ``actor`` partitions the store; it is captured in the tools' closures (not exposed to the model). Call this per request with the identity whose - memory to use (e.g. the signed-in user). + memory to use (e.g. the signed-in user). An actor ending with ``:langgraph`` is refused: the + LangGraph memory store keeps each user's entries under that id. """ from databricks_agentkit.runtime.tool_manifest import resolve_memory_store + check_actor(actor) + store_id = resolve_memory_store(store) if not store_id: return [] diff --git a/integrations/agentbricks/src/databricks_agentkit/langgraph/session_store.py b/integrations/agentbricks/src/databricks_agentkit/langgraph/session_store.py index 69be14bd3..74fccbdb2 100644 --- a/integrations/agentbricks/src/databricks_agentkit/langgraph/session_store.py +++ b/integrations/agentbricks/src/databricks_agentkit/langgraph/session_store.py @@ -40,6 +40,7 @@ ) from langgraph.checkpoint.memory import InMemorySaver +from databricks_agentkit._actors import check_actor from databricks_agentkit.runtime.session_store_client import Session, SessionStoreClient # Discriminators stored inside each session item's `data` JSON. @@ -96,9 +97,10 @@ def thread_config(session_id: str, actor: str | None = None) -> RunnableConfig: Includes ``actor_id`` because the durable saver maps it onto the Session's actor; it's ignored by the in-memory default. ``actor`` partitions the durable store — the caller supplies it (typically the signed-in user), so each user's threads stay separate; it defaults to ``session_id`` (one - actor per conversation). + actor per conversation). An actor ending with ``:langgraph`` is refused: the LangGraph memory + store keeps each user's entries under that id. """ - return {"configurable": {"thread_id": session_id, "actor_id": actor or session_id}} + return {"configurable": {"thread_id": session_id, "actor_id": check_actor(actor or session_id)}} async def history_checkpoint(config: RunnableConfig) -> CheckpointTuple | None: @@ -487,7 +489,7 @@ def _parse_config(config: RunnableConfig | dict[str, Any] | None) -> tuple[str, raise ValueError("config must contain configurable.thread_id") if not actor_id: raise ValueError("config must contain configurable.actor_id") - return thread_id, actor_id, configurable.get("checkpoint_ns", "") or "" + return thread_id, check_actor(actor_id), configurable.get("checkpoint_ns", "") or "" def _session_id(thread_id: str, checkpoint_ns: str) -> str: diff --git a/integrations/agentbricks/src/databricks_agentkit/langgraph/store.py b/integrations/agentbricks/src/databricks_agentkit/langgraph/store.py new file mode 100644 index 000000000..803495ebc --- /dev/null +++ b/integrations/agentbricks/src/databricks_agentkit/langgraph/store.py @@ -0,0 +1,782 @@ +"""LangGraph long-term memory store on Databricks managed memory. + +``DatabricksMemoryStore`` is a LangGraph ``BaseStore`` over the Agent Bricks managed memory API, +reached through the AgentKit client. Pass it as ``store=`` to ``create_agent`` or +``StateGraph.compile`` and any LangGraph or LangMem code that uses the store gets memory that +outlives one conversation, with no Postgres connection, credential rotation or pool tuning:: + + from databricks_agentkit.langgraph import checkpointer, memory_store, thread_config + + agent = create_agent( + model, tools=[...], checkpointer=checkpointer(), store=memory_store() + ) + agent.invoke(inputs, config=thread_config(session_id, actor)) + +Who the user is. Every entry belongs to an actor. Inside a LangGraph run the store reads it per +call from the run config (``configurable["actor_id"]``, which ``thread_config(session_id, actor)`` +sets and the checkpointer already reads), so one store serves many users. ``actor=`` is for calls +outside a run; inside a run it must match the run's actor, or the call raises. The actor +partitions memories but is not an access-control boundary: the memory store is. Set the actor from +a trusted identity (the signed-in user), never from a model-chosen or user-supplied field, and use +separate stores for separate tenants. + +Each item is one cross-session memory (no ``session_id``) at ``{root}/{namespace...}/{key}``, +each segment percent-encoded. Entries are kept under the storage actor ``:langgraph``, not +the plain actor. Memory pipelines find the actors they process through sessions, and no session +uses that id, so automatic pipeline runs leave these entries alone; that's a convention the +session and memory helpers enforce, not the service. ``memory_tools`` and sessions, which use the +plain actor, don't see the entries. The content is a small JSON envelope, ``{"lg": 1, "v": +}``, plus ``"x"`` (expiry, epoch milliseconds) when the item has a TTL. Entries outside +``root``, saved under a session, or not holding an envelope this store wrote are skipped, with a +one-time warning for the last kind. + +Limits: + - Search ranking is the service's. The managed memory docs describe it as full-text (BM25) + combined with semantic matching over each entry's content and description; the response + doesn't say which matching ran. Scores are the service's rank-based values: higher is better, + but they aren't similarities and there's no cutoff, so a query can return entries that share + no words with it. There are no client-side embeddings, so ``index`` configs and per-item + ``index`` field lists are rejected. + - One search returns at most 100 entries, with no pagination; a filter runs after ranking, so + entries ranked below the first 100 can't be reached by search, and ``offset + limit`` past + what the service returns comes back short. Expired entries count toward the 100 until they're + swept. + - ``put`` isn't atomic: concurrent writers to one key overwrite each other (last writer wins). + A batch isn't a transaction, and a call that raised may still have been applied. + - ``search`` without a query reads the namespace's entries, 100 per request, until the page is + full; ``list_namespaces`` reads every entry the actor has. Their cost grows with the entries + stored. + - Values are JSON objects (dicts). Tuples come back as lists and non-string keys as strings. + - TTL (``ttl=`` on ``put``, or ``TTLConfig.default_ttl``) counts from the last write, on this + process's clock. Reads never refresh it. Expired items aren't returned, but they stay in the + service until ``sweep_ttl`` deletes them or they are overwritten; there's no background + sweeper. + - Python 3.10: LangGraph's ``get_config()`` and ``get_store()`` don't work in async code, so + build the store per request with ``actor=`` and use that object directly there. +""" + +from __future__ import annotations + +import asyncio +import builtins +import json +import logging +import math +import os +import re +import threading +import time +import urllib.parse +from dataclasses import dataclass, field +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Any, Iterable, Iterator, Mapping, Optional, Protocol, cast + +from langgraph.config import get_config +from langgraph.store.base import ( + BaseStore, + GetOp, + Item, + ListNamespacesOp, + MatchCondition, + Op, + PutOp, + Result, + SearchItem, + SearchOp, + TTLConfig, +) + +from databricks_agentbricks.errors import AgentCliError +from databricks_agentkit._actors import STORE_ACTOR_SUFFIX, check_actor +from databricks_agentkit.client import AgentKitClient +from databricks_agentkit.runtime.store import RUNTIME_STORE_LOCAL_ENV +from databricks_agentkit.runtime.tool_manifest import resolve_memory_store +from databricks_agentkit.runtime.workspace import workspace_client + +if TYPE_CHECKING: + from databricks_agentkit.memory_store import Memory + +logger = logging.getLogger(__name__) + +_SEARCH_MAX = 100 # the service returns at most this many entries per search, with no pagination +_SEARCH_HEADROOM = 50 # extra results requested to cover entries dropped by client-side filters +_LIST_PAGE_SIZE = 100 +# The service rejects request bodies over 1 MiB, and only after a long wait. The SDK escapes +# non-ASCII text and quotes in the JSON it sends, so measure the encoded request, not the value. +_REQUEST_MAX_BYTES = 900_000 +# The service indexes actor and path together and rejects pairs over about 2.7 KB. +_KEY_MAX_BYTES = 2_400 +_TTL_KEYS = frozenset({"default_ttl", "refresh_on_read", "sweep_interval_minutes", "omit_expired"}) +_TTL_MAX_MINUTES = 100 * 365 * 24 * 60 # a century + + +def _encode_segment(segment: str, what: str) -> str: + if not isinstance(segment, str) or not segment: + raise ValueError(f"{what} must be a non-empty string, got {segment!r}") + try: + encoded = urllib.parse.quote(segment, safe="") + except UnicodeEncodeError as e: # eg a lone surrogate + raise ValueError(f"{what} {segment!r} isn't valid Unicode text: {e.reason}") from e + # quote() leaves dots alone; "." and ".." would read as path navigation, so encode them. + return encoded.replace(".", "%2E") if encoded in (".", "..") else encoded + + +def _prefix_path(root: str, namespace: tuple[str, ...]) -> str: + """The listing prefix for ``namespace``: ends in "/" so ("a",) never matches ("ab",).""" + segments = [_encode_segment(label, "namespace label") for label in namespace] + return "/".join([root, *segments, ""]) + + +def _item_path(root: str, namespace: tuple[str, ...], key: str) -> str: + return _prefix_path(root, namespace) + _encode_segment(key, "key") + + +def _split_path(root: str, path: str) -> Optional[tuple[tuple[str, ...], str]]: + """``(namespace, key)`` for a path this store wrote, else None.""" + if not path.startswith(root + "/"): + return None + segments = path[len(root) + 1 :].split("/") + if any(not segment for segment in segments): + return None + *namespace, key = (urllib.parse.unquote(segment) for segment in segments) + return tuple(namespace), key + + +_ENVELOPE_VERSION = 1 +_ENVELOPE_KEYS = frozenset({"lg", "v", "x"}) +_FOREIGN = ( + "Skipped an entry under the store's root that doesn't hold a value this store wrote; a memory " + "pipeline, the CLI or another writer may have changed it." +) +_NEWER = "Skipped an entry written by a newer version of DatabricksMemoryStore; upgrade to read it." + + +@dataclass(frozen=True) +class _Stored: + value: dict[str, Any] + expires_at_ms: Optional[int] + + +def _wrap(value: Mapping[str, Any], expires_at_ms: Optional[int]) -> str: + envelope: dict[str, Any] = {"lg": _ENVELOPE_VERSION, "v": value} + if expires_at_ms is not None: + envelope["x"] = expires_at_ms + # ensure_ascii=False keeps non-Latin text searchable by the service's full-text matching; + # allow_nan=False keeps the content strict JSON for the CLI and other readers. + return json.dumps(envelope, ensure_ascii=False, allow_nan=False, separators=(",", ":")) + + +def _unwrap(memory: Memory) -> Optional[_Stored]: + """The value and expiry of an entry this store wrote, else None (with a one-time warning).""" + try: + envelope = json.loads(memory.content or "") + except (TypeError, ValueError, RecursionError): # RecursionError: deeply nested JSON + envelope = None + if isinstance(envelope, dict): + version, value, expires = envelope.get("lg"), envelope.get("v"), envelope.get("x") + if type(version) is int and version > _ENVELOPE_VERSION: + _warn_once(_NEWER) + return None + if ( + type(version) is int + and version == _ENVELOPE_VERSION + and set(envelope) <= _ENVELOPE_KEYS + and isinstance(value, dict) + and (expires is None or type(expires) is int) + ): + return _Stored(value, expires) + _warn_once(_FOREIGN) + return None + + +def _now_ms() -> int: + return time.time_ns() // 1_000_000 + + +def _compare_values(item_value: Any, filter_value: Any) -> bool: + """Filter matching with InMemoryStore semantics ($eq/$ne/$gt/$gte/$lt/$lte, nested dicts). + + ``_compare_values`` and ``_apply_operator`` are adapted from the functions of the same names in + ``langgraph.store.memory`` (MIT License, Copyright (c) 2024 LangChain, Inc.). + """ + if isinstance(filter_value, dict): + if any(k.startswith("$") for k in filter_value): + return all( + _apply_operator(item_value, op, operand) for op, operand in filter_value.items() + ) + if not isinstance(item_value, dict): + return False + return all(_compare_values(item_value.get(k), v) for k, v in filter_value.items()) + if isinstance(filter_value, (list, tuple)): + return ( + isinstance(item_value, (list, tuple)) + and len(item_value) == len(filter_value) + and all( + _compare_values(iv, fv) for iv, fv in zip(item_value, filter_value, strict=True) + ) + ) + return item_value == filter_value + + +def _apply_operator(value: Any, operator: str, operand: Any) -> bool: + if operator == "$eq": + return value == operand + if operator == "$ne": + return value != operand + if operator in ("$gt", "$gte", "$lt", "$lte"): + try: + left, right = float(value), float(operand) + except (TypeError, ValueError): + return False + return { + "$gt": left > right, + "$gte": left >= right, + "$lt": left < right, + "$lte": left <= right, + }[operator] + raise ValueError(f"Unsupported filter operator: {operator}") + + +def _matches(value: dict[str, Any], filter_: Optional[dict[str, Any]]) -> bool: + return not filter_ or all(_compare_values(value.get(k), v) for k, v in filter_.items()) + + +def _matches_condition(namespace: tuple[str, ...], condition: MatchCondition) -> bool: + pattern = tuple(condition.path) + if len(namespace) < len(pattern): + return False + window = ( + namespace[: len(pattern)] + if condition.match_type == "prefix" + else namespace[-len(pattern) :] + ) + return all(p == "*" or p == label for p, label in zip(pattern, window, strict=True)) + + +class _Memories(Protocol): + """What the store needs from a managed memory store handle (``MemoryStore``).""" + + def put( + self, *, actor_id: str, path: str, content: str, description: Optional[str] = None + ) -> Any: ... + + def list( + self, + *, + actor_id: str, + path: Optional[str] = None, + path_prefix: Optional[str] = None, + page_size: Optional[int] = None, + ) -> Iterable[Any]: ... + + def search( + self, + *, + actor_id: str, + query: str, + limit: Optional[int] = None, + path_prefix: Optional[str] = None, + ) -> Iterable[Any]: ... + + +@dataclass +class _LocalMemory: + path: str + content: str + create_time: datetime + update_time: datetime + _entries: dict[tuple[str, str], _LocalMemory] = field(repr=False) + _key: tuple[str, str] = field(repr=False) + session_id: Optional[str] = None + + def delete(self) -> None: + self._entries.pop(self._key, None) + + +@dataclass +class _LocalResult: + memory: _LocalMemory + score: float + + +class _LocalMemories: + """An in-process stand-in for a managed memory store, for local development only. + + Same actor partitioning and paths as the service; search ranks by shared words. + """ + + def __init__(self) -> None: + self._entries: dict[tuple[str, str], _LocalMemory] = {} + self._lock = threading.Lock() + + def put( + self, *, actor_id: str, path: str, content: str, description: Optional[str] = None + ) -> _LocalMemory: + now = datetime.now(timezone.utc) + key = (actor_id, path) + with self._lock: + old = self._entries.get(key) + memory = _LocalMemory( + path, content, old.create_time if old else now, now, self._entries, key + ) + self._entries[key] = memory + return memory + + def list( + self, + *, + actor_id: str, + path: Optional[str] = None, + path_prefix: Optional[str] = None, + page_size: Optional[int] = None, + ) -> builtins.list[_LocalMemory]: + with self._lock: + entries = [m for (a, _), m in self._entries.items() if a == actor_id] + return [ + m + for m in entries + if (path is None or m.path == path) + and (path_prefix is None or m.path.startswith(path_prefix)) + ] + + def search( + self, + *, + actor_id: str, + query: str, + limit: Optional[int] = None, + path_prefix: Optional[str] = None, + ) -> builtins.list[_LocalResult]: + terms = set(re.findall(r"\w+", query.lower())) + hits = [] + for memory in self.list(actor_id=actor_id, path_prefix=path_prefix): + score = len(terms & set(re.findall(r"\w+", memory.content.lower()))) + if score: + hits.append(_LocalResult(memory, float(score))) + hits.sort(key=lambda r: -r.score) + return hits[:limit] + + +class DatabricksMemoryStore(BaseStore): + """LangGraph ``BaseStore`` on an Agent Bricks managed memory store. + + Args: + store: Memory store id. Resolves ``store`` → ``AGENT_MEMORY_STORE`` on first use (deployed + Apps set it from the agent.toml binding; ``agentbricks dev`` doesn't, so export it in + your shell to use a store locally). With neither, the first call raises. + actor: The identity whose memories to use, for calls outside a LangGraph run. Inside a + run the actor comes from ``configurable[actor_key]`` (see ``thread_config``); a run + for a different actor than ``actor`` raises rather than mixing users. Where the run + config isn't visible (threads you start, async code on Python 3.10), ``actor`` is + used as is, so build such a store per request, for that request's user. + actor_key: The run-config key holding the actor. + strict_user_key: Optional run-config key that must be present and equal the actor, eg + ``"langgraph_user_id"`` for LangMem namespaces templated on the user id. + root: Path prefix for this store's entries, so they don't mix with other writers. + ttl: A LangGraph ``TTLConfig``. ``default_ttl`` (minutes) applies to puts that don't pass + ``ttl=``. ``omit_expired`` defaults to True here (LangGraph's default is False). + Reads never refresh a TTL, so ``refresh_on_read`` is set to False, and True raises: a + refresh would rewrite the entry and could undo a concurrent put. ``sweep_interval_minutes`` + raises, because the memory API lists one actor at a time, so there's no global sweep; + call :meth:`sweep_ttl` per user instead. Unknown keys raise. + client: An AgentKit client; built from the workspace client on first use when omitted. + """ + + # Plain BaseStore on purpose: LangGraph's batched base pools concurrent callers into one + # batch, where no caller's identity is visible. Each batch() here runs for one actor. + supports_ttl = True + + def __init__( + self, + store: Optional[str] = None, + *, + actor: Optional[str] = None, + actor_key: str = "actor_id", + strict_user_key: Optional[str] = None, + root: str = "/langgraph", + ttl: Optional[TTLConfig] = None, + client: Optional[AgentKitClient] = None, + ) -> None: + if not root.startswith("/") or root.endswith("/"): + raise ValueError(f"root must start with '/' and not end with one, got {root!r}") + if actor is not None: + _explicit_actor(actor) + ttl_opts: dict[str, Any] = dict(cast(Mapping[str, Any], ttl or {})) + unknown = set(ttl_opts) - _TTL_KEYS + if unknown: + raise ValueError(f"Unknown ttl keys {sorted(unknown)}; expected {sorted(_TTL_KEYS)}.") + if ttl_opts.get("sweep_interval_minutes"): + raise ValueError( + "sweep_interval_minutes isn't supported: the memory API lists one actor at a time, " + "so the store can't sweep every user. Call sweep_ttl(actor=...) instead." + ) + if ttl_opts.get("refresh_on_read"): + raise ValueError( + "refresh_on_read isn't supported: a refresh rewrites the entry and could undo a " + "concurrent put. A TTL counts from the last write." + ) + if ttl_opts.get("default_ttl") is not None: + _check_ttl(ttl_opts["default_ttl"]) + if ttl_opts: + ttl_opts["refresh_on_read"] = False # so LangGraph's own TTL defaults match the store + self.ttl_config = cast(TTLConfig, ttl_opts) if ttl_opts else None + self._omit_expired = ttl_opts.get("omit_expired", True) + self._now_ms = _now_ms + self._store = store + self._actor = actor + self._actor_key = actor_key + self._strict_user_key = strict_user_key + self._root = root + self._client = client + self._handle: Optional[_Memories] = None + self._local_fallback = False # memory_store() sets it: no store means a local stand-in + self._lock = threading.Lock() + + @classmethod + def _local(cls, **kwargs: Any) -> DatabricksMemoryStore: + store = cls("local", **kwargs) + store._handle = _LocalMemories() + return store + + def batch(self, ops: Iterable[Op]) -> list[Result]: + return self._run(list(ops), self._resolve_actor()) + + async def abatch(self, ops: Iterable[Op]) -> list[Result]: + # Read the run config here, on the event loop, before hopping to a worker thread. + actor = self._resolve_actor() + return await asyncio.to_thread(self._run, list(ops), actor) + + def _run(self, ops: list[Op], actor: str) -> list[Result]: + handle, partition = self._memory_store(), self.storage_actor(actor) + return [self._apply(handle, partition, op) for op in ops] + + @staticmethod + def storage_actor(actor: str) -> str: + """The actor id the service keeps ``actor``'s entries under, eg for the CLI.""" + return check_actor(actor) + STORE_ACTOR_SUFFIX + + def sweep_ttl(self, *, actor: Optional[str] = None) -> int: + """Delete one user's expired items and return how many were deleted. + + ``actor`` defaults to the run's actor, as for any other call. There's no sweep across + users, because the memory API lists one actor at a time. + """ + user = _explicit_actor(actor) if actor is not None else self._resolve_actor() + return self._sweep(self.storage_actor(user)) + + async def asweep_ttl(self, *, actor: Optional[str] = None) -> int: + user = _explicit_actor(actor) if actor is not None else self._resolve_actor() + return await asyncio.to_thread(self._sweep, self.storage_actor(user)) + + def _sweep(self, partition: str) -> int: + handle, now, deleted = self._memory_store(), self._now_ms(), 0 + for memory in list(self._under(handle, partition, self._root + "/")): + path = memory.path or "" + if _split_path(self._root, path) is None or not _expired(_unwrap(memory), now): + continue + # Re-read just before deleting, so an entry a put renewed in the meantime survives. + for current in self._at(handle, partition, path): + if _expired(_unwrap(current), now): + try: + current.delete() + except AgentCliError as e: + if e.error_code != "NOT_FOUND": + raise + else: + deleted += 1 + return deleted + + def _resolve_actor(self) -> str: + try: + configurable = get_config().get("configurable", {}) + except RuntimeError as e: + if self._actor is not None and self._strict_user_key is None: + return self._actor # outside a LangGraph run, or where its config isn't visible + raise ValueError( + "DatabricksMemoryStore needs a run config: call it inside a LangGraph run with " + "config=thread_config(session_id, actor), or pass actor= for calls outside one. " + "Threads you start yourself don't inherit the run config (copy it with " + "contextvars.copy_context()), and async code on Python 3.10 doesn't see it." + ) from e + run_actor = configurable.get(self._actor_key) + if self._actor is None: + actor = run_actor + # thread_config(session_id) without an actor sets the actor to the session id. + if run_actor is not None and run_actor == configurable.get("thread_id"): + _warn_once( + "The run config's actor is the session id, so memories won't carry over to " + "other conversations; pass the user to thread_config(session_id, actor)." + ) + else: + actor = self._actor + if run_actor is not None and run_actor != actor: + raise ValueError( + f"This DatabricksMemoryStore was built for actor {actor!r} but the run is for " + f"{run_actor!r} (thread_config(session_id) without an actor uses the session " + "id). Build one store without actor= and pass each user through " + "thread_config(session_id, actor)." + ) + if not actor: + raise ValueError( + f"DatabricksMemoryStore needs an actor: the run config has no " + f"{self._actor_key!r}. Use config=thread_config(session_id, actor)." + ) + if not isinstance(actor, str): + raise ValueError( + f"DatabricksMemoryStore needs a string actor, got {type(actor).__name__} " + f"from the run config's {self._actor_key!r}." + ) + if self._strict_user_key is not None: + user = configurable.get(self._strict_user_key) + if user != actor: + raise ValueError( + f"Run config {self._strict_user_key!r}={user!r} doesn't match actor {actor!r}; " + "refusing to read or write another user's memory." + ) + return actor + + def _memory_store(self) -> _Memories: + # Built on first use and reused: the client resolves workspace auth, and the id-only + # handle avoids fetching the store before every call. + with self._lock: + if self._handle is None: + store_id = resolve_memory_store(self._store) + if store_id: + client = self._client or AgentKitClient(workspace_client()) + self._handle = client.memory_stores._handle(store_id) + elif self._local_fallback and not _deployed(): + _warn_once( + "No memory store configured (pass store=, set AGENT_MEMORY_STORE, or run " + "`agentbricks memory bind`); using an in-process store that keeps nothing " + "across restarts." + ) + self._handle = _LocalMemories() + else: + raise ValueError( + "No memory store configured: pass store=, set AGENT_MEMORY_STORE, or bind " + "one with `agentbricks memory bind`. A deployed App needs a store; the " + "in-process stand-in is for local development only." + ) + return self._handle + + def _apply(self, handle: _Memories, partition: str, op: Op) -> Result: + if isinstance(op, GetOp): + return self._get(handle, partition, op) + if isinstance(op, PutOp): + return self._put(handle, partition, op) + if isinstance(op, SearchOp): + return self._search(handle, partition, op) + if isinstance(op, ListNamespacesOp): + return self._list_namespaces(handle, partition, op) + raise TypeError(f"Unsupported store operation: {type(op).__name__}") + + def _at(self, handle: _Memories, partition: str, path: str) -> list[Memory]: + # The API ignores session_id alongside path, so recheck both. + listed = handle.list(actor_id=partition, path=path, page_size=_LIST_PAGE_SIZE) + return [m for m in listed if m.path == path and m.session_id is None] + + def _under(self, handle: _Memories, partition: str, prefix: str) -> Iterator[Memory]: + listed = handle.list(actor_id=partition, path_prefix=prefix, page_size=_LIST_PAGE_SIZE) + for memory in listed: + if memory.session_id is None and (memory.path or "").startswith(prefix): + yield memory + + def _live(self, memory: Memory, now: int) -> Optional[_Stored]: + """The item an entry holds, unless it isn't this store's or has expired (and is omitted).""" + stored = _unwrap(memory) + if stored is None or (self._omit_expired and _expired(stored, now)): + return None + return stored + + def _get(self, handle: _Memories, partition: str, op: GetOp) -> Optional[Item]: + now = self._now_ms() + for memory in self._at(handle, partition, _item_path(self._root, op.namespace, op.key)): + stored = self._live(memory, now) + if stored is not None: + return Item( + value=stored.value, key=op.key, namespace=op.namespace, **_times(memory) + ) + return None + + def _put(self, handle: _Memories, partition: str, op: PutOp) -> None: + if isinstance(op.index, list): + raise NotImplementedError( + "DatabricksMemoryStore doesn't support per-item index fields: the managed memory " + "service ranks search results itself, with no client-side embeddings." + ) + path = _item_path(self._root, op.namespace, op.key) + if op.value is None: + for memory in self._at(handle, partition, path): + try: + memory.delete() + except AgentCliError as e: + if e.error_code != "NOT_FOUND": # already gone, eg a retried delete + raise + return + if not isinstance(op.value, dict): + raise TypeError( + f"Value for key {op.key!r} must be a dict (a JSON object), " + f"got {type(op.value).__name__}." + ) + key_bytes = len(partition.encode()) + len(path.encode()) + if key_bytes > _KEY_MAX_BYTES: + raise ValueError( + f"The actor, namespace and key {op.key!r} take {key_bytes:,} bytes together; " + f"managed memory accepts about {_KEY_MAX_BYTES:,}." + ) + expires = None if op.ttl is None else self._now_ms() + round(_check_ttl(op.ttl) * 60_000) + try: + content = _wrap(op.value, expires) + except TypeError as e: + raise TypeError(f"Value for key {op.key!r} isn't JSON-serializable: {e}") from e + except (ValueError, RecursionError) as e: # NaN or infinity, a huge int, deep nesting + raise ValueError(f"Value for key {op.key!r} can't be stored as JSON: {e}") from e + try: + content.encode() + except UnicodeEncodeError as e: # eg a lone surrogate in a string + raise ValueError( + f"Value for key {op.key!r} isn't valid Unicode text: {e.reason}" + ) from e + description = _describe(op.value) + size = len(json.dumps({"content": content, "description": description})) + if size > _REQUEST_MAX_BYTES: + raise ValueError( + f"Value for key {op.key!r} needs a {size:,}-byte request; managed memory entries " + f"must stay under about {_REQUEST_MAX_BYTES:,} bytes." + ) + handle.put(actor_id=partition, path=path, content=content, description=description) + + def _search(self, handle: _Memories, partition: str, op: SearchOp) -> list[SearchItem]: + prefix, now = _prefix_path(self._root, op.namespace_prefix), self._now_ms() + query, want = (op.query or "").strip(), op.offset + op.limit + found: list[Any] = [] + limit = 0 + candidates: Iterable[tuple[Memory, Optional[float]]] + if query: + # A filter runs after the service ranks, so ask for the most it returns. + limit = _SEARCH_MAX if op.filter else min(_SEARCH_MAX, want + _SEARCH_HEADROOM) + found = list( + handle.search(actor_id=partition, query=query, limit=limit, path_prefix=prefix) + ) + candidates = [(r.memory, r.score) for r in found] + else: + candidates = ((m, None) for m in self._under(handle, partition, prefix)) + items = [] + for memory, score in candidates: + path = memory.path or "" + if memory.session_id is not None or not path.startswith(prefix): + continue + parts = _split_path(self._root, path) + stored = self._live(memory, now) if parts is not None else None + if parts is None or stored is None or not _matches(stored.value, op.filter): + continue + namespace, key = parts + items.append(SearchItem(namespace, key, stored.value, score=score, **_times(memory))) + if not query and len(items) >= want: + break # listing order is fixed, so later pages can't change this one + if query and len(found) >= limit and len(items) < want: + _warn_once( + "A search reached the most results managed memory returns per query; matches " + "ranked lower aren't seen, so the page may be short. Narrow the namespace or query." + ) + return items[op.offset : want] + + def _list_namespaces( + self, handle: _Memories, partition: str, op: ListNamespacesOp + ) -> list[tuple[str, ...]]: + namespaces, now = set(), self._now_ms() + for memory in self._under(handle, partition, self._root + "/"): + parts = _split_path(self._root, memory.path or "") + if parts is None or self._live(memory, now) is None: # same entries get and search see + continue + namespace = parts[0] + if op.match_conditions and not all( + _matches_condition(namespace, c) for c in op.match_conditions + ): + continue + namespaces.add(namespace[: op.max_depth] if op.max_depth is not None else namespace) + return sorted(namespaces)[op.offset : op.offset + op.limit] + + +_DESCRIPTION_MAX = 500 + + +def _describe(value: Mapping[str, Any]) -> str: + """The value's text, for the CLI and UI listings and the service's search over descriptions.""" + texts: list[str] = [] + stack: list[Any] = [value] + while stack: # depth-first, in the value's own order, without recursion + node = stack.pop() + if isinstance(node, str): + texts.append(node) + elif isinstance(node, Mapping): + stack.extend(reversed(list(node.values()))) + elif isinstance(node, (list, tuple)): + stack.extend(reversed(node)) + # The service can't store NUL characters in a description. + return " ".join(" ".join(texts).replace("\x00", " ").split())[:_DESCRIPTION_MAX] + + +def _check_ttl(minutes: Any) -> float: + if ( + isinstance(minutes, bool) + or not isinstance(minutes, (int, float)) + or not math.isfinite(minutes) + or not 0 < minutes <= _TTL_MAX_MINUTES + ): + raise ValueError( + f"ttl must be a positive number of minutes, at most {_TTL_MAX_MINUTES:,}, " + f"got {minutes!r}" + ) + if round(minutes * 60_000) < 1: + raise ValueError(f"ttl must be at least one millisecond, got {minutes!r} minutes") + return float(minutes) + + +def _expired(stored: Optional[_Stored], now: int) -> bool: + return stored is not None and stored.expires_at_ms is not None and stored.expires_at_ms <= now + + +def _explicit_actor(actor: Any) -> str: + if not isinstance(actor, str) or not actor: + raise ValueError(f"actor must be a non-empty string, got {actor!r}") + return check_actor(actor) + + +def _deployed() -> bool: + """True in a deployed Databricks App, where an in-process stand-in would lose memories.""" + local = os.getenv(RUNTIME_STORE_LOCAL_ENV, "").lower() == "true" + return bool(os.getenv("DATABRICKS_APP_NAME")) and not local + + +_warned: set[str] = set() +_warned_lock = threading.Lock() + + +def _warn_once(message: str) -> None: + with _warned_lock: + if message in _warned: + return + _warned.add(message) + logger.warning(message) + + +def _times(memory: Any) -> dict[str, datetime]: + now = datetime.now(timezone.utc) + return {"created_at": memory.create_time or now, "updated_at": memory.update_time or now} + + +def memory_store( + store: Optional[str] = None, *, actor: Optional[str] = None, **kwargs: Any +) -> BaseStore: + """The LangGraph long-term memory store on managed memory. Build it once and reuse it. + + Resolves ``store`` → ``AGENT_MEMORY_STORE`` on first use (deployed Apps set it from the + agent.toml binding; ``agentbricks dev`` doesn't). With neither, a local run logs a warning + and uses an in-process stand-in with the same actor partitioning, which keeps nothing across + restarts; a deployed App raises instead. ``actor`` and ``kwargs`` go to + :class:`DatabricksMemoryStore`. + """ + instance = DatabricksMemoryStore(store, actor=actor, **kwargs) + instance._local_fallback = True + return instance diff --git a/integrations/agentbricks/src/databricks_agentkit/openai/memory.py b/integrations/agentbricks/src/databricks_agentkit/openai/memory.py index ba52d4cfc..479396255 100644 --- a/integrations/agentbricks/src/databricks_agentkit/openai/memory.py +++ b/integrations/agentbricks/src/databricks_agentkit/openai/memory.py @@ -30,6 +30,7 @@ from agents import FunctionTool, function_tool +from databricks_agentkit._actors import check_actor from databricks_agentkit.client import AgentKitClient from databricks_agentkit.runtime.workspace import workspace_client @@ -111,10 +112,13 @@ def memory_tools(actor: str, store: str | None = None) -> list[FunctionTool]: The store resolves ``store`` arg → ``AGENT_MEMORY_STORE`` env → the ``[memory_store]`` binding in agent.toml (`agentbricks memory bind`) → none (no tools). ``actor`` partitions the store; it is captured in the tools' closures (not exposed to the model). Call this per request with the identity whose - memory to use (e.g. the signed-in user). + memory to use (e.g. the signed-in user). An actor ending with ``:langgraph`` is refused: the + LangGraph memory store keeps each user's entries under that id. """ from databricks_agentkit.runtime.tool_manifest import resolve_memory_store + check_actor(actor) + store_id = resolve_memory_store(store) if not store_id: return [] diff --git a/integrations/agentbricks/src/databricks_agentkit/openai/sessions.py b/integrations/agentbricks/src/databricks_agentkit/openai/sessions.py index d9a064e46..922076a99 100644 --- a/integrations/agentbricks/src/databricks_agentkit/openai/sessions.py +++ b/integrations/agentbricks/src/databricks_agentkit/openai/sessions.py @@ -30,6 +30,7 @@ from agents import SQLiteSession, TResponseInputItem from agents.memory import SessionABC +from databricks_agentkit._actors import check_actor from databricks_agentkit.runtime.session_store_client import Session as _StoreSession from databricks_agentkit.runtime.session_store_client import SessionStoreClient @@ -50,13 +51,15 @@ def session_store( is configured. The store name resolves ``store`` arg → ``AGENT_SESSION_STORE`` env → the ``[session_store]`` binding in agent.toml (`agentbricks sessions bind`) → none. ``actor`` partitions the durable store — the caller supplies it (typically the signed-in user), so each user's transcripts - stay separate; it defaults to ``session_id`` and is ignored by the in-memory store. + stay separate; it defaults to ``session_id`` and is ignored by the in-memory store. An actor ending + with ``:langgraph`` is refused: the LangGraph memory store keeps each user's entries under that id. """ from databricks_agentkit.runtime.tool_manifest import resolve_session_store + actor_id = check_actor(actor or session_id) store = resolve_session_store(store) if store: - return DatabricksSessionStore(session_id, store, actor_id=actor or session_id) + return DatabricksSessionStore(session_id, store, actor_id=actor_id) cached = _local_sessions.get(session_id) if cached is None: cached = SQLiteSession(session_id) # ":memory:" — process-local, non-durable diff --git a/integrations/agentbricks/tests/suites/test_langgraph_store.py b/integrations/agentbricks/tests/suites/test_langgraph_store.py new file mode 100644 index 000000000..c5d58ae5d --- /dev/null +++ b/integrations/agentbricks/tests/suites/test_langgraph_store.py @@ -0,0 +1,166 @@ +"""Live check of ``DatabricksMemoryStore`` against a real managed memory store (gated). + +Each test creates its own memory store and registers its deletion before the create, then drives +the public LangGraph store methods. The store uses the suite's workspace client +(``workspace_client.client``), not ambient auth. + +Run with the gate set: + RUN_LANGGRAPH_STORE_TESTS=1 pytest tests/suites/test_langgraph_store.py -v +""" + +from __future__ import annotations + +import os +import uuid + +import pytest + +gated = pytest.mark.skipif( + os.getenv("RUN_LANGGRAPH_STORE_TESTS") != "1", + reason="Set RUN_LANGGRAPH_STORE_TESTS=1 to run DatabricksMemoryStore tests.", +) + + +def _new_store(agentbricks_cli, cleanup) -> str: + """Create a memory store and return its id. The create takes the display name as the id.""" + store_id = f"lg-store-{uuid.uuid4().hex[:10]}" + # Registered first: a create that fails partway can still leave the store behind. + cleanup(f"memory stores delete memory-stores/{store_id} --yes") + agentbricks_cli.run(f"memory stores create --display-name {store_id}") + return store_id + + +@gated +def test_langgraph_store_put_get_overwrite_delete(workspace_client, agentbricks_cli, cleanup): + """DatabricksMemoryStore: put, get, overwrite, search and delete on a live store.""" + from langgraph.store.base import PutOp + + from databricks_agentkit.client import AgentKitClient + from databricks_agentkit.langgraph.store import DatabricksMemoryStore + + store_id = _new_store(agentbricks_cli, cleanup) + store = DatabricksMemoryStore( + store=store_id, + actor=f"e2e-{uuid.uuid4().hex[:8]}", + client=AgentKitClient(workspace_client.client), + ) + + store.put(("notes",), "k1", {"text": "alpha"}) + item = store.get(("notes",), "k1") + assert item is not None and item.value == {"text": "alpha"} + + store.put(("notes",), "k1", {"text": "beta"}) # the second write takes the update path + item = store.get(("notes",), "k1") + assert item is not None and item.value == {"text": "beta"} + assert [hit.key for hit in store.search(("notes",))] == ["k1"] + + store.batch([PutOp(namespace=("notes",), key="k1", value=None)]) # a None value deletes + assert store.get(("notes",), "k1") is None + + +@gated +def test_langgraph_store_isolates_actors(workspace_client, agentbricks_cli, cleanup): + """DatabricksMemoryStore: two actors on one live store see only their own entries.""" + from databricks_agentkit.client import AgentKitClient + from databricks_agentkit.langgraph.store import DatabricksMemoryStore + + store_id = _new_store(agentbricks_cli, cleanup) + client = AgentKitClient(workspace_client.client) + suffix = uuid.uuid4().hex[:8] + alice = DatabricksMemoryStore(store=store_id, actor=f"alice-{suffix}", client=client) + bob = DatabricksMemoryStore(store=store_id, actor=f"bob-{suffix}", client=client) + + # The same key in the same namespace for both actors: no collision, each keeps its own value. + alice.put(("notes",), "shared-key", {"owner": "alice"}) + bob.put(("notes",), "shared-key", {"owner": "bob"}) + bob.put(("bob-only",), "secret", {"owner": "bob"}) + + alice_item = alice.get(("notes",), "shared-key") + bob_item = bob.get(("notes",), "shared-key") + assert alice_item is not None and alice_item.value == {"owner": "alice"} + assert bob_item is not None and bob_item.value == {"owner": "bob"} + assert alice.get(("bob-only",), "secret") is None + assert [hit.key for hit in alice.search(("notes",))] == ["shared-key"] + assert alice.list_namespaces() == [("notes",)] + + +@gated +def test_langgraph_store_keeps_entries_under_the_storage_actor( + workspace_client, agentbricks_cli, cleanup +): + """DatabricksMemoryStore: entries live under ':langgraph', apart from the plain actor.""" + from databricks_agentkit.client import AgentKitClient + from databricks_agentkit.langgraph.store import DatabricksMemoryStore + + store_id = _new_store(agentbricks_cli, cleanup) + client = AgentKitClient(workspace_client.client) + actor = f"tenant:user-{uuid.uuid4().hex[:8]}" # a composite id, as the docs suggest + store = DatabricksMemoryStore(store=store_id, actor=actor, client=client) + store.put(("notes",), "k1", {"text": "alpha"}) + + handle = client.memory_stores._handle(store_id) + storage = DatabricksMemoryStore.storage_actor(actor) + assert [m.path for m in handle.list(actor_id=storage)] == ["/langgraph/notes/k1"] + assert list(handle.list(actor_id=actor)) == [] # the plain actor, as sessions see it + + +@gated +def test_langgraph_store_ttl_hides_expired_items_until_a_sweep( + workspace_client, agentbricks_cli, cleanup +): + """DatabricksMemoryStore: a short TTL hides the item, and sweep_ttl deletes it.""" + import time + + from databricks_agentkit.client import AgentKitClient + from databricks_agentkit.langgraph.store import DatabricksMemoryStore + + store_id = _new_store(agentbricks_cli, cleanup) + store = DatabricksMemoryStore( + store=store_id, + actor=f"e2e-{uuid.uuid4().hex[:8]}", + client=AgentKitClient(workspace_client.client), + ) + store.put(("notes",), "short", {"text": "alpha"}, ttl=0.1) # 6 seconds + store.put(("notes",), "kept", {"text": "beta"}) + assert store.get(("notes",), "short") is not None + + time.sleep(8) + assert store.get(("notes",), "short") is None + assert [hit.key for hit in store.search(("notes",))] == ["kept"] + assert store.sweep_ttl() == 1 + assert store.sweep_ttl() == 0 + assert store.get(("notes",), "kept") is not None + + +@gated +def test_search_mode(workspace_client, agentbricks_cli, cleanup): + """Managed-memory search: detect semantic matching, and check paraphrase recall where it's on. + + A query that shares no word with any entry comes back empty under full-text matching only, so + the test skips there. Where semantic matching is on, the paraphrase must rank its fact first. + """ + import time + + from databricks_agentkit.client import AgentKitClient + from databricks_agentkit.langgraph.store import DatabricksMemoryStore + + store_id = _new_store(agentbricks_cli, cleanup) + store = DatabricksMemoryStore( + store=store_id, + actor=f"e2e-{uuid.uuid4().hex[:8]}", + client=AgentKitClient(workspace_client.client), + ) + store.put(("facts",), "email", {"text": "Prefers to be contacted by email rather than phone."}) + store.put(("facts",), "peanuts", {"text": "Allergic to peanuts."}) + + # The service embeds entries asynchronously after a write, so poll for a bounded time. + deadline, hits = time.monotonic() + 120, [] + while time.monotonic() < deadline: + hits = store.search(("facts",), query="electronic mail correspondence") + if hits: + break + time.sleep(5) + if not hits: + pytest.skip("This workspace's memory search matched no meaning-only query: full-text only.") + assert hits[0].key == "email" + assert store.search(("facts",), query="food allergy")[0].key == "peanuts" diff --git a/integrations/agentbricks/tests/unit_tests/_memory_api_fake.py b/integrations/agentbricks/tests/unit_tests/_memory_api_fake.py new file mode 100644 index 000000000..13c7a066c --- /dev/null +++ b/integrations/agentbricks/tests/unit_tests/_memory_api_fake.py @@ -0,0 +1,353 @@ +"""Wire-level fake of the managed memory entries API, shared by the LangGraph store tests. + +The fake stands in for ``WorkspaceClient.api_client``, the one call ``_AgentBricksApiClient._do`` +makes. Everything above it runs for real: AgentKitClient, _AgentBricksApiClient (route, query and +body building, error mapping), MemoryStores, MemoryStore, Memory and DatabricksMemoryStore. + +Response shapes are the ones the client parses: lists return ``managed_memory_entries`` plus +``next_page_token``, and search returns ``results[].managed_memory_entry`` with a ``score``. +Server behaviour the client code cannot confirm is set by the constructor knobs, and the tests +that depend on it say so. Nothing here talks to a service. +""" + +from __future__ import annotations + +import copy +import re +import threading +import types +from dataclasses import dataclass, field +from datetime import datetime, timezone +from typing import Any, Callable, Optional + +STORE = "langgraph-store" +_API_PREFIX = "/api/2.0/agents/" +_ENTRY_PATH = re.compile( + r"memory-stores/(?P[^/]+)/entries(?P/(?P[^/]+)|:search)?" +) +_SEARCH_MAX = 100 # the server returns at most this many search hits per call +_LIST_DEFAULT_PAGE = 100 + + +class WireError(Exception): + """An HTTP error as the SDK raises it: a message and an ``error_code`` attribute.""" + + def __init__(self, error_code: str, message: str = "", *, status: int = 400) -> None: + super().__init__(message or error_code) + self.error_code = error_code + self.status = status + + +@dataclass +class Call: + route: str # create | list | get | update | delete | search + store: str + entry: Optional[str] + query: dict + body: Any + thread: int = field(default_factory=threading.get_ident) + + @property + def actor_id(self) -> Optional[str]: + return self.query.get("actor_id") or (self.body or {}).get("actor_id") + + +@dataclass +class _Fault: + route: str + nth: int + error: WireError + commit: bool # apply the request's effect before failing, like a lost response or a race + + +class FakeMemoryApi: + """In-memory memory entries API behind the SDK's ``api_client.do`` seam. + + Knobs: + server_page_size: the server's page size; a list page is ``min(page_size, this)``. + ignore_filters: the server ignores ``path``, ``path_prefix`` and ``session_id`` (the + actor is always applied), so the store's own client-side rechecks are exercised. + + Without ``ignore_filters``, an omitted ``session_id`` returns entries of every session. The + server's real default is not confirmed here, so the store must recheck session scope either way. + Search ranks by whole-word matches (case-insensitive) over content and description. + """ + + def __init__(self, *, server_page_size: int = 100, ignore_filters: bool = False) -> None: + self.server_page_size = server_page_size + self.ignore_filters = ignore_filters + self.entries: dict[str, dict[str, Any]] = {} # resource name -> server row + self.calls: list[Call] = [] + # Called with the resource name at the start of update/delete, to simulate a race. + self.hooks: dict[str, Callable[[str], None]] = {} + self._route_counts: dict[str, int] = {} + self._faults: list[_Fault] = [] + self._seq = 0 + self._lock = threading.Lock() + + # --- setup and inspection ------------------------------------------------ + + def seed( + self, + path: str, + content: Optional[str], + *, + actor_id: str = "alice:langgraph", # the store's storage actor for "alice" + session_id: Optional[str] = None, + description: Optional[str] = None, + store: str = STORE, + create_time: Optional[str] = None, + update_time: Optional[str] = None, + ) -> str: + """Write a row straight into server state (test setup, not a logged call).""" + name = self._new_name(store) + now = _rfc3339(datetime.now(timezone.utc)) + row: dict[str, Any] = { + "name": name, + "actor_id": actor_id, + "path": path, + "create_time": create_time or now, + "update_time": update_time or now, + } + if content is not None: + row["content"] = content + if description is not None: + row["description"] = description + if session_id is not None: + row["session_id"] = session_id + self.entries[name] = row + return name + + def fail( + self, + route: str, + error_code: str, + *, + nth: int = 1, + status: int = 400, + commit: bool = False, + ) -> None: + """Fail the ``nth`` call to ``route`` with ``error_code``. + + With ``commit``, the call's effect is applied first, then the error is raised. + """ + self._faults.append(_Fault(route, nth, WireError(error_code, status=status), commit)) + + def rows(self) -> list[dict[str, Any]]: + return [copy.deepcopy(row) for row in self.entries.values()] + + def routes(self) -> list[str]: + return [call.route for call in self.calls] + + def workspace_client(self) -> types.SimpleNamespace: + """A stand-in ``WorkspaceClient`` whose ``api_client`` is this fake.""" + return types.SimpleNamespace(api_client=self, config=types.SimpleNamespace(host=None)) + + # --- the SDK seam --------------------------------------------------------- + + def do(self, method: str, path: str, query: Optional[dict] = None, body: Any = None) -> Any: + with self._lock: + route, store, entry = _route(method, path) + call = Call(route, store, entry, copy.deepcopy(dict(query or {})), copy.deepcopy(body)) + self.calls.append(call) + self._route_counts[route] = self._route_counts.get(route, 0) + 1 + fault = self._fault_for(route) + if fault is not None and not fault.commit: + raise fault.error + result = getattr(self, f"_{route}")(store, entry, call.query, call.body) + if fault is not None: + raise fault.error + return result + + # --- handlers, one per route ----------------------------------------------- + + def _create(self, store: str, _entry, _query: dict, body: dict) -> dict: + if not body.get("actor_id"): + raise _invalid("actor_id is required") + _check_path(body.get("path") or "") + key = (body["actor_id"], body.get("session_id"), body["path"]) + for row in self.entries.values(): + if (row["actor_id"], row.get("session_id"), row["path"]) == key: + raise WireError( + "ALREADY_EXISTS", "a memory with this path already exists", status=409 + ) + name = self._new_name(store) + now = _rfc3339(datetime.now(timezone.utc)) + row: dict[str, Any] = {"name": name, "create_time": now, "update_time": now} + for field_name in ( + "actor_id", + "path", + "content", + "description", + "session_id", + "source_type", + ): + if body.get(field_name) is not None: + row[field_name] = body[field_name] + self.entries[name] = row + return copy.deepcopy(row) + + def _list(self, store: str, _entry, query: dict, _body) -> dict: + rows = [row for row in self._store_rows(store) if self._visible(row, query)] + size = min(_page_size(query.get("page_size"), _LIST_DEFAULT_PAGE), self.server_page_size) + token = query.get("page_token") or "0" + if not str(token).isdigit(): + raise _invalid("invalid page_token") + start = int(token) + page: dict[str, Any] = {"managed_memory_entries": copy.deepcopy(rows[start : start + size])} + if start + size < len(rows): + page["next_page_token"] = str(start + size) + return page + + def _get(self, store: str, entry: str, _query: dict, _body) -> dict: + return copy.deepcopy(self._row(_name(store, entry))) + + def _update(self, store: str, entry: str, query: dict, body: dict) -> dict: + name = _name(store, entry) + self._run_hook("update", name) + row = self._row(name) + if not query.get("update_mask"): + raise _invalid("update_mask is required") + mask = set(query["update_mask"].split(",")) + if "content" in mask and "content" in body: + row["content"] = body["content"] + if "description" in mask and "description" in body: + row["description"] = body["description"] + row["update_time"] = _rfc3339(datetime.now(timezone.utc)) + return copy.deepcopy(row) + + def _delete(self, store: str, entry: str, _query: dict, _body) -> dict: + name = _name(store, entry) + self._run_hook("delete", name) + self._row(name) + del self.entries[name] + return {} + + def _search(self, store: str, _entry, _query: dict, body: dict) -> dict: + if not (body.get("query") or "").strip(): + raise _invalid("query is required") + terms = _tokens(body["query"]) + cap = _page_size(body.get("page_size") or body.get("limit"), _SEARCH_MAX) + scored = [] + for row in self._store_rows(store): + if not self._visible(row, body): + continue + score = _score(row, terms) + if score: + scored.append((score, row)) + scored.sort(key=lambda hit: -hit[0]) # stable: insertion order breaks ties + return { + "results": [ + {"managed_memory_entry": copy.deepcopy(row), "score": float(score)} + for score, row in scored[:cap] + ] + } + + # --- helpers -------------------------------------------------------------- + + def _new_name(self, store: str) -> str: + self._seq += 1 + return f"memory-stores/{store}/entries/e{self._seq}" + + def _store_rows(self, store: str) -> list[dict[str, Any]]: + return [ + row + for row in self.entries.values() + if row["name"].startswith(f"memory-stores/{store}/") + ] + + def _visible(self, row: dict[str, Any], query: dict) -> bool: + """The server's filter. The actor always applies; the rest only without ``ignore_filters``.""" + if row["actor_id"] != query.get("actor_id"): + return False + if self.ignore_filters: + return True + if query.get("path") is not None: # exact path, every session; session_id is ignored + return row["path"] == query["path"] + prefix = query.get("path_prefix") + if prefix is not None and not row["path"].startswith(prefix): + return False + session = query.get("session_id") + return session is None or row.get("session_id") == session + + def _row(self, name: str) -> dict[str, Any]: + if name not in self.entries: + raise WireError("NOT_FOUND", f"{name} not found", status=404) + return self.entries[name] + + def _run_hook(self, route: str, name: str) -> None: + hook = self.hooks.get(route) + if hook is not None: + hook(name) + + def _fault_for(self, route: str) -> Optional[_Fault]: + nth = self._route_counts[route] + for fault in self._faults: + if fault.route == route and fault.nth == nth: + return fault + return None + + +def install(monkeypatch, fake: Optional[FakeMemoryApi] = None) -> FakeMemoryApi: + """Route the store's real AgentKitClient at a fake, by patching its workspace client only.""" + fake = fake or FakeMemoryApi() + monkeypatch.setattr( + "databricks_agentkit.langgraph.store.workspace_client", fake.workspace_client + ) + return fake + + +def _route(method: str, path: str) -> tuple[str, str, Optional[str]]: + assert path.startswith(_API_PREFIX), f"unexpected path {path!r}" + match = _ENTRY_PATH.fullmatch(path[len(_API_PREFIX) :]) + assert match is not None, f"unexpected memory route {method} {path!r}" + store, rest, entry = match["store"], match["rest"] or "", match["entry"] + if rest == ":search" and method == "POST": + return "search", store, None + if rest == "" and method == "POST": + return "create", store, None + if rest == "" and method == "GET": + return "list", store, None + verbs = {"GET": "get", "PATCH": "update", "DELETE": "delete"} + assert entry is not None and method in verbs, f"unexpected memory route {method} {path!r}" + return verbs[method], store, entry + + +def _name(store: str, entry: str) -> str: + return f"memory-stores/{store}/entries/{entry}" + + +def _tokens(text: str) -> list[str]: + return re.findall(r"\w+", text.lower()) + + +def _score(row: dict[str, Any], terms: list[str]) -> int: + """How many distinct query terms appear as whole words in content or description.""" + words = set(_tokens(f"{row.get('content') or ''} {row.get('description') or ''}")) + return len(words & set(terms)) + + +def _invalid(message: str) -> WireError: + """The service's answer to a malformed request.""" + return WireError("INVALID_PARAMETER_VALUE", message, status=400) + + +def _page_size(value: Any, default: int) -> int: + """Page sizes run 1 to 100; the service rejects anything else rather than clamping it.""" + if value is None: + return default + size = int(value) + if not 1 <= size <= _SEARCH_MAX: + raise _invalid(f"page_size must be between 1 and {_SEARCH_MAX}, got {size}") + return size + + +def _check_path(path: str) -> None: + """Paths start with "/" and have no empty, "." or ".." segments; a trailing "/" is allowed.""" + segments = path[1:].rstrip("/").split("/") if path.startswith("/") else [""] + if not path.startswith("/") or any(s in ("", ".", "..") for s in segments): + raise _invalid(f"invalid path {path!r}") + + +def _rfc3339(moment: datetime) -> str: + return moment.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z" # milliseconds, as the service sends diff --git a/integrations/agentbricks/tests/unit_tests/langgraph_store_edge_test.py b/integrations/agentbricks/tests/unit_tests/langgraph_store_edge_test.py new file mode 100644 index 000000000..3a0b19460 --- /dev/null +++ b/integrations/agentbricks/tests/unit_tests/langgraph_store_edge_test.py @@ -0,0 +1,315 @@ +"""DatabricksMemoryStore edge cases: clocks, envelopes, filters, the search window, actors and inputs.""" + +import time + +import pytest + +pytest.importorskip("langgraph") + +from _memory_api_fake import STORE, install # noqa: E402 +from langchain_core.runnables import RunnableConfig # noqa: E402 +from langgraph.config import var_child_runnable_config # noqa: E402 + +from databricks_agentbricks.errors import AgentCliError # noqa: E402 +from databricks_agentkit._actors import check_actor # noqa: E402 +from databricks_agentkit.langgraph import store as store_mod # noqa: E402 + +T0 = 1_800_000_000_000 +MINUTE = 60_000 + + +@pytest.fixture +def api(monkeypatch): + return install(monkeypatch) + + +@pytest.fixture(autouse=True) +def _fresh_warnings(monkeypatch): + monkeypatch.setattr(store_mod, "_warned", set()) + + +class Clock: + def __init__(self, now=T0): + self.now = now + + def __call__(self): + return self.now + + +def _store(actor="alice", clock=None, **kwargs): + store = store_mod.DatabricksMemoryStore(store=STORE, actor=actor, **kwargs) + if clock is not None: + store._now_ms = clock + return store + + +def _keys(items): + return sorted(item.key for item in items) + + +def test_default_clock_is_epoch_milliseconds(): + before = time.time() * 1000 + now = store_mod._now_ms() + assert abs(now - before) < 60_000 + + +def test_an_empty_dict_is_stored_not_deleted(api): + store = _store() + store.put(("ns",), "k", {}) + assert store.get(("ns",), "k").value == {} + assert len(api.rows()) == 1 + + +@pytest.mark.parametrize("content", ['{"lg":1,"v":"tea"}', '{"lg":1}', '{"lg":1,"v":[1]}']) +def test_envelopes_without_an_object_value_are_skipped(api, content): + api.seed("/langgraph/ns/k", content) + store = _store() + assert store.get(("ns",), "k") is None + assert store.search(("ns",)) == [] + + +def test_nested_and_list_filters_compare_structurally(api): + store = _store() + store.put(("ns",), "en", {"meta": {"lang": "en"}, "tags": ["a", "b"]}) + store.put(("ns",), "str", {"meta": "en", "tags": ["a", "b"]}) + store.put(("ns",), "fr", {"meta": {"lang": "fr"}, "tags": ["a"]}) + store.put(("ns",), "partial", {"meta": {"lang": "x"}, "tags": ["a", "x"]}) + + assert _keys(store.search(("ns",), filter={"meta": {"lang": "en"}})) == ["en"] + assert _keys(store.search(("ns",), filter={"tags": ["a", "b"]})) == [ + "en", + "str", + ] # not "partial" + assert _keys(store.search(("ns",), filter={"tags": ["a"]})) == ["fr"] + + +def test_query_search_warns_when_expired_entries_fill_the_service_window(api, caplog): + import logging + + clock = Clock() + store = _store(clock=clock) + for i in range(70): + store.put(("ns",), f"live{i}", {"note": "tea"}) + for i in range(30): + store.put(("ns",), f"old{i}", {"note": "tea"}, ttl=1) + clock.now += 2 * MINUTE + + with caplog.at_level(logging.WARNING, logger=store_mod.logger.name): + assert len(store.search(("ns",), query="tea", limit=100)) == 70 + assert len([r for r in caplog.records if "reached the most results" in r.getMessage()]) == 1 + + +def test_query_search_does_not_warn_on_an_exactly_full_filtered_page(api, caplog): + import logging + + store = _store() + for i in range(90): + store.put(("ns",), f"red{i}", {"color": "red", "note": "tea"}) + for i in range(10): + store.put(("ns",), f"blue{i}", {"color": "blue", "note": "tea"}) + + with caplog.at_level(logging.WARNING, logger=store_mod.logger.name): + assert len(store.search(("ns",), query="tea", filter={"color": "blue"}, limit=10)) == 10 + assert not [r for r in caplog.records if "reached the most results" in r.getMessage()] + + +def test_query_search_requests_offset_plus_limit_plus_headroom(api): + store = _store() + store.put(("ns",), "k", {"note": "tea"}) + api.calls.clear() + store.search(("ns",), query="tea", limit=5, offset=20) + [search] = [call for call in api.calls if call.route == "search"] + assert search.body["page_size"] == 75 + + +def test_a_store_without_an_id_raises_on_first_use(api, monkeypatch): + monkeypatch.delenv("AGENT_MEMORY_STORE", raising=False) + with pytest.raises(ValueError, match="No memory store configured"): + store_mod.DatabricksMemoryStore(actor="alice").put(("ns",), "k", {"v": 1}) + assert api.calls == [] + + +@pytest.mark.parametrize( + ("code", "commit", "expect_error", "count"), + [ + pytest.param("PERMISSION_DENIED", False, True, None, id="other-error-propagates"), + pytest.param("NOT_FOUND", True, False, 0, id="not-found-is-done"), + ], +) +def test_sweep_propagates_delete_errors_except_not_found(api, code, commit, expect_error, count): + clock = Clock() + store = _store(clock=clock) + store.put(("ns",), "k", {"v": 1}, ttl=1) + clock.now += 2 * MINUTE + api.fail("delete", code, status=403 if code != "NOT_FOUND" else 404, commit=commit) + + if expect_error: + with pytest.raises(AgentCliError): + store.sweep_ttl() + assert len(api.rows()) == 1 + else: + assert store.sweep_ttl() == count + + +def test_search_excludes_a_sibling_namespace_when_the_server_ignores_the_prefix(api): + api.ignore_filters = True + store = _store() + store.put(("ns",), "k", {"note": "tea"}) + store.put(("nsx",), "y", {"note": "tea"}) + + assert [item.namespace for item in store.search(("ns",), query="tea")] == [("ns",)] + assert _keys(store.search(("ns",))) == ["k"] + + +def test_an_empty_run_actor_is_refused(api): + config: RunnableConfig = {"configurable": {"thread_id": "s1", "actor_id": ""}} + token = var_child_runnable_config.set(config) + try: + with pytest.raises(ValueError, match="needs an actor"): + store_mod.DatabricksMemoryStore(store=STORE).put(("ns",), "k", {"v": 1}) + finally: + var_child_runnable_config.reset(token) + assert api.calls == [] + + +@pytest.mark.parametrize("actor", ["alice-langgraph", "langgraph", "alice:langgraph2"]) +def test_actors_that_only_resemble_the_suffix_are_accepted(actor): + assert check_actor(actor) == actor + + +def test_paths_with_empty_segments_are_not_read_as_items(api): + api.seed("/langgraph/ns//k", store_mod._wrap({"color": "tea"}, None)) + store = _store() + assert store.search(("ns",)) == [] + assert store.list_namespaces() == [] + + +def test_in_process_search_ranks_by_shared_terms_and_drops_non_matches(monkeypatch): + monkeypatch.delenv("AGENT_MEMORY_STORE", raising=False) + store = store_mod.DatabricksMemoryStore._local(actor="alice") + store.put(("ns",), "one", {"note": "tea"}) + store.put(("ns",), "both", {"note": "tea and coffee"}) + store.put(("ns",), "none", {"note": "water"}) + + assert [item.key for item in store.search(("ns",), query="tea coffee")] == ["both", "one"] + + +# --- inputs the service would reject ------------------------------------------------------- + + +@pytest.mark.parametrize("value", [["tea"], "tea", 42, ("a", "b")]) +def test_a_value_that_is_not_a_dict_is_refused_before_any_call(api, value): + with pytest.raises(TypeError, match="must be a dict"): + _store().put(("ns",), "k", value) + assert api.calls == [] + + +@pytest.mark.parametrize( + ("value", "error"), + [ + pytest.param({"n": float("nan")}, ValueError, id="nan"), + pytest.param({"n": float("inf")}, ValueError, id="infinity"), + pytest.param({"when": time.gmtime(0)}, None, id="struct-time-is-a-tuple"), + pytest.param({"blob": b"x"}, TypeError, id="bytes"), + pytest.param({"set": {1, 2}}, TypeError, id="set"), + ], +) +def test_values_must_be_strict_json(api, value, error): + if error is None: + _store().put(("ns",), "k", value) + return + with pytest.raises(error, match="for key 'k'"): + _store().put(("ns",), "k", value) + assert api.calls == [] + + +def test_an_over_long_actor_and_path_is_refused_before_any_call(api): + with pytest.raises(ValueError, match="take .* bytes together"): + _store().put(("ns",), "k" * 2_400, {"v": 1}) + assert api.calls == [] + + +def test_a_key_that_is_not_valid_unicode_is_refused(api): + with pytest.raises(ValueError, match="isn't valid Unicode"): + _store().put(("ns",), "\ud800", {"v": 1}) + assert api.calls == [] + + +def test_nul_characters_are_kept_in_the_value_but_not_in_the_description(api): + store = _store() + store.put(("ns",), "k", {"note": "a\x00b"}) + [create] = [call for call in api.calls if call.route == "create"] + assert "\x00" not in create.body["description"] + assert store.get(("ns",), "k").value == {"note": "a\x00b"} + + +def test_a_blank_query_lists_instead_of_searching(api): + store = _store() + store.put(("ns",), "k", {"note": "tea"}) + api.calls.clear() + assert _keys(store.search(("ns",), query=" ")) == ["k"] + assert api.routes() == ["list"] + + +def test_a_search_without_a_query_stops_reading_once_the_page_is_full(api): + store = _store() + for i in range(250): + store.put(("ns",), f"k{i:03}", {"n": i}) + api.calls.clear() + + assert len(store.search(("ns",), limit=10)) == 10 + assert api.routes() == ["list"] # one page of 100, not all three + + +# --- explicit actors ----------------------------------------------------------------------- + + +@pytest.mark.parametrize("actor", ["", 7]) +def test_an_explicit_actor_must_be_a_non_empty_string(api, actor): + with pytest.raises(ValueError, match="non-empty string"): + store_mod.DatabricksMemoryStore(store=STORE, actor=actor) + with pytest.raises(ValueError, match="non-empty string"): + _store().sweep_ttl(actor=actor) + assert api.calls == [] + + +# --- TTL configuration --------------------------------------------------------------------- + + +def test_reads_never_refresh_so_the_ttl_config_says_so(): + store = _store(ttl={"default_ttl": 5}) + assert store.ttl_config == {"default_ttl": 5, "refresh_on_read": False} + + +@pytest.mark.parametrize( + ("ttl", "match"), + [ + pytest.param({"default_tll": 5}, "Unknown ttl keys", id="typo"), + pytest.param({"default_ttl": float("inf")}, "positive number", id="infinite"), + pytest.param({"default_ttl": float("nan")}, "positive number", id="nan"), + pytest.param({"default_ttl": 10**12}, "at most", id="absurd"), + pytest.param({"default_ttl": 1e-9}, "one millisecond", id="rounds-to-zero"), + ], +) +def test_bad_ttl_configs_are_refused(ttl, match): + with pytest.raises(ValueError, match=match): + _store(ttl=ttl) + + +# --- no store bound ------------------------------------------------------------------------ + + +def test_memory_store_resolves_the_store_on_first_use_not_when_built(api, monkeypatch): + monkeypatch.delenv("AGENT_MEMORY_STORE", raising=False) + store = store_mod.memory_store(actor="alice") # eg built at import, before the env is set + monkeypatch.setenv("AGENT_MEMORY_STORE", STORE) + store.put(("ns",), "k", {"v": 1}) + assert api.routes() == ["create"] + + +def test_memory_store_in_a_deployed_app_without_a_store_raises(api, monkeypatch): + monkeypatch.delenv("AGENT_MEMORY_STORE", raising=False) + monkeypatch.setenv("DATABRICKS_APP_NAME", "my-agent") + with pytest.raises(ValueError, match="A deployed App needs a store"): + store_mod.memory_store(actor="alice").put(("ns",), "k", {"v": 1}) + assert api.calls == [] diff --git a/integrations/agentbricks/tests/unit_tests/langgraph_store_graph_test.py b/integrations/agentbricks/tests/unit_tests/langgraph_store_graph_test.py new file mode 100644 index 000000000..5794629b7 --- /dev/null +++ b/integrations/agentbricks/tests/unit_tests/langgraph_store_graph_test.py @@ -0,0 +1,341 @@ +"""DatabricksMemoryStore inside real compiled LangGraph runs: identity, concurrency, strict mode. + +Each test compiles a StateGraph with ``store=`` and runs it with ``thread_config``, so the store +reads the actor from the run config the way an agent does. Only the workspace transport is faked +(``_memory_api_fake``). Async run-config tests need Python 3.11+: on 3.10 async code can't read the +run config, so the matching ``actor=`` variants run on every Python. +""" + +# No `from __future__ import annotations` here: LangGraph injects `runtime: Runtime` by annotation. +import asyncio +import contextvars +import logging +import sys +import threading +from concurrent.futures import ThreadPoolExecutor +from typing import Any, TypedDict, cast + +import pytest + +pytest.importorskip("langgraph") + +from _memory_api_fake import STORE, install # noqa: E402 +from langgraph.graph import END, START, StateGraph # noqa: E402 +from langgraph.runtime import Runtime # noqa: E402 +from langgraph.store.base import BaseStore, Item # noqa: E402 + +from databricks_agentkit.langgraph import store as store_mod # noqa: E402 +from databricks_agentkit.langgraph.session_store import thread_config # noqa: E402 + +needs_3_11 = pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.10: async code can't read the run config; the actor= variant covers it", +) + + +class State(TypedDict, total=False): + text: str + found: object + user: str + op: str + + +@pytest.fixture +def api(monkeypatch): + return install(monkeypatch) + + +def _chain(store, save, load): + builder = StateGraph(cast(Any, State)) + builder.add_node("save", save) + builder.add_node("load", load) + builder.add_edge(START, "save") + builder.add_edge("save", "load") + builder.add_edge("load", END) + return builder.compile(store=store) + + +def _store_of(runtime: Runtime) -> BaseStore: + assert runtime.store is not None, "the graph was compiled without a store" + return runtime.store + + +def save(state: State, runtime: Runtime) -> dict: + if state.get("text"): + _store_of(runtime).put(("notes",), "note", {"text": state["text"]}) + return {} + + +def load(state: State, runtime: Runtime) -> dict: + item = _store_of(runtime).get(("notes",), "note") + return {"found": item.value if item is not None else None} + + +async def async_save(state: State, runtime: Runtime) -> dict: + if state.get("text"): + await _store_of(runtime).aput(("notes",), "note", {"text": state["text"]}) + return {} + + +async def async_load(state: State, runtime: Runtime) -> dict: + item = await _store_of(runtime).aget(("notes",), "note") + return {"found": item.value if item is not None else None} + + +def _creates(api): + return [call for call in api.calls if call.route == "create"] + + +# --- two users, one store ---------------------------------------------------------- + + +def test_two_users_share_one_store_and_see_only_their_own_notes(api): + graph = _chain(store_mod.DatabricksMemoryStore(store=STORE), save, load) + + assert graph.invoke({"text": "alice note"}, config=thread_config("alice-1", "alice")) == { + "text": "alice note", + "found": {"text": "alice note"}, + } + assert graph.invoke({}, config=thread_config("bob-1", "bob"))["found"] is None + # A new session for alice still sees her memory: memories outlive one conversation. + assert graph.invoke({}, config=thread_config("alice-2", "alice"))["found"] == { + "text": "alice note" + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "mode", + [pytest.param("run_config", marks=needs_3_11), "actor_kwarg"], +) +async def test_two_users_share_one_store_async(api, mode): + if mode == "run_config": + alice_graph = bob_graph = _chain( + store_mod.DatabricksMemoryStore(store=STORE), async_save, async_load + ) + else: + alice_graph = _chain( + store_mod.DatabricksMemoryStore(store=STORE, actor="alice"), async_save, async_load + ) + bob_graph = _chain( + store_mod.DatabricksMemoryStore(store=STORE, actor="bob"), async_save, async_load + ) + + alice = await alice_graph.ainvoke( + {"text": "alice note"}, config=thread_config("alice-1", "alice") + ) + bob = await bob_graph.ainvoke({}, config=thread_config("bob-1", "bob")) + alice_again = await alice_graph.ainvoke({}, config=thread_config("alice-2", "alice")) + + assert alice["found"] == {"text": "alice note"} + assert bob["found"] is None + assert alice_again["found"] == {"text": "alice note"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "mode", + [pytest.param("run_config", marks=needs_3_11), "actor_kwarg"], +) +async def test_fifty_concurrent_async_runs_do_not_mix_users(api, mode): + users = [f"user-{i:02d}" for i in range(50)] + runs = [] + for user in users: + store = ( + store_mod.DatabricksMemoryStore(store=STORE) + if mode == "run_config" + else store_mod.DatabricksMemoryStore(store=STORE, actor=user) + ) + graph = _chain(store, async_save, async_load) + runs.append( + graph.ainvoke({"text": f"note of {user}"}, config=thread_config(f"s-{user}", user)) + ) + + results = await asyncio.gather(*runs) + + assert [result["found"] for result in results] == [ + {"text": f"note of {user}"} for user in users + ] + assert {call.actor_id for call in _creates(api)} == {f"{user}:langgraph" for user in users} + assert len(api.rows()) == 50 + + +# --- identity checks in a run ----------------------------------------------------------- + + +@pytest.mark.parametrize( + ("run_actor", "expect_error"), + [ + pytest.param("alice", False, id="matching-actor"), + pytest.param("bob", True, id="different-actor"), + pytest.param(None, True, id="session-defaulted-actor"), + ], +) +def test_explicit_actor_is_checked_against_the_run_actor(api, run_actor, expect_error): + graph = _chain(store_mod.DatabricksMemoryStore(store=STORE, actor="alice"), save, load) + config = thread_config("s1", run_actor) # run_actor=None: the actor defaults to "s1" + + if expect_error: + with pytest.raises(ValueError, match="built for actor 'alice'"): + graph.invoke({"text": "x"}, config=config) + assert api.calls == [] + return + graph.invoke({"text": "x"}, config=config) + assert {call.actor_id for call in _creates(api)} == {"alice:langgraph"} + + +def test_warns_once_when_the_actor_defaults_to_the_session_id(api, caplog, monkeypatch): + monkeypatch.setattr(store_mod, "_warned", set()) + graph = _chain(store_mod.DatabricksMemoryStore(store=STORE), save, load) + + with caplog.at_level(logging.WARNING, logger=store_mod.logger.name): + graph.invoke({"text": "a"}, config=thread_config("s9")) + graph.invoke({"text": "b"}, config=thread_config("s9")) + + warned = [r for r in caplog.records if "session id" in r.getMessage()] + assert len(warned) == 1 + assert {call.actor_id for call in _creates(api)} == {"s9:langgraph"} + + +@pytest.mark.parametrize( + ("langgraph_user_id", "error"), + [ + pytest.param(None, "langgraph_user_id", id="key-missing"), + pytest.param("bob", "refusing to read or write another user's memory", id="mismatch"), + pytest.param("alice", None, id="match"), + ], +) +def test_strict_mode_needs_the_user_key_to_match_the_actor(api, langgraph_user_id, error): + store = store_mod.DatabricksMemoryStore(store=STORE, strict_user_key="langgraph_user_id") + graph = _chain(store, save, load) + config = thread_config("s1", "alice") + if langgraph_user_id is not None: + config["configurable"]["langgraph_user_id"] = langgraph_user_id + + if error is not None: + with pytest.raises(ValueError, match=error): + graph.invoke({"text": "mine"}, config=config) + assert api.calls == [] + return + assert graph.invoke({"text": "mine"}, config=config)["found"] == {"text": "mine"} + assert {call.actor_id for call in api.calls} == {"alice:langgraph"} + + +def test_outside_a_run_the_actor_kwarg_is_used_and_no_actor_raises(api): + store = store_mod.DatabricksMemoryStore(store=STORE, actor="alice") + store.put(("ns",), "k", {"v": 1}) + item = store.get(("ns",), "k") + assert item is not None and item.value == {"v": 1} + assert [call.route for call in api.calls] == ["create", "list"] + assert {call.actor_id for call in api.calls} == {"alice:langgraph"} + + with pytest.raises(ValueError, match="needs a run config"): + store_mod.DatabricksMemoryStore(store=STORE).put(("ns",), "k", {"v": 2}) + # Strict mode needs a run config to verify the user, so an actor= alone is not enough. + strict = store_mod.DatabricksMemoryStore( + store=STORE, actor="alice", strict_user_key="langgraph_user_id" + ) + with pytest.raises(ValueError, match="needs a run config"): + strict.get(("ns",), "k") + assert len(api.calls) == 2 + + +# --- threads and context ------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("use_copy_context", "expected"), + [ + pytest.param(False, "copy_context", id="plain-thread-pool"), + pytest.param(True, {"text": "mine"}, id="copied-context"), + ], +) +def test_a_thread_started_in_a_node_needs_a_copied_context(api, use_copy_context, expected): + def fetch_in_pool(state: State, runtime: Runtime) -> dict: + store = _store_of(runtime) + ctx = contextvars.copy_context() # taken on the node's thread, where the run config lives + + def read() -> Item | None: + return store.get(("notes",), "note") + + def read_in_copied_context() -> Item | None: + return ctx.run(read) + + target = read_in_copied_context if use_copy_context else read + with ThreadPoolExecutor(max_workers=1) as pool: + item = pool.submit(target).result() + return {"found": item.value if item is not None else None} + + graph = _chain(store_mod.DatabricksMemoryStore(store=STORE), save, fetch_in_pool) + config = thread_config("s1", "alice") + + if use_copy_context: + assert graph.invoke({"text": "mine"}, config=config)["found"] == expected + return + with pytest.raises(ValueError, match=expected): + graph.invoke({"text": "mine"}, config=config) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "mode", + [pytest.param("run_config", marks=needs_3_11), "actor_kwarg"], +) +async def test_abatch_reads_the_actor_on_the_event_loop_not_in_the_worker(api, monkeypatch, mode): + config_reads = [] + real_get_config = store_mod.get_config + + def recording_get_config(): + config_reads.append(threading.get_ident()) + return real_get_config() + + monkeypatch.setattr(store_mod, "get_config", recording_get_config) + loop_thread = [] + + async def record_and_save(state: State, runtime: Runtime) -> dict: + loop_thread.append(threading.get_ident()) + await _store_of(runtime).aput(("notes",), "note", {"text": "mine"}) + return {} + + store = ( + store_mod.DatabricksMemoryStore(store=STORE) + if mode == "run_config" + else store_mod.DatabricksMemoryStore(store=STORE, actor="alice") + ) + graph = _chain(store, record_and_save, async_load) + result = await graph.ainvoke({}, config=thread_config("s1", "alice")) + + assert result["found"] == {"text": "mine"} + assert set(config_reads) == {loop_thread[0]} + # The service call itself ran on a worker thread, so the actor had to be read before it. + assert loop_thread[0] not in {call.thread for call in api.calls} + + +def test_langmem_style_namespace_is_partitioned_by_actor_not_by_label(api): + store = store_mod.DatabricksMemoryStore(store=STORE, strict_user_key="langgraph_user_id") + + def act(state: State, runtime: Runtime) -> dict: + # LangMem-style: the namespace is templated on the user id. + namespace = ("memories", state["user"]) + if state["op"] == "remember": + _store_of(runtime).put(namespace, "fact", {"text": state["text"]}) + return {} + return {"found": [item.value for item in _store_of(runtime).search(namespace)]} + + builder = StateGraph(cast(Any, State)) + builder.add_node("act", act) + builder.add_edge(START, "act") + builder.add_edge("act", END) + graph = builder.compile(store=store) + + alice_config = thread_config("s-a", "alice") + alice_config["configurable"]["langgraph_user_id"] = "alice" + bob_config = thread_config("s-b", "bob") + bob_config["configurable"]["langgraph_user_id"] = "bob" + + graph.invoke({"user": "alice", "op": "remember", "text": "likes tea"}, config=alice_config) + # Bob asks for alice's namespace label: the actor partitions, so he gets nothing of hers. + assert graph.invoke({"user": "alice", "op": "recall"}, config=bob_config)["found"] == [] + assert graph.invoke({"user": "alice", "op": "recall"}, config=alice_config)["found"] == [ + {"text": "likes tea"} + ] diff --git a/integrations/agentbricks/tests/unit_tests/langgraph_store_isolation_ttl_test.py b/integrations/agentbricks/tests/unit_tests/langgraph_store_isolation_ttl_test.py new file mode 100644 index 000000000..467ec051e --- /dev/null +++ b/integrations/agentbricks/tests/unit_tests/langgraph_store_isolation_ttl_test.py @@ -0,0 +1,325 @@ +"""DatabricksMemoryStore: the storage actor, the stored-value envelope and client-side TTL. + +Same seam as ``langgraph_store_test.py``: the real AgentKit client over the wire fake +(``_memory_api_fake``). The store's clock is replaced per test, so no test sleeps. +""" + +import asyncio +import logging +import uuid +from contextlib import contextmanager +from typing import Any + +import pytest + +pytest.importorskip("langgraph") + +from _memory_api_fake import STORE, install # noqa: E402 +from langchain_core.runnables import RunnableConfig # noqa: E402 +from langgraph.config import var_child_runnable_config # noqa: E402 +from langgraph.store.base import PutOp # noqa: E402 + +from databricks_agentkit.langgraph import memory as memory_mod # noqa: E402 +from databricks_agentkit.langgraph import session_store as session_store_mod # noqa: E402 +from databricks_agentkit.langgraph import store as store_mod # noqa: E402 +from databricks_agentkit.langgraph import thread_config # noqa: E402 + +T0 = 1_800_000_000_000 # an arbitrary "now", in epoch milliseconds +MINUTE = 60_000 + + +@pytest.fixture +def api(monkeypatch): + return install(monkeypatch) + + +@pytest.fixture(autouse=True) +def _fresh_warnings(monkeypatch): + monkeypatch.setattr(store_mod, "_warned", set()) + + +class Clock: + def __init__(self, now: int = T0) -> None: + self.now = now + + def __call__(self) -> int: + return self.now + + +def _store(actor="alice", clock=None, **kwargs): + store = store_mod.DatabricksMemoryStore(store=STORE, actor=actor, **kwargs) + if clock is not None: + store._now_ms = clock + return store + + +@contextmanager +def run_as(actor, **configurable): + config: RunnableConfig = { + "configurable": {"thread_id": "s1", "actor_id": actor, **configurable} + } + token = var_child_runnable_config.set(config) + try: + yield + finally: + var_child_runnable_config.reset(token) + + +def _warnings(caplog, text): + return [r for r in caplog.records if text in r.getMessage()] + + +# --- the storage actor -------------------------------------------------------------- + + +def test_every_operation_goes_to_the_storage_actor_never_the_plain_one(api): + store = _store() + store.put(("ns",), "k", {"note": "tea"}) + store.put(("ns",), "k", {"note": "green tea"}) # create conflict, then update + store.get(("ns",), "k") + store.search(("ns",)) + store.search(("ns",), query="tea") + store.list_namespaces() + store.put(("ns",), "k", None) + + actors = {call.actor_id for call in api.calls if call.actor_id is not None} + assert actors == {"alice:langgraph"} + assert {"create", "list", "update", "search", "delete"} <= set(api.routes()) + + +def test_storage_actor_composes_with_composite_actor_ids(): + assert store_mod.DatabricksMemoryStore.storage_actor("alice") == "alice:langgraph" + assert store_mod.DatabricksMemoryStore.storage_actor("acme:alice") == "acme:alice:langgraph" + + +@pytest.mark.parametrize("actor", ["alice:langgraph", "Alice:LangGraph", "acme:langgraph"]) +def test_an_actor_ending_with_the_reserved_suffix_is_refused_everywhere(api, monkeypatch, actor): + with pytest.raises(ValueError, match="can't end with ':langgraph'"): + store_mod.DatabricksMemoryStore(store=STORE, actor=actor) + store = store_mod.DatabricksMemoryStore(store=STORE) + with run_as(actor): + with pytest.raises(ValueError, match="can't end with ':langgraph'"): + store.put(("ns",), "k", {"v": 1}) + with pytest.raises(ValueError, match="can't end with ':langgraph'"): + store.get(("ns",), "k") + with pytest.raises(ValueError, match="can't end with ':langgraph'"): + store.sweep_ttl(actor=actor) + with pytest.raises(ValueError, match="can't end with ':langgraph'"): + thread_config("s1", actor) + with pytest.raises(ValueError, match="can't end with ':langgraph'"): + thread_config(actor) # the actor defaults to the session id + with pytest.raises(ValueError, match="can't end with ':langgraph'"): # a hand-built config + session_store_mod._parse_config({"configurable": {"thread_id": "s1", "actor_id": actor}}) + monkeypatch.setenv("AGENT_MEMORY_STORE", STORE) + with pytest.raises(ValueError, match="can't end with ':langgraph'"): + memory_mod.memory_tools(actor) + assert api.calls == [] + + +@pytest.mark.parametrize("actor", ["alice:langgraph", "Alice:LangGraph"]) +def test_the_openai_session_and_memory_helpers_refuse_the_reserved_suffix(monkeypatch, actor): + openai_memory = pytest.importorskip("databricks_agentkit.openai.memory") + openai_sessions = pytest.importorskip("databricks_agentkit.openai.sessions") + monkeypatch.setenv("AGENT_MEMORY_STORE", STORE) + with pytest.raises(ValueError, match="can't end with ':langgraph'"): + openai_memory.memory_tools(actor) + with pytest.raises(ValueError, match="can't end with ':langgraph'"): + openai_sessions.session_store("s1", actor) + + +def test_actor_ids_that_are_not_strings_pass_through_the_helpers_unchanged(): + actor = uuid.uuid4() + assert thread_config("s1", actor)["configurable"]["actor_id"] is actor # ty: ignore[invalid-argument-type] + + +def test_memory_tools_for_the_same_user_never_see_store_entries(api, monkeypatch): + monkeypatch.setattr(memory_mod, "workspace_client", api.workspace_client) + monkeypatch.setenv("AGENT_MEMORY_STORE", STORE) + _store().put(("prefs",), "style", {"note": "tea, no sugar"}) + + tools = {t.name: t for t in memory_mod.memory_tools("alice")} + + recall_args: dict[str, Any] = {"query": "tea"} + list_args: dict[str, Any] = {} + assert "tea" not in tools["recall"].invoke(recall_args) + assert "prefs" not in tools["list_memories"].invoke(list_args) + + +# --- the envelope ------------------------------------------------------------------------- + + +def test_a_newer_envelope_version_is_skipped_with_one_warning(api, caplog): + api.seed("/langgraph/ns/k", '{"lg":2,"v":{"note":"tea"},"y":1}') + store = _store() + + with caplog.at_level(logging.WARNING, logger=store_mod.logger.name): + assert store.get(("ns",), "k") is None + assert store.search(("ns",)) == [] + assert store.search(("ns",), query="tea") == [] + assert store.list_namespaces() == [] + + assert len(_warnings(caplog, "newer version")) == 1 + + +def test_foreign_entries_warn_once_without_naming_paths_or_values(api, caplog): + api.seed("/langgraph/users/alice@example.com/a", "User prefers tea") + api.seed("/langgraph/users/alice@example.com/b", '{"note":"tea"}') + store = _store() + + with caplog.at_level(logging.WARNING, logger=store_mod.logger.name): + assert store.search(("users",)) == [] + assert store.list_namespaces() == [] + + [warning] = _warnings(caplog, "this store wrote") + assert "alice" not in warning.getMessage() and "tea" not in warning.getMessage() + + +def test_a_value_that_is_not_valid_unicode_is_refused_before_any_call(api): + with pytest.raises(ValueError, match="isn't valid Unicode"): + _store().put(("ns",), "k", {"text": "\ud800"}) + assert api.calls == [] + + +# --- TTL ---------------------------------------------------------------------------- + + +def test_ttl_sets_an_expiry_from_the_clock_and_expires_at_that_millisecond(api): + clock = Clock() + store = _store(clock=clock) + store.put(("ns",), "k", {"v": 1}, ttl=1.5) + + [row] = api.rows() + assert row["content"] == store_mod._wrap({"v": 1}, T0 + 90_000) + clock.now = T0 + 90_000 - 1 + assert store.get(("ns",), "k").value == {"v": 1} + clock.now = T0 + 90_000 + assert store.get(("ns",), "k") is None + + +def test_expired_items_are_hidden_from_every_read_and_reads_never_write(api): + clock = Clock() + store = _store(clock=clock) + store.put(("old",), "k", {"note": "tea"}, ttl=1) + store.put(("new",), "k", {"note": "tea"}) + clock.now += 2 * MINUTE + api.calls.clear() + + assert store.get(("old",), "k") is None + assert [item.namespace for item in store.search(())] == [("new",)] + assert [item.namespace for item in store.search((), query="tea")] == [("new",)] + assert store.list_namespaces() == [("new",)] + + # get and search default to refresh_ttl=True; this store neither refreshes nor deletes on read. + assert set(api.routes()) <= {"list", "search"} + assert len(api.rows()) == 2 # the expired entry stays until a sweep + + +def test_put_without_ttl_clears_an_earlier_ttl(api): + clock = Clock() + store = _store(clock=clock) + store.put(("ns",), "k", {"v": 1}, ttl=1) + store.put(("ns",), "k", {"v": 2}) + clock.now += 10 * MINUTE + + assert store.get(("ns",), "k").value == {"v": 2} + assert '"x"' not in api.rows()[0]["content"] + + +def test_default_ttl_applies_to_puts_that_do_not_pass_ttl(api): + clock = Clock() + store = _store(clock=clock, ttl={"default_ttl": 2}) + store.put(("ns",), "default", {"v": 1}) + store.put(("ns",), "explicit", {"v": 1}, ttl=5) + store.put(("ns",), "none", {"v": 1}, ttl=None) # an explicit None means no TTL + store.batch([PutOp(("ns",), "direct", {"v": 1})]) # a raw PutOp carries its own ttl (None) + + contents = {row["path"].rsplit("/", 1)[-1]: row["content"] for row in api.rows()} + assert contents == { + "default": store_mod._wrap({"v": 1}, T0 + 2 * MINUTE), + "explicit": store_mod._wrap({"v": 1}, T0 + 5 * MINUTE), + "none": store_mod._wrap({"v": 1}, None), + "direct": store_mod._wrap({"v": 1}, None), + } + + +@pytest.mark.parametrize( + ("config", "message"), + [ + ({"sweep_interval_minutes": 5}, "sweep_interval_minutes isn't supported"), + ({"refresh_on_read": True}, "refresh_on_read isn't supported"), + ({"default_ttl": 0}, "ttl must be a positive number of minutes"), + ({"default_ttl": -1.0}, "ttl must be a positive number of minutes"), + ], +) +def test_unsupported_ttl_settings_are_refused_at_construction(api, config, message): + with pytest.raises(ValueError, match=message): + _store(ttl=config) + assert api.calls == [] + + +def test_supported_ttl_settings_are_kept_and_omit_expired_false_returns_expired_items(api): + clock = Clock() + store = _store(clock=clock, ttl={"refresh_on_read": False, "omit_expired": False}) + assert store.ttl_config == {"refresh_on_read": False, "omit_expired": False} + store.put(("ns",), "k", {"v": 1}, ttl=1) + clock.now += 2 * MINUTE + + assert store.get(("ns",), "k").value == {"v": 1} + assert [item.key for item in store.search(("ns",))] == ["k"] + + +def test_sweep_deletes_only_one_users_expired_items(api): + clock = Clock() + alice, bob = _store("alice", clock=clock), _store("bob", clock=clock) + alice.put(("ns",), "expired", {"v": 1}, ttl=1) + alice.put(("ns",), "live", {"v": 2}, ttl=60) + alice.put(("ns",), "forever", {"v": 3}) + bob.put(("ns",), "expired", {"v": 4}, ttl=1) + api.seed("/langgraph/ns/foreign", "User likes tea") + clock.now += 2 * MINUTE + + assert alice.sweep_ttl() == 1 + remaining = sorted((row["actor_id"], row["path"]) for row in api.rows()) + assert remaining == [ + ("alice:langgraph", "/langgraph/ns/foreign"), + ("alice:langgraph", "/langgraph/ns/forever"), + ("alice:langgraph", "/langgraph/ns/live"), + ("bob:langgraph", "/langgraph/ns/expired"), + ] + assert bob.sweep_ttl(actor="bob") == 1 + assert alice.sweep_ttl() == 0 + + +def test_sweep_keeps_an_entry_that_a_put_renewed_after_the_listing(api, monkeypatch): + clock = Clock() + store = _store(clock=clock) + store.put(("ns",), "k", {"v": 1}, ttl=1) + clock.now += 2 * MINUTE + at = store._at + + def renew_then_reread(handle, partition, path): + _store(clock=clock).put(("ns",), "k", {"v": 2}, ttl=60) # another process renews it + return at(handle, partition, path) + + monkeypatch.setattr(store, "_at", renew_then_reread) + + assert store.sweep_ttl() == 0 + monkeypatch.undo() + assert store.get(("ns",), "k").value == {"v": 2} + + +def test_sweep_uses_the_run_actor_and_has_an_async_form(api): + clock = Clock() + store = _store(actor=None, clock=clock) + with run_as("carol"): + store.put(("ns",), "a", {"v": 1}, ttl=1) + store.put(("ns",), "b", {"v": 1}, ttl=1) + clock.now += 2 * MINUTE + with run_as("carol"): + assert store.sweep_ttl() == 2 + + with run_as("carol"): + store.put(("ns",), "c", {"v": 1}, ttl=1) + clock.now += 2 * MINUTE + assert asyncio.run(store.asweep_ttl(actor="carol")) == 1 + assert api.rows() == [] diff --git a/integrations/agentbricks/tests/unit_tests/langgraph_store_test.py b/integrations/agentbricks/tests/unit_tests/langgraph_store_test.py new file mode 100644 index 000000000..de4126a8b --- /dev/null +++ b/integrations/agentbricks/tests/unit_tests/langgraph_store_test.py @@ -0,0 +1,481 @@ +"""DatabricksMemoryStore: ops, paths, filters, search window, errors and the in-process fallback. + +The real store drives the real AgentKitClient, MemoryStores, MemoryStore and Memory wrappers. Only +the workspace transport is faked (``_memory_api_fake``). Run-config and graph behaviour lives in +``langgraph_store_graph_test.py``. +""" + +import json +import logging +from contextlib import contextmanager +from datetime import datetime, timezone + +import pytest + +pytest.importorskip("langgraph") + +from _memory_api_fake import STORE, install # noqa: E402 +from langchain_core.runnables import RunnableConfig # noqa: E402 +from langgraph.config import var_child_runnable_config # noqa: E402 + +from databricks_agentbricks.errors import AgentCliError # noqa: E402 +from databricks_agentkit.langgraph import store as store_mod # noqa: E402 + + +@pytest.fixture +def api(monkeypatch): + return install(monkeypatch) + + +def _store(actor="alice", **kwargs): + return store_mod.DatabricksMemoryStore(store=STORE, actor=actor, **kwargs) + + +def _env(value, expires_at_ms=None): + """An entry's content as the store writes it.""" + return store_mod._wrap(value, expires_at_ms) + + +@contextmanager +def run_as(actor, **configurable): + """Run the body inside a LangGraph run config for ``actor`` (thread ``s1``).""" + config: RunnableConfig = { + "configurable": {"thread_id": "s1", "actor_id": actor, **configurable} + } + token = var_child_runnable_config.set(config) + try: + yield + finally: + var_child_runnable_config.reset(token) + + +# --- paths and keys ---------------------------------------------------------- + + +@pytest.mark.parametrize( + ("key", "segment"), + [(".", "%2E"), ("..", "%2E%2E"), ("a/b", "a%2Fb"), ("%x", "%25x"), ("café", "caf%C3%A9")], +) +def test_keys_round_trip_as_percent_encoded_paths(api, key, segment): + store = _store() + store.put(("ns",), key, {"v": key}) + + assert [row["path"] for row in api.rows()] == [f"/langgraph/ns/{segment}"] + assert store.get(("ns",), key).value == {"v": key} + assert [(item.namespace, item.key) for item in store.search(("ns",))] == [(("ns",), key)] + + +def test_namespace_labels_with_slash_round_trip(api): + store = _store() + namespace = ("a/b", "c d") + store.put(namespace, "k", {"v": 1}) + + assert [row["path"] for row in api.rows()] == ["/langgraph/a%2Fb/c%20d/k"] + assert store.get(namespace, "k").value == {"v": 1} + [hit] = store.search(("a/b",)) + assert (hit.namespace, hit.key) == (namespace, "k") + assert store.list_namespaces() == [namespace] + + +def test_rejected_inputs_raise_without_touching_the_service(api): + with pytest.raises(ValueError, match="key must be a non-empty string"): + _store().put(("ns",), "", {"v": 1}) + with pytest.raises(ValueError, match="root must start with"): + store_mod.DatabricksMemoryStore(store=STORE, actor="alice", root="/x/") + for ttl in (0, -1, True, "5"): + with pytest.raises(ValueError, match="ttl must be a positive number of minutes"): + _store().put(("ns",), "k", {"v": 1}, ttl=ttl) + with pytest.raises(NotImplementedError, match="index"): + _store().put(("ns",), "k", {"v": 1}, index=["text"]) + assert api.calls == [] + + +# --- put, get, overwrite, delete ----------------------------------------------- + + +def test_put_get_and_overwrite_take_the_create_conflict_then_update_path(api): + store = _store() + store.put(("ns",), "k", {"msg": "héllo", "n": 1}) + item = store.get(("ns",), "k") + + assert (item.namespace, item.key, item.value) == (("ns",), "k", {"msg": "héllo", "n": 1}) + [row] = api.rows() + assert row["content"] == '{"lg":1,"v":{"msg":"héllo","n":1}}' + assert row["actor_id"] == "alice:langgraph" + + api.calls.clear() + store.put(("ns",), "k", {"msg": "v2"}) + + # The create hits ALREADY_EXISTS, then the store looks the entry up and updates it. + assert api.routes() == ["create", "list", "update"] + assert store.get(("ns",), "k").value == {"msg": "v2"} + assert len(api.rows()) == 1 + + +def test_description_is_the_text_of_the_string_values(api): + store = _store() + store.put( + ("ns",), "k", {"title": "Trip", "tags": ["paris", {"note": " spaced\n out "}], "n": 3} + ) + store.put(("ns",), "long", {"text": "x" * 600}) + + descriptions = {row["path"].rsplit("/", 1)[-1]: row["description"] for row in api.rows()} + assert descriptions == {"k": "Trip paris spaced out", "long": "x" * 500} + + +def test_delete_removes_only_the_exact_key_and_a_missing_key_is_a_noop(api): + store = _store() + store.put(("ns",), "k", {"v": 1}) + store.put(("ns",), "k2", {"v": 2}) + + store.put(("ns",), "k", None) + assert [row["path"] for row in api.rows()] == ["/langgraph/ns/k2"] + assert store.get(("ns",), "k") is None + assert store.get(("ns",), "k2").value == {"v": 2} + + api.calls.clear() + store.put(("ns",), "never-written", None) + assert api.routes() == ["list"] # looked up, found nothing, deleted nothing + + +@pytest.mark.parametrize( + ("error_code", "status", "commit", "expect_error"), + [ + # The entry was already removed (a retried or raced delete): the store treats it as done. + pytest.param("NOT_FOUND", 404, True, False, id="not-found-is-noop"), + # Any other error reaches the caller, and the entry is still there. + pytest.param("PERMISSION_DENIED", 403, False, True, id="other-errors-propagate"), + ], +) +def test_delete_treats_not_found_as_done_and_propagates_other_errors( + api, error_code, status, commit, expect_error +): + store = _store() + store.put(("ns",), "k", {"v": 1}) + api.fail("delete", error_code, status=status, commit=commit) + + if not expect_error: + store.put(("ns",), "k", None) + assert store.get(("ns",), "k") is None + return + with pytest.raises(AgentCliError) as exc: + store.put(("ns",), "k", None) + assert exc.value.error_code == error_code + assert store.get(("ns",), "k").value == {"v": 1} + + +def test_item_times_come_from_the_server(api): + api.seed( + "/langgraph/ns/k", + _env({"v": 1}), + create_time="2026-01-02T03:04:05.000000Z", + update_time="2026-02-03T04:05:06.000000Z", + ) + item = _store().get(("ns",), "k") + + assert item.created_at == datetime(2026, 1, 2, 3, 4, 5, tzinfo=timezone.utc) + assert item.updated_at == datetime(2026, 2, 3, 4, 5, 6, tzinfo=timezone.utc) + + +# --- what the store must never return ---------------------------------------- + + +@pytest.mark.parametrize("ignore_filters", [False, True], ids=["server-filters", "server-ignores"]) +def test_session_foreign_and_non_object_entries_are_invisible(api, ignore_filters): + api.ignore_filters = ignore_filters + tea = _env({"color": "tea"}) + api.seed("/langgraph/ns/k", tea) + api.seed("/langgraph/ns/sess", tea, session_id="s1") + api.seed("/langgraph/ns/fact", "User likes tea") # plain text, eg rewritten by a pipeline + api.seed("/langgraph/ns/bare", '{"color":"tea"}') # a JSON object without the envelope + api.seed("/langgraph/ns/arr", '["tea"]') + api.seed("/langgraph/ns/num", "1") + api.seed("/langgraph/ns/empty", None) + api.seed("/langgraph/ns/extra", '{"lg":1,"v":{"color":"tea"},"note":"tea"}') + api.seed("/langgraph/ns/badx", '{"lg":1,"v":{"color":"tea"},"x":"soon"}') + api.seed("/langgraph/ns/boolv", '{"lg":true,"v":{"color":"tea"}}') + api.seed("/langgraph/ns/deep", "[" * 200_000 + "]" * 200_000) # json.loads recursion + api.seed("/other/ns/x", tea) # outside the store's root + api.seed("/langgraph/sessonly/x", tea, session_id="s1") # only a session + api.seed("/langgraph/ns/plain", tea, actor_id="alice") # the plain actor: not the store's + store = _store() + + assert store.get(("ns",), "k").value == {"color": "tea"} + hidden = ("sess", "fact", "bare", "arr", "num", "empty", "extra", "badx", "boolv", "deep") + for key in hidden + ("plain",): + assert store.get(("ns",), key) is None + assert [item.key for item in store.search(("ns",))] == ["k"] + assert [item.key for item in store.search(("ns",), query="tea")] == ["k"] + assert store.list_namespaces() == [("ns",)] + + +# --- search without a query ----------------------------------------------------- + + +def test_search_without_query_returns_the_prefix_with_exact_offset_and_limit(api): + store = _store() + for i in range(5): + store.put(("ns", "sub"), f"k{i}", {"i": i}) + store.put(("ns",), "x", {"i": 9}) + store.put(("nsx",), "y", {"i": 8}) # a sibling namespace must not match the prefix + + assert sorted(item.key for item in store.search(("ns",), limit=10)) == [ + "k0", + "k1", + "k2", + "k3", + "k4", + "x", + ] + assert [item.key for item in store.search(("ns", "sub"), limit=3, offset=1)] == [ + "k1", + "k2", + "k3", + ] + + +def test_search_without_query_follows_the_server_pages(api): + api.server_page_size = 2 + store = _store() + for i in range(5): + store.put(("ns",), f"k{i}", {"i": i}) + api.calls.clear() + + assert [item.key for item in store.search(("ns",), limit=10)] == [ + "k0", + "k1", + "k2", + "k3", + "k4", + ] + assert [call.query.get("page_token") for call in api.calls] == [None, "2", "4"] + + +@pytest.mark.parametrize( + ("filter_", "expected"), + [ + ({"score": {"$gt": 15}}, ["c"]), + ({"score": {"$gte": 15}}, ["b", "c"]), + ({"score": {"$lt": 15}}, ["a"]), + ({"score": {"$lte": 15}}, ["a", "b"]), + ({"score": {"$ne": 15}}, ["a", "c", "d"]), # "high" is not a number, but $ne still matches + ({"score": {"$gte": 15}, "kind": "y"}, ["c"]), # every key must match, not any + ], + ids=["gt", "gte", "lt", "lte", "ne", "multi-key-and"], +) +def test_filter_operators_and_boundaries_on_search_without_query(api, filter_, expected): + store = _store() + store.put(("ns",), "a", {"score": 10, "kind": "x"}) + store.put(("ns",), "b", {"score": 15, "kind": "x"}) + store.put(("ns",), "c", {"score": 20, "kind": "y"}) + store.put(("ns",), "d", {"score": "high", "kind": "y"}) + + assert sorted(item.key for item in store.search(("ns",), filter=filter_)) == expected + + +# --- search with a query -------------------------------------------------------- + + +def test_query_search_matches_whole_words_and_stays_within_the_actor(api): + alice, bob = _store("alice"), _store("bob") + alice.put(("ns",), "tea-house", {"title": "Tea House"}) + alice.put(("ns",), "tealeaf", {"note": "tealeaf"}) + alice.put(("ns",), "green", {"note": "green TEA"}) + bob.put(("ns",), "bobs", {"note": "tea"}) + + hits = alice.search(("ns",), query="tea") + assert {item.key: item.score for item in hits} == {"tea-house": 1.0, "green": 1.0} + assert [item.key for item in bob.search(("ns",), query="tea")] == ["bobs"] + + +def test_query_search_passes_the_path_prefix_so_other_namespaces_do_not_use_the_window(api): + for i in range(150): + api.seed(f"/langgraph/other/k{i}", _env({"note": "tea"})) + api.seed("/langgraph/ns/mine", _env({"note": "tea"})) + + results = _store().search(("ns",), query="tea", limit=5) + + assert [item.key for item in results] == ["mine"] + [search] = [call for call in api.calls if call.route == "search"] + assert search.body["path_prefix"] == "/langgraph/ns/" + assert search.body["page_size"] == 55 # offset + limit + 50 headroom + + +def test_query_search_with_a_filter_requests_the_full_window_and_warns_when_it_is_short( + api, caplog +): + store = _store() + for i in range(10): + store.put(("ns",), f"blue{i}", {"color": "blue", "note": "tea"}) + for i in range(60): # these rank first but fail the filter + store.put(("ns",), f"red{i}", {"color": "red", "note": "tea tea tea"}) + + results = store.search(("ns",), query="tea", filter={"color": "blue"}, limit=10) + + assert sorted(item.key for item in results) == sorted(f"blue{i}" for i in range(10)) + [search] = [call for call in api.calls if call.route == "search"] + assert search.body["page_size"] == 100 + + # 120 matches, none passing the filter: the window of 100 is full and the page is short. + for i in range(120): + store.put(("full",), f"red{i}", {"color": "red", "note": "tea"}) + with caplog.at_level(logging.WARNING, logger=store_mod.logger.name): + assert store.search(("full",), query="tea", filter={"color": "blue"}, limit=10) == [] + warned = [r for r in caplog.records if "reached the most results" in r.getMessage()] + assert len(warned) == 1 + + # 70 matches, no filter: the window of 60 covers the page, so there is nothing to warn about. + caplog.clear() + for i in range(70): + store.put(("quiet",), f"q{i}", {"note": "tea"}) + with caplog.at_level(logging.WARNING, logger=store_mod.logger.name): + assert len(store.search(("quiet",), query="tea", limit=10)) == 10 + assert not [r for r in caplog.records if "reached the most results" in r.getMessage()] + + +# --- list_namespaces -------------------------------------------------------------- + + +def test_list_namespaces_by_prefix_suffix_and_wildcards(api): + store = _store() + for namespace in [("a", "x"), ("b", "x"), ("a", "y", "z"), ("c",)]: + store.put(namespace, "k", {"v": 1}) + + assert store.list_namespaces(prefix=("a",)) == [("a", "x"), ("a", "y", "z")] + assert store.list_namespaces(suffix=("x",)) == [("a", "x"), ("b", "x")] + assert store.list_namespaces(prefix=("*", "x")) == [("a", "x"), ("b", "x")] + assert store.list_namespaces(suffix=("*", "z")) == [("a", "y", "z")] + + +def test_list_namespaces_truncates_depth_dedupes_and_pages(api): + store = _store() + for namespace in [("a", "b"), ("a", "b", "c"), ("a", "d")]: + store.put(namespace, "k", {"v": 1}) + for i in range(5): + store.put(("ns", str(i)), "k", {"v": i}) + + assert store.list_namespaces(prefix=("a",), max_depth=1) == [("a",)] + assert store.list_namespaces(prefix=("a",), max_depth=2) == [("a", "b"), ("a", "d")] + assert store.list_namespaces(prefix=("ns",), limit=2, offset=1) == [("ns", "1"), ("ns", "2")] + assert store.list_namespaces(prefix=("ns",), limit=2, offset=4) == [("ns", "4")] + + +# --- service handle, factory and fallback -------------------------------------------- + + +def test_client_is_built_once_on_first_use(api, monkeypatch): + built = [] + + class CountingClient(store_mod.AgentKitClient): + def __init__(self, *args, **kwargs): + built.append(1) + super().__init__(*args, **kwargs) + + monkeypatch.setattr(store_mod, "AgentKitClient", CountingClient) + store = _store() + assert built == [] + + store.put(("ns",), "k", {"v": 1}) + store.get(("ns",), "k") + store.search(("ns",)) + + assert built == [1] + + +def test_memory_store_factory_uses_the_explicit_id_then_the_environment(api, monkeypatch): + monkeypatch.setenv("AGENT_MEMORY_STORE", "env-store") + + store_mod.memory_store(actor="alice").put(("ns",), "k", {"v": 1}) + store_mod.memory_store("explicit-store", actor="alice").put(("ns",), "j", {"v": 2}) + + creates = [call for call in api.calls if call.route == "create"] + assert [(call.store, call.actor_id) for call in creates] == [ + ("env-store", "alice:langgraph"), + ("explicit-store", "alice:langgraph"), + ] + + +def test_memory_store_without_a_store_is_in_process_partitioned_by_actor_and_strict( + monkeypatch, caplog, api +): + monkeypatch.delenv("AGENT_MEMORY_STORE", raising=False) + monkeypatch.delenv("DATABRICKS_APP_NAME", raising=False) + store = store_mod.memory_store() + with caplog.at_level(logging.WARNING, logger=store_mod.logger.name): + with run_as("alice"): + store.put(("ns",), "k", {"v": "alice"}) + assert "No memory store configured" in caplog.text # warned on first use, not at import + + with run_as("bob"): + assert store.get(("ns",), "k") is None # bob can't read alice's entry + with run_as("alice"): + item = store.get(("ns",), "k") + assert item is not None and item.value == {"v": "alice"} + + strict = store_mod.memory_store(strict_user_key="langgraph_user_id") + with run_as("alice", langgraph_user_id="bob"): + with pytest.raises(ValueError, match="refusing to read or write another user's memory"): + strict.put(("ns",), "k", {"v": "x"}) + with run_as("alice", langgraph_user_id="alice"): + strict.put(("ns",), "k", {"v": "strict"}) + item = strict.get(("ns",), "k") + assert item is not None and item.value == {"v": "strict"} + + assert api.calls == [] # the fallback never calls the service + + +def test_value_over_the_size_cap_is_rejected_before_any_call(api): + store = _store() + cap = store_mod._REQUEST_MAX_BYTES + # The request carries the envelope as a JSON string plus a 500-character description. + base = len(json.dumps({"content": _env({"blob": ""}), "description": "x" * 500})) + + store.put(("ns",), "at-cap", {"blob": "x" * (cap - base)}) + assert api.routes() == ["create"] + + api.calls.clear() + with pytest.raises(ValueError, match="byte request"): + store.put(("ns",), "over", {"blob": "x" * (cap - base + 1)}) + assert api.calls == [] + + +@pytest.mark.parametrize( + "text", + [ + pytest.param("\u5317" * 250_000, id="non-ascii-escaped-on-the-wire"), + pytest.param('"' * 300_000, id="quotes-escaped-twice"), + ], +) +def test_the_size_cap_counts_the_escaped_request_not_the_raw_value(api, text): + assert len(_env({"t": text}).encode()) < store_mod._REQUEST_MAX_BYTES # small as raw bytes + with pytest.raises(ValueError, match="byte request"): + _store().put(("ns",), "k", {"t": text}) + assert api.calls == [] + + +def test_list_namespaces_ignores_namespaces_with_only_foreign_entries(api): + api.seed("/langgraph/ghost/fact", "User likes tea") + + assert _store().list_namespaces() == [] + + +# --- known bugs (strict xfail: they flip to XPASS when fixed) ---------------------- + + +@pytest.mark.xfail( + strict=True, + raises=AgentCliError, + reason=( + "bug: put(value) raises NOT_FOUND when the entry is deleted between the lookup and the " + "update. MemoryStore._put_memory in memory_store.py retries only the create, so the " + "update sits outside the retry loop." + ), +) +def test_put_after_a_concurrent_delete_still_writes_the_value(api): + store = _store() + store.put(("ns",), "k", {"v": 1}) + api.hooks["update"] = lambda name: api.entries.pop(name, None) + + store.put(("ns",), "k", {"v": 2}) + + assert store.get(("ns",), "k").value == {"v": 2} From 8ca5130ef197f052fbbb8feb801a968dea1e73d0 Mon Sep 17 00:00:00 2001 From: pshekhar2308 <124131160+pshekhar2308@users.noreply.github.com> Date: Mon, 2 Nov 2026 06:41:30 -0800 Subject: [PATCH 2/2] [agentbricks] Refuse :langgraph actors in SessionStores add/fork (#2) thread_config, memory_tools and the session helpers reject actors ending in ':langgraph' so automatic memory pipelines never see the LangGraph store's storage actor, but the typed session SDK still let callers create a session under that id directly. A session for :langgraph makes it discoverable to pipelines, which can rewrite the store's JSON envelopes. Apply check_actor in _add_session and _fork_session so the SDK enforces the same rule, and cover both paths with unit tests. Generated with [Devin](https://devin.ai) Co-authored-by: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../src/databricks_agentkit/session_store.py | 5 +++-- .../tests/unit_tests/session_store_test.py | 21 +++++++++++++++++++ 2 files changed, 24 insertions(+), 2 deletions(-) diff --git a/integrations/agentbricks/src/databricks_agentkit/session_store.py b/integrations/agentbricks/src/databricks_agentkit/session_store.py index ee7a7a426..c5c87b253 100644 --- a/integrations/agentbricks/src/databricks_agentkit/session_store.py +++ b/integrations/agentbricks/src/databricks_agentkit/session_store.py @@ -7,6 +7,7 @@ from datetime import datetime from typing import TYPE_CHECKING, Any, List, Optional +from databricks_agentkit._actors import check_actor from databricks_agentkit._pagination import validate_page_size from databricks_agentkit.timefmt import parse_timestamp @@ -245,7 +246,7 @@ def _add_session( ) -> Session: response = self._api.create_session( store.name, - actor_id, + check_actor(actor_id), session_id=session_id, parent_session_id=parent_session_id, metadata=metadata, @@ -310,7 +311,7 @@ def _fork_session( response = self._api.fork_session( session.store_name, session.session_id, - actor_id, + check_actor(actor_id), up_to_item_id=up_to_item_id, session_id=session_id, metadata=metadata, diff --git a/integrations/agentbricks/tests/unit_tests/session_store_test.py b/integrations/agentbricks/tests/unit_tests/session_store_test.py index 80ba6c66e..ff240ec80 100644 --- a/integrations/agentbricks/tests/unit_tests/session_store_test.py +++ b/integrations/agentbricks/tests/unit_tests/session_store_test.py @@ -123,6 +123,16 @@ def test_add_get_and_list_sessions() -> None: api.get_session.assert_called_once_with(SESSION_ID, SESSION_STORE) +def test_add_refuses_reserved_langgraph_storage_actor() -> None: + client, api = resource_client() + api.get_session_store.return_value = session_store_payload() + store = client.session_stores.get(SESSION_STORE) + + with pytest.raises(ValueError, match=":langgraph"): + store.add(actor_id="customer-123:langgraph") + api.create_session.assert_not_called() + + def test_update_and_delete_session() -> None: client, api = resource_client() api.get_session_store.return_value = session_store_payload() @@ -175,6 +185,17 @@ def test_fork_unwraps_session_field() -> None: ) +def test_fork_refuses_reserved_langgraph_storage_actor() -> None: + client, api = resource_client() + api.get_session_store.return_value = session_store_payload() + api.get_session.return_value = session_payload() + session = client.session_stores.get(SESSION_STORE).get(SESSION_ID) + + with pytest.raises(ValueError, match=":langgraph"): + session.fork(actor_id="agent:langgraph") + api.fork_session.assert_not_called() + + def test_item_operations_and_pagination() -> None: client, api = resource_client() api.get_session_store.return_value = session_store_payload()