Build a Crash-Safe AI Execution Engine with PostgreSQL Checkpoints

The worst time for an AI worker to crash is not before it starts. It is after it has decided to do something expensive.

The model returns a tool call. Your worker begins executing it. Then the container is killed, the node reboots, the database failover happens, or the LLM request times out after the side effect has already occurred. When the job restarts, you have two bad options: repeat the work and risk duplicate side effects, or skip ahead and risk losing state.

Most AI applications fail here because they treat execution state as memory. The current step lives in a Python object. The model response lives in a local variable. The tool-call status lives in a queue message that gets acknowledged too early. That works in a demo. It breaks in production.

A crash-safe AI execution engine treats state differently. Every meaningful transition is written to durable storage before the engine moves on. PostgreSQL is a good place to build this because it gives you transactions, row locks, SKIP LOCKED, JSONB, advisory locks, and durable commits in one system.

To be clear: this is not about PostgreSQL’s internal WAL CHECKPOINT command. This is about application-level checkpoints: durable rows that record what your AI engine has already done, what it is about to do, and what it must not do again.

TL;DR

  • Treat AI executions as long-running state machines, not short-lived function calls.
  • Store runs, steps, checkpoints, and side-effect commands in PostgreSQL.
  • Use one database transaction per state transition.
  • Use FOR UPDATE SKIP LOCKED with leases to hand work to workers.
  • Checkpoint LLM output before executing tools.
  • Make external side effects idempotent with command IDs.
  • Run a reconciler to recover stuck leases and expired workers.
  • Use PostgreSQL when durability and transactional consistency matter more than extreme queue throughput.

📋 Table of Contents


The Real Problem

AI workflows are deceptively long-running.

A simple support agent may need to:

  1. Load customer context.
  2. Call a model.
  3. Parse structured output.
  4. Call an internal tool.
  5. Write a result.
  6. Call the model again.
  7. Send a final response.

Each of those steps can fail in a different way. The model can time out. The tool can succeed but return a network error. The worker can be killed during deployment. The database can restart. The user can cancel the run.

If your execution state lives only in process memory, a crash turns a recoverable failure into a correctness problem.

The goal is not to make the LLM deterministic. That is not what this layer controls. The goal is to make the execution engine deterministic enough that a crash does not cause duplicate refunds, missing audit records, lost context, or silently repeated work.

PostgreSQL is not a queue in the same sense as Redis or RabbitMQ, and it is not a workflow engine in the same sense as Temporal. But for many AI products, it is the right checkpoint store because your application already trusts it with durable business data.


1. Design for the Crash That Happens Mid-Action

Scenario: Your worker calls a payment provider, then crashes before it records that the payment was made. On restart, the engine retries the step and charges the customer twice.

Why it matters: The obvious fix — “just retry the step” — is unsafe when steps have external side effects. Retry logic assumes the step either did not happen or can safely happen again. AI workflows often violate both assumptions.

A crash-safe engine starts with a failure model.

The important failures are:

  • worker process killed,
  • machine reboot,
  • deployment rollout,
  • database connection reset,
  • LLM request timeout,
  • tool call succeeds but response is lost,
  • tool call fails after partially succeeding externally,
  • run is cancelled while a worker is still active.

Those failures force a few architectural decisions.

First, execution must be resumable. If a run has completed step 6, a restart should begin at step 6 or 7, not step 1.

Second, side effects must be idempotent. If a tool call may be retried, the external system needs a stable command ID or idempotency key.

Third, state transitions must be durable. If the engine believes a step is complete, that belief must survive a crash.

Fourth, work claiming must be safe under concurrency. Two workers must not lease the same run unless the system explicitly allows that.

This is where PostgreSQL fits well. A committed transaction gives you a durable point of agreement: “this run has reached this state.” If the process dies immediately after the commit, the state remains. If the process dies before the commit, the transition did not happen.

That is the foundation.

🔍 Why this matters: You are not trying to achieve perfect exactly-once execution across arbitrary external services. That is usually impossible. You are trying to achieve safe at-least-once execution with idempotent side effects and durable checkpoints.


2. Make the Run a Row and the Plan a State Machine

Scenario: Your agent loop is a Python function. The current step, model messages, and intermediate results live in local variables. When the worker restarts, there is nothing to recover.

