mirror of
https://github.com/openai/codex.git
synced 2026-11-05 20:10:33 +00:00
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:
committed by
copyberry
parent
5ba12929f8
commit
fcdae21073
Generated
+1
@@ -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",
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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) => {
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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() {
|
||||
|
||||
@@ -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 }
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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)> {
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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<'_, ()>;
|
||||
|
||||
|
||||
Reference in New Issue
Block a user