From 7e310bc7f3c94af86c1d06d6e387054b3685ebaf Mon Sep 17 00:00:00 2001 From: Rasmus Rygaard Date: Tue, 5 May 2026 14:45:29 -0700 Subject: [PATCH] Inject state DB, agent graph store (#20689) ## Why We want the agent graph store to be passed down the stack as a real dependency, the same way we already treat the thread store. This will let us inject the agent graph store as a real dependency and support implementations other than the local SQLite-backed one. Right now most code instantiates a state DB and an agent graph store just-in-time. Ideally, we would not depend on the state DB directly but only read through the higher-level interfaces. This change makes the dependency boundaries explicit and moves state DB initialization to process bootstrap instead of hiding it inside local store implementations. ## What changed - `ThreadManager` now requires a `StateDbHandle` and an `AgentGraphStore` at construction time instead of treating them as optional internals. - The local store constructors no longer lazily initialize SQLite. Callers now initialize the state DB once per process and use that shared handle to build: - `LocalThreadStore` - `LocalAgentGraphStore` - App bootstraps (`app-server`, `mcp-server`, `prompt_debug`, and the thread-manager sample) now initialize the state DB up front and inject the resulting handle down the stack. - `app-server` now consistently uses its process-scoped state DB handle instead of reopening SQLite or trying to recover it from loaded threads. - Device-key storage now reuses the shared state DB handle instead of maintaining its own lazy opener. - The thread archive / descendant traversal paths now use the injected `AgentGraphStore` instead of reaching through local thread-store-specific state. ## Verification - `cargo check -p codex-core -p codex-thread-store -p codex-app-server -p codex-mcp-server -p codex-thread-manager-sample --tests` - `cargo test -p codex-thread-store` - `cargo test -p codex-core thread_manager_accepts_separate_agent_graph_store_and_thread_store -- --nocapture` - `cargo test -p codex-app-server thread_archive_archives_spawned_descendants -- --nocapture` --- codex-rs/Cargo.lock | 1 + codex-rs/app-server-client/src/lib.rs | 2 +- .../app-server/src/bespoke_event_handling.rs | 6 +- codex-rs/app-server/src/in_process.rs | 25 ++- codex-rs/app-server/src/lib.rs | 38 ++--- codex-rs/app-server/src/message_processor.rs | 11 +- .../src/message_processor_tracing_tests.rs | 6 +- .../device_key_processor.rs | 22 +-- .../thread_goal_processor.rs | 31 +--- .../request_processors/thread_processor.rs | 8 +- .../tests/suite/v2/remote_thread_store.rs | 10 +- codex-rs/cli/src/main.rs | 2 +- codex-rs/core-api/src/lib.rs | 2 + codex-rs/core/Cargo.toml | 1 + codex-rs/core/src/agent/control.rs | 48 +++--- codex-rs/core/src/agent/control_tests.rs | 41 ++--- codex-rs/core/src/codex_delegate.rs | 1 + codex-rs/core/src/goals.rs | 12 -- codex-rs/core/src/lib.rs | 2 + codex-rs/core/src/personality_migration.rs | 17 +- .../core/src/personality_migration_tests.rs | 64 +++++++- codex-rs/core/src/prompt_debug.rs | 11 +- codex-rs/core/src/session/mod.rs | 17 +- codex-rs/core/src/session/session.rs | 26 +--- codex-rs/core/src/session/tests.rs | 146 ++++++++---------- .../core/src/session/tests/guardian_tests.rs | 8 +- codex-rs/core/src/session/turn_context.rs | 1 + codex-rs/core/src/stream_events_utils.rs | 9 +- codex-rs/core/src/test_support.rs | 23 +-- codex-rs/core/src/thread_manager.rs | 110 +++++++++---- codex-rs/core/src/thread_manager_tests.rs | 82 +++++++--- .../src/tools/handlers/multi_agents_tests.rs | 107 +++++++------ codex-rs/core/tests/common/test_codex.rs | 17 +- codex-rs/core/tests/suite/client.rs | 19 ++- .../core/tests/suite/personality_migration.rs | 54 ++++--- .../core/tests/suite/prompt_debug_tests.rs | 1 - codex-rs/core/tests/suite/resume_warning.rs | 3 +- codex-rs/core/tests/suite/skills.rs | 12 +- .../tests/suite/unstable_features_warning.rs | 6 +- codex-rs/mcp-server/src/lib.rs | 8 +- codex-rs/mcp-server/src/message_processor.rs | 18 ++- codex-rs/thread-manager-sample/src/main.rs | 11 +- .../thread-store/src/local/archive_thread.rs | 30 ++-- .../thread-store/src/local/create_thread.rs | 4 +- .../thread-store/src/local/list_threads.rs | 68 ++++---- .../thread-store/src/local/live_writer.rs | 4 +- codex-rs/thread-store/src/local/mod.rs | 42 ++--- .../thread-store/src/local/read_thread.rs | 110 ++++--------- .../thread-store/src/local/test_support.rs | 18 ++- .../src/local/unarchive_thread.rs | 30 ++-- .../src/local/update_thread_metadata.rs | 103 ++++-------- codex-rs/tui/src/lib.rs | 62 ++++---- codex-rs/tui/src/resume_picker.rs | 3 + 53 files changed, 768 insertions(+), 745 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 1e1590e6ce..7819fa223a 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2411,6 +2411,7 @@ dependencies = [ "bm25", "chrono", "clap", + "codex-agent-graph-store", "codex-analytics", "codex-api", "codex-app-server-protocol", diff --git a/codex-rs/app-server-client/src/lib.rs b/codex-rs/app-server-client/src/lib.rs index 79141da24e..0911cc448d 100644 --- a/codex-rs/app-server-client/src/lib.rs +++ b/codex-rs/app-server-client/src/lib.rs @@ -29,7 +29,6 @@ pub use codex_app_server::in_process::DEFAULT_IN_PROCESS_CHANNEL_CAPACITY; pub use codex_app_server::in_process::InProcessServerEvent; use codex_app_server::in_process::InProcessStartArgs; use codex_app_server::in_process::LogDbLayer; -pub use codex_app_server::in_process::StateDbHandle; use codex_app_server_protocol::ClientInfo; use codex_app_server_protocol::ClientNotification; use codex_app_server_protocol::ClientRequest; @@ -47,6 +46,7 @@ use codex_config::LoaderOverrides; use codex_config::NoopThreadConfigLoader; use codex_config::RemoteThreadConfigLoader; use codex_config::ThreadConfigLoader; +pub use codex_core::StateDbHandle; use codex_core::config::Config; pub use codex_exec_server::EnvironmentManager; pub use codex_exec_server::EnvironmentManagerArgs; diff --git a/codex-rs/app-server/src/bespoke_event_handling.rs b/codex-rs/app-server/src/bespoke_event_handling.rs index 1907f1c3f4..55b3fea814 100644 --- a/codex-rs/app-server/src/bespoke_event_handling.rs +++ b/codex-rs/app-server/src/bespoke_event_handling.rs @@ -2629,7 +2629,8 @@ mod tests { config.model_provider.clone(), config.codex_home.to_path_buf(), Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - ), + ) + .await, ); let codex_core::NewThread { thread_id: conversation_id, @@ -3214,7 +3215,8 @@ mod tests { config.model_provider.clone(), config.codex_home.to_path_buf(), Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - ), + ) + .await, ); let codex_core::NewThread { thread_id: conversation_id, diff --git a/codex-rs/app-server/src/in_process.rs b/codex-rs/app-server/src/in_process.rs index a8c023e95a..73cfe75a54 100644 --- a/codex-rs/app-server/src/in_process.rs +++ b/codex-rs/app-server/src/in_process.rs @@ -82,11 +82,12 @@ use codex_config::CloudRequirementsLoader; use codex_config::LoaderOverrides; use codex_config::ThreadConfigLoader; use codex_core::config::Config; +use codex_core::init_state_db_from_config; use codex_exec_server::EnvironmentManager; use codex_feedback::CodexFeedback; use codex_login::AuthManager; use codex_protocol::protocol::SessionSource; -pub use codex_rollout::StateDbHandle; +use codex_rollout::state_db::StateDbHandle; pub use codex_state::log_db::LogDbLayer; use tokio::sync::mpsc; use tokio::sync::oneshot; @@ -127,7 +128,7 @@ pub struct InProcessStartArgs { pub feedback: CodexFeedback, /// SQLite tracing layer used to flush recently emitted logs before feedback upload. pub log_db: Option, - /// Process-wide SQLite state handle shared with embedded app-server consumers. + /// Optional state DB handle to use for the in-process runtime. pub state_db: Option, /// Environment manager used by core execution and filesystem operations. pub environment_manager: Arc, @@ -344,7 +345,7 @@ 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 client = start_uninitialized(args).await; let initialize_response = client .request(ClientRequest::Initialize { @@ -364,8 +365,12 @@ pub async fn start(args: InProcessStartArgs) -> IoResult Ok(client) } -fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle { +async fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle { let channel_capacity = args.channel_capacity.max(1); + let state_db = match args.state_db.clone() { + Some(state_db) => Some(state_db), + None => init_state_db_from_config(args.config.as_ref()).await, + }; let (client_tx, mut client_rx) = mpsc::channel::(channel_capacity); let (event_tx, event_rx) = mpsc::channel::(channel_capacity); @@ -414,6 +419,12 @@ fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle { ); let (processor_tx, mut processor_rx) = mpsc::channel::(channel_capacity); let mut processor_handle = tokio::spawn(async move { + let Some(state_db) = state_db else { + warn!( + "in-process app-server state db initialization failed; shutting down processor task" + ); + return; + }; let processor = Arc::new(MessageProcessor::new(MessageProcessorArgs { outgoing: Arc::clone(&processor_outgoing), analytics_events_client, @@ -423,7 +434,7 @@ fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle { environment_manager: args.environment_manager, feedback: args.feedback, log_db: args.log_db, - state_db: args.state_db, + state_db, config_warnings: args.config_warnings, session_source: args.session_source, auth_manager, @@ -761,7 +772,7 @@ mod tests { ) -> InProcessClientHandle { let codex_home = TempDir::new().expect("temp dir"); let config = Arc::new(build_test_config(codex_home.path()).await); - let state_db = codex_rollout::state_db::try_init(config.as_ref()) + let state_db = init_state_db_from_config(config.as_ref()) .await .expect("state db should initialize for in-process test"); let args = InProcessStartArgs { @@ -819,7 +830,7 @@ mod tests { } #[tokio::test] - async fn in_process_allows_device_key_requests_to_reach_device_key_processor() { + async fn in_process_allows_device_key_requests_to_reach_device_key_api() { let client = start_test_client(SessionSource::Cli).await; const MALFORMED_KEY_ID_MESSAGE: &str = concat!( "invalid device key payload: keyId must be dk_hse_, dk_tpm_, or dk_osn_ ", diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index 4b1723d88c..4013bbe76b 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -50,11 +50,11 @@ use codex_config::TextRange as CoreTextRange; use codex_core::ExecPolicyError; use codex_core::check_execpolicy_for_warnings; use codex_core::config::find_codex_home; +use codex_core::init_state_db_from_config; use codex_exec_server::EnvironmentManager; use codex_exec_server::ExecServerRuntimePaths; use codex_feedback::CodexFeedback; use codex_protocol::protocol::SessionSource; -use codex_rollout::state_db as rollout_state_db; use codex_state::log_db; use tokio::sync::mpsc; use tokio::sync::oneshot; @@ -487,9 +487,9 @@ pub async fn run_main_with_transport_options( } }; - let state_db_result = rollout_state_db::try_init(&config).await; - let state_db_init_error = state_db_result.as_ref().err().map(ToString::to_string); - let state_db = state_db_result.ok(); + let state_db = init_state_db_from_config(&config) + .await + .ok_or_else(|| std::io::Error::other("failed to initialize sqlite state db"))?; if should_run_personality_migration { let effective_toml = config.config_layer_stack.effective_config(); @@ -598,10 +598,12 @@ pub async fn run_main_with_transport_options( let feedback_layer = feedback.logger_layer(); let feedback_metadata_layer = feedback.metadata_layer(); - let log_db = state_db.clone().map(log_db::start); - let log_db_layer = log_db - .clone() - .map(|layer| layer.with_filter(Targets::new().with_default(Level::TRACE))); + let log_db = log_db::start(state_db.clone()); + let log_db_layer = Some( + log_db + .clone() + .with_filter(Targets::new().with_default(Level::TRACE)), + ); let otel_logger_layer = otel.as_ref().and_then(|o| o.logger_layer()); let otel_tracing_layer = otel.as_ref().and_then(|o| o.tracing_layer()); let _ = tracing_subscriber::registry() @@ -618,10 +620,6 @@ pub async fn run_main_with_transport_options( None => error!("{}", warning.summary), } } - if let Some(err) = &state_db_init_error { - error!("failed to initialize sqlite state db: {err}"); - } - let transport_shutdown_token = CancellationToken::new(); let mut transport_accept_handles = Vec::>::new(); @@ -666,25 +664,17 @@ pub async fn run_main_with_transport_options( let auth_manager = AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false).await; - let remote_control_config_enabled = config.features.enabled(Feature::RemoteControl); - let remote_control_enabled = remote_control_config_enabled && state_db.is_some(); - if remote_control_config_enabled && state_db.is_none() { - error!("remote control disabled because sqlite state db is unavailable"); - } + let remote_control_enabled = config.features.enabled(Feature::RemoteControl); if transport_accept_handles.is_empty() && !remote_control_enabled { return Err(std::io::Error::new( ErrorKind::InvalidInput, - if remote_control_config_enabled && state_db.is_none() { - "no transport configured; remote control disabled because sqlite state db is unavailable" - } else { - "no transport configured; use --listen or enable remote control" - }, + "no transport configured; use --listen or enable remote control", )); } let (remote_control_accept_handle, remote_control_handle) = start_remote_control( config.chatgpt_base_url.clone(), - state_db.clone(), + Some(state_db.clone()), auth_manager.clone(), transport_event_tx.clone(), transport_shutdown_token.clone(), @@ -768,7 +758,7 @@ pub async fn run_main_with_transport_options( config_manager, environment_manager, feedback: feedback.clone(), - log_db, + log_db: Some(log_db), state_db: state_db.clone(), config_warnings, session_source, diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 8d26cec70f..4c1d16eeac 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -61,6 +61,7 @@ use codex_app_server_protocol::experimental_required_message; use codex_arg0::Arg0DispatchPaths; use codex_chatgpt::workspace_settings; use codex_core::ThreadManager; +use codex_core::agent_graph_store_from_state_db; use codex_core::config::Config; use codex_core::thread_store_from_config; use codex_exec_server::EnvironmentManager; @@ -254,7 +255,7 @@ pub(crate) struct MessageProcessorArgs { pub(crate) environment_manager: Arc, pub(crate) feedback: CodexFeedback, pub(crate) log_db: Option, - pub(crate) state_db: Option, + pub(crate) state_db: StateDbHandle, pub(crate) config_warnings: Vec, pub(crate) session_source: SessionSource, pub(crate) auth_manager: Arc, @@ -291,14 +292,16 @@ impl MessageProcessor { // affect per-thread behavior, but they must not move newly started, // resumed, or forked threads to a different persistence backend/root. let thread_store = thread_store_from_config(config.as_ref(), state_db.clone()); + let agent_graph_store = agent_graph_store_from_state_db(state_db.clone()); let thread_manager = Arc::new(ThreadManager::new( config.as_ref(), auth_manager.clone(), session_source, environment_manager, Some(analytics_events_client.clone()), - Arc::clone(&thread_store), state_db.clone(), + Arc::clone(&thread_store), + agent_graph_store.clone(), )); thread_manager .plugins_manager() @@ -344,7 +347,7 @@ impl MessageProcessor { Arc::clone(&config), feedback, log_db, - state_db.clone(), + Some(state_db.clone()), ); let git_processor = GitRequestProcessor::new(); let initialize_processor = InitializeRequestProcessor::new( @@ -395,7 +398,7 @@ impl MessageProcessor { thread_watch_manager.clone(), Arc::clone(&thread_list_state_permit), thread_goal_processor.clone(), - state_db.clone(), + Some(state_db.clone()), ); let turn_processor = TurnRequestProcessor::new( auth_manager.clone(), 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 45ea709180..67b668bafc 100644 --- a/codex-rs/app-server/src/message_processor_tracing_tests.rs +++ b/codex-rs/app-server/src/message_processor_tracing_tests.rs @@ -32,6 +32,7 @@ use codex_config::CloudRequirementsLoader; use codex_config::LoaderOverrides; use codex_core::config::Config; use codex_core::config::ConfigBuilder; +use codex_core::init_state_db_from_config; use codex_exec_server::EnvironmentManager; use codex_feedback::CodexFeedback; use codex_login::AuthManager; @@ -281,6 +282,9 @@ async fn build_test_processor( outgoing_tx, analytics_events_client.clone(), )); + let state_db = init_state_db_from_config(config.as_ref()) + .await + .expect("tracing test processor requires state db"); let processor = Arc::new(MessageProcessor::new(MessageProcessorArgs { outgoing, analytics_events_client, @@ -290,7 +294,7 @@ async fn build_test_processor( environment_manager: Arc::new(EnvironmentManager::default_for_tests()), feedback: CodexFeedback::new(), log_db: None, - state_db: None, + state_db, config_warnings: Vec::new(), session_source: SessionSource::VSCode, auth_manager, diff --git a/codex-rs/app-server/src/request_processors/device_key_processor.rs b/codex-rs/app-server/src/request_processors/device_key_processor.rs index ea0a96c2af..e1da459a2e 100644 --- a/codex-rs/app-server/src/request_processors/device_key_processor.rs +++ b/codex-rs/app-server/src/request_processors/device_key_processor.rs @@ -33,8 +33,8 @@ use codex_device_key::RemoteControlClientConnectionAudience; use codex_device_key::RemoteControlClientConnectionSignPayload; use codex_device_key::RemoteControlClientEnrollmentAudience; use codex_device_key::RemoteControlClientEnrollmentSignPayload; +use codex_rollout::state_db::StateDbHandle; use codex_state::DeviceKeyBindingRecord; -use codex_state::StateRuntime; #[derive(Clone)] pub(crate) struct DeviceKeyRequestProcessor { @@ -43,10 +43,7 @@ pub(crate) struct DeviceKeyRequestProcessor { } impl DeviceKeyRequestProcessor { - pub(crate) fn new( - outgoing: Arc, - state_db: Option>, - ) -> Self { + pub(crate) fn new(outgoing: Arc, state_db: StateDbHandle) -> Self { Self { outgoing, store: DeviceKeyStore::new(Arc::new(StateDeviceKeyBindingStore::new(state_db))), @@ -170,25 +167,18 @@ async fn sign_device_key( } struct StateDeviceKeyBindingStore { - state_db: Option>, + state_db: StateDbHandle, } impl StateDeviceKeyBindingStore { - fn new(state_db: Option>) -> Self { + fn new(state_db: StateDbHandle) -> Self { Self { state_db } } - - async fn state_db(&self) -> Result, DeviceKeyError> { - self.state_db - .clone() - .ok_or_else(|| DeviceKeyError::Platform("sqlite state db unavailable".to_string())) - } } impl fmt::Debug for StateDeviceKeyBindingStore { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("StateDeviceKeyBindingStore") - .field("has_state_db", &self.state_db.is_some()) .finish_non_exhaustive() } } @@ -196,7 +186,7 @@ impl fmt::Debug for StateDeviceKeyBindingStore { #[async_trait] impl DeviceKeyBindingStore for StateDeviceKeyBindingStore { async fn get_binding(&self, key_id: &str) -> Result, DeviceKeyError> { - let state_db = self.state_db().await?; + let state_db = self.state_db.clone(); state_db .get_device_key_binding(key_id) .await @@ -214,7 +204,7 @@ impl DeviceKeyBindingStore for StateDeviceKeyBindingStore { key_id: &str, binding: &DeviceKeyBinding, ) -> Result<(), DeviceKeyError> { - let state_db = self.state_db().await?; + let state_db = self.state_db.clone(); state_db .upsert_device_key_binding(&DeviceKeyBindingRecord { key_id: key_id.to_string(), diff --git a/codex-rs/app-server/src/request_processors/thread_goal_processor.rs b/codex-rs/app-server/src/request_processors/thread_goal_processor.rs index 0e12e44ce5..ff192e34b2 100644 --- a/codex-rs/app-server/src/request_processors/thread_goal_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_goal_processor.rs @@ -7,7 +7,7 @@ pub(crate) struct ThreadGoalRequestProcessor { outgoing: Arc, config: Arc, thread_state_manager: ThreadStateManager, - state_db: Option, + state_db: StateDbHandle, } impl ThreadGoalRequestProcessor { @@ -16,7 +16,7 @@ impl ThreadGoalRequestProcessor { outgoing: Arc, config: Arc, thread_state_manager: ThreadStateManager, - state_db: Option, + state_db: StateDbHandle, ) -> Self { Self { thread_manager, @@ -72,23 +72,6 @@ impl ThreadGoalRequestProcessor { } } - pub(crate) async fn pending_resume_goal_state( - &self, - thread: &CodexThread, - ) -> (bool, Option) { - let emit_thread_goal_update = self.config.features.enabled(Feature::Goals); - let thread_goal_state_db = if emit_thread_goal_update { - if let Some(state_db) = thread.state_db() { - Some(state_db) - } else { - self.state_db.clone() - } - } else { - None - }; - (emit_thread_goal_update, thread_goal_state_db) - } - async fn thread_goal_set_inner( &self, request_id: ConnectionRequestId, @@ -110,7 +93,7 @@ impl ThreadGoalRequestProcessor { None => find_thread_path_by_id_str( &self.config.codex_home, &thread_id.to_string(), - self.state_db.as_deref(), + Some(self.state_db.as_ref()), ) .await .map_err(|err| { @@ -275,7 +258,7 @@ impl ThreadGoalRequestProcessor { None => find_thread_path_by_id_str( &self.config.codex_home, &thread_id.to_string(), - self.state_db.as_deref(), + Some(self.state_db.as_ref()), ) .await .map_err(|err| { @@ -339,7 +322,7 @@ impl ThreadGoalRequestProcessor { find_thread_path_by_id_str( &self.config.codex_home, &thread_id.to_string(), - self.state_db.as_deref(), + Some(self.state_db.as_ref()), ) .await .map_err(|err| { @@ -348,9 +331,7 @@ impl ThreadGoalRequestProcessor { .ok_or_else(|| invalid_request(format!("thread not found: {thread_id}")))?; } - self.state_db - .clone() - .ok_or_else(|| internal_error("sqlite state db unavailable for thread goals")) + Ok(self.state_db.clone()) } async fn emit_thread_goal_snapshot(&self, thread_id: ThreadId) { diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index 1921fb69b1..fd61182788 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -2611,10 +2611,10 @@ impl ThreadRequestProcessor { ))); }; - let (emit_thread_goal_update, thread_goal_state_db) = self - .thread_goal_processor - .pending_resume_goal_state(existing_thread.as_ref()) - .await; + let emit_thread_goal_update = self.config.features.enabled(Feature::Goals); + let thread_goal_state_db = emit_thread_goal_update + .then(|| self.state_db.clone()) + .flatten(); let command = crate::thread_state::ThreadListenerCommand::SendThreadResumeResponse( Box::new(crate::thread_state::PendingThreadResumeRequest { diff --git a/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs b/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs index b04eb12a45..9363217864 100644 --- a/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs +++ b/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs @@ -42,6 +42,7 @@ use codex_core::config::ConfigBuilder; use codex_exec_server::EnvironmentManager; use codex_feedback::CodexFeedback; use codex_protocol::protocol::SessionSource; +use codex_state::StateRuntime; use codex_thread_store::InMemoryThreadStore; use pretty_assertions::assert_eq; use tempfile::TempDir; @@ -67,6 +68,13 @@ async fn thread_start_with_non_local_thread_store_does_not_create_local_persiste .loader_overrides(loader_overrides.clone()) .build() .await?; + let sqlite_home = TempDir::new()?; + let state_db = StateRuntime::init( + sqlite_home.path().to_path_buf(), + config.model_provider_id.clone(), + ) + .await + .expect("remote thread store regression test should initialize state db"); let thread_store = InMemoryThreadStore::for_id(store_id.clone()); let _in_memory_store = InMemoryThreadStoreId { store_id }; @@ -80,7 +88,7 @@ async fn thread_start_with_non_local_thread_store_does_not_create_local_persiste thread_config_loader: Arc::new(NoopThreadConfigLoader), feedback: CodexFeedback::new(), log_db: None, - state_db: None, + state_db: Some(state_db), environment_manager: Arc::new(EnvironmentManager::default_for_tests()), config_warnings: Vec::new(), session_source: SessionSource::Cli, diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index 48b6886382..5b5ad7a2a3 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -1387,7 +1387,7 @@ async fn run_debug_prompt_input_command( }); } - let prompt_input = codex_core::build_prompt_input(config, input, /*state_db*/ None).await?; + let prompt_input = codex_core::build_prompt_input(config, input).await?; println!("{}", serde_json::to_string_pretty(&prompt_input)?); Ok(()) diff --git a/codex-rs/core-api/src/lib.rs b/codex-rs/core-api/src/lib.rs index e884cfae9b..13c87df9e0 100644 --- a/codex-rs/core-api/src/lib.rs +++ b/codex-rs/core-api/src/lib.rs @@ -31,6 +31,7 @@ pub use codex_core::StartThreadOptions; pub use codex_core::StateDbHandle; pub use codex_core::ThreadManager; pub use codex_core::ThreadShutdownReport; +pub use codex_core::agent_graph_store_from_state_db; pub use codex_core::config::Config; pub use codex_core::config::Constrained; pub use codex_core::config::GhostSnapshotConfig; @@ -40,6 +41,7 @@ pub use codex_core::config::TerminalResizeReflowConfig; pub use codex_core::config::ThreadStoreConfig; pub use codex_core::config::find_codex_home; pub use codex_core::init_state_db; +pub use codex_core::init_state_db_from_config; pub use codex_core::skills::SkillsManager; pub use codex_core::thread_store_from_config; pub use codex_exec_server::EnvironmentManager; diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index 44c6aacac5..49ed304623 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -26,6 +26,7 @@ bm25 = { workspace = true } chrono = { workspace = true, features = ["serde"] } clap = { workspace = true, features = ["derive"] } codex-analytics = { workspace = true } +codex-agent-graph-store = { workspace = true } codex-api = { workspace = true } codex-app-server-protocol = { workspace = true } codex-apply-patch = { workspace = true } diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 1a13c0b00c..3917ad5e95 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -29,7 +29,6 @@ use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::user_input::UserInput; -use codex_state::DirectionalThreadSpawnEdgeStatus; use codex_thread_store::ReadThreadParams; use serde::Serialize; use std::collections::HashMap; @@ -296,7 +295,6 @@ impl AgentControl { state.notify_thread_created(new_thread.thread_id); self.persist_thread_spawn_edge_for_source( - new_thread.thread.as_ref(), new_thread.thread_id, notification_source.as_ref(), ) @@ -449,19 +447,14 @@ impl AgentControl { )) .await?; let state = self.upgrade()?; - let Ok(resumed_thread) = state.get_thread(resumed_thread_id).await else { - return Ok(resumed_thread_id); - }; - let Some(state_db_ctx) = resumed_thread.state_db() else { - return Ok(resumed_thread_id); - }; + let agent_graph_store = state.agent_graph_store(); let mut resume_queue = VecDeque::from([(thread_id, root_depth)]); while let Some((parent_thread_id, parent_depth)) = resume_queue.pop_front() { - let child_ids = match state_db_ctx - .list_thread_spawn_children_with_status( + let child_ids = match agent_graph_store + .list_thread_spawn_children( parent_thread_id, - DirectionalThreadSpawnEdgeStatus::Open, + Some(codex_agent_graph_store::ThreadSpawnEdgeStatus::Open), ) .await { @@ -525,7 +518,6 @@ impl AgentControl { let _ = config.features.disable(Feature::Collab); } let state = self.upgrade()?; - let state_db_ctx = state.state_db(); let mut reservation = self.state.reserve_spawn_slot(config.agent_max_threads)?; let (session_source, agent_metadata) = match session_source { SessionSource::SubAgent(SubAgentSource::ThreadSpawn { @@ -535,14 +527,11 @@ impl AgentControl { agent_role: _, agent_nickname: _, }) => { + let state_db_ctx = state.state_db(); let (resumed_agent_nickname, resumed_agent_role) = - 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), - } - } else { - (None, None) + match state_db_ctx.get_thread(thread_id).await { + Ok(Some(metadata)) => (metadata.agent_nickname, metadata.agent_role), + Ok(None) | Err(_) => (None, None), }; self.prepare_thread_spawn( &mut reservation, @@ -609,7 +598,6 @@ impl AgentControl { ); } self.persist_thread_spawn_edge_for_source( - resumed_thread.thread.as_ref(), resumed_thread.thread_id, Some(¬ification_source), ) @@ -722,11 +710,13 @@ impl AgentControl { /// agent and any live descendants reached from the in-memory tree. pub(crate) async fn close_agent(&self, agent_id: ThreadId) -> CodexResult { let state = self.upgrade()?; - if let Ok(thread) = state.get_thread(agent_id).await - && let Some(state_db_ctx) = thread.state_db() - && let Err(err) = state_db_ctx - .set_thread_spawn_edge_status(agent_id, DirectionalThreadSpawnEdgeStatus::Closed) - .await + if let Err(err) = state + .agent_graph_store() + .set_thread_spawn_edge_status( + agent_id, + codex_agent_graph_store::ThreadSpawnEdgeStatus::Closed, + ) + .await { warn!("failed to persist thread-spawn edge status for {agent_id}: {err}"); } @@ -1142,21 +1132,21 @@ impl AgentControl { async fn persist_thread_spawn_edge_for_source( &self, - thread: &crate::CodexThread, child_thread_id: ThreadId, session_source: Option<&SessionSource>, ) { let Some(parent_thread_id) = session_source.and_then(thread_spawn_parent_thread_id) else { return; }; - let Some(state_db_ctx) = thread.state_db() else { + let Ok(state) = self.upgrade() else { return; }; - if let Err(err) = state_db_ctx + if let Err(err) = state + .agent_graph_store() .upsert_thread_spawn_edge( parent_thread_id, child_thread_id, - DirectionalThreadSpawnEdgeStatus::Open, + codex_agent_graph_store::ThreadSpawnEdgeStatus::Open, ) .await { diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index b95aad4489..a8de427814 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -1,6 +1,5 @@ use super::*; use crate::CodexThread; -use crate::StateDbHandle; use crate::ThreadManager; use crate::agent::agent_status_from_event; use crate::config::AgentRoleConfig; @@ -8,7 +7,6 @@ use crate::config::Config; use crate::config::ConfigBuilder; use crate::context::ContextualUserFragment; use crate::context::SubagentNotification; -use crate::init_state_db; use assert_matches::assert_matches; use codex_features::Feature; use codex_login::CodexAuth; @@ -86,7 +84,6 @@ fn spawn_agent_call(call_id: &str) -> ResponseItem { struct AgentControlHarness { _home: TempDir, config: Config, - state_db: Option, manager: ThreadManager, control: AgentControl, } @@ -94,19 +91,17 @@ struct AgentControlHarness { impl AgentControlHarness { async fn new() -> Self { let (home, config) = test_config().await; - let state_db = init_state_db(&config).await; - let manager = ThreadManager::with_models_provider_home_and_state_for_tests( + 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(), std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - state_db.clone(), - ); + ) + .await; let control = manager.agent_control(); Self { _home: home, config, - state_db, manager, control, } @@ -955,7 +950,8 @@ async fn spawn_agent_respects_max_threads_limit() { config.model_provider.clone(), config.codex_home.to_path_buf(), std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - ); + ) + .await; let control = manager.agent_control(); let _ = manager @@ -1007,7 +1003,8 @@ async fn spawn_agent_releases_slot_after_shutdown() { config.model_provider.clone(), config.codex_home.to_path_buf(), std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - ); + ) + .await; let control = manager.agent_control(); let first_agent_id = control @@ -1050,7 +1047,8 @@ async fn spawn_agent_limit_shared_across_clones() { config.model_provider.clone(), config.codex_home.to_path_buf(), std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - ); + ) + .await; let control = manager.agent_control(); let cloned = control.clone(); @@ -1095,7 +1093,8 @@ async fn resume_agent_respects_max_threads_limit() { config.model_provider.clone(), config.codex_home.to_path_buf(), std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - ); + ) + .await; let control = manager.agent_control(); let resumable_id = control @@ -1151,7 +1150,8 @@ async fn resume_agent_releases_slot_after_resume_failure() { config.model_provider.clone(), config.codex_home.to_path_buf(), std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - ); + ) + .await; let control = manager.agent_control(); let _ = control @@ -1543,19 +1543,17 @@ async fn resume_thread_subagent_restores_stored_nickname_and_role() { .features .enable(Feature::Sqlite) .expect("test config should allow sqlite"); - let state_db = init_state_db(&config).await; - let manager = ThreadManager::with_models_provider_home_and_state_for_tests( + 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(), std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - state_db.clone(), - ); + ) + .await; let control = manager.agent_control(); let harness = AgentControlHarness { _home: home, config, - state_db, manager, control, }; @@ -1706,7 +1704,12 @@ async fn resume_agent_from_rollout_reads_archived_rollout_path() { .expect("child shutdown should succeed"); let store = LocalThreadStore::new( LocalThreadStoreConfig::from_config(&harness.config), - harness.state_db.clone(), + codex_state::StateRuntime::init( + harness.config.sqlite_home.clone(), + harness.config.model_provider_id.clone(), + ) + .await + .expect("state db should initialize"), ); store .archive_thread(ArchiveThreadParams { diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 1aeb125091..9578fa5726 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -95,6 +95,7 @@ pub(crate) async fn run_codex_thread_interactive( parent_trace: None, environment_selections: parent_ctx.environments.clone(), analytics_events_client: Some(parent_session.services.analytics_events_client.clone()), + state_db: parent_session.services.state_db.clone(), thread_store: Arc::clone(&parent_session.services.thread_store), })) .or_cancel(&cancel_token) diff --git a/codex-rs/core/src/goals.rs b/codex-rs/core/src/goals.rs index 3cea25f9b8..4af73e6a01 100644 --- a/codex-rs/core/src/goals.rs +++ b/codex-rs/core/src/goals.rs @@ -29,7 +29,6 @@ use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::validate_thread_goal_objective; use codex_rollout::state_db::reconcile_rollout; -use codex_thread_store::LocalThreadStore; use codex_utils_template::Template; use futures::future::BoxFuture; use std::sync::Arc; @@ -1338,17 +1337,6 @@ impl Session { state_db } else if let Some(state_db) = self.goal_runtime.state_db.lock().await.clone() { state_db - } else if let Some(local_store) = self - .services - .thread_store - .as_any() - .downcast_ref::() - { - local_store.state_db().await.ok_or_else(|| { - anyhow::anyhow!( - "thread goals require a local persisted thread with a state database" - ) - })? } else { anyhow::bail!("thread goals require a local persisted thread with a state database"); }; diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index d8a3fbcf0f..cbada9a26c 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -124,7 +124,9 @@ pub use thread_manager::NewThread; pub use thread_manager::StartThreadOptions; pub use thread_manager::ThreadManager; pub use thread_manager::ThreadShutdownReport; +pub use thread_manager::agent_graph_store_from_state_db; pub use thread_manager::build_models_manager; +pub use thread_manager::init_state_db_from_config; pub use thread_manager::thread_store_from_config; pub use web_search::web_search_action_detail; pub use web_search::web_search_detail; diff --git a/codex-rs/core/src/personality_migration.rs b/codex-rs/core/src/personality_migration.rs index 975aecd4af..bb1cb43587 100644 --- a/codex-rs/core/src/personality_migration.rs +++ b/codex-rs/core/src/personality_migration.rs @@ -25,7 +25,7 @@ pub enum PersonalityMigrationStatus { pub async fn maybe_migrate_personality( codex_home: &Path, config_toml: &ConfigToml, - state_db: Option, + state_db: StateDbHandle, ) -> io::Result { let marker_path = codex_home.join(PERSONALITY_MIGRATION_FILENAME); if tokio::fs::try_exists(&marker_path).await? { @@ -65,16 +65,13 @@ pub async fn maybe_migrate_personality( async fn has_recorded_sessions( codex_home: &Path, default_provider: &str, - state_db: Option, + state_db: StateDbHandle, ) -> io::Result { - let store = LocalThreadStore::new( - LocalThreadStoreConfig { - codex_home: codex_home.to_path_buf(), - sqlite_home: codex_home.to_path_buf(), - default_model_provider_id: default_provider.to_string(), - }, - state_db, - ); + let config = LocalThreadStoreConfig { + codex_home: codex_home.to_path_buf(), + default_model_provider_id: default_provider.to_string(), + }; + let store = LocalThreadStore::new(config, state_db); if has_threads(&store, /*archived*/ false).await? { return Ok(true); } diff --git a/codex-rs/core/src/personality_migration_tests.rs b/codex-rs/core/src/personality_migration_tests.rs index 28d34bc3c9..63faae418a 100644 --- a/codex-rs/core/src/personality_migration_tests.rs +++ b/codex-rs/core/src/personality_migration_tests.rs @@ -8,7 +8,10 @@ use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::UserMessageEvent; use codex_rollout::ARCHIVED_SESSIONS_SUBDIR; +use codex_rollout::RolloutConfig; use codex_rollout::SESSIONS_SUBDIR; +use codex_rollout::state_db::StateDbHandle; +use codex_state::state_db_path; use pretty_assertions::assert_eq; use tempfile::TempDir; use tokio::io::AsyncWriteExt; @@ -20,6 +23,26 @@ async fn read_config_toml(codex_home: &Path) -> io::Result { toml::from_str(&contents).map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err)) } +async fn state_db_for_test(codex_home: &Path) -> io::Result { + state_db_for_test_with_sqlite_home(codex_home, codex_home).await +} + +async fn state_db_for_test_with_sqlite_home( + codex_home: &Path, + sqlite_home: &Path, +) -> io::Result { + let config = RolloutConfig { + codex_home: codex_home.to_path_buf(), + sqlite_home: sqlite_home.to_path_buf(), + cwd: codex_home.to_path_buf(), + model_provider_id: "openai".to_string(), + generate_memories: false, + }; + codex_rollout::state_db::try_init(&config) + .await + .map_err(io::Error::other) +} + async fn write_session_with_user_event(codex_home: &Path) -> io::Result<()> { let thread_id = ThreadId::new(); let dir = codex_home @@ -87,7 +110,8 @@ async fn applies_when_sessions_exist_and_no_personality() -> io::Result<()> { write_session_with_user_event(temp.path()).await?; let config_toml = ConfigToml::default(); - let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?; + let state_db = state_db_for_test(temp.path()).await?; + let status = maybe_migrate_personality(temp.path(), &config_toml, state_db).await?; assert_eq!(status, PersonalityMigrationStatus::Applied); assert!(temp.path().join(PERSONALITY_MIGRATION_FILENAME).exists()); @@ -103,7 +127,8 @@ async fn applies_when_only_archived_sessions_exist_and_no_personality() -> io::R write_archived_session_with_user_event(temp.path()).await?; let config_toml = ConfigToml::default(); - let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?; + let state_db = state_db_for_test(temp.path()).await?; + let status = maybe_migrate_personality(temp.path(), &config_toml, state_db).await?; assert_eq!(status, PersonalityMigrationStatus::Applied); assert!(temp.path().join(PERSONALITY_MIGRATION_FILENAME).exists()); @@ -119,7 +144,8 @@ async fn skips_when_marker_exists() -> io::Result<()> { create_marker(&temp.path().join(PERSONALITY_MIGRATION_FILENAME)).await?; let config_toml = ConfigToml::default(); - let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?; + let state_db = state_db_for_test(temp.path()).await?; + let status = maybe_migrate_personality(temp.path(), &config_toml, state_db).await?; assert_eq!(status, PersonalityMigrationStatus::SkippedMarker); assert!(!temp.path().join("config.toml").exists()); @@ -136,7 +162,8 @@ async fn skips_when_personality_explicit() -> io::Result<()> { .map_err(|err| io::Error::other(format!("failed to write config: {err}")))?; let config_toml = read_config_toml(temp.path()).await?; - let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?; + let state_db = state_db_for_test(temp.path()).await?; + let status = maybe_migrate_personality(temp.path(), &config_toml, state_db).await?; assert_eq!( status, @@ -153,10 +180,37 @@ async fn skips_when_personality_explicit() -> io::Result<()> { async fn skips_when_no_sessions() -> io::Result<()> { let temp = TempDir::new()?; let config_toml = ConfigToml::default(); - let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?; + let state_db = state_db_for_test(temp.path()).await?; + let status = maybe_migrate_personality(temp.path(), &config_toml, state_db).await?; assert_eq!(status, PersonalityMigrationStatus::SkippedNoSessions); assert!(temp.path().join(PERSONALITY_MIGRATION_FILENAME).exists()); assert!(!temp.path().join("config.toml").exists()); Ok(()) } + +#[tokio::test] +async fn uses_configured_sqlite_home_when_checking_for_sessions() -> io::Result<()> { + let codex_home = TempDir::new()?; + let sqlite_home = TempDir::new()?; + write_session_with_user_event(codex_home.path()).await?; + + let config_toml = ConfigToml::default(); + let state_db = + state_db_for_test_with_sqlite_home(codex_home.path(), sqlite_home.path()).await?; + let status = maybe_migrate_personality(codex_home.path(), &config_toml, state_db).await?; + + assert_eq!(status, PersonalityMigrationStatus::Applied); + assert!( + codex_home + .path() + .join(PERSONALITY_MIGRATION_FILENAME) + .exists() + ); + + let persisted = read_config_toml(codex_home.path()).await?; + assert_eq!(persisted.personality, Some(Personality::Pragmatic)); + assert!(!state_db_path(codex_home.path()).exists()); + assert!(state_db_path(sqlite_home.path()).exists()); + Ok(()) +} diff --git a/codex-rs/core/src/prompt_debug.rs b/codex-rs/core/src/prompt_debug.rs index 7c6144c10e..a40fd27942 100644 --- a/codex-rs/core/src/prompt_debug.rs +++ b/codex-rs/core/src/prompt_debug.rs @@ -16,8 +16,9 @@ use crate::config::Config; use crate::session::session::Session; use crate::session::turn::build_prompt; use crate::session::turn::built_tools; -use crate::state_db_bridge::StateDbHandle; use crate::thread_manager::ThreadManager; +use crate::thread_manager::agent_graph_store_from_state_db; +use crate::thread_manager::init_state_db_from_config; use crate::thread_manager::thread_store_from_config; /// Build the model-visible `input` list for a single debug turn. @@ -25,7 +26,6 @@ use crate::thread_manager::thread_store_from_config; pub async fn build_prompt_input( mut config: Config, input: Vec, - state_db: Option, ) -> CodexResult> { config.ephemeral = true; @@ -37,15 +37,20 @@ pub async fn build_prompt_input( config.codex_linux_sandbox_exe.clone(), )?; + let state_db = init_state_db_from_config(&config) + .await + .ok_or_else(|| std::io::Error::other("prompt debug requires state db"))?; let thread_store = thread_store_from_config(&config, state_db.clone()); + let agent_graph_store = agent_graph_store_from_state_db(state_db.clone()); let thread_manager = ThreadManager::new( &config, Arc::clone(&auth_manager), SessionSource::Exec, Arc::new(EnvironmentManager::new(EnvironmentManagerArgs::new(local_runtime_paths)).await), /*analytics_events_client*/ None, + state_db, thread_store, - state_db.clone(), + agent_graph_store, ); let thread = thread_manager.start_thread(config).await?; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index d35bccb96e..85f2ef438b 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -132,7 +132,6 @@ use codex_terminal_detection::user_agent; use codex_thread_store::CreateThreadParams; use codex_thread_store::LiveThread; use codex_thread_store::LiveThreadInitGuard; -use codex_thread_store::LocalThreadStore; use codex_thread_store::ResumeThreadParams; use codex_thread_store::ThreadEventPersistenceMode; use codex_thread_store::ThreadPersistenceMetadata; @@ -356,7 +355,6 @@ use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TokenUsageInfo; use codex_protocol::protocol::WarningEvent; use codex_protocol::user_input::UserInput; -use codex_tools::ToolEnvironmentMode; use codex_tools::ToolsConfig; use codex_tools::ToolsConfigParams; use codex_utils_absolute_path::AbsolutePathBuf; @@ -412,6 +410,7 @@ pub(crate) struct CodexSpawnArgs { pub(crate) parent_trace: Option, pub(crate) environment_selections: ResolvedTurnEnvironments, pub(crate) analytics_events_client: Option, + pub(crate) state_db: Option, pub(crate) thread_store: Arc, } @@ -469,6 +468,7 @@ impl Codex { parent_trace: _, environment_selections, analytics_events_client, + state_db, thread_store, } = args; let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY); @@ -557,15 +557,7 @@ impl Codex { }; match thread_id { Some(thread_id) => { - let state_db_ctx = if config.ephemeral { - None - } else if let Some(local_store) = - thread_store.as_any().downcast_ref::() - { - local_store.state_db().await - } else { - None - }; + let state_db_ctx = state_db.clone(); state_db::get_dynamic_tools(state_db_ctx.as_deref(), thread_id, "codex_spawn") .await } @@ -651,6 +643,7 @@ impl Codex { agent_control, environment_manager, analytics_events_client, + state_db, thread_store, parent_rollout_thread_trace, ) @@ -1311,7 +1304,7 @@ impl Session { self.services.user_shell.as_ref().clone(), self.services.shell_snapshot_tx.clone(), self.services.session_telemetry.clone(), - self.services.state_db.clone(), + self.state_db(), ); } diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 0bb8da6ea6..23621f87cf 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -343,6 +343,7 @@ impl Session { agent_control: AgentControl, environment_manager: Arc, analytics_events_client: Option, + state_db: Option, thread_store: Arc, parent_rollout_thread_trace: ThreadTraceContext, ) -> anyhow::Result> { @@ -441,22 +442,7 @@ impl Session { otel.name = "session_init.thread_persistence", session_init.ephemeral = config.ephemeral, )); - let state_db_fut = async { - if config.ephemeral { - None - } else if let Some(local_store) = - thread_store.as_any().downcast_ref::() - { - local_store.state_db().await - } else { - None - } - } - .instrument(info_span!( - "session_init.state_db", - otel.name = "session_init.state_db", - session_init.ephemeral = config.ephemeral, - )); + let state_db_ctx = if config.ephemeral { None } else { state_db }; let is_subagent = session_configuration.session_source.is_non_root_agent(); let history_meta_fut = async { @@ -495,15 +481,9 @@ impl Session { // Join all independent futures. let ( thread_persistence_result, - state_db_ctx, (history_log_id, history_entry_count), (auth, mcp_servers, auth_statuses), - ) = tokio::join!( - thread_persistence_fut, - state_db_fut, - history_meta_fut, - auth_and_mcp_fut - ); + ) = tokio::join!(thread_persistence_fut, history_meta_fut, auth_and_mcp_fut); let mut live_thread_init = LiveThreadInitGuard::new(thread_persistence_result.map_err(|e| { diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 6995ef8dc9..f3a8d95107 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -867,7 +867,7 @@ async fn danger_full_access_turns_do_not_expose_managed_network_proxy() -> anyho &permission_profile_for_sandbox_policy(&SandboxPolicy::DangerFullAccess), )?; - let session = make_session_with_config(move |config| { + let (session, _codex_home) = make_session_with_config(move |config| { let cwd = config.cwd.clone(); config .permissions @@ -933,7 +933,7 @@ async fn danger_full_access_tool_attempts_do_not_enforce_managed_network() -> an &permission_profile_for_sandbox_policy(&SandboxPolicy::DangerFullAccess), )?; - let session = make_session_with_config(move |config| { + let (session, _codex_home) = make_session_with_config(move |config| { let cwd = config.cwd.clone(); config .permissions @@ -1008,7 +1008,7 @@ async fn workspace_write_turns_continue_to_expose_managed_network_proxy() -> any &permission_profile_for_sandbox_policy(&sandbox_policy), )?; - let session = make_session_with_config(move |config| { + let (session, _codex_home) = make_session_with_config(move |config| { let cwd = config.cwd.clone(); config .permissions @@ -1035,7 +1035,7 @@ async fn user_shell_commands_do_not_inherit_managed_network_proxy() -> anyhow::R &permission_profile_for_sandbox_policy(&sandbox_policy), )?; - let (session, rx) = make_session_with_config_and_rx(move |config| { + let (session, rx, _codex_home) = make_session_with_config_and_rx(move |config| { let cwd = config.cwd.clone(); config .permissions @@ -1165,7 +1165,7 @@ async fn reload_user_config_layer_updates_effective_apps_config() { #[tokio::test] async fn reload_user_config_layer_refreshes_hooks() -> anyhow::Result<()> { - let session = make_session_with_config(|config| { + let (session, _codex_home) = make_session_with_config(|config| { config .features .enable(Feature::CodexHooks) @@ -2982,15 +2982,13 @@ pub(crate) async fn make_session_configuration_for_tests() -> SessionConfigurati fn turn_environments_for_tests( environment: &Arc, cwd: &codex_utils_absolute_path::AbsolutePathBuf, -) -> crate::environment_selection::ResolvedTurnEnvironments { - crate::environment_selection::ResolvedTurnEnvironments { - turn_environments: vec![TurnEnvironment { - environment_id: codex_exec_server::LOCAL_ENVIRONMENT_ID.to_string(), - environment: Arc::clone(environment), - cwd: cwd.clone(), - shell: "bash".to_string(), - }], - } +) -> Vec { + vec![TurnEnvironment { + environment_id: codex_exec_server::LOCAL_ENVIRONMENT_ID.to_string(), + environment: Arc::clone(environment), + cwd: cwd.clone(), + shell: "bash".to_string(), + }] } #[tokio::test] @@ -3528,9 +3526,15 @@ async fn session_new_fails_when_zsh_fork_enabled_without_zsh_path() { AgentControl::default(), Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, + /*state_db*/ None, Arc::new(codex_thread_store::LocalThreadStore::new( codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()), - /*state_db*/ None, + codex_state::StateRuntime::init( + config.sqlite_home.clone(), + config.model_provider_id.clone(), + ) + .await + .expect("state db should initialize"), )), codex_rollout_trace::ThreadTraceContext::disabled(), ) @@ -3678,7 +3682,12 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { live_thread: None, thread_store: Arc::new(codex_thread_store::LocalThreadStore::new( codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()), - /*state_db*/ None, + codex_state::StateRuntime::init( + config.sqlite_home.clone(), + config.model_provider_id.clone(), + ) + .await + .expect("state db should initialize"), )), model_client: ModelClient::new( Some(auth_manager.clone()), @@ -3723,7 +3732,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { model_info, &models_manager, /*network*/ None, - turn_environments, + crate::environment_selection::ResolvedTurnEnvironments { turn_environments }, session_configuration.cwd.clone(), "turn_id".to_string(), skills_outcome, @@ -3756,14 +3765,18 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { async fn make_session_with_config( mutator: impl FnOnce(&mut Config), -) -> anyhow::Result> { - let (session, _rx_event) = make_session_with_config_and_rx(mutator).await?; - Ok(session) +) -> anyhow::Result<(Arc, tempfile::TempDir)> { + let (session, _rx_event, codex_home) = make_session_with_config_and_rx(mutator).await?; + Ok((session, codex_home)) } async fn make_session_with_config_and_rx( mutator: impl FnOnce(&mut Config), -) -> anyhow::Result<(Arc, async_channel::Receiver)> { +) -> anyhow::Result<( + Arc, + async_channel::Receiver, + tempfile::TempDir, +)> { let codex_home = tempfile::tempdir().expect("create temp dir"); let mut config = build_test_config(codex_home.path()).await; mutator(&mut config); @@ -3848,15 +3861,21 @@ async fn make_session_with_config_and_rx( AgentControl::default(), Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, + /*state_db*/ None, Arc::new(codex_thread_store::LocalThreadStore::new( codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()), - /*state_db*/ None, + codex_state::StateRuntime::init( + config.sqlite_home.clone(), + config.model_provider_id.clone(), + ) + .await + .expect("state db should initialize"), )), codex_rollout_trace::ThreadTraceContext::disabled(), ) .await?; - Ok((session, rx_event)) + Ok((session, rx_event, codex_home)) } #[tokio::test] @@ -4518,7 +4537,7 @@ async fn default_turn_honors_empty_stored_thread_environments() { let turn_context = session.new_default_turn().await; - assert!(turn_context.environments.primary().is_none()); + assert!(turn_context.environments.primary_environment().is_none()); assert!(turn_context.environments.turn_environments.is_empty()); assert_eq!(turn_context.cwd, session_cwd); assert_eq!(turn_context.config.cwd, session_cwd); @@ -4537,7 +4556,7 @@ async fn primary_environment_uses_first_turn_environment() { environment_id: "second".to_string(), environment: Arc::clone(&first_environment.environment), cwd: second_cwd.clone(), - shell: first_environment.shell.clone(), + shell: "bash".to_string(), }); assert_eq!( @@ -4580,7 +4599,7 @@ async fn empty_turn_environments_clear_primary_environment() { .await .expect("turn should start"); - assert!(turn_context.environments.primary().is_none()); + assert!(turn_context.environments.primary_environment().is_none()); assert!(turn_context.environments.turn_environments.is_empty()); assert_eq!(turn_context.cwd, session.get_config().await.cwd); assert_eq!(turn_context.config.cwd, session.get_config().await.cwd); @@ -5036,47 +5055,13 @@ async fn make_session_and_context_with_auth_and_config_and_rx( Arc, async_channel::Receiver, ) -where - F: FnOnce(&mut Config), -{ - let codex_home = tempfile::tempdir().expect("create temp dir"); - make_session_and_context_with_auth_config_home_and_rx( - auth, - dynamic_tools, - codex_home.path(), - configure_config, - ) - .await -} - -async fn make_session_and_context_with_auth_config_home_and_rx( - auth: CodexAuth, - dynamic_tools: Vec, - codex_home: &Path, - configure_config: F, -) -> ( - Arc, - Arc, - async_channel::Receiver, -) where F: FnOnce(&mut Config), { let (tx_event, rx_event) = async_channel::unbounded(); - let mut config = build_test_config(codex_home).await; + let codex_home = tempfile::tempdir().expect("create temp dir").keep(); + let mut config = build_test_config(codex_home.as_path()).await; configure_config(&mut config); - let state_db = if config.features.enabled(Feature::Goals) { - Some( - codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.model_provider_id.clone(), - ) - .await - .expect("goal tests should initialize sqlite state db"), - ) - } else { - None - }; let config = Arc::new(config); let conversation_id = ThreadId::default(); let auth_manager = AuthManager::from_auth_for_testing(auth); @@ -5162,6 +5147,12 @@ where .expect("create environment"), ); + let state_db = codex_state::StateRuntime::init( + config.sqlite_home.clone(), + config.model_provider_id.clone(), + ) + .await + .expect("state db should initialize"); let skills_watcher = Arc::new(SkillsWatcher::noop()); let services = SessionServices { mcp_connection_manager: Arc::new(RwLock::new(McpConnectionManager::new_uninitialized( @@ -5202,7 +5193,7 @@ where agent_control, network_proxy: None, network_approval: Arc::clone(&network_approval), - state_db: state_db.clone(), + state_db: Some(state_db.clone()), live_thread: None, thread_store: Arc::new(codex_thread_store::LocalThreadStore::new( codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()), @@ -5251,7 +5242,7 @@ where model_info, &models_manager, /*network*/ None, - turn_environments, + crate::environment_selection::ResolvedTurnEnvironments { turn_environments }, session_configuration.cwd.clone(), "turn_id".to_string(), skills_outcome, @@ -5301,13 +5292,10 @@ async fn make_goal_session_and_context_with_rx() -> ( Arc, Arc, async_channel::Receiver, - tempfile::TempDir, ) { - let codex_home = tempfile::tempdir().expect("create temp dir"); - let (session, turn_context, rx) = make_session_and_context_with_auth_config_home_and_rx( + let (session, turn_context, rx) = make_session_and_context_with_auth_and_config_and_rx( CodexAuth::from_api_key("Test API Key"), Vec::new(), - codex_home.path(), |config| { config .features @@ -5317,14 +5305,14 @@ async fn make_goal_session_and_context_with_rx() -> ( ) .await; upsert_goal_test_thread(session.as_ref()).await; - (session, turn_context, rx, codex_home) + (session, turn_context, rx) } async fn upsert_goal_test_thread(session: &Session) { let config = session.get_config().await; - let state_db = session - .state_db() - .expect("goal test session should have a state db"); + let state_db = goal_test_state_db(session) + .await + .expect("goal test state db should initialize"); let mut builder = codex_state::ThreadMetadataBuilder::new( session.conversation_id, config @@ -7148,7 +7136,7 @@ async fn abort_empty_active_turn_preserves_pending_input() { #[tokio::test] async fn interrupt_accounts_active_goal_before_pausing() -> anyhow::Result<()> { - let (sess, tc, _rx, _codex_home) = make_goal_session_and_context_with_rx().await; + let (sess, tc, _rx) = make_goal_session_and_context_with_rx().await; sess.set_thread_goal( tc.as_ref(), SetGoalRequest { @@ -7412,7 +7400,7 @@ async fn goal_test_state_db(sess: &Session) -> anyhow::Result anyhow::Result<()> { - let (sess, tc, rx, _codex_home) = make_goal_session_and_context_with_rx().await; + let (sess, tc, rx) = make_goal_session_and_context_with_rx().await; sess.set_thread_goal( tc.as_ref(), SetGoalRequest { @@ -7512,7 +7500,7 @@ async fn budget_limited_accounting_steers_active_turn_without_aborting() -> anyh #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn external_goal_mutation_accounts_active_turn_before_status_change() -> anyhow::Result<()> { - let (sess, tc, _rx, _codex_home) = make_goal_session_and_context_with_rx().await; + let (sess, tc, _rx) = make_goal_session_and_context_with_rx().await; sess.set_thread_goal( tc.as_ref(), SetGoalRequest { @@ -7579,7 +7567,7 @@ async fn external_goal_mutation_accounts_active_turn_before_status_change() -> a #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn external_active_goal_set_marks_current_turn_for_accounting() -> anyhow::Result<()> { - let (sess, tc, _rx, _codex_home) = make_goal_session_and_context_with_rx().await; + let (sess, tc, _rx) = make_goal_session_and_context_with_rx().await; sess.spawn_task( Arc::clone(&tc), Vec::new(), @@ -8246,7 +8234,7 @@ async fn sample_rollout( #[tokio::test] async fn create_goal_tool_rejects_existing_goal() { - let (session, turn_context, _rx, _codex_home) = make_goal_session_and_context_with_rx().await; + let (session, turn_context, _rx) = make_goal_session_and_context_with_rx().await; let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); let handler = CreateGoalHandler; @@ -8308,7 +8296,7 @@ async fn create_goal_tool_rejects_existing_goal() { #[tokio::test] async fn update_goal_tool_rejects_pausing_goal() { - let (session, turn_context, _rx, _codex_home) = make_goal_session_and_context_with_rx().await; + let (session, turn_context, _rx) = make_goal_session_and_context_with_rx().await; let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); let create_handler = CreateGoalHandler; let update_handler = UpdateGoalHandler; @@ -8369,7 +8357,7 @@ async fn update_goal_tool_rejects_pausing_goal() { #[tokio::test] async fn update_goal_tool_marks_goal_complete() { - let (session, turn_context, _rx, _codex_home) = make_goal_session_and_context_with_rx().await; + let (session, turn_context, _rx) = make_goal_session_and_context_with_rx().await; let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); let create_handler = CreateGoalHandler; let update_handler = UpdateGoalHandler; diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index 4c61dad18d..d7ec5937d2 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -731,7 +731,12 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { let skills_watcher = Arc::new(SkillsWatcher::noop()); let thread_store = Arc::new(codex_thread_store::LocalThreadStore::new( codex_thread_store::LocalThreadStoreConfig::from_config(&config), - /*state_db*/ None, + codex_state::StateRuntime::init( + config.sqlite_home.clone(), + config.model_provider_id.clone(), + ) + .await + .expect("state db should initialize"), )); let CodexSpawnOk { codex, .. } = Codex::spawn(CodexSpawnArgs { @@ -760,6 +765,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { turn_environments: Vec::new(), }, analytics_events_client: None, + state_db: None, thread_store, }) .await diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index fe54d61546..b4ef59df95 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -8,6 +8,7 @@ use codex_protocol::protocol::TurnEnvironmentSelection; use codex_sandboxing::compatibility_sandbox_policy_for_permission_profile; use codex_sandboxing::policy_transforms::effective_file_system_sandbox_policy; use codex_sandboxing::policy_transforms::effective_network_sandbox_policy; +use codex_tools::ToolEnvironmentMode; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; diff --git a/codex-rs/core/src/stream_events_utils.rs b/codex-rs/core/src/stream_events_utils.rs index 29884da4ae..fc608aed93 100644 --- a/codex-rs/core/src/stream_events_utils.rs +++ b/codex-rs/core/src/stream_events_utils.rs @@ -138,11 +138,8 @@ pub(crate) async fn record_completed_response_item( .await; } mark_thread_memory_mode_polluted_if_external_context(sess, turn_context, item).await; - let has_memory_citation = record_stage1_output_usage_and_detect_memory_citation( - sess.services.state_db.as_ref(), - item, - ) - .await; + let has_memory_citation = + record_stage1_output_usage_and_detect_memory_citation(sess.state_db(), item).await; if has_memory_citation { sess.record_memory_citation_for_turn(&turn_context.sub_id) .await; @@ -177,7 +174,7 @@ pub(crate) async fn mark_thread_memory_mode_polluted_if_external_context( } async fn record_stage1_output_usage_and_detect_memory_citation( - state_db_ctx: Option<&state_db::StateDbHandle>, + state_db_ctx: Option, item: &ResponseItem, ) -> bool { let Some(raw_text) = raw_assistant_output_text_from_item(item) else { diff --git a/codex-rs/core/src/test_support.rs b/codex-rs/core/src/test_support.rs index 48eec66c58..74305fb518 100644 --- a/codex-rs/core/src/test_support.rs +++ b/codex-rs/core/src/test_support.rs @@ -52,14 +52,14 @@ pub fn auth_manager_from_auth_with_home(auth: CodexAuth, codex_home: PathBuf) -> AuthManager::from_auth_for_testing_with_home(auth, codex_home) } -pub fn thread_manager_with_models_provider( +pub async fn thread_manager_with_models_provider( auth: CodexAuth, provider: ModelProviderInfo, ) -> ThreadManager { - ThreadManager::with_models_provider_for_tests(auth, provider) + ThreadManager::with_models_provider_for_tests(auth, provider).await } -pub fn thread_manager_with_models_provider_and_home( +pub async fn thread_manager_with_models_provider_and_home( auth: CodexAuth, provider: ModelProviderInfo, codex_home: PathBuf, @@ -71,22 +71,7 @@ pub fn thread_manager_with_models_provider_and_home( codex_home, environment_manager, ) -} - -pub fn thread_manager_with_models_provider_home_and_state( - auth: CodexAuth, - provider: ModelProviderInfo, - codex_home: PathBuf, - environment_manager: Arc, - state_db: Option, -) -> ThreadManager { - ThreadManager::with_models_provider_home_and_state_for_tests( - auth, - provider, - codex_home, - environment_manager, - state_db, - ) + .await } pub async fn start_thread_with_user_shell_override( diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index d4e36ab3c8..06f8118db7 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -18,6 +18,8 @@ use crate::skills_watcher::SkillsWatcher; use crate::skills_watcher::SkillsWatcherEvent; use crate::tasks::InterruptedTurnHistoryMarker; use crate::tasks::interrupted_turn_history_marker; +use codex_agent_graph_store::AgentGraphStore; +use codex_agent_graph_store::LocalAgentGraphStore; use codex_analytics::AnalyticsEventsClient; use codex_app_server_protocol::ThreadHistoryBuilder; use codex_app_server_protocol::TurnStatus; @@ -50,8 +52,8 @@ use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnAbortedEvent; use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::protocol::W3cTraceContext; +use codex_rollout::state_db; use codex_rollout::state_db::StateDbHandle; -use codex_state::DirectionalThreadSpawnEdgeStatus; use codex_thread_store::InMemoryThreadStore; use codex_thread_store::LocalThreadStore; use codex_thread_store::LocalThreadStoreConfig; @@ -247,9 +249,10 @@ pub(crate) struct ThreadManagerState { mcp_manager: Arc, skills_watcher: Arc, thread_store: Arc, + state_db: StateDbHandle, + agent_graph_store: Arc, session_source: SessionSource, analytics_events_client: Option, - state_db: Option, // Captures submitted ops for testing purpose when test mode is enabled. ops_log: Option, } @@ -265,10 +268,11 @@ pub fn build_models_manager( ) } -pub fn thread_store_from_config( - config: &Config, - state_db: Option, -) -> Arc { +pub async fn init_state_db_from_config(config: &Config) -> Option { + state_db::init(config).await +} + +pub fn thread_store_from_config(config: &Config, state_db: StateDbHandle) -> Arc { match &config.experimental_thread_store { ThreadStoreConfig::Local => Arc::new(LocalThreadStore::new( LocalThreadStoreConfig::from_config(config), @@ -279,15 +283,38 @@ pub fn thread_store_from_config( } } +pub fn agent_graph_store_from_state_db(state_db: StateDbHandle) -> Arc { + Arc::new(LocalAgentGraphStore::new(state_db)) +} + +async fn state_db_from_roots_for_tests( + codex_home: PathBuf, + sqlite_home: PathBuf, + default_model_provider_id: String, +) -> StateDbHandle { + let config = codex_rollout::RolloutConfig { + codex_home: codex_home.clone(), + sqlite_home, + cwd: codex_home, + model_provider_id: default_model_provider_id, + generate_memories: false, + }; + state_db::try_init(&config) + .await + .unwrap_or_else(|err| panic!("test state db should initialize: {err}")) +} + impl ThreadManager { + #[allow(clippy::too_many_arguments)] pub fn new( config: &Config, auth_manager: Arc, session_source: SessionSource, environment_manager: Arc, analytics_events_client: Option, + state_db: StateDbHandle, thread_store: Arc, - state_db: Option, + agent_graph_store: Arc, ) -> Self { let codex_home = config.codex_home.clone(); let restriction_product = session_source.restriction_product(); @@ -314,10 +341,11 @@ impl ThreadManager { mcp_manager, skills_watcher, thread_store, + state_db, + agent_graph_store, auth_manager, session_source, analytics_events_client, - state_db, ops_log: should_use_test_thread_manager_behavior() .then(|| Arc::new(std::sync::Mutex::new(Vec::new()))), }), @@ -327,7 +355,7 @@ impl ThreadManager { /// Construct with a dummy AuthManager containing the provided CodexAuth. /// Used for integration tests: should not be used by ordinary business logic. - pub(crate) fn with_models_provider_for_tests( + pub(crate) async fn with_models_provider_for_tests( auth: CodexAuth, provider: ModelProviderInfo, ) -> Self { @@ -338,11 +366,18 @@ impl ThreadManager { )); std::fs::create_dir_all(&codex_home) .unwrap_or_else(|err| panic!("temp codex home dir create failed: {err}")); - let mut manager = Self::with_models_provider_and_home_for_tests( + let state_db = state_db_from_roots_for_tests( + codex_home.clone(), + codex_home.clone(), + OPENAI_PROVIDER_ID.to_string(), + ) + .await; + let mut manager = Self::with_models_provider_and_home_and_state_db_for_tests( auth, provider, codex_home.clone(), Arc::new(EnvironmentManager::default_for_tests()), + state_db, ); manager._test_codex_home_guard = Some(TempCodexHomeGuard { path: codex_home }); manager @@ -350,27 +385,33 @@ impl ThreadManager { /// Construct with a dummy AuthManager containing the provided CodexAuth and codex home. /// Used for integration tests: should not be used by ordinary business logic. - pub(crate) fn with_models_provider_and_home_for_tests( + pub(crate) async fn with_models_provider_and_home_for_tests( auth: CodexAuth, provider: ModelProviderInfo, codex_home: PathBuf, environment_manager: Arc, ) -> Self { - Self::with_models_provider_home_and_state_for_tests( + let state_db = state_db_from_roots_for_tests( + codex_home.clone(), + codex_home.clone(), + OPENAI_PROVIDER_ID.to_string(), + ) + .await; + Self::with_models_provider_and_home_and_state_db_for_tests( auth, provider, codex_home, environment_manager, - /*state_db*/ None, + state_db, ) } - pub(crate) fn with_models_provider_home_and_state_for_tests( + fn with_models_provider_and_home_and_state_db_for_tests( auth: CodexAuth, provider: ModelProviderInfo, codex_home: PathBuf, environment_manager: Arc, - state_db: Option, + state_db: StateDbHandle, ) -> Self { set_thread_manager_test_mode_for_tests(/*enabled*/ true); let auth_manager = AuthManager::from_auth_for_testing(auth); @@ -396,11 +437,11 @@ impl ThreadManager { let thread_store: Arc = Arc::new(LocalThreadStore::new( LocalThreadStoreConfig { codex_home: codex_home.clone(), - sqlite_home: codex_home.clone(), default_model_provider_id: OPENAI_PROVIDER_ID.to_string(), }, state_db.clone(), )); + let agent_graph_store = agent_graph_store_from_state_db(state_db.clone()); Self { state: Arc::new(ThreadManagerState { threads: Arc::new(RwLock::new(HashMap::new())), @@ -413,10 +454,11 @@ impl ThreadManager { mcp_manager, skills_watcher, thread_store, + state_db, + agent_graph_store, auth_manager, session_source: SessionSource::Exec, analytics_events_client: None, - state_db, ops_log: should_use_test_thread_manager_behavior() .then(|| Arc::new(std::sync::Mutex::new(Vec::new()))), }), @@ -523,22 +565,17 @@ impl ThreadManager { subtree_thread_ids.push(thread_id); seen_thread_ids.insert(thread_id); - if let Some(state_db_ctx) = thread.state_db() { - for status in [ - DirectionalThreadSpawnEdgeStatus::Open, - DirectionalThreadSpawnEdgeStatus::Closed, - ] { - for descendant_id in state_db_ctx - .list_thread_spawn_descendants_with_status(thread_id, status) - .await - .map_err(|err| { - CodexErr::Fatal(format!("failed to load thread-spawn descendants: {err}")) - })? - { - if seen_thread_ids.insert(descendant_id) { - subtree_thread_ids.push(descendant_id); - } - } + for descendant_id in self + .state + .agent_graph_store + .list_thread_spawn_descendants(thread_id, /*status_filter*/ None) + .await + .map_err(|err| { + CodexErr::Fatal(format!("failed to load thread-spawn descendants: {err}")) + })? + { + if seen_thread_ids.insert(descendant_id) { + subtree_thread_ids.push(descendant_id); } } @@ -864,10 +901,14 @@ impl ThreadManager { } impl ThreadManagerState { - pub(crate) fn state_db(&self) -> Option { + pub(crate) fn state_db(&self) -> StateDbHandle { self.state_db.clone() } + pub(crate) fn agent_graph_store(&self) -> Arc { + self.agent_graph_store.clone() + } + pub(crate) async fn list_thread_ids(&self) -> Vec { self.threads .read() @@ -1172,6 +1213,7 @@ impl ThreadManagerState { parent_trace, environment_selections, analytics_events_client: self.analytics_events_client.clone(), + state_db: Some(self.state_db.clone()), thread_store: Arc::clone(&self.thread_store), }) .await?; diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index d50a3ae8f9..643309aac1 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -1,6 +1,5 @@ use super::*; use crate::config::test_config; -use crate::init_state_db; use crate::rollout::RolloutRecorder; use crate::session::session::SessionSettingsUpdate; use crate::session::tests::make_session_and_context; @@ -47,6 +46,21 @@ fn assistant_msg(text: &str) -> ResponseItem { } } +async fn state_backed_stores( + config: &Config, +) -> ( + StateDbHandle, + Arc, + Arc, +) { + let state_db = init_state_db_from_config(config) + .await + .expect("thread manager test requires state db"); + let thread_store = thread_store_from_config(config, state_db.clone()); + let agent_graph_store = agent_graph_store_from_state_db(state_db.clone()); + (state_db, thread_store, agent_graph_store) +} + fn contextual_user_interrupted_marker() -> ResponseItem { interrupted_turn_history_marker(InterruptedTurnHistoryMarker::ContextualUser) .expect("contextual-user interrupted marker should be enabled") @@ -261,7 +275,8 @@ async fn shutdown_all_threads_bounded_submits_shutdown_to_every_thread() { config.model_provider.clone(), config.codex_home.to_path_buf(), Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - ); + ) + .await; let thread_1 = manager .start_thread(config.clone()) .await @@ -310,7 +325,8 @@ async fn start_thread_accepts_explicit_environment_when_default_environment_is_d config.model_provider.clone(), config.codex_home.to_path_buf(), environment_manager, - ); + ) + .await; let thread = manager .start_thread_with_options(StartThreadOptions { @@ -345,7 +361,8 @@ async fn start_thread_keeps_internal_threads_hidden_from_normal_lookups() { config.model_provider.clone(), config.codex_home.to_path_buf(), Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - ); + ) + .await; let thread = manager .start_thread_with_options(StartThreadOptions { config, @@ -384,14 +401,16 @@ async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await; let manager = ThreadManager::new( &config, auth_manager.clone(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, - thread_store_from_config(&config, /*state_db*/ None), - /*state_db*/ None, + state_db, + thread_store, + agent_graph_store, ); let selected_cwd = AbsolutePathBuf::try_from(config.cwd.as_path().join("selected")).expect("absolute path"); @@ -494,14 +513,16 @@ async fn resume_active_thread_from_rollout_returns_running_thread() { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await; let manager = ThreadManager::new( &config, auth_manager.clone(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, - thread_store_from_config(&config, /*state_db*/ None), - /*state_db*/ None, + state_db, + thread_store, + agent_graph_store, ); let source = manager @@ -548,14 +569,16 @@ async fn resume_stopped_thread_from_rollout_spawns_new_thread() { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await; let manager = ThreadManager::new( &config, auth_manager.clone(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, - thread_store_from_config(&config, /*state_db*/ None), - /*state_db*/ None, + state_db, + thread_store, + agent_graph_store, ); let source = manager @@ -612,14 +635,16 @@ async fn new_uses_active_provider_for_model_refresh() { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await; let manager = ThreadManager::new( &config, auth_manager, SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, - thread_store_from_config(&config, /*state_db*/ None), - /*state_db*/ None, + state_db, + thread_store, + agent_graph_store, ); let _ = manager.list_models(RefreshStrategy::Online).await; @@ -824,15 +849,16 @@ async fn interrupted_fork_snapshot_does_not_synthesize_turn_id_for_legacy_histor let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); - let state_db = init_state_db(&config).await; + let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await; let manager = ThreadManager::new( &config, auth_manager.clone(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, - thread_store_from_config(&config, state_db.clone()), - state_db.clone(), + state_db, + thread_store, + agent_graph_store, ); let source = manager @@ -928,15 +954,16 @@ async fn interrupted_fork_snapshot_preserves_explicit_turn_id() { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); - let state_db = init_state_db(&config).await; + let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await; let manager = ThreadManager::new( &config, auth_manager.clone(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, - thread_store_from_config(&config, state_db.clone()), - state_db.clone(), + state_db, + thread_store, + agent_graph_store, ); let source = manager @@ -1021,15 +1048,16 @@ async fn interrupted_fork_snapshot_uses_persisted_mid_turn_history_without_live_ let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); - let state_db = init_state_db(&config).await; + let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await; let manager = ThreadManager::new( &config, auth_manager.clone(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, - thread_store_from_config(&config, state_db.clone()), - state_db.clone(), + state_db, + thread_store, + agent_graph_store, ); let source = manager @@ -1159,15 +1187,16 @@ async fn resumed_thread_keeps_paused_goal_paused() -> anyhow::Result<()> { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); - let state_db = init_state_db(&config).await; + let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await; let manager = ThreadManager::new( &config, auth_manager.clone(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, - thread_store_from_config(&config, state_db.clone()), - state_db.clone(), + state_db, + thread_store, + agent_graph_store, ); let source = manager @@ -1236,6 +1265,11 @@ async fn resumed_thread_keeps_paused_goal_paused() -> anyhow::Result<()> { .await .is_none() ); + let goal = state_db + .get_thread_goal(resumed.thread_id) + .await? + .expect("goal should still exist after resume"); + assert_eq!(codex_state::ThreadGoalStatus::Paused, goal.status); resumed.thread.shutdown_and_wait().await?; Ok(()) 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 88c1698d8f..9a23550d41 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -6,6 +6,7 @@ use crate::function_tool::FunctionCallError; use crate::init_state_db; use crate::session::tests::make_session_and_context; use crate::session_prefix::format_subagent_notification_message; +use crate::thread_manager::agent_graph_store_from_state_db; use crate::thread_manager::thread_store_from_config; use crate::tools::context::ToolOutput; use crate::tools::handlers::multi_agents_v2::CloseAgentHandler as CloseAgentHandlerV2; @@ -90,11 +91,12 @@ fn parse_agent_id(id: &str) -> ThreadId { ThreadId::from_string(id).expect("agent id should be valid") } -fn thread_manager() -> ThreadManager { +async fn thread_manager() -> ThreadManager { ThreadManager::with_models_provider_for_tests( CodexAuth::from_api_key("dummy"), built_in_model_providers(/* openai_base_url */ /*openai_base_url*/ None)["openai"].clone(), ) + .await } async fn install_role_with_model_override(turn: &mut TurnContext) -> String { @@ -241,7 +243,7 @@ async fn spawn_agent_uses_explorer_role_and_preserves_approval_policy() { } let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let mut config = (*turn.config).clone(); let provider_info = @@ -296,7 +298,7 @@ async fn spawn_agent_uses_explorer_role_and_preserves_approval_policy() { async fn spawn_agent_fork_context_rejects_agent_type_override() { let (mut session, mut turn) = make_session_and_context().await; let role_name = install_role_with_model_override(&mut turn).await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -328,7 +330,7 @@ async fn spawn_agent_fork_context_rejects_agent_type_override() { #[tokio::test] async fn spawn_agent_fork_context_rejects_child_model_overrides() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -363,7 +365,7 @@ async fn spawn_agent_fork_context_rejects_child_model_overrides() { async fn multi_agent_v2_spawn_fork_turns_all_rejects_agent_type_override() { let (mut session, mut turn) = make_session_and_context().await; let role_name = install_role_with_model_override(&mut turn).await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -406,7 +408,7 @@ async fn multi_agent_v2_spawn_fork_turns_all_rejects_agent_type_override() { #[tokio::test] async fn multi_agent_v2_spawn_defaults_to_full_fork_and_rejects_child_model_overrides() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -447,7 +449,7 @@ async fn multi_agent_v2_spawn_defaults_to_full_fork_and_rejects_child_model_over async fn multi_agent_v2_spawn_partial_fork_turns_allows_agent_type_override() { let (mut session, mut turn) = make_session_and_context().await; let role_name = install_role_with_model_override(&mut turn).await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -503,7 +505,7 @@ async fn multi_agent_v2_spawn_partial_fork_turns_allows_agent_type_override() { #[tokio::test] async fn spawn_agent_returns_agent_id_without_task_name() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let output = SpawnAgentHandler @@ -530,7 +532,7 @@ async fn spawn_agent_returns_agent_id_without_task_name() { #[tokio::test] async fn multi_agent_v2_spawn_requires_task_name() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -564,7 +566,7 @@ async fn multi_agent_v2_spawn_requires_task_name() { #[tokio::test] async fn multi_agent_v2_spawn_rejects_legacy_items_field() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -624,7 +626,7 @@ async fn multi_agent_v2_spawn_returns_path_and_send_message_accepts_relative_pat } let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -721,7 +723,7 @@ async fn multi_agent_v2_spawn_returns_path_and_send_message_accepts_relative_pat #[tokio::test] async fn multi_agent_v2_spawn_rejects_legacy_fork_context() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -760,7 +762,7 @@ async fn multi_agent_v2_spawn_rejects_legacy_fork_context() { #[tokio::test] async fn multi_agent_v2_spawn_rejects_invalid_fork_turns_string() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -799,7 +801,7 @@ async fn multi_agent_v2_spawn_rejects_invalid_fork_turns_string() { #[tokio::test] async fn multi_agent_v2_spawn_rejects_zero_fork_turns() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -838,7 +840,7 @@ async fn multi_agent_v2_spawn_rejects_zero_fork_turns() { #[tokio::test] async fn multi_agent_v2_send_message_accepts_root_target_from_child() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -914,7 +916,7 @@ async fn multi_agent_v2_send_message_accepts_root_target_from_child() { #[tokio::test] async fn multi_agent_v2_followup_task_rejects_root_target_from_child() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -995,7 +997,7 @@ async fn multi_agent_v2_followup_task_rejects_root_target_from_child() { #[tokio::test] async fn multi_agent_v2_list_agents_returns_completed_status_and_last_task_message() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -1089,7 +1091,7 @@ async fn multi_agent_v2_list_agents_returns_completed_status_and_last_task_messa #[tokio::test] async fn multi_agent_v2_list_agents_filters_by_relative_path_prefix() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -1176,7 +1178,7 @@ async fn multi_agent_v2_list_agents_filters_by_relative_path_prefix() { #[tokio::test] async fn multi_agent_v2_list_agents_omits_closed_agents() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -1240,7 +1242,7 @@ async fn multi_agent_v2_list_agents_omits_closed_agents() { #[tokio::test] async fn multi_agent_v2_send_message_rejects_legacy_items_field() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -1296,7 +1298,7 @@ async fn multi_agent_v2_send_message_rejects_legacy_items_field() { #[tokio::test] async fn multi_agent_v2_send_message_rejects_interrupt_parameter() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -1369,7 +1371,7 @@ async fn multi_agent_v2_send_message_rejects_interrupt_parameter() { #[tokio::test] async fn multi_agent_v2_followup_task_completion_notifies_parent_on_every_turn() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -1504,7 +1506,7 @@ async fn multi_agent_v2_followup_task_completion_notifies_parent_on_every_turn() #[tokio::test] async fn multi_agent_v2_followup_task_rejects_legacy_items_field() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -1557,7 +1559,7 @@ async fn multi_agent_v2_followup_task_rejects_legacy_items_field() { #[tokio::test] async fn multi_agent_v2_interrupted_turn_does_not_notify_parent() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -1634,7 +1636,7 @@ async fn multi_agent_v2_interrupted_turn_does_not_notify_parent() { #[tokio::test] async fn multi_agent_v2_spawn_omits_agent_id_when_named() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -1673,7 +1675,7 @@ async fn multi_agent_v2_spawn_omits_agent_id_when_named() { #[tokio::test] async fn multi_agent_v2_spawn_surfaces_task_name_validation_errors() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -1716,7 +1718,7 @@ async fn spawn_agent_reapplies_runtime_sandbox_after_role_config() { } let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let expected_sandbox = turn.config.legacy_sandbox_policy(); let mut expected_file_system_sandbox_policy = @@ -1797,7 +1799,7 @@ async fn spawn_agent_reapplies_runtime_sandbox_after_role_config() { #[tokio::test] async fn spawn_agent_rejects_when_depth_limit_exceeded() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let max_depth = turn.config.agent_max_depth; @@ -1835,7 +1837,7 @@ async fn spawn_agent_allows_depth_up_to_configured_max_depth() { } let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let mut config = (*turn.config).clone(); @@ -1881,7 +1883,7 @@ async fn multi_agent_v2_spawn_agent_ignores_configured_max_depth() { } let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let mut config = (*turn.config).clone(); config.agent_max_depth = 1; config @@ -1989,7 +1991,7 @@ async fn send_input_rejects_invalid_id() { #[tokio::test] async fn send_input_reports_missing_agent() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let agent_id = ThreadId::new(); let invocation = invocation( @@ -2010,7 +2012,7 @@ async fn send_input_reports_missing_agent() { #[tokio::test] async fn send_input_interrupts_before_prompt() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let config = turn.config.as_ref().clone(); let thread = manager @@ -2052,7 +2054,7 @@ async fn send_input_interrupts_before_prompt() { #[tokio::test] async fn send_input_accepts_structured_items() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let config = turn.config.as_ref().clone(); let thread = manager @@ -2126,7 +2128,7 @@ async fn resume_agent_rejects_invalid_id() { #[tokio::test] async fn resume_agent_reports_missing_agent() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let agent_id = ThreadId::new(); let invocation = invocation( @@ -2147,7 +2149,7 @@ async fn resume_agent_reports_missing_agent() { #[tokio::test] async fn resume_agent_noops_for_active_agent() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let config = turn.config.as_ref().clone(); let thread = manager @@ -2186,7 +2188,7 @@ async fn resume_agent_noops_for_active_agent() { #[tokio::test] async fn resume_agent_restores_closed_agent_and_accepts_send_input() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let config = turn.config.as_ref().clone(); let thread = manager @@ -2265,7 +2267,7 @@ async fn resume_agent_restores_closed_agent_and_accepts_send_input() { #[tokio::test] async fn resume_agent_rejects_when_depth_limit_exceeded() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let max_depth = turn.config.agent_max_depth; @@ -2354,7 +2356,7 @@ async fn wait_agent_rejects_empty_targets() { #[tokio::test] async fn multi_agent_v2_wait_agent_accepts_timeout_only_argument() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -2493,7 +2495,7 @@ async fn multi_agent_v2_wait_agent_uses_configured_min_timeout() { #[tokio::test] async fn wait_agent_returns_not_found_for_missing_agents() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let id_a = ThreadId::new(); let id_b = ThreadId::new(); @@ -2529,7 +2531,7 @@ async fn wait_agent_returns_not_found_for_missing_agents() { #[tokio::test] async fn wait_agent_times_out_when_status_is_not_final() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let config = turn.config.as_ref().clone(); let thread = manager @@ -2572,7 +2574,7 @@ async fn wait_agent_times_out_when_status_is_not_final() { #[tokio::test] async fn wait_agent_clamps_short_timeouts_to_minimum() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let config = turn.config.as_ref().clone(); let thread = manager @@ -2610,7 +2612,7 @@ async fn wait_agent_clamps_short_timeouts_to_minimum() { #[tokio::test] async fn wait_agent_returns_final_status_without_timeout() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let config = turn.config.as_ref().clone(); let thread = manager @@ -2662,7 +2664,7 @@ async fn wait_agent_returns_final_status_without_timeout() { #[tokio::test] async fn multi_agent_v2_wait_agent_returns_summary_for_mailbox_activity() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -2753,7 +2755,7 @@ async fn multi_agent_v2_wait_agent_returns_summary_for_mailbox_activity() { #[tokio::test] async fn multi_agent_v2_wait_agent_returns_for_already_queued_mail() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -2831,7 +2833,7 @@ async fn multi_agent_v2_wait_agent_returns_for_already_queued_mail() { #[tokio::test] async fn multi_agent_v2_wait_agent_wakes_on_any_mailbox_notification() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -2919,7 +2921,7 @@ async fn multi_agent_v2_wait_agent_wakes_on_any_mailbox_notification() { #[tokio::test] async fn multi_agent_v2_wait_agent_does_not_return_completed_content() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -3005,7 +3007,7 @@ async fn multi_agent_v2_wait_agent_does_not_return_completed_content() { #[tokio::test] async fn multi_agent_v2_close_agent_accepts_task_name_target() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -3064,7 +3066,7 @@ async fn multi_agent_v2_close_agent_accepts_task_name_target() { #[tokio::test] async fn multi_agent_v2_close_agent_rejects_root_target_and_id() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -3112,7 +3114,7 @@ async fn multi_agent_v2_close_agent_rejects_root_target_and_id() { #[tokio::test] async fn close_agent_submits_shutdown_and_returns_previous_status() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let config = turn.config.as_ref().clone(); let thread = manager @@ -3157,15 +3159,18 @@ async fn tool_handlers_cascade_close_and_resume_and_keep_explicitly_closed_subtr .features .enable(Feature::Sqlite) .expect("test config should allow sqlite"); - let state_db = init_state_db(&config).await; + let state_db = init_state_db(&config) + .await + .expect("test config should initialize state db"); let manager = ThreadManager::new( &config, AuthManager::from_auth_for_testing(CodexAuth::from_api_key("dummy")), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, - thread_store_from_config(&config, state_db.clone()), state_db.clone(), + thread_store_from_config(&config, state_db.clone()), + agent_graph_store_from_state_db(state_db.clone()), ); let parent = manager diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index 14f652f58b..13dd026654 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -15,7 +15,9 @@ use anyhow::anyhow; use codex_config::CloudRequirementsLoader; use codex_core::CodexThread; use codex_core::ThreadManager; +use codex_core::agent_graph_store_from_state_db; use codex_core::config::Config; +use codex_core::init_state_db_from_config; use codex_core::shell::Shell; use codex_core::shell::get_shell_by_model_provided_path; use codex_core::thread_store_from_config; @@ -423,25 +425,30 @@ impl TestCodexBuilder { environment_manager: Arc, ) -> anyhow::Result { let auth = self.auth.clone(); - let state_db = codex_core::init_state_db(&config).await; let thread_manager = if config.model_catalog.is_some() { + let state_db = init_state_db_from_config(&config) + .await + .expect("test codex requires state db"); + let thread_store = thread_store_from_config(&config, state_db.clone()); + let agent_graph_store = agent_graph_store_from_state_db(state_db.clone()); ThreadManager::new( &config, codex_core::test_support::auth_manager_from_auth(auth.clone()), SessionSource::Exec, Arc::clone(&environment_manager), /*analytics_events_client*/ None, - thread_store_from_config(&config, state_db.clone()), - state_db.clone(), + state_db, + thread_store, + agent_graph_store, ) } else { - codex_core::test_support::thread_manager_with_models_provider_home_and_state( + codex_core::test_support::thread_manager_with_models_provider_and_home( auth.clone(), config.model_provider.clone(), config.codex_home.to_path_buf(), Arc::clone(&environment_manager), - state_db.clone(), ) + .await }; let thread_manager = Arc::new(thread_manager); let user_shell_override = self.user_shell_override.clone(); diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index afa586a158..dafde4ccea 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -5,6 +5,8 @@ use codex_core::NewThread; use codex_core::Prompt; use codex_core::ResponseEvent; use codex_core::ThreadManager; +use codex_core::agent_graph_store_from_state_db; +use codex_core::init_state_db_from_config; use codex_core::thread_store_from_config; use codex_features::Feature; use codex_login::AuthManager; @@ -770,15 +772,10 @@ async fn includes_conversation_id_and_model_headers_in_request() { let installation_id = std::fs::read_to_string(test.codex_home_path().join(INSTALLATION_ID_FILENAME)) .expect("read installation id"); - let session_id_string = session_id.to_string(); - assert_eq!(request_session_id, session_id_string); + assert_eq!(request_session_id, session_id.to_string()); assert_eq!(request_originator, originator().value); assert_eq!(request_authorization, "Bearer Test API Key"); - assert_eq!( - request_body["prompt_cache_key"].as_str(), - Some(session_id_string.as_str()) - ); assert_eq!( request_body["client_metadata"]["x-codex-installation-id"].as_str(), Some(installation_id.as_str()) @@ -1106,14 +1103,20 @@ async fn prefers_apikey_when_config_prefers_apikey_even_with_chatgpt_tokens() { Ok(None) => panic!("No CodexAuth found in codex_home"), Err(e) => panic!("Failed to load CodexAuth: {e}"), }; + let state_db = init_state_db_from_config(&config) + .await + .expect("client test requires state db"); + let thread_store = thread_store_from_config(&config, state_db.clone()); + let agent_graph_store = agent_graph_store_from_state_db(state_db.clone()); let thread_manager = ThreadManager::new( &config, auth_manager, SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, - thread_store_from_config(&config, /*state_db*/ None), - /*state_db*/ None, + state_db, + thread_store, + agent_graph_store, ); let NewThread { thread: codex, .. } = thread_manager .start_thread(config.clone()) diff --git a/codex-rs/core/tests/suite/personality_migration.rs b/codex-rs/core/tests/suite/personality_migration.rs index 68121d2e12..8d33ba1705 100644 --- a/codex-rs/core/tests/suite/personality_migration.rs +++ b/codex-rs/core/tests/suite/personality_migration.rs @@ -13,6 +13,8 @@ use codex_protocol::protocol::SessionMeta; use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::UserMessageEvent; +use codex_rollout::RolloutConfig; +use codex_rollout::state_db::StateDbHandle; use pretty_assertions::assert_eq; use std::io; use std::path::Path; @@ -26,6 +28,27 @@ async fn read_config_toml(codex_home: &Path) -> io::Result { toml::from_str(&contents).map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err)) } +async fn state_db_for_test(codex_home: &Path) -> io::Result { + let config = RolloutConfig { + codex_home: codex_home.to_path_buf(), + sqlite_home: codex_home.to_path_buf(), + cwd: codex_home.to_path_buf(), + model_provider_id: "openai".to_string(), + generate_memories: false, + }; + codex_rollout::state_db::try_init(&config) + .await + .map_err(io::Error::other) +} + +async fn run_migration( + codex_home: &Path, + config_toml: &ConfigToml, +) -> io::Result { + let state_db = state_db_for_test(codex_home).await?; + maybe_migrate_personality(codex_home, config_toml, state_db).await +} + async fn write_session_with_user_event(codex_home: &Path) -> io::Result<()> { let thread_id = ThreadId::new(); let dir = codex_home @@ -141,8 +164,7 @@ async fn migration_marker_exists_no_sessions_no_change() -> io::Result<()> { let marker_path = temp.path().join(PERSONALITY_MIGRATION_FILENAME); tokio::fs::write(&marker_path, "v1\n").await?; - let status = - maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?; + let status = run_migration(temp.path(), &ConfigToml::default()).await?; assert_eq!(status, PersonalityMigrationStatus::SkippedMarker); assert_eq!( @@ -156,8 +178,7 @@ async fn migration_marker_exists_no_sessions_no_change() -> io::Result<()> { async fn no_marker_no_sessions_no_change() -> io::Result<()> { let temp = TempDir::new()?; - let status = - maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?; + let status = run_migration(temp.path(), &ConfigToml::default()).await?; assert_eq!(status, PersonalityMigrationStatus::SkippedNoSessions); assert_eq!( @@ -176,8 +197,7 @@ async fn no_marker_sessions_sets_personality() -> io::Result<()> { let temp = TempDir::new()?; write_session_with_user_event(temp.path()).await?; - let status = - maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?; + let status = run_migration(temp.path(), &ConfigToml::default()).await?; assert_eq!(status, PersonalityMigrationStatus::Applied); assert_eq!( @@ -197,7 +217,7 @@ async fn no_marker_sessions_preserves_existing_config_fields() -> io::Result<()> tokio::fs::write(temp.path().join("config.toml"), "model = \"gpt-5.4\"\n").await?; let config_toml = read_config_toml(temp.path()).await?; - let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?; + let status = run_migration(temp.path(), &config_toml).await?; assert_eq!(status, PersonalityMigrationStatus::Applied); let persisted = read_config_toml(temp.path()).await?; @@ -211,8 +231,7 @@ async fn no_marker_meta_only_rollout_is_treated_as_no_sessions() -> io::Result<( let temp = TempDir::new()?; write_session_with_meta_only(temp.path()).await?; - let status = - maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?; + let status = run_migration(temp.path(), &ConfigToml::default()).await?; assert_eq!(status, PersonalityMigrationStatus::SkippedNoSessions); assert_eq!( @@ -232,7 +251,7 @@ async fn no_marker_explicit_global_personality_skips_migration() -> io::Result<( write_session_with_user_event(temp.path()).await?; let config_toml = parse_config_toml("personality = \"friendly\"\n")?; - let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?; + let status = run_migration(temp.path(), &config_toml).await?; assert_eq!( status, @@ -262,7 +281,7 @@ personality = "friendly" "#, )?; - let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?; + let status = run_migration(temp.path(), &config_toml).await?; assert_eq!( status, @@ -285,7 +304,7 @@ async fn marker_short_circuits_invalid_profile_resolution() -> io::Result<()> { tokio::fs::write(temp.path().join(PERSONALITY_MIGRATION_FILENAME), "v1\n").await?; let config_toml = parse_config_toml("profile = \"missing\"\n")?; - let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?; + let status = run_migration(temp.path(), &config_toml).await?; assert_eq!(status, PersonalityMigrationStatus::SkippedMarker); Ok(()) @@ -296,7 +315,7 @@ async fn invalid_selected_profile_returns_error_and_does_not_write_marker() -> i let temp = TempDir::new()?; let config_toml = parse_config_toml("profile = \"missing\"\n")?; - let err = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None) + let err = run_migration(temp.path(), &config_toml) .await .expect_err("missing profile should fail"); @@ -313,10 +332,8 @@ async fn applied_migration_is_idempotent_on_second_run() -> io::Result<()> { let temp = TempDir::new()?; write_session_with_user_event(temp.path()).await?; - let first_status = - maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?; - let second_status = - maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?; + let first_status = run_migration(temp.path(), &ConfigToml::default()).await?; + let second_status = run_migration(temp.path(), &ConfigToml::default()).await?; assert_eq!(first_status, PersonalityMigrationStatus::Applied); assert_eq!(second_status, PersonalityMigrationStatus::SkippedMarker); @@ -330,8 +347,7 @@ async fn no_marker_archived_sessions_sets_personality() -> io::Result<()> { let temp = TempDir::new()?; write_archived_session_with_user_event(temp.path()).await?; - let status = - maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?; + let status = run_migration(temp.path(), &ConfigToml::default()).await?; assert_eq!(status, PersonalityMigrationStatus::Applied); assert_eq!( diff --git a/codex-rs/core/tests/suite/prompt_debug_tests.rs b/codex-rs/core/tests/suite/prompt_debug_tests.rs index dc506bc474..4fee438261 100644 --- a/codex-rs/core/tests/suite/prompt_debug_tests.rs +++ b/codex-rs/core/tests/suite/prompt_debug_tests.rs @@ -29,7 +29,6 @@ async fn build_prompt_input_includes_context_and_user_message() -> Result<()> { text: "hello from debug prompt".to_string(), text_elements: Vec::new(), }], - /*state_db*/ None, ) .await?; diff --git a/codex-rs/core/tests/suite/resume_warning.rs b/codex-rs/core/tests/suite/resume_warning.rs index cb545df351..02e7626d87 100644 --- a/codex-rs/core/tests/suite/resume_warning.rs +++ b/codex-rs/core/tests/suite/resume_warning.rs @@ -95,7 +95,8 @@ async fn emits_warning_when_resumed_model_differs() { let thread_manager = codex_core::test_support::thread_manager_with_models_provider( CodexAuth::from_api_key("test"), config.model_provider.clone(), - ); + ) + .await; let auth_manager = codex_core::test_support::auth_manager_from_auth(CodexAuth::from_api_key("test")); diff --git a/codex-rs/core/tests/suite/skills.rs b/codex-rs/core/tests/suite/skills.rs index add1ee6ea7..ebf2c3226e 100644 --- a/codex-rs/core/tests/suite/skills.rs +++ b/codex-rs/core/tests/suite/skills.rs @@ -3,6 +3,8 @@ use anyhow::Result; use codex_core::ThreadManager; +use codex_core::agent_graph_store_from_state_db; +use codex_core::init_state_db_from_config; use codex_core::thread_store_from_config; use codex_exec_server::CreateDirectoryOptions; use codex_exec_server::EnvironmentManager; @@ -235,6 +237,11 @@ async fn list_skills_skips_cwd_roots_when_environment_disabled() -> Result<()> { let mut config = load_default_config_for_test(&codex_home).await; config.cwd = AbsolutePathBuf::from_absolute_path_checked(cwd.path())?; + let state_db = init_state_db_from_config(&config) + .await + .expect("skills test requires state db"); + let thread_store = thread_store_from_config(&config, state_db.clone()); + let agent_graph_store = agent_graph_store_from_state_db(state_db.clone()); let thread_manager = ThreadManager::new( &config, codex_core::test_support::auth_manager_from_auth(CodexAuth::from_api_key("dummy")), @@ -246,8 +253,9 @@ async fn list_skills_skips_cwd_roots_when_environment_disabled() -> Result<()> { )?, )), /*analytics_events_client*/ None, - thread_store_from_config(&config, /*state_db*/ None), - /*state_db*/ None, + state_db, + thread_store, + agent_graph_store, ); let new_thread = thread_manager.start_thread(config.clone()).await?; let cwd = config.cwd.to_path_buf(); diff --git a/codex-rs/core/tests/suite/unstable_features_warning.rs b/codex-rs/core/tests/suite/unstable_features_warning.rs index 66a7366589..7a4924cdda 100644 --- a/codex-rs/core/tests/suite/unstable_features_warning.rs +++ b/codex-rs/core/tests/suite/unstable_features_warning.rs @@ -34,7 +34,8 @@ async fn emits_warning_when_unstable_features_enabled_via_config() { let thread_manager = codex_core::test_support::thread_manager_with_models_provider( CodexAuth::from_api_key("test"), config.model_provider.clone(), - ); + ) + .await; let auth_manager = codex_core::test_support::auth_manager_from_auth(CodexAuth::from_api_key("test")); @@ -81,7 +82,8 @@ async fn suppresses_warning_when_configured() { let thread_manager = codex_core::test_support::thread_manager_with_models_provider( CodexAuth::from_api_key("test"), config.model_provider.clone(), - ); + ) + .await; let auth_manager = codex_core::test_support::auth_manager_from_auth(CodexAuth::from_api_key("test")); diff --git a/codex-rs/mcp-server/src/lib.rs b/codex-rs/mcp-server/src/lib.rs index 75372c6e7c..ac764456f5 100644 --- a/codex-rs/mcp-server/src/lib.rs +++ b/codex-rs/mcp-server/src/lib.rs @@ -83,7 +83,6 @@ pub async fn run_main( std::io::Error::new(ErrorKind::InvalidData, format!("error loading config: {e}")) })?; set_default_client_residency_requirement(config.enforce_residency.value()); - let state_db = codex_core::init_state_db(&config).await; let otel = codex_core::otel_init::build_provider( &config, @@ -140,15 +139,18 @@ pub async fn run_main( // Task: process incoming messages. let processor_handle = tokio::spawn({ let outgoing_message_sender = OutgoingMessageSender::new(outgoing_tx); - let mut processor = MessageProcessor::new( + let processor = MessageProcessor::new( outgoing_message_sender, arg0_paths, Arc::new(config), environment_manager, - state_db, ) .await; async move { + let Some(mut processor) = processor else { + error!("failed to initialize MCP processor"); + return; + }; while let Some(msg) = incoming_rx.recv().await { match msg { JsonRpcMessage::Request(r) => processor.process_request(r).await, diff --git a/codex-rs/mcp-server/src/message_processor.rs b/codex-rs/mcp-server/src/message_processor.rs index 9ce46dfeda..2591891e4d 100644 --- a/codex-rs/mcp-server/src/message_processor.rs +++ b/codex-rs/mcp-server/src/message_processor.rs @@ -2,9 +2,10 @@ use std::collections::HashMap; use std::sync::Arc; use codex_arg0::Arg0DispatchPaths; -use codex_core::StateDbHandle; use codex_core::ThreadManager; +use codex_core::agent_graph_store_from_state_db; use codex_core::config::Config; +use codex_core::init_state_db_from_config; use codex_core::thread_store_from_config; use codex_exec_server::EnvironmentManager; use codex_login::AuthManager; @@ -54,30 +55,33 @@ impl MessageProcessor { arg0_paths: Arg0DispatchPaths, config: Arc, environment_manager: Arc, - state_db: Option, - ) -> Self { + ) -> Option { let outgoing = Arc::new(outgoing); let auth_manager = AuthManager::shared_from_config( config.as_ref(), /*enable_codex_api_key_env*/ false, ) .await; + let state_db = init_state_db_from_config(config.as_ref()).await?; + let thread_store = thread_store_from_config(config.as_ref(), state_db.clone()); + let agent_graph_store = agent_graph_store_from_state_db(state_db.clone()); let thread_manager = Arc::new(ThreadManager::new( config.as_ref(), auth_manager, SessionSource::Mcp, environment_manager, /*analytics_events_client*/ None, - thread_store_from_config(config.as_ref(), state_db.clone()), - state_db.clone(), + state_db, + thread_store, + agent_graph_store, )); - Self { + Some(Self { outgoing, initialized: false, arg0_paths, thread_manager, running_requests_id_to_codex_uuid: Arc::new(Mutex::new(HashMap::new())), - } + }) } pub(crate) async fn process_request(&mut self, request: JsonRpcRequest) { diff --git a/codex-rs/thread-manager-sample/src/main.rs b/codex-rs/thread-manager-sample/src/main.rs index 9d84ece674..a27c78d543 100644 --- a/codex-rs/thread-manager-sample/src/main.rs +++ b/codex-rs/thread-manager-sample/src/main.rs @@ -52,10 +52,11 @@ use codex_core_api::TuiNotificationSettings; use codex_core_api::UriBasedFileOpener; use codex_core_api::UserInput; use codex_core_api::WebSearchMode; +use codex_core_api::agent_graph_store_from_state_db; use codex_core_api::arg0_dispatch_or_else; use codex_core_api::built_in_model_providers; use codex_core_api::find_codex_home; -use codex_core_api::init_state_db; +use codex_core_api::init_state_db_from_config; use codex_core_api::item_event_to_server_notification; use codex_core_api::set_default_originator; use codex_core_api::thread_store_from_config; @@ -104,7 +105,6 @@ async fn run_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { }; let config = new_config(args.model, arg0_paths)?; - let state_db = init_state_db(&config).await; let auth_manager = AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false).await; @@ -112,7 +112,11 @@ async fn run_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { config.codex_self_exe.clone(), config.codex_linux_sandbox_exe.clone(), )?; + let Some(state_db) = init_state_db_from_config(&config).await else { + bail!("thread manager sample requires state db"); + }; let thread_store = thread_store_from_config(&config, state_db.clone()); + let agent_graph_store = agent_graph_store_from_state_db(state_db.clone()); let environment_manager = Arc::new(EnvironmentManager::new(EnvironmentManagerArgs::new(local_runtime_paths)).await); let thread_manager = ThreadManager::new( @@ -121,8 +125,9 @@ async fn run_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { SessionSource::Exec, environment_manager, /*analytics_events_client*/ None, - Arc::clone(&thread_store), state_db, + Arc::clone(&thread_store), + agent_graph_store, ); let NewThread { diff --git a/codex-rs/thread-store/src/local/archive_thread.rs b/codex-rs/thread-store/src/local/archive_thread.rs index 8fb214e98c..268a6ffc47 100644 --- a/codex-rs/thread-store/src/local/archive_thread.rs +++ b/codex-rs/thread-store/src/local/archive_thread.rs @@ -13,11 +13,11 @@ pub(super) async fn archive_thread( params: ArchiveThreadParams, ) -> ThreadStoreResult<()> { let thread_id = params.thread_id; - let state_db_ctx = store.state_db().await; + let state_db = store.state_db(); let rollout_path = find_thread_path_by_id_str( store.config.codex_home.as_path(), &thread_id.to_string(), - state_db_ctx.as_deref(), + Some(state_db.as_ref()), ) .await .map_err(|err| ThreadStoreError::InvalidRequest { @@ -52,11 +52,10 @@ pub(super) async fn archive_thread( } })?; - if let Some(ctx) = state_db_ctx { - let _ = ctx - .mark_archived(thread_id, archived_path.as_path(), Utc::now()) - .await; - } + let _ = store + .state_db() + .mark_archived(thread_id, archived_path.as_path(), Utc::now()) + .await; Ok(()) } @@ -75,13 +74,15 @@ mod tests { use crate::ThreadSortKey; use crate::ThreadStore; use crate::local::LocalThreadStore; + use crate::local::test_support::init_test_state_db; use crate::local::test_support::test_config; + use crate::local::test_support::test_store; use crate::local::test_support::write_session_file; #[tokio::test] async fn archive_thread_moves_rollout_to_archived_collection() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(201); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let active_path = @@ -127,21 +128,12 @@ mod tests { async fn archive_thread_updates_sqlite_metadata_when_present() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(202); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let active_path = write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file"); - let runtime = codex_state::StateRuntime::init( - home.path().to_path_buf(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); - runtime - .mark_backfill_complete(/*last_watermark*/ None) - .await - .expect("backfill should be complete"); let mut builder = codex_state::ThreadMetadataBuilder::new( thread_id, active_path.clone(), diff --git a/codex-rs/thread-store/src/local/create_thread.rs b/codex-rs/thread-store/src/local/create_thread.rs index e444e5c91d..cb0f64ea9c 100644 --- a/codex-rs/thread-store/src/local/create_thread.rs +++ b/codex-rs/thread-store/src/local/create_thread.rs @@ -22,12 +22,12 @@ pub(super) async fn create_thread( })?; let config = RolloutConfig { codex_home: store.config.codex_home.clone(), - sqlite_home: store.config.sqlite_home.clone(), + sqlite_home: store.sqlite_home(), cwd, model_provider_id: params.metadata.model_provider.clone(), generate_memories: matches!(params.metadata.memory_mode, ThreadMemoryMode::Enabled), }; - let state_db_ctx = store.state_db().await; + let state_db_ctx = Some(store.state_db()); let recorder = RolloutRecorder::new( &config, RolloutRecorderParams::new( diff --git a/codex-rs/thread-store/src/local/list_threads.rs b/codex-rs/thread-store/src/local/list_threads.rs index e470ad2be9..5cf1f27aa1 100644 --- a/codex-rs/thread-store/src/local/list_threads.rs +++ b/codex-rs/thread-store/src/local/list_threads.rs @@ -39,16 +39,16 @@ pub(super) async fn list_threads( SortDirection::Asc => codex_rollout::SortDirection::Asc, SortDirection::Desc => codex_rollout::SortDirection::Desc, }; - let state_db = store.state_db().await; let rollout_config = RolloutConfig { codex_home: store.config.codex_home.clone(), - sqlite_home: store.config.sqlite_home.clone(), + sqlite_home: store.sqlite_home(), cwd: store.config.codex_home.clone(), model_provider_id: store.config.default_model_provider_id.clone(), generate_memories: false, }; + let state_db_ctx = Some(store.state_db()); let page = list_rollout_threads( - state_db, + state_db_ctx, &rollout_config, store.config.default_model_provider_id.as_str(), ¶ms, @@ -80,14 +80,13 @@ pub(super) async fn list_threads( .map(|thread| thread.thread_id) .collect::>(); let mut names = HashMap::::with_capacity(thread_ids.len()); - if let Some(state_db_ctx) = store.state_db().await { - for &thread_id in &thread_ids { - let Ok(Some(metadata)) = state_db_ctx.get_thread(thread_id).await else { - continue; - }; - if let Some(title) = distinct_thread_metadata_title(&metadata) { - names.insert(thread_id, title); - } + let state_db_ctx = store.state_db(); + for &thread_id in &thread_ids { + let Ok(Some(metadata)) = state_db_ctx.get_thread(thread_id).await else { + continue; + }; + if let Some(title) = distinct_thread_metadata_title(&metadata) { + names.insert(thread_id, title); } } if names.len() < thread_ids.len() @@ -108,9 +107,9 @@ pub(super) async fn list_threads( } async fn list_rollout_threads( - state_db: Option, + state_db_ctx: Option, config: &RolloutConfig, - default_model_provider_id: &str, + default_model_provider: &str, params: &ListThreadsParams, cursor: Option<&codex_rollout::Cursor>, sort_key: codex_rollout::ThreadSortKey, @@ -118,7 +117,7 @@ async fn list_rollout_threads( ) -> ThreadStoreResult { let page = if params.use_state_db_only && params.archived { RolloutRecorder::list_archived_threads_from_state_db( - state_db, + state_db_ctx.clone(), config, params.page_size, cursor, @@ -127,13 +126,13 @@ async fn list_rollout_threads( params.allowed_sources.as_slice(), params.model_providers.as_deref(), params.cwd_filters.as_deref(), - default_model_provider_id, + default_model_provider, params.search_term.as_deref(), ) .await } else if params.use_state_db_only { RolloutRecorder::list_threads_from_state_db( - state_db, + state_db_ctx.clone(), config, params.page_size, cursor, @@ -142,13 +141,13 @@ async fn list_rollout_threads( params.allowed_sources.as_slice(), params.model_providers.as_deref(), params.cwd_filters.as_deref(), - default_model_provider_id, + default_model_provider, params.search_term.as_deref(), ) .await } else if params.archived { RolloutRecorder::list_archived_threads( - state_db, + state_db_ctx.clone(), config, params.page_size, cursor, @@ -157,13 +156,13 @@ async fn list_rollout_threads( params.allowed_sources.as_slice(), params.model_providers.as_deref(), params.cwd_filters.as_deref(), - default_model_provider_id, + default_model_provider, params.search_term.as_deref(), ) .await } else { RolloutRecorder::list_threads( - state_db, + state_db_ctx, config, params.page_size, cursor, @@ -172,7 +171,7 @@ async fn list_rollout_threads( params.allowed_sources.as_slice(), params.model_providers.as_deref(), params.cwd_filters.as_deref(), - default_model_provider_id, + default_model_provider, params.search_term.as_deref(), ) .await @@ -195,7 +194,9 @@ mod tests { use super::*; use crate::ThreadStore; use crate::local::LocalThreadStore; + use crate::local::test_support::init_test_state_db; use crate::local::test_support::test_config; + use crate::local::test_support::test_store; use crate::local::test_support::write_archived_session_file; use crate::local::test_support::write_session_file; use crate::local::test_support::write_session_file_with; @@ -203,7 +204,7 @@ mod tests { #[tokio::test] async fn list_threads_uses_default_provider_when_rollout_omits_provider() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; write_session_file_with( home.path(), home.path().join("sessions/2025/01/03"), @@ -238,22 +239,13 @@ mod tests { async fn list_threads_preserves_sqlite_title_search_results() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(103); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = home.path().join("rollout-title-search.jsonl"); fs::write(&rollout_path, "").expect("placeholder rollout file"); - let runtime = codex_state::StateRuntime::init( - home.path().to_path_buf(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); - runtime - .mark_backfill_complete(/*last_watermark*/ None) - .await - .expect("backfill should be complete"); let created_at = Utc::now(); let mut builder = codex_state::ThreadMetadataBuilder::new( thread_id, @@ -267,6 +259,10 @@ mod tests { let mut metadata = builder.build(config.default_model_provider_id.as_str()); metadata.title = "needle title".to_string(); metadata.first_user_message = Some("plain preview".to_string()); + runtime + .mark_backfill_complete(/*last_watermark*/ None) + .await + .expect("backfill should be complete"); runtime .upsert_thread(&metadata) .await @@ -303,7 +299,7 @@ mod tests { #[tokio::test] async fn list_threads_selects_active_or_archived_collection() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let active_uuid = Uuid::from_u128(105); let archived_uuid = Uuid::from_u128(106); write_session_file(home.path(), "2025-01-03T12-00-00", active_uuid) @@ -372,7 +368,7 @@ mod tests { async fn list_threads_returns_local_rollout_summary() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); - let store = LocalThreadStore::new(config, /*state_db*/ None); + let store = LocalThreadStore::new(config.clone(), init_test_state_db(&config).await); let uuid = Uuid::from_u128(101); let path = write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file"); @@ -411,7 +407,7 @@ mod tests { #[tokio::test] async fn list_threads_rejects_invalid_cursor() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let err = store .list_threads(ListThreadsParams { diff --git a/codex-rs/thread-store/src/local/live_writer.rs b/codex-rs/thread-store/src/local/live_writer.rs index 643207b59d..37d9758973 100644 --- a/codex-rs/thread-store/src/local/live_writer.rs +++ b/codex-rs/thread-store/src/local/live_writer.rs @@ -66,12 +66,12 @@ pub(super) async fn resume_thread( })?; let config = RolloutConfig { codex_home: store.config.codex_home.clone(), - sqlite_home: store.config.sqlite_home.clone(), + sqlite_home: store.sqlite_home(), cwd, model_provider_id: params.metadata.model_provider.clone(), generate_memories: matches!(params.metadata.memory_mode, ThreadMemoryMode::Enabled), }; - let state_db_ctx = store.state_db().await; + let state_db_ctx = Some(store.state_db()); let recorder = RolloutRecorder::new( &config, RolloutRecorderParams::resume( diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index 058c7b2309..e6ea63bba1 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -41,7 +41,7 @@ use crate::UpdateThreadMetadataParams; pub struct LocalThreadStore { pub(super) config: LocalThreadStoreConfig, live_recorders: Arc>>, - state_db: Option, + state_db: StateDbHandle, } /// Process-scoped configuration for local thread storage. @@ -51,7 +51,6 @@ pub struct LocalThreadStore { #[derive(Clone, Debug, PartialEq, Eq)] pub struct LocalThreadStoreConfig { pub codex_home: PathBuf, - pub sqlite_home: PathBuf, /// Provider used only when older local metadata does not contain one. pub default_model_provider_id: String, } @@ -60,7 +59,6 @@ impl LocalThreadStoreConfig { pub fn from_config(config: &impl codex_rollout::RolloutConfigView) -> Self { Self { codex_home: config.codex_home().to_path_buf(), - sqlite_home: config.sqlite_home().to_path_buf(), default_model_provider_id: config.model_provider_id().to_string(), } } @@ -75,8 +73,9 @@ impl std::fmt::Debug for LocalThreadStore { } impl LocalThreadStore { - /// Create a local store using an already initialized state DB handle. - pub fn new(config: LocalThreadStoreConfig, state_db: Option) -> Self { + /// Create a local store from process-scoped local storage configuration and + /// the caller-provided shared state DB handle. + pub fn new(config: LocalThreadStoreConfig, state_db: StateDbHandle) -> Self { Self { config, live_recorders: Arc::new(Mutex::new(HashMap::new())), @@ -85,10 +84,14 @@ impl LocalThreadStore { } /// Return the state DB handle used by local rollout writers. - pub async fn state_db(&self) -> Option { + pub fn state_db(&self) -> StateDbHandle { self.state_db.clone() } + pub(super) fn sqlite_home(&self) -> PathBuf { + self.state_db.codex_home().to_path_buf() + } + /// Read a local rollout-backed thread by path. pub async fn read_thread_by_rollout_path( &self, @@ -282,14 +285,16 @@ mod tests { use super::*; use crate::ThreadEventPersistenceMode; use crate::ThreadPersistenceMetadata; + use crate::local::test_support::init_test_state_db; use crate::local::test_support::test_config; + use crate::local::test_support::test_store; use crate::local::test_support::write_archived_session_file; use crate::local::test_support::write_session_file; #[tokio::test] async fn live_writer_lifecycle_writes_and_closes() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let thread_id = ThreadId::default(); store @@ -338,7 +343,7 @@ mod tests { #[tokio::test] async fn create_thread_rejects_missing_cwd() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let thread_id = ThreadId::default(); let mut params = create_thread_params(thread_id); params.metadata.cwd = None; @@ -358,7 +363,7 @@ mod tests { #[tokio::test] async fn discard_thread_drops_unmaterialized_live_writer() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let thread_id = ThreadId::default(); store @@ -396,8 +401,9 @@ mod tests { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); let thread_id = ThreadId::default(); + let state_db = init_test_state_db(&config).await; - let first_store = LocalThreadStore::new(config.clone(), /*state_db*/ None); + let first_store = LocalThreadStore::new(config.clone(), state_db.clone()); first_store .create_thread(create_thread_params(thread_id)) .await @@ -426,7 +432,7 @@ mod tests { .await .expect("shutdown initial writer"); - let resumed_store = LocalThreadStore::new(config, /*state_db*/ None); + let resumed_store = LocalThreadStore::new(config, state_db); resumed_store .resume_thread(ResumeThreadParams { thread_id, @@ -457,7 +463,7 @@ mod tests { #[tokio::test] async fn create_thread_rejects_duplicate_live_writer() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let thread_id = ThreadId::default(); store @@ -477,7 +483,7 @@ mod tests { #[tokio::test] async fn resume_thread_rejects_duplicate_live_writer() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let thread_id = ThreadId::default(); store @@ -506,7 +512,7 @@ mod tests { #[tokio::test] async fn resume_thread_rejects_missing_cwd() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = uuid::Uuid::from_u128(407); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = @@ -535,7 +541,7 @@ mod tests { async fn load_history_uses_live_writer_rollout_path() { let home = TempDir::new().expect("temp dir"); let external_home = TempDir::new().expect("external temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = uuid::Uuid::from_u128(404); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = write_session_file(external_home.path(), "2025-01-04T10-00-00", uuid) @@ -584,7 +590,7 @@ mod tests { async fn read_thread_uses_live_writer_rollout_path_for_external_resume() { let home = TempDir::new().expect("temp dir"); let external_home = TempDir::new().expect("external temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = uuid::Uuid::from_u128(406); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = write_session_file(external_home.path(), "2025-01-04T11-00-00", uuid) @@ -623,7 +629,7 @@ mod tests { #[tokio::test] async fn load_history_uses_live_writer_rollout_path_for_archived_source() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = uuid::Uuid::from_u128(405); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = write_archived_session_file(home.path(), "2025-01-04T10-30-00", uuid) @@ -691,7 +697,7 @@ mod tests { #[tokio::test] async fn read_thread_by_rollout_path_includes_history() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let thread_id = ThreadId::default(); store diff --git a/codex-rs/thread-store/src/local/read_thread.rs b/codex-rs/thread-store/src/local/read_thread.rs index c538d938a3..5bcf411beb 100644 --- a/codex-rs/thread-store/src/local/read_thread.rs +++ b/codex-rs/thread-store/src/local/read_thread.rs @@ -176,12 +176,12 @@ async fn resolve_rollout_path( return Ok(Some(path)); } - let state_db_ctx = store.state_db().await; + let state_db = store.state_db(); if include_archived { match find_thread_path_by_id_str( store.config.codex_home.as_path(), &thread_id.to_string(), - state_db_ctx.as_deref(), + Some(state_db.as_ref()), ) .await .map_err(|err| ThreadStoreError::InvalidRequest { @@ -191,7 +191,7 @@ async fn resolve_rollout_path( None => find_archived_thread_path_by_id_str( store.config.codex_home.as_path(), &thread_id.to_string(), - state_db_ctx.as_deref(), + Some(state_db.as_ref()), ) .await .map_err(|err| ThreadStoreError::InvalidRequest { @@ -202,7 +202,7 @@ async fn resolve_rollout_path( find_thread_path_by_id_str( store.config.codex_home.as_path(), &thread_id.to_string(), - state_db_ctx.as_deref(), + Some(state_db.as_ref()), ) .await .map_err(|err| ThreadStoreError::InvalidRequest { @@ -260,8 +260,7 @@ async fn read_sqlite_metadata( store: &LocalThreadStore, thread_id: codex_protocol::ThreadId, ) -> Option { - let runtime = store.state_db().await?; - runtime.get_thread(thread_id).await.ok().flatten() + store.state_db().get_thread(thread_id).await.ok().flatten() } async fn stored_thread_from_sqlite_metadata( @@ -412,7 +411,9 @@ mod tests { use super::*; use crate::ThreadStore; use crate::local::LocalThreadStore; + use crate::local::test_support::init_test_state_db; use crate::local::test_support::test_config; + use crate::local::test_support::test_store; use crate::local::test_support::write_archived_session_file; use crate::local::test_support::write_session_file; use crate::local::test_support::write_session_file_with_fork; @@ -420,7 +421,7 @@ mod tests { #[tokio::test] async fn read_thread_returns_active_rollout_summary() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(205); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let active_path = @@ -448,7 +449,7 @@ mod tests { #[tokio::test] async fn read_thread_returns_rollout_path_summary() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(211); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let active_path = @@ -479,17 +480,12 @@ mod tests { async fn read_thread_by_rollout_path_prefers_sqlite_git_info() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(223); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let active_path = write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file"); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); let mut builder = ThreadMetadataBuilder::new( thread_id, active_path.clone(), @@ -527,7 +523,7 @@ mod tests { #[tokio::test] async fn read_thread_returns_archived_rollout_when_requested() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(207); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let archived_path = write_archived_session_file(home.path(), "2025-01-03T12-00-00", uuid) @@ -568,7 +564,7 @@ mod tests { #[tokio::test] async fn read_thread_prefers_active_rollout_over_archived() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(208); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let active_path = @@ -593,7 +589,7 @@ mod tests { #[tokio::test] async fn read_thread_returns_forked_from_id() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(209); let parent_uuid = Uuid::from_u128(210); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); @@ -626,17 +622,12 @@ mod tests { async fn read_thread_applies_sqlite_thread_name() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(212); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file"); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); let mut builder = ThreadMetadataBuilder::new(thread_id, rollout_path, Utc::now(), SessionSource::Cli); builder.model_provider = Some(config.default_model_provider_id.clone()); @@ -666,13 +657,8 @@ mod tests { async fn read_thread_preserves_rollout_cwd_when_sqlite_metadata_exists() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(224); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let day_dir = home.path().join("sessions/2025/01/03"); @@ -741,7 +727,7 @@ mod tests { #[tokio::test] async fn read_thread_uses_legacy_thread_name_when_sqlite_title_is_missing() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(213); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file"); @@ -765,6 +751,8 @@ mod tests { async fn read_thread_uses_sqlite_metadata_for_rollout_without_user_preview() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(217); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let day_dir = home.path().join("sessions/2025/01/03"); @@ -786,13 +774,6 @@ mod tests { }); writeln!(file, "{meta}").expect("write session meta"); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); let mut builder = ThreadMetadataBuilder::new( thread_id, rollout_path.clone(), @@ -835,18 +816,13 @@ mod tests { let home = TempDir::new().expect("temp dir"); let external = TempDir::new().expect("external temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(220); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file"); let stale_path = external.path().join("missing-rollout.jsonl"); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); let mut builder = ThreadMetadataBuilder::new( thread_id, stale_path.clone(), @@ -884,6 +860,8 @@ mod tests { let home = TempDir::new().expect("temp dir"); let external = TempDir::new().expect("external temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(221); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = @@ -891,13 +869,6 @@ mod tests { let other_uuid = Uuid::from_u128(222); let stale_path = write_session_file(external.path(), "2025-01-04T12-00-00", other_uuid) .expect("other session file"); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); let mut builder = ThreadMetadataBuilder::new(thread_id, stale_path, Utc::now(), SessionSource::Cli); builder.model_provider = Some("wrong-sqlite-provider".to_string()); @@ -929,7 +900,7 @@ mod tests { #[tokio::test] async fn read_thread_uses_session_meta_for_rollout_without_user_preview_or_sqlite_metadata() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(218); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let day_dir = home.path().join("sessions/2025/01/03"); @@ -984,18 +955,13 @@ mod tests { let home = TempDir::new().expect("temp dir"); let external = TempDir::new().expect("external temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(214); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = external .path() .join(format!("rollout-2025-01-03T12-00-00-{uuid}.jsonl")); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); let mut builder = ThreadMetadataBuilder::new( thread_id, rollout_path.clone(), @@ -1042,20 +1008,15 @@ mod tests { let home = TempDir::new().expect("temp dir"); let external = TempDir::new().expect("external temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(216); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = external .path() .join(format!("rollout-2025-01-03T12-00-00-{uuid}.jsonl")); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); let mut builder = ThreadMetadataBuilder::new(thread_id, rollout_path, Utc::now(), SessionSource::Cli); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); builder.archived_at = Some(Utc::now()); let mut metadata = builder.build(config.default_model_provider_id.as_str()); metadata.first_user_message = Some("Archived SQLite preview".to_string()); @@ -1098,17 +1059,12 @@ mod tests { async fn read_thread_sqlite_fallback_loads_archived_history() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(219); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let archived_path = write_archived_session_file(home.path(), "2025-01-03T12-00-00", uuid) .expect("archived session file"); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); let mut builder = ThreadMetadataBuilder::new( thread_id, archived_path.clone(), @@ -1144,7 +1100,7 @@ mod tests { #[tokio::test] async fn read_thread_fails_without_rollout() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(206); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); diff --git a/codex-rs/thread-store/src/local/test_support.rs b/codex-rs/thread-store/src/local/test_support.rs index 98321880ff..597014efe9 100644 --- a/codex-rs/thread-store/src/local/test_support.rs +++ b/codex-rs/thread-store/src/local/test_support.rs @@ -4,18 +4,34 @@ use std::path::Path; use std::path::PathBuf; use codex_rollout::ARCHIVED_SESSIONS_SUBDIR; +use codex_rollout::StateDbHandle; use uuid::Uuid; +use super::LocalThreadStore; use super::LocalThreadStoreConfig; pub(super) fn test_config(codex_home: &Path) -> LocalThreadStoreConfig { LocalThreadStoreConfig { codex_home: codex_home.to_path_buf(), - sqlite_home: codex_home.to_path_buf(), default_model_provider_id: "test-provider".to_string(), } } +pub(super) async fn init_test_state_db(config: &LocalThreadStoreConfig) -> StateDbHandle { + codex_state::StateRuntime::init( + config.codex_home.clone(), + config.default_model_provider_id.clone(), + ) + .await + .expect("state db should initialize") +} + +pub(super) async fn test_store(codex_home: &Path) -> LocalThreadStore { + let config = test_config(codex_home); + let state_db = init_test_state_db(&config).await; + LocalThreadStore::new(config, state_db) +} + pub(super) fn write_session_file(root: &Path, ts: &str, uuid: Uuid) -> std::io::Result { write_session_file_with( root, diff --git a/codex-rs/thread-store/src/local/unarchive_thread.rs b/codex-rs/thread-store/src/local/unarchive_thread.rs index ad41db69ac..7ac3f01a0f 100644 --- a/codex-rs/thread-store/src/local/unarchive_thread.rs +++ b/codex-rs/thread-store/src/local/unarchive_thread.rs @@ -17,11 +17,11 @@ pub(super) async fn unarchive_thread( params: ArchiveThreadParams, ) -> ThreadStoreResult { let thread_id = params.thread_id; - let state_db_ctx = store.state_db().await; + let state_db = store.state_db(); let archived_path = find_archived_thread_path_by_id_str( store.config.codex_home.as_path(), &thread_id.to_string(), - state_db_ctx.as_deref(), + Some(state_db.as_ref()), ) .await .map_err(|err| ThreadStoreError::InvalidRequest { @@ -73,11 +73,10 @@ pub(super) async fn unarchive_thread( message: format!("failed to update unarchived thread timestamp: {err}"), })?; - if let Some(ctx) = state_db_ctx { - let _ = ctx - .mark_unarchived(thread_id, restored_path.as_path()) - .await; - } + let _ = store + .state_db() + .mark_unarchived(thread_id, restored_path.as_path()) + .await; let item = read_thread_item_from_rollout(restored_path.clone()) .await @@ -112,13 +111,15 @@ mod tests { use super::*; use crate::ThreadStore; use crate::local::LocalThreadStore; + use crate::local::test_support::init_test_state_db; use crate::local::test_support::test_config; + use crate::local::test_support::test_store; use crate::local::test_support::write_archived_session_file; #[tokio::test] async fn unarchive_thread_restores_rollout_and_returns_updated_thread() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(203); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let archived_path = write_archived_session_file(home.path(), "2025-01-03T13-00-00", uuid) @@ -149,21 +150,12 @@ mod tests { async fn unarchive_thread_updates_sqlite_metadata_when_present() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(204); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let archived_path = write_archived_session_file(home.path(), "2025-01-03T13-00-00", uuid) .expect("archived session file"); - let runtime = codex_state::StateRuntime::init( - home.path().to_path_buf(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); - runtime - .mark_backfill_complete(/*last_watermark*/ None) - .await - .expect("backfill should be complete"); let mut builder = codex_state::ThreadMetadataBuilder::new( thread_id, archived_path.clone(), diff --git a/codex-rs/thread-store/src/local/update_thread_metadata.rs b/codex-rs/thread-store/src/local/update_thread_metadata.rs index 98ee788662..2c70c815ed 100644 --- a/codex-rs/thread-store/src/local/update_thread_metadata.rs +++ b/codex-rs/thread-store/src/local/update_thread_metadata.rs @@ -59,9 +59,8 @@ pub(super) async fn update_thread_metadata( .await?; } - let state_db_ctx = store.state_db().await; codex_rollout::state_db::reconcile_rollout( - state_db_ctx.as_deref(), + Some(store.state_db()).as_deref(), resolved_rollout_path.path.as_path(), store.config.default_model_provider_id.as_str(), /*builder*/ None, @@ -73,11 +72,7 @@ pub(super) async fn update_thread_metadata( let resolved_git_info = match git_info { Some(git_info) => { - let Some(state_db) = store.state_db().await else { - return Err(ThreadStoreError::Internal { - message: format!("sqlite state db unavailable for thread {thread_id}"), - }); - }; + let state_db = store.state_db(); let metadata = state_db .get_thread(thread_id) @@ -157,11 +152,7 @@ async fn apply_thread_git_info( branch: &Option, origin_url: &Option, ) -> ThreadStoreResult<()> { - let Some(state_db) = store.state_db().await else { - return Err(ThreadStoreError::Internal { - message: format!("sqlite state db unavailable for thread {thread_id}"), - }); - }; + let state_db = store.state_db(); let updated = state_db .update_thread_git_info( thread_id, @@ -307,11 +298,11 @@ async fn resolve_rollout_path( return Ok(ResolvedRolloutPath { path, archived }); } - let state_db_ctx = store.state_db().await; + let state_db = store.state_db(); let active_path = find_thread_path_by_id_str( store.config.codex_home.as_path(), &thread_id.to_string(), - state_db_ctx.as_deref(), + Some(state_db.as_ref()), ) .await .map_err(|err| ThreadStoreError::InvalidRequest { @@ -331,7 +322,7 @@ async fn resolve_rollout_path( find_archived_thread_path_by_id_str( store.config.codex_home.as_path(), &thread_id.to_string(), - state_db_ctx.as_deref(), + Some(state_db.as_ref()), ) .await .map_err(|err| ThreadStoreError::InvalidRequest { @@ -366,14 +357,16 @@ mod tests { use crate::ThreadPersistenceMetadata; use crate::ThreadStore; use crate::local::LocalThreadStore; + use crate::local::test_support::init_test_state_db; use crate::local::test_support::test_config; + use crate::local::test_support::test_store; use crate::local::test_support::write_archived_session_file; use crate::local::test_support::write_session_file; #[tokio::test] async fn update_thread_metadata_sets_name_on_active_rollout_and_indexes_name() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(301); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let path = @@ -408,18 +401,12 @@ mod tests { async fn update_thread_metadata_sets_memory_mode_on_active_rollout() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(302); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let path = write_session_file(home.path(), "2025-01-03T14-30-00", uuid).expect("session file"); - let runtime = codex_state::StateRuntime::init( - home.path().to_path_buf(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); - let thread = store .update_thread_metadata(UpdateThreadMetadataParams { thread_id, @@ -452,13 +439,8 @@ mod tests { let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let path = write_session_file(home.path(), "2025-01-03T18-30-00", uuid).expect("session file"); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); store .update_thread_metadata(UpdateThreadMetadataParams { @@ -517,7 +499,7 @@ mod tests { async fn update_thread_metadata_uses_live_rollout_path_for_external_resume() { let home = TempDir::new().expect("temp dir"); let external_home = TempDir::new().expect("external temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(307); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let path = write_session_file(external_home.path(), "2025-01-03T14-45-00", uuid) @@ -558,13 +540,8 @@ mod tests { async fn update_thread_metadata_sets_git_info() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config, Some(runtime)); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config, runtime); let uuid = Uuid::from_u128(309); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); write_session_file(home.path(), "2025-01-03T17-00-00", uuid).expect("session file"); @@ -601,13 +578,8 @@ mod tests { async fn update_thread_metadata_partially_updates_git_info() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config, Some(runtime)); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config, runtime); let uuid = Uuid::from_u128(310); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); write_session_file(home.path(), "2025-01-03T17-30-00", uuid).expect("session file"); @@ -659,13 +631,8 @@ mod tests { async fn update_thread_metadata_clears_git_info_fields() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(311); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let path = @@ -829,7 +796,7 @@ mod tests { #[tokio::test] async fn update_thread_metadata_rejects_mismatched_session_meta_id() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let filename_uuid = Uuid::from_u128(303); let metadata_uuid = Uuid::from_u128(304); let thread_id = ThreadId::from_string(&filename_uuid.to_string()).expect("valid thread id"); @@ -861,7 +828,7 @@ mod tests { #[tokio::test] async fn update_thread_metadata_rejects_multi_field_patch_without_partial_write() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(305); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let path = @@ -896,21 +863,12 @@ mod tests { async fn update_thread_metadata_keeps_archived_thread_archived_in_sqlite() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(306); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let archived_path = write_archived_session_file(home.path(), "2025-01-03T16-00-00", uuid) .expect("archived session file"); - let runtime = codex_state::StateRuntime::init( - home.path().to_path_buf(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); - runtime - .mark_backfill_complete(/*last_watermark*/ None) - .await - .expect("backfill should be complete"); codex_rollout::state_db::reconcile_rollout( Some(runtime.as_ref()), archived_path.as_path(), @@ -959,21 +917,12 @@ mod tests { async fn update_thread_metadata_keeps_live_archived_thread_archived_in_sqlite() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(308); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let archived_path = write_archived_session_file(home.path(), "2025-01-03T16-30-00", uuid) .expect("archived session file"); - let runtime = codex_state::StateRuntime::init( - home.path().to_path_buf(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); - runtime - .mark_backfill_complete(/*last_watermark*/ None) - .await - .expect("backfill should be complete"); codex_rollout::state_db::reconcile_rollout( Some(runtime.as_ref()), archived_path.as_path(), diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index 5622c59f65..b2e92a19b4 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -881,38 +881,40 @@ pub async fn run_main( AppServerTarget::Remote { .. } => state_db::get_state_db(&config).await, }; - let effective_toml = config.config_layer_stack.effective_config(); - match effective_toml.try_into() { - Ok(config_toml) => { - match crate::legacy_core::personality_migration::maybe_migrate_personality( - &config.codex_home, - &config_toml, - state_db.clone(), - ) - .await - { - Ok( - crate::legacy_core::personality_migration::PersonalityMigrationStatus::Applied, - ) => { - config = load_config_or_exit( - cli_kv_overrides.clone(), - overrides.clone(), - cloud_requirements.clone(), - ) - .await; - } - Ok( - crate::legacy_core::personality_migration::PersonalityMigrationStatus::SkippedMarker - | crate::legacy_core::personality_migration::PersonalityMigrationStatus::SkippedExplicitPersonality - | crate::legacy_core::personality_migration::PersonalityMigrationStatus::SkippedNoSessions, - ) => {} - Err(err) => { - tracing::warn!(error = %err, "failed to run personality migration"); + if let Some(state_db) = state_db.clone() { + let effective_toml = config.config_layer_stack.effective_config(); + match effective_toml.try_into() { + Ok(config_toml) => { + match crate::legacy_core::personality_migration::maybe_migrate_personality( + &config.codex_home, + &config_toml, + state_db, + ) + .await + { + Ok( + crate::legacy_core::personality_migration::PersonalityMigrationStatus::Applied, + ) => { + config = load_config_or_exit( + cli_kv_overrides.clone(), + overrides.clone(), + cloud_requirements.clone(), + ) + .await; + } + Ok( + crate::legacy_core::personality_migration::PersonalityMigrationStatus::SkippedMarker + | crate::legacy_core::personality_migration::PersonalityMigrationStatus::SkippedExplicitPersonality + | crate::legacy_core::personality_migration::PersonalityMigrationStatus::SkippedNoSessions, + ) => {} + Err(err) => { + tracing::warn!(error = %err, "failed to run personality migration"); + } } } - } - Err(err) => { - tracing::warn!(error = %err, "failed to deserialize config for personality migration"); + Err(err) => { + tracing::warn!(error = %err, "failed to deserialize config for personality migration"); + } } } diff --git a/codex-rs/tui/src/resume_picker.rs b/codex-rs/tui/src/resume_picker.rs index 16c3d19071..e06ccfd118 100644 --- a/codex-rs/tui/src/resume_picker.rs +++ b/codex-rs/tui/src/resume_picker.rs @@ -5733,6 +5733,7 @@ session_picker_view = "dense" name: None, turns: vec![codex_app_server_protocol::Turn { id: String::from("turn-1"), + items_view: codex_app_server_protocol::TurnItemsView::Full, items: vec![ ThreadItem::UserMessage { id: String::from("user-1"), @@ -5798,6 +5799,7 @@ session_picker_view = "dense" name: None, turns: vec![codex_app_server_protocol::Turn { id: String::from("turn-1"), + items_view: codex_app_server_protocol::TurnItemsView::Full, items: vec![ThreadItem::Reasoning { id: String::from("reasoning-1"), summary: Vec::new(), @@ -5853,6 +5855,7 @@ session_picker_view = "dense" name: None, turns: vec![codex_app_server_protocol::Turn { id: String::from("turn-1"), + items_view: codex_app_server_protocol::TurnItemsView::Full, items: vec![ThreadItem::Reasoning { id: String::from("reasoning-1"), summary: vec![String::from("public summary")],