From bfe7ed3dc88982dc90b5f61f0870a16de1d52f7b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tomasz=20Id=C5=BAkowski?= <idzkowskit@gmail.com>
Date: Tue, 11 Aug 2026 15:37:05 +0200
Subject: [PATCH] Restrict agent_id to just String

---
 CHANGELOG.md                        |  2 +-
 lib/legion.ex                       | 23 +++++++++++++++++++++--
 lib/legion/agent_server.ex          | 13 ++++++++++++-
 lib/legion/eval_guard.ex            |  2 +-
 lib/legion/store.ex                 |  9 +++++----
 lib/legion/store/postgres.ex        | 27 +++++++++++++++++++++++----
 lib/legion/telemetry.ex             | 16 ++++++++--------
 test/legion/agent_server_test.exs   | 25 +++++++++++++++++++++++++
 test/legion/store/postgres_test.exs | 12 ++++++++++++
 9 files changed, 108 insertions(+), 21 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 028386b..fd3d114 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,7 +10,7 @@
 - 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 configurable, generated agent ids, `Legion.get_agent_id/1`, `Legion.lookup/1`, and `Legion.resume/2` for identifying, finding, and restarting persisted conversations
+- Add configurable, generated UTF-8 string agent ids, `Legion.get_agent_id/1`, `Legion.lookup/1`, and `Legion.resume/2` for identifying, finding, and restarting persisted conversations
 - Register live agents cluster-wide by agent id through `:global`, preventing duplicate ownership across connected nodes
 - Remove the `:name` option from `Legion.start_link/2`; use `:agent_id` for identity and `Legion.lookup/1` to resolve a live process
 - Add `Legion.recover/2` and opt-in `:recovery` startup configuration to automatically drive interrupted root-agent runs (`status: :running`) to completion after an application restart. Recovery skips sub-agents to avoid replaying delegated work.
diff --git a/lib/legion.ex b/lib/legion.ex
index 19643fb..ae46670 100644
--- a/lib/legion.ex
+++ b/lib/legion.ex
@@ -78,6 +78,7 @@ defmodule Legion do
 
   ## Options
     - `:store`, `:agent_id` - persist the conversation across restarts; see `Legion.Store`.
+      When supplied, `:agent_id` must be a valid UTF-8 string.
       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`.
@@ -133,8 +134,8 @@ defmodule Legion do
   end
 
   @doc """
-  Returns the id of a running agent. Always set - Legion generates one when
-  none is passed.
+  Returns the valid UTF-8 string 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.
@@ -180,6 +181,8 @@ defmodule Legion do
   @doc """
   Looks up the live process for an agent id.
 
+  `agent_id` must be a valid UTF-8 string. Raises `ArgumentError` otherwise.
+
   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 on any connected node, or `:error` when no live process owns
@@ -191,6 +194,10 @@ defmodule Legion do
       :error = Legion.lookup("missing_agent_id")
   """
   def lookup(agent_id) do
+    if not is_binary(agent_id) or not String.valid?(agent_id) do
+      raise ArgumentError, ":agent_id must be a valid UTF-8 string, got: #{inspect(agent_id)}"
+    end
+
     case AgentIndex.lookup(agent_id) do
       {:ok, pid} -> {:ok, pid}
       _ -> :error
@@ -200,6 +207,8 @@ defmodule Legion do
   @doc """
   Resumes a persisted conversation.
 
+  `agent_id` must be a valid UTF-8 string. Raises `ArgumentError` otherwise.
+
   Loads the persisted conversation with `get/1` to determine its agent module,
   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
@@ -227,6 +236,10 @@ defmodule Legion do
         Legion.resume("missing_chat", store: MyApp.AgentStore)
   """
   def resume(agent_id, opts \\ []) do
+    if not is_binary(agent_id) or not String.valid?(agent_id) do
+      raise ArgumentError, ":agent_id must be a valid UTF-8 string, got: #{inspect(agent_id)}"
+    end
+
     store = store!(opts, :resume)
 
     case store.get(agent_id) do
@@ -248,6 +261,8 @@ defmodule Legion do
   @doc """
   Recovers an interrupted persisted run and waits for it to finish.
 
+  `agent_id` must be a valid UTF-8 string. Raises `ArgumentError` otherwise.
+
   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.
@@ -277,6 +292,10 @@ defmodule Legion do
         Legion.recover("completed_chat", store: MyApp.AgentStore)
   """
   def recover(agent_id, opts \\ []) do
+    if not is_binary(agent_id) or not String.valid?(agent_id) do
+      raise ArgumentError, ":agent_id must be a valid UTF-8 string, got: #{inspect(agent_id)}"
+    end
+
     store = store!(opts, :recover)
 
     case store.get(agent_id) do
diff --git a/lib/legion/agent_server.ex b/lib/legion/agent_server.ex
index 6e1b818..8e6c881 100644
--- a/lib/legion/agent_server.ex
+++ b/lib/legion/agent_server.ex
@@ -58,7 +58,18 @@ defmodule Legion.AgentServer do
             ":agent_id requires a :store - pass one or set `config :legion, :store, MyStore`"
     end
 
-    agent_id = agent_id || generate_id()
+    agent_id =
+      cond do
+        is_nil(agent_id) ->
+          generate_id()
+
+        is_binary(agent_id) and String.valid?(agent_id) ->
+          agent_id
+
+        true ->
+          raise ArgumentError, ":agent_id must be a valid UTF-8 string, got: #{inspect(agent_id)}"
+      end
+
     persistence_frequency = Store.persistence_frequency(store)
     track_usage = Application.get_env(:legion, :track_usage, true)
 
diff --git a/lib/legion/eval_guard.ex b/lib/legion/eval_guard.ex
index 2751259..072af4d 100644
--- a/lib/legion/eval_guard.ex
+++ b/lib/legion/eval_guard.ex
@@ -57,7 +57,7 @@ defmodule Legion.EvalGuard do
 
   alias Legion.Telemetry
 
-  @type context :: %{agent: module(), agent_id: term(), tools: [module()]}
+  @type context :: %{agent: module(), agent_id: Legion.Store.agent_id(), tools: [module()]}
 
   @callback check(code :: String.t(), context :: context()) :: :allow | {:deny, String.t()}
 
diff --git a/lib/legion/store.ex b/lib/legion/store.ex
index f84c5fb..12d9dea 100644
--- a/lib/legion/store.ex
+++ b/lib/legion/store.ex
@@ -67,7 +67,8 @@ defmodule Legion.Store do
 
   `: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:
