# From SQLite to Cosmos DB: How I Built a Production-Grade Checkpoint Saver for LangGraph

*Your AI agent remembers everything — until the server restarts. Here's how I adapted LangGraph's persistence layer from a local SQLite file to a globally distributed Azure Cosmos DB backend, one design decision at a time.*

---

<!-- 🖼️ HERO IMAGE: A wide, cinematic banner showing a futuristic network of connected nodes/neurons with Azure blue and LangGraph green gradients. Think "AI infrastructure meets cloud" aesthetic. Dimensions: 1400×600px -->

![hero-banner-image](placeholder-hero-banner.png)

---

## Why I Built This

If you've spent any time building with [LangGraph](https://github.com/langchain-ai/langgraph) — LangChain's framework for building stateful, multi-step AI agents — you've probably experienced the magic moment where your agent *remembers* a conversation across multiple turns.

That magic comes from **checkpointing**: LangGraph takes a snapshot of graph state after every node execution and saves it somewhere. By default, that "somewhere" is a SQLite file.

And SQLite is *great* — for demos.

<!-- 🖼️ IMAGE: A simple side-by-side comparison graphic. Left side: laptop icon with SQLite logo and text "Great for prototyping". Right side: cloud/globe icon with Cosmos DB logo and text "Built for production". Clean, minimal style. Dimensions: 900×400px -->

![sqlite-vs-cosmos-comparison](placeholder-sqlite-vs-cosmos.png)

But the moment you try to:

- 🔄 **Run multiple instances** of your agent behind a load balancer
- 🚀 **Deploy to Kubernetes** or Azure Container Apps
- 📈 **Handle thousands of concurrent conversations**
- 🌍 **Serve users across geographies** with low latency
- 💾 **Guarantee durability** — your agent's memory must survive hardware failures

...SQLite becomes a bottleneck. It's a single file on a single disk. It doesn't replicate. It doesn't distribute. And when your container restarts, that file might just disappear.

