0Pricing
AI Engineering Academy · Lesson

Checkpointing and Task Resumption

Persist agent state at each completed step so a long-running task can be resumed from the last successful checkpoint rather than restarted from scratch after a failure.

Checkpointing and Task Resumption is a free AI Engineering Academy lesson on CoddyKit — lesson 3 of 4. You can read the complete lesson below for free — then practise it hands-on in the browser with a built-in code editor and a 24/7 AI tutor. It is part of the AI Engineering Academy learning path, one of 4 lessons in the course, and your progress syncs across the web and the CoddyKit app.

The Problem with Long-Running Agents

An agent executing a 50-step research task might run for 30 minutes. If it fails at step 47 due to an API timeout or server restart, restarting from scratch wastes all the work done and costs tokens. Checkpointing persists the agent's state at each completed step so the task can resume from the last successful point rather than the beginning. This is essential for any agent task longer than a few minutes.

What Agent State to Persist

An agent's state consists of: the task definition, the completed steps with their tool calls and observations, the current step index, any accumulated results (files written, data collected), and metadata like start time and total token usage. Persist all of this after each step completes. State must be serializable — prefer JSON over Python objects for portability.

from dataclasses import dataclass, field
from typing import List, Any, Optional

@dataclass
class AgentStep:
    step_index: int
    thought: str
    tool_name: str
    tool_args: dict
    observation: str
    tokens_used: int
    completed_at: str

@dataclass
class AgentCheckpoint:
    task_id: str
    task_description: str
    status: str  # 'running', 'completed', 'failed'
    current_step: int
    completed_steps: List[AgentStep] = field(default_factory=list)
    accumulated_results: dict = field(default_factory=dict)
    total_tokens: int = 0
    final_answer: Optional[str] = None

Saving Checkpoints After Each Step

After each successful step, serialize the checkpoint and write it to a persistent store. Use a database or object storage (S3, Redis) rather than local disk so the checkpoint survives server restarts. Include the task_id in the key so you can retrieve a specific task's checkpoint. Also write a step log entry for every step to maintain an audit trail even if the checkpoint file gets corrupted.

import json
import redis
from dataclasses import asdict

redis_client = redis.Redis()

def save_checkpoint(checkpoint: AgentCheckpoint):
    key = f'agent:checkpoint:{checkpoint.task_id}'
    data = json.dumps(asdict(checkpoint), default=str)
    redis_client.set(key, data, ex=86400)  # 24h TTL
    # Also append to step log
    log_key = f'agent:log:{checkpoint.task_id}'
    if checkpoint.completed_steps:
        last = checkpoint.completed_steps[-1]
        redis_client.rpush(log_key, json.dumps(asdict(last), default=str))

def load_checkpoint(task_id: str) -> AgentCheckpoint | None:
    key = f'agent:checkpoint:{task_id}'
    data = redis_client.get(key)
    if data:
        return AgentCheckpoint(**json.loads(data))
    return None

Resuming from a Checkpoint

When resuming a task, load the checkpoint and reconstruct the agent's message history from the completed steps. Start the execution loop from the next uncompleted step index. The model receives its previous thoughts and observations in the message history so it has full context of what has been done, allowing it to continue without repeating work.

async def resume_or_start(task_id: str, task_description: str) -> str:
    checkpoint = load_checkpoint(task_id)
    if checkpoint and checkpoint.status == 'running':
        print(f'Resuming task {task_id} from step {checkpoint.current_step}')
        messages = rebuild_history(checkpoint)
        start_step = checkpoint.current_step
    else:
        print(f'Starting new task {task_id}')
        checkpoint = AgentCheckpoint(task_id=task_id, task_description=task_description, status='running', current_step=0)
        messages = [{'role': 'user', 'content': task_description}]
        start_step = 0
        save_checkpoint(checkpoint)
    return await run_agent_from(checkpoint, messages, start_step)

Rebuilding Message History from Steps