Why it matters: If the execution is not represented as durable data, you cannot resume it. You can only restart it.

Start with a table for runs:

CREATE TABLE ai_runs (
    id uuid PRIMARY KEY,
    kind text NOT NULL,
    status text NOT NULL CHECK (status IN (
        'pending',
        'running',
        'waiting',
        'succeeded',
        'failed',
        'cancelled'
    )),
    input jsonb NOT NULL,
    state jsonb NOT NULL DEFAULT '{}'::jsonb,
    next_step integer NOT NULL DEFAULT 0,
    attempts integer NOT NULL DEFAULT 0,
    max_attempts integer NOT NULL DEFAULT 5,
    worker_id text,
    lease_expires_at timestamptz,
    created_at timestamptz NOT NULL DEFAULT now(),
    updated_at timestamptz NOT NULL DEFAULT now()
);

CREATE INDEX ai_runs_claim_idx
ON ai_runs (status, lease_expires_at, created_at)
WHERE status IN ('pending', 'waiting');

The important fields are not fancy:

  • status is the current state machine position.
  • next_step tells the engine where to resume.
  • state holds structured progress data.
  • attempts and max_attempts bound retries.
  • worker_id and lease_expires_at support safe worker leasing.

Then add checkpoints:

CREATE TABLE ai_checkpoints (
    run_id uuid NOT NULL REFERENCES ai_runs(id) ON DELETE CASCADE,
    step_index integer NOT NULL,
    step_name text NOT NULL,
    output jsonb NOT NULL,
    created_at timestamptz NOT NULL DEFAULT now(),
    PRIMARY KEY (run_id, step_index)
);

A checkpoint says: “This step produced this output. The engine may use it without redoing the expensive work.”

This design works because the run is no longer a hidden object inside a worker. It is a row that can be claimed, inspected, resumed, cancelled, and audited.

💡 Practical note: Do not put huge blobs into checkpoint rows by default. If a model response is large, store a reference to object storage and keep metadata in PostgreSQL. The checkpoint should be durable and queryable, not a memory dump.


3. Wrap Every State Transition in a Transaction

Scenario: Your worker updates a run to running, calls an external API, then updates the run to waiting. If the external call succeeds but the final update fails, the database no longer reflects reality.

Why it matters: Crash safety comes from precise transaction boundaries. If you mutate state outside a transaction, or if you hold a transaction open while calling slow external services, you create windows where the database and the outside world disagree.

The basic rule:

  • Database state changes happen inside transactions.
  • External side effects happen outside the transaction that changes engine state.
  • External side effects are recorded as commands before or after execution, depending on your delivery semantics.

For example, advancing a run should be transactional:

import json

def advance_run(conn, run_id: str, next_step: int, state_patch: dict) -> None:
    with conn.transaction():
        conn.execute(
            """
            UPDATE ai_runs
            SET
                next_step = %(next_step)s,
                state = state || %(state_patch)s::jsonb,
                status = 'waiting',
                updated_at = now()
            WHERE id = %(run_id)s
            """,
            {
                "run_id": run_id,
                "next_step": next_step,
                "state_patch": json.dumps(state_patch),
            },
        )

This is simple, but the structure matters. The engine does not say, “I think I am at step 4.” It commits a row that says, “This run is now at step 4.”

The mistake to avoid is wrapping an external HTTP call inside the same transaction:

# Bad pattern
with conn.transaction():
    conn.execute("UPDATE ai_runs SET status = 'running' WHERE id = %s", run_id)
    external_api.charge_customer(...)
    conn.execute("UPDATE ai_runs SET status = 'waiting' WHERE id = %s", run_id)

That transaction may hold locks while the external service is slow. Worse, if the transaction rolls back after the external call succeeds, your database says the call did not happen while the outside world disagrees.

Use transactions for engine state. Use idempotent commands for external effects.


4. Claim Work with SKIP LOCKED and Leases

Scenario: You have three workers consuming AI runs. A run becomes available. Without careful locking, two workers can pick it up at the same time and both start executing.

Why it matters: Duplicate execution is not just inefficient. If the run calls tools, sends messages, or writes business data, concurrent duplicate execution can corrupt results.

