diff --git a/CHANGELOG.md b/CHANGELOG.md index ff07d7d..2903ddf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,8 +8,12 @@ - Add `Legion.Sandbox.Lua` - agents write Lua evaluated by [lua](https://hexdocs.pm/lua), a Lua 5.3 VM in pure Elixir. Generated code cannot reach the host BEAM at all (no AST allowlist to escape); tools are bridged in as global Lua tables with values converted at the boundary - Add `Legion.Store` for persisting conversations across process and application restarts; stores exchange partial `Legion.Store.Payload` values containing conversation state and metadata through `get/1` and `save/1` - Persist the user message with `status: :running` before execution and the final conversation with `status: :idle` before replying, so a reply is a commit receipt for the completed turn -- Add optional `persistence_frequency/0`; stores default to `:turn`, while `:step` also checkpoints intermediate eval results, recoverable errors, bindings, and executor progress. Interrupted turns are recorded but are not resumed automatically -- Add globally configured and per-agent stores, generated agent ids, `Legion.get_agent_id/1`, `Legion.lookup/1`, `Legion.running?/1`, and `Legion.resume/2` for identifying, finding, and restarting persisted conversations +- Add optional `persistence_frequency/0`; stores default to `:turn`, while `:step` also checkpoints intermediate eval results, recoverable errors, bindings, and executor progress +- Add configurable, generated agent ids, `Legion.get_agent_id/1`, `Legion.lookup/1`, and `Legion.resume/2` for identifying, finding, and restarting persisted conversations +- Register live agents cluster-wide by agent id through `:global`, preventing duplicate ownership across connected nodes +- Remove the `:name` option from `Legion.start_link/2`; use `:agent_id` for identity and `Legion.lookup/1` to resolve a live process +- Add `Legion.recover/2` and opt-in `:recovery` startup configuration to automatically drive interrupted root-agent runs (`status: :running`) to completion after an application restart. Recovery skips sub-agents to avoid replaying delegated work. +- Add `Legion.resume/2` to restore a persisted conversation under its original agent ID. When the saved conversation ends in a user message, it continues the interrupted executor loop in the background, resuming from a step checkpoint when available; it returns the existing process if that ID is already live. - Propagate stores to sub-agents and persist `parent_agent_id`, `agent_module`, and `started_at` metadata for reconstructing conversation trees - Add `Legion.Store.Postgres`, backed by an existing PostgreSQL Ecto repo, with partial upserts, `get/1`, `list/1`, configurable table names, configurable persistence frequency, and an optional `ecto_sql` dependency - Add versioned, idempotent `Legion.Store.Migration.Postgres` migrations with configurable table names and `pg_notify` notifications for inserts and updates; migration versions are tracked in the agents table comment; generated stores expose `__repo__/0` and `__table__/0` for database-backed consumers such as LegionWeb diff --git a/lib/legion.ex b/lib/legion.ex index 369054a..a5b499a 100644 --- a/lib/legion.ex +++ b/lib/legion.ex @@ -5,7 +5,8 @@ defmodule Legion do |> String.split("") |> Enum.fetch!(1) - alias Legion.AgentServer + alias Legion.{AgentIndex, AgentServer} + alias Legion.Store.Payload @doc """ Runs an agent on a single task and returns the result. @@ -32,17 +33,16 @@ defmodule Legion do Starts a long-lived agent process. ## Options - - `:name` - register the process under a name - `:store`, `:agent_id` - persist the conversation across restarts; see `Legion.Store`. A store set globally with `config :legion, :store, MyApp.AgentStore` applies to every agent, so you need only pass `:agent_id`. If a store is in effect but no `:agent_id` is given, Legion generates one - read it back with `get_agent_id/1`. + An agent ID can belong to at most one live process across connected nodes. - Any config overrides (`:model`, `:max_iterations`, etc.) ## Examples {:ok, pid} = Legion.start_link(AssistantAgent) - {:ok, pid} = Legion.start_link(AssistantAgent, name: MyAssistant, model: "openai:gpt-4o") {:ok, pid} = Legion.start_link(ChatAgent, store: MyApp.AgentStore, agent_id: "user_42:chat_7") {:ok, pid} = Legion.start_link(ChatAgent, agent_id: "user_42:chat_7") # store from app config """ @@ -105,10 +105,14 @@ defmodule Legion do end @doc """ - Returns whether `pid` is alive on this node. Non-pid values return `false`. + Whether `pid` - typically one returned by `lookup/1` or recorded in persisted + run metadata - is alive. Accepts `nil` and returns `false`. - Use `lookup/1` to resolve an `agent_id` to its currently registered process - before checking it. + Checks local processes directly and processes on connected nodes through RPC. + An unreachable remote node returns `false`. + + A stored pid can outlive the VM that wrote it, so after a restart this remains + best-effort: a recycled pid value can collide with an unrelated live process. ## Examples @@ -118,10 +122,12 @@ defmodule Legion do end """ def running?(pid) when is_pid(pid) do - node(pid) == node() and Process.alive?(pid) + if node(pid) == node() do + Process.alive?(pid) + else + :rpc.call(node(pid), :erlang, :is_process_alive, [pid]) == true + end rescue - # A pid deserialized from a previous VM incarnation is not a local pid, - # for which Process.alive?/1 raises. ArgumentError -> false end @@ -130,9 +136,10 @@ defmodule Legion do @doc """ Looks up the live process for an agent id. - Uses `Legion.AgentRegistry` as the runtime source of truth for the + Uses Legion's cluster-wide runtime index as the source of truth for the `agent_id -> pid` mapping. Returns `{:ok, pid}` when the agent is currently - registered, or `:error` when no live process is registered for `agent_id`. + registered on any connected node, or `:error` when no live process owns + `agent_id`. ## Examples @@ -140,9 +147,9 @@ defmodule Legion do :error = Legion.lookup("missing_agent_id") """ def lookup(agent_id) do - case Registry.lookup(Legion.AgentRegistry, agent_id) do - [{pid, _}] -> {:ok, pid} - [] -> :error + case AgentIndex.lookup(agent_id) do + {:ok, pid} -> {:ok, pid} + _ -> :error end end @@ -150,48 +157,115 @@ defmodule Legion do Resumes a persisted conversation. Loads the persisted conversation with `get/1` to determine its agent module, - then returns the process registered for `agent_id` if it is still running. - If no live process is registered, starts the persisted agent module under the - same `agent_id`, so it reloads its conversation state. `opts` are passed - through to `start_link/2`. + then atomically starts it under the same `agent_id`. If a live process already + owns that ID, returns the existing pid instead. A new process restores the + conversation and continues execution in the background. It resumes from a + saved checkpoint when one exists; otherwise it starts a new executor loop + with the restored history. + + `opts` are passed through to `start_link/2`. - Pass `:store` or configure one globally. Raises if `get/1` does not return a - `Legion.Store.Payload` containing an agent module for `agent_id`. + Pass `:store` or configure one globally. Returns: + + - `{:ok, pid}` when a new process starts or one already owns `agent_id` + - `{:error, :not_resumable}` when the store has no payload containing an + agent module for `agent_id` + - `{:error, reason}` when the process cannot start + + Raises when no store is available. ## Examples {:ok, pid} = Legion.resume("user_42:chat_7") {:ok, pid} = Legion.resume("user_42:chat_7", store: MyApp.AgentStore) + + {:error, :not_resumable} = + Legion.resume("missing_chat", store: MyApp.AgentStore) """ def resume(agent_id, opts \\ []) do - store = - Keyword.get(opts, :store) || Vault.get(:store) || Application.get_env(:legion, :store) || - raise ArgumentError, - "resume/2 requires a :store - pass one or set `config :legion, :store, MyStore`" + store = store!(opts, :resume) + + case store.get(agent_id) do + {:ok, %Payload{agent_module: agent_module}} when not is_nil(agent_module) -> + case AgentServer.start_link( + agent_module, + Keyword.merge(opts, agent_id: agent_id, store: store, start_mode: :resume) + ) do + {:ok, pid} -> {:ok, pid} + {:error, {:already_started, pid}} -> {:ok, pid} + {:error, reason} -> {:error, reason} + end - agent_module = - case store.get(agent_id) do - {:ok, %Legion.Store.Payload{agent_module: agent_module}} when not is_nil(agent_module) -> - agent_module + _ -> + {:error, :not_resumable} + end + end - _ -> - raise ArgumentError, - "no run recorded for agent_id #{inspect(agent_id)} in #{inspect(store)}" - end + @doc """ + Recovers an interrupted persisted run and waits for it to finish. - case lookup(agent_id) do - {:ok, pid} -> - if running?(pid) do - {:ok, pid} - else - start_link(agent_module, Keyword.put(opts, :agent_id, agent_id)) + An interrupted run is signaled by a stored payload with `status: :running`. In addition + only runs without `parent_agent_id` are recoverable, since sub-agents are not restarted. Unlike + `resume/2`, a live agent is not returned. + + The agent is only driven to completion, the executor result is not returned. `:ok` means + only that the temporary agent process stopped normally. + + Pass `:store` or configure one globally. Other options are passed through to + `start_link/2`. Returns: + + - `:ok` when the temporary process stops normally + - `{:error, reason}` when the temporary process stops abnormally + - `{:error, :already_running}` when a recoverable run already has a live process + - `{:error, :not_recoverable}` when there is no stored payload with an agent + module, or when the stored payload is not an interrupted run + + Raises when no store is available. + + ## Examples + + :ok = Legion.recover("user_42:chat_7", store: MyApp.AgentStore) + + {:error, :already_running} = + Legion.recover("active_chat", store: MyApp.AgentStore) + + {:error, :not_recoverable} = + Legion.recover("completed_chat", store: MyApp.AgentStore) + """ + def recover(agent_id, opts \\ []) do + store = store!(opts, :recover) + + case store.get(agent_id) do + {:ok, %Payload{agent_module: agent_module, status: :running, parent_agent_id: nil}} + when not is_nil(agent_module) -> + case AgentServer.start_monitor( + agent_module, + Keyword.merge(opts, agent_id: agent_id, store: store, start_mode: :recover) + ) do + {:ok, {pid, ref}} -> + receive do + {:DOWN, ^ref, :process, ^pid, :normal} -> :ok + {:DOWN, ^ref, :process, ^pid, reason} -> {:error, reason} + end + + {:error, {:already_started, _pid}} -> + {:error, :already_running} + + {:error, reason} -> + {:error, reason} end - :error -> - start_link(agent_module, Keyword.put(opts, :agent_id, agent_id)) + _ -> + {:error, :not_recoverable} end end + defp store!(opts, operation) do + Keyword.get(opts, :store) || Vault.get(:store) || Application.get_env(:legion, :store) || + raise ArgumentError, + "#{operation}/2 requires a :store - pass one or set `config :legion, :store, MyStore`" + end + @doc """ Runs multiple agent tasks concurrently and collects results. diff --git a/lib/legion/agent_index.ex b/lib/legion/agent_index.ex new file mode 100644 index 0000000..ad8bc31 --- /dev/null +++ b/lib/legion/agent_index.ex @@ -0,0 +1,44 @@ +defmodule Legion.AgentIndex do + @moduledoc """ + Cluster-wide process index used by Legion's agent servers. + + Implements the `:via` registration callbacks over `:global`. Starting an + agent atomically claims its agent ID across connected nodes. The name is + released automatically when the process stops. + """ + + @doc false + def name(agent_id) do + {:via, __MODULE__, agent_id} + end + + @doc false + def register_name(agent_id, pid) do + :global.register_name(key(agent_id), pid) + end + + @doc false + def unregister_name(agent_id) do + :global.unregister_name(key(agent_id)) + end + + @doc false + def whereis_name(agent_id) do + :global.whereis_name(key(agent_id)) + end + + @doc false + def send(agent_id, message) do + :global.send(key(agent_id), message) + end + + @doc false + def lookup(agent_id) do + case whereis_name(agent_id) do + :undefined -> :error + pid -> {:ok, pid} + end + end + + defp key(agent_id), do: {:legion_agent, agent_id} +end diff --git a/lib/legion/agent_server.ex b/lib/legion/agent_server.ex index 01020f9..411b980 100644 --- a/lib/legion/agent_server.ex +++ b/lib/legion/agent_server.ex @@ -4,6 +4,10 @@ defmodule Legion.AgentServer do Holds the message history across multiple turns. Each `call` or `cast` appends the user message and runs `Executor` to completion (blocking). + + When Legion starts an agent in resume or recovery mode, the server restores + the persisted conversation and executor checkpoint, then continues it during + startup. Recovery stops the temporary process after that execution finishes. """ use GenServer @@ -21,13 +25,27 @@ defmodule Legion.AgentServer do :store, :agent_id, :persistence_frequency, + executor_state: :nonexistent, bindings: [] ] # Client API def start_link(agent_module, opts \\ []) do - {name, opts} = Keyword.pop(opts, :name) + {init_arg, gen_opts} = start_args(agent_module, opts) + + GenServer.start_link(__MODULE__, init_arg, gen_opts) + end + + @doc false + def start_monitor(agent_module, opts \\ []) do + {init_arg, gen_opts} = start_args(agent_module, opts) + {name, gen_opts} = Keyword.pop!(gen_opts, :name) + + :gen_server.start_monitor(name, __MODULE__, init_arg, gen_opts) + end + + defp start_args(agent_module, opts) do {store, opts} = Keyword.pop(opts, :store) {agent_id, opts} = Keyword.pop(opts, :agent_id) @@ -41,14 +59,10 @@ defmodule Legion.AgentServer do agent_id = agent_id || generate_id() persistence_frequency = Store.persistence_frequency(store) - gen_opts = if name, do: [name: name], else: [] + gen_opts = [name: Legion.AgentIndex.name(agent_id)] config = resolve_config(agent_module, opts) - GenServer.start_link( - __MODULE__, - {agent_module, config, store, agent_id, persistence_frequency}, - gen_opts - ) + {{agent_module, config, store, agent_id, persistence_frequency}, gen_opts} end def call(agent, message, timeout \\ :infinity) do @@ -72,17 +86,12 @@ defmodule Legion.AgentServer do @impl true def init({agent_module, config, store, agent_id, persistence_frequency}) do parent_agent_id = Vault.get(:agent_id) + mode = Map.get(config, :start_mode, :normal) Vault.unsafe_put(:agent_id, agent_id) Vault.unsafe_put(:parent_agent_id, parent_agent_id) if store, do: Vault.unsafe_put(:store, store) - {:ok, _} = - Registry.register(Legion.AgentRegistry, agent_id, %{ - parent_agent_id: parent_agent_id, - started_at: NaiveDateTime.utc_now() - }) - for tool <- agent_module.tools() do Vault.unsafe_put(tool, agent_module.tool_config(tool)) end @@ -95,13 +104,20 @@ defmodule Legion.AgentServer do %{agent: agent_module} ) - {saved_messages, saved_bindings} = + {saved_messages, saved_bindings, saved_executor_state} = case store && store.get(agent_id) do - {:ok, %Payload{conversation_state: %{messages: messages, bindings: bindings}}} -> - {messages, bindings} + {:ok, + %Payload{ + conversation_state: %{ + messages: messages, + bindings: bindings, + executor_state: executor_state + } + }} -> + {messages, bindings, executor_state} _no_state -> - {[], []} + {[], [], :nonexistent} end state = %__MODULE__{ @@ -111,7 +127,8 @@ defmodule Legion.AgentServer do store: store, agent_id: agent_id, persistence_frequency: persistence_frequency, - bindings: saved_bindings + bindings: saved_bindings, + executor_state: saved_executor_state } {:ok, @@ -119,7 +136,26 @@ defmodule Legion.AgentServer do agent_module: state.agent_module, parent_agent_id: parent_agent_id, started_at: NaiveDateTime.utc_now() - )} + ), {:continue, %{start_mode: mode, executor_state: saved_executor_state}}} + end + + @impl true + def handle_continue(%{start_mode: :normal}, state), do: {:noreply, state} + + @impl true + def handle_continue(%{start_mode: :resume, executor_state: executor_state}, state) do + if match?(%{role: "user"}, List.last(state.messages)) do + {_reply, state} = do_run(state, executor_state) + {:noreply, state} + else + {:noreply, state} + end + end + + @impl true + def handle_continue(%{start_mode: :recover, executor_state: executor_state}, state) do + {_reply, state} = do_run(state, executor_state) + {:stop, :normal, state} end @impl true @@ -168,7 +204,6 @@ defmodule Legion.AgentServer do """ def handle_message(message, state) do content = stringify(message, state.config[:max_message_length]) - conversation_scope? = Map.get(state.config, :binding_scope, :turn) == :conversation # Persist the user message before the turn runs so store-backed views # (e.g. the legion_web database source) show it without waiting for the @@ -178,6 +213,12 @@ defmodule Legion.AgentServer do |> Map.update!(:messages, &(&1 ++ [Executor.message(:user, content)])) |> persist([:conversation_state, status: :running]) + do_run(state) + end + + defp do_run(state, executor_state \\ :nonexistent) do + conversation_scope? = Map.get(state.config, :binding_scope, :turn) == :conversation + checkpoint = if state.persistence_frequency == :step do fn checkpoint -> @@ -191,7 +232,7 @@ defmodule Legion.AgentServer do {status, value, final_messages, final_bindings} = Telemetry.span( [:legion, :agent, :message], - %{agent: state.agent_module, message: content}, + %{agent: state.agent_module, message: state.messages |> List.last() |> Map.get(:content)}, fn -> messages = state.messages prev_count = Enum.count(messages, &(&1[:role] == "assistant")) @@ -199,7 +240,14 @@ defmodule Legion.AgentServer do initial_bindings = if conversation_scope?, do: state.bindings, else: [] {status, value, messages, bindings} = - result = Executor.run(state.agent_module, messages, executor_config, initial_bindings) + result = + Executor.run( + state.agent_module, + messages, + executor_config, + initial_bindings, + executor_state + ) iterations = Enum.count(messages, &(&1[:role] == "assistant")) - prev_count {result, %{iterations: iterations, status: status, result: value, bindings: bindings}} @@ -242,15 +290,15 @@ defmodule Legion.AgentServer do defp persisted_conversation_state(%__MODULE__{} = state) do [%{role: "system"} | messages] = state.messages - %{messages: messages, bindings: state.bindings} + %{messages: messages, bindings: state.bindings, executor_state: :nonexistent} end defp persisted_conversation_state(%{ messages: [%{role: "system"} | messages], bindings: bindings, - execution: execution + executor_state: executor_state }) do - %{messages: messages, bindings: bindings, execution: execution} + %{messages: messages, bindings: bindings, executor_state: executor_state} end defp generate_id, do: Base.url_encode64(:crypto.strong_rand_bytes(16), padding: false) @@ -275,7 +323,7 @@ defmodule Legion.AgentServer do |> Executor.truncate_content(max_length) end - @known_config_keys ~w(binding_scope eval_guard max_iterations max_message_length max_retries model sandbox sandbox_max_heap sandbox_max_reductions sandbox_priority sandbox_timeout)a + @known_config_keys ~w(binding_scope eval_guard max_iterations max_message_length max_retries model sandbox sandbox_max_heap sandbox_max_reductions sandbox_priority sandbox_timeout start_mode)a defp resolve_config(agent_module, opts) do app_config = Application.get_env(:legion, :config, %{}) diff --git a/lib/legion/application.ex b/lib/legion/application.ex index 3776c97..d5bcbf4 100644 --- a/lib/legion/application.ex +++ b/lib/legion/application.ex @@ -7,10 +7,13 @@ defmodule Legion.Application do @impl true def start(_type, _args) do - children = [ - {Registry, keys: :unique, name: Legion.AgentRegistry} - ] + Supervisor.start_link(children(), strategy: :one_for_one, name: Legion.Supervisor) + end - Supervisor.start_link(children, strategy: :one_for_one, name: Legion.Supervisor) + @doc false + def children do + [ + {Legion.Recovery, Application.fetch_env(:legion, :recovery)} + ] end end diff --git a/lib/legion/executor.ex b/lib/legion/executor.ex index 9209050..7cbcbf9 100644 --- a/lib/legion/executor.ex +++ b/lib/legion/executor.ex @@ -126,12 +126,26 @@ defmodule Legion.Executor do Runs the LLM loop against the given message history. `messages` must already include the system prompt and the current user message. + `bindings` seeds the code-evaluation binding. `executor_state` resumes a step + checkpoint when present: `:awaiting_llm` continues from its saved iteration + and retry counters, while `:completing` finishes without another LLM request. + Pass `:nonexistent` to start a new loop. Returns `{:ok, result, messages, bindings}` or `{:cancel, reason, messages, bindings}`. """ - def run(agent_module, messages, config, bindings \\ []) do + def run(agent_module, messages, config, bindings \\ [], executor_state \\ :nonexistent) do config = Map.merge(@default_config, config) - loop(agent_module, messages, config, 0, 0, bindings) + + case executor_state do + :nonexistent -> + loop(agent_module, messages, config, 0, 0, bindings) + + %{phase: :awaiting_llm, iteration: i, retries: r} -> + loop(agent_module, messages, config, i, r, bindings) + + %{phase: :completing, iteration: _i, retries: _r} -> + {:ok, nil, messages, bindings} + end end defp loop(agent_module, messages, config, iteration, retries, bindings) do @@ -202,7 +216,7 @@ defmodule Legion.Executor do ) end - defp checkpoint!(config, messages, bindings, execution) do + defp checkpoint!(config, messages, bindings, executor_state) do case config[:checkpoint] do nil -> :ok @@ -213,7 +227,7 @@ defmodule Legion.Executor do callback.(%{ messages: messages, bindings: bindings, - execution: execution + executor_state: executor_state }) rescue error -> exit({:checkpoint_persistence_failed, error}) @@ -256,14 +270,14 @@ defmodule Legion.Executor do messages = messages ++ [message(:eval_result, format_result(result, new_bindings, config))] - execution = + executor_state = if eval == "eval_and_continue" do %{phase: :awaiting_llm, iteration: i + 1, retries: 0} else %{phase: :completing, iteration: i, retries: 0} end - checkpoint!(config, messages, new_bindings, execution) + checkpoint!(config, messages, new_bindings, executor_state) if eval == "eval_and_continue", do: loop(agent, messages, config, i + 1, 0, new_bindings), diff --git a/lib/legion/recovery.ex b/lib/legion/recovery.ex new file mode 100644 index 0000000..db2f0bf --- /dev/null +++ b/lib/legion/recovery.ex @@ -0,0 +1,79 @@ +defmodule Legion.Recovery do + @moduledoc """ + Worker which runs the recovery process for interrupted runs. + + `Legion.Application` starts this worker only when `:recovery` is configured: + + config :legion, :recovery, + stores: [MyApp.AgentStore], + store_scan_limit: 100, + concurrent_request_limit: 10 + + The worker calls `list(store_scan_limit)` on every configured store, selects + payloads with `status: :running` and no `parent_agent_id`, then invokes + `Legion.recover/2` for each selected run, with at most + `concurrent_request_limit` recoveries in flight. + + Recovery deliberately does not restart sub-agents. A parent can persist its + state before a code evaluation dispatches a sub-agent, then crash before the + following checkpoint records that dispatch. When the parent recovers, it can + replay that evaluation and dispatch a new sub-agent; recovering the recorded + child independently could execute the same work twice. + + `store_scan_limit` limits payloads read from each store. + `concurrent_request_limit` limits recoveries in flight across all stores and + defaults to `3`. + Recovery runs asynchronously during application startup. The worker performs + the recovery scan, then exits and is not restarted. Configure it through the + application environment rather than starting it directly. + + The temporary agent process drives each interrupted run to completion, then + stops. + """ + + alias Legion.Store.Payload + + @default_concurrent_request_limit 3 + + @doc false + def start_link(:error), do: :ignore + + def start_link({:ok, config}) do + Task.start_link(fn -> run(config) end) + end + + @doc false + def run(config) do + stores = Keyword.fetch!(config, :stores) + store_scan_limit = Keyword.fetch!(config, :store_scan_limit) + + concurrent_request_limit = + Keyword.get(config, :concurrent_request_limit, @default_concurrent_request_limit) + + stores + |> Enum.flat_map(fn store -> + store.list(store_scan_limit) + |> Enum.map(fn payload -> {store, payload} end) + end) + |> Enum.filter(fn {_store, payload} -> running_root?(payload) end) + |> Task.async_stream( + fn {store, %Payload{agent_id: agent_id}} -> Legion.recover(agent_id, store: store) end, + max_concurrency: concurrent_request_limit, + ordered: false, + timeout: :infinity + ) + |> Stream.run() + end + + defp running_root?(%Payload{status: :running, parent_agent_id: nil}), do: true + defp running_root?(_payload), do: false + + @doc false + def child_spec(config) do + %{ + id: __MODULE__, + start: {__MODULE__, :start_link, [config]}, + restart: :temporary + } + end +end diff --git a/lib/legion/store.ex b/lib/legion/store.ex index 31ecc00..7952f3d 100644 --- a/lib/legion/store.ex +++ b/lib/legion/store.ex @@ -60,22 +60,21 @@ defmodule Legion.Store do Legion.start_link(ChatAgent, agent_id: "user_42:chat_7") - Omitting `:agent_id` makes Legion generate one. That - suits a brand-new conversation: read it back with `Legion.get_agent_id/1` and - persist the mapping if you want to resume the chat later. Pass your own id to - resume an existing conversation. Two agents started under the same id race onto - the same row, so route each conversation to a single process. + Omitting `:agent_id` makes Legion generate one. That suits a brand-new conversation: + read it back with `Legion.get_agent_id/1` and persist the mapping if you want to resume + the chat later. ## Required callbacks Stores must implement `get/1`, `list/1`, and `save/1`. `save/1` receives a `Legion.Store.Payload`. Its `:conversation_state` is a - map containing the conversation's `:messages` (without the system prompt) - and `:bindings` from evaluated code. Step snapshots also contain an - `:execution` map with `:phase`, `:iteration`, and `:retries`. `:status` - records whether the agent is mid-turn. The payload also carries the agent - module, parent conversation, and start time when those values are known. + map containing the conversation's `:messages` (without the system prompt), + `:bindings` from evaluated code, and `:executor_state`. `:executor_state` is + `:nonexistent` for ordinary snapshots and is a map with `:phase`, `:iteration`, + and `:retries` for step checkpoints. `:status` records whether the agent is + mid-turn. The payload also carries the agent module, parent conversation, + and start time when those values are known. With `binding_scope: :turn`, active bindings are included in step snapshots while the turn is running and cleared from the final snapshot. Bindings with @@ -96,8 +95,10 @@ defmodule Legion.Store do Step persistence accepts a replay window between an LLM selecting an eval action and the following result or error checkpoint. A crash in that window - can replay the action and any external side effects. Automatic continuation - of an interrupted turn is not currently performed. + can replay the action and any external side effects. Configure + `:recovery` with stores, a store scan limit, and a concurrent request limit + to recover interrupted root turns once when the Legion application starts; + see `Legion.Recovery` and `Legion.recover/2`. ## Reading conversations diff --git a/lib/legion/store/migration/postgres.ex b/lib/legion/store/migration/postgres.ex index d6120a7..b3280cc 100644 --- a/lib/legion/store/migration/postgres.ex +++ b/lib/legion/store/migration/postgres.ex @@ -36,7 +36,7 @@ defmodule Legion.Store.Migration.Postgres do ## Migrating Without Ecto If your application uses something other than Ecto for migrations, be it an external system or - another ORM, it may be helpful to create plain SQL migrations for Oban database schema changes. + another ORM, it may be helpful to create plain SQL migrations for Legion's database schema changes. The simplest mechanism for obtaining the SQL changes is to create the migration locally and run `mix ecto.migrate --log-migrations-sql`. That will log all of the generated SQL, which you can diff --git a/lib/legion/store/payload.ex b/lib/legion/store/payload.ex index 46ac048..4c106b0 100644 --- a/lib/legion/store/payload.ex +++ b/lib/legion/store/payload.ex @@ -1,6 +1,14 @@ defmodule Legion.Store.Payload do @moduledoc """ Data supplied to and returned from a `Legion.Store`. + + Payloads are partial updates: `agent_id` is required, while a `nil` value for + every other field means the store must preserve its existing value. A + `conversation_state` holds the persisted messages, bindings, and executor + checkpoint. Its `:executor_state` value is `:nonexistent` for an ordinary snapshot + or a map when step persistence captures an interrupted turn. + + See `Legion.Store` for the full store contract. """ @enforce_keys [:agent_id] @@ -14,16 +22,19 @@ defmodule Legion.Store.Payload do ] @type status :: :idle | :running - @type execution :: %{ + + @type executor_state :: %{ phase: :awaiting_llm | :completing, iteration: non_neg_integer(), retries: non_neg_integer() } + @type state :: %{ - required(:messages) => [map()], - required(:bindings) => keyword(), - optional(:execution) => execution() + messages: [map()], + bindings: keyword(), + executor_state: :nonexistent | executor_state() } + @type t :: %__MODULE__{ agent_id: Legion.Store.agent_id(), agent_module: module() | nil, diff --git a/lib/legion/store/postgres.ex b/lib/legion/store/postgres.ex index fe66113..70d1ed7 100644 --- a/lib/legion/store/postgres.ex +++ b/lib/legion/store/postgres.ex @@ -185,9 +185,9 @@ defmodule Legion.Store.Postgres do %{ messages: Map.get(state, :messages, []), - bindings: Map.get(state, :bindings, []) + bindings: Map.get(state, :bindings, []), + executor_state: Map.get(state, :executor_state, :nonexistent) } - |> Map.merge(Map.take(state, [:execution])) end defp decode_status("running"), do: :running diff --git a/mix.exs b/mix.exs index eb7376c..7dc7a2a 100644 --- a/mix.exs +++ b/mix.exs @@ -2,7 +2,7 @@ defmodule Legion.MixProject do use Mix.Project @version "0.4.0" - @source_url "https://github.com/dimamik/legion" + @source_url "https://github.com/software-mansion-labs/legion" def project do [ @@ -62,9 +62,14 @@ defmodule Legion.MixProject do Legion.Store.Postgres, Legion.Store.Migration.Postgres ], - Runtime: [Legion.AgentServer, Legion.Executor, ~r/^Legion\.Sandbox/], + Runtime: [Legion.AgentServer, Legion.Executor, Legion.Recovery, ~r/^Legion\.Sandbox/], Tools: [~r/^Legion\.Tools\./], - Internals: [Legion.AgentPrompt, Legion.SourceRegistry, Legion.Telemetry] + Internals: [ + Legion.AgentIndex, + Legion.AgentPrompt, + Legion.SourceRegistry, + Legion.Telemetry + ] ] end diff --git a/test/integration/agent_index_test.exs b/test/integration/agent_index_test.exs new file mode 100644 index 0000000..0ff8539 --- /dev/null +++ b/test/integration/agent_index_test.exs @@ -0,0 +1,77 @@ +defmodule Legion.Integration.AgentIndexTest do + use ExUnit.Case, async: false + + @moduletag :integration + + test "resolves ownership and process exit across connected nodes" do + started_distribution? = start_distribution() + peer_name = String.to_atom("legion_peer_#{System.unique_integer([:positive])}") + {:ok, peer, peer_node} = :peer.start_link(%{name: peer_name}) + + try do + :ok = :global.sync() + :ok = :rpc.call(peer_node, :global, :sync, []) + + {Legion.AgentIndex, binary, filename} = :code.get_object_code(Legion.AgentIndex) + + assert {:module, Legion.AgentIndex} = + :rpc.call( + peer_node, + :code, + :load_binary, + [Legion.AgentIndex, filename, binary] + ) + + agent_id = "remote-owner" + remote_pid = :rpc.call(peer_node, :erlang, :spawn, [:timer, :sleep, [:infinity]]) + + assert :yes = + :rpc.call(peer_node, Legion.AgentIndex, :register_name, [agent_id, remote_pid]) + + assert {:ok, ^remote_pid} = Legion.lookup(agent_id) + assert Legion.running?(remote_pid) + + assert {:error, {:already_started, ^remote_pid}} = + Agent.start_link(fn -> %{} end, name: Legion.AgentIndex.name(agent_id)) + + monitor_ref = Process.monitor(remote_pid) + Process.exit(remote_pid, :kill) + + assert_receive {:DOWN, ^monitor_ref, :process, ^remote_pid, :killed} + assert_eventually(fn -> Legion.lookup(agent_id) == :error end) + after + :peer.stop(peer) + if started_distribution?, do: Node.stop() + end + end + + defp start_distribution do + if Node.alive?() do + false + else + {_output, 0} = System.cmd("epmd", ["-daemon"]) + origin_name = String.to_atom("legion_origin_#{System.unique_integer([:positive])}") + {:ok, _pid} = Node.start(origin_name, :shortnames) + true + end + end + + defp assert_eventually(condition, timeout \\ 1_000) do + deadline = System.monotonic_time(:millisecond) + timeout + wait_for(condition, deadline) + end + + defp wait_for(condition, deadline) do + cond do + condition.() -> + :ok + + System.monotonic_time(:millisecond) >= deadline -> + flunk("condition was not met within the timeout") + + true -> + Process.sleep(10) + wait_for(condition, deadline) + end + end +end diff --git a/test/integration/step_persistence_test.exs b/test/integration/step_persistence_test.exs index 07cabe8..1ba5510 100644 --- a/test/integration/step_persistence_test.exs +++ b/test/integration/step_persistence_test.exs @@ -20,7 +20,7 @@ defmodule Legion.Integration.StepPersistenceTest do :ok end - test "persists the recoverable turn state before execution advances" do + test "persists the recoverable turn state before executor_state advances" do agent_id = "step-persistence-integration" test_pid = self() request_count = :counters.new(1, [:atomics]) @@ -52,7 +52,7 @@ defmodule Legion.Integration.StepPersistenceTest do assert Enum.map(initial_state.messages, & &1.type) == [:user] assert initial_state.bindings == [] - refute Map.has_key?(initial_state, :execution) + assert initial_state.executor_state == :nonexistent assert_received {:before_second_request, {:ok, @@ -64,7 +64,7 @@ defmodule Legion.Integration.StepPersistenceTest do assert Enum.map(checkpoint.messages, & &1.type) == [:user, :assistant, :eval_result] assert checkpoint.bindings == [x: 42] - assert checkpoint.execution == %{ + assert checkpoint.executor_state == %{ phase: :awaiting_llm, iteration: 1, retries: 0 @@ -74,7 +74,7 @@ defmodule Legion.Integration.StepPersistenceTest do StepStore.get(agent_id) assert completed.bindings == [] - refute Map.has_key?(completed, :execution) + assert completed.executor_state == :nonexistent end defp response(action, code, result) do diff --git a/test/legion/agent_server_test.exs b/test/legion/agent_server_test.exs index 18697c4..5378a7c 100644 --- a/test/legion/agent_server_test.exs +++ b/test/legion/agent_server_test.exs @@ -92,6 +92,17 @@ defmodule Legion.AgentServerTest do end end + describe "start_monitor/2" do + test "starts an agent and returns a monitor reference" do + assert {:ok, {pid, monitor_ref}} = Legion.AgentServer.start_monitor(MathAgent) + assert Process.alive?(pid) + + GenServer.stop(pid) + + assert_receive {:DOWN, ^monitor_ref, :process, ^pid, :normal} + end + end + describe "config validation" do test "warns about unknown config keys" do stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> @@ -413,17 +424,6 @@ defmodule Legion.AgentServerTest do end end - describe "named registration" do - test "agent can be started with a registered name" do - stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> - llm_response("42") - end) - - {:ok, _pid} = Legion.start_link(MathAgent, name: :test_named_agent) - assert {:ok, "42"} = Legion.call(:test_named_agent, "What is 42?") - end - end - defmodule MemoryStore do @behaviour Legion.Store @@ -463,9 +463,14 @@ defmodule Legion.AgentServerTest do merged = merge(existing, payload) + state = + state + |> Map.put(payload.agent_id, {:ok, merged}) + |> Map.update({:writes, payload.agent_id}, [payload], &[payload | &1]) + + if watcher = Map.get(state, :save_watcher), do: send(watcher, {:store_saved, payload}) + state - |> Map.put(payload.agent_id, {:ok, merged}) - |> Map.update({:writes, payload.agent_id}, [payload], &[payload | &1]) end) :ok @@ -478,6 +483,10 @@ defmodule Legion.AgentServerTest do |> Enum.reverse() end + def watch_saves(test_pid) do + Agent.update(__MODULE__, &Map.put(&1, :save_watcher, test_pid)) + end + def statuses(agent_id) do agent_id |> writes() @@ -510,6 +519,19 @@ defmodule Legion.AgentServerTest do end end + defmodule EmptyStore do + @behaviour Legion.Store + + @impl Legion.Store + def get(_agent_id), do: :error + + @impl Legion.Store + def list(_limit), do: [] + + @impl Legion.Store + def save(_payload), do: :ok + end + defmodule StepMemoryStore do @behaviour Legion.Store @@ -534,6 +556,13 @@ defmodule Legion.AgentServerTest do :ok end + test "passes the store persistence frequency into the agent server state" do + {:ok, pid} = + Legion.start_link(MathAgent, store: StepMemoryStore, agent_id: "step-frequency") + + assert %{persistence_frequency: :step} = :sys.get_state(pid) + end + test "brackets each turn with :running and :idle status writes" do stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> llm_response("Paris") @@ -691,7 +720,11 @@ defmodule Legion.AgentServerTest do assert %Payload{ status: :running, - conversation_state: %{messages: [%{type: :user}], bindings: []} + conversation_state: %{ + messages: [%{type: :user}], + bindings: [], + executor_state: :nonexistent + } } = running assert %Payload{ @@ -699,13 +732,13 @@ defmodule Legion.AgentServerTest do conversation_state: %{ messages: [%{type: :user}, %{type: :assistant}, %{type: :eval_result}], bindings: [x: 42], - execution: %{phase: :awaiting_llm, iteration: 1, retries: 0} + executor_state: %{phase: :awaiting_llm, iteration: 1, retries: 0} } } = checkpoint assert %Payload{status: :idle, conversation_state: final_state} = completed assert final_state.bindings == [] - refute Map.has_key?(final_state, :execution) + assert final_state.executor_state == :nonexistent end test "a :step store persists eval_and_complete before the final snapshot" do @@ -723,12 +756,12 @@ defmodule Legion.AgentServerTest do assert %Payload{ status: nil, conversation_state: %{ - execution: %{phase: :completing, iteration: 0, retries: 0} + executor_state: %{phase: :completing, iteration: 0, retries: 0} } } = checkpoint assert %Payload{status: :idle, conversation_state: final_state} = completed - refute Map.has_key?(final_state, :execution) + assert final_state.executor_state == :nonexistent end test "a :step store persists retry state after an error message" do @@ -753,7 +786,7 @@ defmodule Legion.AgentServerTest do conversation_state: %{ messages: messages, bindings: [], - execution: %{phase: :awaiting_llm, iteration: 0, retries: 1} + executor_state: %{phase: :awaiting_llm, iteration: 0, retries: 1} } } = checkpoint @@ -907,6 +940,49 @@ defmodule Legion.AgentServerTest do assert {:ok, ^pid} = Legion.lookup("lookup") end + test "allows only one live process to own an agent_id" do + {:ok, pid} = Legion.start_link(MathAgent, store: MemoryStore, agent_id: "unique") + + assert {:error, {:already_started, ^pid}} = + Legion.start_link(MathAgent, store: MemoryStore, agent_id: "unique") + end + + test "concurrent starts atomically choose one owner for an agent_id" do + caller = self() + + contenders = + for _index <- 1..8 do + Task.async(fn -> + send(caller, {:ready, self()}) + + receive do + :start -> Legion.start_link(MathAgent, store: MemoryStore, agent_id: "race") + end + end) + end + + contender_pids = + for _index <- 1..8 do + assert_receive {:ready, contender_pid} + contender_pid + end + + Enum.each(contender_pids, &send(&1, :start)) + results = Task.await_many(contenders) + started_pids = for {:ok, pid} <- results, do: pid + + on_exit(fn -> + Enum.each(started_pids, fn pid -> + if Process.alive?(pid), do: GenServer.stop(pid) + end) + end) + + assert [winner] = started_pids + + assert Enum.count(results, &(&1 == {:error, {:already_started, winner}})) == 7 + assert {:ok, ^winner} = Legion.lookup("race") + end + test "resume/2 returns the recorded process while it is alive" do {:ok, pid} = Legion.start_link(MathAgent, store: MemoryStore, agent_id: "resume-live") @@ -914,6 +990,14 @@ defmodule Legion.AgentServerTest do assert {:ok, ^pid} = Legion.resume("resume-live", store: MemoryStore) end + test "resume/2 validates the requested store before resolving a live process" do + {:ok, _pid} = + Legion.start_link(MathAgent, store: MemoryStore, agent_id: "resume-wrong-store") + + assert {:error, :not_resumable} = + Legion.resume("resume-wrong-store", store: EmptyStore) + end + test "resume/2 restarts a stopped conversation from its run metadata" do stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> llm_response("Paris") @@ -935,12 +1019,198 @@ defmodule Legion.AgentServerTest do ] = Legion.get_messages(revived) end - test "resume/2 raises for an agent_id the store has no run for" do - assert_raise ArgumentError, ~r/no run recorded/, fn -> - Legion.resume("ghost", store: MemoryStore) + test "an awaiting-LLM checkpoint resumes with one request and finishes idle" do + assert :ok = + MemoryStore.save(%Payload{ + agent_id: "resume-awaiting-llm", + agent_module: MathAgent, + status: :running, + conversation_state: %{ + messages: [%{role: "user", type: :user, content: "compute"}], + bindings: [x: 42], + executor_state: %{phase: :awaiting_llm, iteration: 1, retries: 0} + } + }) + + MemoryStore.watch_saves(self()) + test_pid = self() + + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + send(test_pid, :llm_requested) + llm_response("done") + end) + + assert {:ok, _pid} = Legion.resume("resume-awaiting-llm", store: MemoryStore) + assert_receive :llm_requested + + assert_receive {:store_saved, + %Payload{status: :idle, conversation_state: %{executor_state: :nonexistent}}} + + refute_receive :llm_requested, 50 + end + + test "a completing checkpoint resumes without a request and finishes idle" do + assert :ok = + MemoryStore.save(%Payload{ + agent_id: "resume-completing", + agent_module: MathAgent, + status: :running, + conversation_state: %{ + messages: [%{role: "user", type: :user, content: "compute"}], + bindings: [x: 42], + executor_state: %{phase: :completing, iteration: 1, retries: 0} + } + }) + + MemoryStore.watch_saves(self()) + test_pid = self() + + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + send(test_pid, :llm_requested) + llm_response("unexpected") + end) + + assert {:ok, _pid} = Legion.resume("resume-completing", store: MemoryStore) + + assert_receive {:store_saved, + %Payload{status: :idle, conversation_state: %{executor_state: :nonexistent}}} + + refute_receive :llm_requested, 100 + end + + test "resume/2 returns not_resumable for an agent_id the store has no run for" do + assert {:error, :not_resumable} = Legion.resume("ghost", store: MemoryStore) + end + + test "resume/2 returns not_resumable when the stored run has no agent module" do + assert :ok = + MemoryStore.save(%Payload{ + agent_id: "resume-missing-agent-module", + agent_module: nil + }) + + assert {:error, :not_resumable} = + Legion.resume("resume-missing-agent-module", store: MemoryStore) + end + + test "resume/2 identifies itself when no store is configured" do + assert_raise ArgumentError, ~r/resume\/2 requires a :store/, fn -> + Legion.resume("missing-store") + end + end + + test "recover/2 identifies itself when no store is configured" do + assert_raise ArgumentError, ~r/recover\/2 requires a :store/, fn -> + Legion.recover("missing-store") end end + test "recover/2 completes an interrupted run and stops its process" do + assert :ok = + MemoryStore.save(%Payload{ + agent_id: "recover-awaiting-llm", + parent_agent_id: nil, + agent_module: MathAgent, + status: :running, + conversation_state: %{ + messages: [%{role: "user", type: :user, content: "compute"}], + bindings: [x: 42], + executor_state: %{phase: :awaiting_llm, iteration: 1, retries: 0} + } + }) + + assert :ok = Legion.recover("recover-awaiting-llm", store: MemoryStore) + + assert {:ok, %Payload{status: :idle, conversation_state: %{executor_state: :nonexistent}}} = + MemoryStore.get("recover-awaiting-llm") + + assert( + case Legion.lookup("recover-awaiting-llm") do + :error -> true + {:ok, pid} -> not Process.alive?(pid) + end + ) + end + + test "recover/2 returns error when agent is running" do + {:ok, pid} = Legion.start_link(MathAgent, store: MemoryStore, agent_id: "recover-running") + + assert :ok = + MemoryStore.save(%Payload{ + agent_id: "recover-running", + status: :running, + conversation_state: %{ + messages: [%{role: "user", type: :user, content: "recover me"}], + bindings: [], + executor_state: :nonexistent + } + }) + + assert {:error, :already_running} = Legion.recover("recover-running", store: MemoryStore) + + assert Process.alive?(pid) + end + + test "recover/2 validates the requested store before resolving a live process" do + {:ok, _pid} = + Legion.start_link(MathAgent, store: MemoryStore, agent_id: "recover-wrong-store") + + assert {:error, :not_recoverable} = + Legion.recover("recover-wrong-store", store: EmptyStore) + end + + test "recover/2 returns not_recoverable for an agent_id the store has no run for" do + assert {:error, :not_recoverable} = + Legion.recover("recover-missing", store: MemoryStore) + end + + test "recover/2 returns not_recoverable when the stored run has no agent module" do + assert :ok = + MemoryStore.save(%Payload{ + agent_id: "recover-missing-agent-module", + agent_module: nil, + status: :running + }) + + assert {:error, :not_recoverable} = + Legion.recover("recover-missing-agent-module", store: MemoryStore) + end + + test "recover/2 refuses an idle run" do + assert :ok = + MemoryStore.save(%Payload{ + agent_id: "recover-idle-root", + parent_agent_id: nil, + agent_module: MathAgent, + status: :idle, + conversation_state: %{ + messages: [%{role: "user", type: :user, content: "compute"}], + bindings: [x: 42], + executor_state: %{phase: :completing, iteration: 1, retries: 0} + } + }) + + assert {:error, :not_recoverable} = Legion.recover("recover-idle-root", store: MemoryStore) + end + + test "recover/2 refuses a running child run" do + assert :ok = + MemoryStore.save(%Payload{ + agent_id: "recover-running-child", + parent_agent_id: "recover-parent", + agent_module: MathAgent, + status: :running, + conversation_state: %{ + messages: [%{role: "user", type: :user, content: "compute"}], + bindings: [x: 42], + executor_state: %{phase: :completing, iteration: 1, retries: 0} + } + }) + + assert {:error, :not_recoverable} = + Legion.recover("recover-running-child", store: MemoryStore) + end + test "sub-agents inherit the parent store and link to the parent conversation" do stub(ReqLLM, :generate_object, fn _model, messages, _schema -> if Enum.any?(messages, &(&1[:content] == "child task")) do diff --git a/test/legion/agent_test.exs b/test/legion/agent_test.exs index a1876d8..617bbae 100644 --- a/test/legion/agent_test.exs +++ b/test/legion/agent_test.exs @@ -70,10 +70,10 @@ defmodule Legion.AgentTest do end test "passes opts through to start args" do - spec = MinimalAgent.child_spec(name: :my_agent, model: "openai:gpt-4o") + spec = MinimalAgent.child_spec(model: "openai:gpt-4o", max_iterations: 5) assert spec.start == - {Legion, :start_link, [MinimalAgent, [name: :my_agent, model: "openai:gpt-4o"]]} + {Legion, :start_link, [MinimalAgent, [model: "openai:gpt-4o", max_iterations: 5]]} end end diff --git a/test/legion/application_test.exs b/test/legion/application_test.exs new file mode 100644 index 0000000..a235bf9 --- /dev/null +++ b/test/legion/application_test.exs @@ -0,0 +1,36 @@ +defmodule Legion.ApplicationTest do + use ExUnit.Case, async: false + + setup do + previous = Application.fetch_env(:legion, :recovery) + + on_exit(fn -> + case previous do + {:ok, config} -> Application.put_env(:legion, :recovery, config) + :error -> Application.delete_env(:legion, :recovery) + end + end) + end + + test "passes absent recovery configuration to the recovery child" do + Application.delete_env(:legion, :recovery) + + assert {Legion.Recovery, :error} in Legion.Application.children() + end + + test "does not start a local registry for agent names" do + refute Process.whereis(Legion.AgentRegistry) + end + + test "passes configured recovery options to the recovery child" do + config = [ + stores: [RecoveryStore], + store_scan_limit: 3, + concurrent_request_limit: 2 + ] + + Application.put_env(:legion, :recovery, config) + + assert Enum.member?(Legion.Application.children(), {Legion.Recovery, {:ok, config}}) + end +end diff --git a/test/legion/executor_test.exs b/test/legion/executor_test.exs index 20b87af..60a8d9c 100644 --- a/test/legion/executor_test.exs +++ b/test/legion/executor_test.exs @@ -278,7 +278,7 @@ defmodule Legion.ExecutorTest do %{ messages: continuing_messages, bindings: [x: 10], - execution: %{phase: :awaiting_llm, iteration: 1, retries: 0} + executor_state: %{phase: :awaiting_llm, iteration: 1, retries: 0} }} assert List.last(continuing_messages).type == :eval_result @@ -287,7 +287,7 @@ defmodule Legion.ExecutorTest do %{ messages: completing_messages, bindings: [x: 10], - execution: %{phase: :completing, iteration: 1, retries: 0} + executor_state: %{phase: :completing, iteration: 1, retries: 0} }} assert List.last(completing_messages).type == :eval_result @@ -329,7 +329,7 @@ defmodule Legion.ExecutorTest do %{ messages: messages, bindings: [], - execution: %{phase: :awaiting_llm, iteration: 0, retries: 1} + executor_state: %{phase: :awaiting_llm, iteration: 0, retries: 1} }} assert List.last(messages).type == :error diff --git a/test/legion/recovery_test.exs b/test/legion/recovery_test.exs new file mode 100644 index 0000000..592ae74 --- /dev/null +++ b/test/legion/recovery_test.exs @@ -0,0 +1,263 @@ +defmodule Legion.RecoveryTest do + use ExUnit.Case, async: false + use Mimic + + alias Legion.Store.Payload + + defmodule RecoveryAgent do + @moduledoc "Agent used to exercise startup recovery." + use Legion.Agent + end + + defmodule StoreState do + use Agent + + def start_link(test_pid) do + Agent.start_link(fn -> %{stores: %{}, test_pid: test_pid} end, name: __MODULE__) + end + + def put(store, payloads) do + Agent.update(__MODULE__, fn state -> + payloads = Map.new(payloads, &{&1.agent_id, &1}) + put_in(state, [:stores, store], payloads) + end) + end + + def get(store, agent_id) do + {test_pid, result} = + Agent.get(__MODULE__, fn state -> + result = + case get_in(state, [:stores, store, agent_id]) do + nil -> :error + payload -> {:ok, payload} + end + + {state.test_pid, result} + end) + + send(test_pid, {:looked_up, store, agent_id}) + result + end + + def list(store, limit) do + {test_pid, payloads} = + Agent.get(__MODULE__, fn state -> + {state.test_pid, state.stores |> Map.fetch!(store) |> Map.values() |> Enum.take(limit)} + end) + + send(test_pid, {:listed, store, limit}) + payloads + end + + def save(store, %Payload{} = payload) do + Agent.update(__MODULE__, fn state -> + existing = get_in(state, [:stores, store, payload.agent_id]) + + merged = + Enum.reduce(Map.from_struct(payload), existing, fn + {:agent_id, _agent_id}, stored -> stored + {_field, nil}, stored -> stored + {field, value}, stored -> Map.put(stored, field, value) + end) + + put_in(state, [:stores, store, payload.agent_id], merged) + end) + + :ok + end + end + + defmodule RecoveryStoreOne do + @behaviour Legion.Store + + @impl Legion.Store + def get(agent_id), do: StoreState.get(__MODULE__, agent_id) + + @impl Legion.Store + def list(limit), do: StoreState.list(__MODULE__, limit) + + @impl Legion.Store + def save(payload), do: StoreState.save(__MODULE__, payload) + end + + defmodule RecoveryStoreTwo do + @behaviour Legion.Store + + @impl Legion.Store + def get(agent_id), do: StoreState.get(__MODULE__, agent_id) + + @impl Legion.Store + def list(limit), do: StoreState.list(__MODULE__, limit) + + @impl Legion.Store + def save(payload), do: StoreState.save(__MODULE__, payload) + end + + setup :set_mimic_global + + setup do + start_supervised!({StoreState, self()}) + :ok + end + + @moduletag capture_log: true + + test "does not start when recovery config is absent" do + assert :ignore = Legion.Recovery.start_link(:error) + end + + test "uses separate store scan and concurrent request limits" do + stores = [RecoveryStoreOne, RecoveryStoreTwo] + + Enum.each(stores, fn store -> + StoreState.put(store, [ + interrupted_payload("#{store}-one"), + interrupted_payload("#{store}-two") + ]) + end) + + test_pid = self() + + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + send(test_pid, {:recovering, self()}) + + receive do + :complete_recovery -> llm_response("recovered") + end + end) + + assert {:ok, worker} = + Legion.Recovery.start_link( + {:ok, stores: stores, store_scan_limit: 2, concurrent_request_limit: 1} + ) + + monitor_ref = Process.monitor(worker) + + assert_receive {:listed, RecoveryStoreOne, 2} + assert_receive {:listed, RecoveryStoreTwo, 2} + + assert_receive {:recovering, first} + refute_receive {:recovering, _}, 50 + + send(first, :complete_recovery) + + assert_receive {:recovering, second} + refute_receive {:recovering, _}, 50 + + send(second, :complete_recovery) + + assert_receive {:recovering, third} + refute_receive {:recovering, _}, 50 + + send(third, :complete_recovery) + + assert_receive {:recovering, fourth} + refute_receive {:recovering, _}, 50 + + send(fourth, :complete_recovery) + + assert_receive {:DOWN, ^monitor_ref, :process, ^worker, :normal} + + for store <- stores, + suffix <- ["one", "two"] do + assert {:ok, %Payload{status: :idle}} = StoreState.get(store, "#{store}-#{suffix}") + end + end + + test "defaults the concurrent request limit to three" do + StoreState.put(RecoveryStoreOne, [ + interrupted_payload("one"), + interrupted_payload("two"), + interrupted_payload("three"), + interrupted_payload("four") + ]) + + test_pid = self() + + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + send(test_pid, {:recovering, self()}) + + receive do + :complete_recovery -> llm_response("recovered") + end + end) + + assert {:ok, worker} = + Legion.Recovery.start_link({:ok, stores: [RecoveryStoreOne], store_scan_limit: 4}) + + monitor_ref = Process.monitor(worker) + + assert_receive {:recovering, first} + assert_receive {:recovering, second} + assert_receive {:recovering, third} + refute_receive {:recovering, _}, 50 + + send(first, :complete_recovery) + + assert_receive {:recovering, fourth} + + send(second, :complete_recovery) + send(third, :complete_recovery) + send(fourth, :complete_recovery) + + assert_receive {:DOWN, ^monitor_ref, :process, ^worker, :normal} + end + + test "filters out idle and child runs before recovering" do + eligible = interrupted_payload("eligible") + idle_root = %{interrupted_payload("idle-root") | status: :idle} + running_child = %{interrupted_payload("running-child") | parent_agent_id: "parent"} + + StoreState.put(RecoveryStoreOne, [eligible, idle_root, running_child]) + + test_pid = self() + + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + send(test_pid, {:recovering, self()}) + + receive do + :complete_recovery -> llm_response("recovered") + end + end) + + assert {:ok, worker} = + Legion.Recovery.start_link( + {:ok, stores: [RecoveryStoreOne], store_scan_limit: 3, concurrent_request_limit: 3} + ) + + monitor_ref = Process.monitor(worker) + + assert_receive {:listed, RecoveryStoreOne, 3} + assert_receive {:recovering, recovery_pid} + refute_receive {:looked_up, RecoveryStoreOne, "idle-root"}, 50 + refute_receive {:looked_up, RecoveryStoreOne, "running-child"}, 50 + + send(recovery_pid, :complete_recovery) + + assert_receive {:DOWN, ^monitor_ref, :process, ^worker, :normal} + end + + defp interrupted_payload(agent_id) do + %Payload{ + agent_id: agent_id, + parent_agent_id: nil, + agent_module: RecoveryAgent, + status: :running, + conversation_state: %{ + messages: [%{role: "user", type: :user, content: "recover me"}], + bindings: [], + executor_state: %{phase: :awaiting_llm, iteration: 1, retries: 0} + } + } + end + + defp llm_response(result) do + {:ok, + %ReqLLM.Response{ + id: "test", + model: "test", + context: nil, + object: %{"action" => "return", "code" => "", "result" => result} + }} + end +end diff --git a/test/legion/store/postgres_db_test.exs b/test/legion/store/postgres_db_test.exs index 4463a35..92928e1 100644 --- a/test/legion/store/postgres_db_test.exs +++ b/test/legion/store/postgres_db_test.exs @@ -21,7 +21,11 @@ defmodule Legion.Store.PostgresDbTest do parent_agent_id: "parent-1", status: :idle, started_at: ~N[2026-01-01 00:00:00.000000], - conversation_state: %{messages: [%{role: "user", content: "hi"}], bindings: [x: 42]} + conversation_state: %{ + messages: [%{role: "user", content: "hi"}], + bindings: [x: 42], + executor_state: :nonexistent + } } assert :ok = Store.save(payload) @@ -31,7 +35,11 @@ defmodule Legion.Store.PostgresDbTest do test "save/1 partially inserts only the supplied payload fields" do payload = %Payload{ agent_id: "state-only", - conversation_state: %{messages: [%{role: "user", content: "hi"}], bindings: []} + conversation_state: %{ + messages: [%{role: "user", content: "hi"}], + bindings: [], + executor_state: :nonexistent + } } assert :ok = Store.save(payload) @@ -46,7 +54,11 @@ defmodule Legion.Store.PostgresDbTest do parent_agent_id: "parent-1", status: :running, started_at: ~N[2026-01-01 00:00:00.000000], - conversation_state: %{messages: [%{role: "user", content: "hi"}], bindings: [x: 42]} + conversation_state: %{ + messages: [%{role: "user", content: "hi"}], + bindings: [x: 42], + executor_state: :nonexistent + } } assert :ok = Store.save(initial) @@ -63,7 +75,11 @@ defmodule Legion.Store.PostgresDbTest do parent_agent_id: "parent-1", status: :idle, started_at: ~N[2026-01-01 00:00:00.000000], - conversation_state: %{messages: [%{role: "user", content: "hi"}], bindings: [x: 42]} + conversation_state: %{ + messages: [%{role: "user", content: "hi"}], + bindings: [x: 42], + executor_state: :nonexistent + } }} = Store.get("user_42") %{rows: [[updated_at]]} = diff --git a/test/legion/store/postgres_test.exs b/test/legion/store/postgres_test.exs index e1801b9..e255baa 100644 --- a/test/legion/store/postgres_test.exs +++ b/test/legion/store/postgres_test.exs @@ -70,7 +70,11 @@ defmodule Legion.Store.PostgresTest do parent_agent_id: "parent-1", status: :idle, started_at: 123, - conversation_state: %{messages: [%{role: "user", content: "hi"}], bindings: [x: 42]} + conversation_state: %{ + messages: [%{role: "user", content: "hi"}], + bindings: [x: 42], + executor_state: :nonexistent + } } assert :ok = Store.save(payload) @@ -80,7 +84,11 @@ defmodule Legion.Store.PostgresTest do test "save/1 partially inserts only the supplied payload fields" do payload = %Payload{ agent_id: "state-only", - conversation_state: %{messages: [%{role: "user", content: "hi"}], bindings: []} + conversation_state: %{ + messages: [%{role: "user", content: "hi"}], + bindings: [], + executor_state: :nonexistent + } } expected_payload = %{payload | status: :idle} @@ -88,8 +96,8 @@ defmodule Legion.Store.PostgresTest do assert {:ok, ^expected_payload} = Store.get("state-only") end - test "save/1 round trips step execution state" do - execution = %{phase: :awaiting_llm, iteration: 2, retries: 1} + test "save/1 round trips executor_state for a step checkpoint" do + executor_state = %{phase: :awaiting_llm, iteration: 2, retries: 1} payload = %Payload{ agent_id: "step-state", @@ -97,7 +105,7 @@ defmodule Legion.Store.PostgresTest do conversation_state: %{ messages: [%{role: "user", content: "result"}], bindings: [x: 42], - execution: execution + executor_state: executor_state } } @@ -112,7 +120,11 @@ defmodule Legion.Store.PostgresTest do parent_agent_id: "parent-1", status: :running, started_at: 123, - conversation_state: %{messages: [%{role: "user", content: "hi"}], bindings: [x: 42]} + conversation_state: %{ + messages: [%{role: "user", content: "hi"}], + bindings: [x: 42], + executor_state: :nonexistent + } } assert :ok = Store.save(initial) @@ -126,7 +138,11 @@ defmodule Legion.Store.PostgresTest do parent_agent_id: "parent-1", status: :idle, started_at: 123, - conversation_state: %{messages: [%{role: "user", content: "hi"}], bindings: [x: 42]} + conversation_state: %{ + messages: [%{role: "user", content: "hi"}], + bindings: [x: 42], + executor_state: :nonexistent + } }} = Store.get("user_42") assert NaiveDateTime.compare(FakeRepo.run("user_42").updated_at, previous_updated_at) == :gt