Multi-Agent State and Checkpointing with Cosmos DB

Posted on by

Post 4 of 6 on Cosmos DB multi-agent state coordinating what several agents know about the same conversation, without a separate message bus.

Post 3 settled retrieval for a single agent working alone. This post is about what changes once a second agent enters the picture. Coordinating what several agents know about the same conversation turns out to be a different problem from storing and retrieving one agent’s memory, and Cosmos DB multi-agent state ends up resting on two mechanisms this series already covered: hierarchical partitioning from post 2, and change feed from post 1’s retail monitoring callback.

Shared but Separable: What Changes with Multiple Agents

A single agent needs one memory scope. Moreover, a multi-agent system needs two at once: shared memory that every agent can read and write for coordination, and private memory that lets each agent keep its own persona, prompts, and reasoning history separate from the others. Lose the separation, and agents start bleeding into each other’s context. Lose the sharing, and they can’t coordinate at all.

Diagram showing three agents with their own persona connected to shared Cosmos DB multi-agent state
Each agent keeps its own persona and prompts, while all three read and write the same shared memory container.

A triage agent, a product agent, and a specialist agent a common pattern in production multi-agent apps each hold their own scoped state. Still, all three write to the same underlying container, so any of them can pick up where another left off.

LangGraph Checkpointing on Cosmos DB Multi-Agent State

LangGraph’s checkpoint interface persists a graph’s state after every step, and Cosmos DB has more than one implementation of it: langgraph-checkpoint-cosmosdb on PyPI, and the checkpoint saver that ships inside langchain-azure-cosmosdb. Both plug into the same standard LangGraph pattern: compile the graph with a checkpointer, then pass a thread_id on every invocation:

from langgraph. graph import StateGraph
from langgraph_checkpoint_cosmosdb import CosmosDBSaver
checkpointer = CosmosDBSaver(
    endpoint=cosmos_endpoint,
    key=cosmos_key,
    database_name="agentmemory",
    container_name="checkpoints",
)
graph = StateGraph(AgentState)
# add_node / add_edge calls wire up triage -> specialist routing here
app = graph.compile(checkpointer=checkpointer)
config = {"configurable": {"thread_id": "contoso:thread-1234"}}
app.invoke({"messages": [...]}, config=config)

Encode tenantId:threadId into the thread_id string, and the checkpointer’s hierarchical partitioning lines up with the [tenantId, threadId] partition key from post 2 — the same pattern manages per-user, per-session state at scale, this time for graph checkpoints instead of turn-based memory items. Microsoft’s own multi-agent-langgraph sample builds a personal-shopper scenario on exactly this foundation: a triage agent routes requests, and a product agent answers them using retrieval-augmented generation against the same Cosmos DB account.

Change Feed as the Handoff Mechanism

Post 1 covered change feed as the primitive behind a 2023 retail monitoring solution, a new record in Cosmos DB firing a Function that could raise an incident. The same primitive coordinates agent handoffs: one agent writes a turn, a Function listening on the container’s change feed picks it up, and it hands the conversation to whichever agent should act next. No polling loop checks for new work; the write itself is the signal.

Diagram showing change feed triggering a handoff between agents in Cosmos DB multi-agent state
One agent’s write becomes another agent’s signal to act no polling required.

A minimal handoff trigger, using the turn-based schema from post 2:

python

import azure.functions as func
def main(documents: func.DocumentList) -> None:
    for doc in documents:
        if doc.get("targetAgent") == "specialist":
            notify_specialist_agent(doc["threadId"], doc["turnIndex"])

The triage agent sets targetAgent on the turn it writes; the Function reacts to that write and wakes the specialist agent for that thread.

A Second Worked Example: Spring AI for Java Shops

Python and LangGraph aren’t the only path here. Spring AI 2.0 shipped with a Cosmos DB-backed vector store and memory integration for Java, and Microsoft’s multi-agent-spring-ai sample mirrors the LangGraph pattern in Java: multiple agents, one Cosmos DB account, the same shared-but-separable memory shape. Worth a look if the rest of the stack runs on the JVM rather than Python.

Pitfalls

Shared containers without tenant or session isolation. A checkpoint container that mixes every tenant’s graph state leaks context across customers the moment a partition or vector index goes unsharded; the same isolation failure post 3 flagged for vector search is now showing up in agent state instead of retrieved memories. Apply the same [tenantId, threadId] discipline to checkpoints that post 2 applied to turns.

Treating change feed as globally ordered. Change feed guarantees order within a single partition key, not across the whole container. Moreover, a handoff design that assumes “the Function always sees writes in the exact order they happened across every agent” breaks the moment two agents write to different partitions at close to the same time. Design handoffs so each step only depends on ordering within its own thread’s partition, not on a global sequence that Cosmos DB never promised.

Next: Cosmos DB Inside Microsoft Foundry Agent Service

That covers Cosmos DB multi-agent state when you manage the account directly. Post 5 covers the other path: Microsoft Foundry Agent Service’s bring-your-own thread storage, where Cosmos DB still does the work, but Foundry owns the orchestration layer on top of it.

Sources

Leave a Reply