How Checkpoint and Restore Work for Agent Persistence in Jido: A Complete Guide

Jido's checkpoint and restore mechanism captures agent state as serializable maps, externalizes plugin data via callbacks, and rehydrates threads during restoration to enable durable agent persistence across system restarts.

The agentjido/jido repository implements a robust persistence layer that allows Elixir agents to survive process restarts and system crashes. Understanding how checkpoint and restore work for agent persistence in Jido is essential for building reliable long-running applications that maintain state integrity across hibernation cycles.

Understanding the Checkpoint and Restore Architecture

Jido orchestrates persistence through three core modules that separate concerns between agent logic, storage operations, and plugin lifecycle management.

Core Modules Involved

• Jido.Agent – Defines the checkpoint/2 and restore/2 callbacks and supplies default implementations via code generation.
• Jido.Persist – The persistence layer that orchestrates checkpoint creation, storage adapter interaction, and agent rehydration including thread re-attachment.
• Jido.Plugin – Provides the behaviour contract for on_checkpoint/2 and on_restore/2, allowing plugins to externalize large state objects or drop ephemeral data.

How Checkpoint Creation Works in Jido

When an agent is hibernated, Jido constructs a checkpoint through a coordinated sequence of agent callbacks, plugin externalization, and storage adapter writes.

The Agent Callback and Default Implementation

The Persist.hibernate/2 function initiates checkpoint creation by calling create_checkpoint/3, which checks for a custom checkpoint/2 implementation before falling back to the default.

# lib/jido/persist.ex – create_checkpoint/3

defp create_checkpoint(agent_module, agent, thread) do
  ctx = %{}
  result =
    if function_exported?(agent_module, :checkpoint, 2) do
      agent_module.checkpoint(agent, ctx)          # custom logic

    else
      {:ok, default_checkpoint(agent_module, agent, thread)} # default

    end
  # ...

end

The default_checkpoint/3 function in lib/jido/persist.ex constructs a versioned map containing the agent module, ID, state (with the thread removed), and a lightweight thread pointer:

# lib/jido/persist.ex – default_checkpoint/3

defp default_checkpoint(agent_module, agent, thread) do
  thread_pointer =
    case thread do
      nil -> nil
      %Thread{id: id, rev: rev} -> %{id: id, rev: rev}
    end

  %{
    version: 1,
    agent_module: agent_module,
    id: agent.id,
    state: Map.delete(agent.state, :__thread__),
    thread: thread_pointer
  }
end

Plugin State Externalization

During checkpoint creation, the macro-generated checkpoint/2 implementation in lib/jido/agent.ex iterates over all plugin instances, allowing each to decide its persistence strategy:

# lib/jido/agent.ex – quoted checkpoint implementation

def checkpoint(agent, ctx) do
  {state, externalized, externalized_keys} =
    Enum.reduce(@plugin_instances, {agent.state, %{}, %{}}, fn instance,
                                                           {state_acc, ext_acc,
                                                            keys_acc} ->
      plugin_state = Map.get(state_acc, instance.state_key)
      config = instance.config || %{}
      case instance.module.on_checkpoint(plugin_state,
            Map.put(ctx, :config, config)) do
        {:externalize, key, pointer} ->
          {Map.delete(state_acc, instance.state_key),
           Map.put(ext_acc, key, pointer),
           Map.put(keys_acc, key, instance.state_key)}
        :drop -> {Map.delete(state_acc, instance.state_key), ext_acc, keys_acc}
        :keep -> {state_acc, ext_acc, keys_acc}
      end
    end)

  base = %{version: 1, agent_module: __MODULE__, id: agent.id, state: state}
  base = if externalized_keys == %{}, do: base, else: Map.put(base, :externalized_keys, externalized_keys)
  {:ok, Map.merge(base, externalized)}
end

Plugins return one of three directives:

