managed-memory

name

managed-memory

description

Give an agent durable, cross-session long-term memory using Databricks MANAGED memory (the Unity Catalog memory-store REST APIs) as tools — governed by UC with no infra the customer needs to run. This works for either OpenAI Agents SDK or LangGraph templates. Use when: the agent should remember a user's (or a team/org's shared) preferences/facts/decisions across conversations; keywords 'long-term memory', 'managed memory', 'memory store', 'agentic memory'. This is separate from the self-hosted Lakebase memory solution with skills in (agent-openai-memory / agent-langgraph-memory).

Long-Term Memory with Databricks Managed Memory (UC memory-store)

Give your agent durable, cross-session memory about each user, exposed as five tools (save_memory, get_memory, list_memories, update_memory, delete_memory). The tools are thin REST calls to the Unity Catalog memory-store APIs.

Beta. The Databricks memory-store APIs are in beta — APIs and behavior may change.

This is Databricks managed memory — NOT the self-hosted Lakebase memory

A memory store is a governed Unity Catalog securable you read/write purely over REST: no database to provision, no tables to create, no embedding endpoint, and no extra Python dependency (it uses the databricks-sdk already in the template). This is different from the agent-openai-memory / agent-langgraph-memory skills, which persist to a Lakebase instance you run yourself. It's additive to short-term/session memory (the OpenAI AsyncDatabricksSession or the LangGraph checkpointer) — keep that. But it is the agent's long-term memory, and there should be only one: if the template already has a long-term memory system, remove it before adding these tools.

This skill is framework-agnostic and flexible with both the OpenAI Agents SDK and LangGraph; each step notes the small per-SDK difference.

For a pre-existing agent (not built from a default template) — still on Databricks Apps. The core (memory-store REST API, the five tools, the scope-as-isolation rule, the grant calls in Steps 1–2) is identical; only the template specifics differ. Map the agent_server/... paths to your own modules and reuse the Databricks Apps primitives you already have: the forwarded OBO user token for the signed-in user's id (what resolve_scope() reads), config.env for DATABRICKS_MEMORY_STORE, and databricks apps to deploy. Two invariants never change: the tools authenticate via WorkspaceClient() as the app service principal you grant on the store, and you pass the end user's id as scope — fail closed, never the SP.

Prerequisites — this is an add-on

This skill adds long-term memory to an agent that's already set up — it doesn't scaffold one. If there's no .env (auth not configured), run the quickstart skill first — it sets the Databricks profile + MLflow experiment, and on the advanced templates provisions the Lakebase used for short-term session memory (which this skill leaves intact). Then come back here. Verify the app already has everything it needs to run first — the quickstart skill tells you what each template needs set up.

Concepts

Object	What it is
Memory store	A UC securable catalog.schema.name (type MEMORY_STORE) — the governance object you grant on and the container for memories. Read/written over REST, no SQL.
Memory entry	One memory: a path (e.g. /memories/preferences/coffee.md), a one-line description, and optional contents.
Scope	The partition key the caller assigns — decides whose memories you read/write. Per-user (a private partition, the default), a shared constant (org/team-wide), or your own logic (per project/tenant, user×project); see Scope strategy below.

Access is two separate questions:

• Can the caller use the store? → make sure the caller has READ_MEMORY_STORE to retrieve memory entries and WRITE_MEMORY_STORE to write them. When testing locally the tools are called with the developer's credentials; when the agent is deployed on Apps they run with the app's credentials.
• Whose memories? → the explicit scope, set by your code: the end user's id for private per-user memory, or a shared org/team constant for memory common to everyone (see Scope strategy below).

The SP can see every scope, so scope is your isolation boundary: always set it in trusted code (to the end user, or a deliberate shared constant), and never let the model choose it.

Step 1 — Create or choose the memory store

Have your admin or agent developer create a memory store you can read/write memory entries to. First establish workspace creds:

export DATABRICKS_HOST="https://<your-workspace-host>"
export TOKEN="$(databricks auth token -p <profile> | jq -r .access_token)"

Ask the user with AskUserQuestion — two setup choices, in one call:

1. The store — "Do you have an existing memory store you can manage, or should I create one?"

• Use an existing store — you own it, or hold MANAGE / MANAGE_ACCESS_CONTROL on it.
• Create a new store — under a catalog + schema you choose; you become the owner (needs CREATE_MEMORY_STORE on that schema).

2. The scope strategy — "How should memories be partitioned: private per end user, shared across a team/org/project, or by your own logic?" (see Scope strategy below for the tradeoffs)

