Managing Agent State in Long-Running Workflows

Managing Agent State in Long-Running Workflows

When building AI agents that perform complex, multi-step tasks, you quickly run into a fundamental problem: agents don't complete their work in a single request. They think, act, observe, and repeat — sometimes across minutes, hours, or even days. Between each of those steps, the agent needs to remember where it is, what it has learned, and what it plans to do next. That memory is called agent state, and managing it well is one of the most important architectural decisions you will make.

What Is Agent State?

Agent state is the persistent context that defines an agent's situation at any point in time. It includes the conversation history, the current step in a workflow, intermediate results, tool outputs, configuration, and any metadata the agent needs to resume work after an interruption. Think of it as a snapshot that, if captured fully, would let you reconstruct the agent's exact position in its task and continue as if nothing had happened.

In a short-lived workflow — say, a single chat completion — state is trivial. It lives in memory for the duration of the request and disappears when the response is returned. But in a long-running workflow, the agent may pause to wait for external events, hit rate limits, crash, or simply take too long to finish in one process. In all those cases, state must survive beyond a single execution.

Why State Management Matters

Poor state management leads to a cascade of problems. Agents lose track of progress and repeat work. Users see inconsistent behavior when they return to a task. Crashes wipe out hours of computation. Costs balloon because the agent reprocesses context it already handled. And debugging becomes nearly impossible when you cannot inspect what the agent was doing at the moment something went wrong.

Good state management solves these problems by making the agent resumable, observable, and durable. A resumable agent can be stopped and started again without losing progress. An observable agent exposes its internal state for monitoring and debugging. A durable agent survives infrastructure failures because its state is stored outside the process.

Core Components of Agent State

Before diving into implementation, it helps to decompose agent state into its logical components. Most long-running agent workflows need to track the following:

• Conversation history: The full sequence of messages exchanged between the user, the agent, and any tools.
• Workflow step: An identifier or index indicating where the agent is in its multi-step plan.
• Intermediate results: Outputs from tool calls, sub-tasks, or reasoning steps that later steps depend on.
• Execution metadata: Retry counts, timestamps, token usage, and error logs.
• Configuration: Model parameters, tool definitions, and user preferences that should remain constant across steps.
• External references: Pointers to large artifacts stored outside the state object, such as file paths or blob storage URLs.

Approaches to Storing Agent State

In-Memory State

The simplest approach is to keep state in a Python dictionary or dataclass that lives in the process. This works for prototyping and for workflows that complete within a single request, but it fails the moment the process restarts or the workflow spans multiple requests.

from dataclasses import dataclass, field
from typing import Any

@dataclass
class AgentState:
    messages: list[dict] = field(default_factory=list)
    current_step: int = 0
    intermediate_results: dict[str, Any] = field(default_factory=dict)
    retry_count: int = 0

state = AgentState()
state.messages.append({"role": "user", "content": "Research the latest LLM benchmarks."})

Database-Backed State

For production workflows, state should live in a database. Relational databases like PostgreSQL work well when state has a structured schema. Document stores like MongoDB or DynamoDB are a good fit when state is a flexible JSON blob. Redis is excellent for short-lived state with fast access. The key principle is that the database — not the process — is the source of truth.

import json
import sqlite3
from datetime import datetime

def save_state(conn: sqlite3.Connection, run_id: str, state: dict) -> None:
    conn.execute(
        """INSERT OR REPLACE INTO agent_runs (run_id, state, updated_at)
           VALUES (?, ?, ?)""",
        (run_id, json.dumps(state), datetime.utcnow().isoformat())
    )
    conn.commit()

def load_state(conn: sqlite3.Connection, run_id: str) -> dict | None:
    row = conn.execute(
        "SELECT state FROM agent_runs WHERE run_id = ?", (run_id,)
    ).fetchone()
    if row is None:
        return None
    return json.loads(row[0])

Checkpoint-Based State

Checkpointing means saving state at well-defined points in the workflow — typically after each step completes. This gives you granular recovery: if step five fails, you can resume from the step four checkpoint rather than starting over. Frameworks like LangGraph build checkpointing into their execution model, but you can implement it yourself with a simple pattern.