+  chat. Agent ids must be valid UTF-8 strings, but Legion otherwise treats their
+  contents as opaque, so compose them however you like:
 
       Legion.start_link(ChatAgent, agent_id: "user_42:chat_7")
 
@@ -114,8 +115,8 @@ defmodule Legion.Store do
 
   ## Reading conversations
 
-  `get/1` returns the persisted conversation for one `agent_id`, or `:error`
-  when the store has no row for that id.
+  `get/1` receives a valid UTF-8 `agent_id` and returns its persisted
+  conversation, or `:error` when the store has no row for that id.
 
   `list/1` returns persisted conversations newest first for consumers that
   rebuild a view of past conversations from the store alone.
@@ -123,7 +124,7 @@ defmodule Legion.Store do
 
   alias Legion.Store.Payload
 
-  @type agent_id :: term()
+  @type agent_id :: String.t()
   @type status :: Payload.status()
   @type payload :: Payload.t()
   @type persistence_frequency :: :turn | :step
diff --git a/lib/legion/store/postgres.ex b/lib/legion/store/postgres.ex
index cba255c..ac07601 100644
--- a/lib/legion/store/postgres.ex
+++ b/lib/legion/store/postgres.ex
@@ -42,7 +42,7 @@ defmodule Legion.Store.Postgres do
           persistence_frequency: :step
       end
 
-  Agent ids must be strings. Snapshots are stored as compressed
+  Agent ids must be valid UTF-8 strings. Snapshots are stored as compressed
   `:erlang.term_to_binary/2` blobs - readable only from Elixir, one row
   per conversation, upserted on every save. Step snapshots therefore require
   no additional migration.
@@ -102,7 +102,18 @@ defmodule Legion.Store.Postgres do
       def persistence_frequency, do: unquote(persistence_frequency)
 
       @impl Legion.Store
-      def get(agent_id) when is_binary(agent_id) do
+      def get(agent_id) when not is_binary(agent_id), do: :error
+
+      @impl Legion.Store
+      def get(agent_id) do
+        if String.valid?(agent_id) do
+          do_get(agent_id)
+        else
+          :error
+        end
+      end
+
+      defp do_get(agent_id) do
         case unquote(repo).get(Record, agent_id) do
           nil -> :error
           record -> {:ok, Postgres.decode_record(record)}
@@ -119,10 +130,18 @@ defmodule Legion.Store.Postgres do
       def save(map) when is_map(map) and not is_struct(map), do: :error
 
       @impl Legion.Store
-      def save(%Payload{agent_id: nil}), do: :error
+      def save(%Payload{agent_id: agent_id}) when not is_binary(agent_id), do: :error
 
       @impl Legion.Store
-      def save(%Payload{} = payload) do
+      def save(%Payload{agent_id: agent_id} = payload) do
+        if String.valid?(agent_id) do
+          do_save(payload)
+        else
+          :error
+        end
+      end
+
+      defp do_save(payload) do
         attrs =
           payload
           |> Postgres.encode_data()