The key to resumption is faithfully reconstructing the message history from the persisted steps. Each completed step corresponds to an assistant message (thought + tool call) and a tool message (observation). Replay all completed steps as messages before continuing, so the model has the same context as if it had never been interrupted.

def rebuild_history(checkpoint: AgentCheckpoint) -> list:
    messages = [{'role': 'user', 'content': checkpoint.task_description}]
    for step in checkpoint.completed_steps:
        # Reconstruct the agent's reasoning message
        messages.append({
            'role': 'assistant',
            'content': f'Thought: {step.thought}\nAction: {step.tool_name}({step.tool_args})'
        })
        # Reconstruct the tool observation
        messages.append({
            'role': 'user',
            'content': f'Observation: {step.observation}'
        })
    return messages

Idempotent Tool Calls

If a step was partially completed before a crash (the tool was called but the observation was not saved), resumption might call the tool again. Design tools to be idempotent: calling them twice with the same arguments produces the same result as calling once. For write operations (creating files, sending emails), use deduplication keys to prevent duplicate effects even if the tool is called multiple times.

async def idempotent_write_file(content: str, path: str, task_id: str, step: int) -> str:
    dedup_key = f'{task_id}:step_{step}:write:{path}'
    if redis_client.exists(dedup_key):
        return f'File {path} already written (dedup key present)'
    with open(path, 'w') as f:
        f.write(content)
    redis_client.set(dedup_key, '1', ex=3600)
    return f'Successfully wrote {len(content)} chars to {path}'

Checkpoint Cleanup and Retention

Checkpoints consume storage and should not accumulate indefinitely. Set a retention policy: delete completed checkpoints after 24 hours, delete failed checkpoints after 7 days (for post-mortem analysis), and never delete running checkpoints automatically. Implement a background cleanup job that runs hourly and purges expired checkpoints according to the policy.

from datetime import datetime, timedelta

RETENTION = {
    'completed': timedelta(hours=24),
    'failed': timedelta(days=7),
    'running': None  # never auto-delete
}

def cleanup_expired_checkpoints():
    now = datetime.utcnow()
    for key in redis_client.scan_iter('agent:checkpoint:*'):
        data = json.loads(redis_client.get(key))
        status = data.get('status', 'running')
        retention = RETENTION.get(status)
        if retention is None:
            continue
        started = datetime.fromisoformat(data.get('started_at', str(now)))
        if now - started > retention:
            redis_client.delete(key)

Checkpointing in LangGraph

LangGraph has native checkpointing support via its MemorySaver and SqliteSaver classes. Attach a checkpointer to your graph and every node execution is automatically persisted. To resume, call graph.invoke with the same thread_id. LangGraph handles history reconstruction and step tracking, so you do not need to implement the checkpointing logic manually.

from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.graph import StateGraph

# Create graph with persistent checkpointer
checkpointer = SqliteSaver.from_conn_string('/tmp/agent_state.db')
graph = StateGraph(AgentState)
graph.add_node('reason', reason_node)
graph.add_node('act', act_node)
# ... add edges ...
app = graph.compile(checkpointer=checkpointer)

# Run with thread_id - LangGraph auto-checkpoints
config = {'configurable': {'thread_id': 'task_abc123'}}
result = await app.ainvoke({'task': 'Research climate change'}, config)

# Resume same thread - LangGraph loads from checkpoint
result = await app.ainvoke({'task': 'Continue'}, config)

Distributed Checkpointing for Parallel Agents

When multiple agents work in parallel on sub-tasks, each agent needs its own checkpoint namespace. Use a hierarchical key structure: parent_task_id:sub_task_id. The parent agent's checkpoint records which sub-tasks have completed and their results. When the parent resumes, it re-uses completed sub-task results from the checkpoint instead of re-running them.

async def parallel_with_checkpoints(parent_id: str, subtasks: list) -> list:
    results = []
    for i, subtask in enumerate(subtasks):
        sub_id = f'{parent_id}:sub_{i}'
        # Check if subtask already completed
        existing = load_checkpoint(sub_id)
        if existing and existing.status == 'completed':
            print(f'Sub-task {i} already done, using cached result')
            results.append(existing.final_answer)
        else:
            result = await run_agent(sub_id, subtask)
            results.append(result)
    return results

