Add session-scoped stores for stable model-facing IDs (#61339)

## What changed

- Add `SessionModelIdStore` with SQLite and in-memory implementations for looking up and assigning aliases by session, namespace, and external key.
- Commit SQLite assignments before returning them, and recheck under a write transaction so concurrent writers reuse the same alias.
- Support snapshotting and initializing alias maps from generator state and ordered keys, and remove mappings when their thread is deleted.
- Expose alias stores through `ThreadStore` and add distinct types for registered environment IDs and model-facing `env` aliases.

## Testing

Add a test using two SQLite runtime instances to verify that concurrent assignments return the same alias, different keys receive different aliases, and lookups and snapshots preserve the assignments.

GitOrigin-RevId: 2668f605251ca76d4891f70f94bfb4151ba8142c

codex-rs/Cargo.lock

@@ -5084,10 +5084,12 @@ dependencies = [
  "codex-history",
  "codex-protocol",
  "codex-utils-absolute-path",
+ "codex-utils-model-facing-id",
  "file-id",
  "libsqlite3-sys",
  "log",
  "pretty_assertions",
+ "rand 0.9.3",
  "scopeguard",
  "serde",
  "serde_json",

codex-rs/state/Cargo.toml

@@ -11,9 +11,11 @@ codex-file-system = { workspace = true }
 codex-history = { workspace = true }
 codex-protocol = { workspace = true }
 codex-utils-absolute-path = { workspace = true }
+codex-utils-model-facing-id = { workspace = true }
 file-id = { workspace = true }
 libsqlite3-sys = { workspace = true }
 log = { workspace = true }
+rand.workspace = true
 serde = { workspace = true, features = ["derive"] }
 serde_json = { workspace = true }
 sqlx = { workspace = true }

codex-rs/state/migrations/0063_session_model_ids.sql

@@ -0,0 +1,10 @@
+CREATE TABLE session_model_ids (
+    session_id TEXT NOT NULL,
+    kind TEXT NOT NULL,
+    external_key TEXT NOT NULL,
+    ordinal INTEGER NOT NULL CHECK (ordinal >= 0),
+    initial_state TEXT,
+    PRIMARY KEY (session_id, kind, external_key),
+    UNIQUE (session_id, kind, ordinal),
+    CHECK ((ordinal = 0) = (initial_state IS NOT NULL))
+);

codex-rs/state/src/lib.rs

@@ -106,6 +106,9 @@ pub use runtime::ThreadReadState;
 pub use runtime::backup_runtime_db_for_fresh_start;
 pub use runtime::collect_runtime_db_backups;
 pub use runtime::is_sqlite_corruption_error;
+pub use runtime::model_ids::MemoryModelIdStore;
+pub use runtime::model_ids::ModelIdStoreFuture;
+pub use runtime::model_ids::SessionModelIdStore;
 pub use runtime::open_thread_history_db;
 pub use runtime::recovery::RuntimeDbBackup;
 pub use runtime::runtime_db_path_for_corruption_error;

codex-rs/state/src/runtime.rs

@@ -48,6 +48,7 @@ mod logs;
 mod logs_maintenance;
 mod memories;
 mod memory_versions;
+pub(crate) mod model_ids;
 mod projects;
 mod queued_items;
 pub(crate) mod reclamation;

codex-rs/state/src/runtime/model_ids.rs

@@ -0,0 +1,251 @@
+//! Session-scoped aliases. SQLite commits each assignment before the model can see it.
+
+use std::collections::HashMap;
+use std::future::Future;
+use std::pin::Pin;
+use std::sync::Mutex;
+
+use anyhow::Result;
+use codex_protocol::SessionId;
+use codex_utils_model_facing_id::ModelFacingId;
+use codex_utils_model_facing_id::ModelFacingIdGenerator;
+use codex_utils_model_facing_id::ModelFacingIdGeneratorState;
+use codex_utils_model_facing_id::ModelIdKey;
+use codex_utils_model_facing_id::ModelIdMap;
+use sqlx::SqliteConnection;
+use sqlx::types::Json;
+
+use super::StateRuntime;
+use crate::SqlitePoolWriteExt;
+
+pub type ModelIdStoreFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T>> + Send + 'a>>;
+
+/// Stable aliases shared by a session's root and workers. A fork reads the source maps once
+/// and installs them in its new session before it can issue an alias of its own.
+pub trait SessionModelIdStore: Send + Sync {
+    fn lookup(
+        &self,
+        session: SessionId,
+        key: ModelIdKey,
+    ) -> ModelIdStoreFuture<'_, Option<ModelFacingId>>;
+    fn get_or_assign(
+        &self,
+        session: SessionId,
+        key: ModelIdKey,
+    ) -> ModelIdStoreFuture<'_, ModelFacingId>;
+    fn snapshot(&self, session: SessionId) -> ModelIdStoreFuture<'_, Vec<ModelIdMap>>;
+    fn initialize(&self, session: SessionId, maps: Vec<ModelIdMap>) -> ModelIdStoreFuture<'_, ()>;
+}
+
+impl SessionModelIdStore for StateRuntime {
+    fn lookup(
+        &self,
+        session: SessionId,
+        key: ModelIdKey,
+    ) -> ModelIdStoreFuture<'_, Option<ModelFacingId>> {
+        Box::pin(async move {
+            let mut conn = self.pool.acquire().await?;
+            lookup(&mut conn, &session.to_string(), &key).await
+        })
+    }
+
+    fn get_or_assign(
+        &self,
+        session: SessionId,
+        key: ModelIdKey,
+    ) -> ModelIdStoreFuture<'_, ModelFacingId> {
+        Box::pin(async move {
+            let session = session.to_string();
+            {
+                let mut conn = self.pool.acquire().await?;
+                if let Some(id) = lookup(&mut conn, &session, &key).await? {
+                    return Ok(id);
+                }
+            }
+            // Another writer may have assigned the key since the read, so check again under the lock.
+            let mut tx = self.pool.begin_immediate().await?;
+            if let Some(id) = lookup(&mut tx, &session, &key).await? {
+                tx.commit().await?;
+                return Ok(id);
+            }
+            let initial: Option<Json<ModelFacingIdGeneratorState>> = sqlx::query_scalar(
+                "SELECT initial_state FROM session_model_ids WHERE session_id=? AND kind=? AND ordinal=0",
+            )
+            .bind(&session)
+            .bind(key.kind)
+            .fetch_optional(&mut *tx)
+            .await?;
+            let first = initial.is_none();
+            let initial = initial.unwrap_or_else(|| {
+                Json(ModelFacingIdGenerator::new(rand::random(), key.kind).snapshot())
+            });
+            let ordinal: i64 = sqlx::query_scalar(
+                "INSERT INTO session_model_ids (session_id, kind, external_key, ordinal, initial_state)
+                 SELECT ?1, ?2, ?3, COALESCE(MAX(ordinal)+1, 0), ?4 FROM session_model_ids
+                 WHERE session_id=?1 AND kind=?2
+                 RETURNING ordinal",
+            )
+            .bind(&session)
+            .bind(key.kind)
+            .bind(key.key)
+            .bind(first.then_some(&initial))
+            .fetch_one(&mut *tx)
+            .await?;
+            let id = initial.id_at(ordinal as u64)?;
+            tx.commit().await?;
+            Ok(id)
+        })
+    }
+
+    fn snapshot(&self, session: SessionId) -> ModelIdStoreFuture<'_, Vec<ModelIdMap>> {
+        Box::pin(async move {
+            let rows: Vec<(Json<ModelFacingIdGeneratorState>, Json<Vec<String>>)> = sqlx::query_as(
+                "SELECT MAX(initial_state), json_group_array(external_key ORDER BY ordinal)
+                 FROM session_model_ids WHERE session_id=? GROUP BY kind ORDER BY kind",
+            )
+            .bind(session.to_string())
+            .fetch_all(self.pool.as_ref())
+            .await?;
+            Ok(rows
+                .into_iter()
+                .map(|(Json(initial), Json(keys))| ModelIdMap { initial, keys })
+                .collect())
+        })
+    }
+
+    fn initialize(&self, session: SessionId, maps: Vec<ModelIdMap>) -> ModelIdStoreFuture<'_, ()> {
+        Box::pin(async move {
+            let mut tx = self.pool.begin_immediate().await?;
+            let session = session.to_string();
+            for map in maps {
+                sqlx::query(
+                    "INSERT INTO session_model_ids (session_id, kind, external_key, ordinal, initial_state)
+                     SELECT ?1, ?2, value, key, CASE WHEN key = 0 THEN ?3 END FROM json_each(?4)",
+                )
+                .bind(&session)
+                .bind(map.initial.salt())
+                .bind(Json(&map.initial))
+                .bind(Json(map.keys))
+                .execute(&mut *tx)
+                .await?;
+            }
+            tx.commit().await?;
+            Ok(())
+        })
+    }
+}
+
+async fn lookup(
+    conn: &mut SqliteConnection,
+    session: &str,
+    key: &ModelIdKey,
+) -> Result<Option<ModelFacingId>> {
+    let row: Option<(Json<ModelFacingIdGeneratorState>, i64)> = sqlx::query_as(
+        "SELECT first.initial_state, b.ordinal FROM session_model_ids b
+         JOIN session_model_ids first USING (session_id, kind)
+         WHERE b.session_id=? AND b.kind=? AND b.external_key=? AND first.ordinal=0",
+    )
+    .bind(session)
+    .bind(key.kind)
+    .bind(&key.key)
+    .fetch_optional(conn)
+    .await?;
+    let Some((initial, ordinal)) = row else {
+        return Ok(None);
+    };
+    Ok(Some(initial.id_at(ordinal as u64)?))
+}
+
+/// Remove the names with the rest of the session state in the caller's delete transaction.
+pub(super) async fn delete_session(conn: &mut SqliteConnection, session: &str) -> Result<()> {
+    sqlx::query("DELETE FROM session_model_ids WHERE session_id = ?")
+        .bind(session)
+        .execute(conn)
+        .await?;
+    Ok(())
+}
+
+/// Uses the same namespace and fork semantics for ephemeral sessions and in-memory thread stores.
+#[derive(Default, Debug)]
+pub struct MemoryModelIdStore(Mutex<HashMap<SessionId, Vec<ModelIdMap>>>);
+
+impl MemoryModelIdStore {
+    pub fn remove(&self, session: SessionId) {
+        self.0
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner)
+            .remove(&session);
+    }
+}
+
+impl SessionModelIdStore for MemoryModelIdStore {
+    fn lookup(
+        &self,
+        session: SessionId,
+        key: ModelIdKey,
+    ) -> ModelIdStoreFuture<'_, Option<ModelFacingId>> {
+        Box::pin(async move {
+            let sessions = self
+                .0
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner);
+            let Some(map) = sessions
+                .get(&session)
+                .and_then(|maps| maps.iter().find(|map| map.initial.salt() == key.kind))
+            else {
+                return Ok(None);
+            };
+            Ok(map.lookup(&key.key)?)
+        })
+    }
+
+    fn get_or_assign(
+        &self,
+        session: SessionId,
+        key: ModelIdKey,
+    ) -> ModelIdStoreFuture<'_, ModelFacingId> {
+        Box::pin(async move {
+            let mut sessions = self
+                .0
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner);
+            let maps = sessions.entry(session).or_default();
+            let map = match maps.iter().position(|map| map.initial.salt() == key.kind) {
+                Some(index) => &mut maps[index],
+                None => maps.push_mut(ModelIdMap::new(rand::random(), key.kind)),
+            };
+            if let Some(id) = map.lookup(&key.key)? {
+                return Ok(id);
+            }
+            let id = map.initial.id_at(map.keys.len() as u64)?;
+            map.keys.push(key.key);
+            Ok(id)
+        })
+    }
+
+    fn snapshot(&self, session: SessionId) -> ModelIdStoreFuture<'_, Vec<ModelIdMap>> {
+        Box::pin(async move {
+            Ok(self
+                .0
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner)
+                .get(&session)
+                .cloned()
+                .unwrap_or_default())
+        })
+    }
+
+    fn initialize(&self, session: SessionId, maps: Vec<ModelIdMap>) -> ModelIdStoreFuture<'_, ()> {
+        Box::pin(async move {
+            self.0
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner)
+                .insert(session, maps);
+            Ok(())
+        })
+    }
+}
+
+#[cfg(test)]
+#[path = "model_ids_tests.rs"]
+mod tests;