I needed something better. So I built **[langgraph-checkpoint-cosmos](https://github.com/LangModule/checkpoint-cosmos)** — a drop-in replacement that swaps SQLite for Azure Cosmos DB, while preserving every bit of the LangGraph checkpoint interface.

```bash
pip install langgraph-checkpoint-cosmos
```

This article is the story of *how* I adapted the original SQLite implementation to work with a globally distributed document database — every trade-off, every design pattern, every decision.

---

## Understanding the Starting Point: How the SQLite Saver Works

Before we can talk about adaptation, we need to understand what we're adapting *from*. The original `SqliteSaver` from the LangGraph repo is clean, well-designed, and — crucially — simple. Let's break down its architecture.

<!-- 🖼️ IMAGE: A clean architectural diagram showing the SqliteSaver internals: two SQL tables (checkpoints and writes), arrows showing data flow from LangGraph → SqliteSaver → SQLite file. Use a technical blueprint/whiteboard style. Dimensions: 1000×500px -->

![sqlite-architecture-diagram](placeholder-sqlite-architecture.png)

### The Schema: Two Tables

The SQLite saver creates two tables at startup:

```sql
CREATE TABLE IF NOT EXISTS checkpoints (
    thread_id TEXT NOT NULL,
    checkpoint_ns TEXT NOT NULL DEFAULT '',
    checkpoint_id TEXT NOT NULL,
    parent_checkpoint_id TEXT,
    type TEXT,
    checkpoint BLOB,
    metadata BLOB,
    PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id)
);

CREATE TABLE IF NOT EXISTS writes (
    thread_id TEXT NOT NULL,
    checkpoint_ns TEXT NOT NULL DEFAULT '',
    checkpoint_id TEXT NOT NULL,
    task_id TEXT NOT NULL,
    idx INTEGER NOT NULL,
    channel TEXT NOT NULL,
    type TEXT,
    value BLOB,
    PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
);
```

Simple. Elegant. Two tables, primary keys, BLOB columns for serialized data.

### The Core Operations

The SQLite saver implements five key methods from `BaseCheckpointSaver`:

| Method | What It Does |
|--------|-------------|
| `get_tuple(config)` | Fetches the latest (or specific) checkpoint + associated writes |
| `list(config, ...)` | Lists/searches checkpoints with filters and pagination |
| `put(config, checkpoint, metadata, ...)` | Saves a new checkpoint |
| `put_writes(config, writes, task_id)` | Stores intermediate writes for a checkpoint |
| `get_next_version(current, channel)` | Generates monotonic version IDs |

### The "Get Latest" Query

Here's how SQLite finds the latest checkpoint — the most frequent operation in any LangGraph workflow:

```sql
SELECT thread_id, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata
FROM checkpoints
WHERE thread_id = ? AND checkpoint_ns = ?
ORDER BY checkpoint_id DESC
LIMIT 1
```

It queries *all* checkpoints for a thread, sorts them by ID in descending order, and grabs the first result. For SQLite with a small dataset on a local disk, this is fine. But for a distributed database where queries cost money (measured in [Request Units](https://learn.microsoft.com/en-us/azure/cosmos-db/request-units)) and data could be spread across partitions... this pattern is expensive.

**This is where the adaptation story begins.**

---

## The Adaptation: SQLite → Cosmos DB

Adapting a relational checkpoint saver to a document database isn't a line-by-line translation. It requires rethinking data modeling, access patterns, consistency guarantees, and performance characteristics from the ground up.

Here's every major adaptation I made, and *why*.

<!-- 🖼️ IMAGE: A transformation/metamorphosis graphic — left side shows a SQLite database icon morphing/transforming into Azure Cosmos DB icon on the right, with arrows and connecting lines showing the key changes listed below. Style: clean tech illustration with gradients. Dimensions: 1100×450px -->

![transformation-diagram](placeholder-transformation.png)

---

### Adaptation 1: Tables → Documents (Data Modeling)

**SQLite** stores data in *rows* within *tables*. Each checkpoint is a row with typed columns.

**Cosmos DB** stores data as *JSON documents* within *containers*. There are no tables, no columns — just documents with properties.

#### The SQLite Row

```
| thread_id | checkpoint_ns | checkpoint_id | type | checkpoint (BLOB) | metadata (BLOB) |
|-----------|--------------|---------------|------|-------------------|-----------------|
| "user-1"  | ""           | "abc123"     | "..."| <binary>          | <binary>        |
```

#### The Cosmos DB Document

```json
{
    "id": "checkpoint$user-1$$abc123",
    "partition_key": "checkpoint$user-1$",
    "thread_id": "user-1",
    "checkpoint_ns": "",
    "checkpoint_id": "abc123",
    "parent_checkpoint_id": "xyz789",
    "type": "msgpack",
    "checkpoint": "base64-encoded-string...",
    "metadata": { "step": 3, "source": "loop" },
    "metadata_type": "json",
    "metadata_blob": "base64-encoded-string..."
}
```

#### Key Differences

| Aspect | SQLite | Cosmos DB |
|--------|--------|----------|
| **Primary Key** | Composite `(thread_id, ns, id)` | Single `id` field — composite key encoded as `checkpoint$thread$ns$id` |
| **Binary Data** | Native `BLOB` columns | Base64-encoded strings (JSON doesn't support binary) |
| **Partition Key** | Not applicable | `/partition_key` — determines data distribution |
| **Metadata** | Single BLOB | **Dual storage**: searchable JSON + high-fidelity Base64 blob |

That last point — **dual metadata storage** — is worth explaining. SQLite stores metadata as a single BLOB and deserializes it when reading. But Cosmos DB supports SQL queries over JSON properties. I wanted to enable `list()` with metadata filters (e.g., "find all checkpoints where `metadata.source == 'input'`") without deserializing every document.

So I store metadata twice:
1. **`metadata`** (JSON) — the raw dict, queryable via Cosmos DB SQL
2. **`metadata_blob`** (Base64 string) — the serialized bytes, for perfect fidelity during deserialization

This way, you get the best of both worlds: fast filtered queries *and* lossless checkpoint recovery.

```python
# From __init__.py — the dual metadata storage
data = {
    "id": key,
    "partition_key": partition_key,
    # ...
    "metadata": metadata,          # Searchable JSON
    "metadata_type": md_type,
    "metadata_blob": metadata_b64  # High-fidelity blob
}
```

---

### Adaptation 2: `ORDER BY DESC LIMIT 1` → The Tip Document Pattern

This is the single most important architectural change in the entire project.

<!-- 🖼️ IMAGE: A split diagram. Left side: "SQLite Approach" showing a stack of checkpoint rows with an arrow scanning through them top-to-bottom, labeled "ORDER BY ... LIMIT 1 (O(n) scan)". Right side: "Cosmos DB Approach" showing a small "Tip" document pointing directly to the latest checkpoint, labeled "Direct read by ID (O(1))". Use contrasting colors — red/slow for left, green/fast for right. Dimensions: 1100×500px -->

![tip-document-pattern](placeholder-tip-document.png)

#### The SQLite way

```sql
SELECT * FROM checkpoints
WHERE thread_id = ? AND checkpoint_ns = ?
ORDER BY checkpoint_id DESC
LIMIT 1
```

This works because SQLite can sort rows efficiently with local indexes. The cost is negligible.

#### The Cosmos DB problem

In Cosmos DB, the equivalent query would be:

```sql
SELECT TOP 1 * FROM c
WHERE c.thread_id = @thread_id AND c.checkpoint_ns = @ns
ORDER BY c.checkpoint_id DESC
```

This requires:
- A **composite index** on `(thread_id, checkpoint_ns, checkpoint_id DESC)` — which you'd have to configure manually
- A **query** — which costs more Request Units (RU) than a point-read
- **Sorting within the partition** — still `O(n)` relative to checkpoint count

#### The Tip Document solution

Instead, I introduced a **Tip Document** — a lightweight pointer that always references the latest checkpoint:

```json
{
    "id": "tip$user-1$",
    "partition_key": "checkpoint$user-1$",
    "thread_id": "user-1",
    "checkpoint_ns": "",
    "checkpoint_id": "abc123",
    "type": "tip"
}
```

**Reading the latest checkpoint is now two point-reads:**

1. Read the Tip Document by its deterministic ID → get `checkpoint_id`
2. Read the actual checkpoint by its ID → get the full state

<!-- 🖼️ IMAGE: A step diagram showing the two-step process: Step 1: "Read Tip (O(1))" with an arrow pointing to a small Tip document. Step 2: "Read Checkpoint (O(1))" with an arrow from Tip to the full checkpoint. Total: "2 × O(1) = O(1)". Clean flowchart style. Dimensions: 800×350px -->

![tip-read-flow](placeholder-tip-read-flow.png)

Point-reads in Cosmos DB are the cheapest possible operation — guaranteed **1 RU** each, with **single-digit millisecond latency**, regardless of container size. No indexes needed. No sorting. No scanning.

Here's the implementation:

```python
def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
    thread_id = config["configurable"]["thread_id"]
    checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
    checkpoint_id = get_checkpoint_id(config)

    partition_key = _make_cosmosdb_checkpoint_key(thread_id, checkpoint_ns, '')

    if checkpoint_id:
        # Direct access — the caller knows which checkpoint they want
        key = _make_cosmosdb_checkpoint_key(thread_id, checkpoint_ns, checkpoint_id)
        item = self.container.read_item(item=key, partition_key=partition_key)
    else:
        # Tip Document optimization — O(1) lookup of the latest checkpoint
        tip_key = _make_cosmosdb_tip_key(thread_id, checkpoint_ns)
        tip_item = self.container.read_item(item=tip_key, partition_key=partition_key)
        latest_checkpoint_id = tip_item["checkpoint_id"]

        key = _make_cosmosdb_checkpoint_key(
            thread_id, checkpoint_ns, latest_checkpoint_id
        )
        item = self.container.read_item(item=key, partition_key=partition_key)
```

**Compare this to the original SQLite code:**

```python
# SQLite: Query + sort + limit
cur.execute(
    "SELECT ... FROM checkpoints "
    "WHERE thread_id = ? AND checkpoint_ns = ? "
    "ORDER BY checkpoint_id DESC LIMIT 1",
    (thread_id, checkpoint_ns),
)
value = cur.fetchone()
```

Same result. Fundamentally different execution path. The Cosmos DB version scales to *millions* of checkpoints per thread with zero performance degradation.

---

### Adaptation 3: SQL Transactions → Transactional Batches

In SQLite, the `put()` method simply inserts a row. The connection's transaction ensures atomicity.

In Cosmos DB, saving a checkpoint requires **two** document operations:
1. **Upsert** the checkpoint document
2. **Upsert** the Tip Document to point to the new checkpoint

If operation 1 succeeds but operation 2 fails, you have a phantom checkpoint — saved but invisible. If operation 2 succeeds but operation 1 fails, the Tip points to nothing.

<!-- 🖼️ IMAGE: A "danger zone" diagram showing what happens without atomicity: Split-brain scenario where Tip points to Checkpoint B but only Checkpoint A exists. Red warning styling. Then show the solution: both operations in a single atomic batch. Green "safe" styling. Dimensions: 1000×500px -->

![atomic-batch-diagram](placeholder-atomic-batch.png)

#### The solution: Cosmos DB Transactional Batches

Cosmos DB supports **transactional batches** — a set of operations on documents within the *same logical partition* that execute atomically. All succeed, or all roll back.

```python
def put(self, config, checkpoint, metadata, new_versions):
    # ... serialize checkpoint and metadata ...

    # Checkpoint document
    data = {
        "id": key,
        "partition_key": partition_key,
        "checkpoint": checkpoint_b64,
        "metadata": metadata,
        "metadata_blob": metadata_b64,
        # ...
    }

    # Tip Document — always points to the latest
    tip_data = {
        "id": tip_key,
        "partition_key": partition_key,  # Same partition = atomic!
        "checkpoint_id": checkpoint_id,
        "type": "tip"
    }

    # Both operations execute as ONE atomic unit
    batch_operations = [
        ("upsert", (data,), {}),
        ("upsert", (tip_data,), {})
    ]

    self.container.execute_item_batch(batch_operations, partition_key)
```

**The critical insight:** The Tip Document lives in the *same partition* as the checkpoints. This is by design — Cosmos DB transactional batches only work within a single logical partition. By co-locating them, we get atomicity for free.

---

### Adaptation 4: Single Table → Dual-Partition Strategy

In SQLite, checkpoints and writes live in separate *tables* but the same *database file*. Queries to either table are equally fast.

In Cosmos DB, everything lives in *one container*, but the **partition key** determines which logical partition a document belongs to. Documents in different partitions are physically separated and can be on different servers.

I designed a **dual-partition strategy** that separates checkpoints from writes:

```
┌─────────────────────────────────────────────────────┐
│                   Cosmos DB Container                │
│                                                      │
│  ┌──────────────────────────────────┐               │
│  │  Partition: checkpoint$user-1$   │               │
│  │                                   │               │
│  │  📄 tip$user-1$                  │  (Tip)        │
│  │  📄 checkpoint$user-1$$abc123    │  (Checkpoint) │
│  │  📄 checkpoint$user-1$$def456    │  (Checkpoint) │
│  └──────────────────────────────────┘               │
│                                                      │
│  ┌──────────────────────────────────┐               │
│  │  Partition: writes$user-1$       │               │
│  │                                   │               │
│  │  📝 writes$user-1$$abc123$t1$0   │  (Write)      │
│  │  📝 writes$user-1$$abc123$t1$1   │  (Write)      │
│  │  📝 writes$user-1$$abc123$t2$0   │  (Write)      │
│  └──────────────────────────────────┘               │
│                                                      │
└─────────────────────────────────────────────────────┘
```

#### Why split them?

1. **Performance isolation**: Writes can be bursty. An agent making 20 parallel tool calls generates 20 write documents in rapid succession. If these were in the checkpoint partition, they'd compete with the fast point-reads of `get_tuple()`.

2. **Cost optimization**: Cosmos DB charges RU/s (Request Units per second) and throttles at the partition level. Separating workloads means bursty writes don't cause checkpoint reads to be throttled.

3. **Clean deletion**: When you delete a thread, you can target each partition independently using `delete_all_items_by_partition_key()` — a server-side bulk operation that's *dramatically* cheaper than item-by-item deletion.

```python
def _make_cosmosdb_writes_partition_key(thread_id: str, checkpoint_ns: str) -> str:
    """Writes get their own partition: writes$thread-1$namespace"""
    return COSMOSDB_KEY_SEPARATOR.join(["writes", thread_id, checkpoint_ns])
```

---

### Adaptation 5: BLOB Columns → Base64 Encoding

SQLite has native support for `BLOB` (Binary Large Object) columns. You can store raw bytes directly:

```python
# SQLite — binary data stored directly
cur.execute("INSERT INTO checkpoints ... VALUES (?, ?, ?, ?, ?, ?)",
    (thread_id, ns, id, parent_id, type_, checkpoint_bytes))
```

Cosmos DB documents are JSON. JSON has no binary type. Every piece of binary data needs to be encoded as a string.

I chose **Base64 encoding**, which increases size by ~33% but is universally supported and safe for storage:

```python
# Cosmos DB — binary serialized to Base64
type_, checkpoint_bytes = self.serde.dumps_typed(checkpoint)
checkpoint_b64 = base64.b64encode(checkpoint_bytes).decode('utf-8')

# Reading back
c_bytes = base64.b64decode(item["checkpoint"])
checkpoint = self.serde.loads_typed((c_type, c_bytes))
```

<!-- 🖼️ IMAGE: A visual showing the encoding pipeline: "Python Object → Serialize (msgpack) → Raw Bytes → Base64 Encode → JSON String → Cosmos DB Document". Use a horizontal pipeline/conveyor-belt visual style with icons at each stage. Dimensions: 1100×250px -->

![serialization-pipeline](placeholder-serialization-pipeline.png)

This adds a tiny bit of CPU overhead, but it's negligible compared to network latency. And it keeps our documents as standard JSON — no tricks, no hacks.

---

### Adaptation 6: Thread Lock → Concurrency via Document Semantics

SQLite uses a **threading lock** to ensure only one thread accesses the database at a time:

```python
# SQLite saver — manual thread safety
class SqliteSaver:
    def __init__(self, conn):
        self.lock = threading.Lock()

    @contextmanager
    def cursor(self, transaction=True):
        with self.lock:  # Only one caller at a time
            self.setup()
            cur = self.conn.cursor()
            try:
                yield cur
            finally:
                if transaction:
                    self.conn.commit()
                cur.close()
```

In Cosmos DB, there's no shared file. Each operation is an HTTP request. **Concurrency is handled by the database itself**, through document-level optimistic concurrency and partition-level batch atomicity.

My implementation doesn't need a lock for regular operations. But for writes (`put_writes`), I needed to handle a subtle concurrency case:

**What happens when two parallel tasks try to create the same write document?**

```python
def _execute_batch_with_fallback(self, batch_ops, pk, is_upsert):
    """Execute batch with graceful conflict handling."""
    try:
        self.container.execute_item_batch(batch_ops, pk)
    except CosmosBatchOperationError as e:
        if not is_upsert and e.status_code == 409:
            # Batch failed due to conflict — retry each item individually
            for op in batch_ops:
                item = op[1][0]
                try:
                    self.container.create_item(item)
                except CosmosHttpResponseError as e2:
                    if e2.status_code == 409:
                        pass  # Already exists — safe to skip
                    else:
                        raise
        else:
            raise
```

The design here is intentional:

- **System channels** (like `__start__`) use **upsert** — last write wins, idempotent
- **Stream channels** use **create** — if it already exists, we get a `409 Conflict`, which we skip

This gives us exactly the same semantics as SQLite's `INSERT OR REPLACE` and `INSERT OR IGNORE`, but expressed through HTTP status codes instead of SQL constraints.

---

### Adaptation 7: `?` Placeholders → `@` Parameters + Injection Prevention

SQLite uses `?` positional placeholders for parameterized queries. Cosmos DB uses named `@parameters`.

But the adaptation goes deeper than just syntax. In SQLite, you typically trust the schema — the columns are fixed and defined at table creation. In Cosmos DB, metadata filters query into *arbitrary JSON paths*, which opens the door to injection attacks if you're not careful.

I built a multi-layer defense:

<!-- 🖼️ IMAGE: A shield/security layers diagram showing three concentric layers: outer layer "Regex Validation (^[a-zA-Z0-9_.-]+$)", middle layer "Bracket Notation Escaping", inner layer "Parameterized Values (@param)". Security-themed, professional style with a lock icon. Dimensions: 800×500px -->

![security-layers-diagram](placeholder-security-layers.png)

#### Layer 1: Key Validation

Every metadata filter key is validated against a strict regex before it touches any query:

```python
_FILTER_PATTERN = re.compile(r"^[a-zA-Z0-9_.-]+$")

def _validate_filter_key(key: str) -> None:
    if not _FILTER_PATTERN.match(key):
        raise ValueError(
            f"Invalid filter key: '{key}'. "
            "Filter keys must contain only alphanumeric characters, "
            "underscores, dots, and hyphens."
        )
```

#### Layer 2: Bracket Notation

Cosmos DB SQL uses bracket notation for property access — which sidesteps any interpretation of special characters in key names:

```python
# "source.type" becomes c.metadata["source"]["type"]
path_parts = query_key.split(".")
path = "".join([f'["{p}"]' for p in path_parts])
predicates.append(f"c.metadata{path} = {param_name}")
```

#### Layer 3: Parameterized Values

All values are passed as `@parameters`, never interpolated into the query string:

```python
parameters.append({"name": "@metadata_0", "value": query_value})
```

**The SQLite version** relies on `?` placeholders (which are also safe), but doesn't need key validation because column names are fixed at schema time. In a document database with dynamic schemas, you need to be more defensive.

---

### Adaptation 8: Connection String → Azure Identity

SQLite connections are simple — a file path:

```python
# SQLite
conn = sqlite3.connect("checkpoints.sqlite")
saver = SqliteSaver(conn)
```

For Cosmos DB, I designed two authentication paths:

```python
# Development: connection key
with CosmosDBSaver.from_conn_info(
    endpoint="https://account.documents.azure.com:443/",
    credential="your-key-here",
    database_name="langgraph",
    container_name="checkpoints"
) as checkpointer:
    ...

# Production: Azure Identity (keyless)
from azure.identity import DefaultAzureCredential

with CosmosDBSaver.from_conn_info(
    endpoint="https://account.documents.azure.com:443/",
    credential=DefaultAzureCredential(),  # No keys in code!
    database_name="langgraph",
    container_name="checkpoints"
) as checkpointer:
    ...
```

If you pass `None` for `credential`, the library automatically uses `DefaultAzureCredential()`:

```python
@classmethod
@contextmanager
def from_conn_info(cls, *, endpoint, database_name, container_name, credential=None):
    if not credential:
        credential = DefaultAzureCredential()

    client = CosmosClient(endpoint, credential=credential)
    database = client.get_database_client(database_name)
    container = database.get_container_client(container_name)
    yield cls(container)
```

This means **zero secrets in code** when running on Azure — Managed Identity handles everything.

---

## The Full Picture: A Side-by-Side Comparison

Here's every adaptation at a glance:

<!-- 🖼️ IMAGE: A comprehensive comparison table/infographic — styled as a modern dark-mode card layout with two columns: "SQLite (Before)" and "Cosmos DB (After)". Each row is a category (Data Model, Latest Checkpoint, Consistency, etc.) with visual icons. Professional infographic style. Dimensions: 1200×800px -->

![full-comparison-infographic](placeholder-full-comparison.png)

| Area | SQLite (`SqliteSaver`) | Cosmos DB (`CosmosDBSaver`) |
|------|----------------------|---------------------------|
| **Data Model** | Two tables with typed columns | JSON documents with structured IDs |
| **Get Latest** | `ORDER BY DESC LIMIT 1` (scan) | Tip Document → point-read (O(1)) |
| **Binary Storage** | Native BLOB columns | Base64-encoded strings |
| **Atomicity** | SQL transactions + thread lock | Transactional batches (same partition) |
| **Data Isolation** | Separate tables | Dual-partition strategy |
| **Concurrency** | `threading.Lock()` | Optimistic concurrency + batch fallback |
| **Query Params** | `?` positional | `@named` + key validation + bracket notation |
| **Authentication** | File path | Azure Identity / connection key |
| **Scaling** | Single process | Globally distributed, auto-scale |
| **Thread Deletion** | `DELETE WHERE thread_id = ?` | `delete_all_items_by_partition_key()` |
| **Async Support** | Separate `aiosqlite` package | Built-in `AsyncCosmosDBSaver` with sync bridges |

---

## Full Async Support: Not an Afterthought

The SQLite saver has a separate async variant (`AsyncSqliteSaver`) that depends on the `aiosqlite` third-party package. It's a community library that wraps SQLite in an async interface.

For Cosmos DB, async is **first-class**. The Azure Cosmos DB SDK provides a native async client (`azure.cosmos.aio`). I built `AsyncCosmosDBSaver` as a full-featured async implementation — not a wrapper:

```python
from langgraph_checkpoint_cosmos.aio import AsyncCosmosDBSaver

async with AsyncCosmosDBSaver.from_conn_info(
    endpoint="https://your-account.documents.azure.com:443/",
    credential="your-key",
    database_name="langgraph",
    container_name="checkpoints"
) as checkpointer:
    graph = builder.compile(checkpointer=checkpointer)

    config = {"configurable": {"thread_id": "user-456"}}
    result = await graph.ainvoke(1, config)
```

But I went a step further. LangGraph's `BaseCheckpointSaver` requires *both* sync and async methods. So `AsyncCosmosDBSaver` also provides **synchronous bridge methods** that transparently call the async versions:

```python
class AsyncCosmosDBSaver(BaseCheckpointSaver[str]):
    def get_tuple(self, config):
        """Sync bridge — calls async aget_tuple under the hood."""
        # Uses asyncio event loop to run the async version
        ...

    async def aget_tuple(self, config):
        """Native async implementation."""
        # Uses azure.cosmos.aio.ContainerProxy
        ...
```

This means you can use `AsyncCosmosDBSaver` in *any* context — fully async FastAPI handlers, mixed sync/async LangGraph subgraphs, or traditional synchronous code.

---

## In Practice: Using It with LangGraph

Let me show you how this looks in a real-world scenario — a conversational AI agent with multi-turn memory:

<!-- 🖼️ IMAGE: A workflow diagram showing: User Message → LangGraph StateGraph → Agent Node → Tool Calls → Response, with a "Cosmos DB Checkpointer" component on the side connected to every step showing state being saved/loaded. Use a modern flow diagram style with rounded boxes. Dimensions: 1100×500px -->

![langgraph-workflow-diagram](placeholder-langgraph-workflow.png)

```python
from langgraph.graph import StateGraph, MessagesState, START, END
from langgraph_checkpoint_cosmos import CosmosDBSaver

# Define a simple chatbot graph
def chatbot(state: MessagesState):
    # Your LLM logic here
    return {"messages": [ai_response]}

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

# Plug in Cosmos DB persistence
with CosmosDBSaver.from_conn_info(
    endpoint="https://my-account.documents.azure.com:443/",
    credential="my-key",
    database_name="langgraph",
    container_name="checkpoints"
) as checkpointer:
    graph = builder.compile(checkpointer=checkpointer)

    # Every invocation with the same thread_id builds on previous state
    config = {"configurable": {"thread_id": "conversation-42"}}

    # Turn 1
    graph.invoke({"messages": [("user", "Hi, I'm Ram")]}, config)

    # Turn 2 — the agent remembers Turn 1
    graph.invoke({"messages": [("user", "What's my name?")]}, config)

    # Turn 3 — even after a server restart, state persists
    state = graph.get_state(config)
    print(state.values)  # Full conversation history
```

**Nothing changes in your graph code.** The only difference is swapping `SqliteSaver` for `CosmosDBSaver`. That's the beauty of the `BaseCheckpointSaver` abstraction — and exactly why I built strictly against it.

---

## Setting Up Azure Cosmos DB

The setup is intentionally minimal:

<!-- 🖼️ IMAGE: A 4-step setup guide graphic. Step 1: "Create Account" (Azure portal icon). Step 2: "Create Database" (database icon). Step 3: "Create Container" (container icon with /partition_key label). Step 4: "Enable Partition Delete" (optional, settings gear icon). Clean numbered-steps style with Azure blue color scheme. Dimensions: 1100×300px -->

![cosmos-setup-steps](placeholder-cosmos-setup.png)

1. **Create an Azure Cosmos DB account** — NoSQL API
2. **Create a database** — e.g., `langgraph`
3. **Create a container** — partition key must be `/partition_key`
4. *(Optional)* **Enable "Delete All Items By Partition Key"** — for efficient `delete_thread()` support

No schema definitions. No index configuration. No migrations. The library handles everything at the document level.

### Environment Variables (for testing)

```bash
export COSMOS_DB_ENDPOINT="https://your-account.documents.azure.com:443/"
export COSMOS_DB_KEY="your-primary-key"
export COSMOS_DB_NAME="langgraph"
export COSMOS_DB_CONTAINER="checkpoints"
```

---

## Document Key Design: The Hidden Architecture

One thing I haven't fully shown yet is the **key generation system** — the underlying naming convention that makes everything work. This is the glue that enables point-reads, partition isolation, and predictable IDs.

```python
# Composite document IDs — uses '$' as separator
COSMOSDB_KEY_SEPARATOR = "$"

# Checkpoint: checkpoint$thread-1$namespace$abc123
def _make_cosmosdb_checkpoint_key(thread_id, checkpoint_ns, checkpoint_id):
    return "$".join(["checkpoint", thread_id, checkpoint_ns, checkpoint_id])

# Write: writes$thread-1$namespace$abc123$task-1$0
def _make_cosmosdb_checkpoint_writes_key(thread_id, ns, cp_id, task_id, idx):
    return "$".join(["writes", thread_id, ns, cp_id, task_id, str(idx)])

# Tip: tip$thread-1$namespace
def _make_cosmosdb_tip_key(thread_id, checkpoint_ns):
    return "$".join(["tip", thread_id, checkpoint_ns])

# Writes partition: writes$thread-1$namespace
def _make_cosmosdb_writes_partition_key(thread_id, checkpoint_ns):
    return "$".join(["writes", thread_id, checkpoint_ns])
```

<!-- 🖼️ IMAGE: A key anatomy diagram showing a sample document ID "checkpoint$user-123$$abc456" broken down into labeled segments: [prefix]$[thread_id]$[namespace]$[checkpoint_id], with each segment color-coded and labeled. Technical, blueprint style. Dimensions: 1000×250px -->

![key-anatomy-diagram](placeholder-key-anatomy.png)

Every ID is **deterministic**. Given a `thread_id`, `checkpoint_ns`, and `checkpoint_id`, you can compute the exact document ID and read it directly — no query required.

The `$` separator was chosen because it's:
- URL-safe (important for Cosmos DB REST API)
- Unlikely to appear in user-provided IDs
- Visually clear when debugging

---

## Thread Deletion: Server-Side Bulk Operations

In SQLite, deleting a thread is a simple `DELETE WHERE`:

```sql
DELETE FROM checkpoints WHERE thread_id = ?;
DELETE FROM writes WHERE thread_id = ?;
```

In Cosmos DB, deleting items one-by-one is expensive. Each delete costs RUs, and a busy thread could have hundreds of checkpoint and write documents.

Instead, I use Cosmos DB's **"Delete All Items By Partition Key"** feature — a server-side operation that wipes an entire logical partition in one call:

```python
def delete_thread(self, thread_id):
    # Find all partitions for this thread
    query = """
    SELECT DISTINCT c.partition_key FROM c
    WHERE STARTSWITH(c.partition_key, @cp_prefix)
       OR STARTSWITH(c.partition_key, @wr_prefix)
    """
    params = [
        {"name": "@cp_prefix", "value": f"checkpoint${thread_id}$"},
        {"name": "@wr_prefix", "value": f"writes${thread_id}$"}
    ]

    # Server-side bulk delete — no item-by-item loop
    for item in self.container.query_items(query=query, parameters=params, ...):
        self.container.delete_all_items_by_partition_key(item["partition_key"])
```

This is *orders of magnitude* faster and cheaper than item-level deletion, especially for threads with long conversation histories.

---

## What's Next

This is **v0.1.0** — the foundation is solid, battle-tested, and ready for use. Here's where I'm going next:

| Feature | Description | Status |
|---------|-------------|--------|
| 📊 **TTL Support** | Auto-expire old checkpoints to manage costs | Planned |
| 🔍 **Change Feed** | React to checkpoint changes in real-time | Planned |
| 📈 **OpenTelemetry** | Distributed tracing for agent debugging | Planned |
| 🧪 **Chaos Testing** | Fault injection, partition splits, conflict storms | Planned |
| 🌍 **Multi-Region** | Guidance for geo-distributed agent deployments | Planned |

---

## Try It

```bash
pip install langgraph-checkpoint-cosmos
```

📦 **PyPI**: [langgraph-checkpoint-cosmos](https://pypi.org/project/langgraph-checkpoint-cosmos/)
🐙 **GitHub**: [LangModule/checkpoint-cosmos](https://github.com/LangModule/checkpoint-cosmos)
📄 **License**: MIT

The entire codebase is open source, well-documented, and ready for contributions. If you're building AI agents on Azure, this library will save you weeks of infrastructure work.

Drop a ⭐ on GitHub if it's useful, or open an issue if you've got ideas. I'd love to hear what you're building.

---

<!-- 🖼️ AUTHOR CARD IMAGE: A professional headshot-style card with your name, a brief bio, and social links. Clean, modern, dark-mode card style. Dimensions: 800×200px -->

![author-card](placeholder-author-card.png)

*I'm **Ramacharan Reddy Kasireddy** — I build developer tools for AI engineers. This is my first open-source library, and it was built to solve a real problem I faced scaling LangGraph in production. Follow me for more deep dives into LangGraph, Azure, and production AI infrastructure.*

---

> *If you enjoyed this article, consider sharing it with someone building AI agents. The more people stress-test this library, the better it gets for everyone.* 🚀