Testing Resumption Behavior

Write integration tests that deliberately crash the agent mid-task and verify that resumption produces the correct final result. Simulate a crash by raising an exception at a specific step index. After resumption, check that only the steps after the crash are re-executed (not the ones before). Also test that idempotent tool calls do not produce duplicate side effects when a step is replayed.

import pytest

@pytest.mark.asyncio
async def test_resumption_from_step_3():
    task_id = 'test_resume_001'
    # Run until step 3, then crash
    with pytest.raises(SimulatedCrash):
        await run_agent_crashing_at(task_id, 'Research AI trends', crash_at_step=3)

    checkpoint = load_checkpoint(task_id)
    assert checkpoint.current_step == 3
    assert len(checkpoint.completed_steps) == 3

    # Resume and complete
    result = await resume_or_start(task_id, 'Research AI trends')
    assert result is not None
    # Verify only steps 4+ were re-executed
    assert checkpoint_step_was_not_replayed(task_id, step=0)

Checkpoint Versioning for Schema Changes

When you change the AgentCheckpoint schema (adding or renaming fields), old checkpoints in storage become incompatible. Handle this with checkpoint versioning: add a checkpoint_version field and write migration functions to upgrade old checkpoints to the new schema when they are loaded. This prevents crashes when agents resume after a deployment that changed the checkpoint format.

def load_and_migrate_checkpoint(task_id: str) -> AgentCheckpoint:
    raw = json.loads(redis_client.get(f'agent:checkpoint:{task_id}'))
    version = raw.get('checkpoint_version', '1.0')
    if version == '1.0':
        # Migrate: add new fields added in v2.0
        raw['checkpoint_version'] = '2.0'
        raw['accumulated_results'] = raw.get('accumulated_results', {})
        raw['total_tokens'] = raw.get('total_tokens', 0)
    return AgentCheckpoint(**raw)

Quick Check

Test your understanding of checkpointing and task resumption in agents.

Lesson Recap

In this lesson you learned: checkpointing serializes agent state after each step so long-running tasks survive failures without restarting from scratch, rebuilding message history from persisted steps gives the model full context on resumption, and idempotent tools with deduplication keys prevent duplicate side effects when steps are replayed. Next up we design human-in-the-loop escalation triggers.

Frequently asked questions

Is the “Checkpointing and Task Resumption” lesson free?

Yes — the full text of “Checkpointing and Task Resumption” is free to read here on the web, and the AI Engineering Academy course includes 4 lessons in total. To practise it interactively (a built-in code editor and a 24/7 AI tutor) and unlock the rest of the AI Engineering Academy course, upgrade to CoddyKit PRO.

What will I learn in “Checkpointing and Task Resumption”?

Persist agent state at each completed step so a long-running task can be resumed from the last successful checkpoint rather than restarted from scratch after a failure. You practise AI Engineering Academy with hands-on code you run directly in the browser, and a 24/7 AI tutor answers your questions as you work through the lesson.

Do I need any experience to start AI Engineering Academy?

No prior experience is required. AI Engineering Academy on CoddyKit is structured for beginners through advanced learners; this is — lesson 3 of 4, so you can start here or from the beginning and move at your own pace.

How long does the “Checkpointing and Task Resumption” lesson take?

Most CoddyKit lessons take about 5–10 minutes. Each one is bite-sized and interactive, so you make steady progress and pick up exactly where you left off across the web and the app.

Can I write and run code in this AI Engineering Academy lesson?

Yes. Every AI Engineering Academy lesson includes a built-in code editor, so you write and run real code right in your browser and get instant AI feedback — no local setup required.

All lessons in this course

  1. Classifying Agent Failure Modes
  2. Self-Correction and Reflective Prompting
  3. Checkpointing and Task Resumption
  4. Human-in-the-Loop Escalation
← Back to AI Engineering Academy