codex-rs/state/src/runtime/model_ids_tests.rs

@@ -0,0 +1,30 @@
+#![cfg(test)]
+
+use super::*;
+use crate::SqliteConfig;
+use codex_utils_absolute_path::test_support::PathExt;
+use codex_utils_model_facing_id::RegisteredEnvironmentId;
+use pretty_assertions::assert_eq;
+
+/// Concurrent writers reserve a key once, and an independent connection observes the winner.
+#[tokio::test]
+async fn concurrent_assignments_share_a_durable_alias() -> Result<()> {
+    let home = tempfile::tempdir()?;
+    let config = SqliteConfig::new_for_testing(home.path().abs());
+    let a = StateRuntime::init(config.clone(), "test".into()).await?;
+    let b = StateRuntime::init(config, "test".into()).await?;
+    let session = SessionId::new();
+    let key = ModelIdKey::new(&RegisteredEnvironmentId::new("first"))?;
+    let (left, right) = tokio::try_join!(
+        a.get_or_assign(session, key.clone()),
+        b.get_or_assign(session, key.clone()),
+    )?;
+    assert_eq!(left, right);
+    let second = ModelIdKey::new(&RegisteredEnvironmentId::new("second"))?;
+    let next = b.get_or_assign(session, second.clone()).await?;
+    assert_ne!(left, next);
+    assert_eq!(a.lookup(session, key.clone()).await?, Some(left));
+    let maps = a.snapshot(session).await?;
+    assert_eq!(maps[0].keys, vec![key.key, second.key]);
+    Ok(())
+}