• Per-user (recommended) — each user gets a private partition; the default wiring.
• Custom (shared) — fixed scope multiple users can access; if chosen, collect the scope id as a free-text follow-up and set it as a constant in resolve_scope.
• Custom logic — partition some other way (per project/tenant, or user×project). Ask the user to describe their isolation model, then write resolve_scope to it, honoring the contract under Scope strategy → Your own logic.

The scope answer routes resolve_scope (Steps 3–4) and the MEMORY_INSTRUCTIONS framing (Step 5) — wire whichever the user picked.

Then collect the store details as normal chat messages (free-text — not AskUserQuestion) and run the matching API call. Run these yourself so you can see exactly what each does:

# CREATE a store you own. Returns the securable {full_name, owner, memory_store_id, ...}.
curl -sS -X POST "$DATABRICKS_HOST/api/2.1/unity-catalog/memory-stores" \
  -H "Authorization: Bearer $TOKEN" -H "Content-Type: application/json" \
  -d '{"name":"<name>","catalog_name":"<catalog>","schema_name":"<schema>","description":"Long-term memory for my agent"}'

# OR CONFIRM an existing store resolves (catches a typo'd name now instead of as NOT_FOUND at runtime):
curl -sS "$DATABRICKS_HOST/api/2.1/unity-catalog/memory-stores/<catalog.schema.name>" -H "Authorization: Bearer $TOKEN"

Record the full name in the same env var, in two places — .env (read locally) and databricks.yml under the app's config.env (the deployed app doesn't read .env):

      config:
        env:
          - name: DATABRICKS_MEMORY_STORE
            value: "<catalog.schema.name>"

Step 2 — Grant on the store (API calls)

The tools call the API as whatever principal the agent runs as: the app service principal once deployed, and the developer's own user when running locally (the agent's WorkspaceClient() picks up the local profile). Grant three things to both principals: READ_MEMORY_STORE + WRITE_MEMORY_STORE on the store, plus USE_CATALOG on its catalog and USE_SCHEMA on its schema. The last two are easy to miss and non-obvious: without them every entry (and conversation) call fails with User does not have USE CATALOG — Unity Catalog hides a securable whose parent catalog/schema the caller can't traverse. DAB has no MEMORY_STORE grant yet, so these are direct permissions API calls (not databricks.yml) — run the PATCHes below yourself. STORE is the full name:

export STORE="<catalog.schema.name>"
PERM="$DATABRICKS_HOST/api/2.1/unity-catalog/permissions/memory_store/$STORE"

# Grant the DEVELOPER's user (for local testing — the local agent runs as them):
curl -sS -X PATCH "$PERM" -H "Authorization: Bearer $TOKEN" -H "Content-Type: application/json" \
  -d '{"changes":[{"principal":"<[e-mail]>","add":["READ_MEMORY_STORE","WRITE_MEMORY_STORE"]}]}'

# Grant the DEPLOYED app's service principal (run AFTER deploy, once the app + its SP exist):
APP_SP=$(databricks apps get <your-app> -o json | jq -r .service_principal_client_id)
curl -sS -X PATCH "$PERM" -H "Authorization: Bearer $TOKEN" -H "Content-Type: application/json" \
  -d "{\"changes\":[{\"principal\":\"$APP_SP\",\"add\":[\"READ_MEMORY_STORE\",\"WRITE_MEMORY_STORE\"]}]}"

# ALSO grant USE_CATALOG on the catalog and USE_SCHEMA on the schema to the SAME principals.
# Required — without them the calls above fail "User does not have USE CATALOG". Replace <catalog>
# and <catalog.schema>, and run once per principal (developer + $APP_SP):
for PRIN in "<[e-mail]>" "$APP_SP"; do
  curl -sS -X PATCH "$DATABRICKS_HOST/api/2.1/unity-catalog/permissions/catalog/<catalog>" \
    -H "Authorization: Bearer $TOKEN" -H "Content-Type: application/json" \
    -d "{\"changes\":[{\"principal\":\"$PRIN\",\"add\":[\"USE_CATALOG\"]}]}"
  curl -sS -X PATCH "$DATABRICKS_HOST/api/2.1/unity-catalog/permissions/schema/<catalog.schema>" \
    -H "Authorization: Bearer $TOKEN" -H "Content-Type: application/json" \
    -d "{\"changes\":[{\"principal\":\"$PRIN\",\"add\":[\"USE_SCHEMA\"]}]}"
done

# Verify current store grants:
curl -sS "$PERM" -H "Authorization: Bearer $TOKEN"

Step 3 — Add the memory tools