def run_workflow(run_id: str, state: dict, steps: list) -> dict:
    for index, step in enumerate(steps):
        state["current_step"] = index
        try:
            result = step(state)
            state["intermediate_results"][step.__name__] = result
            save_checkpoint(run_id, state)
        except Exception as exc:
            state["retry_count"] += 1
            state["last_error"] = str(exc)
            save_checkpoint(run_id, state)
            if state["retry_count"] >= 3:
                raise
            return run_workflow(run_id, state, steps[index:])
    return state

Designing a State Schema

A well-designed state schema is the foundation of a maintainable agent. The schema should be explicit, versioned, and serializable. Explicit schemas prevent silent bugs where a missing field causes a crash three steps later. Versioning lets you evolve the schema without breaking in-flight workflows. Serializability ensures the state can be stored and transmitted without loss.

Using Pydantic is a popular choice because it gives you validation, serialization, and documentation in one package.

from pydantic import BaseModel, Field
from typing import Literal

class Message(BaseModel):
    role: Literal["user", "assistant", "tool"]
    content: str
    tool_call_id: str | None = None

class WorkflowState(BaseModel):
    version: int = 1
    run_id: str
    messages: list[Message] = Field(default_factory=list)
    current_step: int = 0
    plan: list[str] = Field(default_factory=list)
    results: dict[str, str] = Field(default_factory=dict)
    status: Literal["pending", "running", "paused", "completed", "failed"] = "pending"
    error: str | None = None
    token_usage: int = 0

    def to_storage(self) -> str:
        return self.model_dump_json()

    @classmethod
    def from_storage(cls, data: str) -> "WorkflowState":
        return cls.model_validate_json(data)

Handling Long Contexts

One of the biggest challenges in long-running workflows is context growth. Every step adds messages and tool outputs, and eventually the context exceeds the model's window or becomes prohibitively expensive. State management must include strategies for keeping context manageable without losing important information.

Summarization

A common technique is to periodically summarize older messages and replace them with a compact summary. This preserves the gist of earlier interactions while freeing up context budget for new steps.

def maybe_summarize(state: WorkflowState, llm_client, threshold: int = 20) -> None:
    if len(state.messages) <= threshold:
        return

    old_messages = state.messages[:-10]
    recent_messages = state.messages[-10:]

    summary_prompt = "Summarize the following conversation, preserving key facts and decisions:\n"
    for msg in old_messages:
        summary_prompt += f"{msg.role}: {msg.content}\n"

    summary = llm_client.complete(summary_prompt)
    state.messages = [
        Message(role="assistant", content=f"[Summary of earlier steps]: {summary}"),
        *recent_messages,
    ]

Externalizing Large Artifacts

Tool outputs can be enormous — full web pages, database dumps, or generated files. Storing these directly in state bloats the database and slows down serialization. Instead, store large artifacts in blob storage and keep only a reference in state.

import hashlib

def store_artifact(content: str, storage_client) -> str:
    key = hashlib.sha256(content.encode()).hexdigest()[:16]
    storage_client.put(f"artifacts/{key}.txt", content)
    return f"artifact://{key}"

def retrieve_artifact(reference: str, storage_client) -> str:
    key = reference.replace("artifact://", "")
    return storage_client.get(f"artifacts/{key}.txt")

Making Workflows Resumable

Resumability means that at any point, you can stop the workflow and restart it from exactly where it left off. This requires two things: a persistent checkpoint after every meaningful step, and a runner that can inspect the checkpoint and decide which step to execute next.

from enum import Enum

class StepName(str, Enum):
    PLAN = "plan"
    RESEARCH = "research"
    ANALYZE = "analyze"
    REPORT = "report"
    DONE = "done"

STEP_ORDER = [StepName.PLAN, StepName.RESEARCH, StepName.ANALYZE, StepName.REPORT]

