diff --git a/.bazelrc b/.bazelrc index a5fea40281..911e0802b5 100644 --- a/.bazelrc +++ b/.bazelrc @@ -183,5 +183,15 @@ common:ci-v8 --build_metadata=TAG_os=linux common:ci-v8 --config=remote common:ci-v8 --strategy=remote +# Source-built Bazel V8 artifacts use the in-process sandbox by default. This +# does not affect Cargo's default prebuilt rusty_v8 path. +common --@v8//:v8_enable_pointer_compression=True +common --@v8//:v8_enable_sandbox=True + +# Keep currently published rusty_v8 release artifacts non-sandboxed until the +# artifact migration ships matching Rust feature selection for Cargo consumers. +common:v8-release-compat --@v8//:v8_enable_pointer_compression=False +common:v8-release-compat --@v8//:v8_enable_sandbox=False + # Optional per-user local overrides. try-import %workspace%/user.bazelrc diff --git a/.github/scripts/rusty_v8_bazel.py b/.github/scripts/rusty_v8_bazel.py index ec73e0e5a7..5ad6d3c408 100644 --- a/.github/scripts/rusty_v8_bazel.py +++ b/.github/scripts/rusty_v8_bazel.py @@ -63,8 +63,10 @@ def bazel_output_files( platform: str, labels: list[str], compilation_mode: str = "fastbuild", + bazel_configs: list[str] | None = None, ) -> list[Path]: expression = "set(" + " ".join(labels) + ")" + bazel_configs = bazel_configs or [] result = subprocess.run( [ "bazel", @@ -72,6 +74,7 @@ def bazel_output_files( "-c", compilation_mode, f"--platforms=@llvm//platforms:{platform}", + *[f"--config={config}" for config in bazel_configs], "--output=files", expression, ], @@ -87,7 +90,9 @@ def bazel_build( platform: str, labels: list[str], compilation_mode: str = "fastbuild", + bazel_configs: list[str] | None = None, ) -> None: + bazel_configs = bazel_configs or [] subprocess.run( [ "bazel", @@ -95,6 +100,7 @@ def bazel_build( "-c", compilation_mode, f"--platforms=@llvm//platforms:{platform}", + *[f"--config={config}" for config in bazel_configs], *labels, ], cwd=ROOT, @@ -106,13 +112,14 @@ def ensure_bazel_output_files( platform: str, labels: list[str], compilation_mode: str = "fastbuild", + bazel_configs: list[str] | None = None, ) -> list[Path]: - outputs = bazel_output_files(platform, labels, compilation_mode) + outputs = bazel_output_files(platform, labels, compilation_mode, bazel_configs) if all(path.exists() for path in outputs): return outputs - bazel_build(platform, labels, compilation_mode) - outputs = bazel_output_files(platform, labels, compilation_mode) + bazel_build(platform, labels, compilation_mode, bazel_configs) + outputs = bazel_output_files(platform, labels, compilation_mode, bazel_configs) missing = [str(path) for path in outputs if not path.exists()] if missing: raise SystemExit(f"missing built outputs for {labels}: {missing}") @@ -187,8 +194,9 @@ def single_bazel_output_file( platform: str, label: str, compilation_mode: str = "fastbuild", + bazel_configs: list[str] | None = None, ) -> Path: - outputs = ensure_bazel_output_files(platform, [label], compilation_mode) + outputs = ensure_bazel_output_files(platform, [label], compilation_mode, bazel_configs) if len(outputs) != 1: raise SystemExit(f"expected exactly one output for {label}, found {outputs}") return outputs[0] @@ -198,11 +206,17 @@ def merged_musl_archive( platform: str, lib_path: Path, compilation_mode: str = "fastbuild", + bazel_configs: list[str] | None = None, ) -> Path: - llvm_ar = single_bazel_output_file(platform, LLVM_AR_LABEL, compilation_mode) - llvm_ranlib = single_bazel_output_file(platform, LLVM_RANLIB_LABEL, compilation_mode) + llvm_ar = single_bazel_output_file(platform, LLVM_AR_LABEL, compilation_mode, bazel_configs) + llvm_ranlib = single_bazel_output_file( + platform, + LLVM_RANLIB_LABEL, + compilation_mode, + bazel_configs, + ) runtime_archives = [ - single_bazel_output_file(platform, label, compilation_mode) + single_bazel_output_file(platform, label, compilation_mode, bazel_configs) for label in MUSL_RUNTIME_ARCHIVE_LABELS ] @@ -233,11 +247,13 @@ def stage_release_pair( target: str, output_dir: Path, compilation_mode: str = "fastbuild", + bazel_configs: list[str] | None = None, ) -> None: outputs = ensure_bazel_output_files( platform, [release_pair_label(target)], compilation_mode, + bazel_configs, ) try: @@ -254,7 +270,7 @@ def stage_release_pair( staged_library = output_dir / staged_archive_name(target, lib_path) staged_binding = output_dir / f"src_binding_release_{target}.rs" source_archive = ( - merged_musl_archive(platform, lib_path, compilation_mode) + merged_musl_archive(platform, lib_path, compilation_mode, bazel_configs) if is_musl_archive_target(target, lib_path) else lib_path ) @@ -293,6 +309,12 @@ def parse_args() -> argparse.Namespace: stage_release_pair_parser.add_argument("--platform", required=True) stage_release_pair_parser.add_argument("--target", required=True) stage_release_pair_parser.add_argument("--output-dir", required=True) + stage_release_pair_parser.add_argument( + "--bazel-config", + action="append", + default=[], + dest="bazel_configs", + ) stage_release_pair_parser.add_argument( "--compilation-mode", default="fastbuild", @@ -330,6 +352,7 @@ def main() -> int: target=args.target, output_dir=Path(args.output_dir), compilation_mode=args.compilation_mode, + bazel_configs=args.bazel_configs, ) return 0 if args.command == "resolved-v8-crate-version": diff --git a/.github/scripts/verify_cargo_workspace_manifests.py b/.github/scripts/verify_cargo_workspace_manifests.py index 4812e2428d..93b41ea59f 100644 --- a/.github/scripts/verify_cargo_workspace_manifests.py +++ b/.github/scripts/verify_cargo_workspace_manifests.py @@ -25,7 +25,10 @@ TOP_LEVEL_NAME_EXCEPTIONS = { UTILITY_NAME_EXCEPTIONS = { "path-utils": "codex-utils-path", } -MANIFEST_FEATURE_EXCEPTIONS = {} +MANIFEST_FEATURE_EXCEPTIONS = { + "codex-rs/code-mode/Cargo.toml": {"sandbox": ("v8/v8_enable_sandbox",)}, + "codex-rs/v8-poc/Cargo.toml": {"sandbox": ("v8/v8_enable_sandbox",)}, +} OPTIONAL_DEPENDENCY_EXCEPTIONS = set() INTERNAL_DEPENDENCY_FEATURE_EXCEPTIONS = {} diff --git a/.github/workflows/rusty-v8-release.yml b/.github/workflows/rusty-v8-release.yml index ee92eff4fa..9731948f60 100644 --- a/.github/workflows/rusty-v8-release.yml +++ b/.github/workflows/rusty-v8-release.yml @@ -111,6 +111,7 @@ jobs: -c opt "--platforms=@llvm//platforms:${PLATFORM}" + --config=v8-release-compat "${pair_target}" "${extra_targets[@]}" --build_metadata=COMMIT_SHA=$(git rev-parse HEAD) @@ -134,6 +135,7 @@ jobs: --platform "${PLATFORM}" \ --target "${TARGET}" \ --compilation-mode opt \ + --bazel-config v8-release-compat \ --output-dir "dist/${TARGET}" - name: Upload staged musl artifacts diff --git a/.github/workflows/v8-canary.yml b/.github/workflows/v8-canary.yml index 119d042275..3b59348775 100644 --- a/.github/workflows/v8-canary.yml +++ b/.github/workflows/v8-canary.yml @@ -105,6 +105,7 @@ jobs: bazel_args=( build "--platforms=@llvm//platforms:${PLATFORM}" + --config=v8-release-compat "${pair_target}" "${extra_targets[@]}" --build_metadata=COMMIT_SHA=$(git rev-parse HEAD) @@ -127,6 +128,7 @@ jobs: python3 .github/scripts/rusty_v8_bazel.py stage-release-pair \ --platform "${PLATFORM}" \ --target "${TARGET}" \ + --bazel-config v8-release-compat \ --output-dir "dist/${TARGET}" - name: Upload staged musl artifacts diff --git a/MODULE.bazel b/MODULE.bazel index be00937946..f675038791 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -327,6 +327,18 @@ crate.annotation( "RUSTY_V8_SRC_BINDING_PATH": "$(execpath @v8_targets//:rusty_v8_binding_for_target)", }, crate = "v8", + # Keep the Rust feature aligned with the source-built Bazel artifacts. + # Windows MSVC still consumes upstream non-sandboxed prebuilts. + crate_features_select = { + "aarch64-apple-darwin": ["v8_enable_sandbox"], + "aarch64-pc-windows-gnullvm": ["v8_enable_sandbox"], + "aarch64-unknown-linux-gnu": ["v8_enable_sandbox"], + "aarch64-unknown-linux-musl": ["v8_enable_sandbox"], + "x86_64-apple-darwin": ["v8_enable_sandbox"], + "x86_64-pc-windows-gnullvm": ["v8_enable_sandbox"], + "x86_64-unknown-linux-gnu": ["v8_enable_sandbox"], + "x86_64-unknown-linux-musl": ["v8_enable_sandbox"], + }, gen_build_script = "on", patch_args = ["-p1"], patches = [ diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index ef84ae89d1..1b92b5afd0 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2420,6 +2420,7 @@ dependencies = [ "bm25", "chrono", "clap", + "codex-agent-graph-store", "codex-analytics", "codex-api", "codex-app-server-protocol", @@ -2709,6 +2710,7 @@ dependencies = [ "serde", "serde_json", "serial_test", + "sha2", "tempfile", "test-case", "thiserror 2.0.18", @@ -2717,6 +2719,7 @@ dependencies = [ "tokio-util", "tracing", "uuid", + "wiremock", ] [[package]] diff --git a/codex-rs/app-server-client/src/lib.rs b/codex-rs/app-server-client/src/lib.rs index 79141da24e..0911cc448d 100644 --- a/codex-rs/app-server-client/src/lib.rs +++ b/codex-rs/app-server-client/src/lib.rs @@ -29,7 +29,6 @@ pub use codex_app_server::in_process::DEFAULT_IN_PROCESS_CHANNEL_CAPACITY; pub use codex_app_server::in_process::InProcessServerEvent; use codex_app_server::in_process::InProcessStartArgs; use codex_app_server::in_process::LogDbLayer; -pub use codex_app_server::in_process::StateDbHandle; use codex_app_server_protocol::ClientInfo; use codex_app_server_protocol::ClientNotification; use codex_app_server_protocol::ClientRequest; @@ -47,6 +46,7 @@ use codex_config::LoaderOverrides; use codex_config::NoopThreadConfigLoader; use codex_config::RemoteThreadConfigLoader; use codex_config::ThreadConfigLoader; +pub use codex_core::StateDbHandle; use codex_core::config::Config; pub use codex_exec_server::EnvironmentManager; pub use codex_exec_server::EnvironmentManagerArgs; diff --git a/codex-rs/app-server/src/bespoke_event_handling.rs b/codex-rs/app-server/src/bespoke_event_handling.rs index 1907f1c3f4..55b3fea814 100644 --- a/codex-rs/app-server/src/bespoke_event_handling.rs +++ b/codex-rs/app-server/src/bespoke_event_handling.rs @@ -2629,7 +2629,8 @@ mod tests { config.model_provider.clone(), config.codex_home.to_path_buf(), Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - ), + ) + .await, ); let codex_core::NewThread { thread_id: conversation_id, @@ -3214,7 +3215,8 @@ mod tests { config.model_provider.clone(), config.codex_home.to_path_buf(), Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - ), + ) + .await, ); let codex_core::NewThread { thread_id: conversation_id, diff --git a/codex-rs/app-server/src/in_process.rs b/codex-rs/app-server/src/in_process.rs index a8c023e95a..73cfe75a54 100644 --- a/codex-rs/app-server/src/in_process.rs +++ b/codex-rs/app-server/src/in_process.rs @@ -82,11 +82,12 @@ use codex_config::CloudRequirementsLoader; use codex_config::LoaderOverrides; use codex_config::ThreadConfigLoader; use codex_core::config::Config; +use codex_core::init_state_db_from_config; use codex_exec_server::EnvironmentManager; use codex_feedback::CodexFeedback; use codex_login::AuthManager; use codex_protocol::protocol::SessionSource; -pub use codex_rollout::StateDbHandle; +use codex_rollout::state_db::StateDbHandle; pub use codex_state::log_db::LogDbLayer; use tokio::sync::mpsc; use tokio::sync::oneshot; @@ -127,7 +128,7 @@ pub struct InProcessStartArgs { pub feedback: CodexFeedback, /// SQLite tracing layer used to flush recently emitted logs before feedback upload. pub log_db: Option, - /// Process-wide SQLite state handle shared with embedded app-server consumers. + /// Optional state DB handle to use for the in-process runtime. pub state_db: Option, /// Environment manager used by core execution and filesystem operations. pub environment_manager: Arc, @@ -344,7 +345,7 @@ impl InProcessClientHandle { /// the runtime is shut down and an `InvalidData` error is returned. pub async fn start(args: InProcessStartArgs) -> IoResult { let initialize = args.initialize.clone(); - let client = start_uninitialized(args); + let client = start_uninitialized(args).await; let initialize_response = client .request(ClientRequest::Initialize { @@ -364,8 +365,12 @@ pub async fn start(args: InProcessStartArgs) -> IoResult Ok(client) } -fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle { +async fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle { let channel_capacity = args.channel_capacity.max(1); + let state_db = match args.state_db.clone() { + Some(state_db) => Some(state_db), + None => init_state_db_from_config(args.config.as_ref()).await, + }; let (client_tx, mut client_rx) = mpsc::channel::(channel_capacity); let (event_tx, event_rx) = mpsc::channel::(channel_capacity); @@ -414,6 +419,12 @@ fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle { ); let (processor_tx, mut processor_rx) = mpsc::channel::(channel_capacity); let mut processor_handle = tokio::spawn(async move { + let Some(state_db) = state_db else { + warn!( + "in-process app-server state db initialization failed; shutting down processor task" + ); + return; + }; let processor = Arc::new(MessageProcessor::new(MessageProcessorArgs { outgoing: Arc::clone(&processor_outgoing), analytics_events_client, @@ -423,7 +434,7 @@ fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle { environment_manager: args.environment_manager, feedback: args.feedback, log_db: args.log_db, - state_db: args.state_db, + state_db, config_warnings: args.config_warnings, session_source: args.session_source, auth_manager, @@ -761,7 +772,7 @@ mod tests { ) -> InProcessClientHandle { let codex_home = TempDir::new().expect("temp dir"); let config = Arc::new(build_test_config(codex_home.path()).await); - let state_db = codex_rollout::state_db::try_init(config.as_ref()) + let state_db = init_state_db_from_config(config.as_ref()) .await .expect("state db should initialize for in-process test"); let args = InProcessStartArgs { @@ -819,7 +830,7 @@ mod tests { } #[tokio::test] - async fn in_process_allows_device_key_requests_to_reach_device_key_processor() { + async fn in_process_allows_device_key_requests_to_reach_device_key_api() { let client = start_test_client(SessionSource::Cli).await; const MALFORMED_KEY_ID_MESSAGE: &str = concat!( "invalid device key payload: keyId must be dk_hse_, dk_tpm_, or dk_osn_ ", diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index 4b1723d88c..4013bbe76b 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -50,11 +50,11 @@ use codex_config::TextRange as CoreTextRange; use codex_core::ExecPolicyError; use codex_core::check_execpolicy_for_warnings; use codex_core::config::find_codex_home; +use codex_core::init_state_db_from_config; use codex_exec_server::EnvironmentManager; use codex_exec_server::ExecServerRuntimePaths; use codex_feedback::CodexFeedback; use codex_protocol::protocol::SessionSource; -use codex_rollout::state_db as rollout_state_db; use codex_state::log_db; use tokio::sync::mpsc; use tokio::sync::oneshot; @@ -487,9 +487,9 @@ pub async fn run_main_with_transport_options( } }; - let state_db_result = rollout_state_db::try_init(&config).await; - let state_db_init_error = state_db_result.as_ref().err().map(ToString::to_string); - let state_db = state_db_result.ok(); + let state_db = init_state_db_from_config(&config) + .await + .ok_or_else(|| std::io::Error::other("failed to initialize sqlite state db"))?; if should_run_personality_migration { let effective_toml = config.config_layer_stack.effective_config(); @@ -598,10 +598,12 @@ pub async fn run_main_with_transport_options( let feedback_layer = feedback.logger_layer(); let feedback_metadata_layer = feedback.metadata_layer(); - let log_db = state_db.clone().map(log_db::start); - let log_db_layer = log_db - .clone() - .map(|layer| layer.with_filter(Targets::new().with_default(Level::TRACE))); + let log_db = log_db::start(state_db.clone()); + let log_db_layer = Some( + log_db + .clone() + .with_filter(Targets::new().with_default(Level::TRACE)), + ); let otel_logger_layer = otel.as_ref().and_then(|o| o.logger_layer()); let otel_tracing_layer = otel.as_ref().and_then(|o| o.tracing_layer()); let _ = tracing_subscriber::registry() @@ -618,10 +620,6 @@ pub async fn run_main_with_transport_options( None => error!("{}", warning.summary), } } - if let Some(err) = &state_db_init_error { - error!("failed to initialize sqlite state db: {err}"); - } - let transport_shutdown_token = CancellationToken::new(); let mut transport_accept_handles = Vec::>::new(); @@ -666,25 +664,17 @@ pub async fn run_main_with_transport_options( let auth_manager = AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false).await; - let remote_control_config_enabled = config.features.enabled(Feature::RemoteControl); - let remote_control_enabled = remote_control_config_enabled && state_db.is_some(); - if remote_control_config_enabled && state_db.is_none() { - error!("remote control disabled because sqlite state db is unavailable"); - } + let remote_control_enabled = config.features.enabled(Feature::RemoteControl); if transport_accept_handles.is_empty() && !remote_control_enabled { return Err(std::io::Error::new( ErrorKind::InvalidInput, - if remote_control_config_enabled && state_db.is_none() { - "no transport configured; remote control disabled because sqlite state db is unavailable" - } else { - "no transport configured; use --listen or enable remote control" - }, + "no transport configured; use --listen or enable remote control", )); } let (remote_control_accept_handle, remote_control_handle) = start_remote_control( config.chatgpt_base_url.clone(), - state_db.clone(), + Some(state_db.clone()), auth_manager.clone(), transport_event_tx.clone(), transport_shutdown_token.clone(), @@ -768,7 +758,7 @@ pub async fn run_main_with_transport_options( config_manager, environment_manager, feedback: feedback.clone(), - log_db, + log_db: Some(log_db), state_db: state_db.clone(), config_warnings, session_source, diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 8d26cec70f..4c1d16eeac 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -61,6 +61,7 @@ use codex_app_server_protocol::experimental_required_message; use codex_arg0::Arg0DispatchPaths; use codex_chatgpt::workspace_settings; use codex_core::ThreadManager; +use codex_core::agent_graph_store_from_state_db; use codex_core::config::Config; use codex_core::thread_store_from_config; use codex_exec_server::EnvironmentManager; @@ -254,7 +255,7 @@ pub(crate) struct MessageProcessorArgs { pub(crate) environment_manager: Arc, pub(crate) feedback: CodexFeedback, pub(crate) log_db: Option, - pub(crate) state_db: Option, + pub(crate) state_db: StateDbHandle, pub(crate) config_warnings: Vec, pub(crate) session_source: SessionSource, pub(crate) auth_manager: Arc, @@ -291,14 +292,16 @@ impl MessageProcessor { // affect per-thread behavior, but they must not move newly started, // resumed, or forked threads to a different persistence backend/root. let thread_store = thread_store_from_config(config.as_ref(), state_db.clone()); + let agent_graph_store = agent_graph_store_from_state_db(state_db.clone()); let thread_manager = Arc::new(ThreadManager::new( config.as_ref(), auth_manager.clone(), session_source, environment_manager, Some(analytics_events_client.clone()), - Arc::clone(&thread_store), state_db.clone(), + Arc::clone(&thread_store), + agent_graph_store.clone(), )); thread_manager .plugins_manager() @@ -344,7 +347,7 @@ impl MessageProcessor { Arc::clone(&config), feedback, log_db, - state_db.clone(), + Some(state_db.clone()), ); let git_processor = GitRequestProcessor::new(); let initialize_processor = InitializeRequestProcessor::new( @@ -395,7 +398,7 @@ impl MessageProcessor { thread_watch_manager.clone(), Arc::clone(&thread_list_state_permit), thread_goal_processor.clone(), - state_db.clone(), + Some(state_db.clone()), ); let turn_processor = TurnRequestProcessor::new( auth_manager.clone(), diff --git a/codex-rs/app-server/src/message_processor_tracing_tests.rs b/codex-rs/app-server/src/message_processor_tracing_tests.rs index 45ea709180..67b668bafc 100644 --- a/codex-rs/app-server/src/message_processor_tracing_tests.rs +++ b/codex-rs/app-server/src/message_processor_tracing_tests.rs @@ -32,6 +32,7 @@ use codex_config::CloudRequirementsLoader; use codex_config::LoaderOverrides; use codex_core::config::Config; use codex_core::config::ConfigBuilder; +use codex_core::init_state_db_from_config; use codex_exec_server::EnvironmentManager; use codex_feedback::CodexFeedback; use codex_login::AuthManager; @@ -281,6 +282,9 @@ async fn build_test_processor( outgoing_tx, analytics_events_client.clone(), )); + let state_db = init_state_db_from_config(config.as_ref()) + .await + .expect("tracing test processor requires state db"); let processor = Arc::new(MessageProcessor::new(MessageProcessorArgs { outgoing, analytics_events_client, @@ -290,7 +294,7 @@ async fn build_test_processor( environment_manager: Arc::new(EnvironmentManager::default_for_tests()), feedback: CodexFeedback::new(), log_db: None, - state_db: None, + state_db, config_warnings: Vec::new(), session_source: SessionSource::VSCode, auth_manager, diff --git a/codex-rs/app-server/src/request_processors/device_key_processor.rs b/codex-rs/app-server/src/request_processors/device_key_processor.rs index ea0a96c2af..e1da459a2e 100644 --- a/codex-rs/app-server/src/request_processors/device_key_processor.rs +++ b/codex-rs/app-server/src/request_processors/device_key_processor.rs @@ -33,8 +33,8 @@ use codex_device_key::RemoteControlClientConnectionAudience; use codex_device_key::RemoteControlClientConnectionSignPayload; use codex_device_key::RemoteControlClientEnrollmentAudience; use codex_device_key::RemoteControlClientEnrollmentSignPayload; +use codex_rollout::state_db::StateDbHandle; use codex_state::DeviceKeyBindingRecord; -use codex_state::StateRuntime; #[derive(Clone)] pub(crate) struct DeviceKeyRequestProcessor { @@ -43,10 +43,7 @@ pub(crate) struct DeviceKeyRequestProcessor { } impl DeviceKeyRequestProcessor { - pub(crate) fn new( - outgoing: Arc, - state_db: Option>, - ) -> Self { + pub(crate) fn new(outgoing: Arc, state_db: StateDbHandle) -> Self { Self { outgoing, store: DeviceKeyStore::new(Arc::new(StateDeviceKeyBindingStore::new(state_db))), @@ -170,25 +167,18 @@ async fn sign_device_key( } struct StateDeviceKeyBindingStore { - state_db: Option>, + state_db: StateDbHandle, } impl StateDeviceKeyBindingStore { - fn new(state_db: Option>) -> Self { + fn new(state_db: StateDbHandle) -> Self { Self { state_db } } - - async fn state_db(&self) -> Result, DeviceKeyError> { - self.state_db - .clone() - .ok_or_else(|| DeviceKeyError::Platform("sqlite state db unavailable".to_string())) - } } impl fmt::Debug for StateDeviceKeyBindingStore { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("StateDeviceKeyBindingStore") - .field("has_state_db", &self.state_db.is_some()) .finish_non_exhaustive() } } @@ -196,7 +186,7 @@ impl fmt::Debug for StateDeviceKeyBindingStore { #[async_trait] impl DeviceKeyBindingStore for StateDeviceKeyBindingStore { async fn get_binding(&self, key_id: &str) -> Result, DeviceKeyError> { - let state_db = self.state_db().await?; + let state_db = self.state_db.clone(); state_db .get_device_key_binding(key_id) .await @@ -214,7 +204,7 @@ impl DeviceKeyBindingStore for StateDeviceKeyBindingStore { key_id: &str, binding: &DeviceKeyBinding, ) -> Result<(), DeviceKeyError> { - let state_db = self.state_db().await?; + let state_db = self.state_db.clone(); state_db .upsert_device_key_binding(&DeviceKeyBindingRecord { key_id: key_id.to_string(), diff --git a/codex-rs/app-server/src/request_processors/thread_goal_processor.rs b/codex-rs/app-server/src/request_processors/thread_goal_processor.rs index 0e12e44ce5..ff192e34b2 100644 --- a/codex-rs/app-server/src/request_processors/thread_goal_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_goal_processor.rs @@ -7,7 +7,7 @@ pub(crate) struct ThreadGoalRequestProcessor { outgoing: Arc, config: Arc, thread_state_manager: ThreadStateManager, - state_db: Option, + state_db: StateDbHandle, } impl ThreadGoalRequestProcessor { @@ -16,7 +16,7 @@ impl ThreadGoalRequestProcessor { outgoing: Arc, config: Arc, thread_state_manager: ThreadStateManager, - state_db: Option, + state_db: StateDbHandle, ) -> Self { Self { thread_manager, @@ -72,23 +72,6 @@ impl ThreadGoalRequestProcessor { } } - pub(crate) async fn pending_resume_goal_state( - &self, - thread: &CodexThread, - ) -> (bool, Option) { - let emit_thread_goal_update = self.config.features.enabled(Feature::Goals); - let thread_goal_state_db = if emit_thread_goal_update { - if let Some(state_db) = thread.state_db() { - Some(state_db) - } else { - self.state_db.clone() - } - } else { - None - }; - (emit_thread_goal_update, thread_goal_state_db) - } - async fn thread_goal_set_inner( &self, request_id: ConnectionRequestId, @@ -110,7 +93,7 @@ impl ThreadGoalRequestProcessor { None => find_thread_path_by_id_str( &self.config.codex_home, &thread_id.to_string(), - self.state_db.as_deref(), + Some(self.state_db.as_ref()), ) .await .map_err(|err| { @@ -275,7 +258,7 @@ impl ThreadGoalRequestProcessor { None => find_thread_path_by_id_str( &self.config.codex_home, &thread_id.to_string(), - self.state_db.as_deref(), + Some(self.state_db.as_ref()), ) .await .map_err(|err| { @@ -339,7 +322,7 @@ impl ThreadGoalRequestProcessor { find_thread_path_by_id_str( &self.config.codex_home, &thread_id.to_string(), - self.state_db.as_deref(), + Some(self.state_db.as_ref()), ) .await .map_err(|err| { @@ -348,9 +331,7 @@ impl ThreadGoalRequestProcessor { .ok_or_else(|| invalid_request(format!("thread not found: {thread_id}")))?; } - self.state_db - .clone() - .ok_or_else(|| internal_error("sqlite state db unavailable for thread goals")) + Ok(self.state_db.clone()) } async fn emit_thread_goal_snapshot(&self, thread_id: ThreadId) { diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index 1921fb69b1..fd61182788 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -2611,10 +2611,10 @@ impl ThreadRequestProcessor { ))); }; - let (emit_thread_goal_update, thread_goal_state_db) = self - .thread_goal_processor - .pending_resume_goal_state(existing_thread.as_ref()) - .await; + let emit_thread_goal_update = self.config.features.enabled(Feature::Goals); + let thread_goal_state_db = emit_thread_goal_update + .then(|| self.state_db.clone()) + .flatten(); let command = crate::thread_state::ThreadListenerCommand::SendThreadResumeResponse( Box::new(crate::thread_state::PendingThreadResumeRequest { diff --git a/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs b/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs index b04eb12a45..9363217864 100644 --- a/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs +++ b/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs @@ -42,6 +42,7 @@ use codex_core::config::ConfigBuilder; use codex_exec_server::EnvironmentManager; use codex_feedback::CodexFeedback; use codex_protocol::protocol::SessionSource; +use codex_state::StateRuntime; use codex_thread_store::InMemoryThreadStore; use pretty_assertions::assert_eq; use tempfile::TempDir; @@ -67,6 +68,13 @@ async fn thread_start_with_non_local_thread_store_does_not_create_local_persiste .loader_overrides(loader_overrides.clone()) .build() .await?; + let sqlite_home = TempDir::new()?; + let state_db = StateRuntime::init( + sqlite_home.path().to_path_buf(), + config.model_provider_id.clone(), + ) + .await + .expect("remote thread store regression test should initialize state db"); let thread_store = InMemoryThreadStore::for_id(store_id.clone()); let _in_memory_store = InMemoryThreadStoreId { store_id }; @@ -80,7 +88,7 @@ async fn thread_start_with_non_local_thread_store_does_not_create_local_persiste thread_config_loader: Arc::new(NoopThreadConfigLoader), feedback: CodexFeedback::new(), log_db: None, - state_db: None, + state_db: Some(state_db), environment_manager: Arc::new(EnvironmentManager::default_for_tests()), config_warnings: Vec::new(), session_source: SessionSource::Cli, diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index 48b6886382..5a8ffab36a 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -447,12 +447,20 @@ struct AppServerCommand { #[derive(Debug, Parser)] struct ExecServerCommand { /// Transport endpoint URL. Supported values: `ws://IP:PORT` (default), `stdio`, `stdio://`. - #[arg( - long = "listen", - value_name = "URL", - default_value = "ws://127.0.0.1:0" - )] - listen: String, + #[arg(long = "listen", value_name = "URL", conflicts_with = "remote")] + listen: Option, + + /// Register this exec-server as a remote executor using the given base URL. + #[arg(long = "remote", value_name = "URL", requires = "executor_id")] + remote: Option, + + /// Executor id to attach to when registering remotely. + #[arg(long = "executor-id", value_name = "ID")] + executor_id: Option, + + /// Human-readable executor name. + #[arg(long = "name", value_name = "NAME")] + name: Option, } #[derive(Debug, clap::Subcommand)] @@ -1264,7 +1272,23 @@ async fn run_exec_server_command( codex_self_exe, arg0_paths.codex_linux_sandbox_exe.clone(), )?; - codex_exec_server::run_main(&cmd.listen, runtime_paths) + if let Some(base_url) = cmd.remote { + let executor_id = cmd + .executor_id + .ok_or_else(|| anyhow::anyhow!("--executor-id is required when --remote is set"))?; + let mut remote_config = + codex_exec_server::RemoteExecutorConfig::new(base_url, executor_id)?; + if let Some(name) = cmd.name { + remote_config.name = name; + } + codex_exec_server::run_remote_executor(remote_config, runtime_paths).await?; + return Ok(()); + } + let listen_url = cmd + .listen + .as_deref() + .unwrap_or(codex_exec_server::DEFAULT_LISTEN_URL); + codex_exec_server::run_main(listen_url, runtime_paths) .await .map_err(anyhow::Error::from_boxed) } @@ -1387,7 +1411,7 @@ async fn run_debug_prompt_input_command( }); } - let prompt_input = codex_core::build_prompt_input(config, input, /*state_db*/ None).await?; + let prompt_input = codex_core::build_prompt_input(config, input).await?; println!("{}", serde_json::to_string_pretty(&prompt_input)?); Ok(()) diff --git a/codex-rs/code-mode/Cargo.toml b/codex-rs/code-mode/Cargo.toml index 23b2ce2306..19d6c3ab45 100644 --- a/codex-rs/code-mode/Cargo.toml +++ b/codex-rs/code-mode/Cargo.toml @@ -9,6 +9,9 @@ doctest = false name = "codex_code_mode" path = "src/lib.rs" +[features] +sandbox = ["v8/v8_enable_sandbox"] + [lints] workspace = true diff --git a/codex-rs/core-api/src/lib.rs b/codex-rs/core-api/src/lib.rs index e884cfae9b..13c87df9e0 100644 --- a/codex-rs/core-api/src/lib.rs +++ b/codex-rs/core-api/src/lib.rs @@ -31,6 +31,7 @@ pub use codex_core::StartThreadOptions; pub use codex_core::StateDbHandle; pub use codex_core::ThreadManager; pub use codex_core::ThreadShutdownReport; +pub use codex_core::agent_graph_store_from_state_db; pub use codex_core::config::Config; pub use codex_core::config::Constrained; pub use codex_core::config::GhostSnapshotConfig; @@ -40,6 +41,7 @@ pub use codex_core::config::TerminalResizeReflowConfig; pub use codex_core::config::ThreadStoreConfig; pub use codex_core::config::find_codex_home; pub use codex_core::init_state_db; +pub use codex_core::init_state_db_from_config; pub use codex_core::skills::SkillsManager; pub use codex_core::thread_store_from_config; pub use codex_exec_server::EnvironmentManager; diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index 44c6aacac5..49ed304623 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -26,6 +26,7 @@ bm25 = { workspace = true } chrono = { workspace = true, features = ["serde"] } clap = { workspace = true, features = ["derive"] } codex-analytics = { workspace = true } +codex-agent-graph-store = { workspace = true } codex-api = { workspace = true } codex-app-server-protocol = { workspace = true } codex-apply-patch = { workspace = true } diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 1a13c0b00c..3c98fd15d7 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -29,7 +29,6 @@ use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::user_input::UserInput; -use codex_state::DirectionalThreadSpawnEdgeStatus; use codex_thread_store::ReadThreadParams; use serde::Serialize; use std::collections::HashMap; @@ -296,7 +295,6 @@ impl AgentControl { state.notify_thread_created(new_thread.thread_id); self.persist_thread_spawn_edge_for_source( - new_thread.thread.as_ref(), new_thread.thread_id, notification_source.as_ref(), ) @@ -356,14 +354,10 @@ impl AgentControl { let parent_thread_id = *parent_thread_id; let parent_thread = state.get_thread(parent_thread_id).await.ok(); if let Some(parent_thread) = parent_thread.as_ref() { - // `record_conversation_items` only queues rollout writes asynchronously. - // Flush/materialize the live parent before snapshotting JSONL for a fork. - parent_thread - .codex - .session - .ensure_rollout_materialized() - .await; - parent_thread.codex.session.flush_rollout().await?; + // `record_conversation_items` only queues persistence writes asynchronously. + // Flush before snapshotting store history for a fork. + parent_thread.ensure_rollout_materialized().await; + parent_thread.flush_rollout().await?; } let parent_history = state @@ -449,19 +443,14 @@ impl AgentControl { )) .await?; let state = self.upgrade()?; - let Ok(resumed_thread) = state.get_thread(resumed_thread_id).await else { - return Ok(resumed_thread_id); - }; - let Some(state_db_ctx) = resumed_thread.state_db() else { - return Ok(resumed_thread_id); - }; + let agent_graph_store = state.agent_graph_store(); let mut resume_queue = VecDeque::from([(thread_id, root_depth)]); while let Some((parent_thread_id, parent_depth)) = resume_queue.pop_front() { - let child_ids = match state_db_ctx - .list_thread_spawn_children_with_status( + let child_ids = match agent_graph_store + .list_thread_spawn_children( parent_thread_id, - DirectionalThreadSpawnEdgeStatus::Open, + Some(codex_agent_graph_store::ThreadSpawnEdgeStatus::Open), ) .await { @@ -525,7 +514,6 @@ impl AgentControl { let _ = config.features.disable(Feature::Collab); } let state = self.upgrade()?; - let state_db_ctx = state.state_db(); let mut reservation = self.state.reserve_spawn_slot(config.agent_max_threads)?; let (session_source, agent_metadata) = match session_source { SessionSource::SubAgent(SubAgentSource::ThreadSpawn { @@ -535,14 +523,11 @@ impl AgentControl { agent_role: _, agent_nickname: _, }) => { + let state_db_ctx = state.state_db(); let (resumed_agent_nickname, resumed_agent_role) = - if let Some(state_db_ctx) = state_db_ctx.as_ref() { - match state_db_ctx.get_thread(thread_id).await { - Ok(Some(metadata)) => (metadata.agent_nickname, metadata.agent_role), - Ok(None) | Err(_) => (None, None), - } - } else { - (None, None) + match state_db_ctx.get_thread(thread_id).await { + Ok(Some(metadata)) => (metadata.agent_nickname, metadata.agent_role), + Ok(None) | Err(_) => (None, None), }; self.prepare_thread_spawn( &mut reservation, @@ -609,7 +594,6 @@ impl AgentControl { ); } self.persist_thread_spawn_edge_for_source( - resumed_thread.thread.as_ref(), resumed_thread.thread_id, Some(¬ification_source), ) @@ -722,11 +706,13 @@ impl AgentControl { /// agent and any live descendants reached from the in-memory tree. pub(crate) async fn close_agent(&self, agent_id: ThreadId) -> CodexResult { let state = self.upgrade()?; - if let Ok(thread) = state.get_thread(agent_id).await - && let Some(state_db_ctx) = thread.state_db() - && let Err(err) = state_db_ctx - .set_thread_spawn_edge_status(agent_id, DirectionalThreadSpawnEdgeStatus::Closed) - .await + if let Err(err) = state + .agent_graph_store() + .set_thread_spawn_edge_status( + agent_id, + codex_agent_graph_store::ThreadSpawnEdgeStatus::Closed, + ) + .await { warn!("failed to persist thread-spawn edge status for {agent_id}: {err}"); } @@ -1142,21 +1128,21 @@ impl AgentControl { async fn persist_thread_spawn_edge_for_source( &self, - thread: &crate::CodexThread, child_thread_id: ThreadId, session_source: Option<&SessionSource>, ) { let Some(parent_thread_id) = session_source.and_then(thread_spawn_parent_thread_id) else { return; }; - let Some(state_db_ctx) = thread.state_db() else { + let Ok(state) = self.upgrade() else { return; }; - if let Err(err) = state_db_ctx + if let Err(err) = state + .agent_graph_store() .upsert_thread_spawn_edge( parent_thread_id, child_thread_id, - DirectionalThreadSpawnEdgeStatus::Open, + codex_agent_graph_store::ThreadSpawnEdgeStatus::Open, ) .await { diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index b95aad4489..a8de427814 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -1,6 +1,5 @@ use super::*; use crate::CodexThread; -use crate::StateDbHandle; use crate::ThreadManager; use crate::agent::agent_status_from_event; use crate::config::AgentRoleConfig; @@ -8,7 +7,6 @@ use crate::config::Config; use crate::config::ConfigBuilder; use crate::context::ContextualUserFragment; use crate::context::SubagentNotification; -use crate::init_state_db; use assert_matches::assert_matches; use codex_features::Feature; use codex_login::CodexAuth; @@ -86,7 +84,6 @@ fn spawn_agent_call(call_id: &str) -> ResponseItem { struct AgentControlHarness { _home: TempDir, config: Config, - state_db: Option, manager: ThreadManager, control: AgentControl, } @@ -94,19 +91,17 @@ struct AgentControlHarness { impl AgentControlHarness { async fn new() -> Self { let (home, config) = test_config().await; - let state_db = init_state_db(&config).await; - let manager = ThreadManager::with_models_provider_home_and_state_for_tests( + let manager = ThreadManager::with_models_provider_and_home_for_tests( CodexAuth::from_api_key("dummy"), config.model_provider.clone(), config.codex_home.to_path_buf(), std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - state_db.clone(), - ); + ) + .await; let control = manager.agent_control(); Self { _home: home, config, - state_db, manager, control, } @@ -955,7 +950,8 @@ async fn spawn_agent_respects_max_threads_limit() { config.model_provider.clone(), config.codex_home.to_path_buf(), std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - ); + ) + .await; let control = manager.agent_control(); let _ = manager @@ -1007,7 +1003,8 @@ async fn spawn_agent_releases_slot_after_shutdown() { config.model_provider.clone(), config.codex_home.to_path_buf(), std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - ); + ) + .await; let control = manager.agent_control(); let first_agent_id = control @@ -1050,7 +1047,8 @@ async fn spawn_agent_limit_shared_across_clones() { config.model_provider.clone(), config.codex_home.to_path_buf(), std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - ); + ) + .await; let control = manager.agent_control(); let cloned = control.clone(); @@ -1095,7 +1093,8 @@ async fn resume_agent_respects_max_threads_limit() { config.model_provider.clone(), config.codex_home.to_path_buf(), std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - ); + ) + .await; let control = manager.agent_control(); let resumable_id = control @@ -1151,7 +1150,8 @@ async fn resume_agent_releases_slot_after_resume_failure() { config.model_provider.clone(), config.codex_home.to_path_buf(), std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - ); + ) + .await; let control = manager.agent_control(); let _ = control @@ -1543,19 +1543,17 @@ async fn resume_thread_subagent_restores_stored_nickname_and_role() { .features .enable(Feature::Sqlite) .expect("test config should allow sqlite"); - let state_db = init_state_db(&config).await; - let manager = ThreadManager::with_models_provider_home_and_state_for_tests( + let manager = ThreadManager::with_models_provider_and_home_for_tests( CodexAuth::from_api_key("dummy"), config.model_provider.clone(), config.codex_home.to_path_buf(), std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - state_db.clone(), - ); + ) + .await; let control = manager.agent_control(); let harness = AgentControlHarness { _home: home, config, - state_db, manager, control, }; @@ -1706,7 +1704,12 @@ async fn resume_agent_from_rollout_reads_archived_rollout_path() { .expect("child shutdown should succeed"); let store = LocalThreadStore::new( LocalThreadStoreConfig::from_config(&harness.config), - harness.state_db.clone(), + codex_state::StateRuntime::init( + harness.config.sqlite_home.clone(), + harness.config.model_provider_id.clone(), + ) + .await + .expect("state db should initialize"), ); store .archive_thread(ArchiveThreadParams { diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 1aeb125091..9578fa5726 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -95,6 +95,7 @@ pub(crate) async fn run_codex_thread_interactive( parent_trace: None, environment_selections: parent_ctx.environments.clone(), analytics_events_client: Some(parent_session.services.analytics_events_client.clone()), + state_db: parent_session.services.state_db.clone(), thread_store: Arc::clone(&parent_session.services.thread_store), })) .or_cancel(&cancel_token) diff --git a/codex-rs/core/src/goals.rs b/codex-rs/core/src/goals.rs index 3cea25f9b8..4af73e6a01 100644 --- a/codex-rs/core/src/goals.rs +++ b/codex-rs/core/src/goals.rs @@ -29,7 +29,6 @@ use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::validate_thread_goal_objective; use codex_rollout::state_db::reconcile_rollout; -use codex_thread_store::LocalThreadStore; use codex_utils_template::Template; use futures::future::BoxFuture; use std::sync::Arc; @@ -1338,17 +1337,6 @@ impl Session { state_db } else if let Some(state_db) = self.goal_runtime.state_db.lock().await.clone() { state_db - } else if let Some(local_store) = self - .services - .thread_store - .as_any() - .downcast_ref::() - { - local_store.state_db().await.ok_or_else(|| { - anyhow::anyhow!( - "thread goals require a local persisted thread with a state database" - ) - })? } else { anyhow::bail!("thread goals require a local persisted thread with a state database"); }; diff --git a/codex-rs/core/src/guardian/review_session.rs b/codex-rs/core/src/guardian/review_session.rs index 6fd50219d8..a419d7cbfa 100644 --- a/codex-rs/core/src/guardian/review_session.rs +++ b/codex-rs/core/src/guardian/review_session.rs @@ -35,7 +35,6 @@ use crate::config::NetworkProxySpec; use crate::config::Permissions; use crate::context::ContextualUserFragment; use crate::context::GuardianFollowupReviewReminder; -use crate::rollout::recorder::RolloutRecorder; use crate::session::Codex; use crate::session::session::Session; use crate::session::turn_context::TurnContext; @@ -774,12 +773,11 @@ async fn append_guardian_followup_reminder(review_session: &GuardianReviewSessio async fn load_rollout_items_for_fork( session: &Session, ) -> anyhow::Result>> { + session.try_ensure_rollout_materialized().await?; session.flush_rollout().await?; - let Some(rollout_path) = session.current_rollout_path().await? else { - return Ok(None); - }; - let history = RolloutRecorder::get_rollout_history(rollout_path.as_path()).await?; - Ok(Some(history.get_rollout_items())) + let live_thread = session.live_thread_for_persistence("guardian review fork")?; + let history = live_thread.load_history(/*include_archived*/ true).await?; + Ok(Some(history.items)) } async fn wait_for_guardian_review( diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index d8a3fbcf0f..cbada9a26c 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -124,7 +124,9 @@ pub use thread_manager::NewThread; pub use thread_manager::StartThreadOptions; pub use thread_manager::ThreadManager; pub use thread_manager::ThreadShutdownReport; +pub use thread_manager::agent_graph_store_from_state_db; pub use thread_manager::build_models_manager; +pub use thread_manager::init_state_db_from_config; pub use thread_manager::thread_store_from_config; pub use web_search::web_search_action_detail; pub use web_search::web_search_detail; diff --git a/codex-rs/core/src/personality_migration.rs b/codex-rs/core/src/personality_migration.rs index 975aecd4af..bb1cb43587 100644 --- a/codex-rs/core/src/personality_migration.rs +++ b/codex-rs/core/src/personality_migration.rs @@ -25,7 +25,7 @@ pub enum PersonalityMigrationStatus { pub async fn maybe_migrate_personality( codex_home: &Path, config_toml: &ConfigToml, - state_db: Option, + state_db: StateDbHandle, ) -> io::Result { let marker_path = codex_home.join(PERSONALITY_MIGRATION_FILENAME); if tokio::fs::try_exists(&marker_path).await? { @@ -65,16 +65,13 @@ pub async fn maybe_migrate_personality( async fn has_recorded_sessions( codex_home: &Path, default_provider: &str, - state_db: Option, + state_db: StateDbHandle, ) -> io::Result { - let store = LocalThreadStore::new( - LocalThreadStoreConfig { - codex_home: codex_home.to_path_buf(), - sqlite_home: codex_home.to_path_buf(), - default_model_provider_id: default_provider.to_string(), - }, - state_db, - ); + let config = LocalThreadStoreConfig { + codex_home: codex_home.to_path_buf(), + default_model_provider_id: default_provider.to_string(), + }; + let store = LocalThreadStore::new(config, state_db); if has_threads(&store, /*archived*/ false).await? { return Ok(true); } diff --git a/codex-rs/core/src/personality_migration_tests.rs b/codex-rs/core/src/personality_migration_tests.rs index 28d34bc3c9..63faae418a 100644 --- a/codex-rs/core/src/personality_migration_tests.rs +++ b/codex-rs/core/src/personality_migration_tests.rs @@ -8,7 +8,10 @@ use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::UserMessageEvent; use codex_rollout::ARCHIVED_SESSIONS_SUBDIR; +use codex_rollout::RolloutConfig; use codex_rollout::SESSIONS_SUBDIR; +use codex_rollout::state_db::StateDbHandle; +use codex_state::state_db_path; use pretty_assertions::assert_eq; use tempfile::TempDir; use tokio::io::AsyncWriteExt; @@ -20,6 +23,26 @@ async fn read_config_toml(codex_home: &Path) -> io::Result { toml::from_str(&contents).map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err)) } +async fn state_db_for_test(codex_home: &Path) -> io::Result { + state_db_for_test_with_sqlite_home(codex_home, codex_home).await +} + +async fn state_db_for_test_with_sqlite_home( + codex_home: &Path, + sqlite_home: &Path, +) -> io::Result { + let config = RolloutConfig { + codex_home: codex_home.to_path_buf(), + sqlite_home: sqlite_home.to_path_buf(), + cwd: codex_home.to_path_buf(), + model_provider_id: "openai".to_string(), + generate_memories: false, + }; + codex_rollout::state_db::try_init(&config) + .await + .map_err(io::Error::other) +} + async fn write_session_with_user_event(codex_home: &Path) -> io::Result<()> { let thread_id = ThreadId::new(); let dir = codex_home @@ -87,7 +110,8 @@ async fn applies_when_sessions_exist_and_no_personality() -> io::Result<()> { write_session_with_user_event(temp.path()).await?; let config_toml = ConfigToml::default(); - let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?; + let state_db = state_db_for_test(temp.path()).await?; + let status = maybe_migrate_personality(temp.path(), &config_toml, state_db).await?; assert_eq!(status, PersonalityMigrationStatus::Applied); assert!(temp.path().join(PERSONALITY_MIGRATION_FILENAME).exists()); @@ -103,7 +127,8 @@ async fn applies_when_only_archived_sessions_exist_and_no_personality() -> io::R write_archived_session_with_user_event(temp.path()).await?; let config_toml = ConfigToml::default(); - let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?; + let state_db = state_db_for_test(temp.path()).await?; + let status = maybe_migrate_personality(temp.path(), &config_toml, state_db).await?; assert_eq!(status, PersonalityMigrationStatus::Applied); assert!(temp.path().join(PERSONALITY_MIGRATION_FILENAME).exists()); @@ -119,7 +144,8 @@ async fn skips_when_marker_exists() -> io::Result<()> { create_marker(&temp.path().join(PERSONALITY_MIGRATION_FILENAME)).await?; let config_toml = ConfigToml::default(); - let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?; + let state_db = state_db_for_test(temp.path()).await?; + let status = maybe_migrate_personality(temp.path(), &config_toml, state_db).await?; assert_eq!(status, PersonalityMigrationStatus::SkippedMarker); assert!(!temp.path().join("config.toml").exists()); @@ -136,7 +162,8 @@ async fn skips_when_personality_explicit() -> io::Result<()> { .map_err(|err| io::Error::other(format!("failed to write config: {err}")))?; let config_toml = read_config_toml(temp.path()).await?; - let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?; + let state_db = state_db_for_test(temp.path()).await?; + let status = maybe_migrate_personality(temp.path(), &config_toml, state_db).await?; assert_eq!( status, @@ -153,10 +180,37 @@ async fn skips_when_personality_explicit() -> io::Result<()> { async fn skips_when_no_sessions() -> io::Result<()> { let temp = TempDir::new()?; let config_toml = ConfigToml::default(); - let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?; + let state_db = state_db_for_test(temp.path()).await?; + let status = maybe_migrate_personality(temp.path(), &config_toml, state_db).await?; assert_eq!(status, PersonalityMigrationStatus::SkippedNoSessions); assert!(temp.path().join(PERSONALITY_MIGRATION_FILENAME).exists()); assert!(!temp.path().join("config.toml").exists()); Ok(()) } + +#[tokio::test] +async fn uses_configured_sqlite_home_when_checking_for_sessions() -> io::Result<()> { + let codex_home = TempDir::new()?; + let sqlite_home = TempDir::new()?; + write_session_with_user_event(codex_home.path()).await?; + + let config_toml = ConfigToml::default(); + let state_db = + state_db_for_test_with_sqlite_home(codex_home.path(), sqlite_home.path()).await?; + let status = maybe_migrate_personality(codex_home.path(), &config_toml, state_db).await?; + + assert_eq!(status, PersonalityMigrationStatus::Applied); + assert!( + codex_home + .path() + .join(PERSONALITY_MIGRATION_FILENAME) + .exists() + ); + + let persisted = read_config_toml(codex_home.path()).await?; + assert_eq!(persisted.personality, Some(Personality::Pragmatic)); + assert!(!state_db_path(codex_home.path()).exists()); + assert!(state_db_path(sqlite_home.path()).exists()); + Ok(()) +} diff --git a/codex-rs/core/src/prompt_debug.rs b/codex-rs/core/src/prompt_debug.rs index 7c6144c10e..a40fd27942 100644 --- a/codex-rs/core/src/prompt_debug.rs +++ b/codex-rs/core/src/prompt_debug.rs @@ -16,8 +16,9 @@ use crate::config::Config; use crate::session::session::Session; use crate::session::turn::build_prompt; use crate::session::turn::built_tools; -use crate::state_db_bridge::StateDbHandle; use crate::thread_manager::ThreadManager; +use crate::thread_manager::agent_graph_store_from_state_db; +use crate::thread_manager::init_state_db_from_config; use crate::thread_manager::thread_store_from_config; /// Build the model-visible `input` list for a single debug turn. @@ -25,7 +26,6 @@ use crate::thread_manager::thread_store_from_config; pub async fn build_prompt_input( mut config: Config, input: Vec, - state_db: Option, ) -> CodexResult> { config.ephemeral = true; @@ -37,15 +37,20 @@ pub async fn build_prompt_input( config.codex_linux_sandbox_exe.clone(), )?; + let state_db = init_state_db_from_config(&config) + .await + .ok_or_else(|| std::io::Error::other("prompt debug requires state db"))?; let thread_store = thread_store_from_config(&config, state_db.clone()); + let agent_graph_store = agent_graph_store_from_state_db(state_db.clone()); let thread_manager = ThreadManager::new( &config, Arc::clone(&auth_manager), SessionSource::Exec, Arc::new(EnvironmentManager::new(EnvironmentManagerArgs::new(local_runtime_paths)).await), /*analytics_events_client*/ None, + state_db, thread_store, - state_db.clone(), + agent_graph_store, ); let thread = thread_manager.start_thread(config).await?; diff --git a/codex-rs/core/src/rollout.rs b/codex-rs/core/src/rollout.rs index d4ac5c699a..e0d268dc2d 100644 --- a/codex-rs/core/src/rollout.rs +++ b/codex-rs/core/src/rollout.rs @@ -50,6 +50,7 @@ pub(crate) mod list { pub use codex_rollout::find_thread_path_by_id_str; } +#[cfg(test)] pub(crate) mod recorder { pub use codex_rollout::RolloutRecorder; } diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index d35bccb96e..85f2ef438b 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -132,7 +132,6 @@ use codex_terminal_detection::user_agent; use codex_thread_store::CreateThreadParams; use codex_thread_store::LiveThread; use codex_thread_store::LiveThreadInitGuard; -use codex_thread_store::LocalThreadStore; use codex_thread_store::ResumeThreadParams; use codex_thread_store::ThreadEventPersistenceMode; use codex_thread_store::ThreadPersistenceMetadata; @@ -356,7 +355,6 @@ use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TokenUsageInfo; use codex_protocol::protocol::WarningEvent; use codex_protocol::user_input::UserInput; -use codex_tools::ToolEnvironmentMode; use codex_tools::ToolsConfig; use codex_tools::ToolsConfigParams; use codex_utils_absolute_path::AbsolutePathBuf; @@ -412,6 +410,7 @@ pub(crate) struct CodexSpawnArgs { pub(crate) parent_trace: Option, pub(crate) environment_selections: ResolvedTurnEnvironments, pub(crate) analytics_events_client: Option, + pub(crate) state_db: Option, pub(crate) thread_store: Arc, } @@ -469,6 +468,7 @@ impl Codex { parent_trace: _, environment_selections, analytics_events_client, + state_db, thread_store, } = args; let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY); @@ -557,15 +557,7 @@ impl Codex { }; match thread_id { Some(thread_id) => { - let state_db_ctx = if config.ephemeral { - None - } else if let Some(local_store) = - thread_store.as_any().downcast_ref::() - { - local_store.state_db().await - } else { - None - }; + let state_db_ctx = state_db.clone(); state_db::get_dynamic_tools(state_db_ctx.as_deref(), thread_id, "codex_spawn") .await } @@ -651,6 +643,7 @@ impl Codex { agent_control, environment_manager, analytics_events_client, + state_db, thread_store, parent_rollout_thread_trace, ) @@ -1311,7 +1304,7 @@ impl Session { self.services.user_shell.as_ref().clone(), self.services.shell_snapshot_tx.clone(), self.services.session_telemetry.clone(), - self.services.state_db.clone(), + self.state_db(), ); } diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 0bb8da6ea6..23621f87cf 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -343,6 +343,7 @@ impl Session { agent_control: AgentControl, environment_manager: Arc, analytics_events_client: Option, + state_db: Option, thread_store: Arc, parent_rollout_thread_trace: ThreadTraceContext, ) -> anyhow::Result> { @@ -441,22 +442,7 @@ impl Session { otel.name = "session_init.thread_persistence", session_init.ephemeral = config.ephemeral, )); - let state_db_fut = async { - if config.ephemeral { - None - } else if let Some(local_store) = - thread_store.as_any().downcast_ref::() - { - local_store.state_db().await - } else { - None - } - } - .instrument(info_span!( - "session_init.state_db", - otel.name = "session_init.state_db", - session_init.ephemeral = config.ephemeral, - )); + let state_db_ctx = if config.ephemeral { None } else { state_db }; let is_subagent = session_configuration.session_source.is_non_root_agent(); let history_meta_fut = async { @@ -495,15 +481,9 @@ impl Session { // Join all independent futures. let ( thread_persistence_result, - state_db_ctx, (history_log_id, history_entry_count), (auth, mcp_servers, auth_statuses), - ) = tokio::join!( - thread_persistence_fut, - state_db_fut, - history_meta_fut, - auth_and_mcp_fut - ); + ) = tokio::join!(thread_persistence_fut, history_meta_fut, auth_and_mcp_fut); let mut live_thread_init = LiveThreadInitGuard::new(thread_persistence_result.map_err(|e| { diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 6995ef8dc9..f3a8d95107 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -867,7 +867,7 @@ async fn danger_full_access_turns_do_not_expose_managed_network_proxy() -> anyho &permission_profile_for_sandbox_policy(&SandboxPolicy::DangerFullAccess), )?; - let session = make_session_with_config(move |config| { + let (session, _codex_home) = make_session_with_config(move |config| { let cwd = config.cwd.clone(); config .permissions @@ -933,7 +933,7 @@ async fn danger_full_access_tool_attempts_do_not_enforce_managed_network() -> an &permission_profile_for_sandbox_policy(&SandboxPolicy::DangerFullAccess), )?; - let session = make_session_with_config(move |config| { + let (session, _codex_home) = make_session_with_config(move |config| { let cwd = config.cwd.clone(); config .permissions @@ -1008,7 +1008,7 @@ async fn workspace_write_turns_continue_to_expose_managed_network_proxy() -> any &permission_profile_for_sandbox_policy(&sandbox_policy), )?; - let session = make_session_with_config(move |config| { + let (session, _codex_home) = make_session_with_config(move |config| { let cwd = config.cwd.clone(); config .permissions @@ -1035,7 +1035,7 @@ async fn user_shell_commands_do_not_inherit_managed_network_proxy() -> anyhow::R &permission_profile_for_sandbox_policy(&sandbox_policy), )?; - let (session, rx) = make_session_with_config_and_rx(move |config| { + let (session, rx, _codex_home) = make_session_with_config_and_rx(move |config| { let cwd = config.cwd.clone(); config .permissions @@ -1165,7 +1165,7 @@ async fn reload_user_config_layer_updates_effective_apps_config() { #[tokio::test] async fn reload_user_config_layer_refreshes_hooks() -> anyhow::Result<()> { - let session = make_session_with_config(|config| { + let (session, _codex_home) = make_session_with_config(|config| { config .features .enable(Feature::CodexHooks) @@ -2982,15 +2982,13 @@ pub(crate) async fn make_session_configuration_for_tests() -> SessionConfigurati fn turn_environments_for_tests( environment: &Arc, cwd: &codex_utils_absolute_path::AbsolutePathBuf, -) -> crate::environment_selection::ResolvedTurnEnvironments { - crate::environment_selection::ResolvedTurnEnvironments { - turn_environments: vec![TurnEnvironment { - environment_id: codex_exec_server::LOCAL_ENVIRONMENT_ID.to_string(), - environment: Arc::clone(environment), - cwd: cwd.clone(), - shell: "bash".to_string(), - }], - } +) -> Vec { + vec![TurnEnvironment { + environment_id: codex_exec_server::LOCAL_ENVIRONMENT_ID.to_string(), + environment: Arc::clone(environment), + cwd: cwd.clone(), + shell: "bash".to_string(), + }] } #[tokio::test] @@ -3528,9 +3526,15 @@ async fn session_new_fails_when_zsh_fork_enabled_without_zsh_path() { AgentControl::default(), Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, + /*state_db*/ None, Arc::new(codex_thread_store::LocalThreadStore::new( codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()), - /*state_db*/ None, + codex_state::StateRuntime::init( + config.sqlite_home.clone(), + config.model_provider_id.clone(), + ) + .await + .expect("state db should initialize"), )), codex_rollout_trace::ThreadTraceContext::disabled(), ) @@ -3678,7 +3682,12 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { live_thread: None, thread_store: Arc::new(codex_thread_store::LocalThreadStore::new( codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()), - /*state_db*/ None, + codex_state::StateRuntime::init( + config.sqlite_home.clone(), + config.model_provider_id.clone(), + ) + .await + .expect("state db should initialize"), )), model_client: ModelClient::new( Some(auth_manager.clone()), @@ -3723,7 +3732,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { model_info, &models_manager, /*network*/ None, - turn_environments, + crate::environment_selection::ResolvedTurnEnvironments { turn_environments }, session_configuration.cwd.clone(), "turn_id".to_string(), skills_outcome, @@ -3756,14 +3765,18 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { async fn make_session_with_config( mutator: impl FnOnce(&mut Config), -) -> anyhow::Result> { - let (session, _rx_event) = make_session_with_config_and_rx(mutator).await?; - Ok(session) +) -> anyhow::Result<(Arc, tempfile::TempDir)> { + let (session, _rx_event, codex_home) = make_session_with_config_and_rx(mutator).await?; + Ok((session, codex_home)) } async fn make_session_with_config_and_rx( mutator: impl FnOnce(&mut Config), -) -> anyhow::Result<(Arc, async_channel::Receiver)> { +) -> anyhow::Result<( + Arc, + async_channel::Receiver, + tempfile::TempDir, +)> { let codex_home = tempfile::tempdir().expect("create temp dir"); let mut config = build_test_config(codex_home.path()).await; mutator(&mut config); @@ -3848,15 +3861,21 @@ async fn make_session_with_config_and_rx( AgentControl::default(), Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, + /*state_db*/ None, Arc::new(codex_thread_store::LocalThreadStore::new( codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()), - /*state_db*/ None, + codex_state::StateRuntime::init( + config.sqlite_home.clone(), + config.model_provider_id.clone(), + ) + .await + .expect("state db should initialize"), )), codex_rollout_trace::ThreadTraceContext::disabled(), ) .await?; - Ok((session, rx_event)) + Ok((session, rx_event, codex_home)) } #[tokio::test] @@ -4518,7 +4537,7 @@ async fn default_turn_honors_empty_stored_thread_environments() { let turn_context = session.new_default_turn().await; - assert!(turn_context.environments.primary().is_none()); + assert!(turn_context.environments.primary_environment().is_none()); assert!(turn_context.environments.turn_environments.is_empty()); assert_eq!(turn_context.cwd, session_cwd); assert_eq!(turn_context.config.cwd, session_cwd); @@ -4537,7 +4556,7 @@ async fn primary_environment_uses_first_turn_environment() { environment_id: "second".to_string(), environment: Arc::clone(&first_environment.environment), cwd: second_cwd.clone(), - shell: first_environment.shell.clone(), + shell: "bash".to_string(), }); assert_eq!( @@ -4580,7 +4599,7 @@ async fn empty_turn_environments_clear_primary_environment() { .await .expect("turn should start"); - assert!(turn_context.environments.primary().is_none()); + assert!(turn_context.environments.primary_environment().is_none()); assert!(turn_context.environments.turn_environments.is_empty()); assert_eq!(turn_context.cwd, session.get_config().await.cwd); assert_eq!(turn_context.config.cwd, session.get_config().await.cwd); @@ -5036,47 +5055,13 @@ async fn make_session_and_context_with_auth_and_config_and_rx( Arc, async_channel::Receiver, ) -where - F: FnOnce(&mut Config), -{ - let codex_home = tempfile::tempdir().expect("create temp dir"); - make_session_and_context_with_auth_config_home_and_rx( - auth, - dynamic_tools, - codex_home.path(), - configure_config, - ) - .await -} - -async fn make_session_and_context_with_auth_config_home_and_rx( - auth: CodexAuth, - dynamic_tools: Vec, - codex_home: &Path, - configure_config: F, -) -> ( - Arc, - Arc, - async_channel::Receiver, -) where F: FnOnce(&mut Config), { let (tx_event, rx_event) = async_channel::unbounded(); - let mut config = build_test_config(codex_home).await; + let codex_home = tempfile::tempdir().expect("create temp dir").keep(); + let mut config = build_test_config(codex_home.as_path()).await; configure_config(&mut config); - let state_db = if config.features.enabled(Feature::Goals) { - Some( - codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.model_provider_id.clone(), - ) - .await - .expect("goal tests should initialize sqlite state db"), - ) - } else { - None - }; let config = Arc::new(config); let conversation_id = ThreadId::default(); let auth_manager = AuthManager::from_auth_for_testing(auth); @@ -5162,6 +5147,12 @@ where .expect("create environment"), ); + let state_db = codex_state::StateRuntime::init( + config.sqlite_home.clone(), + config.model_provider_id.clone(), + ) + .await + .expect("state db should initialize"); let skills_watcher = Arc::new(SkillsWatcher::noop()); let services = SessionServices { mcp_connection_manager: Arc::new(RwLock::new(McpConnectionManager::new_uninitialized( @@ -5202,7 +5193,7 @@ where agent_control, network_proxy: None, network_approval: Arc::clone(&network_approval), - state_db: state_db.clone(), + state_db: Some(state_db.clone()), live_thread: None, thread_store: Arc::new(codex_thread_store::LocalThreadStore::new( codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()), @@ -5251,7 +5242,7 @@ where model_info, &models_manager, /*network*/ None, - turn_environments, + crate::environment_selection::ResolvedTurnEnvironments { turn_environments }, session_configuration.cwd.clone(), "turn_id".to_string(), skills_outcome, @@ -5301,13 +5292,10 @@ async fn make_goal_session_and_context_with_rx() -> ( Arc, Arc, async_channel::Receiver, - tempfile::TempDir, ) { - let codex_home = tempfile::tempdir().expect("create temp dir"); - let (session, turn_context, rx) = make_session_and_context_with_auth_config_home_and_rx( + let (session, turn_context, rx) = make_session_and_context_with_auth_and_config_and_rx( CodexAuth::from_api_key("Test API Key"), Vec::new(), - codex_home.path(), |config| { config .features @@ -5317,14 +5305,14 @@ async fn make_goal_session_and_context_with_rx() -> ( ) .await; upsert_goal_test_thread(session.as_ref()).await; - (session, turn_context, rx, codex_home) + (session, turn_context, rx) } async fn upsert_goal_test_thread(session: &Session) { let config = session.get_config().await; - let state_db = session - .state_db() - .expect("goal test session should have a state db"); + let state_db = goal_test_state_db(session) + .await + .expect("goal test state db should initialize"); let mut builder = codex_state::ThreadMetadataBuilder::new( session.conversation_id, config @@ -7148,7 +7136,7 @@ async fn abort_empty_active_turn_preserves_pending_input() { #[tokio::test] async fn interrupt_accounts_active_goal_before_pausing() -> anyhow::Result<()> { - let (sess, tc, _rx, _codex_home) = make_goal_session_and_context_with_rx().await; + let (sess, tc, _rx) = make_goal_session_and_context_with_rx().await; sess.set_thread_goal( tc.as_ref(), SetGoalRequest { @@ -7412,7 +7400,7 @@ async fn goal_test_state_db(sess: &Session) -> anyhow::Result anyhow::Result<()> { - let (sess, tc, rx, _codex_home) = make_goal_session_and_context_with_rx().await; + let (sess, tc, rx) = make_goal_session_and_context_with_rx().await; sess.set_thread_goal( tc.as_ref(), SetGoalRequest { @@ -7512,7 +7500,7 @@ async fn budget_limited_accounting_steers_active_turn_without_aborting() -> anyh #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn external_goal_mutation_accounts_active_turn_before_status_change() -> anyhow::Result<()> { - let (sess, tc, _rx, _codex_home) = make_goal_session_and_context_with_rx().await; + let (sess, tc, _rx) = make_goal_session_and_context_with_rx().await; sess.set_thread_goal( tc.as_ref(), SetGoalRequest { @@ -7579,7 +7567,7 @@ async fn external_goal_mutation_accounts_active_turn_before_status_change() -> a #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn external_active_goal_set_marks_current_turn_for_accounting() -> anyhow::Result<()> { - let (sess, tc, _rx, _codex_home) = make_goal_session_and_context_with_rx().await; + let (sess, tc, _rx) = make_goal_session_and_context_with_rx().await; sess.spawn_task( Arc::clone(&tc), Vec::new(), @@ -8246,7 +8234,7 @@ async fn sample_rollout( #[tokio::test] async fn create_goal_tool_rejects_existing_goal() { - let (session, turn_context, _rx, _codex_home) = make_goal_session_and_context_with_rx().await; + let (session, turn_context, _rx) = make_goal_session_and_context_with_rx().await; let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); let handler = CreateGoalHandler; @@ -8308,7 +8296,7 @@ async fn create_goal_tool_rejects_existing_goal() { #[tokio::test] async fn update_goal_tool_rejects_pausing_goal() { - let (session, turn_context, _rx, _codex_home) = make_goal_session_and_context_with_rx().await; + let (session, turn_context, _rx) = make_goal_session_and_context_with_rx().await; let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); let create_handler = CreateGoalHandler; let update_handler = UpdateGoalHandler; @@ -8369,7 +8357,7 @@ async fn update_goal_tool_rejects_pausing_goal() { #[tokio::test] async fn update_goal_tool_marks_goal_complete() { - let (session, turn_context, _rx, _codex_home) = make_goal_session_and_context_with_rx().await; + let (session, turn_context, _rx) = make_goal_session_and_context_with_rx().await; let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); let create_handler = CreateGoalHandler; let update_handler = UpdateGoalHandler; diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index 4c61dad18d..d7ec5937d2 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -731,7 +731,12 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { let skills_watcher = Arc::new(SkillsWatcher::noop()); let thread_store = Arc::new(codex_thread_store::LocalThreadStore::new( codex_thread_store::LocalThreadStoreConfig::from_config(&config), - /*state_db*/ None, + codex_state::StateRuntime::init( + config.sqlite_home.clone(), + config.model_provider_id.clone(), + ) + .await + .expect("state db should initialize"), )); let CodexSpawnOk { codex, .. } = Codex::spawn(CodexSpawnArgs { @@ -760,6 +765,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { turn_environments: Vec::new(), }, analytics_events_client: None, + state_db: None, thread_store, }) .await diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index fe54d61546..b4ef59df95 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -8,6 +8,7 @@ use codex_protocol::protocol::TurnEnvironmentSelection; use codex_sandboxing::compatibility_sandbox_policy_for_permission_profile; use codex_sandboxing::policy_transforms::effective_file_system_sandbox_policy; use codex_sandboxing::policy_transforms::effective_network_sandbox_policy; +use codex_tools::ToolEnvironmentMode; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; diff --git a/codex-rs/core/src/stream_events_utils.rs b/codex-rs/core/src/stream_events_utils.rs index 29884da4ae..fc608aed93 100644 --- a/codex-rs/core/src/stream_events_utils.rs +++ b/codex-rs/core/src/stream_events_utils.rs @@ -138,11 +138,8 @@ pub(crate) async fn record_completed_response_item( .await; } mark_thread_memory_mode_polluted_if_external_context(sess, turn_context, item).await; - let has_memory_citation = record_stage1_output_usage_and_detect_memory_citation( - sess.services.state_db.as_ref(), - item, - ) - .await; + let has_memory_citation = + record_stage1_output_usage_and_detect_memory_citation(sess.state_db(), item).await; if has_memory_citation { sess.record_memory_citation_for_turn(&turn_context.sub_id) .await; @@ -177,7 +174,7 @@ pub(crate) async fn mark_thread_memory_mode_polluted_if_external_context( } async fn record_stage1_output_usage_and_detect_memory_citation( - state_db_ctx: Option<&state_db::StateDbHandle>, + state_db_ctx: Option, item: &ResponseItem, ) -> bool { let Some(raw_text) = raw_assistant_output_text_from_item(item) else { diff --git a/codex-rs/core/src/test_support.rs b/codex-rs/core/src/test_support.rs index 48eec66c58..74305fb518 100644 --- a/codex-rs/core/src/test_support.rs +++ b/codex-rs/core/src/test_support.rs @@ -52,14 +52,14 @@ pub fn auth_manager_from_auth_with_home(auth: CodexAuth, codex_home: PathBuf) -> AuthManager::from_auth_for_testing_with_home(auth, codex_home) } -pub fn thread_manager_with_models_provider( +pub async fn thread_manager_with_models_provider( auth: CodexAuth, provider: ModelProviderInfo, ) -> ThreadManager { - ThreadManager::with_models_provider_for_tests(auth, provider) + ThreadManager::with_models_provider_for_tests(auth, provider).await } -pub fn thread_manager_with_models_provider_and_home( +pub async fn thread_manager_with_models_provider_and_home( auth: CodexAuth, provider: ModelProviderInfo, codex_home: PathBuf, @@ -71,22 +71,7 @@ pub fn thread_manager_with_models_provider_and_home( codex_home, environment_manager, ) -} - -pub fn thread_manager_with_models_provider_home_and_state( - auth: CodexAuth, - provider: ModelProviderInfo, - codex_home: PathBuf, - environment_manager: Arc, - state_db: Option, -) -> ThreadManager { - ThreadManager::with_models_provider_home_and_state_for_tests( - auth, - provider, - codex_home, - environment_manager, - state_db, - ) + .await } pub async fn start_thread_with_user_shell_override( diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index d4e36ab3c8..06f8118db7 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -18,6 +18,8 @@ use crate::skills_watcher::SkillsWatcher; use crate::skills_watcher::SkillsWatcherEvent; use crate::tasks::InterruptedTurnHistoryMarker; use crate::tasks::interrupted_turn_history_marker; +use codex_agent_graph_store::AgentGraphStore; +use codex_agent_graph_store::LocalAgentGraphStore; use codex_analytics::AnalyticsEventsClient; use codex_app_server_protocol::ThreadHistoryBuilder; use codex_app_server_protocol::TurnStatus; @@ -50,8 +52,8 @@ use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnAbortedEvent; use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::protocol::W3cTraceContext; +use codex_rollout::state_db; use codex_rollout::state_db::StateDbHandle; -use codex_state::DirectionalThreadSpawnEdgeStatus; use codex_thread_store::InMemoryThreadStore; use codex_thread_store::LocalThreadStore; use codex_thread_store::LocalThreadStoreConfig; @@ -247,9 +249,10 @@ pub(crate) struct ThreadManagerState { mcp_manager: Arc, skills_watcher: Arc, thread_store: Arc, + state_db: StateDbHandle, + agent_graph_store: Arc, session_source: SessionSource, analytics_events_client: Option, - state_db: Option, // Captures submitted ops for testing purpose when test mode is enabled. ops_log: Option, } @@ -265,10 +268,11 @@ pub fn build_models_manager( ) } -pub fn thread_store_from_config( - config: &Config, - state_db: Option, -) -> Arc { +pub async fn init_state_db_from_config(config: &Config) -> Option { + state_db::init(config).await +} + +pub fn thread_store_from_config(config: &Config, state_db: StateDbHandle) -> Arc { match &config.experimental_thread_store { ThreadStoreConfig::Local => Arc::new(LocalThreadStore::new( LocalThreadStoreConfig::from_config(config), @@ -279,15 +283,38 @@ pub fn thread_store_from_config( } } +pub fn agent_graph_store_from_state_db(state_db: StateDbHandle) -> Arc { + Arc::new(LocalAgentGraphStore::new(state_db)) +} + +async fn state_db_from_roots_for_tests( + codex_home: PathBuf, + sqlite_home: PathBuf, + default_model_provider_id: String, +) -> StateDbHandle { + let config = codex_rollout::RolloutConfig { + codex_home: codex_home.clone(), + sqlite_home, + cwd: codex_home, + model_provider_id: default_model_provider_id, + generate_memories: false, + }; + state_db::try_init(&config) + .await + .unwrap_or_else(|err| panic!("test state db should initialize: {err}")) +} + impl ThreadManager { + #[allow(clippy::too_many_arguments)] pub fn new( config: &Config, auth_manager: Arc, session_source: SessionSource, environment_manager: Arc, analytics_events_client: Option, + state_db: StateDbHandle, thread_store: Arc, - state_db: Option, + agent_graph_store: Arc, ) -> Self { let codex_home = config.codex_home.clone(); let restriction_product = session_source.restriction_product(); @@ -314,10 +341,11 @@ impl ThreadManager { mcp_manager, skills_watcher, thread_store, + state_db, + agent_graph_store, auth_manager, session_source, analytics_events_client, - state_db, ops_log: should_use_test_thread_manager_behavior() .then(|| Arc::new(std::sync::Mutex::new(Vec::new()))), }), @@ -327,7 +355,7 @@ impl ThreadManager { /// Construct with a dummy AuthManager containing the provided CodexAuth. /// Used for integration tests: should not be used by ordinary business logic. - pub(crate) fn with_models_provider_for_tests( + pub(crate) async fn with_models_provider_for_tests( auth: CodexAuth, provider: ModelProviderInfo, ) -> Self { @@ -338,11 +366,18 @@ impl ThreadManager { )); std::fs::create_dir_all(&codex_home) .unwrap_or_else(|err| panic!("temp codex home dir create failed: {err}")); - let mut manager = Self::with_models_provider_and_home_for_tests( + let state_db = state_db_from_roots_for_tests( + codex_home.clone(), + codex_home.clone(), + OPENAI_PROVIDER_ID.to_string(), + ) + .await; + let mut manager = Self::with_models_provider_and_home_and_state_db_for_tests( auth, provider, codex_home.clone(), Arc::new(EnvironmentManager::default_for_tests()), + state_db, ); manager._test_codex_home_guard = Some(TempCodexHomeGuard { path: codex_home }); manager @@ -350,27 +385,33 @@ impl ThreadManager { /// Construct with a dummy AuthManager containing the provided CodexAuth and codex home. /// Used for integration tests: should not be used by ordinary business logic. - pub(crate) fn with_models_provider_and_home_for_tests( + pub(crate) async fn with_models_provider_and_home_for_tests( auth: CodexAuth, provider: ModelProviderInfo, codex_home: PathBuf, environment_manager: Arc, ) -> Self { - Self::with_models_provider_home_and_state_for_tests( + let state_db = state_db_from_roots_for_tests( + codex_home.clone(), + codex_home.clone(), + OPENAI_PROVIDER_ID.to_string(), + ) + .await; + Self::with_models_provider_and_home_and_state_db_for_tests( auth, provider, codex_home, environment_manager, - /*state_db*/ None, + state_db, ) } - pub(crate) fn with_models_provider_home_and_state_for_tests( + fn with_models_provider_and_home_and_state_db_for_tests( auth: CodexAuth, provider: ModelProviderInfo, codex_home: PathBuf, environment_manager: Arc, - state_db: Option, + state_db: StateDbHandle, ) -> Self { set_thread_manager_test_mode_for_tests(/*enabled*/ true); let auth_manager = AuthManager::from_auth_for_testing(auth); @@ -396,11 +437,11 @@ impl ThreadManager { let thread_store: Arc = Arc::new(LocalThreadStore::new( LocalThreadStoreConfig { codex_home: codex_home.clone(), - sqlite_home: codex_home.clone(), default_model_provider_id: OPENAI_PROVIDER_ID.to_string(), }, state_db.clone(), )); + let agent_graph_store = agent_graph_store_from_state_db(state_db.clone()); Self { state: Arc::new(ThreadManagerState { threads: Arc::new(RwLock::new(HashMap::new())), @@ -413,10 +454,11 @@ impl ThreadManager { mcp_manager, skills_watcher, thread_store, + state_db, + agent_graph_store, auth_manager, session_source: SessionSource::Exec, analytics_events_client: None, - state_db, ops_log: should_use_test_thread_manager_behavior() .then(|| Arc::new(std::sync::Mutex::new(Vec::new()))), }), @@ -523,22 +565,17 @@ impl ThreadManager { subtree_thread_ids.push(thread_id); seen_thread_ids.insert(thread_id); - if let Some(state_db_ctx) = thread.state_db() { - for status in [ - DirectionalThreadSpawnEdgeStatus::Open, - DirectionalThreadSpawnEdgeStatus::Closed, - ] { - for descendant_id in state_db_ctx - .list_thread_spawn_descendants_with_status(thread_id, status) - .await - .map_err(|err| { - CodexErr::Fatal(format!("failed to load thread-spawn descendants: {err}")) - })? - { - if seen_thread_ids.insert(descendant_id) { - subtree_thread_ids.push(descendant_id); - } - } + for descendant_id in self + .state + .agent_graph_store + .list_thread_spawn_descendants(thread_id, /*status_filter*/ None) + .await + .map_err(|err| { + CodexErr::Fatal(format!("failed to load thread-spawn descendants: {err}")) + })? + { + if seen_thread_ids.insert(descendant_id) { + subtree_thread_ids.push(descendant_id); } } @@ -864,10 +901,14 @@ impl ThreadManager { } impl ThreadManagerState { - pub(crate) fn state_db(&self) -> Option { + pub(crate) fn state_db(&self) -> StateDbHandle { self.state_db.clone() } + pub(crate) fn agent_graph_store(&self) -> Arc { + self.agent_graph_store.clone() + } + pub(crate) async fn list_thread_ids(&self) -> Vec { self.threads .read() @@ -1172,6 +1213,7 @@ impl ThreadManagerState { parent_trace, environment_selections, analytics_events_client: self.analytics_events_client.clone(), + state_db: Some(self.state_db.clone()), thread_store: Arc::clone(&self.thread_store), }) .await?; diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index d50a3ae8f9..643309aac1 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -1,6 +1,5 @@ use super::*; use crate::config::test_config; -use crate::init_state_db; use crate::rollout::RolloutRecorder; use crate::session::session::SessionSettingsUpdate; use crate::session::tests::make_session_and_context; @@ -47,6 +46,21 @@ fn assistant_msg(text: &str) -> ResponseItem { } } +async fn state_backed_stores( + config: &Config, +) -> ( + StateDbHandle, + Arc, + Arc, +) { + let state_db = init_state_db_from_config(config) + .await + .expect("thread manager test requires state db"); + let thread_store = thread_store_from_config(config, state_db.clone()); + let agent_graph_store = agent_graph_store_from_state_db(state_db.clone()); + (state_db, thread_store, agent_graph_store) +} + fn contextual_user_interrupted_marker() -> ResponseItem { interrupted_turn_history_marker(InterruptedTurnHistoryMarker::ContextualUser) .expect("contextual-user interrupted marker should be enabled") @@ -261,7 +275,8 @@ async fn shutdown_all_threads_bounded_submits_shutdown_to_every_thread() { config.model_provider.clone(), config.codex_home.to_path_buf(), Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - ); + ) + .await; let thread_1 = manager .start_thread(config.clone()) .await @@ -310,7 +325,8 @@ async fn start_thread_accepts_explicit_environment_when_default_environment_is_d config.model_provider.clone(), config.codex_home.to_path_buf(), environment_manager, - ); + ) + .await; let thread = manager .start_thread_with_options(StartThreadOptions { @@ -345,7 +361,8 @@ async fn start_thread_keeps_internal_threads_hidden_from_normal_lookups() { config.model_provider.clone(), config.codex_home.to_path_buf(), Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - ); + ) + .await; let thread = manager .start_thread_with_options(StartThreadOptions { config, @@ -384,14 +401,16 @@ async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await; let manager = ThreadManager::new( &config, auth_manager.clone(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, - thread_store_from_config(&config, /*state_db*/ None), - /*state_db*/ None, + state_db, + thread_store, + agent_graph_store, ); let selected_cwd = AbsolutePathBuf::try_from(config.cwd.as_path().join("selected")).expect("absolute path"); @@ -494,14 +513,16 @@ async fn resume_active_thread_from_rollout_returns_running_thread() { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await; let manager = ThreadManager::new( &config, auth_manager.clone(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, - thread_store_from_config(&config, /*state_db*/ None), - /*state_db*/ None, + state_db, + thread_store, + agent_graph_store, ); let source = manager @@ -548,14 +569,16 @@ async fn resume_stopped_thread_from_rollout_spawns_new_thread() { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await; let manager = ThreadManager::new( &config, auth_manager.clone(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, - thread_store_from_config(&config, /*state_db*/ None), - /*state_db*/ None, + state_db, + thread_store, + agent_graph_store, ); let source = manager @@ -612,14 +635,16 @@ async fn new_uses_active_provider_for_model_refresh() { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await; let manager = ThreadManager::new( &config, auth_manager, SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, - thread_store_from_config(&config, /*state_db*/ None), - /*state_db*/ None, + state_db, + thread_store, + agent_graph_store, ); let _ = manager.list_models(RefreshStrategy::Online).await; @@ -824,15 +849,16 @@ async fn interrupted_fork_snapshot_does_not_synthesize_turn_id_for_legacy_histor let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); - let state_db = init_state_db(&config).await; + let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await; let manager = ThreadManager::new( &config, auth_manager.clone(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, - thread_store_from_config(&config, state_db.clone()), - state_db.clone(), + state_db, + thread_store, + agent_graph_store, ); let source = manager @@ -928,15 +954,16 @@ async fn interrupted_fork_snapshot_preserves_explicit_turn_id() { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); - let state_db = init_state_db(&config).await; + let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await; let manager = ThreadManager::new( &config, auth_manager.clone(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, - thread_store_from_config(&config, state_db.clone()), - state_db.clone(), + state_db, + thread_store, + agent_graph_store, ); let source = manager @@ -1021,15 +1048,16 @@ async fn interrupted_fork_snapshot_uses_persisted_mid_turn_history_without_live_ let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); - let state_db = init_state_db(&config).await; + let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await; let manager = ThreadManager::new( &config, auth_manager.clone(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, - thread_store_from_config(&config, state_db.clone()), - state_db.clone(), + state_db, + thread_store, + agent_graph_store, ); let source = manager @@ -1159,15 +1187,16 @@ async fn resumed_thread_keeps_paused_goal_paused() -> anyhow::Result<()> { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); - let state_db = init_state_db(&config).await; + let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await; let manager = ThreadManager::new( &config, auth_manager.clone(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, - thread_store_from_config(&config, state_db.clone()), - state_db.clone(), + state_db, + thread_store, + agent_graph_store, ); let source = manager @@ -1236,6 +1265,11 @@ async fn resumed_thread_keeps_paused_goal_paused() -> anyhow::Result<()> { .await .is_none() ); + let goal = state_db + .get_thread_goal(resumed.thread_id) + .await? + .expect("goal should still exist after resume"); + assert_eq!(codex_state::ThreadGoalStatus::Paused, goal.status); resumed.thread.shutdown_and_wait().await?; Ok(()) diff --git a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs index 88c1698d8f..9a23550d41 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -6,6 +6,7 @@ use crate::function_tool::FunctionCallError; use crate::init_state_db; use crate::session::tests::make_session_and_context; use crate::session_prefix::format_subagent_notification_message; +use crate::thread_manager::agent_graph_store_from_state_db; use crate::thread_manager::thread_store_from_config; use crate::tools::context::ToolOutput; use crate::tools::handlers::multi_agents_v2::CloseAgentHandler as CloseAgentHandlerV2; @@ -90,11 +91,12 @@ fn parse_agent_id(id: &str) -> ThreadId { ThreadId::from_string(id).expect("agent id should be valid") } -fn thread_manager() -> ThreadManager { +async fn thread_manager() -> ThreadManager { ThreadManager::with_models_provider_for_tests( CodexAuth::from_api_key("dummy"), built_in_model_providers(/* openai_base_url */ /*openai_base_url*/ None)["openai"].clone(), ) + .await } async fn install_role_with_model_override(turn: &mut TurnContext) -> String { @@ -241,7 +243,7 @@ async fn spawn_agent_uses_explorer_role_and_preserves_approval_policy() { } let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let mut config = (*turn.config).clone(); let provider_info = @@ -296,7 +298,7 @@ async fn spawn_agent_uses_explorer_role_and_preserves_approval_policy() { async fn spawn_agent_fork_context_rejects_agent_type_override() { let (mut session, mut turn) = make_session_and_context().await; let role_name = install_role_with_model_override(&mut turn).await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -328,7 +330,7 @@ async fn spawn_agent_fork_context_rejects_agent_type_override() { #[tokio::test] async fn spawn_agent_fork_context_rejects_child_model_overrides() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -363,7 +365,7 @@ async fn spawn_agent_fork_context_rejects_child_model_overrides() { async fn multi_agent_v2_spawn_fork_turns_all_rejects_agent_type_override() { let (mut session, mut turn) = make_session_and_context().await; let role_name = install_role_with_model_override(&mut turn).await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -406,7 +408,7 @@ async fn multi_agent_v2_spawn_fork_turns_all_rejects_agent_type_override() { #[tokio::test] async fn multi_agent_v2_spawn_defaults_to_full_fork_and_rejects_child_model_overrides() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -447,7 +449,7 @@ async fn multi_agent_v2_spawn_defaults_to_full_fork_and_rejects_child_model_over async fn multi_agent_v2_spawn_partial_fork_turns_allows_agent_type_override() { let (mut session, mut turn) = make_session_and_context().await; let role_name = install_role_with_model_override(&mut turn).await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -503,7 +505,7 @@ async fn multi_agent_v2_spawn_partial_fork_turns_allows_agent_type_override() { #[tokio::test] async fn spawn_agent_returns_agent_id_without_task_name() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let output = SpawnAgentHandler @@ -530,7 +532,7 @@ async fn spawn_agent_returns_agent_id_without_task_name() { #[tokio::test] async fn multi_agent_v2_spawn_requires_task_name() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -564,7 +566,7 @@ async fn multi_agent_v2_spawn_requires_task_name() { #[tokio::test] async fn multi_agent_v2_spawn_rejects_legacy_items_field() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -624,7 +626,7 @@ async fn multi_agent_v2_spawn_returns_path_and_send_message_accepts_relative_pat } let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -721,7 +723,7 @@ async fn multi_agent_v2_spawn_returns_path_and_send_message_accepts_relative_pat #[tokio::test] async fn multi_agent_v2_spawn_rejects_legacy_fork_context() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -760,7 +762,7 @@ async fn multi_agent_v2_spawn_rejects_legacy_fork_context() { #[tokio::test] async fn multi_agent_v2_spawn_rejects_invalid_fork_turns_string() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -799,7 +801,7 @@ async fn multi_agent_v2_spawn_rejects_invalid_fork_turns_string() { #[tokio::test] async fn multi_agent_v2_spawn_rejects_zero_fork_turns() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -838,7 +840,7 @@ async fn multi_agent_v2_spawn_rejects_zero_fork_turns() { #[tokio::test] async fn multi_agent_v2_send_message_accepts_root_target_from_child() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -914,7 +916,7 @@ async fn multi_agent_v2_send_message_accepts_root_target_from_child() { #[tokio::test] async fn multi_agent_v2_followup_task_rejects_root_target_from_child() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -995,7 +997,7 @@ async fn multi_agent_v2_followup_task_rejects_root_target_from_child() { #[tokio::test] async fn multi_agent_v2_list_agents_returns_completed_status_and_last_task_message() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -1089,7 +1091,7 @@ async fn multi_agent_v2_list_agents_returns_completed_status_and_last_task_messa #[tokio::test] async fn multi_agent_v2_list_agents_filters_by_relative_path_prefix() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -1176,7 +1178,7 @@ async fn multi_agent_v2_list_agents_filters_by_relative_path_prefix() { #[tokio::test] async fn multi_agent_v2_list_agents_omits_closed_agents() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -1240,7 +1242,7 @@ async fn multi_agent_v2_list_agents_omits_closed_agents() { #[tokio::test] async fn multi_agent_v2_send_message_rejects_legacy_items_field() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -1296,7 +1298,7 @@ async fn multi_agent_v2_send_message_rejects_legacy_items_field() { #[tokio::test] async fn multi_agent_v2_send_message_rejects_interrupt_parameter() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -1369,7 +1371,7 @@ async fn multi_agent_v2_send_message_rejects_interrupt_parameter() { #[tokio::test] async fn multi_agent_v2_followup_task_completion_notifies_parent_on_every_turn() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -1504,7 +1506,7 @@ async fn multi_agent_v2_followup_task_completion_notifies_parent_on_every_turn() #[tokio::test] async fn multi_agent_v2_followup_task_rejects_legacy_items_field() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -1557,7 +1559,7 @@ async fn multi_agent_v2_followup_task_rejects_legacy_items_field() { #[tokio::test] async fn multi_agent_v2_interrupted_turn_does_not_notify_parent() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -1634,7 +1636,7 @@ async fn multi_agent_v2_interrupted_turn_does_not_notify_parent() { #[tokio::test] async fn multi_agent_v2_spawn_omits_agent_id_when_named() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -1673,7 +1675,7 @@ async fn multi_agent_v2_spawn_omits_agent_id_when_named() { #[tokio::test] async fn multi_agent_v2_spawn_surfaces_task_name_validation_errors() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -1716,7 +1718,7 @@ async fn spawn_agent_reapplies_runtime_sandbox_after_role_config() { } let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let expected_sandbox = turn.config.legacy_sandbox_policy(); let mut expected_file_system_sandbox_policy = @@ -1797,7 +1799,7 @@ async fn spawn_agent_reapplies_runtime_sandbox_after_role_config() { #[tokio::test] async fn spawn_agent_rejects_when_depth_limit_exceeded() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let max_depth = turn.config.agent_max_depth; @@ -1835,7 +1837,7 @@ async fn spawn_agent_allows_depth_up_to_configured_max_depth() { } let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let mut config = (*turn.config).clone(); @@ -1881,7 +1883,7 @@ async fn multi_agent_v2_spawn_agent_ignores_configured_max_depth() { } let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let mut config = (*turn.config).clone(); config.agent_max_depth = 1; config @@ -1989,7 +1991,7 @@ async fn send_input_rejects_invalid_id() { #[tokio::test] async fn send_input_reports_missing_agent() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let agent_id = ThreadId::new(); let invocation = invocation( @@ -2010,7 +2012,7 @@ async fn send_input_reports_missing_agent() { #[tokio::test] async fn send_input_interrupts_before_prompt() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let config = turn.config.as_ref().clone(); let thread = manager @@ -2052,7 +2054,7 @@ async fn send_input_interrupts_before_prompt() { #[tokio::test] async fn send_input_accepts_structured_items() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let config = turn.config.as_ref().clone(); let thread = manager @@ -2126,7 +2128,7 @@ async fn resume_agent_rejects_invalid_id() { #[tokio::test] async fn resume_agent_reports_missing_agent() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let agent_id = ThreadId::new(); let invocation = invocation( @@ -2147,7 +2149,7 @@ async fn resume_agent_reports_missing_agent() { #[tokio::test] async fn resume_agent_noops_for_active_agent() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let config = turn.config.as_ref().clone(); let thread = manager @@ -2186,7 +2188,7 @@ async fn resume_agent_noops_for_active_agent() { #[tokio::test] async fn resume_agent_restores_closed_agent_and_accepts_send_input() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let config = turn.config.as_ref().clone(); let thread = manager @@ -2265,7 +2267,7 @@ async fn resume_agent_restores_closed_agent_and_accepts_send_input() { #[tokio::test] async fn resume_agent_rejects_when_depth_limit_exceeded() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let max_depth = turn.config.agent_max_depth; @@ -2354,7 +2356,7 @@ async fn wait_agent_rejects_empty_targets() { #[tokio::test] async fn multi_agent_v2_wait_agent_accepts_timeout_only_argument() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -2493,7 +2495,7 @@ async fn multi_agent_v2_wait_agent_uses_configured_min_timeout() { #[tokio::test] async fn wait_agent_returns_not_found_for_missing_agents() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let id_a = ThreadId::new(); let id_b = ThreadId::new(); @@ -2529,7 +2531,7 @@ async fn wait_agent_returns_not_found_for_missing_agents() { #[tokio::test] async fn wait_agent_times_out_when_status_is_not_final() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let config = turn.config.as_ref().clone(); let thread = manager @@ -2572,7 +2574,7 @@ async fn wait_agent_times_out_when_status_is_not_final() { #[tokio::test] async fn wait_agent_clamps_short_timeouts_to_minimum() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let config = turn.config.as_ref().clone(); let thread = manager @@ -2610,7 +2612,7 @@ async fn wait_agent_clamps_short_timeouts_to_minimum() { #[tokio::test] async fn wait_agent_returns_final_status_without_timeout() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let config = turn.config.as_ref().clone(); let thread = manager @@ -2662,7 +2664,7 @@ async fn wait_agent_returns_final_status_without_timeout() { #[tokio::test] async fn multi_agent_v2_wait_agent_returns_summary_for_mailbox_activity() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -2753,7 +2755,7 @@ async fn multi_agent_v2_wait_agent_returns_summary_for_mailbox_activity() { #[tokio::test] async fn multi_agent_v2_wait_agent_returns_for_already_queued_mail() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -2831,7 +2833,7 @@ async fn multi_agent_v2_wait_agent_returns_for_already_queued_mail() { #[tokio::test] async fn multi_agent_v2_wait_agent_wakes_on_any_mailbox_notification() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -2919,7 +2921,7 @@ async fn multi_agent_v2_wait_agent_wakes_on_any_mailbox_notification() { #[tokio::test] async fn multi_agent_v2_wait_agent_does_not_return_completed_content() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -3005,7 +3007,7 @@ async fn multi_agent_v2_wait_agent_does_not_return_completed_content() { #[tokio::test] async fn multi_agent_v2_close_agent_accepts_task_name_target() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -3064,7 +3066,7 @@ async fn multi_agent_v2_close_agent_accepts_task_name_target() { #[tokio::test] async fn multi_agent_v2_close_agent_rejects_root_target_and_id() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -3112,7 +3114,7 @@ async fn multi_agent_v2_close_agent_rejects_root_target_and_id() { #[tokio::test] async fn close_agent_submits_shutdown_and_returns_previous_status() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let config = turn.config.as_ref().clone(); let thread = manager @@ -3157,15 +3159,18 @@ async fn tool_handlers_cascade_close_and_resume_and_keep_explicitly_closed_subtr .features .enable(Feature::Sqlite) .expect("test config should allow sqlite"); - let state_db = init_state_db(&config).await; + let state_db = init_state_db(&config) + .await + .expect("test config should initialize state db"); let manager = ThreadManager::new( &config, AuthManager::from_auth_for_testing(CodexAuth::from_api_key("dummy")), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, - thread_store_from_config(&config, state_db.clone()), state_db.clone(), + thread_store_from_config(&config, state_db.clone()), + agent_graph_store_from_state_db(state_db.clone()), ); let parent = manager diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index 14f652f58b..13dd026654 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -15,7 +15,9 @@ use anyhow::anyhow; use codex_config::CloudRequirementsLoader; use codex_core::CodexThread; use codex_core::ThreadManager; +use codex_core::agent_graph_store_from_state_db; use codex_core::config::Config; +use codex_core::init_state_db_from_config; use codex_core::shell::Shell; use codex_core::shell::get_shell_by_model_provided_path; use codex_core::thread_store_from_config; @@ -423,25 +425,30 @@ impl TestCodexBuilder { environment_manager: Arc, ) -> anyhow::Result { let auth = self.auth.clone(); - let state_db = codex_core::init_state_db(&config).await; let thread_manager = if config.model_catalog.is_some() { + let state_db = init_state_db_from_config(&config) + .await + .expect("test codex requires state db"); + let thread_store = thread_store_from_config(&config, state_db.clone()); + let agent_graph_store = agent_graph_store_from_state_db(state_db.clone()); ThreadManager::new( &config, codex_core::test_support::auth_manager_from_auth(auth.clone()), SessionSource::Exec, Arc::clone(&environment_manager), /*analytics_events_client*/ None, - thread_store_from_config(&config, state_db.clone()), - state_db.clone(), + state_db, + thread_store, + agent_graph_store, ) } else { - codex_core::test_support::thread_manager_with_models_provider_home_and_state( + codex_core::test_support::thread_manager_with_models_provider_and_home( auth.clone(), config.model_provider.clone(), config.codex_home.to_path_buf(), Arc::clone(&environment_manager), - state_db.clone(), ) + .await }; let thread_manager = Arc::new(thread_manager); let user_shell_override = self.user_shell_override.clone(); diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index afa586a158..dafde4ccea 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -5,6 +5,8 @@ use codex_core::NewThread; use codex_core::Prompt; use codex_core::ResponseEvent; use codex_core::ThreadManager; +use codex_core::agent_graph_store_from_state_db; +use codex_core::init_state_db_from_config; use codex_core::thread_store_from_config; use codex_features::Feature; use codex_login::AuthManager; @@ -770,15 +772,10 @@ async fn includes_conversation_id_and_model_headers_in_request() { let installation_id = std::fs::read_to_string(test.codex_home_path().join(INSTALLATION_ID_FILENAME)) .expect("read installation id"); - let session_id_string = session_id.to_string(); - assert_eq!(request_session_id, session_id_string); + assert_eq!(request_session_id, session_id.to_string()); assert_eq!(request_originator, originator().value); assert_eq!(request_authorization, "Bearer Test API Key"); - assert_eq!( - request_body["prompt_cache_key"].as_str(), - Some(session_id_string.as_str()) - ); assert_eq!( request_body["client_metadata"]["x-codex-installation-id"].as_str(), Some(installation_id.as_str()) @@ -1106,14 +1103,20 @@ async fn prefers_apikey_when_config_prefers_apikey_even_with_chatgpt_tokens() { Ok(None) => panic!("No CodexAuth found in codex_home"), Err(e) => panic!("Failed to load CodexAuth: {e}"), }; + let state_db = init_state_db_from_config(&config) + .await + .expect("client test requires state db"); + let thread_store = thread_store_from_config(&config, state_db.clone()); + let agent_graph_store = agent_graph_store_from_state_db(state_db.clone()); let thread_manager = ThreadManager::new( &config, auth_manager, SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, - thread_store_from_config(&config, /*state_db*/ None), - /*state_db*/ None, + state_db, + thread_store, + agent_graph_store, ); let NewThread { thread: codex, .. } = thread_manager .start_thread(config.clone()) diff --git a/codex-rs/core/tests/suite/personality_migration.rs b/codex-rs/core/tests/suite/personality_migration.rs index 68121d2e12..8d33ba1705 100644 --- a/codex-rs/core/tests/suite/personality_migration.rs +++ b/codex-rs/core/tests/suite/personality_migration.rs @@ -13,6 +13,8 @@ use codex_protocol::protocol::SessionMeta; use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::UserMessageEvent; +use codex_rollout::RolloutConfig; +use codex_rollout::state_db::StateDbHandle; use pretty_assertions::assert_eq; use std::io; use std::path::Path; @@ -26,6 +28,27 @@ async fn read_config_toml(codex_home: &Path) -> io::Result { toml::from_str(&contents).map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err)) } +async fn state_db_for_test(codex_home: &Path) -> io::Result { + let config = RolloutConfig { + codex_home: codex_home.to_path_buf(), + sqlite_home: codex_home.to_path_buf(), + cwd: codex_home.to_path_buf(), + model_provider_id: "openai".to_string(), + generate_memories: false, + }; + codex_rollout::state_db::try_init(&config) + .await + .map_err(io::Error::other) +} + +async fn run_migration( + codex_home: &Path, + config_toml: &ConfigToml, +) -> io::Result { + let state_db = state_db_for_test(codex_home).await?; + maybe_migrate_personality(codex_home, config_toml, state_db).await +} + async fn write_session_with_user_event(codex_home: &Path) -> io::Result<()> { let thread_id = ThreadId::new(); let dir = codex_home @@ -141,8 +164,7 @@ async fn migration_marker_exists_no_sessions_no_change() -> io::Result<()> { let marker_path = temp.path().join(PERSONALITY_MIGRATION_FILENAME); tokio::fs::write(&marker_path, "v1\n").await?; - let status = - maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?; + let status = run_migration(temp.path(), &ConfigToml::default()).await?; assert_eq!(status, PersonalityMigrationStatus::SkippedMarker); assert_eq!( @@ -156,8 +178,7 @@ async fn migration_marker_exists_no_sessions_no_change() -> io::Result<()> { async fn no_marker_no_sessions_no_change() -> io::Result<()> { let temp = TempDir::new()?; - let status = - maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?; + let status = run_migration(temp.path(), &ConfigToml::default()).await?; assert_eq!(status, PersonalityMigrationStatus::SkippedNoSessions); assert_eq!( @@ -176,8 +197,7 @@ async fn no_marker_sessions_sets_personality() -> io::Result<()> { let temp = TempDir::new()?; write_session_with_user_event(temp.path()).await?; - let status = - maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?; + let status = run_migration(temp.path(), &ConfigToml::default()).await?; assert_eq!(status, PersonalityMigrationStatus::Applied); assert_eq!( @@ -197,7 +217,7 @@ async fn no_marker_sessions_preserves_existing_config_fields() -> io::Result<()> tokio::fs::write(temp.path().join("config.toml"), "model = \"gpt-5.4\"\n").await?; let config_toml = read_config_toml(temp.path()).await?; - let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?; + let status = run_migration(temp.path(), &config_toml).await?; assert_eq!(status, PersonalityMigrationStatus::Applied); let persisted = read_config_toml(temp.path()).await?; @@ -211,8 +231,7 @@ async fn no_marker_meta_only_rollout_is_treated_as_no_sessions() -> io::Result<( let temp = TempDir::new()?; write_session_with_meta_only(temp.path()).await?; - let status = - maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?; + let status = run_migration(temp.path(), &ConfigToml::default()).await?; assert_eq!(status, PersonalityMigrationStatus::SkippedNoSessions); assert_eq!( @@ -232,7 +251,7 @@ async fn no_marker_explicit_global_personality_skips_migration() -> io::Result<( write_session_with_user_event(temp.path()).await?; let config_toml = parse_config_toml("personality = \"friendly\"\n")?; - let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?; + let status = run_migration(temp.path(), &config_toml).await?; assert_eq!( status, @@ -262,7 +281,7 @@ personality = "friendly" "#, )?; - let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?; + let status = run_migration(temp.path(), &config_toml).await?; assert_eq!( status, @@ -285,7 +304,7 @@ async fn marker_short_circuits_invalid_profile_resolution() -> io::Result<()> { tokio::fs::write(temp.path().join(PERSONALITY_MIGRATION_FILENAME), "v1\n").await?; let config_toml = parse_config_toml("profile = \"missing\"\n")?; - let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?; + let status = run_migration(temp.path(), &config_toml).await?; assert_eq!(status, PersonalityMigrationStatus::SkippedMarker); Ok(()) @@ -296,7 +315,7 @@ async fn invalid_selected_profile_returns_error_and_does_not_write_marker() -> i let temp = TempDir::new()?; let config_toml = parse_config_toml("profile = \"missing\"\n")?; - let err = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None) + let err = run_migration(temp.path(), &config_toml) .await .expect_err("missing profile should fail"); @@ -313,10 +332,8 @@ async fn applied_migration_is_idempotent_on_second_run() -> io::Result<()> { let temp = TempDir::new()?; write_session_with_user_event(temp.path()).await?; - let first_status = - maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?; - let second_status = - maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?; + let first_status = run_migration(temp.path(), &ConfigToml::default()).await?; + let second_status = run_migration(temp.path(), &ConfigToml::default()).await?; assert_eq!(first_status, PersonalityMigrationStatus::Applied); assert_eq!(second_status, PersonalityMigrationStatus::SkippedMarker); @@ -330,8 +347,7 @@ async fn no_marker_archived_sessions_sets_personality() -> io::Result<()> { let temp = TempDir::new()?; write_archived_session_with_user_event(temp.path()).await?; - let status = - maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?; + let status = run_migration(temp.path(), &ConfigToml::default()).await?; assert_eq!(status, PersonalityMigrationStatus::Applied); assert_eq!( diff --git a/codex-rs/core/tests/suite/prompt_debug_tests.rs b/codex-rs/core/tests/suite/prompt_debug_tests.rs index dc506bc474..4fee438261 100644 --- a/codex-rs/core/tests/suite/prompt_debug_tests.rs +++ b/codex-rs/core/tests/suite/prompt_debug_tests.rs @@ -29,7 +29,6 @@ async fn build_prompt_input_includes_context_and_user_message() -> Result<()> { text: "hello from debug prompt".to_string(), text_elements: Vec::new(), }], - /*state_db*/ None, ) .await?; diff --git a/codex-rs/core/tests/suite/resume_warning.rs b/codex-rs/core/tests/suite/resume_warning.rs index cb545df351..02e7626d87 100644 --- a/codex-rs/core/tests/suite/resume_warning.rs +++ b/codex-rs/core/tests/suite/resume_warning.rs @@ -95,7 +95,8 @@ async fn emits_warning_when_resumed_model_differs() { let thread_manager = codex_core::test_support::thread_manager_with_models_provider( CodexAuth::from_api_key("test"), config.model_provider.clone(), - ); + ) + .await; let auth_manager = codex_core::test_support::auth_manager_from_auth(CodexAuth::from_api_key("test")); diff --git a/codex-rs/core/tests/suite/skills.rs b/codex-rs/core/tests/suite/skills.rs index add1ee6ea7..ebf2c3226e 100644 --- a/codex-rs/core/tests/suite/skills.rs +++ b/codex-rs/core/tests/suite/skills.rs @@ -3,6 +3,8 @@ use anyhow::Result; use codex_core::ThreadManager; +use codex_core::agent_graph_store_from_state_db; +use codex_core::init_state_db_from_config; use codex_core::thread_store_from_config; use codex_exec_server::CreateDirectoryOptions; use codex_exec_server::EnvironmentManager; @@ -235,6 +237,11 @@ async fn list_skills_skips_cwd_roots_when_environment_disabled() -> Result<()> { let mut config = load_default_config_for_test(&codex_home).await; config.cwd = AbsolutePathBuf::from_absolute_path_checked(cwd.path())?; + let state_db = init_state_db_from_config(&config) + .await + .expect("skills test requires state db"); + let thread_store = thread_store_from_config(&config, state_db.clone()); + let agent_graph_store = agent_graph_store_from_state_db(state_db.clone()); let thread_manager = ThreadManager::new( &config, codex_core::test_support::auth_manager_from_auth(CodexAuth::from_api_key("dummy")), @@ -246,8 +253,9 @@ async fn list_skills_skips_cwd_roots_when_environment_disabled() -> Result<()> { )?, )), /*analytics_events_client*/ None, - thread_store_from_config(&config, /*state_db*/ None), - /*state_db*/ None, + state_db, + thread_store, + agent_graph_store, ); let new_thread = thread_manager.start_thread(config.clone()).await?; let cwd = config.cwd.to_path_buf(); diff --git a/codex-rs/core/tests/suite/unstable_features_warning.rs b/codex-rs/core/tests/suite/unstable_features_warning.rs index 66a7366589..7a4924cdda 100644 --- a/codex-rs/core/tests/suite/unstable_features_warning.rs +++ b/codex-rs/core/tests/suite/unstable_features_warning.rs @@ -34,7 +34,8 @@ async fn emits_warning_when_unstable_features_enabled_via_config() { let thread_manager = codex_core::test_support::thread_manager_with_models_provider( CodexAuth::from_api_key("test"), config.model_provider.clone(), - ); + ) + .await; let auth_manager = codex_core::test_support::auth_manager_from_auth(CodexAuth::from_api_key("test")); @@ -81,7 +82,8 @@ async fn suppresses_warning_when_configured() { let thread_manager = codex_core::test_support::thread_manager_with_models_provider( CodexAuth::from_api_key("test"), config.model_provider.clone(), - ); + ) + .await; let auth_manager = codex_core::test_support::auth_manager_from_auth(CodexAuth::from_api_key("test")); diff --git a/codex-rs/exec-server/Cargo.toml b/codex-rs/exec-server/Cargo.toml index 5f31ca4329..1495397c78 100644 --- a/codex-rs/exec-server/Cargo.toml +++ b/codex-rs/exec-server/Cargo.toml @@ -23,9 +23,10 @@ codex-sandboxing = { workspace = true } codex-utils-absolute-path = { workspace = true } codex-utils-pty = { workspace = true } futures = { workspace = true } -reqwest = { workspace = true, features = ["rustls-tls", "stream"] } +reqwest = { workspace = true, features = ["json", "rustls-tls", "stream"] } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } +sha2 = { workspace = true } thiserror = { workspace = true } tokio = { workspace = true, features = [ "fs", @@ -51,3 +52,4 @@ pretty_assertions = { workspace = true } serial_test = { workspace = true } tempfile = { workspace = true } test-case = "3.3.1" +wiremock = { workspace = true } diff --git a/codex-rs/exec-server/README.md b/codex-rs/exec-server/README.md index 78b92e1a71..81664eaca0 100644 --- a/codex-rs/exec-server/README.md +++ b/codex-rs/exec-server/README.md @@ -22,6 +22,11 @@ the wire. The CLI entrypoint supports: - `ws://IP:PORT` (default) +- `--remote URL --executor-id ID [--name NAME]` + +Remote mode registers the local exec-server with the executor registry, +then reconnects to the service-provided rendezvous websocket as the executor. +It requires a bearer token in `CODEX_EXEC_SERVER_REMOTE_BEARER_TOKEN`. Wire framing: @@ -308,6 +313,8 @@ The crate exports: - `DEFAULT_LISTEN_URL` and `ExecServerListenUrlParseError` - `ExecServerRuntimePaths` - `run_main()` for embedding the websocket server +- `RemoteExecutorConfig` and `run_remote_executor()` for embedding remote + registration mode Callers must pass `ExecServerRuntimePaths` to `run_main()`. The top-level `codex exec-server` command builds these paths from the `codex` arg0 dispatch diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index f26069ac7a..47359393d3 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -254,6 +254,18 @@ pub enum ExecServerError { Protocol(String), #[error("exec-server rejected request ({code}): {message}")] Server { code: i64, message: String }, + #[error("executor registry request failed ({status}{code_suffix}): {message}", code_suffix = .code.as_ref().map(|code| format!(", {code}")).unwrap_or_default())] + ExecutorRegistryHttp { + status: reqwest::StatusCode, + code: Option, + message: String, + }, + #[error("executor registry configuration error: {0}")] + ExecutorRegistryConfig(String), + #[error("executor registry authentication error: {0}")] + ExecutorRegistryAuth(String), + #[error("executor registry request failed: {0}")] + ExecutorRegistryRequest(#[from] reqwest::Error), } impl ExecServerClient { diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index 1550653d94..d860d59aba 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -11,6 +11,7 @@ mod local_process; mod process; mod process_id; mod protocol; +mod remote; mod remote_file_system; mod remote_process; mod rpc; @@ -87,6 +88,9 @@ pub use protocol::TerminateResponse; pub use protocol::WriteParams; pub use protocol::WriteResponse; pub use protocol::WriteStatus; +pub use remote::CODEX_EXEC_SERVER_REMOTE_BEARER_TOKEN_ENV_VAR; +pub use remote::RemoteExecutorConfig; +pub use remote::run_remote_executor; pub use runtime_paths::ExecServerRuntimePaths; pub use server::DEFAULT_LISTEN_URL; pub use server::ExecServerListenUrlParseError; diff --git a/codex-rs/exec-server/src/remote.rs b/codex-rs/exec-server/src/remote.rs new file mode 100644 index 0000000000..b574ced72f --- /dev/null +++ b/codex-rs/exec-server/src/remote.rs @@ -0,0 +1,392 @@ +use std::collections::BTreeMap; +use std::env; +use std::time::Duration; + +use reqwest::StatusCode; +use serde::Deserialize; +use serde::Serialize; +use serde_json::Value; +use sha2::Digest as _; +use tokio::time::sleep; +use tokio_tungstenite::connect_async; +use tracing::warn; +use uuid::Uuid; + +use crate::ExecServerError; +use crate::ExecServerRuntimePaths; +use crate::connection::JsonRpcConnection; +use crate::server::ConnectionProcessor; + +pub const CODEX_EXEC_SERVER_REMOTE_BEARER_TOKEN_ENV_VAR: &str = + "CODEX_EXEC_SERVER_REMOTE_BEARER_TOKEN"; + +const PROTOCOL_VERSION: &str = "codex-exec-server-v1"; +const ERROR_BODY_PREVIEW_BYTES: usize = 4096; + +#[derive(Clone)] +struct ExecutorRegistryClient { + base_url: String, + bearer_token: String, + http: reqwest::Client, +} + +impl std::fmt::Debug for ExecutorRegistryClient { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ExecutorRegistryClient") + .field("base_url", &self.base_url) + .field("bearer_token", &"") + .finish_non_exhaustive() + } +} + +impl ExecutorRegistryClient { + fn new(base_url: String, bearer_token: String) -> Result { + let base_url = normalize_base_url(base_url)?; + Ok(Self { + base_url, + bearer_token, + http: reqwest::Client::new(), + }) + } + + async fn register_executor( + &self, + request: &ExecutorRegistryRegisterExecutorRequest, + ) -> Result { + self.post_json( + &format!("/cloud/executor/{}/register", request.executor_id), + request, + ) + .await + } + + async fn post_json(&self, path: &str, request: &T) -> Result + where + T: Serialize + Sync, + R: for<'de> Deserialize<'de>, + { + let response = self + .http + .post(endpoint_url(&self.base_url, path)) + .bearer_auth(&self.bearer_token) + .json(request) + .send() + .await?; + + if response.status().is_success() { + return response.json::().await.map_err(ExecServerError::from); + } + + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + if matches!(status, StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN) { + return Err(executor_registry_auth_error(status, &body)); + } + + Err(executor_registry_http_error(status, &body)) + } +} + +#[derive(Debug, Clone, Eq, PartialEq, Serialize)] +struct ExecutorRegistryRegisterExecutorRequest { + idempotency_id: String, + executor_id: String, + #[serde(skip_serializing_if = "Option::is_none")] + name: Option, + labels: BTreeMap, + metadata: Value, +} + +#[derive(Debug, Clone, Eq, PartialEq, Deserialize)] +struct ExecutorRegistryExecutorRegistrationResponse { + id: String, + executor_id: String, + url: String, +} + +/// Configuration for registering an exec-server for remote use. +#[derive(Clone, Eq, PartialEq)] +pub struct RemoteExecutorConfig { + pub base_url: String, + pub executor_id: String, + pub name: String, + bearer_token: String, +} + +impl std::fmt::Debug for RemoteExecutorConfig { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("RemoteExecutorConfig") + .field("base_url", &self.base_url) + .field("executor_id", &self.executor_id) + .field("name", &self.name) + .field("bearer_token", &"") + .finish() + } +} + +impl RemoteExecutorConfig { + pub fn new(base_url: String, executor_id: String) -> Result { + Self::with_bearer_token(base_url, executor_id, read_remote_bearer_token_from_env()?) + } + + fn with_bearer_token( + base_url: String, + executor_id: String, + bearer_token: String, + ) -> Result { + let executor_id = normalize_executor_id(executor_id)?; + let bearer_token = normalize_bearer_token(bearer_token)?; + Ok(Self { + base_url, + executor_id, + name: "codex-exec-server".to_string(), + bearer_token, + }) + } + + fn registration_request( + &self, + registration_id: Uuid, + ) -> ExecutorRegistryRegisterExecutorRequest { + ExecutorRegistryRegisterExecutorRequest { + idempotency_id: self.default_idempotency_id(registration_id), + executor_id: self.executor_id.clone(), + name: Some(self.name.clone()), + labels: BTreeMap::new(), + metadata: Value::Object(Default::default()), + } + } + + fn default_idempotency_id(&self, registration_id: Uuid) -> String { + let mut hasher = sha2::Sha256::new(); + hasher.update(self.executor_id.as_bytes()); + hasher.update(b"\0"); + hasher.update(self.name.as_bytes()); + hasher.update(b"\0"); + hasher.update(PROTOCOL_VERSION); + hasher.update(b"\0"); + hasher.update(registration_id.as_bytes()); + let digest = hasher.finalize(); + format!("codex-exec-server-{digest:x}") + } +} + +/// Register an exec-server for remote use and serve requests over the returned +/// rendezvous websocket. +pub async fn run_remote_executor( + config: RemoteExecutorConfig, + runtime_paths: ExecServerRuntimePaths, +) -> Result<(), ExecServerError> { + let client = ExecutorRegistryClient::new(config.base_url.clone(), config.bearer_token.clone())?; + let processor = ConnectionProcessor::new(runtime_paths); + let registration_id = Uuid::new_v4(); + let mut backoff = Duration::from_secs(1); + + loop { + let request = config.registration_request(registration_id); + let response = client.register_executor(&request).await?; + eprintln!( + "codex exec-server remote executor {} registered with executor_id {}", + response.id, response.executor_id + ); + + match connect_async(response.url.as_str()).await { + Ok((websocket, _)) => { + backoff = Duration::from_secs(1); + processor + .run_connection(JsonRpcConnection::from_websocket( + websocket, + "remote exec-server websocket".to_string(), + )) + .await; + } + Err(err) => { + warn!("failed to connect remote exec-server websocket: {err}"); + } + } + + sleep(backoff).await; + backoff = (backoff * 2).min(Duration::from_secs(30)); + } +} + +fn read_remote_bearer_token_from_env() -> Result { + read_remote_bearer_token_from_env_with(|name| env::var(name)) +} + +fn read_remote_bearer_token_from_env_with(get_var: F) -> Result +where + F: FnOnce(&str) -> Result, +{ + let bearer_token = get_var(CODEX_EXEC_SERVER_REMOTE_BEARER_TOKEN_ENV_VAR).map_err(|_| { + ExecServerError::ExecutorRegistryAuth(format!( + "executor registry bearer token environment variable `{CODEX_EXEC_SERVER_REMOTE_BEARER_TOKEN_ENV_VAR}` is not set" + )) + })?; + normalize_bearer_token(bearer_token) +} + +fn normalize_bearer_token(bearer_token: String) -> Result { + let bearer_token = bearer_token.trim().to_string(); + if bearer_token.is_empty() { + return Err(ExecServerError::ExecutorRegistryAuth(format!( + "executor registry bearer token environment variable `{CODEX_EXEC_SERVER_REMOTE_BEARER_TOKEN_ENV_VAR}` is empty" + ))); + } + Ok(bearer_token) +} + +fn normalize_executor_id(executor_id: String) -> Result { + let executor_id = executor_id.trim().to_string(); + if executor_id.is_empty() { + return Err(ExecServerError::ExecutorRegistryConfig( + "executor id is required for remote exec-server registration".to_string(), + )); + } + Ok(executor_id) +} + +#[derive(Deserialize)] +struct RegistryErrorBody { + error: Option, +} + +#[derive(Deserialize)] +struct RegistryError { + code: Option, + message: Option, +} + +fn normalize_base_url(base_url: String) -> Result { + let trimmed = base_url.trim().trim_end_matches('/').to_string(); + if trimmed.is_empty() { + return Err(ExecServerError::ExecutorRegistryConfig( + "executor registry base URL is required".to_string(), + )); + } + Ok(trimmed) +} + +fn endpoint_url(base_url: &str, path: &str) -> String { + format!("{base_url}/{}", path.trim_start_matches('/')) +} + +fn executor_registry_auth_error(status: StatusCode, body: &str) -> ExecServerError { + let message = registry_error_message(body).unwrap_or_else(|| "empty error body".to_string()); + ExecServerError::ExecutorRegistryAuth(format!( + "executor registry authentication failed ({status}): {message}" + )) +} + +fn executor_registry_http_error(status: StatusCode, body: &str) -> ExecServerError { + let parsed = serde_json::from_str::(body).ok(); + let (code, message) = parsed + .and_then(|body| body.error) + .map(|error| { + ( + error.code, + error.message.unwrap_or_else(|| { + preview_error_body(body).unwrap_or_else(|| "empty error body".to_string()) + }), + ) + }) + .unwrap_or_else(|| { + ( + None, + preview_error_body(body) + .unwrap_or_else(|| "empty or malformed error body".to_string()), + ) + }); + ExecServerError::ExecutorRegistryHttp { + status, + code, + message, + } +} + +fn registry_error_message(body: &str) -> Option { + serde_json::from_str::(body) + .ok() + .and_then(|body| body.error) + .and_then(|error| error.message) + .or_else(|| preview_error_body(body)) +} + +fn preview_error_body(body: &str) -> Option { + let trimmed = body.trim(); + if trimmed.is_empty() { + return None; + } + Some(trimmed.chars().take(ERROR_BODY_PREVIEW_BYTES).collect()) +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + use serde_json::json; + use wiremock::Mock; + use wiremock::MockServer; + use wiremock::ResponseTemplate; + use wiremock::matchers::body_json; + use wiremock::matchers::header; + use wiremock::matchers::method; + use wiremock::matchers::path; + + use super::*; + + #[tokio::test] + async fn register_executor_posts_with_bearer_token_header() { + let server = MockServer::start().await; + let registration_id = Uuid::from_u128(1); + let config = RemoteExecutorConfig::with_bearer_token( + server.uri(), + "exec-requested".to_string(), + "registry-token".to_string(), + ) + .expect("config"); + let request = config.registration_request(registration_id); + let expected_request = serde_json::to_value(&request).expect("serialize request"); + Mock::given(method("POST")) + .and(path("/cloud/executor/exec-requested/register")) + .and(header("authorization", "Bearer registry-token")) + .and(body_json(expected_request)) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "id": "registration-1", + "executor_id": "exec-1", + "url": "wss://rendezvous.test/executor/exec-1?role=executor&sig=abc" + }))) + .mount(&server) + .await; + let client = ExecutorRegistryClient::new(server.uri(), "registry-token".to_string()) + .expect("client"); + + let response = client + .register_executor(&request) + .await + .expect("register executor"); + + assert_eq!( + response, + ExecutorRegistryExecutorRegistrationResponse { + id: "registration-1".to_string(), + executor_id: "exec-1".to_string(), + url: "wss://rendezvous.test/executor/exec-1?role=executor&sig=abc".to_string(), + } + ); + } + + #[test] + fn debug_output_redacts_bearer_token() { + let config = RemoteExecutorConfig::with_bearer_token( + "https://registry.example".to_string(), + "exec-1".to_string(), + "secret-token".to_string(), + ) + .expect("config"); + + let debug = format!("{config:?}"); + + assert!(debug.contains("")); + assert!(!debug.contains("secret-token")); + } +} diff --git a/codex-rs/exec-server/src/server.rs b/codex-rs/exec-server/src/server.rs index 62c1787381..bf33eb77ba 100644 --- a/codex-rs/exec-server/src/server.rs +++ b/codex-rs/exec-server/src/server.rs @@ -7,6 +7,7 @@ mod session_registry; mod transport; pub(crate) use handler::ExecServerHandler; +pub(crate) use processor::ConnectionProcessor; pub use transport::DEFAULT_LISTEN_URL; pub use transport::ExecServerListenUrlParseError; diff --git a/codex-rs/mcp-server/src/lib.rs b/codex-rs/mcp-server/src/lib.rs index 75372c6e7c..ac764456f5 100644 --- a/codex-rs/mcp-server/src/lib.rs +++ b/codex-rs/mcp-server/src/lib.rs @@ -83,7 +83,6 @@ pub async fn run_main( std::io::Error::new(ErrorKind::InvalidData, format!("error loading config: {e}")) })?; set_default_client_residency_requirement(config.enforce_residency.value()); - let state_db = codex_core::init_state_db(&config).await; let otel = codex_core::otel_init::build_provider( &config, @@ -140,15 +139,18 @@ pub async fn run_main( // Task: process incoming messages. let processor_handle = tokio::spawn({ let outgoing_message_sender = OutgoingMessageSender::new(outgoing_tx); - let mut processor = MessageProcessor::new( + let processor = MessageProcessor::new( outgoing_message_sender, arg0_paths, Arc::new(config), environment_manager, - state_db, ) .await; async move { + let Some(mut processor) = processor else { + error!("failed to initialize MCP processor"); + return; + }; while let Some(msg) = incoming_rx.recv().await { match msg { JsonRpcMessage::Request(r) => processor.process_request(r).await, diff --git a/codex-rs/mcp-server/src/message_processor.rs b/codex-rs/mcp-server/src/message_processor.rs index 9ce46dfeda..2591891e4d 100644 --- a/codex-rs/mcp-server/src/message_processor.rs +++ b/codex-rs/mcp-server/src/message_processor.rs @@ -2,9 +2,10 @@ use std::collections::HashMap; use std::sync::Arc; use codex_arg0::Arg0DispatchPaths; -use codex_core::StateDbHandle; use codex_core::ThreadManager; +use codex_core::agent_graph_store_from_state_db; use codex_core::config::Config; +use codex_core::init_state_db_from_config; use codex_core::thread_store_from_config; use codex_exec_server::EnvironmentManager; use codex_login::AuthManager; @@ -54,30 +55,33 @@ impl MessageProcessor { arg0_paths: Arg0DispatchPaths, config: Arc, environment_manager: Arc, - state_db: Option, - ) -> Self { + ) -> Option { let outgoing = Arc::new(outgoing); let auth_manager = AuthManager::shared_from_config( config.as_ref(), /*enable_codex_api_key_env*/ false, ) .await; + let state_db = init_state_db_from_config(config.as_ref()).await?; + let thread_store = thread_store_from_config(config.as_ref(), state_db.clone()); + let agent_graph_store = agent_graph_store_from_state_db(state_db.clone()); let thread_manager = Arc::new(ThreadManager::new( config.as_ref(), auth_manager, SessionSource::Mcp, environment_manager, /*analytics_events_client*/ None, - thread_store_from_config(config.as_ref(), state_db.clone()), - state_db.clone(), + state_db, + thread_store, + agent_graph_store, )); - Self { + Some(Self { outgoing, initialized: false, arg0_paths, thread_manager, running_requests_id_to_codex_uuid: Arc::new(Mutex::new(HashMap::new())), - } + }) } pub(crate) async fn process_request(&mut self, request: JsonRpcRequest) { diff --git a/codex-rs/thread-manager-sample/src/main.rs b/codex-rs/thread-manager-sample/src/main.rs index 9d84ece674..a27c78d543 100644 --- a/codex-rs/thread-manager-sample/src/main.rs +++ b/codex-rs/thread-manager-sample/src/main.rs @@ -52,10 +52,11 @@ use codex_core_api::TuiNotificationSettings; use codex_core_api::UriBasedFileOpener; use codex_core_api::UserInput; use codex_core_api::WebSearchMode; +use codex_core_api::agent_graph_store_from_state_db; use codex_core_api::arg0_dispatch_or_else; use codex_core_api::built_in_model_providers; use codex_core_api::find_codex_home; -use codex_core_api::init_state_db; +use codex_core_api::init_state_db_from_config; use codex_core_api::item_event_to_server_notification; use codex_core_api::set_default_originator; use codex_core_api::thread_store_from_config; @@ -104,7 +105,6 @@ async fn run_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { }; let config = new_config(args.model, arg0_paths)?; - let state_db = init_state_db(&config).await; let auth_manager = AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false).await; @@ -112,7 +112,11 @@ async fn run_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { config.codex_self_exe.clone(), config.codex_linux_sandbox_exe.clone(), )?; + let Some(state_db) = init_state_db_from_config(&config).await else { + bail!("thread manager sample requires state db"); + }; let thread_store = thread_store_from_config(&config, state_db.clone()); + let agent_graph_store = agent_graph_store_from_state_db(state_db.clone()); let environment_manager = Arc::new(EnvironmentManager::new(EnvironmentManagerArgs::new(local_runtime_paths)).await); let thread_manager = ThreadManager::new( @@ -121,8 +125,9 @@ async fn run_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { SessionSource::Exec, environment_manager, /*analytics_events_client*/ None, - Arc::clone(&thread_store), state_db, + Arc::clone(&thread_store), + agent_graph_store, ); let NewThread { diff --git a/codex-rs/thread-store/src/local/archive_thread.rs b/codex-rs/thread-store/src/local/archive_thread.rs index 8fb214e98c..268a6ffc47 100644 --- a/codex-rs/thread-store/src/local/archive_thread.rs +++ b/codex-rs/thread-store/src/local/archive_thread.rs @@ -13,11 +13,11 @@ pub(super) async fn archive_thread( params: ArchiveThreadParams, ) -> ThreadStoreResult<()> { let thread_id = params.thread_id; - let state_db_ctx = store.state_db().await; + let state_db = store.state_db(); let rollout_path = find_thread_path_by_id_str( store.config.codex_home.as_path(), &thread_id.to_string(), - state_db_ctx.as_deref(), + Some(state_db.as_ref()), ) .await .map_err(|err| ThreadStoreError::InvalidRequest { @@ -52,11 +52,10 @@ pub(super) async fn archive_thread( } })?; - if let Some(ctx) = state_db_ctx { - let _ = ctx - .mark_archived(thread_id, archived_path.as_path(), Utc::now()) - .await; - } + let _ = store + .state_db() + .mark_archived(thread_id, archived_path.as_path(), Utc::now()) + .await; Ok(()) } @@ -75,13 +74,15 @@ mod tests { use crate::ThreadSortKey; use crate::ThreadStore; use crate::local::LocalThreadStore; + use crate::local::test_support::init_test_state_db; use crate::local::test_support::test_config; + use crate::local::test_support::test_store; use crate::local::test_support::write_session_file; #[tokio::test] async fn archive_thread_moves_rollout_to_archived_collection() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(201); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let active_path = @@ -127,21 +128,12 @@ mod tests { async fn archive_thread_updates_sqlite_metadata_when_present() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(202); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let active_path = write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file"); - let runtime = codex_state::StateRuntime::init( - home.path().to_path_buf(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); - runtime - .mark_backfill_complete(/*last_watermark*/ None) - .await - .expect("backfill should be complete"); let mut builder = codex_state::ThreadMetadataBuilder::new( thread_id, active_path.clone(), diff --git a/codex-rs/thread-store/src/local/create_thread.rs b/codex-rs/thread-store/src/local/create_thread.rs index e444e5c91d..cb0f64ea9c 100644 --- a/codex-rs/thread-store/src/local/create_thread.rs +++ b/codex-rs/thread-store/src/local/create_thread.rs @@ -22,12 +22,12 @@ pub(super) async fn create_thread( })?; let config = RolloutConfig { codex_home: store.config.codex_home.clone(), - sqlite_home: store.config.sqlite_home.clone(), + sqlite_home: store.sqlite_home(), cwd, model_provider_id: params.metadata.model_provider.clone(), generate_memories: matches!(params.metadata.memory_mode, ThreadMemoryMode::Enabled), }; - let state_db_ctx = store.state_db().await; + let state_db_ctx = Some(store.state_db()); let recorder = RolloutRecorder::new( &config, RolloutRecorderParams::new( diff --git a/codex-rs/thread-store/src/local/list_threads.rs b/codex-rs/thread-store/src/local/list_threads.rs index e470ad2be9..5cf1f27aa1 100644 --- a/codex-rs/thread-store/src/local/list_threads.rs +++ b/codex-rs/thread-store/src/local/list_threads.rs @@ -39,16 +39,16 @@ pub(super) async fn list_threads( SortDirection::Asc => codex_rollout::SortDirection::Asc, SortDirection::Desc => codex_rollout::SortDirection::Desc, }; - let state_db = store.state_db().await; let rollout_config = RolloutConfig { codex_home: store.config.codex_home.clone(), - sqlite_home: store.config.sqlite_home.clone(), + sqlite_home: store.sqlite_home(), cwd: store.config.codex_home.clone(), model_provider_id: store.config.default_model_provider_id.clone(), generate_memories: false, }; + let state_db_ctx = Some(store.state_db()); let page = list_rollout_threads( - state_db, + state_db_ctx, &rollout_config, store.config.default_model_provider_id.as_str(), ¶ms, @@ -80,14 +80,13 @@ pub(super) async fn list_threads( .map(|thread| thread.thread_id) .collect::>(); let mut names = HashMap::::with_capacity(thread_ids.len()); - if let Some(state_db_ctx) = store.state_db().await { - for &thread_id in &thread_ids { - let Ok(Some(metadata)) = state_db_ctx.get_thread(thread_id).await else { - continue; - }; - if let Some(title) = distinct_thread_metadata_title(&metadata) { - names.insert(thread_id, title); - } + let state_db_ctx = store.state_db(); + for &thread_id in &thread_ids { + let Ok(Some(metadata)) = state_db_ctx.get_thread(thread_id).await else { + continue; + }; + if let Some(title) = distinct_thread_metadata_title(&metadata) { + names.insert(thread_id, title); } } if names.len() < thread_ids.len() @@ -108,9 +107,9 @@ pub(super) async fn list_threads( } async fn list_rollout_threads( - state_db: Option, + state_db_ctx: Option, config: &RolloutConfig, - default_model_provider_id: &str, + default_model_provider: &str, params: &ListThreadsParams, cursor: Option<&codex_rollout::Cursor>, sort_key: codex_rollout::ThreadSortKey, @@ -118,7 +117,7 @@ async fn list_rollout_threads( ) -> ThreadStoreResult { let page = if params.use_state_db_only && params.archived { RolloutRecorder::list_archived_threads_from_state_db( - state_db, + state_db_ctx.clone(), config, params.page_size, cursor, @@ -127,13 +126,13 @@ async fn list_rollout_threads( params.allowed_sources.as_slice(), params.model_providers.as_deref(), params.cwd_filters.as_deref(), - default_model_provider_id, + default_model_provider, params.search_term.as_deref(), ) .await } else if params.use_state_db_only { RolloutRecorder::list_threads_from_state_db( - state_db, + state_db_ctx.clone(), config, params.page_size, cursor, @@ -142,13 +141,13 @@ async fn list_rollout_threads( params.allowed_sources.as_slice(), params.model_providers.as_deref(), params.cwd_filters.as_deref(), - default_model_provider_id, + default_model_provider, params.search_term.as_deref(), ) .await } else if params.archived { RolloutRecorder::list_archived_threads( - state_db, + state_db_ctx.clone(), config, params.page_size, cursor, @@ -157,13 +156,13 @@ async fn list_rollout_threads( params.allowed_sources.as_slice(), params.model_providers.as_deref(), params.cwd_filters.as_deref(), - default_model_provider_id, + default_model_provider, params.search_term.as_deref(), ) .await } else { RolloutRecorder::list_threads( - state_db, + state_db_ctx, config, params.page_size, cursor, @@ -172,7 +171,7 @@ async fn list_rollout_threads( params.allowed_sources.as_slice(), params.model_providers.as_deref(), params.cwd_filters.as_deref(), - default_model_provider_id, + default_model_provider, params.search_term.as_deref(), ) .await @@ -195,7 +194,9 @@ mod tests { use super::*; use crate::ThreadStore; use crate::local::LocalThreadStore; + use crate::local::test_support::init_test_state_db; use crate::local::test_support::test_config; + use crate::local::test_support::test_store; use crate::local::test_support::write_archived_session_file; use crate::local::test_support::write_session_file; use crate::local::test_support::write_session_file_with; @@ -203,7 +204,7 @@ mod tests { #[tokio::test] async fn list_threads_uses_default_provider_when_rollout_omits_provider() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; write_session_file_with( home.path(), home.path().join("sessions/2025/01/03"), @@ -238,22 +239,13 @@ mod tests { async fn list_threads_preserves_sqlite_title_search_results() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(103); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = home.path().join("rollout-title-search.jsonl"); fs::write(&rollout_path, "").expect("placeholder rollout file"); - let runtime = codex_state::StateRuntime::init( - home.path().to_path_buf(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); - runtime - .mark_backfill_complete(/*last_watermark*/ None) - .await - .expect("backfill should be complete"); let created_at = Utc::now(); let mut builder = codex_state::ThreadMetadataBuilder::new( thread_id, @@ -267,6 +259,10 @@ mod tests { let mut metadata = builder.build(config.default_model_provider_id.as_str()); metadata.title = "needle title".to_string(); metadata.first_user_message = Some("plain preview".to_string()); + runtime + .mark_backfill_complete(/*last_watermark*/ None) + .await + .expect("backfill should be complete"); runtime .upsert_thread(&metadata) .await @@ -303,7 +299,7 @@ mod tests { #[tokio::test] async fn list_threads_selects_active_or_archived_collection() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let active_uuid = Uuid::from_u128(105); let archived_uuid = Uuid::from_u128(106); write_session_file(home.path(), "2025-01-03T12-00-00", active_uuid) @@ -372,7 +368,7 @@ mod tests { async fn list_threads_returns_local_rollout_summary() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); - let store = LocalThreadStore::new(config, /*state_db*/ None); + let store = LocalThreadStore::new(config.clone(), init_test_state_db(&config).await); let uuid = Uuid::from_u128(101); let path = write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file"); @@ -411,7 +407,7 @@ mod tests { #[tokio::test] async fn list_threads_rejects_invalid_cursor() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let err = store .list_threads(ListThreadsParams { diff --git a/codex-rs/thread-store/src/local/live_writer.rs b/codex-rs/thread-store/src/local/live_writer.rs index 643207b59d..37d9758973 100644 --- a/codex-rs/thread-store/src/local/live_writer.rs +++ b/codex-rs/thread-store/src/local/live_writer.rs @@ -66,12 +66,12 @@ pub(super) async fn resume_thread( })?; let config = RolloutConfig { codex_home: store.config.codex_home.clone(), - sqlite_home: store.config.sqlite_home.clone(), + sqlite_home: store.sqlite_home(), cwd, model_provider_id: params.metadata.model_provider.clone(), generate_memories: matches!(params.metadata.memory_mode, ThreadMemoryMode::Enabled), }; - let state_db_ctx = store.state_db().await; + let state_db_ctx = Some(store.state_db()); let recorder = RolloutRecorder::new( &config, RolloutRecorderParams::resume( diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index 058c7b2309..e6ea63bba1 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -41,7 +41,7 @@ use crate::UpdateThreadMetadataParams; pub struct LocalThreadStore { pub(super) config: LocalThreadStoreConfig, live_recorders: Arc>>, - state_db: Option, + state_db: StateDbHandle, } /// Process-scoped configuration for local thread storage. @@ -51,7 +51,6 @@ pub struct LocalThreadStore { #[derive(Clone, Debug, PartialEq, Eq)] pub struct LocalThreadStoreConfig { pub codex_home: PathBuf, - pub sqlite_home: PathBuf, /// Provider used only when older local metadata does not contain one. pub default_model_provider_id: String, } @@ -60,7 +59,6 @@ impl LocalThreadStoreConfig { pub fn from_config(config: &impl codex_rollout::RolloutConfigView) -> Self { Self { codex_home: config.codex_home().to_path_buf(), - sqlite_home: config.sqlite_home().to_path_buf(), default_model_provider_id: config.model_provider_id().to_string(), } } @@ -75,8 +73,9 @@ impl std::fmt::Debug for LocalThreadStore { } impl LocalThreadStore { - /// Create a local store using an already initialized state DB handle. - pub fn new(config: LocalThreadStoreConfig, state_db: Option) -> Self { + /// Create a local store from process-scoped local storage configuration and + /// the caller-provided shared state DB handle. + pub fn new(config: LocalThreadStoreConfig, state_db: StateDbHandle) -> Self { Self { config, live_recorders: Arc::new(Mutex::new(HashMap::new())), @@ -85,10 +84,14 @@ impl LocalThreadStore { } /// Return the state DB handle used by local rollout writers. - pub async fn state_db(&self) -> Option { + pub fn state_db(&self) -> StateDbHandle { self.state_db.clone() } + pub(super) fn sqlite_home(&self) -> PathBuf { + self.state_db.codex_home().to_path_buf() + } + /// Read a local rollout-backed thread by path. pub async fn read_thread_by_rollout_path( &self, @@ -282,14 +285,16 @@ mod tests { use super::*; use crate::ThreadEventPersistenceMode; use crate::ThreadPersistenceMetadata; + use crate::local::test_support::init_test_state_db; use crate::local::test_support::test_config; + use crate::local::test_support::test_store; use crate::local::test_support::write_archived_session_file; use crate::local::test_support::write_session_file; #[tokio::test] async fn live_writer_lifecycle_writes_and_closes() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let thread_id = ThreadId::default(); store @@ -338,7 +343,7 @@ mod tests { #[tokio::test] async fn create_thread_rejects_missing_cwd() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let thread_id = ThreadId::default(); let mut params = create_thread_params(thread_id); params.metadata.cwd = None; @@ -358,7 +363,7 @@ mod tests { #[tokio::test] async fn discard_thread_drops_unmaterialized_live_writer() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let thread_id = ThreadId::default(); store @@ -396,8 +401,9 @@ mod tests { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); let thread_id = ThreadId::default(); + let state_db = init_test_state_db(&config).await; - let first_store = LocalThreadStore::new(config.clone(), /*state_db*/ None); + let first_store = LocalThreadStore::new(config.clone(), state_db.clone()); first_store .create_thread(create_thread_params(thread_id)) .await @@ -426,7 +432,7 @@ mod tests { .await .expect("shutdown initial writer"); - let resumed_store = LocalThreadStore::new(config, /*state_db*/ None); + let resumed_store = LocalThreadStore::new(config, state_db); resumed_store .resume_thread(ResumeThreadParams { thread_id, @@ -457,7 +463,7 @@ mod tests { #[tokio::test] async fn create_thread_rejects_duplicate_live_writer() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let thread_id = ThreadId::default(); store @@ -477,7 +483,7 @@ mod tests { #[tokio::test] async fn resume_thread_rejects_duplicate_live_writer() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let thread_id = ThreadId::default(); store @@ -506,7 +512,7 @@ mod tests { #[tokio::test] async fn resume_thread_rejects_missing_cwd() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = uuid::Uuid::from_u128(407); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = @@ -535,7 +541,7 @@ mod tests { async fn load_history_uses_live_writer_rollout_path() { let home = TempDir::new().expect("temp dir"); let external_home = TempDir::new().expect("external temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = uuid::Uuid::from_u128(404); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = write_session_file(external_home.path(), "2025-01-04T10-00-00", uuid) @@ -584,7 +590,7 @@ mod tests { async fn read_thread_uses_live_writer_rollout_path_for_external_resume() { let home = TempDir::new().expect("temp dir"); let external_home = TempDir::new().expect("external temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = uuid::Uuid::from_u128(406); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = write_session_file(external_home.path(), "2025-01-04T11-00-00", uuid) @@ -623,7 +629,7 @@ mod tests { #[tokio::test] async fn load_history_uses_live_writer_rollout_path_for_archived_source() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = uuid::Uuid::from_u128(405); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = write_archived_session_file(home.path(), "2025-01-04T10-30-00", uuid) @@ -691,7 +697,7 @@ mod tests { #[tokio::test] async fn read_thread_by_rollout_path_includes_history() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let thread_id = ThreadId::default(); store diff --git a/codex-rs/thread-store/src/local/read_thread.rs b/codex-rs/thread-store/src/local/read_thread.rs index c538d938a3..5bcf411beb 100644 --- a/codex-rs/thread-store/src/local/read_thread.rs +++ b/codex-rs/thread-store/src/local/read_thread.rs @@ -176,12 +176,12 @@ async fn resolve_rollout_path( return Ok(Some(path)); } - let state_db_ctx = store.state_db().await; + let state_db = store.state_db(); if include_archived { match find_thread_path_by_id_str( store.config.codex_home.as_path(), &thread_id.to_string(), - state_db_ctx.as_deref(), + Some(state_db.as_ref()), ) .await .map_err(|err| ThreadStoreError::InvalidRequest { @@ -191,7 +191,7 @@ async fn resolve_rollout_path( None => find_archived_thread_path_by_id_str( store.config.codex_home.as_path(), &thread_id.to_string(), - state_db_ctx.as_deref(), + Some(state_db.as_ref()), ) .await .map_err(|err| ThreadStoreError::InvalidRequest { @@ -202,7 +202,7 @@ async fn resolve_rollout_path( find_thread_path_by_id_str( store.config.codex_home.as_path(), &thread_id.to_string(), - state_db_ctx.as_deref(), + Some(state_db.as_ref()), ) .await .map_err(|err| ThreadStoreError::InvalidRequest { @@ -260,8 +260,7 @@ async fn read_sqlite_metadata( store: &LocalThreadStore, thread_id: codex_protocol::ThreadId, ) -> Option { - let runtime = store.state_db().await?; - runtime.get_thread(thread_id).await.ok().flatten() + store.state_db().get_thread(thread_id).await.ok().flatten() } async fn stored_thread_from_sqlite_metadata( @@ -412,7 +411,9 @@ mod tests { use super::*; use crate::ThreadStore; use crate::local::LocalThreadStore; + use crate::local::test_support::init_test_state_db; use crate::local::test_support::test_config; + use crate::local::test_support::test_store; use crate::local::test_support::write_archived_session_file; use crate::local::test_support::write_session_file; use crate::local::test_support::write_session_file_with_fork; @@ -420,7 +421,7 @@ mod tests { #[tokio::test] async fn read_thread_returns_active_rollout_summary() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(205); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let active_path = @@ -448,7 +449,7 @@ mod tests { #[tokio::test] async fn read_thread_returns_rollout_path_summary() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(211); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let active_path = @@ -479,17 +480,12 @@ mod tests { async fn read_thread_by_rollout_path_prefers_sqlite_git_info() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(223); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let active_path = write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file"); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); let mut builder = ThreadMetadataBuilder::new( thread_id, active_path.clone(), @@ -527,7 +523,7 @@ mod tests { #[tokio::test] async fn read_thread_returns_archived_rollout_when_requested() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(207); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let archived_path = write_archived_session_file(home.path(), "2025-01-03T12-00-00", uuid) @@ -568,7 +564,7 @@ mod tests { #[tokio::test] async fn read_thread_prefers_active_rollout_over_archived() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(208); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let active_path = @@ -593,7 +589,7 @@ mod tests { #[tokio::test] async fn read_thread_returns_forked_from_id() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(209); let parent_uuid = Uuid::from_u128(210); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); @@ -626,17 +622,12 @@ mod tests { async fn read_thread_applies_sqlite_thread_name() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(212); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file"); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); let mut builder = ThreadMetadataBuilder::new(thread_id, rollout_path, Utc::now(), SessionSource::Cli); builder.model_provider = Some(config.default_model_provider_id.clone()); @@ -666,13 +657,8 @@ mod tests { async fn read_thread_preserves_rollout_cwd_when_sqlite_metadata_exists() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(224); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let day_dir = home.path().join("sessions/2025/01/03"); @@ -741,7 +727,7 @@ mod tests { #[tokio::test] async fn read_thread_uses_legacy_thread_name_when_sqlite_title_is_missing() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(213); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file"); @@ -765,6 +751,8 @@ mod tests { async fn read_thread_uses_sqlite_metadata_for_rollout_without_user_preview() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(217); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let day_dir = home.path().join("sessions/2025/01/03"); @@ -786,13 +774,6 @@ mod tests { }); writeln!(file, "{meta}").expect("write session meta"); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); let mut builder = ThreadMetadataBuilder::new( thread_id, rollout_path.clone(), @@ -835,18 +816,13 @@ mod tests { let home = TempDir::new().expect("temp dir"); let external = TempDir::new().expect("external temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(220); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file"); let stale_path = external.path().join("missing-rollout.jsonl"); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); let mut builder = ThreadMetadataBuilder::new( thread_id, stale_path.clone(), @@ -884,6 +860,8 @@ mod tests { let home = TempDir::new().expect("temp dir"); let external = TempDir::new().expect("external temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(221); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = @@ -891,13 +869,6 @@ mod tests { let other_uuid = Uuid::from_u128(222); let stale_path = write_session_file(external.path(), "2025-01-04T12-00-00", other_uuid) .expect("other session file"); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); let mut builder = ThreadMetadataBuilder::new(thread_id, stale_path, Utc::now(), SessionSource::Cli); builder.model_provider = Some("wrong-sqlite-provider".to_string()); @@ -929,7 +900,7 @@ mod tests { #[tokio::test] async fn read_thread_uses_session_meta_for_rollout_without_user_preview_or_sqlite_metadata() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(218); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let day_dir = home.path().join("sessions/2025/01/03"); @@ -984,18 +955,13 @@ mod tests { let home = TempDir::new().expect("temp dir"); let external = TempDir::new().expect("external temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(214); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = external .path() .join(format!("rollout-2025-01-03T12-00-00-{uuid}.jsonl")); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); let mut builder = ThreadMetadataBuilder::new( thread_id, rollout_path.clone(), @@ -1042,20 +1008,15 @@ mod tests { let home = TempDir::new().expect("temp dir"); let external = TempDir::new().expect("external temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(216); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = external .path() .join(format!("rollout-2025-01-03T12-00-00-{uuid}.jsonl")); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); let mut builder = ThreadMetadataBuilder::new(thread_id, rollout_path, Utc::now(), SessionSource::Cli); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); builder.archived_at = Some(Utc::now()); let mut metadata = builder.build(config.default_model_provider_id.as_str()); metadata.first_user_message = Some("Archived SQLite preview".to_string()); @@ -1098,17 +1059,12 @@ mod tests { async fn read_thread_sqlite_fallback_loads_archived_history() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(219); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let archived_path = write_archived_session_file(home.path(), "2025-01-03T12-00-00", uuid) .expect("archived session file"); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); let mut builder = ThreadMetadataBuilder::new( thread_id, archived_path.clone(), @@ -1144,7 +1100,7 @@ mod tests { #[tokio::test] async fn read_thread_fails_without_rollout() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(206); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); diff --git a/codex-rs/thread-store/src/local/test_support.rs b/codex-rs/thread-store/src/local/test_support.rs index 98321880ff..597014efe9 100644 --- a/codex-rs/thread-store/src/local/test_support.rs +++ b/codex-rs/thread-store/src/local/test_support.rs @@ -4,18 +4,34 @@ use std::path::Path; use std::path::PathBuf; use codex_rollout::ARCHIVED_SESSIONS_SUBDIR; +use codex_rollout::StateDbHandle; use uuid::Uuid; +use super::LocalThreadStore; use super::LocalThreadStoreConfig; pub(super) fn test_config(codex_home: &Path) -> LocalThreadStoreConfig { LocalThreadStoreConfig { codex_home: codex_home.to_path_buf(), - sqlite_home: codex_home.to_path_buf(), default_model_provider_id: "test-provider".to_string(), } } +pub(super) async fn init_test_state_db(config: &LocalThreadStoreConfig) -> StateDbHandle { + codex_state::StateRuntime::init( + config.codex_home.clone(), + config.default_model_provider_id.clone(), + ) + .await + .expect("state db should initialize") +} + +pub(super) async fn test_store(codex_home: &Path) -> LocalThreadStore { + let config = test_config(codex_home); + let state_db = init_test_state_db(&config).await; + LocalThreadStore::new(config, state_db) +} + pub(super) fn write_session_file(root: &Path, ts: &str, uuid: Uuid) -> std::io::Result { write_session_file_with( root, diff --git a/codex-rs/thread-store/src/local/unarchive_thread.rs b/codex-rs/thread-store/src/local/unarchive_thread.rs index ad41db69ac..7ac3f01a0f 100644 --- a/codex-rs/thread-store/src/local/unarchive_thread.rs +++ b/codex-rs/thread-store/src/local/unarchive_thread.rs @@ -17,11 +17,11 @@ pub(super) async fn unarchive_thread( params: ArchiveThreadParams, ) -> ThreadStoreResult { let thread_id = params.thread_id; - let state_db_ctx = store.state_db().await; + let state_db = store.state_db(); let archived_path = find_archived_thread_path_by_id_str( store.config.codex_home.as_path(), &thread_id.to_string(), - state_db_ctx.as_deref(), + Some(state_db.as_ref()), ) .await .map_err(|err| ThreadStoreError::InvalidRequest { @@ -73,11 +73,10 @@ pub(super) async fn unarchive_thread( message: format!("failed to update unarchived thread timestamp: {err}"), })?; - if let Some(ctx) = state_db_ctx { - let _ = ctx - .mark_unarchived(thread_id, restored_path.as_path()) - .await; - } + let _ = store + .state_db() + .mark_unarchived(thread_id, restored_path.as_path()) + .await; let item = read_thread_item_from_rollout(restored_path.clone()) .await @@ -112,13 +111,15 @@ mod tests { use super::*; use crate::ThreadStore; use crate::local::LocalThreadStore; + use crate::local::test_support::init_test_state_db; use crate::local::test_support::test_config; + use crate::local::test_support::test_store; use crate::local::test_support::write_archived_session_file; #[tokio::test] async fn unarchive_thread_restores_rollout_and_returns_updated_thread() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(203); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let archived_path = write_archived_session_file(home.path(), "2025-01-03T13-00-00", uuid) @@ -149,21 +150,12 @@ mod tests { async fn unarchive_thread_updates_sqlite_metadata_when_present() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(204); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let archived_path = write_archived_session_file(home.path(), "2025-01-03T13-00-00", uuid) .expect("archived session file"); - let runtime = codex_state::StateRuntime::init( - home.path().to_path_buf(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); - runtime - .mark_backfill_complete(/*last_watermark*/ None) - .await - .expect("backfill should be complete"); let mut builder = codex_state::ThreadMetadataBuilder::new( thread_id, archived_path.clone(), diff --git a/codex-rs/thread-store/src/local/update_thread_metadata.rs b/codex-rs/thread-store/src/local/update_thread_metadata.rs index 98ee788662..2c70c815ed 100644 --- a/codex-rs/thread-store/src/local/update_thread_metadata.rs +++ b/codex-rs/thread-store/src/local/update_thread_metadata.rs @@ -59,9 +59,8 @@ pub(super) async fn update_thread_metadata( .await?; } - let state_db_ctx = store.state_db().await; codex_rollout::state_db::reconcile_rollout( - state_db_ctx.as_deref(), + Some(store.state_db()).as_deref(), resolved_rollout_path.path.as_path(), store.config.default_model_provider_id.as_str(), /*builder*/ None, @@ -73,11 +72,7 @@ pub(super) async fn update_thread_metadata( let resolved_git_info = match git_info { Some(git_info) => { - let Some(state_db) = store.state_db().await else { - return Err(ThreadStoreError::Internal { - message: format!("sqlite state db unavailable for thread {thread_id}"), - }); - }; + let state_db = store.state_db(); let metadata = state_db .get_thread(thread_id) @@ -157,11 +152,7 @@ async fn apply_thread_git_info( branch: &Option, origin_url: &Option, ) -> ThreadStoreResult<()> { - let Some(state_db) = store.state_db().await else { - return Err(ThreadStoreError::Internal { - message: format!("sqlite state db unavailable for thread {thread_id}"), - }); - }; + let state_db = store.state_db(); let updated = state_db .update_thread_git_info( thread_id, @@ -307,11 +298,11 @@ async fn resolve_rollout_path( return Ok(ResolvedRolloutPath { path, archived }); } - let state_db_ctx = store.state_db().await; + let state_db = store.state_db(); let active_path = find_thread_path_by_id_str( store.config.codex_home.as_path(), &thread_id.to_string(), - state_db_ctx.as_deref(), + Some(state_db.as_ref()), ) .await .map_err(|err| ThreadStoreError::InvalidRequest { @@ -331,7 +322,7 @@ async fn resolve_rollout_path( find_archived_thread_path_by_id_str( store.config.codex_home.as_path(), &thread_id.to_string(), - state_db_ctx.as_deref(), + Some(state_db.as_ref()), ) .await .map_err(|err| ThreadStoreError::InvalidRequest { @@ -366,14 +357,16 @@ mod tests { use crate::ThreadPersistenceMetadata; use crate::ThreadStore; use crate::local::LocalThreadStore; + use crate::local::test_support::init_test_state_db; use crate::local::test_support::test_config; + use crate::local::test_support::test_store; use crate::local::test_support::write_archived_session_file; use crate::local::test_support::write_session_file; #[tokio::test] async fn update_thread_metadata_sets_name_on_active_rollout_and_indexes_name() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(301); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let path = @@ -408,18 +401,12 @@ mod tests { async fn update_thread_metadata_sets_memory_mode_on_active_rollout() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(302); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let path = write_session_file(home.path(), "2025-01-03T14-30-00", uuid).expect("session file"); - let runtime = codex_state::StateRuntime::init( - home.path().to_path_buf(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); - let thread = store .update_thread_metadata(UpdateThreadMetadataParams { thread_id, @@ -452,13 +439,8 @@ mod tests { let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let path = write_session_file(home.path(), "2025-01-03T18-30-00", uuid).expect("session file"); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); store .update_thread_metadata(UpdateThreadMetadataParams { @@ -517,7 +499,7 @@ mod tests { async fn update_thread_metadata_uses_live_rollout_path_for_external_resume() { let home = TempDir::new().expect("temp dir"); let external_home = TempDir::new().expect("external temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(307); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let path = write_session_file(external_home.path(), "2025-01-03T14-45-00", uuid) @@ -558,13 +540,8 @@ mod tests { async fn update_thread_metadata_sets_git_info() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config, Some(runtime)); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config, runtime); let uuid = Uuid::from_u128(309); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); write_session_file(home.path(), "2025-01-03T17-00-00", uuid).expect("session file"); @@ -601,13 +578,8 @@ mod tests { async fn update_thread_metadata_partially_updates_git_info() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config, Some(runtime)); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config, runtime); let uuid = Uuid::from_u128(310); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); write_session_file(home.path(), "2025-01-03T17-30-00", uuid).expect("session file"); @@ -659,13 +631,8 @@ mod tests { async fn update_thread_metadata_clears_git_info_fields() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(311); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let path = @@ -829,7 +796,7 @@ mod tests { #[tokio::test] async fn update_thread_metadata_rejects_mismatched_session_meta_id() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let filename_uuid = Uuid::from_u128(303); let metadata_uuid = Uuid::from_u128(304); let thread_id = ThreadId::from_string(&filename_uuid.to_string()).expect("valid thread id"); @@ -861,7 +828,7 @@ mod tests { #[tokio::test] async fn update_thread_metadata_rejects_multi_field_patch_without_partial_write() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(305); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let path = @@ -896,21 +863,12 @@ mod tests { async fn update_thread_metadata_keeps_archived_thread_archived_in_sqlite() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(306); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let archived_path = write_archived_session_file(home.path(), "2025-01-03T16-00-00", uuid) .expect("archived session file"); - let runtime = codex_state::StateRuntime::init( - home.path().to_path_buf(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); - runtime - .mark_backfill_complete(/*last_watermark*/ None) - .await - .expect("backfill should be complete"); codex_rollout::state_db::reconcile_rollout( Some(runtime.as_ref()), archived_path.as_path(), @@ -959,21 +917,12 @@ mod tests { async fn update_thread_metadata_keeps_live_archived_thread_archived_in_sqlite() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(308); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let archived_path = write_archived_session_file(home.path(), "2025-01-03T16-30-00", uuid) .expect("archived session file"); - let runtime = codex_state::StateRuntime::init( - home.path().to_path_buf(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); - runtime - .mark_backfill_complete(/*last_watermark*/ None) - .await - .expect("backfill should be complete"); codex_rollout::state_db::reconcile_rollout( Some(runtime.as_ref()), archived_path.as_path(), diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index 5622c59f65..b2e92a19b4 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -881,38 +881,40 @@ pub async fn run_main( AppServerTarget::Remote { .. } => state_db::get_state_db(&config).await, }; - let effective_toml = config.config_layer_stack.effective_config(); - match effective_toml.try_into() { - Ok(config_toml) => { - match crate::legacy_core::personality_migration::maybe_migrate_personality( - &config.codex_home, - &config_toml, - state_db.clone(), - ) - .await - { - Ok( - crate::legacy_core::personality_migration::PersonalityMigrationStatus::Applied, - ) => { - config = load_config_or_exit( - cli_kv_overrides.clone(), - overrides.clone(), - cloud_requirements.clone(), - ) - .await; - } - Ok( - crate::legacy_core::personality_migration::PersonalityMigrationStatus::SkippedMarker - | crate::legacy_core::personality_migration::PersonalityMigrationStatus::SkippedExplicitPersonality - | crate::legacy_core::personality_migration::PersonalityMigrationStatus::SkippedNoSessions, - ) => {} - Err(err) => { - tracing::warn!(error = %err, "failed to run personality migration"); + if let Some(state_db) = state_db.clone() { + let effective_toml = config.config_layer_stack.effective_config(); + match effective_toml.try_into() { + Ok(config_toml) => { + match crate::legacy_core::personality_migration::maybe_migrate_personality( + &config.codex_home, + &config_toml, + state_db, + ) + .await + { + Ok( + crate::legacy_core::personality_migration::PersonalityMigrationStatus::Applied, + ) => { + config = load_config_or_exit( + cli_kv_overrides.clone(), + overrides.clone(), + cloud_requirements.clone(), + ) + .await; + } + Ok( + crate::legacy_core::personality_migration::PersonalityMigrationStatus::SkippedMarker + | crate::legacy_core::personality_migration::PersonalityMigrationStatus::SkippedExplicitPersonality + | crate::legacy_core::personality_migration::PersonalityMigrationStatus::SkippedNoSessions, + ) => {} + Err(err) => { + tracing::warn!(error = %err, "failed to run personality migration"); + } } } - } - Err(err) => { - tracing::warn!(error = %err, "failed to deserialize config for personality migration"); + Err(err) => { + tracing::warn!(error = %err, "failed to deserialize config for personality migration"); + } } } diff --git a/codex-rs/tui/src/resume_picker.rs b/codex-rs/tui/src/resume_picker.rs index 16c3d19071..06ad0a61a7 100644 --- a/codex-rs/tui/src/resume_picker.rs +++ b/codex-rs/tui/src/resume_picker.rs @@ -5733,6 +5733,7 @@ session_picker_view = "dense" name: None, turns: vec![codex_app_server_protocol::Turn { id: String::from("turn-1"), + items_view: codex_app_server_protocol::TurnItemsView::Full, items: vec![ ThreadItem::UserMessage { id: String::from("user-1"), @@ -5752,7 +5753,6 @@ session_picker_view = "dense" text: String::from("1. Do the thing"), }, ], - items_view: codex_app_server_protocol::TurnItemsView::Full, status: codex_app_server_protocol::TurnStatus::Completed, error: None, started_at: None, @@ -5798,12 +5798,12 @@ session_picker_view = "dense" name: None, turns: vec![codex_app_server_protocol::Turn { id: String::from("turn-1"), + items_view: codex_app_server_protocol::TurnItemsView::Full, items: vec![ThreadItem::Reasoning { id: String::from("reasoning-1"), summary: Vec::new(), content: vec![String::from("private raw chain of thought")], }], - items_view: codex_app_server_protocol::TurnItemsView::Full, status: codex_app_server_protocol::TurnStatus::Completed, error: None, started_at: None, @@ -5853,12 +5853,12 @@ session_picker_view = "dense" name: None, turns: vec![codex_app_server_protocol::Turn { id: String::from("turn-1"), + items_view: codex_app_server_protocol::TurnItemsView::Full, items: vec![ThreadItem::Reasoning { id: String::from("reasoning-1"), summary: vec![String::from("public summary")], content: vec![String::from("raw reasoning content")], }], - items_view: codex_app_server_protocol::TurnItemsView::Full, status: codex_app_server_protocol::TurnStatus::Completed, error: None, started_at: None, diff --git a/codex-rs/v8-poc/BUILD.bazel b/codex-rs/v8-poc/BUILD.bazel index 0cadbd518d..05a46451a9 100644 --- a/codex-rs/v8-poc/BUILD.bazel +++ b/codex-rs/v8-poc/BUILD.bazel @@ -3,6 +3,10 @@ load("//:defs.bzl", "codex_rust_crate") codex_rust_crate( name = "v8-poc", crate_name = "codex_v8_poc", + crate_features = select({ + "@rules_rs//rs/experimental/platforms/constraints:windows_msvc": [], + "//conditions:default": ["sandbox"], + }), deps_extra = ["@crates//:v8"], ) diff --git a/codex-rs/v8-poc/Cargo.toml b/codex-rs/v8-poc/Cargo.toml index 4bf008c095..cbcb7ca404 100644 --- a/codex-rs/v8-poc/Cargo.toml +++ b/codex-rs/v8-poc/Cargo.toml @@ -8,6 +8,9 @@ license.workspace = true name = "codex_v8_poc" path = "src/lib.rs" +[features] +sandbox = ["v8/v8_enable_sandbox"] + [lints] workspace = true diff --git a/codex-rs/v8-poc/src/lib.rs b/codex-rs/v8-poc/src/lib.rs index 0e9faaab10..f43a9a1253 100644 --- a/codex-rs/v8-poc/src/lib.rs +++ b/codex-rs/v8-poc/src/lib.rs @@ -12,6 +12,17 @@ pub fn embedded_v8_version() -> &'static str { v8::V8::get_version() } +/// Returns whether the linked V8 library was built with the in-process sandbox. +#[must_use] +pub fn linked_v8_has_sandbox() -> bool { + unsafe extern "C" { + fn v8__V8__IsSandboxEnabled() -> bool; + } + + // `rusty_v8` exposes this symbol for its own sandbox verification tests. + unsafe { v8__V8__IsSandboxEnabled() } +} + #[cfg(test)] mod tests { use pretty_assertions::assert_eq; @@ -53,6 +64,11 @@ mod tests { assert!(!super::embedded_v8_version().is_empty()); } + #[test] + fn sandbox_feature_matches_linked_v8() { + assert_eq!(super::linked_v8_has_sandbox(), cfg!(feature = "sandbox")); + } + #[test] fn evaluates_integer_addition() { assert_eq!(evaluate_expression("1 + 2"), "3"); diff --git a/defs.bzl b/defs.bzl index fc8b0e0da6..4af81ecef1 100644 --- a/defs.bzl +++ b/defs.bzl @@ -331,6 +331,7 @@ def codex_rust_crate( rust_test( name = unit_test_binary, crate = name, + crate_features = crate_features, deps = all_crate_deps(normal = True, normal_dev = True) + maybe_deps + deps_extra, # Unit tests also compile to standalone Windows executables, so # keep their stack reserve aligned with binaries and integration diff --git a/patches/v8_bazel_rules.patch b/patches/v8_bazel_rules.patch index df845939d0..70ab440d85 100644 --- a/patches/v8_bazel_rules.patch +++ b/patches/v8_bazel_rules.patch @@ -65,10 +65,19 @@ index 9648e4a..88efd41 100644 ":should_add_rdynamic": ["-rdynamic"], "//conditions:default": [], diff --git a/orig/v8-14.6.202.11/BUILD.bazel b/mod/v8-14.6.202.11/BUILD.bazel -index 85f31b7..7314584 100644 +index 85f31b7..bbc351b 100644 --- a/orig/v8-14.6.202.11/BUILD.bazel +++ b/mod/v8-14.6.202.11/BUILD.bazel -@@ -303,7 +303,7 @@ v8_int( +@@ -148,6 +148,8 @@ v8_flag(name = "v8_enable_trace_maps") + + v8_flag(name = "v8_enable_v8_checks") + ++v8_flag(name = "v8_enable_sandbox") ++ + v8_flag(name = "v8_enable_verify_csa") + + v8_flag(name = "v8_enable_verify_heap") +@@ -303,7 +305,7 @@ v8_int( # If no explicit value for v8_enable_pointer_compression, we set it to 'none'. v8_string( name = "v8_enable_pointer_compression", @@ -77,7 +86,15 @@ index 85f31b7..7314584 100644 ) # Default setting for v8_enable_pointer_compression. -@@ -4077,28 +4077,14 @@ filegroup( +@@ -503,6 +505,7 @@ v8_config( + "v8_enable_slow_dchecks": "ENABLE_SLOW_DCHECKS", + "v8_enable_runtime_call_stats": "V8_RUNTIME_CALL_STATS", + "v8_enable_snapshot_native_code_counters": "V8_SNAPSHOT_NATIVE_CODE_COUNTERS", ++ "v8_enable_sandbox": "V8_ENABLE_SANDBOX", + "v8_enable_trace_maps": "V8_TRACE_MAPS", + "v8_enable_turbofan": "V8_ENABLE_TURBOFAN", + "v8_enable_v8_checks": "V8_ENABLE_CHECKS", +@@ -4077,28 +4080,14 @@ filegroup( }), ) @@ -112,7 +129,7 @@ index 85f31b7..7314584 100644 ) filegroup( -@@ -4405,6 +4391,20 @@ genrule( +@@ -4405,6 +4394,20 @@ genrule( srcs = [ "include/js_protocol.pdl", "src/inspector/inspector_protocol_config.json", @@ -133,7 +150,7 @@ index 85f31b7..7314584 100644 ], outs = [ "include/inspector/Debugger.h", -@@ -4426,15 +4426,19 @@ genrule( +@@ -4426,15 +4429,19 @@ genrule( "src/inspector/protocol/Schema.cpp", "src/inspector/protocol/Schema.h", ], @@ -157,7 +174,7 @@ index 85f31b7..7314584 100644 ], ) -@@ -4448,6 +4451,15 @@ filegroup( +@@ -4448,6 +4455,15 @@ filegroup( ], ) @@ -173,7 +190,7 @@ index 85f31b7..7314584 100644 filegroup( name = "d8_files", srcs = [ -@@ -4567,16 +4579,9 @@ cc_library( +@@ -4567,16 +4583,9 @@ cc_library( ], ) @@ -192,7 +209,8 @@ index 85f31b7..7314584 100644 + actual = "@simdutf//:simdutf", ) -@@ -4593,7 +4598,7 @@ v8_library( + v8_library( +@@ -4593,7 +4602,7 @@ v8_library( copts = ["-Wno-implicit-fallthrough"], icu_deps = [ ":icu/generated_torque_definitions_headers", @@ -201,7 +219,7 @@ index 85f31b7..7314584 100644 ], icu_srcs = [ ":generated_regexp_special_case", -@@ -4608,7 +4613,7 @@ v8_library( +@@ -4608,7 +4617,7 @@ v8_library( ], deps = [ ":lib_dragonbox", @@ -210,7 +228,7 @@ index 85f31b7..7314584 100644 ":lib_fp16", ":simdutf", ":v8_libbase", -@@ -4664,6 +4669,7 @@ alias( +@@ -4664,6 +4673,7 @@ alias( alias( name = "core_lib_icu", actual = "icu/v8", @@ -218,7 +236,7 @@ index 85f31b7..7314584 100644 ) v8_library( -@@ -4715,7 +4721,7 @@ v8_binary( +@@ -4715,7 +4725,7 @@ v8_binary( ], deps = [ ":v8_libbase", diff --git a/third_party/v8/README.md b/third_party/v8/README.md index 6b85ba66c2..f0961df1d0 100644 --- a/third_party/v8/README.md +++ b/third_party/v8/README.md @@ -3,14 +3,21 @@ This directory wires the `v8` crate to exact-version Bazel inputs. Bazel consumer builds use: -- upstream `denoland/rusty_v8` release archives on Windows -- source-built V8 archives on Darwin, GNU Linux, and musl Linux +- upstream `denoland/rusty_v8` release archives on Windows MSVC +- source-built V8 archives on Darwin, GNU Linux, musl Linux, and Windows GNU - `openai/codex` release assets for published musl release pairs Cargo builds still use prebuilt `rusty_v8` archives by default. Only Bazel overrides `RUSTY_V8_ARCHIVE`/`RUSTY_V8_SRC_BINDING_PATH` in `MODULE.bazel` to select source-built local archives for its consumer builds. +Source-built Bazel V8 artifacts enable V8's in-process sandbox by default, and +the Bazel `v8` crate feature selection tracks those targets. A full consumer +rollout still needs matching sandbox-enabled archives for every non-source-built +target. Until that artifact migration lands, the rusty_v8 publishing workflows +use `--config=v8-release-compat` to preserve the current non-sandboxed release +artifact contract. + Current pinned versions: - Rust crate: `v8 = =146.4.0`