• :keep – State remains in the checkpoint map.
• :drop – State is excluded from persistence (useful for ephemeral data).
• {:externalize, key, pointer} – State is removed and replaced with a lightweight pointer, enabling zero-copy persistence for large objects like threads.

Storage Adapter Integration

After constructing the checkpoint map, Persist.hibernate/2 writes it via Storage.put_checkpoint/3. The default Jido.Storage.File adapter performs atomic writes to disk, while Jido.Storage.ETS provides in-memory storage for testing.

How Agent Restoration Works in Jido

Restoration reverses the checkpoint process: fetching serialized data, rebuilding the agent, re-attaching threads, and restoring plugin state.

Fetching and Rebuilding the Agent

Persist.thaw/3 retrieves the checkpoint using the storage adapter and initiates restoration:

# lib/jido/persist.ex – thaw flow (excerpt)

checkpoint_key = make_checkpoint_key(agent_module, key)
case Jido.Storage.fetch_checkpoint(adapter, checkpoint_key, opts) do
  {:ok, checkpoint} -> restore_from_checkpoint(adapter, opts, agent_module, checkpoint)
  ...
end

The restore_agent/3 function checks for a custom restore/2 callback before applying the default restoration logic:

# lib/jido/persist.ex – restore_agent/3

if function_exported?(agent_module, :restore, 2) do
  agent_module.restore(checkpoint, ctx)
else
  default_restore(agent_module, checkpoint)
end

The default_restore/2 function creates a fresh agent via new/1 and merges the persisted state:

# lib/jido/persist.ex – default_restore/2

defp default_restore(agent_module, checkpoint) do
  case agent_module.new(id: checkpoint.id) do
    {:ok, agent} ->
      merged_state = Map.merge(agent.state, checkpoint.state || %{})
      {:ok, %{agent | state: merged_state}}
    # also handles the case where new/1 returns a struct directly

    agent when is_struct(agent) ->
      merged_state = Map.merge(agent.state, checkpoint.state || %{})
      {:ok, %{agent | state: merged_state}}
  end
end

Thread Rehydration

If the checkpoint contains a thread pointer, Persist.rehydrate_thread/4 loads the full thread from storage and calls Agent.attach_thread/2. A revision mismatch between the pointer and the stored thread aborts the restoration to prevent stale state corruption.

Plugin State Restoration

After rebuilding the base agent, Persist.restore_from_checkpoint/4 iterates over plugin instances to restore externalized state:

# lib/jido/agent.ex – restore/2 default implementation (calls Persist)

# (the actual iteration lives in Persist.restore_from_checkpoint)

# See lines 96‑123 in persist.ex for the loop:

Enum.reduce_while(@plugin_instances, {:ok, agent}, fn instance, {:ok, acc} ->
  ext_key = Enum.find_value(externalized_keys, fn {k, v} -> if v == instance.state_key, do: k end)
  pointer = if ext_key, do: data[ext_key]
  if pointer do
    case instance.module.on_restore(pointer, restore_ctx) do
      {:ok, nil} -> {:cont, {:ok, acc}}                 # externalised plugin will rehydrate later

      {:ok, restored_state} ->
        {:cont, {:ok, %{acc | state: Map.put(acc.state, instance.state_key, restored_state)}}}
      {:error, reason} -> {:halt, {:error, reason}}
    end
  else
    {:cont, {:ok, acc}}
  end
end)

Plugins receive their stored pointer via on_restore/2 and return:

• {:ok, nil} – Indicates the plugin will rehydrate independently (e.g., threads loaded separately by Persist).
• {:ok, restored_state} – Returns the actual state to merge into the agent.
• {:error, reason} – Aborts the entire restoration.

Complete Implementation Example

Here is a practical implementation demonstrating custom checkpoint and restore logic with plugin externalization:

