diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index 00f855156c..93a535f8f8 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -210,16 +210,17 @@ use codex_core::exec::ExecCapturePolicy; use codex_core::exec::ExecExpiration; use codex_core::exec::ExecParams; use codex_core::exec_env::create_env; -use codex_core::find_archived_thread_path_by_id_str_with_state_db; +use codex_core::find_archived_thread_path_by_id_str; use codex_core::find_thread_name_by_id; use codex_core::find_thread_names_by_ids; -use codex_core::find_thread_path_by_id_str_with_state_db; +use codex_core::find_thread_path_by_id_str; use codex_core::mcp::auth::discover_supported_scopes; use codex_core::mcp::auth::resolve_oauth_scopes; use codex_core::mcp::collect_mcp_snapshot; use codex_core::mcp::group_tools_by_server; use codex_core::models_manager::collaboration_mode_presets::CollaborationModesConfig; use codex_core::parse_cursor; +use codex_core::path_utils::normalize_for_path_comparison; use codex_core::plugins::MarketplaceError; use codex_core::plugins::MarketplacePluginSource; use codex_core::plugins::OPENAI_CURATED_MARKETPLACE_NAME; @@ -234,7 +235,6 @@ use codex_core::read_session_meta_line; use codex_core::rollout_date_parts; use codex_core::sandboxing::SandboxPermissions; use codex_core::state_db::StateDbHandle; -use codex_core::state_db::get_state_db; use codex_core::state_db::reconcile_rollout; use codex_core::windows_sandbox::WindowsSandboxLevelExt; use codex_core::windows_sandbox::WindowsSandboxSetupMode as CoreWindowsSandboxSetupMode; @@ -300,7 +300,6 @@ use std::sync::atomic::Ordering; use std::time::Duration; use std::time::SystemTime; use tokio::sync::Mutex; -use tokio::sync::OnceCell; use tokio::sync::broadcast; use tokio::sync::oneshot; use tokio::sync::watch; @@ -385,7 +384,7 @@ pub(crate) struct CodexMessageProcessor { pending_fuzzy_searches: Arc>>>, fuzzy_search_sessions: Arc>>, background_tasks: TaskTracker, - app_server_state_db_cell: OnceCell, + state_db: StateDbHandle, feedback: CodexFeedback, log_db: Option, } @@ -432,6 +431,7 @@ pub(crate) struct CodexMessageProcessorArgs { pub(crate) cloud_requirements: Arc>, pub(crate) feedback: CodexFeedback, pub(crate) log_db: Option, + pub(crate) state_db: StateDbHandle, } impl CodexMessageProcessor { @@ -485,60 +485,6 @@ impl CodexMessageProcessor { Ok((thread_id, thread)) } - async fn app_server_state_db(&self) -> Option { - let config = Arc::clone(&self.config); - self.app_server_state_db_cell - .get_or_try_init(|| async move { get_state_db(&config).await.ok_or(()) }) - .await - .ok() - .cloned() - } - - async fn find_thread_path_by_id_with_state_db( - &self, - thread_id: ThreadId, - state_db_ctx: Option, - ) -> std::io::Result> { - find_thread_path_by_id_str_with_state_db( - &self.config.codex_home, - &thread_id.to_string(), - state_db_ctx, - ) - .await - } - - async fn find_archived_thread_path_by_id_with_state_db( - &self, - thread_id: ThreadId, - state_db_ctx: Option, - ) -> std::io::Result> { - find_archived_thread_path_by_id_str_with_state_db( - &self.config.codex_home, - &thread_id.to_string(), - state_db_ctx, - ) - .await - } - - async fn find_thread_path_by_id( - &self, - thread_id: ThreadId, - ) -> std::io::Result> { - self.find_thread_path_by_id_with_state_db(thread_id, self.app_server_state_db().await) - .await - } - - async fn find_archived_thread_path_by_id( - &self, - thread_id: ThreadId, - ) -> std::io::Result> { - self.find_archived_thread_path_by_id_with_state_db( - thread_id, - self.app_server_state_db().await, - ) - .await - } - pub fn new(args: CodexMessageProcessorArgs) -> Self { let CodexMessageProcessorArgs { auth_manager, @@ -551,6 +497,7 @@ impl CodexMessageProcessor { cloud_requirements, feedback, log_db, + state_db, } = args; Self { auth_manager, @@ -569,7 +516,7 @@ impl CodexMessageProcessor { pending_fuzzy_searches: Arc::new(Mutex::new(HashMap::new())), fuzzy_search_sessions: Arc::new(Mutex::new(HashMap::new())), background_tasks: TaskTracker::new(), - app_server_state_db_cell: OnceCell::new(), + state_db, feedback, log_db, } @@ -2294,7 +2241,13 @@ impl CodexMessageProcessor { } }; - let rollout_path = match self.find_thread_path_by_id(thread_id).await { + let rollout_path = match find_thread_path_by_id_str( + &self.config.codex_home, + &thread_id.to_string(), + Some(self.state_db.clone()), + ) + .await + { Ok(Some(p)) => p, Ok(None) => { let error = JSONRPCErrorError { @@ -2442,7 +2395,13 @@ impl CodexMessageProcessor { return; } - let thread_exists = match self.find_thread_path_by_id(thread_id).await { + let thread_exists = match find_thread_path_by_id_str( + &self.config.codex_home, + &thread_id.to_string(), + Some(self.state_db.clone()), + ) + .await + { Ok(Some(_)) => true, Ok(None) => false, Err(err) => { @@ -2526,7 +2485,7 @@ impl CodexMessageProcessor { let loaded_thread = self.thread_manager.get_thread(thread_uuid).await.ok(); let mut state_db_ctx = loaded_thread.as_ref().and_then(|thread| thread.state_db()); if state_db_ctx.is_none() { - state_db_ctx = self.app_server_state_db().await; + state_db_ctx = Some(self.state_db.clone()); } let Some(state_db_ctx) = state_db_ctx else { self.send_internal_error( @@ -2729,17 +2688,20 @@ impl CodexMessageProcessor { return Ok(()); } - let rollout_path = match self - .find_thread_path_by_id_with_state_db(thread_uuid, Some(state_db_ctx.clone())) - .await + let rollout_path = match find_thread_path_by_id_str( + &self.config.codex_home, + &thread_uuid.to_string(), + Some(state_db_ctx.clone()), + ) + .await { Ok(Some(path)) => path, - Ok(None) => match self - .find_archived_thread_path_by_id_with_state_db( - thread_uuid, - Some(state_db_ctx.clone()), - ) - .await + Ok(None) => match find_archived_thread_path_by_id_str( + &self.config.codex_home, + &thread_uuid.to_string(), + Some(state_db_ctx.clone()), + ) + .await { Ok(Some(path)) => path, Ok(None) => { @@ -2799,7 +2761,13 @@ impl CodexMessageProcessor { } }; - let archived_path = match self.find_archived_thread_path_by_id(thread_id).await { + let archived_path = match find_archived_thread_path_by_id_str( + &self.config.codex_home, + &thread_id.to_string(), + Some(self.state_db.clone()), + ) + .await + { Ok(Some(path)) => path, Ok(None) => { let error = JSONRPCErrorError { @@ -2823,7 +2791,7 @@ impl CodexMessageProcessor { let rollout_path_display = archived_path.display().to_string(); let fallback_provider = self.config.model_provider_id.clone(); - let state_db_ctx = self.app_server_state_db().await; + let state_db_ctx = Some(self.state_db.clone()); let archived_folder = self .config .codex_home @@ -3303,24 +3271,25 @@ impl CodexMessageProcessor { let loaded_thread = self.thread_manager.get_thread(thread_uuid).await.ok(); let loaded_thread_state_db = loaded_thread.as_ref().and_then(|thread| thread.state_db()); - let app_server_state_db = if loaded_thread_state_db.is_some() { + let state_db = if loaded_thread_state_db.is_some() { None } else { - self.app_server_state_db().await + Some(self.state_db.clone()) }; let db_summary = if let Some(state_db_ctx) = loaded_thread_state_db.as_ref() { read_summary_from_state_db_context_by_thread_id(Some(state_db_ctx), thread_uuid).await } else { - read_summary_from_state_db_by_thread_id(app_server_state_db.as_ref(), thread_uuid).await + read_summary_from_state_db_by_thread_id(state_db.as_ref(), thread_uuid).await }; let mut rollout_path = db_summary.as_ref().map(|summary| summary.path.clone()); if rollout_path.is_none() || include_turns { - let state_db_ctx = loaded_thread_state_db - .clone() - .or_else(|| app_server_state_db.clone()); - rollout_path = match self - .find_thread_path_by_id_with_state_db(thread_uuid, state_db_ctx) - .await + let state_db_ctx = loaded_thread_state_db.clone().or_else(|| state_db.clone()); + rollout_path = match find_thread_path_by_id_str( + &self.config.codex_home, + &thread_uuid.to_string(), + state_db_ctx, + ) + .await { Ok(Some(path)) => Some(path), Ok(None) => { @@ -3709,7 +3678,7 @@ impl CodexMessageProcessor { let InitialHistory::Resumed(resumed_history) = thread_history else { return None; }; - let state_db_ctx = self.app_server_state_db().await?; + let state_db_ctx = &self.state_db; let persisted_metadata = state_db_ctx .get_thread(resumed_history.conversation_id) .await @@ -3742,7 +3711,15 @@ impl CodexMessageProcessor { if path.exists() { path } else { - match self.find_thread_path_by_id(existing_thread_id).await { + match find_thread_path_by_id_str( + &self.config.codex_home, + &existing_thread_id.to_string(), + existing_thread + .state_db() + .or_else(|| Some(self.state_db.clone())), + ) + .await + { Ok(Some(path)) => path, Ok(None) => { self.send_invalid_request_error( @@ -3763,7 +3740,15 @@ impl CodexMessageProcessor { } } } else { - match self.find_thread_path_by_id(existing_thread_id).await { + match find_thread_path_by_id_str( + &self.config.codex_home, + &existing_thread_id.to_string(), + existing_thread + .state_db() + .or_else(|| Some(self.state_db.clone())), + ) + .await + { Ok(Some(path)) => path, Ok(None) => { self.send_invalid_request_error( @@ -3917,7 +3902,13 @@ impl CodexMessageProcessor { } }; - match self.find_thread_path_by_id(existing_thread_id).await { + match find_thread_path_by_id_str( + &self.config.codex_home, + &existing_thread_id.to_string(), + Some(self.state_db.clone()), + ) + .await + { Ok(Some(path)) => path, Ok(None) => { self.send_invalid_request_error( @@ -4044,7 +4035,13 @@ impl CodexMessageProcessor { } }; - match self.find_thread_path_by_id(existing_thread_id).await { + match find_thread_path_by_id_str( + &self.config.codex_home, + &existing_thread_id.to_string(), + Some(self.state_db.clone()), + ) + .await + { Ok(Some(p)) => (p, Some(existing_thread_id)), Ok(None) => { self.send_invalid_request_error( @@ -4065,9 +4062,9 @@ impl CodexMessageProcessor { } }; - let app_server_state_db = self.app_server_state_db().await; + let state_db = &self.state_db; let history_cwd = read_history_cwd_from_state_db( - app_server_state_db.as_ref(), + Some(state_db), source_thread_id, rollout_path.as_path(), ) @@ -4293,13 +4290,10 @@ impl CodexMessageProcessor { request_id: ConnectionRequestId, params: GetConversationSummaryParams, ) { - let app_server_state_db = self.app_server_state_db().await; + let state_db = &self.state_db; if let GetConversationSummaryParams::ThreadId { conversation_id } = ¶ms - && let Some(summary) = read_summary_from_state_db_by_thread_id( - app_server_state_db.as_ref(), - *conversation_id, - ) - .await + && let Some(summary) = + read_summary_from_state_db_by_thread_id(Some(state_db), *conversation_id).await { let response = GetConversationSummaryResponse { summary }; self.outgoing.send_response(request_id, response).await; @@ -4315,12 +4309,12 @@ impl CodexMessageProcessor { } } GetConversationSummaryParams::ThreadId { conversation_id } => { - match self - .find_thread_path_by_id_with_state_db( - conversation_id, - app_server_state_db.clone(), - ) - .await + match find_thread_path_by_id_str( + &self.config.codex_home, + &conversation_id.to_string(), + Some(self.state_db.clone()), + ) + .await { Ok(Some(p)) => p, _ => { @@ -4401,7 +4395,11 @@ impl CodexMessageProcessor { let fallback_provider = self.config.model_provider_id.clone(); let (allowed_sources_vec, source_kind_filter) = compute_source_filters(source_kinds); let allowed_sources = allowed_sources_vec.as_slice(); - let state_db_ctx = self.app_server_state_db().await; + let state_db_ctx = Some(self.state_db.clone()); + let normalized_cwd = cwd + .as_ref() + .and_then(|cwd| normalize_for_path_comparison(cwd).ok()); + let cwd_filter = normalized_cwd.as_ref().or(cwd.as_ref()); while remaining > 0 { let page_size = remaining.min(THREAD_LIST_MAX_LIMIT); @@ -4443,7 +4441,7 @@ impl CodexMessageProcessor { let mut filtered = Vec::with_capacity(page.items.len()); for it in page.items { - let Some(summary) = summary_from_thread_list_item( + let Some(mut summary) = summary_from_thread_list_item( it, fallback_provider.as_str(), state_db_ctx.as_ref(), @@ -4455,10 +4453,11 @@ impl CodexMessageProcessor { if source_kind_filter .as_ref() .is_none_or(|filter| source_kind_matches(&summary.source, filter)) - && cwd - .as_ref() - .is_none_or(|expected_cwd| &summary.cwd == expected_cwd) + && cwd_filter.is_none_or(|expected_cwd| &summary.cwd == expected_cwd) { + if let Some(cwd) = cwd.as_ref() { + summary.cwd = cwd.clone(); + } filtered.push(summary); if filtered.len() >= remaining { break; @@ -5254,7 +5253,7 @@ impl CodexMessageProcessor { self.finalize_thread_teardown(thread_id).await; if state_db_ctx.is_none() { - state_db_ctx = self.app_server_state_db().await; + state_db_ctx = Some(self.state_db.clone()); } // Move the rollout file to archived. @@ -6605,18 +6604,24 @@ impl CodexMessageProcessor { let rollout_path = if let Some(path) = parent_thread.rollout_path() { path } else { - self.find_thread_path_by_id(parent_thread_id) - .await - .map_err(|err| JSONRPCErrorError { - code: INTERNAL_ERROR_CODE, - message: format!("failed to locate thread id {parent_thread_id}: {err}"), - data: None, - })? - .ok_or_else(|| JSONRPCErrorError { - code: INVALID_REQUEST_ERROR_CODE, - message: format!("no rollout found for thread id {parent_thread_id}"), - data: None, - })? + find_thread_path_by_id_str( + &self.config.codex_home, + &parent_thread_id.to_string(), + parent_thread + .state_db() + .or_else(|| Some(self.state_db.clone())), + ) + .await + .map_err(|err| JSONRPCErrorError { + code: INTERNAL_ERROR_CODE, + message: format!("failed to locate thread id {parent_thread_id}: {err}"), + data: None, + })? + .ok_or_else(|| JSONRPCErrorError { + code: INVALID_REQUEST_ERROR_CODE, + message: format!("no rollout found for thread id {parent_thread_id}"), + data: None, + })? }; let mut config = self.config.as_ref().clone(); @@ -7209,11 +7214,10 @@ impl CodexMessageProcessor { if let Some(log_db) = self.log_db.as_ref() { log_db.flush().await; } - let state_db_ctx = self.app_server_state_db().await; - match (state_db_ctx.as_ref(), conversation_id) { - (Some(state_db_ctx), Some(conversation_id)) => { + match conversation_id { + Some(conversation_id) => { let thread_id_text = conversation_id.to_string(); - match state_db_ctx.query_feedback_logs(&thread_id_text).await { + match self.state_db.query_feedback_logs(&thread_id_text).await { Ok(logs) if logs.is_empty() => None, Ok(logs) => Some(logs), Err(err) => { @@ -7224,7 +7228,7 @@ impl CodexMessageProcessor { } } } - _ => None, + None => None, } } else { None @@ -7362,11 +7366,14 @@ impl CodexMessageProcessor { async fn resolve_rollout_path(&self, conversation_id: ThreadId) -> Option { match self.thread_manager.get_thread(conversation_id).await { Ok(conv) => conv.rollout_path(), - Err(_) => self - .find_thread_path_by_id(conversation_id) - .await - .ok() - .flatten(), + Err(_) => find_thread_path_by_id_str( + &self.config.codex_home, + &conversation_id.to_string(), + Some(self.state_db.clone()), + ) + .await + .ok() + .flatten(), } } } diff --git a/codex-rs/app-server/src/in_process.rs b/codex-rs/app-server/src/in_process.rs index 79f32c9083..2048f8b607 100644 --- a/codex-rs/app-server/src/in_process.rs +++ b/codex-rs/app-server/src/in_process.rs @@ -77,6 +77,8 @@ use codex_arg0::Arg0DispatchPaths; use codex_core::config::Config; use codex_core::config_loader::CloudRequirementsLoader; use codex_core::config_loader::LoaderOverrides; +use codex_core::state_db::StateDbHandle; +use codex_core::state_db::init; use codex_feedback::CodexFeedback; use codex_protocol::protocol::SessionSource; use tokio::sync::mpsc; @@ -325,7 +327,13 @@ impl InProcessClientHandle { /// the runtime is shut down and an `InvalidData` error is returned. pub async fn start(args: InProcessStartArgs) -> IoResult { let initialize = args.initialize.clone(); - let client = start_uninitialized(args); + let state_db = init(args.config.as_ref()).await.ok_or_else(|| { + IoError::other(format!( + "failed to initialize state db at {}", + args.config.sqlite_home.display() + )) + })?; + let client = start_uninitialized(args, state_db); let initialize_response = client .request(ClientRequest::Initialize { @@ -345,7 +353,7 @@ pub async fn start(args: InProcessStartArgs) -> IoResult Ok(client) } -fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle { +fn start_uninitialized(args: InProcessStartArgs, state_db: StateDbHandle) -> InProcessClientHandle { let channel_capacity = args.channel_capacity.max(1); let (client_tx, mut client_rx) = mpsc::channel::(channel_capacity); let (event_tx, event_rx) = mpsc::channel::(channel_capacity); @@ -388,6 +396,7 @@ fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle { cloud_requirements: args.cloud_requirements, feedback: args.feedback, log_db: None, + state_db, config_warnings: args.config_warnings, session_source: args.session_source, enable_codex_api_key_env: args.enable_codex_api_key_env, @@ -693,13 +702,17 @@ mod tests { use codex_app_server_protocol::TurnStatus; use codex_core::config::ConfigBuilder; use pretty_assertions::assert_eq; + use tempfile::TempDir; async fn build_test_config() -> Config { - match ConfigBuilder::default().build().await { - Ok(config) => config, - Err(_) => Config::load_default_with_cli_overrides(Vec::new()) - .expect("default config should load"), - } + let codex_home = TempDir::new().expect("create temp dir"); + let config = ConfigBuilder::default() + .codex_home(codex_home.path().to_path_buf()) + .build() + .await + .expect("test config should build"); + std::mem::forget(codex_home); + config } async fn start_test_client_with_capacity( diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index dc2512aaec..aa90de7ef4 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -8,9 +8,11 @@ use codex_core::config::ConfigBuilder; use codex_core::config_loader::CloudRequirementsLoader; use codex_core::config_loader::ConfigLayerStackOrdering; use codex_core::config_loader::LoaderOverrides; +use codex_core::state_db::init; use codex_utils_cli::CliConfigOverrides; use std::collections::HashMap; use std::collections::HashSet; +use std::io::Error as IoError; use std::io::ErrorKind; use std::io::Result as IoResult; use std::sync::Arc; @@ -489,6 +491,12 @@ pub async fn run_main_with_transport( } let feedback = CodexFeedback::new(); + let state_db = init(&config).await.ok_or_else(|| { + IoError::other(format!( + "failed to initialize state db at {}", + config.sqlite_home.display() + )) + })?; let otel = codex_core::otel_init::build_provider( &config, @@ -522,13 +530,7 @@ pub async fn run_main_with_transport( let feedback_layer = feedback.logger_layer(); let feedback_metadata_layer = feedback.metadata_layer(); - let log_db = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.model_provider_id.clone(), - ) - .await - .ok() - .map(log_db::start); + let log_db = Some(log_db::start(Arc::clone(&state_db))); let log_db_layer = log_db .clone() .map(|layer| layer.with_filter(Targets::new().with_default(Level::TRACE))); @@ -618,6 +620,7 @@ pub async fn run_main_with_transport( cloud_requirements: cloud_requirements.clone(), feedback: feedback.clone(), log_db, + state_db, config_warnings, session_source, enable_codex_api_key_env: false, diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index d070738fba..7232b3d4ab 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -66,6 +66,7 @@ use codex_core::default_client::get_codex_user_agent; use codex_core::default_client::set_default_client_residency_requirement; use codex_core::default_client::set_default_originator; use codex_core::models_manager::collaboration_mode_presets::CollaborationModesConfig; +use codex_core::state_db::StateDbHandle; use codex_features::Feature; use codex_feedback::CodexFeedback; use codex_login::auth::ExternalAuthRefreshContext; @@ -181,6 +182,7 @@ pub(crate) struct MessageProcessorArgs { pub(crate) cloud_requirements: CloudRequirementsLoader, pub(crate) feedback: CodexFeedback, pub(crate) log_db: Option, + pub(crate) state_db: StateDbHandle, pub(crate) config_warnings: Vec, pub(crate) session_source: SessionSource, pub(crate) enable_codex_api_key_env: bool, @@ -199,6 +201,7 @@ impl MessageProcessor { cloud_requirements, feedback, log_db, + state_db, config_warnings, session_source, enable_codex_api_key_env, @@ -242,6 +245,7 @@ impl MessageProcessor { cloud_requirements: cloud_requirements.clone(), feedback, log_db, + state_db, }); // Keep plugin startup warmups aligned at app-server startup. // TODO(xl): Move into PluginManager once this no longer depends on config feature gating. diff --git a/codex-rs/app-server/src/message_processor/tracing_tests.rs b/codex-rs/app-server/src/message_processor/tracing_tests.rs index aa66f79f29..9128f906c9 100644 --- a/codex-rs/app-server/src/message_processor/tracing_tests.rs +++ b/codex-rs/app-server/src/message_processor/tracing_tests.rs @@ -24,6 +24,7 @@ use codex_core::config::Config; use codex_core::config::ConfigBuilder; use codex_core::config_loader::CloudRequirementsLoader; use codex_core::config_loader::LoaderOverrides; +use codex_core::state_db::init; use codex_feedback::CodexFeedback; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::W3cTraceContext; @@ -117,7 +118,7 @@ impl TracingHarness { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; let config = Arc::new(build_test_config(codex_home.path(), &server.uri()).await?); - let (processor, outgoing_rx) = build_test_processor(config); + let (processor, outgoing_rx) = build_test_processor(config).await; let tracing = init_test_tracing(); tracing.exporter.reset(); tracing::callsite::rebuild_interest_cache(); @@ -224,7 +225,7 @@ async fn build_test_config(codex_home: &Path, server_uri: &str) -> Result, ) -> ( MessageProcessor, @@ -232,6 +233,9 @@ fn build_test_processor( ) { let (outgoing_tx, outgoing_rx) = mpsc::channel(16); let outgoing = Arc::new(OutgoingMessageSender::new(outgoing_tx)); + let state_db = init(config.as_ref()) + .await + .expect("state db should initialize in tracing tests"); let processor = MessageProcessor::new(MessageProcessorArgs { outgoing, arg0_paths: Arg0DispatchPaths::default(), @@ -241,6 +245,7 @@ fn build_test_processor( cloud_requirements: CloudRequirementsLoader::default(), feedback: CodexFeedback::new(), log_db: None, + state_db, config_warnings: Vec::new(), session_source: SessionSource::VSCode, enable_codex_api_key_env: false, diff --git a/codex-rs/app-server/tests/suite/v2/thread_archive.rs b/codex-rs/app-server/tests/suite/v2/thread_archive.rs index 20fd6fd884..3f58e35c92 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_archive.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_archive.rs @@ -58,7 +58,7 @@ async fn thread_archive_requires_materialized_rollout() -> Result<()> { rollout_path.display() ); assert!( - find_thread_path_by_id_str(codex_home.path(), &thread.id) + find_thread_path_by_id_str(codex_home.path(), &thread.id, /*state_db_ctx*/ None) .await? .is_none(), "thread id should not be discoverable before rollout materialization" @@ -113,9 +113,10 @@ async fn thread_archive_requires_materialized_rollout() -> Result<()> { rollout_path.display() ); - let discovered_path = find_thread_path_by_id_str(codex_home.path(), &thread.id) - .await? - .expect("expected rollout path for thread id to exist after materialization"); + let discovered_path = + find_thread_path_by_id_str(codex_home.path(), &thread.id, /*state_db_ctx*/ None) + .await? + .expect("expected rollout path for thread id to exist after materialization"); assert_paths_match_on_disk(&discovered_path, &rollout_path)?; let archive_id = mcp diff --git a/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs b/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs index b2ae60ae35..588764edb8 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs @@ -75,9 +75,10 @@ async fn thread_unarchive_moves_rollout_back_into_sessions_directory() -> Result ) .await??; - let found_rollout_path = find_thread_path_by_id_str(codex_home.path(), &thread.id) - .await? - .expect("expected rollout path for thread id to exist"); + let found_rollout_path = + find_thread_path_by_id_str(codex_home.path(), &thread.id, /*state_db_ctx*/ None) + .await? + .expect("expected rollout path for thread id to exist"); assert_paths_match_on_disk(&found_rollout_path, &rollout_path)?; let archive_id = mcp @@ -92,9 +93,13 @@ async fn thread_unarchive_moves_rollout_back_into_sessions_directory() -> Result .await??; let _: ThreadArchiveResponse = to_response::(archive_resp)?; - let archived_path = find_archived_thread_path_by_id_str(codex_home.path(), &thread.id) - .await? - .expect("expected archived rollout path for thread id to exist"); + let archived_path = find_archived_thread_path_by_id_str( + codex_home.path(), + &thread.id, + /*state_db_ctx*/ None, + ) + .await? + .expect("expected archived rollout path for thread id to exist"); let archived_path_display = archived_path.display(); assert!( archived_path.exists(), diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index a3b9897688..87d25bfd7e 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -205,6 +205,9 @@ impl AgentControl { .or(find_thread_path_by_id_str( config.codex_home.as_path(), &parent_thread_id.to_string(), + parent_thread + .as_ref() + .and_then(|parent_thread| parent_thread.state_db()), ) .await?) .ok_or_else(|| { @@ -377,6 +380,7 @@ impl AgentControl { } let state = self.upgrade()?; let mut reservation = self.state.reserve_spawn_slot(config.agent_max_threads)?; + let state_db_ctx = state_db::get_state_db(&config).await; let (session_source, agent_metadata) = match session_source { SessionSource::SubAgent(SubAgentSource::ThreadSpawn { parent_thread_id, @@ -386,7 +390,7 @@ impl AgentControl { agent_nickname: _, }) => { let (resumed_agent_nickname, resumed_agent_role) = - if let Some(state_db_ctx) = state_db::get_state_db(&config).await { + if let Some(state_db_ctx) = state_db_ctx.as_ref() { match state_db_ctx.get_thread(thread_id).await { Ok(Some(metadata)) => (metadata.agent_nickname, metadata.agent_role), Ok(None) | Err(_) => (None, None), @@ -413,18 +417,22 @@ impl AgentControl { let inherited_exec_policy = self .inherited_exec_policy_for_source(&state, Some(&session_source), &config) .await; - let rollout_path = - match find_thread_path_by_id_str(config.codex_home.as_path(), &thread_id.to_string()) - .await? - { - Some(rollout_path) => rollout_path, - None => find_archived_thread_path_by_id_str( - config.codex_home.as_path(), - &thread_id.to_string(), - ) - .await? - .ok_or_else(|| CodexErr::ThreadNotFound(thread_id))?, - }; + let rollout_path = match find_thread_path_by_id_str( + config.codex_home.as_path(), + &thread_id.to_string(), + state_db_ctx.clone(), + ) + .await? + { + Some(rollout_path) => rollout_path, + None => find_archived_thread_path_by_id_str( + config.codex_home.as_path(), + &thread_id.to_string(), + state_db_ctx, + ) + .await? + .ok_or_else(|| CodexErr::ThreadNotFound(thread_id))?, + }; let resumed_thread = state .resume_thread_from_rollout_with_source( diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 5008b60849..71f4328aed 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -144,12 +144,10 @@ pub use rollout::SESSIONS_SUBDIR; pub use rollout::SessionMeta; pub use rollout::append_thread_name; pub use rollout::find_archived_thread_path_by_id_str; -pub use rollout::find_archived_thread_path_by_id_str_with_state_db; #[deprecated(note = "use find_thread_path_by_id_str")] pub use rollout::find_conversation_path_by_id_str; pub use rollout::find_thread_name_by_id; pub use rollout::find_thread_path_by_id_str; -pub use rollout::find_thread_path_by_id_str_with_state_db; pub use rollout::find_thread_path_by_name_str; pub use rollout::list::Cursor; pub use rollout::list::ThreadItem; diff --git a/codex-rs/core/src/shell_snapshot.rs b/codex-rs/core/src/shell_snapshot.rs index 29b50cb9e8..48199a45ea 100644 --- a/codex-rs/core/src/shell_snapshot.rs +++ b/codex-rs/core/src/shell_snapshot.rs @@ -518,7 +518,8 @@ pub async fn cleanup_stale_snapshots(codex_home: &Path, active_session_id: Threa continue; } - let rollout_path = find_thread_path_by_id_str(codex_home, session_id).await?; + let rollout_path = + find_thread_path_by_id_str(codex_home, session_id, /*state_db_ctx*/ None).await?; let Some(rollout_path) = rollout_path else { remove_snapshot_file(&path).await; continue; diff --git a/codex-rs/core/tests/suite/rollout_list_find.rs b/codex-rs/core/tests/suite/rollout_list_find.rs index a9ea2b9c8a..7c0ce6e598 100644 --- a/codex-rs/core/tests/suite/rollout_list_find.rs +++ b/codex-rs/core/tests/suite/rollout_list_find.rs @@ -78,9 +78,10 @@ async fn find_locates_rollout_file_by_id() { let id = Uuid::new_v4(); let expected = write_minimal_rollout_with_id(home.path(), id); - let found = find_thread_path_by_id_str(home.path(), &id.to_string()) - .await - .unwrap(); + let found = + find_thread_path_by_id_str(home.path(), &id.to_string(), /*state_db_ctx*/ None) + .await + .unwrap(); assert_eq!(found.unwrap(), expected); } @@ -94,9 +95,10 @@ async fn find_handles_gitignore_covering_codex_home_directory() { let id = Uuid::new_v4(); let expected = write_minimal_rollout_with_id(&codex_home, id); - let found = find_thread_path_by_id_str(&codex_home, &id.to_string()) - .await - .unwrap(); + let found = + find_thread_path_by_id_str(&codex_home, &id.to_string(), /*state_db_ctx*/ None) + .await + .unwrap(); assert_eq!(found, Some(expected)); } @@ -114,9 +116,10 @@ async fn find_prefers_sqlite_path_by_id() { write_minimal_rollout_with_id(home.path(), id); upsert_thread_metadata(home.path(), thread_id, db_path.clone()).await; - let found = find_thread_path_by_id_str(home.path(), &id.to_string()) - .await - .unwrap(); + let found = + find_thread_path_by_id_str(home.path(), &id.to_string(), /*state_db_ctx*/ None) + .await + .unwrap(); assert_eq!(found, Some(db_path)); } @@ -133,9 +136,10 @@ async fn find_falls_back_to_filesystem_when_sqlite_has_no_match() { .join("sessions/2030/12/30/rollout-2030-12-30T00-00-00-unrelated.jsonl"); upsert_thread_metadata(home.path(), unrelated_thread_id, unrelated_path).await; - let found = find_thread_path_by_id_str(home.path(), &id.to_string()) - .await - .unwrap(); + let found = + find_thread_path_by_id_str(home.path(), &id.to_string(), /*state_db_ctx*/ None) + .await + .unwrap(); assert_eq!(found, Some(expected)); } @@ -147,9 +151,10 @@ async fn find_ignores_granular_gitignore_rules() { let expected = write_minimal_rollout_with_id(home.path(), id); std::fs::write(home.path().join("sessions/.gitignore"), "*.jsonl\n").unwrap(); - let found = find_thread_path_by_id_str(home.path(), &id.to_string()) - .await - .unwrap(); + let found = + find_thread_path_by_id_str(home.path(), &id.to_string(), /*state_db_ctx*/ None) + .await + .unwrap(); assert_eq!(found, Some(expected)); } @@ -210,9 +215,13 @@ async fn find_archived_locates_rollout_file_by_id() { let id = Uuid::new_v4(); let expected = write_minimal_rollout_with_id_in_subdir(home.path(), "archived_sessions", id); - let found = find_archived_thread_path_by_id_str(home.path(), &id.to_string()) - .await - .unwrap(); + let found = find_archived_thread_path_by_id_str( + home.path(), + &id.to_string(), + /*state_db_ctx*/ None, + ) + .await + .unwrap(); assert_eq!(found, Some(expected)); } diff --git a/codex-rs/rollout/src/lib.rs b/codex-rs/rollout/src/lib.rs index 3c2c6e1e52..160792a390 100644 --- a/codex-rs/rollout/src/lib.rs +++ b/codex-rs/rollout/src/lib.rs @@ -33,11 +33,9 @@ pub use codex_protocol::protocol::SessionMeta; pub use config::RolloutConfig; pub use config::RolloutConfigView; pub use list::find_archived_thread_path_by_id_str; -pub use list::find_archived_thread_path_by_id_str_with_state_db; pub use list::find_thread_path_by_id_str; #[deprecated(note = "use find_thread_path_by_id_str")] pub use list::find_thread_path_by_id_str as find_conversation_path_by_id_str; -pub use list::find_thread_path_by_id_str_with_state_db; pub use list::rollout_date_parts; pub use policy::EventPersistenceMode; pub use recorder::RolloutRecorder; diff --git a/codex-rs/rollout/src/list.rs b/codex-rs/rollout/src/list.rs index e1242e37f9..8bef6d2b6b 100644 --- a/codex-rs/rollout/src/list.rs +++ b/codex-rs/rollout/src/list.rs @@ -1251,51 +1251,21 @@ async fn find_thread_path_by_id_str_in_subdir( } /// Locate a recorded thread rollout file by its UUID string using the existing -/// paginated listing implementation. Returns `Ok(Some(path))` if found, `Ok(None)` if not present -/// or the id is invalid. +/// paginated listing implementation. Uses `state_db_ctx` when available to avoid reopening +/// SQLite. Returns `Ok(Some(path))` if found, `Ok(None)` if not present or the id is invalid. pub async fn find_thread_path_by_id_str( codex_home: &Path, id_str: &str, -) -> io::Result> { - find_thread_path_by_id_str_in_subdir( - codex_home, - SESSIONS_SUBDIR, - id_str, - /*state_db_ctx*/ None, - ) - .await -} - -/// Locate a recorded thread rollout file by its UUID string using an already-open SQLite handle -/// when available. Returns `Ok(Some(path))` if found, `Ok(None)` if not present or the id is -/// invalid. -pub async fn find_thread_path_by_id_str_with_state_db( - codex_home: &Path, - id_str: &str, state_db_ctx: Option, ) -> io::Result> { find_thread_path_by_id_str_in_subdir(codex_home, SESSIONS_SUBDIR, id_str, state_db_ctx).await } -/// Locate an archived thread rollout file by its UUID string. +/// Locate an archived thread rollout file by its UUID string. Uses `state_db_ctx` when available +/// to avoid reopening SQLite. pub async fn find_archived_thread_path_by_id_str( codex_home: &Path, id_str: &str, -) -> io::Result> { - find_thread_path_by_id_str_in_subdir( - codex_home, - ARCHIVED_SESSIONS_SUBDIR, - id_str, - /*state_db_ctx*/ None, - ) - .await -} - -/// Locate an archived thread rollout file by its UUID string using an already-open SQLite handle -/// when available. -pub async fn find_archived_thread_path_by_id_str_with_state_db( - codex_home: &Path, - id_str: &str, state_db_ctx: Option, ) -> io::Result> { find_thread_path_by_id_str_in_subdir(codex_home, ARCHIVED_SESSIONS_SUBDIR, id_str, state_db_ctx) diff --git a/codex-rs/rollout/src/recorder.rs b/codex-rs/rollout/src/recorder.rs index f39c38af60..0c281c337b 100644 --- a/codex-rs/rollout/src/recorder.rs +++ b/codex-rs/rollout/src/recorder.rs @@ -287,6 +287,12 @@ impl RolloutRecorder { ) .await { + if db_page.items.is_empty() && !fs_page.items.is_empty() { + tracing::warn!( + "state db discrepancy during list_threads_with_db_fallback: empty_db_page" + ); + return Ok(truncate_fs_page(fs_page, page_size, sort_key)); + } return Ok(db_page.into()); } // If SQLite listing still fails, return the filesystem page rather than failing the list. diff --git a/codex-rs/rollout/src/session_index.rs b/codex-rs/rollout/src/session_index.rs index 8c88dd39ad..cde33cf811 100644 --- a/codex-rs/rollout/src/session_index.rs +++ b/codex-rs/rollout/src/session_index.rs @@ -139,7 +139,12 @@ pub async fn find_thread_path_by_name_str( let Some(thread_id) = find_thread_id_by_name(codex_home, name).await? else { return Ok(None); }; - super::list::find_thread_path_by_id_str(codex_home, &thread_id.to_string()).await + super::list::find_thread_path_by_id_str( + codex_home, + &thread_id.to_string(), + /*state_db_ctx*/ None, + ) + .await } fn session_index_path(codex_home: &Path) -> PathBuf { diff --git a/codex-rs/rollout/src/state_db.rs b/codex-rs/rollout/src/state_db.rs index 6367a27ca9..6b1957775b 100644 --- a/codex-rs/rollout/src/state_db.rs +++ b/codex-rs/rollout/src/state_db.rs @@ -25,6 +25,9 @@ use uuid::Uuid; pub type StateDbHandle = Arc; /// Initialize the state runtime for thread state persistence and backfill checks. +/// +/// Callers that keep a shared runtime alive for repeated state reads should use +/// this helper instead of reopening SQLite on each lookup. pub async fn init(config: &impl RolloutConfigView) -> Option { let config = RolloutConfig::from_view(config); let runtime = match codex_state::StateRuntime::init( diff --git a/codex-rs/rollout/src/tests.rs b/codex-rs/rollout/src/tests.rs index 6a5f25286d..d0fd23643e 100644 --- a/codex-rs/rollout/src/tests.rs +++ b/codex-rs/rollout/src/tests.rs @@ -231,7 +231,7 @@ async fn find_thread_path_falls_back_when_db_path_is_stale() { )); insert_state_db_thread(home, thread_id, stale_db_path.as_path(), false).await; - let found = find_thread_path_by_id_str(home, &uuid.to_string()) + let found = find_thread_path_by_id_str(home, &uuid.to_string(), /*state_db_ctx*/ None) .await .expect("lookup should succeed"); assert_eq!(found, Some(fs_rollout_path.clone())); @@ -257,7 +257,7 @@ async fn find_thread_path_repairs_missing_db_row_after_filesystem_fallback() { .await .expect("backfill should be complete"); - let found = find_thread_path_by_id_str(home, &uuid.to_string()) + let found = find_thread_path_by_id_str(home, &uuid.to_string(), /*state_db_ctx*/ None) .await .expect("lookup should succeed"); assert_eq!(found, Some(fs_rollout_path.clone())); diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index 039148a718..0e0b55880f 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -721,7 +721,8 @@ async fn run_ratatui_app( if let Some(id_str) = cli.fork_session_id.as_deref() { let is_uuid = Uuid::parse_str(id_str).is_ok(); let path = if is_uuid { - find_thread_path_by_id_str(&config.codex_home, id_str).await? + find_thread_path_by_id_str(&config.codex_home, id_str, /*state_db_ctx*/ None) + .await? } else { find_thread_path_by_name_str(&config.codex_home, id_str).await? }; @@ -814,7 +815,7 @@ async fn run_ratatui_app( } else if let Some(id_str) = cli.resume_session_id.as_deref() { let is_uuid = Uuid::parse_str(id_str).is_ok(); let path = if is_uuid { - find_thread_path_by_id_str(&config.codex_home, id_str).await? + find_thread_path_by_id_str(&config.codex_home, id_str, /*state_db_ctx*/ None).await? } else { find_thread_path_by_name_str(&config.codex_home, id_str).await? };