def resume(run_id: str, state_store, step_handlers: dict) -> WorkflowState:
    state = state_store.load(run_id)
    if state is None:
        raise ValueError(f"No workflow found for run_id={run_id}")

    if state.status in ("completed", "failed"):
        return state

    state.status = "running"
    state_store.save(run_id, state)

    current_index = state.current_step
    for index in range(current_index, len(STEP_ORDER)):
        step_name = STEP_ORDER[index]
        handler = step_handlers[step_name]

        state.current_step = index
        state_store.save(run_id, state)

        try:
            handler(state)
        except Exception as exc:
            state.status = "failed"
            state.error = str(exc)
            state_store.save(run_id, state)
            return state

    state.status = "completed"
    state.current_step = len(STEP_ORDER)
    state_store.save(run_id, state)
    return state

Concurrency and State Consistency

When multiple agents or multiple instances of the same agent run concurrently, state consistency becomes critical. Two common problems arise: lost updates, where one writer overwrites another's changes, and race conditions, where two steps read stale state and make conflicting decisions.

Optimistic locking is a lightweight solution. Each state record carries a version number. When a writer saves, it includes the version it read. If the version in storage has advanced, the save is rejected and the writer must reload and retry.

def save_with_optimistic_lock(conn, run_id: str, state: dict, expected_version: int) -> bool:
    new_version = expected_version + 1
    cursor = conn.execute(
        """UPDATE agent_runs
           SET state = ?, version = ?, updated_at = ?
           WHERE run_id = ? AND version = ?""",
        (json.dumps(state), new_version, datetime.utcnow().isoformat(),
         run_id, expected_version)
    )
    conn.commit()
    return cursor.rowcount == 1

def update_state_safely(conn, run_id: str, transform_fn, max_retries: int = 5) -> dict:
    for attempt in range(max_retries):
        row = conn.execute(
            "SELECT state, version FROM agent_runs WHERE run_id = ?", (run_id,)
        ).fetchone()
        if row is None:
            raise ValueError("Run not found")

        state = json.loads(row[0])
        version = row[1]
        new_state = transform_fn(state)

        if save_with_optimistic_lock(conn, run_id, new_state, version):
            return new_state

    raise RuntimeError(f"Failed to update state after {max_retries} retries")

Observability and Debugging

State is not just for the agent — it is also for the humans who build and operate the system. A well-structured state object makes it easy to answer questions like: what step is this run on? How many tokens has it consumed? What did the last tool call return? Where did it fail?

Logging state transitions at each checkpoint creates an audit trail that is invaluable for debugging. Even better, storing the full state history — not just the latest snapshot — lets you replay a workflow step by step.

def save_checkpoint(run_id: str, state: WorkflowState, conn) -> None:
    conn.execute(
        """INSERT INTO state_history (run_id, step, state, timestamp)
           VALUES (?, ?, ?, ?)""",
        (run_id, state.current_step, state.to_storage(),
         datetime.utcnow().isoformat())
    )
    conn.execute(
        """INSERT OR REPLACE INTO agent_runs (run_id, state, updated_at)
           VALUES (?, ?, ?)""",
        (run_id, state.to_storage(), datetime.utcnow().isoformat())
    )
    conn.commit()

def get_state_history(run_id: str, conn) -> list[dict]:
    rows = conn.execute(
        "SELECT step, state, timestamp FROM state_history WHERE run_id = ? ORDER BY id",
        (run_id,)
    ).fetchall()
    return [{"step": r[0], "state": json.loads(r[1]), "timestamp": r[2]} for r in rows]

Best Practices

• Make state explicit. Use typed schemas rather than ad-hoc dictionaries. The cost of defining a schema is tiny compared to the cost of debugging untyped state.
• Checkpoint after every meaningful step. The more granular your checkpoints, the less work you lose on failure and the easier debugging becomes.
• Keep state serializable. Avoid storing file handles, database connections, or other non-serializable objects in state. Store references instead.
• Version your schema. Include a version field in state and write migration logic for each version bump. This prevents old in-flight workflows from breaking when you deploy changes.
• Externalize large data. Keep state objects small by storing large artifacts in blob storage and keeping only references in state.
• Use optimistic locking for concurrent writes. This prevents lost updates without the overhead of distributed locks.
• Separate state from logic. Step handlers should be pure functions that take state and return updated state. This makes them testable and replayable.
• Plan for failure from the start. Every step should be idempotent or at least safe to retry. Assume that any step can fail and that the workflow will be resumed.
• Log state transitions. An audit trail of state changes is the single most useful debugging tool in a long-running agent system.