defmodule MyApp.CounterAgent do
  use Jido.Agent, plugins: [MyApp.ThreadPlugin]

  # ---- custom checkpoint -------------------------------------------------

  @impl true
  def checkpoint(agent, _ctx) do
    # keep normal state but externalise the heavy thread plugin

    {:ok,
     %{
       version: 1,
       agent_module: __MODULE__,
       id: agent.id,
       state: agent.state,
       externalized_keys: %{thread: :thread_state}
     }
     |> Map.merge(%{thread: %{id: agent.state.thread.id, rev: agent.state.thread.rev}})
    }
  end

  # ---- custom restore ----------------------------------------------------

  @impl true
  def restore(data, _ctx) do
    # Re‑create a fresh agent and merge persisted state

    with {:ok, agent} <- new(id: data.id) do
      merged = Map.merge(agent.state, data.state || %{})
      {:ok, %{agent | state: merged}}
    end
  end
end

defmodule MyApp.ThreadPlugin do
  use Jido.Plugin

  @impl true
  def on_checkpoint(_state, _ctx) do
    # Store only a lightweight pointer; Persist will later reattach the full thread

    {:externalize, :thread, %{id: :my_thread, rev: 42}}
  end

  @impl true
  def on_restore(pointer, _ctx) do
    # Persist will load the thread using the pointer; we just signal “OK”

    {:ok, nil}
  end
end

This example demonstrates:

• The agent’s checkpoint/2 building a map with externalized_keys linking the key :thread to the plugin’s internal state key.
• The plugin’s on_checkpoint/2 returning {:externalize, :thread, pointer} to store only a lightweight reference.
• During thaw, Persist loading the checkpoint, calling the agent’s restore/2, then invoking MyApp.ThreadPlugin.on_restore/2 with the pointer, which returns {:ok, nil} because Persist handles full thread rehydration separately.

Summary

• Checkpoint creation occurs in Jido.Persist.create_checkpoint/3, which calls the agent’s checkpoint/2 callback or falls back to default_checkpoint/3 in lib/jido/persist.ex.
• Plugin externalization allows plugins to return {:externalize, key, pointer} from on_checkpoint/2, storing lightweight pointers instead of full state.
• Storage adapters like Jido.Storage.File and Jido.Storage.ETS handle the actual serialization via Storage.put_checkpoint/3.
• Restoration in Jido.Persist.thaw/3 fetches the checkpoint, rebuilds the agent via restore/2 or default_restore/2, reattaches threads via rehydrate_thread/4, and restores plugin state through on_restore/2 callbacks.

Frequently Asked Questions

What is the difference between hibernation and checkpointing in Jido?

Hibernation refers to the high-level operation of pausing an agent and saving its state, while checkpointing is the specific mechanism that serializes the agent's state into a storable map. When you call Persist.hibernate/2, it internally invokes create_checkpoint/3 to build the checkpoint before writing it to storage.

How do plugins handle large state objects during checkpoint?

Plugins handle large state objects by returning {:externalize, key, pointer} from their on_checkpoint/2 callback. This instructs Jido.Persist to remove the bulky state from the checkpoint map and store only a lightweight pointer (such as a thread ID and revision). The actual heavy object is persisted separately, and the plugin receives the pointer during on_restore/2 to rehydrate the data.

Can I use custom storage adapters with Jido's persistence layer?

Yes, Jido's persistence layer is adapter-based. You can implement the storage behaviour defined in the codebase to create custom adapters. The repository includes Jido.Storage.File for atomic file-based persistence and Jido.Storage.ETS for in-memory testing. Your custom adapter must implement put_checkpoint/3 and fetch_checkpoint/3 to integrate with Jido.Persist.

What happens if the thread revision mismatches during restore?

If the thread pointer in the checkpoint contains a revision that does not match the current revision of the stored thread, Jido.Persist.rehydrate_thread/4 aborts the restoration process. This revision-checking mechanism prevents stale state corruption by ensuring that the agent only attaches to threads that haven't been modified since the checkpoint was created.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}