codex-rs/state/src/runtime/threads.rs

@@ -1298,6 +1298,7 @@ ON CONFLICT(id) DO UPDATE SET
         }
         let mut rows_affected = 0;
         for thread_id_string in &thread_id_strings {
+            super::model_ids::delete_session(&mut tx, thread_id_string).await?;
             sqlx::query("DELETE FROM thread_read_receipts WHERE thread_id = ?")
                 .bind(thread_id_string)
                 .execute(&mut *tx)

codex-rs/thread-store/src/in_memory.rs

@@ -525,6 +525,7 @@ pub struct InMemoryThreadStoreCalls {
 pub struct InMemoryThreadStore {
     state: Arc<tokio::sync::Mutex<InMemoryThreadStoreState>>,
     omit_metadata_update_result: Arc<AtomicBool>,
+    model_ids: Arc<codex_state::MemoryModelIdStore>,
     state_db: Option<codex_rollout::StateDbHandle>,
 }
 
@@ -558,6 +559,7 @@ impl InMemoryThreadStore {
         Self {
             state: Arc::clone(&self.state),
             omit_metadata_update_result: Arc::clone(&self.omit_metadata_update_result),
+            model_ids: Arc::clone(&self.model_ids),
             state_db,
         }
     }
@@ -906,6 +908,7 @@ impl InMemoryThreadStore {
             0
         };
         let mut state = self.state.lock().await;
+        self.model_ids.remove(params.thread_id.into());
         let existed = state.histories.remove(&params.thread_id).is_some();
         state.created_threads.remove(&params.thread_id);
         state.names.remove(&params.thread_id);
@@ -927,6 +930,10 @@ impl InMemoryThreadStore {
 }
 
 impl ThreadStore for InMemoryThreadStore {
+    fn model_ids(&self) -> Option<Arc<dyn codex_state::SessionModelIdStore>> {
+        Some(self.model_ids.clone())
+    }
+
     fn as_any(&self) -> &dyn std::any::Any {
         self
     }

codex-rs/thread-store/src/local/mod.rs

@@ -516,6 +516,12 @@ impl LocalThreadStore {
 }
 
 impl ThreadStore for LocalThreadStore {
+    fn model_ids(&self) -> Option<Arc<dyn codex_state::SessionModelIdStore>> {
+        self.state_db
+            .clone()
+            .map(|db| db as Arc<dyn codex_state::SessionModelIdStore>)
+    }
+
     fn default_history_mode(&self) -> ThreadHistoryMode {
         ThreadHistoryMode::Paginated
     }

codex-rs/thread-store/src/store.rs

@@ -146,6 +146,11 @@ pub trait ThreadStore: Any + Send + Sync {
     /// Return this store as [`Any`] for implementation-owned escape hatches.
     fn as_any(&self) -> &dyn Any;
 
+    /// Returns the authority for session aliases when this backend supports them.
+    fn model_ids(&self) -> Option<std::sync::Arc<dyn codex_state::SessionModelIdStore>> {
+        None
+    }
+
     /// Returns the history mode to use when history does not carry a persisted mode.
     ///
     /// The default is legacy so existing stores stay compatible. Stores whose durable contract is

codex-rs/utils/model-facing-id/Cargo.toml

@@ -15,9 +15,9 @@ test = false
 blake3.workspace = true
 schemars.workspace = true
 serde = { workspace = true, features = ["derive"] }
+serde_json.workspace = true
 thiserror.workspace = true
 tracing.workspace = true
 
 [dev-dependencies]
-serde_json.workspace = true
 pretty_assertions.workspace = true

codex-rs/utils/model-facing-id/src/bindings.rs

@@ -0,0 +1,59 @@
+//! Stable names for external objects in low-cardinality ID spaces.
+
+use std::hash::Hash;
+
+use serde::Serialize;
+use serde::de::DeserializeOwned;
+
+use crate::ModelFacingId;
+use crate::ModelFacingIdGenerator;
+use crate::ModelFacingIdGeneratorState;
+use crate::RestoreError;
+
+/// The key type selects the alias type and a unique, stable prefix. Its JSON encoding must
+/// distinguish unequal keys and remain stable across releases.
+pub trait ModelIdKind: Eq + Hash + Serialize + DeserializeOwned {
+    const PREFIX: &'static str;
+    type Alias: From<ModelFacingId>;
+}
+
+#[derive(Clone)]
+pub struct ModelIdKey {
+    pub kind: &'static str,
+    pub key: String,
+}
+
+impl ModelIdKey {
+    pub fn new<K: ModelIdKind>(key: &K) -> serde_json::Result<Self> {
+        let mut key = serde_json::to_value(key)?;
+        key.sort_all_objects();
+        Ok(Self {
+            kind: K::PREFIX,
+            key: key.to_string(),
+        })
+    }
+}
+
+/// A namespace needs only its initial generator state and the external keys in allocation order.
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct ModelIdMap {
+    pub initial: ModelFacingIdGeneratorState,
+    pub keys: Vec<String>,
+}
+
+impl ModelIdMap {
+    pub fn new(seed: u64, kind: &'static str) -> Self {
+        Self {
+            initial: ModelFacingIdGenerator::new(seed, kind).snapshot(),
+            keys: Vec::new(),
+        }
+    }
+
+    pub fn lookup(&self, key: &str) -> Result<Option<ModelFacingId>, RestoreError> {
+        self.keys
+            .iter()
+            .position(|bound| bound == key)
+            .map(|ordinal| self.initial.id_at(ordinal as u64))
+            .transpose()
+    }
+}

codex-rs/utils/model-facing-id/src/environment.rs

@@ -0,0 +1,56 @@
+//! Distinguishes client executor IDs from aliases accepted in model-facing environment arguments.
+
+use std::fmt;
+
+use serde::Deserialize;
+use serde::Serialize;
+
+use crate::ModelFacingId;
+use crate::ModelIdKind;
+
+/// An ID assigned to an executor by the client.
+#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
+#[serde(transparent)]
+pub struct RegisteredEnvironmentId(String);
+
+impl RegisteredEnvironmentId {
+    pub fn new(id: impl Into<String>) -> Self {
+        Self(id.into())
+    }
+
+    pub fn as_str(&self) -> &str {
+        &self.0
+    }
+}
+
+impl fmt::Display for RegisteredEnvironmentId {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_str(self.as_str())
+    }
+}
+
+/// The environment namespace supplies the prefix when this value is displayed.
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+pub struct ModelEnvironmentId(u64);
+
+impl fmt::Display for ModelEnvironmentId {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "{}{}",
+            RegisteredEnvironmentId::PREFIX,
+            crate::padded_digits(self.0)
+        )
+    }
+}
+
+impl ModelIdKind for RegisteredEnvironmentId {
+    const PREFIX: &'static str = "env";
+    type Alias = ModelEnvironmentId;
+}
+
+impl From<ModelFacingId> for ModelEnvironmentId {
+    fn from(id: ModelFacingId) -> Self {
+        Self(id.value)
+    }
+}

codex-rs/utils/model-facing-id/src/lib.rs

@@ -11,6 +11,14 @@ use serde::Deserialize;
 use serde::Serialize;
 use thiserror::Error;
 
+mod bindings;
+mod environment;
+pub use bindings::ModelIdKey;
+pub use bindings::ModelIdKind;
+pub use bindings::ModelIdMap;
+pub use environment::ModelEnvironmentId;
+pub use environment::RegisteredEnvironmentId;
+
 const VERSION: u32 = 1;
 // Our tokenizers encode every decimal ID of up to three digits as a single token.
 const GROUP_SIZE: u64 = 1_000;
@@ -47,17 +55,14 @@ impl ModelFacingId {
 
     /// Formats the padded digits for protocols that supply their own type prefix.
     pub fn digits(&self) -> impl fmt::Display + '_ {
-        fmt::from_fn(|f| {
-            write!(
-                f,
-                "{:0width$}",
-                self.value,
-                width = decimal_width(self.value)
-            )
-        })
+        padded_digits(self.value)
     }
 }
 
+fn padded_digits(value: u64) -> impl fmt::Display {
+    fmt::from_fn(move |f| write!(f, "{value:0width$}", width = decimal_width(value)))
+}
+
 fn decimal_width(value: u64) -> usize {
     (value.checked_ilog10().unwrap_or(0) as usize + 1).next_multiple_of(3)
 }
@@ -153,6 +158,13 @@ pub struct ModelFacingIdGeneratorState {
 }
 
 impl ModelFacingIdGeneratorState {
+    /// Recovers the ID at a stored binding's ordinal in this version of the sequence.
+    pub fn id_at(&self, ordinal: u64) -> Result<ModelFacingId, RestoreError> {
+        let mut state = self.clone();
+        state.next_ordinal = ordinal;
+        Ok(ModelFacingIdGenerator::restore(state)?.fetch_next())
+    }
+
     /// Returns the namespace this snapshot belongs to.
     pub fn salt(&self) -> &str {
         &self.salt