LangGraph Checkpointer — CosmosDB¶

A LangGraph BaseCheckpointSaver for Azure CosmosDB with built-in message pruning. It persists agent state between runs so your graphs can resume from any prior checkpoint, and it can automatically cap your message history before each write — no changes to your graph code or state annotations required.

Current version

langgraph-checkpoint-cosmosdb 0.2.7 · Requires Python 3.10+

What it is¶

CosmosDBSaver implements the LangGraph checkpointer interface (put / get_tuple / list plus async counterparts) backed by an Azure CosmosDB container. Unlike other CosmosDB checkpointers, it accepts an optional MessageReducer that prunes the message list at the persistence layer.

Full checkpoint persistence — save, retrieve, and list checkpoints
Built-in message pruning at save time (see Built-in message pruning)
Sync and async API
Subgraph support — parent and subgraph state checkpointed independently
Flexible auth — key-based or Azure RBAC
Auto-creates the database and container under key-based auth

Installation¶

BaseWith pruning

pip install langgraph-checkpoint-cosmosdb

pip install "langgraph-checkpoint-cosmosdb[reducer]"

The [reducer] extra pulls in agentstate-reducer, required only if you pass a reducer.

Database and container setup¶

Auth mode	Database	Container	Partition key
Key-based (`COSMOSDB_KEY` set)	Created automatically if absent	Created automatically if absent	`/partition_key` (set by saver)
RBAC / Managed Identity (no key)	Must pre-exist	Must pre-exist	`/partition_key` (must be pre-configured)

Key-based auth is the easiest way to get started — point the saver at an existing CosmosDB account and it provisions everything.

For RBAC, the saver only calls get_database_client / get_container_client (no setup-time write permissions), so the database and container must already exist:

az cosmosdb sql database create --account-name <account> --name <db>
az cosmosdb sql container create \
  --account-name <account> --database-name <db> --name <container> \
  --partition-key-path "/partition_key"

Partition key path

The partition key path must be /partition_key regardless of how the container is created.

Authentication¶

Key-based (dev)Azure RBAC (prod)User-assigned MIService principal

export COSMOSDB_ENDPOINT="https://<account>.documents.azure.com:443/"
export COSMOSDB_KEY="<your-key>"

Set only the endpoint — no key. The saver uses DefaultAzureCredential, resolving in order: environment service principal → managed identity → az login.

export COSMOSDB_ENDPOINT="https://<account>.documents.azure.com:443/"
# COSMOSDB_KEY not set → DefaultAzureCredential is used

export AZURE_CLIENT_ID="<managed-identity-client-id>"

export AZURE_TENANT_ID="<tenant-id>"
export AZURE_CLIENT_ID="<client-id>"
export AZURE_CLIENT_SECRET="<client-secret>"

Quick start¶

from langgraph.graph import StateGraph, MessagesState, START
from langchain_openai import ChatOpenAI
from langgraph_checkpoint_cosmosdb import CosmosDBSaver

model = ChatOpenAI(model="gpt-4o-mini")

def call_model(state: MessagesState):
    return {"messages": model.invoke(state["messages"])}

builder = StateGraph(MessagesState)
builder.add_node("call_model", call_model)
builder.add_edge(START, "call_model")

checkpointer = CosmosDBSaver(database_name="mydb", container_name="checkpoints")
graph = builder.compile(checkpointer=checkpointer)

config = {"configurable": {"thread_id": "user-123"}}

# First run — state is saved to CosmosDB
graph.invoke({"messages": [{"role": "user", "content": "Hi, I'm Kamal"}]}, config)

# Second run — picks up where it left off
graph.invoke({"messages": [{"role": "user", "content": "What's my name?"}]}, config)

API reference¶

`CosmosDBSaver(database_name, container_name, reducer=None, messages_key="messages")`¶

Parameter	Type	Default	Description
`database_name`	`str`	required	CosmosDB database name
`container_name`	`str`	required	CosmosDB container name
`reducer`	`MessageReducer`	`None`	Optional pruner — see Built-in message pruning
`messages_key`	`str`	`"messages"`	State channel name that holds the message list

Sync methods¶

Method	Description
`put(config, checkpoint, metadata, new_versions)`	Save a checkpoint
`put_writes(config, writes, task_id)`	Save pending writes for a checkpoint
`get_tuple(config)`	Retrieve the latest (or a specific) checkpoint
`list(config, *, before, limit)`	Iterate checkpoints for a thread

Async methods¶

All sync methods have async counterparts: aput, aput_writes, aget_tuple, alist, and adelete.

checkpoint = await saver.aget_tuple(config)
await saver.adelete(thread_id="user-123", checkpoint_namespace="", checkpoint_id="<id>")

list filtering

list only supports filtering by thread_id. The filter parameter (filtering by metadata) is not yet implemented.

Built-in message pruning¶

Long-running agents accumulate message history with every turn, inflating checkpoint size, increasing storage costs, and eventually blowing past LLM context limits.

Pass a MessageReducer and the saver prunes the message list inside put() before the checkpoint is serialised and written to CosmosDB. Your graph code, state definition, and node logic stay untouched. This is an alternative to — or complement of — the LangGraph Annotated[list, reducer_fn] pattern; use it when:

You don't own the graph or state definition (e.g. a pre-built LangGraph agent)
You want pruning at every save, regardless of which node triggered it
You want in-memory state intact and only prune what gets persisted

from agentstate_reducer import MessageReducer
from langgraph_checkpoint_cosmosdb import CosmosDBSaver

reducer = MessageReducer(min_messages=10, max_messages=20)

checkpointer = CosmosDBSaver(
    database_name="mydb",
    container_name="checkpoints",
    reducer=reducer,        # prune before each checkpoint save
    messages_key="messages" # state channel holding the message list (default)
)

When len(messages) > max_messages, the oldest human/ai messages are removed until min_messages remain. System-prompt index 0, system/function messages, and tool messages (unless their parent ai message is pruned) are preserved.

Full reducer configuration

For preserve_first, cascade_tool_messages, summarize_fn, token budgeting, and role aliases, see the reducer overview and token budget docs.

Data model¶

Checkpoints and writes are stored as separate items in the same container, differentiated by a key prefix and partition key:

Item type	Partition key format	Item id format
Checkpoint	`checkpoint$<thread_id>$<ns>$`	`checkpoint$<thread_id>$<ns>$<checkpoint_id>`
Pending write	`writes$<thread_id>$<ns>$<checkpoint_id>$`	`writes$<thread_id>$<ns>$<checkpoint_id>$<task_id>$<idx>`

The container requires a partition key path of /partition_key.