Add usage tracking

CHANGELOG.md

@@ -18,6 +18,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`.; 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`
 

lib/legion.ex

@@ -5,9 +5,53 @@ defmodule Legion do
              |> String.split("<!-- MDOC -->")
              |> 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
 

lib/legion/agent_server.ex

@@ -25,6 +25,8 @@ defmodule Legion.AgentServer do
     :store,
     :agent_id,
     :persistence_frequency,
+    :track_usage,
+    usage: nil,
     executor_state: :nonexistent,
     bindings: []
   ]
@@ -58,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 = [name: Legion.AgentIndex.name(agent_id)]
     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
@@ -84,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)
 
@@ -104,20 +107,21 @@ defmodule Legion.AgentServer do
       %{agent: agent_module}
     )
 
-    {saved_messages, saved_bindings, saved_executor_state} =
+    {saved_messages, saved_bindings, saved_executor_state, saved_usage} =
       case store && store.get(agent_id) do
         {:ok,
          %Payload{
            conversation_state: %{
              messages: messages,
              bindings: bindings,
              executor_state: executor_state
-           }
+           },
+           usage: usage
          }} ->
-          {messages, bindings, executor_state}
+          {messages, bindings, executor_state, if(track_usage, do: usage || [], else: nil)}
 
         _no_state ->
-          {[], [], :nonexistent}
+          {[], [], :nonexistent, if(track_usage, do: [], else: nil)}
       end
 
     state = %__MODULE__{
@@ -128,14 +132,17 @@ 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
     }
 
     {:ok,
      persist(state,
        agent_module: state.agent_module,
        parent_agent_id: parent_agent_id,
-       started_at: NaiveDateTime.utc_now()
+       started_at: NaiveDateTime.utc_now(),
+       usage: state.usage
      ), {:continue, %{start_mode: mode, executor_state: saved_executor_state}}}
   end
 
@@ -229,7 +236,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_usage} =
       Telemetry.span(
         [:legion, :agent, :message],
         %{agent: state.agent_module, message: state.messages |> List.last() |> Map.get(:content)},
@@ -239,7 +246,7 @@ defmodule Legion.AgentServer do
 
           initial_bindings = if conversation_scope?, do: state.bindings, else: []
 
-          {status, value, messages, bindings} =
+          {status, value, messages, bindings, _turn_usage} =
             result =
             Executor.run(
               state.agent_module,
@@ -250,15 +257,31 @@ 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
+           }}
         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}
-      |> persist([:conversation_state, status: :idle])
+      %{
+        state
+        | messages: final_messages,
+          bindings: kept_bindings,
+          usage: usage
+      }
+
+    fields = [:conversation_state, status: :idle, usage: usage]
+    state = persist(state, fields)
 
     {{status, value}, state}
   end
@@ -279,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] ->
+      when field in [:agent_module, :parent_agent_id, :status, :started_at, :usage] ->
         Map.put(payload, field, value)
 
       unknown, _payload ->