PostgreSQL has a queue-friendly pattern: FOR UPDATE SKIP LOCKED. It lets workers claim rows without blocking each other.

A claim query can look like this:

WITH claimable AS (
    SELECT id
    FROM ai_runs
    WHERE status IN ('pending', 'waiting')
      AND (lease_expires_at IS NULL OR lease_expires_at < now())
      AND attempts < max_attempts
    ORDER BY created_at
    LIMIT 1
    FOR UPDATE SKIP LOCKED
)
UPDATE ai_runs r
SET
    status = 'running',
    worker_id = %(worker_id)s,
    attempts = attempts + 1,
    lease_expires_at = now() + interval '5 minutes',
    updated_at = now()
FROM claimable c
WHERE r.id = c.id
RETURNING
    r.id,
    r.kind,
    r.input,
    r.state,
    r.next_step;

In Python with psycopg:

import psycopg
from psycopg.rows import dict_row

def claim_next_run(conn: psycopg.Connection, worker_id: str):
    with conn.transaction():
        cur = conn.execute(
            """
            WITH claimable AS (
                SELECT id
                FROM ai_runs
                WHERE status IN ('pending', 'waiting')
                  AND (lease_expires_at IS NULL OR lease_expires_at < now())
                  AND attempts < max_attempts
                ORDER BY created_at
                LIMIT 1
                FOR UPDATE SKIP LOCKED
            )
            UPDATE ai_runs r
            SET
                status = 'running',
                worker_id = %(worker_id)s,
                attempts = attempts + 1,
                lease_expires_at = now() + interval '5 minutes',
                updated_at = now()
            FROM claimable c
            WHERE r.id = c.id
            RETURNING
                r.id,
                r.kind,
                r.input,
                r.state,
                r.next_step
            """,
            {"worker_id": worker_id},
        )

        return cur.fetchone()

Use a dict row factory if you want named access:

conn = psycopg.connect(dsn, row_factory=dict_row)

The lease is important. If a worker crashes while a run is marked running, the run would otherwise stay stuck forever. The lease_expires_at column lets a reconciler reclaim it later.

Leases introduce a subtle problem: long-running steps can outlive the lease. If your LLM call or tool call can take longer than the lease period, the worker needs to extend the lease.

def extend_lease(conn, run_id: str, worker_id: str) -> None:
    with conn.transaction():
        conn.execute(
            """
            UPDATE ai_runs
            SET lease_expires_at = now() + interval '5 minutes'
            WHERE id = %(run_id)s
              AND worker_id = %(worker_id)s
              AND status = 'running'
            """,
            {"run_id": run_id, "worker_id": worker_id},
        )

⚠️ Gotcha: Do not make the lease too short. If the lease expires while a worker is still executing a dangerous side effect, another worker may claim the run and repeat the step. Prefer idempotent side effects plus lease extension over very aggressive lease timeouts.


5. Checkpoint Model Outputs Before You Execute Tools

Scenario: Your agent calls a model. The model returns a tool call: refund_order. Before the worker executes the refund, the process crashes. On restart, do you call the model again?

Why it matters: Model calls are expensive, slow, and sometimes rate-limited. Worse, a second model call may return a different plan. If the first response was meaningful, you want to recover it, not replace it.

Checkpoint the model output before acting on it.

import json

def save_checkpoint(
    conn,
    run_id: str,
    step_index: int,
    step_name: str,
    output: dict,
) -> None:
    with conn.transaction():
        conn.execute(
            """
            INSERT INTO ai_checkpoints (
                run_id,
                step_index,
                step_name,
                output
            )
            VALUES (
                %(run_id)s,
                %(step_index)s,
                %(step_name)s,
                %(output)s::jsonb
            )
            ON CONFLICT (run_id, step_index)
            DO UPDATE SET
                step_name = EXCLUDED.step_name,
                output = EXCLUDED.output
            """,
            {
                "run_id": run_id,
                "step_index": step_index,
                "step_name": step_name,
                "output": json.dumps(output),
            },
        )

Then load the checkpoint before calling the model:

def get_checkpoint(conn, run_id: str, step_index: int):
    cur = conn.execute(
        """
        SELECT output
        FROM ai_checkpoints
        WHERE run_id = %(run_id)s
          AND step_index = %(step_index)s
        """,
        {"run_id": run_id, "step_index": step_index},
    )

    row = cur.fetchone()

    if row is None:
        return None

    return row["output"]

