AIdeazz Blog About Portfolio

Three LangGraph Rewrites: Checkpointing Stateful Agents in Production

· by

My first LangGraph agent, a simple job router for a Telegram bot, silently dropped 100% of its tasks for two weeks. The StateSchema was defined as TypedDict, but I was passing a Pydantic model instance. LangGraph's MemorySaver just swallowed the type mismatch, logging no error. The agent appeared to run, but the state never persisted. No jobs were ever picked up by the next node. This was the first of three full rewrites before I shipped a stable, stateful LangGraph agent to production on Oracle Cloud.

The Silent StateSchema Mismatch

The initial agent was meant to take incoming user requests from Telegram, classify them, and route them to specialized LLM agents (Groq for fast, simple queries; Claude 3 Opus for complex, multi-turn tasks). The core idea was a TypedDict for the state:

class AgentState(TypedDict):
    chat_id: int
    user_input: str
    classification: Optional[str]
    response: Optional[str]
    history: List[Tuple[str, str]]

When a new message arrived, I'd initialize this state and pass it to the graph. The problem arose when I tried to enrich the state with Pydantic models representing more complex objects, like a UserSession or JobDetails. Instead of updating AgentState['user_input'] = job_details.description, I was trying to store the JobDetails Pydantic object directly into a field that expected a str.

LangGraph's MemorySaver (and by extension, the SqliteSaver I later tried) performs a shallow copy and update. If the incoming state dict has keys that don't match the TypedDict schema, or if the types of the values for matching keys don't align, it doesn't raise an error. It just ignores the update for the mismatched fields. My job_details Pydantic object, despite being assigned to user_input, was never actually stored. The user_input field remained None or its initial value, effectively discarding the core information needed for routing.

The fix was to strictly serialize Pydantic models to dictionaries or JSON strings before updating the LangGraph state. For example:

# Incorrect:
# state['job_details'] = job_details_pydantic_instance

# Correct:
state['job_details_json'] = job_details_pydantic_instance.model_dump_json()
# Or, if the schema expects a dict:
state['job_details_dict'] = job_details_pydantic_instance.model_dump()

This forced me to treat the AgentState as a strict data transfer object, not a flexible container for arbitrary Python objects.

Checkpoint Corruption and Oracle Object Storage

After fixing the state schema, I moved to a more robust checkpointing mechanism. SqliteSaver is fine for local development, but not for a production agent running on Oracle Cloud Infrastructure (OCI) Container Instances. My agents are stateless in terms of their compute, but stateful in terms of their graph execution. I needed external persistence.

I initially tried to use OCI Object Storage directly. LangGraph's MemorySaver can be extended, and I wrote a custom OCISaver that would serialize the state to JSON and upload it to an S3-compatible bucket on OCI Object Storage.

The problem: concurrent updates. LangGraph's default MemorySaver (and its derivatives) assumes a single writer or handles concurrency internally for local files. When multiple instances of my agent (e.g., handling different chat_ids) tried to update their respective checkpoints in Object Storage, I ran into race conditions. An agent would read the state, process it, and then try to write it back. If another agent had written an update in between, the first agent's write would overwrite the newer state, leading to lost turns or corrupted history.

OCI Object Storage, like S3, is eventually consistent. While it offers strong read-after-write consistency for new objects, overwrites can exhibit eventual consistency. More critically, it doesn't provide atomic compare-and-swap operations for object content directly. You can use object versioning, but that just creates new versions; it doesn't prevent overwrites or provide a mechanism for optimistic locking without additional logic.

My custom OCISaver looked something like this:

