From f0797149a492ff9449613a1651bd027e671a2320 Mon Sep 17 00:00:00 2001 From: Dima Mikielewicz Date: Mon, 6 Jul 2026 22:27:36 +0200 Subject: [PATCH 01/30] Add Legion.Store conversation persistence with a built-in Postgres adapter --- CHANGELOG.md | 7 +++ README.md | 15 +++++- lib/legion.ex | 2 + lib/legion/agent_server.ex | 37 ++++++++++--- lib/legion/store.ex | 68 +++++++++++++++++++++++ lib/legion/store/postgres.ex | 75 ++++++++++++++++++++++++++ mix.exs | 2 +- test/legion/agent_server_test.exs | 84 +++++++++++++++++++++++++++++ test/legion/store/postgres_test.exs | 62 +++++++++++++++++++++ 9 files changed, 343 insertions(+), 9 deletions(-) create mode 100644 lib/legion/store.ex create mode 100644 lib/legion/store/postgres.ex create mode 100644 test/legion/store/postgres_test.exs diff --git a/CHANGELOG.md b/CHANGELOG.md index 8d190b3..a91fcbf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,12 @@ # Changelog +## Unreleased + +### Changes + +- Add `Legion.Store` behaviour for persisting conversations across restarts. Pass `store:` and `agent_id:` to `Legion.start_link/2`; snapshots (messages + bindings) are saved after every completed turn, before the caller receives its reply +- Add `Legion.Store.Postgres`, a ready-made store adapter that reuses your Ecto repo (`use Legion.Store.Postgres, repo: MyApp.Repo`) without adding Ecto as a dependency + ## v0.4.0 - 2026-05-17 ### Security diff --git a/README.md b/README.md index ebe353c..d6521e5 100644 --- a/README.md +++ b/README.md @@ -7,9 +7,17 @@ -An Elixir framework for building AI agents that write and execute code instead of making function calls. +An Elixir framework for AI agents that live inside your application and act by writing code. -Traditional agents call tools one at a time - fetch, wait, decide, fetch again - burning tokens and latency on every round-trip. Legion agents write Elixir code that fetches, filters, decides, and acts in a single step, running safely in a sandbox. Fewer LLM calls, smarter behavior, full language expressivity. [Why code execution beats function calling.](https://www.anthropic.com/engineering/code-execution-with-mcp) +Legion is a runtime, not a coding assistant. No human reads, reviews, or commits the code its agents write. An agent writes Elixir the way you would write a shell one-liner: to do something, right now. Given a task, it reads your tools' source, writes a snippet that fetches, filters, decides, and acts, executes it, observes the result, and continues - all inside your running application. + +That single move replaces function calling. A traditional agent pays an LLM round-trip for every tool call: fetch, wait, decide, call the next one. A Legion agent collapses that loop into one evaluation, with pipelines, pattern matching, and the standard library at its disposal. Fewer LLM calls, lower latency, smarter behavior. [Why code execution beats function calling.](https://www.anthropic.com/engineering/code-execution-with-mcp) + +The whole framework fits in three ideas: + +- **Tools are plain Elixir modules.** The LLM reads their source directly - no schemas, no wrappers, no glue. +- **Agents are BEAM processes.** Supervise them, pool them, `call` and `cast` them like GenServers. +- **Generated code runs in a sandbox.** Dangerous constructs are blocked at the AST level; module access is allowlisted. ## Quick Start @@ -85,6 +93,7 @@ A traditional agent would need a separate LLM call for each filter decision and - **Tools are just modules** - `use Legion.Tool` on any module to expose it. The LLM reads your source code and calls your functions. No schemas to write, no wrappers - reuse existing app logic directly. - **Authorization via Vault** - Set auth context before the agent starts, validate inside tools at runtime. LLM-generated code never touches credentials. See [Vault](https://github.com/dimamik/vault). - **Long-lived agents** - Start agents with `Legion.start_link/2` and message them with `call/2` and `cast/2`, just like a GenServer. Variables can persist across turns with `binding_scope: :conversation`. +- **Persistence** - Make conversations survive crashes, restarts, and deploys with `use Legion.Store.Postgres, repo: MyApp.Repo` (reuses your Ecto repo, one table), or implement the two-callback `Legion.Store` behaviour for any other storage. Snapshots are saved after every completed turn, before the caller sees its reply - a reply is a commit receipt. - **Multi-agent orchestration** - Agents delegate to other agents via the built-in `AgentTool`. Fan out with `parallel/2`, chain with `pipeline/1`. Sub-agents are linked processes - when a parent dies, children stop too. - **Human in the loop** - The built-in `HumanTool` pauses agent execution until a human responds. It's just message passing - your handler receives a question and sends back an answer. - **Structured output** - Define a JSON Schema via `output_schema/0` to get typed, validated responses. Or skip it and work with plain text. @@ -215,8 +224,10 @@ config :legion, :config, %{ | Option | Description | | -------------------- | -------------------------------------------------------------------------------------------------------------------------- | +| `model` | LLM model string passed to [ReqLLM](https://hexdocs.pm/req_llm), e.g. `"openai:gpt-4o-mini"`. | | `max_iterations` | Successful execution steps before the agent is stopped. | | `max_retries` | Consecutive failures (bad code, tool errors) before giving up. Resets after each success. | +| `sandbox_timeout` | Milliseconds a single code evaluation may run before it is killed. | | `binding_scope` | `:iteration` (fresh each step), `:turn` (persist within a message, default), or `:conversation` (persist across messages). | | `max_message_length` | Byte limit for any single message. Longer content is truncated. Set to `:infinity` to disable. | diff --git a/lib/legion.ex b/lib/legion.ex index bd3f604..33b9ba6 100644 --- a/lib/legion.ex +++ b/lib/legion.ex @@ -29,12 +29,14 @@ defmodule Legion do ## Options - `:name` - register the process under a name + - `:store`, `:agent_id` - persist the conversation across restarts; see `Legion.Store` - 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(AssistantAgent, store: MyApp.AgentStore, agent_id: "user_42") """ def start_link(agent_module, opts \\ []) do AgentServer.start_link(agent_module, opts) diff --git a/lib/legion/agent_server.ex b/lib/legion/agent_server.ex index 80dde8f..8ebb86f 100644 --- a/lib/legion/agent_server.ex +++ b/lib/legion/agent_server.ex @@ -13,15 +13,22 @@ defmodule Legion.AgentServer do alias Legion.{Executor, Telemetry} alias ReqLLM.Message.ContentPart - defstruct [:agent_module, :messages, :config, bindings: []] + defstruct [:agent_module, :messages, :config, :store, :agent_id, bindings: []] # Client API def start_link(agent_module, opts \\ []) do {name, opts} = Keyword.pop(opts, :name) + {store, opts} = Keyword.pop(opts, :store) + {agent_id, opts} = Keyword.pop(opts, :agent_id) + + if is_nil(store) != is_nil(agent_id) do + raise ArgumentError, ":store and :agent_id must be given together" + end + gen_opts = if name, do: [name: name], else: [] config = resolve_config(agent_module, opts) - GenServer.start_link(__MODULE__, {agent_module, config}, gen_opts) + GenServer.start_link(__MODULE__, {agent_module, config, store, agent_id}, gen_opts) end def call(agent, message, timeout \\ :infinity) do @@ -39,7 +46,7 @@ defmodule Legion.AgentServer do # Server callbacks @impl true - def init({agent_module, config}) do + def init({agent_module, config, store, agent_id}) do parent_run_id = Vault.get(:run_id) run_id = make_ref() @@ -58,10 +65,19 @@ defmodule Legion.AgentServer do %{agent: agent_module} ) + {saved_messages, saved_bindings} = + case store && store.load(agent_id) do + {:ok, %{messages: messages, bindings: bindings}} -> {messages, bindings} + _no_snapshot -> {[], []} + end + state = %__MODULE__{ agent_module: agent_module, - messages: [%{role: "system", content: system_prompt}], - config: config + messages: [%{role: "system", content: system_prompt} | saved_messages], + config: config, + store: store, + agent_id: agent_id, + bindings: saved_bindings } {:ok, state} @@ -130,7 +146,16 @@ defmodule Legion.AgentServer do end ) - {{status, value}, %{state | messages: final_messages, bindings: final_bindings}} + state = persist(%{state | messages: final_messages, bindings: final_bindings}) + {{status, value}, state} + end + + defp persist(%{store: nil} = state), do: state + + defp persist(state) do + [%{role: "system"} | messages] = state.messages + :ok = state.store.save(state.agent_id, %{messages: messages, bindings: state.bindings}) + state end defp stringify(message, max_length) when is_binary(message), diff --git a/lib/legion/store.ex b/lib/legion/store.ex new file mode 100644 index 0000000..3ae56d6 --- /dev/null +++ b/lib/legion/store.ex @@ -0,0 +1,68 @@ +defmodule Legion.Store do + @moduledoc """ + Behaviour for persisting agent conversations across restarts. + + Legion decides *when* to persist; your store decides *where*. Pass a store + module and an agent id when starting an agent: + + {:ok, pid} = Legion.start_link(AssistantAgent, store: MyApp.AgentStore, agent_id: "user_42") + + On start, the agent calls `c:load/1` and resumes from the snapshot if one + exists. The system prompt is regenerated fresh on every start, so prompt or + tool changes apply to restored conversations. + + After every completed turn, the agent calls `c:save/2` **before** replying + to the caller. A reply is a commit receipt: any turn a caller observed + survives a crash, restart, or deploy. A crash mid-turn rolls back to the + last completed turn. + + For Postgres users there is a ready-made adapter - see `Legion.Store.Postgres`: + + defmodule MyApp.AgentStore do + use Legion.Store.Postgres, repo: MyApp.Repo + end + + Or implement the two callbacks against any storage you like: + + ## Example: hand-rolled Ecto store + + defmodule MyApp.AgentStore do + @behaviour Legion.Store + + def load(agent_id) do + case MyApp.Repo.get(MyApp.AgentSnapshot, agent_id) do + nil -> :error + row -> {:ok, :erlang.binary_to_term(row.snapshot)} + end + end + + def save(agent_id, snapshot) do + MyApp.Repo.insert!( + %MyApp.AgentSnapshot{id: agent_id, snapshot: :erlang.term_to_binary(snapshot)}, + on_conflict: {:replace, [:snapshot]}, + conflict_target: :id + ) + + :ok + end + end + + ## What is persisted + + The snapshot holds the conversation `:messages` (without the system prompt) + and the `:bindings` from evaluated code (relevant with + `binding_scope: :conversation`). Bindings are arbitrary Elixir terms - + values like pids, references, or functions will not survive + serialization, so keep conversation-scoped variables to plain data if you + persist agents. + """ + + @type agent_id :: term() + @type snapshot :: %{messages: [map()], bindings: keyword()} + + @doc "Returns the last saved snapshot for `agent_id`, or `:error` if none exists." + @callback load(agent_id()) :: {:ok, snapshot()} | :error + + @doc "Saves the snapshot for `agent_id`. Raise on failure - the turn is not acked until this returns." + @callback save(agent_id(), snapshot()) :: :ok +end diff --git a/lib/legion/store/postgres.ex b/lib/legion/store/postgres.ex new file mode 100644 index 0000000..7012b73 --- /dev/null +++ b/lib/legion/store/postgres.ex @@ -0,0 +1,75 @@ +defmodule Legion.Store.Postgres do + @moduledoc """ + A ready-made `Legion.Store` backed by Postgres, through your existing Ecto repo. + + Legion does not depend on Ecto - the generated store only calls + `repo.query!/2` at runtime, so it works with any `Ecto.Repo` on + `Ecto.Adapters.Postgres` that your application already runs. + + ## Usage + + Define a store module: + + defmodule MyApp.AgentStore do + use Legion.Store.Postgres, repo: MyApp.Repo + end + + Create the table in a migration: + + defmodule MyApp.Repo.Migrations.AddLegionAgents do + use Ecto.Migration + + def change do + create table(:legion_agents, primary_key: false) do + add :agent_id, :text, primary_key: true + add :snapshot, :binary, null: false + timestamps(type: :timestamptz) + end + end + end + + Then start agents with it: + + {:ok, pid} = Legion.start_link(AssistantAgent, store: MyApp.AgentStore, agent_id: "user_42") + + ## Options + + - `:repo` (required) - your Ecto repo module + - `:table` - the table name, defaults to `"legion_agents"` + + Agent ids must be strings. Snapshots are stored as `:erlang.term_to_binary/1` + blobs - readable only from Elixir, one row per agent, upserted on every turn. + """ + + defmacro __using__(opts) do + repo = Keyword.fetch!(opts, :repo) + table = Keyword.get(opts, :table, "legion_agents") + + select_sql = "SELECT snapshot FROM #{table} WHERE agent_id = $1" + + upsert_sql = """ + INSERT INTO #{table} (agent_id, snapshot, inserted_at, updated_at) + VALUES ($1, $2, now(), now()) + ON CONFLICT (agent_id) DO UPDATE SET snapshot = EXCLUDED.snapshot, updated_at = now() + """ + + quote do + @behaviour Legion.Store + + @impl Legion.Store + # sobelow_skip ["Misc.BinToTerm"] + def load(agent_id) when is_binary(agent_id) do + case unquote(repo).query!(unquote(select_sql), [agent_id]) do + %{rows: [[snapshot]]} -> {:ok, :erlang.binary_to_term(snapshot)} + %{rows: []} -> :error + end + end + + @impl Legion.Store + def save(agent_id, snapshot) when is_binary(agent_id) do + unquote(repo).query!(unquote(upsert_sql), [agent_id, :erlang.term_to_binary(snapshot)]) + :ok + end + end + end +end diff --git a/mix.exs b/mix.exs index 7b6af18..c2f1999 100644 --- a/mix.exs +++ b/mix.exs @@ -48,7 +48,7 @@ defmodule Legion.MixProject do defp groups_for_modules do [ - Core: [Legion, Legion.Agent, Legion.Tool], + Core: [Legion, Legion.Agent, Legion.Tool, Legion.Store, Legion.Store.Postgres], Runtime: [Legion.AgentServer, Legion.Executor, ~r/^Legion\.Sandbox/], Tools: [~r/^Legion\.Tools\./], Internals: [Legion.AgentPrompt, Legion.SourceRegistry, Legion.Telemetry] diff --git a/test/legion/agent_server_test.exs b/test/legion/agent_server_test.exs index 3c77dbd..f08ba22 100644 --- a/test/legion/agent_server_test.exs +++ b/test/legion/agent_server_test.exs @@ -397,6 +397,90 @@ defmodule Legion.AgentServerTest do end end + defmodule MemoryStore do + @behaviour Legion.Store + + def start_link, do: Agent.start_link(fn -> %{} end, name: __MODULE__) + + def load(agent_id), do: Agent.get(__MODULE__, &Map.fetch(&1, agent_id)) + + def save(agent_id, snapshot) do + Agent.update(__MODULE__, &Map.put(&1, agent_id, snapshot)) + end + end + + describe "persistence" do + setup do + start_supervised!(%{id: MemoryStore, start: {MemoryStore, :start_link, []}}) + :ok + end + + test "saves a snapshot before the caller receives its reply" do + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + llm_response("Paris") + end) + + {:ok, pid} = Legion.start_link(MathAgent, store: MemoryStore, agent_id: "receipt") + {:ok, _} = Legion.call(pid, "What is the capital of France?") + + assert {:ok, %{messages: messages, bindings: []}} = MemoryStore.load("receipt") + + assert [ + %{role: "user", content: "What is the capital of France?"}, + %{role: "assistant"} | _ + ] = messages + + refute Enum.any?(messages, &(&1.role == "system")) + end + + test "restores the conversation under a fresh system prompt after a restart" do + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + llm_response("Paris") + end) + + {:ok, pid} = Legion.start_link(MathAgent, store: MemoryStore, agent_id: "restore") + {:ok, _} = Legion.call(pid, "What is the capital of France?") + GenServer.stop(pid) + + {:ok, revived} = Legion.start_link(MathAgent, store: MemoryStore, agent_id: "restore") + + assert [ + %{role: "system"}, + %{role: "user", content: "What is the capital of France?"}, + %{role: "assistant"} | _ + ] = Legion.get_messages(revived) + end + + test "restores conversation-scoped bindings after a restart" do + stub(ReqLLM, :generate_object, fn _model, messages, _schema -> + assistant_count = Enum.count(messages, &(&1[:role] == "assistant")) + + if assistant_count == 0 do + llm_eval_response("x = 42") + else + llm_eval_response("x + 1") + end + end) + + {:ok, pid} = + Legion.start_link(ConversationBindingsAgent, store: MemoryStore, agent_id: "bindings") + + {:ok, 42} = Legion.call(pid, "set x") + GenServer.stop(pid) + + {:ok, revived} = + Legion.start_link(ConversationBindingsAgent, store: MemoryStore, agent_id: "bindings") + + assert {:ok, 43} = Legion.call(revived, "use x") + end + + test "raises when :store is given without :agent_id" do + assert_raise ArgumentError, ":store and :agent_id must be given together", fn -> + Legion.start_link(MathAgent, store: MemoryStore) + end + end + end + describe "binding_scope" do test "bindings do not persist across turns by default (:turn)" do stub(ReqLLM, :generate_object, fn _model, messages, _schema -> diff --git a/test/legion/store/postgres_test.exs b/test/legion/store/postgres_test.exs new file mode 100644 index 0000000..583fd3b --- /dev/null +++ b/test/legion/store/postgres_test.exs @@ -0,0 +1,62 @@ +defmodule Legion.Store.PostgresTest do + use ExUnit.Case, async: true + + defmodule FakeRepo do + @moduledoc "Emulates repo.query!/2 for the two statements the store issues." + + def start_link, do: Agent.start_link(fn -> %{} end, name: __MODULE__) + + def query!("SELECT snapshot FROM " <> _rest, [agent_id]) do + case Agent.get(__MODULE__, &Map.fetch(&1, agent_id)) do + {:ok, snapshot} -> %{rows: [[snapshot]]} + :error -> %{rows: []} + end + end + + def query!("INSERT INTO " <> _rest, [agent_id, snapshot]) do + Agent.update(__MODULE__, &Map.put(&1, agent_id, snapshot)) + %{num_rows: 1} + end + end + + defmodule Store do + use Legion.Store.Postgres, repo: Legion.Store.PostgresTest.FakeRepo + end + + setup do + start_supervised!(%{id: FakeRepo, start: {FakeRepo, :start_link, []}}) + :ok + end + + test "save/2 then load/1 round-trips the snapshot through term_to_binary" do + snapshot = %{messages: [%{role: "user", content: "hi"}], bindings: [x: 42]} + + assert :ok = Store.save("user_42", snapshot) + assert {:ok, ^snapshot} = Store.load("user_42") + end + + test "load/1 returns :error when no snapshot exists" do + assert :error = Store.load("missing") + end + + test "ids must be strings" do + assert_raise FunctionClauseError, fn -> Store.load(42) end + assert_raise FunctionClauseError, fn -> Store.save(42, %{messages: [], bindings: []}) end + end + + test "a custom table name is interpolated into the statements" do + defmodule TableCapturingRepo do + def query!(sql, _params), do: send(self(), {:sql, sql}) && %{rows: []} + end + + defmodule CustomTableStore do + use Legion.Store.Postgres, + repo: Legion.Store.PostgresTest.TableCapturingRepo, + table: "my_agents" + end + + CustomTableStore.load("user_42") + + assert_received {:sql, "SELECT snapshot FROM my_agents WHERE agent_id = $1"} + end +end From be6b8737e594a33681afaf43637bece4466dbf34 Mon Sep 17 00:00:00 2001 From: Dima Mikielewicz Date: Thu, 9 Jul 2026 09:49:18 +0200 Subject: [PATCH 02/30] Review feedback and minor flow improvements --- .github/workflows/ci.yml | 12 +++++ lib/legion.ex | 30 ++++++++++-- lib/legion/agent_server.ex | 29 +++++++++--- lib/legion/store.ex | 22 +++++++++ mix.exs | 3 +- mix.lock | 3 ++ test/legion/agent_server_test.exs | 64 ++++++++++++++++++++++++-- test/legion/store/postgres_db_test.exs | 39 ++++++++++++++++ test/test_helper.exs | 25 ++++++++++ 9 files changed, 212 insertions(+), 15 deletions(-) create mode 100644 test/legion/store/postgres_db_test.exs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 07cab64..68f3090 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -7,6 +7,18 @@ jobs: lint_and_test: name: Lint and Test runs-on: ubuntu-latest + services: + postgres: + image: postgres:17 + env: + POSTGRES_PASSWORD: postgres + ports: + - 5432:5432 + options: >- + --health-cmd pg_isready + --health-interval 10s + --health-timeout 5s + --health-retries 5 steps: - uses: actions/checkout@v4 - uses: erlef/setup-beam@v1 diff --git a/lib/legion.ex b/lib/legion.ex index 33b9ba6..d71585d 100644 --- a/lib/legion.ex +++ b/lib/legion.ex @@ -11,14 +11,17 @@ defmodule Legion do Runs an agent on a single task and returns the result. Starts a temporary agent process, blocks until the task completes, then stops it. + Accepts the same `opts` as `start_link/2`, so a one-off run can resume and + persist a conversation by passing `:store` and `:agent_id`. ## Examples {:ok, summary} = Legion.execute(ResearchAgent, "Summarize the Elixir getting started guide") {:cancel, :reached_max_iterations} = Legion.execute(ResearchAgent, "impossible task") + {:ok, reply} = Legion.execute(ChatAgent, "next question", store: MyApp.AgentStore, agent_id: "user_42:chat_7") """ - def execute(agent_module, task) do - {:ok, pid} = AgentServer.start_link(agent_module) + def execute(agent_module, task, opts \\ []) do + {:ok, pid} = AgentServer.start_link(agent_module, opts) result = AgentServer.call(pid, task) GenServer.stop(pid) result @@ -29,14 +32,18 @@ defmodule Legion do ## Options - `:name` - register the process under a name - - `:store`, `:agent_id` - persist the conversation across restarts; see `Legion.Store` + - `: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`. - 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(AssistantAgent, store: MyApp.AgentStore, agent_id: "user_42") + {: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 """ def start_link(agent_module, opts \\ []) do AgentServer.start_link(agent_module, opts) @@ -80,6 +87,21 @@ defmodule Legion do AgentServer.get_messages(pid) end + @doc """ + Returns the persistence id of a running agent, or `nil` if it has no store. + + When a store is configured but you let Legion generate the id, capture it + with this to resume the same conversation on a later start. + + ## Examples + + {:ok, pid} = Legion.start_link(ChatAgent) + agent_id = Legion.get_agent_id(pid) + """ + def get_agent_id(pid) do + AgentServer.get_agent_id(pid) + end + @doc """ Runs multiple agent tasks concurrently and collects results. diff --git a/lib/legion/agent_server.ex b/lib/legion/agent_server.ex index 8ebb86f..b8a6bb4 100644 --- a/lib/legion/agent_server.ex +++ b/lib/legion/agent_server.ex @@ -22,10 +22,15 @@ defmodule Legion.AgentServer do {store, opts} = Keyword.pop(opts, :store) {agent_id, opts} = Keyword.pop(opts, :agent_id) - if is_nil(store) != is_nil(agent_id) do - raise ArgumentError, ":store and :agent_id must be given together" + store = store || Application.get_env(:legion, :store) + + if is_nil(store) and not is_nil(agent_id) do + raise ArgumentError, + ":agent_id requires a :store - pass one or set `config :legion, :store, MyStore`" end + agent_id = if store, do: agent_id || generate_agent_id(), else: nil + gen_opts = if name, do: [name: name], else: [] config = resolve_config(agent_module, opts) GenServer.start_link(__MODULE__, {agent_module, config, store, agent_id}, gen_opts) @@ -43,6 +48,10 @@ defmodule Legion.AgentServer do GenServer.call(agent, :get_messages) end + def get_agent_id(agent) do + GenServer.call(agent, :get_agent_id) + end + # Server callbacks @impl true @@ -97,6 +106,11 @@ defmodule Legion.AgentServer do {:reply, state.messages, state} end + @impl true + def handle_call(:get_agent_id, _from, state) do + {:reply, state.agent_id, state} + end + @impl true def handle_call({:message, message}, _from, state) do {reply, state} = handle_message(message, state) @@ -124,6 +138,7 @@ 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 {status, value, final_messages, final_bindings} = Telemetry.span( @@ -133,10 +148,7 @@ defmodule Legion.AgentServer do messages = state.messages ++ [%{role: "user", content: content}] prev_count = Enum.count(messages, &(&1[:role] == "assistant")) - initial_bindings = - if Map.get(state.config, :binding_scope, :turn) == :conversation, - do: state.bindings, - else: [] + initial_bindings = if conversation_scope?, do: state.bindings, else: [] {status, value, messages, bindings} = result = Executor.run(state.agent_module, messages, state.config, initial_bindings) @@ -146,7 +158,8 @@ defmodule Legion.AgentServer do end ) - state = persist(%{state | messages: final_messages, bindings: final_bindings}) + kept_bindings = if conversation_scope?, do: final_bindings, else: [] + state = persist(%{state | messages: final_messages, bindings: kept_bindings}) {{status, value}, state} end @@ -158,6 +171,8 @@ defmodule Legion.AgentServer do state end + defp generate_agent_id, do: Base.url_encode64(:crypto.strong_rand_bytes(16), padding: false) + defp stringify(message, max_length) when is_binary(message), do: Executor.truncate_content(message, max_length) diff --git a/lib/legion/store.ex b/lib/legion/store.ex index 3ae56d6..4c158c3 100644 --- a/lib/legion/store.ex +++ b/lib/legion/store.ex @@ -22,6 +22,28 @@ defmodule Legion.Store do use Legion.Store.Postgres, repo: MyApp.Repo end + ## Configuring the store + + Pass `:store` per agent, or set one globally so every agent persists by default: + + config :legion, :store, MyApp.AgentStore + + A `:store` given to `start_link/2` overrides the global one. + + ## Identifying a conversation + + `:agent_id` is the key a snapshot is saved under - it names one conversation, + not one user. A chat app with many chats per user keys by the chat; compose the + id however you like, since Legion treats it as opaque: + + Legion.start_link(ChatAgent, agent_id: "user_42:chat_7") + + With a store in effect, 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. + Or implement the two callbacks against any storage you like: ## Example: hand-rolled Ecto store diff --git a/mix.exs b/mix.exs index c2f1999..b36d0c1 100644 --- a/mix.exs +++ b/mix.exs @@ -64,7 +64,8 @@ defmodule Legion.MixProject do {:ex_doc, ">= 0.0.0", only: :dev, runtime: false}, {:credo, ">= 0.0.0", only: [:dev, :test], runtime: false}, {:sobelow, ">= 0.0.0", only: [:dev, :test], runtime: false}, - {:mimic, "~> 1.7", only: :test} + {:mimic, "~> 1.7", only: :test}, + {:postgrex, "~> 0.22", only: :test} ] end diff --git a/mix.lock b/mix.lock index beca5c7..a019df6 100644 --- a/mix.lock +++ b/mix.lock @@ -2,6 +2,8 @@ "abnf_parsec": {:hex, :abnf_parsec, "2.1.0", "c4e88d5d089f1698297c0daced12be1fb404e6e577ecf261313ebba5477941f9", [:mix], [{:nimble_parsec, "~> 1.4", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm", "e0ed6290c7cc7e5020c006d1003520390c9bdd20f7c3f776bd49bfe3c5cd362a"}, "bunt": {:hex, :bunt, "1.0.0", "081c2c665f086849e6d57900292b3a161727ab40431219529f13c4ddcf3e7a44", [:mix], [], "hexpm", "dc5f86aa08a5f6fa6b8096f0735c4e76d54ae5c9fa2c143e5a1fc7c1cd9bb6b5"}, "credo": {:hex, :credo, "1.7.18", "5c5596bf7aedf9c8c227f13272ac499fe8eae6237bd326f2f07dfc173786f042", [:mix], [{:bunt, "~> 0.2.1 or ~> 1.0", [hex: :bunt, repo: "hexpm", optional: false]}, {:file_system, "~> 0.2 or ~> 1.0", [hex: :file_system, repo: "hexpm", optional: false]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: false]}], "hexpm", "a189d164685fd945809e862fe76a7420c4398fa288d76257662aecb909d6b3e5"}, + "db_connection": {:hex, :db_connection, "2.10.2", "ae391e803a5adff104da913c2fc1c0c14a37f8b10001dcef568796e1fb7bf95c", [:mix], [{:telemetry, "~> 0.4 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "510b14482330f1af6490a2fa0efd8d4f1435d1529b165647df22ac0f2df0fa93"}, + "decimal": {:hex, :decimal, "2.4.1", "6c0fbede12fb122ba685e9ab41c6a40c129e322b3aa192f9e072e61f3a6ffaf2", [:mix], [], "hexpm", "7e618897933a8455f19a727d7c5e50a2c071a544b700e5e724298ecb4340187f"}, "deep_merge": {:hex, :deep_merge, "1.0.0", "b4aa1a0d1acac393bdf38b2291af38cb1d4a52806cf7a4906f718e1feb5ee961", [:mix], [], "hexpm", "ce708e5f094b9cd4e8f2be4f00d2f4250c4095be93f8cd6d018c753894885430"}, "dotenvy": {:hex, :dotenvy, "1.1.1", "00e318f3c51de9fafc4b48598447e386f19204dc18ca69886905bb8f8b08b667", [:mix], [], "hexpm", "c8269471b5701e9e56dc86509c1199ded2b33dce088c3471afcfef7839766d8e"}, "earmark_parser": {:hex, :earmark_parser, "1.4.44", "f20830dd6b5c77afe2b063777ddbbff09f9759396500cdbe7523efd58d7a339c", [:mix], [], "hexpm", "4778ac752b4701a5599215f7030989c989ffdc4f6df457c5f36938cc2d2a2750"}, @@ -24,6 +26,7 @@ "nimble_options": {:hex, :nimble_options, "1.1.1", "e3a492d54d85fc3fd7c5baf411d9d2852922f66e69476317787a7b2bb000a61b", [:mix], [], "hexpm", "821b2470ca9442c4b6984882fe9bb0389371b8ddec4d45a9504f00a66f650b44"}, "nimble_parsec": {:hex, :nimble_parsec, "1.4.2", "8efba0122db06df95bfaa78f791344a89352ba04baedd3849593bfce4d0dc1c6", [:mix], [], "hexpm", "4b21398942dda052b403bbe1da991ccd03a053668d147d53fb8c4e0efe09c973"}, "nimble_pool": {:hex, :nimble_pool, "1.1.0", "bf9c29fbdcba3564a8b800d1eeb5a3c58f36e1e11d7b7fb2e084a643f645f06b", [:mix], [], "hexpm", "af2e4e6b34197db81f7aad230c1118eac993acc0dae6bc83bac0126d4ae0813a"}, + "postgrex": {:hex, :postgrex, "0.22.2", "4aec14df2a72722aee92492566edbeeb44e233ecb86b1915d03136297ef1385d", [:mix], [{:db_connection, "~> 2.9", [hex: :db_connection, repo: "hexpm", optional: false]}, {:decimal, "~> 1.5 or ~> 2.0 or ~> 3.0", [hex: :decimal, repo: "hexpm", optional: false]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: true]}, {:table, "~> 0.1.0", [hex: :table, repo: "hexpm", optional: true]}], "hexpm", "8946382ddb06294f56026ac4278b3cc212bac8a2c82ed68b4087819ed1abc53b"}, "process_tree": {:hex, :process_tree, "0.3.0", "0eb58ec68f3d22a4f36040b0374469464c95fae5465c011b55bea01649b0bd70", [:mix], [], "hexpm", "6cb3b7be9c7d74b28a9f6e0f03115d5953e6bbda44be2b6fd9e667020870eb86"}, "req": {:hex, :req, "0.5.17", "0096ddd5b0ed6f576a03dde4b158a0c727215b15d2795e59e0916c6971066ede", [:mix], [{:brotli, "~> 0.3.1", [hex: :brotli, repo: "hexpm", optional: true]}, {:ezstd, "~> 1.0", [hex: :ezstd, repo: "hexpm", optional: true]}, {:finch, "~> 0.17", [hex: :finch, repo: "hexpm", optional: false]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: false]}, {:mime, "~> 2.0.6 or ~> 2.1", [hex: :mime, repo: "hexpm", optional: false]}, {:nimble_csv, "~> 1.0", [hex: :nimble_csv, repo: "hexpm", optional: true]}, {:plug, "~> 1.0", [hex: :plug, repo: "hexpm", optional: true]}], "hexpm", "0b8bc6ffdfebbc07968e59d3ff96d52f2202d0536f10fef4dc11dc02a2a43e39"}, "req_llm": {:hex, :req_llm, "1.11.0", "56c823e40e1409ef7f8d972301058277671161262ca9201db8c5b531a33f25ae", [:mix], [{:dotenvy, "~> 1.1", [hex: :dotenvy, repo: "hexpm", optional: false]}, {:ex_aws_auth, "~> 1.3", [hex: :ex_aws_auth, repo: "hexpm", optional: false]}, {:igniter, "~> 0.7", [hex: :igniter, repo: "hexpm", optional: true]}, {:jason, "~> 1.4", [hex: :jason, repo: "hexpm", optional: false]}, {:jsv, "~> 0.11", [hex: :jsv, repo: "hexpm", optional: false]}, {:llm_db, "~> 2026.4.0", [hex: :llm_db, repo: "hexpm", optional: false]}, {:nimble_options, "~> 1.1", [hex: :nimble_options, repo: "hexpm", optional: false]}, {:req, "~> 0.5", [hex: :req, repo: "hexpm", optional: false]}, {:server_sent_events, "~> 1.0.0", [hex: :server_sent_events, repo: "hexpm", optional: false]}, {:splode, "~> 0.3.0", [hex: :splode, repo: "hexpm", optional: false]}, {:uniq, "~> 0.6", [hex: :uniq, repo: "hexpm", optional: false]}, {:websockex, "~> 0.5.1", [hex: :websockex, repo: "hexpm", optional: false]}, {:zoi, "~> 0.14", [hex: :zoi, repo: "hexpm", optional: false]}], "hexpm", "81118f324632e60d7641911eb97d01798129fb480d04a860ba6007507b866671"}, diff --git a/test/legion/agent_server_test.exs b/test/legion/agent_server_test.exs index f08ba22..9550278 100644 --- a/test/legion/agent_server_test.exs +++ b/test/legion/agent_server_test.exs @@ -433,6 +433,32 @@ defmodule Legion.AgentServerTest do refute Enum.any?(messages, &(&1.role == "system")) end + test "a one-off execute/3 persists its snapshot before stopping" do + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + llm_response("Paris") + end) + + {:ok, _} = + Legion.execute(MathAgent, "What is the capital of France?", + store: MemoryStore, + agent_id: "one-off" + ) + + assert {:ok, %{messages: [%{role: "user"}, %{role: "assistant"} | _]}} = + MemoryStore.load("one-off") + end + + test "does not persist bindings under the default :turn scope" do + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + llm_eval_response("x = 42") + end) + + {:ok, pid} = Legion.start_link(MathAgent, store: MemoryStore, agent_id: "turn-bindings") + {:ok, 42} = Legion.call(pid, "set x") + + assert {:ok, %{bindings: []}} = MemoryStore.load("turn-bindings") + end + test "restores the conversation under a fresh system prompt after a restart" do stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> llm_response("Paris") @@ -474,9 +500,41 @@ defmodule Legion.AgentServerTest do assert {:ok, 43} = Legion.call(revived, "use x") end - test "raises when :store is given without :agent_id" do - assert_raise ArgumentError, ":store and :agent_id must be given together", fn -> - Legion.start_link(MathAgent, store: MemoryStore) + test "generates an agent_id when a store is given without one" do + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + llm_response("Paris") + end) + + {:ok, pid} = Legion.start_link(MathAgent, store: MemoryStore) + agent_id = Legion.get_agent_id(pid) + + assert is_binary(agent_id) + {:ok, _} = Legion.call(pid, "What is the capital of France?") + assert {:ok, _snapshot} = MemoryStore.load(agent_id) + end + + test "uses a store configured globally, needing only an agent_id" do + Application.put_env(:legion, :store, MemoryStore) + on_exit(fn -> Application.delete_env(:legion, :store) end) + + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + llm_response("Paris") + end) + + {:ok, pid} = Legion.start_link(MathAgent, agent_id: "global-store") + {:ok, _} = Legion.call(pid, "What is the capital of France?") + + assert {:ok, _snapshot} = MemoryStore.load("global-store") + end + + test "get_agent_id/1 returns nil without a store" do + {:ok, pid} = Legion.start_link(MathAgent) + assert Legion.get_agent_id(pid) == nil + end + + test "raises when :agent_id is given without a :store" do + assert_raise ArgumentError, ~r/:agent_id requires a :store/, fn -> + Legion.start_link(MathAgent, agent_id: "orphan") end end end diff --git a/test/legion/store/postgres_db_test.exs b/test/legion/store/postgres_db_test.exs new file mode 100644 index 0000000..630848c --- /dev/null +++ b/test/legion/store/postgres_db_test.exs @@ -0,0 +1,39 @@ +defmodule Legion.Store.PostgresDbTest do + @moduledoc """ + Exercises the generated Postgres store against a real database, so the SQL it + issues - the `ON CONFLICT` upsert in particular - is verified for real rather + than shape-matched against a fake. + """ + use ExUnit.Case, async: false + + defmodule Repo do + def query!(sql, params), do: Postgrex.query!(:legion_store_test, sql, params) + end + + defmodule Store do + use Legion.Store.Postgres, repo: Legion.Store.PostgresDbTest.Repo + end + + setup do + Postgrex.query!(:legion_store_test, "TRUNCATE legion_agents", []) + :ok + end + + test "round-trips a snapshot through a real bytea column" do + snapshot = %{messages: [%{role: "user", content: "hi"}], bindings: [x: 42]} + + assert :ok = Store.save("user_42", snapshot) + assert {:ok, ^snapshot} = Store.load("user_42") + end + + test "load/1 returns :error when the row is absent" do + assert :error = Store.load("missing") + end + + test "save/2 upserts on conflict - the latest snapshot wins" do + assert :ok = Store.save("user_42", %{messages: [], bindings: [v: 1]}) + assert :ok = Store.save("user_42", %{messages: [], bindings: [v: 2]}) + + assert {:ok, %{bindings: [v: 2]}} = Store.load("user_42") + end +end diff --git a/test/test_helper.exs b/test/test_helper.exs index 253ab14..5cae099 100644 --- a/test/test_helper.exs +++ b/test/test_helper.exs @@ -1,4 +1,29 @@ Mimic.copy(ReqLLM) Legion.Telemetry.attach_default_logger() + +# Shared connection and schema for the Legion.Store.Postgres database tests. +{:ok, _} = + Postgrex.start_link( + name: :legion_store_test, + hostname: System.get_env("POSTGRES_HOST", "localhost"), + port: String.to_integer(System.get_env("POSTGRES_PORT", "5432")), + username: System.get_env("POSTGRES_USER", "postgres"), + password: System.get_env("POSTGRES_PASSWORD", "postgres"), + database: System.get_env("POSTGRES_DB", "postgres") + ) + +Postgrex.query!( + :legion_store_test, + """ + CREATE TABLE IF NOT EXISTS legion_agents ( + agent_id text PRIMARY KEY, + snapshot bytea NOT NULL, + inserted_at timestamptz NOT NULL DEFAULT now(), + updated_at timestamptz NOT NULL DEFAULT now() + ) + """, + [] +) + ExUnit.start(exclude: [:integration]) From 328ef14caddd224963036303643cdd8dc862aa5c Mon Sep 17 00:00:00 2001 From: Dima Mikielewicz Date: Wed, 15 Jul 2026 18:21:05 +0200 Subject: [PATCH 03/30] WiP --- CHANGELOG.md | 4 + README.md | 4 +- lib/legion.ex | 62 ++++++++++- lib/legion/agent.ex | 2 +- lib/legion/agent_server.ex | 59 ++++++++-- lib/legion/executor.ex | 47 ++++++-- lib/legion/prompts/system_prompt.eex | 2 + lib/legion/store.ex | 73 +++++++++++- lib/legion/store/postgres.ex | 116 +++++++++++++++++-- lib/legion/store/postgres/migration.ex | 143 ++++++++++++++++++++++++ lib/legion/telemetry.ex | 40 +++---- lib/legion/tools/human_tool.ex | 14 ++- mix.exs | 9 +- test/legion/agent_server_test.exs | 147 ++++++++++++++++++++++++- test/legion/store/postgres_db_test.exs | 81 ++++++++++++++ test/legion/store/postgres_test.exs | 124 +++++++++++++++++++-- test/legion/tools/human_tool_test.exs | 22 +++- test/test_helper.exs | 19 ++-- 18 files changed, 882 insertions(+), 86 deletions(-) create mode 100644 lib/legion/store/postgres/migration.ex diff --git a/CHANGELOG.md b/CHANGELOG.md index a91fcbf..74e9c0d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,10 @@ - Add `Legion.Store` behaviour for persisting conversations across restarts. Pass `store:` and `agent_id:` to `Legion.start_link/2`; snapshots (messages + bindings) are saved after every completed turn, before the caller receives its reply - Add `Legion.Store.Postgres`, a ready-made store adapter that reuses your Ecto repo (`use Legion.Store.Postgres, repo: MyApp.Repo`) without adding Ecto as a dependency +- `Legion.Store.Postgres.Migration` installs a trigger that `pg_notify`s the table's channel with the agent_id on every write, and generated stores expose `__repo__/0` and `__table__/0`, so consumers (LegionWeb) can follow store changes live +- Add optional `Legion.Store.save_status/2` callback, called with `:running` when a turn starts and `:idle` after its snapshot is saved; `Legion.Store.Postgres` implements it via a `status` column, so a `'running'` status under a dead pid identifies a conversation that crashed mid-turn +- Bump the default model from `openai:gpt-4o-mini` to `openai:gpt-5.4` +- `Legion.Tools.HumanTool.ask/1` now raises when called under `eval_and_complete` - the turn would end as soon as the code returns, silently discarding the human's answer; the error feeds back to the model, which retries under `eval_and_continue` ## v0.4.0 - 2026-05-17 diff --git a/README.md b/README.md index d6521e5..a484bae 100644 --- a/README.md +++ b/README.md @@ -213,7 +213,7 @@ Your handler receives `{:human_request, ref, from_pid, question, meta}` and repl ```elixir config :legion, :config, %{ - model: "openai:gpt-4o-mini", + model: "openai:gpt-5.4", max_iterations: 10, max_retries: 3, sandbox_timeout: 60_000, @@ -224,7 +224,7 @@ config :legion, :config, %{ | Option | Description | | -------------------- | -------------------------------------------------------------------------------------------------------------------------- | -| `model` | LLM model string passed to [ReqLLM](https://hexdocs.pm/req_llm), e.g. `"openai:gpt-4o-mini"`. | +| `model` | LLM model string passed to [ReqLLM](https://hexdocs.pm/req_llm), e.g. `"openai:gpt-5.4"`. | | `max_iterations` | Successful execution steps before the agent is stopped. | | `max_retries` | Consecutive failures (bad code, tool errors) before giving up. Resets after each success. | | `sandbox_timeout` | Milliseconds a single code evaluation may run before it is killed. | diff --git a/lib/legion.ex b/lib/legion.ex index d71585d..1545836 100644 --- a/lib/legion.ex +++ b/lib/legion.ex @@ -12,7 +12,8 @@ defmodule Legion do Starts a temporary agent process, blocks until the task completes, then stops it. Accepts the same `opts` as `start_link/2`, so a one-off run can resume and - persist a conversation by passing `:store` and `:agent_id`. + persist a conversation by passing `:store` and `:agent_id`. Passing `:store` + overrides the globally configured store for this agent; see `Legion.Store`. ## Examples @@ -88,7 +89,8 @@ defmodule Legion do end @doc """ - Returns the persistence id of a running agent, or `nil` if it has no store. + Returns the id of a running agent. Always set - Legion generates one when + none is passed. When a store is configured but you let Legion generate the id, capture it with this to resume the same conversation on a later start. @@ -102,6 +104,62 @@ defmodule Legion do AgentServer.get_agent_id(pid) end + @doc """ + Whether `pid` - typically the one recorded in a persisted run's metadata - + is alive on this node. Accepts `nil` and returns `false`. + + A stored pid outlives the VM that wrote it, so after a restart the check is + best-effort: a recycled pid value can collide with an unrelated live + process. + + ## Examples + + run = MyApp.AgentStore.get_run("user_42:chat_7") + Legion.running?(run.pid) + """ + def running?(pid) when is_pid(pid) do + node(pid) == node() and Process.alive?(pid) + rescue + # A pid deserialized from a previous VM incarnation is not a local pid, + # for which Process.alive?/1 raises. + ArgumentError -> false + end + + def running?(_other), do: false + + @doc """ + Resumes a persisted conversation. + + Returns the recorded process if it is still `running?/1`; otherwise starts + the agent again under the same `agent_id`, so it reloads its snapshot from + the store. `opts` are passed through to `start_link/2`. + + Requires a store implementing `c:Legion.Store.get_run/1` - pass `:store` or + configure one globally. Raises if the store has no run for `agent_id`. + + ## Examples + + {:ok, pid} = Legion.resume("user_42:chat_7") + {:ok, pid} = Legion.resume("user_42:chat_7", 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`" + + run = + store.get_run(agent_id) || + raise ArgumentError, + "no run recorded for agent_id #{inspect(agent_id)} in #{inspect(store)}" + + if running?(run[:pid]) do + {:ok, run[:pid]} + else + start_link(run.agent_module, Keyword.put(opts, :agent_id, agent_id)) + end + end + @doc """ Runs multiple agent tasks concurrently and collects results. diff --git a/lib/legion/agent.ex b/lib/legion/agent.ex index 0fa9f7f..05bb105 100644 --- a/lib/legion/agent.ex +++ b/lib/legion/agent.ex @@ -37,7 +37,7 @@ defmodule Legion.Agent do - `config/0` — agent-level configuration merged with application config and call-time opts. Defaults to `%{}`. Available keys: - - `model` — LLM model identifier (default: `"openai:gpt-4o-mini"`) + - `model` — LLM model identifier (default: `"openai:gpt-5.4"`) - `max_iterations` — max successful execution steps per turn (default: `10`) - `max_retries` — max consecutive failures before giving up (default: `3`) - `sandbox_timeout` — timeout in ms for code execution (default: `60_000`) diff --git a/lib/legion/agent_server.ex b/lib/legion/agent_server.ex index b8a6bb4..2a9bf72 100644 --- a/lib/legion/agent_server.ex +++ b/lib/legion/agent_server.ex @@ -22,14 +22,14 @@ defmodule Legion.AgentServer do {store, opts} = Keyword.pop(opts, :store) {agent_id, opts} = Keyword.pop(opts, :agent_id) - store = store || Application.get_env(:legion, :store) + store = store || Vault.get(:store) || Application.get_env(:legion, :store) if is_nil(store) and not is_nil(agent_id) do raise ArgumentError, ":agent_id requires a :store - pass one or set `config :legion, :store, MyStore`" end - agent_id = if store, do: agent_id || generate_agent_id(), else: nil + agent_id = agent_id || generate_id() gen_opts = if name, do: [name: name], else: [] config = resolve_config(agent_module, opts) @@ -56,11 +56,11 @@ defmodule Legion.AgentServer do @impl true def init({agent_module, config, store, agent_id}) do - parent_run_id = Vault.get(:run_id) - run_id = make_ref() + parent_agent_id = Vault.get(:agent_id) - Vault.unsafe_put(:run_id, run_id) - Vault.unsafe_put(:parent_run_id, parent_run_id) + Vault.unsafe_put(:agent_id, agent_id) + Vault.unsafe_put(:parent_agent_id, parent_agent_id) + if store, do: Vault.unsafe_put(:store, store) for tool <- agent_module.tools() do Vault.unsafe_put(tool, agent_module.tool_config(tool)) @@ -74,6 +74,13 @@ defmodule Legion.AgentServer do %{agent: agent_module} ) + save_run(store, agent_id, %{ + agent_module: agent_module, + parent_agent_id: parent_agent_id, + pid: self(), + started_at: System.system_time(:millisecond) + }) + {saved_messages, saved_bindings} = case store && store.load(agent_id) do {:ok, %{messages: messages, bindings: bindings}} -> {messages, bindings} @@ -82,7 +89,7 @@ defmodule Legion.AgentServer do state = %__MODULE__{ agent_module: agent_module, - messages: [%{role: "system", content: system_prompt} | saved_messages], + messages: [Executor.message(:system, system_prompt) | saved_messages], config: config, store: store, agent_id: agent_id, @@ -140,12 +147,18 @@ defmodule Legion.AgentServer 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 + # response. + state = persist(%{state | messages: state.messages ++ [Executor.message(:user, content)]}) + save_status(state, :running) + {status, value, final_messages, final_bindings} = Telemetry.span( [:legion, :agent, :message], %{agent: state.agent_module, message: content}, fn -> - messages = state.messages ++ [%{role: "user", content: content}] + messages = state.messages prev_count = Enum.count(messages, &(&1[:role] == "assistant")) initial_bindings = if conversation_scope?, do: state.bindings, else: [] @@ -160,6 +173,7 @@ defmodule Legion.AgentServer do kept_bindings = if conversation_scope?, do: final_bindings, else: [] state = persist(%{state | messages: final_messages, bindings: kept_bindings}) + save_status(state, :idle) {{status, value}, state} end @@ -171,7 +185,34 @@ defmodule Legion.AgentServer do state end - defp generate_agent_id, do: Base.url_encode64(:crypto.strong_rand_bytes(16), padding: false) + # Status is written outside the turn's save/persist path: a crash between + # the :running write and the :idle write leaves 'running' in the store, + # which consumers read as "crashed mid-turn" under a dead pid. + defp save_status(%{store: nil}, _status), do: :ok + + defp save_status(state, status) do + if Code.ensure_loaded?(state.store) and function_exported?(state.store, :save_status, 2) do + state.store.save_status(state.agent_id, status) + end + + :ok + end + + defp save_run(nil, _agent_id, _metadata), do: :ok + + defp save_run(store, agent_id, metadata) do + if Code.ensure_loaded?(store) and function_exported?(store, :save_run, 2) do + store.save_run(agent_id, metadata) + else + Logger.warning( + "Store #{inspect(store)} does not implement save_run/2; run metadata not persisted" + ) + end + + :ok + end + + defp generate_id, do: Base.url_encode64(:crypto.strong_rand_bytes(16), padding: false) defp stringify(message, max_length) when is_binary(message), do: Executor.truncate_content(message, max_length) diff --git a/lib/legion/executor.ex b/lib/legion/executor.ex index d046f9b..c041c40 100644 --- a/lib/legion/executor.ex +++ b/lib/legion/executor.ex @@ -12,7 +12,7 @@ defmodule Legion.Executor do alias Legion.{Sandbox, Telemetry} @default_config %{ - model: "openai:gpt-4o-mini", + model: "openai:gpt-5.4", max_iterations: 10, max_retries: 3, sandbox_timeout: 60_000, @@ -23,14 +23,38 @@ defmodule Legion.Executor do @doc false def default_config, do: @default_config + @message_roles %{ + system: "system", + user: "user", + assistant: "assistant", + eval_result: "user", + error: "user" + } + + @doc """ + Builds a conversation message stamped with its `:type` and creation time + (`:at`, milliseconds). The extra keys ride along into persisted snapshots so + consumers (e.g. LegionWeb) can classify messages without parsing content; + ReqLLM ignores them. + """ + def message(type, content) do + %{ + role: Map.fetch!(@message_roles, type), + type: type, + content: content, + at: System.system_time(:millisecond) + } + end + @action_descriptions %{ "eval_and_continue" => "Execute code and continue the turn. Use when you need the result before deciding the next step.", "eval_and_complete" => "Finish the turn with the code's result. Use when the final answer comes from executing code.", "return" => - "Finish the turn with a structured result and no code execution. Only use when the task is fully done - not to report in-progress work or bail out of execution errors (fix the code and re-run instead).", - "done" => "Task complete with no result to return." + "Finish the turn with a structured result and no code execution. Only use when the task is fully done - not to report in-progress work, ask the user something, or bail out of execution errors (fix the code and re-run instead). The result is a final answer, not a chat message: never return a status like \"starting\" or \"working on it\" - do the work first.", + "done" => + "Finish the turn with no result to return. This ends your run - nothing executes after it, so only use it when everything you were asked to do is fully done. Never announce upcoming work and then pick this action; do the work first (eval_and_continue)." } defp action_schema(agent_module) do @@ -146,7 +170,7 @@ defmodule Legion.Executor do case ReqLLM.generate_object(config.model, messages, action_schema(agent_module)) do {:ok, response} -> action = extract_object(response) - messages = messages ++ [%{role: "assistant", content: Jason.encode!(action)}] + messages = messages ++ [message(:assistant, Jason.encode!(action))] {{:ok, action, messages}, %{object: action}} {:error, reason} -> @@ -180,12 +204,16 @@ defmodule Legion.Executor do bindings ) when eval in ["eval_and_continue", "eval_and_complete"] and code != "" do + # Tools that must see the answer come back to the model (e.g. HumanTool) + # read this to reject running under a turn-ending action. + Vault.unsafe_put(:current_action, eval) + case eval_in_span(agent, code, config, bindings) do {:ok, {result, new_bindings}} -> new_bindings = if config.binding_scope == :iteration, do: [], else: new_bindings messages = - messages ++ [%{role: "user", content: format_result(result, new_bindings, config)}] + messages ++ [message(:eval_result, format_result(result, new_bindings, config))] if eval == "eval_and_continue", do: loop(agent, messages, config, i + 1, 0, new_bindings), @@ -241,11 +269,10 @@ defmodule Legion.Executor do messages = messages ++ [ - %{ - role: "user", - content: - "Code execution failed:\n\n#{error_text}\n\nPlease fix the error and try again." - } + message( + :error, + "Code execution failed:\n\n#{error_text}\n\nPlease fix the error and try again." + ) ] loop(agent_module, messages, config, iteration, retries + 1, bindings) diff --git a/lib/legion/prompts/system_prompt.eex b/lib/legion/prompts/system_prompt.eex index 80833a7..5763211 100644 --- a/lib/legion/prompts/system_prompt.eex +++ b/lib/legion/prompts/system_prompt.eex @@ -9,6 +9,8 @@ For each step, you respond with a JSON object containing: **Prefer writing code to accomplish your task.** Use the available tools by calling their functions in your code. Examine the tool source code below to understand what functions are available and how to use them. +**Finishing actions end your run.** After `return`, `done`, or `eval_and_complete`, nothing executes until the user sends a new message, and the returned result is a final answer - not a chat message the user is expected to reply to. Never finish just to announce what you are about to do; do it now with `eval_and_continue`. If you need input from the user mid-task, ask through a human-input tool (such as `HumanTool.ask`) when one is listed below - that is the only way to reach the user before your run ends. + <%= if binding_scope == :iteration do %>**Variables do not persist.** Each code execution starts with a clean slate - variables defined in one execution are not available in the next.<% else %>**Variables persist.** When you use `eval_and_continue`, any variables you define are available in subsequent code executions within the same turn. After each execution, you'll see which variables are available.<%= if binding_scope == :conversation do %> Variables also persist across turns - a variable you assigned while handling an earlier user message is still in scope when handling the next one, so you can reference it directly instead of recomputing it.<% end %><% end %> **Constraints:** diff --git a/lib/legion/store.ex b/lib/legion/store.ex index 4c158c3..aacb154 100644 --- a/lib/legion/store.ex +++ b/lib/legion/store.ex @@ -28,7 +28,11 @@ defmodule Legion.Store do config :legion, :store, MyApp.AgentStore - A `:store` given to `start_link/2` overrides the global one. + A `:store` given to `start_link/2` overrides the global one. Sub-agents + spawned from a running agent (e.g. via `Legion.Tools.AgentTool`) inherit + the parent's store automatically. The inheritance is ambient: *any* agent + started from within an agent's process tree picks up that store unless + given an explicit `:store` of its own. ## Identifying a conversation @@ -38,7 +42,7 @@ defmodule Legion.Store do Legion.start_link(ChatAgent, agent_id: "user_42:chat_7") - With a store in effect, omitting `:agent_id` makes Legion generate one. That + 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 @@ -73,18 +77,81 @@ defmodule Legion.Store do The snapshot holds the conversation `:messages` (without the system prompt) and the `:bindings` from evaluated code (relevant with - `binding_scope: :conversation`). Bindings are arbitrary Elixir terms - + `binding_scope: :conversation`). Each message carries a `:type` + (`:user`, `:assistant`, `:eval_result`, or `:error`) and an `:at` timestamp + in milliseconds, so consumers can classify and order messages without + parsing content. Bindings are arbitrary Elixir terms - values like pids, references, or functions will not survive serialization, so keep conversation-scoped variables to plain data if you persist agents. + + ## Run metadata + + Implement the optional `c:save_run/2` to also record each conversation's + identity: which agent module ran, under which parent conversation (for + sub-agents), when, and under which `pid`. Legion calls it once per agent + start, so persisted conversations stay attributable to the agent tree that + produced them. Restarting a conversation under the same `agent_id` calls it + again with a fresh `started_at` and `pid` - treat it as an upsert. The + stored pid always names the newest process for the conversation, so + `Legion.running?/1` and `Legion.resume/2` can tell whether it is still + live (a pid outlives the VM that wrote it, so never trust it blindly). On conflict, keep the stored + `parent_agent_id`: the callback reports where the agent was started *this + time*, so a conversation resumed from iex or another agent's tree would + otherwise be silently reparented. + + ## Turn status + + Implement the optional `c:save_status/2` to also record whether a + conversation is mid-turn: Legion calls it with `:running` when a message + starts and `:idle` after the turn's snapshot is saved. Combined with the + stored pid this lets a consumer distinguish a live agent that is working + from one waiting for input, and a conversation that crashed mid-turn + (status still `:running` under a dead pid) from one that finished. + + ## Reading conversations back + + The optional `c:list_runs/1` and `c:get_run/1` callbacks expose persisted + runs for consumers that rebuild a view of past conversations from the store + alone - `LegionWeb` uses them when configured with + `config :legion_web, agents_source: :database`. `Legion.Store.Postgres` + implements both. """ @type agent_id :: term() @type snapshot :: %{messages: [map()], bindings: keyword()} + @type run_metadata :: %{ + agent_module: module(), + parent_agent_id: agent_id() | nil, + pid: pid(), + started_at: integer() + } + @type run :: %{ + agent_id: agent_id(), + agent_module: module() | nil, + parent_agent_id: agent_id() | nil, + pid: pid() | nil, + status: :running | :idle | nil, + started_at: integer() | nil + } @doc "Returns the last saved snapshot for `agent_id`, or `:error` if none exists." @callback load(agent_id()) :: {:ok, snapshot()} | :error @doc "Saves the snapshot for `agent_id`. Raise on failure - the turn is not acked until this returns." @callback save(agent_id(), snapshot()) :: :ok + + @doc "Records run metadata when an agent starts. Optional. `started_at` is in milliseconds." + @callback save_run(agent_id(), run_metadata()) :: :ok + + @doc "Records whether the conversation is mid-turn. Optional." + @callback save_status(agent_id(), :running | :idle) :: :ok + + @doc "Returns the newest `limit` persisted runs, newest first. Optional." + @callback list_runs(limit :: pos_integer()) :: [run()] + + @doc "Returns the persisted run for `agent_id`, or `nil` if none exists. Optional." + @callback get_run(agent_id()) :: run() | nil + + @optional_callbacks save_run: 2, save_status: 2, list_runs: 1, get_run: 1 end diff --git a/lib/legion/store/postgres.ex b/lib/legion/store/postgres.ex index 7012b73..94cb0f7 100644 --- a/lib/legion/store/postgres.ex +++ b/lib/legion/store/postgres.ex @@ -14,18 +14,13 @@ defmodule Legion.Store.Postgres do use Legion.Store.Postgres, repo: MyApp.Repo end - Create the table in a migration: + Create the table in a migration with `Legion.Store.Postgres.Migration`: defmodule MyApp.Repo.Migrations.AddLegionAgents do use Ecto.Migration - def change do - create table(:legion_agents, primary_key: false) do - add :agent_id, :text, primary_key: true - add :snapshot, :binary, null: false - timestamps(type: :timestamptz) - end - end + def up, do: Legion.Store.Postgres.Migration.up() + def down, do: Legion.Store.Postgres.Migration.down() end Then start agents with it: @@ -38,7 +33,29 @@ defmodule Legion.Store.Postgres do - `:table` - the table name, defaults to `"legion_agents"` Agent ids must be strings. Snapshots are stored as `:erlang.term_to_binary/1` - blobs - readable only from Elixir, one row per agent, upserted on every turn. + blobs - readable only from Elixir, one row per conversation, upserted on + every turn. + + The store also implements `c:Legion.Store.save_run/2`, so the same row + carries the conversation's identity: `agent_module` (in `inspect/1` form, + e.g. `"MyApp.ResearchAgent"`), `parent_agent_id` linking a sub-agent to the + conversation that spawned it, `pid` of the newest process that ran the + conversation (a `:erlang.term_to_binary/1` blob, like the snapshot), and + `started_at` in milliseconds (last start wins). `snapshot` is null until the conversation's first turn completes; + `parent_agent_id` is kept once set, so resuming from elsewhere does not + reparent the conversation. + + `c:Legion.Store.save_status/2` is implemented as well: the row's `status` + flips to `'running'` when a turn starts and back to `'idle'` when it + completes (`save_run` resets it to `'idle'` on start), so a `'running'` + status under a dead pid identifies a conversation that crashed mid-turn. + + `c:Legion.Store.list_runs/1` and `c:Legion.Store.get_run/1` are implemented + too, so persisted conversations can be read back into a view of past runs. + + The migration also installs a trigger that `pg_notify`s the table's channel + (the table name) with the `agent_id` on every insert or update, so + consumers can follow store changes live without polling. """ defmacro __using__(opts) do @@ -53,6 +70,27 @@ defmodule Legion.Store.Postgres do ON CONFLICT (agent_id) DO UPDATE SET snapshot = EXCLUDED.snapshot, updated_at = now() """ + save_run_sql = """ + INSERT INTO #{table} (agent_id, agent_module, parent_agent_id, pid, status, started_at, inserted_at, updated_at) + VALUES ($1, $2, $3, $4, 'idle', $5, now(), now()) + ON CONFLICT (agent_id) DO UPDATE + SET agent_module = EXCLUDED.agent_module, + parent_agent_id = COALESCE(#{table}.parent_agent_id, EXCLUDED.parent_agent_id), + pid = EXCLUDED.pid, + status = 'idle', + started_at = EXCLUDED.started_at, + updated_at = now() + """ + + save_status_sql = "UPDATE #{table} SET status = $2, updated_at = now() WHERE agent_id = $1" + + run_columns = "agent_id, agent_module, parent_agent_id, pid, status, started_at" + + list_runs_sql = + "SELECT #{run_columns} FROM #{table} ORDER BY started_at DESC NULLS LAST LIMIT $1" + + get_run_sql = "SELECT #{run_columns} FROM #{table} WHERE agent_id = $1" + quote do @behaviour Legion.Store @@ -60,6 +98,7 @@ defmodule Legion.Store.Postgres do # sobelow_skip ["Misc.BinToTerm"] def load(agent_id) when is_binary(agent_id) do case unquote(repo).query!(unquote(select_sql), [agent_id]) do + %{rows: [[nil]]} -> :error %{rows: [[snapshot]]} -> {:ok, :erlang.binary_to_term(snapshot)} %{rows: []} -> :error end @@ -70,6 +109,65 @@ defmodule Legion.Store.Postgres do unquote(repo).query!(unquote(upsert_sql), [agent_id, :erlang.term_to_binary(snapshot)]) :ok end + + @impl Legion.Store + def save_run(agent_id, metadata) when is_binary(agent_id) do + unquote(repo).query!(unquote(save_run_sql), [ + agent_id, + inspect(metadata.agent_module), + metadata.parent_agent_id, + metadata[:pid] && :erlang.term_to_binary(metadata.pid), + metadata.started_at + ]) + + :ok + end + + @impl Legion.Store + def save_status(agent_id, status) when is_binary(agent_id) do + unquote(repo).query!(unquote(save_status_sql), [agent_id, Atom.to_string(status)]) + :ok + end + + @impl Legion.Store + def list_runs(limit) do + %{rows: rows} = unquote(repo).query!(unquote(list_runs_sql), [limit]) + Enum.map(rows, &Legion.Store.Postgres.decode_run_row/1) + end + + @impl Legion.Store + def get_run(agent_id) when is_binary(agent_id) do + case unquote(repo).query!(unquote(get_run_sql), [agent_id]) do + %{rows: [row]} -> Legion.Store.Postgres.decode_run_row(row) + %{rows: []} -> nil + end + end + + @doc false + def __repo__, do: unquote(repo) + + @doc false + def __table__, do: unquote(table) end end + + @doc false + # A pid stored by a previous VM run decodes with its original creation + # number, so `Legion.running?/1` reads it as not running instead of + # colliding with a recycled pid value. + # sobelow_skip ["Misc.BinToTerm"] + def decode_run_row([agent_id, agent_module, parent_agent_id, pid, status, started_at]) do + %{ + agent_id: agent_id, + agent_module: agent_module && Module.concat([agent_module]), + parent_agent_id: parent_agent_id, + pid: pid && :erlang.binary_to_term(pid), + status: decode_status(status), + started_at: started_at + } + end + + defp decode_status("running"), do: :running + defp decode_status("idle"), do: :idle + defp decode_status(nil), do: nil end diff --git a/lib/legion/store/postgres/migration.ex b/lib/legion/store/postgres/migration.ex new file mode 100644 index 0000000..fe54336 --- /dev/null +++ b/lib/legion/store/postgres/migration.ex @@ -0,0 +1,143 @@ +defmodule Legion.Store.Postgres.Migration do + @moduledoc """ + Migration helpers for `Legion.Store.Postgres`. + + ## Usage + + defmodule MyApp.Repo.Migrations.AddLegionAgents do + use Ecto.Migration + + def up, do: Legion.Store.Postgres.Migration.up() + def down, do: Legion.Store.Postgres.Migration.down() + end + + Migrations are versioned and idempotent - `up/1` only runs the versions the + database hasn't seen yet, so when a new Legion release ships schema changes + you generate another migration with the same two calls (optionally pinning + `version:`): + + defmodule MyApp.Repo.Migrations.UpgradeLegionAgentsToV2 do + use Ecto.Migration + + def up, do: Legion.Store.Postgres.Migration.up(version: 2) + def down, do: Legion.Store.Postgres.Migration.down(version: 2) + end + + ## Options + + - `:table` - the table name, defaults to `"legion_agents"`. Must match + the `:table` given to `use Legion.Store.Postgres`. + - `:version` - the target version. `up/1` defaults to the latest version, + `down/1` to rolling everything back. + """ + + # Legion does not depend on Ecto - these run inside the host app's + # migrations, where Ecto.Migration is available. + @compile {:no_warn_undefined, Ecto.Migration} + + @default_table "legion_agents" + @initial_version 1 + @current_version 1 + + @doc "Migrates the agents table up to `:version`, defaulting to the latest." + def up(opts \\ []) do + table = table(opts) + version = Keyword.get(opts, :version, @current_version) + migrated = migrated_version(opts) + + if migrated < version do + for step <- (migrated + 1)..version, sql <- List.wrap(up_sql(step, table)) do + Ecto.Migration.execute(sql) + end + + record_version(table, version) + end + + :ok + end + + @doc "Rolls the agents table back down to and including `:version`." + def down(opts \\ []) do + table = table(opts) + version = Keyword.get(opts, :version, @initial_version) + migrated = migrated_version(opts) + + if migrated >= version do + for step <- migrated..version//-1, sql <- List.wrap(down_sql(step, table)) do + Ecto.Migration.execute(sql) + end + + record_version(table, version - 1) + end + + :ok + end + + @doc "Returns the version the database is migrated to, `0` when the table is absent." + def migrated_version(opts \\ []) do + # The version lives in the table's comment, read directly so it reflects + # the database before this migration's queued commands run. + query = """ + SELECT pg_catalog.obj_description(pg_class.oid, 'pg_class') + FROM pg_class + WHERE pg_class.relname = '#{table(opts)}' + """ + + case Ecto.Migration.repo().query(query, [], log: false) do + {:ok, %{rows: [[version]]}} when is_binary(version) -> String.to_integer(version) + _ -> 0 + end + end + + defp record_version(_table, 0), do: :ok + + defp record_version(table, version) do + Ecto.Migration.execute("COMMENT ON TABLE #{table} IS '#{version}'") + end + + # The trigger notifies the table's channel with the agent_id on every + # write, so consumers (e.g. LegionWeb.Source.Listener) can refresh from the + # database instead of capturing telemetry. Postgres has no CREATE TRIGGER + # IF NOT EXISTS, so the trigger is dropped first to keep the step idempotent. + @doc false + def up_sql(1, table) do + [ + """ + CREATE TABLE IF NOT EXISTS #{table} ( + agent_id text PRIMARY KEY, + agent_module text, + parent_agent_id text, + pid bytea, + status text, + started_at bigint, + snapshot bytea, + inserted_at timestamptz NOT NULL DEFAULT now(), + updated_at timestamptz NOT NULL DEFAULT now() + ) + """, + """ + CREATE OR REPLACE FUNCTION #{table}_notify() RETURNS trigger AS $$ + BEGIN + PERFORM pg_notify('#{table}', NEW.agent_id); + RETURN NEW; + END; + $$ LANGUAGE plpgsql + """, + "DROP TRIGGER IF EXISTS #{table}_notify ON #{table}", + """ + CREATE TRIGGER #{table}_notify AFTER INSERT OR UPDATE ON #{table} + FOR EACH ROW EXECUTE FUNCTION #{table}_notify() + """ + ] + end + + @doc false + def down_sql(1, table) do + [ + "DROP TABLE IF EXISTS #{table}", + "DROP FUNCTION IF EXISTS #{table}_notify()" + ] + end + + defp table(opts), do: Keyword.get(opts, :table, @default_table) +end diff --git a/lib/legion/telemetry.ex b/lib/legion/telemetry.ex index ac6872f..dcc9fff 100644 --- a/lib/legion/telemetry.ex +++ b/lib/legion/telemetry.ex @@ -8,21 +8,23 @@ defmodule Legion.Telemetry do - `[:legion, :agent, :started]` — agent process finished `init/1` - Measurements: `%{system_time: integer}` - - Metadata: `%{agent: module, run_id: reference, parent_run_id: reference}` - - `parent_run_id` is only present when the agent was started inside another + - Metadata: `%{agent: module, agent_id: term, parent_agent_id: term}` + - `agent_id` names the conversation — stable across restarts, so a resumed + conversation emits under the same id. + - `parent_agent_id` is only present when the agent was started inside another agent's run. Not emitted if `init/1` itself crashes (e.g. while building the system prompt) — in that case `:stopped` is not emitted either, since GenServer does not call `terminate/2` on init failure. - `[:legion, :agent, :stopped]` — agent process terminated via `terminate/2` - Measurements: `%{system_time: integer}` - - Metadata: `%{agent: module, run_id: reference}` (plus `parent_run_id` + - Metadata: `%{agent: module, agent_id: term}` (plus `parent_agent_id` when the parent's run is still on the process Vault) ## Agent Message Events (spans) - `[:legion, :agent, :message, :start | :stop | :exception]` — agent handling a message - - Metadata includes: `agent`, `run_id`, `message` + - Metadata includes: `agent`, `agent_id`, `message` - Stop adds: `iterations` (count of assistant turns in this message), `status` (`:ok` or `:cancel`), `result` (the value returned, or the cancellation reason such as `:reached_max_iterations`), and `bindings` @@ -31,18 +33,18 @@ defmodule Legion.Telemetry do ## Iteration Events (spans) - `[:legion, :iteration, :start | :stop | :exception]` - - Metadata includes: `agent`, `run_id`, `iteration` + - Metadata includes: `agent`, `agent_id`, `iteration` - Stop adds: `action` ## LLM Request Events (spans) - `[:legion, :llm, :request, :start | :stop | :exception]` - - Metadata includes: `agent`, `run_id`, `model`, `message_count`, `iteration` + - Metadata includes: `agent`, `agent_id`, `model`, `message_count`, `iteration` ## Sandbox Eval Events (spans) - `[:legion, :sandbox, :eval, :start | :stop | :exception]` - - Metadata includes: `agent`, `run_id`, `code` + - Metadata includes: `agent`, `agent_id`, `code` - Stop adds: `success`, `result` or `error` ## Default Logger @@ -114,10 +116,10 @@ defmodule Legion.Telemetry do Wraps a function with `:start` / `:stop` / `:exception` telemetry events. The function should return `{result, extra_stop_metadata}`. - `run_id` is automatically injected from the process dictionary. + `agent_id` is automatically injected from the process dictionary. """ def span(event_prefix, metadata, fun) when is_function(fun, 0) do - metadata = with_run_id(metadata) + metadata = with_agent_id(metadata) start_time = System.monotonic_time() :telemetry.execute(event_prefix ++ [:start], %{system_time: System.system_time()}, metadata) @@ -159,16 +161,16 @@ defmodule Legion.Telemetry do end @doc """ - Emits a single telemetry event. Injects `run_id` from the process dictionary. + Emits a single telemetry event. Injects `agent_id` from the process dictionary. """ def emit(event, measurements \\ %{}, metadata) do - :telemetry.execute(event, measurements, with_run_id(metadata)) + :telemetry.execute(event, measurements, with_agent_id(metadata)) end - defp with_run_id(metadata) do + defp with_agent_id(metadata) do metadata - |> put_from_vault(:run_id) - |> put_from_vault(:parent_run_id) + |> put_from_vault(:agent_id) + |> put_from_vault(:parent_agent_id) end defp put_from_vault(metadata, key) do @@ -283,7 +285,7 @@ defmodule Legion.Telemetry do Logger.log(level, "#{prefix} #{message}") end - # ANSI colors indexed by run_id hash for visual grouping + # ANSI colors indexed by agent_id hash for visual grouping @colors [ IO.ANSI.cyan(), IO.ANSI.green(), @@ -298,15 +300,15 @@ defmodule Legion.Telemetry do ] defp run_prefix(meta) do - color = color_for(meta[:run_id]) - marker = if meta[:parent_run_id], do: "┃▸", else: "┃ " + color = color_for(meta[:agent_id]) + marker = if meta[:parent_agent_id], do: "┃▸", else: "┃ " "#{color}#{marker}#{IO.ANSI.reset()}" end defp color_for(nil), do: IO.ANSI.white() - defp color_for(run_id) do - index = :erlang.phash2(run_id, length(@colors)) + defp color_for(agent_id) do + index = :erlang.phash2(agent_id, length(@colors)) Enum.at(@colors, index) end diff --git a/lib/legion/tools/human_tool.ex b/lib/legion/tools/human_tool.ex index fc0ced1..09eb189 100644 --- a/lib/legion/tools/human_tool.ex +++ b/lib/legion/tools/human_tool.ex @@ -28,9 +28,19 @@ defmodule Legion.Tools.HumanTool do Asks a human a question and blocks until they respond. Returns the human's answer as a string. - Use this tool with `eval_and_continue`. + + Only works under `eval_and_continue` - the answer comes back to you as the + eval result so you can act on it. Under `eval_and_complete` the turn would + end the moment this code returns and the answer would be discarded, so + `ask/1` raises there. """ def ask(question) when is_binary(question) do + if Vault.get(:current_action) == "eval_and_complete" do + raise "HumanTool.ask/1 must run under eval_and_continue: with eval_and_complete " <> + "the turn ends when this code returns and the human's answer is discarded. " <> + "Re-run this code with action eval_and_continue, then act on the answer." + end + config = Vault.get(__MODULE__, []) handler = config[:handler] @@ -41,7 +51,7 @@ defmodule Legion.Tools.HumanTool do timeout = config[:timeout] || :infinity ref = make_ref() - send(handler, {:human_request, ref, self(), question, %{run_id: Vault.get(:run_id)}}) + send(handler, {:human_request, ref, self(), question, %{agent_id: Vault.get(:agent_id)}}) receive do {:human_response, ^ref, answer} -> answer diff --git a/mix.exs b/mix.exs index b36d0c1..a5ffd1e 100644 --- a/mix.exs +++ b/mix.exs @@ -48,7 +48,14 @@ defmodule Legion.MixProject do defp groups_for_modules do [ - Core: [Legion, Legion.Agent, Legion.Tool, Legion.Store, Legion.Store.Postgres], + Core: [ + Legion, + Legion.Agent, + Legion.Tool, + Legion.Store, + Legion.Store.Postgres, + Legion.Store.Postgres.Migration + ], Runtime: [Legion.AgentServer, Legion.Executor, ~r/^Legion\.Sandbox/], Tools: [~r/^Legion\.Tools\./], Internals: [Legion.AgentPrompt, Legion.SourceRegistry, Legion.Telemetry] diff --git a/test/legion/agent_server_test.exs b/test/legion/agent_server_test.exs index 9550278..a896a77 100644 --- a/test/legion/agent_server_test.exs +++ b/test/legion/agent_server_test.exs @@ -22,6 +22,20 @@ defmodule Legion.AgentServerTest do def tools, do: [Legion.Test.Support.MathTool] end + defmodule ChildAgent do + @moduledoc "Sub-agent invoked through AgentTool." + use Legion.Agent + end + + defmodule DelegatingAgent do + @moduledoc "Agent that delegates work to ChildAgent." + use Legion.Agent + + def tools, do: [Legion.Tools.AgentTool] + def tool_config(Legion.Tools.AgentTool), do: [agents: [ChildAgent]] + def tool_config(_tool), do: [] + end + setup :set_mimic_global @moduletag capture_log: true @@ -58,10 +72,12 @@ defmodule Legion.AgentServerTest do messages = Legion.get_messages(pid) assert [ - %{role: "system", content: _system}, - %{role: "user", content: "What is the capital of France?"}, - %{role: "assistant"} | _ + %{role: "system", type: :system, content: _system}, + %{role: "user", type: :user, content: "What is the capital of France?", at: at}, + %{role: "assistant", type: :assistant} | _ ] = messages + + assert is_integer(at) end end @@ -407,6 +423,34 @@ defmodule Legion.AgentServerTest do def save(agent_id, snapshot) do Agent.update(__MODULE__, &Map.put(&1, agent_id, snapshot)) end + + def save_run(agent_id, metadata) do + Agent.update(__MODULE__, &Map.put(&1, {:run, agent_id}, metadata)) + end + + def save_status(agent_id, status) do + Agent.update(__MODULE__, fn state -> + Map.update(state, {:statuses, agent_id}, [status], &[status | &1]) + end) + end + + def statuses(agent_id) do + Agent.get(__MODULE__, &Map.get(&1, {:statuses, agent_id}, [])) + |> Enum.reverse() + end + + def get_run(agent_id) do + case Agent.get(__MODULE__, &Map.get(&1, {:run, agent_id})) do + nil -> nil + metadata -> Map.put(metadata, :agent_id, agent_id) + end + end + + def runs do + Agent.get(__MODULE__, fn state -> + for {{:run, agent_id}, metadata} <- state, do: Map.put(metadata, :agent_id, agent_id) + end) + end end describe "persistence" do @@ -415,6 +459,20 @@ defmodule Legion.AgentServerTest do :ok end + test "brackets each turn with :running and :idle status writes" do + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + llm_response("Paris") + end) + + {:ok, pid} = Legion.start_link(MathAgent, store: MemoryStore, agent_id: "statuses") + {:ok, _} = Legion.call(pid, "What is the capital of France?") + + assert MemoryStore.statuses("statuses") == [:running, :idle] + + {:ok, _} = Legion.call(pid, "And of Germany?") + assert MemoryStore.statuses("statuses") == [:running, :idle, :running, :idle] + end + test "saves a snapshot before the caller receives its reply" do stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> llm_response("Paris") @@ -433,6 +491,21 @@ defmodule Legion.AgentServerTest do refute Enum.any?(messages, &(&1.role == "system")) end + test "persists the user message before the turn runs" do + test_process = self() + + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + send(test_process, {:snapshot_during_turn, MemoryStore.load("early-save")}) + llm_response("Paris") + end) + + {:ok, pid} = Legion.start_link(MathAgent, store: MemoryStore, agent_id: "early-save") + {:ok, _} = Legion.call(pid, "What is the capital of France?") + + assert_received {:snapshot_during_turn, {:ok, %{messages: messages}}} + assert [%{role: "user", content: "What is the capital of France?"}] = messages + end + test "a one-off execute/3 persists its snapshot before stopping" do stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> llm_response("Paris") @@ -527,9 +600,9 @@ defmodule Legion.AgentServerTest do assert {:ok, _snapshot} = MemoryStore.load("global-store") end - test "get_agent_id/1 returns nil without a store" do + test "get_agent_id/1 returns a generated id without a store" do {:ok, pid} = Legion.start_link(MathAgent) - assert Legion.get_agent_id(pid) == nil + assert is_binary(Legion.get_agent_id(pid)) end test "raises when :agent_id is given without a :store" do @@ -537,6 +610,70 @@ defmodule Legion.AgentServerTest do Legion.start_link(MathAgent, agent_id: "orphan") end end + + test "records run metadata on start" do + {:ok, pid} = Legion.start_link(MathAgent, store: MemoryStore, agent_id: "meta") + + assert [run] = MemoryStore.runs() + assert run.agent_id == "meta" + assert run.agent_module == MathAgent + assert run.parent_agent_id == nil + assert run.pid == pid + assert is_integer(run.started_at) + 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") + + assert Legion.running?(pid) + assert {:ok, ^pid} = Legion.resume("resume-live", store: MemoryStore) + end + + test "resume/2 restarts a stopped conversation from its run metadata" do + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + llm_response("Paris") + end) + + {:ok, pid} = Legion.start_link(MathAgent, store: MemoryStore, agent_id: "resume-dead") + {:ok, _} = Legion.call(pid, "What is the capital of France?") + GenServer.stop(pid) + refute Legion.running?(pid) + + {:ok, revived} = Legion.resume("resume-dead", store: MemoryStore) + + assert revived != pid + + assert [ + %{role: "system"}, + %{role: "user", content: "What is the capital of France?"}, + %{role: "assistant"} | _ + ] = 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) + end + 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 + llm_response("child done") + else + llm_eval_response(~s|AgentTool.call(ChildAgent, "child task")|) + end + end) + + {:ok, _} = Legion.execute(DelegatingAgent, "parent task", store: MemoryStore) + + runs = MemoryStore.runs() + parent = Enum.find(runs, &(&1.agent_module == DelegatingAgent)) + child = Enum.find(runs, &(&1.agent_module == ChildAgent)) + + assert child.parent_agent_id == parent.agent_id + assert {:ok, %{messages: [%{content: "child task"} | _]}} = MemoryStore.load(child.agent_id) + end end describe "binding_scope" do diff --git a/test/legion/store/postgres_db_test.exs b/test/legion/store/postgres_db_test.exs index 630848c..c38ae9a 100644 --- a/test/legion/store/postgres_db_test.exs +++ b/test/legion/store/postgres_db_test.exs @@ -36,4 +36,85 @@ defmodule Legion.Store.PostgresDbTest do assert {:ok, %{bindings: [v: 2]}} = Store.load("user_42") end + + test "save_run/2 records conversation identity and keeps the parent on conflict" do + metadata = %{ + agent_module: MyApp.Worker, + parent_agent_id: "parent-1", + pid: self(), + started_at: 100 + } + + assert :ok = Store.save_run("child-1", metadata) + + # Resumed from elsewhere: no parent this time, later started_at. + assert :ok = Store.save_run("child-1", %{metadata | parent_agent_id: nil, started_at: 200}) + + %{rows: [[agent_module, parent_agent_id, started_at]]} = + Postgrex.query!( + :legion_store_test, + "SELECT agent_module, parent_agent_id, started_at FROM legion_agents WHERE agent_id = $1", + ["child-1"] + ) + + assert agent_module == "MyApp.Worker" + assert parent_agent_id == "parent-1" + assert started_at == 200 + assert %{pid: pid} = Store.get_run("child-1") + assert pid == self() + end + + test "save_status/2 flips the stored status and save_run/2 resets it to idle" do + :ok = Store.save_run("s1", %{agent_module: A, parent_agent_id: nil, started_at: 1}) + assert %{status: :idle} = Store.get_run("s1") + + :ok = Store.save_status("s1", :running) + assert %{status: :running} = Store.get_run("s1") + + :ok = Store.save_run("s1", %{agent_module: A, parent_agent_id: nil, started_at: 2}) + assert %{status: :idle} = Store.get_run("s1") + end + + test "the trigger notifies the table's channel with the agent_id on every write" do + {:ok, notifications} = + Postgrex.Notifications.start_link( + hostname: System.get_env("POSTGRES_HOST", "localhost"), + port: String.to_integer(System.get_env("POSTGRES_PORT", "5432")), + username: System.get_env("POSTGRES_USER", "postgres"), + password: System.get_env("POSTGRES_PASSWORD", "postgres"), + database: System.get_env("POSTGRES_DB", "postgres") + ) + + {:ok, _ref} = Postgrex.Notifications.listen(notifications, "legion_agents") + + :ok = Store.save_run("notify-1", %{agent_module: A, parent_agent_id: nil, started_at: 1}) + assert_receive {:notification, _pid, _ref, "legion_agents", "notify-1"}, 1_000 + + :ok = Store.save_status("notify-1", :running) + assert_receive {:notification, _pid, _ref, "legion_agents", "notify-1"}, 1_000 + + :ok = Store.save("notify-1", %{messages: [], bindings: []}) + assert_receive {:notification, _pid, _ref, "legion_agents", "notify-1"}, 1_000 + + GenServer.stop(notifications) + end + + test "save_run/2 then save/2 share one row and load/1 sees the snapshot" do + assert :ok = Store.save_run("both", %{agent_module: A, parent_agent_id: nil, started_at: 1}) + assert :error = Store.load("both") + + snapshot = %{messages: [%{role: "user", content: "hi"}], bindings: []} + assert :ok = Store.save("both", snapshot) + + assert {:ok, ^snapshot} = Store.load("both") + + %{rows: [[count]]} = + Postgrex.query!( + :legion_store_test, + "SELECT COUNT(*) FROM legion_agents WHERE agent_id = $1", + ["both"] + ) + + assert count == 1 + end end diff --git a/test/legion/store/postgres_test.exs b/test/legion/store/postgres_test.exs index 583fd3b..2b20b42 100644 --- a/test/legion/store/postgres_test.exs +++ b/test/legion/store/postgres_test.exs @@ -2,21 +2,61 @@ defmodule Legion.Store.PostgresTest do use ExUnit.Case, async: true defmodule FakeRepo do - @moduledoc "Emulates repo.query!/2 for the two statements the store issues." + @moduledoc "Emulates repo.query!/2 for the statements the store issues." - def start_link, do: Agent.start_link(fn -> %{} end, name: __MODULE__) + def start_link, do: Agent.start_link(fn -> %{snapshots: %{}, runs: %{}} end, name: __MODULE__) def query!("SELECT snapshot FROM " <> _rest, [agent_id]) do - case Agent.get(__MODULE__, &Map.fetch(&1, agent_id)) do - {:ok, snapshot} -> %{rows: [[snapshot]]} - :error -> %{rows: []} - end + Agent.get(__MODULE__, fn state -> + case {state.snapshots, state.runs} do + {%{^agent_id => snapshot}, _runs} -> %{rows: [[snapshot]]} + {_snapshots, %{^agent_id => _run}} -> %{rows: [[nil]]} + _neither -> %{rows: []} + end + end) end def query!("INSERT INTO " <> _rest, [agent_id, snapshot]) do - Agent.update(__MODULE__, &Map.put(&1, agent_id, snapshot)) + Agent.update(__MODULE__, &put_in(&1.snapshots[agent_id], snapshot)) + %{num_rows: 1} + end + + def query!("INSERT INTO " <> _rest, [agent_id, agent_module, parent_agent_id, pid, started_at]) do + Agent.update( + __MODULE__, + &put_in(&1.runs[agent_id], %{ + agent_module: agent_module, + parent_agent_id: parent_agent_id, + pid: pid, + status: "idle", + started_at: started_at + }) + ) + %{num_rows: 1} end + + def query!("UPDATE " <> _rest, [agent_id, status]) do + Agent.update(__MODULE__, &put_in(&1.runs[agent_id].status, status)) + %{num_rows: 1} + end + + def query!("SELECT agent_id" <> _rest = sql, [param]) do + rows = + Agent.get(__MODULE__, fn state -> + for {agent_id, run} <- state.runs do + [agent_id, run.agent_module, run.parent_agent_id, run.pid, run.status, run.started_at] + end + end) + + if String.contains?(sql, "WHERE agent_id") do + %{rows: Enum.filter(rows, fn [agent_id | _rest] -> agent_id == param end)} + else + %{rows: rows |> Enum.sort_by(&List.last/1, :desc) |> Enum.take(param)} + end + end + + def run(agent_id), do: Agent.get(__MODULE__, & &1.runs[agent_id]) end defmodule Store do @@ -44,6 +84,76 @@ defmodule Legion.Store.PostgresTest do assert_raise FunctionClauseError, fn -> Store.save(42, %{messages: [], bindings: []}) end end + test "save_run/2 stores the module in inspect form with parent, pid, and start time" do + metadata = %{ + agent_module: Legion.Test.Support.MathAgent, + parent_agent_id: "p1", + pid: self(), + started_at: 123 + } + + assert :ok = Store.save_run("user_42", metadata) + + assert FakeRepo.run("user_42") == %{ + agent_module: "Legion.Test.Support.MathAgent", + parent_agent_id: "p1", + pid: :erlang.term_to_binary(self()), + status: "idle", + started_at: 123 + } + end + + test "save_status/2 flips the run status and get_run/1 decodes it" do + :ok = Store.save_run("s1", %{agent_module: SomeAgent, parent_agent_id: nil, started_at: 1}) + assert %{status: :idle} = Store.get_run("s1") + + :ok = Store.save_status("s1", :running) + assert %{status: :running} = Store.get_run("s1") + + :ok = Store.save_status("s1", :idle) + assert %{status: :idle} = Store.get_run("s1") + end + + test "get_run/1 round-trips the pid" do + metadata = %{ + agent_module: SomeAgent, + parent_agent_id: nil, + pid: self(), + started_at: 123 + } + + :ok = Store.save_run("with-pid", metadata) + + assert %{pid: pid} = Store.get_run("with-pid") + assert pid == self() + end + + test "list_runs/1 returns decoded runs newest first" do + :ok = Store.save_run("a", %{agent_module: SomeAgent, parent_agent_id: nil, started_at: 1}) + :ok = Store.save_run("b", %{agent_module: OtherAgent, parent_agent_id: "a", started_at: 2}) + + assert [ + %{agent_id: "b", agent_module: OtherAgent, parent_agent_id: "a", started_at: 2}, + %{agent_id: "a", agent_module: SomeAgent, parent_agent_id: nil, started_at: 1} + ] = Store.list_runs(10) + + assert [%{agent_id: "b"}] = Store.list_runs(1) + end + + test "get_run/1 returns the decoded run, or nil when missing" do + :ok = Store.save_run("a", %{agent_module: SomeAgent, parent_agent_id: nil, started_at: 1}) + + assert %{agent_id: "a", agent_module: SomeAgent} = Store.get_run("a") + assert Store.get_run("missing") == nil + end + + test "load/1 returns :error when only run metadata exists (snapshot still null)" do + metadata = %{agent_module: SomeAgent, parent_agent_id: nil, started_at: 1} + + assert :ok = Store.save_run("started-only", metadata) + assert :error = Store.load("started-only") + end + test "a custom table name is interpolated into the statements" do defmodule TableCapturingRepo do def query!(sql, _params), do: send(self(), {:sql, sql}) && %{rows: []} diff --git a/test/legion/tools/human_tool_test.exs b/test/legion/tools/human_tool_test.exs index b2937f3..0156629 100644 --- a/test/legion/tools/human_tool_test.exs +++ b/test/legion/tools/human_tool_test.exs @@ -38,8 +38,22 @@ defmodule Legion.Tools.HumanToolTest do assert HumanTool.ask("What is your name?") == "fake answer" end - test "ask/1 passes question and run_id metadata to handler" do - run_id = make_ref() + test "ask/1 works under eval_and_continue" do + Vault.unsafe_merge(%{current_action: "eval_and_continue"}) + + assert HumanTool.ask("What is your name?") == "fake answer" + end + + test "ask/1 raises under eval_and_complete so the answer is not discarded" do + Vault.unsafe_merge(%{current_action: "eval_and_complete"}) + + assert_raise RuntimeError, ~r/must run under eval_and_continue/, fn -> + HumanTool.ask("What is your name?") + end + end + + test "ask/1 passes question and agent_id metadata to handler" do + agent_id = "human-tool-test" test_pid = self() spy_pid = @@ -51,10 +65,10 @@ defmodule Legion.Tools.HumanToolTest do end end) - Vault.unsafe_merge(%{Legion.Tools.HumanTool => [handler: spy_pid], run_id: run_id}) + Vault.unsafe_merge(%{Legion.Tools.HumanTool => [handler: spy_pid], agent_id: agent_id}) assert HumanTool.ask("tell me something") == "spy answer" - assert_receive {:captured, "tell me something", %{run_id: ^run_id}} + assert_receive {:captured, "tell me something", %{agent_id: ^agent_id}} end test "ask/1 raises when no handler is configured" do diff --git a/test/test_helper.exs b/test/test_helper.exs index 5cae099..efed6d6 100644 --- a/test/test_helper.exs +++ b/test/test_helper.exs @@ -13,17 +13,12 @@ Legion.Telemetry.attach_default_logger() database: System.get_env("POSTGRES_DB", "postgres") ) -Postgrex.query!( - :legion_store_test, - """ - CREATE TABLE IF NOT EXISTS legion_agents ( - agent_id text PRIMARY KEY, - snapshot bytea NOT NULL, - inserted_at timestamptz NOT NULL DEFAULT now(), - updated_at timestamptz NOT NULL DEFAULT now() - ) - """, - [] -) +for sql <- Legion.Store.Postgres.Migration.down_sql(1, "legion_agents") do + Postgrex.query!(:legion_store_test, sql, []) +end + +for sql <- Legion.Store.Postgres.Migration.up_sql(1, "legion_agents") do + Postgrex.query!(:legion_store_test, sql, []) +end ExUnit.start(exclude: [:integration]) From 07e837d3bedf2ffa54e8a72a179c325e4de147e3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tomasz=20Id=C5=BAkowski?= Date: Thu, 16 Jul 2026 14:29:07 +0200 Subject: [PATCH 04/30] Move handling AgentServer PID from Store to Registry --- lib/legion.ex | 33 ++++++++++++++++++++------ lib/legion/agent_server.ex | 3 ++- lib/legion/application.ex | 12 ++++++++++ lib/legion/store.ex | 18 ++++---------- lib/legion/store/postgres.ex | 25 +++++++------------ lib/legion/store/postgres/migration.ex | 1 - mix.exs | 1 + test/legion/agent_server_test.exs | 9 +++++-- test/legion/store/postgres_db_test.exs | 3 --- test/legion/store/postgres_test.exs | 23 +++--------------- 10 files changed, 65 insertions(+), 63 deletions(-) create mode 100644 lib/legion/application.ex diff --git a/lib/legion.ex b/lib/legion.ex index 1545836..da05521 100644 --- a/lib/legion.ex +++ b/lib/legion.ex @@ -127,12 +127,32 @@ defmodule Legion do def running?(_other), do: false + @doc """ + Looks up the live process for an agent id. + + Uses `Legion.AgentRegistry` as the runtime 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`. + + ## Examples + + {:ok, pid} = Legion.lookup("user_42:chat_7") + :error = Legion.lookup("missing_agent_id") + """ + def lookup(agent_id) do + case Registry.lookup(Legion.AgentRegistry, agent_id) do + [{pid, _}] -> {:ok, pid} + [] -> :error + end + end + @doc """ Resumes a persisted conversation. - Returns the recorded process if it is still `running?/1`; otherwise starts - the agent again under the same `agent_id`, so it reloads its snapshot from - the store. `opts` are passed through to `start_link/2`. + Checks the persisted run exists, then returns the process registered for + `agent_id` if the agent is already running. If no process is registered, + starts the agent again under the same `agent_id`, so it reloads its snapshot + from the store. `opts` are passed through to `start_link/2`. Requires a store implementing `c:Legion.Store.get_run/1` - pass `:store` or configure one globally. Raises if the store has no run for `agent_id`. @@ -153,10 +173,9 @@ defmodule Legion do raise ArgumentError, "no run recorded for agent_id #{inspect(agent_id)} in #{inspect(store)}" - if running?(run[:pid]) do - {:ok, run[:pid]} - else - start_link(run.agent_module, Keyword.put(opts, :agent_id, agent_id)) + case lookup(agent_id) do + {:ok, pid} -> {:ok, pid} + :error -> start_link(run.agent_module, Keyword.put(opts, :agent_id, agent_id)) end end diff --git a/lib/legion/agent_server.ex b/lib/legion/agent_server.ex index 2a9bf72..d2ddf1c 100644 --- a/lib/legion/agent_server.ex +++ b/lib/legion/agent_server.ex @@ -62,6 +62,8 @@ defmodule Legion.AgentServer do Vault.unsafe_put(:parent_agent_id, parent_agent_id) if store, do: Vault.unsafe_put(:store, store) + Registry.register(Legion.AgentRegistry, agent_id, self()) + for tool <- agent_module.tools() do Vault.unsafe_put(tool, agent_module.tool_config(tool)) end @@ -77,7 +79,6 @@ defmodule Legion.AgentServer do save_run(store, agent_id, %{ agent_module: agent_module, parent_agent_id: parent_agent_id, - pid: self(), started_at: System.system_time(:millisecond) }) diff --git a/lib/legion/application.ex b/lib/legion/application.ex new file mode 100644 index 0000000..0d541ac --- /dev/null +++ b/lib/legion/application.ex @@ -0,0 +1,12 @@ +defmodule Legion.Application do + use Application + + @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 +end diff --git a/lib/legion/store.ex b/lib/legion/store.ex index aacb154..bfbd600 100644 --- a/lib/legion/store.ex +++ b/lib/legion/store.ex @@ -89,13 +89,10 @@ defmodule Legion.Store do Implement the optional `c:save_run/2` to also record each conversation's identity: which agent module ran, under which parent conversation (for - sub-agents), when, and under which `pid`. Legion calls it once per agent - start, so persisted conversations stay attributable to the agent tree that - produced them. Restarting a conversation under the same `agent_id` calls it - again with a fresh `started_at` and `pid` - treat it as an upsert. The - stored pid always names the newest process for the conversation, so - `Legion.running?/1` and `Legion.resume/2` can tell whether it is still - live (a pid outlives the VM that wrote it, so never trust it blindly). On conflict, keep the stored + sub-agents), and when. Legion calls it once per agent start, so persisted + conversations stay attributable to the agent tree that produced them. + Restarting a conversation under the same `agent_id` calls it again with a + fresh `started_at` - treat it as an upsert. On conflict, keep the stored `parent_agent_id`: the callback reports where the agent was started *this time*, so a conversation resumed from iex or another agent's tree would otherwise be silently reparented. @@ -104,10 +101,7 @@ defmodule Legion.Store do Implement the optional `c:save_status/2` to also record whether a conversation is mid-turn: Legion calls it with `:running` when a message - starts and `:idle` after the turn's snapshot is saved. Combined with the - stored pid this lets a consumer distinguish a live agent that is working - from one waiting for input, and a conversation that crashed mid-turn - (status still `:running` under a dead pid) from one that finished. + starts and `:idle` after the turn's snapshot is saved. ## Reading conversations back @@ -123,14 +117,12 @@ defmodule Legion.Store do @type run_metadata :: %{ agent_module: module(), parent_agent_id: agent_id() | nil, - pid: pid(), started_at: integer() } @type run :: %{ agent_id: agent_id(), agent_module: module() | nil, parent_agent_id: agent_id() | nil, - pid: pid() | nil, status: :running | :idle | nil, started_at: integer() | nil } diff --git a/lib/legion/store/postgres.ex b/lib/legion/store/postgres.ex index 94cb0f7..626642f 100644 --- a/lib/legion/store/postgres.ex +++ b/lib/legion/store/postgres.ex @@ -39,16 +39,16 @@ defmodule Legion.Store.Postgres do The store also implements `c:Legion.Store.save_run/2`, so the same row carries the conversation's identity: `agent_module` (in `inspect/1` form, e.g. `"MyApp.ResearchAgent"`), `parent_agent_id` linking a sub-agent to the - conversation that spawned it, `pid` of the newest process that ran the - conversation (a `:erlang.term_to_binary/1` blob, like the snapshot), and - `started_at` in milliseconds (last start wins). `snapshot` is null until the conversation's first turn completes; + conversation that spawned it, and `started_at` in milliseconds (last start + wins). `snapshot` is null until the conversation's first turn completes; `parent_agent_id` is kept once set, so resuming from elsewhere does not reparent the conversation. `c:Legion.Store.save_status/2` is implemented as well: the row's `status` flips to `'running'` when a turn starts and back to `'idle'` when it - completes (`save_run` resets it to `'idle'` on start), so a `'running'` - status under a dead pid identifies a conversation that crashed mid-turn. + completes (`save_run` resets it to `'idle'` on start), so consumers can + identify conversations that were mid-turn when persistence last observed + them. `c:Legion.Store.list_runs/1` and `c:Legion.Store.get_run/1` are implemented too, so persisted conversations can be read back into a view of past runs. @@ -71,12 +71,11 @@ defmodule Legion.Store.Postgres do """ save_run_sql = """ - INSERT INTO #{table} (agent_id, agent_module, parent_agent_id, pid, status, started_at, inserted_at, updated_at) - VALUES ($1, $2, $3, $4, 'idle', $5, now(), now()) + INSERT INTO #{table} (agent_id, agent_module, parent_agent_id, status, started_at, inserted_at, updated_at) + VALUES ($1, $2, $3, 'idle', $4, now(), now()) ON CONFLICT (agent_id) DO UPDATE SET agent_module = EXCLUDED.agent_module, parent_agent_id = COALESCE(#{table}.parent_agent_id, EXCLUDED.parent_agent_id), - pid = EXCLUDED.pid, status = 'idle', started_at = EXCLUDED.started_at, updated_at = now() @@ -84,7 +83,7 @@ defmodule Legion.Store.Postgres do save_status_sql = "UPDATE #{table} SET status = $2, updated_at = now() WHERE agent_id = $1" - run_columns = "agent_id, agent_module, parent_agent_id, pid, status, started_at" + run_columns = "agent_id, agent_module, parent_agent_id, status, started_at" list_runs_sql = "SELECT #{run_columns} FROM #{table} ORDER BY started_at DESC NULLS LAST LIMIT $1" @@ -116,7 +115,6 @@ defmodule Legion.Store.Postgres do agent_id, inspect(metadata.agent_module), metadata.parent_agent_id, - metadata[:pid] && :erlang.term_to_binary(metadata.pid), metadata.started_at ]) @@ -152,16 +150,11 @@ defmodule Legion.Store.Postgres do end @doc false - # A pid stored by a previous VM run decodes with its original creation - # number, so `Legion.running?/1` reads it as not running instead of - # colliding with a recycled pid value. - # sobelow_skip ["Misc.BinToTerm"] - def decode_run_row([agent_id, agent_module, parent_agent_id, pid, status, started_at]) do + def decode_run_row([agent_id, agent_module, parent_agent_id, status, started_at]) do %{ agent_id: agent_id, agent_module: agent_module && Module.concat([agent_module]), parent_agent_id: parent_agent_id, - pid: pid && :erlang.binary_to_term(pid), status: decode_status(status), started_at: started_at } diff --git a/lib/legion/store/postgres/migration.ex b/lib/legion/store/postgres/migration.ex index fe54336..0ad30f3 100644 --- a/lib/legion/store/postgres/migration.ex +++ b/lib/legion/store/postgres/migration.ex @@ -107,7 +107,6 @@ defmodule Legion.Store.Postgres.Migration do agent_id text PRIMARY KEY, agent_module text, parent_agent_id text, - pid bytea, status text, started_at bigint, snapshot bytea, diff --git a/mix.exs b/mix.exs index a5ffd1e..a17ac37 100644 --- a/mix.exs +++ b/mix.exs @@ -42,6 +42,7 @@ defmodule Legion.MixProject do def application do [ + mod: {Legion.Application, []}, extra_applications: [:logger] ] end diff --git a/test/legion/agent_server_test.exs b/test/legion/agent_server_test.exs index a896a77..918fae5 100644 --- a/test/legion/agent_server_test.exs +++ b/test/legion/agent_server_test.exs @@ -612,16 +612,21 @@ defmodule Legion.AgentServerTest do end test "records run metadata on start" do - {:ok, pid} = Legion.start_link(MathAgent, store: MemoryStore, agent_id: "meta") + {:ok, _pid} = Legion.start_link(MathAgent, store: MemoryStore, agent_id: "meta") assert [run] = MemoryStore.runs() assert run.agent_id == "meta" assert run.agent_module == MathAgent assert run.parent_agent_id == nil - assert run.pid == pid assert is_integer(run.started_at) end + test "registers the agent pid by agent_id" do + {:ok, pid} = Legion.start_link(MathAgent, store: MemoryStore, agent_id: "lookup") + + assert {:ok, ^pid} = Legion.lookup("lookup") + 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") diff --git a/test/legion/store/postgres_db_test.exs b/test/legion/store/postgres_db_test.exs index c38ae9a..f7e90bb 100644 --- a/test/legion/store/postgres_db_test.exs +++ b/test/legion/store/postgres_db_test.exs @@ -41,7 +41,6 @@ defmodule Legion.Store.PostgresDbTest do metadata = %{ agent_module: MyApp.Worker, parent_agent_id: "parent-1", - pid: self(), started_at: 100 } @@ -60,8 +59,6 @@ defmodule Legion.Store.PostgresDbTest do assert agent_module == "MyApp.Worker" assert parent_agent_id == "parent-1" assert started_at == 200 - assert %{pid: pid} = Store.get_run("child-1") - assert pid == self() end test "save_status/2 flips the stored status and save_run/2 resets it to idle" do diff --git a/test/legion/store/postgres_test.exs b/test/legion/store/postgres_test.exs index 2b20b42..8e14c61 100644 --- a/test/legion/store/postgres_test.exs +++ b/test/legion/store/postgres_test.exs @@ -21,13 +21,12 @@ defmodule Legion.Store.PostgresTest do %{num_rows: 1} end - def query!("INSERT INTO " <> _rest, [agent_id, agent_module, parent_agent_id, pid, started_at]) do + def query!("INSERT INTO " <> _rest, [agent_id, agent_module, parent_agent_id, started_at]) do Agent.update( __MODULE__, &put_in(&1.runs[agent_id], %{ agent_module: agent_module, parent_agent_id: parent_agent_id, - pid: pid, status: "idle", started_at: started_at }) @@ -45,7 +44,7 @@ defmodule Legion.Store.PostgresTest do rows = Agent.get(__MODULE__, fn state -> for {agent_id, run} <- state.runs do - [agent_id, run.agent_module, run.parent_agent_id, run.pid, run.status, run.started_at] + [agent_id, run.agent_module, run.parent_agent_id, run.status, run.started_at] end end) @@ -84,11 +83,10 @@ defmodule Legion.Store.PostgresTest do assert_raise FunctionClauseError, fn -> Store.save(42, %{messages: [], bindings: []}) end end - test "save_run/2 stores the module in inspect form with parent, pid, and start time" do + test "save_run/2 stores the module in inspect form with parent, and start time" do metadata = %{ agent_module: Legion.Test.Support.MathAgent, parent_agent_id: "p1", - pid: self(), started_at: 123 } @@ -97,7 +95,6 @@ defmodule Legion.Store.PostgresTest do assert FakeRepo.run("user_42") == %{ agent_module: "Legion.Test.Support.MathAgent", parent_agent_id: "p1", - pid: :erlang.term_to_binary(self()), status: "idle", started_at: 123 } @@ -114,20 +111,6 @@ defmodule Legion.Store.PostgresTest do assert %{status: :idle} = Store.get_run("s1") end - test "get_run/1 round-trips the pid" do - metadata = %{ - agent_module: SomeAgent, - parent_agent_id: nil, - pid: self(), - started_at: 123 - } - - :ok = Store.save_run("with-pid", metadata) - - assert %{pid: pid} = Store.get_run("with-pid") - assert pid == self() - end - test "list_runs/1 returns decoded runs newest first" do :ok = Store.save_run("a", %{agent_module: SomeAgent, parent_agent_id: nil, started_at: 1}) :ok = Store.save_run("b", %{agent_module: OtherAgent, parent_agent_id: "a", started_at: 2}) From a6d2026a70724db3bc14680a35a24a1046487909 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tomasz=20Id=C5=BAkowski?= Date: Thu, 16 Jul 2026 16:22:05 +0200 Subject: [PATCH 05/30] Resolve Credo style warnings --- lib/legion/application.ex | 4 ++++ lib/legion/store/postgres.ex | 6 ++++-- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/lib/legion/application.ex b/lib/legion/application.ex index 0d541ac..3776c97 100644 --- a/lib/legion/application.ex +++ b/lib/legion/application.ex @@ -1,4 +1,8 @@ defmodule Legion.Application do + @moduledoc """ + The OTP application for Legion. + """ + use Application @impl true diff --git a/lib/legion/store/postgres.ex b/lib/legion/store/postgres.ex index 626642f..138ca85 100644 --- a/lib/legion/store/postgres.ex +++ b/lib/legion/store/postgres.ex @@ -93,6 +93,8 @@ defmodule Legion.Store.Postgres do quote do @behaviour Legion.Store + alias Legion.Store.Postgres + @impl Legion.Store # sobelow_skip ["Misc.BinToTerm"] def load(agent_id) when is_binary(agent_id) do @@ -130,13 +132,13 @@ defmodule Legion.Store.Postgres do @impl Legion.Store def list_runs(limit) do %{rows: rows} = unquote(repo).query!(unquote(list_runs_sql), [limit]) - Enum.map(rows, &Legion.Store.Postgres.decode_run_row/1) + Enum.map(rows, &Postgres.decode_run_row/1) end @impl Legion.Store def get_run(agent_id) when is_binary(agent_id) do case unquote(repo).query!(unquote(get_run_sql), [agent_id]) do - %{rows: [row]} -> Legion.Store.Postgres.decode_run_row(row) + %{rows: [row]} -> Postgres.decode_run_row(row) %{rows: []} -> nil end end From d361d1e3af37fe2616c87880a24b0fa666d3c1b6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tomasz=20Id=C5=BAkowski?= Date: Fri, 17 Jul 2026 12:01:37 +0200 Subject: [PATCH 06/30] Modify Store contract - merge saves and loads --- lib/legion/store.ex | 142 ++++++++-------------- lib/legion/store/conversation_metadata.ex | 18 +++ lib/legion/store/conversation_state.ex | 17 +++ 3 files changed, 88 insertions(+), 89 deletions(-) create mode 100644 lib/legion/store/conversation_metadata.ex create mode 100644 lib/legion/store/conversation_state.ex diff --git a/lib/legion/store.ex b/lib/legion/store.ex index bfbd600..35f2e84 100644 --- a/lib/legion/store.ex +++ b/lib/legion/store.ex @@ -7,14 +7,15 @@ defmodule Legion.Store do {:ok, pid} = Legion.start_link(AssistantAgent, store: MyApp.AgentStore, agent_id: "user_42") - On start, the agent calls `c:load/1` and resumes from the snapshot if one - exists. The system prompt is regenerated fresh on every start, so prompt or - tool changes apply to restored conversations. + On start, the agent calls `c:get/1` and resumes from the returned + `Legion.Store.ConversationState` if one exists. The system prompt is + regenerated fresh on every start, so prompt or tool changes apply to + restored conversations. - After every completed turn, the agent calls `c:save/2` **before** replying - to the caller. A reply is a commit receipt: any turn a caller observed - survives a crash, restart, or deploy. A crash mid-turn rolls back to the - last completed turn. + After every completed turn, the agent calls `c:save/2` with a + `Legion.Store.ConversationState` **before** replying to the caller. A reply + is a commit receipt: any turn a caller observed survives a crash, restart, + or deploy. A crash mid-turn rolls back to the last completed turn. For Postgres users there is a ready-made adapter - see `Legion.Store.Postgres`: @@ -37,8 +38,8 @@ defmodule Legion.Store do ## Identifying a conversation `:agent_id` is the key a snapshot is saved under - it names one conversation, - not one user. A chat app with many chats per user keys by the chat; compose the - id however you like, since Legion treats it as opaque: + not one user. A chat app with many chats per user keys by the chat; + compose the id however you like, since Legion treats it as opaque: Legion.start_link(ChatAgent, agent_id: "user_42:chat_7") @@ -48,102 +49,65 @@ defmodule Legion.Store do resume an existing conversation. Two agents started under the same id race onto the same row, so route each conversation to a single process. - Or implement the two callbacks against any storage you like: + ## Required persistence - ## Example: hand-rolled Ecto store + Stores must implement `c:get/1` and `c:save/2`. - defmodule MyApp.AgentStore do - @behaviour Legion.Store - - def load(agent_id) do - case MyApp.Repo.get(MyApp.AgentSnapshot, agent_id) do - nil -> :error - row -> {:ok, :erlang.binary_to_term(row.snapshot)} - end - end - - def save(agent_id, snapshot) do - MyApp.Repo.insert!( - %MyApp.AgentSnapshot{id: agent_id, snapshot: :erlang.term_to_binary(snapshot)}, - on_conflict: {:replace, [:snapshot]}, - conflict_target: :id - ) - - :ok - end - end - - ## What is persisted - - The snapshot holds the conversation `:messages` (without the system prompt) - and the `:bindings` from evaluated code (relevant with - `binding_scope: :conversation`). Each message carries a `:type` - (`:user`, `:assistant`, `:eval_result`, or `:error`) and an `:at` timestamp - in milliseconds, so consumers can classify and order messages without - parsing content. Bindings are arbitrary Elixir terms - - values like pids, references, or functions will not survive + The required save payload is `Legion.Store.ConversationState`, which holds + the conversation `:messages` (without the system prompt) and the `:bindings` + from evaluated code (relevant with `binding_scope: :conversation`). Each + message carries a `:type` (`:user`, `:assistant`, `:eval_result`, or + `:error`) and an `:at` timestamp in milliseconds, so consumers can classify + and order messages without parsing content. Bindings are arbitrary Elixir + terms - values like pids, references, or functions will not survive serialization, so keep conversation-scoped variables to plain data if you persist agents. - ## Run metadata + ## Optional persistence + + Stores may also accept `Legion.Store.ConversationMetadata` through + `c:save/2` to record which agent module ran, under which parent + conversation, and when. - Implement the optional `c:save_run/2` to also record each conversation's - identity: which agent module ran, under which parent conversation (for - sub-agents), and when. Legion calls it once per agent start, so persisted - conversations stay attributable to the agent tree that produced them. - Restarting a conversation under the same `agent_id` calls it again with a - fresh `started_at` - treat it as an upsert. On conflict, keep the stored - `parent_agent_id`: the callback reports where the agent was started *this - time*, so a conversation resumed from iex or another agent's tree would - otherwise be silently reparented. + Stores may also accept `{:status, :running | :idle}` through `c:save/2` to + record whether the agent is mid-turn. - ## Turn status + Because metadata and status persistence are optional, snapshots returned + from `c:get/1` or `c:list/1` may have `metadata: nil` or `status: nil`. + Because a store may record optional data before any conversation state is + saved, `state` may also be nil. - Implement the optional `c:save_status/2` to also record whether a - conversation is mid-turn: Legion calls it with `:running` when a message - starts and `:idle` after the turn's snapshot is saved. + ## Reading conversations - ## Reading conversations back + `c:get/1` returns the persisted snapshot for one `agent_id`, or `:error` + when the store has no row for that id. - The optional `c:list_runs/1` and `c:get_run/1` callbacks expose persisted - runs for consumers that rebuild a view of past conversations from the store - alone - `LegionWeb` uses them when configured with - `config :legion_web, agents_source: :database`. `Legion.Store.Postgres` - implements both. + The optional `c:list/1` callback returns persisted snapshots newest first + for consumers that rebuild a view of past conversations from the store + alone. """ @type agent_id :: term() - @type snapshot :: %{messages: [map()], bindings: keyword()} - @type run_metadata :: %{ - agent_module: module(), - parent_agent_id: agent_id() | nil, - started_at: integer() + @type status :: :idle | :running + @type snapshot :: { + agent_id :: agent_id(), + metadata :: Legion.Store.ConversationMetadata.t() | nil, + status :: status() | nil, + state :: Legion.Store.ConversationState.t() | nil } - @type run :: %{ - agent_id: agent_id(), - agent_module: module() | nil, - parent_agent_id: agent_id() | nil, - status: :running | :idle | nil, - started_at: integer() | nil - } - - @doc "Returns the last saved snapshot for `agent_id`, or `:error` if none exists." - @callback load(agent_id()) :: {:ok, snapshot()} | :error - - @doc "Saves the snapshot for `agent_id`. Raise on failure - the turn is not acked until this returns." - @callback save(agent_id(), snapshot()) :: :ok - - @doc "Records run metadata when an agent starts. Optional. `started_at` is in milliseconds." - @callback save_run(agent_id(), run_metadata()) :: :ok + @type payload :: + Legion.Store.ConversationState.t() + | Legion.Store.ConversationMetadata.t() + | {:status, status()} - @doc "Records whether the conversation is mid-turn. Optional." - @callback save_status(agent_id(), :running | :idle) :: :ok + @doc "Returns the persisted snapshot for `agent_id`, or `:error` if none exists." + @callback get(agent_id()) :: {:ok, snapshot()} | :error - @doc "Returns the newest `limit` persisted runs, newest first. Optional." - @callback list_runs(limit :: pos_integer()) :: [run()] + @doc "Returns the newest `limit` persisted snapshots, newest first." + @callback list(limit :: pos_integer()) :: [snapshot()] - @doc "Returns the persisted run for `agent_id`, or `nil` if none exists. Optional." - @callback get_run(agent_id()) :: run() | nil + @doc "Saves a conversation state, conversation metadata, or status payload for `agent_id`." + @callback save(agent_id(), payload()) :: :ok - @optional_callbacks save_run: 2, save_status: 2, list_runs: 1, get_run: 1 + @optional_callbacks list: 1 end diff --git a/lib/legion/store/conversation_metadata.ex b/lib/legion/store/conversation_metadata.ex new file mode 100644 index 0000000..75d3ae8 --- /dev/null +++ b/lib/legion/store/conversation_metadata.ex @@ -0,0 +1,18 @@ +defmodule Legion.Store.ConversationMetadata do + @moduledoc """ + Identity metadata for a persisted agent conversation. + + Stores can save this payload to record which agent module owns a + conversation, which parent conversation spawned it, and when it started. + """ + + @enforce_keys [:agent_module, :parent_agent_id, :started_at] + defstruct [:agent_module, :parent_agent_id, :started_at] + + @typedoc "Metadata describing a persisted conversation's agent identity." + @type t() :: %__MODULE__{ + agent_module: module(), + parent_agent_id: Legion.Store.agent_id() | nil, + started_at: integer() + } +end diff --git a/lib/legion/store/conversation_state.ex b/lib/legion/store/conversation_state.ex new file mode 100644 index 0000000..39e3da9 --- /dev/null +++ b/lib/legion/store/conversation_state.ex @@ -0,0 +1,17 @@ +defmodule Legion.Store.ConversationState do + @moduledoc """ + Replayable state for a persisted agent conversation. + + Stores save this payload after turns so a restarted agent can restore its + conversation messages and conversation-scoped bindings. + """ + + @enforce_keys [:messages, :bindings] + defstruct [:messages, :bindings] + + @typedoc "Messages and bindings needed to restore a conversation." + @type t() :: %__MODULE__{ + messages: [map()], + bindings: keyword() + } +end From 31dabdf9c8a87550dbc7997e46cd457e05243e08 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tomasz=20Id=C5=BAkowski?= Date: Fri, 17 Jul 2026 14:39:49 +0200 Subject: [PATCH 07/30] Finalize Store behavior: Snapshot -> Conversation; add Conversation struct; add default save/2 implementations --- lib/legion/store.ex | 106 ++++++++++++------ lib/legion/store/conversation.ex | 25 +++++ .../metadata.ex} | 2 +- .../state.ex} | 2 +- test/legion/store_test.exs | 55 +++++++++ 5 files changed, 155 insertions(+), 35 deletions(-) create mode 100644 lib/legion/store/conversation.ex rename lib/legion/store/{conversation_metadata.ex => conversation/metadata.ex} (92%) rename lib/legion/store/{conversation_state.ex => conversation/state.ex} (91%) create mode 100644 test/legion/store_test.exs diff --git a/lib/legion/store.ex b/lib/legion/store.ex index 35f2e84..617f31c 100644 --- a/lib/legion/store.ex +++ b/lib/legion/store.ex @@ -8,12 +8,12 @@ defmodule Legion.Store do {:ok, pid} = Legion.start_link(AssistantAgent, store: MyApp.AgentStore, agent_id: "user_42") On start, the agent calls `c:get/1` and resumes from the returned - `Legion.Store.ConversationState` if one exists. The system prompt is + `Legion.Store.Conversation.State` if one exists. The system prompt is regenerated fresh on every start, so prompt or tool changes apply to restored conversations. After every completed turn, the agent calls `c:save/2` with a - `Legion.Store.ConversationState` **before** replying to the caller. A reply + `Legion.Store.Conversation.State` **before** replying to the caller. A reply is a commit receipt: any turn a caller observed survives a crash, restart, or deploy. A crash mid-turn rolls back to the last completed turn. @@ -37,9 +37,9 @@ defmodule Legion.Store do ## Identifying a conversation - `:agent_id` is the key a snapshot is saved under - it names one conversation, - not one user. A chat app with many chats per user keys by the chat; - compose the id however you like, since Legion treats it as opaque: + `:agent_id` is the key a conversation is saved under - it names one + conversation, not one user. A chat app with many chats per user keys by the + chat; compose the id however you like, since Legion treats it as opaque: Legion.start_link(ChatAgent, agent_id: "user_42:chat_7") @@ -53,10 +53,15 @@ defmodule Legion.Store do Stores must implement `c:get/1` and `c:save/2`. - The required save payload is `Legion.Store.ConversationState`, which holds - the conversation `:messages` (without the system prompt) and the `:bindings` - from evaluated code (relevant with `binding_scope: :conversation`). Each - message carries a `:type` (`:user`, `:assistant`, `:eval_result`, or + Stores must accept two required save payloads: + + - `Legion.Store.Conversation.State`, which holds the conversation + `:messages` (without the system prompt) and the `:bindings` from + evaluated code (relevant with `binding_scope: :conversation`) + - `{:status, :running | :idle}`, which records whether the agent is + mid-turn + + Each message carries a `:type` (`:user`, `:assistant`, `:eval_result`, or `:error`) and an `:at` timestamp in milliseconds, so consumers can classify and order messages without parsing content. Bindings are arbitrary Elixir terms - values like pids, references, or functions will not survive @@ -65,49 +70,84 @@ defmodule Legion.Store do ## Optional persistence - Stores may also accept `Legion.Store.ConversationMetadata` through + Stores may also accept `Legion.Store.Conversation.Metadata` through `c:save/2` to record which agent module ran, under which parent conversation, and when. - Stores may also accept `{:status, :running | :idle}` through `c:save/2` to - record whether the agent is mid-turn. + `use Legion.Store` provides default no-op `save/2` clauses for state, + metadata, and status payloads. They log a warning and return `:ok`, so a + store can opt into only the payloads it persists without breaking agent + execution. Override `save/2` for durable persistence. - Because metadata and status persistence are optional, snapshots returned - from `c:get/1` or `c:list/1` may have `metadata: nil` or `status: nil`. - Because a store may record optional data before any conversation state is - saved, `state` may also be nil. + Because metadata persistence is optional, conversations returned from + `c:get/1` or `c:list/1` may have `metadata: nil`. `status` may also be nil + for persisted conversations created before status was recorded. Because a + store may record metadata or status before any conversation state is saved, + `state` may also be nil. ## Reading conversations - `c:get/1` returns the persisted snapshot for one `agent_id`, or `:error` + `c:get/1` returns the persisted conversation for one `agent_id`, or `:error` when the store has no row for that id. - The optional `c:list/1` callback returns persisted snapshots newest first + The optional `c:list/1` callback returns persisted conversations newest first for consumers that rebuild a view of past conversations from the store alone. """ + alias Legion.Store.Conversation + alias Legion.Store.Conversation.{Metadata, State} + @type agent_id :: term() - @type status :: :idle | :running - @type snapshot :: { - agent_id :: agent_id(), - metadata :: Legion.Store.ConversationMetadata.t() | nil, - status :: status() | nil, - state :: Legion.Store.ConversationState.t() | nil - } + @type status :: Conversation.status() + @type conversation :: Conversation.t() @type payload :: - Legion.Store.ConversationState.t() - | Legion.Store.ConversationMetadata.t() + State.t() + | Metadata.t() | {:status, status()} - @doc "Returns the persisted snapshot for `agent_id`, or `:error` if none exists." - @callback get(agent_id()) :: {:ok, snapshot()} | :error + @doc "Returns the persisted conversation for `agent_id`, or `:error` if none exists." + @callback get(agent_id()) :: {:ok, conversation()} | :error - @doc "Returns the newest `limit` persisted snapshots, newest first." - @callback list(limit :: pos_integer()) :: [snapshot()] + @doc "Returns the newest `limit` persisted conversations, newest first." + @callback list(limit :: pos_integer()) :: [conversation()] - @doc "Saves a conversation state, conversation metadata, or status payload for `agent_id`." - @callback save(agent_id(), payload()) :: :ok + @doc "Saves a conversation state, status, or optional conversation metadata for `agent_id`." + @callback save(agent_id(), payload()) :: :ok | :error @optional_callbacks list: 1 + + defmacro __using__(_opts) do + quote do + @behaviour Legion.Store + + require Logger + + def save(_agent_id, %Conversation.State{}) do + Logger.warning( + "Store #{inspect(__MODULE__)} does not persist conversation state; override save/2 to persist this payload" + ) + + :ok + end + + def save(_agent_id, %Conversation.Metadata{}) do + Logger.warning( + "Store #{inspect(__MODULE__)} does not persist conversation metadata; override save/2 to persist this payload" + ) + + :ok + end + + def save(_agent_id, {:status, status}) when status in [:running, :idle] do + Logger.warning( + "Store #{inspect(__MODULE__)} does not persist conversation status; override save/2 to persist this payload" + ) + + :ok + end + + defoverridable save: 2 + end + end end diff --git a/lib/legion/store/conversation.ex b/lib/legion/store/conversation.ex new file mode 100644 index 0000000..a0a77db --- /dev/null +++ b/lib/legion/store/conversation.ex @@ -0,0 +1,25 @@ +defmodule Legion.Store.Conversation do + @moduledoc """ + One persisted conversation record in a `Legion.Store`. + + A conversation combines the opaque store key with the persisted data a store + may know about that conversation: identity metadata, current run status, and + replayable state. + """ + + alias Legion.Store.Conversation.{Metadata, State} + + @enforce_keys [:agent_id] + defstruct [:agent_id, :metadata, :status, :state] + + @typedoc "Whether the persisted conversation is idle or mid-turn." + @type status :: :idle | :running + + @typedoc "One persisted conversation record." + @type t :: %__MODULE__{ + agent_id: Legion.Store.agent_id(), + metadata: Metadata.t() | nil, + status: status() | nil, + state: State.t() | nil + } +end diff --git a/lib/legion/store/conversation_metadata.ex b/lib/legion/store/conversation/metadata.ex similarity index 92% rename from lib/legion/store/conversation_metadata.ex rename to lib/legion/store/conversation/metadata.ex index 75d3ae8..d155eb6 100644 --- a/lib/legion/store/conversation_metadata.ex +++ b/lib/legion/store/conversation/metadata.ex @@ -1,4 +1,4 @@ -defmodule Legion.Store.ConversationMetadata do +defmodule Legion.Store.Conversation.Metadata do @moduledoc """ Identity metadata for a persisted agent conversation. diff --git a/lib/legion/store/conversation_state.ex b/lib/legion/store/conversation/state.ex similarity index 91% rename from lib/legion/store/conversation_state.ex rename to lib/legion/store/conversation/state.ex index 39e3da9..47bcbf3 100644 --- a/lib/legion/store/conversation_state.ex +++ b/lib/legion/store/conversation/state.ex @@ -1,4 +1,4 @@ -defmodule Legion.Store.ConversationState do +defmodule Legion.Store.Conversation.State do @moduledoc """ Replayable state for a persisted agent conversation. diff --git a/test/legion/store_test.exs b/test/legion/store_test.exs new file mode 100644 index 0000000..9aba1f1 --- /dev/null +++ b/test/legion/store_test.exs @@ -0,0 +1,55 @@ +defmodule Legion.StoreTest do + use ExUnit.Case, async: true + + import ExUnit.CaptureLog + + alias Legion.Store.Conversation.{Metadata, State} + + defmodule DefaultStore do + use Legion.Store + + def get(_agent_id), do: :error + end + + defmodule CustomStore do + use Legion.Store + + def get(_agent_id), do: :error + def save(_agent_id, _payload), do: :custom + end + + test "default save/2 warns and returns :ok for conversation state" do + payload = %State{messages: [], bindings: []} + + log = + capture_log(fn -> + assert :ok = DefaultStore.save("agent-1", payload) + end) + + assert log =~ "Store #{inspect(DefaultStore)} does not persist conversation state" + end + + test "default save/2 warns and returns :ok for conversation metadata" do + payload = %Metadata{agent_module: SomeAgent, parent_agent_id: nil, started_at: 1} + + log = + capture_log(fn -> + assert :ok = DefaultStore.save("agent-1", payload) + end) + + assert log =~ "Store #{inspect(DefaultStore)} does not persist conversation metadata" + end + + test "default save/2 warns and returns :ok for conversation status" do + log = + capture_log(fn -> + assert :ok = DefaultStore.save("agent-1", {:status, :running}) + end) + + assert log =~ "Store #{inspect(DefaultStore)} does not persist conversation status" + end + + test "default save/2 can be overridden" do + assert :custom = CustomStore.save("agent-1", %State{messages: [], bindings: []}) + end +end From be36ac89be3dfe02972d333f1050005107fc05e6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tomasz=20Id=C5=BAkowski?= Date: Fri, 17 Jul 2026 15:53:53 +0200 Subject: [PATCH 08/30] Adjust Postgres adapter to follow new Store contract --- lib/legion/store/postgres.ex | 126 ++++++++++++++++--------- lib/legion/store/postgres/migration.ex | 2 +- mix.exs | 2 + 3 files changed, 82 insertions(+), 48 deletions(-) diff --git a/lib/legion/store/postgres.ex b/lib/legion/store/postgres.ex index 138ca85..2b0babb 100644 --- a/lib/legion/store/postgres.ex +++ b/lib/legion/store/postgres.ex @@ -62,86 +62,81 @@ defmodule Legion.Store.Postgres do repo = Keyword.fetch!(opts, :repo) table = Keyword.get(opts, :table, "legion_agents") - select_sql = "SELECT snapshot FROM #{table} WHERE agent_id = $1" + table_columns = + "agent_id, agent_module, parent_agent_id, status, started_at, conversation_state" - upsert_sql = """ - INSERT INTO #{table} (agent_id, snapshot, inserted_at, updated_at) + select_sql = "SELECT #{table_columns} FROM #{table} WHERE agent_id = $1" + + list_sql = + "SELECT #{table_columns} FROM #{table} ORDER BY updated_at DESC NULLS LAST LIMIT $1" + + save_state_sql = """ + INSERT INTO #{table} (agent_id, conversation_state, inserted_at, updated_at) VALUES ($1, $2, now(), now()) - ON CONFLICT (agent_id) DO UPDATE SET snapshot = EXCLUDED.snapshot, updated_at = now() + ON CONFLICT (agent_id) DO UPDATE SET conversation_state = EXCLUDED.conversation_state, updated_at = now() """ - save_run_sql = """ - INSERT INTO #{table} (agent_id, agent_module, parent_agent_id, status, started_at, inserted_at, updated_at) - VALUES ($1, $2, $3, 'idle', $4, now(), now()) + save_metadata_sql = """ + INSERT INTO #{table} (agent_id, agent_module, parent_agent_id, started_at, inserted_at, updated_at) + VALUES ($1, $2, $3, $4, now(), now()) ON CONFLICT (agent_id) DO UPDATE SET agent_module = EXCLUDED.agent_module, parent_agent_id = COALESCE(#{table}.parent_agent_id, EXCLUDED.parent_agent_id), - status = 'idle', started_at = EXCLUDED.started_at, updated_at = now() """ - save_status_sql = "UPDATE #{table} SET status = $2, updated_at = now() WHERE agent_id = $1" - - run_columns = "agent_id, agent_module, parent_agent_id, status, started_at" - - list_runs_sql = - "SELECT #{run_columns} FROM #{table} ORDER BY started_at DESC NULLS LAST LIMIT $1" - - get_run_sql = "SELECT #{run_columns} FROM #{table} WHERE agent_id = $1" + save_status_sql = """ + INSERT INTO #{table} (agent_id, status, inserted_at, updated_at) + VALUES ($1, $2, now(), now()) + ON CONFLICT (agent_id) DO UPDATE + SET status = EXCLUDED.status, updated_at = now() + """ quote do - @behaviour Legion.Store + use Legion.Store alias Legion.Store.Postgres @impl Legion.Store - # sobelow_skip ["Misc.BinToTerm"] - def load(agent_id) when is_binary(agent_id) do + def get(agent_id) when is_binary(agent_id) do case unquote(repo).query!(unquote(select_sql), [agent_id]) do - %{rows: [[nil]]} -> :error - %{rows: [[snapshot]]} -> {:ok, :erlang.binary_to_term(snapshot)} %{rows: []} -> :error + %{rows: [row]} -> {:ok, Postgres.decode_conversation(row)} end end @impl Legion.Store - def save(agent_id, snapshot) when is_binary(agent_id) do - unquote(repo).query!(unquote(upsert_sql), [agent_id, :erlang.term_to_binary(snapshot)]) + def list(limit) when is_integer(limit) and limit > 0 do + %{rows: rows} = unquote(repo).query!(unquote(list_sql), [limit]) + Enum.map(rows, &Postgres.decode_conversation/1) + end + + @impl Legion.Store + def save(agent_id, %Legion.Store.Conversation.State{} = state) when is_binary(agent_id) do + unquote(repo).query!(unquote(save_state_sql), [agent_id, :erlang.term_to_binary(state)]) :ok end @impl Legion.Store - def save_run(agent_id, metadata) when is_binary(agent_id) do - unquote(repo).query!(unquote(save_run_sql), [ - agent_id, - inspect(metadata.agent_module), - metadata.parent_agent_id, - metadata.started_at - ]) + def save(agent_id, %Legion.Store.Conversation.Metadata{} = metadata) + when is_binary(agent_id) do + unquote(repo).query!( + unquote(save_metadata_sql), + [agent_id] ++ Postgres.encode_metadata(metadata) + ) :ok end @impl Legion.Store - def save_status(agent_id, status) when is_binary(agent_id) do + def save(agent_id, {:status, status}) + when is_binary(agent_id) and status in [:running, :idle] do unquote(repo).query!(unquote(save_status_sql), [agent_id, Atom.to_string(status)]) :ok end - @impl Legion.Store - def list_runs(limit) do - %{rows: rows} = unquote(repo).query!(unquote(list_runs_sql), [limit]) - Enum.map(rows, &Postgres.decode_run_row/1) - end - - @impl Legion.Store - def get_run(agent_id) when is_binary(agent_id) do - case unquote(repo).query!(unquote(get_run_sql), [agent_id]) do - %{rows: [row]} -> Postgres.decode_run_row(row) - %{rows: []} -> nil - end - end + def save(_agent_id, _payload), do: :error @doc false def __repo__, do: unquote(repo) @@ -152,16 +147,53 @@ defmodule Legion.Store.Postgres do end @doc false - def decode_run_row([agent_id, agent_module, parent_agent_id, status, started_at]) do - %{ + def decode_conversation([ + agent_id, + agent_module, + parent_agent_id, + status, + started_at, + conversation_state | _ + ]) do + %Legion.Store.Conversation{ agent_id: agent_id, + metadata: decode_metadata(agent_module, parent_agent_id, started_at), + status: decode_status(status), + state: decode_state(conversation_state) + } + end + + @doc false + def encode_metadata(%Legion.Store.Conversation.Metadata{} = metadata) do + [ + metadata.agent_module && inspect(metadata.agent_module), + metadata.parent_agent_id, + metadata.started_at + ] + end + + defp decode_metadata(nil, nil, nil), do: nil + + defp decode_metadata(agent_module, parent_agent_id, started_at) do + %Legion.Store.Conversation.Metadata{ agent_module: agent_module && Module.concat([agent_module]), parent_agent_id: parent_agent_id, - status: decode_status(status), started_at: started_at } end + defp decode_state(nil), do: nil + + # sobelow_skip ["Misc.BinToTerm"] + defp decode_state(binary) when is_binary(binary) do + state = :erlang.binary_to_term(binary) + + %Legion.Store.Conversation.State{ + messages: Map.get(state, :messages, []), + bindings: Map.get(state, :bindings, []) + } + end + defp decode_status("running"), do: :running defp decode_status("idle"), do: :idle defp decode_status(nil), do: nil diff --git a/lib/legion/store/postgres/migration.ex b/lib/legion/store/postgres/migration.ex index 0ad30f3..51e7947 100644 --- a/lib/legion/store/postgres/migration.ex +++ b/lib/legion/store/postgres/migration.ex @@ -109,7 +109,7 @@ defmodule Legion.Store.Postgres.Migration do parent_agent_id text, status text, started_at bigint, - snapshot bytea, + conversation_state bytea, inserted_at timestamptz NOT NULL DEFAULT now(), updated_at timestamptz NOT NULL DEFAULT now() ) diff --git a/mix.exs b/mix.exs index a17ac37..749f30c 100644 --- a/mix.exs +++ b/mix.exs @@ -54,6 +54,8 @@ defmodule Legion.MixProject do Legion.Agent, Legion.Tool, Legion.Store, + Legion.Store.Conversation, + ~r/^Legion\.Store\.Conversation\./, Legion.Store.Postgres, Legion.Store.Postgres.Migration ], From d109205fc859e5ffadbf18ce724952bd35c164be Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tomasz=20Id=C5=BAkowski?= Date: Fri, 17 Jul 2026 17:47:36 +0200 Subject: [PATCH 09/30] Wire new store contract into AgentServer --- lib/legion/agent_server.ex | 78 +++---- lib/legion/executor.ex | 2 +- test/legion/agent_server_test.exs | 62 +++++- test/legion/store/postgres_db_test.exs | 86 +++++--- test/legion/store/postgres_test.exs | 294 ++++++++++++++++++------- 5 files changed, 363 insertions(+), 159 deletions(-) diff --git a/lib/legion/agent_server.ex b/lib/legion/agent_server.ex index d2ddf1c..7e26c5d 100644 --- a/lib/legion/agent_server.ex +++ b/lib/legion/agent_server.ex @@ -11,6 +11,8 @@ defmodule Legion.AgentServer do require Logger alias Legion.{Executor, Telemetry} + alias Legion.Store.Conversation + alias Legion.Store.Conversation.{Metadata, State} alias ReqLLM.Message.ContentPart defstruct [:agent_module, :messages, :config, :store, :agent_id, bindings: []] @@ -76,16 +78,13 @@ defmodule Legion.AgentServer do %{agent: agent_module} ) - save_run(store, agent_id, %{ - agent_module: agent_module, - parent_agent_id: parent_agent_id, - started_at: System.system_time(:millisecond) - }) - {saved_messages, saved_bindings} = - case store && store.load(agent_id) do - {:ok, %{messages: messages, bindings: bindings}} -> {messages, bindings} - _no_snapshot -> {[], []} + case store && store.get(agent_id) do + {:ok, %Conversation{state: %State{messages: messages, bindings: bindings}}} -> + {messages, bindings} + + _no_state -> + {[], []} end state = %__MODULE__{ @@ -97,7 +96,7 @@ defmodule Legion.AgentServer do bindings: saved_bindings } - {:ok, state} + {:ok, persist(state, {:metadata, parent_agent_id})} end @impl true @@ -151,8 +150,11 @@ defmodule Legion.AgentServer do # 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 # response. - state = persist(%{state | messages: state.messages ++ [Executor.message(:user, content)]}) - save_status(state, :running) + state = + state + |> Map.update!(:messages, &(&1 ++ [Executor.message(:user, content)])) + |> persist(:state) + |> persist({:status, :running}) {status, value, final_messages, final_bindings} = Telemetry.span( @@ -173,44 +175,42 @@ defmodule Legion.AgentServer do ) kept_bindings = if conversation_scope?, do: final_bindings, else: [] - state = persist(%{state | messages: final_messages, bindings: kept_bindings}) - save_status(state, :idle) + + state = + %{state | messages: final_messages, bindings: kept_bindings} + |> persist(:state) + |> persist({:status, :idle}) + {{status, value}, state} end - defp persist(%{store: nil} = state), do: state + defp persist(%{store: nil} = state, _payload), do: state - defp persist(state) do + defp persist(state, :state) do [%{role: "system"} | messages] = state.messages - :ok = state.store.save(state.agent_id, %{messages: messages, bindings: state.bindings}) + :ok = state.store.save(state.agent_id, %State{messages: messages, bindings: state.bindings}) state end - # Status is written outside the turn's save/persist path: a crash between - # the :running write and the :idle write leaves 'running' in the store, - # which consumers read as "crashed mid-turn" under a dead pid. - defp save_status(%{store: nil}, _status), do: :ok - - defp save_status(state, status) do - if Code.ensure_loaded?(state.store) and function_exported?(state.store, :save_status, 2) do - state.store.save_status(state.agent_id, status) - end + defp persist(state, {:metadata, parent_agent_id}) do + :ok = + state.store.save( + state.agent_id, + %Metadata{ + agent_module: state.agent_module, + parent_agent_id: parent_agent_id, + started_at: System.system_time(:millisecond) + } + ) - :ok + state end - defp save_run(nil, _agent_id, _metadata), do: :ok - - defp save_run(store, agent_id, metadata) do - if Code.ensure_loaded?(store) and function_exported?(store, :save_run, 2) do - store.save_run(agent_id, metadata) - else - Logger.warning( - "Store #{inspect(store)} does not implement save_run/2; run metadata not persisted" - ) - end - - :ok + # A crash between the :running write and :idle write leaves `:running` in + # the store, which consumers read as "crashed mid-turn" under a dead pid. + defp persist(state, {:status, status}) when status in [:running, :idle] do + :ok = state.store.save(state.agent_id, {:status, status}) + state end defp generate_id, do: Base.url_encode64(:crypto.strong_rand_bytes(16), padding: false) diff --git a/lib/legion/executor.ex b/lib/legion/executor.ex index c041c40..52efab2 100644 --- a/lib/legion/executor.ex +++ b/lib/legion/executor.ex @@ -33,7 +33,7 @@ defmodule Legion.Executor do @doc """ Builds a conversation message stamped with its `:type` and creation time - (`:at`, milliseconds). The extra keys ride along into persisted snapshots so + (`:at`, milliseconds). The extra keys ride along into persisted state so consumers (e.g. LegionWeb) can classify messages without parsing content; ReqLLM ignores them. """ diff --git a/test/legion/agent_server_test.exs b/test/legion/agent_server_test.exs index 918fae5..94dd7d5 100644 --- a/test/legion/agent_server_test.exs +++ b/test/legion/agent_server_test.exs @@ -4,6 +4,8 @@ defmodule Legion.AgentServerTest do import ExUnit.CaptureLog + alias Legion.Store.Conversation + alias Legion.Store.Conversation.{Metadata, State} alias Legion.Test.Support.MathAgent alias ReqLLM.Message.ContentPart @@ -418,20 +420,46 @@ defmodule Legion.AgentServerTest do def start_link, do: Agent.start_link(fn -> %{} end, name: __MODULE__) - def load(agent_id), do: Agent.get(__MODULE__, &Map.fetch(&1, agent_id)) + @impl Legion.Store + def get(agent_id) do + Agent.get(__MODULE__, fn state -> + case Map.fetch(state, agent_id) do + {:ok, conversation} -> {:ok, conversation} + :error -> :error + end + end) + end - def save(agent_id, snapshot) do - Agent.update(__MODULE__, &Map.put(&1, agent_id, snapshot)) + def load(agent_id) do + case get(agent_id) do + {:ok, %Conversation{state: state}} when not is_nil(state) -> {:ok, state} + _ -> :error + end + end + + @impl Legion.Store + def save(agent_id, %State{} = snapshot) do + update_conversation(agent_id, &%{&1 | state: snapshot}) end - def save_run(agent_id, metadata) do - Agent.update(__MODULE__, &Map.put(&1, {:run, agent_id}, metadata)) + def save(agent_id, %Metadata{} = metadata) do + update_conversation(agent_id, &%{&1 | metadata: metadata}) end - def save_status(agent_id, status) do + def save(agent_id, {:status, status}) do Agent.update(__MODULE__, fn state -> - Map.update(state, {:statuses, agent_id}, [status], &[status | &1]) + state + |> Map.update({:statuses, agent_id}, [status], &[status | &1]) + |> Map.update( + agent_id, + %Conversation{agent_id: agent_id, status: status}, + fn conversation -> + %{conversation | status: status} + end + ) end) + + :ok end def statuses(agent_id) do @@ -440,16 +468,28 @@ defmodule Legion.AgentServerTest do end def get_run(agent_id) do - case Agent.get(__MODULE__, &Map.get(&1, {:run, agent_id})) do - nil -> nil - metadata -> Map.put(metadata, :agent_id, agent_id) + case get(agent_id) do + {:ok, %Conversation{metadata: nil}} -> nil + {:ok, %Conversation{metadata: metadata}} -> Map.put(metadata, :agent_id, agent_id) + :error -> nil end end def runs do Agent.get(__MODULE__, fn state -> - for {{:run, agent_id}, metadata} <- state, do: Map.put(metadata, :agent_id, agent_id) + for {agent_id, %Conversation{metadata: metadata}} <- state, + is_binary(agent_id), + not is_nil(metadata), + do: Map.put(metadata, :agent_id, agent_id) + end) + end + + defp update_conversation(agent_id, update) do + Agent.update(__MODULE__, fn state -> + Map.update(state, agent_id, update.(%Conversation{agent_id: agent_id}), update) end) + + :ok end end diff --git a/test/legion/store/postgres_db_test.exs b/test/legion/store/postgres_db_test.exs index f7e90bb..8ed709c 100644 --- a/test/legion/store/postgres_db_test.exs +++ b/test/legion/store/postgres_db_test.exs @@ -6,6 +6,9 @@ defmodule Legion.Store.PostgresDbTest do """ use ExUnit.Case, async: false + alias Legion.Store.Conversation + alias Legion.Store.Conversation.{Metadata, State} + defmodule Repo do def query!(sql, params), do: Postgrex.query!(:legion_store_test, sql, params) end @@ -19,35 +22,35 @@ defmodule Legion.Store.PostgresDbTest do :ok end - test "round-trips a snapshot through a real bytea column" do - snapshot = %{messages: [%{role: "user", content: "hi"}], bindings: [x: 42]} + test "save/2 with State round-trips through a real bytea column" do + state = %State{messages: [%{type: :user, content: "hi"}], bindings: [x: 42]} - assert :ok = Store.save("user_42", snapshot) - assert {:ok, ^snapshot} = Store.load("user_42") + assert :ok = Store.save("user_42", state) + assert {:ok, %Conversation{agent_id: "user_42", state: ^state}} = Store.get("user_42") end - test "load/1 returns :error when the row is absent" do - assert :error = Store.load("missing") + test "get/1 returns :error when the row is absent" do + assert :error = Store.get("missing") end - test "save/2 upserts on conflict - the latest snapshot wins" do - assert :ok = Store.save("user_42", %{messages: [], bindings: [v: 1]}) - assert :ok = Store.save("user_42", %{messages: [], bindings: [v: 2]}) + test "save/2 with State upserts on conflict - the latest state wins" do + assert :ok = Store.save("user_42", %State{messages: [], bindings: [v: 1]}) + assert :ok = Store.save("user_42", %State{messages: [], bindings: [v: 2]}) - assert {:ok, %{bindings: [v: 2]}} = Store.load("user_42") + assert {:ok, %Conversation{state: %State{bindings: [v: 2]}}} = Store.get("user_42") end - test "save_run/2 records conversation identity and keeps the parent on conflict" do - metadata = %{ + test "save/2 with Metadata records conversation identity and keeps the parent on conflict" do + metadata = %Metadata{ agent_module: MyApp.Worker, parent_agent_id: "parent-1", started_at: 100 } - assert :ok = Store.save_run("child-1", metadata) + assert :ok = Store.save("child-1", metadata) # Resumed from elsewhere: no parent this time, later started_at. - assert :ok = Store.save_run("child-1", %{metadata | parent_agent_id: nil, started_at: 200}) + assert :ok = Store.save("child-1", %{metadata | parent_agent_id: nil, started_at: 200}) %{rows: [[agent_module, parent_agent_id, started_at]]} = Postgrex.query!( @@ -61,15 +64,36 @@ defmodule Legion.Store.PostgresDbTest do assert started_at == 200 end - test "save_status/2 flips the stored status and save_run/2 resets it to idle" do - :ok = Store.save_run("s1", %{agent_module: A, parent_agent_id: nil, started_at: 1}) - assert %{status: :idle} = Store.get_run("s1") + test "save/2 with status payload flips the stored status" do + :ok = Store.save("s1", {:status, :running}) + assert {:ok, %Conversation{status: :running, metadata: nil, state: nil}} = Store.get("s1") + + :ok = Store.save("s1", {:status, :idle}) + assert {:ok, %Conversation{status: :idle}} = Store.get("s1") + end - :ok = Store.save_status("s1", :running) - assert %{status: :running} = Store.get_run("s1") + test "save/2 with Metadata does not reset an existing status" do + metadata = %Metadata{agent_module: A, parent_agent_id: nil, started_at: 2} - :ok = Store.save_run("s1", %{agent_module: A, parent_agent_id: nil, started_at: 2}) - assert %{status: :idle} = Store.get_run("s1") + :ok = Store.save("s1", {:status, :running}) + :ok = Store.save("s1", metadata) + + assert {:ok, %Conversation{status: :running, metadata: ^metadata}} = Store.get("s1") + end + + test "list/1 returns newest conversations including partial rows" do + state = %State{messages: [], bindings: []} + metadata = %Metadata{agent_module: A, parent_agent_id: nil, started_at: 1} + + :ok = Store.save("state-only", state) + :ok = Store.save("metadata-only", metadata) + :ok = Store.save("status-only", {:status, :idle}) + + assert [ + %Conversation{agent_id: "status-only", status: :idle, metadata: nil, state: nil}, + %Conversation{agent_id: "metadata-only", metadata: ^metadata, state: nil}, + %Conversation{agent_id: "state-only", metadata: nil, state: ^state} + ] = Store.list(10) end test "the trigger notifies the table's channel with the agent_id on every write" do @@ -84,26 +108,28 @@ defmodule Legion.Store.PostgresDbTest do {:ok, _ref} = Postgrex.Notifications.listen(notifications, "legion_agents") - :ok = Store.save_run("notify-1", %{agent_module: A, parent_agent_id: nil, started_at: 1}) + :ok = Store.save("notify-1", %Metadata{agent_module: A, parent_agent_id: nil, started_at: 1}) assert_receive {:notification, _pid, _ref, "legion_agents", "notify-1"}, 1_000 - :ok = Store.save_status("notify-1", :running) + :ok = Store.save("notify-1", {:status, :running}) assert_receive {:notification, _pid, _ref, "legion_agents", "notify-1"}, 1_000 - :ok = Store.save("notify-1", %{messages: [], bindings: []}) + :ok = Store.save("notify-1", %State{messages: [], bindings: []}) assert_receive {:notification, _pid, _ref, "legion_agents", "notify-1"}, 1_000 GenServer.stop(notifications) end - test "save_run/2 then save/2 share one row and load/1 sees the snapshot" do - assert :ok = Store.save_run("both", %{agent_module: A, parent_agent_id: nil, started_at: 1}) - assert :error = Store.load("both") + test "metadata then state share one row and get/1 sees both" do + metadata = %Metadata{agent_module: A, parent_agent_id: nil, started_at: 1} + + assert :ok = Store.save("both", metadata) + assert {:ok, %Conversation{metadata: ^metadata, state: nil}} = Store.get("both") - snapshot = %{messages: [%{role: "user", content: "hi"}], bindings: []} - assert :ok = Store.save("both", snapshot) + state = %State{messages: [%{type: :user, content: "hi"}], bindings: []} + assert :ok = Store.save("both", state) - assert {:ok, ^snapshot} = Store.load("both") + assert {:ok, %Conversation{metadata: ^metadata, state: ^state}} = Store.get("both") %{rows: [[count]]} = Postgrex.query!( diff --git a/test/legion/store/postgres_test.exs b/test/legion/store/postgres_test.exs index 8e14c61..f00b8e9 100644 --- a/test/legion/store/postgres_test.exs +++ b/test/legion/store/postgres_test.exs @@ -1,61 +1,125 @@ defmodule Legion.Store.PostgresTest do use ExUnit.Case, async: true + alias Legion.Store.Conversation + alias Legion.Store.Conversation.{Metadata, State} + defmodule FakeRepo do @moduledoc "Emulates repo.query!/2 for the statements the store issues." - def start_link, do: Agent.start_link(fn -> %{snapshots: %{}, runs: %{}} end, name: __MODULE__) + def start_link, do: Agent.start_link(fn -> %{rows: %{}, clock: 0} end, name: __MODULE__) - def query!("SELECT snapshot FROM " <> _rest, [agent_id]) do - Agent.get(__MODULE__, fn state -> - case {state.snapshots, state.runs} do - {%{^agent_id => snapshot}, _runs} -> %{rows: [[snapshot]]} - {_snapshots, %{^agent_id => _run}} -> %{rows: [[nil]]} - _neither -> %{rows: []} - end - end) - end + def query!("SELECT agent_id" <> _rest = sql, [param]) do + rows = + Agent.get(__MODULE__, fn state -> + state.rows + |> select_rows(sql, param) + |> Enum.map(&row/1) + end) - def query!("INSERT INTO " <> _rest, [agent_id, snapshot]) do - Agent.update(__MODULE__, &put_in(&1.snapshots[agent_id], snapshot)) - %{num_rows: 1} + %{rows: rows} end - def query!("INSERT INTO " <> _rest, [agent_id, agent_module, parent_agent_id, started_at]) do - Agent.update( - __MODULE__, - &put_in(&1.runs[agent_id], %{ - agent_module: agent_module, - parent_agent_id: parent_agent_id, - status: "idle", - started_at: started_at - }) - ) - - %{num_rows: 1} - end + def query!("INSERT INTO " <> _rest = sql, [agent_id, value]) do + cond do + String.contains?(sql, "(agent_id, conversation_state") -> + update_row(agent_id, %{conversation_state: value}) + %{num_rows: 1} + + String.contains?(sql, "(agent_id, status") -> + update_row(agent_id, %{status: value}) + %{num_rows: 1} - def query!("UPDATE " <> _rest, [agent_id, status]) do - Agent.update(__MODULE__, &put_in(&1.runs[agent_id].status, status)) - %{num_rows: 1} + true -> + raise ArgumentError, "unexpected two-argument insert: #{sql}" + end end - def query!("SELECT agent_id" <> _rest = sql, [param]) do - rows = - Agent.get(__MODULE__, fn state -> - for {agent_id, run} <- state.runs do - [agent_id, run.agent_module, run.parent_agent_id, run.status, run.started_at] - end + def query!("INSERT INTO " <> _rest = sql, [ + agent_id, + agent_module, + parent_agent_id, + started_at + ]) do + if String.contains?(sql, "(agent_id, agent_module, parent_agent_id, started_at") do + Agent.update(__MODULE__, fn state -> + row = Map.get(state.rows, agent_id, empty_row(agent_id)) + + put_in(state.rows[agent_id], %{ + row + | agent_module: agent_module, + parent_agent_id: row.parent_agent_id || parent_agent_id, + started_at: started_at, + updated_at: state.clock + 1 + }) + |> Map.update!(:clock, &(&1 + 1)) end) + %{num_rows: 1} + else + raise ArgumentError, "unexpected metadata insert: #{sql}" + end + end + + def run(agent_id), do: Agent.get(__MODULE__, & &1.rows[agent_id]) + + defp select_rows(rows, sql, param) do if String.contains?(sql, "WHERE agent_id") do - %{rows: Enum.filter(rows, fn [agent_id | _rest] -> agent_id == param end)} + select_row(rows, param) else - %{rows: rows |> Enum.sort_by(&List.last/1, :desc) |> Enum.take(param)} + select_recent_rows(rows, param) end end - def run(agent_id), do: Agent.get(__MODULE__, & &1.runs[agent_id]) + defp select_row(rows, agent_id) do + rows + |> Map.values() + |> Enum.filter(&(&1.agent_id == agent_id)) + end + + defp select_recent_rows(rows, limit) do + rows + |> Map.values() + |> Enum.sort_by(& &1.updated_at, :desc) + |> Enum.take(limit) + end + + defp update_row(agent_id, attrs) do + Agent.update(__MODULE__, fn state -> + row = + state.rows + |> Map.get(agent_id, empty_row(agent_id)) + |> Map.merge(attrs) + |> Map.put(:updated_at, state.clock + 1) + + state + |> put_in([:rows, agent_id], row) + |> Map.update!(:clock, &(&1 + 1)) + end) + end + + defp empty_row(agent_id) do + %{ + agent_id: agent_id, + agent_module: nil, + parent_agent_id: nil, + status: nil, + started_at: nil, + conversation_state: nil, + updated_at: 0 + } + end + + defp row(row) do + [ + row.agent_id, + row.agent_module, + row.parent_agent_id, + row.status, + row.started_at, + row.conversation_state + ] + end end defmodule Store do @@ -67,74 +131,147 @@ defmodule Legion.Store.PostgresTest do :ok end - test "save/2 then load/1 round-trips the snapshot through term_to_binary" do - snapshot = %{messages: [%{role: "user", content: "hi"}], bindings: [x: 42]} + test "save/2 with State then get/1 returns a conversation with decoded state" do + state = %State{messages: [%{type: :user, content: "hi"}], bindings: [x: 42]} - assert :ok = Store.save("user_42", snapshot) - assert {:ok, ^snapshot} = Store.load("user_42") + assert :ok = Store.save("user_42", state) + + assert {:ok, + %Conversation{ + agent_id: "user_42", + metadata: nil, + status: nil, + state: ^state + }} = Store.get("user_42") end - test "load/1 returns :error when no snapshot exists" do - assert :error = Store.load("missing") + test "get/1 returns :error when the row is absent" do + assert :error = Store.get("missing") end test "ids must be strings" do - assert_raise FunctionClauseError, fn -> Store.load(42) end - assert_raise FunctionClauseError, fn -> Store.save(42, %{messages: [], bindings: []}) end + assert_raise FunctionClauseError, fn -> Store.get(42) end + assert :error = Store.save(42, %State{messages: [], bindings: []}) end - test "save_run/2 stores the module in inspect form with parent, and start time" do - metadata = %{ + test "save/2 with Metadata stores the module in inspect form with parent and start time" do + metadata = %Metadata{ agent_module: Legion.Test.Support.MathAgent, parent_agent_id: "p1", started_at: 123 } - assert :ok = Store.save_run("user_42", metadata) + assert :ok = Store.save("user_42", metadata) - assert FakeRepo.run("user_42") == %{ + assert %{ agent_module: "Legion.Test.Support.MathAgent", parent_agent_id: "p1", - status: "idle", - started_at: 123 - } + status: nil, + started_at: 123, + conversation_state: nil + } = FakeRepo.run("user_42") end - test "save_status/2 flips the run status and get_run/1 decodes it" do - :ok = Store.save_run("s1", %{agent_module: SomeAgent, parent_agent_id: nil, started_at: 1}) - assert %{status: :idle} = Store.get_run("s1") - - :ok = Store.save_status("s1", :running) - assert %{status: :running} = Store.get_run("s1") + test "save/2 with status payload flips the stored status and get/1 decodes it" do + :ok = Store.save("s1", {:status, :running}) + assert {:ok, %Conversation{status: :running, metadata: nil, state: nil}} = Store.get("s1") - :ok = Store.save_status("s1", :idle) - assert %{status: :idle} = Store.get_run("s1") + :ok = Store.save("s1", {:status, :idle}) + assert {:ok, %Conversation{status: :idle}} = Store.get("s1") end - test "list_runs/1 returns decoded runs newest first" do - :ok = Store.save_run("a", %{agent_module: SomeAgent, parent_agent_id: nil, started_at: 1}) - :ok = Store.save_run("b", %{agent_module: OtherAgent, parent_agent_id: "a", started_at: 2}) + test "list/1 returns decoded conversations newest first" do + :ok = + Store.save("a", %Metadata{agent_module: SomeAgent, parent_agent_id: nil, started_at: 1}) + + :ok = + Store.save("b", %Metadata{agent_module: OtherAgent, parent_agent_id: "a", started_at: 2}) assert [ - %{agent_id: "b", agent_module: OtherAgent, parent_agent_id: "a", started_at: 2}, - %{agent_id: "a", agent_module: SomeAgent, parent_agent_id: nil, started_at: 1} - ] = Store.list_runs(10) + %Conversation{ + agent_id: "b", + metadata: %Metadata{ + agent_module: OtherAgent, + parent_agent_id: "a", + started_at: 2 + } + }, + %Conversation{ + agent_id: "a", + metadata: %Metadata{ + agent_module: SomeAgent, + parent_agent_id: nil, + started_at: 1 + } + } + ] = Store.list(10) + + assert [%Conversation{agent_id: "b"}] = Store.list(1) + end + + test "get/1 returns decoded conversation, or :error when missing" do + :ok = Store.save("a", %Metadata{agent_module: SomeAgent, parent_agent_id: nil, started_at: 1}) + + assert {:ok, %Conversation{agent_id: "a", metadata: %Metadata{agent_module: SomeAgent}}} = + Store.get("a") + + assert Store.get("missing") == :error + end + + test "get/1 returns conversation with nil state when only metadata exists" do + metadata = %Metadata{agent_module: SomeAgent, parent_agent_id: nil, started_at: 1} - assert [%{agent_id: "b"}] = Store.list_runs(1) + assert :ok = Store.save("started-only", metadata) + assert {:ok, %Conversation{metadata: ^metadata, state: nil}} = Store.get("started-only") end - test "get_run/1 returns the decoded run, or nil when missing" do - :ok = Store.save_run("a", %{agent_module: SomeAgent, parent_agent_id: nil, started_at: 1}) + test "get/1 returns conversation with nil metadata when only state exists" do + state = %State{messages: [], bindings: [answer: 42]} - assert %{agent_id: "a", agent_module: SomeAgent} = Store.get_run("a") - assert Store.get_run("missing") == nil + assert :ok = Store.save("state-only", state) + assert {:ok, %Conversation{metadata: nil, state: ^state}} = Store.get("state-only") end - test "load/1 returns :error when only run metadata exists (snapshot still null)" do - metadata = %{agent_module: SomeAgent, parent_agent_id: nil, started_at: 1} + test "save/2 with Metadata does not reset an existing status" do + metadata = %Metadata{agent_module: SomeAgent, parent_agent_id: nil, started_at: 1} + + assert :ok = Store.save("status-metadata", {:status, :running}) + assert :ok = Store.save("status-metadata", metadata) + + assert {:ok, %Conversation{status: :running, metadata: ^metadata}} = + Store.get("status-metadata") + end + + test "save/2 with Metadata preserves an existing parent on conflict" do + metadata = %Metadata{agent_module: SomeAgent, parent_agent_id: "parent-1", started_at: 1} + + assert :ok = Store.save("child", metadata) + assert :ok = Store.save("child", %{metadata | parent_agent_id: nil, started_at: 2}) + + assert {:ok, + %Conversation{ + metadata: %Metadata{parent_agent_id: "parent-1", started_at: 2} + }} = Store.get("child") + end + + test "list/1 includes state-only metadata-only and status-only conversations" do + state = %State{messages: [], bindings: []} + metadata = %Metadata{agent_module: SomeAgent, parent_agent_id: nil, started_at: 1} + + :ok = Store.save("state-only", state) + :ok = Store.save("metadata-only", metadata) + :ok = Store.save("status-only", {:status, :idle}) + + assert [ + %Conversation{agent_id: "status-only", status: :idle, metadata: nil, state: nil}, + %Conversation{agent_id: "metadata-only", metadata: ^metadata, state: nil}, + %Conversation{agent_id: "state-only", metadata: nil, state: ^state} + ] = Store.list(10) + end - assert :ok = Store.save_run("started-only", metadata) - assert :error = Store.load("started-only") + test "save/2 returns :error for unsupported payloads" do + assert :error = Store.save("bad-payload", %{messages: [], bindings: []}) + assert :error = Store.save("bad-status", {:status, :paused}) end test "a custom table name is interpolated into the statements" do @@ -148,8 +285,9 @@ defmodule Legion.Store.PostgresTest do table: "my_agents" end - CustomTableStore.load("user_42") + CustomTableStore.get("user_42") - assert_received {:sql, "SELECT snapshot FROM my_agents WHERE agent_id = $1"} + assert_received {:sql, + "SELECT agent_id, agent_module, parent_agent_id, status, started_at, conversation_state FROM my_agents WHERE agent_id = $1"} end end From 9103db6de82cbc40918bc686067ea3f93df1b53e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tomasz=20Id=C5=BAkowski?= Date: Tue, 21 Jul 2026 09:21:03 +0200 Subject: [PATCH 10/30] Update store contract\nStore no longer uses specific column groups that can only be updated together, it now supports updating any combination of columns at once through the Payload wrapper. To achieve this Ecto was added as a dependency, so Postgres could use it for DB operations instead of raw sql queries (main benefit was in the unconstrained save/1 function where we could skip tedious string ops) --- lib/legion/store.ex | 103 ++------ lib/legion/store/conversation.ex | 25 -- lib/legion/store/conversation/metadata.ex | 18 -- lib/legion/store/conversation/state.ex | 17 -- lib/legion/store/payload.ex | 29 +++ lib/legion/store/postgres.ex | 161 ++++++------ mix.exs | 1 + mix.lock | 1 + test/legion/store/postgres_db_test.exs | 193 +++++++------- test/legion/store/postgres_test.exs | 294 +++++----------------- test/legion/store_test.exs | 55 ---- 11 files changed, 272 insertions(+), 625 deletions(-) delete mode 100644 lib/legion/store/conversation.ex delete mode 100644 lib/legion/store/conversation/metadata.ex delete mode 100644 lib/legion/store/conversation/state.ex create mode 100644 lib/legion/store/payload.ex delete mode 100644 test/legion/store_test.exs diff --git a/lib/legion/store.ex b/lib/legion/store.ex index 617f31c..943e564 100644 --- a/lib/legion/store.ex +++ b/lib/legion/store.ex @@ -8,14 +8,14 @@ defmodule Legion.Store do {:ok, pid} = Legion.start_link(AssistantAgent, store: MyApp.AgentStore, agent_id: "user_42") On start, the agent calls `c:get/1` and resumes from the returned - `Legion.Store.Conversation.State` if one exists. The system prompt is - regenerated fresh on every start, so prompt or tool changes apply to - restored conversations. + `Legion.Store.Payload` when it has a `:conversation_state`. The system prompt + is regenerated for every start, so prompt and tool changes apply to restored + conversations. - After every completed turn, the agent calls `c:save/2` with a - `Legion.Store.Conversation.State` **before** replying to the caller. A reply - is a commit receipt: any turn a caller observed survives a crash, restart, - or deploy. A crash mid-turn rolls back to the last completed turn. + After every completed turn, the agent calls `c:save/1` with a payload before + replying to the caller. A reply is a commit receipt: any observed turn + survives a crash, restart, or deploy. A crash mid-turn rolls back to the last + completed turn. For Postgres users there is a ready-made adapter - see `Legion.Store.Postgres`: @@ -51,15 +51,14 @@ defmodule Legion.Store do ## Required persistence - Stores must implement `c:get/1` and `c:save/2`. + Stores must implement `c:get/1` and `c:save/1`. - Stores must accept two required save payloads: - - - `Legion.Store.Conversation.State`, which holds the conversation - `:messages` (without the system prompt) and the `:bindings` from - evaluated code (relevant with `binding_scope: :conversation`) - - `{:status, :running | :idle}`, which records whether the agent is - mid-turn + `c: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 (relevant with + `binding_scope: :conversation`). `: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. Each message carries a `:type` (`:user`, `:assistant`, `:eval_result`, or `:error`) and an `:at` timestamp in milliseconds, so consumers can classify @@ -68,22 +67,11 @@ defmodule Legion.Store do serialization, so keep conversation-scoped variables to plain data if you persist agents. - ## Optional persistence - - Stores may also accept `Legion.Store.Conversation.Metadata` through - `c:save/2` to record which agent module ran, under which parent - conversation, and when. + `use Legion.Store` provides a default no-op `save/1` that logs a warning. + Override it for durable persistence. - `use Legion.Store` provides default no-op `save/2` clauses for state, - metadata, and status payloads. They log a warning and return `:ok`, so a - store can opt into only the payloads it persists without breaking agent - execution. Override `save/2` for durable persistence. - - Because metadata persistence is optional, conversations returned from - `c:get/1` or `c:list/1` may have `metadata: nil`. `status` may also be nil - for persisted conversations created before status was recorded. Because a - store may record metadata or status before any conversation state is saved, - `state` may also be nil. + Payload fields other than `:agent_id` may be nil. A store can therefore + persist identity or status before a conversation state exists. ## Reading conversations @@ -95,59 +83,18 @@ defmodule Legion.Store do alone. """ - alias Legion.Store.Conversation - alias Legion.Store.Conversation.{Metadata, State} + alias Legion.Store.Payload @type agent_id :: term() - @type status :: Conversation.status() - @type conversation :: Conversation.t() - @type payload :: - State.t() - | Metadata.t() - | {:status, status()} + @type status :: Payload.status() + @type payload :: Payload.t() @doc "Returns the persisted conversation for `agent_id`, or `:error` if none exists." - @callback get(agent_id()) :: {:ok, conversation()} | :error + @callback get(agent_id()) :: {:ok, payload()} | :error @doc "Returns the newest `limit` persisted conversations, newest first." - @callback list(limit :: pos_integer()) :: [conversation()] - - @doc "Saves a conversation state, status, or optional conversation metadata for `agent_id`." - @callback save(agent_id(), payload()) :: :ok | :error - - @optional_callbacks list: 1 - - defmacro __using__(_opts) do - quote do - @behaviour Legion.Store - - require Logger - - def save(_agent_id, %Conversation.State{}) do - Logger.warning( - "Store #{inspect(__MODULE__)} does not persist conversation state; override save/2 to persist this payload" - ) - - :ok - end - - def save(_agent_id, %Conversation.Metadata{}) do - Logger.warning( - "Store #{inspect(__MODULE__)} does not persist conversation metadata; override save/2 to persist this payload" - ) - - :ok - end - - def save(_agent_id, {:status, status}) when status in [:running, :idle] do - Logger.warning( - "Store #{inspect(__MODULE__)} does not persist conversation status; override save/2 to persist this payload" - ) - - :ok - end + @callback list(limit :: pos_integer()) :: [payload()] - defoverridable save: 2 - end - end + @doc "Saves a conversation payload." + @callback save(payload()) :: :ok | :error end diff --git a/lib/legion/store/conversation.ex b/lib/legion/store/conversation.ex deleted file mode 100644 index a0a77db..0000000 --- a/lib/legion/store/conversation.ex +++ /dev/null @@ -1,25 +0,0 @@ -defmodule Legion.Store.Conversation do - @moduledoc """ - One persisted conversation record in a `Legion.Store`. - - A conversation combines the opaque store key with the persisted data a store - may know about that conversation: identity metadata, current run status, and - replayable state. - """ - - alias Legion.Store.Conversation.{Metadata, State} - - @enforce_keys [:agent_id] - defstruct [:agent_id, :metadata, :status, :state] - - @typedoc "Whether the persisted conversation is idle or mid-turn." - @type status :: :idle | :running - - @typedoc "One persisted conversation record." - @type t :: %__MODULE__{ - agent_id: Legion.Store.agent_id(), - metadata: Metadata.t() | nil, - status: status() | nil, - state: State.t() | nil - } -end diff --git a/lib/legion/store/conversation/metadata.ex b/lib/legion/store/conversation/metadata.ex deleted file mode 100644 index d155eb6..0000000 --- a/lib/legion/store/conversation/metadata.ex +++ /dev/null @@ -1,18 +0,0 @@ -defmodule Legion.Store.Conversation.Metadata do - @moduledoc """ - Identity metadata for a persisted agent conversation. - - Stores can save this payload to record which agent module owns a - conversation, which parent conversation spawned it, and when it started. - """ - - @enforce_keys [:agent_module, :parent_agent_id, :started_at] - defstruct [:agent_module, :parent_agent_id, :started_at] - - @typedoc "Metadata describing a persisted conversation's agent identity." - @type t() :: %__MODULE__{ - agent_module: module(), - parent_agent_id: Legion.Store.agent_id() | nil, - started_at: integer() - } -end diff --git a/lib/legion/store/conversation/state.ex b/lib/legion/store/conversation/state.ex deleted file mode 100644 index 47bcbf3..0000000 --- a/lib/legion/store/conversation/state.ex +++ /dev/null @@ -1,17 +0,0 @@ -defmodule Legion.Store.Conversation.State do - @moduledoc """ - Replayable state for a persisted agent conversation. - - Stores save this payload after turns so a restarted agent can restore its - conversation messages and conversation-scoped bindings. - """ - - @enforce_keys [:messages, :bindings] - defstruct [:messages, :bindings] - - @typedoc "Messages and bindings needed to restore a conversation." - @type t() :: %__MODULE__{ - messages: [map()], - bindings: keyword() - } -end diff --git a/lib/legion/store/payload.ex b/lib/legion/store/payload.ex new file mode 100644 index 0000000..00d69bf --- /dev/null +++ b/lib/legion/store/payload.ex @@ -0,0 +1,29 @@ +defmodule Legion.Store.Payload do + @moduledoc """ + Data supplied to and returned from a `Legion.Store`. + """ + + @enforce_keys [:agent_id] + defstruct [ + :agent_id, + :agent_module, + :parent_agent_id, + :status, + :started_at, + :conversation_state + ] + + @type status :: :idle | :running + @type state :: %{ + messages: [map()], + bindings: keyword() + } + @type t :: %__MODULE__{ + agent_id: Legion.Store.agent_id(), + agent_module: module() | nil, + parent_agent_id: Legion.Store.agent_id() | nil, + status: status() | nil, + started_at: integer() | nil, + conversation_state: state() | nil + } +end diff --git a/lib/legion/store/postgres.ex b/lib/legion/store/postgres.ex index 2b0babb..60d76c8 100644 --- a/lib/legion/store/postgres.ex +++ b/lib/legion/store/postgres.ex @@ -58,86 +58,77 @@ defmodule Legion.Store.Postgres do consumers can follow store changes live without polling. """ + alias Legion.Store.Payload + defmacro __using__(opts) do repo = Keyword.fetch!(opts, :repo) table = Keyword.get(opts, :table, "legion_agents") - table_columns = - "agent_id, agent_module, parent_agent_id, status, started_at, conversation_state" - - select_sql = "SELECT #{table_columns} FROM #{table} WHERE agent_id = $1" - - list_sql = - "SELECT #{table_columns} FROM #{table} ORDER BY updated_at DESC NULLS LAST LIMIT $1" - - save_state_sql = """ - INSERT INTO #{table} (agent_id, conversation_state, inserted_at, updated_at) - VALUES ($1, $2, now(), now()) - ON CONFLICT (agent_id) DO UPDATE SET conversation_state = EXCLUDED.conversation_state, updated_at = now() - """ - - save_metadata_sql = """ - INSERT INTO #{table} (agent_id, agent_module, parent_agent_id, started_at, inserted_at, updated_at) - VALUES ($1, $2, $3, $4, now(), now()) - ON CONFLICT (agent_id) DO UPDATE - SET agent_module = EXCLUDED.agent_module, - parent_agent_id = COALESCE(#{table}.parent_agent_id, EXCLUDED.parent_agent_id), - started_at = EXCLUDED.started_at, - updated_at = now() - """ - - save_status_sql = """ - INSERT INTO #{table} (agent_id, status, inserted_at, updated_at) - VALUES ($1, $2, now(), now()) - ON CONFLICT (agent_id) DO UPDATE - SET status = EXCLUDED.status, updated_at = now() - """ - quote do - use Legion.Store + @behaviour Legion.Store + + import Ecto.Query, only: [from: 2] + alias Legion.Store.Payload alias Legion.Store.Postgres + defmodule Record do + use Ecto.Schema + + @primary_key {:agent_id, :string, autogenerate: false} + + schema unquote(table) do + field(:agent_module, :string) + field(:parent_agent_id, :string) + field(:status, :string) + field(:started_at, :integer) + field(:conversation_state, :binary) + field(:inserted_at, :utc_datetime_usec) + field(:updated_at, :utc_datetime_usec) + end + end + @impl Legion.Store def get(agent_id) when is_binary(agent_id) do - case unquote(repo).query!(unquote(select_sql), [agent_id]) do - %{rows: []} -> :error - %{rows: [row]} -> {:ok, Postgres.decode_conversation(row)} + case unquote(repo).get(Record, agent_id) do + nil -> :error + record -> {:ok, Postgres.decode_record(record)} end end @impl Legion.Store def list(limit) when is_integer(limit) and limit > 0 do - %{rows: rows} = unquote(repo).query!(unquote(list_sql), [limit]) - Enum.map(rows, &Postgres.decode_conversation/1) + from(record in Record, order_by: [desc: record.updated_at], limit: ^limit) + |> unquote(repo).all() + |> Enum.map(&Postgres.decode_record/1) end - @impl Legion.Store - def save(agent_id, %Legion.Store.Conversation.State{} = state) when is_binary(agent_id) do - unquote(repo).query!(unquote(save_state_sql), [agent_id, :erlang.term_to_binary(state)]) - :ok - end + def save(map) when is_map(map) and not is_struct(map), do: :error @impl Legion.Store - def save(agent_id, %Legion.Store.Conversation.Metadata{} = metadata) - when is_binary(agent_id) do - unquote(repo).query!( - unquote(save_metadata_sql), - [agent_id] ++ Postgres.encode_metadata(metadata) - ) - - :ok - end + def save(%Payload{agent_id: nil}), do: :error @impl Legion.Store - def save(agent_id, {:status, status}) - when is_binary(agent_id) and status in [:running, :idle] do - unquote(repo).query!(unquote(save_status_sql), [agent_id, Atom.to_string(status)]) - :ok + def save(%Payload{} = payload) do + attrs = + payload + |> Postgres.encode_data() + |> Map.reject(fn {_k, v} -> is_nil(v) end) + |> Map.put(:updated_at, DateTime.utc_now()) + + update_columns = attrs |> Map.delete(:agent_id) |> Map.keys() + + case unquote(repo).insert_all( + Record, + [attrs], + conflict_target: :agent_id, + on_conflict: {:replace, update_columns} + ) do + {1, _} -> :ok + _ -> :error + end end - def save(_agent_id, _payload), do: :error - @doc false def __repo__, do: unquote(repo) @@ -146,49 +137,39 @@ defmodule Legion.Store.Postgres do end end - @doc false - def decode_conversation([ - agent_id, - agent_module, - parent_agent_id, - status, - started_at, - conversation_state | _ - ]) do - %Legion.Store.Conversation{ - agent_id: agent_id, - metadata: decode_metadata(agent_module, parent_agent_id, started_at), - status: decode_status(status), - state: decode_state(conversation_state) - } + def encode_data(%Payload{} = payload) do + payload + |> Map.from_struct() + |> Map.update(:conversation_state, nil, fn state -> + if is_nil(state), do: nil, else: :erlang.term_to_binary(state) + end) + |> Map.update(:agent_module, nil, fn module -> + if is_nil(module), do: nil, else: inspect(module) + end) + |> Map.update(:status, nil, fn status -> + if is_nil(status), do: nil, else: Atom.to_string(status) + end) end @doc false - def encode_metadata(%Legion.Store.Conversation.Metadata{} = metadata) do - [ - metadata.agent_module && inspect(metadata.agent_module), - metadata.parent_agent_id, - metadata.started_at - ] - end - - defp decode_metadata(nil, nil, nil), do: nil - - defp decode_metadata(agent_module, parent_agent_id, started_at) do - %Legion.Store.Conversation.Metadata{ - agent_module: agent_module && Module.concat([agent_module]), - parent_agent_id: parent_agent_id, - started_at: started_at + def decode_record(record) do + %Payload{ + agent_id: record.agent_id, + agent_module: record.agent_module && Module.concat([record.agent_module]), + parent_agent_id: record.parent_agent_id, + status: decode_status(record.status), + started_at: record.started_at, + conversation_state: decode_conversation_state(record.conversation_state) } end - defp decode_state(nil), do: nil + defp decode_conversation_state(nil), do: nil # sobelow_skip ["Misc.BinToTerm"] - defp decode_state(binary) when is_binary(binary) do + defp decode_conversation_state(binary) when is_binary(binary) do state = :erlang.binary_to_term(binary) - %Legion.Store.Conversation.State{ + %{ messages: Map.get(state, :messages, []), bindings: Map.get(state, :bindings, []) } diff --git a/mix.exs b/mix.exs index 749f30c..05ad39d 100644 --- a/mix.exs +++ b/mix.exs @@ -67,6 +67,7 @@ defmodule Legion.MixProject do defp deps do [ + {:ecto, "~> 3.13"}, {:req_llm, "~> 1.2"}, {:vault, "~> 0.2"}, {:jason, "~> 1.4"}, diff --git a/mix.lock b/mix.lock index a019df6..e7fb010 100644 --- a/mix.lock +++ b/mix.lock @@ -7,6 +7,7 @@ "deep_merge": {:hex, :deep_merge, "1.0.0", "b4aa1a0d1acac393bdf38b2291af38cb1d4a52806cf7a4906f718e1feb5ee961", [:mix], [], "hexpm", "ce708e5f094b9cd4e8f2be4f00d2f4250c4095be93f8cd6d018c753894885430"}, "dotenvy": {:hex, :dotenvy, "1.1.1", "00e318f3c51de9fafc4b48598447e386f19204dc18ca69886905bb8f8b08b667", [:mix], [], "hexpm", "c8269471b5701e9e56dc86509c1199ded2b33dce088c3471afcfef7839766d8e"}, "earmark_parser": {:hex, :earmark_parser, "1.4.44", "f20830dd6b5c77afe2b063777ddbbff09f9759396500cdbe7523efd58d7a339c", [:mix], [], "hexpm", "4778ac752b4701a5599215f7030989c989ffdc4f6df457c5f36938cc2d2a2750"}, + "ecto": {:hex, :ecto, "3.13.6", "352135b474f91d1ab99a1b502171d207e9db60421c9e3d0ecab4c7ab96b24d14", [:mix], [{:decimal, "~> 2.0 or ~> 3.0", [hex: :decimal, repo: "hexpm", optional: false]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: true]}, {:telemetry, "~> 0.4 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "8afa059bc16cd2c94739ec0a11e3e5df69d828125119109bef35f20a21a76af2"}, "ex_aws_auth": {:hex, :ex_aws_auth, "1.3.1", "3963992d6f7cb251b53573603c3615cec70c3f4d86199fdb865ff440295ef7a4", [:mix], [{:jason, "~> 1.4", [hex: :jason, repo: "hexpm", optional: true]}, {:req, "~> 0.5", [hex: :req, repo: "hexpm", optional: true]}], "hexpm", "025793aa08fa419aabdb652db60edbdb2e12346bd447988a1bb5854c4dd64903"}, "ex_doc": {:hex, :ex_doc, "0.40.2", "f50edec428c4b0a457a167de42414c461122a3585a99515a69d09fff19e5597e", [:mix], [{:earmark_parser, "~> 1.4.44", [hex: :earmark_parser, repo: "hexpm", optional: false]}, {:makeup_c, ">= 0.1.0", [hex: :makeup_c, repo: "hexpm", optional: true]}, {:makeup_elixir, "~> 0.14 or ~> 1.0", [hex: :makeup_elixir, repo: "hexpm", optional: false]}, {:makeup_erlang, "~> 0.1 or ~> 1.0", [hex: :makeup_erlang, repo: "hexpm", optional: false]}, {:makeup_html, ">= 0.1.0", [hex: :makeup_html, repo: "hexpm", optional: true]}], "hexpm", "4fa426e2beb47854a162e2c488727fdec51cd4692e319b23810c2804cb1a40fe"}, "file_system": {:hex, :file_system, "1.1.1", "31864f4685b0148f25bd3fbef2b1228457c0c89024ad67f7a81a3ffbc0bbad3a", [:mix], [], "hexpm", "7a15ff97dfe526aeefb090a7a9d3d03aa907e100e262a0f8f7746b78f8f87a5d"}, diff --git a/test/legion/store/postgres_db_test.exs b/test/legion/store/postgres_db_test.exs index 8ed709c..2e76fde 100644 --- a/test/legion/store/postgres_db_test.exs +++ b/test/legion/store/postgres_db_test.exs @@ -1,16 +1,47 @@ defmodule Legion.Store.PostgresDbTest do - @moduledoc """ - Exercises the generated Postgres store against a real database, so the SQL it - issues - the `ON CONFLICT` upsert in particular - is verified for real rather - than shape-matched against a fake. - """ + @moduledoc "Exercises the generated Postgres store against a real database." use ExUnit.Case, async: false - alias Legion.Store.Conversation - alias Legion.Store.Conversation.{Metadata, State} + alias Legion.Store.Payload defmodule Repo do + @columns ~w(agent_id agent_module parent_agent_id status started_at conversation_state inserted_at updated_at)a + + def get(_schema, agent_id) do + case query!("SELECT #{columns_sql()} FROM legion_agents WHERE agent_id = $1", [agent_id]).rows do + [] -> nil + [row] -> Map.new(Enum.zip(@columns, row)) + end + end + + def insert_all(_schema, [attrs], conflict_target: :agent_id, on_conflict: {:replace, columns}) do + attrs = Map.take(attrs, [:agent_id | columns]) + fields = Map.keys(attrs) + values = Enum.map(fields, &Map.fetch!(attrs, &1)) + + sql = """ + INSERT INTO legion_agents (#{Enum.join(fields, ", ")}) + VALUES (#{placeholders(length(fields))}) + ON CONFLICT (agent_id) DO UPDATE + SET #{updates_sql(columns)} + """ + + %{num_rows: count} = query!(sql, values) + {count, nil} + end + def query!(sql, params), do: Postgrex.query!(:legion_store_test, sql, params) + + defp columns_sql, do: Enum.join(@columns, ", ") + + defp placeholders(count) do + 1..count + |> Enum.map_join(", ", &"$#{&1}") + end + + defp updates_sql(columns) do + Enum.map_join(columns, ", ", fn column -> "#{column} = EXCLUDED.#{column}" end) + end end defmodule Store do @@ -18,126 +49,78 @@ defmodule Legion.Store.PostgresDbTest do end setup do - Postgrex.query!(:legion_store_test, "TRUNCATE legion_agents", []) + Repo.query!("TRUNCATE legion_agents", []) :ok end - test "save/2 with State round-trips through a real bytea column" do - state = %State{messages: [%{type: :user, content: "hi"}], bindings: [x: 42]} - - assert :ok = Store.save("user_42", state) - assert {:ok, %Conversation{agent_id: "user_42", state: ^state}} = Store.get("user_42") - end + test "save/1 fully inserts every payload field" do + payload = %Payload{ + agent_id: "user_42", + agent_module: Legion.Test.Support.MathAgent, + parent_agent_id: "parent-1", + status: :idle, + started_at: 123, + conversation_state: %{messages: [%{role: "user", content: "hi"}], bindings: [x: 42]} + } - test "get/1 returns :error when the row is absent" do - assert :error = Store.get("missing") + assert :ok = Store.save(payload) + assert {:ok, ^payload} = Store.get("user_42") end - test "save/2 with State upserts on conflict - the latest state wins" do - assert :ok = Store.save("user_42", %State{messages: [], bindings: [v: 1]}) - assert :ok = Store.save("user_42", %State{messages: [], bindings: [v: 2]}) + 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: []} + } - assert {:ok, %Conversation{state: %State{bindings: [v: 2]}}} = Store.get("user_42") + assert :ok = Store.save(payload) + assert {:ok, ^payload} = Store.get("state-only") end - test "save/2 with Metadata records conversation identity and keeps the parent on conflict" do - metadata = %Metadata{ - agent_module: MyApp.Worker, + test "save/1 partial upsert preserves omitted fields and advances updated_at" do + initial = %Payload{ + agent_id: "user_42", + agent_module: Legion.Test.Support.MathAgent, parent_agent_id: "parent-1", - started_at: 100 + status: :running, + started_at: 123, + conversation_state: %{messages: [%{role: "user", content: "hi"}], bindings: [x: 42]} } - assert :ok = Store.save("child-1", metadata) - - # Resumed from elsewhere: no parent this time, later started_at. - assert :ok = Store.save("child-1", %{metadata | parent_agent_id: nil, started_at: 200}) - - %{rows: [[agent_module, parent_agent_id, started_at]]} = - Postgrex.query!( - :legion_store_test, - "SELECT agent_module, parent_agent_id, started_at FROM legion_agents WHERE agent_id = $1", - ["child-1"] - ) - - assert agent_module == "MyApp.Worker" - assert parent_agent_id == "parent-1" - assert started_at == 200 - end - - test "save/2 with status payload flips the stored status" do - :ok = Store.save("s1", {:status, :running}) - assert {:ok, %Conversation{status: :running, metadata: nil, state: nil}} = Store.get("s1") - - :ok = Store.save("s1", {:status, :idle}) - assert {:ok, %Conversation{status: :idle}} = Store.get("s1") - end + assert :ok = Store.save(initial) - test "save/2 with Metadata does not reset an existing status" do - metadata = %Metadata{agent_module: A, parent_agent_id: nil, started_at: 2} + %{rows: [[previous_updated_at]]} = + Repo.query!("SELECT updated_at FROM legion_agents WHERE agent_id = $1", ["user_42"]) - :ok = Store.save("s1", {:status, :running}) - :ok = Store.save("s1", metadata) + Process.sleep(1) + assert :ok = Store.save(%Payload{agent_id: "user_42", status: :idle}) - assert {:ok, %Conversation{status: :running, metadata: ^metadata}} = Store.get("s1") - end + assert {:ok, + %Payload{ + agent_module: Legion.Test.Support.MathAgent, + parent_agent_id: "parent-1", + status: :idle, + started_at: 123, + conversation_state: %{messages: [%{role: "user", content: "hi"}], bindings: [x: 42]} + }} = Store.get("user_42") - test "list/1 returns newest conversations including partial rows" do - state = %State{messages: [], bindings: []} - metadata = %Metadata{agent_module: A, parent_agent_id: nil, started_at: 1} + %{rows: [[updated_at]]} = + Repo.query!("SELECT updated_at FROM legion_agents WHERE agent_id = $1", ["user_42"]) - :ok = Store.save("state-only", state) - :ok = Store.save("metadata-only", metadata) - :ok = Store.save("status-only", {:status, :idle}) - - assert [ - %Conversation{agent_id: "status-only", status: :idle, metadata: nil, state: nil}, - %Conversation{agent_id: "metadata-only", metadata: ^metadata, state: nil}, - %Conversation{agent_id: "state-only", metadata: nil, state: ^state} - ] = Store.list(10) + assert DateTime.compare(updated_at, previous_updated_at) == :gt end - test "the trigger notifies the table's channel with the agent_id on every write" do - {:ok, notifications} = - Postgrex.Notifications.start_link( - hostname: System.get_env("POSTGRES_HOST", "localhost"), - port: String.to_integer(System.get_env("POSTGRES_PORT", "5432")), - username: System.get_env("POSTGRES_USER", "postgres"), - password: System.get_env("POSTGRES_PASSWORD", "postgres"), - database: System.get_env("POSTGRES_DB", "postgres") - ) - - {:ok, _ref} = Postgrex.Notifications.listen(notifications, "legion_agents") - - :ok = Store.save("notify-1", %Metadata{agent_module: A, parent_agent_id: nil, started_at: 1}) - assert_receive {:notification, _pid, _ref, "legion_agents", "notify-1"}, 1_000 + test "a payload cannot be constructed without agent_id" do + assert_raise ArgumentError, fn -> struct!(Payload, %{}) end - :ok = Store.save("notify-1", {:status, :running}) - assert_receive {:notification, _pid, _ref, "legion_agents", "notify-1"}, 1_000 - - :ok = Store.save("notify-1", %State{messages: [], bindings: []}) - assert_receive {:notification, _pid, _ref, "legion_agents", "notify-1"}, 1_000 - - GenServer.stop(notifications) + %{rows: [[count]]} = Repo.query!("SELECT COUNT(*) FROM legion_agents", []) + assert count == 0 end - test "metadata then state share one row and get/1 sees both" do - metadata = %Metadata{agent_module: A, parent_agent_id: nil, started_at: 1} - - assert :ok = Store.save("both", metadata) - assert {:ok, %Conversation{metadata: ^metadata, state: nil}} = Store.get("both") - - state = %State{messages: [%{type: :user, content: "hi"}], bindings: []} - assert :ok = Store.save("both", state) - - assert {:ok, %Conversation{metadata: ^metadata, state: ^state}} = Store.get("both") - - %{rows: [[count]]} = - Postgrex.query!( - :legion_store_test, - "SELECT COUNT(*) FROM legion_agents WHERE agent_id = $1", - ["both"] - ) + test "save/1 rejects unknown payload keys without inserting a row" do + assert :error = Store.save(%{agent_id: "user_42", unexpected: "value"}) - assert count == 1 + %{rows: [[count]]} = Repo.query!("SELECT COUNT(*) FROM legion_agents", []) + assert count == 0 end end diff --git a/test/legion/store/postgres_test.exs b/test/legion/store/postgres_test.exs index f00b8e9..02b580e 100644 --- a/test/legion/store/postgres_test.exs +++ b/test/legion/store/postgres_test.exs @@ -1,103 +1,32 @@ defmodule Legion.Store.PostgresTest do use ExUnit.Case, async: true - alias Legion.Store.Conversation - alias Legion.Store.Conversation.{Metadata, State} + alias Legion.Store.Payload defmodule FakeRepo do - @moduledoc "Emulates repo.query!/2 for the statements the store issues." + @moduledoc "Emulates the Ecto repository calls made by the generated store." - def start_link, do: Agent.start_link(fn -> %{rows: %{}, clock: 0} end, name: __MODULE__) + def start_link, do: Agent.start_link(fn -> %{rows: %{}} end, name: __MODULE__) - def query!("SELECT agent_id" <> _rest = sql, [param]) do - rows = - Agent.get(__MODULE__, fn state -> - state.rows - |> select_rows(sql, param) - |> Enum.map(&row/1) - end) - - %{rows: rows} - end - - def query!("INSERT INTO " <> _rest = sql, [agent_id, value]) do - cond do - String.contains?(sql, "(agent_id, conversation_state") -> - update_row(agent_id, %{conversation_state: value}) - %{num_rows: 1} - - String.contains?(sql, "(agent_id, status") -> - update_row(agent_id, %{status: value}) - %{num_rows: 1} - - true -> - raise ArgumentError, "unexpected two-argument insert: #{sql}" - end - end - - def query!("INSERT INTO " <> _rest = sql, [ - agent_id, - agent_module, - parent_agent_id, - started_at - ]) do - if String.contains?(sql, "(agent_id, agent_module, parent_agent_id, started_at") do - Agent.update(__MODULE__, fn state -> - row = Map.get(state.rows, agent_id, empty_row(agent_id)) - - put_in(state.rows[agent_id], %{ - row - | agent_module: agent_module, - parent_agent_id: row.parent_agent_id || parent_agent_id, - started_at: started_at, - updated_at: state.clock + 1 - }) - |> Map.update!(:clock, &(&1 + 1)) - end) - - %{num_rows: 1} - else - raise ArgumentError, "unexpected metadata insert: #{sql}" - end - end - - def run(agent_id), do: Agent.get(__MODULE__, & &1.rows[agent_id]) - - defp select_rows(rows, sql, param) do - if String.contains?(sql, "WHERE agent_id") do - select_row(rows, param) - else - select_recent_rows(rows, param) - end - end - - defp select_row(rows, agent_id) do - rows - |> Map.values() - |> Enum.filter(&(&1.agent_id == agent_id)) + def get(_schema, agent_id) do + Agent.get(__MODULE__, &Map.get(&1.rows, agent_id)) end - defp select_recent_rows(rows, limit) do - rows - |> Map.values() - |> Enum.sort_by(& &1.updated_at, :desc) - |> Enum.take(limit) - end - - defp update_row(agent_id, attrs) do + def insert_all(_schema, [attrs], conflict_target: :agent_id, on_conflict: {:replace, columns}) do Agent.update(__MODULE__, fn state -> row = state.rows - |> Map.get(agent_id, empty_row(agent_id)) - |> Map.merge(attrs) - |> Map.put(:updated_at, state.clock + 1) + |> Map.get(attrs.agent_id, empty_row(attrs.agent_id)) + |> Map.merge(Map.take(attrs, [:agent_id | columns])) - state - |> put_in([:rows, agent_id], row) - |> Map.update!(:clock, &(&1 + 1)) + put_in(state.rows[attrs.agent_id], row) end) + + {1, nil} end + def run(agent_id), do: Agent.get(__MODULE__, &Map.get(&1.rows, agent_id)) + defp empty_row(agent_id) do %{ agent_id: agent_id, @@ -106,20 +35,10 @@ defmodule Legion.Store.PostgresTest do status: nil, started_at: nil, conversation_state: nil, - updated_at: 0 + inserted_at: nil, + updated_at: nil } end - - defp row(row) do - [ - row.agent_id, - row.agent_module, - row.parent_agent_id, - row.status, - row.started_at, - row.conversation_state - ] - end end defmodule Store do @@ -131,163 +50,64 @@ defmodule Legion.Store.PostgresTest do :ok end - test "save/2 with State then get/1 returns a conversation with decoded state" do - state = %State{messages: [%{type: :user, content: "hi"}], bindings: [x: 42]} - - assert :ok = Store.save("user_42", state) - - assert {:ok, - %Conversation{ - agent_id: "user_42", - metadata: nil, - status: nil, - state: ^state - }} = Store.get("user_42") - end - - test "get/1 returns :error when the row is absent" do - assert :error = Store.get("missing") - end - - test "ids must be strings" do - assert_raise FunctionClauseError, fn -> Store.get(42) end - assert :error = Store.save(42, %State{messages: [], bindings: []}) - end - - test "save/2 with Metadata stores the module in inspect form with parent and start time" do - metadata = %Metadata{ + test "save/1 fully inserts every payload field" do + payload = %Payload{ + agent_id: "user_42", agent_module: Legion.Test.Support.MathAgent, - parent_agent_id: "p1", - started_at: 123 + parent_agent_id: "parent-1", + status: :idle, + started_at: 123, + conversation_state: %{messages: [%{role: "user", content: "hi"}], bindings: [x: 42]} } - assert :ok = Store.save("user_42", metadata) - - assert %{ - agent_module: "Legion.Test.Support.MathAgent", - parent_agent_id: "p1", - status: nil, - started_at: 123, - conversation_state: nil - } = FakeRepo.run("user_42") - end - - test "save/2 with status payload flips the stored status and get/1 decodes it" do - :ok = Store.save("s1", {:status, :running}) - assert {:ok, %Conversation{status: :running, metadata: nil, state: nil}} = Store.get("s1") - - :ok = Store.save("s1", {:status, :idle}) - assert {:ok, %Conversation{status: :idle}} = Store.get("s1") - end - - test "list/1 returns decoded conversations newest first" do - :ok = - Store.save("a", %Metadata{agent_module: SomeAgent, parent_agent_id: nil, started_at: 1}) - - :ok = - Store.save("b", %Metadata{agent_module: OtherAgent, parent_agent_id: "a", started_at: 2}) - - assert [ - %Conversation{ - agent_id: "b", - metadata: %Metadata{ - agent_module: OtherAgent, - parent_agent_id: "a", - started_at: 2 - } - }, - %Conversation{ - agent_id: "a", - metadata: %Metadata{ - agent_module: SomeAgent, - parent_agent_id: nil, - started_at: 1 - } - } - ] = Store.list(10) - - assert [%Conversation{agent_id: "b"}] = Store.list(1) - end - - test "get/1 returns decoded conversation, or :error when missing" do - :ok = Store.save("a", %Metadata{agent_module: SomeAgent, parent_agent_id: nil, started_at: 1}) - - assert {:ok, %Conversation{agent_id: "a", metadata: %Metadata{agent_module: SomeAgent}}} = - Store.get("a") - - assert Store.get("missing") == :error - end - - test "get/1 returns conversation with nil state when only metadata exists" do - metadata = %Metadata{agent_module: SomeAgent, parent_agent_id: nil, started_at: 1} - - assert :ok = Store.save("started-only", metadata) - assert {:ok, %Conversation{metadata: ^metadata, state: nil}} = Store.get("started-only") + assert :ok = Store.save(payload) + assert {:ok, ^payload} = Store.get("user_42") end - test "get/1 returns conversation with nil metadata when only state exists" do - state = %State{messages: [], bindings: [answer: 42]} + 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: []} + } - assert :ok = Store.save("state-only", state) - assert {:ok, %Conversation{metadata: nil, state: ^state}} = Store.get("state-only") + assert :ok = Store.save(payload) + assert {:ok, ^payload} = Store.get("state-only") end - test "save/2 with Metadata does not reset an existing status" do - metadata = %Metadata{agent_module: SomeAgent, parent_agent_id: nil, started_at: 1} - - assert :ok = Store.save("status-metadata", {:status, :running}) - assert :ok = Store.save("status-metadata", metadata) - - assert {:ok, %Conversation{status: :running, metadata: ^metadata}} = - Store.get("status-metadata") - end + test "save/1 partial upsert preserves omitted fields and advances updated_at" do + initial = %Payload{ + agent_id: "user_42", + agent_module: Legion.Test.Support.MathAgent, + parent_agent_id: "parent-1", + status: :running, + started_at: 123, + conversation_state: %{messages: [%{role: "user", content: "hi"}], bindings: [x: 42]} + } - test "save/2 with Metadata preserves an existing parent on conflict" do - metadata = %Metadata{agent_module: SomeAgent, parent_agent_id: "parent-1", started_at: 1} + assert :ok = Store.save(initial) + previous_updated_at = FakeRepo.run("user_42").updated_at - assert :ok = Store.save("child", metadata) - assert :ok = Store.save("child", %{metadata | parent_agent_id: nil, started_at: 2}) + assert :ok = Store.save(%Payload{agent_id: "user_42", status: :idle}) assert {:ok, - %Conversation{ - metadata: %Metadata{parent_agent_id: "parent-1", started_at: 2} - }} = Store.get("child") - end - - test "list/1 includes state-only metadata-only and status-only conversations" do - state = %State{messages: [], bindings: []} - metadata = %Metadata{agent_module: SomeAgent, parent_agent_id: nil, started_at: 1} - - :ok = Store.save("state-only", state) - :ok = Store.save("metadata-only", metadata) - :ok = Store.save("status-only", {:status, :idle}) + %Payload{ + agent_module: Legion.Test.Support.MathAgent, + parent_agent_id: "parent-1", + status: :idle, + started_at: 123, + conversation_state: %{messages: [%{role: "user", content: "hi"}], bindings: [x: 42]} + }} = Store.get("user_42") - assert [ - %Conversation{agent_id: "status-only", status: :idle, metadata: nil, state: nil}, - %Conversation{agent_id: "metadata-only", metadata: ^metadata, state: nil}, - %Conversation{agent_id: "state-only", metadata: nil, state: ^state} - ] = Store.list(10) + assert DateTime.compare(FakeRepo.run("user_42").updated_at, previous_updated_at) == :gt end - test "save/2 returns :error for unsupported payloads" do - assert :error = Store.save("bad-payload", %{messages: [], bindings: []}) - assert :error = Store.save("bad-status", {:status, :paused}) + test "a payload cannot be constructed without agent_id" do + assert_raise ArgumentError, fn -> struct!(Payload, %{}) end + assert FakeRepo.run("missing") == nil end - test "a custom table name is interpolated into the statements" do - defmodule TableCapturingRepo do - def query!(sql, _params), do: send(self(), {:sql, sql}) && %{rows: []} - end - - defmodule CustomTableStore do - use Legion.Store.Postgres, - repo: Legion.Store.PostgresTest.TableCapturingRepo, - table: "my_agents" - end - - CustomTableStore.get("user_42") - - assert_received {:sql, - "SELECT agent_id, agent_module, parent_agent_id, status, started_at, conversation_state FROM my_agents WHERE agent_id = $1"} + test "save/1 rejects unknown payload keys without inserting a row" do + assert :error = Store.save(%{agent_id: "user_42", unexpected: "value"}) + assert FakeRepo.run("user_42") == nil end end diff --git a/test/legion/store_test.exs b/test/legion/store_test.exs deleted file mode 100644 index 9aba1f1..0000000 --- a/test/legion/store_test.exs +++ /dev/null @@ -1,55 +0,0 @@ -defmodule Legion.StoreTest do - use ExUnit.Case, async: true - - import ExUnit.CaptureLog - - alias Legion.Store.Conversation.{Metadata, State} - - defmodule DefaultStore do - use Legion.Store - - def get(_agent_id), do: :error - end - - defmodule CustomStore do - use Legion.Store - - def get(_agent_id), do: :error - def save(_agent_id, _payload), do: :custom - end - - test "default save/2 warns and returns :ok for conversation state" do - payload = %State{messages: [], bindings: []} - - log = - capture_log(fn -> - assert :ok = DefaultStore.save("agent-1", payload) - end) - - assert log =~ "Store #{inspect(DefaultStore)} does not persist conversation state" - end - - test "default save/2 warns and returns :ok for conversation metadata" do - payload = %Metadata{agent_module: SomeAgent, parent_agent_id: nil, started_at: 1} - - log = - capture_log(fn -> - assert :ok = DefaultStore.save("agent-1", payload) - end) - - assert log =~ "Store #{inspect(DefaultStore)} does not persist conversation metadata" - end - - test "default save/2 warns and returns :ok for conversation status" do - log = - capture_log(fn -> - assert :ok = DefaultStore.save("agent-1", {:status, :running}) - end) - - assert log =~ "Store #{inspect(DefaultStore)} does not persist conversation status" - end - - test "default save/2 can be overridden" do - assert :custom = CustomStore.save("agent-1", %State{messages: [], bindings: []}) - end -end From 9aa4605b2ba8f6247c0e515a374f419cd37c723c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tomasz=20Id=C5=BAkowski?= Date: Tue, 21 Jul 2026 10:32:24 +0200 Subject: [PATCH 11/30] Wire new store contract into AgentServer --- lib/legion/agent_server.ex | 56 ++++++------ test/legion/agent_server_test.exs | 146 ++++++++++++++++++++++-------- 2 files changed, 136 insertions(+), 66 deletions(-) diff --git a/lib/legion/agent_server.ex b/lib/legion/agent_server.ex index 7e26c5d..228e839 100644 --- a/lib/legion/agent_server.ex +++ b/lib/legion/agent_server.ex @@ -11,8 +11,7 @@ defmodule Legion.AgentServer do require Logger alias Legion.{Executor, Telemetry} - alias Legion.Store.Conversation - alias Legion.Store.Conversation.{Metadata, State} + alias Legion.Store.Payload alias ReqLLM.Message.ContentPart defstruct [:agent_module, :messages, :config, :store, :agent_id, bindings: []] @@ -80,7 +79,7 @@ defmodule Legion.AgentServer do {saved_messages, saved_bindings} = case store && store.get(agent_id) do - {:ok, %Conversation{state: %State{messages: messages, bindings: bindings}}} -> + {:ok, %Payload{conversation_state: %{messages: messages, bindings: bindings}}} -> {messages, bindings} _no_state -> @@ -96,7 +95,12 @@ defmodule Legion.AgentServer do bindings: saved_bindings } - {:ok, persist(state, {:metadata, parent_agent_id})} + {:ok, + persist(state, + agent_module: state.agent_module, + parent_agent_id: parent_agent_id, + started_at: System.system_time(:millisecond) + )} end @impl true @@ -153,8 +157,7 @@ defmodule Legion.AgentServer do state = state |> Map.update!(:messages, &(&1 ++ [Executor.message(:user, content)])) - |> persist(:state) - |> persist({:status, :running}) + |> persist([:conversation_state, status: :running]) {status, value, final_messages, final_bindings} = Telemetry.span( @@ -178,39 +181,36 @@ defmodule Legion.AgentServer do state = %{state | messages: final_messages, bindings: kept_bindings} - |> persist(:state) - |> persist({:status, :idle}) + |> persist([:conversation_state, status: :idle]) {{status, value}, state} end - defp persist(%{store: nil} = state, _payload), do: state + defp persist(%{store: nil} = state, _fields), do: state - defp persist(state, :state) do - [%{role: "system"} | messages] = state.messages - :ok = state.store.save(state.agent_id, %State{messages: messages, bindings: state.bindings}) + defp persist(state, fields) do + :ok = state.store.save(payload(state, fields)) state end - defp persist(state, {:metadata, parent_agent_id}) do - :ok = - state.store.save( - state.agent_id, - %Metadata{ - agent_module: state.agent_module, - parent_agent_id: parent_agent_id, - started_at: System.system_time(:millisecond) - } - ) + defp payload(state, fields) do + Enum.reduce(fields, %Payload{agent_id: state.agent_id}, fn + :conversation_state, payload -> + %{payload | conversation_state: persisted_conversation_state(state)} - state + {field, value}, payload + when field in [:agent_module, :parent_agent_id, :status, :started_at] -> + Map.put(payload, field, value) + + unknown, _payload -> + raise ArgumentError, "unsupported persistence field: #{inspect(unknown)}" + end) end - # A crash between the :running write and :idle write leaves `:running` in - # the store, which consumers read as "crashed mid-turn" under a dead pid. - defp persist(state, {:status, status}) when status in [:running, :idle] do - :ok = state.store.save(state.agent_id, {:status, status}) - state + defp persisted_conversation_state(state) do + [%{role: "system"} | messages] = state.messages + + %{messages: messages, bindings: state.bindings} end defp generate_id, do: Base.url_encode64(:crypto.strong_rand_bytes(16), padding: false) diff --git a/test/legion/agent_server_test.exs b/test/legion/agent_server_test.exs index 94dd7d5..9d3cc00 100644 --- a/test/legion/agent_server_test.exs +++ b/test/legion/agent_server_test.exs @@ -4,8 +4,7 @@ defmodule Legion.AgentServerTest do import ExUnit.CaptureLog - alias Legion.Store.Conversation - alias Legion.Store.Conversation.{Metadata, State} + alias Legion.Store.Payload alias Legion.Test.Support.MathAgent alias ReqLLM.Message.ContentPart @@ -421,75 +420,83 @@ defmodule Legion.AgentServerTest do def start_link, do: Agent.start_link(fn -> %{} end, name: __MODULE__) @impl Legion.Store - def get(agent_id) do + def get(agent_id), do: Agent.get(__MODULE__, &Map.get(&1, agent_id, :error)) + + @impl Legion.Store + def list(limit) do Agent.get(__MODULE__, fn state -> - case Map.fetch(state, agent_id) do - {:ok, conversation} -> {:ok, conversation} - :error -> :error - end + state + |> Map.values() + |> Enum.flat_map(fn + {:ok, %Payload{} = payload} -> [payload] + _ -> [] + end) + |> Enum.take(limit) end) end def load(agent_id) do case get(agent_id) do - {:ok, %Conversation{state: state}} when not is_nil(state) -> {:ok, state} + {:ok, %Payload{conversation_state: state}} when not is_nil(state) -> {:ok, state} _ -> :error end end @impl Legion.Store - def save(agent_id, %State{} = snapshot) do - update_conversation(agent_id, &%{&1 | state: snapshot}) - end + def save(%Payload{} = payload) do + Agent.update(__MODULE__, fn state -> + existing = + case Map.get(state, payload.agent_id) do + {:ok, stored} -> stored + nil -> %Payload{agent_id: payload.agent_id} + end - def save(agent_id, %Metadata{} = metadata) do - update_conversation(agent_id, &%{&1 | metadata: metadata}) - end + merged = merge(existing, payload) - def save(agent_id, {:status, status}) do - Agent.update(__MODULE__, fn state -> state - |> Map.update({:statuses, agent_id}, [status], &[status | &1]) - |> Map.update( - agent_id, - %Conversation{agent_id: agent_id, status: status}, - fn conversation -> - %{conversation | status: status} - end - ) + |> Map.put(payload.agent_id, {:ok, merged}) + |> Map.update({:writes, payload.agent_id}, [payload], &[payload | &1]) end) :ok end - def statuses(agent_id) do - Agent.get(__MODULE__, &Map.get(&1, {:statuses, agent_id}, [])) + def save(_invalid), do: :error + + def writes(agent_id) do + Agent.get(__MODULE__, &Map.get(&1, {:writes, agent_id}, [])) |> Enum.reverse() end + def statuses(agent_id) do + agent_id + |> writes() + |> Enum.map(& &1.status) + |> Enum.reject(&is_nil/1) + end + def get_run(agent_id) do case get(agent_id) do - {:ok, %Conversation{metadata: nil}} -> nil - {:ok, %Conversation{metadata: metadata}} -> Map.put(metadata, :agent_id, agent_id) + {:ok, %Payload{agent_module: nil}} -> nil + {:ok, %Payload{} = payload} -> payload :error -> nil end end def runs do Agent.get(__MODULE__, fn state -> - for {agent_id, %Conversation{metadata: metadata}} <- state, - is_binary(agent_id), - not is_nil(metadata), - do: Map.put(metadata, :agent_id, agent_id) + for {_agent_id, {:ok, %Payload{agent_module: agent_module} = payload}} <- state, + not is_nil(agent_module), + do: payload end) end - defp update_conversation(agent_id, update) do - Agent.update(__MODULE__, fn state -> - Map.update(state, agent_id, update.(%Conversation{agent_id: agent_id}), update) + defp merge(existing, incoming) do + Enum.reduce(Map.from_struct(incoming), existing, fn + {:agent_id, _agent_id}, payload -> payload + {_field, nil}, payload -> payload + {field, value}, payload -> Map.put(payload, field, value) end) - - :ok end end @@ -513,6 +520,50 @@ defmodule Legion.AgentServerTest do assert MemoryStore.statuses("statuses") == [:running, :idle, :running, :idle] end + test "writes the new Store payloads for a completed message" do + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + llm_response("Paris") + end) + + {:ok, pid} = Legion.start_link(MathAgent, store: MemoryStore, agent_id: "payloads") + {:ok, "Paris"} = Legion.call(pid, "What is the capital of France?") + + [started, running, completed] = MemoryStore.writes("payloads") + + assert %Payload{ + agent_id: "payloads", + agent_module: MathAgent, + parent_agent_id: nil, + started_at: started_at, + status: nil, + conversation_state: nil + } = started + + assert is_integer(started_at) + + assert %Payload{ + agent_id: "payloads", + status: :running, + conversation_state: %{ + messages: [%{role: "user", content: "What is the capital of France?"}], + bindings: [] + } + } = running + + assert %Payload{ + agent_id: "payloads", + status: :idle, + conversation_state: %{messages: messages, bindings: []} + } = completed + + assert [ + %{role: "user", content: "What is the capital of France?"}, + %{role: "assistant"} | _ + ] = messages + + refute Enum.any?(messages, &(&1.role == "system")) + end + test "saves a snapshot before the caller receives its reply" do stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> llm_response("Paris") @@ -572,6 +623,25 @@ defmodule Legion.AgentServerTest do assert {:ok, %{bindings: []}} = MemoryStore.load("turn-bindings") end + test "persists conversation-scoped bindings in the completed payload" do + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + llm_eval_response("x = 42") + end) + + {:ok, pid} = + Legion.start_link(ConversationBindingsAgent, + store: MemoryStore, + agent_id: "conversation-bindings" + ) + + assert {:ok, 42} = Legion.call(pid, "set x") + + assert {:ok, + %Payload{ + conversation_state: %{bindings: [x: 42]} + }} = MemoryStore.get("conversation-bindings") + end + test "restores the conversation under a fresh system prompt after a restart" do stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> llm_response("Paris") @@ -686,7 +756,7 @@ defmodule Legion.AgentServerTest do {:ok, revived} = Legion.resume("resume-dead", store: MemoryStore) - assert revived != pid + assert Legion.running?(revived) assert [ %{role: "system"}, From 4f3353f6c569ef51aa24c2298c0d4e88a7d9f3e5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tomasz=20Id=C5=BAkowski?= Date: Tue, 21 Jul 2026 13:00:56 +0200 Subject: [PATCH 12/30] Fix race condition in resume/2 Besides looking in the registry (lookup/1) resume/2 also checks whether the process is actually alive using running?/1 --- lib/legion.ex | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/lib/legion.ex b/lib/legion.ex index da05521..dd63219 100644 --- a/lib/legion.ex +++ b/lib/legion.ex @@ -174,8 +174,15 @@ defmodule Legion do "no run recorded for agent_id #{inspect(agent_id)} in #{inspect(store)}" case lookup(agent_id) do - {:ok, pid} -> {:ok, pid} - :error -> start_link(run.agent_module, Keyword.put(opts, :agent_id, agent_id)) + {:ok, pid} -> + if running?(pid) do + {:ok, pid} + else + start_link(run.agent_module, Keyword.put(opts, :agent_id, agent_id)) + end + + :error -> + start_link(run.agent_module, Keyword.put(opts, :agent_id, agent_id)) end end From f71080b24f361e1f8fb53529c3aeca0110a87e0c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tomasz=20Id=C5=BAkowski?= Date: Tue, 21 Jul 2026 14:15:34 +0200 Subject: [PATCH 13/30] Add ability to store intra-turn updates Legion now supports emitting live updates from within a single LLM turn. It's configurable through implementing the optional persistence_frequency/0 callback and providing either :turn or :step, default value is :turn and it persists the state at the end of a turn (that does inclde the intra state history, but the refresh rate is slower). When using :step writes also occur after each intermediate eval result or recovorable error to ensure full recoverability is achievable in the future --- lib/legion/agent_server.ex | 48 +++++++- lib/legion/executor.ex | 38 ++++++- lib/legion/store.ex | 76 +++++++++++-- lib/legion/store/payload.ex | 10 +- lib/legion/store/postgres.ex | 54 +++++---- test/legion/agent_server_test.exs | 170 ++++++++++++++++++++++++++++ test/legion/executor_test.exs | 145 ++++++++++++++++++++++++ test/legion/store/postgres_test.exs | 28 +++++ test/legion/store_test.exs | 63 +++++++++++ 9 files changed, 592 insertions(+), 40 deletions(-) create mode 100644 test/legion/store_test.exs diff --git a/lib/legion/agent_server.ex b/lib/legion/agent_server.ex index 228e839..1ab38ba 100644 --- a/lib/legion/agent_server.ex +++ b/lib/legion/agent_server.ex @@ -10,11 +10,19 @@ defmodule Legion.AgentServer do require Logger - alias Legion.{Executor, Telemetry} + alias Legion.{Executor, Store, Telemetry} alias Legion.Store.Payload alias ReqLLM.Message.ContentPart - defstruct [:agent_module, :messages, :config, :store, :agent_id, bindings: []] + defstruct [ + :agent_module, + :messages, + :config, + :store, + :agent_id, + :persistence_frequency, + bindings: [] + ] # Client API @@ -31,10 +39,16 @@ defmodule Legion.AgentServer do end agent_id = agent_id || generate_id() + persistence_frequency = Store.persistence_frequency(store) gen_opts = if name, do: [name: name], else: [] config = resolve_config(agent_module, opts) - GenServer.start_link(__MODULE__, {agent_module, config, store, agent_id}, gen_opts) + + GenServer.start_link( + __MODULE__, + {agent_module, config, store, agent_id, persistence_frequency}, + gen_opts + ) end def call(agent, message, timeout \\ :infinity) do @@ -56,7 +70,7 @@ defmodule Legion.AgentServer do # Server callbacks @impl true - def init({agent_module, config, store, agent_id}) do + def init({agent_module, config, store, agent_id, persistence_frequency}) do parent_agent_id = Vault.get(:agent_id) Vault.unsafe_put(:agent_id, agent_id) @@ -92,6 +106,7 @@ defmodule Legion.AgentServer do config: config, store: store, agent_id: agent_id, + persistence_frequency: persistence_frequency, bindings: saved_bindings } @@ -159,6 +174,16 @@ defmodule Legion.AgentServer do |> Map.update!(:messages, &(&1 ++ [Executor.message(:user, content)])) |> persist([:conversation_state, status: :running]) + checkpoint = + if state.persistence_frequency == :step do + fn checkpoint -> + persist(state, [{:conversation_state, checkpoint}]) + :ok + end + end + + executor_config = Map.put(state.config, :checkpoint, checkpoint) + {status, value, final_messages, final_bindings} = Telemetry.span( [:legion, :agent, :message], @@ -170,7 +195,7 @@ defmodule Legion.AgentServer do initial_bindings = if conversation_scope?, do: state.bindings, else: [] {status, value, messages, bindings} = - result = Executor.run(state.agent_module, messages, state.config, initial_bindings) + result = Executor.run(state.agent_module, messages, executor_config, initial_bindings) iterations = Enum.count(messages, &(&1[:role] == "assistant")) - prev_count {result, %{iterations: iterations, status: status, result: value, bindings: bindings}} @@ -198,6 +223,9 @@ defmodule Legion.AgentServer do :conversation_state, payload -> %{payload | conversation_state: persisted_conversation_state(state)} + {:conversation_state, checkpoint}, payload -> + %{payload | conversation_state: persisted_conversation_state(checkpoint)} + {field, value}, payload when field in [:agent_module, :parent_agent_id, :status, :started_at] -> Map.put(payload, field, value) @@ -207,12 +235,20 @@ defmodule Legion.AgentServer do end) end - defp persisted_conversation_state(state) do + defp persisted_conversation_state(%__MODULE__{} = state) do [%{role: "system"} | messages] = state.messages %{messages: messages, bindings: state.bindings} end + defp persisted_conversation_state(%{ + messages: [%{role: "system"} | messages], + bindings: bindings, + execution: execution + }) do + %{messages: messages, bindings: bindings, execution: execution} + end + defp generate_id, do: Base.url_encode64(:crypto.strong_rand_bytes(16), padding: false) defp stringify(message, max_length) when is_binary(message), diff --git a/lib/legion/executor.ex b/lib/legion/executor.ex index 52efab2..8ea9bed 100644 --- a/lib/legion/executor.ex +++ b/lib/legion/executor.ex @@ -180,6 +180,25 @@ defmodule Legion.Executor do ) end + defp checkpoint!(config, messages, bindings, execution) do + case config[:checkpoint] do + nil -> + :ok + + callback -> + try do + :ok = + callback.(%{ + messages: messages, + bindings: bindings, + execution: execution + }) + rescue + error -> exit({:checkpoint_persistence_failed, error}) + end + end + end + defp handle_action( _agent, messages, @@ -215,6 +234,15 @@ defmodule Legion.Executor do messages = messages ++ [message(:eval_result, format_result(result, new_bindings, config))] + execution = + 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) + if eval == "eval_and_continue", do: loop(agent, messages, config, i + 1, 0, new_bindings), else: {:ok, result, messages, new_bindings} @@ -275,7 +303,15 @@ defmodule Legion.Executor do ) ] - loop(agent_module, messages, config, iteration, retries + 1, bindings) + next_retries = retries + 1 + + checkpoint!(config, messages, bindings, %{ + phase: :awaiting_llm, + iteration: iteration, + retries: next_retries + }) + + loop(agent_module, messages, config, iteration, next_retries, bindings) end end diff --git a/lib/legion/store.ex b/lib/legion/store.ex index 943e564..af3e74e 100644 --- a/lib/legion/store.ex +++ b/lib/legion/store.ex @@ -12,10 +12,16 @@ defmodule Legion.Store do is regenerated for every start, so prompt and tool changes apply to restored conversations. - After every completed turn, the agent calls `c:save/1` with a payload before - replying to the caller. A reply is a commit receipt: any observed turn - survives a crash, restart, or deploy. A crash mid-turn rolls back to the last - completed turn. + By default, Legion saves after appending the external user message and again + after completing the turn. The final save happens before replying to the + caller, so a reply is a commit receipt: any observed turn survives a crash, + restart, or deploy. + + A store can opt into step persistence by implementing + `c:persistence_frequency/0` and returning `:step`. Legion then also saves + after intermediate user-role messages, including eval results and recoverable + errors. Each step save contains the complete conversation and executor state + at that checkpoint. For Postgres users there is a ready-made adapter - see `Legion.Store.Postgres`: @@ -29,6 +35,17 @@ defmodule Legion.Store do config :legion, :store, MyApp.AgentStore + Stores default to `:turn` persistence. To persist intermediate executor + steps, return `:step`: + + defmodule MyApp.AgentStore do + @behaviour Legion.Store + + def persistence_frequency, do: :step + + # Implement get/1, list/1, and save/1... + end + A `:store` given to `start_link/2` overrides the global one. Sub-agents spawned from a running agent (e.g. via `Legion.Tools.AgentTool`) inherit the parent's store automatically. The inheritance is ambient: *any* agent @@ -55,10 +72,15 @@ defmodule Legion.Store do `c: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 (relevant with - `binding_scope: :conversation`). `: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. + 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. + + With `binding_scope: :turn`, active bindings are included in step snapshots + while the turn is running and cleared from the final snapshot. Bindings with + `binding_scope: :conversation` remain in the final snapshot, while + iteration-scoped bindings are cleared before a step is saved. Each message carries a `:type` (`:user`, `:assistant`, `:eval_result`, or `:error`) and an `:at` timestamp in milliseconds, so consumers can classify @@ -67,11 +89,15 @@ defmodule Legion.Store do serialization, so keep conversation-scoped variables to plain data if you persist agents. - `use Legion.Store` provides a default no-op `save/1` that logs a warning. - Override it for durable persistence. - Payload fields other than `:agent_id` may be nil. A store can therefore - persist identity or status before a conversation state exists. + persist identity or status before a conversation state exists. Stores must + treat nil fields as omitted partial updates so later state-only saves preserve + metadata and the running status. + + 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. ## Reading conversations @@ -88,6 +114,7 @@ defmodule Legion.Store do @type agent_id :: term() @type status :: Payload.status() @type payload :: Payload.t() + @type persistence_frequency :: :turn | :step @doc "Returns the persisted conversation for `agent_id`, or `:error` if none exists." @callback get(agent_id()) :: {:ok, payload()} | :error @@ -97,4 +124,29 @@ defmodule Legion.Store do @doc "Saves a conversation payload." @callback save(payload()) :: :ok | :error + + @doc "Returns how frequently Legion persists conversation state for this store." + @callback persistence_frequency() :: persistence_frequency() + + @optional_callbacks persistence_frequency: 0 + + @doc false + def persistence_frequency(nil), do: :turn + + def persistence_frequency(store) do + frequency = + if Code.ensure_loaded?(store) and function_exported?(store, :persistence_frequency, 0) do + store.persistence_frequency() + else + :turn + end + + if frequency in [:turn, :step] do + frequency + else + raise ArgumentError, + "#{inspect(store)}.persistence_frequency/0 must return " <> + ":turn or :step, got: #{inspect(frequency)}" + end + end end diff --git a/lib/legion/store/payload.ex b/lib/legion/store/payload.ex index 00d69bf..0f9dbf9 100644 --- a/lib/legion/store/payload.ex +++ b/lib/legion/store/payload.ex @@ -14,9 +14,15 @@ defmodule Legion.Store.Payload do ] @type status :: :idle | :running + @type execution :: %{ + phase: :awaiting_llm | :completing, + iteration: non_neg_integer(), + retries: non_neg_integer() + } @type state :: %{ - messages: [map()], - bindings: keyword() + required(:messages) => [map()], + required(:bindings) => keyword(), + optional(:execution) => execution() } @type t :: %__MODULE__{ agent_id: Legion.Store.agent_id(), diff --git a/lib/legion/store/postgres.ex b/lib/legion/store/postgres.ex index 60d76c8..524acfd 100644 --- a/lib/legion/store/postgres.ex +++ b/lib/legion/store/postgres.ex @@ -2,13 +2,19 @@ defmodule Legion.Store.Postgres do @moduledoc """ A ready-made `Legion.Store` backed by Postgres, through your existing Ecto repo. - Legion does not depend on Ecto - the generated store only calls - `repo.query!/2` at runtime, so it works with any `Ecto.Repo` on - `Ecto.Adapters.Postgres` that your application already runs. + Legion depends on Ecto. The generated store defines an `Ecto.Schema` and + uses your application's existing `Ecto.Repo` for reads and partial upserts. + The repo must use `Ecto.Adapters.Postgres`. ## Usage - Define a store module: + Define a Postgres-backed Ecto repo and a store module that uses it: + + defmodule MyApp.Repo do + use Ecto.Repo, + otp_app: :my_app, + adapter: Ecto.Adapters.Postgres + end defmodule MyApp.AgentStore do use Legion.Store.Postgres, repo: MyApp.Repo @@ -31,27 +37,32 @@ defmodule Legion.Store.Postgres do - `:repo` (required) - your Ecto repo module - `:table` - the table name, defaults to `"legion_agents"` + - `:persistence_frequency` - `:turn` (default) or `:step` + + To persist intermediate eval results and recoverable errors: + + defmodule MyApp.AgentStore do + use Legion.Store.Postgres, + repo: MyApp.Repo, + persistence_frequency: :step + end Agent ids must be strings. Snapshots are stored as `:erlang.term_to_binary/1` blobs - readable only from Elixir, one row per conversation, upserted on - every turn. + every save. Step snapshots therefore require no additional migration. - The store also implements `c:Legion.Store.save_run/2`, so the same row - carries the conversation's identity: `agent_module` (in `inspect/1` form, - e.g. `"MyApp.ResearchAgent"`), `parent_agent_id` linking a sub-agent to the - conversation that spawned it, and `started_at` in milliseconds (last start - wins). `snapshot` is null until the conversation's first turn completes; - `parent_agent_id` is kept once set, so resuming from elsewhere does not - reparent the conversation. + `c:Legion.Store.save/1` performs partial upserts, so the same row carries the + conversation state and identity: `agent_module` (in `inspect/1` form, e.g. + `"MyApp.ResearchAgent"`), `parent_agent_id` linking a sub-agent to the + conversation that spawned it, and `started_at` in milliseconds. Omitted + payload fields preserve their existing values. - `c:Legion.Store.save_status/2` is implemented as well: the row's `status` - flips to `'running'` when a turn starts and back to `'idle'` when it - completes (`save_run` resets it to `'idle'` on start), so consumers can - identify conversations that were mid-turn when persistence last observed - them. + The row's `status` flips to `'running'` when a turn starts and back to + `'idle'` when it completes. Step writes update only the conversation state, + leaving the running status unchanged. - `c:Legion.Store.list_runs/1` and `c:Legion.Store.get_run/1` are implemented - too, so persisted conversations can be read back into a view of past runs. + `c:Legion.Store.list/1` and `c:Legion.Store.get/1` read persisted + conversations back from the same table. The migration also installs a trigger that `pg_notify`s the table's channel (the table name) with the `agent_id` on every insert or update, so @@ -63,6 +74,7 @@ defmodule Legion.Store.Postgres do defmacro __using__(opts) do repo = Keyword.fetch!(opts, :repo) table = Keyword.get(opts, :table, "legion_agents") + persistence_frequency = Keyword.get(opts, :persistence_frequency, :turn) quote do @behaviour Legion.Store @@ -88,6 +100,9 @@ defmodule Legion.Store.Postgres do end end + @impl Legion.Store + def persistence_frequency, do: unquote(persistence_frequency) + @impl Legion.Store def get(agent_id) when is_binary(agent_id) do case unquote(repo).get(Record, agent_id) do @@ -173,6 +188,7 @@ defmodule Legion.Store.Postgres do messages: Map.get(state, :messages, []), bindings: Map.get(state, :bindings, []) } + |> Map.merge(Map.take(state, [:execution])) end defp decode_status("running"), do: :running diff --git a/test/legion/agent_server_test.exs b/test/legion/agent_server_test.exs index 9d3cc00..0b7c816 100644 --- a/test/legion/agent_server_test.exs +++ b/test/legion/agent_server_test.exs @@ -61,6 +61,16 @@ defmodule Legion.AgentServerTest do }} end + defp llm_eval_continue_response(code) do + {:ok, + %ReqLLM.Response{ + id: "test", + model: "test", + context: nil, + object: %{"action" => "eval_and_continue", "code" => code, "result" => ""} + }} + end + describe "get_messages/1" do test "returns conversation history from a running agent" do stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> @@ -500,6 +510,24 @@ defmodule Legion.AgentServerTest do end end + defmodule StepMemoryStore do + @behaviour Legion.Store + + alias Legion.AgentServerTest.MemoryStore + + @impl Legion.Store + def persistence_frequency, do: :step + + @impl Legion.Store + def get(agent_id), do: MemoryStore.get(agent_id) + + @impl Legion.Store + def list(limit), do: MemoryStore.list(limit) + + @impl Legion.Store + def save(payload), do: MemoryStore.save(payload) + end + describe "persistence" do setup do start_supervised!(%{id: MemoryStore, start: {MemoryStore, :start_link, []}}) @@ -642,6 +670,148 @@ defmodule Legion.AgentServerTest do }} = MemoryStore.get("conversation-bindings") end + test "a :step store persists a complete eval_and_continue checkpoint" do + call_count = :counters.new(1, [:atomics]) + + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + :counters.add(call_count, 1, 1) + + case :counters.get(call_count, 1) do + 1 -> llm_eval_continue_response("x = 42") + 2 -> llm_response("done") + end + end) + + {:ok, pid} = + Legion.start_link(MathAgent, store: StepMemoryStore, agent_id: "step-continue") + + assert {:ok, "done"} = Legion.call(pid, "compute") + + [_started, running, checkpoint, completed] = MemoryStore.writes("step-continue") + + assert %Payload{ + status: :running, + conversation_state: %{messages: [%{type: :user}], bindings: []} + } = running + + assert %Payload{ + status: nil, + conversation_state: %{ + messages: [%{type: :user}, %{type: :assistant}, %{type: :eval_result}], + bindings: [x: 42], + execution: %{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) + end + + test "a :step store persists eval_and_complete before the final snapshot" do + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + llm_eval_response("1 + 1") + end) + + {:ok, pid} = + Legion.start_link(MathAgent, store: StepMemoryStore, agent_id: "step-complete") + + assert {:ok, 2} = Legion.call(pid, "compute") + + [_started, _running, checkpoint, completed] = MemoryStore.writes("step-complete") + + assert %Payload{ + status: nil, + conversation_state: %{ + execution: %{phase: :completing, iteration: 0, retries: 0} + } + } = checkpoint + + assert %Payload{status: :idle, conversation_state: final_state} = completed + refute Map.has_key?(final_state, :execution) + end + + test "a :step store persists retry state after an error message" do + call_count = :counters.new(1, [:atomics]) + + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + :counters.add(call_count, 1, 1) + + case :counters.get(call_count, 1) do + 1 -> llm_eval_response("raise \"boom\"") + 2 -> llm_response("recovered") + end + end) + + {:ok, pid} = Legion.start_link(MathAgent, store: StepMemoryStore, agent_id: "step-retry") + assert {:ok, "recovered"} = Legion.call(pid, "compute") + + [_started, _running, checkpoint, _completed] = MemoryStore.writes("step-retry") + + assert %Payload{ + status: nil, + conversation_state: %{ + messages: messages, + bindings: [], + execution: %{phase: :awaiting_llm, iteration: 0, retries: 1} + } + } = checkpoint + + assert List.last(messages).type == :error + end + + test "a :step store retains conversation bindings in the final snapshot" do + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + llm_eval_response("x = 42") + end) + + {:ok, pid} = + Legion.start_link(ConversationBindingsAgent, + store: StepMemoryStore, + agent_id: "step-conversation-bindings" + ) + + assert {:ok, 42} = Legion.call(pid, "compute") + + [_started, _running, checkpoint, completed] = + MemoryStore.writes("step-conversation-bindings") + + assert checkpoint.conversation_state.bindings == [x: 42] + assert completed.conversation_state.bindings == [x: 42] + end + + test "a :step store persists empty iteration-scoped bindings" do + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + llm_eval_response("x = 42") + end) + + {:ok, pid} = + Legion.start_link(MathAgent, + store: StepMemoryStore, + agent_id: "step-iteration-bindings", + binding_scope: :iteration + ) + + assert {:ok, 42} = Legion.call(pid, "compute") + + [_started, _running, checkpoint, completed] = + MemoryStore.writes("step-iteration-bindings") + + assert checkpoint.conversation_state.bindings == [] + assert completed.conversation_state.bindings == [] + end + + test "a :step store does not add a checkpoint for return" do + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + llm_response("done") + end) + + {:ok, pid} = Legion.start_link(MathAgent, store: StepMemoryStore, agent_id: "step-return") + assert {:ok, "done"} = Legion.call(pid, "compute") + + assert [_started, _running, _completed] = MemoryStore.writes("step-return") + end + test "restores the conversation under a fresh system prompt after a restart" do stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> llm_response("Paris") diff --git a/test/legion/executor_test.exs b/test/legion/executor_test.exs index e230439..307f403 100644 --- a/test/legion/executor_test.exs +++ b/test/legion/executor_test.exs @@ -76,6 +76,13 @@ defmodule Legion.ExecutorTest do {:ok, %ReqLLM.Response{id: "test", model: "test", context: nil, object: object}} end + defp executor_messages(message) do + [ + Legion.Executor.message(:system, "system"), + Legion.Executor.message(:user, message) + ] + end + describe "run/4" do test "returns result for return action" do stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> @@ -180,6 +187,144 @@ defmodule Legion.ExecutorTest do end end + describe "checkpoints" do + test "emits complete checkpoints for continuing and completing eval results" do + test_pid = self() + call_count = :counters.new(1, [:atomics]) + + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + :counters.add(call_count, 1, 1) + + case :counters.get(call_count, 1) do + 1 -> + response(%{ + "action" => "eval_and_continue", + "code" => "x = 10", + "result" => "" + }) + + 2 -> + response(%{ + "action" => "eval_and_complete", + "code" => "x * 2", + "result" => "" + }) + end + end) + + checkpoint = fn state -> + send(test_pid, {:checkpoint, state}) + :ok + end + + assert {:ok, 20, _messages, _bindings} = + Legion.Executor.run( + MathAgent, + executor_messages("compute"), + %{checkpoint: checkpoint} + ) + + assert_received {:checkpoint, + %{ + messages: continuing_messages, + bindings: [x: 10], + execution: %{phase: :awaiting_llm, iteration: 1, retries: 0} + }} + + assert List.last(continuing_messages).type == :eval_result + + assert_received {:checkpoint, + %{ + messages: completing_messages, + bindings: [x: 10], + execution: %{phase: :completing, iteration: 1, retries: 0} + }} + + assert List.last(completing_messages).type == :eval_result + end + + test "emits the current counters after a recoverable error" do + test_pid = self() + call_count = :counters.new(1, [:atomics]) + + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + :counters.add(call_count, 1, 1) + + case :counters.get(call_count, 1) do + 1 -> + response(%{ + "action" => "eval_and_complete", + "code" => "raise \"boom\"", + "result" => "" + }) + + 2 -> + response(%{"action" => "return", "code" => "", "result" => "recovered"}) + end + end) + + checkpoint = fn state -> + send(test_pid, {:checkpoint, state}) + :ok + end + + assert {:ok, "recovered", _messages, []} = + Legion.Executor.run( + MathAgent, + executor_messages("recover"), + %{checkpoint: checkpoint} + ) + + assert_received {:checkpoint, + %{ + messages: messages, + bindings: [], + execution: %{phase: :awaiting_llm, iteration: 0, retries: 1} + }} + + assert List.last(messages).type == :error + refute_received {:checkpoint, _other} + end + + test "does not checkpoint a return action" do + test_pid = self() + + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + response(%{"action" => "return", "code" => "", "result" => "done"}) + end) + + assert {:ok, "done", _messages, []} = + Legion.Executor.run( + MathAgent, + executor_messages("finish"), + %{checkpoint: fn state -> send(test_pid, {:checkpoint, state}) end} + ) + + refute_received {:checkpoint, _state} + end + + test "checkpoint failure exits before the next LLM request" do + call_count = :counters.new(1, [:atomics]) + + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + :counters.add(call_count, 1, 1) + response(%{"action" => "eval_and_continue", "code" => "1 + 1", "result" => ""}) + end) + + reason = + catch_exit( + Legion.Executor.run( + MathAgent, + executor_messages("compute"), + %{checkpoint: fn _state -> :error end} + ) + ) + + assert {:checkpoint_persistence_failed, %MatchError{term: :error}} = reason + assert :counters.get(call_count, 1) == 1 + end + end + describe "result formatting" do test "available variables are listed in the result message" do {:ok, counter} = Agent.start_link(fn -> 0 end) diff --git a/test/legion/store/postgres_test.exs b/test/legion/store/postgres_test.exs index 02b580e..e6b2281 100644 --- a/test/legion/store/postgres_test.exs +++ b/test/legion/store/postgres_test.exs @@ -45,11 +45,22 @@ defmodule Legion.Store.PostgresTest do use Legion.Store.Postgres, repo: Legion.Store.PostgresTest.FakeRepo end + defmodule StepStore do + use Legion.Store.Postgres, + repo: Legion.Store.PostgresTest.FakeRepo, + persistence_frequency: :step + end + setup do start_supervised!(%{id: FakeRepo, start: {FakeRepo, :start_link, []}}) :ok end + test "generated stores expose their configured persistence frequency" do + assert Legion.Store.persistence_frequency(Store) == :turn + assert Legion.Store.persistence_frequency(StepStore) == :step + end + test "save/1 fully inserts every payload field" do payload = %Payload{ agent_id: "user_42", @@ -74,6 +85,23 @@ defmodule Legion.Store.PostgresTest do assert {:ok, ^payload} = Store.get("state-only") end + test "save/1 round trips step execution state" do + execution = %{phase: :awaiting_llm, iteration: 2, retries: 1} + + payload = %Payload{ + agent_id: "step-state", + status: :running, + conversation_state: %{ + messages: [%{role: "user", content: "result"}], + bindings: [x: 42], + execution: execution + } + } + + assert :ok = Store.save(payload) + assert {:ok, ^payload} = Store.get("step-state") + end + test "save/1 partial upsert preserves omitted fields and advances updated_at" do initial = %Payload{ agent_id: "user_42", diff --git a/test/legion/store_test.exs b/test/legion/store_test.exs new file mode 100644 index 0000000..7a97d62 --- /dev/null +++ b/test/legion/store_test.exs @@ -0,0 +1,63 @@ +defmodule Legion.StoreTest do + use ExUnit.Case, async: true + + defmodule DefaultStore 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 StepStore do + @behaviour Legion.Store + + @impl Legion.Store + def persistence_frequency, do: :step + + @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 InvalidStore do + @behaviour Legion.Store + + @impl Legion.Store + def persistence_frequency, do: :often + + @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 + + test "defaults stores without a frequency callback to :turn" do + assert Legion.Store.persistence_frequency(DefaultStore) == :turn + assert Legion.Store.persistence_frequency(nil) == :turn + end + + test "returns a store's declared :step frequency" do + assert Legion.Store.persistence_frequency(StepStore) == :step + end + + test "rejects unsupported frequencies" do + assert_raise ArgumentError, ~r/must return :turn or :step, got: :often/, fn -> + Legion.Store.persistence_frequency(InvalidStore) + end + end +end From ed2f72a517599bdd05cec41050ca2c1c3f6e7da6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tomasz=20Id=C5=BAkowski?= Date: Tue, 21 Jul 2026 15:53:07 +0200 Subject: [PATCH 14/30] Add parent_agent_id and started_at time to AgentRegistry --- lib/legion/agent_server.ex | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/lib/legion/agent_server.ex b/lib/legion/agent_server.ex index 1ab38ba..b92d123 100644 --- a/lib/legion/agent_server.ex +++ b/lib/legion/agent_server.ex @@ -77,7 +77,10 @@ defmodule Legion.AgentServer do Vault.unsafe_put(:parent_agent_id, parent_agent_id) if store, do: Vault.unsafe_put(:store, store) - Registry.register(Legion.AgentRegistry, agent_id, self()) + Registry.register(Legion.AgentRegistry, agent_id, %{ + parent_agent_id: parent_agent_id, + started_at: System.system_time(:millisecond) + }) for tool <- agent_module.tools() do Vault.unsafe_put(tool, agent_module.tool_config(tool)) From afcff80acd1941aaacbdd472baa01d3377706df4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tomasz=20Id=C5=BAkowski?= Date: Thu, 23 Jul 2026 11:31:00 +0200 Subject: [PATCH 15/30] Apply review suggestions and update docs --- .formatter.exs | 3 +- CHANGELOG.md | 11 ++- README.md | 67 +++++++++++++++- lib/legion.ex | 44 ++++++----- lib/legion/agent_server.ex | 15 ++-- lib/legion/store.ex | 17 ++--- lib/legion/store/payload.ex | 2 +- lib/legion/store/postgres.ex | 45 +++++------ lib/legion/store/postgres/migration.ex | 12 +-- lib/legion/telemetry.ex | 4 +- mix.exs | 5 +- mix.lock | 1 + test/integration/step_persistence_test.exs | 89 ++++++++++++++++++++++ test/legion/agent_server_test.exs | 4 +- test/legion/store/postgres_db_test.exs | 54 ++----------- test/legion/store/postgres_test.exs | 9 ++- test/support/postgres_repo.ex | 7 ++ test/test_helper.exs | 7 +- 18 files changed, 263 insertions(+), 133 deletions(-) create mode 100644 test/integration/step_persistence_test.exs create mode 100644 test/support/postgres_repo.ex diff --git a/.formatter.exs b/.formatter.exs index d2cda26..a9a7f0f 100644 --- a/.formatter.exs +++ b/.formatter.exs @@ -1,4 +1,5 @@ # Used by "mix format" [ - inputs: ["{mix,.formatter}.exs", "{config,lib,test}/**/*.{ex,exs}"] + inputs: ["{mix,.formatter}.exs", "{config,lib,test}/**/*.{ex,exs}"], + locals_without_parens: [field: 2, field: 3] ] diff --git a/CHANGELOG.md b/CHANGELOG.md index 74e9c0d..4340afc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,10 +4,13 @@ ### Changes -- Add `Legion.Store` behaviour for persisting conversations across restarts. Pass `store:` and `agent_id:` to `Legion.start_link/2`; snapshots (messages + bindings) are saved after every completed turn, before the caller receives its reply -- Add `Legion.Store.Postgres`, a ready-made store adapter that reuses your Ecto repo (`use Legion.Store.Postgres, repo: MyApp.Repo`) without adding Ecto as a dependency -- `Legion.Store.Postgres.Migration` installs a trigger that `pg_notify`s the table's channel with the agent_id on every write, and generated stores expose `__repo__/0` and `__table__/0`, so consumers (LegionWeb) can follow store changes live -- Add optional `Legion.Store.save_status/2` callback, called with `:running` when a turn starts and `:idle` after its snapshot is saved; `Legion.Store.Postgres` implements it via a `status` column, so a `'running'` status under a dead pid identifies a conversation that crashed mid-turn +- 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 +- 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.Postgres.Migration` helpers and `pg_notify` notifications for inserts and updates; generated stores expose `__repo__/0` and `__table__/0` for database-backed consumers such as LegionWeb - Bump the default model from `openai:gpt-4o-mini` to `openai:gpt-5.4` - `Legion.Tools.HumanTool.ask/1` now raises when called under `eval_and_complete` - the turn would end as soon as the code returns, silently discarding the human's answer; the error feeds back to the model, which retries under `eval_and_continue` diff --git a/README.md b/README.md index a484bae..64329f9 100644 --- a/README.md +++ b/README.md @@ -93,7 +93,7 @@ A traditional agent would need a separate LLM call for each filter decision and - **Tools are just modules** - `use Legion.Tool` on any module to expose it. The LLM reads your source code and calls your functions. No schemas to write, no wrappers - reuse existing app logic directly. - **Authorization via Vault** - Set auth context before the agent starts, validate inside tools at runtime. LLM-generated code never touches credentials. See [Vault](https://github.com/dimamik/vault). - **Long-lived agents** - Start agents with `Legion.start_link/2` and message them with `call/2` and `cast/2`, just like a GenServer. Variables can persist across turns with `binding_scope: :conversation`. -- **Persistence** - Make conversations survive crashes, restarts, and deploys with `use Legion.Store.Postgres, repo: MyApp.Repo` (reuses your Ecto repo, one table), or implement the two-callback `Legion.Store` behaviour for any other storage. Snapshots are saved after every completed turn, before the caller sees its reply - a reply is a commit receipt. +- **Persistence** - Persist conversations across process and application restarts with `use Legion.Store.Postgres, repo: MyApp.Repo`, or implement `Legion.Store` for any other storage. Legion saves the incoming user message before execution and the final state before replying; stores can also opt into step checkpoints for intermediate results and recoverable errors. - **Multi-agent orchestration** - Agents delegate to other agents via the built-in `AgentTool`. Fan out with `parallel/2`, chain with `pipeline/1`. Sub-agents are linked processes - when a parent dies, children stop too. - **Human in the loop** - The built-in `HumanTool` pauses agent execution until a human responds. It's just message passing - your handler receives a question and sends back an answer. - **Structured output** - Define a JSON Schema via `output_schema/0` to get typed, validated responses. Or skip it and work with plain text. @@ -138,6 +138,61 @@ config :req_llm, openai_api_key: System.get_env("OPENAI_API_KEY") Legion.cast(pid, "Also check the reviews") ``` +## Conversation Persistence + +Persist conversations with the built-in Postgres adapter or your own implementation of `Legion.Store`. The Postgres adapter reuses your application's Ecto repo: + +```elixir +defmodule MyApp.AgentStore do + use Legion.Store.Postgres, repo: MyApp.Repo +end +``` + +Create its table with the migration helper: + +```elixir +defmodule MyApp.Repo.Migrations.AddLegionAgents do + use Ecto.Migration + + def up, do: Legion.Store.Postgres.Migration.up() + def down, do: Legion.Store.Postgres.Migration.down() +end +``` + +Use a stable `agent_id` to continue the same conversation after a process or application restart: + +```elixir +{:ok, pid} = + Legion.start_link(MyApp.AssistantAgent, + store: MyApp.AgentStore, + agent_id: "user_42:chat_7" + ) + +{:ok, response} = Legion.call(pid, "Remember that my budget is $100") + +# Later, after the original process has stopped +{:ok, pid} = Legion.resume("user_42:chat_7", store: MyApp.AgentStore) +``` + +If you omit `agent_id`, Legion generates one. Save it if you want to resume the conversation later: + +```elixir +{:ok, pid} = Legion.start_link(MyApp.AssistantAgent, store: MyApp.AgentStore) +agent_id = Legion.get_agent_id(pid) +``` + +Stores persist at turn boundaries by default. To also checkpoint intermediate eval results, recoverable errors, bindings, and executor progress: + +```elixir +defmodule MyApp.AgentStore do + use Legion.Store.Postgres, + repo: MyApp.Repo, + persistence_frequency: :step +end +``` + +Step persistence records the latest recoverable state, but Legion does not automatically continue an interrupted turn. + ## Multi-Agent Systems Agents orchestrate other agents through the built-in `AgentTool`: @@ -211,6 +266,16 @@ Your handler receives `{:human_request, ref, from_pid, question, meta}` and repl ## Configuration +Set a global store to persist agents by default: + +```elixir +config :legion, :store, MyApp.AgentStore +``` + +A `store:` passed to `Legion.start_link/2` overrides the global store. With a global store configured, pass only `agent_id:` to select an existing conversation; if you omit it, Legion generates one. + +Configure model and runtime options separately: + ```elixir config :legion, :config, %{ model: "openai:gpt-5.4", diff --git a/lib/legion.ex b/lib/legion.ex index dd63219..369054a 100644 --- a/lib/legion.ex +++ b/lib/legion.ex @@ -105,17 +105,17 @@ defmodule Legion do end @doc """ - Whether `pid` - typically the one recorded in a persisted run's metadata - - is alive on this node. Accepts `nil` and returns `false`. + Returns whether `pid` is alive on this node. Non-pid values return `false`. - A stored pid outlives the VM that wrote it, so after a restart the check is - best-effort: a recycled pid value can collide with an unrelated live - process. + Use `lookup/1` to resolve an `agent_id` to its currently registered process + before checking it. ## Examples - run = MyApp.AgentStore.get_run("user_42:chat_7") - Legion.running?(run.pid) + case Legion.lookup("user_42:chat_7") do + {:ok, pid} -> Legion.running?(pid) + :error -> false + end """ def running?(pid) when is_pid(pid) do node(pid) == node() and Process.alive?(pid) @@ -149,13 +149,14 @@ defmodule Legion do @doc """ Resumes a persisted conversation. - Checks the persisted run exists, then returns the process registered for - `agent_id` if the agent is already running. If no process is registered, - starts the agent again under the same `agent_id`, so it reloads its snapshot - from the store. `opts` are passed through to `start_link/2`. + 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`. - Requires a store implementing `c:Legion.Store.get_run/1` - pass `:store` or - configure one globally. Raises if the store has no run for `agent_id`. + Pass `:store` or configure one globally. Raises if `get/1` does not return a + `Legion.Store.Payload` containing an agent module for `agent_id`. ## Examples @@ -168,21 +169,26 @@ defmodule Legion do raise ArgumentError, "resume/2 requires a :store - pass one or set `config :legion, :store, MyStore`" - run = - store.get_run(agent_id) || - raise ArgumentError, - "no run recorded for agent_id #{inspect(agent_id)} in #{inspect(store)}" + agent_module = + case store.get(agent_id) do + {:ok, %Legion.Store.Payload{agent_module: agent_module}} when not is_nil(agent_module) -> + agent_module + + _ -> + raise ArgumentError, + "no run recorded for agent_id #{inspect(agent_id)} in #{inspect(store)}" + end case lookup(agent_id) do {:ok, pid} -> if running?(pid) do {:ok, pid} else - start_link(run.agent_module, Keyword.put(opts, :agent_id, agent_id)) + start_link(agent_module, Keyword.put(opts, :agent_id, agent_id)) end :error -> - start_link(run.agent_module, Keyword.put(opts, :agent_id, agent_id)) + start_link(agent_module, Keyword.put(opts, :agent_id, agent_id)) end end diff --git a/lib/legion/agent_server.ex b/lib/legion/agent_server.ex index b92d123..10a2053 100644 --- a/lib/legion/agent_server.ex +++ b/lib/legion/agent_server.ex @@ -77,10 +77,11 @@ defmodule Legion.AgentServer do Vault.unsafe_put(:parent_agent_id, parent_agent_id) if store, do: Vault.unsafe_put(:store, store) - Registry.register(Legion.AgentRegistry, agent_id, %{ - parent_agent_id: parent_agent_id, - started_at: System.system_time(:millisecond) - }) + {: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)) @@ -90,7 +91,7 @@ defmodule Legion.AgentServer do Telemetry.emit( [:legion, :agent, :started], - %{system_time: System.system_time()}, + %{system_time: NaiveDateTime.utc_now()}, %{agent: agent_module} ) @@ -117,7 +118,7 @@ defmodule Legion.AgentServer do persist(state, agent_module: state.agent_module, parent_agent_id: parent_agent_id, - started_at: System.system_time(:millisecond) + started_at: NaiveDateTime.utc_now() )} end @@ -125,7 +126,7 @@ defmodule Legion.AgentServer do def terminate(_reason, state) do Telemetry.emit( [:legion, :agent, :stopped], - %{system_time: System.system_time()}, + %{system_time: NaiveDateTime.utc_now()}, %{agent: state.agent_module} ) end diff --git a/lib/legion/store.ex b/lib/legion/store.ex index af3e74e..31ecc00 100644 --- a/lib/legion/store.ex +++ b/lib/legion/store.ex @@ -7,7 +7,7 @@ defmodule Legion.Store do {:ok, pid} = Legion.start_link(AssistantAgent, store: MyApp.AgentStore, agent_id: "user_42") - On start, the agent calls `c:get/1` and resumes from the returned + On start, the agent calls `get/1` and resumes from the returned `Legion.Store.Payload` when it has a `:conversation_state`. The system prompt is regenerated for every start, so prompt and tool changes apply to restored conversations. @@ -18,7 +18,7 @@ defmodule Legion.Store do restart, or deploy. A store can opt into step persistence by implementing - `c:persistence_frequency/0` and returning `:step`. Legion then also saves + `persistence_frequency/0` and returning `:step`. Legion then also saves after intermediate user-role messages, including eval results and recoverable errors. Each step save contains the complete conversation and executor state at that checkpoint. @@ -66,11 +66,11 @@ defmodule Legion.Store do resume an existing conversation. Two agents started under the same id race onto the same row, so route each conversation to a single process. - ## Required persistence + ## Required callbacks - Stores must implement `c:get/1` and `c:save/1`. + Stores must implement `get/1`, `list/1`, and `save/1`. - `c:save/1` receives a `Legion.Store.Payload`. Its `:conversation_state` is a + `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` @@ -101,12 +101,11 @@ defmodule Legion.Store do ## Reading conversations - `c:get/1` returns the persisted conversation for one `agent_id`, or `:error` + `get/1` returns the persisted conversation for one `agent_id`, or `:error` when the store has no row for that id. - The optional `c:list/1` callback returns persisted conversations newest first - for consumers that rebuild a view of past conversations from the store - alone. + `list/1` returns persisted conversations newest first for consumers that + rebuild a view of past conversations from the store alone. """ alias Legion.Store.Payload diff --git a/lib/legion/store/payload.ex b/lib/legion/store/payload.ex index 0f9dbf9..46ac048 100644 --- a/lib/legion/store/payload.ex +++ b/lib/legion/store/payload.ex @@ -29,7 +29,7 @@ defmodule Legion.Store.Payload do agent_module: module() | nil, parent_agent_id: Legion.Store.agent_id() | nil, status: status() | nil, - started_at: integer() | nil, + started_at: NaiveDateTime.t() | nil, conversation_state: state() | nil } end diff --git a/lib/legion/store/postgres.ex b/lib/legion/store/postgres.ex index 524acfd..0ac301d 100644 --- a/lib/legion/store/postgres.ex +++ b/lib/legion/store/postgres.ex @@ -2,20 +2,15 @@ defmodule Legion.Store.Postgres do @moduledoc """ A ready-made `Legion.Store` backed by Postgres, through your existing Ecto repo. - Legion depends on Ecto. The generated store defines an `Ecto.Schema` and - uses your application's existing `Ecto.Repo` for reads and partial upserts. - The repo must use `Ecto.Adapters.Postgres`. + The generated store uses Ecto for its schema, reads, and partial upserts. + Legion declares `ecto_sql` as an optional dependency; applications using + this adapter must provide an `Ecto.Repo` backed by `Ecto.Adapters.Postgres` + and include Postgrex. ## Usage Define a Postgres-backed Ecto repo and a store module that uses it: - defmodule MyApp.Repo do - use Ecto.Repo, - otp_app: :my_app, - adapter: Ecto.Adapters.Postgres - end - defmodule MyApp.AgentStore do use Legion.Store.Postgres, repo: MyApp.Repo end @@ -51,17 +46,18 @@ defmodule Legion.Store.Postgres do blobs - readable only from Elixir, one row per conversation, upserted on every save. Step snapshots therefore require no additional migration. - `c:Legion.Store.save/1` performs partial upserts, so the same row carries the + `save/1` performs partial upserts, so the same row carries the conversation state and identity: `agent_module` (in `inspect/1` form, e.g. `"MyApp.ResearchAgent"`), `parent_agent_id` linking a sub-agent to the - conversation that spawned it, and `started_at` in milliseconds. Omitted - payload fields preserve their existing values. + conversation that spawned it, and `started_at` as a UTC `NaiveDateTime` + stored with microsecond precision. Omitted payload fields preserve their + existing values. The row's `status` flips to `'running'` when a turn starts and back to `'idle'` when it completes. Step writes update only the conversation state, leaving the running status unchanged. - `c:Legion.Store.list/1` and `c:Legion.Store.get/1` read persisted + `list/1` and `get/1` read persisted conversations back from the same table. The migration also installs a trigger that `pg_notify`s the table's channel @@ -90,13 +86,13 @@ defmodule Legion.Store.Postgres do @primary_key {:agent_id, :string, autogenerate: false} schema unquote(table) do - field(:agent_module, :string) - field(:parent_agent_id, :string) - field(:status, :string) - field(:started_at, :integer) - field(:conversation_state, :binary) - field(:inserted_at, :utc_datetime_usec) - field(:updated_at, :utc_datetime_usec) + field :agent_module, :string + field :parent_agent_id, :string + field :status, :string + field :started_at, :naive_datetime_usec + field :conversation_state, :binary + field :inserted_at, :naive_datetime_usec + field :updated_at, :naive_datetime_usec end end @@ -129,7 +125,7 @@ defmodule Legion.Store.Postgres do payload |> Postgres.encode_data() |> Map.reject(fn {_k, v} -> is_nil(v) end) - |> Map.put(:updated_at, DateTime.utc_now()) + |> Map.put(:updated_at, NaiveDateTime.utc_now()) update_columns = attrs |> Map.delete(:agent_id) |> Map.keys() @@ -156,13 +152,13 @@ defmodule Legion.Store.Postgres do payload |> Map.from_struct() |> Map.update(:conversation_state, nil, fn state -> - if is_nil(state), do: nil, else: :erlang.term_to_binary(state) + if not is_nil(state), do: :erlang.term_to_binary(state) end) |> Map.update(:agent_module, nil, fn module -> - if is_nil(module), do: nil, else: inspect(module) + if not is_nil(module), do: inspect(module) end) |> Map.update(:status, nil, fn status -> - if is_nil(status), do: nil, else: Atom.to_string(status) + if not is_nil(status), do: Atom.to_string(status) end) end @@ -193,5 +189,4 @@ defmodule Legion.Store.Postgres do defp decode_status("running"), do: :running defp decode_status("idle"), do: :idle - defp decode_status(nil), do: nil end diff --git a/lib/legion/store/postgres/migration.ex b/lib/legion/store/postgres/migration.ex index 51e7947..586a1d0 100644 --- a/lib/legion/store/postgres/migration.ex +++ b/lib/legion/store/postgres/migration.ex @@ -31,8 +31,8 @@ defmodule Legion.Store.Postgres.Migration do `down/1` to rolling everything back. """ - # Legion does not depend on Ecto - these run inside the host app's - # migrations, where Ecto.Migration is available. + # Ecto SQL is optional for Legion. These helpers run inside the host + # application's migrations, where Ecto.Migration is available. @compile {:no_warn_undefined, Ecto.Migration} @default_table "legion_agents" @@ -107,11 +107,11 @@ defmodule Legion.Store.Postgres.Migration do agent_id text PRIMARY KEY, agent_module text, parent_agent_id text, - status text, - started_at bigint, + status text NOT NULL DEFAULT 'idle', + started_at timestamp, conversation_state bytea, - inserted_at timestamptz NOT NULL DEFAULT now(), - updated_at timestamptz NOT NULL DEFAULT now() + inserted_at timestamp NOT NULL DEFAULT now(), + updated_at timestamp NOT NULL DEFAULT now() ) """, """ diff --git a/lib/legion/telemetry.ex b/lib/legion/telemetry.ex index dcc9fff..4c41174 100644 --- a/lib/legion/telemetry.ex +++ b/lib/legion/telemetry.ex @@ -7,7 +7,7 @@ defmodule Legion.Telemetry do ## Agent Lifecycle Events - `[:legion, :agent, :started]` — agent process finished `init/1` - - Measurements: `%{system_time: integer}` + - Measurements: `%{system_time: NaiveDateTime.t()}` - Metadata: `%{agent: module, agent_id: term, parent_agent_id: term}` - `agent_id` names the conversation — stable across restarts, so a resumed conversation emits under the same id. @@ -17,7 +17,7 @@ defmodule Legion.Telemetry do since GenServer does not call `terminate/2` on init failure. - `[:legion, :agent, :stopped]` — agent process terminated via `terminate/2` - - Measurements: `%{system_time: integer}` + - Measurements: `%{system_time: NaiveDateTime.t()}` - Metadata: `%{agent: module, agent_id: term}` (plus `parent_agent_id` when the parent's run is still on the process Vault) diff --git a/mix.exs b/mix.exs index 05ad39d..de587b4 100644 --- a/mix.exs +++ b/mix.exs @@ -54,8 +54,7 @@ defmodule Legion.MixProject do Legion.Agent, Legion.Tool, Legion.Store, - Legion.Store.Conversation, - ~r/^Legion\.Store\.Conversation\./, + Legion.Store.Payload, Legion.Store.Postgres, Legion.Store.Postgres.Migration ], @@ -67,7 +66,7 @@ defmodule Legion.MixProject do defp deps do [ - {:ecto, "~> 3.13"}, + {:ecto_sql, "~> 3.13", optional: true}, {:req_llm, "~> 1.2"}, {:vault, "~> 0.2"}, {:jason, "~> 1.4"}, diff --git a/mix.lock b/mix.lock index e7fb010..f98a619 100644 --- a/mix.lock +++ b/mix.lock @@ -8,6 +8,7 @@ "dotenvy": {:hex, :dotenvy, "1.1.1", "00e318f3c51de9fafc4b48598447e386f19204dc18ca69886905bb8f8b08b667", [:mix], [], "hexpm", "c8269471b5701e9e56dc86509c1199ded2b33dce088c3471afcfef7839766d8e"}, "earmark_parser": {:hex, :earmark_parser, "1.4.44", "f20830dd6b5c77afe2b063777ddbbff09f9759396500cdbe7523efd58d7a339c", [:mix], [], "hexpm", "4778ac752b4701a5599215f7030989c989ffdc4f6df457c5f36938cc2d2a2750"}, "ecto": {:hex, :ecto, "3.13.6", "352135b474f91d1ab99a1b502171d207e9db60421c9e3d0ecab4c7ab96b24d14", [:mix], [{:decimal, "~> 2.0 or ~> 3.0", [hex: :decimal, repo: "hexpm", optional: false]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: true]}, {:telemetry, "~> 0.4 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "8afa059bc16cd2c94739ec0a11e3e5df69d828125119109bef35f20a21a76af2"}, + "ecto_sql": {:hex, :ecto_sql, "3.13.5", "2f8282b2ad97bf0f0d3217ea0a6fff320ead9e2f8770f810141189d182dc304e", [:mix], [{:db_connection, "~> 2.4.1 or ~> 2.5", [hex: :db_connection, repo: "hexpm", optional: false]}, {:ecto, "~> 3.13.0", [hex: :ecto, repo: "hexpm", optional: false]}, {:myxql, "~> 0.7", [hex: :myxql, repo: "hexpm", optional: true]}, {:postgrex, "~> 0.19 or ~> 1.0", [hex: :postgrex, repo: "hexpm", optional: true]}, {:tds, "~> 2.1.1 or ~> 2.2", [hex: :tds, repo: "hexpm", optional: true]}, {:telemetry, "~> 0.4.0 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "aa36751f4e6a2b56ae79efb0e088042e010ff4935fc8684e74c23b1f49e25fdc"}, "ex_aws_auth": {:hex, :ex_aws_auth, "1.3.1", "3963992d6f7cb251b53573603c3615cec70c3f4d86199fdb865ff440295ef7a4", [:mix], [{:jason, "~> 1.4", [hex: :jason, repo: "hexpm", optional: true]}, {:req, "~> 0.5", [hex: :req, repo: "hexpm", optional: true]}], "hexpm", "025793aa08fa419aabdb652db60edbdb2e12346bd447988a1bb5854c4dd64903"}, "ex_doc": {:hex, :ex_doc, "0.40.2", "f50edec428c4b0a457a167de42414c461122a3585a99515a69d09fff19e5597e", [:mix], [{:earmark_parser, "~> 1.4.44", [hex: :earmark_parser, repo: "hexpm", optional: false]}, {:makeup_c, ">= 0.1.0", [hex: :makeup_c, repo: "hexpm", optional: true]}, {:makeup_elixir, "~> 0.14 or ~> 1.0", [hex: :makeup_elixir, repo: "hexpm", optional: false]}, {:makeup_erlang, "~> 0.1 or ~> 1.0", [hex: :makeup_erlang, repo: "hexpm", optional: false]}, {:makeup_html, ">= 0.1.0", [hex: :makeup_html, repo: "hexpm", optional: true]}], "hexpm", "4fa426e2beb47854a162e2c488727fdec51cd4692e319b23810c2804cb1a40fe"}, "file_system": {:hex, :file_system, "1.1.1", "31864f4685b0148f25bd3fbef2b1228457c0c89024ad67f7a81a3ffbc0bbad3a", [:mix], [], "hexpm", "7a15ff97dfe526aeefb090a7a9d3d03aa907e100e262a0f8f7746b78f8f87a5d"}, diff --git a/test/integration/step_persistence_test.exs b/test/integration/step_persistence_test.exs new file mode 100644 index 0000000..07cabe8 --- /dev/null +++ b/test/integration/step_persistence_test.exs @@ -0,0 +1,89 @@ +defmodule Legion.Integration.StepPersistenceTest do + use ExUnit.Case, async: false + use Mimic + + alias Legion.Store.Payload + alias Legion.Test.Support.{MathAgent, PostgresRepo} + + @moduletag :integration + + defmodule StepStore do + use Legion.Store.Postgres, + repo: Legion.Test.Support.PostgresRepo, + persistence_frequency: :step + end + + setup :set_mimic_global + + setup do + PostgresRepo.query!("TRUNCATE legion_agents") + :ok + end + + test "persists the recoverable turn state before execution advances" do + agent_id = "step-persistence-integration" + test_pid = self() + request_count = :counters.new(1, [:atomics]) + + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + :counters.add(request_count, 1, 1) + + case :counters.get(request_count, 1) do + 1 -> + send(test_pid, {:before_first_request, StepStore.get(agent_id)}) + response("eval_and_continue", "x = 42", "") + + 2 -> + send(test_pid, {:before_second_request, StepStore.get(agent_id)}) + response("return", "", "done") + end + end) + + {:ok, pid} = Legion.start_link(MathAgent, store: StepStore, agent_id: agent_id) + + assert {:ok, "done"} = Legion.call(pid, "compute") + + assert_received {:before_first_request, + {:ok, + %Payload{ + status: :running, + conversation_state: initial_state + }}} + + assert Enum.map(initial_state.messages, & &1.type) == [:user] + assert initial_state.bindings == [] + refute Map.has_key?(initial_state, :execution) + + assert_received {:before_second_request, + {:ok, + %Payload{ + status: :running, + conversation_state: checkpoint + }}} + + assert Enum.map(checkpoint.messages, & &1.type) == [:user, :assistant, :eval_result] + assert checkpoint.bindings == [x: 42] + + assert checkpoint.execution == %{ + phase: :awaiting_llm, + iteration: 1, + retries: 0 + } + + assert {:ok, %Payload{status: :idle, conversation_state: completed}} = + StepStore.get(agent_id) + + assert completed.bindings == [] + refute Map.has_key?(completed, :execution) + end + + defp response(action, code, result) do + {:ok, + %ReqLLM.Response{ + id: "test", + model: "test", + context: nil, + object: %{"action" => action, "code" => code, "result" => result} + }} + end +end diff --git a/test/legion/agent_server_test.exs b/test/legion/agent_server_test.exs index 0b7c816..98fb8e6 100644 --- a/test/legion/agent_server_test.exs +++ b/test/legion/agent_server_test.exs @@ -567,7 +567,7 @@ defmodule Legion.AgentServerTest do conversation_state: nil } = started - assert is_integer(started_at) + assert is_struct(started_at, NaiveDateTime) assert %Payload{ agent_id: "payloads", @@ -898,7 +898,7 @@ defmodule Legion.AgentServerTest do assert run.agent_id == "meta" assert run.agent_module == MathAgent assert run.parent_agent_id == nil - assert is_integer(run.started_at) + assert is_struct(run.started_at, NaiveDateTime) end test "registers the agent pid by agent_id" do diff --git a/test/legion/store/postgres_db_test.exs b/test/legion/store/postgres_db_test.exs index 2e76fde..4463a35 100644 --- a/test/legion/store/postgres_db_test.exs +++ b/test/legion/store/postgres_db_test.exs @@ -3,49 +3,10 @@ defmodule Legion.Store.PostgresDbTest do use ExUnit.Case, async: false alias Legion.Store.Payload - - defmodule Repo do - @columns ~w(agent_id agent_module parent_agent_id status started_at conversation_state inserted_at updated_at)a - - def get(_schema, agent_id) do - case query!("SELECT #{columns_sql()} FROM legion_agents WHERE agent_id = $1", [agent_id]).rows do - [] -> nil - [row] -> Map.new(Enum.zip(@columns, row)) - end - end - - def insert_all(_schema, [attrs], conflict_target: :agent_id, on_conflict: {:replace, columns}) do - attrs = Map.take(attrs, [:agent_id | columns]) - fields = Map.keys(attrs) - values = Enum.map(fields, &Map.fetch!(attrs, &1)) - - sql = """ - INSERT INTO legion_agents (#{Enum.join(fields, ", ")}) - VALUES (#{placeholders(length(fields))}) - ON CONFLICT (agent_id) DO UPDATE - SET #{updates_sql(columns)} - """ - - %{num_rows: count} = query!(sql, values) - {count, nil} - end - - def query!(sql, params), do: Postgrex.query!(:legion_store_test, sql, params) - - defp columns_sql, do: Enum.join(@columns, ", ") - - defp placeholders(count) do - 1..count - |> Enum.map_join(", ", &"$#{&1}") - end - - defp updates_sql(columns) do - Enum.map_join(columns, ", ", fn column -> "#{column} = EXCLUDED.#{column}" end) - end - end + alias Legion.Test.Support.PostgresRepo, as: Repo defmodule Store do - use Legion.Store.Postgres, repo: Legion.Store.PostgresDbTest.Repo + use Legion.Store.Postgres, repo: Legion.Test.Support.PostgresRepo end setup do @@ -59,7 +20,7 @@ defmodule Legion.Store.PostgresDbTest do agent_module: Legion.Test.Support.MathAgent, parent_agent_id: "parent-1", status: :idle, - started_at: 123, + started_at: ~N[2026-01-01 00:00:00.000000], conversation_state: %{messages: [%{role: "user", content: "hi"}], bindings: [x: 42]} } @@ -74,7 +35,8 @@ defmodule Legion.Store.PostgresDbTest do } assert :ok = Store.save(payload) - assert {:ok, ^payload} = Store.get("state-only") + assert {:ok, stored} = Store.get("state-only") + assert stored == %{payload | status: :idle} end test "save/1 partial upsert preserves omitted fields and advances updated_at" do @@ -83,7 +45,7 @@ defmodule Legion.Store.PostgresDbTest do agent_module: Legion.Test.Support.MathAgent, parent_agent_id: "parent-1", status: :running, - started_at: 123, + started_at: ~N[2026-01-01 00:00:00.000000], conversation_state: %{messages: [%{role: "user", content: "hi"}], bindings: [x: 42]} } @@ -100,14 +62,14 @@ defmodule Legion.Store.PostgresDbTest do agent_module: Legion.Test.Support.MathAgent, parent_agent_id: "parent-1", status: :idle, - started_at: 123, + started_at: ~N[2026-01-01 00:00:00.000000], conversation_state: %{messages: [%{role: "user", content: "hi"}], bindings: [x: 42]} }} = Store.get("user_42") %{rows: [[updated_at]]} = Repo.query!("SELECT updated_at FROM legion_agents WHERE agent_id = $1", ["user_42"]) - assert DateTime.compare(updated_at, previous_updated_at) == :gt + assert NaiveDateTime.compare(updated_at, previous_updated_at) == :gt end test "a payload cannot be constructed without agent_id" do diff --git a/test/legion/store/postgres_test.exs b/test/legion/store/postgres_test.exs index e6b2281..e1801b9 100644 --- a/test/legion/store/postgres_test.exs +++ b/test/legion/store/postgres_test.exs @@ -12,6 +12,8 @@ defmodule Legion.Store.PostgresTest do Agent.get(__MODULE__, &Map.get(&1.rows, agent_id)) end + def all(_query), do: Agent.get(__MODULE__, &Map.values(&1.rows)) + def insert_all(_schema, [attrs], conflict_target: :agent_id, on_conflict: {:replace, columns}) do Agent.update(__MODULE__, fn state -> row = @@ -32,7 +34,7 @@ defmodule Legion.Store.PostgresTest do agent_id: agent_id, agent_module: nil, parent_agent_id: nil, - status: nil, + status: "idle", started_at: nil, conversation_state: nil, inserted_at: nil, @@ -81,8 +83,9 @@ defmodule Legion.Store.PostgresTest do conversation_state: %{messages: [%{role: "user", content: "hi"}], bindings: []} } + expected_payload = %{payload | status: :idle} assert :ok = Store.save(payload) - assert {:ok, ^payload} = Store.get("state-only") + assert {:ok, ^expected_payload} = Store.get("state-only") end test "save/1 round trips step execution state" do @@ -126,7 +129,7 @@ defmodule Legion.Store.PostgresTest do conversation_state: %{messages: [%{role: "user", content: "hi"}], bindings: [x: 42]} }} = Store.get("user_42") - assert DateTime.compare(FakeRepo.run("user_42").updated_at, previous_updated_at) == :gt + assert NaiveDateTime.compare(FakeRepo.run("user_42").updated_at, previous_updated_at) == :gt end test "a payload cannot be constructed without agent_id" do diff --git a/test/support/postgres_repo.ex b/test/support/postgres_repo.ex new file mode 100644 index 0000000..297a81a --- /dev/null +++ b/test/support/postgres_repo.ex @@ -0,0 +1,7 @@ +defmodule Legion.Test.Support.PostgresRepo do + @moduledoc false + + use Ecto.Repo, + otp_app: :legion, + adapter: Ecto.Adapters.Postgres +end diff --git a/test/test_helper.exs b/test/test_helper.exs index efed6d6..404c3ef 100644 --- a/test/test_helper.exs +++ b/test/test_helper.exs @@ -4,8 +4,7 @@ Legion.Telemetry.attach_default_logger() # Shared connection and schema for the Legion.Store.Postgres database tests. {:ok, _} = - Postgrex.start_link( - name: :legion_store_test, + Legion.Test.Support.PostgresRepo.start_link( hostname: System.get_env("POSTGRES_HOST", "localhost"), port: String.to_integer(System.get_env("POSTGRES_PORT", "5432")), username: System.get_env("POSTGRES_USER", "postgres"), @@ -14,11 +13,11 @@ Legion.Telemetry.attach_default_logger() ) for sql <- Legion.Store.Postgres.Migration.down_sql(1, "legion_agents") do - Postgrex.query!(:legion_store_test, sql, []) + Legion.Test.Support.PostgresRepo.query!(sql) end for sql <- Legion.Store.Postgres.Migration.up_sql(1, "legion_agents") do - Postgrex.query!(:legion_store_test, sql, []) + Legion.Test.Support.PostgresRepo.query!(sql) end ExUnit.start(exclude: [:integration]) From cee0be8c36c8d098a40ed4c6b1b64c64f2428bc9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tomasz=20Id=C5=BAkowski?= Date: Thu, 23 Jul 2026 13:45:57 +0200 Subject: [PATCH 16/30] Move migrations closer to Oban style Change folder structure and make migrations use Ecto --- .formatter.exs | 31 ++++- CHANGELOG.md | 2 +- lib/legion/store/migration/postgres.ex | 104 +++++++++++++++ lib/legion/store/migration/postgres/v01.ex | 49 +++++++ lib/legion/store/postgres/migration.ex | 142 --------------------- mix.exs | 2 +- test/test_helper.exs | 22 +++- 7 files changed, 200 insertions(+), 152 deletions(-) create mode 100644 lib/legion/store/migration/postgres.ex create mode 100644 lib/legion/store/migration/postgres/v01.ex delete mode 100644 lib/legion/store/postgres/migration.ex diff --git a/.formatter.exs b/.formatter.exs index a9a7f0f..da6c471 100644 --- a/.formatter.exs +++ b/.formatter.exs @@ -1,5 +1,34 @@ # Used by "mix format" [ + import_deps: [:ecto_sql], inputs: ["{mix,.formatter}.exs", "{config,lib,test}/**/*.{ex,exs}"], - locals_without_parens: [field: 2, field: 3] + locals_without_parens: [ + add: 2, + add: 3, + add_if_not_exists: 2, + add_if_not_exists: 3, + alter: 2, + create: 1, + create: 2, + create_if_not_exists: 1, + create_if_not_exists: 2, + drop: 1, + drop: 2, + drop_if_exists: 1, + drop_if_exists: 2, + execute: 1, + execute: 2, + field: 2, + field: 3, + modify: 2, + modify: 3, + remove: 1, + remove: 2, + remove: 3, + remove_if_exists: 1, + remove_if_exists: 2, + rename: 2, + rename: 3, + timestamps: 1 + ] ] diff --git a/CHANGELOG.md b/CHANGELOG.md index 4340afc..4b13e4e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,7 +10,7 @@ - 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 - 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.Postgres.Migration` helpers and `pg_notify` notifications for inserts and updates; generated stores expose `__repo__/0` and `__table__/0` for database-backed consumers such as LegionWeb +- 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 - Bump the default model from `openai:gpt-4o-mini` to `openai:gpt-5.4` - `Legion.Tools.HumanTool.ask/1` now raises when called under `eval_and_complete` - the turn would end as soon as the code returns, silently discarding the human's answer; the error feeds back to the model, which retries under `eval_and_continue` diff --git a/lib/legion/store/migration/postgres.ex b/lib/legion/store/migration/postgres.ex new file mode 100644 index 0000000..d56a3c7 --- /dev/null +++ b/lib/legion/store/migration/postgres.ex @@ -0,0 +1,104 @@ +defmodule Legion.Store.Migration.Postgres do + @moduledoc """ + Runs versioned PostgreSQL migrations for `Legion.Store.Postgres`. + + ## Usage + + defmodule MyApp.Repo.Migrations.AddLegionAgents do + use Ecto.Migration + + def up, do: Legion.Store.Migration.Postgres.up() + def down, do: Legion.Store.Migration.Postgres.down() + end + + Migrations are versioned and idempotent. `up/1` applies only versions that + haven't already run and records the resulting version in the table comment. + + When a new Legion release adds a schema version, generate another migration + using the same calls. You can pin that migration to a specific version: + + defmodule MyApp.Repo.Migrations.UpgradeLegionAgentsToV2 do + use Ecto.Migration + + def up, do: Legion.Store.Migration.Postgres.up(version: 2) + def down, do: Legion.Store.Migration.Postgres.down(version: 2) + end + + Rolling back the example above removes version 2 and leaves version 1 applied. + + ## Options + + * `:table` - the table name, defaults to `"legion_agents"`. It must match + the `:table` given to `use Legion.Store.Postgres`. + * `:version` - the target version. `up/1` defaults to the latest version; + `down/1` defaults to rolling back all versions. + """ + + use Ecto.Migration + + @default_table "legion_agents" + @initial_version 1 + @current_version 1 + + def up(opts \\ []) do + opts = Keyword.put_new(opts, :table, @default_table) + initial = migrated_version(opts) + migrated = Keyword.get(opts, :version, @current_version) + + if initial < migrated do + change((initial + 1)..migrated, :up, opts) + end + + :ok + end + + def down(opts \\ []) do + opts = Keyword.put_new(opts, :table, @default_table) + initial = max(migrated_version(opts), @initial_version) + migrated = Keyword.get(opts, :version, @initial_version) + + if initial >= migrated do + change(initial..migrated//-1, :down, opts) + end + + :ok + end + + def migrated_version(opts \\ []) do + opts = Keyword.put_new(opts, :table, @default_table) + table = Keyword.fetch!(opts, :table) + + query = """ + SELECT pg_catalog.obj_description(pg_class.oid, 'pg_class') + FROM pg_class + WHERE pg_class.relname = '#{table}' + """ + + case Ecto.Migration.repo().query(query, [], log: false) do + {:ok, %{rows: [[version]]}} when is_binary(version) -> String.to_integer(version) + _ -> 0 + end + end + + defp change(range, direction, opts) do + for index <- range do + pad_idx = String.pad_leading(to_string(index), 2, "0") + + [__MODULE__, "V#{pad_idx}"] + |> Module.concat() + |> apply(direction, [opts]) + end + + case direction do + :up -> record_version(opts, Enum.max(range)) + :down -> record_version(opts, Enum.min(range) - 1) + end + end + + defp record_version(_opts, 0), do: :ok + + defp record_version(opts, version) do + table = Keyword.get(opts, :table, @default_table) + Ecto.Migration.execute("COMMENT ON TABLE #{table} IS '#{version}'") + end +end diff --git a/lib/legion/store/migration/postgres/v01.ex b/lib/legion/store/migration/postgres/v01.ex new file mode 100644 index 0000000..8d863be --- /dev/null +++ b/lib/legion/store/migration/postgres/v01.ex @@ -0,0 +1,49 @@ +defmodule Legion.Store.Migration.Postgres.V01 do + @moduledoc false + + use Ecto.Migration + + def up(opts) do + table = Keyword.fetch!(opts, :table) + + create_if_not_exists table(table, primary_key: false) do + add :agent_id, :text, primary_key: true + add :agent_module, :text + add :parent_agent_id, :text + add :status, :text, null: false, default: "idle" + add :started_at, :naive_datetime_usec + add :conversation_state, :binary + + add :inserted_at, :naive_datetime_usec, + null: false, + default: fragment("now()") + + add :updated_at, :naive_datetime_usec, + null: false, + default: fragment("now()") + end + + execute """ + CREATE OR REPLACE FUNCTION #{table}_notify() RETURNS trigger AS $$ + BEGIN + PERFORM pg_notify('#{table}', NEW.agent_id); + RETURN NEW; + END; + $$ LANGUAGE plpgsql + """ + + execute "DROP TRIGGER IF EXISTS #{table}_notify ON #{table}" + + execute """ + CREATE TRIGGER #{table}_notify AFTER INSERT OR UPDATE ON #{table} + FOR EACH ROW EXECUTE FUNCTION #{table}_notify() + """ + end + + def down(opts) do + table = Keyword.fetch!(opts, :table) + + drop_if_exists table(table) + execute "DROP FUNCTION IF EXISTS #{table}_notify()" + end +end diff --git a/lib/legion/store/postgres/migration.ex b/lib/legion/store/postgres/migration.ex deleted file mode 100644 index 586a1d0..0000000 --- a/lib/legion/store/postgres/migration.ex +++ /dev/null @@ -1,142 +0,0 @@ -defmodule Legion.Store.Postgres.Migration do - @moduledoc """ - Migration helpers for `Legion.Store.Postgres`. - - ## Usage - - defmodule MyApp.Repo.Migrations.AddLegionAgents do - use Ecto.Migration - - def up, do: Legion.Store.Postgres.Migration.up() - def down, do: Legion.Store.Postgres.Migration.down() - end - - Migrations are versioned and idempotent - `up/1` only runs the versions the - database hasn't seen yet, so when a new Legion release ships schema changes - you generate another migration with the same two calls (optionally pinning - `version:`): - - defmodule MyApp.Repo.Migrations.UpgradeLegionAgentsToV2 do - use Ecto.Migration - - def up, do: Legion.Store.Postgres.Migration.up(version: 2) - def down, do: Legion.Store.Postgres.Migration.down(version: 2) - end - - ## Options - - - `:table` - the table name, defaults to `"legion_agents"`. Must match - the `:table` given to `use Legion.Store.Postgres`. - - `:version` - the target version. `up/1` defaults to the latest version, - `down/1` to rolling everything back. - """ - - # Ecto SQL is optional for Legion. These helpers run inside the host - # application's migrations, where Ecto.Migration is available. - @compile {:no_warn_undefined, Ecto.Migration} - - @default_table "legion_agents" - @initial_version 1 - @current_version 1 - - @doc "Migrates the agents table up to `:version`, defaulting to the latest." - def up(opts \\ []) do - table = table(opts) - version = Keyword.get(opts, :version, @current_version) - migrated = migrated_version(opts) - - if migrated < version do - for step <- (migrated + 1)..version, sql <- List.wrap(up_sql(step, table)) do - Ecto.Migration.execute(sql) - end - - record_version(table, version) - end - - :ok - end - - @doc "Rolls the agents table back down to and including `:version`." - def down(opts \\ []) do - table = table(opts) - version = Keyword.get(opts, :version, @initial_version) - migrated = migrated_version(opts) - - if migrated >= version do - for step <- migrated..version//-1, sql <- List.wrap(down_sql(step, table)) do - Ecto.Migration.execute(sql) - end - - record_version(table, version - 1) - end - - :ok - end - - @doc "Returns the version the database is migrated to, `0` when the table is absent." - def migrated_version(opts \\ []) do - # The version lives in the table's comment, read directly so it reflects - # the database before this migration's queued commands run. - query = """ - SELECT pg_catalog.obj_description(pg_class.oid, 'pg_class') - FROM pg_class - WHERE pg_class.relname = '#{table(opts)}' - """ - - case Ecto.Migration.repo().query(query, [], log: false) do - {:ok, %{rows: [[version]]}} when is_binary(version) -> String.to_integer(version) - _ -> 0 - end - end - - defp record_version(_table, 0), do: :ok - - defp record_version(table, version) do - Ecto.Migration.execute("COMMENT ON TABLE #{table} IS '#{version}'") - end - - # The trigger notifies the table's channel with the agent_id on every - # write, so consumers (e.g. LegionWeb.Source.Listener) can refresh from the - # database instead of capturing telemetry. Postgres has no CREATE TRIGGER - # IF NOT EXISTS, so the trigger is dropped first to keep the step idempotent. - @doc false - def up_sql(1, table) do - [ - """ - CREATE TABLE IF NOT EXISTS #{table} ( - agent_id text PRIMARY KEY, - agent_module text, - parent_agent_id text, - status text NOT NULL DEFAULT 'idle', - started_at timestamp, - conversation_state bytea, - inserted_at timestamp NOT NULL DEFAULT now(), - updated_at timestamp NOT NULL DEFAULT now() - ) - """, - """ - CREATE OR REPLACE FUNCTION #{table}_notify() RETURNS trigger AS $$ - BEGIN - PERFORM pg_notify('#{table}', NEW.agent_id); - RETURN NEW; - END; - $$ LANGUAGE plpgsql - """, - "DROP TRIGGER IF EXISTS #{table}_notify ON #{table}", - """ - CREATE TRIGGER #{table}_notify AFTER INSERT OR UPDATE ON #{table} - FOR EACH ROW EXECUTE FUNCTION #{table}_notify() - """ - ] - end - - @doc false - def down_sql(1, table) do - [ - "DROP TABLE IF EXISTS #{table}", - "DROP FUNCTION IF EXISTS #{table}_notify()" - ] - end - - defp table(opts), do: Keyword.get(opts, :table, @default_table) -end diff --git a/mix.exs b/mix.exs index de587b4..89117df 100644 --- a/mix.exs +++ b/mix.exs @@ -56,7 +56,7 @@ defmodule Legion.MixProject do Legion.Store, Legion.Store.Payload, Legion.Store.Postgres, - Legion.Store.Postgres.Migration + Legion.Store.Migration.Postgres ], Runtime: [Legion.AgentServer, Legion.Executor, ~r/^Legion\.Sandbox/], Tools: [~r/^Legion\.Tools\./], diff --git a/test/test_helper.exs b/test/test_helper.exs index 404c3ef..452188c 100644 --- a/test/test_helper.exs +++ b/test/test_helper.exs @@ -2,6 +2,15 @@ Mimic.copy(ReqLLM) Legion.Telemetry.attach_default_logger() +defmodule Legion.Test.Support.LegionAgentsMigration do + use Ecto.Migration + + alias Legion.Store.Migration.Postgres + + def up, do: Postgres.up() + def down, do: Postgres.down() +end + # Shared connection and schema for the Legion.Store.Postgres database tests. {:ok, _} = Legion.Test.Support.PostgresRepo.start_link( @@ -12,12 +21,11 @@ Legion.Telemetry.attach_default_logger() database: System.get_env("POSTGRES_DB", "postgres") ) -for sql <- Legion.Store.Postgres.Migration.down_sql(1, "legion_agents") do - Legion.Test.Support.PostgresRepo.query!(sql) -end - -for sql <- Legion.Store.Postgres.Migration.up_sql(1, "legion_agents") do - Legion.Test.Support.PostgresRepo.query!(sql) -end +Ecto.Migrator.up( + Legion.Test.Support.PostgresRepo, + 20_260_723_100_815, + Legion.Test.Support.LegionAgentsMigration, + log: false +) ExUnit.start(exclude: [:integration]) From 1a2a24a36206b2997bc173bbe3d7e06a3887ef74 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tomasz=20Id=C5=BAkowski?= Date: Thu, 23 Jul 2026 13:52:57 +0200 Subject: [PATCH 17/30] Add 'Migrating Without Ecto' section --- lib/legion/store/migration/postgres.ex | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/lib/legion/store/migration/postgres.ex b/lib/legion/store/migration/postgres.ex index d56a3c7..d6120a7 100644 --- a/lib/legion/store/migration/postgres.ex +++ b/lib/legion/store/migration/postgres.ex @@ -32,6 +32,15 @@ defmodule Legion.Store.Migration.Postgres do the `:table` given to `use Legion.Store.Postgres`. * `:version` - the target version. `up/1` defaults to the latest version; `down/1` defaults to rolling back all versions. + + ## 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. + + 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 + then paste into your migration system of choice. """ use Ecto.Migration From 4b08fab75ff7e111605293ccf0d5d5bc9d42eba3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tomasz=20Id=C5=BAkowski?= Date: Wed, 29 Jul 2026 12:29:23 +0200 Subject: [PATCH 18/30] Revise resume/2 wiring Complete the resume/2 loop by creating a mechanism where not only is an agent's history recovered, but also it actually RESUMES executing, if it was interrupted before completion (switching back to :idle). --- lib/legion.ex | 4 +- lib/legion/agent_server.ex | 49 ++++++++++--- lib/legion/executor.ex | 14 +++- lib/legion/store.ex | 9 +-- lib/legion/store/payload.ex | 6 +- test/integration/step_persistence_test.exs | 4 +- test/legion/agent_server_test.exs | 82 ++++++++++++++++++++-- test/legion/store/postgres_db_test.exs | 24 +++++-- test/legion/store/postgres_test.exs | 24 +++++-- 9 files changed, 179 insertions(+), 37 deletions(-) diff --git a/lib/legion.ex b/lib/legion.ex index 369054a..0e39aa3 100644 --- a/lib/legion.ex +++ b/lib/legion.ex @@ -184,11 +184,11 @@ defmodule Legion do if running?(pid) do {:ok, pid} else - start_link(agent_module, Keyword.put(opts, :agent_id, agent_id)) + start_link(agent_module, Keyword.merge(opts, agent_id: agent_id, start_mode: :resume)) end :error -> - start_link(agent_module, Keyword.put(opts, :agent_id, agent_id)) + start_link(agent_module, Keyword.merge(opts, agent_id: agent_id, start_mode: :resume)) end end diff --git a/lib/legion/agent_server.ex b/lib/legion/agent_server.ex index 10a2053..24e1cf8 100644 --- a/lib/legion/agent_server.ex +++ b/lib/legion/agent_server.ex @@ -21,6 +21,7 @@ defmodule Legion.AgentServer do :store, :agent_id, :persistence_frequency, + :execution, bindings: [] ] @@ -72,6 +73,7 @@ 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) @@ -95,13 +97,16 @@ defmodule Legion.AgentServer do %{agent: agent_module} ) - {saved_messages, saved_bindings} = + {saved_messages, saved_bindings, saved_execution} = 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, execution: execution} + }} -> + {messages, bindings, execution} _no_state -> - {[], []} + {[], [], nil} end state = %__MODULE__{ @@ -111,7 +116,8 @@ defmodule Legion.AgentServer do store: store, agent_id: agent_id, persistence_frequency: persistence_frequency, - bindings: saved_bindings + bindings: saved_bindings, + execution: saved_execution } {:ok, @@ -119,7 +125,16 @@ defmodule Legion.AgentServer do agent_module: state.agent_module, parent_agent_id: parent_agent_id, started_at: NaiveDateTime.utc_now() - )} + ), {:continue, %{start_mode: mode, execution: saved_execution}}} + end + + @impl true + def handle_continue(%{start_mode: :normal}, state), do: {:noreply, state} + + @impl true + def handle_continue(%{start_mode: :resume, execution: execution}, state) do + {_reply, state} = perform_run(state, execution) + {:noreply, state} end @impl true @@ -168,7 +183,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 +192,12 @@ defmodule Legion.AgentServer do |> Map.update!(:messages, &(&1 ++ [Executor.message(:user, content)])) |> persist([:conversation_state, status: :running]) + perform_run(state) + end + + defp perform_run(state, execution \\ nil) do + conversation_scope? = Map.get(state.config, :binding_scope, :turn) == :conversation + checkpoint = if state.persistence_frequency == :step do fn checkpoint -> @@ -191,7 +211,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 +219,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, + execution + ) iterations = Enum.count(messages, &(&1[:role] == "assistant")) - prev_count {result, %{iterations: iterations, status: status, result: value, bindings: bindings}} @@ -242,7 +269,7 @@ 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, execution: nil} end defp persisted_conversation_state(%{ @@ -275,7 +302,7 @@ defmodule Legion.AgentServer do |> Executor.truncate_content(max_length) end - @known_config_keys ~w(binding_scope max_iterations max_message_length max_retries model sandbox_timeout)a + @known_config_keys ~w(binding_scope max_iterations max_message_length max_retries model sandbox_timeout start_mode)a defp resolve_config(agent_module, opts) do app_config = Application.get_env(:legion, :config, %{}) diff --git a/lib/legion/executor.ex b/lib/legion/executor.ex index 8ea9bed..4472f41 100644 --- a/lib/legion/executor.ex +++ b/lib/legion/executor.ex @@ -107,9 +107,19 @@ defmodule Legion.Executor do 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 \\ [], execution \\ nil) do config = Map.merge(@default_config, config) - loop(agent_module, messages, config, 0, 0, bindings) + + case execution do + nil -> + 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 diff --git a/lib/legion/store.ex b/lib/legion/store.ex index 31ecc00..e0a7076 100644 --- a/lib/legion/store.ex +++ b/lib/legion/store.ex @@ -71,10 +71,11 @@ defmodule Legion.Store do 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 + map containing the conversation's `:messages` (without the system prompt), + `:bindings` from evaluated code, and `:execution`. `:execution` is `nil` + 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 diff --git a/lib/legion/store/payload.ex b/lib/legion/store/payload.ex index 46ac048..a026356 100644 --- a/lib/legion/store/payload.ex +++ b/lib/legion/store/payload.ex @@ -20,9 +20,9 @@ defmodule Legion.Store.Payload do retries: non_neg_integer() } @type state :: %{ - required(:messages) => [map()], - required(:bindings) => keyword(), - optional(:execution) => execution() + messages: [map()], + bindings: keyword(), + execution: execution() | nil } @type t :: %__MODULE__{ agent_id: Legion.Store.agent_id(), diff --git a/test/integration/step_persistence_test.exs b/test/integration/step_persistence_test.exs index 07cabe8..5650f14 100644 --- a/test/integration/step_persistence_test.exs +++ b/test/integration/step_persistence_test.exs @@ -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.execution == nil assert_received {:before_second_request, {:ok, @@ -74,7 +74,7 @@ defmodule Legion.Integration.StepPersistenceTest do StepStore.get(agent_id) assert completed.bindings == [] - refute Map.has_key?(completed, :execution) + assert completed.execution == nil end defp response(action, code, result) do diff --git a/test/legion/agent_server_test.exs b/test/legion/agent_server_test.exs index 98fb8e6..4c9161d 100644 --- a/test/legion/agent_server_test.exs +++ b/test/legion/agent_server_test.exs @@ -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() @@ -691,7 +700,11 @@ defmodule Legion.AgentServerTest do assert %Payload{ status: :running, - conversation_state: %{messages: [%{type: :user}], bindings: []} + conversation_state: %{ + messages: [%{type: :user}], + bindings: [], + execution: nil + } } = running assert %Payload{ @@ -705,7 +718,7 @@ defmodule Legion.AgentServerTest do assert %Payload{status: :idle, conversation_state: final_state} = completed assert final_state.bindings == [] - refute Map.has_key?(final_state, :execution) + assert final_state.execution == nil end test "a :step store persists eval_and_complete before the final snapshot" do @@ -728,7 +741,7 @@ defmodule Legion.AgentServerTest do } = checkpoint assert %Payload{status: :idle, conversation_state: final_state} = completed - refute Map.has_key?(final_state, :execution) + assert final_state.execution == nil end test "a :step store persists retry state after an error message" do @@ -935,6 +948,65 @@ defmodule Legion.AgentServerTest do ] = Legion.get_messages(revived) end + 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], + execution: %{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: %{execution: nil}}} + + 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], + execution: %{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: %{execution: nil}}} + + refute_receive :llm_requested, 100 + 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) diff --git a/test/legion/store/postgres_db_test.exs b/test/legion/store/postgres_db_test.exs index 4463a35..3eb4866 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], + execution: nil + } } 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: [], + execution: nil + } } 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], + execution: nil + } } 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], + execution: nil + } }} = 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..bf8821b 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], + execution: nil + } } 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: [], + execution: nil + } } expected_payload = %{payload | status: :idle} @@ -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], + execution: nil + } } 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], + execution: nil + } }} = Store.get("user_42") assert NaiveDateTime.compare(FakeRepo.run("user_42").updated_at, previous_updated_at) == :gt From 84b4865b22a7c538a5cd60ece7c6175ce5c7fe5d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tomasz=20Id=C5=BAkowski?= Date: Thu, 30 Jul 2026 15:29:06 +0200 Subject: [PATCH 19/30] Add recovery mechanics Users can now opt in for Legion to automatically drive interrupted (:running without PID) agents to completion. --- README.md | 14 ++ lib/legion.ex | 77 +++++++++-- lib/legion/agent_server.ex | 25 +++- lib/legion/application.ex | 12 +- lib/legion/recovery.ex | 41 ++++++ lib/legion/store.ex | 5 +- test/legion/agent_server_test.exs | 94 +++++++++++++ test/legion/application_test.exs | 27 ++++ test/legion/recovery_test.exs | 216 ++++++++++++++++++++++++++++++ 9 files changed, 490 insertions(+), 21 deletions(-) create mode 100644 lib/legion/recovery.ex create mode 100644 test/legion/application_test.exs create mode 100644 test/legion/recovery_test.exs diff --git a/README.md b/README.md index 64329f9..6ab5cf1 100644 --- a/README.md +++ b/README.md @@ -274,6 +274,20 @@ config :legion, :store, MyApp.AgentStore A `store:` passed to `Legion.start_link/2` overrides the global store. With a global store configured, pass only `agent_id:` to select an existing conversation; if you omit it, Legion generates one. +To recover interrupted root runs when the application starts, configure the +stores to scan and a shared scan/concurrency limit: + +```elixir +config :legion, :recovery, + stores: [MyApp.AgentStore], + limit: 10 +``` + +Recovery is a one-shot asynchronous startup worker, so it does not delay the +application starting. It calls `list(10)` on each configured store, filters the +returned runs to interrupted roots, then calls `Legion.recover/2` with no more +than 10 recoveries in flight across all stores. Omit `:recovery` to disable it. + Configure model and runtime options separately: ```elixir diff --git a/lib/legion.ex b/lib/legion.ex index 0e39aa3..70de199 100644 --- a/lib/legion.ex +++ b/lib/legion.ex @@ -164,10 +164,7 @@ defmodule Legion do {:ok, pid} = Legion.resume("user_42:chat_7", 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) agent_module = case store.get(agent_id) do @@ -179,16 +176,76 @@ defmodule Legion do "no run recorded for agent_id #{inspect(agent_id)} in #{inspect(store)}" end - case lookup(agent_id) do - {:ok, pid} -> - if running?(pid) do - {:ok, pid} + with {:ok, pid} <- lookup(agent_id), true <- running?(pid) do + {:ok, pid} + else + _ -> start_link(agent_module, Keyword.merge(opts, agent_id: agent_id, start_mode: :resume)) + end + end + + @doc """ + Recovers an interrupted persisted root run and waits for it to finish. + + Loads the run from its store and resumes it only when it is a root run with + `:running` status. The recovered agent process stops after completing the + interrupted execution, so this is suitable for one-shot startup recovery. + + Pass `:store` or configure one globally. Returns `:ok` after a recovered run + finishes, `{:error, :already_running}` when that `agent_id` is live, and + `{:error, :not_recoverable}` for persisted runs that are not interrupted root + runs. Raises when the store has no recorded run for `agent_id`. + """ + def recover(agent_id, opts \\ []) do + store = store!(opts, :recover) + + with {:ok, pid} <- lookup(agent_id), true <- running?(pid) do + {:error, :already_running} + else + _ -> recover_stored_agent(store, agent_id, opts) + 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 + + defp recover_stored_agent(store, agent_id, opts) do + case store.get(agent_id) do + {:ok, + %Legion.Store.Payload{ + agent_module: agent_module, + status: :running, + parent_agent_id: nil + }} + when not is_nil(agent_module) -> + with {:ok, pid} <- lookup(agent_id), true <- running?(pid) do + {:error, :already_running} else - start_link(agent_module, Keyword.merge(opts, agent_id: agent_id, start_mode: :resume)) + _ -> + {:ok, {_pid, ref}} = + AgentServer.start_monitor( + agent_module, + Keyword.merge(opts, agent_id: agent_id, store: store, start_mode: :recover) + ) + + receive do + {:DOWN, ^ref, :process, _pid, :normal} -> :ok + {:DOWN, ^ref, :process, _pid, reason} -> {:error, reason} + end end + {:ok, %Legion.Store.Payload{agent_module: nil}} -> + raise ArgumentError, + "no run recorded for agent_id #{inspect(agent_id)} in #{inspect(store)}" + + {:ok, %Legion.Store.Payload{}} -> + {:error, :not_recoverable} + :error -> - start_link(agent_module, Keyword.merge(opts, agent_id: agent_id, start_mode: :resume)) + raise ArgumentError, + "no run recorded for agent_id #{inspect(agent_id)} in #{inspect(store)}" end end diff --git a/lib/legion/agent_server.ex b/lib/legion/agent_server.ex index 24e1cf8..a963bbc 100644 --- a/lib/legion/agent_server.ex +++ b/lib/legion/agent_server.ex @@ -28,6 +28,19 @@ defmodule Legion.AgentServer do # Client API def start_link(agent_module, opts \\ []) do + {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) + + :gen_server.start_monitor(__MODULE__, init_arg, gen_opts) + end + + defp start_args(agent_module, opts) do {name, opts} = Keyword.pop(opts, :name) {store, opts} = Keyword.pop(opts, :store) {agent_id, opts} = Keyword.pop(opts, :agent_id) @@ -45,11 +58,7 @@ defmodule Legion.AgentServer do gen_opts = if name, do: [name: name], else: [] 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 @@ -137,6 +146,12 @@ defmodule Legion.AgentServer do {:noreply, state} end + @impl true + def handle_continue(%{start_mode: :recover, execution: execution}, state) do + {_reply, state} = perform_run(state, execution) + {:stop, :normal, state} + end + @impl true def terminate(_reason, state) do Telemetry.emit( diff --git a/lib/legion/application.ex b/lib/legion/application.ex index 3776c97..2247ee6 100644 --- a/lib/legion/application.ex +++ b/lib/legion/application.ex @@ -7,10 +7,14 @@ 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 + [ + {Registry, keys: :unique, name: Legion.AgentRegistry}, + {Legion.Recovery, Application.fetch_env(:legion, :recovery)} + ] end end diff --git a/lib/legion/recovery.ex b/lib/legion/recovery.ex new file mode 100644 index 0000000..f2b1a74 --- /dev/null +++ b/lib/legion/recovery.ex @@ -0,0 +1,41 @@ +defmodule Legion.Recovery do + @moduledoc false + + alias Legion.Store.Payload + + def start_link(:error), do: :ignore + + def start_link({:ok, config}) do + Task.start_link(fn -> run(config) end) + end + + def run(config) do + stores = Keyword.fetch!(config, :stores) + limit = Keyword.fetch!(config, :limit) + + stores + |> Enum.flat_map(fn store -> + store.list(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: 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 + + 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 e0a7076..a5367ce 100644 --- a/lib/legion/store.ex +++ b/lib/legion/store.ex @@ -97,8 +97,9 @@ 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 and a limit to recover interrupted root turns once + when the Legion application starts. ## Reading conversations diff --git a/test/legion/agent_server_test.exs b/test/legion/agent_server_test.exs index 4c9161d..fc64fba 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 -> @@ -543,6 +554,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") @@ -1013,6 +1031,82 @@ defmodule Legion.AgentServerTest do end 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 completes an interrupted root 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], + execution: %{phase: :awaiting_llm, iteration: 1, retries: 0} + } + }) + + assert :ok = Legion.recover("recover-awaiting-llm", store: MemoryStore) + + assert {:ok, %Payload{status: :idle, conversation_state: %{execution: nil}}} = + 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 {:error, :already_running} = Legion.recover("recover-running", store: MemoryStore) + + assert Legion.running?(pid) + end + + test "recover/2 refuses an idle root 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], + execution: %{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], + execution: %{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/application_test.exs b/test/legion/application_test.exs new file mode 100644 index 0000000..3a8902c --- /dev/null +++ b/test/legion/application_test.exs @@ -0,0 +1,27 @@ +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 "passes configured recovery options to the recovery child" do + config = [stores: [RecoveryStore], limit: 3] + Application.put_env(:legion, :recovery, config) + + assert Enum.member?(Legion.Application.children(), {Legion.Recovery, {:ok, config}}) + end +end diff --git a/test/legion/recovery_test.exs b/test/legion/recovery_test.exs new file mode 100644 index 0000000..7ee0c2b --- /dev/null +++ b/test/legion/recovery_test.exs @@ -0,0 +1,216 @@ +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 "lists every store and recovers runs with the configured concurrency limit" 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, limit: 2}) + + monitor_ref = Process.monitor(worker) + + assert_receive {:listed, RecoveryStoreOne, 2} + assert_receive {:listed, RecoveryStoreTwo, 2} + + assert_receive {:recovering, first} + assert_receive {:recovering, second} + refute first == second + refute_receive {:recovering, _}, 50 + + send(first, :complete_recovery) + send(second, :complete_recovery) + + assert_receive {:recovering, third} + assert_receive {:recovering, fourth} + refute third == fourth + refute_receive {:recovering, _}, 50 + + send(third, :complete_recovery) + 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 "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], 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: [], + execution: %{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 From 5fb7d1457ad53e953d7da8652d00d20cf81e278d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tomasz=20Id=C5=BAkowski?= Date: Thu, 30 Jul 2026 17:24:08 +0200 Subject: [PATCH 20/30] Update docs for newly added recovery flow and consistency --- CHANGELOG.md | 3 +- README.md | 27 +++++++++++++---- lib/legion.ex | 42 ++++++++++++++++++++------ lib/legion/agent_server.ex | 4 +++ lib/legion/executor.ex | 4 +++ lib/legion/recovery.ex | 32 +++++++++++++++++++- lib/legion/store.ex | 3 +- lib/legion/store/migration/postgres.ex | 2 +- lib/legion/store/payload.ex | 11 +++++++ mix.exs | 4 +-- 10 files changed, 111 insertions(+), 21 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4b13e4e..e22497a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,8 +6,9 @@ - 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 optional `persistence_frequency/0`; stores default to `:turn`, while `:step` also checkpoints intermediate eval results, recoverable errors, bindings, and executor progress - 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 `Legion.recover/2` and opt-in `:recovery` startup configuration for recovering interrupted root runs - 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/README.md b/README.md index 5a9fab1..5e0b665 100644 --- a/README.md +++ b/README.md @@ -191,7 +191,9 @@ defmodule MyApp.AgentStore do end ``` -Step persistence records the latest recoverable state, but Legion does not automatically continue an interrupted turn. +Step persistence records the latest recoverable state. When you configure +`:recovery` below, Legion scans persisted runs once at application startup and +automatically attempts to recover eligible interrupted root runs. ## Multi-Agent Systems @@ -275,7 +277,7 @@ config :legion, :store, MyApp.AgentStore A `store:` passed to `Legion.start_link/2` overrides the global store. With a global store configured, pass only `agent_id:` to select an existing conversation; if you omit it, Legion generates one. To recover interrupted root runs when the application starts, configure the -stores to scan and a shared scan/concurrency limit: +stores to scan and a recovery limit: ```elixir config :legion, :recovery, @@ -283,10 +285,23 @@ config :legion, :recovery, limit: 10 ``` -Recovery is a one-shot asynchronous startup worker, so it does not delay the -application starting. It calls `list(10)` on each configured store, filters the -returned runs to interrupted roots, then calls `Legion.recover/2` with no more -than 10 recoveries in flight across all stores. Omit `:recovery` to disable it. +Recovery starts a temporary worker asynchronously, so it does not delay +application startup. The worker calls `list(limit)` on every configured store, +selects interrupted root runs, and calls `Legion.recover/2` for each. `limit` +is both the maximum number of runs read from each store and the maximum number +of recoveries in flight across all stores. The worker performs the recovery +scan, then exits and is not restarted. Omit `:recovery` to disable it. + +To recover a known interrupted root run directly: + +```elixir +case Legion.recover("user_42:chat_7", store: MyApp.AgentStore) do + :ok -> :recovered + {:error, :already_running} -> :already_running + {:error, :not_recoverable} -> :not_an_interrupted_root + {:error, reason} -> {:recovery_failed, reason} +end +``` Configure model and runtime options separately: diff --git a/lib/legion.ex b/lib/legion.ex index 70de199..4a2357e 100644 --- a/lib/legion.ex +++ b/lib/legion.ex @@ -152,8 +152,12 @@ defmodule Legion do 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`. + same `agent_id` and returns its pid immediately. The 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`. @@ -186,14 +190,34 @@ defmodule Legion do @doc """ Recovers an interrupted persisted root run and waits for it to finish. - Loads the run from its store and resumes it only when it is a root run with - `:running` status. The recovered agent process stops after completing the - interrupted execution, so this is suitable for one-shot startup recovery. + The stored payload must have `status: :running` and no `parent_agent_id`. + `recover/2` starts a temporary agent process, waits without a timeout while + it drives the interrupted execution to completion, then stops that process. + Unlike `resume/2`, it does not revive the conversation as a live agent. + + The executor result is not returned. `:ok` means only that the temporary + process stopped normally, including when the executor cancelled the run. + + 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 live process is registered for `agent_id` + - `{:error, :not_recoverable}` when the stored payload is not an interrupted root run + + Raises when no store is available, when the store has no payload for + `agent_id`, or when the payload has no `agent_module`. + + ## Examples + + :ok = Legion.recover("user_42:chat_7", store: MyApp.AgentStore) + + {:error, :already_running} = + Legion.recover("active_chat", store: MyApp.AgentStore) - Pass `:store` or configure one globally. Returns `:ok` after a recovered run - finishes, `{:error, :already_running}` when that `agent_id` is live, and - `{:error, :not_recoverable}` for persisted runs that are not interrupted root - runs. Raises when the store has no recorded run for `agent_id`. + {:error, :not_recoverable} = + Legion.recover("completed_chat", store: MyApp.AgentStore) """ def recover(agent_id, opts \\ []) do store = store!(opts, :recover) diff --git a/lib/legion/agent_server.ex b/lib/legion/agent_server.ex index a963bbc..dae1d7e 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 diff --git a/lib/legion/executor.ex b/lib/legion/executor.ex index 4472f41..5a263e2 100644 --- a/lib/legion/executor.ex +++ b/lib/legion/executor.ex @@ -104,6 +104,10 @@ 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. `execution` resumes a step + checkpoint when present: `:awaiting_llm` continues from its saved iteration + and retry counters, while `:completing` finishes without another LLM request. + Pass `nil` to start a new loop. Returns `{:ok, result, messages, bindings}` or `{:cancel, reason, messages, bindings}`. """ diff --git a/lib/legion/recovery.ex b/lib/legion/recovery.ex index f2b1a74..99a4072 100644 --- a/lib/legion/recovery.ex +++ b/lib/legion/recovery.ex @@ -1,14 +1,43 @@ defmodule Legion.Recovery do - @moduledoc false + @moduledoc """ + Startup worker for persisted interrupted root runs. + + `Legion.Application` starts this worker only when `:recovery` is configured: + + config :legion, :recovery, + stores: [MyApp.AgentStore], + limit: 10 + + The worker calls `list(limit)` on every configured store, selects payloads + with `status: :running` and no `parent_agent_id`, then invokes + `Legion.recover/2` for each selected run. + + 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. + + `limit` is both the maximum number of payloads read from each store and the + maximum number of recoveries in flight across all stores. 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 root run to completion, + then stops. + """ alias Legion.Store.Payload + @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) limit = Keyword.fetch!(config, :limit) @@ -31,6 +60,7 @@ defmodule Legion.Recovery do 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__, diff --git a/lib/legion/store.ex b/lib/legion/store.ex index a5367ce..b47c7b7 100644 --- a/lib/legion/store.ex +++ b/lib/legion/store.ex @@ -99,7 +99,8 @@ defmodule Legion.Store do action and the following result or error checkpoint. A crash in that window can replay the action and any external side effects. Configure `:recovery` with stores and a limit to recover interrupted root turns once - when the Legion application starts. + 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 a026356..0d2bedc 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 `:execution` value is `nil` 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 :: %{ phase: :awaiting_llm | :completing, iteration: non_neg_integer(), retries: non_neg_integer() } + @type state :: %{ messages: [map()], bindings: keyword(), execution: execution() | nil } + @type t :: %__MODULE__{ agent_id: Legion.Store.agent_id(), agent_module: module() | nil, diff --git a/mix.exs b/mix.exs index 89117df..95eda6c 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 [ @@ -58,7 +58,7 @@ 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] ] From 0bbd68ffa0ecb049555c556383b390d0b8e1807d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tomasz=20Id=C5=BAkowski?= Date: Fri, 31 Jul 2026 11:12:54 +0200 Subject: [PATCH 21/30] Resolve merge conflicts --- lib/legion/executor.ex | 117 ++++++++++++++++----- lib/legion/store/migration/postgres/v01.ex | 1 + lib/legion/store/payload.ex | 6 +- lib/legion/store/postgres.ex | 4 +- test/legion/store/postgres_db_test.exs | 23 +++- test/legion/store/postgres_test.exs | 28 ++++- 6 files changed, 139 insertions(+), 40 deletions(-) diff --git a/lib/legion/executor.ex b/lib/legion/executor.ex index 8ea9bed..6eae132 100644 --- a/lib/legion/executor.ex +++ b/lib/legion/executor.ex @@ -104,33 +104,61 @@ 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. - - Returns `{:ok, result, messages, bindings}` or `{:cancel, reason, messages, bindings}`. + `bindings` seeds the code-evaluation binding. `execution` resumes a step + checkpoint when present: `:awaiting_llm` continues from its saved iteration + and retry counters, while `:completing` finishes without another LLM request. + Pass `nil` to start a new loop. `:turn_tokens` is the number of tokens used + in the current turn so far. + + Returns `{:ok, result, messages, bindings, turn_tokens}` or + `{:cancel, reason, messages, bindings, turn_tokens}`. """ - def run(agent_module, messages, config, bindings \\ []) do + def run(agent_module, messages, config, bindings \\ [], execution \\ nil, turn_tokens \\ 0) do config = Map.merge(@default_config, config) - loop(agent_module, messages, config, 0, 0, bindings) + + case execution do + nil -> + loop(agent_module, messages, config, 0, 0, bindings, turn_tokens) + + %{phase: :awaiting_llm, iteration: i, retries: r} -> + loop(agent_module, messages, config, i, r, bindings, turn_tokens) + + %{phase: :completing, iteration: _i, retries: _r} -> + {:ok, nil, messages, bindings, turn_tokens} + end end - defp loop(agent_module, messages, config, iteration, retries, bindings) do + defp loop(agent_module, messages, config, iteration, retries, bindings, turn_tokens) do if iteration >= config.max_iterations do - {:cancel, :reached_max_iterations, messages, bindings} + {:cancel, :reached_max_iterations, messages, bindings, turn_tokens} else Telemetry.span( [:legion, :iteration], %{agent: agent_module, iteration: iteration}, - fn -> iterate(agent_module, messages, config, iteration, retries, bindings) end + fn -> + iterate(agent_module, messages, config, iteration, retries, bindings, turn_tokens) + end ) end end - defp iterate(agent_module, messages, config, iteration, retries, bindings) do + defp iterate(agent_module, messages, config, iteration, retries, bindings, turn_tokens) do # credo:disable-for-next-line try do - with {:ok, action, messages} <- call_llm(agent_module, messages, config, iteration), + with {:ok, action, messages, turn_tokens} <- + call_llm(agent_module, messages, config, iteration, turn_tokens), :ok <- validate_action_type(agent_module, action) do result = - handle_action(agent_module, messages, config, action, iteration, retries, bindings) + handle_action( + agent_module, + messages, + config, + action, + iteration, + retries, + bindings, + turn_tokens + ) {result, %{action: action["action"]}} else @@ -143,7 +171,8 @@ defmodule Legion.Executor do reason, iteration, retries, - bindings + bindings, + turn_tokens ) {result, %{action: nil}} @@ -151,27 +180,38 @@ defmodule Legion.Executor do rescue e -> result = - handle_execution_error(agent_module, messages, config, e, iteration, retries, bindings) + handle_execution_error( + agent_module, + messages, + config, + e, + iteration, + retries, + bindings, + turn_tokens + ) {result, %{action: nil}} end end - defp call_llm(agent_module, messages, config, iteration) do + defp call_llm(agent_module, messages, config, iteration, turn_tokens) do Telemetry.span( [:legion, :llm, :request], %{ agent: agent_module, model: config.model, message_count: length(messages), - iteration: iteration + iteration: iteration, + turn_tokens: turn_tokens }, fn -> case ReqLLM.generate_object(config.model, messages, action_schema(agent_module)) do {:ok, response} -> action = extract_object(response) messages = messages ++ [message(:assistant, Jason.encode!(action))] - {{:ok, action, messages}, %{object: action}} + turn_tokens = turn_tokens + response.usage.total_tokens + {{:ok, action, messages, turn_tokens}, %{object: action}} {:error, reason} -> {{:error, "LLM request failed: #{inspect(reason)}"}, %{error: reason}} @@ -206,12 +246,22 @@ defmodule Legion.Executor do %{"action" => "return", "result" => result}, _i, _r, - bindings + bindings, + turn_tokens ), - do: {:ok, result, messages, bindings} + do: {:ok, result, messages, bindings, turn_tokens} - defp handle_action(_agent, messages, _config, %{"action" => "done"}, _i, _r, bindings), - do: {:ok, nil, messages, bindings} + defp handle_action( + _agent, + messages, + _config, + %{"action" => "done"}, + _i, + _r, + bindings, + turn_tokens + ), + do: {:ok, nil, messages, bindings, turn_tokens} defp handle_action( agent, @@ -220,7 +270,8 @@ defmodule Legion.Executor do %{"action" => eval, "code" => code}, i, retries, - bindings + bindings, + turn_tokens ) when eval in ["eval_and_continue", "eval_and_complete"] and code != "" do # Tools that must see the answer come back to the model (e.g. HumanTool) @@ -244,15 +295,15 @@ defmodule Legion.Executor do checkpoint!(config, messages, new_bindings, execution) if eval == "eval_and_continue", - do: loop(agent, messages, config, i + 1, 0, new_bindings), - else: {:ok, result, messages, new_bindings} + do: loop(agent, messages, config, i + 1, 0, new_bindings, turn_tokens), + else: {:ok, result, messages, new_bindings, turn_tokens} {:error, error} -> - handle_execution_error(agent, messages, config, error, i, retries, bindings) + handle_execution_error(agent, messages, config, error, i, retries, bindings, turn_tokens) end end - defp handle_action(agent, messages, config, action, i, retries, bindings), + defp handle_action(agent, messages, config, action, i, retries, bindings, turn_tokens), do: handle_execution_error( agent, @@ -261,7 +312,8 @@ defmodule Legion.Executor do "Unexpected action: #{inspect(action)}", i, retries, - bindings + bindings, + turn_tokens ) defp eval_in_span(agent_module, code, config, bindings) do @@ -288,9 +340,18 @@ defmodule Legion.Executor do end end - defp handle_execution_error(agent_module, messages, config, error, iteration, retries, bindings) do + defp handle_execution_error( + agent_module, + messages, + config, + error, + iteration, + retries, + bindings, + turn_tokens + ) do if retries >= config.max_retries do - {:cancel, :reached_max_retries, messages, bindings} + {:cancel, :reached_max_retries, messages, bindings, turn_tokens} else error_text = error |> format_error() |> truncate_content(config[:max_message_length]) @@ -311,7 +372,7 @@ defmodule Legion.Executor do retries: next_retries }) - loop(agent_module, messages, config, iteration, next_retries, bindings) + loop(agent_module, messages, config, iteration, next_retries, bindings, turn_tokens) end end diff --git a/lib/legion/store/migration/postgres/v01.ex b/lib/legion/store/migration/postgres/v01.ex index 8d863be..1b20881 100644 --- a/lib/legion/store/migration/postgres/v01.ex +++ b/lib/legion/store/migration/postgres/v01.ex @@ -13,6 +13,7 @@ defmodule Legion.Store.Migration.Postgres.V01 do add :status, :text, null: false, default: "idle" add :started_at, :naive_datetime_usec add :conversation_state, :binary + add :total_tokens, :bigint, null: false, default: 0 add :inserted_at, :naive_datetime_usec, null: false, diff --git a/lib/legion/store/payload.ex b/lib/legion/store/payload.ex index 46ac048..aa89490 100644 --- a/lib/legion/store/payload.ex +++ b/lib/legion/store/payload.ex @@ -10,7 +10,8 @@ defmodule Legion.Store.Payload do :parent_agent_id, :status, :started_at, - :conversation_state + :conversation_state, + :total_tokens ] @type status :: :idle | :running @@ -30,6 +31,7 @@ defmodule Legion.Store.Payload do parent_agent_id: Legion.Store.agent_id() | nil, status: status() | nil, started_at: NaiveDateTime.t() | nil, - conversation_state: state() | nil + conversation_state: state() | nil, + total_tokens: non_neg_integer() | nil } end diff --git a/lib/legion/store/postgres.ex b/lib/legion/store/postgres.ex index cc8dbf4..88042e1 100644 --- a/lib/legion/store/postgres.ex +++ b/lib/legion/store/postgres.ex @@ -91,6 +91,7 @@ defmodule Legion.Store.Postgres do field :status, :string field :started_at, :naive_datetime_usec field :conversation_state, :binary + field :total_tokens, :integer field :inserted_at, :naive_datetime_usec field :updated_at, :naive_datetime_usec end @@ -170,7 +171,8 @@ defmodule Legion.Store.Postgres do parent_agent_id: record.parent_agent_id, status: decode_status(record.status), started_at: record.started_at, - conversation_state: decode_conversation_state(record.conversation_state) + conversation_state: decode_conversation_state(record.conversation_state), + total_tokens: record.total_tokens } end diff --git a/test/legion/store/postgres_db_test.exs b/test/legion/store/postgres_db_test.exs index 4463a35..0aa7093 100644 --- a/test/legion/store/postgres_db_test.exs +++ b/test/legion/store/postgres_db_test.exs @@ -21,7 +21,12 @@ 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], + execution: nil + }, + total_tokens: 100 } assert :ok = Store.save(payload) @@ -36,7 +41,7 @@ defmodule Legion.Store.PostgresDbTest do assert :ok = Store.save(payload) assert {:ok, stored} = Store.get("state-only") - assert stored == %{payload | status: :idle} + assert stored == %{payload | status: :idle, total_tokens: 0} end test "save/1 partial upsert preserves omitted fields and advances updated_at" do @@ -46,7 +51,12 @@ 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], + execution: nil + }, + total_tokens: 100 } assert :ok = Store.save(initial) @@ -63,7 +73,12 @@ 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], + execution: nil + }, + total_tokens: 100 }} = 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..8c1b7b4 100644 --- a/test/legion/store/postgres_test.exs +++ b/test/legion/store/postgres_test.exs @@ -37,6 +37,7 @@ defmodule Legion.Store.PostgresTest do status: "idle", started_at: nil, conversation_state: nil, + total_tokens: 0, inserted_at: nil, updated_at: nil } @@ -70,7 +71,12 @@ 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], + execution: nil + }, + total_tokens: 100 } assert :ok = Store.save(payload) @@ -83,7 +89,7 @@ defmodule Legion.Store.PostgresTest do conversation_state: %{messages: [%{role: "user", content: "hi"}], bindings: []} } - expected_payload = %{payload | status: :idle} + expected_payload = %{payload | status: :idle, total_tokens: 0} assert :ok = Store.save(payload) assert {:ok, ^expected_payload} = Store.get("state-only") end @@ -102,7 +108,9 @@ defmodule Legion.Store.PostgresTest do } assert :ok = Store.save(payload) - assert {:ok, ^payload} = Store.get("step-state") + + expected_payload = %{payload | total_tokens: 0} + assert {:ok, ^expected_payload} = Store.get("step-state") end test "save/1 partial upsert preserves omitted fields and advances updated_at" do @@ -112,7 +120,12 @@ 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], + execution: nil + }, + total_tokens: 100 } assert :ok = Store.save(initial) @@ -126,7 +139,12 @@ 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], + execution: nil + }, + total_tokens: 100 }} = Store.get("user_42") assert NaiveDateTime.compare(FakeRepo.run("user_42").updated_at, previous_updated_at) == :gt From a3f4847108d6b29ef295e510976d3c96542c9b4e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tomasz=20Id=C5=BAkowski?= Date: Fri, 31 Jul 2026 13:54:50 +0200 Subject: [PATCH 22/30] Add token tracking to Executor and AgentServer --- lib/legion/agent_server.ex | 45 +++++--- lib/legion/executor.ex | 116 ++++++++++++--------- test/integration/step_persistence_test.exs | 3 +- test/legion/agent_server_test.exs | 77 ++++++++++---- test/legion/executor_test.exs | 96 +++++++++++++++-- test/legion/parallel_and_pipeline_test.exs | 9 +- test/legion/recovery_test.exs | 4 +- 7 files changed, 261 insertions(+), 89 deletions(-) diff --git a/lib/legion/agent_server.ex b/lib/legion/agent_server.ex index dae1d7e..2f84ef7 100644 --- a/lib/legion/agent_server.ex +++ b/lib/legion/agent_server.ex @@ -26,6 +26,7 @@ defmodule Legion.AgentServer do :agent_id, :persistence_frequency, :execution, + total_tokens: 0, bindings: [] ] @@ -110,16 +111,17 @@ defmodule Legion.AgentServer do %{agent: agent_module} ) - {saved_messages, saved_bindings, saved_execution} = + {saved_messages, saved_bindings, saved_execution, saved_total_tokens} = case store && store.get(agent_id) do {:ok, %Payload{ - conversation_state: %{messages: messages, bindings: bindings, execution: execution} + conversation_state: %{messages: messages, bindings: bindings, execution: execution}, + total_tokens: total_tokens }} -> - {messages, bindings, execution} + {messages, bindings, execution, total_tokens} _no_state -> - {[], [], nil} + {[], [], nil, 0} end state = %__MODULE__{ @@ -130,14 +132,16 @@ defmodule Legion.AgentServer do agent_id: agent_id, persistence_frequency: persistence_frequency, bindings: saved_bindings, - execution: saved_execution + execution: saved_execution, + total_tokens: saved_total_tokens } {:ok, persist(state, agent_module: state.agent_module, parent_agent_id: parent_agent_id, - started_at: NaiveDateTime.utc_now() + started_at: NaiveDateTime.utc_now(), + total_tokens: state.total_tokens ), {:continue, %{start_mode: mode, execution: saved_execution}}} end @@ -227,7 +231,7 @@ defmodule Legion.AgentServer do executor_config = Map.put(state.config, :checkpoint, checkpoint) - {status, value, final_messages, final_bindings} = + {status, value, final_messages, final_bindings, turn_tokens} = Telemetry.span( [:legion, :agent, :message], %{agent: state.agent_module, message: state.messages |> List.last() |> Map.get(:content)}, @@ -237,7 +241,7 @@ defmodule Legion.AgentServer do initial_bindings = if conversation_scope?, do: state.bindings, else: [] - {status, value, messages, bindings} = + {status, value, messages, bindings, turn_tokens} = result = Executor.run( state.agent_module, @@ -248,15 +252,32 @@ defmodule Legion.AgentServer do ) iterations = Enum.count(messages, &(&1[:role] == "assistant")) - prev_count - {result, %{iterations: iterations, status: status, result: value, bindings: bindings}} + + {result, + %{ + iterations: iterations, + status: status, + result: value, + bindings: bindings, + turn_tokens: turn_tokens + }} end ) kept_bindings = if conversation_scope?, do: final_bindings, else: [] state = - %{state | messages: final_messages, bindings: kept_bindings} - |> persist([:conversation_state, status: :idle]) + %{ + state + | messages: final_messages, + bindings: kept_bindings, + total_tokens: state.total_tokens + turn_tokens + } + |> persist([ + :conversation_state, + status: :idle, + total_tokens: state.total_tokens + turn_tokens + ]) {{status, value}, state} end @@ -277,7 +298,7 @@ defmodule Legion.AgentServer do %{payload | conversation_state: persisted_conversation_state(checkpoint)} {field, value}, payload - when field in [:agent_module, :parent_agent_id, :status, :started_at] -> + when field in [:agent_module, :parent_agent_id, :status, :started_at, :total_tokens] -> Map.put(payload, field, value) unknown, _payload -> diff --git a/lib/legion/executor.ex b/lib/legion/executor.ex index 6eae132..e8bcd63 100644 --- a/lib/legion/executor.ex +++ b/lib/legion/executor.ex @@ -10,6 +10,7 @@ defmodule Legion.Executor do """ alias Legion.{Sandbox, Telemetry} + alias ReqLLM.Usage @default_config %{ model: "openai:gpt-5.4", @@ -113,18 +114,18 @@ defmodule Legion.Executor do Returns `{:ok, result, messages, bindings, turn_tokens}` or `{:cancel, reason, messages, bindings, turn_tokens}`. """ - def run(agent_module, messages, config, bindings \\ [], execution \\ nil, turn_tokens \\ 0) do + def run(agent_module, messages, config, bindings \\ [], execution \\ nil) do config = Map.merge(@default_config, config) case execution do nil -> - loop(agent_module, messages, config, 0, 0, bindings, turn_tokens) + loop(agent_module, messages, config, 0, 0, bindings, 0) %{phase: :awaiting_llm, iteration: i, retries: r} -> - loop(agent_module, messages, config, i, r, bindings, turn_tokens) + loop(agent_module, messages, config, i, r, bindings, 0) %{phase: :completing, iteration: _i, retries: _r} -> - {:ok, nil, messages, bindings, turn_tokens} + {:ok, nil, messages, bindings, 0} end end @@ -144,47 +145,54 @@ defmodule Legion.Executor do defp iterate(agent_module, messages, config, iteration, retries, bindings, turn_tokens) do # credo:disable-for-next-line - try do - with {:ok, action, messages, turn_tokens} <- - call_llm(agent_module, messages, config, iteration, turn_tokens), - :ok <- validate_action_type(agent_module, action) do - result = - handle_action( - agent_module, - messages, - config, - action, - iteration, - retries, - bindings, - turn_tokens - ) + llm_result = + try do + call_llm(agent_module, messages, config, iteration, turn_tokens) + rescue + error -> {:error, error, turn_tokens} + end - {result, %{action: action["action"]}} - else - {:error, reason} -> - result = - handle_execution_error( - agent_module, - messages, - config, - reason, - iteration, - retries, - bindings, - turn_tokens - ) + case llm_result do + {:ok, action, messages, turn_tokens} -> + case validate_action_type(agent_module, action) do + :ok -> + result = + handle_action( + agent_module, + messages, + config, + action, + iteration, + retries, + bindings, + turn_tokens + ) + + {result, %{action: action["action"]}} - {result, %{action: nil}} - end - rescue - e -> + {:error, reason} -> + result = + handle_execution_error( + agent_module, + messages, + config, + reason, + iteration, + retries, + bindings, + turn_tokens + ) + + {result, %{action: nil}} + end + + {:error, reason, turn_tokens} -> result = handle_execution_error( agent_module, messages, config, - e, + reason, iteration, retries, bindings, @@ -208,18 +216,30 @@ defmodule Legion.Executor do fn -> case ReqLLM.generate_object(config.model, messages, action_schema(agent_module)) do {:ok, response} -> - action = extract_object(response) - messages = messages ++ [message(:assistant, Jason.encode!(action))] - turn_tokens = turn_tokens + response.usage.total_tokens - {{:ok, action, messages, turn_tokens}, %{object: action}} + handle_llm_response(response, messages, turn_tokens) {:error, reason} -> - {{:error, "LLM request failed: #{inspect(reason)}"}, %{error: reason}} + {{:error, "LLM request failed: #{inspect(reason)}", turn_tokens}, %{error: reason}} end end ) end + defp handle_llm_response(response, messages, turn_tokens) do + usage = Usage.normalize(response.usage) + turn_tokens = turn_tokens + usage.total_tokens + + case extract_object(response) do + {:ok, action} when is_map(action) -> + messages = messages ++ [message(:assistant, Jason.encode!(action))] + {{:ok, action, messages, turn_tokens}, %{object: action}} + + {:error, reason} -> + {{:error, "LLM response object invalid: #{inspect(reason)}", turn_tokens}, + %{error: reason}} + end + end + defp checkpoint!(config, messages, bindings, execution) do case config[:checkpoint] do nil -> @@ -392,14 +412,16 @@ defmodule Legion.Executor do {:error, "Response missing required 'action' field, got: #{inspect(action)}"} end - defp extract_object(%{object: object}) when is_map(object), do: object + defp extract_object(%{object: object}) when is_map(object), do: {:ok, object} defp extract_object(%{message: %{tool_calls: tool_calls}}) when is_list(tool_calls) do - ReqLLM.ToolCall.find_args(tool_calls, "structured_output") || - raise "LLM response contained no structured object" + case ReqLLM.ToolCall.find_args(tool_calls, "structured_output") do + args when is_map(args) -> {:ok, args} + _ -> {:error, "LLM response contained no structured object"} + end end - defp extract_object(_response), do: raise("LLM response contained no structured object") + defp extract_object(_response), do: {:error, "LLM response contained no structured object"} defp format_result(result, bindings, config) do variable_names = bindings |> Keyword.keys() |> Enum.map(&"`#{&1}`") diff --git a/test/integration/step_persistence_test.exs b/test/integration/step_persistence_test.exs index 5650f14..b39f123 100644 --- a/test/integration/step_persistence_test.exs +++ b/test/integration/step_persistence_test.exs @@ -83,7 +83,8 @@ defmodule Legion.Integration.StepPersistenceTest do id: "test", model: "test", context: nil, - object: %{"action" => action, "code" => code, "result" => result} + object: %{"action" => action, "code" => code, "result" => result}, + usage: %{total_tokens: 0} }} end end diff --git a/test/legion/agent_server_test.exs b/test/legion/agent_server_test.exs index fc64fba..176650d 100644 --- a/test/legion/agent_server_test.exs +++ b/test/legion/agent_server_test.exs @@ -41,33 +41,29 @@ defmodule Legion.AgentServerTest do @moduletag capture_log: true - defp llm_response(result) do - {:ok, - %ReqLLM.Response{ - id: "test", - model: "test", - context: nil, - object: %{"action" => "return", "code" => "", "result" => result} - }} + defp llm_response(result, total_tokens \\ 0) do + llm_object(%{"action" => "return", "code" => "", "result" => result}, total_tokens) end - defp llm_eval_response(code) do - {:ok, - %ReqLLM.Response{ - id: "test", - model: "test", - context: nil, - object: %{"action" => "eval_and_complete", "code" => code, "result" => ""} - }} + defp llm_eval_response(code, total_tokens \\ 0) do + llm_object(%{"action" => "eval_and_complete", "code" => code, "result" => ""}, total_tokens) + end + + defp llm_eval_continue_response(code, total_tokens \\ 0) do + llm_object( + %{"action" => "eval_and_continue", "code" => code, "result" => ""}, + total_tokens + ) end - defp llm_eval_continue_response(code) do + defp llm_object(object, total_tokens) do {:ok, %ReqLLM.Response{ id: "test", model: "test", context: nil, - object: %{"action" => "eval_and_continue", "code" => code, "result" => ""} + object: object, + usage: %{total_tokens: total_tokens} }} end @@ -591,7 +587,8 @@ defmodule Legion.AgentServerTest do parent_agent_id: nil, started_at: started_at, status: nil, - conversation_state: nil + conversation_state: nil, + total_tokens: 0 } = started assert is_struct(started_at, NaiveDateTime) @@ -619,6 +616,43 @@ defmodule Legion.AgentServerTest do refute Enum.any?(messages, &(&1.role == "system")) end + test "accumulates token totals across turns" do + call_count = :counters.new(1, [:atomics]) + + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + :counters.add(call_count, 1, 1) + + case :counters.get(call_count, 1) do + 1 -> llm_response("first", 7) + 2 -> llm_response("second", 11) + end + end) + + {:ok, pid} = Legion.start_link(MathAgent, store: MemoryStore, agent_id: "usage-turns") + assert {:ok, "first"} = Legion.call(pid, "first turn") + assert {:ok, "second"} = Legion.call(pid, "second turn") + + assert {:ok, %Payload{total_tokens: 18}} = MemoryStore.get("usage-turns") + end + + test "restored conversations add only new invocation usage" do + assert :ok = + MemoryStore.save(%Payload{ + agent_id: "usage-restore", + total_tokens: 100, + conversation_state: %{messages: [], bindings: [], execution: nil} + }) + + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + llm_response("new work", 20) + end) + + {:ok, pid} = Legion.start_link(MathAgent, store: MemoryStore, agent_id: "usage-restore") + assert {:ok, "new work"} = Legion.call(pid, "continue") + + assert {:ok, %Payload{total_tokens: 120}} = MemoryStore.get("usage-restore") + end + test "saves a snapshot before the caller receives its reply" do stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> llm_response("Paris") @@ -972,6 +1006,7 @@ defmodule Legion.AgentServerTest do agent_id: "resume-awaiting-llm", agent_module: MathAgent, status: :running, + total_tokens: 0, conversation_state: %{ messages: [%{role: "user", type: :user, content: "compute"}], bindings: [x: 42], @@ -1002,6 +1037,7 @@ defmodule Legion.AgentServerTest do agent_id: "resume-completing", agent_module: MathAgent, status: :running, + total_tokens: 0, conversation_state: %{ messages: [%{role: "user", type: :user, content: "compute"}], bindings: [x: 42], @@ -1044,6 +1080,7 @@ defmodule Legion.AgentServerTest do parent_agent_id: nil, agent_module: MathAgent, status: :running, + total_tokens: 0, conversation_state: %{ messages: [%{role: "user", type: :user, content: "compute"}], bindings: [x: 42], @@ -1079,6 +1116,7 @@ defmodule Legion.AgentServerTest do parent_agent_id: nil, agent_module: MathAgent, status: :idle, + total_tokens: 0, conversation_state: %{ messages: [%{role: "user", type: :user, content: "compute"}], bindings: [x: 42], @@ -1096,6 +1134,7 @@ defmodule Legion.AgentServerTest do parent_agent_id: "recover-parent", agent_module: MathAgent, status: :running, + total_tokens: 0, conversation_state: %{ messages: [%{role: "user", type: :user, content: "compute"}], bindings: [x: 42], diff --git a/test/legion/executor_test.exs b/test/legion/executor_test.exs index 307f403..2709ff7 100644 --- a/test/legion/executor_test.exs +++ b/test/legion/executor_test.exs @@ -72,8 +72,15 @@ defmodule Legion.ExecutorTest do @moduletag capture_log: true - defp response(object) do - {:ok, %ReqLLM.Response{id: "test", model: "test", context: nil, object: object}} + defp response(object, total_tokens \\ 0) do + {:ok, + %ReqLLM.Response{ + id: "test", + model: "test", + context: nil, + object: object, + usage: %{total_tokens: total_tokens} + }} end defp executor_messages(message) do @@ -83,7 +90,68 @@ defmodule Legion.ExecutorTest do ] end - describe "run/4" do + describe "run/3-5" do + test "returns tokens used by a single LLM request" do + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + response(%{"action" => "return", "code" => "", "result" => "42"}, 17) + end) + + assert {:ok, "42", _messages, [], 17} = + Legion.Executor.run(MathAgent, executor_messages("what is 42?"), %{}) + end + + test "normalizes provider usage before adding turn tokens" do + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + {:ok, + %ReqLLM.Response{ + id: "test", + model: "test", + context: nil, + object: %{"action" => "return", "code" => "", "result" => "42"}, + usage: %{input_tokens: 12, output_tokens: 5} + }} + end) + + assert {:ok, "42", _messages, [], 17} = + Legion.Executor.run(MathAgent, executor_messages("what is 42?"), %{}) + end + + test "adds each LLM response total across a multi-response turn" do + call_count = :counters.new(1, [:atomics]) + + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + :counters.add(call_count, 1, 1) + + case :counters.get(call_count, 1) do + 1 -> response(%{"action" => "eval_and_continue", "code" => "x = 10", "result" => ""}, 7) + 2 -> response(%{"action" => "return", "code" => "", "result" => "done"}, 11) + end + end) + + assert {:ok, "done", _messages, [x: 10], 18} = + Legion.Executor.run( + MathAgent, + executor_messages("compute"), + %{} + ) + end + + test "retains tokens from an invalid response while retrying" do + call_count = :counters.new(1, [:atomics]) + + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + :counters.add(call_count, 1, 1) + + case :counters.get(call_count, 1) do + 1 -> response(nil, 7) + 2 -> response(%{"action" => "return", "code" => "", "result" => "recovered"}, 11) + end + end) + + assert {:ok, "recovered", _messages, [], 18} = + Legion.Executor.run(MathAgent, executor_messages("recover"), %{}) + end + test "returns result for return action" do stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> response(%{"action" => "return", "code" => "", "result" => "42"}) @@ -159,6 +227,22 @@ defmodule Legion.ExecutorTest do assert {:ok, "recovered"} = Legion.execute(MathAgent, "retry me") end + test "raised LLM exception triggers retry without adding usage" do + call_count = :counters.new(1, [:atomics]) + + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + :counters.add(call_count, 1, 1) + + case :counters.get(call_count, 1) do + 1 -> raise "provider exploded" + 2 -> response(%{"action" => "return", "code" => "", "result" => "recovered"}, 11) + end + end) + + assert {:ok, "recovered", _messages, [], 11} = + Legion.Executor.run(MathAgent, executor_messages("retry raised error"), %{}) + end + test "third-party tool module without extra_allowed_modules/0 does not crash eval" do stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> response(%{ @@ -217,7 +301,7 @@ defmodule Legion.ExecutorTest do :ok end - assert {:ok, 20, _messages, _bindings} = + assert {:ok, 20, _messages, _bindings, _turn_tokens} = Legion.Executor.run( MathAgent, executor_messages("compute"), @@ -268,7 +352,7 @@ defmodule Legion.ExecutorTest do :ok end - assert {:ok, "recovered", _messages, []} = + assert {:ok, "recovered", _messages, [], _turn_tokens} = Legion.Executor.run( MathAgent, executor_messages("recover"), @@ -293,7 +377,7 @@ defmodule Legion.ExecutorTest do response(%{"action" => "return", "code" => "", "result" => "done"}) end) - assert {:ok, "done", _messages, []} = + assert {:ok, "done", _messages, [], _turn_tokens} = Legion.Executor.run( MathAgent, executor_messages("finish"), diff --git a/test/legion/parallel_and_pipeline_test.exs b/test/legion/parallel_and_pipeline_test.exs index ce06ba5..3373bd9 100644 --- a/test/legion/parallel_and_pipeline_test.exs +++ b/test/legion/parallel_and_pipeline_test.exs @@ -19,7 +19,8 @@ defmodule Legion.ParallelAndPipelineTest do id: "test", model: "test", context: nil, - object: %{"action" => "return", "code" => "", "result" => result} + object: %{"action" => "return", "code" => "", "result" => result}, + usage: %{total_tokens: 0} }} end @@ -53,7 +54,8 @@ defmodule Legion.ParallelAndPipelineTest do id: "cancel", model: "test", context: nil, - object: %{"action" => "eval_and_continue", "code" => "1 + 1", "result" => ""} + object: %{"action" => "eval_and_continue", "code" => "1 + 1", "result" => ""}, + usage: %{total_tokens: 0} }} end end) @@ -126,7 +128,8 @@ defmodule Legion.ParallelAndPipelineTest do id: "cancel", model: "test", context: nil, - object: %{"action" => "eval_and_continue", "code" => "1 + 1", "result" => ""} + object: %{"action" => "eval_and_continue", "code" => "1 + 1", "result" => ""}, + usage: %{total_tokens: 0} }} end) diff --git a/test/legion/recovery_test.exs b/test/legion/recovery_test.exs index 7ee0c2b..a1f1f8e 100644 --- a/test/legion/recovery_test.exs +++ b/test/legion/recovery_test.exs @@ -196,6 +196,7 @@ defmodule Legion.RecoveryTest do parent_agent_id: nil, agent_module: RecoveryAgent, status: :running, + total_tokens: 0, conversation_state: %{ messages: [%{role: "user", type: :user, content: "recover me"}], bindings: [], @@ -210,7 +211,8 @@ defmodule Legion.RecoveryTest do id: "test", model: "test", context: nil, - object: %{"action" => "return", "code" => "", "result" => result} + object: %{"action" => "return", "code" => "", "result" => result}, + usage: %{total_tokens: 0} }} end end From 161da2d077bce74484e50a4378198c421937b0f1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tomasz=20Id=C5=BAkowski?= Date: Sun, 2 Aug 2026 14:21:21 +0200 Subject: [PATCH 23/30] Update tracking contract: token -> usage Stores full ReqLLM.Response.Usage struct (as jsonb[] of each consequent message) instead of just total token count --- CHANGELOG.md | 1 + README.md | 17 +++++ lib/legion/agent_server.ex | 56 ++++++++------- lib/legion/executor.ex | 81 +++++++++++----------- lib/legion/store.ex | 13 ++++ lib/legion/store/migration/postgres/v01.ex | 2 +- lib/legion/store/payload.ex | 4 +- lib/legion/store/postgres.ex | 7 +- test/legion/agent_server_test.exs | 50 ++++++++++--- test/legion/executor_test.exs | 41 ++++++++--- test/legion/recovery_test.exs | 2 +- test/legion/store/postgres_db_test.exs | 42 +++++++++-- test/legion/store/postgres_test.exs | 12 ++-- 13 files changed, 224 insertions(+), 104 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e22497a..82bc0cc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,7 @@ - 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 +- Add configurable per-request LLM usage persistence, enabled by default and disabled globally with `config :legion, :track_usage, false`. `Legion.Store.Postgres` stores usage as `jsonb[]`. - Bump the default model from `openai:gpt-4o-mini` to `openai:gpt-5.4` - `Legion.Tools.HumanTool.ask/1` now raises when called under `eval_and_complete` - the turn would end as soon as the code returns, silently discarding the human's answer; the error feeds back to the model, which retries under `eval_and_continue` diff --git a/README.md b/README.md index 5e0b665..b7ec3d6 100644 --- a/README.md +++ b/README.md @@ -195,6 +195,23 @@ Step persistence records the latest recoverable state. When you configure `:recovery` below, Legion scans persisted runs once at application startup and automatically attempts to recover eligible interrupted root runs. +### Usage tracking + +Legion stores one complete `ReqLLM.Response.usage` map for each successful LLM +request, in request order. Usage tracking is enabled by default. + +Disable it globally before starting agents: + +```elixir +config :legion, :track_usage, false +``` + +The setting is read when an agent starts. Disabled agents do not load or update +usage, so existing stored usage remains unchanged. + +With `Legion.Store.Postgres`, usage is stored as a `jsonb[]`. PostgreSQL reloads +JSON maps with string keys. + ## Multi-Agent Systems Agents orchestrate other agents through the built-in `AgentTool`: diff --git a/lib/legion/agent_server.ex b/lib/legion/agent_server.ex index 2f84ef7..5b2756d 100644 --- a/lib/legion/agent_server.ex +++ b/lib/legion/agent_server.ex @@ -26,7 +26,8 @@ defmodule Legion.AgentServer do :agent_id, :persistence_frequency, :execution, - total_tokens: 0, + :track_usage, + usage: nil, bindings: [] ] @@ -59,11 +60,12 @@ defmodule Legion.AgentServer do agent_id = agent_id || generate_id() persistence_frequency = Store.persistence_frequency(store) + track_usage = Application.get_env(:legion, :track_usage, true) gen_opts = if name, do: [name: name], else: [] config = resolve_config(agent_module, opts) - {{agent_module, config, store, agent_id, persistence_frequency}, gen_opts} + {{agent_module, config, store, agent_id, persistence_frequency, track_usage}, gen_opts} end def call(agent, message, timeout \\ :infinity) do @@ -85,7 +87,7 @@ defmodule Legion.AgentServer do # Server callbacks @impl true - def init({agent_module, config, store, agent_id, persistence_frequency}) do + def init({agent_module, config, store, agent_id, persistence_frequency, track_usage}) do parent_agent_id = Vault.get(:agent_id) mode = Map.get(config, :start_mode, :normal) @@ -111,17 +113,17 @@ defmodule Legion.AgentServer do %{agent: agent_module} ) - {saved_messages, saved_bindings, saved_execution, saved_total_tokens} = + {saved_messages, saved_bindings, saved_execution, saved_usage} = case store && store.get(agent_id) do {:ok, %Payload{ conversation_state: %{messages: messages, bindings: bindings, execution: execution}, - total_tokens: total_tokens + usage: usage }} -> - {messages, bindings, execution, total_tokens} + {messages, bindings, execution, if(track_usage, do: usage || [], else: nil)} _no_state -> - {[], [], nil, 0} + {[], [], nil, if(track_usage, do: [], else: nil)} end state = %__MODULE__{ @@ -133,16 +135,19 @@ defmodule Legion.AgentServer do persistence_frequency: persistence_frequency, bindings: saved_bindings, execution: saved_execution, - total_tokens: saved_total_tokens + track_usage: track_usage, + usage: saved_usage } - {:ok, - persist(state, - agent_module: state.agent_module, - parent_agent_id: parent_agent_id, - started_at: NaiveDateTime.utc_now(), - total_tokens: state.total_tokens - ), {:continue, %{start_mode: mode, execution: saved_execution}}} + initial_fields = [ + agent_module: state.agent_module, + parent_agent_id: parent_agent_id, + started_at: NaiveDateTime.utc_now(), + usage: state.usage + ] + + {:ok, persist(state, initial_fields), + {:continue, %{start_mode: mode, execution: saved_execution}}} end @impl true @@ -231,7 +236,7 @@ defmodule Legion.AgentServer do executor_config = Map.put(state.config, :checkpoint, checkpoint) - {status, value, final_messages, final_bindings, turn_tokens} = + {status, value, final_messages, final_bindings, turn_usage} = Telemetry.span( [:legion, :agent, :message], %{agent: state.agent_module, message: state.messages |> List.last() |> Map.get(:content)}, @@ -241,7 +246,7 @@ defmodule Legion.AgentServer do initial_bindings = if conversation_scope?, do: state.bindings, else: [] - {status, value, messages, bindings, turn_tokens} = + {status, value, messages, bindings, _turn_usage} = result = Executor.run( state.agent_module, @@ -258,26 +263,25 @@ defmodule Legion.AgentServer do iterations: iterations, status: status, result: value, - bindings: bindings, - turn_tokens: turn_tokens + bindings: bindings }} end ) kept_bindings = if conversation_scope?, do: final_bindings, else: [] + usage = if state.track_usage, do: state.usage ++ turn_usage + state = %{ state | messages: final_messages, bindings: kept_bindings, - total_tokens: state.total_tokens + turn_tokens + usage: usage } - |> persist([ - :conversation_state, - status: :idle, - total_tokens: state.total_tokens + turn_tokens - ]) + + fields = [:conversation_state, status: :idle, usage: usage] + state = persist(state, fields) {{status, value}, state} end @@ -298,7 +302,7 @@ defmodule Legion.AgentServer do %{payload | conversation_state: persisted_conversation_state(checkpoint)} {field, value}, payload - when field in [:agent_module, :parent_agent_id, :status, :started_at, :total_tokens] -> + when field in [:agent_module, :parent_agent_id, :status, :started_at, :usage] -> Map.put(payload, field, value) unknown, _payload -> diff --git a/lib/legion/executor.ex b/lib/legion/executor.ex index e8bcd63..9884ee9 100644 --- a/lib/legion/executor.ex +++ b/lib/legion/executor.ex @@ -10,7 +10,6 @@ defmodule Legion.Executor do """ alias Legion.{Sandbox, Telemetry} - alias ReqLLM.Usage @default_config %{ model: "openai:gpt-5.4", @@ -108,52 +107,52 @@ defmodule Legion.Executor do `bindings` seeds the code-evaluation binding. `execution` resumes a step checkpoint when present: `:awaiting_llm` continues from its saved iteration and retry counters, while `:completing` finishes without another LLM request. - Pass `nil` to start a new loop. `:turn_tokens` is the number of tokens used - in the current turn so far. + Pass `nil` to start a new loop. `:turn_usage` is the complete, ordered list + of usage maps returned by LLM requests in the current turn. - Returns `{:ok, result, messages, bindings, turn_tokens}` or - `{:cancel, reason, messages, bindings, turn_tokens}`. + Returns `{:ok, result, messages, bindings, turn_usage}` or + `{:cancel, reason, messages, bindings, turn_usage}`. """ def run(agent_module, messages, config, bindings \\ [], execution \\ nil) do config = Map.merge(@default_config, config) case execution do nil -> - loop(agent_module, messages, config, 0, 0, bindings, 0) + loop(agent_module, messages, config, 0, 0, bindings, []) %{phase: :awaiting_llm, iteration: i, retries: r} -> - loop(agent_module, messages, config, i, r, bindings, 0) + loop(agent_module, messages, config, i, r, bindings, []) %{phase: :completing, iteration: _i, retries: _r} -> - {:ok, nil, messages, bindings, 0} + {:ok, nil, messages, bindings, []} end end - defp loop(agent_module, messages, config, iteration, retries, bindings, turn_tokens) do + defp loop(agent_module, messages, config, iteration, retries, bindings, turn_usage) do if iteration >= config.max_iterations do - {:cancel, :reached_max_iterations, messages, bindings, turn_tokens} + {:cancel, :reached_max_iterations, messages, bindings, turn_usage} else Telemetry.span( [:legion, :iteration], %{agent: agent_module, iteration: iteration}, fn -> - iterate(agent_module, messages, config, iteration, retries, bindings, turn_tokens) + iterate(agent_module, messages, config, iteration, retries, bindings, turn_usage) end ) end end - defp iterate(agent_module, messages, config, iteration, retries, bindings, turn_tokens) do + defp iterate(agent_module, messages, config, iteration, retries, bindings, turn_usage) do # credo:disable-for-next-line llm_result = try do - call_llm(agent_module, messages, config, iteration, turn_tokens) + call_llm(agent_module, messages, config, iteration, turn_usage) rescue - error -> {:error, error, turn_tokens} + error -> {:error, error, turn_usage} end case llm_result do - {:ok, action, messages, turn_tokens} -> + {:ok, action, messages, turn_usage} -> case validate_action_type(agent_module, action) do :ok -> result = @@ -165,7 +164,7 @@ defmodule Legion.Executor do iteration, retries, bindings, - turn_tokens + turn_usage ) {result, %{action: action["action"]}} @@ -180,13 +179,13 @@ defmodule Legion.Executor do iteration, retries, bindings, - turn_tokens + turn_usage ) {result, %{action: nil}} end - {:error, reason, turn_tokens} -> + {:error, reason, turn_usage} -> result = handle_execution_error( agent_module, @@ -196,46 +195,44 @@ defmodule Legion.Executor do iteration, retries, bindings, - turn_tokens + turn_usage ) {result, %{action: nil}} end end - defp call_llm(agent_module, messages, config, iteration, turn_tokens) do + defp call_llm(agent_module, messages, config, iteration, turn_usage) do Telemetry.span( [:legion, :llm, :request], %{ agent: agent_module, model: config.model, message_count: length(messages), - iteration: iteration, - turn_tokens: turn_tokens + iteration: iteration }, fn -> case ReqLLM.generate_object(config.model, messages, action_schema(agent_module)) do {:ok, response} -> - handle_llm_response(response, messages, turn_tokens) + handle_llm_response(response, messages, turn_usage) {:error, reason} -> - {{:error, "LLM request failed: #{inspect(reason)}", turn_tokens}, %{error: reason}} + {{:error, "LLM request failed: #{inspect(reason)}", turn_usage}, %{error: reason}} end end ) end - defp handle_llm_response(response, messages, turn_tokens) do - usage = Usage.normalize(response.usage) - turn_tokens = turn_tokens + usage.total_tokens + defp handle_llm_response(response, messages, turn_usage) do + turn_usage = turn_usage ++ [response.usage] case extract_object(response) do {:ok, action} when is_map(action) -> messages = messages ++ [message(:assistant, Jason.encode!(action))] - {{:ok, action, messages, turn_tokens}, %{object: action}} + {{:ok, action, messages, turn_usage}, %{object: action}} {:error, reason} -> - {{:error, "LLM response object invalid: #{inspect(reason)}", turn_tokens}, + {{:error, "LLM response object invalid: #{inspect(reason)}", turn_usage}, %{error: reason}} end end @@ -267,9 +264,9 @@ defmodule Legion.Executor do _i, _r, bindings, - turn_tokens + turn_usage ), - do: {:ok, result, messages, bindings, turn_tokens} + do: {:ok, result, messages, bindings, turn_usage} defp handle_action( _agent, @@ -279,9 +276,9 @@ defmodule Legion.Executor do _i, _r, bindings, - turn_tokens + turn_usage ), - do: {:ok, nil, messages, bindings, turn_tokens} + do: {:ok, nil, messages, bindings, turn_usage} defp handle_action( agent, @@ -291,7 +288,7 @@ defmodule Legion.Executor do i, retries, bindings, - turn_tokens + turn_usage ) when eval in ["eval_and_continue", "eval_and_complete"] and code != "" do # Tools that must see the answer come back to the model (e.g. HumanTool) @@ -315,15 +312,15 @@ defmodule Legion.Executor do checkpoint!(config, messages, new_bindings, execution) if eval == "eval_and_continue", - do: loop(agent, messages, config, i + 1, 0, new_bindings, turn_tokens), - else: {:ok, result, messages, new_bindings, turn_tokens} + do: loop(agent, messages, config, i + 1, 0, new_bindings, turn_usage), + else: {:ok, result, messages, new_bindings, turn_usage} {:error, error} -> - handle_execution_error(agent, messages, config, error, i, retries, bindings, turn_tokens) + handle_execution_error(agent, messages, config, error, i, retries, bindings, turn_usage) end end - defp handle_action(agent, messages, config, action, i, retries, bindings, turn_tokens), + defp handle_action(agent, messages, config, action, i, retries, bindings, turn_usage), do: handle_execution_error( agent, @@ -333,7 +330,7 @@ defmodule Legion.Executor do i, retries, bindings, - turn_tokens + turn_usage ) defp eval_in_span(agent_module, code, config, bindings) do @@ -368,10 +365,10 @@ defmodule Legion.Executor do iteration, retries, bindings, - turn_tokens + turn_usage ) do if retries >= config.max_retries do - {:cancel, :reached_max_retries, messages, bindings, turn_tokens} + {:cancel, :reached_max_retries, messages, bindings, turn_usage} else error_text = error |> format_error() |> truncate_content(config[:max_message_length]) @@ -392,7 +389,7 @@ defmodule Legion.Executor do retries: next_retries }) - loop(agent_module, messages, config, iteration, next_retries, bindings, turn_tokens) + loop(agent_module, messages, config, iteration, next_retries, bindings, turn_usage) end end diff --git a/lib/legion/store.ex b/lib/legion/store.ex index b47c7b7..98f4754 100644 --- a/lib/legion/store.ex +++ b/lib/legion/store.ex @@ -52,6 +52,17 @@ defmodule Legion.Store do started from within an agent's process tree picks up that store unless given an explicit `:store` of its own. + ## Usage tracking + + Legion persists the complete `ReqLLM.Response.usage` map for every LLM + request in a conversation, ordered by request. Tracking is enabled by + default. Disable it globally before starting an agent: + + config :legion, :track_usage, false + + The setting is read when an agent starts. Disabled agents neither restore nor + update usage; existing usage in a store is preserved. + ## Identifying a conversation `:agent_id` is the key a conversation is saved under - it names one @@ -77,6 +88,8 @@ defmodule Legion.Store do `: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. + Its `:usage` field is the ordered list of raw LLM usage maps when tracking is + enabled. With `binding_scope: :turn`, active bindings are included in step snapshots while the turn is running and cleared from the final snapshot. Bindings with diff --git a/lib/legion/store/migration/postgres/v01.ex b/lib/legion/store/migration/postgres/v01.ex index 1b20881..45a099b 100644 --- a/lib/legion/store/migration/postgres/v01.ex +++ b/lib/legion/store/migration/postgres/v01.ex @@ -13,7 +13,7 @@ defmodule Legion.Store.Migration.Postgres.V01 do add :status, :text, null: false, default: "idle" add :started_at, :naive_datetime_usec add :conversation_state, :binary - add :total_tokens, :bigint, null: false, default: 0 + add :usage, {:array, :map}, null: false, default: [] add :inserted_at, :naive_datetime_usec, null: false, diff --git a/lib/legion/store/payload.ex b/lib/legion/store/payload.ex index f40770a..65fc5e4 100644 --- a/lib/legion/store/payload.ex +++ b/lib/legion/store/payload.ex @@ -19,7 +19,7 @@ defmodule Legion.Store.Payload do :status, :started_at, :conversation_state, - :total_tokens + :usage ] @type status :: :idle | :running @@ -43,6 +43,6 @@ defmodule Legion.Store.Payload do status: status() | nil, started_at: NaiveDateTime.t() | nil, conversation_state: state() | nil, - total_tokens: non_neg_integer() | nil + usage: [map()] | nil } end diff --git a/lib/legion/store/postgres.ex b/lib/legion/store/postgres.ex index 88042e1..af3c9b3 100644 --- a/lib/legion/store/postgres.ex +++ b/lib/legion/store/postgres.ex @@ -57,6 +57,9 @@ defmodule Legion.Store.Postgres do `'idle'` when it completes. Step writes update only the conversation state, leaving the running status unchanged. + Usage is stored as a `jsonb[]`: each element contains one complete LLM usage + map. PostgreSQL JSON decoding returns string-keyed maps. + `list/1` and `get/1` read persisted conversations back from the same table. @@ -91,7 +94,7 @@ defmodule Legion.Store.Postgres do field :status, :string field :started_at, :naive_datetime_usec field :conversation_state, :binary - field :total_tokens, :integer + field :usage, {:array, :map} field :inserted_at, :naive_datetime_usec field :updated_at, :naive_datetime_usec end @@ -172,7 +175,7 @@ defmodule Legion.Store.Postgres do status: decode_status(record.status), started_at: record.started_at, conversation_state: decode_conversation_state(record.conversation_state), - total_tokens: record.total_tokens + usage: record.usage } end diff --git a/test/legion/agent_server_test.exs b/test/legion/agent_server_test.exs index 176650d..51e6ebf 100644 --- a/test/legion/agent_server_test.exs +++ b/test/legion/agent_server_test.exs @@ -588,7 +588,7 @@ defmodule Legion.AgentServerTest do started_at: started_at, status: nil, conversation_state: nil, - total_tokens: 0 + usage: [] } = started assert is_struct(started_at, NaiveDateTime) @@ -616,7 +616,7 @@ defmodule Legion.AgentServerTest do refute Enum.any?(messages, &(&1.role == "system")) end - test "accumulates token totals across turns" do + test "accumulates raw usage across turns" do call_count = :counters.new(1, [:atomics]) stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> @@ -632,14 +632,15 @@ defmodule Legion.AgentServerTest do assert {:ok, "first"} = Legion.call(pid, "first turn") assert {:ok, "second"} = Legion.call(pid, "second turn") - assert {:ok, %Payload{total_tokens: 18}} = MemoryStore.get("usage-turns") + assert {:ok, payload} = MemoryStore.get("usage-turns") + assert Map.get(payload, :usage) == [%{total_tokens: 7}, %{total_tokens: 11}] end test "restored conversations add only new invocation usage" do assert :ok = MemoryStore.save(%Payload{ agent_id: "usage-restore", - total_tokens: 100, + usage: [%{total_tokens: 100}], conversation_state: %{messages: [], bindings: [], execution: nil} }) @@ -650,7 +651,36 @@ defmodule Legion.AgentServerTest do {:ok, pid} = Legion.start_link(MathAgent, store: MemoryStore, agent_id: "usage-restore") assert {:ok, "new work"} = Legion.call(pid, "continue") - assert {:ok, %Payload{total_tokens: 120}} = MemoryStore.get("usage-restore") + assert {:ok, %Payload{usage: [%{total_tokens: 100}, %{total_tokens: 20}]}} = + MemoryStore.get("usage-restore") + end + + test "does not update usage when globally disabled" do + previous = Application.get_env(:legion, :track_usage, :unset) + Application.put_env(:legion, :track_usage, false) + + on_exit(fn -> + if previous == :unset, + do: Application.delete_env(:legion, :track_usage), + else: Application.put_env(:legion, :track_usage, previous) + end) + + assert :ok = + MemoryStore.save(%Payload{ + agent_id: "usage-disabled", + usage: [%{total_tokens: 100}], + conversation_state: %{messages: [], bindings: [], execution: nil} + }) + + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + llm_response("new work", 20) + end) + + {:ok, pid} = Legion.start_link(MathAgent, store: MemoryStore, agent_id: "usage-disabled") + assert {:ok, "new work"} = Legion.call(pid, "continue") + + assert {:ok, %Payload{usage: [%{total_tokens: 100}]}} = + MemoryStore.get("usage-disabled") end test "saves a snapshot before the caller receives its reply" do @@ -1006,7 +1036,7 @@ defmodule Legion.AgentServerTest do agent_id: "resume-awaiting-llm", agent_module: MathAgent, status: :running, - total_tokens: 0, + usage: [], conversation_state: %{ messages: [%{role: "user", type: :user, content: "compute"}], bindings: [x: 42], @@ -1037,7 +1067,7 @@ defmodule Legion.AgentServerTest do agent_id: "resume-completing", agent_module: MathAgent, status: :running, - total_tokens: 0, + usage: [], conversation_state: %{ messages: [%{role: "user", type: :user, content: "compute"}], bindings: [x: 42], @@ -1080,7 +1110,7 @@ defmodule Legion.AgentServerTest do parent_agent_id: nil, agent_module: MathAgent, status: :running, - total_tokens: 0, + usage: [], conversation_state: %{ messages: [%{role: "user", type: :user, content: "compute"}], bindings: [x: 42], @@ -1116,7 +1146,7 @@ defmodule Legion.AgentServerTest do parent_agent_id: nil, agent_module: MathAgent, status: :idle, - total_tokens: 0, + usage: [], conversation_state: %{ messages: [%{role: "user", type: :user, content: "compute"}], bindings: [x: 42], @@ -1134,7 +1164,7 @@ defmodule Legion.AgentServerTest do parent_agent_id: "recover-parent", agent_module: MathAgent, status: :running, - total_tokens: 0, + usage: [], conversation_state: %{ messages: [%{role: "user", type: :user, content: "compute"}], bindings: [x: 42], diff --git a/test/legion/executor_test.exs b/test/legion/executor_test.exs index 2709ff7..b72fc14 100644 --- a/test/legion/executor_test.exs +++ b/test/legion/executor_test.exs @@ -91,16 +91,39 @@ defmodule Legion.ExecutorTest do end describe "run/3-5" do - test "returns tokens used by a single LLM request" do + test "returns raw usage from a single LLM request" do + usage = %{ + input_tokens: 12, + output_tokens: 5, + total_tokens: 17, + tool_usage: %{web_search: 1} + } + + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + {:ok, + %ReqLLM.Response{ + id: "test", + model: "test", + context: nil, + object: %{"action" => "return", "code" => "", "result" => "42"}, + usage: usage + }} + end) + + assert {:ok, "42", _messages, [], [^usage]} = + Legion.Executor.run(MathAgent, executor_messages("what is 42?"), %{}) + end + + test "returns usage list from a single LLM request" do stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> response(%{"action" => "return", "code" => "", "result" => "42"}, 17) end) - assert {:ok, "42", _messages, [], 17} = + assert {:ok, "42", _messages, [], [%{total_tokens: 17}]} = Legion.Executor.run(MathAgent, executor_messages("what is 42?"), %{}) end - test "normalizes provider usage before adding turn tokens" do + test "does not normalize provider usage" do stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> {:ok, %ReqLLM.Response{ @@ -112,11 +135,11 @@ defmodule Legion.ExecutorTest do }} end) - assert {:ok, "42", _messages, [], 17} = + assert {:ok, "42", _messages, [], [%{input_tokens: 12, output_tokens: 5}]} = Legion.Executor.run(MathAgent, executor_messages("what is 42?"), %{}) end - test "adds each LLM response total across a multi-response turn" do + test "preserves usage order across a multi-response turn" do call_count = :counters.new(1, [:atomics]) stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> @@ -128,7 +151,7 @@ defmodule Legion.ExecutorTest do end end) - assert {:ok, "done", _messages, [x: 10], 18} = + assert {:ok, "done", _messages, [x: 10], [%{total_tokens: 7}, %{total_tokens: 11}]} = Legion.Executor.run( MathAgent, executor_messages("compute"), @@ -136,7 +159,7 @@ defmodule Legion.ExecutorTest do ) end - test "retains tokens from an invalid response while retrying" do + test "retains usage from an invalid response while retrying" do call_count = :counters.new(1, [:atomics]) stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> @@ -148,7 +171,7 @@ defmodule Legion.ExecutorTest do end end) - assert {:ok, "recovered", _messages, [], 18} = + assert {:ok, "recovered", _messages, [], [%{total_tokens: 7}, %{total_tokens: 11}]} = Legion.Executor.run(MathAgent, executor_messages("recover"), %{}) end @@ -239,7 +262,7 @@ defmodule Legion.ExecutorTest do end end) - assert {:ok, "recovered", _messages, [], 11} = + assert {:ok, "recovered", _messages, [], [%{total_tokens: 11}]} = Legion.Executor.run(MathAgent, executor_messages("retry raised error"), %{}) end diff --git a/test/legion/recovery_test.exs b/test/legion/recovery_test.exs index a1f1f8e..afc03a0 100644 --- a/test/legion/recovery_test.exs +++ b/test/legion/recovery_test.exs @@ -196,7 +196,7 @@ defmodule Legion.RecoveryTest do parent_agent_id: nil, agent_module: RecoveryAgent, status: :running, - total_tokens: 0, + usage: [], conversation_state: %{ messages: [%{role: "user", type: :user, content: "recover me"}], bindings: [], diff --git a/test/legion/store/postgres_db_test.exs b/test/legion/store/postgres_db_test.exs index 8e0f150..b41eb3a 100644 --- a/test/legion/store/postgres_db_test.exs +++ b/test/legion/store/postgres_db_test.exs @@ -14,6 +14,36 @@ defmodule Legion.Store.PostgresDbTest do :ok end + test "stores usage as a jsonb array" do + payload = %Payload{ + agent_id: "usage-jsonb", + usage: [ + %{input_tokens: 12, output_tokens: 5, total_tokens: 17, tool_usage: %{web_search: 1}}, + %{input_tokens: 7, output_tokens: 3, total_tokens: 10} + ] + } + + assert :ok = Store.save(payload) + + assert {:ok, + %Payload{ + usage: [ + %{ + "input_tokens" => 12, + "output_tokens" => 5, + "total_tokens" => 17, + "tool_usage" => %{"web_search" => 1} + }, + %{"input_tokens" => 7, "output_tokens" => 3, "total_tokens" => 10} + ] + }} = Store.get("usage-jsonb") + + assert %{rows: [["jsonb[]"]]} = + Repo.query!("SELECT pg_typeof(usage)::text FROM legion_agents WHERE agent_id = $1", [ + "usage-jsonb" + ]) + end + test "save/1 fully inserts every payload field" do payload = %Payload{ agent_id: "user_42", @@ -26,11 +56,13 @@ defmodule Legion.Store.PostgresDbTest do bindings: [x: 42], execution: nil }, - total_tokens: 100 + usage: [%{total_tokens: 100}] } + expected_payload = %{payload | usage: [%{"total_tokens" => 100}]} + assert :ok = Store.save(payload) - assert {:ok, ^payload} = Store.get("user_42") + assert {:ok, ^expected_payload} = Store.get("user_42") end test "save/1 partially inserts only the supplied payload fields" do @@ -45,7 +77,7 @@ defmodule Legion.Store.PostgresDbTest do assert :ok = Store.save(payload) assert {:ok, stored} = Store.get("state-only") - assert stored == %{payload | status: :idle, total_tokens: 0} + assert stored == %{payload | status: :idle, usage: []} end test "save/1 partial upsert preserves omitted fields and advances updated_at" do @@ -60,7 +92,7 @@ defmodule Legion.Store.PostgresDbTest do bindings: [x: 42], execution: nil }, - total_tokens: 100 + usage: [%{total_tokens: 100}] } assert :ok = Store.save(initial) @@ -82,7 +114,7 @@ defmodule Legion.Store.PostgresDbTest do bindings: [x: 42], execution: nil }, - total_tokens: 100 + usage: [%{"total_tokens" => 100}] }} = Store.get("user_42") %{rows: [[updated_at]]} = diff --git a/test/legion/store/postgres_test.exs b/test/legion/store/postgres_test.exs index 56e87a9..bda580f 100644 --- a/test/legion/store/postgres_test.exs +++ b/test/legion/store/postgres_test.exs @@ -37,7 +37,7 @@ defmodule Legion.Store.PostgresTest do status: "idle", started_at: nil, conversation_state: nil, - total_tokens: 0, + usage: [], inserted_at: nil, updated_at: nil } @@ -76,7 +76,7 @@ defmodule Legion.Store.PostgresTest do bindings: [x: 42], execution: nil }, - total_tokens: 100 + usage: [%{total_tokens: 100}] } assert :ok = Store.save(payload) @@ -93,7 +93,7 @@ defmodule Legion.Store.PostgresTest do } } - expected_payload = %{payload | status: :idle, total_tokens: 0} + expected_payload = %{payload | status: :idle, usage: []} assert :ok = Store.save(payload) assert {:ok, ^expected_payload} = Store.get("state-only") end @@ -113,7 +113,7 @@ defmodule Legion.Store.PostgresTest do assert :ok = Store.save(payload) - expected_payload = %{payload | total_tokens: 0} + expected_payload = %{payload | usage: []} assert {:ok, ^expected_payload} = Store.get("step-state") end @@ -129,7 +129,7 @@ defmodule Legion.Store.PostgresTest do bindings: [x: 42], execution: nil }, - total_tokens: 100 + usage: [%{total_tokens: 100}] } assert :ok = Store.save(initial) @@ -148,7 +148,7 @@ defmodule Legion.Store.PostgresTest do bindings: [x: 42], execution: nil }, - total_tokens: 100 + usage: [%{total_tokens: 100}] }} = Store.get("user_42") assert NaiveDateTime.compare(FakeRepo.run("user_42").updated_at, previous_updated_at) == :gt From e57cc8d3aa49bbae7919e7d993ed36ffbc89f690 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tomasz=20Id=C5=BAkowski?= Date: Mon, 3 Aug 2026 12:43:23 +0200 Subject: [PATCH 24/30] Address revive concerns from diff branch --- lib/legion/agent_server.ex | 8 ++++++-- lib/legion/store/postgres.ex | 18 +++++++++--------- 2 files changed, 15 insertions(+), 11 deletions(-) diff --git a/lib/legion/agent_server.ex b/lib/legion/agent_server.ex index dae1d7e..1d900c7 100644 --- a/lib/legion/agent_server.ex +++ b/lib/legion/agent_server.ex @@ -146,8 +146,12 @@ defmodule Legion.AgentServer do @impl true def handle_continue(%{start_mode: :resume, execution: execution}, state) do - {_reply, state} = perform_run(state, execution) - {:noreply, state} + if match?(%{role: "user"}, List.last(state.messages)) do + {_reply, state} = perform_run(state, execution) + {:noreply, state} + else + {:noreply, state} + end end @impl true diff --git a/lib/legion/store/postgres.ex b/lib/legion/store/postgres.ex index cc8dbf4..a9ba088 100644 --- a/lib/legion/store/postgres.ex +++ b/lib/legion/store/postgres.ex @@ -86,13 +86,13 @@ defmodule Legion.Store.Postgres do @primary_key {:agent_id, :string, autogenerate: false} schema unquote(table) do - field :agent_module, :string - field :parent_agent_id, :string - field :status, :string - field :started_at, :naive_datetime_usec - field :conversation_state, :binary - field :inserted_at, :naive_datetime_usec - field :updated_at, :naive_datetime_usec + field(:agent_module, :string) + field(:parent_agent_id, :string) + field(:status, :string) + field(:started_at, :naive_datetime_usec) + field(:conversation_state, :binary) + field(:inserted_at, :naive_datetime_usec) + field(:updated_at, :naive_datetime_usec) end end @@ -182,9 +182,9 @@ defmodule Legion.Store.Postgres do %{ messages: Map.get(state, :messages, []), - bindings: Map.get(state, :bindings, []) + bindings: Map.get(state, :bindings, []), + execution: Map.get(state, :execution, nil) } - |> Map.merge(Map.take(state, [:execution])) end defp decode_status("running"), do: :running From 8930fbe5067609a28e620b008ece91dd4c41063f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tomasz=20Id=C5=BAkowski?= Date: Tue, 4 Aug 2026 14:48:09 +0200 Subject: [PATCH 25/30] Address review comments --- README.md | 24 +++++---- lib/legion.ex | 39 +++++++------- lib/legion/agent_server.ex | 34 ++++++------ lib/legion/executor.ex | 10 ++-- lib/legion/recovery.ex | 38 +++++++------ lib/legion/store.ex | 8 +-- lib/legion/store/payload.ex | 6 +-- lib/legion/store/postgres.ex | 16 +++--- test/integration/step_persistence_test.exs | 8 +-- test/legion/agent_server_test.exs | 28 +++++----- test/legion/application_test.exs | 7 ++- test/legion/executor_test.exs | 10 ++-- test/legion/recovery_test.exs | 63 +++++++++++++++++++--- test/legion/store/postgres_db_test.exs | 8 +-- test/legion/store/postgres_test.exs | 14 ++--- 15 files changed, 188 insertions(+), 125 deletions(-) diff --git a/README.md b/README.md index 5e0b665..ee8b402 100644 --- a/README.md +++ b/README.md @@ -193,7 +193,7 @@ end Step persistence records the latest recoverable state. When you configure `:recovery` below, Legion scans persisted runs once at application startup and -automatically attempts to recover eligible interrupted root runs. +automatically attempts to recover eligible interrupted runs. ## Multi-Agent Systems @@ -276,29 +276,31 @@ config :legion, :store, MyApp.AgentStore A `store:` passed to `Legion.start_link/2` overrides the global store. With a global store configured, pass only `agent_id:` to select an existing conversation; if you omit it, Legion generates one. -To recover interrupted root runs when the application starts, configure the -stores to scan and a recovery limit: +To recover interrupted runs when the application starts, configure the +stores to scan, how many runs to read from each, and how many recovery requests +may run concurrently: ```elixir config :legion, :recovery, stores: [MyApp.AgentStore], - limit: 10 + store_scan_limit: 100, + concurrent_request_limit: 10 ``` Recovery starts a temporary worker asynchronously, so it does not delay -application startup. The worker calls `list(limit)` on every configured store, -selects interrupted root runs, and calls `Legion.recover/2` for each. `limit` -is both the maximum number of runs read from each store and the maximum number -of recoveries in flight across all stores. The worker performs the recovery -scan, then exits and is not restarted. Omit `:recovery` to disable it. +application startup. The worker calls `list(store_scan_limit)` on every +configured store, selects interrupted runs, and calls `Legion.recover/2` +for each. `concurrent_request_limit` caps recoveries in flight across all +stores and defaults to `3`. The worker performs the recovery scan, then exits +and is not restarted. Omit `:recovery` to disable it. -To recover a known interrupted root run directly: +To recover a known interrupted run directly: ```elixir case Legion.recover("user_42:chat_7", store: MyApp.AgentStore) do :ok -> :recovered {:error, :already_running} -> :already_running - {:error, :not_recoverable} -> :not_an_interrupted_root + {:error, :not_recoverable} -> :not_interrupted {:error, reason} -> {:recovery_failed, reason} end ``` diff --git a/lib/legion.ex b/lib/legion.ex index 4a2357e..c98301b 100644 --- a/lib/legion.ex +++ b/lib/legion.ex @@ -32,11 +32,11 @@ 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`. + - `:agent_id` - Any config overrides (`:model`, `:max_iterations`, etc.) ## Examples @@ -188,15 +188,14 @@ defmodule Legion do end @doc """ - Recovers an interrupted persisted root run and waits for it to finish. + Recovers an interrupted persisted run and waits for it to finish. - The stored payload must have `status: :running` and no `parent_agent_id`. - `recover/2` starts a temporary agent process, waits without a timeout while - it drives the interrupted execution to completion, then stops that process. - Unlike `resume/2`, it does not revive the conversation as a live agent. + 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 executor result is not returned. `:ok` means only that the temporary - process stopped normally, including when the executor cancelled the run. + 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: @@ -204,7 +203,7 @@ defmodule Legion do - `:ok` when the temporary process stops normally - `{:error, reason}` when the temporary process stops abnormally - `{:error, :already_running}` when a live process is registered for `agent_id` - - `{:error, :not_recoverable}` when the stored payload is not an interrupted root run + - `{:error, :not_recoverable}` when the stored payload is not an interrupted run Raises when no store is available, when the store has no payload for `agent_id`, or when the payload has no `agent_module`. @@ -244,20 +243,18 @@ defmodule Legion do parent_agent_id: nil }} when not is_nil(agent_module) -> - with {:ok, pid} <- lookup(agent_id), true <- running?(pid) do - {:error, :already_running} - else - _ -> - {:ok, {_pid, ref}} = - AgentServer.start_monitor( - agent_module, - Keyword.merge(opts, agent_id: agent_id, store: store, start_mode: :recover) - ) - + 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} + {:DOWN, ^ref, :process, ^pid, :normal} -> :ok + {:DOWN, ^ref, :process, ^pid, reason} -> {:error, reason} end + + {:error, reason} -> + {:error, reason} end {:ok, %Legion.Store.Payload{agent_module: nil}} -> diff --git a/lib/legion/agent_server.ex b/lib/legion/agent_server.ex index 1d900c7..ad5bbc3 100644 --- a/lib/legion/agent_server.ex +++ b/lib/legion/agent_server.ex @@ -25,7 +25,7 @@ defmodule Legion.AgentServer do :store, :agent_id, :persistence_frequency, - :execution, + :executor_state, bindings: [] ] @@ -110,13 +110,17 @@ defmodule Legion.AgentServer do %{agent: agent_module} ) - {saved_messages, saved_bindings, saved_execution} = + {saved_messages, saved_bindings, saved_executor_state} = case store && store.get(agent_id) do {:ok, %Payload{ - conversation_state: %{messages: messages, bindings: bindings, execution: execution} + conversation_state: %{ + messages: messages, + bindings: bindings, + executor_state: executor_state + } }} -> - {messages, bindings, execution} + {messages, bindings, executor_state} _no_state -> {[], [], nil} @@ -130,7 +134,7 @@ defmodule Legion.AgentServer do agent_id: agent_id, persistence_frequency: persistence_frequency, bindings: saved_bindings, - execution: saved_execution + executor_state: saved_executor_state } {:ok, @@ -138,16 +142,16 @@ defmodule Legion.AgentServer do agent_module: state.agent_module, parent_agent_id: parent_agent_id, started_at: NaiveDateTime.utc_now() - ), {:continue, %{start_mode: mode, execution: saved_execution}}} + ), {: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, execution: execution}, state) do + def handle_continue(%{start_mode: :resume, executor_state: executor_state}, state) do if match?(%{role: "user"}, List.last(state.messages)) do - {_reply, state} = perform_run(state, execution) + {_reply, state} = perform_run(state, executor_state) {:noreply, state} else {:noreply, state} @@ -155,8 +159,8 @@ defmodule Legion.AgentServer do end @impl true - def handle_continue(%{start_mode: :recover, execution: execution}, state) do - {_reply, state} = perform_run(state, execution) + def handle_continue(%{start_mode: :recover, executor_state: executor_state}, state) do + {_reply, state} = perform_run(state, executor_state) {:stop, :normal, state} end @@ -218,7 +222,7 @@ defmodule Legion.AgentServer do perform_run(state) end - defp perform_run(state, execution \\ nil) do + defp perform_run(state, executor_state \\ nil) do conversation_scope? = Map.get(state.config, :binding_scope, :turn) == :conversation checkpoint = @@ -248,7 +252,7 @@ defmodule Legion.AgentServer do messages, executor_config, initial_bindings, - execution + executor_state ) iterations = Enum.count(messages, &(&1[:role] == "assistant")) - prev_count @@ -292,15 +296,15 @@ defmodule Legion.AgentServer do defp persisted_conversation_state(%__MODULE__{} = state) do [%{role: "system"} | messages] = state.messages - %{messages: messages, bindings: state.bindings, execution: nil} + %{messages: messages, bindings: state.bindings, executor_state: nil} 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) diff --git a/lib/legion/executor.ex b/lib/legion/executor.ex index 5a263e2..655da27 100644 --- a/lib/legion/executor.ex +++ b/lib/legion/executor.ex @@ -104,17 +104,17 @@ 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. `execution` resumes a step + `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 `nil` to start a new loop. Returns `{:ok, result, messages, bindings}` or `{:cancel, reason, messages, bindings}`. """ - def run(agent_module, messages, config, bindings \\ [], execution \\ nil) do + def run(agent_module, messages, config, bindings \\ [], executor_state \\ nil) do config = Map.merge(@default_config, config) - case execution do + case executor_state do nil -> loop(agent_module, messages, config, 0, 0, bindings) @@ -194,7 +194,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 @@ -205,7 +205,7 @@ defmodule Legion.Executor do callback.(%{ messages: messages, bindings: bindings, - execution: execution + executor_state: executor_state }) rescue error -> exit({:checkpoint_persistence_failed, error}) diff --git a/lib/legion/recovery.ex b/lib/legion/recovery.ex index 99a4072..db2f0bf 100644 --- a/lib/legion/recovery.ex +++ b/lib/legion/recovery.ex @@ -1,16 +1,18 @@ defmodule Legion.Recovery do @moduledoc """ - Startup worker for persisted interrupted root runs. + 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], - limit: 10 + store_scan_limit: 100, + concurrent_request_limit: 10 - The worker calls `list(limit)` on every configured store, selects payloads - with `status: :running` and no `parent_agent_id`, then invokes - `Legion.recover/2` for each selected run. + 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 @@ -18,18 +20,21 @@ defmodule Legion.Recovery do replay that evaluation and dispatch a new sub-agent; recovering the recorded child independently could execute the same work twice. - `limit` is both the maximum number of payloads read from each store and the - maximum number of recoveries in flight across all stores. 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. + `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 root run to completion, - then stops. + 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 @@ -40,17 +45,20 @@ defmodule Legion.Recovery do @doc false def run(config) do stores = Keyword.fetch!(config, :stores) - limit = Keyword.fetch!(config, :limit) + 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(limit) + 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: limit, + max_concurrency: concurrent_request_limit, ordered: false, timeout: :infinity ) diff --git a/lib/legion/store.ex b/lib/legion/store.ex index b47c7b7..43c6fb3 100644 --- a/lib/legion/store.ex +++ b/lib/legion/store.ex @@ -72,7 +72,7 @@ defmodule Legion.Store do `save/1` receives a `Legion.Store.Payload`. Its `:conversation_state` is a map containing the conversation's `:messages` (without the system prompt), - `:bindings` from evaluated code, and `:execution`. `:execution` is `nil` + `:bindings` from evaluated code, and `:executor_state`. `:executor_state` is `nil` 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 @@ -98,9 +98,9 @@ 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. Configure - `:recovery` with stores and a limit to recover interrupted root turns once - when the Legion application starts; see `Legion.Recovery` and - `Legion.recover/2`. + `: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/payload.ex b/lib/legion/store/payload.ex index 0d2bedc..580a13f 100644 --- a/lib/legion/store/payload.ex +++ b/lib/legion/store/payload.ex @@ -5,7 +5,7 @@ defmodule Legion.Store.Payload do 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 `:execution` value is `nil` for an ordinary snapshot or a + checkpoint. Its `:executor_state` value is `nil` for an ordinary snapshot or a map when step persistence captures an interrupted turn. See `Legion.Store` for the full store contract. @@ -23,7 +23,7 @@ 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() @@ -32,7 +32,7 @@ defmodule Legion.Store.Payload do @type state :: %{ messages: [map()], bindings: keyword(), - execution: execution() | nil + executor_state: executor_state() | nil } @type t :: %__MODULE__{ diff --git a/lib/legion/store/postgres.ex b/lib/legion/store/postgres.ex index a9ba088..e22b608 100644 --- a/lib/legion/store/postgres.ex +++ b/lib/legion/store/postgres.ex @@ -86,13 +86,13 @@ defmodule Legion.Store.Postgres do @primary_key {:agent_id, :string, autogenerate: false} schema unquote(table) do - field(:agent_module, :string) - field(:parent_agent_id, :string) - field(:status, :string) - field(:started_at, :naive_datetime_usec) - field(:conversation_state, :binary) - field(:inserted_at, :naive_datetime_usec) - field(:updated_at, :naive_datetime_usec) + field :agent_module, :string + field :parent_agent_id, :string + field :status, :string + field :started_at, :naive_datetime_usec + field :conversation_state, :binary + field :inserted_at, :naive_datetime_usec + field :updated_at, :naive_datetime_usec end end @@ -183,7 +183,7 @@ defmodule Legion.Store.Postgres do %{ messages: Map.get(state, :messages, []), bindings: Map.get(state, :bindings, []), - execution: Map.get(state, :execution, nil) + executor_state: Map.get(state, :executor_state, nil) } end diff --git a/test/integration/step_persistence_test.exs b/test/integration/step_persistence_test.exs index 5650f14..bbc748b 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 == [] - assert initial_state.execution == nil + assert initial_state.executor_state == nil 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 == [] - assert completed.execution == nil + assert completed.executor_state == nil end defp response(action, code, result) do diff --git a/test/legion/agent_server_test.exs b/test/legion/agent_server_test.exs index fc64fba..fe5ef1f 100644 --- a/test/legion/agent_server_test.exs +++ b/test/legion/agent_server_test.exs @@ -721,7 +721,7 @@ defmodule Legion.AgentServerTest do conversation_state: %{ messages: [%{type: :user}], bindings: [], - execution: nil + executor_state: nil } } = running @@ -730,13 +730,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 == [] - assert final_state.execution == nil + assert final_state.executor_state == nil end test "a :step store persists eval_and_complete before the final snapshot" do @@ -754,12 +754,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 - assert final_state.execution == nil + assert final_state.executor_state == nil end test "a :step store persists retry state after an error message" do @@ -784,7 +784,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 @@ -975,7 +975,7 @@ defmodule Legion.AgentServerTest do conversation_state: %{ messages: [%{role: "user", type: :user, content: "compute"}], bindings: [x: 42], - execution: %{phase: :awaiting_llm, iteration: 1, retries: 0} + executor_state: %{phase: :awaiting_llm, iteration: 1, retries: 0} } }) @@ -991,7 +991,7 @@ defmodule Legion.AgentServerTest do assert_receive :llm_requested assert_receive {:store_saved, - %Payload{status: :idle, conversation_state: %{execution: nil}}} + %Payload{status: :idle, conversation_state: %{executor_state: nil}}} refute_receive :llm_requested, 50 end @@ -1005,7 +1005,7 @@ defmodule Legion.AgentServerTest do conversation_state: %{ messages: [%{role: "user", type: :user, content: "compute"}], bindings: [x: 42], - execution: %{phase: :completing, iteration: 1, retries: 0} + executor_state: %{phase: :completing, iteration: 1, retries: 0} } }) @@ -1020,7 +1020,7 @@ defmodule Legion.AgentServerTest do assert {:ok, _pid} = Legion.resume("resume-completing", store: MemoryStore) assert_receive {:store_saved, - %Payload{status: :idle, conversation_state: %{execution: nil}}} + %Payload{status: :idle, conversation_state: %{executor_state: nil}}} refute_receive :llm_requested, 100 end @@ -1047,13 +1047,13 @@ defmodule Legion.AgentServerTest do conversation_state: %{ messages: [%{role: "user", type: :user, content: "compute"}], bindings: [x: 42], - execution: %{phase: :awaiting_llm, iteration: 1, retries: 0} + 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: %{execution: nil}}} = + assert {:ok, %Payload{status: :idle, conversation_state: %{executor_state: nil}}} = MemoryStore.get("recover-awaiting-llm") assert( @@ -1082,7 +1082,7 @@ defmodule Legion.AgentServerTest do conversation_state: %{ messages: [%{role: "user", type: :user, content: "compute"}], bindings: [x: 42], - execution: %{phase: :completing, iteration: 1, retries: 0} + executor_state: %{phase: :completing, iteration: 1, retries: 0} } }) @@ -1099,7 +1099,7 @@ defmodule Legion.AgentServerTest do conversation_state: %{ messages: [%{role: "user", type: :user, content: "compute"}], bindings: [x: 42], - execution: %{phase: :completing, iteration: 1, retries: 0} + executor_state: %{phase: :completing, iteration: 1, retries: 0} } }) diff --git a/test/legion/application_test.exs b/test/legion/application_test.exs index 3a8902c..d0c040a 100644 --- a/test/legion/application_test.exs +++ b/test/legion/application_test.exs @@ -19,7 +19,12 @@ defmodule Legion.ApplicationTest do end test "passes configured recovery options to the recovery child" do - config = [stores: [RecoveryStore], limit: 3] + 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}}) diff --git a/test/legion/executor_test.exs b/test/legion/executor_test.exs index 307f403..e0830f9 100644 --- a/test/legion/executor_test.exs +++ b/test/legion/executor_test.exs @@ -132,7 +132,7 @@ defmodule Legion.ExecutorTest do Legion.execute(MathAgent, "loop forever") end - test "retries on code execution error and cancels after max_retries" do + test "retries on code executor_state error and cancels after max_retries" do stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> response(%{ "action" => "eval_and_complete", @@ -228,7 +228,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 @@ -237,7 +237,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 @@ -279,7 +279,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 @@ -360,7 +360,7 @@ defmodule Legion.ExecutorTest do end describe "max_message_length in result/error feedback" do - test "truncates large code execution results in the feedback message" do + test "truncates large code executor_state results in the feedback message" do test_pid = self() {:ok, counter} = Agent.start_link(fn -> 0 end) diff --git a/test/legion/recovery_test.exs b/test/legion/recovery_test.exs index 7ee0c2b..592ae74 100644 --- a/test/legion/recovery_test.exs +++ b/test/legion/recovery_test.exs @@ -106,7 +106,7 @@ defmodule Legion.RecoveryTest do assert :ignore = Legion.Recovery.start_link(:error) end - test "lists every store and recovers runs with the configured concurrency limit" do + test "uses separate store scan and concurrent request limits" do stores = [RecoveryStoreOne, RecoveryStoreTwo] Enum.each(stores, fn store -> @@ -127,7 +127,9 @@ defmodule Legion.RecoveryTest do end) assert {:ok, worker} = - Legion.Recovery.start_link({:ok, stores: stores, limit: 2}) + Legion.Recovery.start_link( + {:ok, stores: stores, store_scan_limit: 2, concurrent_request_limit: 1} + ) monitor_ref = Process.monitor(worker) @@ -135,19 +137,23 @@ defmodule Legion.RecoveryTest do assert_receive {:listed, RecoveryStoreTwo, 2} assert_receive {:recovering, first} - assert_receive {:recovering, second} - refute first == second refute_receive {:recovering, _}, 50 send(first, :complete_recovery) + + assert_receive {:recovering, second} + refute_receive {:recovering, _}, 50 + send(second, :complete_recovery) assert_receive {:recovering, third} - assert_receive {:recovering, fourth} - refute third == fourth 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} @@ -158,6 +164,45 @@ defmodule Legion.RecoveryTest do 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} @@ -176,7 +221,9 @@ defmodule Legion.RecoveryTest do end) assert {:ok, worker} = - Legion.Recovery.start_link({:ok, stores: [RecoveryStoreOne], limit: 3}) + Legion.Recovery.start_link( + {:ok, stores: [RecoveryStoreOne], store_scan_limit: 3, concurrent_request_limit: 3} + ) monitor_ref = Process.monitor(worker) @@ -199,7 +246,7 @@ defmodule Legion.RecoveryTest do conversation_state: %{ messages: [%{role: "user", type: :user, content: "recover me"}], bindings: [], - execution: %{phase: :awaiting_llm, iteration: 1, retries: 0} + executor_state: %{phase: :awaiting_llm, iteration: 1, retries: 0} } } end diff --git a/test/legion/store/postgres_db_test.exs b/test/legion/store/postgres_db_test.exs index 3eb4866..1029f12 100644 --- a/test/legion/store/postgres_db_test.exs +++ b/test/legion/store/postgres_db_test.exs @@ -24,7 +24,7 @@ defmodule Legion.Store.PostgresDbTest do conversation_state: %{ messages: [%{role: "user", content: "hi"}], bindings: [x: 42], - execution: nil + executor_state: nil } } @@ -38,7 +38,7 @@ defmodule Legion.Store.PostgresDbTest do conversation_state: %{ messages: [%{role: "user", content: "hi"}], bindings: [], - execution: nil + executor_state: nil } } @@ -57,7 +57,7 @@ defmodule Legion.Store.PostgresDbTest do conversation_state: %{ messages: [%{role: "user", content: "hi"}], bindings: [x: 42], - execution: nil + executor_state: nil } } @@ -78,7 +78,7 @@ defmodule Legion.Store.PostgresDbTest do conversation_state: %{ messages: [%{role: "user", content: "hi"}], bindings: [x: 42], - execution: nil + executor_state: nil } }} = Store.get("user_42") diff --git a/test/legion/store/postgres_test.exs b/test/legion/store/postgres_test.exs index bf8821b..eb99a48 100644 --- a/test/legion/store/postgres_test.exs +++ b/test/legion/store/postgres_test.exs @@ -73,7 +73,7 @@ defmodule Legion.Store.PostgresTest do conversation_state: %{ messages: [%{role: "user", content: "hi"}], bindings: [x: 42], - execution: nil + executor_state: nil } } @@ -87,7 +87,7 @@ defmodule Legion.Store.PostgresTest do conversation_state: %{ messages: [%{role: "user", content: "hi"}], bindings: [], - execution: nil + executor_state: nil } } @@ -96,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 step executor_state state" do + executor_state = %{phase: :awaiting_llm, iteration: 2, retries: 1} payload = %Payload{ agent_id: "step-state", @@ -105,7 +105,7 @@ defmodule Legion.Store.PostgresTest do conversation_state: %{ messages: [%{role: "user", content: "result"}], bindings: [x: 42], - execution: execution + executor_state: executor_state } } @@ -123,7 +123,7 @@ defmodule Legion.Store.PostgresTest do conversation_state: %{ messages: [%{role: "user", content: "hi"}], bindings: [x: 42], - execution: nil + executor_state: nil } } @@ -141,7 +141,7 @@ defmodule Legion.Store.PostgresTest do conversation_state: %{ messages: [%{role: "user", content: "hi"}], bindings: [x: 42], - execution: nil + executor_state: nil } }} = Store.get("user_42") From 582b857684aec45d1ba9c3d3c986149bd25c6a51 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tomasz=20Id=C5=BAkowski?= Date: Tue, 4 Aug 2026 17:55:45 +0200 Subject: [PATCH 26/30] Move from Registry to :global, refactor resume/2 and recover/2 --- CHANGELOG.md | 6 +- README.md | 16 +++ lib/legion.ex | 131 +++++++++-------------- lib/legion/agent_index.ex | 44 ++++++++ lib/legion/agent_server.ex | 12 +-- lib/legion/application.ex | 1 - lib/legion/store.ex | 8 +- mix.exs | 7 +- test/integration/agent_index_test.exs | 76 ++++++++++++++ test/legion/agent_server_test.exs | 146 ++++++++++++++++++++++---- test/legion/agent_test.exs | 4 +- test/legion/application_test.exs | 4 + 12 files changed, 332 insertions(+), 123 deletions(-) create mode 100644 lib/legion/agent_index.ex create mode 100644 test/integration/agent_index_test.exs diff --git a/CHANGELOG.md b/CHANGELOG.md index e22497a..f4c8b88 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,8 +7,10 @@ - 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 -- 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 `Legion.recover/2` and opt-in `:recovery` startup configuration for recovering interrupted root runs +- 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 for recovering interrupted runs - 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/README.md b/README.md index ee8b402..a40968c 100644 --- a/README.md +++ b/README.md @@ -174,6 +174,18 @@ Use a stable `agent_id` to continue the same conversation after a process or app {:ok, pid} = Legion.resume("user_42:chat_7", store: MyApp.AgentStore) ``` +An `agent_id` also names its live process across connected nodes. Only one +process can own an id at a time. Concurrent `start_link/2` calls return +`{:error, {:already_started, pid}}` to the loser, while `resume/2` returns +`{:ok, pid}` for either the newly started or existing process. + +Disconnected network partitions can temporarily own the same id. When nodes +reconnect, `:global` resolves the conflict by terminating one owner. + +`resume/2` returns `{:error, :not_resumable}` when the selected store has no +payload containing an agent module. It raises when no store was passed or +configured. + If you omit `agent_id`, Legion generates one. Save it if you want to resume the conversation later: ```elixir @@ -305,6 +317,10 @@ case Legion.recover("user_42:chat_7", store: MyApp.AgentStore) do end ``` +`recover/2` validates the selected store before claiming the agent id. Missing +payloads, payloads without an agent module, idle runs, and child runs return +`{:error, :not_recoverable}`. It raises when no store was passed or configured. + Configure model and runtime options separately: ```elixir diff --git a/lib/legion.ex b/lib/legion.ex index c98301b..69655d6 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. @@ -36,13 +37,12 @@ defmodule Legion do 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`. - - `:agent_id` + 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 """ @@ -104,35 +104,13 @@ defmodule Legion do AgentServer.get_agent_id(pid) end - @doc """ - Returns whether `pid` is alive on this node. Non-pid values return `false`. - - Use `lookup/1` to resolve an `agent_id` to its currently registered process - before checking it. - - ## Examples - - case Legion.lookup("user_42:chat_7") do - {:ok, pid} -> Legion.running?(pid) - :error -> false - end - """ - def running?(pid) when is_pid(pid) do - node(pid) == node() and Process.alive?(pid) - rescue - # A pid deserialized from a previous VM incarnation is not a local pid, - # for which Process.alive?/1 raises. - ArgumentError -> false - end - - def running?(_other), do: false - @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 +118,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,40 +128,47 @@ 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` and returns its pid immediately. The new process restores the + 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 = store!(opts, :resume) - agent_module = - case store.get(agent_id) do - {:ok, %Legion.Store.Payload{agent_module: agent_module}} when not is_nil(agent_module) -> - agent_module - - _ -> - raise ArgumentError, - "no run recorded for agent_id #{inspect(agent_id)} in #{inspect(store)}" - end + 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 - with {:ok, pid} <- lookup(agent_id), true <- running?(pid) do - {:ok, pid} - else - _ -> start_link(agent_module, Keyword.merge(opts, agent_id: agent_id, start_mode: :resume)) + _ -> + {:error, :not_resumable} end end @@ -202,11 +187,11 @@ defmodule Legion do - `:ok` when the temporary process stops normally - `{:error, reason}` when the temporary process stops abnormally - - `{:error, :already_running}` when a live process is registered for `agent_id` - - `{:error, :not_recoverable}` when the stored payload is not an interrupted run + - `{: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, when the store has no payload for - `agent_id`, or when the payload has no `agent_module`. + Raises when no store is available. ## Examples @@ -221,27 +206,8 @@ defmodule Legion do def recover(agent_id, opts \\ []) do store = store!(opts, :recover) - with {:ok, pid} <- lookup(agent_id), true <- running?(pid) do - {:error, :already_running} - else - _ -> recover_stored_agent(store, agent_id, opts) - 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 - - defp recover_stored_agent(store, agent_id, opts) do case store.get(agent_id) do - {:ok, - %Legion.Store.Payload{ - agent_module: agent_module, - status: :running, - parent_agent_id: nil - }} + {:ok, %Payload{agent_module: agent_module, status: :running, parent_agent_id: nil}} when not is_nil(agent_module) -> case AgentServer.start_monitor( agent_module, @@ -253,23 +219,24 @@ defmodule Legion do {:DOWN, ^ref, :process, ^pid, reason} -> {:error, reason} end + {:error, {:already_started, _pid}} -> + {:error, :already_running} + {:error, reason} -> {:error, reason} end - {:ok, %Legion.Store.Payload{agent_module: nil}} -> - raise ArgumentError, - "no run recorded for agent_id #{inspect(agent_id)} in #{inspect(store)}" - - {:ok, %Legion.Store.Payload{}} -> + _ -> {:error, :not_recoverable} - - :error -> - raise ArgumentError, - "no run recorded for agent_id #{inspect(agent_id)} in #{inspect(store)}" 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 ad5bbc3..3c7c102 100644 --- a/lib/legion/agent_server.ex +++ b/lib/legion/agent_server.ex @@ -40,12 +40,12 @@ defmodule Legion.AgentServer do @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(__MODULE__, init_arg, gen_opts) + :gen_server.start_monitor(name, __MODULE__, init_arg, gen_opts) end defp start_args(agent_module, opts) do - {name, opts} = Keyword.pop(opts, :name) {store, opts} = Keyword.pop(opts, :store) {agent_id, opts} = Keyword.pop(opts, :agent_id) @@ -59,7 +59,7 @@ 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) {{agent_module, config, store, agent_id, persistence_frequency}, gen_opts} @@ -92,12 +92,6 @@ defmodule Legion.AgentServer do 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 diff --git a/lib/legion/application.ex b/lib/legion/application.ex index 2247ee6..d5bcbf4 100644 --- a/lib/legion/application.ex +++ b/lib/legion/application.ex @@ -13,7 +13,6 @@ defmodule Legion.Application do @doc false def children do [ - {Registry, keys: :unique, name: Legion.AgentRegistry}, {Legion.Recovery, Application.fetch_env(:legion, :recovery)} ] end diff --git a/lib/legion/store.ex b/lib/legion/store.ex index 43c6fb3..c994adb 100644 --- a/lib/legion/store.ex +++ b/lib/legion/store.ex @@ -60,11 +60,9 @@ 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 diff --git a/mix.exs b/mix.exs index 95eda6c..9d64091 100644 --- a/mix.exs +++ b/mix.exs @@ -60,7 +60,12 @@ defmodule Legion.MixProject do ], 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..b5e7bb6 --- /dev/null +++ b/test/integration/agent_index_test.exs @@ -0,0 +1,76 @@ +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 {: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/legion/agent_server_test.exs b/test/legion/agent_server_test.exs index fe5ef1f..e67f4b6 100644 --- a/test/legion/agent_server_test.exs +++ b/test/legion/agent_server_test.exs @@ -424,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 @@ -530,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 @@ -938,13 +940,64 @@ 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") - assert Legion.running?(pid) + assert Process.alive?(pid) 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") @@ -953,11 +1006,11 @@ defmodule Legion.AgentServerTest do {:ok, pid} = Legion.start_link(MathAgent, store: MemoryStore, agent_id: "resume-dead") {:ok, _} = Legion.call(pid, "What is the capital of France?") GenServer.stop(pid) - refute Legion.running?(pid) + refute Process.alive?(pid) {:ok, revived} = Legion.resume("resume-dead", store: MemoryStore) - assert Legion.running?(revived) + assert Process.alive?(revived) assert [ %{role: "system"}, @@ -1025,10 +1078,19 @@ defmodule Legion.AgentServerTest do refute_receive :llm_requested, 100 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) - 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 @@ -1037,7 +1099,13 @@ defmodule Legion.AgentServerTest do end end - test "recover/2 completes an interrupted root run and stops its process" do + 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", @@ -1067,12 +1135,48 @@ defmodule Legion.AgentServerTest do 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: nil + } + }) + assert {:error, :already_running} = Legion.recover("recover-running", store: MemoryStore) - assert Legion.running?(pid) + 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 root run" do + test "recover/2 refuses an idle run" do assert :ok = MemoryStore.save(%Payload{ agent_id: "recover-idle-root", 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 index d0c040a..a235bf9 100644 --- a/test/legion/application_test.exs +++ b/test/legion/application_test.exs @@ -18,6 +18,10 @@ defmodule Legion.ApplicationTest do 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], From 4e218e2da253369c854fbde9dca23388430cbe83 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tomasz=20Id=C5=BAkowski?= Date: Wed, 5 Aug 2026 12:27:51 +0200 Subject: [PATCH 27/30] Fix docs and test descriptions --- lib/legion/executor.ex | 4 ++-- test/legion/executor_test.exs | 4 ++-- test/legion/store/postgres_test.exs | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/lib/legion/executor.ex b/lib/legion/executor.ex index 655da27..39d9db9 100644 --- a/lib/legion/executor.ex +++ b/lib/legion/executor.ex @@ -248,14 +248,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/test/legion/executor_test.exs b/test/legion/executor_test.exs index e0830f9..03fa511 100644 --- a/test/legion/executor_test.exs +++ b/test/legion/executor_test.exs @@ -132,7 +132,7 @@ defmodule Legion.ExecutorTest do Legion.execute(MathAgent, "loop forever") end - test "retries on code executor_state error and cancels after max_retries" do + test "retries on code execution error and cancels after max_retries" do stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> response(%{ "action" => "eval_and_complete", @@ -360,7 +360,7 @@ defmodule Legion.ExecutorTest do end describe "max_message_length in result/error feedback" do - test "truncates large code executor_state results in the feedback message" do + test "truncates large code execution results in the feedback message" do test_pid = self() {:ok, counter} = Agent.start_link(fn -> 0 end) diff --git a/test/legion/store/postgres_test.exs b/test/legion/store/postgres_test.exs index eb99a48..7437f82 100644 --- a/test/legion/store/postgres_test.exs +++ b/test/legion/store/postgres_test.exs @@ -96,7 +96,7 @@ defmodule Legion.Store.PostgresTest do assert {:ok, ^expected_payload} = Store.get("state-only") end - test "save/1 round trips step executor_state state" do + test "save/1 round trips executor_state for a step checkpoint" do executor_state = %{phase: :awaiting_llm, iteration: 2, retries: 1} payload = %Payload{ From 109472baa20d9dd7ad527b398527b4a5f1401f7a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tomasz=20Id=C5=BAkowski?= <80714708+tom-ehh@users.noreply.github.com> Date: Thu, 6 Aug 2026 12:06:01 +0200 Subject: [PATCH 28/30] Revise Legion startup (#14) Transfer Legion startup to Oban style supervisor to ensure Recovery worker and rest of Legion starts AFTER the repo. ```elixir defmodule MyApp.Application do use Application def start(_type, _args) do children = [ MyApp.Repo, {Legion, []} ] Supervisor.start_link(children, strategy: :one_for_one, name: MyApp.Supervisor) end end ``` Instead of Legion autostarting through `application.ex` --- README.md | 11 +++++ lib/legion.ex | 46 ++++++++++++++++++- lib/legion/application.ex | 19 -------- lib/legion/recovery.ex | 2 +- mix.exs | 1 - test/legion/application_test.exs | 36 --------------- test/legion/supervisor_test.exs | 76 ++++++++++++++++++++++++++++++++ 7 files changed, 133 insertions(+), 58 deletions(-) delete mode 100644 lib/legion/application.ex delete mode 100644 test/legion/application_test.exs create mode 100644 test/legion/supervisor_test.exs diff --git a/README.md b/README.md index a40968c..a781ee8 100644 --- a/README.md +++ b/README.md @@ -288,6 +288,17 @@ config :legion, :store, MyApp.AgentStore A `store:` passed to `Legion.start_link/2` overrides the global store. With a global store configured, pass only `agent_id:` to select an existing conversation; if you omit it, Legion generates one. +Legion is an embeddable supervisor. Add it to your application's supervision +tree after dependencies required by its configured recovery stores. For a +Repo-backed store, place it after your Repo: + +```elixir +children = [ + MyApp.Repo, + {Legion, []} +] +``` + To recover interrupted runs when the application starts, configure the stores to scan, how many runs to read from each, and how many recovery requests may run concurrently: diff --git a/lib/legion.ex b/lib/legion.ex index 69655d6..05ebe84 100644 --- a/lib/legion.ex +++ b/lib/legion.ex @@ -5,9 +5,53 @@ defmodule Legion do |> String.split("") |> Enum.fetch!(1) + use Supervisor + alias Legion.{AgentIndex, AgentServer} alias Legion.Store.Payload + @doc """ + Starts Legion's supervisor. + + Add it to your application's supervision tree after dependencies required by + its configured recovery stores. For a Repo-backed store, place it after your + Repo. + + ## Examples + + defmodule MyApp.Application do + use Application + + def start(_type, _args) do + children = [ + MyApp.Repo, + {Legion, []} + ] + + Supervisor.start_link(children, strategy: :one_for_one, name: MyApp.Supervisor) + end + end + """ + def start_link(_opts) when is_list(_opts) do + Supervisor.start_link(__MODULE__, nil, name: __MODULE__) + end + + @doc false + def child_spec(opts) do + %{ + id: __MODULE__, + start: {__MODULE__, :start_link, [opts]}, + type: :supervisor + } + end + + @impl Supervisor + def init(_opts) do + children = [{Legion.Recovery, Application.fetch_env(:legion, :recovery)}] + + Supervisor.init(children, strategy: :one_for_one, name: Legion.Supervisor) + end + @doc """ Runs an agent on a single task and returns the result. @@ -46,7 +90,7 @@ defmodule Legion do {: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 """ - def start_link(agent_module, opts \\ []) do + def start_link(agent_module, opts \\ []) when is_atom(agent_module) do AgentServer.start_link(agent_module, opts) end diff --git a/lib/legion/application.ex b/lib/legion/application.ex deleted file mode 100644 index d5bcbf4..0000000 --- a/lib/legion/application.ex +++ /dev/null @@ -1,19 +0,0 @@ -defmodule Legion.Application do - @moduledoc """ - The OTP application for Legion. - """ - - use Application - - @impl true - def start(_type, _args) do - Supervisor.start_link(children(), strategy: :one_for_one, name: Legion.Supervisor) - end - - @doc false - def children do - [ - {Legion.Recovery, Application.fetch_env(:legion, :recovery)} - ] - end -end diff --git a/lib/legion/recovery.ex b/lib/legion/recovery.ex index db2f0bf..451b723 100644 --- a/lib/legion/recovery.ex +++ b/lib/legion/recovery.ex @@ -2,7 +2,7 @@ defmodule Legion.Recovery do @moduledoc """ Worker which runs the recovery process for interrupted runs. - `Legion.Application` starts this worker only when `:recovery` is configured: + Legion's supervisor starts this worker only when `:recovery` is configured: config :legion, :recovery, stores: [MyApp.AgentStore], diff --git a/mix.exs b/mix.exs index 9d64091..8b1b10d 100644 --- a/mix.exs +++ b/mix.exs @@ -42,7 +42,6 @@ defmodule Legion.MixProject do def application do [ - mod: {Legion.Application, []}, extra_applications: [:logger] ] end diff --git a/test/legion/application_test.exs b/test/legion/application_test.exs deleted file mode 100644 index a235bf9..0000000 --- a/test/legion/application_test.exs +++ /dev/null @@ -1,36 +0,0 @@ -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/supervisor_test.exs b/test/legion/supervisor_test.exs new file mode 100644 index 0000000..5ebe26f --- /dev/null +++ b/test/legion/supervisor_test.exs @@ -0,0 +1,76 @@ +defmodule Legion.SupervisorTest do + use ExUnit.Case, async: false + + defmodule FakeRepo do + use GenServer + + def start_link(_opts), do: GenServer.start_link(__MODULE__, :ok, name: __MODULE__) + + @impl GenServer + def init(:ok), do: {:ok, :ready} + end + + defmodule RecoveryStore do + def list(_limit) do + send( + Process.whereis(:legion_supervisor_test), + {:recovery_scanned, Process.whereis(FakeRepo)} + ) + + [] + end + end + + defmodule TestAgent do + @moduledoc "Agent used to verify Legion's public startup function." + use Legion.Agent + end + + setup do + previous = Application.fetch_env(:legion, :recovery) + Process.register(self(), :legion_supervisor_test) + + on_exit(fn -> + case previous do + {:ok, config} -> Application.put_env(:legion, :recovery, config) + :error -> Application.delete_env(:legion, :recovery) + end + end) + end + + test "does not auto-start a Legion supervisor" do + assert [] = Application.spec(:legion, :mod) + end + + test "uses start_link when embedded as a child" do + assert %{id: Legion, start: {Legion, :start_link, [[]]}, type: :supervisor} = + Legion.child_spec([]) + end + + test "starts an agent when start_link receives an agent module" do + assert {:ok, pid} = Legion.start_link(TestAgent) + assert is_binary(Legion.get_agent_id(pid)) + end + + test "adds recovery worker with configured options" do + config = [stores: [RecoveryStore], store_scan_limit: 3, concurrent_request_limit: 2] + Application.put_env(:legion, :recovery, config) + + assert {:ok, + {_supervisor_flags, + [%{id: Legion.Recovery, start: {Legion.Recovery, :start_link, [{:ok, ^config}]}}]}} = + Legion.init([]) + end + + test "starts recovery after a client repo" do + Application.put_env(:legion, :recovery, stores: [RecoveryStore], store_scan_limit: 1) + + start_supervised!(%{ + id: :client_supervisor, + start: {Supervisor, :start_link, [[FakeRepo, {Legion, []}], [strategy: :one_for_one]]} + }) + + assert_receive {:recovery_scanned, repo_pid} + assert is_pid(repo_pid) + end +end From 328d028b2de6ae1320278072bb5fdad4ae0e0a6f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tomasz=20Id=C5=BAkowski?= Date: Thu, 6 Aug 2026 12:27:14 +0200 Subject: [PATCH 29/30] Resolve merge artefacts --- lib/legion.ex | 2 +- lib/legion/agent_server.ex | 2 +- test/legion/agent_server_test.exs | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/lib/legion.ex b/lib/legion.ex index 05ebe84..4d04ea4 100644 --- a/lib/legion.ex +++ b/lib/legion.ex @@ -32,7 +32,7 @@ defmodule Legion do end end """ - def start_link(_opts) when is_list(_opts) do + def start_link(opts) when is_list(opts) do Supervisor.start_link(__MODULE__, nil, name: __MODULE__) end diff --git a/lib/legion/agent_server.ex b/lib/legion/agent_server.ex index 15f1eea..6a0cfb4 100644 --- a/lib/legion/agent_server.ex +++ b/lib/legion/agent_server.ex @@ -132,7 +132,7 @@ defmodule Legion.AgentServer do agent_id: agent_id, persistence_frequency: persistence_frequency, bindings: saved_bindings, - executor_state: saved_executor_state + executor_state: saved_executor_state, track_usage: track_usage, usage: saved_usage } diff --git a/test/legion/agent_server_test.exs b/test/legion/agent_server_test.exs index df27975..8ecb0fb 100644 --- a/test/legion/agent_server_test.exs +++ b/test/legion/agent_server_test.exs @@ -643,7 +643,7 @@ defmodule Legion.AgentServerTest do MemoryStore.save(%Payload{ agent_id: "usage-restore", usage: [%{total_tokens: 100}], - conversation_state: %{messages: [], bindings: [], execution: nil} + conversation_state: %{messages: [], bindings: [], executor_state: nil} }) stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> @@ -671,7 +671,7 @@ defmodule Legion.AgentServerTest do MemoryStore.save(%Payload{ agent_id: "usage-disabled", usage: [%{total_tokens: 100}], - conversation_state: %{messages: [], bindings: [], execution: nil} + conversation_state: %{messages: [], bindings: [], executor_state: nil} }) stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> From 9f57184f6f973acfbcc46ee9d580a01554571173 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tomasz=20Id=C5=BAkowski?= Date: Mon, 10 Aug 2026 19:20:35 +0200 Subject: [PATCH 30/30] Switch usage tracking to string keyed map Thix fixes an issue where usage which was persisted, and then recovered would be string-keyed while new records would be atom keyed --- CHANGELOG.md | 2 +- lib/legion/executor.ex | 14 +++++++- lib/legion/store.ex | 6 ++-- lib/legion/store/postgres.ex | 17 ++++------ test/integration/step_persistence_test.exs | 2 +- test/legion/agent_server_test.exs | 28 ++++++++-------- test/legion/executor_test.exs | 37 +++++++++++++--------- test/legion/parallel_and_pipeline_test.exs | 6 ++-- test/legion/recovery_test.exs | 2 +- test/legion/store/postgres_db_test.exs | 16 +++++----- test/legion/store/postgres_test.exs | 6 ++-- 11 files changed, 76 insertions(+), 60 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 418a79c..eb2e892 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,7 +14,7 @@ - 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 -- Add configurable per-request LLM usage persistence, enabled by default and disabled globally with `config :legion, :track_usage, false`. `Legion.Store.Postgres` stores usage as `jsonb[]`. +- Add configurable per-request LLM usage persistence, enabled by default and disabled globally with `config :legion, :track_usage, false`.; each usage map is recursively string-keyed and `Legion.Store.Postgres` stores them as jsonb[]. - Bump the default model from `openai:gpt-4o-mini` to `openai:gpt-5.4` - `Legion.Tools.HumanTool.ask/1` now raises when called under `eval_and_complete` - the turn would end as soon as the code returns, silently discarding the human's answer; the error feeds back to the model, which retries under `eval_and_continue` diff --git a/lib/legion/executor.ex b/lib/legion/executor.ex index 43a2081..1e3dcb3 100644 --- a/lib/legion/executor.ex +++ b/lib/legion/executor.ex @@ -224,7 +224,7 @@ defmodule Legion.Executor do end defp handle_llm_response(response, messages, turn_usage) do - turn_usage = turn_usage ++ [response.usage] + turn_usage = turn_usage ++ [normalize_usage(response.usage)] case extract_object(response) do {:ok, action} when is_map(action) -> @@ -237,6 +237,18 @@ defmodule Legion.Executor do end end + defp normalize_usage(usage) when is_map(usage) do + Map.new(usage, fn {key, value} -> + {normalize_usage_key(key), normalize_usage(value)} + end) + end + + defp normalize_usage(usage) when is_list(usage), do: Enum.map(usage, &normalize_usage/1) + defp normalize_usage(usage), do: usage + + defp normalize_usage_key(key) when is_atom(key), do: Atom.to_string(key) + defp normalize_usage_key(key), do: key + defp checkpoint!(config, messages, bindings, executor_state) do case config[:checkpoint] do nil -> diff --git a/lib/legion/store.ex b/lib/legion/store.ex index 72ade97..fe95996 100644 --- a/lib/legion/store.ex +++ b/lib/legion/store.ex @@ -54,7 +54,7 @@ defmodule Legion.Store do ## Usage tracking - Legion persists the complete `ReqLLM.Response.usage` map for every LLM + Legion persists a string-keyed copy of `ReqLLM.Response.usage` for every LLM request in a conversation, ordered by request. Tracking is enabled by default. Disable it globally before starting an agent: @@ -86,8 +86,8 @@ defmodule Legion.Store do `: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. - Its `:usage` field is the ordered list of raw LLM usage maps when tracking is - enabled. + Its `:usage` field is the ordered list of string-keyed LLM usage maps when + tracking is enabled. With `binding_scope: :turn`, active bindings are included in step snapshots while the turn is running and cleared from the final snapshot. Bindings with diff --git a/lib/legion/store/postgres.ex b/lib/legion/store/postgres.ex index ae10bd1..b3beea2 100644 --- a/lib/legion/store/postgres.ex +++ b/lib/legion/store/postgres.ex @@ -46,22 +46,19 @@ defmodule Legion.Store.Postgres do blobs - readable only from Elixir, one row per conversation, upserted on every save. Step snapshots therefore require no additional migration. - `save/1` performs partial upserts, so the same row carries the - conversation state and identity: `agent_module` (in `inspect/1` form, e.g. - `"MyApp.ResearchAgent"`), `parent_agent_id` linking a sub-agent to the - conversation that spawned it, and `started_at` as a UTC `NaiveDateTime` - stored with microsecond precision. Omitted payload fields preserve their - existing values. + `save/1` performs partial upserts, so a row carries the conversation state + and identity. Omitted payload fields preserve their existing values. The + `updated_at` timestamp is automatically set to the current UTC time on every save. + Only `agent_id` and `inserted_at` are never updated. The row's `status` flips to `'running'` when a turn starts and back to `'idle'` when it completes. Step writes update only the conversation state, leaving the running status unchanged. - Usage is stored as a `jsonb[]`: each element contains one complete LLM usage - map. PostgreSQL JSON decoding returns string-keyed maps. + Usage is stored as a `jsonb[]`: each element contains one complete, + string-keyed LLM usage map. - `list/1` and `get/1` read persisted - conversations back from the same table. + `list/1` and `get/1` read persisted conversations back from the same table. The migration also installs a trigger that `pg_notify`s the table's channel (the table name) with the `agent_id` on every insert or update, so diff --git a/test/integration/step_persistence_test.exs b/test/integration/step_persistence_test.exs index 4813811..31adeb6 100644 --- a/test/integration/step_persistence_test.exs +++ b/test/integration/step_persistence_test.exs @@ -84,7 +84,7 @@ defmodule Legion.Integration.StepPersistenceTest do model: "test", context: nil, object: %{"action" => action, "code" => code, "result" => result}, - usage: %{total_tokens: 0} + usage: %{turn_usage: 0} }} end end diff --git a/test/legion/agent_server_test.exs b/test/legion/agent_server_test.exs index 8ecb0fb..3966e69 100644 --- a/test/legion/agent_server_test.exs +++ b/test/legion/agent_server_test.exs @@ -41,29 +41,29 @@ defmodule Legion.AgentServerTest do @moduletag capture_log: true - defp llm_response(result, total_tokens \\ 0) do - llm_object(%{"action" => "return", "code" => "", "result" => result}, total_tokens) + defp llm_response(result, turn_usage \\ 0) do + llm_object(%{"action" => "return", "code" => "", "result" => result}, turn_usage) end - defp llm_eval_response(code, total_tokens \\ 0) do - llm_object(%{"action" => "eval_and_complete", "code" => code, "result" => ""}, total_tokens) + defp llm_eval_response(code, turn_usage \\ 0) do + llm_object(%{"action" => "eval_and_complete", "code" => code, "result" => ""}, turn_usage) end - defp llm_eval_continue_response(code, total_tokens \\ 0) do + defp llm_eval_continue_response(code, turn_usage \\ 0) do llm_object( %{"action" => "eval_and_continue", "code" => code, "result" => ""}, - total_tokens + turn_usage ) end - defp llm_object(object, total_tokens) do + defp llm_object(object, turn_usage) do {:ok, %ReqLLM.Response{ id: "test", model: "test", context: nil, object: object, - usage: %{total_tokens: total_tokens} + usage: %{turn_usage: turn_usage} }} end @@ -618,7 +618,7 @@ defmodule Legion.AgentServerTest do refute Enum.any?(messages, &(&1.role == "system")) end - test "accumulates raw usage across turns" do + test "accumulates string-keyed usage across turns" do call_count = :counters.new(1, [:atomics]) stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> @@ -635,14 +635,14 @@ defmodule Legion.AgentServerTest do assert {:ok, "second"} = Legion.call(pid, "second turn") assert {:ok, payload} = MemoryStore.get("usage-turns") - assert Map.get(payload, :usage) == [%{total_tokens: 7}, %{total_tokens: 11}] + assert Map.get(payload, :usage) == [%{"turn_usage" => 7}, %{"turn_usage" => 11}] end test "restored conversations add only new invocation usage" do assert :ok = MemoryStore.save(%Payload{ agent_id: "usage-restore", - usage: [%{total_tokens: 100}], + usage: [%{turn_usage: 100}], conversation_state: %{messages: [], bindings: [], executor_state: nil} }) @@ -653,7 +653,7 @@ defmodule Legion.AgentServerTest do {:ok, pid} = Legion.start_link(MathAgent, store: MemoryStore, agent_id: "usage-restore") assert {:ok, "new work"} = Legion.call(pid, "continue") - assert {:ok, %Payload{usage: [%{total_tokens: 100}, %{total_tokens: 20}]}} = + assert {:ok, %Payload{usage: [%{turn_usage: 100}, %{"turn_usage" => 20}]}} = MemoryStore.get("usage-restore") end @@ -670,7 +670,7 @@ defmodule Legion.AgentServerTest do assert :ok = MemoryStore.save(%Payload{ agent_id: "usage-disabled", - usage: [%{total_tokens: 100}], + usage: [%{turn_usage: 100}], conversation_state: %{messages: [], bindings: [], executor_state: nil} }) @@ -681,7 +681,7 @@ defmodule Legion.AgentServerTest do {:ok, pid} = Legion.start_link(MathAgent, store: MemoryStore, agent_id: "usage-disabled") assert {:ok, "new work"} = Legion.call(pid, "continue") - assert {:ok, %Payload{usage: [%{total_tokens: 100}]}} = + assert {:ok, %Payload{usage: [%{turn_usage: 100}]}} = MemoryStore.get("usage-disabled") end diff --git a/test/legion/executor_test.exs b/test/legion/executor_test.exs index 9b1f6ce..0c28459 100644 --- a/test/legion/executor_test.exs +++ b/test/legion/executor_test.exs @@ -72,14 +72,14 @@ defmodule Legion.ExecutorTest do @moduletag capture_log: true - defp response(object, total_tokens \\ 0) do + defp response(object, turn_usage \\ 0) do {:ok, %ReqLLM.Response{ id: "test", model: "test", context: nil, object: object, - usage: %{total_tokens: total_tokens} + usage: %{turn_usage: turn_usage} }} end @@ -91,11 +91,11 @@ defmodule Legion.ExecutorTest do end describe "run/3-5" do - test "returns raw usage from a single LLM request" do + test "returns recursively string-keyed usage from a single LLM request" do usage = %{ input_tokens: 12, output_tokens: 5, - total_tokens: 17, + turn_usage: 17, tool_usage: %{web_search: 1} } @@ -110,8 +110,15 @@ defmodule Legion.ExecutorTest do }} end) - assert {:ok, "42", _messages, [], [^usage]} = - Legion.Executor.run(MathAgent, executor_messages("what is 42?"), %{}) + assert {:ok, "42", _messages, [], + [ + %{ + "input_tokens" => 12, + "output_tokens" => 5, + "turn_usage" => 17, + "tool_usage" => %{"web_search" => 1} + } + ]} = Legion.Executor.run(MathAgent, executor_messages("what is 42?"), %{}) end test "returns usage list from a single LLM request" do @@ -119,11 +126,11 @@ defmodule Legion.ExecutorTest do response(%{"action" => "return", "code" => "", "result" => "42"}, 17) end) - assert {:ok, "42", _messages, [], [%{total_tokens: 17}]} = + assert {:ok, "42", _messages, [], [%{"turn_usage" => 17}]} = Legion.Executor.run(MathAgent, executor_messages("what is 42?"), %{}) end - test "does not normalize provider usage" do + test "preserves provider usage values while stringifying keys" do stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> {:ok, %ReqLLM.Response{ @@ -135,7 +142,7 @@ defmodule Legion.ExecutorTest do }} end) - assert {:ok, "42", _messages, [], [%{input_tokens: 12, output_tokens: 5}]} = + assert {:ok, "42", _messages, [], [%{"input_tokens" => 12, "output_tokens" => 5}]} = Legion.Executor.run(MathAgent, executor_messages("what is 42?"), %{}) end @@ -151,7 +158,7 @@ defmodule Legion.ExecutorTest do end end) - assert {:ok, "done", _messages, [x: 10], [%{total_tokens: 7}, %{total_tokens: 11}]} = + assert {:ok, "done", _messages, [x: 10], [%{"turn_usage" => 7}, %{"turn_usage" => 11}]} = Legion.Executor.run( MathAgent, executor_messages("compute"), @@ -171,7 +178,7 @@ defmodule Legion.ExecutorTest do end end) - assert {:ok, "recovered", _messages, [], [%{total_tokens: 7}, %{total_tokens: 11}]} = + assert {:ok, "recovered", _messages, [], [%{"turn_usage" => 7}, %{"turn_usage" => 11}]} = Legion.Executor.run(MathAgent, executor_messages("recover"), %{}) end @@ -262,7 +269,7 @@ defmodule Legion.ExecutorTest do end end) - assert {:ok, "recovered", _messages, [], [%{total_tokens: 11}]} = + assert {:ok, "recovered", _messages, [], [%{"turn_usage" => 11}]} = Legion.Executor.run(MathAgent, executor_messages("retry raised error"), %{}) end @@ -324,7 +331,7 @@ defmodule Legion.ExecutorTest do :ok end - assert {:ok, 20, _messages, _bindings, _turn_tokens} = + assert {:ok, 20, _messages, _bindings, _turn_usage} = Legion.Executor.run( MathAgent, executor_messages("compute"), @@ -375,7 +382,7 @@ defmodule Legion.ExecutorTest do :ok end - assert {:ok, "recovered", _messages, [], _turn_tokens} = + assert {:ok, "recovered", _messages, [], _turn_usage} = Legion.Executor.run( MathAgent, executor_messages("recover"), @@ -400,7 +407,7 @@ defmodule Legion.ExecutorTest do response(%{"action" => "return", "code" => "", "result" => "done"}) end) - assert {:ok, "done", _messages, [], _turn_tokens} = + assert {:ok, "done", _messages, [], _turn_usage} = Legion.Executor.run( MathAgent, executor_messages("finish"), diff --git a/test/legion/parallel_and_pipeline_test.exs b/test/legion/parallel_and_pipeline_test.exs index 3373bd9..e576c3b 100644 --- a/test/legion/parallel_and_pipeline_test.exs +++ b/test/legion/parallel_and_pipeline_test.exs @@ -20,7 +20,7 @@ defmodule Legion.ParallelAndPipelineTest do model: "test", context: nil, object: %{"action" => "return", "code" => "", "result" => result}, - usage: %{total_tokens: 0} + usage: %{turn_usage: 0} }} end @@ -55,7 +55,7 @@ defmodule Legion.ParallelAndPipelineTest do model: "test", context: nil, object: %{"action" => "eval_and_continue", "code" => "1 + 1", "result" => ""}, - usage: %{total_tokens: 0} + usage: %{turn_usage: 0} }} end end) @@ -129,7 +129,7 @@ defmodule Legion.ParallelAndPipelineTest do model: "test", context: nil, object: %{"action" => "eval_and_continue", "code" => "1 + 1", "result" => ""}, - usage: %{total_tokens: 0} + usage: %{turn_usage: 0} }} end) diff --git a/test/legion/recovery_test.exs b/test/legion/recovery_test.exs index 579b841..478c147 100644 --- a/test/legion/recovery_test.exs +++ b/test/legion/recovery_test.exs @@ -259,7 +259,7 @@ defmodule Legion.RecoveryTest do model: "test", context: nil, object: %{"action" => "return", "code" => "", "result" => result}, - usage: %{total_tokens: 0} + usage: %{turn_usage: 0} }} end end diff --git a/test/legion/store/postgres_db_test.exs b/test/legion/store/postgres_db_test.exs index 940016f..7ced963 100644 --- a/test/legion/store/postgres_db_test.exs +++ b/test/legion/store/postgres_db_test.exs @@ -18,8 +18,8 @@ defmodule Legion.Store.PostgresDbTest do payload = %Payload{ agent_id: "usage-jsonb", usage: [ - %{input_tokens: 12, output_tokens: 5, total_tokens: 17, tool_usage: %{web_search: 1}}, - %{input_tokens: 7, output_tokens: 3, total_tokens: 10} + %{input_tokens: 12, output_tokens: 5, turn_usage: 17, tool_usage: %{web_search: 1}}, + %{input_tokens: 7, output_tokens: 3, turn_usage: 10} ] } @@ -31,10 +31,10 @@ defmodule Legion.Store.PostgresDbTest do %{ "input_tokens" => 12, "output_tokens" => 5, - "total_tokens" => 17, + "turn_usage" => 17, "tool_usage" => %{"web_search" => 1} }, - %{"input_tokens" => 7, "output_tokens" => 3, "total_tokens" => 10} + %{"input_tokens" => 7, "output_tokens" => 3, "turn_usage" => 10} ] }} = Store.get("usage-jsonb") @@ -56,10 +56,10 @@ defmodule Legion.Store.PostgresDbTest do bindings: [x: 42], executor_state: nil }, - usage: [%{total_tokens: 100}] + usage: [%{turn_usage: 100}] } - expected_payload = %{payload | usage: [%{"total_tokens" => 100}]} + expected_payload = %{payload | usage: [%{"turn_usage" => 100}]} assert :ok = Store.save(payload) assert {:ok, ^expected_payload} = Store.get("user_42") @@ -92,7 +92,7 @@ defmodule Legion.Store.PostgresDbTest do bindings: [x: 42], executor_state: nil }, - usage: [%{total_tokens: 100}] + usage: [%{turn_usage: 100}] } assert :ok = Store.save(initial) @@ -114,7 +114,7 @@ defmodule Legion.Store.PostgresDbTest do bindings: [x: 42], executor_state: nil }, - usage: [%{"total_tokens" => 100}] + usage: [%{"turn_usage" => 100}] }} = Store.get("user_42") %{rows: [[updated_at]]} = diff --git a/test/legion/store/postgres_test.exs b/test/legion/store/postgres_test.exs index 9409652..0645279 100644 --- a/test/legion/store/postgres_test.exs +++ b/test/legion/store/postgres_test.exs @@ -76,7 +76,7 @@ defmodule Legion.Store.PostgresTest do bindings: [x: 42], executor_state: nil }, - usage: [%{total_tokens: 100}] + usage: [%{turn_usage: 100}] } assert :ok = Store.save(payload) @@ -129,7 +129,7 @@ defmodule Legion.Store.PostgresTest do bindings: [x: 42], executor_state: nil }, - usage: [%{total_tokens: 100}] + usage: [%{turn_usage: 100}] } assert :ok = Store.save(initial) @@ -148,7 +148,7 @@ defmodule Legion.Store.PostgresTest do bindings: [x: 42], executor_state: nil }, - usage: [%{total_tokens: 100}] + usage: [%{turn_usage: 100}] }} = Store.get("user_42") assert NaiveDateTime.compare(FakeRepo.run("user_42").updated_at, previous_updated_at) == :gt