diff --git a/codex-rs/app-server/src/mcp_refresh.rs b/codex-rs/app-server/src/mcp_refresh.rs index 7fe1e32af4..863d86af74 100644 --- a/codex-rs/app-server/src/mcp_refresh.rs +++ b/codex-rs/app-server/src/mcp_refresh.rs @@ -178,9 +178,9 @@ mod tests { use codex_config::ThreadConfigSource; use codex_config::types::AuthKeyringBackendKind; use codex_config::types::McpServerConfig; + use codex_core::ThreadStoreBuilder; use codex_core::config::ConfigOverrides; use codex_core::init_state_db; - use codex_core::thread_store_from_config; use codex_exec_server::EnvironmentManager; use codex_extension_api::NoopExtensionEventSink; use codex_home::CodexHomeUserInstructionsProvider; @@ -522,7 +522,8 @@ enabled = false let state_db = init_state_db(&good_config) .await .expect("refresh tests require state db"); - let thread_store = thread_store_from_config(&good_config, Some(state_db.clone())); + let thread_store = + ThreadStoreBuilder::from_config(&good_config, Some(state_db.clone())).build(); let environment_manager = Arc::new(EnvironmentManager::default_for_tests()); let executor_skill_provider: Arc = Arc::new( codex_skills_extension::ExecutorSkillProvider::new_with_restriction_product( diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index d296ac1f78..977c6c2735 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -80,6 +80,7 @@ use codex_app_server_protocol::experimental_required_message; use codex_arg0::Arg0DispatchPaths; use codex_code_mode::CodeModeSessionProvider; use codex_core::ThreadManager; +use codex_core::ThreadStoreBuilder; use codex_core::config::Config; use codex_core::config::ThreadStoreConfig; use codex_exec_server::EnvironmentManager; @@ -317,7 +318,8 @@ impl MessageProcessor { // The thread store is intentionally process-scoped. Config reloads can // affect per-thread behavior, but they must not move newly started, // resumed, or forked threads to a different persistence backend/root. - let thread_store = codex_core::thread_store_from_config(config.as_ref(), state_db.clone()); + let thread_store = + ThreadStoreBuilder::from_config(config.as_ref(), state_db.clone()).build(); // Queue persistence requires SQLite, so in-memory thread stores and // app servers without a state database do not have a queue backend. let queue_store: Option> = match &config.experimental_thread_store { diff --git a/codex-rs/core-api/src/lib.rs b/codex-rs/core-api/src/lib.rs index 5934637065..4255f0a3e5 100644 --- a/codex-rs/core-api/src/lib.rs +++ b/codex-rs/core-api/src/lib.rs @@ -58,6 +58,7 @@ pub use codex_core::SteerSubmission; pub use codex_core::SuspendTurnOutcome; pub use codex_core::ThreadManager; pub use codex_core::ThreadShutdownReport; +pub use codex_core::ThreadStoreBuilder; pub use codex_core::TurnInput; pub use codex_core::TurnInputRequest; pub use codex_core::TurnInputSubmission; @@ -82,7 +83,6 @@ pub use codex_core::init_state_db; pub use codex_core::local_agent_graph_store_from_state_db; pub use codex_core::passthrough_image_store; pub use codex_core::resolve_installation_id; -pub use codex_core::thread_store_from_config; pub use codex_exec_server::EnvironmentManager; pub use codex_exec_server::EnvironmentRegistryConnectRequest; pub use codex_exec_server::EnvironmentRegistryConnectResponse; diff --git a/codex-rs/core/src/agent/control/spawn_publication_tests.rs b/codex-rs/core/src/agent/control/spawn_publication_tests.rs index 3df67dbbdc..ce679c37df 100644 --- a/codex-rs/core/src/agent/control/spawn_publication_tests.rs +++ b/codex-rs/core/src/agent/control/spawn_publication_tests.rs @@ -122,7 +122,7 @@ async fn publication_harness( Arc::new(crate::test_support::EmptyUserInstructionsProvider), /*analytics_events_client*/ None, harness.manager.image_store(), - crate::thread_manager::thread_store_from_config(&harness.config, harness.state_db.clone()), + crate::ThreadStoreBuilder::from_config(&harness.config, harness.state_db.clone()).build(), Some(graph.clone()), "spawn-publication-cleanup-test".to_string(), /*attestation_provider*/ None, diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 87247661d2..3a5fe5d8f6 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -3793,10 +3793,8 @@ async fn spawn_agent_full_fork_legacy_compaction_rebuilds_child_instructions_onc Arc::new(crate::test_support::EmptyUserInstructionsProvider), /*analytics_events_client*/ None, crate::thread_manager::passthrough_image_store(), - crate::thread_manager::thread_store_from_config( - &parent_config, - harness.state_db.clone(), - ), + crate::ThreadStoreBuilder::from_config(&parent_config, harness.state_db.clone()) + .build(), crate::thread_manager::local_agent_graph_store_from_state_db(harness.state_db.as_ref()), uuid::Uuid::new_v4().to_string(), /*attestation_provider*/ None, diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 51573406cb..76e0853fb3 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -155,6 +155,7 @@ mod event_mapping; pub use codex_prompts as review_prompts; mod thread_manager; mod thread_spawn; +mod thread_store_builder; pub use agent::host::AgentMcpMetadataRecipient; pub use agent::host::HostThreadSpawnReservation; pub use agent::host::HostToolInvocationLookup; @@ -181,7 +182,7 @@ pub use thread_manager::ThreadShutdownReport; pub use thread_manager::build_models_manager; pub use thread_manager::local_agent_graph_store_from_state_db; pub use thread_manager::passthrough_image_store; -pub use thread_manager::thread_store_from_config; +pub use thread_store_builder::ThreadStoreBuilder; pub use tools::handlers::WaitForEnvironmentToolConfig; pub use web_search::web_search_action_detail; pub use windows_sandbox_read_grants::grant_read_root_non_elevated; diff --git a/codex-rs/core/src/prompt_debug.rs b/codex-rs/core/src/prompt_debug.rs index 8d1af241c6..05c99d2b11 100644 --- a/codex-rs/core/src/prompt_debug.rs +++ b/codex-rs/core/src/prompt_debug.rs @@ -12,6 +12,7 @@ use codex_protocol::protocol::SessionSource; use codex_protocol::user_input::UserInput; use tokio_util::sync::CancellationToken; +use crate::ThreadStoreBuilder; use crate::config::Config; use crate::resolve_installation_id; use crate::session::session::Session; @@ -19,7 +20,6 @@ use crate::session::turn::build_prompt; use crate::state_db_bridge::StateDbHandle; use crate::thread_manager::StartThreadOptions; use crate::thread_manager::ThreadManager; -use crate::thread_manager::thread_store_from_config; /// Build the model-visible `input` list for a single debug turn. #[doc(hidden)] @@ -42,7 +42,7 @@ pub async fn build_prompt_input( config.codex_linux_sandbox_exe.clone(), )?; - let thread_store = thread_store_from_config(&config, state_db.clone()); + let thread_store = ThreadStoreBuilder::from_config(&config, state_db.clone()).build(); let installation_id = resolve_installation_id(&config.codex_home).await?; let thread_manager = ThreadManager::new( &config, diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 1e2bcbdc9d..a30a2b8ab7 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -19,7 +19,6 @@ use crate::agents_md_manager::SessionInstructions; use crate::attestation::AttestationProvider; use crate::codex_thread::CodexThread; use crate::config::Config; -use crate::config::ThreadStoreConfig; use crate::current_time::TimeProvider; use crate::environment_selection::TurnEnvironmentSnapshot; use crate::environment_selection::default_thread_environment_requests; @@ -41,7 +40,6 @@ use crate::tasks::interrupted_turn_history_marker; use crate::thread_startup_metadata::ThreadStartupMetadata; use codex_agent_graph_store::AgentGraphStore; use codex_agent_graph_store::LocalAgentGraphStore; -use codex_agent_message_board_extension::local_board_options; use codex_analytics::AnalyticsEventsClient; use codex_app_server_protocol::ThreadHistoryBuilder; use codex_app_server_protocol::TurnStatus; @@ -73,7 +71,6 @@ use codex_login::AuthManager; use codex_login::CodexAuth; use codex_login::default_client::CODEX_INTERNAL_ORIGINATOR_OVERRIDE_ENV_VAR; use codex_login::default_client::originator; -use codex_message_board::LocalAgentMessageBoard; use codex_model_provider::create_model_provider; use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::OPENAI_PROVIDER_ID; @@ -103,7 +100,6 @@ use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::protocol::W3cTraceContext; use codex_rollout::state_db::StateDbHandle; use codex_skills_extension::HostSkillsService; -use codex_thread_store::InMemoryThreadStore; use codex_thread_store::LoadThreadHistoryParams; use codex_thread_store::LocalThreadStore; use codex_thread_store::LocalThreadStoreConfig; @@ -563,70 +559,6 @@ pub fn build_models_manager( manager } -pub fn thread_store_from_config( - config: &Config, - state_db: Option, -) -> Arc { - match &config.experimental_thread_store { - ThreadStoreConfig::Local => { - let compression_enabled = config - .features - .enabled(Feature::LocalThreadStoreCompression); - let background_migration_enabled = config - .features - .enabled(Feature::BackgroundPaginatedRolloutMigration); - let has_state_db = state_db.is_some(); - let sqlite = config.sqlite_config().clone(); - let store = Arc::new( - LocalThreadStore::new(LocalThreadStoreConfig::from_config(config), state_db) - .with_thread_data_cleanup(move |thread_ids| { - let sqlite = sqlite.clone(); - Box::pin(async move { - let boards = thread_ids.into_iter().map(Into::into).collect::>(); - LocalAgentMessageBoard::delete_boards( - &sqlite, - &local_board_options(), - &boards, - ) - .await - .map_err(|err| { - ThreadStoreError::Internal { - message: format!( - "failed to delete agent message boards: {err}" - ), - } - }) - }) - }), - ); - if has_state_db && background_migration_enabled { - let startup_store = Arc::clone(&store); - let codex_home = config.codex_home.to_path_buf(); - tokio::spawn(async move { - if let Err(err) = startup_store.migrate_rollouts_on_startup().await { - warn!("failed to migrate legacy rollouts on startup: {err}"); - } - if compression_enabled { - codex_rollout::spawn_rollout_compression_worker( - codex_home, - codex_rollout::RolloutCompressionTrigger::Startup, - ); - } - }); - } else if compression_enabled { - codex_rollout::spawn_rollout_compression_worker( - config.codex_home.to_path_buf(), - codex_rollout::RolloutCompressionTrigger::Startup, - ); - } - store - } - ThreadStoreConfig::InMemory { id } => { - Arc::new(InMemoryThreadStore::for_id(id).with_state_db(state_db)) - } - } -} - /// Constructs the default image store that preserves inline images. pub fn passthrough_image_store() -> Arc { Arc::new(InlineAttachmentStore) @@ -967,7 +899,7 @@ impl ThreadManager { /// Starts the local rollout migration path after a runtime feature enablement. /// - /// Startup config handles the initial launch in [`thread_store_from_config`]. This covers + /// Startup config handles the initial launch in [`crate::ThreadStoreBuilder::build`]. This covers /// clients that decide to enable background migration after constructing the app-server. pub fn start_background_rollout_migration(&self) { let Some(store) = self diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 54e06778a3..c4601f0e59 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -1,8 +1,10 @@ #![cfg(test)] use super::*; +use crate::ThreadStoreBuilder; use crate::agent::types::SpawnAgentOptions; use crate::config::RolloutBudgetConfig; +use crate::config::ThreadStoreConfig; use crate::config::test_config; use crate::init_state_db; use crate::installation_id::INSTALLATION_ID_FILENAME; @@ -52,6 +54,7 @@ use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TurnStartedEvent; use codex_protocol::protocol::UserMessageEvent; use codex_protocol::user_input::UserInput; +use codex_thread_store::InMemoryThreadStore; use codex_utils_path_uri::PathUri; use core_test_support::PathBufExt; use core_test_support::PathExt; @@ -1256,7 +1259,7 @@ async fn mcp_invalidation_refreshes_threads_that_are_still_starting() { Arc::new(crate::test_support::EmptyUserInstructionsProvider), /*analytics_events_client*/ None, passthrough_image_store(), - thread_store_from_config(&config, /*state_db*/ None), + ThreadStoreBuilder::from_config(&config, /*state_db*/ None).build(), /*agent_graph_store*/ None, TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, @@ -1924,7 +1927,7 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() Arc::new(crate::test_support::EmptyUserInstructionsProvider), /*analytics_events_client*/ None, passthrough_image_store(), - thread_store_from_config(&config, /*state_db*/ None), + ThreadStoreBuilder::from_config(&config, /*state_db*/ None).build(), /*agent_graph_store*/ None, TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, @@ -2371,7 +2374,7 @@ async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() { Arc::new(crate::test_support::EmptyUserInstructionsProvider), /*analytics_events_client*/ None, passthrough_image_store(), - thread_store_from_config(&config, /*state_db*/ None), + ThreadStoreBuilder::from_config(&config, /*state_db*/ None).build(), /*agent_graph_store*/ None, TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, @@ -2519,7 +2522,7 @@ async fn explicit_installation_id_skips_codex_home_file() { AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); let installation_id = uuid::Uuid::new_v4().to_string(); let state_db = init_state_db(&config).await; - let thread_store = thread_store_from_config(&config, state_db.clone()); + let thread_store = ThreadStoreBuilder::from_config(&config, state_db.clone()).build(); let manager = ThreadManager::new( &config, auth_manager.clone(), @@ -2575,7 +2578,7 @@ async fn resume_active_thread_from_rollout_returns_running_thread() { Arc::new(crate::test_support::EmptyUserInstructionsProvider), /*analytics_events_client*/ None, passthrough_image_store(), - thread_store_from_config(&config, /*state_db*/ None), + ThreadStoreBuilder::from_config(&config, /*state_db*/ None).build(), /*agent_graph_store*/ None, TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, @@ -2641,7 +2644,7 @@ async fn resume_stopped_thread_from_rollout_spawns_new_thread() { Arc::new(crate::test_support::EmptyUserInstructionsProvider), /*analytics_events_client*/ None, passthrough_image_store(), - thread_store_from_config(&config, /*state_db*/ None), + ThreadStoreBuilder::from_config(&config, /*state_db*/ None).build(), /*agent_graph_store*/ None, TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, @@ -2702,7 +2705,7 @@ async fn resume_stopped_thread_from_rollout_preserves_thread_source() { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); let state_db = init_state_db(&config).await; - let thread_store = thread_store_from_config(&config, state_db.clone()); + let thread_store = ThreadStoreBuilder::from_config(&config, state_db.clone()).build(); let manager = ThreadManager::new( &config, auth_manager.clone(), @@ -2801,7 +2804,7 @@ async fn subtree_listing_uses_injected_graph_store_without_state_db() { Arc::new(crate::test_support::EmptyUserInstructionsProvider), /*analytics_events_client*/ None, passthrough_image_store(), - thread_store_from_config(&config, /*state_db*/ None), + ThreadStoreBuilder::from_config(&config, /*state_db*/ None).build(), Some(agent_graph_store), TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, @@ -2833,7 +2836,7 @@ async fn rollout_path_resume_and_fork_read_history_through_thread_store() { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); let state_db = init_state_db(&config).await; - let thread_store = thread_store_from_config(&config, state_db.clone()); + let thread_store = ThreadStoreBuilder::from_config(&config, state_db.clone()).build(); let in_memory_store = thread_store .as_any() .downcast_ref::() @@ -2943,7 +2946,7 @@ async fn metadata_update_without_result_reads_only_when_the_caller_needs_the_thr let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); - let thread_store = thread_store_from_config(&config, /*state_db*/ None); + let thread_store = ThreadStoreBuilder::from_config(&config, /*state_db*/ None).build(); let in_memory_store = thread_store .as_any() .downcast_ref::() @@ -3086,7 +3089,7 @@ async fn new_uses_active_provider_for_model_refresh() { Arc::new(crate::test_support::EmptyUserInstructionsProvider), /*analytics_events_client*/ None, passthrough_image_store(), - thread_store_from_config(&config, /*state_db*/ None), + ThreadStoreBuilder::from_config(&config, /*state_db*/ None).build(), /*agent_graph_store*/ None, TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, @@ -3185,7 +3188,7 @@ async fn injected_models_manager_controls_refresh_policy() { Arc::new(crate::test_support::EmptyUserInstructionsProvider), /*analytics_events_client*/ None, passthrough_image_store(), - thread_store_from_config(&config, /*state_db*/ None), + ThreadStoreBuilder::from_config(&config, /*state_db*/ None).build(), /*agent_graph_store*/ None, TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, @@ -3466,7 +3469,7 @@ async fn interrupted_fork_snapshot_does_not_synthesize_turn_id_for_legacy_histor Arc::new(crate::test_support::EmptyUserInstructionsProvider), /*analytics_events_client*/ None, passthrough_image_store(), - thread_store_from_config(&config, state_db.clone()), + ThreadStoreBuilder::from_config(&config, state_db.clone()).build(), local_agent_graph_store_from_state_db(state_db.as_ref()), TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, @@ -3578,7 +3581,7 @@ async fn interrupted_fork_snapshot_preserves_explicit_turn_id() { Arc::new(crate::test_support::EmptyUserInstructionsProvider), /*analytics_events_client*/ None, passthrough_image_store(), - thread_store_from_config(&config, state_db.clone()), + ThreadStoreBuilder::from_config(&config, state_db.clone()).build(), local_agent_graph_store_from_state_db(state_db.as_ref()), TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, @@ -3683,7 +3686,7 @@ async fn interrupted_fork_snapshot_uses_persisted_mid_turn_history_without_live_ Arc::new(crate::test_support::EmptyUserInstructionsProvider), /*analytics_events_client*/ None, passthrough_image_store(), - thread_store_from_config(&config, state_db.clone()), + ThreadStoreBuilder::from_config(&config, state_db.clone()).build(), local_agent_graph_store_from_state_db(state_db.as_ref()), TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, diff --git a/codex-rs/core/src/thread_store_builder.rs b/codex-rs/core/src/thread_store_builder.rs new file mode 100644 index 0000000000..52d6acbade --- /dev/null +++ b/codex-rs/core/src/thread_store_builder.rs @@ -0,0 +1,131 @@ +//! Configured thread-store construction and permanent-deletion cleanup. +//! Host cleanup runs after agent-board cleanup under the store's existing deletion locks. + +use crate::config::Config; +use crate::config::ThreadStoreConfig; +use codex_agent_message_board_extension::local_board_options; +use codex_features::Feature; +use codex_message_board::LocalAgentMessageBoard; +use codex_protocol::ThreadId; +use codex_rollout::state_db::StateDbHandle; +use codex_thread_store::InMemoryThreadStore; +use codex_thread_store::LocalThreadStore; +use codex_thread_store::LocalThreadStoreConfig; +use codex_thread_store::ThreadStore; +use codex_thread_store::ThreadStoreError; +use codex_thread_store::ThreadStoreFuture; +use std::sync::Arc; +use tracing::warn; + +type AdditionalCleanup = dyn Fn(Vec) -> ThreadStoreFuture<'static, ()> + Send + Sync; + +/// Builds the configured thread store with its existing agent-board cleanup. +/// Hosts can supply additional cleanup before the store is shared with callers. +pub struct ThreadStoreBuilder<'config> { + config: &'config Config, + state_db: Option, + additional_cleanup: Option>, +} + +impl<'config> ThreadStoreBuilder<'config> { + pub fn from_config(config: &'config Config, state_db: Option) -> Self { + Self { + config, + state_db, + additional_cleanup: None, + } + } + + /// Sets host-owned cleanup for permanent deletion in local stores. + /// Runs after successful agent-board cleanup and before deleting rollouts, under the same locks. + /// Cleanup must be idempotent because later deletion steps can fail and be retried. + /// + /// Cleanup must not call, or await work that calls, thread-store operations that reacquire + /// lifecycle or writer locks for any supplied thread ID (for example, flush or shutdown). + /// Deletion holds those locks until cleanup returns, so such calls would deadlock. + /// Flush or shut down threads before starting deletion. + /// + /// In-memory stores do not run persistent-data cleanup. + pub fn with_additional_cleanup( + mut self, + cleanup: impl Fn(Vec) -> ThreadStoreFuture<'static, ()> + Send + Sync + 'static, + ) -> Self { + self.additional_cleanup = Some(Arc::new(cleanup)); + self + } + + pub fn build(self) -> Arc { + let Self { + config, + state_db, + additional_cleanup, + } = self; + match &config.experimental_thread_store { + ThreadStoreConfig::Local => { + let compression_enabled = config + .features + .enabled(Feature::LocalThreadStoreCompression); + let background_migration_enabled = config + .features + .enabled(Feature::BackgroundPaginatedRolloutMigration); + let has_state_db = state_db.is_some(); + let sqlite = config.sqlite_config().clone(); + let store = Arc::new( + LocalThreadStore::new(LocalThreadStoreConfig::from_config(config), state_db) + .with_thread_data_cleanup(move |thread_ids| { + let sqlite = sqlite.clone(); + let additional_cleanup = additional_cleanup.clone(); + Box::pin(async move { + let boards = thread_ids + .iter() + .copied() + .map(Into::into) + .collect::>(); + LocalAgentMessageBoard::delete_boards( + &sqlite, + &local_board_options(), + &boards, + ) + .await + .map_err(|err| { + ThreadStoreError::Internal { + message: format!( + "failed to delete agent message boards: {err}" + ), + } + })?; + if let Some(cleanup) = additional_cleanup { + cleanup(thread_ids).await?; + } + Ok(()) + }) + }), + ); + if has_state_db && background_migration_enabled { + let startup_store = Arc::clone(&store); + let codex_home = config.codex_home.to_path_buf(); + tokio::spawn(async move { + if let Err(err) = startup_store.migrate_rollouts_on_startup().await { + warn!("failed to migrate legacy rollouts on startup: {err}"); + } + if compression_enabled { + codex_rollout::spawn_rollout_compression_worker( + codex_home, + codex_rollout::RolloutCompressionTrigger::Startup, + ); + } + }); + } else if compression_enabled { + codex_rollout::spawn_rollout_compression_worker( + config.codex_home.to_path_buf(), + codex_rollout::RolloutCompressionTrigger::Startup, + ); + } + store + } + ThreadStoreConfig::InMemory { id } => { + Arc::new(InMemoryThreadStore::for_id(id).with_state_db(state_db)) + } + } + } +} diff --git a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs index 0a9d78cd5c..808fe812a5 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -1,6 +1,7 @@ use super::*; use crate::StartThreadOptions; use crate::ThreadManager; +use crate::ThreadStoreBuilder; use crate::agent::child_config::apply_spawn_agent_service_tier; use crate::agent::child_config::build_agent_resume_config; use crate::agent::child_config::build_agent_spawn_config; @@ -19,7 +20,6 @@ use crate::session::tests::update_selected_settings_for_test; use crate::session::tests::update_turn_settings_for_test; use crate::session::turn_context::TurnContext; use crate::session_prefix::format_inter_agent_completion_message; -use crate::thread_manager::thread_store_from_config; use crate::tools::context::ToolOutput; use crate::tools::handlers::multi_agents_v2::FollowupTaskHandler as FollowupTaskHandlerV2; use crate::tools::handlers::multi_agents_v2::InterruptAgentHandler; @@ -4323,7 +4323,7 @@ async fn tool_handlers_cascade_close_and_resume_and_keep_explicitly_closed_subtr Arc::new(crate::test_support::EmptyUserInstructionsProvider), /*analytics_events_client*/ None, crate::thread_manager::passthrough_image_store(), - thread_store_from_config(&config, state_db.clone()), + ThreadStoreBuilder::from_config(&config, state_db.clone()).build(), local_agent_graph_store_from_state_db(state_db.as_ref()), "11111111-1111-4111-8111-111111111111".to_string(), /*attestation_provider*/ None, diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index 89081ae5ca..3cb6b51c1b 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -22,13 +22,13 @@ use codex_config::CloudConfigBundleLoader; use codex_core::CodexThread; pub use codex_core::StartThreadOptions; use codex_core::ThreadManager; +use codex_core::ThreadStoreBuilder; use codex_core::TimeProvider; pub use codex_core::TurnInputRequest; use codex_core::config::Config; use codex_core::resolve_installation_id; use codex_core::shell::Shell; use codex_core::shell::get_shell_by_model_provided_path; -use codex_core::thread_store_from_config; use codex_core::windows_sandbox::WindowsSandboxLevelExt; use codex_exec_server::CreateDirectoryOptions; use codex_exec_server::ExecutorFileSystem; @@ -847,7 +847,7 @@ impl TestCodexBuilder { let thread_store = self .thread_store .clone() - .unwrap_or_else(|| thread_store_from_config(&config, state_db.clone())); + .unwrap_or_else(|| ThreadStoreBuilder::from_config(&config, state_db.clone()).build()); let installation_id = resolve_installation_id(&config.codex_home).await?; let user_instructions_provider = self.user_instructions_provider.clone().unwrap_or_else(|| { diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index 630c515384..6967284e64 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -6,10 +6,10 @@ use codex_core::Prompt; use codex_core::ResponseEvent; use codex_core::StartThreadOptions; use codex_core::ThreadManager; +use codex_core::ThreadStoreBuilder; use codex_core::TurnInputRequest; use codex_core::X_CODEX_ROUTING_HINT_HEADER; use codex_core::resolve_installation_id; -use codex_core::thread_store_from_config; use codex_extension_api::empty_extension_registry; use codex_features::Feature; use codex_history::RolloutItem; @@ -2083,7 +2083,7 @@ async fn prefers_apikey_when_config_prefers_apikey_even_with_chatgpt_tokens() { Arc::new(codex_core::test_support::EmptyUserInstructionsProvider), /*analytics_events_client*/ None, codex_core::passthrough_image_store(), - thread_store_from_config(&config, /*state_db*/ None), + ThreadStoreBuilder::from_config(&config, /*state_db*/ None).build(), /*agent_graph_store*/ None, installation_id, /*attestation_provider*/ None, diff --git a/codex-rs/core/tests/suite/pending_input_persistence.rs b/codex-rs/core/tests/suite/pending_input_persistence.rs index c7dc52e36f..6312e5342a 100644 --- a/codex-rs/core/tests/suite/pending_input_persistence.rs +++ b/codex-rs/core/tests/suite/pending_input_persistence.rs @@ -662,10 +662,11 @@ async fn local_preparation_is_durable_before_first_input_and_survives_restart( let thread_id = initial.session_configured.thread_id; initial.thread_manager.remove_thread(&thread_id).await; // Reopen persisted context by ID, which supports both local history modes. - let store = codex_core::thread_store_from_config( + let store = codex_core::ThreadStoreBuilder::from_config( &initial.config, codex_core::init_state_db(&initial.config).await, - ); + ) + .build(); let context = store .load_latest_model_context(LoadThreadHistoryParams { thread_id, diff --git a/codex-rs/core/tests/suite/rollout_compression.rs b/codex-rs/core/tests/suite/rollout_compression.rs index 03f01be0b4..cf542007a9 100644 --- a/codex-rs/core/tests/suite/rollout_compression.rs +++ b/codex-rs/core/tests/suite/rollout_compression.rs @@ -137,8 +137,11 @@ async fn compressed_shared_fork_resume_preserves_checkpoint_and_frozen_history() .features .enable(Feature::LocalThreadStoreCompression)?; // Use production feature wiring, rather than manually writing zstd files or calling the worker. - let store = - codex_core::thread_store_from_config(&config, codex_core::init_state_db(&config).await); + let store = codex_core::ThreadStoreBuilder::from_config( + &config, + codex_core::init_state_db(&config).await, + ) + .build(); tokio::time::timeout(Duration::from_secs(30), async { loop { if paths diff --git a/codex-rs/core/tests/suite/subagent_notifications.rs b/codex-rs/core/tests/suite/subagent_notifications.rs index ec628fd960..44c26e15ec 100644 --- a/codex-rs/core/tests/suite/subagent_notifications.rs +++ b/codex-rs/core/tests/suite/subagent_notifications.rs @@ -1131,7 +1131,8 @@ async fn spawned_child_receives_forked_parent_context( .expect("legacy child thread id"), )?; // Read the acknowledged fork from storage without flushing the live child first. - let reopened_store = codex_core::thread_store_from_config(&test.config, /*state_db*/ None); + let reopened_store = + codex_core::ThreadStoreBuilder::from_config(&test.config, /*state_db*/ None).build(); let persisted = reopened_store .load_latest_model_context(LoadThreadHistoryParams { thread_id: child_thread_id, diff --git a/codex-rs/thread-manager-sample/src/main.rs b/codex-rs/thread-manager-sample/src/main.rs index 1ced86ef11..d07f21f02c 100644 --- a/codex-rs/thread-manager-sample/src/main.rs +++ b/codex-rs/thread-manager-sample/src/main.rs @@ -52,6 +52,7 @@ use codex_core_api::StartIfIdleSubmission; use codex_core_api::StartThreadOptions; use codex_core_api::TerminalResizeReflowConfig; use codex_core_api::ThreadManager; +use codex_core_api::ThreadStoreBuilder; use codex_core_api::ThreadStoreConfig; use codex_core_api::ToolSuggestConfig; use codex_core_api::TuiKeymap; @@ -75,7 +76,6 @@ use codex_core_api::passthrough_image_store; use codex_core_api::resolve_bootstrap_respect_system_proxy; use codex_core_api::resolve_installation_id; use codex_core_api::set_default_originator; -use codex_core_api::thread_store_from_config; #[derive(Debug, Parser)] #[command( @@ -129,7 +129,7 @@ async fn run_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { config.codex_self_exe.clone(), config.codex_linux_sandbox_exe.clone(), )?; - let thread_store = thread_store_from_config(&config, state_db.clone()); + let thread_store = ThreadStoreBuilder::from_config(&config, state_db.clone()).build(); let environment_manager = Arc::new( EnvironmentManager::from_codex_home( config.codex_home.clone(),