From 33b182d805f9ca3dfadf0897cca7511a738997fd Mon Sep 17 00:00:00 2001 From: Owen Lin Date: Thu, 7 May 2026 17:54:29 -0700 Subject: [PATCH] Add SQLite init and fallback telemetry --- codex-rs/agent-graph-store/src/local.rs | 11 +- codex-rs/app-server-client/src/lib.rs | 15 +- .../src/transport/remote_control/enroll.rs | 10 +- .../src/transport/remote_control/tests.rs | 10 +- .../src/transport/remote_control/websocket.rs | 10 +- codex-rs/app-server/src/in_process.rs | 12 +- codex-rs/app-server/src/lib.rs | 38 +- codex-rs/app-server/src/mcp_refresh.rs | 7 +- codex-rs/app-server/src/message_processor.rs | 13 +- .../src/message_processor_tracing_tests.rs | 2 +- codex-rs/app-server/src/request_processors.rs | 1 + .../thread_goal_processor.rs | 18 +- .../tests/suite/conversation_summary.rs | 2 +- .../app-server/tests/suite/v2/mcp_resource.rs | 2 +- .../app-server/tests/suite/v2/memory_reset.rs | 7 +- .../tests/suite/v2/remote_thread_store.rs | 2 +- .../tests/suite/v2/thread_archive.rs | 76 ++-- .../app-server/tests/suite/v2/thread_list.rs | 20 +- .../tests/suite/v2/thread_memory_mode_set.rs | 7 +- .../tests/suite/v2/thread_metadata_update.rs | 7 +- .../app-server/tests/suite/v2/thread_read.rs | 6 +- .../tests/suite/v2/thread_resume.rs | 8 +- .../tests/suite/v2/thread_unarchive.rs | 15 +- codex-rs/cli/src/main.rs | 12 +- codex-rs/cli/tests/debug_clear_memories.rs | 8 +- codex-rs/core/src/agent/control_tests.rs | 6 +- codex-rs/core/src/codex_delegate.rs | 1 + codex-rs/core/src/lib.rs | 2 + codex-rs/core/src/mcp_tool_call.rs | 2 +- codex-rs/core/src/otel_init.rs | 32 +- codex-rs/core/src/personality_migration.rs | 10 +- .../core/src/personality_migration_tests.rs | 15 +- codex-rs/core/src/prompt_debug.rs | 8 +- codex-rs/core/src/session/mod.rs | 15 +- codex-rs/core/src/session/session.rs | 14 +- codex-rs/core/src/session/tests.rs | 45 ++- .../core/src/session/tests/guardian_tests.rs | 3 +- codex-rs/core/src/shell_snapshot.rs | 22 +- codex-rs/core/src/shell_snapshot_tests.rs | 12 +- codex-rs/core/src/state/service.rs | 4 +- codex-rs/core/src/state_db_bridge.rs | 9 +- codex-rs/core/src/stream_events_utils.rs | 6 +- codex-rs/core/src/thread_manager.rs | 19 +- codex-rs/core/src/thread_manager_tests.rs | 69 ++-- .../src/tools/handlers/multi_agents_tests.rs | 36 +- codex-rs/core/tests/common/test_codex.rs | 11 +- codex-rs/core/tests/suite/client.rs | 4 +- .../core/tests/suite/personality_migration.rs | 36 +- .../core/tests/suite/prompt_debug_tests.rs | 3 +- .../core/tests/suite/rollout_list_find.rs | 68 ++-- codex-rs/exec/src/lib.rs | 20 +- codex-rs/mcp-server/src/lib.rs | 48 ++- codex-rs/mcp-server/src/message_processor.rs | 8 +- codex-rs/memories/write/src/startup_tests.rs | 8 +- codex-rs/otel/src/metrics/client.rs | 11 +- codex-rs/otel/src/metrics/mod.rs | 4 + codex-rs/otel/src/metrics/names.rs | 1 + codex-rs/otel/src/metrics/process.rs | 27 ++ codex-rs/otel/src/metrics/tags.rs | 26 ++ codex-rs/rollout/src/lib.rs | 2 + codex-rs/rollout/src/list.rs | 109 +++-- codex-rs/rollout/src/metadata_tests.rs | 30 +- codex-rs/rollout/src/recorder.rs | 53 ++- codex-rs/rollout/src/recorder_tests.rs | 164 ++++++-- codex-rs/rollout/src/session_index.rs | 3 +- codex-rs/rollout/src/session_index_tests.rs | 9 +- codex-rs/rollout/src/sqlite_metrics.rs | 43 ++ codex-rs/rollout/src/state_db.rs | 203 ++++++++-- codex-rs/rollout/src/state_db_tests.rs | 18 +- codex-rs/rollout/src/tests.rs | 71 +++- codex-rs/state/src/bin/logs_client.rs | 3 +- codex-rs/state/src/lib.rs | 27 +- codex-rs/state/src/log_db.rs | 40 +- codex-rs/state/src/runtime.rs | 143 ++++++- codex-rs/state/src/runtime/agent_jobs.rs | 14 +- codex-rs/state/src/runtime/backfill.rs | 31 +- codex-rs/state/src/runtime/goals.rs | 10 +- codex-rs/state/src/runtime/logs.rs | 220 +++++++--- codex-rs/state/src/runtime/memories.rs | 380 ++++++++++++------ codex-rs/state/src/runtime/remote_control.rs | 20 +- codex-rs/state/src/runtime/threads.rs | 140 +++++-- codex-rs/state/src/telemetry.rs | 263 ++++++++++++ codex-rs/thread-manager-sample/src/main.rs | 8 +- .../thread-store/src/local/archive_thread.rs | 17 +- .../thread-store/src/local/list_threads.rs | 39 +- codex-rs/thread-store/src/local/mod.rs | 66 ++- .../thread-store/src/local/read_thread.rs | 102 ++++- .../src/local/unarchive_thread.rs | 17 +- .../src/local/update_thread_metadata.rs | 62 ++- codex-rs/tui/src/app.rs | 8 +- codex-rs/tui/src/app/event_dispatch.rs | 2 +- codex-rs/tui/src/app/session_lifecycle.rs | 2 +- codex-rs/tui/src/app/test_support.rs | 2 +- codex-rs/tui/src/app/tests.rs | 5 +- codex-rs/tui/src/app/thread_routing.rs | 8 +- codex-rs/tui/src/app/thread_session_state.rs | 8 +- codex-rs/tui/src/lib.rs | 110 ++--- codex-rs/tui/src/onboarding/auth.rs | 2 +- 98 files changed, 2505 insertions(+), 873 deletions(-) create mode 100644 codex-rs/otel/src/metrics/process.rs create mode 100644 codex-rs/rollout/src/sqlite_metrics.rs create mode 100644 codex-rs/state/src/telemetry.rs diff --git a/codex-rs/agent-graph-store/src/local.rs b/codex-rs/agent-graph-store/src/local.rs index f45874855c..a446b4289a 100644 --- a/codex-rs/agent-graph-store/src/local.rs +++ b/codex-rs/agent-graph-store/src/local.rs @@ -125,10 +125,13 @@ mod tests { async fn state_runtime() -> TestRuntime { let codex_home = TempDir::new().expect("tempdir should be created"); - let state_db = - StateRuntime::init(codex_home.path().to_path_buf(), "test-provider".to_string()) - .await - .expect("state db should initialize"); + let state_db = StateRuntime::init( + codex_home.path().to_path_buf(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("state db should initialize"); TestRuntime { state_db, _codex_home: codex_home, diff --git a/codex-rs/app-server-client/src/lib.rs b/codex-rs/app-server-client/src/lib.rs index 3b386ff3ce..d735b91e27 100644 --- a/codex-rs/app-server-client/src/lib.rs +++ b/codex-rs/app-server-client/src/lib.rs @@ -29,6 +29,7 @@ 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::StateDbAccess; pub use codex_app_server::in_process::StateDbHandle; use codex_app_server_protocol::ClientInfo; use codex_app_server_protocol::ClientNotification; @@ -340,8 +341,8 @@ pub struct InProcessClientStartArgs { 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 the embedded app-server. - pub state_db: Option, + /// Process-wide SQLite access shared with the embedded app-server. + pub state_db_access: StateDbAccess, /// Environment manager used by core execution and filesystem operations. pub environment_manager: Arc, /// Startup warnings emitted after initialize succeeds. @@ -404,7 +405,7 @@ impl InProcessClientStartArgs { thread_config_loader, feedback: self.feedback, log_db: self.log_db, - state_db: self.state_db, + state_db_access: self.state_db_access, environment_manager: self.environment_manager, config_warnings: self.config_warnings, session_source: self.session_source, @@ -1017,7 +1018,7 @@ mod tests { ) -> TestClient { let codex_home = TempDir::new().expect("temp dir"); let config = Arc::new(build_test_config_for_codex_home(codex_home.path()).await); - let state_db = init_state_db(config.as_ref()) + let state_db = init_state_db(config.as_ref(), /*metrics*/ None) .await .expect("state db should initialize for in-process test"); let client = InProcessAppServerClient::start(InProcessClientStartArgs { @@ -1028,7 +1029,7 @@ mod tests { cloud_requirements: CloudRequirementsLoader::default(), feedback: CodexFeedback::new(), log_db: None, - state_db: Some(state_db), + state_db_access: StateDbAccess::new(Some(state_db)), environment_manager: Arc::new(EnvironmentManager::default_for_tests()), config_warnings: Vec::new(), session_source, @@ -2112,7 +2113,7 @@ mod tests { cloud_requirements: CloudRequirementsLoader::default(), feedback: CodexFeedback::new(), log_db: None, - state_db: None, + state_db_access: StateDbAccess::none(), environment_manager: environment_manager.clone(), config_warnings: Vec::new(), session_source: SessionSource::Exec, @@ -2152,7 +2153,7 @@ mod tests { cloud_requirements: CloudRequirementsLoader::default(), feedback: CodexFeedback::new(), log_db: None, - state_db: None, + state_db_access: StateDbAccess::none(), environment_manager: Arc::new(EnvironmentManager::default_for_tests()), config_warnings: Vec::new(), session_source: SessionSource::Exec, diff --git a/codex-rs/app-server-transport/src/transport/remote_control/enroll.rs b/codex-rs/app-server-transport/src/transport/remote_control/enroll.rs index 60dfc3d845..1d4957768e 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/enroll.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/enroll.rs @@ -277,9 +277,13 @@ mod tests { use tokio::time::timeout; async fn remote_control_state_runtime(codex_home: &TempDir) -> Arc { - StateRuntime::init(codex_home.path().to_path_buf(), "test-provider".to_string()) - .await - .expect("state runtime should initialize") + StateRuntime::init( + codex_home.path().to_path_buf(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("state runtime should initialize") } #[tokio::test] diff --git a/codex-rs/app-server-transport/src/transport/remote_control/tests.rs b/codex-rs/app-server-transport/src/transport/remote_control/tests.rs index 88b7798827..4c7ee75065 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/tests.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/tests.rs @@ -109,9 +109,13 @@ fn remote_control_auth_dot_json(account_id: Option<&str>) -> AuthDotJson { } async fn remote_control_state_runtime(codex_home: &TempDir) -> Arc { - StateRuntime::init(codex_home.path().to_path_buf(), "test-provider".to_string()) - .await - .expect("state runtime should initialize") + StateRuntime::init( + codex_home.path().to_path_buf(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("state runtime should initialize") } fn remote_control_url_for_listener(listener: &TcpListener) -> String { diff --git a/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs b/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs index 472639bc68..1fca441964 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs @@ -1286,9 +1286,13 @@ mod tests { } async fn remote_control_state_runtime(codex_home: &TempDir) -> Arc { - StateRuntime::init(codex_home.path().to_path_buf(), "test-provider".to_string()) - .await - .expect("state runtime should initialize") + StateRuntime::init( + codex_home.path().to_path_buf(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("state runtime should initialize") } fn remote_control_auth_manager() -> Arc { diff --git a/codex-rs/app-server/src/in_process.rs b/codex-rs/app-server/src/in_process.rs index d812888e62..3df93be0de 100644 --- a/codex-rs/app-server/src/in_process.rs +++ b/codex-rs/app-server/src/in_process.rs @@ -86,7 +86,9 @@ use codex_exec_server::EnvironmentManager; use codex_feedback::CodexFeedback; use codex_login::AuthManager; use codex_protocol::protocol::SessionSource; +pub use codex_rollout::StateDbAccess; pub use codex_rollout::StateDbHandle; +pub use codex_state::DbMetricsRecorderHandle; pub use codex_state::log_db::LogDbLayer; use tokio::sync::mpsc; use tokio::sync::oneshot; @@ -127,8 +129,8 @@ 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. - pub state_db: Option, + /// Process-wide SQLite access shared with embedded app-server consumers. + pub state_db_access: StateDbAccess, /// Environment manager used by core execution and filesystem operations. pub environment_manager: Arc, /// Startup warnings emitted after initialize succeeds. @@ -424,7 +426,7 @@ async fn start_uninitialized(args: InProcessStartArgs) -> IoResult 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 = codex_rollout::state_db::try_init(config.as_ref(), /*metrics*/ None) .await .expect("state db should initialize for in-process test"); let args = InProcessStartArgs { @@ -769,7 +771,7 @@ mod tests { thread_config_loader: Arc::new(codex_config::NoopThreadConfigLoader), feedback: CodexFeedback::new(), log_db: None, - state_db: Some(state_db), + state_db_access: StateDbAccess::new(Some(state_db)), environment_manager: Arc::new(EnvironmentManager::default_for_tests()), config_warnings: Vec::new(), session_source, diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index af95800b24..0bf8a48c76 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -107,6 +107,7 @@ pub use crate::transport::auth::AppServerWebsocketAuthSettings; pub use crate::transport::auth::WebsocketAuthCliMode; const LOG_FORMAT_ENV_VAR: &str = "LOG_FORMAT"; +const OTEL_SERVICE_NAME: &str = "codex-app-server"; #[derive(Clone, Copy, Debug, Eq, PartialEq)] enum LogFormat { @@ -492,9 +493,27 @@ pub async fn run_main_with_transport_options( } }; - let state_db_result = rollout_state_db::try_init(&config).await; + let otel = codex_core::otel_init::build_provider( + &config, + env!("CARGO_PKG_VERSION"), + Some(OTEL_SERVICE_NAME), + default_analytics_enabled, + ) + .map_err(|e| { + std::io::Error::new( + ErrorKind::InvalidData, + format!("error loading otel config: {e}"), + ) + })?; + codex_core::otel_init::record_process_start(otel.as_ref(), OTEL_SERVICE_NAME); + let state_db_metrics = + codex_core::otel_init::sqlite_metrics_recorder(otel.as_ref(), OTEL_SERVICE_NAME); + + let state_db_result = rollout_state_db::try_init(&config, state_db_metrics.clone()).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_access = codex_rollout::StateDbAccess::new(state_db.clone()) + .with_standalone_metrics(state_db_metrics); if should_run_personality_migration { let effective_toml = config.config_layer_stack.effective_config(); @@ -503,7 +522,7 @@ pub async fn run_main_with_transport_options( match codex_core::personality_migration::maybe_migrate_personality( &config.codex_home, &config_toml, - state_db.clone(), + state_db_access.clone(), ) .await { @@ -571,19 +590,6 @@ pub async fn run_main_with_transport_options( let feedback = CodexFeedback::new(); - let otel = codex_core::otel_init::build_provider( - &config, - env!("CARGO_PKG_VERSION"), - Some("codex-app-server"), - default_analytics_enabled, - ) - .map_err(|e| { - std::io::Error::new( - ErrorKind::InvalidData, - format!("error loading otel config: {e}"), - ) - })?; - // Install a simple subscriber so `tracing` output is visible. Users can // control the log level with `RUST_LOG` and switch to JSON logs with // `LOG_FORMAT=json`. @@ -778,7 +784,7 @@ pub async fn run_main_with_transport_options( environment_manager, feedback: feedback.clone(), log_db, - state_db: state_db.clone(), + state_db_access, config_warnings, session_source, auth_manager, diff --git a/codex-rs/app-server/src/mcp_refresh.rs b/codex-rs/app-server/src/mcp_refresh.rs index a327b4b125..258820d6a5 100644 --- a/codex-rs/app-server/src/mcp_refresh.rs +++ b/codex-rs/app-server/src/mcp_refresh.rs @@ -174,10 +174,11 @@ mod tests { .await?; let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("dummy")); - let state_db = init_state_db(&good_config) + let state_db = init_state_db(&good_config, /*metrics*/ None) .await .expect("refresh tests require state db"); - let thread_store = thread_store_from_config(&good_config, Some(state_db.clone())); + let state_db_access = codex_core::StateDbAccess::new(Some(state_db.clone())); + let thread_store = thread_store_from_config(&good_config, state_db_access.clone()); let thread_manager = Arc::new(ThreadManager::new( &good_config, auth_manager, @@ -185,7 +186,7 @@ mod tests { Arc::new(EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, thread_store, - Some(state_db.clone()), + state_db_access, "11111111-1111-4111-8111-111111111111".to_string(), /*attestation_provider*/ None, )); diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 062127d5a7..1525c7fca3 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -73,7 +73,7 @@ use codex_login::auth::ExternalAuthTokens; use codex_protocol::ThreadId; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::W3cTraceContext; -use codex_rollout::StateDbHandle; +use codex_rollout::StateDbAccess; use codex_state::log_db::LogDbLayer; use tokio::sync::Mutex; use tokio::sync::Semaphore; @@ -255,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_access: StateDbAccess, pub(crate) config_warnings: Vec, pub(crate) session_source: SessionSource, pub(crate) auth_manager: Arc, @@ -278,7 +278,7 @@ impl MessageProcessor { environment_manager, feedback, log_db, - state_db, + state_db_access, config_warnings, session_source, auth_manager, @@ -294,7 +294,8 @@ impl MessageProcessor { // The thread store is intentionally process-scoped. Config reloads can // affect per-thread behavior, but they must not move newly started, // resumed, or forked threads to a different persistence backend/root. - let thread_store = thread_store_from_config(config.as_ref(), state_db.clone()); + let state_db = state_db_access.state_db(); + let thread_store = thread_store_from_config(config.as_ref(), state_db_access.clone()); let thread_manager = Arc::new(ThreadManager::new( config.as_ref(), auth_manager.clone(), @@ -302,7 +303,7 @@ impl MessageProcessor { environment_manager, Some(analytics_events_client.clone()), Arc::clone(&thread_store), - state_db.clone(), + state_db_access.clone(), installation_id, Some(app_server_attestation_provider( outgoing.clone(), @@ -387,7 +388,7 @@ impl MessageProcessor { outgoing.clone(), Arc::clone(&config), thread_state_manager.clone(), - state_db.clone(), + state_db_access, ); let thread_processor = ThreadRequestProcessor::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 516e042301..43b2051436 100644 --- a/codex-rs/app-server/src/message_processor_tracing_tests.rs +++ b/codex-rs/app-server/src/message_processor_tracing_tests.rs @@ -258,7 +258,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_access: codex_rollout::StateDbAccess::none(), config_warnings: Vec::new(), session_source: SessionSource::VSCode, auth_manager, diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index 5bfe97d305..c7d28aac2b 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -262,6 +262,7 @@ use codex_core::NewThread; #[cfg(test)] use codex_core::SessionMeta; use codex_core::StartThreadOptions; +use codex_core::StateDbAccess; use codex_core::SteerInputError; use codex_core::ThreadConfigSnapshot; use codex_core::ThreadManager; 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..451772c600 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_access: StateDbAccess, } impl ThreadGoalRequestProcessor { @@ -16,14 +16,14 @@ impl ThreadGoalRequestProcessor { outgoing: Arc, config: Arc, thread_state_manager: ThreadStateManager, - state_db: Option, + state_db_access: StateDbAccess, ) -> Self { Self { thread_manager, outgoing, config, thread_state_manager, - state_db, + state_db_access, } } @@ -81,7 +81,7 @@ impl ThreadGoalRequestProcessor { if let Some(state_db) = thread.state_db() { Some(state_db) } else { - self.state_db.clone() + self.state_db_access.state_db() } } else { None @@ -110,7 +110,7 @@ impl ThreadGoalRequestProcessor { None => find_thread_path_by_id_str( &self.config.codex_home, &thread_id.to_string(), - self.state_db.as_deref(), + &self.state_db_access, ) .await .map_err(|err| { @@ -275,7 +275,7 @@ impl ThreadGoalRequestProcessor { None => find_thread_path_by_id_str( &self.config.codex_home, &thread_id.to_string(), - self.state_db.as_deref(), + &self.state_db_access, ) .await .map_err(|err| { @@ -339,7 +339,7 @@ impl ThreadGoalRequestProcessor { find_thread_path_by_id_str( &self.config.codex_home, &thread_id.to_string(), - self.state_db.as_deref(), + &self.state_db_access, ) .await .map_err(|err| { @@ -348,8 +348,8 @@ impl ThreadGoalRequestProcessor { .ok_or_else(|| invalid_request(format!("thread not found: {thread_id}")))?; } - self.state_db - .clone() + self.state_db_access + .state_db() .ok_or_else(|| internal_error("sqlite state db unavailable for thread goals")) } diff --git a/codex-rs/app-server/tests/suite/conversation_summary.rs b/codex-rs/app-server/tests/suite/conversation_summary.rs index 754d1f9467..135287847e 100644 --- a/codex-rs/app-server/tests/suite/conversation_summary.rs +++ b/codex-rs/app-server/tests/suite/conversation_summary.rs @@ -153,7 +153,7 @@ async fn get_conversation_summary_by_thread_id_reads_pathless_store_thread() -> thread_config_loader: Arc::new(codex_config::NoopThreadConfigLoader), feedback: CodexFeedback::new(), log_db: None, - state_db: None, + state_db_access: in_process::StateDbAccess::none(), environment_manager: Arc::new(EnvironmentManager::default_for_tests()), config_warnings: Vec::new(), session_source: SessionSource::Cli, diff --git a/codex-rs/app-server/tests/suite/v2/mcp_resource.rs b/codex-rs/app-server/tests/suite/v2/mcp_resource.rs index a51f4bbd4e..69e3b50107 100644 --- a/codex-rs/app-server/tests/suite/v2/mcp_resource.rs +++ b/codex-rs/app-server/tests/suite/v2/mcp_resource.rs @@ -204,7 +204,7 @@ async fn mcp_resource_read_returns_error_for_unknown_thread() -> Result<()> { thread_config_loader: Arc::new(codex_config::NoopThreadConfigLoader), feedback: CodexFeedback::new(), log_db: None, - state_db: None, + state_db_access: in_process::StateDbAccess::none(), environment_manager: Arc::new(EnvironmentManager::default_for_tests()), config_warnings: Vec::new(), session_source: SessionSource::Cli, diff --git a/codex-rs/app-server/tests/suite/v2/memory_reset.rs b/codex-rs/app-server/tests/suite/v2/memory_reset.rs index 3c7ae38671..097247ff00 100644 --- a/codex-rs/app-server/tests/suite/v2/memory_reset.rs +++ b/codex-rs/app-server/tests/suite/v2/memory_reset.rs @@ -113,7 +113,12 @@ async fn seed_stage1_output(state_db: &Arc, codex_home: &Path) -> } async fn init_state_db(codex_home: &Path) -> Result> { - let state_db = StateRuntime::init(codex_home.to_path_buf(), "mock_provider".into()).await?; + let state_db = StateRuntime::init( + codex_home.to_path_buf(), + "mock_provider".into(), + /*metrics*/ None, + ) + .await?; state_db .mark_backfill_complete(/*last_watermark*/ None) .await?; 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 e5c0b2c53f..5ae6a68749 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 @@ -79,7 +79,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_access: in_process::StateDbAccess::none(), environment_manager: Arc::new(EnvironmentManager::default_for_tests()), config_warnings: Vec::new(), session_source: SessionSource::Cli, diff --git a/codex-rs/app-server/tests/suite/v2/thread_archive.rs b/codex-rs/app-server/tests/suite/v2/thread_archive.rs index b441a23cb6..50380d1f1a 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_archive.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_archive.rs @@ -63,9 +63,13 @@ async fn thread_archive_requires_materialized_rollout() -> Result<()> { rollout_path.display() ); assert!( - find_thread_path_by_id_str(codex_home.path(), &thread.id, /*state_db_ctx*/ None) - .await? - .is_none(), + find_thread_path_by_id_str( + codex_home.path(), + &thread.id, + &codex_core::StateDbAccess::none() + ) + .await? + .is_none(), "thread id should not be discoverable before rollout materialization" ); @@ -118,10 +122,13 @@ async fn thread_archive_requires_materialized_rollout() -> Result<()> { rollout_path.display() ); - let discovered_path = - find_thread_path_by_id_str(codex_home.path(), &thread.id, /*state_db_ctx*/ None) - .await? - .expect("expected rollout path for thread id to exist after materialization"); + let discovered_path = find_thread_path_by_id_str( + codex_home.path(), + &thread.id, + &codex_core::StateDbAccess::none(), + ) + .await? + .expect("expected rollout path for thread id to exist after materialization"); assert_paths_match_on_disk(&discovered_path, &rollout_path)?; let archive_id = mcp @@ -200,8 +207,12 @@ async fn thread_archive_archives_spawned_descendants() -> Result<()> { let parent_thread_id = ThreadId::from_string(&parent_id)?; let child_thread_id = ThreadId::from_string(&child_id)?; let grandchild_thread_id = ThreadId::from_string(&grandchild_id)?; - let state_db = - StateRuntime::init(codex_home.path().to_path_buf(), "mock_provider".into()).await?; + let state_db = StateRuntime::init( + codex_home.path().to_path_buf(), + "mock_provider".into(), + /*metrics*/ None, + ) + .await?; state_db .mark_backfill_complete(/*last_watermark*/ None) .await?; @@ -256,7 +267,7 @@ async fn thread_archive_archives_spawned_descendants() -> Result<()> { find_thread_path_by_id_str( codex_home.path(), &thread_id.to_string(), - /*state_db_ctx*/ None, + &codex_core::StateDbAccess::none(), ) .await? .is_none(), @@ -266,7 +277,7 @@ async fn thread_archive_archives_spawned_descendants() -> Result<()> { find_archived_thread_path_by_id_str( codex_home.path(), &thread_id.to_string(), - /*state_db_ctx*/ None, + &codex_core::StateDbAccess::none(), ) .await? .is_some(), @@ -311,8 +322,12 @@ async fn thread_archive_succeeds_when_descendant_archive_fails() -> Result<()> { let parent_thread_id = ThreadId::from_string(&parent_id)?; let child_thread_id = ThreadId::from_string(&child_id)?; let grandchild_thread_id = ThreadId::from_string(&grandchild_id)?; - let state_db = - StateRuntime::init(codex_home.path().to_path_buf(), "mock_provider".into()).await?; + let state_db = StateRuntime::init( + codex_home.path().to_path_buf(), + "mock_provider".into(), + /*metrics*/ None, + ) + .await?; state_db .mark_backfill_complete(/*last_watermark*/ None) .await?; @@ -331,10 +346,13 @@ async fn thread_archive_succeeds_when_descendant_archive_fails() -> Result<()> { ) .await?; - let child_rollout_path = - find_thread_path_by_id_str(codex_home.path(), &child_id, /*state_db_ctx*/ None) - .await? - .expect("child rollout path"); + let child_rollout_path = find_thread_path_by_id_str( + codex_home.path(), + &child_id, + &codex_core::StateDbAccess::none(), + ) + .await? + .expect("child rollout path"); let archived_child_path = codex_home .path() .join(ARCHIVED_SESSIONS_SUBDIR) @@ -394,7 +412,7 @@ async fn thread_archive_succeeds_when_descendant_archive_fails() -> Result<()> { find_thread_path_by_id_str( codex_home.path(), &thread_id.to_string(), - /*state_db_ctx*/ None, + &codex_core::StateDbAccess::none(), ) .await? .is_none(), @@ -404,7 +422,7 @@ async fn thread_archive_succeeds_when_descendant_archive_fails() -> Result<()> { find_archived_thread_path_by_id_str( codex_home.path(), &thread_id.to_string(), - /*state_db_ctx*/ None, + &codex_core::StateDbAccess::none(), ) .await? .is_some(), @@ -432,8 +450,12 @@ async fn thread_archive_succeeds_when_spawned_descendant_is_missing() -> Result< let parent_thread_id = ThreadId::from_string(&parent_id)?; let missing_child_thread_id = ThreadId::from_string("00000000-0000-0000-0000-000000000901")?; - let state_db = - StateRuntime::init(codex_home.path().to_path_buf(), "mock_provider".into()).await?; + let state_db = StateRuntime::init( + codex_home.path().to_path_buf(), + "mock_provider".into(), + /*metrics*/ None, + ) + .await?; state_db .mark_backfill_complete(/*last_watermark*/ None) .await?; @@ -473,16 +495,20 @@ async fn thread_archive_succeeds_when_spawned_descendant_is_missing() -> Result< assert_eq!(archived_notification.thread_id, parent_id); assert!( - find_thread_path_by_id_str(codex_home.path(), &parent_id, /*state_db_ctx*/ None) - .await? - .is_none(), + find_thread_path_by_id_str( + codex_home.path(), + &parent_id, + &codex_core::StateDbAccess::none() + ) + .await? + .is_none(), "parent should be archived even when a descendant is missing" ); assert!( find_archived_thread_path_by_id_str( codex_home.path(), &parent_id, - /*state_db_ctx*/ None, + &codex_core::StateDbAccess::none(), ) .await? .is_some(), diff --git a/codex-rs/app-server/tests/suite/v2/thread_list.rs b/codex-rs/app-server/tests/suite/v2/thread_list.rs index 80254d8f47..68edbb9641 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_list.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_list.rs @@ -600,9 +600,12 @@ sqlite = true // `thread/list` applies `search_term` on the sqlite fast path. This test creates // rollouts manually, so mark the DB backfill complete and then run an unsearched // list large enough to repair every rollout the searched list should find. - let state_db = - codex_state::StateRuntime::init(codex_home.path().to_path_buf(), "mock_provider".into()) - .await?; + let state_db = codex_state::StateRuntime::init( + codex_home.path().to_path_buf(), + "mock_provider".into(), + /*metrics*/ None, + ) + .await?; state_db .mark_backfill_complete(/*last_watermark*/ None) .await?; @@ -614,7 +617,7 @@ sqlite = true generate_memories: false, }; let repaired_page = codex_core::RolloutRecorder::list_threads( - Some(state_db.clone()), + codex_core::StateDbAccess::new(Some(state_db.clone())), &rollout_config, /*page_size*/ 10, /*cursor*/ None, @@ -683,9 +686,12 @@ sqlite = true Some("mock_provider"), /*git_info*/ None, )?; - let state_db = - codex_state::StateRuntime::init(codex_home.path().to_path_buf(), "mock_provider".into()) - .await?; + let state_db = codex_state::StateRuntime::init( + codex_home.path().to_path_buf(), + "mock_provider".into(), + /*metrics*/ None, + ) + .await?; state_db .mark_backfill_complete(/*last_watermark*/ None) .await?; diff --git a/codex-rs/app-server/tests/suite/v2/thread_memory_mode_set.rs b/codex-rs/app-server/tests/suite/v2/thread_memory_mode_set.rs index bf9bba7b2f..bb97650ad2 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_memory_mode_set.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_memory_mode_set.rs @@ -103,7 +103,12 @@ async fn thread_memory_mode_set_updates_stored_thread_state() -> Result<()> { } async fn init_state_db(codex_home: &Path) -> Result> { - let state_db = StateRuntime::init(codex_home.to_path_buf(), "mock_provider".into()).await?; + let state_db = StateRuntime::init( + codex_home.to_path_buf(), + "mock_provider".into(), + /*metrics*/ None, + ) + .await?; state_db .mark_backfill_complete(/*last_watermark*/ None) .await?; diff --git a/codex-rs/app-server/tests/suite/v2/thread_metadata_update.rs b/codex-rs/app-server/tests/suite/v2/thread_metadata_update.rs index c78e9b8152..73eca2de02 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_metadata_update.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_metadata_update.rs @@ -486,7 +486,12 @@ async fn thread_metadata_update_can_clear_stored_git_fields() -> Result<()> { } async fn init_state_db(codex_home: &Path) -> Result> { - let state_db = StateRuntime::init(codex_home.to_path_buf(), "mock_provider".into()).await?; + let state_db = StateRuntime::init( + codex_home.to_path_buf(), + "mock_provider".into(), + /*metrics*/ None, + ) + .await?; state_db .mark_backfill_complete(/*last_watermark*/ None) .await?; diff --git a/codex-rs/app-server/tests/suite/v2/thread_read.rs b/codex-rs/app-server/tests/suite/v2/thread_read.rs index 52420c0c80..161a68c919 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_read.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_read.rs @@ -379,7 +379,7 @@ async fn thread_turns_list_reads_store_history_without_rollout_path() -> Result< thread_config_loader: Arc::new(codex_config::NoopThreadConfigLoader), feedback: CodexFeedback::new(), log_db: None, - state_db: None, + state_db_access: in_process::StateDbAccess::none(), environment_manager: Arc::new(EnvironmentManager::default_for_tests()), config_warnings: Vec::new(), session_source: SessionSource::Cli.into(), @@ -444,7 +444,7 @@ async fn thread_read_loaded_include_turns_reads_store_history_without_rollout_pa thread_config_loader: Arc::new(codex_config::NoopThreadConfigLoader), feedback: CodexFeedback::new(), log_db: None, - state_db: None, + state_db_access: in_process::StateDbAccess::none(), environment_manager: Arc::new(EnvironmentManager::default_for_tests()), config_warnings: Vec::new(), session_source: SessionSource::Cli.into(), @@ -529,7 +529,7 @@ async fn thread_list_includes_store_thread_without_rollout_path() -> Result<()> thread_config_loader: Arc::new(codex_config::NoopThreadConfigLoader), feedback: CodexFeedback::new(), log_db: None, - state_db: None, + state_db_access: in_process::StateDbAccess::none(), environment_manager: Arc::new(EnvironmentManager::default_for_tests()), config_warnings: Vec::new(), session_source: SessionSource::Cli.into(), diff --git a/codex-rs/app-server/tests/suite/v2/thread_resume.rs b/codex-rs/app-server/tests/suite/v2/thread_resume.rs index 2b0eafd00a..6acd4ab93e 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_resume.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_resume.rs @@ -1223,8 +1223,12 @@ stream_max_retries = 0 .join("\n") + "\n", )?; - let state_db = - StateRuntime::init(codex_home.path().to_path_buf(), "mock_provider".into()).await?; + let state_db = StateRuntime::init( + codex_home.path().to_path_buf(), + "mock_provider".into(), + /*metrics*/ None, + ) + .await?; state_db .mark_backfill_complete(/*last_watermark*/ None) .await?; diff --git a/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs b/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs index 5b421dcec5..5f31b88d64 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs @@ -100,10 +100,13 @@ async fn thread_unarchive_moves_rollout_back_into_sessions_directory() -> Result ) .await??; - let found_rollout_path = - find_thread_path_by_id_str(codex_home.path(), &thread.id, /*state_db_ctx*/ None) - .await? - .expect("expected rollout path for thread id to exist"); + let found_rollout_path = find_thread_path_by_id_str( + codex_home.path(), + &thread.id, + &codex_core::StateDbAccess::none(), + ) + .await? + .expect("expected rollout path for thread id to exist"); assert_paths_match_on_disk(&found_rollout_path, &rollout_path)?; let archive_id = mcp @@ -121,7 +124,7 @@ async fn thread_unarchive_moves_rollout_back_into_sessions_directory() -> Result let archived_path = find_archived_thread_path_by_id_str( codex_home.path(), &thread.id, - /*state_db_ctx*/ None, + &codex_core::StateDbAccess::none(), ) .await? .expect("expected archived rollout path for thread id to exist"); @@ -249,7 +252,7 @@ async fn thread_unarchive_preserves_pathless_store_metadata() -> Result<()> { thread_config_loader: Arc::new(codex_config::NoopThreadConfigLoader), feedback: CodexFeedback::new(), log_db: None, - state_db: None, + state_db_access: in_process::StateDbAccess::none(), 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 dbe6b7605c..c465cc47f4 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -1440,7 +1440,8 @@ 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, codex_core::StateDbAccess::none()).await?; println!("{}", serde_json::to_string_pretty(&prompt_input)?); Ok(()) @@ -1487,9 +1488,12 @@ async fn run_debug_clear_memories_command( let state_path = state_db_path(config.sqlite_home.as_path()); let mut cleared_state_db = false; if tokio::fs::try_exists(&state_path).await? { - let state_db = - StateRuntime::init(config.sqlite_home.clone(), config.model_provider_id.clone()) - .await?; + let state_db = StateRuntime::init( + config.sqlite_home.clone(), + config.model_provider_id.clone(), + /*metrics*/ None, + ) + .await?; state_db.clear_memory_data().await?; cleared_state_db = true; } diff --git a/codex-rs/cli/tests/debug_clear_memories.rs b/codex-rs/cli/tests/debug_clear_memories.rs index 9d5e114dbf..c2dbd398f9 100644 --- a/codex-rs/cli/tests/debug_clear_memories.rs +++ b/codex-rs/cli/tests/debug_clear_memories.rs @@ -16,8 +16,12 @@ fn codex_command(codex_home: &Path) -> Result { #[tokio::test] async fn debug_clear_memories_resets_state_and_removes_memory_dir() -> Result<()> { let codex_home = TempDir::new()?; - let runtime = - StateRuntime::init(codex_home.path().to_path_buf(), "test-provider".to_string()).await?; + let runtime = StateRuntime::init( + codex_home.path().to_path_buf(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await?; drop(runtime); let thread_id = "00000000-0000-0000-0000-000000000123"; diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index b95aad4489..4946faf4b5 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -94,7 +94,7 @@ struct AgentControlHarness { impl AgentControlHarness { async fn new() -> Self { let (home, config) = test_config().await; - let state_db = init_state_db(&config).await; + let state_db = init_state_db(&config, /*metrics*/ None).await; let manager = ThreadManager::with_models_provider_home_and_state_for_tests( CodexAuth::from_api_key("dummy"), config.model_provider.clone(), @@ -1543,7 +1543,7 @@ 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 state_db = init_state_db(&config, /*metrics*/ None).await; let manager = ThreadManager::with_models_provider_home_and_state_for_tests( CodexAuth::from_api_key("dummy"), config.model_provider.clone(), @@ -1706,7 +1706,7 @@ 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(), + crate::StateDbAccess::new(harness.state_db.clone()), ); store .archive_thread(ArchiveThreadParams { diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 2503764904..991b197b74 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -100,6 +100,7 @@ pub(crate) async fn run_codex_thread_interactive( analytics_events_client: Some(parent_session.services.analytics_events_client.clone()), thread_store: Arc::clone(&parent_session.services.thread_store), attestation_provider: parent_session.services.attestation_provider.clone(), + state_db_access: parent_session.services.state_db_access.clone(), })) .or_cancel(&cancel_token) .await??; diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 57ac9b4a59..2a4412bd34 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -140,6 +140,8 @@ pub mod shell; pub(crate) mod shell_snapshot; pub mod spawn; pub(crate) mod state_db_bridge; +pub use state_db_bridge::DbMetricsRecorderHandle; +pub use state_db_bridge::StateDbAccess; pub use state_db_bridge::StateDbHandle; pub use state_db_bridge::init_state_db; mod thread_rollout_truncation; diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index 5dd16acb9a..3bb4d2898e 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -770,7 +770,7 @@ async fn maybe_mark_thread_memory_mode_polluted( return; } state_db::mark_thread_memory_mode_polluted( - sess.services.state_db.as_deref(), + sess.services.state_db_access.as_deref(), sess.conversation_id, "mcp_tool_call", ) diff --git a/codex-rs/core/src/otel_init.rs b/codex-rs/core/src/otel_init.rs index 0cd1f06994..b212953b4a 100644 --- a/codex-rs/core/src/otel_init.rs +++ b/codex-rs/core/src/otel_init.rs @@ -2,7 +2,6 @@ use crate::config::Config; use codex_config::types::OtelExporterKind as Kind; use codex_config::types::OtelHttpProtocol as Protocol; use codex_features::Feature; -use codex_login::default_client::originator; use codex_otel::OtelExporter; use codex_otel::OtelHttpProtocol; use codex_otel::OtelProvider; @@ -76,7 +75,7 @@ pub fn build_provider( OtelExporter::None }; - let originator = originator(); + let originator = codex_login::default_client::originator(); let service_name = service_name_override.unwrap_or(originator.value.as_str()); let runtime_metrics = config.features.enabled(Feature::RuntimeMetrics); @@ -94,6 +93,35 @@ pub fn build_provider( }) } +/// Record the process-start denominator using the caller's explicit process identity. +pub fn record_process_start(provider: Option<&OtelProvider>, originator: &str) { + if let Some(provider) = provider + && let Some(metrics) = provider.metrics() + { + let _ = codex_otel::record_process_start_once(metrics, originator); + } +} + +/// Build a SQLite metrics recorder from the caller's explicit process identity. +pub fn sqlite_metrics_recorder( + provider: Option<&OtelProvider>, + originator: &str, +) -> Option { + provider.and_then(|provider| { + provider.metrics().map(|metrics| { + codex_rollout::state_db::otel_db_metrics_recorder(metrics.clone(), originator) + }) + }) +} + +/// Build a SQLite metrics recorder from the process-global OTEL metrics client. +pub fn sqlite_metrics_recorder_from_global( + originator: &str, +) -> Option { + codex_otel::global() + .map(|metrics| codex_rollout::state_db::otel_db_metrics_recorder(metrics, originator)) +} + /// Filter predicate for exporting only Codex-owned events via OTEL. /// Keeps events that originated from codex_otel module pub fn codex_export_filter(meta: &tracing::Metadata<'_>) -> bool { diff --git a/codex-rs/core/src/personality_migration.rs b/codex-rs/core/src/personality_migration.rs index 975aecd4af..33070ade68 100644 --- a/codex-rs/core/src/personality_migration.rs +++ b/codex-rs/core/src/personality_migration.rs @@ -1,7 +1,7 @@ use crate::config::edit::ConfigEditsBuilder; use codex_config::config_toml::ConfigToml; use codex_protocol::config_types::Personality; -use codex_rollout::state_db::StateDbHandle; +use codex_rollout::StateDbAccess; use codex_thread_store::ListThreadsParams; use codex_thread_store::LocalThreadStore; use codex_thread_store::LocalThreadStoreConfig; @@ -25,7 +25,7 @@ pub enum PersonalityMigrationStatus { pub async fn maybe_migrate_personality( codex_home: &Path, config_toml: &ConfigToml, - state_db: Option, + state_db_access: StateDbAccess, ) -> io::Result { let marker_path = codex_home.join(PERSONALITY_MIGRATION_FILENAME); if tokio::fs::try_exists(&marker_path).await? { @@ -45,7 +45,7 @@ pub async fn maybe_migrate_personality( .or_else(|| config_toml.model_provider.clone()) .unwrap_or_else(|| "openai".to_string()); - if !has_recorded_sessions(codex_home, model_provider_id.as_str(), state_db).await? { + if !has_recorded_sessions(codex_home, model_provider_id.as_str(), state_db_access).await? { create_marker(&marker_path).await?; return Ok(PersonalityMigrationStatus::SkippedNoSessions); } @@ -65,7 +65,7 @@ pub async fn maybe_migrate_personality( async fn has_recorded_sessions( codex_home: &Path, default_provider: &str, - state_db: Option, + state_db_access: StateDbAccess, ) -> io::Result { let store = LocalThreadStore::new( LocalThreadStoreConfig { @@ -73,7 +73,7 @@ async fn has_recorded_sessions( sqlite_home: codex_home.to_path_buf(), default_model_provider_id: default_provider.to_string(), }, - state_db, + state_db_access, ); 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 699e06fe67..779835c647 100644 --- a/codex-rs/core/src/personality_migration_tests.rs +++ b/codex-rs/core/src/personality_migration_tests.rs @@ -88,7 +88,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 status = + maybe_migrate_personality(temp.path(), &config_toml, StateDbAccess::none()).await?; assert_eq!(status, PersonalityMigrationStatus::Applied); assert!(temp.path().join(PERSONALITY_MIGRATION_FILENAME).exists()); @@ -104,7 +105,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 status = + maybe_migrate_personality(temp.path(), &config_toml, StateDbAccess::none()).await?; assert_eq!(status, PersonalityMigrationStatus::Applied); assert!(temp.path().join(PERSONALITY_MIGRATION_FILENAME).exists()); @@ -120,7 +122,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 status = + maybe_migrate_personality(temp.path(), &config_toml, StateDbAccess::none()).await?; assert_eq!(status, PersonalityMigrationStatus::SkippedMarker); assert!(!temp.path().join("config.toml").exists()); @@ -137,7 +140,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 status = + maybe_migrate_personality(temp.path(), &config_toml, StateDbAccess::none()).await?; assert_eq!( status, @@ -154,7 +158,8 @@ 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 status = + maybe_migrate_personality(temp.path(), &config_toml, StateDbAccess::none()).await?; assert_eq!(status, PersonalityMigrationStatus::SkippedNoSessions); assert!(temp.path().join(PERSONALITY_MIGRATION_FILENAME).exists()); diff --git a/codex-rs/core/src/prompt_debug.rs b/codex-rs/core/src/prompt_debug.rs index 688ce27508..70fcbb35c8 100644 --- a/codex-rs/core/src/prompt_debug.rs +++ b/codex-rs/core/src/prompt_debug.rs @@ -17,7 +17,7 @@ use crate::resolve_installation_id; 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::state_db_bridge::StateDbAccess; use crate::thread_manager::ThreadManager; use crate::thread_manager::thread_store_from_config; @@ -26,7 +26,7 @@ use crate::thread_manager::thread_store_from_config; pub async fn build_prompt_input( mut config: Config, input: Vec, - state_db: Option, + state_db_access: StateDbAccess, ) -> CodexResult> { config.ephemeral = true; @@ -38,7 +38,7 @@ pub async fn build_prompt_input( config.codex_linux_sandbox_exe.clone(), )?; - let thread_store = thread_store_from_config(&config, state_db.clone()); + let thread_store = thread_store_from_config(&config, state_db_access.clone()); let installation_id = resolve_installation_id(&config.codex_home).await?; let thread_manager = ThreadManager::new( &config, @@ -51,7 +51,7 @@ pub async fn build_prompt_input( ), /*analytics_events_client*/ None, thread_store, - state_db.clone(), + state_db_access, installation_id, /*attestation_provider*/ None, ); diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index c6acf16f67..e0095c2fe3 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::ReadThreadParams; use codex_thread_store::ResumeThreadParams; use codex_thread_store::ThreadEventPersistenceMode; @@ -273,6 +272,7 @@ use crate::SkillLoadOutcome; #[cfg(test)] use crate::SkillMetadata; use crate::SkillsManager; +use crate::StateDbAccess; use crate::agents_md::AgentsMdManager; use crate::context::UserInstructions; use crate::exec_policy::ExecPolicyUpdateError; @@ -414,6 +414,7 @@ pub(crate) struct CodexSpawnArgs { pub(crate) analytics_events_client: Option, pub(crate) thread_store: Arc, pub(crate) attestation_provider: Option>, + pub(crate) state_db_access: StateDbAccess, } pub(crate) const INITIAL_SUBMIT_ID: &str = ""; @@ -474,6 +475,7 @@ impl Codex { analytics_events_client, thread_store, attestation_provider, + state_db_access, } = args; let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY); let (tx_event, rx_event) = async_channel::unbounded(); @@ -563,12 +565,8 @@ impl Codex { 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 + state_db_access.state_db() }; state_db::get_dynamic_tools(state_db_ctx.as_deref(), thread_id, "codex_spawn") .await @@ -660,6 +658,7 @@ impl Codex { thread_store, parent_rollout_thread_trace, attestation_provider, + state_db_access, ) .await .map_err(|e| { @@ -1049,7 +1048,7 @@ impl Session { } pub(crate) fn state_db(&self) -> Option { - self.services.state_db.clone() + self.services.state_db_access.state_db() } pub(crate) fn live_thread_for_persistence( @@ -1328,7 +1327,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.services.state_db_access.clone(), ); } diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 1a790314d5..bb39c9e4be 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1,4 +1,5 @@ use super::*; +use crate::StateDbAccess; use crate::goals::GoalRuntimeState; use codex_protocol::SessionId; use codex_protocol::config_types::ServiceTier; @@ -371,6 +372,7 @@ impl Session { thread_store: Arc, parent_rollout_thread_trace: ThreadTraceContext, attestation_provider: Option>, + state_db_access: StateDbAccess, ) -> anyhow::Result> { debug!( "Configuring session: model={}; provider={:?}", @@ -471,12 +473,8 @@ impl Session { 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 + state_db_access.state_db() } } .instrument(info_span!( @@ -507,7 +505,7 @@ impl Session { )); // Join all independent futures. - let (thread_persistence_result, state_db_ctx, (auth, mcp_servers, auth_statuses)) = + let (thread_persistence_result, _state_db_ctx, (auth, mcp_servers, auth_statuses)) = tokio::join!(thread_persistence_fut, state_db_fut, auth_and_mcp_fut); let mut live_thread_init = @@ -709,7 +707,7 @@ impl Session { session_configuration.cwd.clone(), &mut default_shell, session_telemetry.clone(), - state_db_ctx.clone(), + state_db_access.clone(), ) } } else { @@ -850,7 +848,7 @@ impl Session { agent_control, network_proxy, network_approval: Arc::clone(&network_approval), - state_db: state_db_ctx.clone(), + state_db_access: state_db_access.clone(), live_thread: live_thread_init.as_ref().cloned(), thread_store: Arc::clone(&thread_store), attestation_provider: attestation_provider.clone(), diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 5090fc3f07..39c0aa4f77 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -3731,10 +3731,11 @@ async fn session_new_fails_when_zsh_fork_enabled_without_zsh_path() { /*analytics_events_client*/ None, Arc::new(codex_thread_store::LocalThreadStore::new( codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()), - /*state_db*/ None, + crate::StateDbAccess::none(), )), codex_rollout_trace::ThreadTraceContext::disabled(), /*attestation_provider*/ None, + crate::StateDbAccess::none(), ) .await; @@ -3877,11 +3878,11 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { agent_control, network_proxy: None, network_approval: Arc::clone(&network_approval), - state_db: None, + state_db_access: crate::StateDbAccess::none(), live_thread: None, thread_store: Arc::new(codex_thread_store::LocalThreadStore::new( codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()), - /*state_db*/ None, + crate::StateDbAccess::none(), )), attestation_provider: None, model_client: ModelClient::new( @@ -4070,10 +4071,11 @@ async fn make_session_with_config_and_rx( /*analytics_events_client*/ None, Arc::new(codex_thread_store::LocalThreadStore::new( codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()), - /*state_db*/ None, + crate::StateDbAccess::none(), )), codex_rollout_trace::ThreadTraceContext::disabled(), /*attestation_provider*/ None, + crate::StateDbAccess::none(), ) .await?; @@ -4153,6 +4155,16 @@ async fn make_session_with_history_source_and_agent_control_and_rx( /*bundled_skills_enabled*/ true, )); + let state_db_access = crate::StateDbAccess::new(Some( + codex_state::StateRuntime::init( + config.sqlite_home.clone(), + config.model_provider_id.clone(), + /*metrics*/ None, + ) + .await + .expect("state db should initialize"), + )); + let session = Session::new( session_configuration, Arc::clone(&config), @@ -4173,17 +4185,11 @@ async fn make_session_with_history_source_and_agent_control_and_rx( /*analytics_events_client*/ None, Arc::new(codex_thread_store::LocalThreadStore::new( codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()), - Some( - codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.model_provider_id.clone(), - ) - .await - .expect("state db should initialize"), - ), + state_db_access.clone(), )), codex_rollout_trace::ThreadTraceContext::disabled(), /*attestation_provider*/ None, + state_db_access, ) .await?; @@ -5463,6 +5469,7 @@ where codex_state::StateRuntime::init( config.sqlite_home.clone(), config.model_provider_id.clone(), + /*metrics*/ None, ) .await .expect("goal tests should initialize sqlite state db"), @@ -5470,6 +5477,7 @@ where } else { None }; + let state_db_access = crate::StateDbAccess::new(state_db.clone()); let config = Arc::new(config); let thread_id = ThreadId::default(); let auth_manager = AuthManager::from_auth_for_testing(auth); @@ -5596,11 +5604,11 @@ where agent_control, network_proxy: None, network_approval: Arc::clone(&network_approval), - state_db: state_db.clone(), + state_db_access: state_db_access.clone(), live_thread: None, thread_store: Arc::new(codex_thread_store::LocalThreadStore::new( codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()), - state_db, + state_db_access, )), attestation_provider: None, model_client: ModelClient::new( @@ -7846,8 +7854,12 @@ async fn goal_test_state_db(sess: &Session) -> anyhow::Result a let state_db = codex_state::StateRuntime::init( test.config.sqlite_home.clone(), test.config.model_provider_id.clone(), + /*metrics*/ None, ) .await?; let persisted_goal = state_db diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index 37a5e6bb09..b0cb253125 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -731,7 +731,7 @@ 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, + crate::StateDbAccess::none(), )); let CodexSpawnOk { codex, .. } = Codex::spawn(CodexSpawnArgs { @@ -764,6 +764,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { analytics_events_client: None, thread_store, attestation_provider: None, + state_db_access: crate::StateDbAccess::none(), }) .await .expect("spawn guardian subagent"); diff --git a/codex-rs/core/src/shell_snapshot.rs b/codex-rs/core/src/shell_snapshot.rs index b328a977d7..ad92db5d08 100644 --- a/codex-rs/core/src/shell_snapshot.rs +++ b/codex-rs/core/src/shell_snapshot.rs @@ -5,7 +5,7 @@ use std::sync::Arc; use std::time::Duration; use std::time::SystemTime; -use crate::StateDbHandle; +use crate::StateDbAccess; use crate::rollout::list::find_thread_path_by_id_str; use crate::shell::Shell; use crate::shell::ShellType; @@ -42,7 +42,7 @@ impl ShellSnapshot { session_cwd: AbsolutePathBuf, shell: &mut Shell, session_telemetry: SessionTelemetry, - state_db: Option, + state_db_access: StateDbAccess, ) -> watch::Sender>> { let (shell_snapshot_tx, shell_snapshot_rx) = watch::channel(None); shell.shell_snapshot = shell_snapshot_rx; @@ -54,7 +54,7 @@ impl ShellSnapshot { shell.clone(), shell_snapshot_tx.clone(), session_telemetry, - state_db, + state_db_access, ); shell_snapshot_tx @@ -67,7 +67,7 @@ impl ShellSnapshot { shell: Shell, shell_snapshot_tx: watch::Sender>>, session_telemetry: SessionTelemetry, - state_db: Option, + state_db_access: StateDbAccess, ) { Self::spawn_snapshot_task( codex_home, @@ -76,7 +76,7 @@ impl ShellSnapshot { shell, shell_snapshot_tx, session_telemetry, - state_db, + state_db_access, ); } @@ -87,7 +87,7 @@ impl ShellSnapshot { snapshot_shell: Shell, shell_snapshot_tx: watch::Sender>>, session_telemetry: SessionTelemetry, - state_db: Option, + state_db_access: StateDbAccess, ) { let snapshot_span = info_span!("shell_snapshot", thread_id = %session_id); tokio::spawn( @@ -98,7 +98,7 @@ impl ShellSnapshot { session_id, &session_cwd, &snapshot_shell, - state_db, + state_db_access, ) .await .map(Arc::new); @@ -121,7 +121,7 @@ impl ShellSnapshot { session_id: ThreadId, session_cwd: &AbsolutePathBuf, shell: &Shell, - state_db: Option, + state_db_access: StateDbAccess, ) -> std::result::Result { // File to store the snapshot let extension = match shell.shell_type { @@ -144,7 +144,7 @@ impl ShellSnapshot { let cleanup_session_id = session_id; tokio::spawn(async move { if let Err(err) = - cleanup_stale_snapshots(&codex_home, cleanup_session_id, state_db).await + cleanup_stale_snapshots(&codex_home, cleanup_session_id, state_db_access).await { tracing::warn!("Failed to clean up shell snapshots: {err:?}"); } @@ -500,7 +500,7 @@ $envVars | ForEach-Object { pub async fn cleanup_stale_snapshots( codex_home: &AbsolutePathBuf, active_session_id: ThreadId, - state_db: Option, + state_db_access: StateDbAccess, ) -> Result<()> { let snapshot_dir = codex_home.join(SNAPSHOT_DIR); @@ -531,7 +531,7 @@ pub async fn cleanup_stale_snapshots( } let rollout_path = - find_thread_path_by_id_str(codex_home, session_id, state_db.as_deref()).await?; + find_thread_path_by_id_str(codex_home, session_id, &state_db_access).await?; let Some(rollout_path) = rollout_path else { remove_snapshot_file(&path).await; continue; diff --git a/codex-rs/core/src/shell_snapshot_tests.rs b/codex-rs/core/src/shell_snapshot_tests.rs index 0199347b4e..e29b03f236 100644 --- a/codex-rs/core/src/shell_snapshot_tests.rs +++ b/codex-rs/core/src/shell_snapshot_tests.rs @@ -202,7 +202,7 @@ async fn try_new_creates_and_deletes_snapshot_file() -> Result<()> { ThreadId::new(), &dir.path().abs(), &shell, - /*state_db*/ None, + crate::StateDbAccess::none(), ) .await .expect("snapshot should be created"); @@ -233,7 +233,7 @@ async fn try_new_uses_distinct_generation_paths() -> Result<()> { session_id, &dir.path().abs(), &shell, - /*state_db*/ None, + crate::StateDbAccess::none(), ) .await .expect("initial snapshot should be created"); @@ -242,7 +242,7 @@ async fn try_new_uses_distinct_generation_paths() -> Result<()> { session_id, &dir.path().abs(), &shell, - /*state_db*/ None, + crate::StateDbAccess::none(), ) .await .expect("refreshed snapshot should be created"); @@ -439,7 +439,7 @@ async fn cleanup_stale_snapshots_removes_orphans_and_keeps_live() -> Result<()> fs::write(&orphan_snapshot, "orphan").await?; fs::write(&invalid_snapshot, "invalid").await?; - cleanup_stale_snapshots(&codex_home, ThreadId::new(), /*state_db*/ None).await?; + cleanup_stale_snapshots(&codex_home, ThreadId::new(), crate::StateDbAccess::none()).await?; assert_eq!(live_snapshot.exists(), true); assert_eq!(orphan_snapshot.exists(), false); @@ -462,7 +462,7 @@ async fn cleanup_stale_snapshots_removes_stale_rollouts() -> Result<()> { set_file_mtime(&rollout_path, SNAPSHOT_RETENTION + Duration::from_secs(60))?; - cleanup_stale_snapshots(&codex_home, ThreadId::new(), /*state_db*/ None).await?; + cleanup_stale_snapshots(&codex_home, ThreadId::new(), crate::StateDbAccess::none()).await?; assert_eq!(stale_snapshot.exists(), false); Ok(()) @@ -483,7 +483,7 @@ async fn cleanup_stale_snapshots_skips_active_session() -> Result<()> { set_file_mtime(&rollout_path, SNAPSHOT_RETENTION + Duration::from_secs(60))?; - cleanup_stale_snapshots(&codex_home, active_session, /*state_db*/ None).await?; + cleanup_stale_snapshots(&codex_home, active_session, crate::StateDbAccess::none()).await?; assert_eq!(active_snapshot.exists(), true); Ok(()) diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 0dba931296..fa9824a52b 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -24,7 +24,7 @@ use codex_login::AuthManager; use codex_mcp::McpConnectionManager; use codex_models_manager::manager::SharedModelsManager; use codex_otel::SessionTelemetry; -use codex_rollout::state_db::StateDbHandle; +use codex_rollout::StateDbAccess; use codex_rollout_trace::ThreadTraceContext; use codex_thread_store::LiveThread; use codex_thread_store::ThreadStore; @@ -64,7 +64,7 @@ pub(crate) struct SessionServices { pub(crate) agent_control: AgentControl, pub(crate) network_proxy: Option, pub(crate) network_approval: Arc, - pub(crate) state_db: Option, + pub(crate) state_db_access: StateDbAccess, pub(crate) live_thread: Option, pub(crate) thread_store: Arc, pub(crate) attestation_provider: Option>, diff --git a/codex-rs/core/src/state_db_bridge.rs b/codex-rs/core/src/state_db_bridge.rs index 78d3cb11f9..bc03479a5b 100644 --- a/codex-rs/core/src/state_db_bridge.rs +++ b/codex-rs/core/src/state_db_bridge.rs @@ -1,8 +1,13 @@ use codex_rollout::state_db as rollout_state_db; +pub use codex_rollout::state_db::StateDbAccess; pub use codex_rollout::state_db::StateDbHandle; +pub use codex_state::DbMetricsRecorderHandle; use crate::config::Config; -pub async fn init_state_db(config: &Config) -> Option { - rollout_state_db::init(config).await +pub async fn init_state_db( + config: &Config, + metrics: Option, +) -> Option { + rollout_state_db::init(config, metrics).await } diff --git a/codex-rs/core/src/stream_events_utils.rs b/codex-rs/core/src/stream_events_utils.rs index 14a4e975ae..f5cec804cc 100644 --- a/codex-rs/core/src/stream_events_utils.rs +++ b/codex-rs/core/src/stream_events_utils.rs @@ -139,7 +139,7 @@ pub(crate) async fn record_completed_response_item( } 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(), + sess.services.state_db_access.as_deref(), item, ) .await; @@ -169,7 +169,7 @@ pub(crate) async fn mark_thread_memory_mode_polluted_if_external_context( return; } state_db::mark_thread_memory_mode_polluted( - sess.services.state_db.as_deref(), + sess.services.state_db_access.as_deref(), sess.conversation_id, "record_completed_response_item", ) @@ -177,7 +177,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<&codex_state::StateRuntime>, item: &ResponseItem, ) -> bool { let Some(raw_text) = raw_assistant_output_text_from_item(item) else { diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index de4114189b..c85e549a13 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -51,6 +51,7 @@ use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnAbortedEvent; use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::protocol::W3cTraceContext; +use codex_rollout::StateDbAccess; use codex_rollout::state_db::StateDbHandle; use codex_state::DirectionalThreadSpawnEdgeStatus; use codex_thread_store::InMemoryThreadStore; @@ -253,7 +254,7 @@ pub(crate) struct ThreadManagerState { session_source: SessionSource, installation_id: String, analytics_events_client: Option, - state_db: Option, + state_db_access: StateDbAccess, // Captures submitted ops for testing purpose when test mode is enabled. ops_log: Option, } @@ -271,12 +272,12 @@ pub fn build_models_manager( pub fn thread_store_from_config( config: &Config, - state_db: Option, + state_db_ctx: StateDbAccess, ) -> Arc { match &config.experimental_thread_store { ThreadStoreConfig::Local => Arc::new(LocalThreadStore::new( LocalThreadStoreConfig::from_config(config), - state_db, + state_db_ctx, )), ThreadStoreConfig::InMemory { id } => InMemoryThreadStore::for_id(id), } @@ -291,7 +292,7 @@ impl ThreadManager { environment_manager: Arc, analytics_events_client: Option, thread_store: Arc, - state_db: Option, + state_db_access: StateDbAccess, installation_id: String, attestation_provider: Option>, ) -> Self { @@ -325,7 +326,7 @@ impl ThreadManager { session_source, installation_id, analytics_events_client, - state_db, + state_db_access, ops_log: should_use_test_thread_manager_behavior() .then(|| Arc::new(std::sync::Mutex::new(Vec::new()))), }), @@ -402,13 +403,14 @@ impl ThreadManager { let skills_watcher = build_skills_watcher(Arc::clone(&skills_manager)); // This test constructor has no Config input. Tests that need a non-local // process store should construct ThreadManager::new with an explicit store. + let state_db_access = StateDbAccess::new(state_db.clone()); 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(), + state_db_access.clone(), )); Self { state: Arc::new(ThreadManagerState { @@ -427,7 +429,7 @@ impl ThreadManager { session_source: SessionSource::Exec, installation_id, analytics_events_client: None, - state_db, + state_db_access, ops_log: should_use_test_thread_manager_behavior() .then(|| Arc::new(std::sync::Mutex::new(Vec::new()))), }), @@ -889,7 +891,7 @@ impl ThreadManager { impl ThreadManagerState { pub(crate) fn state_db(&self) -> Option { - self.state_db.clone() + self.state_db_access.state_db() } pub(crate) async fn list_thread_ids(&self) -> Vec { @@ -1210,6 +1212,7 @@ impl ThreadManagerState { analytics_events_client: self.analytics_events_client.clone(), thread_store: Arc::clone(&self.thread_store), attestation_provider: self.attestation_provider.clone(), + state_db_access: self.state_db_access.clone(), }) .await?; let new_thread = self diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 3bc5c77841..b622d29387 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -486,14 +486,15 @@ 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_access = StateDbAccess::none(); 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, + thread_store_from_config(&config, state_db_access.clone()), + state_db_access, TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, ); @@ -601,8 +602,9 @@ async fn explicit_installation_id_skips_codex_home_file() { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); let installation_id = uuid::Uuid::new_v4().to_string(); - let state_db = init_state_db(&config).await; - let thread_store = thread_store_from_config(&config, state_db.clone()); + let state_db = init_state_db(&config, /*metrics*/ None).await; + let state_db_access = StateDbAccess::new(state_db.clone()); + let thread_store = thread_store_from_config(&config, state_db_access.clone()); let manager = ThreadManager::new( &config, auth_manager, @@ -610,7 +612,7 @@ async fn explicit_installation_id_skips_codex_home_file() { Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, thread_store, - state_db.clone(), + state_db_access, installation_id.clone(), /*attestation_provider*/ None, ); @@ -641,14 +643,15 @@ 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_access = StateDbAccess::none(); 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, + thread_store_from_config(&config, state_db_access.clone()), + state_db_access, TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, ); @@ -697,14 +700,15 @@ 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_access = StateDbAccess::none(); 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, + thread_store_from_config(&config, state_db_access.clone()), + state_db_access, TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, ); @@ -758,8 +762,9 @@ async fn resume_stopped_thread_from_rollout_preserves_thread_source() { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); - let state_db = init_state_db(&config).await; - let thread_store = thread_store_from_config(&config, state_db.clone()); + let state_db = init_state_db(&config, /*metrics*/ None).await; + let state_db_access = StateDbAccess::new(state_db.clone()); + let thread_store = thread_store_from_config(&config, state_db_access.clone()); let manager = ThreadManager::new( &config, auth_manager.clone(), @@ -767,7 +772,7 @@ async fn resume_stopped_thread_from_rollout_preserves_thread_source() { Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, thread_store, - state_db.clone(), + state_db_access, TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, ); @@ -843,8 +848,9 @@ async fn rollout_path_resume_and_fork_read_history_through_thread_store() { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); - let state_db = init_state_db(&config).await; - let thread_store = thread_store_from_config(&config, state_db.clone()); + let state_db = init_state_db(&config, /*metrics*/ None).await; + let state_db_access = StateDbAccess::new(state_db.clone()); + let thread_store = thread_store_from_config(&config, state_db_access.clone()); let in_memory_store = thread_store .as_any() .downcast_ref::() @@ -856,7 +862,7 @@ async fn rollout_path_resume_and_fork_read_history_through_thread_store() { Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, thread_store.clone(), - state_db, + state_db_access, TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, ); @@ -951,14 +957,15 @@ 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_access = StateDbAccess::none(); 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, + thread_store_from_config(&config, state_db_access.clone()), + state_db_access, TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, ); @@ -1165,15 +1172,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 = init_state_db(&config, /*metrics*/ None).await; + let state_db_access = StateDbAccess::new(state_db.clone()); 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(), + thread_store_from_config(&config, state_db_access.clone()), + state_db_access, TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, ); @@ -1272,15 +1280,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 = init_state_db(&config, /*metrics*/ None).await; + let state_db_access = StateDbAccess::new(state_db.clone()); 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(), + thread_store_from_config(&config, state_db_access.clone()), + state_db_access, TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, ); @@ -1368,15 +1377,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 = init_state_db(&config, /*metrics*/ None).await; + let state_db_access = StateDbAccess::new(state_db.clone()); 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(), + thread_store_from_config(&config, state_db_access.clone()), + state_db_access, TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, ); @@ -1510,15 +1520,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 = init_state_db(&config, /*metrics*/ None).await; + let state_db_access = StateDbAccess::new(state_db.clone()); 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(), + thread_store_from_config(&config, state_db_access.clone()), + state_db_access, TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, ); 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 e7625fa44f..483cc7baea 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -3148,8 +3148,33 @@ async fn close_agent_submits_shutdown_and_returns_previous_status() { assert_eq!(status_after, AgentStatus::NotFound); } -#[tokio::test] -async fn tool_handlers_cascade_close_and_resume_and_keep_explicitly_closed_subtrees_closed() { +#[test] +fn tool_handlers_cascade_close_and_resume_and_keep_explicitly_closed_subtrees_closed() { + const TEST_STACK_SIZE_BYTES: usize = 8 * 1024 * 1024; + + let handle = std::thread::Builder::new() + .name( + "tool_handlers_cascade_close_and_resume_and_keep_explicitly_closed_subtrees_closed" + .to_string(), + ) + .stack_size(TEST_STACK_SIZE_BYTES) + .spawn(|| { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("test runtime should build"); + runtime.block_on(Box::pin( + tool_handlers_cascade_close_and_resume_and_keep_explicitly_closed_subtrees_closed_impl(), + )); + }) + .expect("test thread should spawn"); + + if let Err(payload) = handle.join() { + std::panic::resume_unwind(payload); + } +} + +async fn tool_handlers_cascade_close_and_resume_and_keep_explicitly_closed_subtrees_closed_impl() { let (_session, turn) = make_session_and_context().await; let mut config = turn.config.as_ref().clone(); config.agent_max_depth = 3; @@ -3157,15 +3182,16 @@ 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, /*metrics*/ None).await; + let state_db_access = crate::StateDbAccess::new(state_db.clone()); 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_access.clone()), + state_db_access, "11111111-1111-4111-8111-111111111111".to_string(), /*attestation_provider*/ None, ); diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index c326d84ebf..a743359811 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -66,6 +66,7 @@ type WorkspaceSetup = dyn FnOnce(AbsolutePathBuf, Arc) - + Send; const TEST_MODEL_WITH_EXPERIMENTAL_TOOLS: &str = "test-gpt-5.1-codex"; const REMOTE_EXEC_SERVER_URL_ENV_VAR: &str = "CODEX_TEST_REMOTE_EXEC_SERVER_URL"; +const OTEL_SERVICE_NAME: &str = "codex_exec"; static REMOTE_TEST_INSTANCE_COUNTER: AtomicU64 = AtomicU64::new(0); const SUBMIT_TURN_COMPLETE_TIMEOUT: Duration = Duration::from_secs(30); @@ -423,8 +424,12 @@ impl TestCodexBuilder { environment_manager: Arc, ) -> anyhow::Result { let auth = self.auth.clone(); - let state_db = codex_core::init_state_db(&config).await; - let thread_store = thread_store_from_config(&config, state_db.clone()); + let state_db_metrics = + codex_core::otel_init::sqlite_metrics_recorder_from_global(OTEL_SERVICE_NAME); + let state_db = codex_core::init_state_db(&config, state_db_metrics.clone()).await; + let state_db_access = codex_core::StateDbAccess::new(state_db.clone()) + .with_standalone_metrics(state_db_metrics); + let thread_store = thread_store_from_config(&config, state_db_access.clone()); let installation_id = resolve_installation_id(&config.codex_home).await?; let thread_manager = ThreadManager::new( &config, @@ -433,7 +438,7 @@ impl TestCodexBuilder { Arc::clone(&environment_manager), /*analytics_events_client*/ None, thread_store, - state_db.clone(), + state_db_access, installation_id, /*attestation_provider*/ None, ); diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index 83a7865d16..5787b93036 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -1138,8 +1138,8 @@ async fn prefers_apikey_when_config_prefers_apikey_even_with_chatgpt_tokens() { 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, + thread_store_from_config(&config, codex_core::StateDbAccess::none()), + codex_core::StateDbAccess::none(), installation_id, /*attestation_provider*/ None, ); diff --git a/codex-rs/core/tests/suite/personality_migration.rs b/codex-rs/core/tests/suite/personality_migration.rs index 0b89a9cfba..bf7a9e64af 100644 --- a/codex-rs/core/tests/suite/personality_migration.rs +++ b/codex-rs/core/tests/suite/personality_migration.rs @@ -1,6 +1,7 @@ use codex_config::config_toml::ConfigToml; use codex_core::ARCHIVED_SESSIONS_SUBDIR; use codex_core::SESSIONS_SUBDIR; +use codex_core::StateDbAccess; use codex_core::personality_migration::PERSONALITY_MIGRATION_FILENAME; use codex_core::personality_migration::PersonalityMigrationStatus; use codex_core::personality_migration::maybe_migrate_personality; @@ -144,7 +145,8 @@ async fn migration_marker_exists_no_sessions_no_change() -> io::Result<()> { tokio::fs::write(&marker_path, "v1\n").await?; let status = - maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?; + maybe_migrate_personality(temp.path(), &ConfigToml::default(), StateDbAccess::none()) + .await?; assert_eq!(status, PersonalityMigrationStatus::SkippedMarker); assert_eq!( @@ -159,7 +161,8 @@ 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?; + maybe_migrate_personality(temp.path(), &ConfigToml::default(), StateDbAccess::none()) + .await?; assert_eq!(status, PersonalityMigrationStatus::SkippedNoSessions); assert_eq!( @@ -179,7 +182,8 @@ async fn no_marker_sessions_sets_personality() -> io::Result<()> { write_session_with_user_event(temp.path()).await?; let status = - maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?; + maybe_migrate_personality(temp.path(), &ConfigToml::default(), StateDbAccess::none()) + .await?; assert_eq!(status, PersonalityMigrationStatus::Applied); assert_eq!( @@ -199,7 +203,8 @@ 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 = + maybe_migrate_personality(temp.path(), &config_toml, StateDbAccess::none()).await?; assert_eq!(status, PersonalityMigrationStatus::Applied); let persisted = read_config_toml(temp.path()).await?; @@ -214,7 +219,8 @@ async fn no_marker_meta_only_rollout_is_treated_as_no_sessions() -> io::Result<( write_session_with_meta_only(temp.path()).await?; let status = - maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?; + maybe_migrate_personality(temp.path(), &ConfigToml::default(), StateDbAccess::none()) + .await?; assert_eq!(status, PersonalityMigrationStatus::SkippedNoSessions); assert_eq!( @@ -234,7 +240,8 @@ 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 = + maybe_migrate_personality(temp.path(), &config_toml, StateDbAccess::none()).await?; assert_eq!( status, @@ -264,7 +271,8 @@ personality = "friendly" "#, )?; - let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?; + let status = + maybe_migrate_personality(temp.path(), &config_toml, StateDbAccess::none()).await?; assert_eq!( status, @@ -287,7 +295,8 @@ 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 = + maybe_migrate_personality(temp.path(), &config_toml, StateDbAccess::none()).await?; assert_eq!(status, PersonalityMigrationStatus::SkippedMarker); Ok(()) @@ -298,7 +307,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 = maybe_migrate_personality(temp.path(), &config_toml, StateDbAccess::none()) .await .expect_err("missing profile should fail"); @@ -316,9 +325,11 @@ async fn applied_migration_is_idempotent_on_second_run() -> io::Result<()> { write_session_with_user_event(temp.path()).await?; let first_status = - maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?; + maybe_migrate_personality(temp.path(), &ConfigToml::default(), StateDbAccess::none()) + .await?; let second_status = - maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?; + maybe_migrate_personality(temp.path(), &ConfigToml::default(), StateDbAccess::none()) + .await?; assert_eq!(first_status, PersonalityMigrationStatus::Applied); assert_eq!(second_status, PersonalityMigrationStatus::SkippedMarker); @@ -333,7 +344,8 @@ async fn no_marker_archived_sessions_sets_personality() -> io::Result<()> { write_archived_session_with_user_event(temp.path()).await?; let status = - maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?; + maybe_migrate_personality(temp.path(), &ConfigToml::default(), StateDbAccess::none()) + .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..c1f3aeec67 100644 --- a/codex-rs/core/tests/suite/prompt_debug_tests.rs +++ b/codex-rs/core/tests/suite/prompt_debug_tests.rs @@ -1,4 +1,5 @@ use anyhow::Result; +use codex_core::StateDbAccess; use codex_core::build_prompt_input; use codex_core::config::ConfigBuilder; use codex_core::config::ConfigOverrides; @@ -29,7 +30,7 @@ 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, + StateDbAccess::none(), ) .await?; diff --git a/codex-rs/core/tests/suite/rollout_list_find.rs b/codex-rs/core/tests/suite/rollout_list_find.rs index 51b3f2db1a..62c9d12432 100644 --- a/codex-rs/core/tests/suite/rollout_list_find.rs +++ b/codex-rs/core/tests/suite/rollout_list_find.rs @@ -66,9 +66,13 @@ async fn upsert_thread_metadata( thread_id: ThreadId, rollout_path: PathBuf, ) -> StateDbHandle { - let runtime = StateRuntime::init(codex_home.to_path_buf(), "test-provider".to_string()) - .await - .unwrap(); + let runtime = StateRuntime::init( + codex_home.to_path_buf(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .unwrap(); runtime .mark_backfill_complete(/*last_watermark*/ None) .await @@ -91,10 +95,13 @@ async fn find_locates_rollout_file_by_id() { let id = Uuid::new_v4(); let expected = write_minimal_rollout_with_id(home.path(), id); - let found = - find_thread_path_by_id_str(home.path(), &id.to_string(), /*state_db_ctx*/ None) - .await - .unwrap(); + let found = find_thread_path_by_id_str( + home.path(), + &id.to_string(), + &codex_core::StateDbAccess::none(), + ) + .await + .unwrap(); assert_eq!(found.unwrap(), expected); } @@ -108,10 +115,13 @@ async fn find_handles_gitignore_covering_codex_home_directory() { let id = Uuid::new_v4(); let expected = write_minimal_rollout_with_id(&codex_home, id); - let found = - find_thread_path_by_id_str(&codex_home, &id.to_string(), /*state_db_ctx*/ None) - .await - .unwrap(); + let found = find_thread_path_by_id_str( + &codex_home, + &id.to_string(), + &codex_core::StateDbAccess::none(), + ) + .await + .unwrap(); assert_eq!(found, Some(expected)); } @@ -129,9 +139,13 @@ async fn find_prefers_sqlite_path_by_id() { write_minimal_rollout_with_id(home.path(), id); let state_db = upsert_thread_metadata(home.path(), thread_id, db_path.clone()).await; - let found = find_thread_path_by_id_str(home.path(), &id.to_string(), Some(&state_db)) - .await - .unwrap(); + let found = find_thread_path_by_id_str( + home.path(), + &id.to_string(), + &codex_core::StateDbAccess::new(Some(state_db.clone())), + ) + .await + .unwrap(); assert_eq!(found, Some(db_path)); } @@ -148,9 +162,13 @@ async fn find_falls_back_to_filesystem_when_sqlite_has_no_match() { .join("sessions/2030/12/30/rollout-2030-12-30T00-00-00-unrelated.jsonl"); let state_db = upsert_thread_metadata(home.path(), unrelated_thread_id, unrelated_path).await; - let found = find_thread_path_by_id_str(home.path(), &id.to_string(), Some(&state_db)) - .await - .unwrap(); + let found = find_thread_path_by_id_str( + home.path(), + &id.to_string(), + &codex_core::StateDbAccess::new(Some(state_db.clone())), + ) + .await + .unwrap(); assert_eq!(found, Some(expected)); } @@ -162,10 +180,13 @@ async fn find_ignores_granular_gitignore_rules() { let expected = write_minimal_rollout_with_id(home.path(), id); std::fs::write(home.path().join("sessions/.gitignore"), "*.jsonl\n").unwrap(); - let found = - find_thread_path_by_id_str(home.path(), &id.to_string(), /*state_db_ctx*/ None) - .await - .unwrap(); + let found = find_thread_path_by_id_str( + home.path(), + &id.to_string(), + &codex_core::StateDbAccess::none(), + ) + .await + .unwrap(); assert_eq!(found, Some(expected)); } @@ -212,7 +233,8 @@ async fn find_locates_rollout_file_written_by_recorder() -> std::io::Result<()> )?; let found = - find_thread_meta_by_name_str(home.path(), thread_name, /*state_db_ctx*/ None).await?; + find_thread_meta_by_name_str(home.path(), thread_name, &codex_core::StateDbAccess::none()) + .await?; let (path, session_meta) = found.expect("expected rollout path to be found"); assert_eq!(session_meta.meta.id, thread_id); @@ -232,7 +254,7 @@ async fn find_archived_locates_rollout_file_by_id() { let found = find_archived_thread_path_by_id_str( home.path(), &id.to_string(), - /*state_db_ctx*/ None, + &codex_core::StateDbAccess::none(), ) .await .unwrap(); diff --git a/codex-rs/exec/src/lib.rs b/codex-rs/exec/src/lib.rs index e68c96d00b..d88c2a27bb 100644 --- a/codex-rs/exec/src/lib.rs +++ b/codex-rs/exec/src/lib.rs @@ -56,6 +56,7 @@ use codex_cloud_requirements::cloud_requirements_loader_for_storage; use codex_config::ConfigLoadError; use codex_config::LoaderOverrides; use codex_config::format_config_error_with_source; +use codex_core::StateDbAccess; use codex_core::StateDbHandle; use codex_core::check_execpolicy_for_warnings; use codex_core::config::Config; @@ -156,6 +157,7 @@ use crate::event_processor::EventProcessor; const DEFAULT_ANALYTICS_ENABLED: bool = true; const EXEC_DEFAULT_LOG_FILTER: &str = "error,opentelemetry_sdk=off,opentelemetry_otlp=off"; +const OTEL_SERVICE_NAME: &str = "codex_exec"; enum InitialOperation { UserTurn { @@ -482,6 +484,9 @@ pub async fn run_main(cli: Cli, arg0_paths: Arg0DispatchPaths) -> anyhow::Result 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()); + codex_core::otel_init::record_process_start(otel.as_ref(), OTEL_SERVICE_NAME); + let state_db_metrics = + codex_core::otel_init::sqlite_metrics_recorder(otel.as_ref(), OTEL_SERVICE_NAME); let _ = tracing_subscriber::registry() .with(fmt_layer) @@ -507,7 +512,9 @@ pub async fn run_main(cli: Cli, arg0_paths: Arg0DispatchPaths) -> anyhow::Result arg0_paths.codex_self_exe.clone(), arg0_paths.codex_linux_sandbox_exe.clone(), )?; - let state_db = codex_core::init_state_db(&config).await; + let state_db = codex_core::init_state_db(&config, state_db_metrics.clone()).await; + let state_db_access = + StateDbAccess::new(state_db.clone()).with_standalone_metrics(state_db_metrics); let environment_manager = if run_loader_overrides.ignore_user_config { EnvironmentManager::from_env(local_runtime_paths).await? } else { @@ -521,7 +528,7 @@ pub async fn run_main(cli: Cli, arg0_paths: Arg0DispatchPaths) -> anyhow::Result cloud_requirements: run_cloud_requirements, feedback: CodexFeedback::new(), log_db: None, - state_db: state_db.clone(), + state_db_access, environment_manager: std::sync::Arc::new(environment_manager), config_warnings, session_source: SessionSource::Exec, @@ -1398,9 +1405,12 @@ async fn resolve_resume_thread_id( if let Some(thread) = resolved { return Ok(Some(thread.id.to_string())); } - if let Some((_, session_meta)) = - find_thread_meta_by_name_str(&config.codex_home, session_id, Some(state_db.as_ref())) - .await? + if let Some((_, session_meta)) = find_thread_meta_by_name_str( + &config.codex_home, + session_id, + &StateDbAccess::new(Some((*state_db).clone())), + ) + .await? && (args.all || cwds_match(config.cwd.as_path(), &session_meta.meta.cwd)) { return Ok(Some(session_meta.meta.id.to_string())); diff --git a/codex-rs/mcp-server/src/lib.rs b/codex-rs/mcp-server/src/lib.rs index aa560bbe6a..8eaa5312d1 100644 --- a/codex-rs/mcp-server/src/lib.rs +++ b/codex-rs/mcp-server/src/lib.rs @@ -74,19 +74,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 environment_manager = Arc::new( - EnvironmentManager::from_codex_home( - config.codex_home.clone(), - ExecServerRuntimePaths::from_optional_paths( - arg0_paths.codex_self_exe.clone(), - arg0_paths.codex_linux_sandbox_exe.clone(), - )?, - ) - .await - .map_err(std::io::Error::other)?, - ); - let otel = codex_core::otel_init::build_provider( &config, env!("CARGO_PKG_VERSION"), @@ -99,6 +86,23 @@ pub async fn run_main( format!("error loading otel config: {e}"), ) })?; + codex_core::otel_init::record_process_start(otel.as_ref(), OTEL_SERVICE_NAME); + let state_db_metrics = + codex_core::otel_init::sqlite_metrics_recorder(otel.as_ref(), OTEL_SERVICE_NAME); + let state_db = codex_core::init_state_db(&config, state_db_metrics.clone()).await; + let state_db_access = + codex_core::StateDbAccess::new(state_db.clone()).with_standalone_metrics(state_db_metrics); + let environment_manager = Arc::new( + EnvironmentManager::from_codex_home( + config.codex_home.clone(), + ExecServerRuntimePaths::from_optional_paths( + arg0_paths.codex_self_exe.clone(), + arg0_paths.codex_linux_sandbox_exe.clone(), + )?, + ) + .await + .map_err(std::io::Error::other)?, + ); let fmt_layer = tracing_subscriber::fmt::layer() .with_writer(std::io::stderr) @@ -148,7 +152,7 @@ pub async fn run_main( arg0_paths, Arc::new(config), environment_manager, - state_db, + state_db_access, installation_id, ) .await; @@ -201,10 +205,15 @@ pub async fn run_main( mod tests { use super::*; use codex_config::types::OtelExporterKind; + use codex_config::types::OtelHttpProtocol; use codex_core::config::ConfigBuilder; use pretty_assertions::assert_eq; use std::collections::HashMap; use tempfile::TempDir; + use wiremock::Mock; + use wiremock::MockServer; + use wiremock::ResponseTemplate; + use wiremock::matchers::method; #[test] fn mcp_server_defaults_analytics_to_enabled() { @@ -213,14 +222,21 @@ mod tests { #[tokio::test] async fn mcp_server_builds_otel_provider_with_logs_traces_and_metrics() -> anyhow::Result<()> { + let collector = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200)) + .mount(&collector) + .await; + let codex_home = TempDir::new()?; let mut config = ConfigBuilder::default() .codex_home(codex_home.path().to_path_buf()) .build() .await?; - let exporter = OtelExporterKind::OtlpGrpc { - endpoint: "http://localhost:4317".to_string(), + let exporter = OtelExporterKind::OtlpHttp { + endpoint: collector.uri(), headers: HashMap::new(), + protocol: OtelHttpProtocol::Binary, tls: None, }; config.otel.exporter = exporter.clone(); diff --git a/codex-rs/mcp-server/src/message_processor.rs b/codex-rs/mcp-server/src/message_processor.rs index f963747e87..65d28c18a3 100644 --- a/codex-rs/mcp-server/src/message_processor.rs +++ b/codex-rs/mcp-server/src/message_processor.rs @@ -2,7 +2,7 @@ use std::collections::HashMap; use std::sync::Arc; use codex_arg0::Arg0DispatchPaths; -use codex_core::StateDbHandle; +use codex_core::StateDbAccess; use codex_core::ThreadManager; use codex_core::config::Config; use codex_core::thread_store_from_config; @@ -54,7 +54,7 @@ impl MessageProcessor { arg0_paths: Arg0DispatchPaths, config: Arc, environment_manager: Arc, - state_db: Option, + state_db_access: StateDbAccess, installation_id: String, ) -> Self { let outgoing = Arc::new(outgoing); @@ -69,8 +69,8 @@ impl MessageProcessor { SessionSource::Mcp, environment_manager, /*analytics_events_client*/ None, - thread_store_from_config(config.as_ref(), state_db.clone()), - state_db.clone(), + thread_store_from_config(config.as_ref(), state_db_access.clone()), + state_db_access, installation_id, /*attestation_provider*/ None, )); diff --git a/codex-rs/memories/write/src/startup_tests.rs b/codex-rs/memories/write/src/startup_tests.rs index 4fcb1d409b..fa5ff16216 100644 --- a/codex-rs/memories/write/src/startup_tests.rs +++ b/codex-rs/memories/write/src/startup_tests.rs @@ -308,8 +308,12 @@ async fn build_test_codex( } async fn init_state_db(home: &Arc) -> anyhow::Result> { - let db = - codex_state::StateRuntime::init(home.path().to_path_buf(), "test-provider".into()).await?; + let db = codex_state::StateRuntime::init( + home.path().to_path_buf(), + "test-provider".into(), + /*metrics*/ None, + ) + .await?; db.mark_backfill_complete(/*last_watermark*/ None).await?; Ok(db) } diff --git a/codex-rs/otel/src/metrics/client.rs b/codex-rs/otel/src/metrics/client.rs index 417c1f4bd9..32b050ae84 100644 --- a/codex-rs/otel/src/metrics/client.rs +++ b/codex-rs/otel/src/metrics/client.rs @@ -41,6 +41,7 @@ use std::time::Duration; use tracing::debug; const ENV_ATTRIBUTE: &str = "env"; +const ARCH_ATTRIBUTE: &str = "arch"; const METER_NAME: &str = "codex"; const DURATION_UNIT: &str = "ms"; const DURATION_DESCRIPTION: &str = "Duration in milliseconds."; @@ -198,13 +199,13 @@ impl MetricsClient { validate_tags(&default_tags)?; - let mut resource_attributes = Vec::with_capacity(4); + let mut resource_attributes = Vec::with_capacity(5); resource_attributes.push(KeyValue::new( semconv::attribute::SERVICE_VERSION, service_version, )); resource_attributes.push(KeyValue::new(ENV_ATTRIBUTE, environment)); - resource_attributes.extend(os_resource_attributes()); + resource_attributes.extend(platform_resource_attributes()); let resource = Resource::builder() .with_service_name(service_name) @@ -290,12 +291,13 @@ impl MetricsClient { } } -fn os_resource_attributes() -> Vec { +fn platform_resource_attributes() -> Vec { let os_info = os_info::get(); let os_type_raw = os_info.os_type().to_string(); let os_type = sanitize_metric_tag_value(os_type_raw.as_str()); let os_version_raw = os_info.version().to_string(); let os_version = sanitize_metric_tag_value(os_version_raw.as_str()); + let arch = sanitize_metric_tag_value(std::env::consts::ARCH); let mut attributes = Vec::new(); if os_type != "unspecified" { attributes.push(KeyValue::new("os", os_type)); @@ -303,6 +305,9 @@ fn os_resource_attributes() -> Vec { if os_version != "unspecified" { attributes.push(KeyValue::new("os_version", os_version)); } + if arch != "unspecified" { + attributes.push(KeyValue::new(ARCH_ATTRIBUTE, arch)); + } attributes } diff --git a/codex-rs/otel/src/metrics/mod.rs b/codex-rs/otel/src/metrics/mod.rs index e75840bf37..07b7554224 100644 --- a/codex-rs/otel/src/metrics/mod.rs +++ b/codex-rs/otel/src/metrics/mod.rs @@ -2,6 +2,7 @@ mod client; mod config; mod error; pub(crate) mod names; +mod process; pub(crate) mod runtime_metrics; pub(crate) mod tags; pub(crate) mod timer; @@ -13,9 +14,12 @@ pub use crate::metrics::config::MetricsConfig; pub use crate::metrics::config::MetricsExporter; pub use crate::metrics::error::MetricsError; pub use crate::metrics::error::Result; +pub use crate::metrics::process::record_process_start_once; pub use names::*; use std::sync::OnceLock; +pub use tags::ORIGINATOR_TAG; pub use tags::SessionMetricTagValues; +pub use tags::bounded_originator_tag_value; static GLOBAL_METRICS: OnceLock = OnceLock::new(); static GLOBAL_STATSIG_METRICS_SETTINGS: OnceLock = OnceLock::new(); diff --git a/codex-rs/otel/src/metrics/names.rs b/codex-rs/otel/src/metrics/names.rs index b2f2d639ce..67ac6e19ed 100644 --- a/codex-rs/otel/src/metrics/names.rs +++ b/codex-rs/otel/src/metrics/names.rs @@ -1,6 +1,7 @@ pub const TOOL_CALL_COUNT_METRIC: &str = "codex.tool.call"; pub const TOOL_CALL_DURATION_METRIC: &str = "codex.tool.call.duration_ms"; pub const TOOL_CALL_UNIFIED_EXEC_METRIC: &str = "codex.tool.unified_exec"; +pub const PROCESS_START_METRIC: &str = "codex.process.start"; pub const API_CALL_COUNT_METRIC: &str = "codex.api_request"; pub const API_CALL_DURATION_METRIC: &str = "codex.api_request.duration_ms"; pub const SSE_EVENT_COUNT_METRIC: &str = "codex.sse_event"; diff --git a/codex-rs/otel/src/metrics/process.rs b/codex-rs/otel/src/metrics/process.rs new file mode 100644 index 0000000000..6c7cd18b4e --- /dev/null +++ b/codex-rs/otel/src/metrics/process.rs @@ -0,0 +1,27 @@ +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; + +use super::client::MetricsClient; +use super::error::Result; +use super::names::PROCESS_START_METRIC; +use super::tags::ORIGINATOR_TAG; +use super::tags::bounded_originator_tag_value; + +static PROCESS_START_RECORDED: AtomicBool = AtomicBool::new(false); + +/// Record the process start counter at most once for this process. +pub fn record_process_start_once(metrics: &MetricsClient, originator: &str) -> Result { + if PROCESS_START_RECORDED + .compare_exchange(false, true, Ordering::Relaxed, Ordering::Relaxed) + .is_err() + { + return Ok(false); + } + + metrics.counter( + PROCESS_START_METRIC, + /*inc*/ 1, + &[(ORIGINATOR_TAG, bounded_originator_tag_value(originator))], + )?; + Ok(true) +} diff --git a/codex-rs/otel/src/metrics/tags.rs b/codex-rs/otel/src/metrics/tags.rs index 5c166e307e..96a3157170 100644 --- a/codex-rs/otel/src/metrics/tags.rs +++ b/codex-rs/otel/src/metrics/tags.rs @@ -1,6 +1,7 @@ use crate::metrics::Result; use crate::metrics::validation::validate_tag_key; use crate::metrics::validation::validate_tag_value; +use codex_utils_string::sanitize_metric_tag_value; pub const APP_VERSION_TAG: &str = "app.version"; pub const AUTH_MODE_TAG: &str = "auth_mode"; @@ -9,6 +10,31 @@ pub const ORIGINATOR_TAG: &str = "originator"; pub const SERVICE_NAME_TAG: &str = "service_name"; pub const SESSION_SOURCE_TAG: &str = "session_source"; +const OTHER_ORIGINATOR_TAG_VALUE: &str = "other"; +const KNOWN_ORIGINATOR_TAG_VALUES: &[&str] = &[ + "codex_desktop", + "codex-app-server", + "codex_mcp_server", + "codex_cli_rs", + "codex-tui", + "codex_vscode", + "none", + "codex_exec", + "codex-cli", + "codex_sdk_ts", + "codex-app-server-sdk", +]; + +/// Returns a known low-cardinality originator tag value, or `other`. +pub fn bounded_originator_tag_value(originator: &str) -> &'static str { + let sanitized = sanitize_metric_tag_value(originator); + KNOWN_ORIGINATOR_TAG_VALUES + .iter() + .copied() + .find(|known| *known == sanitized.as_str()) + .unwrap_or(OTHER_ORIGINATOR_TAG_VALUE) +} + pub struct SessionMetricTagValues<'a> { pub auth_mode: Option<&'a str>, pub session_source: &'a str, diff --git a/codex-rs/rollout/src/lib.rs b/codex-rs/rollout/src/lib.rs index d65ddd3d5b..497b9263b6 100644 --- a/codex-rs/rollout/src/lib.rs +++ b/codex-rs/rollout/src/lib.rs @@ -10,6 +10,7 @@ pub(crate) mod metadata; pub(crate) mod policy; pub(crate) mod recorder; pub(crate) mod session_index; +pub(crate) mod sqlite_metrics; pub mod state_db; pub(crate) mod default_client { @@ -62,6 +63,7 @@ pub use session_index::append_thread_name; pub use session_index::find_thread_meta_by_name_str; pub use session_index::find_thread_name_by_id; pub use session_index::find_thread_names_by_ids; +pub use state_db::StateDbAccess; pub use state_db::StateDbHandle; #[cfg(test)] diff --git a/codex-rs/rollout/src/list.rs b/codex-rs/rollout/src/list.rs index 8ba63a713e..c586dc2abc 100644 --- a/codex-rs/rollout/src/list.rs +++ b/codex-rs/rollout/src/list.rs @@ -20,6 +20,7 @@ use super::ARCHIVED_SESSIONS_SUBDIR; use super::SESSIONS_SUBDIR; use crate::protocol::EventMsg; use crate::state_db; +use crate::state_db::StateDbAccess; use codex_file_search as file_search; use codex_protocol::ThreadId; use codex_protocol::protocol::RolloutItem; @@ -1240,7 +1241,9 @@ async fn find_thread_path_by_id_str_in_subdir( subdir: &str, id_str: &str, state_db_ctx: Option<&codex_state::StateRuntime>, + metrics: Option<&dyn codex_state::DbMetricsRecorder>, ) -> io::Result> { + let metrics = metrics.or_else(|| state_db_ctx.and_then(codex_state::StateRuntime::metrics)); // Validate UUID format early. if Uuid::parse_str(id_str).is_err() { return Ok(None); @@ -1255,47 +1258,67 @@ async fn find_thread_path_by_id_str_in_subdir( }; let thread_id = ThreadId::from_string(id_str).ok(); let mut unverified_db_path = None; + let mut fallback_reason = state_db_ctx.is_none().then_some("db_unavailable"); if let Some(state_db_ctx) = state_db_ctx && let Some(thread_id) = thread_id - && let Some(db_path) = state_db::find_rollout_path_by_id( - Some(state_db_ctx), - thread_id, - archived_only, - "find_path_query", - ) - .await { - if tokio::fs::try_exists(&db_path).await.unwrap_or(false) { - match read_session_meta_line(&db_path).await { - Ok(meta_line) if meta_line.meta.id == thread_id => { - return Ok(Some(db_path)); - } - Ok(meta_line) => { + match state_db_ctx + .find_rollout_path_by_id(thread_id, archived_only) + .await + { + Ok(Some(db_path)) => { + if tokio::fs::try_exists(&db_path).await.unwrap_or(false) { + match read_session_meta_line(&db_path).await { + Ok(meta_line) if meta_line.meta.id == thread_id => { + return Ok(Some(db_path)); + } + Ok(meta_line) => { + tracing::error!( + "state db returned rollout path for thread {id_str} but file belongs to thread {}: {}", + meta_line.meta.id, + db_path.display() + ); + tracing::warn!( + "state db discrepancy during find_thread_path_by_id_str_in_subdir: mismatched_db_path" + ); + codex_state::record_db_fallback_metric( + metrics, + "find_thread_path", + "mismatch", + ); + } + Err(err) => { + tracing::debug!( + "state db returned rollout path for thread {id_str} that could not be verified: {}: {err}", + db_path.display() + ); + unverified_db_path = Some(db_path); + } + } + } else { tracing::error!( - "state db returned rollout path for thread {id_str} but file belongs to thread {}: {}", - meta_line.meta.id, + "state db returned stale rollout path for thread {id_str}: {}", db_path.display() ); tracing::warn!( - "state db discrepancy during find_thread_path_by_id_str_in_subdir: mismatched_db_path" + "state db discrepancy during find_thread_path_by_id_str_in_subdir: stale_db_path" ); - } - Err(err) => { - tracing::debug!( - "state db returned rollout path for thread {id_str} that could not be verified: {}: {err}", - db_path.display() + codex_state::record_db_fallback_metric( + metrics, + "find_thread_path", + "stale_path", ); - unverified_db_path = Some(db_path); } } - } else { - tracing::error!( - "state db returned stale rollout path for thread {id_str}: {}", - db_path.display() - ); - tracing::warn!( - "state db discrepancy during find_thread_path_by_id_str_in_subdir: stale_db_path" - ); + Ok(None) => { + fallback_reason = Some("missing_row"); + } + Err(err) => { + tracing::warn!( + "state db find_rollout_path_by_id failed during find_path_query: {err}" + ); + fallback_reason = Some("db_error"); + } } } @@ -1323,6 +1346,9 @@ async fn find_thread_path_by_id_str_in_subdir( tracing::warn!( "state db discrepancy during find_thread_path_by_id_str_in_subdir: falling_back" ); + if let Some(reason) = fallback_reason { + codex_state::record_db_fallback_metric(metrics, "find_thread_path", reason); + } state_db::read_repair_rollout_path( state_db_ctx, thread_id, @@ -1341,19 +1367,32 @@ async fn find_thread_path_by_id_str_in_subdir( pub async fn find_thread_path_by_id_str( codex_home: &Path, id_str: &str, - state_db_ctx: Option<&codex_state::StateRuntime>, + state_db_ctx: &StateDbAccess, ) -> io::Result> { - find_thread_path_by_id_str_in_subdir(codex_home, SESSIONS_SUBDIR, id_str, state_db_ctx).await + find_thread_path_by_id_str_in_subdir( + codex_home, + SESSIONS_SUBDIR, + id_str, + state_db_ctx.as_deref(), + state_db_ctx.metrics(), + ) + .await } /// Locate an archived thread rollout file by its UUID string. pub async fn find_archived_thread_path_by_id_str( codex_home: &Path, id_str: &str, - state_db_ctx: Option<&codex_state::StateRuntime>, + state_db_ctx: &StateDbAccess, ) -> io::Result> { - find_thread_path_by_id_str_in_subdir(codex_home, ARCHIVED_SESSIONS_SUBDIR, id_str, state_db_ctx) - .await + find_thread_path_by_id_str_in_subdir( + codex_home, + ARCHIVED_SESSIONS_SUBDIR, + id_str, + state_db_ctx.as_deref(), + state_db_ctx.metrics(), + ) + .await } /// Extract the `YYYY/MM/DD` directory components from a rollout filename. diff --git a/codex-rs/rollout/src/metadata_tests.rs b/codex-rs/rollout/src/metadata_tests.rs index 45db758c65..54c9a7740f 100644 --- a/codex-rs/rollout/src/metadata_tests.rs +++ b/codex-rs/rollout/src/metadata_tests.rs @@ -187,9 +187,13 @@ async fn backfill_sessions_resumes_from_watermark_and_marks_complete() { /*git*/ None, ); - let runtime = codex_state::StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = codex_state::StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let first_watermark = backfill_watermark_for_path(codex_home.as_path(), first_path.as_path()); runtime.mark_backfill_running().await.expect("mark running"); runtime @@ -252,9 +256,13 @@ async fn backfill_sessions_preserves_existing_git_branch_and_fills_missing_git_f }), ); - let runtime = codex_state::StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = codex_state::StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let thread_id = ThreadId::from_string(&thread_uuid.to_string()).expect("thread id"); let mut existing = extract_metadata_from_rollout(&rollout_path, "test-provider") .await @@ -298,9 +306,13 @@ async fn backfill_sessions_normalizes_cwd_before_upsert() { /*git*/ None, ); - let runtime = codex_state::StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = codex_state::StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); backfill_sessions(runtime.as_ref(), codex_home.as_path(), "test-provider").await; diff --git a/codex-rs/rollout/src/recorder.rs b/codex-rs/rollout/src/recorder.rs index 6b549b745c..70eedd7036 100644 --- a/codex-rs/rollout/src/recorder.rs +++ b/codex-rs/rollout/src/recorder.rs @@ -53,6 +53,7 @@ use super::session_index::find_thread_names_by_ids; use crate::config::RolloutConfigView; use crate::default_client::originator; use crate::state_db; +use crate::state_db::StateDbAccess; use crate::state_db::StateDbHandle; use codex_git_utils::collect_git_info; use codex_protocol::protocol::EventMsg; @@ -236,7 +237,7 @@ impl RolloutRecorder { /// List threads (rollout files) under the provided Codex home directory. #[allow(clippy::too_many_arguments)] pub async fn list_threads( - state_db_ctx: Option, + state_db_ctx: StateDbAccess, config: &impl RolloutConfigView, page_size: usize, cursor: Option<&Cursor>, @@ -268,7 +269,7 @@ impl RolloutRecorder { #[allow(clippy::too_many_arguments)] pub async fn list_threads_from_state_db( - state_db_ctx: Option, + state_db_ctx: StateDbAccess, config: &impl RolloutConfigView, page_size: usize, cursor: Option<&Cursor>, @@ -301,7 +302,7 @@ impl RolloutRecorder { /// List archived threads (rollout files) under the archived sessions directory. #[allow(clippy::too_many_arguments)] pub async fn list_archived_threads( - state_db_ctx: Option, + state_db_ctx: StateDbAccess, config: &impl RolloutConfigView, page_size: usize, cursor: Option<&Cursor>, @@ -333,7 +334,7 @@ impl RolloutRecorder { #[allow(clippy::too_many_arguments)] pub async fn list_archived_threads_from_state_db( - state_db_ctx: Option, + state_db_ctx: StateDbAccess, config: &impl RolloutConfigView, page_size: usize, cursor: Option<&Cursor>, @@ -365,7 +366,7 @@ impl RolloutRecorder { #[allow(clippy::too_many_arguments)] async fn list_threads_with_db_fallback( - state_db_ctx: Option, + state_db_ctx: StateDbAccess, config: &impl RolloutConfigView, page_size: usize, cursor: Option<&Cursor>, @@ -447,9 +448,14 @@ impl RolloutRecorder { } }; - if state_db_ctx.is_none() { + if !state_db_ctx.is_available() { // Keep legacy behavior when SQLite is unavailable: return filesystem results // at the requested page size. + codex_state::record_db_fallback_metric( + state_db_ctx.metrics(), + "list_threads", + "db_unavailable", + ); return Ok(page_from_filesystem_scan( fs_page, sort_direction, @@ -558,6 +564,11 @@ impl RolloutRecorder { ) .await; } + codex_state::record_db_fallback_metric( + state_db_ctx.metrics(), + "list_threads", + "metadata_filter", + ); let page = page_from_filesystem_scan(fs_page, sort_direction, page_size, sort_key); return Ok(fill_missing_thread_item_metadata_from_state_db( state_db_ctx.as_deref(), @@ -569,6 +580,11 @@ impl RolloutRecorder { } if listing_has_metadata_filters { let page = page_from_filesystem_scan(fs_page, sort_direction, page_size, sort_key); + codex_state::record_db_fallback_metric( + state_db_ctx.metrics(), + "list_threads", + "db_error", + ); return Ok(fill_missing_thread_item_metadata_from_state_db( state_db_ctx.as_deref(), page, @@ -578,6 +594,7 @@ impl RolloutRecorder { // If SQLite listing still fails, return the filesystem page rather than failing the list. tracing::error!("Falling back on rollout system"); tracing::warn!("state db discrepancy during list_threads_with_db_fallback: falling_back"); + codex_state::record_db_fallback_metric(state_db_ctx.metrics(), "list_threads", "db_error"); Ok(page_from_filesystem_scan( fs_page, sort_direction, @@ -589,7 +606,7 @@ impl RolloutRecorder { /// Find the newest recorded thread path, optionally filtering to a matching cwd. #[allow(clippy::too_many_arguments)] pub async fn find_latest_thread_path( - state_db_ctx: Option, + state_db_ctx: StateDbAccess, config: &impl RolloutConfigView, page_size: usize, cursor: Option<&Cursor>, @@ -601,10 +618,11 @@ impl RolloutRecorder { ) -> std::io::Result> { let codex_home = config.codex_home(); let cwd_filter = filter_cwd.map(Path::to_path_buf); - if state_db_ctx.is_some() { + let mut fallback_reason = (!state_db_ctx.is_available()).then_some("db_unavailable"); + if state_db_ctx.is_available() { let mut db_cursor = cursor.cloned(); loop { - let Some(db_page) = state_db::list_threads_db( + let db_page = match state_db::list_threads_db( state_db_ctx.as_deref(), codex_home, page_size, @@ -618,8 +636,12 @@ impl RolloutRecorder { /*search_term*/ None, ) .await - else { - break; + { + Some(db_page) => db_page, + None => { + fallback_reason = Some("db_error"); + break; + } }; if let Some(path) = select_resume_path_from_db_page(&db_page, filter_cwd, default_provider).await @@ -628,11 +650,20 @@ impl RolloutRecorder { } db_cursor = db_page.next_anchor.map(Into::into); if db_cursor.is_none() { + fallback_reason = Some("missing_row"); break; } } } + if let Some(reason) = fallback_reason { + codex_state::record_db_fallback_metric( + state_db_ctx.metrics(), + "find_latest_thread_path", + reason, + ); + } + let mut cursor = cursor.cloned(); loop { let page = get_threads( diff --git a/codex-rs/rollout/src/recorder_tests.rs b/codex-rs/rollout/src/recorder_tests.rs index a8b946934b..5a1f633938 100644 --- a/codex-rs/rollout/src/recorder_tests.rs +++ b/codex-rs/rollout/src/recorder_tests.rs @@ -18,11 +18,14 @@ use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::TurnContextItem; use codex_protocol::protocol::UserMessageEvent; use pretty_assertions::assert_eq; +use std::collections::BTreeMap; use std::fs; use std::fs::File; use std::io::Write; use std::path::Path; use std::path::PathBuf; +use std::sync::Arc; +use std::sync::Mutex; use std::time::Duration; use tempfile::TempDir; use uuid::Uuid; @@ -37,6 +40,45 @@ fn test_config(codex_home: &Path) -> RolloutConfig { } } +#[derive(Default)] +struct TestMetrics { + events: Mutex>, +} + +#[derive(Debug, Eq, PartialEq)] +struct MetricEvent { + name: String, + tags: BTreeMap, +} + +impl TestMetrics { + fn has_fallback(&self, caller: &str, reason: &str) -> bool { + self.events + .lock() + .expect("metrics lock") + .iter() + .any(|event| { + event.name == codex_state::DB_FALLBACK_METRIC + && event.tags.get("caller").map(String::as_str) == Some(caller) + && event.tags.get("reason").map(String::as_str) == Some(reason) + }) + } +} + +impl codex_state::DbMetricsRecorder for TestMetrics { + fn counter(&self, name: &str, _inc: i64, tags: &[(&str, &str)]) { + self.events.lock().expect("metrics lock").push(MetricEvent { + name: name.to_string(), + tags: tags + .iter() + .map(|(key, value)| ((*key).to_string(), (*value).to_string())) + .collect(), + }); + } + + fn record_duration(&self, _name: &str, _duration: Duration, _tags: &[(&str, &str)]) {} +} + fn write_session_file(root: &Path, ts: &str, uuid: Uuid) -> std::io::Result { let day_dir = root.join("sessions/2025/01/03"); fs::create_dir_all(&day_dir)?; @@ -124,7 +166,7 @@ async fn state_db_init_backfills_before_returning() -> anyhow::Result<()> { .join("\n"); fs::write(&rollout_path, format!("{jsonl}\n"))?; - let runtime = crate::state_db::init(&test_config(home.path())) + let runtime = crate::state_db::init(&test_config(home.path()), /*metrics*/ None) .await .expect("state db should initialize"); @@ -535,9 +577,13 @@ async fn metadata_irrelevant_events_coalesce_state_db_updated_at() -> std::io::R let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); - let state_db = StateRuntime::init(home.path().to_path_buf(), config.model_provider_id.clone()) - .await - .expect("state db should initialize"); + let state_db = StateRuntime::init( + home.path().to_path_buf(), + config.model_provider_id.clone(), + /*metrics*/ None, + ) + .await + .expect("state db should initialize"); state_db .mark_backfill_complete(/*last_watermark*/ None) .await @@ -639,9 +685,13 @@ async fn shutdown_flushes_pending_metadata_irrelevant_updated_at() -> std::io::R let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); - let state_db = StateRuntime::init(home.path().to_path_buf(), config.model_provider_id.clone()) - .await - .expect("state db should initialize"); + let state_db = StateRuntime::init( + home.path().to_path_buf(), + config.model_provider_id.clone(), + /*metrics*/ None, + ) + .await + .expect("state db should initialize"); state_db .mark_backfill_complete(/*last_watermark*/ None) .await @@ -699,9 +749,13 @@ async fn metadata_irrelevant_events_fall_back_to_upsert_when_thread_missing() -> let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); - let state_db = StateRuntime::init(home.path().to_path_buf(), config.model_provider_id.clone()) - .await - .expect("state db should initialize"); + let state_db = StateRuntime::init( + home.path().to_path_buf(), + config.model_provider_id.clone(), + /*metrics*/ None, + ) + .await + .expect("state db should initialize"); let thread_id = ThreadId::new(); let rollout_path = home.path().join("rollout.jsonl"); let builder = ThreadMetadataBuilder::new( @@ -751,7 +805,7 @@ async fn list_threads_db_disabled_does_not_skip_paginated_items() -> std::io::Re let default_provider = config.model_provider_id.clone(); let page1 = RolloutRecorder::list_threads( - /*state_db_ctx*/ None, + StateDbAccess::none(), &config, /*page_size*/ 1, /*cursor*/ None, @@ -769,7 +823,7 @@ async fn list_threads_db_disabled_does_not_skip_paginated_items() -> std::io::Re let cursor = page1.next_cursor.clone().expect("cursor should be present"); let page2 = RolloutRecorder::list_threads( - /*state_db_ctx*/ None, + StateDbAccess::none(), &config, /*page_size*/ 1, Some(&cursor), @@ -787,6 +841,62 @@ async fn list_threads_db_disabled_does_not_skip_paginated_items() -> std::io::Re Ok(()) } +#[tokio::test] +async fn list_threads_records_db_unavailable_fallback_metric() -> std::io::Result<()> { + let home = TempDir::new().expect("temp dir"); + let config = test_config(home.path()); + let _path = write_session_file(home.path(), "2025-01-03T12-00-00", Uuid::from_u128(9101))?; + let metrics = Arc::new(TestMetrics::default()); + let metrics_handle: codex_state::DbMetricsRecorderHandle = metrics.clone(); + let state_db_ctx = StateDbAccess::none().with_standalone_metrics(Some(metrics_handle)); + + let page = RolloutRecorder::list_threads( + state_db_ctx, + &config, + /*page_size*/ 10, + /*cursor*/ None, + ThreadSortKey::CreatedAt, + SortDirection::Desc, + &[], + /*model_providers*/ None, + /*cwd_filters*/ None, + config.model_provider_id.as_str(), + /*search_term*/ None, + ) + .await?; + + assert_eq!(page.items.len(), 1); + assert!(metrics.has_fallback("list_threads", "db_unavailable")); + Ok(()) +} + +#[tokio::test] +async fn find_latest_thread_path_records_db_unavailable_fallback_metric() -> std::io::Result<()> { + let home = TempDir::new().expect("temp dir"); + let config = test_config(home.path()); + let path = write_session_file(home.path(), "2025-01-03T12-00-00", Uuid::from_u128(9103))?; + let metrics = Arc::new(TestMetrics::default()); + let metrics_handle: codex_state::DbMetricsRecorderHandle = metrics.clone(); + let state_db_ctx = StateDbAccess::none().with_standalone_metrics(Some(metrics_handle)); + + let found = RolloutRecorder::find_latest_thread_path( + state_db_ctx, + &config, + /*page_size*/ 10, + /*cursor*/ None, + ThreadSortKey::CreatedAt, + &[], + /*model_providers*/ None, + config.model_provider_id.as_str(), + /*filter_cwd*/ None, + ) + .await?; + + assert_eq!(found, Some(path)); + assert!(metrics.has_fallback("find_latest_thread_path", "db_unavailable")); + Ok(()) +} + #[tokio::test] async fn list_threads_db_enabled_drops_missing_rollout_paths() -> std::io::Result<()> { let home = TempDir::new().expect("temp dir"); @@ -801,6 +911,7 @@ async fn list_threads_db_enabled_drops_missing_rollout_paths() -> std::io::Resul let runtime = codex_state::StateRuntime::init( home.path().to_path_buf(), config.model_provider_id.clone(), + /*metrics*/ None, ) .await .expect("state db should initialize"); @@ -829,7 +940,7 @@ async fn list_threads_db_enabled_drops_missing_rollout_paths() -> std::io::Resul let default_provider = config.model_provider_id.clone(); let page = RolloutRecorder::list_threads( - Some(runtime.clone()), + StateDbAccess::new(Some(runtime.clone())), &config, /*page_size*/ 10, /*cursor*/ None, @@ -866,6 +977,7 @@ async fn list_threads_db_enabled_repairs_stale_rollout_paths() -> std::io::Resul let runtime = codex_state::StateRuntime::init( home.path().to_path_buf(), config.model_provider_id.clone(), + /*metrics*/ None, ) .await .expect("state db should initialize"); @@ -894,7 +1006,7 @@ async fn list_threads_db_enabled_repairs_stale_rollout_paths() -> std::io::Resul let default_provider = config.model_provider_id.clone(); let page = RolloutRecorder::list_threads( - Some(runtime.clone()), + StateDbAccess::new(Some(runtime.clone())), &config, /*page_size*/ 1, /*cursor*/ None, @@ -926,6 +1038,7 @@ async fn list_threads_state_db_only_skips_jsonl_repair_scan() -> std::io::Result let runtime = codex_state::StateRuntime::init( home.path().to_path_buf(), config.model_provider_id.clone(), + /*metrics*/ None, ) .await .expect("state db should initialize"); @@ -967,7 +1080,7 @@ async fn list_threads_state_db_only_skips_jsonl_repair_scan() -> std::io::Result let cwd_filters = [home.path().to_path_buf()]; let state_db_only_page = RolloutRecorder::list_threads_from_state_db( - Some(runtime.clone()), + StateDbAccess::new(Some(runtime.clone())), &config, /*page_size*/ 10, /*cursor*/ None, @@ -983,7 +1096,7 @@ async fn list_threads_state_db_only_skips_jsonl_repair_scan() -> std::io::Result assert_eq!(state_db_only_page.items.len(), 0); let repaired_page = RolloutRecorder::list_threads( - Some(runtime.clone()), + StateDbAccess::new(Some(runtime.clone())), &config, /*page_size*/ 10, /*cursor*/ None, @@ -999,7 +1112,7 @@ async fn list_threads_state_db_only_skips_jsonl_repair_scan() -> std::io::Result assert_eq!(repaired_page.items.len(), 1); let repaired_state_db_only_page = RolloutRecorder::list_threads_from_state_db( - Some(runtime.clone()), + StateDbAccess::new(Some(runtime.clone())), &config, /*page_size*/ 10, /*cursor*/ None, @@ -1029,6 +1142,7 @@ async fn list_threads_default_filter_returns_filesystem_scan_results() -> std::i let runtime = codex_state::StateRuntime::init( home.path().to_path_buf(), config.model_provider_id.clone(), + /*metrics*/ None, ) .await .expect("state db should initialize"); @@ -1057,7 +1171,7 @@ async fn list_threads_default_filter_returns_filesystem_scan_results() -> std::i let cwd_filters = [stale_cwd]; let state_db_only_page = RolloutRecorder::list_threads_from_state_db( - Some(runtime.clone()), + StateDbAccess::new(Some(runtime.clone())), &config, /*page_size*/ 10, /*cursor*/ None, @@ -1073,7 +1187,7 @@ async fn list_threads_default_filter_returns_filesystem_scan_results() -> std::i assert_eq!(state_db_only_page.items.len(), 1); let scanned_page = RolloutRecorder::list_threads( - Some(runtime.clone()), + StateDbAccess::new(Some(runtime.clone())), &config, /*page_size*/ 10, /*cursor*/ None, @@ -1089,7 +1203,7 @@ async fn list_threads_default_filter_returns_filesystem_scan_results() -> std::i assert_eq!(scanned_page.items.len(), 0); let repaired_state_db_only_page = RolloutRecorder::list_threads_from_state_db( - Some(runtime.clone()), + StateDbAccess::new(Some(runtime.clone())), &config, /*page_size*/ 10, /*cursor*/ None, @@ -1118,6 +1232,7 @@ async fn list_threads_metadata_filter_overlays_state_db_list_metadata() -> std:: let runtime = codex_state::StateRuntime::init( home.path().to_path_buf(), config.model_provider_id.clone(), + /*metrics*/ None, ) .await .expect("state db should initialize"); @@ -1148,7 +1263,7 @@ async fn list_threads_metadata_filter_overlays_state_db_list_metadata() -> std:: .expect("state db upsert should succeed"); let page = RolloutRecorder::list_threads( - Some(runtime.clone()), + StateDbAccess::new(Some(runtime.clone())), &config, /*page_size*/ 10, /*cursor*/ None, @@ -1247,6 +1362,7 @@ async fn list_threads_search_repairs_stale_state_db_hits_before_returning() -> s let runtime = codex_state::StateRuntime::init( home.path().to_path_buf(), config.model_provider_id.clone(), + /*metrics*/ None, ) .await .expect("state db should initialize"); @@ -1275,7 +1391,7 @@ async fn list_threads_search_repairs_stale_state_db_hits_before_returning() -> s .expect("state db upsert should succeed"); let stale_state_db_only_page = RolloutRecorder::list_threads_from_state_db( - Some(runtime.clone()), + StateDbAccess::new(Some(runtime.clone())), &config, /*page_size*/ 10, /*cursor*/ None, @@ -1291,7 +1407,7 @@ async fn list_threads_search_repairs_stale_state_db_hits_before_returning() -> s assert_eq!(stale_state_db_only_page.items.len(), 1); let scanned_page = RolloutRecorder::list_threads( - Some(runtime.clone()), + StateDbAccess::new(Some(runtime.clone())), &config, /*page_size*/ 10, /*cursor*/ None, @@ -1307,7 +1423,7 @@ async fn list_threads_search_repairs_stale_state_db_hits_before_returning() -> s assert_eq!(scanned_page.items.len(), 0); let repaired_state_db_only_page = RolloutRecorder::list_threads_from_state_db( - Some(runtime.clone()), + StateDbAccess::new(Some(runtime.clone())), &config, /*page_size*/ 10, /*cursor*/ None, diff --git a/codex-rs/rollout/src/session_index.rs b/codex-rs/rollout/src/session_index.rs index e227515357..2bc8390b32 100644 --- a/codex-rs/rollout/src/session_index.rs +++ b/codex-rs/rollout/src/session_index.rs @@ -7,6 +7,7 @@ use std::io::SeekFrom; use std::path::Path; use std::path::PathBuf; +use crate::state_db::StateDbAccess; use codex_protocol::ThreadId; use codex_protocol::protocol::SessionMetaLine; use serde::Deserialize; @@ -117,7 +118,7 @@ pub async fn find_thread_names_by_ids( pub async fn find_thread_meta_by_name_str( codex_home: &Path, name: &str, - state_db_ctx: Option<&codex_state::StateRuntime>, + state_db_ctx: &StateDbAccess, ) -> std::io::Result> { if name.trim().is_empty() { return Ok(None); diff --git a/codex-rs/rollout/src/session_index_tests.rs b/codex-rs/rollout/src/session_index_tests.rs index 757b08b4d4..c3753b06ed 100644 --- a/codex-rs/rollout/src/session_index_tests.rs +++ b/codex-rs/rollout/src/session_index_tests.rs @@ -100,7 +100,8 @@ async fn find_thread_meta_by_name_str_skips_newest_entry_without_rollout() -> st ]; write_index(&path, &lines)?; - let found = find_thread_meta_by_name_str(temp.path(), "same", /*state_db_ctx*/ None).await?; + let found = + find_thread_meta_by_name_str(temp.path(), "same", &crate::StateDbAccess::none()).await?; assert_eq!( found.map(|(path, session_meta)| (path, session_meta.meta.id)), @@ -137,7 +138,8 @@ async fn find_thread_meta_by_name_str_skips_partial_rollout() -> std::io::Result ]; write_index(&path, &lines)?; - let found = find_thread_meta_by_name_str(temp.path(), "same", /*state_db_ctx*/ None).await?; + let found = + find_thread_meta_by_name_str(temp.path(), "same", &crate::StateDbAccess::none()).await?; assert_eq!(found.map(|(path, _)| path), Some(saved_rollout_path)); Ok(()) @@ -175,7 +177,8 @@ async fn find_thread_meta_by_name_str_ignores_historical_name_after_rename() -> ]; write_index(&path, &lines)?; - let found = find_thread_meta_by_name_str(temp.path(), "same", /*state_db_ctx*/ None).await?; + let found = + find_thread_meta_by_name_str(temp.path(), "same", &crate::StateDbAccess::none()).await?; assert_eq!(found.map(|(path, _)| path), Some(current_rollout_path)); Ok(()) diff --git a/codex-rs/rollout/src/sqlite_metrics.rs b/codex-rs/rollout/src/sqlite_metrics.rs new file mode 100644 index 0000000000..3f83306a1f --- /dev/null +++ b/codex-rs/rollout/src/sqlite_metrics.rs @@ -0,0 +1,43 @@ +use std::sync::Arc; +use std::time::Duration; + +use codex_otel::ORIGINATOR_TAG; +use codex_otel::bounded_originator_tag_value; +use codex_state::DbMetricsRecorder; +use codex_state::DbMetricsRecorderHandle; + +struct OtelDbMetrics { + metrics: codex_otel::MetricsClient, + originator: &'static str, +} + +impl DbMetricsRecorder for OtelDbMetrics { + fn counter(&self, name: &str, inc: i64, tags: &[(&str, &str)]) { + let tags = sqlite_originator_tags(tags, self.originator); + let _ = self.metrics.counter(name, inc, &tags); + } + + fn record_duration(&self, name: &str, duration: Duration, tags: &[(&str, &str)]) { + let tags = sqlite_originator_tags(tags, self.originator); + let _ = self.metrics.record_duration(name, duration, &tags); + } +} + +pub(crate) fn recorder( + metrics: codex_otel::MetricsClient, + originator: &str, +) -> DbMetricsRecorderHandle { + Arc::new(OtelDbMetrics { + metrics, + originator: bounded_originator_tag_value(originator), + }) +} + +fn sqlite_originator_tags<'a>( + tags: &[(&'a str, &'a str)], + originator: &'static str, +) -> Vec<(&'a str, &'a str)> { + let mut tags = tags.to_vec(); + tags.push((ORIGINATOR_TAG, originator)); + tags +} diff --git a/codex-rs/rollout/src/state_db.rs b/codex-rs/rollout/src/state_db.rs index d039e16d68..156ef6df46 100644 --- a/codex-rs/rollout/src/state_db.rs +++ b/codex-rs/rollout/src/state_db.rs @@ -4,6 +4,7 @@ use crate::list::Cursor; use crate::list::SortDirection; use crate::list::ThreadSortKey; use crate::metadata; +use crate::sqlite_metrics; use chrono::DateTime; use chrono::Utc; use codex_protocol::ThreadId; @@ -25,6 +26,61 @@ use tracing::warn; /// Core-facing handle to the SQLite-backed state runtime. pub type StateDbHandle = Arc; +/// Access to SQLite-backed state plus fallback telemetry for local rollout operations. +/// +/// Some fallback metrics need to be recorded after SQLite initialization fails, +/// so callers must carry the metrics sink independently from the runtime handle. +#[derive(Clone)] +pub struct StateDbAccess { + state_db: Option, + standalone_metrics: Option, +} + +impl StateDbAccess { + /// Construct an access bundle with no DB handle and no fallback metrics recorder. + pub fn none() -> Self { + Self { + state_db: None, + standalone_metrics: None, + } + } + + pub fn new(state_db: Option) -> Self { + Self { + state_db, + standalone_metrics: None, + } + } + + pub fn with_standalone_metrics( + mut self, + standalone_metrics: Option, + ) -> Self { + self.standalone_metrics = standalone_metrics; + self + } + + pub fn state_db(&self) -> Option { + self.state_db.clone() + } + + pub fn as_deref(&self) -> Option<&codex_state::StateRuntime> { + self.state_db.as_deref() + } + + pub fn is_available(&self) -> bool { + self.state_db.is_some() + } + + pub fn metrics(&self) -> Option<&dyn codex_state::DbMetricsRecorder> { + self.standalone_metrics.as_deref().or_else(|| { + self.state_db + .as_deref() + .and_then(codex_state::StateRuntime::metrics) + }) + } +} + #[cfg(not(test))] const STARTUP_BACKFILL_POLL_INTERVAL: Duration = Duration::from_secs(1); #[cfg(test)] @@ -39,12 +95,16 @@ const STARTUP_BACKFILL_WAIT_TIMEOUT: Duration = Duration::from_secs(2); /// This is the process entry point for local state: it opens the SQLite-backed /// runtime, applies rollout metadata backfills as needed, and returns the /// initialized handle. -pub async fn init(config: &impl RolloutConfigView) -> Option { +pub async fn init( + config: &impl RolloutConfigView, + metrics: Option, +) -> Option { let config = RolloutConfig::from_view(config); match try_init_with_roots( config.codex_home, config.sqlite_home, config.model_provider_id, + metrics, ) .await { @@ -60,12 +120,16 @@ pub async fn init(config: &impl RolloutConfigView) -> Option { /// /// Prefer [`init`] unless the caller needs to surface the exact failure after /// tracing or UI setup has completed. -pub async fn try_init(config: &impl RolloutConfigView) -> anyhow::Result { +pub async fn try_init( + config: &impl RolloutConfigView, + metrics: Option, +) -> anyhow::Result { let config = RolloutConfig::from_view(config); try_init_with_roots( config.codex_home, config.sqlite_home, config.model_provider_id, + metrics, ) .await } @@ -74,11 +138,13 @@ async fn try_init_with_roots( codex_home: PathBuf, sqlite_home: PathBuf, default_model_provider_id: String, + metrics: Option, ) -> anyhow::Result { try_init_with_roots_inner( codex_home, sqlite_home, default_model_provider_id, + metrics, /*backfill_lease_seconds*/ None, ) .await @@ -95,6 +161,7 @@ async fn try_init_with_roots_and_backfill_lease( codex_home, sqlite_home, default_model_provider_id, + /*metrics*/ None, Some(backfill_lease_seconds), ) .await @@ -104,54 +171,82 @@ async fn try_init_with_roots_inner( codex_home: PathBuf, sqlite_home: PathBuf, default_model_provider_id: String, + metrics: Option, backfill_lease_seconds: Option, ) -> anyhow::Result { - let runtime = - codex_state::StateRuntime::init(sqlite_home.clone(), default_model_provider_id.clone()) - .await - .map_err(|err| { - anyhow::anyhow!( - "failed to initialize state runtime at {}: {err}", - sqlite_home.display() - ) - })?; + let runtime = codex_state::StateRuntime::init( + sqlite_home.clone(), + default_model_provider_id.clone(), + metrics.clone(), + ) + .await + .map_err(|err| { + anyhow::anyhow!( + "failed to initialize state runtime at {}: {err}", + sqlite_home.display() + ) + })?; + let backfill_gate_started = Instant::now(); + let backfill_gate_result = wait_for_startup_backfill( + runtime.as_ref(), + codex_home.as_path(), + default_model_provider_id.as_str(), + backfill_lease_seconds, + ) + .await; + codex_state::record_db_init_backfill_gate_metric( + metrics.as_deref(), + backfill_gate_started.elapsed(), + &backfill_gate_result, + ); + backfill_gate_result?; + Ok(runtime) +} + +async fn wait_for_startup_backfill( + runtime: &codex_state::StateRuntime, + codex_home: &Path, + default_model_provider_id: &str, + backfill_lease_seconds: Option, +) -> anyhow::Result<()> { let wait_started = Instant::now(); let mut reported_wait = false; loop { - let backfill_state = runtime.get_backfill_state().await.map_err(|err| { - anyhow::anyhow!( - "failed to read backfill state at {}: {err}", - codex_home.display() - ) - })?; + let backfill_state = match runtime.get_backfill_state().await { + Ok(state) => state, + Err(err) => { + return Err(anyhow::anyhow!( + "failed to read backfill state at {}: {err}", + codex_home.display() + )); + } + }; if backfill_state.status == codex_state::BackfillStatus::Complete { - return Ok(runtime); + return Ok(()); } if let Some(backfill_lease_seconds) = backfill_lease_seconds { metadata::backfill_sessions_with_lease( - runtime.as_ref(), - codex_home.as_path(), - default_model_provider_id.as_str(), + runtime, + codex_home, + default_model_provider_id, backfill_lease_seconds, ) .await; } else { - metadata::backfill_sessions( - runtime.as_ref(), - codex_home.as_path(), - default_model_provider_id.as_str(), - ) - .await; + metadata::backfill_sessions(runtime, codex_home, default_model_provider_id).await; } - let backfill_state = runtime.get_backfill_state().await.map_err(|err| { - anyhow::anyhow!( - "failed to read backfill state at {} after startup backfill: {err}", - codex_home.display() - ) - })?; + let backfill_state = match runtime.get_backfill_state().await { + Ok(state) => state, + Err(err) => { + return Err(anyhow::anyhow!( + "failed to read backfill state at {} after startup backfill: {err}", + codex_home.display() + )); + } + }; if backfill_state.status == codex_state::BackfillStatus::Complete { - return Ok(runtime); + return Ok(()); } if wait_started.elapsed() >= STARTUP_BACKFILL_WAIT_TIMEOUT { return Err(anyhow::anyhow!( @@ -192,39 +287,65 @@ fn emit_startup_warning(message: &str) { /// /// Unlike [`init`], this helper does not run rollout backfill. It is for /// optional local reads from non-owning contexts such as remote app-server mode. -pub async fn get_state_db(config: &impl RolloutConfigView) -> Option { +pub async fn get_state_db( + config: &impl RolloutConfigView, + metrics: Option, +) -> Option { let state_path = codex_state::state_db_path(config.sqlite_home()); if !tokio::fs::try_exists(&state_path).await.unwrap_or(false) { + codex_state::record_db_fallback_metric( + metrics.as_deref(), + "get_state_db", + "db_unavailable", + ); return None; } - let runtime = codex_state::StateRuntime::init( + let runtime = match codex_state::StateRuntime::init( config.sqlite_home().to_path_buf(), config.model_provider_id().to_string(), + metrics.clone(), ) .await - .ok()?; - require_backfill_complete(runtime, config.sqlite_home()).await + { + Ok(runtime) => runtime, + Err(_) => { + codex_state::record_db_fallback_metric(metrics.as_deref(), "get_state_db", "db_error"); + return None; + } + }; + require_backfill_complete(runtime, config.sqlite_home(), metrics.as_deref()).await +} + +/// Build a SQLite metrics recorder backed by an OTEL metrics client. +pub fn otel_db_metrics_recorder( + metrics: codex_otel::MetricsClient, + originator: &str, +) -> codex_state::DbMetricsRecorderHandle { + sqlite_metrics::recorder(metrics, originator) } async fn require_backfill_complete( runtime: StateDbHandle, - codex_home: &Path, + sqlite_home: &Path, + metrics: Option<&dyn codex_state::DbMetricsRecorder>, ) -> Option { match runtime.get_backfill_state().await { Ok(state) if state.status == codex_state::BackfillStatus::Complete => Some(runtime), Ok(state) => { warn!( "state db backfill not complete at {} (status: {})", - codex_home.display(), + sqlite_home.display(), state.status.as_str() ); + codex_state::record_db_fallback_metric(metrics, "get_state_db", "backfill_incomplete"); None } Err(err) => { warn!( "failed to read backfill state at {}: {err}", - codex_home.display() + sqlite_home.display() ); + codex_state::record_db_fallback_metric(metrics, "get_state_db", "db_error"); None } } diff --git a/codex-rs/rollout/src/state_db_tests.rs b/codex-rs/rollout/src/state_db_tests.rs index 10a9a3da13..eb98bb081b 100644 --- a/codex-rs/rollout/src/state_db_tests.rs +++ b/codex-rs/rollout/src/state_db_tests.rs @@ -27,9 +27,12 @@ fn cursor_to_anchor_normalizes_timestamp_format() { #[tokio::test] async fn try_init_waits_for_concurrent_startup_backfill() -> anyhow::Result<()> { let home = TempDir::new().expect("temp dir"); - let runtime = - codex_state::StateRuntime::init(home.path().to_path_buf(), "test-provider".to_string()) - .await?; + let runtime = codex_state::StateRuntime::init( + home.path().to_path_buf(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await?; let claimed = runtime.try_claim_backfill(/*lease_seconds*/ 60).await?; assert!(claimed); let runtime_for_completion = runtime.clone(); @@ -59,9 +62,12 @@ async fn try_init_waits_for_concurrent_startup_backfill() -> anyhow::Result<()> #[tokio::test] async fn try_init_times_out_waiting_for_stuck_startup_backfill() -> anyhow::Result<()> { let home = TempDir::new().expect("temp dir"); - let runtime = - codex_state::StateRuntime::init(home.path().to_path_buf(), "test-provider".to_string()) - .await?; + let runtime = codex_state::StateRuntime::init( + home.path().to_path_buf(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await?; let claimed = runtime.try_claim_backfill(/*lease_seconds*/ 60).await?; assert!(claimed); diff --git a/codex-rs/rollout/src/tests.rs b/codex-rs/rollout/src/tests.rs index fdfed3dadc..40114296c4 100644 --- a/codex-rs/rollout/src/tests.rs +++ b/codex-rs/rollout/src/tests.rs @@ -19,6 +19,7 @@ use time::macros::format_description; use uuid::Uuid; use crate::INTERACTIVE_SESSION_SOURCES; +use crate::StateDbAccess; use crate::find_thread_path_by_id_str; use crate::list::Cursor; use crate::list::ThreadItem; @@ -59,9 +60,13 @@ async fn insert_state_db_thread( rollout_path: &Path, archived: bool, ) -> crate::state_db::StateDbHandle { - let runtime = codex_state::StateRuntime::init(home.to_path_buf(), TEST_PROVIDER.to_string()) - .await - .expect("state db should initialize"); + let runtime = codex_state::StateRuntime::init( + home.to_path_buf(), + TEST_PROVIDER.to_string(), + /*metrics*/ None, + ) + .await + .expect("state db should initialize"); runtime .mark_backfill_complete(/*last_watermark*/ None) .await @@ -245,9 +250,13 @@ async fn find_thread_path_falls_back_when_db_path_is_stale() { ) .await; - let found = find_thread_path_by_id_str(home, &uuid.to_string(), Some(runtime.as_ref())) - .await - .expect("lookup should succeed"); + let found = find_thread_path_by_id_str( + home, + &uuid.to_string(), + &StateDbAccess::new(Some(runtime.clone())), + ) + .await + .expect("lookup should succeed"); assert_eq!(found, Some(fs_rollout_path.clone())); assert_state_db_rollout_path(home, thread_id, Some(fs_rollout_path.as_path())).await; } @@ -290,9 +299,13 @@ async fn find_thread_path_falls_back_when_db_path_points_to_another_thread() { ) .await; - let found = find_thread_path_by_id_str(home, &uuid.to_string(), Some(runtime.as_ref())) - .await - .expect("lookup should succeed"); + let found = find_thread_path_by_id_str( + home, + &uuid.to_string(), + &StateDbAccess::new(Some(runtime.clone())), + ) + .await + .expect("lookup should succeed"); assert_eq!(found, Some(fs_rollout_path.clone())); assert_state_db_rollout_path(home, thread_id, Some(fs_rollout_path.as_path())).await; } @@ -315,17 +328,25 @@ async fn find_thread_path_repairs_missing_db_row_after_filesystem_fallback() { let fs_rollout_path = home.join(format!("sessions/2025/01/03/rollout-{ts}-{uuid}.jsonl")); // Create an empty state DB so lookup takes the DB-first path and then falls back to files. - let runtime = codex_state::StateRuntime::init(home.to_path_buf(), TEST_PROVIDER.to_string()) - .await - .expect("state db should initialize"); + let runtime = codex_state::StateRuntime::init( + home.to_path_buf(), + TEST_PROVIDER.to_string(), + /*metrics*/ None, + ) + .await + .expect("state db should initialize"); runtime .mark_backfill_complete(/*last_watermark*/ None) .await .expect("backfill should be complete"); - let found = find_thread_path_by_id_str(home, &uuid.to_string(), Some(runtime.as_ref())) - .await - .expect("lookup should succeed"); + let found = find_thread_path_by_id_str( + home, + &uuid.to_string(), + &StateDbAccess::new(Some(runtime.clone())), + ) + .await + .expect("lookup should succeed"); assert_eq!(found, Some(fs_rollout_path.clone())); assert_state_db_rollout_path(home, thread_id, Some(fs_rollout_path.as_path())).await; } @@ -347,9 +368,13 @@ async fn find_thread_path_accepts_existing_state_db_path_without_canonical_filen ) .await; - let found = find_thread_path_by_id_str(home, &uuid.to_string(), Some(runtime.as_ref())) - .await - .expect("lookup should succeed"); + let found = find_thread_path_by_id_str( + home, + &uuid.to_string(), + &StateDbAccess::new(Some(runtime.clone())), + ) + .await + .expect("lookup should succeed"); assert_eq!(found, Some(db_rollout_path)); } @@ -368,9 +393,13 @@ async fn assert_state_db_rollout_path( thread_id: ThreadId, expected_path: Option<&Path>, ) { - let runtime = codex_state::StateRuntime::init(home.to_path_buf(), TEST_PROVIDER.to_string()) - .await - .expect("state db should initialize"); + let runtime = codex_state::StateRuntime::init( + home.to_path_buf(), + TEST_PROVIDER.to_string(), + /*metrics*/ None, + ) + .await + .expect("state db should initialize"); let path = runtime .find_rollout_path_by_id(thread_id, Some(false)) .await diff --git a/codex-rs/state/src/bin/logs_client.rs b/codex-rs/state/src/bin/logs_client.rs index bc45c24eaa..d7c0d4ca29 100644 --- a/codex-rs/state/src/bin/logs_client.rs +++ b/codex-rs/state/src/bin/logs_client.rs @@ -111,7 +111,8 @@ async fn main() -> anyhow::Result<()> { .parent() .map(ToOwned::to_owned) .unwrap_or_else(|| PathBuf::from(".")); - let runtime = StateRuntime::init(codex_home, "logs-client".to_string()).await?; + let runtime = + StateRuntime::init(codex_home, "logs-client".to_string(), /*metrics*/ None).await?; let mut last_id = print_backfill(runtime.as_ref(), &filter, args.backfill, args.compact).await?; diff --git a/codex-rs/state/src/lib.rs b/codex-rs/state/src/lib.rs index 84582370a5..fd171ce63a 100644 --- a/codex-rs/state/src/lib.rs +++ b/codex-rs/state/src/lib.rs @@ -10,12 +10,13 @@ mod migrations; mod model; mod paths; mod runtime; +mod telemetry; pub use model::LogEntry; pub use model::LogQuery; pub use model::LogRow; pub use model::Phase2JobClaimOutcome; -/// Preferred entrypoint: owns configuration and metrics. +/// Preferred entrypoint: owns SQLite configuration and optional metrics injection. pub use runtime::StateRuntime; /// Low-level storage engine: useful for focused tests. @@ -56,6 +57,8 @@ pub use runtime::logs_db_filename; pub use runtime::logs_db_path; pub use runtime::state_db_filename; pub use runtime::state_db_path; +pub use telemetry::DbMetricsRecorder; +pub use telemetry::DbMetricsRecorderHandle; /// Environment variable for overriding the SQLite state database home directory. pub const SQLITE_HOME_ENV: &str = "CODEX_SQLITE_HOME"; @@ -71,3 +74,25 @@ pub const DB_ERROR_METRIC: &str = "codex.db.error"; pub const DB_METRIC_BACKFILL: &str = "codex.db.backfill"; /// Metrics on backfill duration. Tags: [status] pub const DB_METRIC_BACKFILL_DURATION_MS: &str = "codex.db.backfill.duration_ms"; +/// SQLite startup initialization attempts. Tags: [status, phase, db, error] +pub const DB_INIT_METRIC: &str = "codex.sqlite.init.count"; +/// SQLite startup initialization duration. Tags: [status, phase, db, error] +pub const DB_INIT_DURATION_METRIC: &str = "codex.sqlite.init.duration_ms"; +/// Filesystem fallback after SQLite could not serve a request. Tags: [caller, reason] +pub const DB_FALLBACK_METRIC: &str = "codex.sqlite.fallback.count"; + +pub fn record_db_fallback_metric( + metrics: Option<&dyn DbMetricsRecorder>, + caller: &'static str, + reason: &'static str, +) { + telemetry::record_fallback(metrics, caller, reason); +} + +pub fn record_db_init_backfill_gate_metric( + metrics: Option<&dyn DbMetricsRecorder>, + duration: std::time::Duration, + result: &anyhow::Result<()>, +) { + telemetry::record_init_backfill_gate(metrics, duration, result); +} diff --git a/codex-rs/state/src/log_db.rs b/codex-rs/state/src/log_db.rs index 6024d6bdaa..24d07b01fb 100644 --- a/codex-rs/state/src/log_db.rs +++ b/codex-rs/state/src/log_db.rs @@ -547,9 +547,13 @@ mod tests { #[tokio::test] async fn sqlite_feedback_logs_match_feedback_formatter_shape() { let codex_home = temp_codex_home(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let writer = SharedWriter::default(); let layer = start(runtime.clone()); @@ -605,9 +609,13 @@ mod tests { #[tokio::test] async fn flush_persists_logs_for_query() { let codex_home = temp_codex_home(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let layer = start(runtime.clone()); let guard = tracing_subscriber::registry() @@ -636,9 +644,13 @@ mod tests { #[tokio::test] async fn configured_batch_size_flushes_without_explicit_flush() { let codex_home = temp_codex_home(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let layer = LogDbLayer::start_with_config( runtime.clone(), LogSinkQueueConfig { @@ -685,9 +697,13 @@ mod tests { #[tokio::test] async fn configured_flush_interval_persists_buffered_logs() { let codex_home = temp_codex_home(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let layer = LogDbLayer::start_with_config( runtime.clone(), LogSinkQueueConfig { diff --git a/codex-rs/state/src/runtime.rs b/codex-rs/state/src/runtime.rs index c8b4e7b98e..0ce73d1521 100644 --- a/codex-rs/state/src/runtime.rs +++ b/codex-rs/state/src/runtime.rs @@ -27,6 +27,9 @@ use crate::model::datetime_to_epoch_millis; use crate::model::datetime_to_epoch_seconds; use crate::model::epoch_millis_to_datetime; use crate::paths::file_modified_time_utc; +use crate::telemetry::DbKind; +use crate::telemetry::DbMetricsRecorder; +use crate::telemetry::DbMetricsRecorderHandle; use chrono::DateTime; use chrono::Utc; use codex_protocol::ThreadId; @@ -52,6 +55,7 @@ use std::path::PathBuf; use std::sync::Arc; use std::sync::atomic::AtomicI64; use std::time::Duration; +use std::time::Instant; use tracing::warn; mod agent_jobs; @@ -85,6 +89,7 @@ pub struct StateRuntime { default_provider: String, pool: Arc, logs_pool: Arc, + metrics: Option, thread_updated_at_millis: Arc, } @@ -93,8 +98,13 @@ impl StateRuntime { /// /// This opens (and migrates) the SQLite databases under `codex_home`, /// keeping logs in a dedicated file to reduce lock contention with the - /// rest of the state store. - pub async fn init(codex_home: PathBuf, default_provider: String) -> anyhow::Result> { + /// rest of the state store. Callers may provide a metrics sink for SQLite + /// startup telemetry. + pub async fn init( + codex_home: PathBuf, + default_provider: String, + metrics: Option, + ) -> anyhow::Result> { tokio::fs::create_dir_all(&codex_home).await?; let state_migrator = runtime_state_migrator(); let logs_migrator = runtime_logs_migrator(); @@ -116,28 +126,50 @@ impl StateRuntime { .await; let state_path = state_db_path(codex_home.as_path()); let logs_path = logs_db_path(codex_home.as_path()); - let pool = match open_state_sqlite(&state_path, &state_migrator).await { + let pool = match open_state_sqlite(&state_path, &state_migrator, metrics.as_deref()).await { Ok(db) => Arc::new(db), Err(err) => { warn!("failed to open state db at {}: {err}", state_path.display()); return Err(err); } }; - let logs_pool = match open_logs_sqlite(&logs_path, &logs_migrator).await { + let logs_pool = match open_logs_sqlite(&logs_path, &logs_migrator, metrics.as_deref()).await + { Ok(db) => Arc::new(db), Err(err) => { warn!("failed to open logs db at {}: {err}", logs_path.display()); return Err(err); } }; - let thread_updated_at_millis: Option = + let started = Instant::now(); + let backfill_state_result = ensure_backfill_state_row_in_pool(pool.as_ref()).await; + crate::telemetry::record_init_result( + metrics.as_deref(), + DbKind::State, + "ensure_backfill_state", + started.elapsed(), + &backfill_state_result, + ); + backfill_state_result?; + let started = Instant::now(); + let thread_updated_at_millis_result: anyhow::Result> = sqlx::query_scalar("SELECT MAX(threads.updated_at_ms) FROM threads") .fetch_one(pool.as_ref()) - .await?; + .await + .map_err(anyhow::Error::from); + crate::telemetry::record_init_result( + metrics.as_deref(), + DbKind::State, + "post_init_query", + started.elapsed(), + &thread_updated_at_millis_result, + ); + let thread_updated_at_millis = thread_updated_at_millis_result?; let thread_updated_at_millis = thread_updated_at_millis.unwrap_or(0); let runtime = Arc::new(Self { pool, logs_pool, + metrics, codex_home, default_provider, thread_updated_at_millis: Arc::new(AtomicI64::new(thread_updated_at_millis)), @@ -155,6 +187,11 @@ impl StateRuntime { pub fn codex_home(&self) -> &Path { self.codex_home.as_path() } + + /// Return the explicit metrics recorder attached to this runtime, if any. + pub fn metrics(&self) -> Option<&dyn DbMetricsRecorder> { + self.metrics.as_deref() + } } fn base_sqlite_options(path: &Path) -> SqliteConnectOptions { @@ -167,27 +204,87 @@ fn base_sqlite_options(path: &Path) -> SqliteConnectOptions { .log_statements(LevelFilter::Off) } -async fn open_state_sqlite(path: &Path, migrator: &Migrator) -> anyhow::Result { +async fn open_state_sqlite( + path: &Path, + migrator: &Migrator, + metrics: Option<&dyn DbMetricsRecorder>, +) -> anyhow::Result { // New state DBs should use incremental auto-vacuum, but retrofitting an // existing DB requires a full VACUUM. Do not attempt that during process // startup: it is maintenance work that can contend with foreground writers. + open_sqlite( + path, + migrator, + metrics, + DbKind::State, + "open_state", + "migrate_state", + ) + .await +} + +async fn open_logs_sqlite( + path: &Path, + migrator: &Migrator, + metrics: Option<&dyn DbMetricsRecorder>, +) -> anyhow::Result { + open_sqlite( + path, + migrator, + metrics, + DbKind::Logs, + "open_logs", + "migrate_logs", + ) + .await +} + +async fn open_sqlite( + path: &Path, + migrator: &Migrator, + metrics: Option<&dyn DbMetricsRecorder>, + db: DbKind, + open_phase: &'static str, + migrate_phase: &'static str, +) -> anyhow::Result { let options = base_sqlite_options(path).auto_vacuum(SqliteAutoVacuum::Incremental); - let pool = SqlitePoolOptions::new() + let started = Instant::now(); + let pool_result = SqlitePoolOptions::new() .max_connections(5) + .acquire_slow_level(LevelFilter::Warn) + .acquire_slow_threshold(Duration::from_millis(250)) .connect_with(options) - .await?; - migrator.run(&pool).await?; + .await + .map_err(anyhow::Error::from); + crate::telemetry::record_init_result(metrics, db, open_phase, started.elapsed(), &pool_result); + let pool = pool_result?; + let started = Instant::now(); + let migrate_result = migrator.run(&pool).await.map_err(anyhow::Error::from); + crate::telemetry::record_init_result( + metrics, + db, + migrate_phase, + started.elapsed(), + &migrate_result, + ); + migrate_result?; Ok(pool) } -async fn open_logs_sqlite(path: &Path, migrator: &Migrator) -> anyhow::Result { - let options = base_sqlite_options(path).auto_vacuum(SqliteAutoVacuum::Incremental); - let pool = SqlitePoolOptions::new() - .max_connections(5) - .connect_with(options) - .await?; - migrator.run(&pool).await?; - Ok(pool) +async fn ensure_backfill_state_row_in_pool(pool: &sqlx::SqlitePool) -> anyhow::Result<()> { + sqlx::query( + r#" +INSERT INTO backfill_state (id, status, last_watermark, last_success_at, updated_at) +VALUES (?, ?, NULL, NULL, ?) +ON CONFLICT(id) DO NOTHING + "#, + ) + .bind(1_i64) + .bind(crate::BackfillStatus::Pending.as_str()) + .bind(Utc::now().timestamp()) + .execute(pool) + .await?; + Ok(()) } fn db_filename(base_name: &str, version: u32) -> String { @@ -355,9 +452,13 @@ mod tests { strict_pool.close().await; let tolerant_migrator = runtime_state_migrator(); - let tolerant_pool = open_state_sqlite(state_path.as_path(), &tolerant_migrator) - .await - .expect("runtime migrator should tolerate newer applied migrations"); + let tolerant_pool = open_state_sqlite( + state_path.as_path(), + &tolerant_migrator, + /*metrics*/ None, + ) + .await + .expect("runtime migrator should tolerate newer applied migrations"); tolerant_pool.close().await; let _ = tokio::fs::remove_dir_all(codex_home).await; diff --git a/codex-rs/state/src/runtime/agent_jobs.rs b/codex-rs/state/src/runtime/agent_jobs.rs index 3f5526c58d..3a60269f63 100644 --- a/codex-rs/state/src/runtime/agent_jobs.rs +++ b/codex-rs/state/src/runtime/agent_jobs.rs @@ -615,7 +615,12 @@ mod tests { #[tokio::test] async fn report_agent_job_item_result_completes_item_atomically() -> anyhow::Result<()> { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home, "test-provider".to_string()).await?; + let runtime = StateRuntime::init( + codex_home, + "test-provider".to_string(), + /*metrics*/ None, + ) + .await?; let (job_id, item_id, thread_id) = create_running_single_item_job(runtime.as_ref()).await?; let accepted = runtime @@ -655,7 +660,12 @@ mod tests { #[tokio::test] async fn report_agent_job_item_result_rejects_late_reports() -> anyhow::Result<()> { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home, "test-provider".to_string()).await?; + let runtime = StateRuntime::init( + codex_home, + "test-provider".to_string(), + /*metrics*/ None, + ) + .await?; let (job_id, item_id, thread_id) = create_running_single_item_job(runtime.as_ref()).await?; let marked_failed = runtime diff --git a/codex-rs/state/src/runtime/backfill.rs b/codex-rs/state/src/runtime/backfill.rs index cac3163f0e..066abb2ba6 100644 --- a/codex-rs/state/src/runtime/backfill.rs +++ b/codex-rs/state/src/runtime/backfill.rs @@ -2,7 +2,6 @@ use super::*; impl StateRuntime { pub async fn get_backfill_state(&self) -> anyhow::Result { - self.ensure_backfill_state_row().await?; let row = sqlx::query( r#" SELECT status, last_watermark, last_success_at @@ -160,9 +159,13 @@ mod tests { .await .expect("write numeric"); - let _runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let _runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); for suffix in ["", "-wal", "-shm", "-journal"] { let legacy_path = codex_home.join(format!("{unversioned_name}{suffix}")); @@ -207,9 +210,13 @@ mod tests { #[tokio::test] async fn backfill_state_persists_progress_and_completion() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let initial = runtime .get_backfill_state() @@ -260,9 +267,13 @@ mod tests { #[tokio::test] async fn backfill_claim_is_singleton_until_stale_and_blocked_when_complete() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let claimed = runtime .try_claim_backfill(/*lease_seconds*/ 3600) diff --git a/codex-rs/state/src/runtime/goals.rs b/codex-rs/state/src/runtime/goals.rs index 56613b7d92..57088ad6b5 100644 --- a/codex-rs/state/src/runtime/goals.rs +++ b/codex-rs/state/src/runtime/goals.rs @@ -431,9 +431,13 @@ mod tests { use pretty_assertions::assert_eq; async fn test_runtime() -> std::sync::Arc { - StateRuntime::init(unique_temp_dir(), "test-provider".to_string()) - .await - .expect("state db should initialize") + StateRuntime::init( + unique_temp_dir(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("state db should initialize") } fn test_thread_id() -> ThreadId { diff --git a/codex-rs/state/src/runtime/logs.rs b/codex-rs/state/src/runtime/logs.rs index 6c878db624..d2dd7df765 100644 --- a/codex-rs/state/src/runtime/logs.rs +++ b/codex-rs/state/src/runtime/logs.rs @@ -573,9 +573,13 @@ mod tests { #[tokio::test] async fn insert_logs_use_dedicated_log_database() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); runtime .insert_logs(&[LogEntry { @@ -645,9 +649,13 @@ mod tests { pool.close().await; drop(pool); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let rows = runtime .query_logs(&LogQuery::default()) @@ -736,9 +744,13 @@ mod tests { pool.close().await; drop(pool); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); assert!( !legacy_logs_path.exists(), @@ -762,9 +774,13 @@ mod tests { #[tokio::test] async fn init_configures_logs_db_with_incremental_auto_vacuum() { let codex_home = unique_temp_dir(); - let _runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let _runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let pool = open_db_pool(logs_db_path(codex_home.as_path()).as_path()).await; let auto_vacuum = sqlx::query_scalar::<_, i64>("PRAGMA auto_vacuum") @@ -806,9 +822,13 @@ mod tests { #[tokio::test] async fn query_logs_with_search_matches_rendered_body_substring() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); runtime .insert_logs(&[ @@ -859,9 +879,13 @@ mod tests { #[tokio::test] async fn query_logs_filters_level_set_without_rewriting_stored_level() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); runtime .insert_logs(&[ @@ -944,9 +968,13 @@ mod tests { #[tokio::test] async fn insert_logs_prunes_old_rows_when_thread_exceeds_size_limit() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let six_mebibytes = "a".repeat(6 * 1024 * 1024); runtime @@ -998,9 +1026,13 @@ mod tests { #[tokio::test] async fn insert_logs_prunes_single_thread_row_when_it_exceeds_size_limit() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let eleven_mebibytes = "d".repeat(11 * 1024 * 1024); runtime @@ -1036,9 +1068,13 @@ mod tests { #[tokio::test] async fn insert_logs_prunes_threadless_rows_per_process_uuid_only() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let six_mebibytes = "b".repeat(6 * 1024 * 1024); runtime @@ -1105,9 +1141,13 @@ mod tests { #[tokio::test] async fn insert_logs_prunes_single_threadless_process_row_when_it_exceeds_size_limit() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let eleven_mebibytes = "e".repeat(11 * 1024 * 1024); runtime @@ -1143,9 +1183,13 @@ mod tests { #[tokio::test] async fn insert_logs_prunes_threadless_rows_with_null_process_uuid() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let six_mebibytes = "c".repeat(6 * 1024 * 1024); runtime @@ -1211,9 +1255,13 @@ mod tests { #[tokio::test] async fn insert_logs_prunes_single_threadless_null_process_row_when_it_exceeds_limit() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let eleven_mebibytes = "f".repeat(11 * 1024 * 1024); runtime @@ -1249,9 +1297,13 @@ mod tests { #[tokio::test] async fn insert_logs_prunes_old_rows_when_thread_exceeds_row_limit() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let entries: Vec = (1..=1_001) .map(|ts| LogEntry { @@ -1292,9 +1344,13 @@ mod tests { #[tokio::test] async fn insert_logs_prunes_old_threadless_rows_when_process_exceeds_row_limit() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let entries: Vec = (1..=1_001) .map(|ts| LogEntry { @@ -1339,9 +1395,13 @@ mod tests { #[tokio::test] async fn insert_logs_prunes_old_threadless_null_process_rows_when_row_limit_exceeded() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let entries: Vec = (1..=1_001) .map(|ts| LogEntry { @@ -1386,9 +1446,13 @@ mod tests { #[tokio::test] async fn query_feedback_logs_returns_newest_lines_within_limit_in_order() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); runtime .insert_logs(&[ @@ -1456,9 +1520,13 @@ mod tests { #[tokio::test] async fn query_feedback_logs_excludes_oversized_newest_row() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let eleven_mebibytes = "z".repeat(11 * 1024 * 1024); runtime @@ -1506,9 +1574,13 @@ mod tests { #[tokio::test] async fn query_feedback_logs_includes_threadless_rows_from_same_process() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); runtime .insert_logs(&[ @@ -1604,9 +1676,13 @@ mod tests { #[tokio::test] async fn query_feedback_logs_excludes_threadless_rows_from_prior_processes() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); runtime .insert_logs(&[ @@ -1702,9 +1778,13 @@ mod tests { #[tokio::test] async fn query_feedback_logs_keeps_newest_suffix_across_thread_and_threadless_logs() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let thread_marker = "thread-scoped-oldest"; let threadless_older_marker = "threadless-older"; let threadless_newer_marker = "threadless-newer"; @@ -1777,9 +1857,13 @@ mod tests { #[tokio::test] async fn query_feedback_logs_for_threads_merges_requested_threads_and_threadless_rows() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); runtime .insert_logs(&[ @@ -1897,9 +1981,13 @@ mod tests { #[tokio::test] async fn query_feedback_logs_for_threads_returns_empty_for_empty_thread_list() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let bytes = runtime .query_feedback_logs_for_threads(&[]) diff --git a/codex-rs/state/src/runtime/memories.rs b/codex-rs/state/src/runtime/memories.rs index 186f2dd341..929fd3719c 100644 --- a/codex-rs/state/src/runtime/memories.rs +++ b/codex-rs/state/src/runtime/memories.rs @@ -1308,9 +1308,13 @@ mod tests { #[tokio::test] async fn stage1_claim_skips_when_up_to_date() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); let metadata = test_thread_metadata(&codex_home, thread_id, codex_home.join("a")); @@ -1376,9 +1380,13 @@ mod tests { #[tokio::test] async fn stage1_running_stale_can_be_stolen_but_fresh_running_is_skipped() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); let owner_a = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); @@ -1431,9 +1439,13 @@ mod tests { #[tokio::test] async fn stage1_concurrent_claim_for_same_thread_is_conflict_safe() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); runtime @@ -1499,9 +1511,13 @@ mod tests { #[tokio::test] async fn stage1_concurrent_claims_respect_running_cap() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let thread_a = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); let thread_b = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); @@ -1567,9 +1583,13 @@ mod tests { #[tokio::test] async fn claim_stage1_jobs_filters_by_age_idle_and_current_thread() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let now = Utc::now(); let fresh_at = now - Duration::hours(1); @@ -1657,9 +1677,13 @@ mod tests { #[tokio::test] async fn claim_stage1_jobs_prefilters_threads_with_up_to_date_memory() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let now = Utc::now(); let eligible_newer_at = now - Duration::hours(13); @@ -1756,9 +1780,13 @@ mod tests { #[tokio::test] async fn claim_stage1_jobs_skips_threads_with_disabled_memory_mode() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let now = Utc::now(); let eligible_at = now - Duration::hours(13); @@ -1827,9 +1855,13 @@ mod tests { #[tokio::test] async fn clear_memory_data_clears_rows_and_preserves_thread_memory_modes() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let now = Utc::now() - Duration::hours(13); let worker_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("worker id"); @@ -1936,9 +1968,13 @@ mod tests { #[tokio::test] async fn claim_stage1_jobs_enforces_global_running_cap() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let current_thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("current thread id"); @@ -2063,9 +2099,13 @@ WHERE kind = 'memory_stage1' #[tokio::test] async fn claim_stage1_jobs_processes_two_full_batches_across_startup_passes() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let current_thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("current thread id"); @@ -2150,9 +2190,13 @@ WHERE kind = 'memory_stage1' #[tokio::test] async fn stage1_output_cascades_on_thread_delete() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); @@ -2220,9 +2264,13 @@ WHERE kind = 'memory_stage1' #[tokio::test] async fn mark_stage1_job_succeeded_no_output_skips_phase2_when_output_was_already_absent() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); @@ -2295,9 +2343,13 @@ WHERE kind = 'memory_stage1' #[tokio::test] async fn mark_stage1_job_succeeded_no_output_enqueues_phase2_when_deleting_output() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); @@ -2420,9 +2472,13 @@ WHERE kind = 'memory_stage1' #[tokio::test] async fn stage1_retry_exhaustion_does_not_block_newer_watermark() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); @@ -2516,9 +2572,13 @@ WHERE kind = 'memory_stage1' #[tokio::test] async fn phase2_global_lock_respects_success_cooldown() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); @@ -2578,9 +2638,13 @@ WHERE kind = 'memory_stage1' #[tokio::test] async fn phase2_global_lock_can_be_claimed_after_retry_budget_is_exhausted() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); runtime .enqueue_global_consolidation(/*input_watermark*/ 100) @@ -2648,9 +2712,13 @@ WHERE kind = 'memory_stage1' #[tokio::test] async fn list_stage1_outputs_for_global_returns_latest_outputs() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let thread_id_a = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); let thread_id_b = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); @@ -2751,9 +2819,13 @@ WHERE kind = 'memory_stage1' #[tokio::test] async fn list_stage1_outputs_for_global_skips_empty_payloads() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let thread_id_non_empty = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); @@ -2820,9 +2892,13 @@ VALUES (?, ?, ?, ?, ?) #[tokio::test] async fn list_stage1_outputs_for_global_skips_polluted_threads() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let thread_id_enabled = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); @@ -2888,9 +2964,13 @@ VALUES (?, ?, ?, ?, ?) #[tokio::test] async fn get_phase2_input_selection_returns_current_selected_rows() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let thread_id_a = stable_thread_id("00000000-0000-4000-8000-000000000001"); let thread_id_b = stable_thread_id("00000000-0000-4000-8000-000000000002"); @@ -3003,9 +3083,13 @@ VALUES (?, ?, ?, ?, ?) #[tokio::test] async fn get_phase2_input_selection_excludes_polluted_previous_selection() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let thread_id_enabled = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); @@ -3096,9 +3180,13 @@ VALUES (?, ?, ?, ?, ?) #[tokio::test] async fn mark_thread_memory_mode_polluted_enqueues_phase2_for_selected_threads() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); @@ -3185,9 +3273,13 @@ VALUES (?, ?, ?, ?, ?) #[tokio::test] async fn get_phase2_input_selection_returns_regenerated_selected_rows() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); @@ -3304,9 +3396,13 @@ VALUES (?, ?, ?, ?, ?) #[tokio::test] async fn get_phase2_input_selection_uses_current_ranking_after_refreshes() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let thread_id_a = stable_thread_id("00000000-0000-4000-8000-000000000001"); let thread_id_b = stable_thread_id("00000000-0000-4000-8000-000000000002"); @@ -3447,9 +3543,13 @@ VALUES (?, ?, ?, ?, ?) #[tokio::test] async fn mark_global_phase2_job_succeeded_updates_selected_snapshot_timestamp() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); @@ -3597,9 +3697,13 @@ VALUES (?, ?, ?, ?, ?) #[tokio::test] async fn mark_global_phase2_job_succeeded_only_marks_exact_selected_snapshots() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); @@ -3717,9 +3821,13 @@ VALUES (?, ?, ?, ?, ?) #[tokio::test] async fn record_stage1_output_usage_updates_usage_metadata() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let thread_a = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id a"); let thread_b = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id b"); @@ -3836,9 +3944,13 @@ VALUES (?, ?, ?, ?, ?) #[tokio::test] async fn get_phase2_input_selection_prioritizes_usage_count_then_recent_usage() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let now = Utc::now(); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); @@ -3932,9 +4044,13 @@ VALUES (?, ?, ?, ?, ?) #[tokio::test] async fn get_phase2_input_selection_excludes_stale_used_memories_but_keeps_fresh_never_used() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let now = Utc::now(); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); @@ -4028,9 +4144,13 @@ VALUES (?, ?, ?, ?, ?) #[tokio::test] async fn get_phase2_input_selection_prefers_recent_thread_updates_over_recent_generation() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); let older_thread = @@ -4114,9 +4234,13 @@ VALUES (?, ?, ?, ?, ?) #[tokio::test] async fn prune_stage1_outputs_for_retention_prunes_stale_unselected_rows_only() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); let stale_unused = @@ -4256,9 +4380,13 @@ VALUES (?, ?, ?, ?, ?) #[tokio::test] async fn prune_stage1_outputs_for_retention_respects_batch_limit() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id"); let thread_a = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread a"); @@ -4334,9 +4462,13 @@ VALUES (?, ?, ?, ?, ?) #[tokio::test] async fn mark_stage1_job_succeeded_enqueues_global_consolidation() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let thread_a = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id a"); let thread_b = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("thread id b"); @@ -4429,9 +4561,13 @@ VALUES (?, ?, ?, ?, ?) #[tokio::test] async fn phase2_global_lock_allows_only_one_fresh_runner() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); runtime .enqueue_global_consolidation(/*input_watermark*/ 200) @@ -4462,9 +4598,13 @@ VALUES (?, ?, ?, ?, ?) #[tokio::test] async fn phase2_global_lock_creates_missing_job_row() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); let owner_a = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner a"); let owner_b = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner b"); @@ -4512,9 +4652,13 @@ VALUES (?, ?, ?, ?, ?) #[tokio::test] async fn phase2_global_lock_stale_lease_allows_takeover() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); runtime .enqueue_global_consolidation(/*input_watermark*/ 300) @@ -4587,9 +4731,13 @@ VALUES (?, ?, ?, ?, ?) #[tokio::test] async fn enqueue_global_consolidation_keeps_phase2_input_watermark_monotonic() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); runtime .enqueue_global_consolidation(/*input_watermark*/ 500) @@ -4651,9 +4799,13 @@ VALUES (?, ?, ?, ?, ?) #[tokio::test] async fn phase2_failure_fallback_updates_unowned_running_job() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); runtime .enqueue_global_consolidation(/*input_watermark*/ 400) diff --git a/codex-rs/state/src/runtime/remote_control.rs b/codex-rs/state/src/runtime/remote_control.rs index fa0b1823f8..7c8ada7b2c 100644 --- a/codex-rs/state/src/runtime/remote_control.rs +++ b/codex-rs/state/src/runtime/remote_control.rs @@ -130,9 +130,13 @@ mod tests { #[tokio::test] async fn remote_control_enrollment_round_trips_by_target_and_account() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); runtime .upsert_remote_control_enrollment(&RemoteControlEnrollmentRecord { @@ -207,9 +211,13 @@ mod tests { #[tokio::test] async fn delete_remote_control_enrollment_removes_only_matching_entry() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("initialize runtime"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("initialize runtime"); runtime .upsert_remote_control_enrollment(&RemoteControlEnrollmentRecord { diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index 1795a864d8..2a49b88d3b 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -1156,9 +1156,13 @@ mod tests { #[tokio::test] async fn upsert_thread_keeps_creation_memory_mode_for_existing_rows() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("state db should initialize"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("state db should initialize"); let thread_id = ThreadId::from_string("00000000-0000-0000-0000-000000000123").expect("valid thread id"); let mut metadata = test_thread_metadata(&codex_home, thread_id, codex_home.clone()); @@ -1194,9 +1198,13 @@ mod tests { #[tokio::test] async fn list_threads_updated_after_returns_oldest_changes_first() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("state db should initialize"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("state db should initialize"); let older_id = ThreadId::from_string("00000000-0000-0000-0000-000000000001").expect("valid thread id"); let middle_id = @@ -1278,9 +1286,13 @@ mod tests { #[tokio::test] async fn list_threads_filters_by_cwd() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("state db should initialize"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("state db should initialize"); let first_id = ThreadId::from_string("00000000-0000-0000-0000-000000000101").expect("valid thread id"); let second_id = @@ -1349,9 +1361,13 @@ mod tests { #[tokio::test] async fn apply_rollout_items_restores_memory_mode_from_session_meta() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("state db should initialize"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("state db should initialize"); let thread_id = ThreadId::from_string("00000000-0000-0000-0000-000000000456").expect("valid thread id"); let metadata = test_thread_metadata(&codex_home, thread_id, codex_home.clone()); @@ -1406,9 +1422,13 @@ mod tests { #[tokio::test] async fn apply_rollout_items_preserves_existing_git_branch_and_fills_missing_git_fields() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("state db should initialize"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("state db should initialize"); let thread_id = ThreadId::from_string("00000000-0000-0000-0000-000000000457").expect("valid thread id"); let mut metadata = test_thread_metadata(&codex_home, thread_id, codex_home.clone()); @@ -1475,9 +1495,13 @@ mod tests { #[tokio::test] async fn upsert_thread_preserves_existing_git_fields_atomically() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("state db should initialize"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("state db should initialize"); let thread_id = ThreadId::from_string("00000000-0000-0000-0000-000000000458").expect("valid thread id"); let mut metadata = test_thread_metadata(&codex_home, thread_id, codex_home.clone()); @@ -1516,9 +1540,13 @@ mod tests { #[tokio::test] async fn update_thread_git_info_preserves_newer_non_git_metadata() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("state db should initialize"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("state db should initialize"); let thread_id = ThreadId::from_string("00000000-0000-0000-0000-000000000789").expect("valid thread id"); let metadata = test_thread_metadata(&codex_home, thread_id, codex_home.clone()); @@ -1576,9 +1604,13 @@ mod tests { #[tokio::test] async fn insert_thread_if_absent_preserves_existing_metadata() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("state db should initialize"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("state db should initialize"); let thread_id = ThreadId::from_string("00000000-0000-0000-0000-000000000791").expect("valid thread id"); @@ -1621,9 +1653,13 @@ mod tests { #[tokio::test] async fn update_thread_git_info_can_clear_fields() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("state db should initialize"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("state db should initialize"); let thread_id = ThreadId::from_string("00000000-0000-0000-0000-000000000790").expect("valid thread id"); let mut metadata = test_thread_metadata(&codex_home, thread_id, codex_home.clone()); @@ -1655,9 +1691,13 @@ mod tests { #[tokio::test] async fn touch_thread_updated_at_updates_only_updated_at() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("state db should initialize"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("state db should initialize"); let thread_id = ThreadId::from_string("00000000-0000-0000-0000-000000000791").expect("valid thread id"); let mut metadata = test_thread_metadata(&codex_home, thread_id, codex_home.clone()); @@ -1692,9 +1732,13 @@ mod tests { #[tokio::test] async fn thread_updated_at_uses_unique_epoch_millis_and_reads_legacy_seconds() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("state db should initialize"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("state db should initialize"); let first_id = ThreadId::from_string("00000000-0000-0000-0000-000000000901").expect("valid thread id"); let second_id = @@ -1790,9 +1834,13 @@ mod tests { #[tokio::test] async fn apply_rollout_items_uses_override_updated_at_when_provided() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("state db should initialize"); + let runtime = StateRuntime::init( + codex_home.clone(), + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("state db should initialize"); let thread_id = ThreadId::from_string("00000000-0000-0000-0000-000000000792").expect("valid thread id"); let metadata = test_thread_metadata(&codex_home, thread_id, codex_home.clone()); @@ -1849,9 +1897,13 @@ mod tests { #[tokio::test] async fn thread_spawn_edges_track_directional_status() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home, "test-provider".to_string()) - .await - .expect("state db should initialize"); + let runtime = StateRuntime::init( + codex_home, + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("state db should initialize"); let parent_thread_id = ThreadId::from_string("00000000-0000-0000-0000-000000000900").expect("valid thread id"); let child_thread_id = @@ -1945,9 +1997,13 @@ mod tests { #[tokio::test] async fn thread_spawn_children_without_status_filter_lists_all_statuses() { let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home, "test-provider".to_string()) - .await - .expect("state db should initialize"); + let runtime = StateRuntime::init( + codex_home, + "test-provider".to_string(), + /*metrics*/ None, + ) + .await + .expect("state db should initialize"); let parent_thread_id = ThreadId::from_string("00000000-0000-0000-0000-000000000910").expect("valid thread id"); let open_child_thread_id = diff --git a/codex-rs/state/src/telemetry.rs b/codex-rs/state/src/telemetry.rs new file mode 100644 index 0000000000..68e17123eb --- /dev/null +++ b/codex-rs/state/src/telemetry.rs @@ -0,0 +1,263 @@ +use std::borrow::Cow; +use std::sync::Arc; +use std::time::Duration; + +use crate::DB_FALLBACK_METRIC; +use crate::DB_INIT_DURATION_METRIC; +use crate::DB_INIT_METRIC; + +/// Low-cardinality metrics sink used by the SQLite state runtime. +/// +/// Implementations should ignore recording errors locally. Database operations +/// must never fail because telemetry delivery failed. +pub trait DbMetricsRecorder: Send + Sync + 'static { + /// Increment a counter metric by `inc` with low-cardinality tags. + fn counter(&self, name: &str, inc: i64, tags: &[(&str, &str)]); + + /// Record an elapsed duration metric with low-cardinality tags. + fn record_duration(&self, name: &str, duration: Duration, tags: &[(&str, &str)]); +} + +/// Shared recorder handle stored by rollout SQLite telemetry plumbing. +pub type DbMetricsRecorderHandle = Arc; + +#[derive(Clone, Copy)] +pub(crate) enum DbKind { + State, + Logs, +} + +impl DbKind { + fn as_str(self) -> &'static str { + match self { + Self::State => "state", + Self::Logs => "logs", + } + } +} + +pub(crate) fn record_init_result( + metrics: Option<&dyn DbMetricsRecorder>, + db: DbKind, + phase: &'static str, + duration: Duration, + result: &anyhow::Result, +) { + let outcome = DbOutcomeTags::from_result(result); + let tags = [ + ("status", outcome.status), + ("phase", phase), + ("db", db.as_str()), + ("error", outcome.error), + ]; + record_counter(metrics, DB_INIT_METRIC, &tags); + record_duration(metrics, DB_INIT_DURATION_METRIC, duration, &tags); +} + +pub fn record_fallback( + metrics: Option<&dyn DbMetricsRecorder>, + caller: &'static str, + reason: &'static str, +) { + let tags = [("caller", caller), ("reason", reason)]; + record_counter(metrics, DB_FALLBACK_METRIC, &tags); +} + +pub fn record_init_backfill_gate( + metrics: Option<&dyn DbMetricsRecorder>, + duration: Duration, + result: &anyhow::Result<()>, +) { + record_init_result(metrics, DbKind::State, "backfill_gate", duration, result); +} + +pub(crate) fn classify_error(err: &anyhow::Error) -> &'static str { + for cause in err.chain() { + if let Some(sqlx_err) = cause.downcast_ref::() { + return classify_sqlx_error(sqlx_err); + } + if cause + .downcast_ref::() + .is_some() + { + return "migration"; + } + if cause.downcast_ref::().is_some() { + return "serde"; + } + if cause.downcast_ref::().is_some() { + return "io"; + } + } + + "unknown" +} + +pub(crate) fn classify_sqlite_code(code: &str) -> &'static str { + // SQLite result codes are documented at https://www.sqlite.org/rescode.html. + // Extended codes preserve the primary code in the low byte. + let primary_code = code.parse::().ok().map(|code| code & 0xff); + match primary_code { + Some(5) => "busy", + Some(6) => "locked", + Some(8) => "readonly", + Some(10) => "io", + Some(11) => "corrupt", + Some(13) => "full", + Some(14) => "cantopen", + Some(19) => "constraint", + Some(17) => "schema", + _ => "unknown", + } +} + +struct DbOutcomeTags { + status: &'static str, + error: &'static str, +} + +impl DbOutcomeTags { + fn from_result(result: &anyhow::Result) -> Self { + match result { + Ok(_) => Self { + status: "success", + error: "none", + }, + Err(err) => Self { + status: "failed", + error: classify_error(err), + }, + } + } +} + +fn classify_sqlx_error(err: &sqlx::Error) -> &'static str { + match err { + sqlx::Error::Database(database_error) => { + let code = database_error + .code() + .unwrap_or(Cow::Borrowed("none")) + .to_string(); + classify_sqlite_code(code.as_str()) + } + sqlx::Error::PoolTimedOut => "pool_timeout", + sqlx::Error::Io(_) => "io", + sqlx::Error::ColumnDecode { source, .. } if source.is::() => "serde", + sqlx::Error::Decode(source) if source.is::() => "serde", + _ => "unknown", + } +} + +fn record_counter(metrics: Option<&dyn DbMetricsRecorder>, name: &str, tags: &[(&str, &str)]) { + if let Some(metrics) = metrics { + metrics.counter(name, /*inc*/ 1, tags); + } +} + +fn record_duration( + metrics: Option<&dyn DbMetricsRecorder>, + name: &str, + duration: Duration, + tags: &[(&str, &str)], +) { + if let Some(metrics) = metrics { + metrics.record_duration(name, duration, tags); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::DB_FALLBACK_METRIC; + use pretty_assertions::assert_eq; + use std::collections::BTreeMap; + use std::sync::Mutex; + + #[derive(Default)] + struct TestMetrics { + events: Mutex>, + } + + #[derive(Debug, Eq, PartialEq)] + struct MetricEvent { + name: String, + tags: BTreeMap, + } + + impl TestMetrics { + fn events(&self) -> Vec { + self.events + .lock() + .expect("metrics lock") + .iter() + .map(|event| MetricEvent { + name: event.name.clone(), + tags: event.tags.clone(), + }) + .collect() + } + } + + impl DbMetricsRecorder for TestMetrics { + fn counter(&self, name: &str, _inc: i64, tags: &[(&str, &str)]) { + self.events.lock().expect("metrics lock").push(MetricEvent { + name: name.to_string(), + tags: tags_to_map(tags), + }); + } + + fn record_duration(&self, _name: &str, _duration: Duration, _tags: &[(&str, &str)]) {} + } + + fn tags_to_map(tags: &[(&str, &str)]) -> BTreeMap { + tags.iter() + .map(|(key, value)| ((*key).to_string(), (*value).to_string())) + .collect() + } + + #[test] + fn classifies_sqlite_primary_codes() { + assert_eq!(classify_sqlite_code("5"), "busy"); + assert_eq!(classify_sqlite_code("6"), "locked"); + assert_eq!(classify_sqlite_code("14"), "cantopen"); + assert_eq!(classify_sqlite_code("2067"), "constraint"); + } + + #[test] + fn classifies_non_sqlite_errors() { + let io_error = + anyhow::Error::new(std::io::Error::new(std::io::ErrorKind::NotFound, "missing")); + assert_eq!(classify_error(&io_error), "io"); + + let serde_error = + anyhow::Error::new(serde_json::from_str::("not-json").unwrap_err()); + assert_eq!(classify_error(&serde_error), "serde"); + + let unknown_error = anyhow::anyhow!("plain failure"); + assert_eq!(classify_error(&unknown_error), "unknown"); + } + + #[test] + fn classifies_sqlx_pool_timeout() { + let err = anyhow::Error::new(sqlx::Error::PoolTimedOut); + assert_eq!(classify_error(&err), "pool_timeout"); + } + + #[test] + fn records_fallback_metric_with_reason() { + let metrics = TestMetrics::default(); + + record_fallback(Some(&metrics), "list_threads", "db_error"); + + assert_eq!( + metrics.events(), + vec![MetricEvent { + name: DB_FALLBACK_METRIC.to_string(), + tags: BTreeMap::from([ + ("caller".to_string(), "list_threads".to_string()), + ("reason".to_string(), "db_error".to_string()), + ]), + }] + ); + } +} diff --git a/codex-rs/thread-manager-sample/src/main.rs b/codex-rs/thread-manager-sample/src/main.rs index 68a95bd2d9..d116ccb7dc 100644 --- a/codex-rs/thread-manager-sample/src/main.rs +++ b/codex-rs/thread-manager-sample/src/main.rs @@ -42,6 +42,7 @@ use codex_core_api::RealtimeConfig; use codex_core_api::SessionPickerViewMode; use codex_core_api::SessionSource; use codex_core_api::ShellEnvironmentPolicy; +use codex_core_api::StateDbAccess; use codex_core_api::TerminalResizeReflowConfig; use codex_core_api::ThreadManager; use codex_core_api::ThreadStoreConfig; @@ -104,7 +105,8 @@ 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 state_db = init_state_db(&config, /*metrics*/ None).await; + let state_db_access = StateDbAccess::new(state_db.clone()); let auth_manager = AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false).await; @@ -112,7 +114,7 @@ async fn run_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { config.codex_self_exe.clone(), config.codex_linux_sandbox_exe.clone(), )?; - let thread_store = thread_store_from_config(&config, state_db.clone()); + let thread_store = thread_store_from_config(&config, state_db_access.clone()); let environment_manager = Arc::new( EnvironmentManager::from_codex_home(config.codex_home.clone(), local_runtime_paths).await?, ); @@ -124,7 +126,7 @@ async fn run_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { environment_manager, /*analytics_events_client*/ None, Arc::clone(&thread_store), - state_db, + state_db_access, installation_id, /*attestation_provider*/ None, ); diff --git a/codex-rs/thread-store/src/local/archive_thread.rs b/codex-rs/thread-store/src/local/archive_thread.rs index 8fb214e98c..08ce9ede28 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_access = store.state_db_access(); let rollout_path = find_thread_path_by_id_str( store.config.codex_home.as_path(), &thread_id.to_string(), - state_db_ctx.as_deref(), + &state_db_access, ) .await .map_err(|err| ThreadStoreError::InvalidRequest { @@ -52,7 +52,7 @@ pub(super) async fn archive_thread( } })?; - if let Some(ctx) = state_db_ctx { + if let Some(ctx) = state_db_access.state_db() { let _ = ctx .mark_archived(thread_id, archived_path.as_path(), Utc::now()) .await; @@ -81,7 +81,10 @@ mod tests { #[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 = LocalThreadStore::new( + test_config(home.path()), + codex_rollout::StateDbAccess::none(), + ); let uuid = Uuid::from_u128(201); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let active_path = @@ -134,10 +137,14 @@ mod tests { let runtime = codex_state::StateRuntime::init( home.path().to_path_buf(), config.default_model_provider_id.clone(), + /*metrics*/ None, ) .await .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); + let store = LocalThreadStore::new( + config.clone(), + codex_rollout::StateDbAccess::new(Some(runtime.clone())), + ); runtime .mark_backfill_complete(/*last_watermark*/ None) .await diff --git a/codex-rs/thread-store/src/local/list_threads.rs b/codex-rs/thread-store/src/local/list_threads.rs index e470ad2be9..99bd7e5408 100644 --- a/codex-rs/thread-store/src/local/list_threads.rs +++ b/codex-rs/thread-store/src/local/list_threads.rs @@ -39,7 +39,7 @@ 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 state_db_access = store.state_db_access(); let rollout_config = RolloutConfig { codex_home: store.config.codex_home.clone(), sqlite_home: store.config.sqlite_home.clone(), @@ -48,7 +48,7 @@ pub(super) async fn list_threads( generate_memories: false, }; let page = list_rollout_threads( - state_db, + state_db_access.clone(), &rollout_config, store.config.default_model_provider_id.as_str(), ¶ms, @@ -80,7 +80,7 @@ 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 { + if let Some(state_db_ctx) = state_db_access.state_db() { for &thread_id in &thread_ids { let Ok(Some(metadata)) = state_db_ctx.get_thread(thread_id).await else { continue; @@ -108,7 +108,7 @@ pub(super) async fn list_threads( } async fn list_rollout_threads( - state_db: Option, + state_db_ctx: codex_rollout::StateDbAccess, config: &RolloutConfig, default_model_provider_id: &str, params: &ListThreadsParams, @@ -118,7 +118,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, config, params.page_size, cursor, @@ -133,7 +133,7 @@ async fn list_rollout_threads( .await } else if params.use_state_db_only { RolloutRecorder::list_threads_from_state_db( - state_db, + state_db_ctx, config, params.page_size, cursor, @@ -148,7 +148,7 @@ async fn list_rollout_threads( .await } else if params.archived { RolloutRecorder::list_archived_threads( - state_db, + state_db_ctx, config, params.page_size, cursor, @@ -163,7 +163,7 @@ async fn list_rollout_threads( .await } else { RolloutRecorder::list_threads( - state_db, + state_db_ctx, config, params.page_size, cursor, @@ -203,7 +203,10 @@ 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 = LocalThreadStore::new( + test_config(home.path()), + codex_rollout::StateDbAccess::none(), + ); write_session_file_with( home.path(), home.path().join("sessions/2025/01/03"), @@ -246,10 +249,14 @@ mod tests { let runtime = codex_state::StateRuntime::init( home.path().to_path_buf(), config.default_model_provider_id.clone(), + /*metrics*/ None, ) .await .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); + let store = LocalThreadStore::new( + config.clone(), + codex_rollout::StateDbAccess::new(Some(runtime.clone())), + ); runtime .mark_backfill_complete(/*last_watermark*/ None) .await @@ -303,7 +310,10 @@ 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 = LocalThreadStore::new( + test_config(home.path()), + codex_rollout::StateDbAccess::none(), + ); 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 +382,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, codex_rollout::StateDbAccess::none()); let uuid = Uuid::from_u128(101); let path = write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file"); @@ -411,7 +421,10 @@ 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 = LocalThreadStore::new( + test_config(home.path()), + codex_rollout::StateDbAccess::none(), + ); let err = store .list_threads(ListThreadsParams { diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index 07aa5e925f..0172199f34 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -13,6 +13,7 @@ mod test_support; use async_trait::async_trait; use codex_protocol::ThreadId; use codex_rollout::RolloutRecorder; +use codex_rollout::StateDbAccess; use codex_rollout::StateDbHandle; use std::collections::HashMap; use std::collections::hash_map::Entry; @@ -41,7 +42,7 @@ use crate::UpdateThreadMetadataParams; pub struct LocalThreadStore { pub(super) config: LocalThreadStoreConfig, live_recorders: Arc>>, - state_db: Option, + state_db: StateDbAccess, } /// Process-scoped configuration for local thread storage. @@ -75,8 +76,8 @@ 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 with explicit SQLite access. + pub fn new(config: LocalThreadStoreConfig, state_db: StateDbAccess) -> Self { Self { config, live_recorders: Arc::new(Mutex::new(HashMap::new())), @@ -86,6 +87,10 @@ impl LocalThreadStore { /// Return the state DB handle used by local rollout writers. pub async fn state_db(&self) -> Option { + self.state_db.state_db() + } + + pub(super) fn state_db_access(&self) -> StateDbAccess { self.state_db.clone() } @@ -289,7 +294,10 @@ mod tests { #[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 = LocalThreadStore::new( + test_config(home.path()), + codex_rollout::StateDbAccess::none(), + ); let thread_id = ThreadId::default(); store @@ -338,7 +346,10 @@ 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 = LocalThreadStore::new( + test_config(home.path()), + codex_rollout::StateDbAccess::none(), + ); let thread_id = ThreadId::default(); let mut params = create_thread_params(thread_id); params.metadata.cwd = None; @@ -358,7 +369,10 @@ 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 = LocalThreadStore::new( + test_config(home.path()), + codex_rollout::StateDbAccess::none(), + ); let thread_id = ThreadId::default(); store @@ -397,7 +411,8 @@ mod tests { let config = test_config(home.path()); let thread_id = ThreadId::default(); - let first_store = LocalThreadStore::new(config.clone(), /*state_db*/ None); + let first_store = + LocalThreadStore::new(config.clone(), codex_rollout::StateDbAccess::none()); first_store .create_thread(create_thread_params(thread_id)) .await @@ -426,7 +441,7 @@ mod tests { .await .expect("shutdown initial writer"); - let resumed_store = LocalThreadStore::new(config, /*state_db*/ None); + let resumed_store = LocalThreadStore::new(config, codex_rollout::StateDbAccess::none()); resumed_store .resume_thread(ResumeThreadParams { thread_id, @@ -457,7 +472,10 @@ 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 = LocalThreadStore::new( + test_config(home.path()), + codex_rollout::StateDbAccess::none(), + ); let thread_id = ThreadId::default(); store @@ -477,7 +495,10 @@ 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 = LocalThreadStore::new( + test_config(home.path()), + codex_rollout::StateDbAccess::none(), + ); let thread_id = ThreadId::default(); store @@ -506,7 +527,10 @@ 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 = LocalThreadStore::new( + test_config(home.path()), + codex_rollout::StateDbAccess::none(), + ); 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 +559,10 @@ 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 = LocalThreadStore::new( + test_config(home.path()), + codex_rollout::StateDbAccess::none(), + ); 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 +611,10 @@ 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 = LocalThreadStore::new( + test_config(home.path()), + codex_rollout::StateDbAccess::none(), + ); 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 +653,10 @@ 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 = LocalThreadStore::new( + test_config(home.path()), + codex_rollout::StateDbAccess::none(), + ); 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 +724,10 @@ 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 = LocalThreadStore::new( + test_config(home.path()), + codex_rollout::StateDbAccess::none(), + ); 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 9d685aace6..3ac6b08b36 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_access = store.state_db_access(); 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(), + &state_db_access, ) .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(), + &state_db_access, ) .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(), + &state_db_access, ) .await .map_err(|err| ThreadStoreError::InvalidRequest { @@ -423,7 +423,10 @@ 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 = LocalThreadStore::new( + test_config(home.path()), + codex_rollout::StateDbAccess::none(), + ); let uuid = Uuid::from_u128(205); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let active_path = @@ -451,7 +454,10 @@ 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 = LocalThreadStore::new( + test_config(home.path()), + codex_rollout::StateDbAccess::none(), + ); let uuid = Uuid::from_u128(211); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let active_path = @@ -489,10 +495,14 @@ mod tests { let runtime = codex_state::StateRuntime::init( config.sqlite_home.clone(), config.default_model_provider_id.clone(), + /*metrics*/ None, ) .await .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); + let store = LocalThreadStore::new( + config.clone(), + codex_rollout::StateDbAccess::new(Some(runtime.clone())), + ); let mut builder = ThreadMetadataBuilder::new( thread_id, active_path.clone(), @@ -530,7 +540,10 @@ 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 = LocalThreadStore::new( + test_config(home.path()), + codex_rollout::StateDbAccess::none(), + ); 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) @@ -571,7 +584,10 @@ 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 = LocalThreadStore::new( + test_config(home.path()), + codex_rollout::StateDbAccess::none(), + ); let uuid = Uuid::from_u128(208); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let active_path = @@ -596,7 +612,10 @@ 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 = LocalThreadStore::new( + test_config(home.path()), + codex_rollout::StateDbAccess::none(), + ); 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"); @@ -636,10 +655,14 @@ mod tests { let runtime = codex_state::StateRuntime::init( config.sqlite_home.clone(), config.default_model_provider_id.clone(), + /*metrics*/ None, ) .await .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); + let store = LocalThreadStore::new( + config.clone(), + codex_rollout::StateDbAccess::new(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()); @@ -672,10 +695,14 @@ mod tests { let runtime = codex_state::StateRuntime::init( config.sqlite_home.clone(), config.default_model_provider_id.clone(), + /*metrics*/ None, ) .await .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); + let store = LocalThreadStore::new( + config.clone(), + codex_rollout::StateDbAccess::new(Some(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"); @@ -744,7 +771,10 @@ 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 = LocalThreadStore::new( + test_config(home.path()), + codex_rollout::StateDbAccess::none(), + ); 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"); @@ -792,10 +822,14 @@ mod tests { let runtime = codex_state::StateRuntime::init( config.sqlite_home.clone(), config.default_model_provider_id.clone(), + /*metrics*/ None, ) .await .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); + let store = LocalThreadStore::new( + config.clone(), + codex_rollout::StateDbAccess::new(Some(runtime.clone())), + ); let mut builder = ThreadMetadataBuilder::new( thread_id, rollout_path.clone(), @@ -846,10 +880,14 @@ mod tests { let runtime = codex_state::StateRuntime::init( config.sqlite_home.clone(), config.default_model_provider_id.clone(), + /*metrics*/ None, ) .await .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); + let store = LocalThreadStore::new( + config.clone(), + codex_rollout::StateDbAccess::new(Some(runtime.clone())), + ); let mut builder = ThreadMetadataBuilder::new( thread_id, stale_path.clone(), @@ -897,10 +935,14 @@ mod tests { let runtime = codex_state::StateRuntime::init( config.sqlite_home.clone(), config.default_model_provider_id.clone(), + /*metrics*/ None, ) .await .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); + let store = LocalThreadStore::new( + config.clone(), + codex_rollout::StateDbAccess::new(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()); @@ -932,7 +974,10 @@ 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 = LocalThreadStore::new( + test_config(home.path()), + codex_rollout::StateDbAccess::none(), + ); 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"); @@ -995,10 +1040,14 @@ mod tests { let runtime = codex_state::StateRuntime::init( config.sqlite_home.clone(), config.default_model_provider_id.clone(), + /*metrics*/ None, ) .await .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); + let store = LocalThreadStore::new( + config.clone(), + codex_rollout::StateDbAccess::new(Some(runtime.clone())), + ); let mut builder = ThreadMetadataBuilder::new( thread_id, rollout_path.clone(), @@ -1053,12 +1102,16 @@ mod tests { let runtime = codex_state::StateRuntime::init( config.sqlite_home.clone(), config.default_model_provider_id.clone(), + /*metrics*/ None, ) .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())); + let store = LocalThreadStore::new( + config.clone(), + codex_rollout::StateDbAccess::new(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()); @@ -1108,10 +1161,14 @@ mod tests { let runtime = codex_state::StateRuntime::init( config.sqlite_home.clone(), config.default_model_provider_id.clone(), + /*metrics*/ None, ) .await .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); + let store = LocalThreadStore::new( + config.clone(), + codex_rollout::StateDbAccess::new(Some(runtime.clone())), + ); let mut builder = ThreadMetadataBuilder::new( thread_id, archived_path.clone(), @@ -1147,7 +1204,10 @@ 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 = LocalThreadStore::new( + test_config(home.path()), + codex_rollout::StateDbAccess::none(), + ); 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/unarchive_thread.rs b/codex-rs/thread-store/src/local/unarchive_thread.rs index ad41db69ac..c7c3a6ada8 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_access = store.state_db_access(); 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(), + &state_db_access, ) .await .map_err(|err| ThreadStoreError::InvalidRequest { @@ -73,7 +73,7 @@ pub(super) async fn unarchive_thread( message: format!("failed to update unarchived thread timestamp: {err}"), })?; - if let Some(ctx) = state_db_ctx { + if let Some(ctx) = state_db_access.state_db() { let _ = ctx .mark_unarchived(thread_id, restored_path.as_path()) .await; @@ -118,7 +118,10 @@ mod tests { #[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 = LocalThreadStore::new( + test_config(home.path()), + codex_rollout::StateDbAccess::none(), + ); 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) @@ -156,10 +159,14 @@ mod tests { let runtime = codex_state::StateRuntime::init( home.path().to_path_buf(), config.default_model_provider_id.clone(), + /*metrics*/ None, ) .await .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); + let store = LocalThreadStore::new( + config.clone(), + codex_rollout::StateDbAccess::new(Some(runtime.clone())), + ); runtime .mark_backfill_complete(/*last_watermark*/ None) .await 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 ef69cfa8b8..55ad7d0643 100644 --- a/codex-rs/thread-store/src/local/update_thread_metadata.rs +++ b/codex-rs/thread-store/src/local/update_thread_metadata.rs @@ -310,11 +310,11 @@ async fn resolve_rollout_path( return Ok(ResolvedRolloutPath { path, archived }); } - let state_db_ctx = store.state_db().await; + let state_db_access = store.state_db_access(); let active_path = find_thread_path_by_id_str( store.config.codex_home.as_path(), &thread_id.to_string(), - state_db_ctx.as_deref(), + &state_db_access, ) .await .map_err(|err| ThreadStoreError::InvalidRequest { @@ -334,7 +334,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(), + &state_db_access, ) .await .map_err(|err| ThreadStoreError::InvalidRequest { @@ -376,7 +376,10 @@ mod tests { #[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 = LocalThreadStore::new( + test_config(home.path()), + codex_rollout::StateDbAccess::none(), + ); let uuid = Uuid::from_u128(301); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); write_session_file(home.path(), "2025-01-03T14-00-00", uuid).expect("session file"); @@ -411,10 +414,14 @@ mod tests { let runtime = codex_state::StateRuntime::init( home.path().to_path_buf(), config.default_model_provider_id.clone(), + /*metrics*/ None, ) .await .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); + let store = LocalThreadStore::new( + config.clone(), + codex_rollout::StateDbAccess::new(Some(runtime.clone())), + ); let thread = store .update_thread_metadata(UpdateThreadMetadataParams { @@ -451,10 +458,14 @@ mod tests { let runtime = codex_state::StateRuntime::init( config.sqlite_home.clone(), config.default_model_provider_id.clone(), + /*metrics*/ None, ) .await .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); + let store = LocalThreadStore::new( + config.clone(), + codex_rollout::StateDbAccess::new(Some(runtime.clone())), + ); store .update_thread_metadata(UpdateThreadMetadataParams { @@ -513,7 +524,10 @@ 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 = LocalThreadStore::new( + test_config(home.path()), + codex_rollout::StateDbAccess::none(), + ); 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) @@ -557,10 +571,11 @@ mod tests { let runtime = codex_state::StateRuntime::init( config.sqlite_home.clone(), config.default_model_provider_id.clone(), + /*metrics*/ None, ) .await .expect("state db should initialize"); - let store = LocalThreadStore::new(config, Some(runtime)); + let store = LocalThreadStore::new(config, codex_rollout::StateDbAccess::new(Some(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"); @@ -600,10 +615,11 @@ mod tests { let runtime = codex_state::StateRuntime::init( config.sqlite_home.clone(), config.default_model_provider_id.clone(), + /*metrics*/ None, ) .await .expect("state db should initialize"); - let store = LocalThreadStore::new(config, Some(runtime)); + let store = LocalThreadStore::new(config, codex_rollout::StateDbAccess::new(Some(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"); @@ -658,10 +674,14 @@ mod tests { let runtime = codex_state::StateRuntime::init( config.sqlite_home.clone(), config.default_model_provider_id.clone(), + /*metrics*/ None, ) .await .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); + let store = LocalThreadStore::new( + config.clone(), + codex_rollout::StateDbAccess::new(Some(runtime.clone())), + ); let uuid = Uuid::from_u128(311); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let path = @@ -825,7 +845,10 @@ 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 = LocalThreadStore::new( + test_config(home.path()), + codex_rollout::StateDbAccess::none(), + ); 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"); @@ -857,7 +880,10 @@ 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 = LocalThreadStore::new( + test_config(home.path()), + codex_rollout::StateDbAccess::none(), + ); let uuid = Uuid::from_u128(305); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let path = @@ -899,10 +925,14 @@ mod tests { let runtime = codex_state::StateRuntime::init( home.path().to_path_buf(), config.default_model_provider_id.clone(), + /*metrics*/ None, ) .await .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); + let store = LocalThreadStore::new( + config.clone(), + codex_rollout::StateDbAccess::new(Some(runtime.clone())), + ); runtime .mark_backfill_complete(/*last_watermark*/ None) .await @@ -962,10 +992,14 @@ mod tests { let runtime = codex_state::StateRuntime::init( home.path().to_path_buf(), config.default_model_provider_id.clone(), + /*metrics*/ None, ) .await .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); + let store = LocalThreadStore::new( + config.clone(), + codex_rollout::StateDbAccess::new(Some(runtime.clone())), + ); runtime .mark_backfill_complete(/*last_watermark*/ None) .await diff --git a/codex-rs/tui/src/app.rs b/codex-rs/tui/src/app.rs index dff288493c..093a77d674 100644 --- a/codex-rs/tui/src/app.rs +++ b/codex-rs/tui/src/app.rs @@ -145,7 +145,7 @@ use codex_protocol::openai_models::ModelUpgrade; use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig; #[cfg(target_os = "windows")] use codex_protocol::permissions::FileSystemSandboxKind; -use codex_rollout::StateDbHandle; +use codex_rollout::StateDbAccess; use codex_terminal_detection::user_agent; use codex_utils_absolute_path::AbsolutePathBuf; use color_eyre::eyre::Result; @@ -437,7 +437,7 @@ pub(crate) struct App { workspace_command_runner: Option, /// Config is stored here so we can recreate ChatWidgets as needed. pub(crate) config: Config, - pub(crate) state_db: Option, + pub(crate) state_db_access: StateDbAccess, pub(crate) active_profile: Option, cli_kv_overrides: Vec<(String, TomlValue)>, harness_overrides: ConfigOverrides, @@ -616,7 +616,7 @@ impl App { should_prompt_windows_sandbox_nux_at_startup: bool, remote_app_server_url: Option, remote_app_server_auth_token: Option, - state_db: Option, + state_db_access: StateDbAccess, environment_manager: Arc, ) -> Result { use tokio_stream::StreamExt; @@ -872,7 +872,7 @@ See the Codex keymap documentation for supported actions and examples." chat_widget, workspace_command_runner: Some(workspace_command_runner), config, - state_db, + state_db_access, active_profile, cli_kv_overrides, harness_overrides, diff --git a/codex-rs/tui/src/app/event_dispatch.rs b/codex-rs/tui/src/app/event_dispatch.rs index b2261f155a..ccff95474e 100644 --- a/codex-rs/tui/src/app/event_dispatch.rs +++ b/codex-rs/tui/src/app/event_dispatch.rs @@ -64,7 +64,7 @@ impl App { }, None => crate::AppServerTarget::Embedded, }, - self.state_db.clone(), + self.state_db_access.clone(), self.environment_manager.clone(), ) .await diff --git a/codex-rs/tui/src/app/session_lifecycle.rs b/codex-rs/tui/src/app/session_lifecycle.rs index 05aba144c8..402b1b2231 100644 --- a/codex-rs/tui/src/app/session_lifecycle.rs +++ b/codex-rs/tui/src/app/session_lifecycle.rs @@ -638,7 +638,7 @@ impl App { } else { match crate::session_resume::resolve_cwd_for_resume_or_fork( tui, - self.state_db.as_deref(), + self.state_db_access.as_deref(), ¤t_cwd, target_session.thread_id, target_session.path.as_deref(), diff --git a/codex-rs/tui/src/app/test_support.rs b/codex-rs/tui/src/app/test_support.rs index f34e22203f..499c3d6a30 100644 --- a/codex-rs/tui/src/app/test_support.rs +++ b/codex-rs/tui/src/app/test_support.rs @@ -21,7 +21,7 @@ pub(super) async fn make_test_app() -> App { chat_widget, workspace_command_runner: None, config, - state_db: None, + state_db_access: StateDbAccess::none(), active_profile: None, cli_kv_overrides: Vec::new(), harness_overrides: ConfigOverrides::default(), diff --git a/codex-rs/tui/src/app/tests.rs b/codex-rs/tui/src/app/tests.rs index b74bfab8f6..9ecb6860f0 100644 --- a/codex-rs/tui/src/app/tests.rs +++ b/codex-rs/tui/src/app/tests.rs @@ -1656,6 +1656,7 @@ fn update_memory_settings_updates_current_thread_memory_mode() -> Result<()> { let state_db = codex_state::StateRuntime::init( codex_home.path().to_path_buf(), app.config.model_provider_id.clone(), + /*metrics*/ None, ) .await .expect("state db should initialize"); @@ -3867,7 +3868,7 @@ async fn make_test_app() -> App { chat_widget, workspace_command_runner: None, config, - state_db: None, + state_db_access: StateDbAccess::none(), active_profile: None, cli_kv_overrides: Vec::new(), harness_overrides: ConfigOverrides::default(), @@ -3930,7 +3931,7 @@ async fn make_test_app_with_channels() -> ( chat_widget, workspace_command_runner: None, config, - state_db: None, + state_db_access: StateDbAccess::none(), active_profile: None, cli_kv_overrides: Vec::new(), harness_overrides: ConfigOverrides::default(), diff --git a/codex-rs/tui/src/app/thread_routing.rs b/codex-rs/tui/src/app/thread_routing.rs index 9a76b54168..77d08ebd6b 100644 --- a/codex-rs/tui/src/app/thread_routing.rs +++ b/codex-rs/tui/src/app/thread_routing.rs @@ -916,8 +916,12 @@ impl App { session.model_provider_id = notification.thread.model_provider.clone(); session.cwd = notification.thread.cwd.clone(); let rollout_path = notification.thread.path.clone(); - if let Some(model) = - read_session_model(self.state_db.as_deref(), thread_id, rollout_path.as_deref()).await + if let Some(model) = read_session_model( + self.state_db_access.as_deref(), + thread_id, + rollout_path.as_deref(), + ) + .await { session.model = model; } else if rollout_path.is_some() { diff --git a/codex-rs/tui/src/app/thread_session_state.rs b/codex-rs/tui/src/app/thread_session_state.rs index 4f5f48cfc8..89f4cd8c1f 100644 --- a/codex-rs/tui/src/app/thread_session_state.rs +++ b/codex-rs/tui/src/app/thread_session_state.rs @@ -85,8 +85,12 @@ impl App { session.active_permission_profile = active_permission_profile; session.instruction_source_paths = Vec::new(); session.rollout_path = thread.path.clone(); - if let Some(model) = - read_session_model(self.state_db.as_deref(), thread_id, thread.path.as_deref()).await + if let Some(model) = read_session_model( + self.state_db_access.as_deref(), + thread_id, + thread.path.as_deref(), + ) + .await { session.model = model; } else if thread.path.is_some() { diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index ac5b489af1..d836fba509 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -42,13 +42,14 @@ use codex_config::format_config_error_with_source; use codex_exec_server::EnvironmentManager; use codex_exec_server::ExecServerRuntimePaths; use codex_login::AuthConfig; +use codex_login::default_client::originator; use codex_login::default_client::set_default_client_residency_requirement; use codex_login::enforce_login_restrictions; use codex_protocol::ThreadId; use codex_protocol::config_types::AltScreenMode; use codex_protocol::config_types::SandboxMode; use codex_protocol::config_types::WindowsSandboxLevel; -use codex_rollout::StateDbHandle; +use codex_rollout::StateDbAccess; use codex_rollout::state_db; use codex_state::log_db; use codex_terminal_detection::terminal_info; @@ -273,7 +274,7 @@ async fn start_embedded_app_server( cloud_requirements: CloudRequirementsLoader, feedback: codex_feedback::CodexFeedback, log_db: Option, - state_db: Option, + state_db_access: StateDbAccess, environment_manager: Arc, ) -> color_eyre::Result { start_embedded_app_server_with( @@ -284,7 +285,7 @@ async fn start_embedded_app_server( cloud_requirements, feedback, log_db, - state_db, + state_db_access, environment_manager, InProcessAppServerClient::start, ) @@ -402,7 +403,7 @@ async fn start_app_server( cloud_requirements: CloudRequirementsLoader, feedback: codex_feedback::CodexFeedback, log_db: Option, - state_db: Option, + state_db_access: StateDbAccess, environment_manager: Arc, ) -> color_eyre::Result { match target { @@ -414,7 +415,7 @@ async fn start_app_server( cloud_requirements, feedback, log_db, - state_db, + state_db_access, environment_manager, ) .await @@ -429,7 +430,7 @@ async fn start_app_server( pub(crate) async fn start_app_server_for_picker( config: &Config, target: &AppServerTarget, - state_db: Option, + state_db_access: StateDbAccess, environment_manager: Arc, ) -> color_eyre::Result { let app_server = start_app_server( @@ -441,7 +442,7 @@ pub(crate) async fn start_app_server_for_picker( CloudRequirementsLoader::default(), codex_feedback::CodexFeedback::new(), /*log_db*/ None, - state_db, + state_db_access, environment_manager, ) .await?; @@ -452,11 +453,11 @@ pub(crate) async fn start_app_server_for_picker( pub(crate) async fn start_embedded_app_server_for_picker( config: &Config, ) -> color_eyre::Result { - let state_db = state_db::init(config).await; + let state_db = state_db::init(config, /*metrics*/ None).await; start_app_server_for_picker( config, &AppServerTarget::Embedded, - state_db, + StateDbAccess::new(state_db), Arc::new(EnvironmentManager::default_for_tests()), ) .await @@ -471,7 +472,7 @@ async fn start_embedded_app_server_with( cloud_requirements: CloudRequirementsLoader, feedback: codex_feedback::CodexFeedback, log_db: Option, - state_db: Option, + state_db_access: StateDbAccess, environment_manager: Arc, start_client: F, ) -> color_eyre::Result @@ -497,7 +498,7 @@ where cloud_requirements, feedback, log_db, - state_db, + state_db_access, environment_manager, config_warnings, session_source: serde_json::from_value(serde_json::json!("cli")) @@ -887,10 +888,45 @@ pub async fn run_main( ) .await; - let state_db = match &app_server_target { - AppServerTarget::Embedded => state_db::init(&config).await, - AppServerTarget::Remote { .. } => state_db::get_state_db(&config).await, + let otel_originator = originator().value; + let otel = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + crate::legacy_core::otel_init::build_provider( + &config, + env!("CARGO_PKG_VERSION"), + /*service_name_override*/ None, + /*default_analytics_enabled*/ true, + ) + })) { + Ok(Ok(otel)) => otel, + Ok(Err(e)) => { + #[allow(clippy::print_stderr)] + { + eprintln!("Could not create otel exporter: {e}"); + } + None + } + Err(_) => { + #[allow(clippy::print_stderr)] + { + eprintln!("Could not create otel exporter: panicked during initialization"); + } + None + } }; + crate::legacy_core::otel_init::record_process_start(otel.as_ref(), otel_originator.as_str()); + let state_db_metrics = crate::legacy_core::otel_init::sqlite_metrics_recorder( + otel.as_ref(), + otel_originator.as_str(), + ); + + let state_db = match &app_server_target { + AppServerTarget::Embedded => state_db::init(&config, state_db_metrics.clone()).await, + AppServerTarget::Remote { .. } => { + state_db::get_state_db(&config, state_db_metrics.clone()).await + } + }; + let state_db_access = + StateDbAccess::new(state_db.clone()).with_standalone_metrics(state_db_metrics.clone()); let effective_toml = config.config_layer_stack.effective_config(); match effective_toml.try_into() { @@ -898,7 +934,7 @@ pub async fn run_main( match crate::legacy_core::personality_migration::maybe_migrate_personality( &config.codex_home, &config_toml, - state_db.clone(), + state_db_access.clone(), ) .await { @@ -1029,31 +1065,6 @@ pub async fn run_main( ensure_oss_provider_ready(provider_id, &config).await?; } - let otel = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - crate::legacy_core::otel_init::build_provider( - &config, - env!("CARGO_PKG_VERSION"), - /*service_name_override*/ None, - /*default_analytics_enabled*/ true, - ) - })) { - Ok(Ok(otel)) => otel, - Ok(Err(e)) => { - #[allow(clippy::print_stderr)] - { - eprintln!("Could not create otel exporter: {e}"); - } - None - } - Err(_) => { - #[allow(clippy::print_stderr)] - { - eprintln!("Could not create otel exporter: panicked during initialization"); - } - None - } - }; - 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()); @@ -1084,7 +1095,7 @@ pub async fn run_main( cloud_requirements, feedback, log_db, - state_db, + state_db_access, remote_url, remote_auth_token, environment_manager, @@ -1106,7 +1117,7 @@ async fn run_ratatui_app( mut cloud_requirements: CloudRequirementsLoader, feedback: codex_feedback::CodexFeedback, log_db: Option, - state_db: Option, + state_db_access: StateDbAccess, remote_url: Option, remote_auth_token: Option, environment_manager: Arc, @@ -1166,7 +1177,7 @@ async fn run_ratatui_app( cloud_requirements.clone(), feedback.clone(), log_db.clone(), - state_db.clone(), + state_db_access.clone(), environment_manager.clone(), ) .await @@ -1414,7 +1425,7 @@ async fn run_ratatui_app( } else { match resolve_cwd_for_resume_or_fork( &mut tui, - state_db.as_deref(), + state_db_access.as_deref(), ¤t_cwd, target_session.thread_id, target_session.path.as_deref(), @@ -1505,7 +1516,7 @@ async fn run_ratatui_app( cloud_requirements.clone(), feedback.clone(), log_db.clone(), - state_db.clone(), + state_db_access.clone(), environment_manager.clone(), ) .await @@ -1536,7 +1547,7 @@ async fn run_ratatui_app( should_prompt_windows_sandbox_nux_at_startup, remote_url, remote_auth_token, - state_db, + state_db_access, environment_manager, ) .await; @@ -1741,7 +1752,7 @@ mod tests { async fn start_test_embedded_app_server( config: Config, ) -> color_eyre::Result { - let state_db = state_db::init(&config).await; + let state_db = state_db::init(&config, /*metrics*/ None).await; start_embedded_app_server( Arg0DispatchPaths::default(), config, @@ -1750,7 +1761,7 @@ mod tests { CloudRequirementsLoader::default(), codex_feedback::CodexFeedback::new(), /*log_db*/ None, - state_db, + StateDbAccess::new(state_db), Arc::new(EnvironmentManager::default_for_tests()), ) .await @@ -2214,6 +2225,7 @@ mod tests { let state_runtime = codex_state::StateRuntime::init( config.codex_home.to_path_buf(), config.model_provider_id.clone(), + /*metrics*/ None, ) .await .map_err(std::io::Error::other)?; @@ -2272,7 +2284,7 @@ mod tests { CloudRequirementsLoader::default(), codex_feedback::CodexFeedback::new(), /*log_db*/ None, - /*state_db*/ None, + StateDbAccess::none(), Arc::new(EnvironmentManager::default_for_tests()), |_args| async { Err(std::io::Error::other("boom")) }, ) diff --git a/codex-rs/tui/src/onboarding/auth.rs b/codex-rs/tui/src/onboarding/auth.rs index 91362bf5a5..d0b4bc7276 100644 --- a/codex-rs/tui/src/onboarding/auth.rs +++ b/codex-rs/tui/src/onboarding/auth.rs @@ -1065,7 +1065,7 @@ mod tests { .await, feedback: codex_feedback::CodexFeedback::new(), log_db: None, - state_db: None, + state_db_access: codex_app_server_client::StateDbAccess::none(), environment_manager: Arc::new( codex_app_server_client::EnvironmentManager::default_for_tests(), ),