A simplified execution step becomes:

def execute_model_step(conn, run_id: str, step_index: int, step_name: str, messages: list):
    existing = get_checkpoint(conn, run_id, step_index)

    if existing is not None:
        return existing

    output = call_llm(messages)

    save_checkpoint(conn, run_id, step_index, step_name, output)

    return output

This gives you a useful recovery property. If the worker crashes after checkpointing the model output but before executing the tool, the restarted worker can reuse the saved output.

There is still a gap: if the worker crashes after the model call returns but before the checkpoint commits, the engine may call the model again. That is usually acceptable. Model calls are usually safer to repeat than side effects. But if your provider gives you request IDs or idempotency mechanisms, store them. If not, accept that model calls may be at-least-once.

The important part is that the checkpoint becomes the boundary between “the model proposed this” and “the engine acted on this.”

🧠 The important part: Do not execute tools directly from a transient model response. Persist the response first. The checkpoint is the engine’s source of truth, not the worker’s local variable.


6. Turn Side Effects into Idempotent Commands

Scenario: Your AI agent decides to send an email. The worker sends it, then crashes before recording success. On retry, the engine sends the email again.

Why it matters: External side effects are where crash safety becomes hard. PostgreSQL can make its own state consistent, but it cannot roll back an email, a payment, or a webhook that already reached another system.

The practical pattern is to treat side effects as commands.

Create a command table:

CREATE TABLE ai_commands (
    idempotency_key text PRIMARY KEY,
    run_id uuid NOT NULL REFERENCES ai_runs(id) ON DELETE CASCADE,
    step_index integer NOT NULL,
    command_name text NOT NULL,
    payload jsonb NOT NULL,
    status text NOT NULL CHECK (status IN (
        'pending',
        'succeeded',
        'failed'
    )) DEFAULT 'pending',
    result jsonb,
    created_at timestamptz NOT NULL DEFAULT now(),
    updated_at timestamptz NOT NULL DEFAULT now()
);

CREATE INDEX ai_commands_run_idx
ON ai_commands (run_id, step_index);

The idempotency key should be deterministic. It should come from the run, step, and tool-call identity, not from a random value generated during execution.

For example:

run_8f2c:step_4:tool_call_ab12:send_email

Reserve the command before executing it:

import json

def reserve_command(
    conn,
    run_id: str,
    step_index: int,
    command_name: str,
    payload: dict,
    idempotency_key: str,
):
    with conn.transaction():
        cur = conn.execute(
            """
            INSERT INTO ai_commands (
                idempotency_key,
                run_id,
                step_index,
                command_name,
                payload,
                status
            )
            VALUES (
                %(idempotency_key)s,
                %(run_id)s,
                %(step_index)s,
                %(command_name)s,
                %(payload)s::jsonb,
                'pending'
            )
            ON CONFLICT (idempotency_key) DO NOTHING
            RETURNING status, result
            """,
            {
                "idempotency_key": idempotency_key,
                "run_id": run_id,
                "step_index": step_index,
                "command_name": command_name,
                "payload": json.dumps(payload),
            },
        )

        row = cur.fetchone()

        if row is not None:
            return row

        cur = conn.execute(
            """
            SELECT status, result
            FROM ai_commands
            WHERE idempotency_key = %(idempotency_key)s
            """,
            {"idempotency_key": idempotency_key},
        )

        return cur.fetchone()

Then execute only if needed:

def execute_command(conn, run_id: str, step_index: int, tool_call: dict, gateway):
    idempotency_key = (
        f"{run_id}:step_{step_index}:tool_call_{tool_call['id']}:{tool_call['name']}"
    )

    existing = reserve_command(
        conn,
        run_id=run_id,
        step_index=step_index,
        command_name=tool_call["name"],
        payload=tool_call["arguments"],
        idempotency_key=idempotency_key,
    )

    if existing["status"] == "succeeded":
        return existing["result"]

    result = gateway.execute(
        tool_call["name"],
        tool_call["arguments"],
        idempotency_key=idempotency_key,
    )

    with conn.transaction():
        conn.execute(
            """
            UPDATE ai_commands
            SET
                status = 'succeeded',
                result = %(result)s::jsonb,
                updated_at = now()
            WHERE idempotency_key = %(idempotency_key)s
            """,
            {
                "result": json.dumps(result),
                "idempotency_key": idempotency_key,
            },
        )

    return result

