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
This commit is contained in:
Anton Panasenko
2026-08-16 00:57:28 +00:00
committed by copyberry
parent 5ba12929f8
commit fcdae21073
17 changed files with 650 additions and 21 deletions
+1
View File
@@ -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",
+1
View File
@@ -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(),
+3
View File
@@ -422,6 +422,7 @@ pub(crate) struct SessionSpawnArgs {
pub(crate) environment_selections: Vec<TurnEnvironmentSelection>,
pub(crate) thread_extension_init: ExtensionDataInit,
pub(crate) client_mcp_extensions: ClientMcpExtensions,
pub(crate) reserved_thread_id: Option<ThreadId>,
pub(crate) analytics_events_client: Option<AnalyticsEventsClient>,
pub(crate) thread_store: Arc<dyn ThreadStore>,
pub(crate) attestation_provider: Option<Arc<dyn AttestationProvider>>,
@@ -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,
+13 -3
View File
@@ -605,6 +605,7 @@ impl Session {
mut thread_extension_init: ExtensionDataInit,
client_mcp_extensions: ClientMcpExtensions,
agent_control: AgentControl,
reserved_thread_id: Option<ThreadId>,
environment_manager: Arc<EnvironmentManager>,
inherited_environments: Option<TurnEnvironmentSnapshot>,
analytics_events_client: Option<AnalyticsEventsClient>,
@@ -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) => {
+3
View File
@@ -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,
@@ -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,
+15
View File
@@ -231,6 +231,8 @@ pub struct StartThreadOptions {
pub environments: Option<Vec<TurnEnvironmentSelection>>,
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<ThreadId>,
}
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(),
+56
View File
@@ -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() {
+1
View File
@@ -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 }
+4 -1
View File
@@ -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<ThreadManager>,
pub thread_store: Arc<dyn ThreadStore>,
_test_env: TestEnv,
}
+60
View File
@@ -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;
+20 -2
View File
@@ -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)> {
+35
View File
@@ -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 })
}
@@ -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)
}
}
@@ -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()
}))
}
@@ -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;
+26
View File
@@ -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<'_, ()>;