State Management

from typing import Literal
from pydantic import BaseModel

class RunState(BaseModel):
    # identity
    run_id: str
    user_id: str
    # agent state: what the model reads
    messages: list[dict] = []
    plan: list[str] = []
    # workflow state: what the code tracks
    status: Literal["running", "waiting_approval", "done", "failed", "handed_over"] = "running"
    step: int = 0
    cost_rupees: float = 0.0
    pending_approval: dict | None = None
    completed_actions: list[str] = []      # idempotency keys of writes already done

def run(run_id: str) -> RunState:
    state = store.load(run_id)                       # resume, or start fresh
    while state.status == "running" and state.step < MAX_STEPS:
        reply = llm(state.messages, tools=TOOL_DEFS)
        state.messages.append(reply.as_message())

        if not reply.tool_calls:
            state.status = "done"
        for call in reply.tool_calls:
            key = f"{run_id}:{state.step}:{call.name}"       # same key on any retry
            result = execute(call, state.user_id, idempotency_key=key)   # safe to repeat
            if key not in state.completed_actions:
                state.completed_actions.append(key)
            state.messages.append(tool_result(call.id, result))

        state.step += 1
        store.save(state)                            # the checkpoint
    return state

def request_approval(state: RunState, call) -> RunState:
    """Pause the run. Nothing stays running while a person decides."""
    state.status = "waiting_approval"
    state.pending_approval = {"tool": call.name, "input": call.input, "call_id": call.id}
    store.save(state)
    notify_supervisor(state.run_id, state.pending_approval)
    return state

def on_decision(run_id: str, approved: bool, reviewer: str) -> RunState:
    """Called hours later by the approval screen. Loads state and carries on."""
    state = store.load(run_id)
    pending, state.pending_approval = state.pending_approval, None
    result = (execute_pending(pending, state.user_id) if approved
              else f"Declined by {reviewer}. Do not retry; explain to the customer.")
    state.messages.append(tool_result(pending["call_id"], result))
    state.status = "running"
    store.save(state)
    return run(run_id)

TRANSITIONS = {
    "collecting":        {"details_complete": "verifying"},
    "verifying":         {"verified": "deciding", "cannot_verify": "handed_over"},
    "deciding":          {"small_refund": "refunding", "large_refund": "awaiting_approval",
                          "no_refund": "replying"},
    "awaiting_approval": {"approved": "refunding", "declined": "replying"},
    "refunding":         {"refund_ok": "replying", "refund_failed": "handed_over"},
    "replying":          {"sent": "done"},
}

TOOLS_BY_STATE = {
    "collecting": ["get_order"],
    "verifying":  ["get_order", "get_policy"],
    "refunding":  ["refund_item"],          # the only state in which money can move
}

def advance(current: str, event: str) -> str:
    try:
        return TRANSITIONS[current][event]
    except KeyError:
        raise ValueError(f"'{event}' is not allowed in state '{current}'")

Q: Stateless versus stateful agents?

A stateless agent keeps nothing between requests; a stateful one keeps the task's record so it can continue. In production the servers are stateless and the state is in a database keyed by run id, so any server can resume any task.

Q: What is in an agent's state?

The messages and tool results the model reads, the plan, and workflow fields the code tracks: status, step count, cost, pending approvals and the idempotency keys of completed actions.

Q: What is checkpointing?

Saving the full state after each step. A failure then loses at most one step, the run can be resumed on any machine, and old checkpoints can be loaded to inspect or replay.

Q: How does an agent wait two days for an approval?

It does not wait. It saves state with status 'waiting', notifies the approver, and stops. The approval event loads the state, appends the decision as a tool result, and restarts the loop.

Q: The process crashes mid-tool-call. What happens on resume?

The step is run again from the last checkpoint. Because each write carries an idempotency key derived from the run and step, the downstream service performs it once.

Q: Why use a state machine around an agent?

To restrict which tools and transitions are possible at each stage, so forbidden sequences cannot occur whatever the model outputs, and so the whole process can be read, tested and audited.

Q: How is state different from memory?

State is the exact, complete record of one run, kept until it ends. Memory is selected information kept across runs and retrieved by relevance.

Your task. Add checkpointing and resumption to a simple agent loop, then prove it by killing the process mid-task.

Steps

1. Define a RunState class with messages, status, step, and completed actions.
2. Write save(state) and load(run_id) using SQLite or one JSON file per run.
3. Change your loop to load state at the start and save after every step.
4. Give each write action an idempotency key of run id, step and tool name, and make the fake refund tool ignore a key it has seen.
5. Add raise SystemExit inside the loop at step 3, run the task, then remove the line and run again with the same run id.

You are done when

• The second run continues from step 3, not from step 0.
• The fake refund was issued exactly once across both runs.
• You can load the checkpoint for any step and print what the model had seen at that point.

Stretch. Add a waiting_approval status: the loop stops at a refund above ₹500, and a separate script approves it and resumes the run.