This does not magically give you exactly-once delivery to every external service. If the external service does not honor idempotency keys, a crash between the external call and the success update can still cause repetition. But it gives you the structure to make side effects safe:

  • the command is visible before execution,
  • the command has a stable key,
  • retries can detect prior success,
  • external services can deduplicate using the key,
  • support tooling can inspect pending and failed commands.

For high-risk operations, add another layer. Require approval before dispatching commands like refunds, deletions, or bulk notifications.

🚨 Production warning: If a command triggers a payment, deletion, or customer-visible message, the command ID is not optional. “We’ll just retry” is not a safety strategy.


7. Sweep Stuck Runs with a Reconciler

Scenario: A worker claims a run, sets status = 'running', then dies. The run stays in running forever because no one is advancing it.

Why it matters: Crash safety is not only about avoiding duplicate work. It is also about recovering work that has stopped. A queue without a sweeper eventually accumulates stuck jobs.

Run a separate reconciler process. Its job is to find expired leases and return runs to a retryable state.

WITH expired AS (
    SELECT id
    FROM ai_runs
    WHERE status = 'running'
      AND lease_expires_at < now()
    LIMIT 100
    FOR UPDATE SKIP LOCKED
)
UPDATE ai_runs r
SET
    status = CASE
        WHEN attempts >= max_attempts THEN 'failed'
        ELSE 'waiting'
    END,
    worker_id = NULL,
    lease_expires_at = NULL,
    updated_at = now()
FROM expired e
WHERE r.id = e.id;

This query is safe to run periodically because it also uses FOR UPDATE SKIP LOCKED. It will not fight with active workers over rows that are currently locked.

The reconciler should also handle poisoned runs. Some runs fail repeatedly because the input is invalid, the model refuses, the tool schema is wrong, or a dependency permanently rejects the request. Those should not retry forever.

A useful policy:

  • transient failures: retry with backoff,
  • validation failures: fail immediately,
  • repeated failures: move to dead-letter status,
  • dangerous side-effect failures: alert,
  • expired leases with pending commands: reconcile carefully.

The last one matters. If a run expires while commands are pending, do not blindly restart the step. Check the command table. If a command is pending, determine whether the external system may have processed it. Use the idempotency key. If the external system supports query-by-key, query before retrying.

You can add a retry delay by adding available_at to ai_runs:

ALTER TABLE ai_runs
ADD COLUMN available_at timestamptz NOT NULL DEFAULT now();

Then update the claim query:

AND available_at <= now()

And set backoff when returning a run to waiting:

UPDATE ai_runs
SET
    status = 'waiting',
    available_at = now() + interval '30 seconds',
    worker_id = NULL,
    lease_expires_at = NULL,
    updated_at = now()
WHERE id = %(run_id)s;

This keeps retries from hammering a model provider or external API immediately after a failure.


8. Add an Append-Only Event Trail for Debugging and Replay

Scenario: A user says the assistant did something unexpected. The run row says succeeded. The checkpoint table says the model returned a tool call. But nobody can tell what happened between steps, when the lease expired, or why a command was retried.

Why it matters: Checkpoints tell you where the engine is. Events tell you how it got there. For AI systems, that difference is critical because the model’s behavior is not always obvious from the final state.

Add an append-only event table:

CREATE TABLE ai_events (
    id bigint GENERATED ALWAYS AS IDENTITY PRIMARY KEY,
    run_id uuid NOT NULL REFERENCES ai_runs(id) ON DELETE CASCADE,
    event_type text NOT NULL,
    payload jsonb NOT NULL DEFAULT '{}'::jsonb,
    created_at timestamptz NOT NULL DEFAULT now()
);

CREATE INDEX ai_events_run_idx
ON ai_events (run_id, id);

Log meaningful transitions:

def record_event(conn, run_id: str, event_type: str, payload: dict) -> None:
    with conn.transaction():
        conn.execute(
            """
            INSERT INTO ai_events (run_id, event_type, payload)
            VALUES (%(run_id)s, %(event_type)s, %(payload)s::jsonb)
            """,
            {
                "run_id": run_id,
                "event_type": event_type,
                "payload": json.dumps(payload),
            },
        )