class OCISaver(BaseCheckpointSaver):
    def __init__(self, bucket_name: str, namespace: str, object_storage_client):
        self.bucket_name = bucket_name
        self.namespace = namespace
        self.client = object_storage_client

    def get_tuple(self, config: RunnableConfig) -> Optional[Tuple[Checkpoint, int]]:
        thread_id = config["configurable"]["thread_id"]
        try:
            response = self.client.get_object(self.namespace, self.bucket_name, f"{thread_id}.json")
            data = json.loads(response.data.content.decode('utf-8'))
            return Checkpoint(**data), data["v"] # Assuming 'v' for version
        except Exception: # Object not found or other error
            return None

    def put_tuple(self, config: RunnableConfig, checkpoint: Checkpoint) -> None:
        thread_id = config["configurable"]["thread_id"]
        # This is where the race condition happens
        self.client.put_object(
            self.namespace, self.bucket_name, f"{thread_id}.json",
            json.dumps(checkpoint.dict()).encode('utf-8')
        )

The put_tuple method was the culprit. It was a blind write.

The solution was to introduce a proper database for checkpointing. I chose Oracle Autonomous Database (ADB) Serverless, specifically its JSON Document Store capabilities. Each thread_id (which maps to a user's chat session) became a document. ADB provides transactional consistency. I could perform a read-modify-write operation within a transaction, ensuring that if two agents tried to update the same thread_id's checkpoint, one would succeed and the other would either wait or fail, allowing for retry logic.

This required a full rewrite of the BaseCheckpointSaver implementation to use SQL UPDATE statements with a WHERE clause checking the version number, or leveraging ADB's native JSON update capabilities.

# Simplified conceptual example for ADB
class ADBSaver(BaseCheckpointSaver):
    def __init__(self, db_connection_pool):
        self.pool = db_connection_pool

    def get_tuple(self, config: RunnableConfig) -> Optional[Tuple[Checkpoint, int]]:
        thread_id = config["configurable"]["thread_id"]
        with self.pool.acquire() as connection:
            cursor = connection.cursor()
            cursor.execute("SELECT checkpoint_data, version FROM checkpoints WHERE thread_id = :1", [thread_id])
            row = cursor.fetchone()
            if row:
                return Checkpoint(**json.loads(row[0])), row[1]
            return None

    def put_tuple(self, config: RunnableConfig, checkpoint: Checkpoint) -> None:
        thread_id = config["configurable"]["thread_id"]
        current_version = checkpoint.v # LangGraph's internal version
        with self.pool.acquire() as connection:
            cursor = connection.cursor()
            # Attempt to update only if the version matches (optimistic locking)
            cursor.execute(
                """
                UPDATE checkpoints
                SET checkpoint_data = :1, version = :2
                WHERE thread_id = :3 AND version = :4
                """,
                [json.dumps(checkpoint.dict()), current_version + 1, thread_id, current_version]
            )
            if cursor.rowcount == 0:
                # This means another process updated it already. Handle conflict.
                # For LangGraph, this typically means reloading the state and retrying the step.
                raise CheckpointConflictError(f"Checkpoint for {thread_id} updated by another process.")
            connection.commit()
            # If it's an insert (first time), handle that too
            # INSERT ... ON CONFLICT DO UPDATE ... or separate INSERT logic

This pattern, using optimistic locking with a version column, is critical for any shared mutable state in a distributed system.

The One Pattern: Explicit State Transitions and Agent Orchestration

Even with robust checkpointing, my multi-agent system was still brittle. The core issue was implicit state transitions and tightly coupled agent logic. An agent would decide the next step, and the graph would just execute it. If an agent made a bad decision (e.g., misclassified a query, or got stuck in a loop), the entire graph would follow.

My agents were:
1. Router Agent: Classifies user input (Groq).
2. Search Agent: Performs web search if needed (Claude 3 Haiku).
3. Summarizer Agent: Summarizes search results (Groq).
4. Response Agent: Generates final response (Claude 3 Opus).

The initial graph was a series of conditional edges based on the classification field in the state. If classification == "search", go to search_agent. If search_agent failed, it would just return an empty result, and summarizer_agent would get nothing, leading to a poor final response. There was no explicit error handling or retry mechanism within the graph structure itself.

The third rewrite introduced a more explicit orchestration pattern:


Supervisor Agent: A dedicated LLM (Groq, for speed) that acts as a meta-agent, reviewing the current state and deciding the next node to execute*. This agent is always the target of a conditional edge.
Tool-like Agents: Each specialized agent (Router, Search, Summarizer, Response) is treated more like a tool. It takes the current state, performs its specific task, and returns its result* and a suggested next_step to the Supervisor.

The AgentState evolved to include:

class AgentState(TypedDict):
    chat_id: int
    user_input: str
    current_task: str # e.g., "classify", "search", "summarize", "respond"
    task_output: Optional[str] # Output of the last completed task
    error_message: Optional[str] # If a task failed
    history: List[Tuple[str, str]]
    # ... other fields for search results, classification, etc.

The graph structure became:

graph = StateGraph(AgentState)

graph.add_node("router", router_agent)
graph.add_node("search", search_agent)
graph.add_node("summarize", summarize_agent)
graph.add_node("respond", response_agent)
graph.add_node("supervisor", supervisor_agent) # The new orchestrator

graph.set_entry_point("router") # Start with classification

# All agents return to the supervisor
graph.add_edge("router", "supervisor")
graph.add_edge("search", "supervisor")
graph.add_edge("summarize", "supervisor")
graph.add_edge("respond", "supervisor")

# Supervisor decides next step
graph.add_conditional_edges(
    "supervisor",
    lambda state: state["next_step"], # Supervisor's output dictates next node
    {
        "router": "router",
        "search": "search",
        "summarize": "summarize",
        "respond": "respond",
        "end": END,
        "error": "error_handler_node" # Or a dedicated error handling path
    }
)

The supervisor_agent's prompt is crucial. It takes the full AgentState and is instructed to output one of the predefined next_step values. For example:

You are an AI orchestrator. Review the current state of the conversation and the last task's output.
Decide the next logical step.
Current State: {state}
Last Task Output: {state['task_output']}
Possible next steps: router, search, summarize, respond, end, error

If the user input needs classification, choose 'router'.
If a search query was generated and needs execution, choose 'search'.
If search results are available and need summarization, choose 'summarize'.
If a final response is ready to be generated, choose 'respond'.
If the conversation is complete, choose 'end'.
If there was an error in the last step, choose 'error'.

Your output MUST be one of these words.

This pattern made the graph significantly more robust. The supervisor could:

This explicit orchestration, combined with transactional checkpointing and strict state schema adherence, finally led to a stable, production-ready multi-agent system. It's more verbose, but the clarity of control flow and error handling is invaluable when debugging live agents.

Frequently Asked Questions

Q: Why not just use a single large LLM for orchestration instead of a supervisor agent?
A: A dedicated, smaller LLM (like Groq's Llama 3 8B) for orchestration is faster and cheaper than a large model. It focuses solely on control flow decisions, reducing the chance of hallucination on the logic of the pipeline, while larger models (Claude 3 Opus) handle complex content generation.

Q: What specific database did you use for checkpointing on OCI?
A: Oracle Autonomous Database (ADB) Serverless, configured for JSON Document Store. It provides ACID transactions and scales automatically, which is critical for unpredictable agent load.

Q: How do you handle schema evolution for AgentState in production?
A: We use Pydantic for AgentState internally, which allows for default values and optional fields. When deploying a new version, we ensure backward compatibility by adding new fields as optional or with defaults. For breaking changes, a migration script is required to update existing checkpoint documents in ADB.

Q: What's the typical latency overhead of the supervisor agent?
A: With Groq, the supervisor agent adds about 100-200ms per decision, which is acceptable for most conversational agents. This is a small price for increased stability and debuggability compared to the potential for agents getting stuck in loops or failing silently.

Q: How do you monitor agent performance and identify issues in this setup?
A: Each node logs its entry, exit, and any errors, along with the chat_id and current_task. We push these logs to OCI Logging Analytics. The supervisor_agent's decisions are also logged, providing a clear trace of the agent's thought process and control flow, making it easier to pinpoint where an agent went off-track.

— Elena Revicheva · AIdeazz · Portfolio