diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 8f4ac8ffea..deffcfbb9b 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -5030,6 +5030,7 @@ dependencies = [ "codex-model-provider-info", "codex-models-manager", "codex-protocol", + "codex-thread-store", "codex-utils-absolute-path", "codex-utils-cargo-bin", "codex-utils-path-uri", diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 1cec86480d..0b4ff1ac46 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -120,6 +120,7 @@ pub(crate) async fn run_codex_thread_interactive( environment_selections: parent_environments.to_selections(), thread_extension_init: codex_extension_api::ExtensionDataInit::default(), client_mcp_extensions: parent_session.services.client_mcp_extensions.clone(), + reserved_thread_id: None, analytics_events_client: Some(parent_session.services.analytics_events_client.clone()), thread_store: Arc::clone(&parent_session.services.thread_store), attestation_provider: parent_session.services.attestation_provider.clone(), diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 33b77fadf4..03ce90bca9 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -422,6 +422,7 @@ pub(crate) struct SessionSpawnArgs { pub(crate) environment_selections: Vec, pub(crate) thread_extension_init: ExtensionDataInit, pub(crate) client_mcp_extensions: ClientMcpExtensions, + pub(crate) reserved_thread_id: Option, pub(crate) analytics_events_client: Option, pub(crate) thread_store: Arc, pub(crate) attestation_provider: Option>, @@ -514,6 +515,7 @@ impl Session { environment_selections, thread_extension_init, client_mcp_extensions, + reserved_thread_id, analytics_events_client, thread_store, attestation_provider, @@ -742,6 +744,7 @@ impl Session { thread_extension_init, client_mcp_extensions, agent_control, + reserved_thread_id, environment_manager, inherited_environments, analytics_events_client, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index dfd6ca1bb9..cb792abf0d 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -605,6 +605,7 @@ impl Session { mut thread_extension_init: ExtensionDataInit, client_mcp_extensions: ClientMcpExtensions, agent_control: AgentControl, + reserved_thread_id: Option, environment_manager: Arc, inherited_environments: Option, analytics_events_client: Option, @@ -663,11 +664,20 @@ impl Session { let multi_agent_version = multi_agent_version.map(OnceLock::from).unwrap_or_default(); let initial_multi_agent_version = multi_agent_version.get().copied(); - let thread_id = match &initial_history { - InitialHistory::New | InitialHistory::Cleared | InitialHistory::Forked(_) => { + let thread_id = match (&initial_history, reserved_thread_id) { + ( + InitialHistory::New | InitialHistory::Cleared | InitialHistory::Forked(_), + Some(thread_id), + ) => thread_id, + (InitialHistory::New | InitialHistory::Cleared | InitialHistory::Forked(_), None) => { agent_control.generate_thread_id() } - InitialHistory::Resumed(resumed_history) => resumed_history.conversation_id, + (InitialHistory::Resumed(resumed_history), None) => resumed_history.conversation_id, + (InitialHistory::Resumed(_), Some(_)) => { + return Err(anyhow::anyhow!( + "reserved thread ID cannot be used when resuming a thread" + )); + } }; let resumed_session_id = match &initial_history { InitialHistory::Resumed(resumed) => { diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 9a90659060..3005aae736 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -5695,6 +5695,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() { codex_extension_api::ExtensionDataInit::default(), ClientMcpExtensions::default(), AgentControl::default(), + /*reserved_thread_id*/ None, environment_manager, /*inherited_environments*/ None, /*analytics_events_client*/ None, @@ -6122,6 +6123,7 @@ async fn make_session_with_config_and_rx( codex_extension_api::ExtensionDataInit::default(), ClientMcpExtensions::default(), AgentControl::default(), + /*reserved_thread_id*/ None, environment_manager, /*inherited_environments*/ None, /*analytics_events_client*/ None, @@ -6240,6 +6242,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx( codex_extension_api::ExtensionDataInit::default(), ClientMcpExtensions::default(), agent_control, + /*reserved_thread_id*/ None, environment_manager, /*inherited_environments*/ None, /*analytics_events_client*/ None, diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index 89cb35129d..9a2dd7c2ed 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -832,6 +832,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { environment_selections: Vec::new(), thread_extension_init: codex_extension_api::ExtensionDataInit::default(), client_mcp_extensions: ClientMcpExtensions::default(), + reserved_thread_id: None, analytics_events_client: None, thread_store, attestation_provider: None, diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 3c373f9fe0..52e87218b9 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -231,6 +231,8 @@ pub struct StartThreadOptions { pub environments: Option>, pub thread_extension_init: ExtensionDataInit, pub client_mcp_extensions: ClientMcpExtensions, + /// Thread ID reserved before startup so the caller can associate host-owned state with it. + pub reserved_thread_id: Option, } impl StartThreadOptions { @@ -248,6 +250,7 @@ impl StartThreadOptions { environments: None, thread_extension_init: ExtensionDataInit::default(), client_mcp_extensions: ClientMcpExtensions::default(), + reserved_thread_id: None, } } } @@ -891,6 +894,11 @@ impl ThreadManager { Box::pin(self.start_thread_inner(options, /*forked_from_thread_id*/ None)).await } + /// Allocates a thread ID before startup so a caller can associate host-owned state with it. + pub fn reserve_thread_id(&self) -> ThreadId { + self.state.thread_id_generator.as_ref()() + } + async fn start_thread_inner( &self, mut options: StartThreadOptions, @@ -1755,6 +1763,7 @@ impl ThreadManagerState { environments, thread_extension_init, client_mcp_extensions, + reserved_thread_id, } = options; let session_source = session_source.unwrap_or_else(|| self.session_source.clone()); let environments = environments.unwrap_or_else(|| { @@ -1765,6 +1774,11 @@ impl ThreadManagerState { ) }); let is_resumed_thread = matches!(&initial_history, InitialHistory::Resumed(_)); + if reserved_thread_id.is_some() && matches!(&initial_history, InitialHistory::Resumed(_)) { + return Err(CodexErr::InvalidRequest( + "reserved thread ID cannot be used when resuming a thread".to_string(), + )); + } if let InitialHistory::Resumed(resumed) = &initial_history { let mut threads = self.threads.write().await; if let Some(thread) = threads.get(&resumed.conversation_id).cloned() { @@ -1851,6 +1865,7 @@ impl ThreadManagerState { environment_selections: environments, thread_extension_init, client_mcp_extensions, + reserved_thread_id, analytics_events_client: self.analytics_events_client.clone(), thread_store: Arc::clone(&self.thread_store), attestation_provider: self.attestation_provider.clone(), diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 525ab9ef08..9ad69d8e92 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -58,6 +58,62 @@ fn thread_id_generator_defaults_to_standard_ids() { ); } +#[tokio::test] +async fn reserved_thread_id_is_used_without_changing_normal_id_generation() { + let temp_dir = tempdir().expect("tempdir"); + let mut config = test_config().await; + config.codex_home = temp_dir.path().join("codex-home").abs(); + config.cwd = config.codex_home.abs(); + std::fs::create_dir_all(&config.codex_home).expect("create codex home"); + + let generated_ids = [ + ThreadId::from_u128(/*value*/ 0x018f_0000_0000_7000_8000_0000_0000_0001), + ThreadId::from_u128(/*value*/ 0x018f_0000_0000_7000_8000_0000_0000_0002), + ThreadId::from_u128(/*value*/ 0x018f_0000_0000_7000_8000_0000_0000_0003), + ]; + let next_id = std::sync::atomic::AtomicUsize::new(0); + let manager = ThreadManager::with_models_provider_and_home_for_tests( + CodexAuth::from_api_key("dummy"), + config.model_provider.clone(), + config.codex_home.to_path_buf(), + Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + ) + .with_thread_id_generator(move || generated_ids[next_id.fetch_add(1, Ordering::Relaxed)]); + + let reserved_id = manager.reserve_thread_id(); + let mut reserved_options = StartThreadOptions::new(config.clone()); + reserved_options.reserved_thread_id = Some(reserved_id); + let reserved = manager + .start_thread(reserved_options) + .await + .expect("start reserved thread"); + let mut resumed_options = StartThreadOptions::new(config.clone()); + resumed_options.initial_history = InitialHistory::Resumed(ResumedHistory { + conversation_id: reserved.thread_id, + history: Arc::new(Vec::new()), + rollout_path: None, + }); + let resumed_id = manager.reserve_thread_id(); + resumed_options.reserved_thread_id = Some(resumed_id); + let resume_error = manager + .start_thread(resumed_options) + .await + .err() + .expect("reject reserved ID for resume"); + let generated = manager + .start_thread(StartThreadOptions::new(config.clone())) + .await + .expect("start generated thread"); + + assert_eq!(reserved.thread_id, generated_ids[0]); + assert!(matches!( + resume_error.details(), + codex_protocol::error::CodexErrorDetails::InvalidRequest(message) + if message == "reserved thread ID cannot be used when resuming a thread" + )); + assert_eq!(generated.thread_id, generated_ids[2]); +} + /// One custom ID factory supplies identifiers for roots, actual child agents, and forks. #[tokio::test] async fn thread_id_generator_applies_to_roots_children_and_forks() { diff --git a/codex-rs/core/tests/common/Cargo.toml b/codex-rs/core/tests/common/Cargo.toml index 05adeb044e..a4851eb1d9 100644 --- a/codex-rs/core/tests/common/Cargo.toml +++ b/codex-rs/core/tests/common/Cargo.toml @@ -27,6 +27,7 @@ codex-login = { workspace = true } codex-model-provider-info = { workspace = true } codex-models-manager = { workspace = true } codex-protocol = { workspace = true } +codex-thread-store = { workspace = true } codex-utils-absolute-path = { workspace = true } codex-utils-cargo-bin = { workspace = true } codex-utils-path-uri = { workspace = true } diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index 01ef0af3f5..63ddaab8bc 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -58,6 +58,7 @@ use codex_protocol::protocol::ThreadSettingsOverrides; use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::protocol::TurnEnvironmentSelections; use codex_protocol::user_input::UserInput; +use codex_thread_store::ThreadStore; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_path_uri::PathUri; use futures::future::BoxFuture; @@ -671,7 +672,7 @@ impl TestCodexBuilder { Arc::clone(&self.extensions), user_instructions_provider, /*analytics_events_client*/ None, - thread_store, + Arc::clone(&thread_store), codex_core::local_agent_graph_store_from_state_db(state_db.as_ref()), installation_id, /*attestation_provider*/ None, @@ -755,6 +756,7 @@ impl TestCodexBuilder { codex: new_conversation.thread, session_configured: new_conversation.session_configured, thread_manager, + thread_store, _test_env: test_env, }) } @@ -845,6 +847,7 @@ pub struct TestCodex { pub session_configured: SessionConfiguredEvent, pub config: Config, pub thread_manager: Arc, + pub thread_store: Arc, _test_env: TestEnv, } diff --git a/codex-rs/core/tests/suite/sqlite_state.rs b/codex-rs/core/tests/suite/sqlite_state.rs index 9fabd70bcd..d4ad40410a 100644 --- a/codex-rs/core/tests/suite/sqlite_state.rs +++ b/codex-rs/core/tests/suite/sqlite_state.rs @@ -28,6 +28,7 @@ use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::ThreadSettingsOverrides; use codex_protocol::protocol::UserMessageEvent; use codex_protocol::user_input::UserInput; +use codex_thread_store::ThreadMetadataPatch; use codex_web_search_extension::install as install_web_search_extension; use core_test_support::responses; use core_test_support::responses::ev_completed; @@ -114,6 +115,65 @@ async fn new_thread_is_recorded_in_state_db() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn staged_metadata_is_persisted_on_first_turn() -> Result<()> { + let server = start_mock_server().await; + let _mock = mount_sse_once( + &server, + responses::sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; + let mut builder = test_codex().with_config(|config| { + config + .features + .enable(Feature::Sqlite) + .expect("test config should allow feature update"); + }); + let test = builder.build_with_auto_env(&server).await?; + let thread_id = test.thread_manager.reserve_thread_id(); + test.thread_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?; + let mut options = StartThreadOptions::new(test.config.clone()); + options.reserved_thread_id = Some(thread_id); + let started = test.thread_manager.start_thread(options).await?; + + started + .thread + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "persist staged metadata".to_string(), + text_elements: Vec::new(), + }])) + .await?; + wait_for_event(&started.thread, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + let db = test.codex.state_db().expect("state db enabled"); + let metadata = db + .get_thread(started.thread_id) + .await? + .expect("thread should exist in state db"); + let thread = started + .thread + .read_thread( + /*include_archived*/ false, /*include_history*/ false, + ) + .await?; + assert_eq!(thread.name.as_deref(), Some("staged-name")); + assert_eq!(metadata.model_provider, test.config.model_provider_id); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn resume_restores_dynamic_tools_from_rollout_with_sqlite_enabled() -> Result<()> { let server = start_mock_server().await; diff --git a/codex-rs/thread-store/src/local/live_writer.rs b/codex-rs/thread-store/src/local/live_writer.rs index 8472027401..3c48136272 100644 --- a/codex-rs/thread-store/src/local/live_writer.rs +++ b/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)> { diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index e5c12dd11f..78fc188438 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/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>>, + pending_thread_metadata: pending_thread_metadata::PendingThreadMetadataRegistry, live_writer_locks: Arc, writer_lock_coordinator: Arc, state_db: Option, @@ -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 }) } diff --git a/codex-rs/thread-store/src/local/pending_thread_metadata.rs b/codex-rs/thread-store/src/local/pending_thread_metadata.rs new file mode 100644 index 0000000000..5eb28660d4 --- /dev/null +++ b/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>>; + +/// 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>>, +} + +pub(super) type LockedPendingThreadMetadata = OwnedMutexGuard>; + +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 { + let entry = self.entries.lock().await.get(&thread_id).cloned()?; + Some(entry.lock_owned().await) + } +} diff --git a/codex-rs/thread-store/src/local/pending_thread_metadata_tests.rs b/codex-rs/thread-store/src/local/pending_thread_metadata_tests.rs new file mode 100644 index 0000000000..4cc4d040b2 --- /dev/null +++ b/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, + Arc, +) { + 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() + })) +} diff --git a/codex-rs/thread-store/src/local/update_thread_metadata.rs b/codex-rs/thread-store/src/local/update_thread_metadata.rs index 12195734fc..61e6b9c1b6 100644 --- a/codex-rs/thread-store/src/local/update_thread_metadata.rs +++ b/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 { 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, +) { + 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; diff --git a/codex-rs/thread-store/src/store.rs b/codex-rs/thread-store/src/store.rs index 90b53b38ed..af2c9a6f08 100644 --- a/codex-rs/thread-store/src/store.rs +++ b/codex-rs/thread-store/src/store.rs @@ -33,6 +33,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; @@ -69,6 +70,31 @@ pub trait ThreadStore: Any + Send + Sync { /// Creates a new live thread. fn create_thread(&self, params: CreateThreadParams) -> ThreadStoreFuture<'_, ()>; + /// Stages host-owned metadata for a thread ID reserved before Core starts the thread. + /// + /// The entry remains in memory until the first successful metadata update for that thread. + /// Callers must remove it if startup fails before the store opens a live thread. + fn stage_pending_thread_metadata( + &self, + _thread_id: ThreadId, + _patch: ThreadMetadataPatch, + ) -> ThreadStoreFuture<'_, ()> { + Box::pin(async { + Err(ThreadStoreError::Unsupported { + operation: "stage_pending_thread_metadata", + }) + }) + } + + /// Removes host-owned metadata staged for a reserved thread ID. + fn remove_pending_thread_metadata(&self, _thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> { + Box::pin(async { + Err(ThreadStoreError::Unsupported { + operation: "remove_pending_thread_metadata", + }) + }) + } + /// Reopens an existing thread for live appends. fn resume_thread(&self, params: ResumeThreadParams) -> ThreadStoreFuture<'_, ()>;