Support metadata staging for reserved thread IDs (#38819)

## Why

Callers need to associate host-owned state with a thread before Core starts it.

## What changed

- Add `ThreadManager::reserve_thread_id` and allow new threads to start with the reserved ID. Reject reserved IDs when resuming an existing thread.
- Let thread stores stage metadata for a reserved ID. The local store merges staged values into the first successful metadata update and then removes them.
- Clear staged metadata when an unmaterialized thread shuts down or is discarded, and require the state database for staging.

## Testing

- Cover reserved ID generation and resume validation.
- Cover metadata persistence, merging, rollout-compatible updates, cleanup, and invalid staging requests.

GitOrigin-RevId: 3f9bc48f0f7764011cc90ea01349e497dda0a9e3

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

@@ -156,6 +156,7 @@ pub(super) async fn shutdown_thread(
     store: &LocalThreadStore,
     thread_id: ThreadId,
 ) -> ThreadStoreResult<()> {
+    let mut pending_metadata = store.pending_thread_metadata.lock(thread_id).await;
     let _live_writer_guard = store.live_writer_locks.lock(thread_id).await;
     let (recorder, rollout_id, history_mode) = live_writer_parts(store, thread_id).await?;
     let rollout_path = recorder.rollout_path().to_path_buf();
@@ -175,12 +176,25 @@ pub(super) async fn shutdown_thread(
     }
     sync_materialized_rollout_path(store, thread_id, rollout_path.as_path()).await?;
     if let Some(metrics) = codex_otel::global()
-        && let Ok(metadata) = tokio::fs::metadata(rollout_path).await
+        && let Ok(metadata) = tokio::fs::metadata(&rollout_path).await
     {
         let size_bytes = i64::try_from(metadata.len()).unwrap_or(i64::MAX);
         let _ = metrics.histogram(ROLLOUT_SIZE_BYTES_METRIC, size_bytes, &[]);
     }
+    let rollout_exists = match tokio::fs::try_exists(&rollout_path).await {
+        Ok(rollout_exists) => rollout_exists,
+        Err(err) => {
+            warn!(
+                "failed to check rollout path after shutdown for {thread_id}; preserving pending metadata: {err}"
+            );
+            true
+        }
+    };
     store.live_recorders.lock().await.remove(&thread_id);
+    drop(_live_writer_guard);
+    if !rollout_exists && pending_metadata.take().is_some() {
+        store.pending_thread_metadata.remove(thread_id).await;
+    }
     Ok(())
 }
 
@@ -188,7 +202,11 @@ pub(super) async fn discard_thread(
     store: &LocalThreadStore,
     thread_id: ThreadId,
 ) -> ThreadStoreResult<()> {
+    let mut pending_metadata = store.pending_thread_metadata.lock(thread_id).await;
     let _live_writer_guard = store.live_writer_locks.lock(thread_id).await;
+    if pending_metadata.take().is_some() {
+        store.pending_thread_metadata.remove(thread_id).await;
+    }
     store
         .live_recorders
         .lock()
@@ -277,7 +295,7 @@ enum RolloutWriteOp {
     Flush,
 }
 
-async fn live_writer_parts(
+pub(super) async fn live_writer_parts(
     store: &LocalThreadStore,
     thread_id: ThreadId,
 ) -> ThreadStoreResult<(RolloutRecorder, ThreadId, ThreadHistoryMode)> {

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

@@ -7,6 +7,7 @@ mod live_writer;
 mod model_context;
 mod move_thread_to_section;
 mod paginated_fork;
+mod pending_thread_metadata;
 mod read_thread;
 mod revert_thread;
 mod rollout_migration;
@@ -22,6 +23,9 @@ mod unarchive_thread;
 mod update_thread_metadata;
 mod writer_lock;
 
+#[cfg(test)]
+#[path = "pending_thread_metadata_tests.rs"]
+mod pending_thread_metadata_tests;
 #[cfg(test)]
 mod test_support;
 
@@ -71,6 +75,7 @@ use crate::StoredThread;
 use crate::StoredThreadHistory;
 use crate::StoredThreadSection;
 use crate::StoredThreadSectionsPage;
+use crate::ThreadMetadataPatch;
 use crate::ThreadOccurrenceSearchPage;
 use crate::ThreadPage;
 use crate::ThreadSearchPage;
@@ -107,6 +112,7 @@ pub use rollout_migration::RolloutMigrationStatus;
 pub struct LocalThreadStore {
     pub(super) config: LocalThreadStoreConfig,
     live_recorders: Arc<Mutex<HashMap<ThreadId, LiveRecorderEntry>>>,
+    pending_thread_metadata: pending_thread_metadata::PendingThreadMetadataRegistry,
     live_writer_locks: Arc<LiveWriterLocks>,
     writer_lock_coordinator: Arc<WriterLockCoordinator>,
     state_db: Option<StateDbHandle>,
@@ -218,6 +224,8 @@ impl LocalThreadStore {
         Self {
             config,
             live_recorders: Arc::new(Mutex::new(HashMap::new())),
+            pending_thread_metadata:
+                pending_thread_metadata::PendingThreadMetadataRegistry::default(),
             live_writer_locks: Arc::new(LiveWriterLocks::default()),
             writer_lock_coordinator,
             state_db,
@@ -401,6 +409,33 @@ impl ThreadStore for LocalThreadStore {
         Box::pin(async move { live_writer::create_thread(self, params).await })
     }
 
+    fn stage_pending_thread_metadata(
+        &self,
+        thread_id: ThreadId,
+        patch: ThreadMetadataPatch,
+    ) -> ThreadStoreFuture<'_, ()> {
+        Box::pin(async move {
+            if self.state_db.is_none() {
+                return Err(ThreadStoreError::InvalidRequest {
+                    message: "pending thread metadata requires a state db".to_string(),
+                });
+            }
+            if patch.rollout_path.is_some() {
+                return Err(ThreadStoreError::InvalidRequest {
+                    message: "pending thread metadata cannot set rollout_path".to_string(),
+                });
+            }
+            self.pending_thread_metadata.stage(thread_id, patch).await
+        })
+    }
+
+    fn remove_pending_thread_metadata(&self, thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> {
+        Box::pin(async move {
+            self.pending_thread_metadata.remove(thread_id).await;
+            Ok(())
+        })
+    }
+
     fn resume_thread(&self, params: ResumeThreadParams) -> ThreadStoreFuture<'_, ()> {
         Box::pin(async move { live_writer::resume_thread(self, params).await })
     }

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

@@ -0,0 +1,58 @@
+use std::collections::HashMap;
+use std::collections::hash_map::Entry;
+use std::sync::Arc;
+
+use codex_protocol::ThreadId;
+use tokio::sync::Mutex;
+use tokio::sync::OwnedMutexGuard;
+
+use crate::ThreadMetadataPatch;
+use crate::ThreadStoreError;
+use crate::ThreadStoreResult;
+
+type PendingThreadMetadataEntry = Arc<Mutex<Option<ThreadMetadataPatch>>>;
+
+/// In-memory metadata staged before a reserved thread starts.
+///
+/// The registry lock is held only while looking up or removing an entry. Never await an entry
+/// lock while holding the registry lock. Callers may hold an entry lock before the live-writer
+/// lock.
+#[derive(Clone, Default)]
+pub(super) struct PendingThreadMetadataRegistry {
+    entries: Arc<Mutex<HashMap<ThreadId, PendingThreadMetadataEntry>>>,
+}
+
+pub(super) type LockedPendingThreadMetadata = OwnedMutexGuard<Option<ThreadMetadataPatch>>;
+
+impl PendingThreadMetadataRegistry {
+    pub(super) async fn stage(
+        &self,
+        thread_id: ThreadId,
+        patch: ThreadMetadataPatch,
+    ) -> ThreadStoreResult<()> {
+        if patch.is_empty() {
+            return Err(ThreadStoreError::InvalidRequest {
+                message: "pending thread metadata cannot be empty".to_string(),
+            });
+        }
+        let mut entries = self.entries.lock().await;
+        match entries.entry(thread_id) {
+            Entry::Vacant(entry) => {
+                entry.insert(Arc::new(Mutex::new(Some(patch))));
+                Ok(())
+            }
+            Entry::Occupied(_) => Err(ThreadStoreError::InvalidRequest {
+                message: format!("pending thread metadata already exists: {thread_id}"),
+            }),
+        }
+    }
+
+    pub(super) async fn remove(&self, thread_id: ThreadId) {
+        self.entries.lock().await.remove(&thread_id);
+    }
+
+    pub(super) async fn lock(&self, thread_id: ThreadId) -> Option<LockedPendingThreadMetadata> {
+        let entry = self.entries.lock().await.get(&thread_id).cloned()?;
+        Some(entry.lock_owned().await)
+    }
+}

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

@@ -0,0 +1,298 @@
+use std::sync::Arc;
+use std::time::Duration;
+
+use codex_protocol::ThreadId;
+use codex_protocol::models::BaseInstructions;
+use codex_protocol::protocol::EventMsg;
+use codex_protocol::protocol::SessionSource;
+use codex_protocol::protocol::ThreadHistoryMode;
+use codex_protocol::protocol::ThreadMemoryMode;
+use codex_protocol::protocol::UserMessageEvent;
+use codex_rollout::RolloutItem;
+use codex_rollout::RolloutRecorder;
+use pretty_assertions::assert_eq;
+use tempfile::TempDir;
+
+use super::*;
+use crate::CreateThreadParams;
+use crate::LiveThread;
+use crate::ThreadPersistenceMetadata;
+use crate::UpdateThreadMetadataParams;
+use crate::local::test_support::test_config;
+
+#[tokio::test]
+async fn pending_thread_metadata_is_consumed_by_first_metadata_update() {
+    let (_home, store, runtime) = store_with_runtime().await;
+    let thread_id = ThreadId::new();
+    let pending_patch = staged_model_patch();
+
+    store
+        .stage_pending_thread_metadata(thread_id, pending_patch.clone())
+        .await
+        .expect("stage pending metadata");
+
+    let live_thread = LiveThread::create(store.clone(), create_thread_params(thread_id))
+        .await
+        .expect("create live thread");
+    live_thread
+        .append_items(&[user_message_item("materialize pending metadata")])
+        .await
+        .expect("append observed item");
+
+    let metadata = runtime
+        .get_thread(thread_id)
+        .await
+        .expect("sqlite metadata read")
+        .expect("sqlite metadata");
+    assert_eq!(metadata.model.as_deref(), Some("staged-model"));
+    assert_eq!(metadata.model_provider, "test-provider");
+    assert_eq!(metadata.source, "exec");
+    assert_eq!(
+        metadata.cwd.canonicalize().expect("canonical metadata cwd"),
+        std::env::current_dir()
+            .expect("current directory")
+            .canonicalize()
+            .expect("canonical current directory")
+    );
+    assert_pending_metadata_consumed(&store, thread_id).await;
+}
+
+#[tokio::test]
+async fn pending_rollout_compatible_metadata_does_not_deadlock() {
+    let (_home, store, runtime) = store_with_runtime().await;
+    let thread_id = ThreadId::new();
+
+    store
+        .stage_pending_thread_metadata(
+            thread_id,
+            ThreadMetadataPatch {
+                name: Some(Some("staged-name".to_string())),
+                model_provider: Some("staged-provider".to_string()),
+                ..Default::default()
+            },
+        )
+        .await
+        .expect("stage pending metadata");
+    let live_thread = LiveThread::create(store.clone(), create_thread_params(thread_id))
+        .await
+        .expect("create live thread");
+    tokio::time::timeout(
+        Duration::from_secs(5),
+        live_thread.update_metadata(
+            ThreadMetadataPatch {
+                model: Some("updated-model".to_string()),
+                ..Default::default()
+            },
+            /*include_archived*/ false,
+        ),
+    )
+    .await
+    .expect("pending metadata update should not deadlock")
+    .expect("update metadata before rollout materialization");
+
+    let metadata = runtime
+        .get_thread(thread_id)
+        .await
+        .expect("sqlite metadata read")
+        .expect("sqlite metadata");
+    assert_eq!(metadata.title, "staged-name");
+    assert_eq!(metadata.model_provider, "test-provider");
+    assert_eq!(metadata.model.as_deref(), Some("updated-model"));
+    assert_pending_metadata_consumed(&store, thread_id).await;
+}
+
+#[tokio::test]
+async fn pending_rollout_compatible_metadata_is_written_with_observed_update() {
+    let (_home, store, _runtime) = store_with_runtime().await;
+    let thread_id = ThreadId::new();
+
+    store
+        .stage_pending_thread_metadata(
+            thread_id,
+            ThreadMetadataPatch {
+                memory_mode: Some(ThreadMemoryMode::Disabled),
+                ..Default::default()
+            },
+        )
+        .await
+        .expect("stage pending metadata");
+    let _live_thread = LiveThread::create(store.clone(), create_thread_params(thread_id))
+        .await
+        .expect("create live thread");
+    store
+        .update_thread_metadata(UpdateThreadMetadataParams {
+            thread_id,
+            patch: ThreadMetadataPatch {
+                model: Some("observed-model".to_string()),
+                ..Default::default()
+            },
+            include_archived: false,
+        })
+        .await
+        .expect("apply observed metadata");
+
+    let rollout_path = store
+        .live_rollout_path(thread_id)
+        .await
+        .expect("rollout path");
+    let (items, _, _) = RolloutRecorder::load_rollout_items(rollout_path.as_path())
+        .await
+        .expect("load durable rollout");
+    let memory_mode = items.iter().rev().find_map(|item| match item {
+        RolloutItem::SessionMeta(session_meta) => session_meta.meta.memory_mode.as_deref(),
+        _ => None,
+    });
+
+    assert_eq!(memory_mode, Some("disabled"));
+    assert_pending_metadata_consumed(&store, thread_id).await;
+}
+
+#[tokio::test]
+async fn idle_shutdown_clears_pending_thread_metadata_without_materializing() {
+    let (_home, store, runtime) = store_with_runtime().await;
+    let thread_id = ThreadId::new();
+
+    store
+        .stage_pending_thread_metadata(thread_id, staged_model_patch())
+        .await
+        .expect("stage pending metadata");
+    let live_thread = LiveThread::create(store.clone(), create_thread_params(thread_id))
+        .await
+        .expect("create live thread");
+    live_thread.shutdown().await.expect("shutdown idle thread");
+
+    assert_eq!(
+        runtime
+            .get_thread(thread_id)
+            .await
+            .expect("sqlite metadata read"),
+        None
+    );
+    assert_pending_metadata_consumed(&store, thread_id).await;
+}
+
+#[tokio::test]
+async fn discarding_live_thread_clears_pending_metadata() {
+    let (_home, store, _runtime) = store_with_runtime().await;
+    let thread_id = ThreadId::new();
+
+    store
+        .stage_pending_thread_metadata(thread_id, staged_model_patch())
+        .await
+        .expect("stage pending metadata");
+    let live_thread = LiveThread::create(store.clone(), create_thread_params(thread_id))
+        .await
+        .expect("create live thread");
+    live_thread.discard().await.expect("discard live thread");
+
+    assert_pending_metadata_consumed(&store, thread_id).await;
+}
+
+#[tokio::test]
+async fn pending_thread_metadata_requires_state_db() {
+    let home = TempDir::new().expect("temp dir");
+    let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
+
+    let err = store
+        .stage_pending_thread_metadata(ThreadId::new(), staged_model_patch())
+        .await
+        .expect_err("reject pending metadata without state db");
+    assert!(matches!(
+        err,
+        ThreadStoreError::InvalidRequest { message }
+            if message == "pending thread metadata requires a state db"
+    ));
+}
+
+#[tokio::test]
+async fn pending_thread_metadata_rejects_rollout_path() {
+    let (home, store, _runtime) = store_with_runtime().await;
+
+    let err = store
+        .stage_pending_thread_metadata(
+            ThreadId::new(),
+            ThreadMetadataPatch {
+                rollout_path: Some(home.path().join("wrong-rollout.jsonl")),
+                ..Default::default()
+            },
+        )
+        .await
+        .expect_err("reject host-owned rollout path");
+    assert!(matches!(
+        err,
+        ThreadStoreError::InvalidRequest { message }
+            if message == "pending thread metadata cannot set rollout_path"
+    ));
+}
+
+async fn store_with_runtime() -> (
+    TempDir,
+    Arc<LocalThreadStore>,
+    Arc<codex_state::StateRuntime>,
+) {
+    let home = TempDir::new().expect("temp dir");
+    let config = test_config(home.path());
+    let runtime = codex_state::StateRuntime::init(
+        config.sqlite.clone(),
+        config.default_model_provider_id.clone(),
+    )
+    .await
+    .expect("state db should initialize");
+    let store = Arc::new(LocalThreadStore::new(config, Some(runtime.clone())));
+    (home, store, runtime)
+}
+
+async fn assert_pending_metadata_consumed(store: &LocalThreadStore, thread_id: ThreadId) {
+    store
+        .stage_pending_thread_metadata(thread_id, staged_model_patch())
+        .await
+        .expect("restage metadata after consumption");
+    store
+        .remove_pending_thread_metadata(thread_id)
+        .await
+        .expect("remove restaged metadata");
+}
+
+fn staged_model_patch() -> ThreadMetadataPatch {
+    ThreadMetadataPatch {
+        model: Some("staged-model".to_string()),
+        ..Default::default()
+    }
+}
+
+fn create_thread_params(thread_id: ThreadId) -> CreateThreadParams {
+    CreateThreadParams {
+        session_id: thread_id.into(),
+        thread_id,
+        extra_config: None,
+        forked_from_id: None,
+        parent_thread_id: None,
+        source: SessionSource::Exec,
+        thread_source: None,
+        originator: "test_originator".to_string(),
+        base_instructions: BaseInstructions::default(),
+        dynamic_tools: Vec::new(),
+        selected_capability_roots: Vec::new(),
+        multi_agent_version: None,
+        history_mode: ThreadHistoryMode::Legacy,
+        history_base: None,
+        subagent_history_start_ordinal: None,
+        initial_window_id: uuid::Uuid::now_v7().to_string(),
+        metadata: ThreadPersistenceMetadata {
+            cwd: Some(std::env::current_dir().expect("cwd")),
+            model_provider: "test-provider".to_string(),
+            memory_mode: ThreadMemoryMode::Enabled,
+        },
+    }
+}
+
+fn user_message_item(message: &str) -> RolloutItem {
+    RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent {
+        client_id: None,
+        message: message.to_string(),
+        images: None,
+        local_images: Vec::new(),
+        text_elements: Vec::new(),
+        ..Default::default()
+    }))
+}

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

@@ -18,6 +18,7 @@ use super::LocalThreadStore;
 use super::helpers::git_info_from_parts;
 use super::helpers::permission_profile_to_metadata_value;
 use super::live_writer;
+use super::pending_thread_metadata;
 use super::thread_rollout_resolver;
 use super::thread_rollout_resolver::ResolvedThreadRollout;
 use super::thread_rollout_resolver::RolloutLocation;
@@ -35,7 +36,16 @@ pub(super) async fn update_thread_metadata(
     params: UpdateThreadMetadataParams,
 ) -> ThreadStoreResult<StoredThread> {
     let thread_id = params.thread_id;
-    let patch = params.patch;
+    let mut pending_metadata = store.pending_thread_metadata.lock(thread_id).await;
+    let pending_patch = pending_metadata
+        .as_ref()
+        .and_then(|metadata| metadata.as_ref().cloned());
+    let mut patch = params.patch;
+    if let Some(staged_patch) = pending_patch.as_ref() {
+        let mut merged_patch = staged_patch.clone();
+        merged_patch.merge(patch);
+        patch = merged_patch;
+    }
     if patch.is_empty() {
         return read_thread::read_thread(
             store,
@@ -48,26 +58,35 @@ pub(super) async fn update_thread_metadata(
         .await;
     }
 
-    let requires_rollout_compat = requires_rollout_compatibility_update(&patch);
+    let staged_requires_rollout_compat = pending_patch
+        .as_ref()
+        .is_some_and(|patch| patch.memory_mode.is_some() || patch.git_info.is_some());
+    let requires_rollout_compat =
+        staged_requires_rollout_compat || requires_rollout_compatibility_update(&patch);
     let has_explicit_metadata = patch.name.is_some() || requires_rollout_compat;
     let history_mode = if has_explicit_metadata {
-        Some(
-            read_thread::read_thread(
-                store,
-                ReadThreadParams {
-                    thread_id,
-                    include_archived: params.include_archived,
-                    include_history: false,
-                },
-            )
-            .await?
-            .history_mode,
-        )
+        match live_writer::live_writer_parts(store, thread_id).await {
+            Ok((_recorder, _rollout_id, history_mode)) => Some(history_mode),
+            Err(ThreadStoreError::ThreadNotFound { .. }) => Some(
+                read_thread::read_thread(
+                    store,
+                    ReadThreadParams {
+                        thread_id,
+                        include_archived: params.include_archived,
+                        include_history: false,
+                    },
+                )
+                .await?
+                .history_mode,
+            ),
+            Err(err) => return Err(err),
+        }
     } else {
         None
     };
     let paginated = matches!(history_mode, Some(ThreadHistoryMode::Paginated));
-    let require_sqlite_write = sqlite_write_failure_should_block(&patch) || paginated;
+    let require_sqlite_write =
+        pending_patch.is_some() || sqlite_write_failure_should_block(&patch) || paginated;
     let mut updated = apply_metadata_update(
         store,
         thread_id,
@@ -112,10 +131,16 @@ pub(super) async fn update_thread_metadata(
         {
             warn!("failed to index paginated thread name for {thread_id}: {err}");
         }
+        if pending_patch.is_some() {
+            remove_pending_thread_metadata(store, thread_id, &mut pending_metadata).await;
+        }
         return Ok(updated);
     }
     let needs_rollout_compat = requires_rollout_compat || patch.name.is_some();
     if !needs_rollout_compat {
+        if pending_patch.is_some() {
+            remove_pending_thread_metadata(store, thread_id, &mut pending_metadata).await;
+        }
         return Ok(updated);
     }
 
@@ -239,9 +264,24 @@ pub(super) async fn update_thread_metadata(
     if let Some(((sha, branch, origin_url), _memory_mode)) = resolved_git_info {
         thread.git_info = git_info_from_parts(sha, branch, origin_url);
     }
+    if pending_patch.is_some() {
+        remove_pending_thread_metadata(store, thread_id, &mut pending_metadata).await;
+    }
     Ok(thread)
 }
 
+async fn remove_pending_thread_metadata(
+    store: &LocalThreadStore,
+    thread_id: ThreadId,
+    pending_metadata: &mut Option<pending_thread_metadata::LockedPendingThreadMetadata>,
+) {
+    if let Some(mut metadata) = pending_metadata.take() {
+        *metadata = None;
+        drop(metadata);
+        store.pending_thread_metadata.remove(thread_id).await;
+    }
+}
+
 async fn refresh_resolved_rollout_path(resolved: &mut ResolvedThreadRollout) {
     if let Some(path) = codex_rollout::existing_rollout_path(resolved.path.as_path()).await {
         resolved.path = path;