Useful events include:

  • run.created
  • run.claimed
  • run.lease_extended
  • run.checkpoint_saved
  • model.call_started
  • model.call_failed
  • tool.command_reserved
  • tool.command_succeeded
  • tool.command_failed
  • run.reconciled
  • run.cancelled
  • run.succeeded
  • run.failed

Do not treat this table as the primary state machine. The run and checkpoint tables remain the authoritative execution state. The event table is for diagnosis, replay, auditing, and operational dashboards.

It also helps with a subtle AI problem: explaining why the agent did something. If the checkpoint stores the model output and the event trail stores the transition sequence, you can answer questions like:

  • Which step produced the tool call?
  • Was the tool call from the original model response or a retry?
  • Did the worker crash after the command was reserved?
  • Did the reconciler restart the run?
  • Was the command already marked succeeded externally?

Be careful with sensitive data. Model prompts and outputs can contain personal information. If your retention policy or privacy requirements demand it, redact fields before logging, store only hashes, or keep full payloads in encrypted object storage with a pointer in PostgreSQL.

💡 Practical note: The event table will grow. If you expect high volume, plan for partitioning by time or periodic archiving. Append-only tables are useful until they become the largest table in the database.


9. Know Where PostgreSQL Stops Being the Right Tool

PostgreSQL is a strong foundation for this pattern, but it is not universally ideal.

The architecture works especially well when:

  • your application already uses PostgreSQL,
  • AI executions are business-critical but not massive in throughput,
  • you need transactional consistency with business data,
  • you want durable checkpoints without adding a new distributed system,
  • your side effects need strong auditability,
  • your team can operate a worker/reconciler pattern safely.

It becomes less attractive when:

  • you need extremely high-throughput queueing,
  • you need complex distributed saga orchestration,
  • you need built-in durable timers, retries, and workflow versioning,
  • your execution graph is highly dynamic and very large,
  • you need multi-region active-active coordination,
  • your team already has a dedicated workflow platform.

A rough comparison:

ApproachDurabilityOperational complexityThroughputBest use
In-memory stateLowLowHighprototypes, non-critical tasks
Redis queueMedium to high depending on configMediumHighfast job distribution
PostgreSQL checkpoint engineHighMediumMediumtransactional AI workflows
Dedicated workflow engineHighHighMedium to highcomplex long-running orchestration

The PostgreSQL approach gives you a lot with a small surface area. You are not introducing a new source of truth. Your runs, checkpoints, commands, and events live next to your business data.

But if your AI system becomes a large orchestration platform, you may eventually want a dedicated engine. That is not a failure. It means your requirements changed.

Even then, PostgreSQL can remain useful as the durable store for business-level checkpoints and command records.


What I Would Ship

If I were building a crash-safe AI execution engine for a production product, I would start with PostgreSQL and keep the design deliberately boring.

The baseline would look like this:

  1. Every AI execution is a row in ai_runs.
  2. Every meaningful step has a checkpoint row.
  3. Workers claim runs using FOR UPDATE SKIP LOCKED.
  4. Running rows have leases and worker IDs.
  5. Long-running steps extend their leases.
  6. Model outputs are checkpointed before tool execution.
  7. Tool calls become commands with deterministic idempotency keys.
  8. External services must accept those keys for dangerous operations.
  9. A reconciler sweeps expired leases.
  10. An append-only event table records transitions.
  11. Failed runs have bounded retries and explicit dead-letter states.
  12. The team regularly tests crash recovery by killing workers.

That last point matters more than it sounds. The real test is not whether the engine works when the model returns a clean response. The real test is what happens when you kill the worker immediately after the model returns, immediately after the tool call succeeds, or immediately after the checkpoint commits.

A crash-safe AI engine is not built by trusting the model, the network, or the worker process. It is built by assuming all of them will fail, then making the next step obvious from durable state.

PostgreSQL cannot stop your AI worker from crashing. But it can make the crash uninteresting: the run is still there, the checkpoint is still there, the command is still traceable, and the engine knows exactly where to resume.

Post a Comment

Previous Post Next Post