Putting It All Together

Here is a minimal but complete example that ties together the concepts we have covered: a typed state schema, a database-backed store with checkpointing, a resumable runner, and a simple multi-step workflow.

import json
import sqlite3
from datetime import datetime
from pydantic import BaseModel, Field
from typing import Literal, Callable

# --- State schema ---
class Message(BaseModel):
    role: Literal["user", "assistant", "tool"]
    content: str

class AgentState(BaseModel):
    version: int = 1
    run_id: str
    messages: list[Message] = Field(default_factory=list)
    current_step: int = 0
    results: dict[str, str] = Field(default_factory=dict)
    status: Literal["pending", "running", "completed", "failed"] = "pending"
    error: str | None = None

# --- State store ---
class StateStore:
    def __init__(self, db_path: str):
        self.conn = sqlite3.connect(db_path)
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS agent_runs (
                run_id TEXT PRIMARY KEY,
                state TEXT NOT NULL,
                updated_at TEXT NOT NULL
            )
        """)
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS state_history (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                run_id TEXT NOT NULL,
                step INTEGER NOT NULL,
                state TEXT NOT NULL,
                timestamp TEXT NOT NULL
            )
        """)
        self.conn.commit()

    def save(self, state: AgentState) -> None:
        data = state.model_dump_json()
        ts = datetime.utcnow().isoformat()
        self.conn.execute(
            "INSERT OR REPLACE INTO agent_runs VALUES (?, ?, ?)",
            (state.run_id, data, ts)
        )
        self.conn.execute(
            "INSERT INTO state_history (run_id, step, state, timestamp) VALUES (?, ?, ?, ?)",
            (state.run_id, state.current_step, data, ts)
        )
        self.conn.commit()

    def load(self, run_id: str) -> AgentState | None:
        row = self.conn.execute(
            "SELECT state FROM agent_runs WHERE run_id = ?", (run_id,)
        ).fetchone()
        if row is None:
            return None
        return AgentState.model_validate_json(row[0])

# --- Runner ---
def run_workflow(
    run_id: str,
    store: StateStore,
    steps: list[Callable[[AgentState], None]],
) -> AgentState:
    state = store.load(run_id)
    if state is None:
        state = AgentState(run_id=run_id, status="running")
    elif state.status in ("completed", "failed"):
        return state
    else:
        state.status = "running"

    store.save(state)

    for index in range(state.current_step, len(steps)):
        state.current_step = index
        store.save(state)
        try:
            steps[index](state)
        except Exception as exc:
            state.status = "failed"
            state.error = str(exc)
            store.save(state)
            return state

    state.status = "completed"
    state.current_step = len(steps)
    store.save(state)
    return state

# --- Example steps ---
def step_plan(state: AgentState) -> None:
    state.results["plan"] = "1. Gather data 2. Analyze 3. Report"
    state.messages.append(Message(role="assistant", content="Plan created."))

def step_gather(state: AgentState) -> None:
    state.results["data"] = "sample data payload"
    state.messages.append(Message(role="tool", content="Data gathered."))

def step_report(state: AgentState) -> None:
    summary = f"Report based on: {state.results['data']}"
    state.results["report"] = summary
    state.messages.append(Message(role="assistant", content=summary))

# --- Usage ---
if __name__ == "__main__":
    store = StateStore("agent.db")
    steps = [step_plan, step_gather, step_report]
    final = run_workflow("run-001", store, steps)
    print(f"Status: {final.status}")
    print(f"Results: {final.results}")

Conclusion

Managing agent state in long-running workflows is not a peripheral concern — it is the architectural backbone that determines whether your agent system is reliable, debuggable, and scalable. By treating state as an explicit, versioned, serializable object, checkpointing after every meaningful step, externalizing large artifacts, and using concurrency controls like optimistic locking, you transform fragile prototypes into production-grade systems. The patterns in this tutorial are framework-agnostic and work equally well whether you are building on LangGraph, CrewAI, or a custom agent runtime. The investment you make in state management early on will pay dividends every time your agent encounters a failure, a timeout, or a user who returns to a task hours later expecting it to pick up right where it left off.