Put these in agent_server/utils_memory.py — use (a) the shared core + the block for your SDK ((b) for the OpenAI Agents SDK or (c) for LangGraph; not both — they each define _scope their own way). No new dependency — it uses the databricks-sdk already in the template. Most templates have no utils_memory.py — create it (note the OpenAI advanced template keeps its session plumbing in utils.py, not here, so you still create a fresh utils_memory.py). The one exception is agent-langgraph-advanced: its existing utils_memory.py already holds the Lakebase plumbing — short-term checkpointer and a long-term AsyncDatabricksStore + memory_tools(). There, add these functions to that same file (don't create a second one), keep the checkpointer, and replace the long-term store — only one long-term system (see the intro and Step 4).

(a) Shared core — the REST calls and scope resolution (SDK-agnostic):

import os
from databricks.sdk import WorkspaceClient
from databricks.sdk.errors import DatabricksError
from mlflow.genai.agent_server import get_request_headers
from agent_server.utils import get_user_workspace_client

# API: BASE = /api/2.1/unity-catalog/memory-stores/{DATABRICKS_MEMORY_STORE}
#   create  POST {BASE}/entries?scope=…   {path,contents,description,creation_reason,creation_source}  (flat body; scope is a query param)
#   get     GET  {BASE}/entries:get       ?scope,path        -> {contents, description, ...}
#   list    GET  {BASE}/entries           ?scope             -> {entries:[{path,description,has_contents}]}  (key omitted entirely when empty)
#   update  PATCH{BASE}/entries          {scope, path, [description], [one contents edit op]}  (>=1 of the two)
#   delete  DELETE {BASE}/entries         ?scope,path

_client: WorkspaceClient | None = None

def _ws() -> WorkspaceClient:
    """The memory caller — the app SP when deployed, the developer when local. Per-user isolation is via
    `scope`, NEVER this identity (the SP can see every scope)."""
    global _client
    if _client is None:
        _client = WorkspaceClient()
    return _client

def _entries(suffix: str = "") -> str:
    store = os.getenv("DATABRICKS_MEMORY_STORE")
    if not store:
        raise RuntimeError("DATABRICKS_MEMORY_STORE is not set — it must be the full catalog.schema.name.")
    return f"/api/2.1/unity-catalog/memory-stores/{store}/entries{suffix}"

def resolve_scope(request=None) -> str | None:
    """The end user's id used as `scope`, or None if it can't be determined (the handler MUST fail
    closed). Deployed: the OBO forwarded token -> current_user.me().id — the ONLY trusted source.
    Local: an X-Forwarded-User header, the request's custom_inputs.user_id (what the bundled chat UI /
    preflight send). NEVER the app's own identity, and
    NEVER a client-supplied value (X-Forwarded-User / custom_inputs) when deployed — those are spoofable."""
    headers = get_request_headers() or {}
    if headers.get("x-forwarded-access-token"):
        obo = get_user_workspace_client()          # None if the OBO client can't be built (see auth_type note)
        return obo.current_user.me().id if obo else None   # fail closed rather than AttributeError
    # Deployed -> the verified OBO token above is the only trusted source. DATABRICKS_APP_NAME is set by
    # the Apps runtime when deployed, so the client-supplied fallbacks below are LOCAL-DEV ONLY:
    if os.getenv("DATABRICKS_APP_NAME"):
        return None
    ci = dict(getattr(request, "custom_inputs", None) or {})
    return headers.get("x-forwarded-user") or ci.get("user_id")

# The five operations. `scope` is passed in (never model-supplied). Each returns a short string.
def _save(scope, path, description, contents=""):
    try:
        _ws().api_client.do("POST", _entries(), query={"scope": scope}, body={
            "path": path, "contents": contents, "description": description,
            "creation_reason": "CREATION_REASON_AGENT_INFERRED",
            "creation_source": "CREATION_SOURCE_ONLINE_AGENT"})
    except DatabricksError as e:
        if e.error_code == "ALREADY_EXISTS":
            return f"A memory already exists at {path}; use update_memory to revise it."
        return f"Could not save {path}: {getattr(e, 'message', str(e))}"
    return f"Saved memory at {path}."

def _get(scope, path):
    try:
        entry = _ws().api_client.do("GET", _entries(":get"), query={"scope": scope, "path": path})
    except DatabricksError as e:
        if e.error_code == "NOT_FOUND":
            return f"No memory at {path}."
        return f"Could not read {path}: {getattr(e, 'message', str(e))}"
    # A brief memory may have empty contents — its description is then the memory.
    return entry.get("contents") or entry.get("description") or f"(empty memory at {path})"

def _list(scope):
    try: