AIdeazz Blog About Portfolio

LangGraph Checkpointing: Three Rewrites to Production Stability

· by

My first LangGraph agent silently discarded every job for weeks. The agent was supposed to process incoming Telegram messages, enrich them with external data, and then route them to a specific LLM (Groq for fast, simple tasks; Claude for complex reasoning). Instead, it would receive a message, acknowledge it, and then… nothing. No error, no log, just a black hole. The problem wasn't in the LLM calls or the routing logic. It was a state schema mismatch, a subtle bug that only surfaced when I tried to re-hydrate the agent's state from a checkpoint.

The Silent State Schema Mismatch

The initial LangGraph state was a simple TypedDict. Something like this:

class AgentState(TypedDict):
    input_message: str
    enriched_data: Optional[Dict]
    llm_response: Optional[str]
    route: Optional[str]

This worked fine for a single-pass execution. The issue arose when I introduced checkpointing. I was using SqliteSaver for local development and PostgresSaver for production on Oracle Cloud. The agent would process the input_message, then call a tool to fetch enriched_data. This enriched_data was sometimes a complex JSON object, sometimes a simple string, depending on the external API response.

My initial thought was to just dump the dict directly into the enriched_data field. Python's dict is flexible. PostgreSQL's JSONB type is also flexible. What could go wrong?

The problem was in the deserialization. When LangGraph tried to load a checkpoint where enriched_data was a string, but the AgentState expected a Dict, it didn't throw an error. It simply assigned None to enriched_data. The subsequent steps, which depended on enriched_data being present, would then silently exit or take a default path, effectively dropping the job. My logs showed the initial message received, but no further processing. It took a full week of debugging, stepping through the SqliteSaver's internal get_state method, to realize that the pydantic validation (or lack thereof in TypedDict for complex types) was the culprit.

The fix was to explicitly define the schema for enriched_data as a str and handle the JSON serialization/deserialization within the agent's nodes.

import json
from typing import Optional, Dict, Any

class AgentState(TypedDict):
    input_message: str
    enriched_data_json: Optional[str] # Store as JSON string
    llm_response: Optional[str]
    route: Optional[str]

# Inside a node:
def enrich_data_node(state: AgentState) -> AgentState:
    input_msg = state["input_message"]
    # Call external API
    raw_data = {"key": "value", "another_key": [1, 2, 3]} # Example
    return {"enriched_data_json": json.dumps(raw_data)}

# Inside another node:
def process_enriched_data_node(state: AgentState) -> AgentState:
    if state["enriched_data_json"]:
        enriched_data = json.loads(state["enriched_data_json"])
        # Process enriched_data (now a dict)
        # ...
    return {}

This explicit serialization ensured that the state could be reliably saved and loaded, regardless of the actual content of enriched_data. The PostgresSaver now correctly stored and retrieved the string, and the agent's internal logic handled the json.loads/json.dumps.

Checkpoint Corruption and Race Conditions

The second major rewrite came after a series of inexplicable agent failures. Sometimes, an agent would just hang. Other times, it would restart from an earlier state, re-processing messages it had already handled. This was happening in production, with multiple Telegram agents running concurrently, each managing its own conversation state.

My initial setup used PostgresSaver with a simple thread_id as the config["configurable"]["thread_id"]. Each Telegram chat had a unique thread_id. The problem was that my agent was sometimes called multiple times for the same thread_id in rapid succession, especially when users sent multiple messages quickly.

LangGraph's PostgresSaver uses optimistic locking. When update_checkpoint is called, it checks the version column. If the version in the database doesn't match the version the saver thinks it has, it raises an IntegrityError. This is good for preventing lost updates. However, my agent wasn't always handling these IntegrityErrors gracefully. Sometimes, the entire process would crash, leaving the checkpoint in an inconsistent state.

The core issue was a race condition. If two concurrent calls to the agent tried to update the same thread_id's state, one would succeed, and the other would fail. The failing one would often leave the agent in a partially updated state or, worse, cause the entire worker process to restart without proper state recovery.

The solution involved two parts:

1. Robust Error Handling and Retries: Wrap the agent invocation in a try-except block specifically for IntegrityError and implement a retry mechanism with exponential backoff. If a checkpoint update fails due to a version mismatch, it means another process updated it. The agent should then reload the latest state from the database and re-attempt its current step.

    from tenacity import retry, wait_exponential, stop_after_attempt, retry_if_exception_type
    from sqlalchemy.exc import IntegrityError

    @retry(wait=wait_exponential(multiplier=1, min=4, max=10), stop=stop_after_attempt(5),
           retry=retry_if_exception_type(IntegrityError))
    def run_agent_with_retries(agent, input_data, config):
        try:
            return agent.invoke(input_data, config)
        except IntegrityError as e:
            print(f"Checkpoint update failed due to race condition: {e}. Retrying...")
            # Crucially, reload the state before retrying
            # LangGraph's invoke will handle loading the latest state if the config is the same
            raise # Re-raise to trigger tenacity retry
    

2. External Locking (for critical sections): While LangGraph handles optimistic locking for its state, certain external operations (like calling a rate-limited API or updating an external database) might need stronger guarantees. For these, I introduced a simple Redis-based distributed lock. Before a critical node executes, it acquires a lock for the thread_id. If it can't acquire the lock, it waits or signals that the operation is already in progress. This prevented multiple agents from, for example, simultaneously fetching the same enriched_data or sending duplicate messages to the user.

    import redis
    import time

    redis_client = redis.Redis(host='localhost', port=6379, db=0)

    def acquire_lock(lock_name, timeout=10):
        end_time = time.time() + timeout
        while time.time() < end_time:
            if redis_client.setnx(lock_name, 1):
                redis_client.expire(lock_name, 30) # Lock expires in 30 seconds
                return True
            time.sleep(0.1)
        return False

    def release_lock(lock_name):
        redis_client.delete(lock_name)

    # Inside a LangGraph node:
    def critical_external_api_call_node(state: AgentState) -> AgentState:
        thread_id = state["thread_id"] # Assuming thread_id is part of state
        lock_name = f"lock:thread:{thread_id}"
        if acquire_lock(lock_name):
            try:
                # Perform critical API call
                print(f"Acquired lock for {thread_id}, calling API...")
                time.sleep(2) # Simulate API call
                return {"api_result": "success"}
            finally:
                release_lock(lock_name)
        else:
            print(f"Could not acquire lock for {thread_id}, another process is working.")
            # Handle case where lock couldn't be acquired (e.g., wait, or return early)
            return {"api_result": "skipped_due_to_lock"}
    

This combination made the multi-agent system significantly more resilient. The retry mechanism handled transient database conflicts, and the external locking prevented race conditions on shared external resources.

The One Pattern That Made Multi-Step Pipelines Stable: Explicit Step Management

My final rewrite, and the one that truly stabilized complex multi-step agents, involved moving away from implicit state transitions and towards explicit step management within the LangGraph state.

Initially, I relied heavily on LangGraph's conditional edges to route between nodes. A node would return a "next_step" string, and the graph would transition accordingly. This worked for simple flows. However, when an agent needed to perform a sequence of actions, sometimes involving human intervention (e.g., "ask user for clarification," "wait for user input"), the state became hard to reason about.

The problem was that the current step was implicit. If an agent crashed and restarted, it would re-evaluate the conditional edges, potentially leading to a different path if external conditions changed or if a previous step's output was ambiguous.

The stable pattern I adopted was to explicitly store the current_step and history of steps in the agent's state.

from typing import List

class AgentState(TypedDict):
    input_message: str
    current_step: str # e.g., "INITIAL", "ENRICHING_DATA", "LLM_CALL", "WAITING_FOR_USER"
    step_history: List[str] # To track the path taken
    # ... other state variables

Each node in the LangGraph would then be responsible for:

1. Checking current_step: Only execute if current_step matches the node's expected state.
2. Updating current_step: Set the current_step to the next logical step upon completion.
3. Appending to step_history: Record the node's execution.

The LangGraph conditional edges then became much simpler: they would primarily check current_step and route to the next node responsible for that step.

Example:

from langgraph.graph import StateGraph, END

# Define nodes
def initial_processing_node(state: AgentState) -> AgentState:
    if state["current_step"] == "INITIAL":
        print("Executing initial processing...")
        return {"current_step": "ENRICHING_DATA", "step_history": state["step_history"] + ["initial_processing"]}
    return {} # No change if not the current step

def enrich_data_node(state: AgentState) -> AgentState:
    if state["current_step"] == "ENRICHING_DATA":
        print("Executing data enrichment...")
        # ... call external API ...
        return {"current_step": "LLM_CALL", "step_history": state["step_history"] + ["enrich_data"]}
    return {}

def llm_call_node(state: AgentState) -> AgentState:
    if state["current_step"] == "LLM_CALL":
        print("Executing LLM call...")
        # ... call Groq/Claude ...
        return {"current_step": "FINAL_RESPONSE", "step_history": state["step_history"] + ["llm_call"]}
    return {}

def final_response_node(state: AgentState) -> AgentState:
    if state["current_step"] == "FINAL_RESPONSE":
        print("Executing final response...")
        return {"current_step": "COMPLETED", "step_history": state["step_history"] + ["final_response"]}
    return {}

# Build the graph
graph_builder = StateGraph(AgentState)
graph_builder.add_node("initial_processing", initial_processing_node)
graph_builder.add_node("enrich_data", enrich_data_node)
graph_builder.add_node("llm_call", llm_call_node)
graph_builder.add_node("final_response", final_response_node)

graph_builder.set_entry_point("initial_processing")

# Conditional edges based on current_step
graph_builder.add_conditional_edges(
    "initial_processing",
    lambda state: state["current_step"],
    {
        "ENRICHING_DATA": "enrich_data",
        "COMPLETED": END, # If for some reason initial processing completes the flow
    },
)
graph_builder.add_conditional_edges(
    "enrich_data",
    lambda state: state["current_step"],
    {
        "LLM_CALL": "llm_call",
        "COMPLETED": END,
    },
)
graph_builder.add_conditional_edges(
    "llm_call",
    lambda state: state["current_step"],
    {
        "FINAL_RESPONSE": "final_response",
        "COMPLETED": END,
    },
)
graph_builder.add_conditional_edges(
    "final_response",
    lambda state: state["current_step"],
    {
        "COMPLETED": END,
    },
)

graph = graph_builder.compile()

# Initial state for a new conversation
initial_state = {"input_message": "Hello", "current_step": "INITIAL", "step_history": []}
# Invoke the agent
# graph.invoke(initial_state, config={"configurable": {"thread_id": "user123"}})

This pattern has several advantages:

Shipping production AI agents with zero VC funding means every line of code needs to be robust. These three lessons, learned through weeks of silent failures and corrupted checkpoints, have been fundamental to building stable, stateful LangGraph agents on Oracle Cloud.

Frequently Asked Questions

Q: Why not just use a database transaction for LangGraph state updates instead of optimistic locking?
A: LangGraph's PostgresSaver does use transactions for its updates. The optimistic locking (version column) is an additional layer to handle concurrent updates to the same checkpoint from different application processes, which a single transaction wouldn't prevent if they started concurrently.

Q: How do you handle schema migrations for AgentState in production?
A: This is tricky. For minor changes, I add new optional fields. For major overhauls, I create a new AgentStateV2 and a migration node. This node, triggered by a current_step like "MIGRATE_STATE", transforms old state schemas to the new one, then updates current_step to proceed with the new schema.

Q: Is SqliteSaver ever production-ready?
A: No. SqliteSaver is for local development and testing only. It's not designed for concurrent access from multiple processes or for high availability. For production, always use PostgresSaver or a custom saver backed by a robust, distributed database.

Q: What's the overhead of storing step_history in the state?
A: For typical conversational agents, the step_history list of strings is small, usually a few dozen entries at most, resulting in negligible storage overhead (a few KB). If steps are very granular or conversations extremely long, consider truncating the history or storing only critical milestones to prevent state bloat.

— Elena Revicheva · AIdeazz · Portfolio