diff --git a/lib/legion/telemetry.ex b/lib/legion/telemetry.ex
index 7b1edec..2a88e75 100644
--- a/lib/legion/telemetry.ex
+++ b/lib/legion/telemetry.ex
@@ -8,18 +8,18 @@ defmodule Legion.Telemetry do
 
   - `[:legion, :agent, :started]` — agent process finished `init/1`
     - Measurements: `%{system_time: NaiveDateTime.t()}`
-    - Metadata: `%{agent: module, agent_id: term, parent_agent_id: term}`
+    - Metadata: `%{agent: module, agent_id: String.t()}` (plus `parent_agent_id: String.t()`
+      when the agent was started inside another agent's run)
     - `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.
+    - 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: NaiveDateTime.t()}`
-    - Metadata: `%{agent: module, agent_id: term}` (plus `parent_agent_id`
-      when the parent's run is still on the process Vault)
+    - Metadata: `%{agent: module, agent_id: String.t()}` (plus `parent_agent_id: String.t()`
+      when the agent was started inside another agent's run)
 
   ## Agent Message Events (spans)
 
@@ -50,7 +50,7 @@ defmodule Legion.Telemetry do
   ## Eval Guard Events
 
   - `[:legion, :eval_guard, :denied]` — a guard refused generated code
-    - Metadata: `%{agent: module, agent_id: term, guard: module, code: String.t(), reason: String.t()}`
+    - Metadata: `%{agent: module, agent_id: String.t(), guard: module, code: String.t(), reason: String.t()}`
 
   ## Default Logger
 
diff --git a/test/legion/agent_server_test.exs b/test/legion/agent_server_test.exs
index 17bc6ae..b8ad4c4 100644
--- a/test/legion/agent_server_test.exs
+++ b/test/legion/agent_server_test.exs
@@ -965,6 +965,7 @@ defmodule Legion.AgentServerTest do
       agent_id = Legion.get_agent_id(pid)
 
       assert is_binary(agent_id)
+      assert String.valid?(agent_id)
       {:ok, _} = Legion.call(pid, "What is the capital of France?")
       assert {:ok, _snapshot} = MemoryStore.load(agent_id)
     end
@@ -994,6 +995,30 @@ defmodule Legion.AgentServerTest do
       end
     end
 
+    test "rejects non-string explicit agent IDs" do
+      for agent_id <- [:agent, make_ref(), <<0xFF>>] do
+        assert_raise ArgumentError, ~r/:agent_id must be a valid UTF-8 string/, fn ->
+          Legion.start_link(MathAgent, store: MemoryStore, agent_id: agent_id)
+        end
+      end
+    end
+
+    test "public identity operations reject invalid non-string agent IDs" do
+      non_binary_agent_id = :agent
+      non_utf8_agent_id = <<0xFF>>
+
+      for operation <- [
+            fn -> Legion.lookup(non_binary_agent_id) end,
+            fn -> Legion.resume(non_binary_agent_id, store: MemoryStore) end,
+            fn -> Legion.recover(non_binary_agent_id, store: MemoryStore) end,
+            fn -> Legion.lookup(non_utf8_agent_id) end,
+            fn -> Legion.resume(non_utf8_agent_id, store: MemoryStore) end,
+            fn -> Legion.recover(non_utf8_agent_id, store: MemoryStore) end
+          ] do
+        assert_raise ArgumentError, ~r/:agent_id must be a valid UTF-8 string/, operation
+      end
+    end
+
     test "records run metadata on start" do
       {:ok, _pid} = Legion.start_link(MathAgent, store: MemoryStore, agent_id: "meta")
 
diff --git a/test/legion/store/postgres_test.exs b/test/legion/store/postgres_test.exs
index 2c0225b..608c76b 100644
--- a/test/legion/store/postgres_test.exs
+++ b/test/legion/store/postgres_test.exs
@@ -163,4 +163,16 @@ defmodule Legion.Store.PostgresTest do
     assert :error = Store.save(%{agent_id: "user_42", unexpected: "value"})
     assert FakeRepo.run("user_42") == nil
   end
+
+  test "save/1 rejects a payload with non-string agent_id without inserting a row" do
+    assert :error = Store.save(%Payload{agent_id: 42})
+    assert FakeRepo.run(42) == nil
+    assert :error = Store.save(%Payload{agent_id: <<0xFF>>})
+    assert FakeRepo.run(<<0xFF>>) == nil
+  end
+
+  test "get/1 rejects a non-string agent_id" do
+    assert :error = Store.get(42)
+    assert :error = Store.get(<<0xFF>>)
+  end
 end