agentmux_srv/migrations/
m0026_registry_agent_id_rekey.rs

1// Copyright 2026, AgentMux Corp.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Re-key the named-agent registry from launch ids to agent ids.
5//!
6//! Agent-concept consolidation, Phase 3b PR 2
7//! (`SPEC_AGENT_ARCHITECTURE_2026_05_27.md`): the cross-version/cross-channel
8//! registry under `<home>/shared/agents/registry/` names one file per named
9//! agent. It used to be keyed by the LAUNCH id (`db_agent_instances.id`);
10//! now that an agent is one `db_agents` row, it is keyed by the AGENT id.
11//!
12//! Without this pass an upgraded install keeps its old files under the old
13//! keys, and every one of them is a duplicate the picker shows next to the
14//! agent's real entry — and worse, "Forget agent" retires the NEW key while
15//! the stale file stays active, so the forgotten agent reappears.
16//!
17//! Mapping (read-only against the legacy table, which still exists until
18//! Phase 3c drops it): a record keyed by launch `L` maps to the agent row
19//! that launch folded into — the launch's definition when that definition is
20//! a user agent, otherwise the root of `L`'s continuation chain. Records
21//! whose launch id is not in THIS channel's `db_agent_instances` are left
22//! alone: the registry is shared across channels, so another channel owns
23//! them and will re-key them when it upgrades.
24//!
25//! Global scope: the registry is one host-global tree, and re-keying it twice
26//! (once per channel) would be wasted work — but each channel only resolves
27//! its own launches, so the pass is written to be safely repeatable and is
28//! recorded per install.
29
30use std::collections::HashMap;
31
32use rusqlite::Connection;
33
34use super::{Migration, MigrationContext, MigrationError, MigrationScope};
35use crate::registry::Registry;
36
37pub struct M0026RegistryAgentIdRekey;
38
39/// One legacy launch row, reduced to what the mapping needs.
40struct LaunchRow {
41    definition_id: String,
42    parent_instance_id: String,
43}
44
45/// The agent id a launch folded into, or `None` when this channel cannot
46/// resolve it. Mirrors `instance_create`'s rule: a launch of a USER agent is
47/// that agent; a launch of a template is its own agent, and a continuation
48/// belongs to the launch at the root of its chain.
49fn agent_id_for_launch(
50    launch_id: &str,
51    launches: &HashMap<String, LaunchRow>,
52    user_agent_defs: &dyn Fn(&str) -> bool,
53) -> Option<String> {
54    let mut id = launch_id.to_string();
55    let mut seen = std::collections::HashSet::new();
56    loop {
57        let row = launches.get(&id)?;
58        if user_agent_defs(&row.definition_id) {
59            return Some(row.definition_id.clone());
60        }
61        if row.parent_instance_id.is_empty() || !launches.contains_key(&row.parent_instance_id) {
62            return Some(id);
63        }
64        if !seen.insert(id.clone()) {
65            return Some(id); // cyclic chain — stop where we are
66        }
67        id = row.parent_instance_id.clone();
68    }
69}
70
71/// How many records moved, and how many were left for another channel.
72#[derive(Debug, Default, PartialEq, Eq)]
73pub(super) struct RekeyStats {
74    pub rekeyed: usize,
75    pub already_current: usize,
76    pub unresolved: usize,
77}
78
79pub(super) fn rekey_registry(conn: &Connection, registry: &Registry) -> Result<RekeyStats, String> {
80    let has_legacy: bool = conn
81        .query_row(
82            "SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'db_agent_instances')",
83            [],
84            |row| row.get(0),
85        )
86        .map_err(|e| format!("probe db_agent_instances: {e}"))?;
87    if !has_legacy {
88        return Ok(RekeyStats::default());
89    }
90
91    let mut stmt = conn
92        .prepare("SELECT id, definition_id, parent_instance_id FROM db_agent_instances")
93        .map_err(|e| format!("read db_agent_instances: {e}"))?;
94    let launches: HashMap<String, LaunchRow> = stmt
95        .query_map([], |row| {
96            Ok((
97                row.get::<_, String>(0)?,
98                LaunchRow { definition_id: row.get(1)?, parent_instance_id: row.get(2)? },
99            ))
100        })
101        .and_then(|rows| rows.collect::<Result<HashMap<_, _>, _>>())
102        .map_err(|e| format!("read db_agent_instances: {e}"))?;
103    drop(stmt);
104    if launches.is_empty() {
105        return Ok(RekeyStats::default());
106    }
107
108    let mut user_stmt = conn
109        .prepare("SELECT id FROM db_agents WHERE is_template = 0")
110        .map_err(|e| format!("read db_agents: {e}"))?;
111    let user_agents: std::collections::HashSet<String> = user_stmt
112        .query_map([], |row| row.get::<_, String>(0))
113        .and_then(|rows| rows.collect::<Result<std::collections::HashSet<_>, _>>())
114        .map_err(|e| format!("read db_agents: {e}"))?;
115    drop(user_stmt);
116    let is_user_agent = |id: &str| user_agents.contains(id);
117
118    let mut stats = RekeyStats::default();
119    let records = registry
120        .list_active()
121        .map_err(|e| format!("list registry: {e}"))?;
122
123    // Group by the agent every record re-keys to BEFORE writing anything.
124    // Several legacy head launches of one user agent each have their own
125    // active registry file, and they all collapse onto that agent's single
126    // id. Upserting them in `list_active()`'s filesystem order would let
127    // whichever ran last win, so an older launch could overwrite a newer
128    // `session_id` and the next continuation would `--resume` stale
129    // conversation state (codex P1 on #3080). The newest launch wins
130    // explicitly instead.
131    let mut groups: HashMap<String, Vec<_>> = HashMap::new();
132    for rec in records {
133        let old_id = rec.data.instance_id.clone();
134        let Some(agent_id) = agent_id_for_launch(&old_id, &launches, &is_user_agent) else {
135            stats.unresolved += 1;
136            continue;
137        };
138        groups.entry(agent_id).or_default().push(rec);
139    }
140
141    for (agent_id, mut group) in groups {
142        // Newest launch wins; the launch id breaks a tie so the outcome is
143        // deterministic rather than filesystem-order dependent.
144        group.sort_by(|a, b| {
145            a.data
146                .last_launched_at_ms
147                .cmp(&b.data.last_launched_at_ms)
148                .then_with(|| a.data.instance_id.cmp(&b.data.instance_id))
149        });
150        let mut winner = group.pop().expect("groups only holds non-empty vecs");
151        let mut stale: Vec<String> = group.into_iter().map(|r| r.data.instance_id).collect();
152        if winner.data.instance_id == agent_id {
153            stats.already_current += 1;
154        } else {
155            // The winner's own file is keyed by its launch id and has to go
156            // with the rest once its content lands under the agent id.
157            stale.push(winner.data.instance_id.clone());
158            // Write the record under the agent id, then drop the old files. An
159            // interrupted run leaves both (the next boot re-runs this and the
160            // duplicate resolves), never neither.
161            winner.data.instance_id = agent_id.clone();
162            winner.data.definition_id = agent_id.clone();
163            registry
164                .upsert(&winner)
165                .map_err(|e| format!("write registry record {agent_id}: {e}"))?;
166            stats.rekeyed += 1;
167        }
168        for old_id in stale {
169            if old_id == agent_id {
170                continue;
171            }
172            registry
173                .hard_delete(&old_id)
174                .map_err(|e| format!("remove stale registry record {old_id}: {e}"))?;
175        }
176    }
177    Ok(stats)
178}
179
180impl Migration for M0026RegistryAgentIdRekey {
181    fn id(&self) -> &'static str { "0026_registry_agent_id_rekey" }
182    fn scope(&self) -> MigrationScope { MigrationScope::Channel }
183    fn description(&self) -> &'static str {
184        "Re-key the named-agent registry from launch ids to agent ids"
185    }
186
187    fn up(&self, ctx: &MigrationContext) -> Result<(), MigrationError> {
188        if !ctx.channel_store_path.exists() {
189            return Ok(());
190        }
191        let Some(registry_dir) = crate::registry::resolve_shared_registry_dir() else {
192            // No shared root (CI, a bare `cargo run`) — nothing to re-key.
193            return Ok(());
194        };
195        if !registry_dir.exists() {
196            return Ok(());
197        }
198        let registry = Registry::open(registry_dir)
199            .map_err(|e| MigrationError(format!("registry_agent_id_rekey: open registry: {e}")))?;
200        let conn = Connection::open(&ctx.channel_store_path)
201            .map_err(|e| MigrationError(format!("registry_agent_id_rekey: open store: {e}")))?;
202        let stats = rekey_registry(&conn, &registry)
203            .map_err(|e| MigrationError(format!("registry_agent_id_rekey: {e}")))?;
204        tracing::info!(
205            rekeyed = stats.rekeyed,
206            already_current = stats.already_current,
207            unresolved = stats.unresolved,
208            "m0026_registry_agent_id_rekey: complete"
209        );
210        Ok(())
211    }
212
213    // No `verify()`: the post-condition is "no active record is keyed by a
214    // launch id", which needs the legacy table to state — and that table is
215    // what Phase 3c removes, so the check could only ever run before its own
216    // precondition disappeared.
217}
218
219#[cfg(test)]
220mod tests {
221    use super::*;
222    use crate::backend::storage::store::Store;
223    use crate::registry::{NamedAgentRecord, NamedAgentRecordV1, MAX_SUPPORTED_SCHEMA};
224
225    fn record(instance_id: &str, definition_id: &str, name: &str) -> NamedAgentRecord {
226        NamedAgentRecord {
227            schema_version: MAX_SUPPORTED_SCHEMA,
228            data: NamedAgentRecordV1 {
229                instance_id: instance_id.to_string(),
230                instance_name: name.to_string(),
231                definition_id: definition_id.to_string(),
232                identity_id: None,
233                memory_id: None,
234                session_id: Some("sess".to_string()),
235                working_dir: name.to_string(),
236                source_agents_base: None,
237                created_at_ms: 1,
238                last_launched_at_ms: 1,
239                created_by_version: "0.55.0".to_string(),
240                last_launched_by_version: "0.55.0".to_string(),
241            },
242        }
243    }
244
245    /// A channel store carrying the pre-consolidation launch rows plus the
246    /// consolidated agent rows they map onto.
247    fn store_with_launches() -> (tempfile::TempDir, Connection) {
248        let dir = tempfile::tempdir().unwrap();
249        let path = dir.path().join("objects.db");
250        drop(Store::open(&path).unwrap());
251        let conn = Connection::open(&path).unwrap();
252        conn.execute_batch(
253            // Neither legacy table exists at all any more (v32,
254            // `OBJECT_SCHEMA_VERSION`'s v32 doc comment) — `rekey_registry`
255            // itself never reads `db_agent_definitions` (grep confirms it),
256            // so this fixture only needs enough of it to satisfy
257            // `db_agent_instances.definition_id`'s FK below.
258            "CREATE TABLE db_agent_definitions (
259                id TEXT PRIMARY KEY, slug TEXT NOT NULL DEFAULT '', name TEXT NOT NULL DEFAULT '',
260                icon TEXT NOT NULL DEFAULT '', provider TEXT NOT NULL DEFAULT '', description TEXT NOT NULL DEFAULT '',
261                working_directory TEXT NOT NULL DEFAULT '', shell TEXT NOT NULL DEFAULT '',
262                provider_flags TEXT NOT NULL DEFAULT '', auto_start INTEGER NOT NULL DEFAULT 0,
263                restart_on_crash INTEGER NOT NULL DEFAULT 0, idle_timeout_minutes INTEGER NOT NULL DEFAULT 0,
264                created_at INTEGER NOT NULL DEFAULT 0, agent_type TEXT NOT NULL DEFAULT '',
265                environment TEXT NOT NULL DEFAULT '', agent_bus_id TEXT NOT NULL DEFAULT '',
266                is_seeded INTEGER NOT NULL DEFAULT 0, accounts TEXT NOT NULL DEFAULT '',
267                parent_id TEXT NOT NULL DEFAULT '', branch_label TEXT NOT NULL DEFAULT '',
268                updated_at INTEGER NOT NULL DEFAULT 0
269             );
270             CREATE TABLE db_agent_instances (
271                id TEXT PRIMARY KEY, definition_id TEXT NOT NULL,
272                parent_instance_id TEXT NOT NULL DEFAULT '', block_id TEXT NOT NULL DEFAULT '',
273                session_id TEXT NOT NULL DEFAULT '', status TEXT NOT NULL DEFAULT '',
274                github_context TEXT NOT NULL DEFAULT '', started_at INTEGER NOT NULL DEFAULT 0,
275                ended_at INTEGER NOT NULL DEFAULT 0, created_at INTEGER NOT NULL DEFAULT 0,
276                identity_id TEXT NOT NULL DEFAULT '', memory_id TEXT NOT NULL DEFAULT '',
277                instance_name TEXT NOT NULL DEFAULT '', working_directory TEXT NOT NULL DEFAULT '',
278                display_hidden INTEGER NOT NULL DEFAULT 0,
279                FOREIGN KEY (definition_id) REFERENCES db_agent_definitions(id) ON DELETE CASCADE
280             );
281             INSERT INTO db_agents (id, name, provider, is_template, parent_template_id, is_seeded, created_at, updated_at)
282             VALUES ('user-agent', 'Maks', 'claude', 0, 'tpl', 0, 1, 1),
283                    ('tpl', 'Coder', 'claude', 1, '', 1, 1, 1),
284                    ('launch-head', 'FromTemplate', 'claude', 0, 'tpl', 0, 1, 1);
285             -- The legacy definition rows the launch rows' FK points at.
286             INSERT INTO db_agent_definitions
287                (id, slug, name, icon, provider, description, working_directory, shell,
288                 provider_flags, auto_start, restart_on_crash, idle_timeout_minutes,
289                 created_at, agent_type, environment, agent_bus_id, is_seeded, accounts,
290                 parent_id, branch_label, updated_at)
291             VALUES ('user-agent', 'maks', 'Maks', '', 'claude', '', '', 'bash', '', 0, 0, 0,
292                     1, 'standalone', '', '', 0, '', 'tpl', '', 1),
293                    ('tpl', 'coder', 'Coder', '', 'claude', '', '', 'bash', '', 0, 0, 0,
294                     1, 'standalone', '', '', 1, '', '', '', 1);
295             INSERT INTO db_agent_instances
296                (id, definition_id, parent_instance_id, block_id, session_id, status,
297                 github_context, started_at, ended_at, created_at,
298                 identity_id, memory_id, instance_name, working_directory, display_hidden)
299             VALUES ('launch-on-user', 'user-agent', '', '', '', '', '', 1, 0, 1, '', '', 'Maks', '', 0),
300                    ('launch-head', 'tpl', '', '', '', '', '', 1, 0, 1, '', '', 'FromTemplate', '', 0),
301                    ('launch-cont', 'tpl', 'launch-head', '', '', '', '', 2, 0, 2, '', '', 'FromTemplate', '', 0);",
302        )
303        .unwrap();
304        (dir, conn)
305    }
306
307    fn open_registry(dir: &tempfile::TempDir) -> Registry {
308        Registry::open(dir.path().join("registry")).unwrap()
309    }
310
311    #[test]
312    fn a_launch_of_a_user_agent_rekeys_to_that_agent() {
313        let (dir, conn) = store_with_launches();
314        let reg = open_registry(&dir);
315        reg.upsert(&record("launch-on-user", "user-agent", "Maks")).unwrap();
316
317        let stats = rekey_registry(&conn, &reg).unwrap();
318        assert_eq!(stats.rekeyed, 1, "{stats:?}");
319        assert!(reg.get("launch-on-user").unwrap().is_none(), "the stale file is gone");
320        let moved = reg.get("user-agent").unwrap().expect("re-keyed record");
321        assert_eq!(moved.data.instance_name, "Maks");
322        assert_eq!(moved.data.session_id.as_deref(), Some("sess"), "content survives the move");
323    }
324
325    #[test]
326    fn a_continuation_rekeys_to_the_chain_head_and_a_head_is_already_current() {
327        let (dir, conn) = store_with_launches();
328        let reg = open_registry(&dir);
329        reg.upsert(&record("launch-cont", "tpl", "FromTemplate")).unwrap();
330        let stats = rekey_registry(&conn, &reg).unwrap();
331        assert_eq!((stats.rekeyed, stats.already_current), (1, 0), "{stats:?}");
332        assert!(reg.get("launch-head").unwrap().is_some());
333
334        // The head's own record is already keyed by the agent id.
335        let stats = rekey_registry(&conn, &reg).unwrap();
336        assert_eq!((stats.rekeyed, stats.already_current), (0, 1), "idempotent: {stats:?}");
337    }
338
339    /// Codex P1 on #3080: several legacy launches of ONE user agent each
340    /// carry their own active registry file, and every one of them re-keys
341    /// to that agent's single id. `list_active()` is filesystem-ordered, so
342    /// without an explicit newest-wins rule an older launch's `session_id`
343    /// could be written last and the next continuation would `--resume` a
344    /// conversation the user has since moved on from.
345    #[test]
346    fn several_launches_of_one_agent_collapse_to_the_newest_record() {
347        let (dir, conn) = store_with_launches();
348        conn.execute_batch(
349            "INSERT INTO db_agent_instances
350                (id, definition_id, parent_instance_id, block_id, session_id, status,
351                 github_context, started_at, ended_at, created_at,
352                 identity_id, memory_id, instance_name, working_directory, display_hidden)
353             VALUES ('launch-on-user-2', 'user-agent', '', '', '', '', '', 3, 0, 3, '', '', 'Maks', '', 0);",
354        )
355        .unwrap();
356        let reg = open_registry(&dir);
357        let mut older = record("launch-on-user", "user-agent", "Maks");
358        older.data.last_launched_at_ms = 100;
359        older.data.session_id = Some("stale".to_string());
360        let mut newer = record("launch-on-user-2", "user-agent", "Maks");
361        newer.data.last_launched_at_ms = 200;
362        newer.data.session_id = Some("current".to_string());
363        reg.upsert(&older).unwrap();
364        reg.upsert(&newer).unwrap();
365
366        let stats = rekey_registry(&conn, &reg).unwrap();
367        assert_eq!(stats.rekeyed, 1, "one agent keeps one record: {stats:?}");
368        assert!(reg.get("launch-on-user").unwrap().is_none());
369        assert!(reg.get("launch-on-user-2").unwrap().is_none());
370        let moved = reg.get("user-agent").unwrap().expect("re-keyed record");
371        assert_eq!(
372            moved.data.session_id.as_deref(),
373            Some("current"),
374            "the newest launch's session must win, not whichever file was listed last"
375        );
376    }
377
378    #[test]
379    fn a_record_this_channel_does_not_own_is_left_for_the_channel_that_does() {
380        let (dir, conn) = store_with_launches();
381        let reg = open_registry(&dir);
382        reg.upsert(&record("launch-from-another-channel", "somewhere", "Foreign")).unwrap();
383        let stats = rekey_registry(&conn, &reg).unwrap();
384        assert_eq!(stats.unresolved, 1, "{stats:?}");
385        assert!(
386            reg.get("launch-from-another-channel").unwrap().is_some(),
387            "another channel's record must survive untouched"
388        );
389    }
390
391    #[test]
392    fn a_store_without_the_legacy_table_is_a_no_op() {
393        // A plain `Store::open` already has neither legacy table as of v32
394        // (`OBJECT_SCHEMA_VERSION`'s v32 doc comment) — no DROP needed.
395        let dir = tempfile::tempdir().unwrap();
396        let path = dir.path().join("objects.db");
397        drop(Store::open(&path).unwrap());
398        let conn = Connection::open(&path).unwrap();
399        let reg = open_registry(&dir);
400        reg.upsert(&record("whatever", "whatever", "Whatever")).unwrap();
401        assert_eq!(rekey_registry(&conn, &reg).unwrap(), RekeyStats::default());
402        assert!(reg.get("whatever").unwrap().is_some());
403    }
404}
agentmux_srv/migrations/m0026_registry_agent_id_rekey.rs

agentmux_srv/migrations/
m0026_registry_agent_id_rekey.rs
