From 36460387ec8bd5a676a568466d6a93d0efafaf12 Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Tue, 5 May 2026 14:36:37 -0700 Subject: [PATCH 1/7] Enable V8 sandboxing for source-built builds (#21146) ## Summary This is the first PR in the V8 in-process sandboxing rollout. It adds the build-system and Rust feature plumbing needed to support sandboxed V8 builds, then enables sandboxing by default for the source-built Bazel V8 path that we control directly. It deliberately keeps the published `rusty_v8` artifact workflows on their current non-sandboxed contract so this PR can land and ship independently before we change any released artifacts. ## Rollout plan - [x] **PR 1: land sandbox plumbing and default source-built Bazel V8 to sandboxed mode** - [ ] **PR 2: publish sandbox-enabled release artifacts and add compatibility validation** - Produce sandboxed artifact pairs for every released Cargo target that does not already use the source-built Bazel path. - Add CI coverage that consumes those sandboxed artifacts and verifies: - `codex-v8-poc` reports sandbox enabled - `codex-code-mode` builds/tests against the sandboxed path - [ ] **PR 3: switch release consumers to sandboxed artifacts by default** - Update released artifact selectors/checksums. - Enable the Rust `v8_enable_sandbox` feature in the default release path. - Make the sandboxed artifact family the normal path for published builds. - [ ] **PR 4: remove rollout-only compatibility paths** - Remove the temporary non-sandbox release compatibility config once the new default has shipped and baked. - Keep the invariant tests permanently. --- .bazelrc | 10 +++++ .github/scripts/rusty_v8_bazel.py | 39 ++++++++++++++---- .../verify_cargo_workspace_manifests.py | 5 ++- .github/workflows/rusty-v8-release.yml | 2 + .github/workflows/v8-canary.yml | 2 + MODULE.bazel | 12 ++++++ codex-rs/code-mode/Cargo.toml | 3 ++ codex-rs/v8-poc/BUILD.bazel | 4 ++ codex-rs/v8-poc/Cargo.toml | 3 ++ codex-rs/v8-poc/src/lib.rs | 16 ++++++++ defs.bzl | 1 + patches/v8_bazel_rules.patch | 40 ++++++++++++++----- third_party/v8/README.md | 11 ++++- 13 files changed, 126 insertions(+), 22 deletions(-) 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/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/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` From 7e310bc7f3c94af86c1d06d6e387054b3685ebaf Mon Sep 17 00:00:00 2001 From: Rasmus Rygaard Date: Tue, 5 May 2026 14:45:29 -0700 Subject: [PATCH 2/7] Inject state DB, agent graph store (#20689) ## Why We want the agent graph store to be passed down the stack as a real dependency, the same way we already treat the thread store. This will let us inject the agent graph store as a real dependency and support implementations other than the local SQLite-backed one. Right now most code instantiates a state DB and an agent graph store just-in-time. Ideally, we would not depend on the state DB directly but only read through the higher-level interfaces. This change makes the dependency boundaries explicit and moves state DB initialization to process bootstrap instead of hiding it inside local store implementations. ## What changed - `ThreadManager` now requires a `StateDbHandle` and an `AgentGraphStore` at construction time instead of treating them as optional internals. - The local store constructors no longer lazily initialize SQLite. Callers now initialize the state DB once per process and use that shared handle to build: - `LocalThreadStore` - `LocalAgentGraphStore` - App bootstraps (`app-server`, `mcp-server`, `prompt_debug`, and the thread-manager sample) now initialize the state DB up front and inject the resulting handle down the stack. - `app-server` now consistently uses its process-scoped state DB handle instead of reopening SQLite or trying to recover it from loaded threads. - Device-key storage now reuses the shared state DB handle instead of maintaining its own lazy opener. - The thread archive / descendant traversal paths now use the injected `AgentGraphStore` instead of reaching through local thread-store-specific state. ## Verification - `cargo check -p codex-core -p codex-thread-store -p codex-app-server -p codex-mcp-server -p codex-thread-manager-sample --tests` - `cargo test -p codex-thread-store` - `cargo test -p codex-core thread_manager_accepts_separate_agent_graph_store_and_thread_store -- --nocapture` - `cargo test -p codex-app-server thread_archive_archives_spawned_descendants -- --nocapture` --- codex-rs/Cargo.lock | 1 + codex-rs/app-server-client/src/lib.rs | 2 +- .../app-server/src/bespoke_event_handling.rs | 6 +- codex-rs/app-server/src/in_process.rs | 25 ++- codex-rs/app-server/src/lib.rs | 38 ++--- codex-rs/app-server/src/message_processor.rs | 11 +- .../src/message_processor_tracing_tests.rs | 6 +- .../device_key_processor.rs | 22 +-- .../thread_goal_processor.rs | 31 +--- .../request_processors/thread_processor.rs | 8 +- .../tests/suite/v2/remote_thread_store.rs | 10 +- codex-rs/cli/src/main.rs | 2 +- codex-rs/core-api/src/lib.rs | 2 + codex-rs/core/Cargo.toml | 1 + codex-rs/core/src/agent/control.rs | 48 +++--- codex-rs/core/src/agent/control_tests.rs | 41 ++--- codex-rs/core/src/codex_delegate.rs | 1 + codex-rs/core/src/goals.rs | 12 -- codex-rs/core/src/lib.rs | 2 + codex-rs/core/src/personality_migration.rs | 17 +- .../core/src/personality_migration_tests.rs | 64 +++++++- codex-rs/core/src/prompt_debug.rs | 11 +- codex-rs/core/src/session/mod.rs | 17 +- codex-rs/core/src/session/session.rs | 26 +--- codex-rs/core/src/session/tests.rs | 146 ++++++++---------- .../core/src/session/tests/guardian_tests.rs | 8 +- codex-rs/core/src/session/turn_context.rs | 1 + codex-rs/core/src/stream_events_utils.rs | 9 +- codex-rs/core/src/test_support.rs | 23 +-- codex-rs/core/src/thread_manager.rs | 110 +++++++++---- codex-rs/core/src/thread_manager_tests.rs | 82 +++++++--- .../src/tools/handlers/multi_agents_tests.rs | 107 +++++++------ codex-rs/core/tests/common/test_codex.rs | 17 +- codex-rs/core/tests/suite/client.rs | 19 ++- .../core/tests/suite/personality_migration.rs | 54 ++++--- .../core/tests/suite/prompt_debug_tests.rs | 1 - codex-rs/core/tests/suite/resume_warning.rs | 3 +- codex-rs/core/tests/suite/skills.rs | 12 +- .../tests/suite/unstable_features_warning.rs | 6 +- codex-rs/mcp-server/src/lib.rs | 8 +- codex-rs/mcp-server/src/message_processor.rs | 18 ++- codex-rs/thread-manager-sample/src/main.rs | 11 +- .../thread-store/src/local/archive_thread.rs | 30 ++-- .../thread-store/src/local/create_thread.rs | 4 +- .../thread-store/src/local/list_threads.rs | 68 ++++---- .../thread-store/src/local/live_writer.rs | 4 +- codex-rs/thread-store/src/local/mod.rs | 42 ++--- .../thread-store/src/local/read_thread.rs | 110 ++++--------- .../thread-store/src/local/test_support.rs | 18 ++- .../src/local/unarchive_thread.rs | 30 ++-- .../src/local/update_thread_metadata.rs | 103 ++++-------- codex-rs/tui/src/lib.rs | 62 ++++---- codex-rs/tui/src/resume_picker.rs | 3 + 53 files changed, 768 insertions(+), 745 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 1e1590e6ce..7819fa223a 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2411,6 +2411,7 @@ dependencies = [ "bm25", "chrono", "clap", + "codex-agent-graph-store", "codex-analytics", "codex-api", "codex-app-server-protocol", diff --git a/codex-rs/app-server-client/src/lib.rs b/codex-rs/app-server-client/src/lib.rs index 79141da24e..0911cc448d 100644 --- a/codex-rs/app-server-client/src/lib.rs +++ b/codex-rs/app-server-client/src/lib.rs @@ -29,7 +29,6 @@ pub use codex_app_server::in_process::DEFAULT_IN_PROCESS_CHANNEL_CAPACITY; pub use codex_app_server::in_process::InProcessServerEvent; use codex_app_server::in_process::InProcessStartArgs; use codex_app_server::in_process::LogDbLayer; -pub use codex_app_server::in_process::StateDbHandle; use codex_app_server_protocol::ClientInfo; use codex_app_server_protocol::ClientNotification; use codex_app_server_protocol::ClientRequest; @@ -47,6 +46,7 @@ use codex_config::LoaderOverrides; use codex_config::NoopThreadConfigLoader; use codex_config::RemoteThreadConfigLoader; use codex_config::ThreadConfigLoader; +pub use codex_core::StateDbHandle; use codex_core::config::Config; pub use codex_exec_server::EnvironmentManager; pub use codex_exec_server::EnvironmentManagerArgs; diff --git a/codex-rs/app-server/src/bespoke_event_handling.rs b/codex-rs/app-server/src/bespoke_event_handling.rs index 1907f1c3f4..55b3fea814 100644 --- a/codex-rs/app-server/src/bespoke_event_handling.rs +++ b/codex-rs/app-server/src/bespoke_event_handling.rs @@ -2629,7 +2629,8 @@ mod tests { config.model_provider.clone(), config.codex_home.to_path_buf(), Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - ), + ) + .await, ); let codex_core::NewThread { thread_id: conversation_id, @@ -3214,7 +3215,8 @@ mod tests { config.model_provider.clone(), config.codex_home.to_path_buf(), Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - ), + ) + .await, ); let codex_core::NewThread { thread_id: conversation_id, diff --git a/codex-rs/app-server/src/in_process.rs b/codex-rs/app-server/src/in_process.rs index a8c023e95a..73cfe75a54 100644 --- a/codex-rs/app-server/src/in_process.rs +++ b/codex-rs/app-server/src/in_process.rs @@ -82,11 +82,12 @@ use codex_config::CloudRequirementsLoader; use codex_config::LoaderOverrides; use codex_config::ThreadConfigLoader; use codex_core::config::Config; +use codex_core::init_state_db_from_config; use codex_exec_server::EnvironmentManager; use codex_feedback::CodexFeedback; use codex_login::AuthManager; use codex_protocol::protocol::SessionSource; -pub use codex_rollout::StateDbHandle; +use codex_rollout::state_db::StateDbHandle; pub use codex_state::log_db::LogDbLayer; use tokio::sync::mpsc; use tokio::sync::oneshot; @@ -127,7 +128,7 @@ pub struct InProcessStartArgs { pub feedback: CodexFeedback, /// SQLite tracing layer used to flush recently emitted logs before feedback upload. pub log_db: Option, - /// Process-wide SQLite state handle shared with embedded app-server consumers. + /// Optional state DB handle to use for the in-process runtime. pub state_db: Option, /// Environment manager used by core execution and filesystem operations. pub environment_manager: Arc, @@ -344,7 +345,7 @@ impl InProcessClientHandle { /// the runtime is shut down and an `InvalidData` error is returned. pub async fn start(args: InProcessStartArgs) -> IoResult { let initialize = args.initialize.clone(); - let client = start_uninitialized(args); + let client = start_uninitialized(args).await; let initialize_response = client .request(ClientRequest::Initialize { @@ -364,8 +365,12 @@ pub async fn start(args: InProcessStartArgs) -> IoResult Ok(client) } -fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle { +async fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle { let channel_capacity = args.channel_capacity.max(1); + let state_db = match args.state_db.clone() { + Some(state_db) => Some(state_db), + None => init_state_db_from_config(args.config.as_ref()).await, + }; let (client_tx, mut client_rx) = mpsc::channel::(channel_capacity); let (event_tx, event_rx) = mpsc::channel::(channel_capacity); @@ -414,6 +419,12 @@ fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle { ); let (processor_tx, mut processor_rx) = mpsc::channel::(channel_capacity); let mut processor_handle = tokio::spawn(async move { + let Some(state_db) = state_db else { + warn!( + "in-process app-server state db initialization failed; shutting down processor task" + ); + return; + }; let processor = Arc::new(MessageProcessor::new(MessageProcessorArgs { outgoing: Arc::clone(&processor_outgoing), analytics_events_client, @@ -423,7 +434,7 @@ fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle { environment_manager: args.environment_manager, feedback: args.feedback, log_db: args.log_db, - state_db: args.state_db, + state_db, config_warnings: args.config_warnings, session_source: args.session_source, auth_manager, @@ -761,7 +772,7 @@ mod tests { ) -> InProcessClientHandle { let codex_home = TempDir::new().expect("temp dir"); let config = Arc::new(build_test_config(codex_home.path()).await); - let state_db = codex_rollout::state_db::try_init(config.as_ref()) + let state_db = init_state_db_from_config(config.as_ref()) .await .expect("state db should initialize for in-process test"); let args = InProcessStartArgs { @@ -819,7 +830,7 @@ mod tests { } #[tokio::test] - async fn in_process_allows_device_key_requests_to_reach_device_key_processor() { + async fn in_process_allows_device_key_requests_to_reach_device_key_api() { let client = start_test_client(SessionSource::Cli).await; const MALFORMED_KEY_ID_MESSAGE: &str = concat!( "invalid device key payload: keyId must be dk_hse_, dk_tpm_, or dk_osn_ ", diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index 4b1723d88c..4013bbe76b 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -50,11 +50,11 @@ use codex_config::TextRange as CoreTextRange; use codex_core::ExecPolicyError; use codex_core::check_execpolicy_for_warnings; use codex_core::config::find_codex_home; +use codex_core::init_state_db_from_config; use codex_exec_server::EnvironmentManager; use codex_exec_server::ExecServerRuntimePaths; use codex_feedback::CodexFeedback; use codex_protocol::protocol::SessionSource; -use codex_rollout::state_db as rollout_state_db; use codex_state::log_db; use tokio::sync::mpsc; use tokio::sync::oneshot; @@ -487,9 +487,9 @@ pub async fn run_main_with_transport_options( } }; - let state_db_result = rollout_state_db::try_init(&config).await; - let state_db_init_error = state_db_result.as_ref().err().map(ToString::to_string); - let state_db = state_db_result.ok(); + let state_db = init_state_db_from_config(&config) + .await + .ok_or_else(|| std::io::Error::other("failed to initialize sqlite state db"))?; if should_run_personality_migration { let effective_toml = config.config_layer_stack.effective_config(); @@ -598,10 +598,12 @@ pub async fn run_main_with_transport_options( let feedback_layer = feedback.logger_layer(); let feedback_metadata_layer = feedback.metadata_layer(); - let log_db = state_db.clone().map(log_db::start); - let log_db_layer = log_db - .clone() - .map(|layer| layer.with_filter(Targets::new().with_default(Level::TRACE))); + let log_db = log_db::start(state_db.clone()); + let log_db_layer = Some( + log_db + .clone() + .with_filter(Targets::new().with_default(Level::TRACE)), + ); let otel_logger_layer = otel.as_ref().and_then(|o| o.logger_layer()); let otel_tracing_layer = otel.as_ref().and_then(|o| o.tracing_layer()); let _ = tracing_subscriber::registry() @@ -618,10 +620,6 @@ pub async fn run_main_with_transport_options( None => error!("{}", warning.summary), } } - if let Some(err) = &state_db_init_error { - error!("failed to initialize sqlite state db: {err}"); - } - let transport_shutdown_token = CancellationToken::new(); let mut transport_accept_handles = Vec::>::new(); @@ -666,25 +664,17 @@ pub async fn run_main_with_transport_options( let auth_manager = AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false).await; - let remote_control_config_enabled = config.features.enabled(Feature::RemoteControl); - let remote_control_enabled = remote_control_config_enabled && state_db.is_some(); - if remote_control_config_enabled && state_db.is_none() { - error!("remote control disabled because sqlite state db is unavailable"); - } + let remote_control_enabled = config.features.enabled(Feature::RemoteControl); if transport_accept_handles.is_empty() && !remote_control_enabled { return Err(std::io::Error::new( ErrorKind::InvalidInput, - if remote_control_config_enabled && state_db.is_none() { - "no transport configured; remote control disabled because sqlite state db is unavailable" - } else { - "no transport configured; use --listen or enable remote control" - }, + "no transport configured; use --listen or enable remote control", )); } let (remote_control_accept_handle, remote_control_handle) = start_remote_control( config.chatgpt_base_url.clone(), - state_db.clone(), + Some(state_db.clone()), auth_manager.clone(), transport_event_tx.clone(), transport_shutdown_token.clone(), @@ -768,7 +758,7 @@ pub async fn run_main_with_transport_options( config_manager, environment_manager, feedback: feedback.clone(), - log_db, + log_db: Some(log_db), state_db: state_db.clone(), config_warnings, session_source, diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 8d26cec70f..4c1d16eeac 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -61,6 +61,7 @@ use codex_app_server_protocol::experimental_required_message; use codex_arg0::Arg0DispatchPaths; use codex_chatgpt::workspace_settings; use codex_core::ThreadManager; +use codex_core::agent_graph_store_from_state_db; use codex_core::config::Config; use codex_core::thread_store_from_config; use codex_exec_server::EnvironmentManager; @@ -254,7 +255,7 @@ pub(crate) struct MessageProcessorArgs { pub(crate) environment_manager: Arc, pub(crate) feedback: CodexFeedback, pub(crate) log_db: Option, - pub(crate) state_db: Option, + pub(crate) state_db: StateDbHandle, pub(crate) config_warnings: Vec, pub(crate) session_source: SessionSource, pub(crate) auth_manager: Arc, @@ -291,14 +292,16 @@ impl MessageProcessor { // affect per-thread behavior, but they must not move newly started, // resumed, or forked threads to a different persistence backend/root. let thread_store = thread_store_from_config(config.as_ref(), state_db.clone()); + let agent_graph_store = agent_graph_store_from_state_db(state_db.clone()); let thread_manager = Arc::new(ThreadManager::new( config.as_ref(), auth_manager.clone(), session_source, environment_manager, Some(analytics_events_client.clone()), - Arc::clone(&thread_store), state_db.clone(), + Arc::clone(&thread_store), + agent_graph_store.clone(), )); thread_manager .plugins_manager() @@ -344,7 +347,7 @@ impl MessageProcessor { Arc::clone(&config), feedback, log_db, - state_db.clone(), + Some(state_db.clone()), ); let git_processor = GitRequestProcessor::new(); let initialize_processor = InitializeRequestProcessor::new( @@ -395,7 +398,7 @@ impl MessageProcessor { thread_watch_manager.clone(), Arc::clone(&thread_list_state_permit), thread_goal_processor.clone(), - state_db.clone(), + Some(state_db.clone()), ); let turn_processor = TurnRequestProcessor::new( auth_manager.clone(), diff --git a/codex-rs/app-server/src/message_processor_tracing_tests.rs b/codex-rs/app-server/src/message_processor_tracing_tests.rs index 45ea709180..67b668bafc 100644 --- a/codex-rs/app-server/src/message_processor_tracing_tests.rs +++ b/codex-rs/app-server/src/message_processor_tracing_tests.rs @@ -32,6 +32,7 @@ use codex_config::CloudRequirementsLoader; use codex_config::LoaderOverrides; use codex_core::config::Config; use codex_core::config::ConfigBuilder; +use codex_core::init_state_db_from_config; use codex_exec_server::EnvironmentManager; use codex_feedback::CodexFeedback; use codex_login::AuthManager; @@ -281,6 +282,9 @@ async fn build_test_processor( outgoing_tx, analytics_events_client.clone(), )); + let state_db = init_state_db_from_config(config.as_ref()) + .await + .expect("tracing test processor requires state db"); let processor = Arc::new(MessageProcessor::new(MessageProcessorArgs { outgoing, analytics_events_client, @@ -290,7 +294,7 @@ async fn build_test_processor( environment_manager: Arc::new(EnvironmentManager::default_for_tests()), feedback: CodexFeedback::new(), log_db: None, - state_db: None, + state_db, config_warnings: Vec::new(), session_source: SessionSource::VSCode, auth_manager, diff --git a/codex-rs/app-server/src/request_processors/device_key_processor.rs b/codex-rs/app-server/src/request_processors/device_key_processor.rs index ea0a96c2af..e1da459a2e 100644 --- a/codex-rs/app-server/src/request_processors/device_key_processor.rs +++ b/codex-rs/app-server/src/request_processors/device_key_processor.rs @@ -33,8 +33,8 @@ use codex_device_key::RemoteControlClientConnectionAudience; use codex_device_key::RemoteControlClientConnectionSignPayload; use codex_device_key::RemoteControlClientEnrollmentAudience; use codex_device_key::RemoteControlClientEnrollmentSignPayload; +use codex_rollout::state_db::StateDbHandle; use codex_state::DeviceKeyBindingRecord; -use codex_state::StateRuntime; #[derive(Clone)] pub(crate) struct DeviceKeyRequestProcessor { @@ -43,10 +43,7 @@ pub(crate) struct DeviceKeyRequestProcessor { } impl DeviceKeyRequestProcessor { - pub(crate) fn new( - outgoing: Arc, - state_db: Option>, - ) -> Self { + pub(crate) fn new(outgoing: Arc, state_db: StateDbHandle) -> Self { Self { outgoing, store: DeviceKeyStore::new(Arc::new(StateDeviceKeyBindingStore::new(state_db))), @@ -170,25 +167,18 @@ async fn sign_device_key( } struct StateDeviceKeyBindingStore { - state_db: Option>, + state_db: StateDbHandle, } impl StateDeviceKeyBindingStore { - fn new(state_db: Option>) -> Self { + fn new(state_db: StateDbHandle) -> Self { Self { state_db } } - - async fn state_db(&self) -> Result, DeviceKeyError> { - self.state_db - .clone() - .ok_or_else(|| DeviceKeyError::Platform("sqlite state db unavailable".to_string())) - } } impl fmt::Debug for StateDeviceKeyBindingStore { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("StateDeviceKeyBindingStore") - .field("has_state_db", &self.state_db.is_some()) .finish_non_exhaustive() } } @@ -196,7 +186,7 @@ impl fmt::Debug for StateDeviceKeyBindingStore { #[async_trait] impl DeviceKeyBindingStore for StateDeviceKeyBindingStore { async fn get_binding(&self, key_id: &str) -> Result, DeviceKeyError> { - let state_db = self.state_db().await?; + let state_db = self.state_db.clone(); state_db .get_device_key_binding(key_id) .await @@ -214,7 +204,7 @@ impl DeviceKeyBindingStore for StateDeviceKeyBindingStore { key_id: &str, binding: &DeviceKeyBinding, ) -> Result<(), DeviceKeyError> { - let state_db = self.state_db().await?; + let state_db = self.state_db.clone(); state_db .upsert_device_key_binding(&DeviceKeyBindingRecord { key_id: key_id.to_string(), diff --git a/codex-rs/app-server/src/request_processors/thread_goal_processor.rs b/codex-rs/app-server/src/request_processors/thread_goal_processor.rs index 0e12e44ce5..ff192e34b2 100644 --- a/codex-rs/app-server/src/request_processors/thread_goal_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_goal_processor.rs @@ -7,7 +7,7 @@ pub(crate) struct ThreadGoalRequestProcessor { outgoing: Arc, config: Arc, thread_state_manager: ThreadStateManager, - state_db: Option, + state_db: StateDbHandle, } impl ThreadGoalRequestProcessor { @@ -16,7 +16,7 @@ impl ThreadGoalRequestProcessor { outgoing: Arc, config: Arc, thread_state_manager: ThreadStateManager, - state_db: Option, + state_db: StateDbHandle, ) -> Self { Self { thread_manager, @@ -72,23 +72,6 @@ impl ThreadGoalRequestProcessor { } } - pub(crate) async fn pending_resume_goal_state( - &self, - thread: &CodexThread, - ) -> (bool, Option) { - let emit_thread_goal_update = self.config.features.enabled(Feature::Goals); - let thread_goal_state_db = if emit_thread_goal_update { - if let Some(state_db) = thread.state_db() { - Some(state_db) - } else { - self.state_db.clone() - } - } else { - None - }; - (emit_thread_goal_update, thread_goal_state_db) - } - async fn thread_goal_set_inner( &self, request_id: ConnectionRequestId, @@ -110,7 +93,7 @@ impl ThreadGoalRequestProcessor { None => find_thread_path_by_id_str( &self.config.codex_home, &thread_id.to_string(), - self.state_db.as_deref(), + Some(self.state_db.as_ref()), ) .await .map_err(|err| { @@ -275,7 +258,7 @@ impl ThreadGoalRequestProcessor { None => find_thread_path_by_id_str( &self.config.codex_home, &thread_id.to_string(), - self.state_db.as_deref(), + Some(self.state_db.as_ref()), ) .await .map_err(|err| { @@ -339,7 +322,7 @@ impl ThreadGoalRequestProcessor { find_thread_path_by_id_str( &self.config.codex_home, &thread_id.to_string(), - self.state_db.as_deref(), + Some(self.state_db.as_ref()), ) .await .map_err(|err| { @@ -348,9 +331,7 @@ impl ThreadGoalRequestProcessor { .ok_or_else(|| invalid_request(format!("thread not found: {thread_id}")))?; } - self.state_db - .clone() - .ok_or_else(|| internal_error("sqlite state db unavailable for thread goals")) + Ok(self.state_db.clone()) } async fn emit_thread_goal_snapshot(&self, thread_id: ThreadId) { diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index 1921fb69b1..fd61182788 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -2611,10 +2611,10 @@ impl ThreadRequestProcessor { ))); }; - let (emit_thread_goal_update, thread_goal_state_db) = self - .thread_goal_processor - .pending_resume_goal_state(existing_thread.as_ref()) - .await; + let emit_thread_goal_update = self.config.features.enabled(Feature::Goals); + let thread_goal_state_db = emit_thread_goal_update + .then(|| self.state_db.clone()) + .flatten(); let command = crate::thread_state::ThreadListenerCommand::SendThreadResumeResponse( Box::new(crate::thread_state::PendingThreadResumeRequest { diff --git a/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs b/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs index b04eb12a45..9363217864 100644 --- a/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs +++ b/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs @@ -42,6 +42,7 @@ use codex_core::config::ConfigBuilder; use codex_exec_server::EnvironmentManager; use codex_feedback::CodexFeedback; use codex_protocol::protocol::SessionSource; +use codex_state::StateRuntime; use codex_thread_store::InMemoryThreadStore; use pretty_assertions::assert_eq; use tempfile::TempDir; @@ -67,6 +68,13 @@ async fn thread_start_with_non_local_thread_store_does_not_create_local_persiste .loader_overrides(loader_overrides.clone()) .build() .await?; + let sqlite_home = TempDir::new()?; + let state_db = StateRuntime::init( + sqlite_home.path().to_path_buf(), + config.model_provider_id.clone(), + ) + .await + .expect("remote thread store regression test should initialize state db"); let thread_store = InMemoryThreadStore::for_id(store_id.clone()); let _in_memory_store = InMemoryThreadStoreId { store_id }; @@ -80,7 +88,7 @@ async fn thread_start_with_non_local_thread_store_does_not_create_local_persiste thread_config_loader: Arc::new(NoopThreadConfigLoader), feedback: CodexFeedback::new(), log_db: None, - state_db: None, + state_db: Some(state_db), environment_manager: Arc::new(EnvironmentManager::default_for_tests()), config_warnings: Vec::new(), session_source: SessionSource::Cli, diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index 48b6886382..5b5ad7a2a3 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -1387,7 +1387,7 @@ async fn run_debug_prompt_input_command( }); } - let prompt_input = codex_core::build_prompt_input(config, input, /*state_db*/ None).await?; + let prompt_input = codex_core::build_prompt_input(config, input).await?; println!("{}", serde_json::to_string_pretty(&prompt_input)?); Ok(()) diff --git a/codex-rs/core-api/src/lib.rs b/codex-rs/core-api/src/lib.rs index e884cfae9b..13c87df9e0 100644 --- a/codex-rs/core-api/src/lib.rs +++ b/codex-rs/core-api/src/lib.rs @@ -31,6 +31,7 @@ pub use codex_core::StartThreadOptions; pub use codex_core::StateDbHandle; pub use codex_core::ThreadManager; pub use codex_core::ThreadShutdownReport; +pub use codex_core::agent_graph_store_from_state_db; pub use codex_core::config::Config; pub use codex_core::config::Constrained; pub use codex_core::config::GhostSnapshotConfig; @@ -40,6 +41,7 @@ pub use codex_core::config::TerminalResizeReflowConfig; pub use codex_core::config::ThreadStoreConfig; pub use codex_core::config::find_codex_home; pub use codex_core::init_state_db; +pub use codex_core::init_state_db_from_config; pub use codex_core::skills::SkillsManager; pub use codex_core::thread_store_from_config; pub use codex_exec_server::EnvironmentManager; diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index 44c6aacac5..49ed304623 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -26,6 +26,7 @@ bm25 = { workspace = true } chrono = { workspace = true, features = ["serde"] } clap = { workspace = true, features = ["derive"] } codex-analytics = { workspace = true } +codex-agent-graph-store = { workspace = true } codex-api = { workspace = true } codex-app-server-protocol = { workspace = true } codex-apply-patch = { workspace = true } diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 1a13c0b00c..3917ad5e95 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -29,7 +29,6 @@ use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::user_input::UserInput; -use codex_state::DirectionalThreadSpawnEdgeStatus; use codex_thread_store::ReadThreadParams; use serde::Serialize; use std::collections::HashMap; @@ -296,7 +295,6 @@ impl AgentControl { state.notify_thread_created(new_thread.thread_id); self.persist_thread_spawn_edge_for_source( - new_thread.thread.as_ref(), new_thread.thread_id, notification_source.as_ref(), ) @@ -449,19 +447,14 @@ impl AgentControl { )) .await?; let state = self.upgrade()?; - let Ok(resumed_thread) = state.get_thread(resumed_thread_id).await else { - return Ok(resumed_thread_id); - }; - let Some(state_db_ctx) = resumed_thread.state_db() else { - return Ok(resumed_thread_id); - }; + let agent_graph_store = state.agent_graph_store(); let mut resume_queue = VecDeque::from([(thread_id, root_depth)]); while let Some((parent_thread_id, parent_depth)) = resume_queue.pop_front() { - let child_ids = match state_db_ctx - .list_thread_spawn_children_with_status( + let child_ids = match agent_graph_store + .list_thread_spawn_children( parent_thread_id, - DirectionalThreadSpawnEdgeStatus::Open, + Some(codex_agent_graph_store::ThreadSpawnEdgeStatus::Open), ) .await { @@ -525,7 +518,6 @@ impl AgentControl { let _ = config.features.disable(Feature::Collab); } let state = self.upgrade()?; - let state_db_ctx = state.state_db(); let mut reservation = self.state.reserve_spawn_slot(config.agent_max_threads)?; let (session_source, agent_metadata) = match session_source { SessionSource::SubAgent(SubAgentSource::ThreadSpawn { @@ -535,14 +527,11 @@ impl AgentControl { agent_role: _, agent_nickname: _, }) => { + let state_db_ctx = state.state_db(); let (resumed_agent_nickname, resumed_agent_role) = - if let Some(state_db_ctx) = state_db_ctx.as_ref() { - match state_db_ctx.get_thread(thread_id).await { - Ok(Some(metadata)) => (metadata.agent_nickname, metadata.agent_role), - Ok(None) | Err(_) => (None, None), - } - } else { - (None, None) + match state_db_ctx.get_thread(thread_id).await { + Ok(Some(metadata)) => (metadata.agent_nickname, metadata.agent_role), + Ok(None) | Err(_) => (None, None), }; self.prepare_thread_spawn( &mut reservation, @@ -609,7 +598,6 @@ impl AgentControl { ); } self.persist_thread_spawn_edge_for_source( - resumed_thread.thread.as_ref(), resumed_thread.thread_id, Some(¬ification_source), ) @@ -722,11 +710,13 @@ impl AgentControl { /// agent and any live descendants reached from the in-memory tree. pub(crate) async fn close_agent(&self, agent_id: ThreadId) -> CodexResult { let state = self.upgrade()?; - if let Ok(thread) = state.get_thread(agent_id).await - && let Some(state_db_ctx) = thread.state_db() - && let Err(err) = state_db_ctx - .set_thread_spawn_edge_status(agent_id, DirectionalThreadSpawnEdgeStatus::Closed) - .await + if let Err(err) = state + .agent_graph_store() + .set_thread_spawn_edge_status( + agent_id, + codex_agent_graph_store::ThreadSpawnEdgeStatus::Closed, + ) + .await { warn!("failed to persist thread-spawn edge status for {agent_id}: {err}"); } @@ -1142,21 +1132,21 @@ impl AgentControl { async fn persist_thread_spawn_edge_for_source( &self, - thread: &crate::CodexThread, child_thread_id: ThreadId, session_source: Option<&SessionSource>, ) { let Some(parent_thread_id) = session_source.and_then(thread_spawn_parent_thread_id) else { return; }; - let Some(state_db_ctx) = thread.state_db() else { + let Ok(state) = self.upgrade() else { return; }; - if let Err(err) = state_db_ctx + if let Err(err) = state + .agent_graph_store() .upsert_thread_spawn_edge( parent_thread_id, child_thread_id, - DirectionalThreadSpawnEdgeStatus::Open, + codex_agent_graph_store::ThreadSpawnEdgeStatus::Open, ) .await { diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index b95aad4489..a8de427814 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -1,6 +1,5 @@ use super::*; use crate::CodexThread; -use crate::StateDbHandle; use crate::ThreadManager; use crate::agent::agent_status_from_event; use crate::config::AgentRoleConfig; @@ -8,7 +7,6 @@ use crate::config::Config; use crate::config::ConfigBuilder; use crate::context::ContextualUserFragment; use crate::context::SubagentNotification; -use crate::init_state_db; use assert_matches::assert_matches; use codex_features::Feature; use codex_login::CodexAuth; @@ -86,7 +84,6 @@ fn spawn_agent_call(call_id: &str) -> ResponseItem { struct AgentControlHarness { _home: TempDir, config: Config, - state_db: Option, manager: ThreadManager, control: AgentControl, } @@ -94,19 +91,17 @@ struct AgentControlHarness { impl AgentControlHarness { async fn new() -> Self { let (home, config) = test_config().await; - let state_db = init_state_db(&config).await; - let manager = ThreadManager::with_models_provider_home_and_state_for_tests( + let manager = ThreadManager::with_models_provider_and_home_for_tests( CodexAuth::from_api_key("dummy"), config.model_provider.clone(), config.codex_home.to_path_buf(), std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - state_db.clone(), - ); + ) + .await; let control = manager.agent_control(); Self { _home: home, config, - state_db, manager, control, } @@ -955,7 +950,8 @@ async fn spawn_agent_respects_max_threads_limit() { config.model_provider.clone(), config.codex_home.to_path_buf(), std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - ); + ) + .await; let control = manager.agent_control(); let _ = manager @@ -1007,7 +1003,8 @@ async fn spawn_agent_releases_slot_after_shutdown() { config.model_provider.clone(), config.codex_home.to_path_buf(), std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - ); + ) + .await; let control = manager.agent_control(); let first_agent_id = control @@ -1050,7 +1047,8 @@ async fn spawn_agent_limit_shared_across_clones() { config.model_provider.clone(), config.codex_home.to_path_buf(), std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - ); + ) + .await; let control = manager.agent_control(); let cloned = control.clone(); @@ -1095,7 +1093,8 @@ async fn resume_agent_respects_max_threads_limit() { config.model_provider.clone(), config.codex_home.to_path_buf(), std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - ); + ) + .await; let control = manager.agent_control(); let resumable_id = control @@ -1151,7 +1150,8 @@ async fn resume_agent_releases_slot_after_resume_failure() { config.model_provider.clone(), config.codex_home.to_path_buf(), std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - ); + ) + .await; let control = manager.agent_control(); let _ = control @@ -1543,19 +1543,17 @@ async fn resume_thread_subagent_restores_stored_nickname_and_role() { .features .enable(Feature::Sqlite) .expect("test config should allow sqlite"); - let state_db = init_state_db(&config).await; - let manager = ThreadManager::with_models_provider_home_and_state_for_tests( + let manager = ThreadManager::with_models_provider_and_home_for_tests( CodexAuth::from_api_key("dummy"), config.model_provider.clone(), config.codex_home.to_path_buf(), std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - state_db.clone(), - ); + ) + .await; let control = manager.agent_control(); let harness = AgentControlHarness { _home: home, config, - state_db, manager, control, }; @@ -1706,7 +1704,12 @@ async fn resume_agent_from_rollout_reads_archived_rollout_path() { .expect("child shutdown should succeed"); let store = LocalThreadStore::new( LocalThreadStoreConfig::from_config(&harness.config), - harness.state_db.clone(), + codex_state::StateRuntime::init( + harness.config.sqlite_home.clone(), + harness.config.model_provider_id.clone(), + ) + .await + .expect("state db should initialize"), ); store .archive_thread(ArchiveThreadParams { diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 1aeb125091..9578fa5726 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -95,6 +95,7 @@ pub(crate) async fn run_codex_thread_interactive( parent_trace: None, environment_selections: parent_ctx.environments.clone(), analytics_events_client: Some(parent_session.services.analytics_events_client.clone()), + state_db: parent_session.services.state_db.clone(), thread_store: Arc::clone(&parent_session.services.thread_store), })) .or_cancel(&cancel_token) diff --git a/codex-rs/core/src/goals.rs b/codex-rs/core/src/goals.rs index 3cea25f9b8..4af73e6a01 100644 --- a/codex-rs/core/src/goals.rs +++ b/codex-rs/core/src/goals.rs @@ -29,7 +29,6 @@ use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::validate_thread_goal_objective; use codex_rollout::state_db::reconcile_rollout; -use codex_thread_store::LocalThreadStore; use codex_utils_template::Template; use futures::future::BoxFuture; use std::sync::Arc; @@ -1338,17 +1337,6 @@ impl Session { state_db } else if let Some(state_db) = self.goal_runtime.state_db.lock().await.clone() { state_db - } else if let Some(local_store) = self - .services - .thread_store - .as_any() - .downcast_ref::() - { - local_store.state_db().await.ok_or_else(|| { - anyhow::anyhow!( - "thread goals require a local persisted thread with a state database" - ) - })? } else { anyhow::bail!("thread goals require a local persisted thread with a state database"); }; diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index d8a3fbcf0f..cbada9a26c 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -124,7 +124,9 @@ pub use thread_manager::NewThread; pub use thread_manager::StartThreadOptions; pub use thread_manager::ThreadManager; pub use thread_manager::ThreadShutdownReport; +pub use thread_manager::agent_graph_store_from_state_db; pub use thread_manager::build_models_manager; +pub use thread_manager::init_state_db_from_config; pub use thread_manager::thread_store_from_config; pub use web_search::web_search_action_detail; pub use web_search::web_search_detail; diff --git a/codex-rs/core/src/personality_migration.rs b/codex-rs/core/src/personality_migration.rs index 975aecd4af..bb1cb43587 100644 --- a/codex-rs/core/src/personality_migration.rs +++ b/codex-rs/core/src/personality_migration.rs @@ -25,7 +25,7 @@ pub enum PersonalityMigrationStatus { pub async fn maybe_migrate_personality( codex_home: &Path, config_toml: &ConfigToml, - state_db: Option, + state_db: StateDbHandle, ) -> io::Result { let marker_path = codex_home.join(PERSONALITY_MIGRATION_FILENAME); if tokio::fs::try_exists(&marker_path).await? { @@ -65,16 +65,13 @@ pub async fn maybe_migrate_personality( async fn has_recorded_sessions( codex_home: &Path, default_provider: &str, - state_db: Option, + state_db: StateDbHandle, ) -> io::Result { - let store = LocalThreadStore::new( - LocalThreadStoreConfig { - codex_home: codex_home.to_path_buf(), - sqlite_home: codex_home.to_path_buf(), - default_model_provider_id: default_provider.to_string(), - }, - state_db, - ); + let config = LocalThreadStoreConfig { + codex_home: codex_home.to_path_buf(), + default_model_provider_id: default_provider.to_string(), + }; + let store = LocalThreadStore::new(config, state_db); if has_threads(&store, /*archived*/ false).await? { return Ok(true); } diff --git a/codex-rs/core/src/personality_migration_tests.rs b/codex-rs/core/src/personality_migration_tests.rs index 28d34bc3c9..63faae418a 100644 --- a/codex-rs/core/src/personality_migration_tests.rs +++ b/codex-rs/core/src/personality_migration_tests.rs @@ -8,7 +8,10 @@ use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::UserMessageEvent; use codex_rollout::ARCHIVED_SESSIONS_SUBDIR; +use codex_rollout::RolloutConfig; use codex_rollout::SESSIONS_SUBDIR; +use codex_rollout::state_db::StateDbHandle; +use codex_state::state_db_path; use pretty_assertions::assert_eq; use tempfile::TempDir; use tokio::io::AsyncWriteExt; @@ -20,6 +23,26 @@ async fn read_config_toml(codex_home: &Path) -> io::Result { toml::from_str(&contents).map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err)) } +async fn state_db_for_test(codex_home: &Path) -> io::Result { + state_db_for_test_with_sqlite_home(codex_home, codex_home).await +} + +async fn state_db_for_test_with_sqlite_home( + codex_home: &Path, + sqlite_home: &Path, +) -> io::Result { + let config = RolloutConfig { + codex_home: codex_home.to_path_buf(), + sqlite_home: sqlite_home.to_path_buf(), + cwd: codex_home.to_path_buf(), + model_provider_id: "openai".to_string(), + generate_memories: false, + }; + codex_rollout::state_db::try_init(&config) + .await + .map_err(io::Error::other) +} + async fn write_session_with_user_event(codex_home: &Path) -> io::Result<()> { let thread_id = ThreadId::new(); let dir = codex_home @@ -87,7 +110,8 @@ async fn applies_when_sessions_exist_and_no_personality() -> io::Result<()> { write_session_with_user_event(temp.path()).await?; let config_toml = ConfigToml::default(); - let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?; + let state_db = state_db_for_test(temp.path()).await?; + let status = maybe_migrate_personality(temp.path(), &config_toml, state_db).await?; assert_eq!(status, PersonalityMigrationStatus::Applied); assert!(temp.path().join(PERSONALITY_MIGRATION_FILENAME).exists()); @@ -103,7 +127,8 @@ async fn applies_when_only_archived_sessions_exist_and_no_personality() -> io::R write_archived_session_with_user_event(temp.path()).await?; let config_toml = ConfigToml::default(); - let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?; + let state_db = state_db_for_test(temp.path()).await?; + let status = maybe_migrate_personality(temp.path(), &config_toml, state_db).await?; assert_eq!(status, PersonalityMigrationStatus::Applied); assert!(temp.path().join(PERSONALITY_MIGRATION_FILENAME).exists()); @@ -119,7 +144,8 @@ async fn skips_when_marker_exists() -> io::Result<()> { create_marker(&temp.path().join(PERSONALITY_MIGRATION_FILENAME)).await?; let config_toml = ConfigToml::default(); - let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?; + let state_db = state_db_for_test(temp.path()).await?; + let status = maybe_migrate_personality(temp.path(), &config_toml, state_db).await?; assert_eq!(status, PersonalityMigrationStatus::SkippedMarker); assert!(!temp.path().join("config.toml").exists()); @@ -136,7 +162,8 @@ async fn skips_when_personality_explicit() -> io::Result<()> { .map_err(|err| io::Error::other(format!("failed to write config: {err}")))?; let config_toml = read_config_toml(temp.path()).await?; - let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?; + let state_db = state_db_for_test(temp.path()).await?; + let status = maybe_migrate_personality(temp.path(), &config_toml, state_db).await?; assert_eq!( status, @@ -153,10 +180,37 @@ async fn skips_when_personality_explicit() -> io::Result<()> { async fn skips_when_no_sessions() -> io::Result<()> { let temp = TempDir::new()?; let config_toml = ConfigToml::default(); - let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?; + let state_db = state_db_for_test(temp.path()).await?; + let status = maybe_migrate_personality(temp.path(), &config_toml, state_db).await?; assert_eq!(status, PersonalityMigrationStatus::SkippedNoSessions); assert!(temp.path().join(PERSONALITY_MIGRATION_FILENAME).exists()); assert!(!temp.path().join("config.toml").exists()); Ok(()) } + +#[tokio::test] +async fn uses_configured_sqlite_home_when_checking_for_sessions() -> io::Result<()> { + let codex_home = TempDir::new()?; + let sqlite_home = TempDir::new()?; + write_session_with_user_event(codex_home.path()).await?; + + let config_toml = ConfigToml::default(); + let state_db = + state_db_for_test_with_sqlite_home(codex_home.path(), sqlite_home.path()).await?; + let status = maybe_migrate_personality(codex_home.path(), &config_toml, state_db).await?; + + assert_eq!(status, PersonalityMigrationStatus::Applied); + assert!( + codex_home + .path() + .join(PERSONALITY_MIGRATION_FILENAME) + .exists() + ); + + let persisted = read_config_toml(codex_home.path()).await?; + assert_eq!(persisted.personality, Some(Personality::Pragmatic)); + assert!(!state_db_path(codex_home.path()).exists()); + assert!(state_db_path(sqlite_home.path()).exists()); + Ok(()) +} diff --git a/codex-rs/core/src/prompt_debug.rs b/codex-rs/core/src/prompt_debug.rs index 7c6144c10e..a40fd27942 100644 --- a/codex-rs/core/src/prompt_debug.rs +++ b/codex-rs/core/src/prompt_debug.rs @@ -16,8 +16,9 @@ use crate::config::Config; use crate::session::session::Session; use crate::session::turn::build_prompt; use crate::session::turn::built_tools; -use crate::state_db_bridge::StateDbHandle; use crate::thread_manager::ThreadManager; +use crate::thread_manager::agent_graph_store_from_state_db; +use crate::thread_manager::init_state_db_from_config; use crate::thread_manager::thread_store_from_config; /// Build the model-visible `input` list for a single debug turn. @@ -25,7 +26,6 @@ use crate::thread_manager::thread_store_from_config; pub async fn build_prompt_input( mut config: Config, input: Vec, - state_db: Option, ) -> CodexResult> { config.ephemeral = true; @@ -37,15 +37,20 @@ pub async fn build_prompt_input( config.codex_linux_sandbox_exe.clone(), )?; + let state_db = init_state_db_from_config(&config) + .await + .ok_or_else(|| std::io::Error::other("prompt debug requires state db"))?; let thread_store = thread_store_from_config(&config, state_db.clone()); + let agent_graph_store = agent_graph_store_from_state_db(state_db.clone()); let thread_manager = ThreadManager::new( &config, Arc::clone(&auth_manager), SessionSource::Exec, Arc::new(EnvironmentManager::new(EnvironmentManagerArgs::new(local_runtime_paths)).await), /*analytics_events_client*/ None, + state_db, thread_store, - state_db.clone(), + agent_graph_store, ); let thread = thread_manager.start_thread(config).await?; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index d35bccb96e..85f2ef438b 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -132,7 +132,6 @@ use codex_terminal_detection::user_agent; use codex_thread_store::CreateThreadParams; use codex_thread_store::LiveThread; use codex_thread_store::LiveThreadInitGuard; -use codex_thread_store::LocalThreadStore; use codex_thread_store::ResumeThreadParams; use codex_thread_store::ThreadEventPersistenceMode; use codex_thread_store::ThreadPersistenceMetadata; @@ -356,7 +355,6 @@ use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TokenUsageInfo; use codex_protocol::protocol::WarningEvent; use codex_protocol::user_input::UserInput; -use codex_tools::ToolEnvironmentMode; use codex_tools::ToolsConfig; use codex_tools::ToolsConfigParams; use codex_utils_absolute_path::AbsolutePathBuf; @@ -412,6 +410,7 @@ pub(crate) struct CodexSpawnArgs { pub(crate) parent_trace: Option, pub(crate) environment_selections: ResolvedTurnEnvironments, pub(crate) analytics_events_client: Option, + pub(crate) state_db: Option, pub(crate) thread_store: Arc, } @@ -469,6 +468,7 @@ impl Codex { parent_trace: _, environment_selections, analytics_events_client, + state_db, thread_store, } = args; let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY); @@ -557,15 +557,7 @@ impl Codex { }; match thread_id { Some(thread_id) => { - let state_db_ctx = if config.ephemeral { - None - } else if let Some(local_store) = - thread_store.as_any().downcast_ref::() - { - local_store.state_db().await - } else { - None - }; + let state_db_ctx = state_db.clone(); state_db::get_dynamic_tools(state_db_ctx.as_deref(), thread_id, "codex_spawn") .await } @@ -651,6 +643,7 @@ impl Codex { agent_control, environment_manager, analytics_events_client, + state_db, thread_store, parent_rollout_thread_trace, ) @@ -1311,7 +1304,7 @@ impl Session { self.services.user_shell.as_ref().clone(), self.services.shell_snapshot_tx.clone(), self.services.session_telemetry.clone(), - self.services.state_db.clone(), + self.state_db(), ); } diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 0bb8da6ea6..23621f87cf 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -343,6 +343,7 @@ impl Session { agent_control: AgentControl, environment_manager: Arc, analytics_events_client: Option, + state_db: Option, thread_store: Arc, parent_rollout_thread_trace: ThreadTraceContext, ) -> anyhow::Result> { @@ -441,22 +442,7 @@ impl Session { otel.name = "session_init.thread_persistence", session_init.ephemeral = config.ephemeral, )); - let state_db_fut = async { - if config.ephemeral { - None - } else if let Some(local_store) = - thread_store.as_any().downcast_ref::() - { - local_store.state_db().await - } else { - None - } - } - .instrument(info_span!( - "session_init.state_db", - otel.name = "session_init.state_db", - session_init.ephemeral = config.ephemeral, - )); + let state_db_ctx = if config.ephemeral { None } else { state_db }; let is_subagent = session_configuration.session_source.is_non_root_agent(); let history_meta_fut = async { @@ -495,15 +481,9 @@ impl Session { // Join all independent futures. let ( thread_persistence_result, - state_db_ctx, (history_log_id, history_entry_count), (auth, mcp_servers, auth_statuses), - ) = tokio::join!( - thread_persistence_fut, - state_db_fut, - history_meta_fut, - auth_and_mcp_fut - ); + ) = tokio::join!(thread_persistence_fut, history_meta_fut, auth_and_mcp_fut); let mut live_thread_init = LiveThreadInitGuard::new(thread_persistence_result.map_err(|e| { diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 6995ef8dc9..f3a8d95107 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -867,7 +867,7 @@ async fn danger_full_access_turns_do_not_expose_managed_network_proxy() -> anyho &permission_profile_for_sandbox_policy(&SandboxPolicy::DangerFullAccess), )?; - let session = make_session_with_config(move |config| { + let (session, _codex_home) = make_session_with_config(move |config| { let cwd = config.cwd.clone(); config .permissions @@ -933,7 +933,7 @@ async fn danger_full_access_tool_attempts_do_not_enforce_managed_network() -> an &permission_profile_for_sandbox_policy(&SandboxPolicy::DangerFullAccess), )?; - let session = make_session_with_config(move |config| { + let (session, _codex_home) = make_session_with_config(move |config| { let cwd = config.cwd.clone(); config .permissions @@ -1008,7 +1008,7 @@ async fn workspace_write_turns_continue_to_expose_managed_network_proxy() -> any &permission_profile_for_sandbox_policy(&sandbox_policy), )?; - let session = make_session_with_config(move |config| { + let (session, _codex_home) = make_session_with_config(move |config| { let cwd = config.cwd.clone(); config .permissions @@ -1035,7 +1035,7 @@ async fn user_shell_commands_do_not_inherit_managed_network_proxy() -> anyhow::R &permission_profile_for_sandbox_policy(&sandbox_policy), )?; - let (session, rx) = make_session_with_config_and_rx(move |config| { + let (session, rx, _codex_home) = make_session_with_config_and_rx(move |config| { let cwd = config.cwd.clone(); config .permissions @@ -1165,7 +1165,7 @@ async fn reload_user_config_layer_updates_effective_apps_config() { #[tokio::test] async fn reload_user_config_layer_refreshes_hooks() -> anyhow::Result<()> { - let session = make_session_with_config(|config| { + let (session, _codex_home) = make_session_with_config(|config| { config .features .enable(Feature::CodexHooks) @@ -2982,15 +2982,13 @@ pub(crate) async fn make_session_configuration_for_tests() -> SessionConfigurati fn turn_environments_for_tests( environment: &Arc, cwd: &codex_utils_absolute_path::AbsolutePathBuf, -) -> crate::environment_selection::ResolvedTurnEnvironments { - crate::environment_selection::ResolvedTurnEnvironments { - turn_environments: vec![TurnEnvironment { - environment_id: codex_exec_server::LOCAL_ENVIRONMENT_ID.to_string(), - environment: Arc::clone(environment), - cwd: cwd.clone(), - shell: "bash".to_string(), - }], - } +) -> Vec { + vec![TurnEnvironment { + environment_id: codex_exec_server::LOCAL_ENVIRONMENT_ID.to_string(), + environment: Arc::clone(environment), + cwd: cwd.clone(), + shell: "bash".to_string(), + }] } #[tokio::test] @@ -3528,9 +3526,15 @@ async fn session_new_fails_when_zsh_fork_enabled_without_zsh_path() { AgentControl::default(), Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, + /*state_db*/ None, Arc::new(codex_thread_store::LocalThreadStore::new( codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()), - /*state_db*/ None, + codex_state::StateRuntime::init( + config.sqlite_home.clone(), + config.model_provider_id.clone(), + ) + .await + .expect("state db should initialize"), )), codex_rollout_trace::ThreadTraceContext::disabled(), ) @@ -3678,7 +3682,12 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { live_thread: None, thread_store: Arc::new(codex_thread_store::LocalThreadStore::new( codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()), - /*state_db*/ None, + codex_state::StateRuntime::init( + config.sqlite_home.clone(), + config.model_provider_id.clone(), + ) + .await + .expect("state db should initialize"), )), model_client: ModelClient::new( Some(auth_manager.clone()), @@ -3723,7 +3732,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { model_info, &models_manager, /*network*/ None, - turn_environments, + crate::environment_selection::ResolvedTurnEnvironments { turn_environments }, session_configuration.cwd.clone(), "turn_id".to_string(), skills_outcome, @@ -3756,14 +3765,18 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { async fn make_session_with_config( mutator: impl FnOnce(&mut Config), -) -> anyhow::Result> { - let (session, _rx_event) = make_session_with_config_and_rx(mutator).await?; - Ok(session) +) -> anyhow::Result<(Arc, tempfile::TempDir)> { + let (session, _rx_event, codex_home) = make_session_with_config_and_rx(mutator).await?; + Ok((session, codex_home)) } async fn make_session_with_config_and_rx( mutator: impl FnOnce(&mut Config), -) -> anyhow::Result<(Arc, async_channel::Receiver)> { +) -> anyhow::Result<( + Arc, + async_channel::Receiver, + tempfile::TempDir, +)> { let codex_home = tempfile::tempdir().expect("create temp dir"); let mut config = build_test_config(codex_home.path()).await; mutator(&mut config); @@ -3848,15 +3861,21 @@ async fn make_session_with_config_and_rx( AgentControl::default(), Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, + /*state_db*/ None, Arc::new(codex_thread_store::LocalThreadStore::new( codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()), - /*state_db*/ None, + codex_state::StateRuntime::init( + config.sqlite_home.clone(), + config.model_provider_id.clone(), + ) + .await + .expect("state db should initialize"), )), codex_rollout_trace::ThreadTraceContext::disabled(), ) .await?; - Ok((session, rx_event)) + Ok((session, rx_event, codex_home)) } #[tokio::test] @@ -4518,7 +4537,7 @@ async fn default_turn_honors_empty_stored_thread_environments() { let turn_context = session.new_default_turn().await; - assert!(turn_context.environments.primary().is_none()); + assert!(turn_context.environments.primary_environment().is_none()); assert!(turn_context.environments.turn_environments.is_empty()); assert_eq!(turn_context.cwd, session_cwd); assert_eq!(turn_context.config.cwd, session_cwd); @@ -4537,7 +4556,7 @@ async fn primary_environment_uses_first_turn_environment() { environment_id: "second".to_string(), environment: Arc::clone(&first_environment.environment), cwd: second_cwd.clone(), - shell: first_environment.shell.clone(), + shell: "bash".to_string(), }); assert_eq!( @@ -4580,7 +4599,7 @@ async fn empty_turn_environments_clear_primary_environment() { .await .expect("turn should start"); - assert!(turn_context.environments.primary().is_none()); + assert!(turn_context.environments.primary_environment().is_none()); assert!(turn_context.environments.turn_environments.is_empty()); assert_eq!(turn_context.cwd, session.get_config().await.cwd); assert_eq!(turn_context.config.cwd, session.get_config().await.cwd); @@ -5036,47 +5055,13 @@ async fn make_session_and_context_with_auth_and_config_and_rx( Arc, async_channel::Receiver, ) -where - F: FnOnce(&mut Config), -{ - let codex_home = tempfile::tempdir().expect("create temp dir"); - make_session_and_context_with_auth_config_home_and_rx( - auth, - dynamic_tools, - codex_home.path(), - configure_config, - ) - .await -} - -async fn make_session_and_context_with_auth_config_home_and_rx( - auth: CodexAuth, - dynamic_tools: Vec, - codex_home: &Path, - configure_config: F, -) -> ( - Arc, - Arc, - async_channel::Receiver, -) where F: FnOnce(&mut Config), { let (tx_event, rx_event) = async_channel::unbounded(); - let mut config = build_test_config(codex_home).await; + let codex_home = tempfile::tempdir().expect("create temp dir").keep(); + let mut config = build_test_config(codex_home.as_path()).await; configure_config(&mut config); - let state_db = if config.features.enabled(Feature::Goals) { - Some( - codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.model_provider_id.clone(), - ) - .await - .expect("goal tests should initialize sqlite state db"), - ) - } else { - None - }; let config = Arc::new(config); let conversation_id = ThreadId::default(); let auth_manager = AuthManager::from_auth_for_testing(auth); @@ -5162,6 +5147,12 @@ where .expect("create environment"), ); + let state_db = codex_state::StateRuntime::init( + config.sqlite_home.clone(), + config.model_provider_id.clone(), + ) + .await + .expect("state db should initialize"); let skills_watcher = Arc::new(SkillsWatcher::noop()); let services = SessionServices { mcp_connection_manager: Arc::new(RwLock::new(McpConnectionManager::new_uninitialized( @@ -5202,7 +5193,7 @@ where agent_control, network_proxy: None, network_approval: Arc::clone(&network_approval), - state_db: state_db.clone(), + state_db: Some(state_db.clone()), live_thread: None, thread_store: Arc::new(codex_thread_store::LocalThreadStore::new( codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()), @@ -5251,7 +5242,7 @@ where model_info, &models_manager, /*network*/ None, - turn_environments, + crate::environment_selection::ResolvedTurnEnvironments { turn_environments }, session_configuration.cwd.clone(), "turn_id".to_string(), skills_outcome, @@ -5301,13 +5292,10 @@ async fn make_goal_session_and_context_with_rx() -> ( Arc, Arc, async_channel::Receiver, - tempfile::TempDir, ) { - let codex_home = tempfile::tempdir().expect("create temp dir"); - let (session, turn_context, rx) = make_session_and_context_with_auth_config_home_and_rx( + let (session, turn_context, rx) = make_session_and_context_with_auth_and_config_and_rx( CodexAuth::from_api_key("Test API Key"), Vec::new(), - codex_home.path(), |config| { config .features @@ -5317,14 +5305,14 @@ async fn make_goal_session_and_context_with_rx() -> ( ) .await; upsert_goal_test_thread(session.as_ref()).await; - (session, turn_context, rx, codex_home) + (session, turn_context, rx) } async fn upsert_goal_test_thread(session: &Session) { let config = session.get_config().await; - let state_db = session - .state_db() - .expect("goal test session should have a state db"); + let state_db = goal_test_state_db(session) + .await + .expect("goal test state db should initialize"); let mut builder = codex_state::ThreadMetadataBuilder::new( session.conversation_id, config @@ -7148,7 +7136,7 @@ async fn abort_empty_active_turn_preserves_pending_input() { #[tokio::test] async fn interrupt_accounts_active_goal_before_pausing() -> anyhow::Result<()> { - let (sess, tc, _rx, _codex_home) = make_goal_session_and_context_with_rx().await; + let (sess, tc, _rx) = make_goal_session_and_context_with_rx().await; sess.set_thread_goal( tc.as_ref(), SetGoalRequest { @@ -7412,7 +7400,7 @@ async fn goal_test_state_db(sess: &Session) -> anyhow::Result anyhow::Result<()> { - let (sess, tc, rx, _codex_home) = make_goal_session_and_context_with_rx().await; + let (sess, tc, rx) = make_goal_session_and_context_with_rx().await; sess.set_thread_goal( tc.as_ref(), SetGoalRequest { @@ -7512,7 +7500,7 @@ async fn budget_limited_accounting_steers_active_turn_without_aborting() -> anyh #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn external_goal_mutation_accounts_active_turn_before_status_change() -> anyhow::Result<()> { - let (sess, tc, _rx, _codex_home) = make_goal_session_and_context_with_rx().await; + let (sess, tc, _rx) = make_goal_session_and_context_with_rx().await; sess.set_thread_goal( tc.as_ref(), SetGoalRequest { @@ -7579,7 +7567,7 @@ async fn external_goal_mutation_accounts_active_turn_before_status_change() -> a #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn external_active_goal_set_marks_current_turn_for_accounting() -> anyhow::Result<()> { - let (sess, tc, _rx, _codex_home) = make_goal_session_and_context_with_rx().await; + let (sess, tc, _rx) = make_goal_session_and_context_with_rx().await; sess.spawn_task( Arc::clone(&tc), Vec::new(), @@ -8246,7 +8234,7 @@ async fn sample_rollout( #[tokio::test] async fn create_goal_tool_rejects_existing_goal() { - let (session, turn_context, _rx, _codex_home) = make_goal_session_and_context_with_rx().await; + let (session, turn_context, _rx) = make_goal_session_and_context_with_rx().await; let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); let handler = CreateGoalHandler; @@ -8308,7 +8296,7 @@ async fn create_goal_tool_rejects_existing_goal() { #[tokio::test] async fn update_goal_tool_rejects_pausing_goal() { - let (session, turn_context, _rx, _codex_home) = make_goal_session_and_context_with_rx().await; + let (session, turn_context, _rx) = make_goal_session_and_context_with_rx().await; let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); let create_handler = CreateGoalHandler; let update_handler = UpdateGoalHandler; @@ -8369,7 +8357,7 @@ async fn update_goal_tool_rejects_pausing_goal() { #[tokio::test] async fn update_goal_tool_marks_goal_complete() { - let (session, turn_context, _rx, _codex_home) = make_goal_session_and_context_with_rx().await; + let (session, turn_context, _rx) = make_goal_session_and_context_with_rx().await; let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); let create_handler = CreateGoalHandler; let update_handler = UpdateGoalHandler; diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index 4c61dad18d..d7ec5937d2 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -731,7 +731,12 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { let skills_watcher = Arc::new(SkillsWatcher::noop()); let thread_store = Arc::new(codex_thread_store::LocalThreadStore::new( codex_thread_store::LocalThreadStoreConfig::from_config(&config), - /*state_db*/ None, + codex_state::StateRuntime::init( + config.sqlite_home.clone(), + config.model_provider_id.clone(), + ) + .await + .expect("state db should initialize"), )); let CodexSpawnOk { codex, .. } = Codex::spawn(CodexSpawnArgs { @@ -760,6 +765,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { turn_environments: Vec::new(), }, analytics_events_client: None, + state_db: None, thread_store, }) .await diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index fe54d61546..b4ef59df95 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -8,6 +8,7 @@ use codex_protocol::protocol::TurnEnvironmentSelection; use codex_sandboxing::compatibility_sandbox_policy_for_permission_profile; use codex_sandboxing::policy_transforms::effective_file_system_sandbox_policy; use codex_sandboxing::policy_transforms::effective_network_sandbox_policy; +use codex_tools::ToolEnvironmentMode; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; diff --git a/codex-rs/core/src/stream_events_utils.rs b/codex-rs/core/src/stream_events_utils.rs index 29884da4ae..fc608aed93 100644 --- a/codex-rs/core/src/stream_events_utils.rs +++ b/codex-rs/core/src/stream_events_utils.rs @@ -138,11 +138,8 @@ pub(crate) async fn record_completed_response_item( .await; } mark_thread_memory_mode_polluted_if_external_context(sess, turn_context, item).await; - let has_memory_citation = record_stage1_output_usage_and_detect_memory_citation( - sess.services.state_db.as_ref(), - item, - ) - .await; + let has_memory_citation = + record_stage1_output_usage_and_detect_memory_citation(sess.state_db(), item).await; if has_memory_citation { sess.record_memory_citation_for_turn(&turn_context.sub_id) .await; @@ -177,7 +174,7 @@ pub(crate) async fn mark_thread_memory_mode_polluted_if_external_context( } async fn record_stage1_output_usage_and_detect_memory_citation( - state_db_ctx: Option<&state_db::StateDbHandle>, + state_db_ctx: Option, item: &ResponseItem, ) -> bool { let Some(raw_text) = raw_assistant_output_text_from_item(item) else { diff --git a/codex-rs/core/src/test_support.rs b/codex-rs/core/src/test_support.rs index 48eec66c58..74305fb518 100644 --- a/codex-rs/core/src/test_support.rs +++ b/codex-rs/core/src/test_support.rs @@ -52,14 +52,14 @@ pub fn auth_manager_from_auth_with_home(auth: CodexAuth, codex_home: PathBuf) -> AuthManager::from_auth_for_testing_with_home(auth, codex_home) } -pub fn thread_manager_with_models_provider( +pub async fn thread_manager_with_models_provider( auth: CodexAuth, provider: ModelProviderInfo, ) -> ThreadManager { - ThreadManager::with_models_provider_for_tests(auth, provider) + ThreadManager::with_models_provider_for_tests(auth, provider).await } -pub fn thread_manager_with_models_provider_and_home( +pub async fn thread_manager_with_models_provider_and_home( auth: CodexAuth, provider: ModelProviderInfo, codex_home: PathBuf, @@ -71,22 +71,7 @@ pub fn thread_manager_with_models_provider_and_home( codex_home, environment_manager, ) -} - -pub fn thread_manager_with_models_provider_home_and_state( - auth: CodexAuth, - provider: ModelProviderInfo, - codex_home: PathBuf, - environment_manager: Arc, - state_db: Option, -) -> ThreadManager { - ThreadManager::with_models_provider_home_and_state_for_tests( - auth, - provider, - codex_home, - environment_manager, - state_db, - ) + .await } pub async fn start_thread_with_user_shell_override( diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index d4e36ab3c8..06f8118db7 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -18,6 +18,8 @@ use crate::skills_watcher::SkillsWatcher; use crate::skills_watcher::SkillsWatcherEvent; use crate::tasks::InterruptedTurnHistoryMarker; use crate::tasks::interrupted_turn_history_marker; +use codex_agent_graph_store::AgentGraphStore; +use codex_agent_graph_store::LocalAgentGraphStore; use codex_analytics::AnalyticsEventsClient; use codex_app_server_protocol::ThreadHistoryBuilder; use codex_app_server_protocol::TurnStatus; @@ -50,8 +52,8 @@ use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnAbortedEvent; use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::protocol::W3cTraceContext; +use codex_rollout::state_db; use codex_rollout::state_db::StateDbHandle; -use codex_state::DirectionalThreadSpawnEdgeStatus; use codex_thread_store::InMemoryThreadStore; use codex_thread_store::LocalThreadStore; use codex_thread_store::LocalThreadStoreConfig; @@ -247,9 +249,10 @@ pub(crate) struct ThreadManagerState { mcp_manager: Arc, skills_watcher: Arc, thread_store: Arc, + state_db: StateDbHandle, + agent_graph_store: Arc, session_source: SessionSource, analytics_events_client: Option, - state_db: Option, // Captures submitted ops for testing purpose when test mode is enabled. ops_log: Option, } @@ -265,10 +268,11 @@ pub fn build_models_manager( ) } -pub fn thread_store_from_config( - config: &Config, - state_db: Option, -) -> Arc { +pub async fn init_state_db_from_config(config: &Config) -> Option { + state_db::init(config).await +} + +pub fn thread_store_from_config(config: &Config, state_db: StateDbHandle) -> Arc { match &config.experimental_thread_store { ThreadStoreConfig::Local => Arc::new(LocalThreadStore::new( LocalThreadStoreConfig::from_config(config), @@ -279,15 +283,38 @@ pub fn thread_store_from_config( } } +pub fn agent_graph_store_from_state_db(state_db: StateDbHandle) -> Arc { + Arc::new(LocalAgentGraphStore::new(state_db)) +} + +async fn state_db_from_roots_for_tests( + codex_home: PathBuf, + sqlite_home: PathBuf, + default_model_provider_id: String, +) -> StateDbHandle { + let config = codex_rollout::RolloutConfig { + codex_home: codex_home.clone(), + sqlite_home, + cwd: codex_home, + model_provider_id: default_model_provider_id, + generate_memories: false, + }; + state_db::try_init(&config) + .await + .unwrap_or_else(|err| panic!("test state db should initialize: {err}")) +} + impl ThreadManager { + #[allow(clippy::too_many_arguments)] pub fn new( config: &Config, auth_manager: Arc, session_source: SessionSource, environment_manager: Arc, analytics_events_client: Option, + state_db: StateDbHandle, thread_store: Arc, - state_db: Option, + agent_graph_store: Arc, ) -> Self { let codex_home = config.codex_home.clone(); let restriction_product = session_source.restriction_product(); @@ -314,10 +341,11 @@ impl ThreadManager { mcp_manager, skills_watcher, thread_store, + state_db, + agent_graph_store, auth_manager, session_source, analytics_events_client, - state_db, ops_log: should_use_test_thread_manager_behavior() .then(|| Arc::new(std::sync::Mutex::new(Vec::new()))), }), @@ -327,7 +355,7 @@ impl ThreadManager { /// Construct with a dummy AuthManager containing the provided CodexAuth. /// Used for integration tests: should not be used by ordinary business logic. - pub(crate) fn with_models_provider_for_tests( + pub(crate) async fn with_models_provider_for_tests( auth: CodexAuth, provider: ModelProviderInfo, ) -> Self { @@ -338,11 +366,18 @@ impl ThreadManager { )); std::fs::create_dir_all(&codex_home) .unwrap_or_else(|err| panic!("temp codex home dir create failed: {err}")); - let mut manager = Self::with_models_provider_and_home_for_tests( + let state_db = state_db_from_roots_for_tests( + codex_home.clone(), + codex_home.clone(), + OPENAI_PROVIDER_ID.to_string(), + ) + .await; + let mut manager = Self::with_models_provider_and_home_and_state_db_for_tests( auth, provider, codex_home.clone(), Arc::new(EnvironmentManager::default_for_tests()), + state_db, ); manager._test_codex_home_guard = Some(TempCodexHomeGuard { path: codex_home }); manager @@ -350,27 +385,33 @@ impl ThreadManager { /// Construct with a dummy AuthManager containing the provided CodexAuth and codex home. /// Used for integration tests: should not be used by ordinary business logic. - pub(crate) fn with_models_provider_and_home_for_tests( + pub(crate) async fn with_models_provider_and_home_for_tests( auth: CodexAuth, provider: ModelProviderInfo, codex_home: PathBuf, environment_manager: Arc, ) -> Self { - Self::with_models_provider_home_and_state_for_tests( + let state_db = state_db_from_roots_for_tests( + codex_home.clone(), + codex_home.clone(), + OPENAI_PROVIDER_ID.to_string(), + ) + .await; + Self::with_models_provider_and_home_and_state_db_for_tests( auth, provider, codex_home, environment_manager, - /*state_db*/ None, + state_db, ) } - pub(crate) fn with_models_provider_home_and_state_for_tests( + fn with_models_provider_and_home_and_state_db_for_tests( auth: CodexAuth, provider: ModelProviderInfo, codex_home: PathBuf, environment_manager: Arc, - state_db: Option, + state_db: StateDbHandle, ) -> Self { set_thread_manager_test_mode_for_tests(/*enabled*/ true); let auth_manager = AuthManager::from_auth_for_testing(auth); @@ -396,11 +437,11 @@ impl ThreadManager { let thread_store: Arc = Arc::new(LocalThreadStore::new( LocalThreadStoreConfig { codex_home: codex_home.clone(), - sqlite_home: codex_home.clone(), default_model_provider_id: OPENAI_PROVIDER_ID.to_string(), }, state_db.clone(), )); + let agent_graph_store = agent_graph_store_from_state_db(state_db.clone()); Self { state: Arc::new(ThreadManagerState { threads: Arc::new(RwLock::new(HashMap::new())), @@ -413,10 +454,11 @@ impl ThreadManager { mcp_manager, skills_watcher, thread_store, + state_db, + agent_graph_store, auth_manager, session_source: SessionSource::Exec, analytics_events_client: None, - state_db, ops_log: should_use_test_thread_manager_behavior() .then(|| Arc::new(std::sync::Mutex::new(Vec::new()))), }), @@ -523,22 +565,17 @@ impl ThreadManager { subtree_thread_ids.push(thread_id); seen_thread_ids.insert(thread_id); - if let Some(state_db_ctx) = thread.state_db() { - for status in [ - DirectionalThreadSpawnEdgeStatus::Open, - DirectionalThreadSpawnEdgeStatus::Closed, - ] { - for descendant_id in state_db_ctx - .list_thread_spawn_descendants_with_status(thread_id, status) - .await - .map_err(|err| { - CodexErr::Fatal(format!("failed to load thread-spawn descendants: {err}")) - })? - { - if seen_thread_ids.insert(descendant_id) { - subtree_thread_ids.push(descendant_id); - } - } + for descendant_id in self + .state + .agent_graph_store + .list_thread_spawn_descendants(thread_id, /*status_filter*/ None) + .await + .map_err(|err| { + CodexErr::Fatal(format!("failed to load thread-spawn descendants: {err}")) + })? + { + if seen_thread_ids.insert(descendant_id) { + subtree_thread_ids.push(descendant_id); } } @@ -864,10 +901,14 @@ impl ThreadManager { } impl ThreadManagerState { - pub(crate) fn state_db(&self) -> Option { + pub(crate) fn state_db(&self) -> StateDbHandle { self.state_db.clone() } + pub(crate) fn agent_graph_store(&self) -> Arc { + self.agent_graph_store.clone() + } + pub(crate) async fn list_thread_ids(&self) -> Vec { self.threads .read() @@ -1172,6 +1213,7 @@ impl ThreadManagerState { parent_trace, environment_selections, analytics_events_client: self.analytics_events_client.clone(), + state_db: Some(self.state_db.clone()), thread_store: Arc::clone(&self.thread_store), }) .await?; diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index d50a3ae8f9..643309aac1 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -1,6 +1,5 @@ use super::*; use crate::config::test_config; -use crate::init_state_db; use crate::rollout::RolloutRecorder; use crate::session::session::SessionSettingsUpdate; use crate::session::tests::make_session_and_context; @@ -47,6 +46,21 @@ fn assistant_msg(text: &str) -> ResponseItem { } } +async fn state_backed_stores( + config: &Config, +) -> ( + StateDbHandle, + Arc, + Arc, +) { + let state_db = init_state_db_from_config(config) + .await + .expect("thread manager test requires state db"); + let thread_store = thread_store_from_config(config, state_db.clone()); + let agent_graph_store = agent_graph_store_from_state_db(state_db.clone()); + (state_db, thread_store, agent_graph_store) +} + fn contextual_user_interrupted_marker() -> ResponseItem { interrupted_turn_history_marker(InterruptedTurnHistoryMarker::ContextualUser) .expect("contextual-user interrupted marker should be enabled") @@ -261,7 +275,8 @@ async fn shutdown_all_threads_bounded_submits_shutdown_to_every_thread() { config.model_provider.clone(), config.codex_home.to_path_buf(), Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - ); + ) + .await; let thread_1 = manager .start_thread(config.clone()) .await @@ -310,7 +325,8 @@ async fn start_thread_accepts_explicit_environment_when_default_environment_is_d config.model_provider.clone(), config.codex_home.to_path_buf(), environment_manager, - ); + ) + .await; let thread = manager .start_thread_with_options(StartThreadOptions { @@ -345,7 +361,8 @@ async fn start_thread_keeps_internal_threads_hidden_from_normal_lookups() { config.model_provider.clone(), config.codex_home.to_path_buf(), Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - ); + ) + .await; let thread = manager .start_thread_with_options(StartThreadOptions { config, @@ -384,14 +401,16 @@ async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await; let manager = ThreadManager::new( &config, auth_manager.clone(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, - thread_store_from_config(&config, /*state_db*/ None), - /*state_db*/ None, + state_db, + thread_store, + agent_graph_store, ); let selected_cwd = AbsolutePathBuf::try_from(config.cwd.as_path().join("selected")).expect("absolute path"); @@ -494,14 +513,16 @@ async fn resume_active_thread_from_rollout_returns_running_thread() { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await; let manager = ThreadManager::new( &config, auth_manager.clone(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, - thread_store_from_config(&config, /*state_db*/ None), - /*state_db*/ None, + state_db, + thread_store, + agent_graph_store, ); let source = manager @@ -548,14 +569,16 @@ async fn resume_stopped_thread_from_rollout_spawns_new_thread() { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await; let manager = ThreadManager::new( &config, auth_manager.clone(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, - thread_store_from_config(&config, /*state_db*/ None), - /*state_db*/ None, + state_db, + thread_store, + agent_graph_store, ); let source = manager @@ -612,14 +635,16 @@ async fn new_uses_active_provider_for_model_refresh() { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await; let manager = ThreadManager::new( &config, auth_manager, SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, - thread_store_from_config(&config, /*state_db*/ None), - /*state_db*/ None, + state_db, + thread_store, + agent_graph_store, ); let _ = manager.list_models(RefreshStrategy::Online).await; @@ -824,15 +849,16 @@ async fn interrupted_fork_snapshot_does_not_synthesize_turn_id_for_legacy_histor let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); - let state_db = init_state_db(&config).await; + let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await; let manager = ThreadManager::new( &config, auth_manager.clone(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, - thread_store_from_config(&config, state_db.clone()), - state_db.clone(), + state_db, + thread_store, + agent_graph_store, ); let source = manager @@ -928,15 +954,16 @@ async fn interrupted_fork_snapshot_preserves_explicit_turn_id() { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); - let state_db = init_state_db(&config).await; + let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await; let manager = ThreadManager::new( &config, auth_manager.clone(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, - thread_store_from_config(&config, state_db.clone()), - state_db.clone(), + state_db, + thread_store, + agent_graph_store, ); let source = manager @@ -1021,15 +1048,16 @@ async fn interrupted_fork_snapshot_uses_persisted_mid_turn_history_without_live_ let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); - let state_db = init_state_db(&config).await; + let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await; let manager = ThreadManager::new( &config, auth_manager.clone(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, - thread_store_from_config(&config, state_db.clone()), - state_db.clone(), + state_db, + thread_store, + agent_graph_store, ); let source = manager @@ -1159,15 +1187,16 @@ async fn resumed_thread_keeps_paused_goal_paused() -> anyhow::Result<()> { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); - let state_db = init_state_db(&config).await; + let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await; let manager = ThreadManager::new( &config, auth_manager.clone(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, - thread_store_from_config(&config, state_db.clone()), - state_db.clone(), + state_db, + thread_store, + agent_graph_store, ); let source = manager @@ -1236,6 +1265,11 @@ async fn resumed_thread_keeps_paused_goal_paused() -> anyhow::Result<()> { .await .is_none() ); + let goal = state_db + .get_thread_goal(resumed.thread_id) + .await? + .expect("goal should still exist after resume"); + assert_eq!(codex_state::ThreadGoalStatus::Paused, goal.status); resumed.thread.shutdown_and_wait().await?; Ok(()) diff --git a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs index 88c1698d8f..9a23550d41 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -6,6 +6,7 @@ use crate::function_tool::FunctionCallError; use crate::init_state_db; use crate::session::tests::make_session_and_context; use crate::session_prefix::format_subagent_notification_message; +use crate::thread_manager::agent_graph_store_from_state_db; use crate::thread_manager::thread_store_from_config; use crate::tools::context::ToolOutput; use crate::tools::handlers::multi_agents_v2::CloseAgentHandler as CloseAgentHandlerV2; @@ -90,11 +91,12 @@ fn parse_agent_id(id: &str) -> ThreadId { ThreadId::from_string(id).expect("agent id should be valid") } -fn thread_manager() -> ThreadManager { +async fn thread_manager() -> ThreadManager { ThreadManager::with_models_provider_for_tests( CodexAuth::from_api_key("dummy"), built_in_model_providers(/* openai_base_url */ /*openai_base_url*/ None)["openai"].clone(), ) + .await } async fn install_role_with_model_override(turn: &mut TurnContext) -> String { @@ -241,7 +243,7 @@ async fn spawn_agent_uses_explorer_role_and_preserves_approval_policy() { } let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let mut config = (*turn.config).clone(); let provider_info = @@ -296,7 +298,7 @@ async fn spawn_agent_uses_explorer_role_and_preserves_approval_policy() { async fn spawn_agent_fork_context_rejects_agent_type_override() { let (mut session, mut turn) = make_session_and_context().await; let role_name = install_role_with_model_override(&mut turn).await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -328,7 +330,7 @@ async fn spawn_agent_fork_context_rejects_agent_type_override() { #[tokio::test] async fn spawn_agent_fork_context_rejects_child_model_overrides() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -363,7 +365,7 @@ async fn spawn_agent_fork_context_rejects_child_model_overrides() { async fn multi_agent_v2_spawn_fork_turns_all_rejects_agent_type_override() { let (mut session, mut turn) = make_session_and_context().await; let role_name = install_role_with_model_override(&mut turn).await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -406,7 +408,7 @@ async fn multi_agent_v2_spawn_fork_turns_all_rejects_agent_type_override() { #[tokio::test] async fn multi_agent_v2_spawn_defaults_to_full_fork_and_rejects_child_model_overrides() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -447,7 +449,7 @@ async fn multi_agent_v2_spawn_defaults_to_full_fork_and_rejects_child_model_over async fn multi_agent_v2_spawn_partial_fork_turns_allows_agent_type_override() { let (mut session, mut turn) = make_session_and_context().await; let role_name = install_role_with_model_override(&mut turn).await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -503,7 +505,7 @@ async fn multi_agent_v2_spawn_partial_fork_turns_allows_agent_type_override() { #[tokio::test] async fn spawn_agent_returns_agent_id_without_task_name() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let output = SpawnAgentHandler @@ -530,7 +532,7 @@ async fn spawn_agent_returns_agent_id_without_task_name() { #[tokio::test] async fn multi_agent_v2_spawn_requires_task_name() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -564,7 +566,7 @@ async fn multi_agent_v2_spawn_requires_task_name() { #[tokio::test] async fn multi_agent_v2_spawn_rejects_legacy_items_field() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -624,7 +626,7 @@ async fn multi_agent_v2_spawn_returns_path_and_send_message_accepts_relative_pat } let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -721,7 +723,7 @@ async fn multi_agent_v2_spawn_returns_path_and_send_message_accepts_relative_pat #[tokio::test] async fn multi_agent_v2_spawn_rejects_legacy_fork_context() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -760,7 +762,7 @@ async fn multi_agent_v2_spawn_rejects_legacy_fork_context() { #[tokio::test] async fn multi_agent_v2_spawn_rejects_invalid_fork_turns_string() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -799,7 +801,7 @@ async fn multi_agent_v2_spawn_rejects_invalid_fork_turns_string() { #[tokio::test] async fn multi_agent_v2_spawn_rejects_zero_fork_turns() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -838,7 +840,7 @@ async fn multi_agent_v2_spawn_rejects_zero_fork_turns() { #[tokio::test] async fn multi_agent_v2_send_message_accepts_root_target_from_child() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -914,7 +916,7 @@ async fn multi_agent_v2_send_message_accepts_root_target_from_child() { #[tokio::test] async fn multi_agent_v2_followup_task_rejects_root_target_from_child() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -995,7 +997,7 @@ async fn multi_agent_v2_followup_task_rejects_root_target_from_child() { #[tokio::test] async fn multi_agent_v2_list_agents_returns_completed_status_and_last_task_message() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -1089,7 +1091,7 @@ async fn multi_agent_v2_list_agents_returns_completed_status_and_last_task_messa #[tokio::test] async fn multi_agent_v2_list_agents_filters_by_relative_path_prefix() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -1176,7 +1178,7 @@ async fn multi_agent_v2_list_agents_filters_by_relative_path_prefix() { #[tokio::test] async fn multi_agent_v2_list_agents_omits_closed_agents() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -1240,7 +1242,7 @@ async fn multi_agent_v2_list_agents_omits_closed_agents() { #[tokio::test] async fn multi_agent_v2_send_message_rejects_legacy_items_field() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -1296,7 +1298,7 @@ async fn multi_agent_v2_send_message_rejects_legacy_items_field() { #[tokio::test] async fn multi_agent_v2_send_message_rejects_interrupt_parameter() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -1369,7 +1371,7 @@ async fn multi_agent_v2_send_message_rejects_interrupt_parameter() { #[tokio::test] async fn multi_agent_v2_followup_task_completion_notifies_parent_on_every_turn() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -1504,7 +1506,7 @@ async fn multi_agent_v2_followup_task_completion_notifies_parent_on_every_turn() #[tokio::test] async fn multi_agent_v2_followup_task_rejects_legacy_items_field() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -1557,7 +1559,7 @@ async fn multi_agent_v2_followup_task_rejects_legacy_items_field() { #[tokio::test] async fn multi_agent_v2_interrupted_turn_does_not_notify_parent() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -1634,7 +1636,7 @@ async fn multi_agent_v2_interrupted_turn_does_not_notify_parent() { #[tokio::test] async fn multi_agent_v2_spawn_omits_agent_id_when_named() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -1673,7 +1675,7 @@ async fn multi_agent_v2_spawn_omits_agent_id_when_named() { #[tokio::test] async fn multi_agent_v2_spawn_surfaces_task_name_validation_errors() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -1716,7 +1718,7 @@ async fn spawn_agent_reapplies_runtime_sandbox_after_role_config() { } let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let expected_sandbox = turn.config.legacy_sandbox_policy(); let mut expected_file_system_sandbox_policy = @@ -1797,7 +1799,7 @@ async fn spawn_agent_reapplies_runtime_sandbox_after_role_config() { #[tokio::test] async fn spawn_agent_rejects_when_depth_limit_exceeded() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let max_depth = turn.config.agent_max_depth; @@ -1835,7 +1837,7 @@ async fn spawn_agent_allows_depth_up_to_configured_max_depth() { } let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let mut config = (*turn.config).clone(); @@ -1881,7 +1883,7 @@ async fn multi_agent_v2_spawn_agent_ignores_configured_max_depth() { } let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let mut config = (*turn.config).clone(); config.agent_max_depth = 1; config @@ -1989,7 +1991,7 @@ async fn send_input_rejects_invalid_id() { #[tokio::test] async fn send_input_reports_missing_agent() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let agent_id = ThreadId::new(); let invocation = invocation( @@ -2010,7 +2012,7 @@ async fn send_input_reports_missing_agent() { #[tokio::test] async fn send_input_interrupts_before_prompt() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let config = turn.config.as_ref().clone(); let thread = manager @@ -2052,7 +2054,7 @@ async fn send_input_interrupts_before_prompt() { #[tokio::test] async fn send_input_accepts_structured_items() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let config = turn.config.as_ref().clone(); let thread = manager @@ -2126,7 +2128,7 @@ async fn resume_agent_rejects_invalid_id() { #[tokio::test] async fn resume_agent_reports_missing_agent() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let agent_id = ThreadId::new(); let invocation = invocation( @@ -2147,7 +2149,7 @@ async fn resume_agent_reports_missing_agent() { #[tokio::test] async fn resume_agent_noops_for_active_agent() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let config = turn.config.as_ref().clone(); let thread = manager @@ -2186,7 +2188,7 @@ async fn resume_agent_noops_for_active_agent() { #[tokio::test] async fn resume_agent_restores_closed_agent_and_accepts_send_input() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let config = turn.config.as_ref().clone(); let thread = manager @@ -2265,7 +2267,7 @@ async fn resume_agent_restores_closed_agent_and_accepts_send_input() { #[tokio::test] async fn resume_agent_rejects_when_depth_limit_exceeded() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let max_depth = turn.config.agent_max_depth; @@ -2354,7 +2356,7 @@ async fn wait_agent_rejects_empty_targets() { #[tokio::test] async fn multi_agent_v2_wait_agent_accepts_timeout_only_argument() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -2493,7 +2495,7 @@ async fn multi_agent_v2_wait_agent_uses_configured_min_timeout() { #[tokio::test] async fn wait_agent_returns_not_found_for_missing_agents() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let id_a = ThreadId::new(); let id_b = ThreadId::new(); @@ -2529,7 +2531,7 @@ async fn wait_agent_returns_not_found_for_missing_agents() { #[tokio::test] async fn wait_agent_times_out_when_status_is_not_final() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let config = turn.config.as_ref().clone(); let thread = manager @@ -2572,7 +2574,7 @@ async fn wait_agent_times_out_when_status_is_not_final() { #[tokio::test] async fn wait_agent_clamps_short_timeouts_to_minimum() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let config = turn.config.as_ref().clone(); let thread = manager @@ -2610,7 +2612,7 @@ async fn wait_agent_clamps_short_timeouts_to_minimum() { #[tokio::test] async fn wait_agent_returns_final_status_without_timeout() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let config = turn.config.as_ref().clone(); let thread = manager @@ -2662,7 +2664,7 @@ async fn wait_agent_returns_final_status_without_timeout() { #[tokio::test] async fn multi_agent_v2_wait_agent_returns_summary_for_mailbox_activity() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -2753,7 +2755,7 @@ async fn multi_agent_v2_wait_agent_returns_summary_for_mailbox_activity() { #[tokio::test] async fn multi_agent_v2_wait_agent_returns_for_already_queued_mail() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -2831,7 +2833,7 @@ async fn multi_agent_v2_wait_agent_returns_for_already_queued_mail() { #[tokio::test] async fn multi_agent_v2_wait_agent_wakes_on_any_mailbox_notification() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -2919,7 +2921,7 @@ async fn multi_agent_v2_wait_agent_wakes_on_any_mailbox_notification() { #[tokio::test] async fn multi_agent_v2_wait_agent_does_not_return_completed_content() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -3005,7 +3007,7 @@ async fn multi_agent_v2_wait_agent_does_not_return_completed_content() { #[tokio::test] async fn multi_agent_v2_close_agent_accepts_task_name_target() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -3064,7 +3066,7 @@ async fn multi_agent_v2_close_agent_accepts_task_name_target() { #[tokio::test] async fn multi_agent_v2_close_agent_rejects_root_target_and_id() { let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; let root = manager .start_thread((*turn.config).clone()) .await @@ -3112,7 +3114,7 @@ async fn multi_agent_v2_close_agent_rejects_root_target_and_id() { #[tokio::test] async fn close_agent_submits_shutdown_and_returns_previous_status() { let (mut session, turn) = make_session_and_context().await; - let manager = thread_manager(); + let manager = thread_manager().await; session.services.agent_control = manager.agent_control(); let config = turn.config.as_ref().clone(); let thread = manager @@ -3157,15 +3159,18 @@ async fn tool_handlers_cascade_close_and_resume_and_keep_explicitly_closed_subtr .features .enable(Feature::Sqlite) .expect("test config should allow sqlite"); - let state_db = init_state_db(&config).await; + let state_db = init_state_db(&config) + .await + .expect("test config should initialize state db"); let manager = ThreadManager::new( &config, AuthManager::from_auth_for_testing(CodexAuth::from_api_key("dummy")), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, - thread_store_from_config(&config, state_db.clone()), state_db.clone(), + thread_store_from_config(&config, state_db.clone()), + agent_graph_store_from_state_db(state_db.clone()), ); let parent = manager diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index 14f652f58b..13dd026654 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -15,7 +15,9 @@ use anyhow::anyhow; use codex_config::CloudRequirementsLoader; use codex_core::CodexThread; use codex_core::ThreadManager; +use codex_core::agent_graph_store_from_state_db; use codex_core::config::Config; +use codex_core::init_state_db_from_config; use codex_core::shell::Shell; use codex_core::shell::get_shell_by_model_provided_path; use codex_core::thread_store_from_config; @@ -423,25 +425,30 @@ impl TestCodexBuilder { environment_manager: Arc, ) -> anyhow::Result { let auth = self.auth.clone(); - let state_db = codex_core::init_state_db(&config).await; let thread_manager = if config.model_catalog.is_some() { + let state_db = init_state_db_from_config(&config) + .await + .expect("test codex requires state db"); + let thread_store = thread_store_from_config(&config, state_db.clone()); + let agent_graph_store = agent_graph_store_from_state_db(state_db.clone()); ThreadManager::new( &config, codex_core::test_support::auth_manager_from_auth(auth.clone()), SessionSource::Exec, Arc::clone(&environment_manager), /*analytics_events_client*/ None, - thread_store_from_config(&config, state_db.clone()), - state_db.clone(), + state_db, + thread_store, + agent_graph_store, ) } else { - codex_core::test_support::thread_manager_with_models_provider_home_and_state( + codex_core::test_support::thread_manager_with_models_provider_and_home( auth.clone(), config.model_provider.clone(), config.codex_home.to_path_buf(), Arc::clone(&environment_manager), - state_db.clone(), ) + .await }; let thread_manager = Arc::new(thread_manager); let user_shell_override = self.user_shell_override.clone(); diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index afa586a158..dafde4ccea 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -5,6 +5,8 @@ use codex_core::NewThread; use codex_core::Prompt; use codex_core::ResponseEvent; use codex_core::ThreadManager; +use codex_core::agent_graph_store_from_state_db; +use codex_core::init_state_db_from_config; use codex_core::thread_store_from_config; use codex_features::Feature; use codex_login::AuthManager; @@ -770,15 +772,10 @@ async fn includes_conversation_id_and_model_headers_in_request() { let installation_id = std::fs::read_to_string(test.codex_home_path().join(INSTALLATION_ID_FILENAME)) .expect("read installation id"); - let session_id_string = session_id.to_string(); - assert_eq!(request_session_id, session_id_string); + assert_eq!(request_session_id, session_id.to_string()); assert_eq!(request_originator, originator().value); assert_eq!(request_authorization, "Bearer Test API Key"); - assert_eq!( - request_body["prompt_cache_key"].as_str(), - Some(session_id_string.as_str()) - ); assert_eq!( request_body["client_metadata"]["x-codex-installation-id"].as_str(), Some(installation_id.as_str()) @@ -1106,14 +1103,20 @@ async fn prefers_apikey_when_config_prefers_apikey_even_with_chatgpt_tokens() { Ok(None) => panic!("No CodexAuth found in codex_home"), Err(e) => panic!("Failed to load CodexAuth: {e}"), }; + let state_db = init_state_db_from_config(&config) + .await + .expect("client test requires state db"); + let thread_store = thread_store_from_config(&config, state_db.clone()); + let agent_graph_store = agent_graph_store_from_state_db(state_db.clone()); let thread_manager = ThreadManager::new( &config, auth_manager, SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, - thread_store_from_config(&config, /*state_db*/ None), - /*state_db*/ None, + state_db, + thread_store, + agent_graph_store, ); let NewThread { thread: codex, .. } = thread_manager .start_thread(config.clone()) diff --git a/codex-rs/core/tests/suite/personality_migration.rs b/codex-rs/core/tests/suite/personality_migration.rs index 68121d2e12..8d33ba1705 100644 --- a/codex-rs/core/tests/suite/personality_migration.rs +++ b/codex-rs/core/tests/suite/personality_migration.rs @@ -13,6 +13,8 @@ use codex_protocol::protocol::SessionMeta; use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::UserMessageEvent; +use codex_rollout::RolloutConfig; +use codex_rollout::state_db::StateDbHandle; use pretty_assertions::assert_eq; use std::io; use std::path::Path; @@ -26,6 +28,27 @@ async fn read_config_toml(codex_home: &Path) -> io::Result { toml::from_str(&contents).map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err)) } +async fn state_db_for_test(codex_home: &Path) -> io::Result { + let config = RolloutConfig { + codex_home: codex_home.to_path_buf(), + sqlite_home: codex_home.to_path_buf(), + cwd: codex_home.to_path_buf(), + model_provider_id: "openai".to_string(), + generate_memories: false, + }; + codex_rollout::state_db::try_init(&config) + .await + .map_err(io::Error::other) +} + +async fn run_migration( + codex_home: &Path, + config_toml: &ConfigToml, +) -> io::Result { + let state_db = state_db_for_test(codex_home).await?; + maybe_migrate_personality(codex_home, config_toml, state_db).await +} + async fn write_session_with_user_event(codex_home: &Path) -> io::Result<()> { let thread_id = ThreadId::new(); let dir = codex_home @@ -141,8 +164,7 @@ async fn migration_marker_exists_no_sessions_no_change() -> io::Result<()> { let marker_path = temp.path().join(PERSONALITY_MIGRATION_FILENAME); tokio::fs::write(&marker_path, "v1\n").await?; - let status = - maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?; + let status = run_migration(temp.path(), &ConfigToml::default()).await?; assert_eq!(status, PersonalityMigrationStatus::SkippedMarker); assert_eq!( @@ -156,8 +178,7 @@ async fn migration_marker_exists_no_sessions_no_change() -> io::Result<()> { async fn no_marker_no_sessions_no_change() -> io::Result<()> { let temp = TempDir::new()?; - let status = - maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?; + let status = run_migration(temp.path(), &ConfigToml::default()).await?; assert_eq!(status, PersonalityMigrationStatus::SkippedNoSessions); assert_eq!( @@ -176,8 +197,7 @@ async fn no_marker_sessions_sets_personality() -> io::Result<()> { let temp = TempDir::new()?; write_session_with_user_event(temp.path()).await?; - let status = - maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?; + let status = run_migration(temp.path(), &ConfigToml::default()).await?; assert_eq!(status, PersonalityMigrationStatus::Applied); assert_eq!( @@ -197,7 +217,7 @@ async fn no_marker_sessions_preserves_existing_config_fields() -> io::Result<()> tokio::fs::write(temp.path().join("config.toml"), "model = \"gpt-5.4\"\n").await?; let config_toml = read_config_toml(temp.path()).await?; - let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?; + let status = run_migration(temp.path(), &config_toml).await?; assert_eq!(status, PersonalityMigrationStatus::Applied); let persisted = read_config_toml(temp.path()).await?; @@ -211,8 +231,7 @@ async fn no_marker_meta_only_rollout_is_treated_as_no_sessions() -> io::Result<( let temp = TempDir::new()?; write_session_with_meta_only(temp.path()).await?; - let status = - maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?; + let status = run_migration(temp.path(), &ConfigToml::default()).await?; assert_eq!(status, PersonalityMigrationStatus::SkippedNoSessions); assert_eq!( @@ -232,7 +251,7 @@ async fn no_marker_explicit_global_personality_skips_migration() -> io::Result<( write_session_with_user_event(temp.path()).await?; let config_toml = parse_config_toml("personality = \"friendly\"\n")?; - let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?; + let status = run_migration(temp.path(), &config_toml).await?; assert_eq!( status, @@ -262,7 +281,7 @@ personality = "friendly" "#, )?; - let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?; + let status = run_migration(temp.path(), &config_toml).await?; assert_eq!( status, @@ -285,7 +304,7 @@ async fn marker_short_circuits_invalid_profile_resolution() -> io::Result<()> { tokio::fs::write(temp.path().join(PERSONALITY_MIGRATION_FILENAME), "v1\n").await?; let config_toml = parse_config_toml("profile = \"missing\"\n")?; - let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?; + let status = run_migration(temp.path(), &config_toml).await?; assert_eq!(status, PersonalityMigrationStatus::SkippedMarker); Ok(()) @@ -296,7 +315,7 @@ async fn invalid_selected_profile_returns_error_and_does_not_write_marker() -> i let temp = TempDir::new()?; let config_toml = parse_config_toml("profile = \"missing\"\n")?; - let err = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None) + let err = run_migration(temp.path(), &config_toml) .await .expect_err("missing profile should fail"); @@ -313,10 +332,8 @@ async fn applied_migration_is_idempotent_on_second_run() -> io::Result<()> { let temp = TempDir::new()?; write_session_with_user_event(temp.path()).await?; - let first_status = - maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?; - let second_status = - maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?; + let first_status = run_migration(temp.path(), &ConfigToml::default()).await?; + let second_status = run_migration(temp.path(), &ConfigToml::default()).await?; assert_eq!(first_status, PersonalityMigrationStatus::Applied); assert_eq!(second_status, PersonalityMigrationStatus::SkippedMarker); @@ -330,8 +347,7 @@ async fn no_marker_archived_sessions_sets_personality() -> io::Result<()> { let temp = TempDir::new()?; write_archived_session_with_user_event(temp.path()).await?; - let status = - maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?; + let status = run_migration(temp.path(), &ConfigToml::default()).await?; assert_eq!(status, PersonalityMigrationStatus::Applied); assert_eq!( diff --git a/codex-rs/core/tests/suite/prompt_debug_tests.rs b/codex-rs/core/tests/suite/prompt_debug_tests.rs index dc506bc474..4fee438261 100644 --- a/codex-rs/core/tests/suite/prompt_debug_tests.rs +++ b/codex-rs/core/tests/suite/prompt_debug_tests.rs @@ -29,7 +29,6 @@ async fn build_prompt_input_includes_context_and_user_message() -> Result<()> { text: "hello from debug prompt".to_string(), text_elements: Vec::new(), }], - /*state_db*/ None, ) .await?; diff --git a/codex-rs/core/tests/suite/resume_warning.rs b/codex-rs/core/tests/suite/resume_warning.rs index cb545df351..02e7626d87 100644 --- a/codex-rs/core/tests/suite/resume_warning.rs +++ b/codex-rs/core/tests/suite/resume_warning.rs @@ -95,7 +95,8 @@ async fn emits_warning_when_resumed_model_differs() { let thread_manager = codex_core::test_support::thread_manager_with_models_provider( CodexAuth::from_api_key("test"), config.model_provider.clone(), - ); + ) + .await; let auth_manager = codex_core::test_support::auth_manager_from_auth(CodexAuth::from_api_key("test")); diff --git a/codex-rs/core/tests/suite/skills.rs b/codex-rs/core/tests/suite/skills.rs index add1ee6ea7..ebf2c3226e 100644 --- a/codex-rs/core/tests/suite/skills.rs +++ b/codex-rs/core/tests/suite/skills.rs @@ -3,6 +3,8 @@ use anyhow::Result; use codex_core::ThreadManager; +use codex_core::agent_graph_store_from_state_db; +use codex_core::init_state_db_from_config; use codex_core::thread_store_from_config; use codex_exec_server::CreateDirectoryOptions; use codex_exec_server::EnvironmentManager; @@ -235,6 +237,11 @@ async fn list_skills_skips_cwd_roots_when_environment_disabled() -> Result<()> { let mut config = load_default_config_for_test(&codex_home).await; config.cwd = AbsolutePathBuf::from_absolute_path_checked(cwd.path())?; + let state_db = init_state_db_from_config(&config) + .await + .expect("skills test requires state db"); + let thread_store = thread_store_from_config(&config, state_db.clone()); + let agent_graph_store = agent_graph_store_from_state_db(state_db.clone()); let thread_manager = ThreadManager::new( &config, codex_core::test_support::auth_manager_from_auth(CodexAuth::from_api_key("dummy")), @@ -246,8 +253,9 @@ async fn list_skills_skips_cwd_roots_when_environment_disabled() -> Result<()> { )?, )), /*analytics_events_client*/ None, - thread_store_from_config(&config, /*state_db*/ None), - /*state_db*/ None, + state_db, + thread_store, + agent_graph_store, ); let new_thread = thread_manager.start_thread(config.clone()).await?; let cwd = config.cwd.to_path_buf(); diff --git a/codex-rs/core/tests/suite/unstable_features_warning.rs b/codex-rs/core/tests/suite/unstable_features_warning.rs index 66a7366589..7a4924cdda 100644 --- a/codex-rs/core/tests/suite/unstable_features_warning.rs +++ b/codex-rs/core/tests/suite/unstable_features_warning.rs @@ -34,7 +34,8 @@ async fn emits_warning_when_unstable_features_enabled_via_config() { let thread_manager = codex_core::test_support::thread_manager_with_models_provider( CodexAuth::from_api_key("test"), config.model_provider.clone(), - ); + ) + .await; let auth_manager = codex_core::test_support::auth_manager_from_auth(CodexAuth::from_api_key("test")); @@ -81,7 +82,8 @@ async fn suppresses_warning_when_configured() { let thread_manager = codex_core::test_support::thread_manager_with_models_provider( CodexAuth::from_api_key("test"), config.model_provider.clone(), - ); + ) + .await; let auth_manager = codex_core::test_support::auth_manager_from_auth(CodexAuth::from_api_key("test")); diff --git a/codex-rs/mcp-server/src/lib.rs b/codex-rs/mcp-server/src/lib.rs index 75372c6e7c..ac764456f5 100644 --- a/codex-rs/mcp-server/src/lib.rs +++ b/codex-rs/mcp-server/src/lib.rs @@ -83,7 +83,6 @@ pub async fn run_main( std::io::Error::new(ErrorKind::InvalidData, format!("error loading config: {e}")) })?; set_default_client_residency_requirement(config.enforce_residency.value()); - let state_db = codex_core::init_state_db(&config).await; let otel = codex_core::otel_init::build_provider( &config, @@ -140,15 +139,18 @@ pub async fn run_main( // Task: process incoming messages. let processor_handle = tokio::spawn({ let outgoing_message_sender = OutgoingMessageSender::new(outgoing_tx); - let mut processor = MessageProcessor::new( + let processor = MessageProcessor::new( outgoing_message_sender, arg0_paths, Arc::new(config), environment_manager, - state_db, ) .await; async move { + let Some(mut processor) = processor else { + error!("failed to initialize MCP processor"); + return; + }; while let Some(msg) = incoming_rx.recv().await { match msg { JsonRpcMessage::Request(r) => processor.process_request(r).await, diff --git a/codex-rs/mcp-server/src/message_processor.rs b/codex-rs/mcp-server/src/message_processor.rs index 9ce46dfeda..2591891e4d 100644 --- a/codex-rs/mcp-server/src/message_processor.rs +++ b/codex-rs/mcp-server/src/message_processor.rs @@ -2,9 +2,10 @@ use std::collections::HashMap; use std::sync::Arc; use codex_arg0::Arg0DispatchPaths; -use codex_core::StateDbHandle; use codex_core::ThreadManager; +use codex_core::agent_graph_store_from_state_db; use codex_core::config::Config; +use codex_core::init_state_db_from_config; use codex_core::thread_store_from_config; use codex_exec_server::EnvironmentManager; use codex_login::AuthManager; @@ -54,30 +55,33 @@ impl MessageProcessor { arg0_paths: Arg0DispatchPaths, config: Arc, environment_manager: Arc, - state_db: Option, - ) -> Self { + ) -> Option { let outgoing = Arc::new(outgoing); let auth_manager = AuthManager::shared_from_config( config.as_ref(), /*enable_codex_api_key_env*/ false, ) .await; + let state_db = init_state_db_from_config(config.as_ref()).await?; + let thread_store = thread_store_from_config(config.as_ref(), state_db.clone()); + let agent_graph_store = agent_graph_store_from_state_db(state_db.clone()); let thread_manager = Arc::new(ThreadManager::new( config.as_ref(), auth_manager, SessionSource::Mcp, environment_manager, /*analytics_events_client*/ None, - thread_store_from_config(config.as_ref(), state_db.clone()), - state_db.clone(), + state_db, + thread_store, + agent_graph_store, )); - Self { + Some(Self { outgoing, initialized: false, arg0_paths, thread_manager, running_requests_id_to_codex_uuid: Arc::new(Mutex::new(HashMap::new())), - } + }) } pub(crate) async fn process_request(&mut self, request: JsonRpcRequest) { diff --git a/codex-rs/thread-manager-sample/src/main.rs b/codex-rs/thread-manager-sample/src/main.rs index 9d84ece674..a27c78d543 100644 --- a/codex-rs/thread-manager-sample/src/main.rs +++ b/codex-rs/thread-manager-sample/src/main.rs @@ -52,10 +52,11 @@ use codex_core_api::TuiNotificationSettings; use codex_core_api::UriBasedFileOpener; use codex_core_api::UserInput; use codex_core_api::WebSearchMode; +use codex_core_api::agent_graph_store_from_state_db; use codex_core_api::arg0_dispatch_or_else; use codex_core_api::built_in_model_providers; use codex_core_api::find_codex_home; -use codex_core_api::init_state_db; +use codex_core_api::init_state_db_from_config; use codex_core_api::item_event_to_server_notification; use codex_core_api::set_default_originator; use codex_core_api::thread_store_from_config; @@ -104,7 +105,6 @@ async fn run_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { }; let config = new_config(args.model, arg0_paths)?; - let state_db = init_state_db(&config).await; let auth_manager = AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false).await; @@ -112,7 +112,11 @@ async fn run_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { config.codex_self_exe.clone(), config.codex_linux_sandbox_exe.clone(), )?; + let Some(state_db) = init_state_db_from_config(&config).await else { + bail!("thread manager sample requires state db"); + }; let thread_store = thread_store_from_config(&config, state_db.clone()); + let agent_graph_store = agent_graph_store_from_state_db(state_db.clone()); let environment_manager = Arc::new(EnvironmentManager::new(EnvironmentManagerArgs::new(local_runtime_paths)).await); let thread_manager = ThreadManager::new( @@ -121,8 +125,9 @@ async fn run_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { SessionSource::Exec, environment_manager, /*analytics_events_client*/ None, - Arc::clone(&thread_store), state_db, + Arc::clone(&thread_store), + agent_graph_store, ); let NewThread { diff --git a/codex-rs/thread-store/src/local/archive_thread.rs b/codex-rs/thread-store/src/local/archive_thread.rs index 8fb214e98c..268a6ffc47 100644 --- a/codex-rs/thread-store/src/local/archive_thread.rs +++ b/codex-rs/thread-store/src/local/archive_thread.rs @@ -13,11 +13,11 @@ pub(super) async fn archive_thread( params: ArchiveThreadParams, ) -> ThreadStoreResult<()> { let thread_id = params.thread_id; - let state_db_ctx = store.state_db().await; + let state_db = store.state_db(); let rollout_path = find_thread_path_by_id_str( store.config.codex_home.as_path(), &thread_id.to_string(), - state_db_ctx.as_deref(), + Some(state_db.as_ref()), ) .await .map_err(|err| ThreadStoreError::InvalidRequest { @@ -52,11 +52,10 @@ pub(super) async fn archive_thread( } })?; - if let Some(ctx) = state_db_ctx { - let _ = ctx - .mark_archived(thread_id, archived_path.as_path(), Utc::now()) - .await; - } + let _ = store + .state_db() + .mark_archived(thread_id, archived_path.as_path(), Utc::now()) + .await; Ok(()) } @@ -75,13 +74,15 @@ mod tests { use crate::ThreadSortKey; use crate::ThreadStore; use crate::local::LocalThreadStore; + use crate::local::test_support::init_test_state_db; use crate::local::test_support::test_config; + use crate::local::test_support::test_store; use crate::local::test_support::write_session_file; #[tokio::test] async fn archive_thread_moves_rollout_to_archived_collection() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(201); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let active_path = @@ -127,21 +128,12 @@ mod tests { async fn archive_thread_updates_sqlite_metadata_when_present() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(202); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let active_path = write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file"); - let runtime = codex_state::StateRuntime::init( - home.path().to_path_buf(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); - runtime - .mark_backfill_complete(/*last_watermark*/ None) - .await - .expect("backfill should be complete"); let mut builder = codex_state::ThreadMetadataBuilder::new( thread_id, active_path.clone(), diff --git a/codex-rs/thread-store/src/local/create_thread.rs b/codex-rs/thread-store/src/local/create_thread.rs index e444e5c91d..cb0f64ea9c 100644 --- a/codex-rs/thread-store/src/local/create_thread.rs +++ b/codex-rs/thread-store/src/local/create_thread.rs @@ -22,12 +22,12 @@ pub(super) async fn create_thread( })?; let config = RolloutConfig { codex_home: store.config.codex_home.clone(), - sqlite_home: store.config.sqlite_home.clone(), + sqlite_home: store.sqlite_home(), cwd, model_provider_id: params.metadata.model_provider.clone(), generate_memories: matches!(params.metadata.memory_mode, ThreadMemoryMode::Enabled), }; - let state_db_ctx = store.state_db().await; + let state_db_ctx = Some(store.state_db()); let recorder = RolloutRecorder::new( &config, RolloutRecorderParams::new( diff --git a/codex-rs/thread-store/src/local/list_threads.rs b/codex-rs/thread-store/src/local/list_threads.rs index e470ad2be9..5cf1f27aa1 100644 --- a/codex-rs/thread-store/src/local/list_threads.rs +++ b/codex-rs/thread-store/src/local/list_threads.rs @@ -39,16 +39,16 @@ pub(super) async fn list_threads( SortDirection::Asc => codex_rollout::SortDirection::Asc, SortDirection::Desc => codex_rollout::SortDirection::Desc, }; - let state_db = store.state_db().await; let rollout_config = RolloutConfig { codex_home: store.config.codex_home.clone(), - sqlite_home: store.config.sqlite_home.clone(), + sqlite_home: store.sqlite_home(), cwd: store.config.codex_home.clone(), model_provider_id: store.config.default_model_provider_id.clone(), generate_memories: false, }; + let state_db_ctx = Some(store.state_db()); let page = list_rollout_threads( - state_db, + state_db_ctx, &rollout_config, store.config.default_model_provider_id.as_str(), ¶ms, @@ -80,14 +80,13 @@ pub(super) async fn list_threads( .map(|thread| thread.thread_id) .collect::>(); let mut names = HashMap::::with_capacity(thread_ids.len()); - if let Some(state_db_ctx) = store.state_db().await { - for &thread_id in &thread_ids { - let Ok(Some(metadata)) = state_db_ctx.get_thread(thread_id).await else { - continue; - }; - if let Some(title) = distinct_thread_metadata_title(&metadata) { - names.insert(thread_id, title); - } + let state_db_ctx = store.state_db(); + for &thread_id in &thread_ids { + let Ok(Some(metadata)) = state_db_ctx.get_thread(thread_id).await else { + continue; + }; + if let Some(title) = distinct_thread_metadata_title(&metadata) { + names.insert(thread_id, title); } } if names.len() < thread_ids.len() @@ -108,9 +107,9 @@ pub(super) async fn list_threads( } async fn list_rollout_threads( - state_db: Option, + state_db_ctx: Option, config: &RolloutConfig, - default_model_provider_id: &str, + default_model_provider: &str, params: &ListThreadsParams, cursor: Option<&codex_rollout::Cursor>, sort_key: codex_rollout::ThreadSortKey, @@ -118,7 +117,7 @@ async fn list_rollout_threads( ) -> ThreadStoreResult { let page = if params.use_state_db_only && params.archived { RolloutRecorder::list_archived_threads_from_state_db( - state_db, + state_db_ctx.clone(), config, params.page_size, cursor, @@ -127,13 +126,13 @@ async fn list_rollout_threads( params.allowed_sources.as_slice(), params.model_providers.as_deref(), params.cwd_filters.as_deref(), - default_model_provider_id, + default_model_provider, params.search_term.as_deref(), ) .await } else if params.use_state_db_only { RolloutRecorder::list_threads_from_state_db( - state_db, + state_db_ctx.clone(), config, params.page_size, cursor, @@ -142,13 +141,13 @@ async fn list_rollout_threads( params.allowed_sources.as_slice(), params.model_providers.as_deref(), params.cwd_filters.as_deref(), - default_model_provider_id, + default_model_provider, params.search_term.as_deref(), ) .await } else if params.archived { RolloutRecorder::list_archived_threads( - state_db, + state_db_ctx.clone(), config, params.page_size, cursor, @@ -157,13 +156,13 @@ async fn list_rollout_threads( params.allowed_sources.as_slice(), params.model_providers.as_deref(), params.cwd_filters.as_deref(), - default_model_provider_id, + default_model_provider, params.search_term.as_deref(), ) .await } else { RolloutRecorder::list_threads( - state_db, + state_db_ctx, config, params.page_size, cursor, @@ -172,7 +171,7 @@ async fn list_rollout_threads( params.allowed_sources.as_slice(), params.model_providers.as_deref(), params.cwd_filters.as_deref(), - default_model_provider_id, + default_model_provider, params.search_term.as_deref(), ) .await @@ -195,7 +194,9 @@ mod tests { use super::*; use crate::ThreadStore; use crate::local::LocalThreadStore; + use crate::local::test_support::init_test_state_db; use crate::local::test_support::test_config; + use crate::local::test_support::test_store; use crate::local::test_support::write_archived_session_file; use crate::local::test_support::write_session_file; use crate::local::test_support::write_session_file_with; @@ -203,7 +204,7 @@ mod tests { #[tokio::test] async fn list_threads_uses_default_provider_when_rollout_omits_provider() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; write_session_file_with( home.path(), home.path().join("sessions/2025/01/03"), @@ -238,22 +239,13 @@ mod tests { async fn list_threads_preserves_sqlite_title_search_results() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(103); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = home.path().join("rollout-title-search.jsonl"); fs::write(&rollout_path, "").expect("placeholder rollout file"); - let runtime = codex_state::StateRuntime::init( - home.path().to_path_buf(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); - runtime - .mark_backfill_complete(/*last_watermark*/ None) - .await - .expect("backfill should be complete"); let created_at = Utc::now(); let mut builder = codex_state::ThreadMetadataBuilder::new( thread_id, @@ -267,6 +259,10 @@ mod tests { let mut metadata = builder.build(config.default_model_provider_id.as_str()); metadata.title = "needle title".to_string(); metadata.first_user_message = Some("plain preview".to_string()); + runtime + .mark_backfill_complete(/*last_watermark*/ None) + .await + .expect("backfill should be complete"); runtime .upsert_thread(&metadata) .await @@ -303,7 +299,7 @@ mod tests { #[tokio::test] async fn list_threads_selects_active_or_archived_collection() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let active_uuid = Uuid::from_u128(105); let archived_uuid = Uuid::from_u128(106); write_session_file(home.path(), "2025-01-03T12-00-00", active_uuid) @@ -372,7 +368,7 @@ mod tests { async fn list_threads_returns_local_rollout_summary() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); - let store = LocalThreadStore::new(config, /*state_db*/ None); + let store = LocalThreadStore::new(config.clone(), init_test_state_db(&config).await); let uuid = Uuid::from_u128(101); let path = write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file"); @@ -411,7 +407,7 @@ mod tests { #[tokio::test] async fn list_threads_rejects_invalid_cursor() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let err = store .list_threads(ListThreadsParams { diff --git a/codex-rs/thread-store/src/local/live_writer.rs b/codex-rs/thread-store/src/local/live_writer.rs index 643207b59d..37d9758973 100644 --- a/codex-rs/thread-store/src/local/live_writer.rs +++ b/codex-rs/thread-store/src/local/live_writer.rs @@ -66,12 +66,12 @@ pub(super) async fn resume_thread( })?; let config = RolloutConfig { codex_home: store.config.codex_home.clone(), - sqlite_home: store.config.sqlite_home.clone(), + sqlite_home: store.sqlite_home(), cwd, model_provider_id: params.metadata.model_provider.clone(), generate_memories: matches!(params.metadata.memory_mode, ThreadMemoryMode::Enabled), }; - let state_db_ctx = store.state_db().await; + let state_db_ctx = Some(store.state_db()); let recorder = RolloutRecorder::new( &config, RolloutRecorderParams::resume( diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index 058c7b2309..e6ea63bba1 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -41,7 +41,7 @@ use crate::UpdateThreadMetadataParams; pub struct LocalThreadStore { pub(super) config: LocalThreadStoreConfig, live_recorders: Arc>>, - state_db: Option, + state_db: StateDbHandle, } /// Process-scoped configuration for local thread storage. @@ -51,7 +51,6 @@ pub struct LocalThreadStore { #[derive(Clone, Debug, PartialEq, Eq)] pub struct LocalThreadStoreConfig { pub codex_home: PathBuf, - pub sqlite_home: PathBuf, /// Provider used only when older local metadata does not contain one. pub default_model_provider_id: String, } @@ -60,7 +59,6 @@ impl LocalThreadStoreConfig { pub fn from_config(config: &impl codex_rollout::RolloutConfigView) -> Self { Self { codex_home: config.codex_home().to_path_buf(), - sqlite_home: config.sqlite_home().to_path_buf(), default_model_provider_id: config.model_provider_id().to_string(), } } @@ -75,8 +73,9 @@ impl std::fmt::Debug for LocalThreadStore { } impl LocalThreadStore { - /// Create a local store using an already initialized state DB handle. - pub fn new(config: LocalThreadStoreConfig, state_db: Option) -> Self { + /// Create a local store from process-scoped local storage configuration and + /// the caller-provided shared state DB handle. + pub fn new(config: LocalThreadStoreConfig, state_db: StateDbHandle) -> Self { Self { config, live_recorders: Arc::new(Mutex::new(HashMap::new())), @@ -85,10 +84,14 @@ impl LocalThreadStore { } /// Return the state DB handle used by local rollout writers. - pub async fn state_db(&self) -> Option { + pub fn state_db(&self) -> StateDbHandle { self.state_db.clone() } + pub(super) fn sqlite_home(&self) -> PathBuf { + self.state_db.codex_home().to_path_buf() + } + /// Read a local rollout-backed thread by path. pub async fn read_thread_by_rollout_path( &self, @@ -282,14 +285,16 @@ mod tests { use super::*; use crate::ThreadEventPersistenceMode; use crate::ThreadPersistenceMetadata; + use crate::local::test_support::init_test_state_db; use crate::local::test_support::test_config; + use crate::local::test_support::test_store; use crate::local::test_support::write_archived_session_file; use crate::local::test_support::write_session_file; #[tokio::test] async fn live_writer_lifecycle_writes_and_closes() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let thread_id = ThreadId::default(); store @@ -338,7 +343,7 @@ mod tests { #[tokio::test] async fn create_thread_rejects_missing_cwd() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let thread_id = ThreadId::default(); let mut params = create_thread_params(thread_id); params.metadata.cwd = None; @@ -358,7 +363,7 @@ mod tests { #[tokio::test] async fn discard_thread_drops_unmaterialized_live_writer() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let thread_id = ThreadId::default(); store @@ -396,8 +401,9 @@ mod tests { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); let thread_id = ThreadId::default(); + let state_db = init_test_state_db(&config).await; - let first_store = LocalThreadStore::new(config.clone(), /*state_db*/ None); + let first_store = LocalThreadStore::new(config.clone(), state_db.clone()); first_store .create_thread(create_thread_params(thread_id)) .await @@ -426,7 +432,7 @@ mod tests { .await .expect("shutdown initial writer"); - let resumed_store = LocalThreadStore::new(config, /*state_db*/ None); + let resumed_store = LocalThreadStore::new(config, state_db); resumed_store .resume_thread(ResumeThreadParams { thread_id, @@ -457,7 +463,7 @@ mod tests { #[tokio::test] async fn create_thread_rejects_duplicate_live_writer() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let thread_id = ThreadId::default(); store @@ -477,7 +483,7 @@ mod tests { #[tokio::test] async fn resume_thread_rejects_duplicate_live_writer() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let thread_id = ThreadId::default(); store @@ -506,7 +512,7 @@ mod tests { #[tokio::test] async fn resume_thread_rejects_missing_cwd() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = uuid::Uuid::from_u128(407); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = @@ -535,7 +541,7 @@ mod tests { async fn load_history_uses_live_writer_rollout_path() { let home = TempDir::new().expect("temp dir"); let external_home = TempDir::new().expect("external temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = uuid::Uuid::from_u128(404); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = write_session_file(external_home.path(), "2025-01-04T10-00-00", uuid) @@ -584,7 +590,7 @@ mod tests { async fn read_thread_uses_live_writer_rollout_path_for_external_resume() { let home = TempDir::new().expect("temp dir"); let external_home = TempDir::new().expect("external temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = uuid::Uuid::from_u128(406); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = write_session_file(external_home.path(), "2025-01-04T11-00-00", uuid) @@ -623,7 +629,7 @@ mod tests { #[tokio::test] async fn load_history_uses_live_writer_rollout_path_for_archived_source() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = uuid::Uuid::from_u128(405); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = write_archived_session_file(home.path(), "2025-01-04T10-30-00", uuid) @@ -691,7 +697,7 @@ mod tests { #[tokio::test] async fn read_thread_by_rollout_path_includes_history() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let thread_id = ThreadId::default(); store diff --git a/codex-rs/thread-store/src/local/read_thread.rs b/codex-rs/thread-store/src/local/read_thread.rs index c538d938a3..5bcf411beb 100644 --- a/codex-rs/thread-store/src/local/read_thread.rs +++ b/codex-rs/thread-store/src/local/read_thread.rs @@ -176,12 +176,12 @@ async fn resolve_rollout_path( return Ok(Some(path)); } - let state_db_ctx = store.state_db().await; + let state_db = store.state_db(); if include_archived { match find_thread_path_by_id_str( store.config.codex_home.as_path(), &thread_id.to_string(), - state_db_ctx.as_deref(), + Some(state_db.as_ref()), ) .await .map_err(|err| ThreadStoreError::InvalidRequest { @@ -191,7 +191,7 @@ async fn resolve_rollout_path( None => find_archived_thread_path_by_id_str( store.config.codex_home.as_path(), &thread_id.to_string(), - state_db_ctx.as_deref(), + Some(state_db.as_ref()), ) .await .map_err(|err| ThreadStoreError::InvalidRequest { @@ -202,7 +202,7 @@ async fn resolve_rollout_path( find_thread_path_by_id_str( store.config.codex_home.as_path(), &thread_id.to_string(), - state_db_ctx.as_deref(), + Some(state_db.as_ref()), ) .await .map_err(|err| ThreadStoreError::InvalidRequest { @@ -260,8 +260,7 @@ async fn read_sqlite_metadata( store: &LocalThreadStore, thread_id: codex_protocol::ThreadId, ) -> Option { - let runtime = store.state_db().await?; - runtime.get_thread(thread_id).await.ok().flatten() + store.state_db().get_thread(thread_id).await.ok().flatten() } async fn stored_thread_from_sqlite_metadata( @@ -412,7 +411,9 @@ mod tests { use super::*; use crate::ThreadStore; use crate::local::LocalThreadStore; + use crate::local::test_support::init_test_state_db; use crate::local::test_support::test_config; + use crate::local::test_support::test_store; use crate::local::test_support::write_archived_session_file; use crate::local::test_support::write_session_file; use crate::local::test_support::write_session_file_with_fork; @@ -420,7 +421,7 @@ mod tests { #[tokio::test] async fn read_thread_returns_active_rollout_summary() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(205); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let active_path = @@ -448,7 +449,7 @@ mod tests { #[tokio::test] async fn read_thread_returns_rollout_path_summary() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(211); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let active_path = @@ -479,17 +480,12 @@ mod tests { async fn read_thread_by_rollout_path_prefers_sqlite_git_info() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(223); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let active_path = write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file"); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); let mut builder = ThreadMetadataBuilder::new( thread_id, active_path.clone(), @@ -527,7 +523,7 @@ mod tests { #[tokio::test] async fn read_thread_returns_archived_rollout_when_requested() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(207); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let archived_path = write_archived_session_file(home.path(), "2025-01-03T12-00-00", uuid) @@ -568,7 +564,7 @@ mod tests { #[tokio::test] async fn read_thread_prefers_active_rollout_over_archived() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(208); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let active_path = @@ -593,7 +589,7 @@ mod tests { #[tokio::test] async fn read_thread_returns_forked_from_id() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(209); let parent_uuid = Uuid::from_u128(210); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); @@ -626,17 +622,12 @@ mod tests { async fn read_thread_applies_sqlite_thread_name() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(212); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file"); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); let mut builder = ThreadMetadataBuilder::new(thread_id, rollout_path, Utc::now(), SessionSource::Cli); builder.model_provider = Some(config.default_model_provider_id.clone()); @@ -666,13 +657,8 @@ mod tests { async fn read_thread_preserves_rollout_cwd_when_sqlite_metadata_exists() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(224); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let day_dir = home.path().join("sessions/2025/01/03"); @@ -741,7 +727,7 @@ mod tests { #[tokio::test] async fn read_thread_uses_legacy_thread_name_when_sqlite_title_is_missing() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(213); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file"); @@ -765,6 +751,8 @@ mod tests { async fn read_thread_uses_sqlite_metadata_for_rollout_without_user_preview() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(217); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let day_dir = home.path().join("sessions/2025/01/03"); @@ -786,13 +774,6 @@ mod tests { }); writeln!(file, "{meta}").expect("write session meta"); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); let mut builder = ThreadMetadataBuilder::new( thread_id, rollout_path.clone(), @@ -835,18 +816,13 @@ mod tests { let home = TempDir::new().expect("temp dir"); let external = TempDir::new().expect("external temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(220); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file"); let stale_path = external.path().join("missing-rollout.jsonl"); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); let mut builder = ThreadMetadataBuilder::new( thread_id, stale_path.clone(), @@ -884,6 +860,8 @@ mod tests { let home = TempDir::new().expect("temp dir"); let external = TempDir::new().expect("external temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(221); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = @@ -891,13 +869,6 @@ mod tests { let other_uuid = Uuid::from_u128(222); let stale_path = write_session_file(external.path(), "2025-01-04T12-00-00", other_uuid) .expect("other session file"); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); let mut builder = ThreadMetadataBuilder::new(thread_id, stale_path, Utc::now(), SessionSource::Cli); builder.model_provider = Some("wrong-sqlite-provider".to_string()); @@ -929,7 +900,7 @@ mod tests { #[tokio::test] async fn read_thread_uses_session_meta_for_rollout_without_user_preview_or_sqlite_metadata() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(218); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let day_dir = home.path().join("sessions/2025/01/03"); @@ -984,18 +955,13 @@ mod tests { let home = TempDir::new().expect("temp dir"); let external = TempDir::new().expect("external temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(214); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = external .path() .join(format!("rollout-2025-01-03T12-00-00-{uuid}.jsonl")); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); let mut builder = ThreadMetadataBuilder::new( thread_id, rollout_path.clone(), @@ -1042,20 +1008,15 @@ mod tests { let home = TempDir::new().expect("temp dir"); let external = TempDir::new().expect("external temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(216); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = external .path() .join(format!("rollout-2025-01-03T12-00-00-{uuid}.jsonl")); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); let mut builder = ThreadMetadataBuilder::new(thread_id, rollout_path, Utc::now(), SessionSource::Cli); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); builder.archived_at = Some(Utc::now()); let mut metadata = builder.build(config.default_model_provider_id.as_str()); metadata.first_user_message = Some("Archived SQLite preview".to_string()); @@ -1098,17 +1059,12 @@ mod tests { async fn read_thread_sqlite_fallback_loads_archived_history() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(219); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let archived_path = write_archived_session_file(home.path(), "2025-01-03T12-00-00", uuid) .expect("archived session file"); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); let mut builder = ThreadMetadataBuilder::new( thread_id, archived_path.clone(), @@ -1144,7 +1100,7 @@ mod tests { #[tokio::test] async fn read_thread_fails_without_rollout() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(206); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); diff --git a/codex-rs/thread-store/src/local/test_support.rs b/codex-rs/thread-store/src/local/test_support.rs index 98321880ff..597014efe9 100644 --- a/codex-rs/thread-store/src/local/test_support.rs +++ b/codex-rs/thread-store/src/local/test_support.rs @@ -4,18 +4,34 @@ use std::path::Path; use std::path::PathBuf; use codex_rollout::ARCHIVED_SESSIONS_SUBDIR; +use codex_rollout::StateDbHandle; use uuid::Uuid; +use super::LocalThreadStore; use super::LocalThreadStoreConfig; pub(super) fn test_config(codex_home: &Path) -> LocalThreadStoreConfig { LocalThreadStoreConfig { codex_home: codex_home.to_path_buf(), - sqlite_home: codex_home.to_path_buf(), default_model_provider_id: "test-provider".to_string(), } } +pub(super) async fn init_test_state_db(config: &LocalThreadStoreConfig) -> StateDbHandle { + codex_state::StateRuntime::init( + config.codex_home.clone(), + config.default_model_provider_id.clone(), + ) + .await + .expect("state db should initialize") +} + +pub(super) async fn test_store(codex_home: &Path) -> LocalThreadStore { + let config = test_config(codex_home); + let state_db = init_test_state_db(&config).await; + LocalThreadStore::new(config, state_db) +} + pub(super) fn write_session_file(root: &Path, ts: &str, uuid: Uuid) -> std::io::Result { write_session_file_with( root, diff --git a/codex-rs/thread-store/src/local/unarchive_thread.rs b/codex-rs/thread-store/src/local/unarchive_thread.rs index ad41db69ac..7ac3f01a0f 100644 --- a/codex-rs/thread-store/src/local/unarchive_thread.rs +++ b/codex-rs/thread-store/src/local/unarchive_thread.rs @@ -17,11 +17,11 @@ pub(super) async fn unarchive_thread( params: ArchiveThreadParams, ) -> ThreadStoreResult { let thread_id = params.thread_id; - let state_db_ctx = store.state_db().await; + let state_db = store.state_db(); let archived_path = find_archived_thread_path_by_id_str( store.config.codex_home.as_path(), &thread_id.to_string(), - state_db_ctx.as_deref(), + Some(state_db.as_ref()), ) .await .map_err(|err| ThreadStoreError::InvalidRequest { @@ -73,11 +73,10 @@ pub(super) async fn unarchive_thread( message: format!("failed to update unarchived thread timestamp: {err}"), })?; - if let Some(ctx) = state_db_ctx { - let _ = ctx - .mark_unarchived(thread_id, restored_path.as_path()) - .await; - } + let _ = store + .state_db() + .mark_unarchived(thread_id, restored_path.as_path()) + .await; let item = read_thread_item_from_rollout(restored_path.clone()) .await @@ -112,13 +111,15 @@ mod tests { use super::*; use crate::ThreadStore; use crate::local::LocalThreadStore; + use crate::local::test_support::init_test_state_db; use crate::local::test_support::test_config; + use crate::local::test_support::test_store; use crate::local::test_support::write_archived_session_file; #[tokio::test] async fn unarchive_thread_restores_rollout_and_returns_updated_thread() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(203); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let archived_path = write_archived_session_file(home.path(), "2025-01-03T13-00-00", uuid) @@ -149,21 +150,12 @@ mod tests { async fn unarchive_thread_updates_sqlite_metadata_when_present() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(204); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let archived_path = write_archived_session_file(home.path(), "2025-01-03T13-00-00", uuid) .expect("archived session file"); - let runtime = codex_state::StateRuntime::init( - home.path().to_path_buf(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); - runtime - .mark_backfill_complete(/*last_watermark*/ None) - .await - .expect("backfill should be complete"); let mut builder = codex_state::ThreadMetadataBuilder::new( thread_id, archived_path.clone(), diff --git a/codex-rs/thread-store/src/local/update_thread_metadata.rs b/codex-rs/thread-store/src/local/update_thread_metadata.rs index 98ee788662..2c70c815ed 100644 --- a/codex-rs/thread-store/src/local/update_thread_metadata.rs +++ b/codex-rs/thread-store/src/local/update_thread_metadata.rs @@ -59,9 +59,8 @@ pub(super) async fn update_thread_metadata( .await?; } - let state_db_ctx = store.state_db().await; codex_rollout::state_db::reconcile_rollout( - state_db_ctx.as_deref(), + Some(store.state_db()).as_deref(), resolved_rollout_path.path.as_path(), store.config.default_model_provider_id.as_str(), /*builder*/ None, @@ -73,11 +72,7 @@ pub(super) async fn update_thread_metadata( let resolved_git_info = match git_info { Some(git_info) => { - let Some(state_db) = store.state_db().await else { - return Err(ThreadStoreError::Internal { - message: format!("sqlite state db unavailable for thread {thread_id}"), - }); - }; + let state_db = store.state_db(); let metadata = state_db .get_thread(thread_id) @@ -157,11 +152,7 @@ async fn apply_thread_git_info( branch: &Option, origin_url: &Option, ) -> ThreadStoreResult<()> { - let Some(state_db) = store.state_db().await else { - return Err(ThreadStoreError::Internal { - message: format!("sqlite state db unavailable for thread {thread_id}"), - }); - }; + let state_db = store.state_db(); let updated = state_db .update_thread_git_info( thread_id, @@ -307,11 +298,11 @@ async fn resolve_rollout_path( return Ok(ResolvedRolloutPath { path, archived }); } - let state_db_ctx = store.state_db().await; + let state_db = store.state_db(); let active_path = find_thread_path_by_id_str( store.config.codex_home.as_path(), &thread_id.to_string(), - state_db_ctx.as_deref(), + Some(state_db.as_ref()), ) .await .map_err(|err| ThreadStoreError::InvalidRequest { @@ -331,7 +322,7 @@ async fn resolve_rollout_path( find_archived_thread_path_by_id_str( store.config.codex_home.as_path(), &thread_id.to_string(), - state_db_ctx.as_deref(), + Some(state_db.as_ref()), ) .await .map_err(|err| ThreadStoreError::InvalidRequest { @@ -366,14 +357,16 @@ mod tests { use crate::ThreadPersistenceMetadata; use crate::ThreadStore; use crate::local::LocalThreadStore; + use crate::local::test_support::init_test_state_db; use crate::local::test_support::test_config; + use crate::local::test_support::test_store; use crate::local::test_support::write_archived_session_file; use crate::local::test_support::write_session_file; #[tokio::test] async fn update_thread_metadata_sets_name_on_active_rollout_and_indexes_name() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(301); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let path = @@ -408,18 +401,12 @@ mod tests { async fn update_thread_metadata_sets_memory_mode_on_active_rollout() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(302); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let path = write_session_file(home.path(), "2025-01-03T14-30-00", uuid).expect("session file"); - let runtime = codex_state::StateRuntime::init( - home.path().to_path_buf(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); - let thread = store .update_thread_metadata(UpdateThreadMetadataParams { thread_id, @@ -452,13 +439,8 @@ mod tests { let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let path = write_session_file(home.path(), "2025-01-03T18-30-00", uuid).expect("session file"); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); store .update_thread_metadata(UpdateThreadMetadataParams { @@ -517,7 +499,7 @@ mod tests { async fn update_thread_metadata_uses_live_rollout_path_for_external_resume() { let home = TempDir::new().expect("temp dir"); let external_home = TempDir::new().expect("external temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(307); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let path = write_session_file(external_home.path(), "2025-01-03T14-45-00", uuid) @@ -558,13 +540,8 @@ mod tests { async fn update_thread_metadata_sets_git_info() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config, Some(runtime)); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config, runtime); let uuid = Uuid::from_u128(309); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); write_session_file(home.path(), "2025-01-03T17-00-00", uuid).expect("session file"); @@ -601,13 +578,8 @@ mod tests { async fn update_thread_metadata_partially_updates_git_info() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config, Some(runtime)); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config, runtime); let uuid = Uuid::from_u128(310); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); write_session_file(home.path(), "2025-01-03T17-30-00", uuid).expect("session file"); @@ -659,13 +631,8 @@ mod tests { async fn update_thread_metadata_clears_git_info_fields() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(311); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let path = @@ -829,7 +796,7 @@ mod tests { #[tokio::test] async fn update_thread_metadata_rejects_mismatched_session_meta_id() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let filename_uuid = Uuid::from_u128(303); let metadata_uuid = Uuid::from_u128(304); let thread_id = ThreadId::from_string(&filename_uuid.to_string()).expect("valid thread id"); @@ -861,7 +828,7 @@ mod tests { #[tokio::test] async fn update_thread_metadata_rejects_multi_field_patch_without_partial_write() { let home = TempDir::new().expect("temp dir"); - let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let store = test_store(home.path()).await; let uuid = Uuid::from_u128(305); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let path = @@ -896,21 +863,12 @@ mod tests { async fn update_thread_metadata_keeps_archived_thread_archived_in_sqlite() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(306); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let archived_path = write_archived_session_file(home.path(), "2025-01-03T16-00-00", uuid) .expect("archived session file"); - let runtime = codex_state::StateRuntime::init( - home.path().to_path_buf(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); - runtime - .mark_backfill_complete(/*last_watermark*/ None) - .await - .expect("backfill should be complete"); codex_rollout::state_db::reconcile_rollout( Some(runtime.as_ref()), archived_path.as_path(), @@ -959,21 +917,12 @@ mod tests { async fn update_thread_metadata_keeps_live_archived_thread_archived_in_sqlite() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); + let runtime = init_test_state_db(&config).await; + let store = LocalThreadStore::new(config.clone(), runtime.clone()); let uuid = Uuid::from_u128(308); let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let archived_path = write_archived_session_file(home.path(), "2025-01-03T16-30-00", uuid) .expect("archived session file"); - let runtime = codex_state::StateRuntime::init( - home.path().to_path_buf(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); - runtime - .mark_backfill_complete(/*last_watermark*/ None) - .await - .expect("backfill should be complete"); codex_rollout::state_db::reconcile_rollout( Some(runtime.as_ref()), archived_path.as_path(), diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index 5622c59f65..b2e92a19b4 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -881,38 +881,40 @@ pub async fn run_main( AppServerTarget::Remote { .. } => state_db::get_state_db(&config).await, }; - let effective_toml = config.config_layer_stack.effective_config(); - match effective_toml.try_into() { - Ok(config_toml) => { - match crate::legacy_core::personality_migration::maybe_migrate_personality( - &config.codex_home, - &config_toml, - state_db.clone(), - ) - .await - { - Ok( - crate::legacy_core::personality_migration::PersonalityMigrationStatus::Applied, - ) => { - config = load_config_or_exit( - cli_kv_overrides.clone(), - overrides.clone(), - cloud_requirements.clone(), - ) - .await; - } - Ok( - crate::legacy_core::personality_migration::PersonalityMigrationStatus::SkippedMarker - | crate::legacy_core::personality_migration::PersonalityMigrationStatus::SkippedExplicitPersonality - | crate::legacy_core::personality_migration::PersonalityMigrationStatus::SkippedNoSessions, - ) => {} - Err(err) => { - tracing::warn!(error = %err, "failed to run personality migration"); + if let Some(state_db) = state_db.clone() { + let effective_toml = config.config_layer_stack.effective_config(); + match effective_toml.try_into() { + Ok(config_toml) => { + match crate::legacy_core::personality_migration::maybe_migrate_personality( + &config.codex_home, + &config_toml, + state_db, + ) + .await + { + Ok( + crate::legacy_core::personality_migration::PersonalityMigrationStatus::Applied, + ) => { + config = load_config_or_exit( + cli_kv_overrides.clone(), + overrides.clone(), + cloud_requirements.clone(), + ) + .await; + } + Ok( + crate::legacy_core::personality_migration::PersonalityMigrationStatus::SkippedMarker + | crate::legacy_core::personality_migration::PersonalityMigrationStatus::SkippedExplicitPersonality + | crate::legacy_core::personality_migration::PersonalityMigrationStatus::SkippedNoSessions, + ) => {} + Err(err) => { + tracing::warn!(error = %err, "failed to run personality migration"); + } } } - } - Err(err) => { - tracing::warn!(error = %err, "failed to deserialize config for personality migration"); + Err(err) => { + tracing::warn!(error = %err, "failed to deserialize config for personality migration"); + } } } diff --git a/codex-rs/tui/src/resume_picker.rs b/codex-rs/tui/src/resume_picker.rs index 16c3d19071..e06ccfd118 100644 --- a/codex-rs/tui/src/resume_picker.rs +++ b/codex-rs/tui/src/resume_picker.rs @@ -5733,6 +5733,7 @@ session_picker_view = "dense" name: None, turns: vec![codex_app_server_protocol::Turn { id: String::from("turn-1"), + items_view: codex_app_server_protocol::TurnItemsView::Full, items: vec![ ThreadItem::UserMessage { id: String::from("user-1"), @@ -5798,6 +5799,7 @@ session_picker_view = "dense" name: None, turns: vec![codex_app_server_protocol::Turn { id: String::from("turn-1"), + items_view: codex_app_server_protocol::TurnItemsView::Full, items: vec![ThreadItem::Reasoning { id: String::from("reasoning-1"), summary: Vec::new(), @@ -5853,6 +5855,7 @@ session_picker_view = "dense" name: None, turns: vec![codex_app_server_protocol::Turn { id: String::from("turn-1"), + items_view: codex_app_server_protocol::TurnItemsView::Full, items: vec![ThreadItem::Reasoning { id: String::from("reasoning-1"), summary: vec![String::from("public summary")], From d0f9d5eba2004c9a2057bc66baff5fe67ff60ff6 Mon Sep 17 00:00:00 2001 From: Michael Zeng Date: Tue, 5 May 2026 15:01:48 -0700 Subject: [PATCH 3/7] Add cloud executor registration to exec-server (#19575) ## Summary This PR adds the first `codex-rs` milestone for remote-exec e2e: a local `codex exec-server` can now register itself with `codex-cloud-environments` and attach to the returned rendezvous websocket. At a high level, `codex exec-server --cloud ...` now: - loads ChatGPT auth from normal Codex config - registers an executor with `codex-cloud-environments` - receives a signed rendezvous websocket URL - serves the existing exec-server JSON-RPC protocol over that websocket ## What Changed - Added `--cloud`, `--cloud-base-url`, `--cloud-environment-id`, and `--cloud-name` to `codex exec-server` - Added a new `exec-server/src/cloud.rs` module that handles: - registration requests - auth/header setup - bounded auth retry on `401/403` - reconnect/backoff after websocket disconnects - Reused the existing `ConnectionProcessor` / `ExecServerHandler` path so cloud mode serves the same exec/filesystem RPC surface as local websocket mode - Added cloud-specific error variants and minimal docs for the new mode ## Testing Manual e2e test that fully goes through exec server flow with our codex cloud agent as orchestrator --- codex-rs/Cargo.lock | 2 + codex-rs/cli/src/main.rs | 38 ++- codex-rs/exec-server/Cargo.toml | 4 +- codex-rs/exec-server/README.md | 7 + codex-rs/exec-server/src/client.rs | 12 + codex-rs/exec-server/src/lib.rs | 4 + codex-rs/exec-server/src/remote.rs | 392 +++++++++++++++++++++++++++++ codex-rs/exec-server/src/server.rs | 1 + 8 files changed, 452 insertions(+), 8 deletions(-) create mode 100644 codex-rs/exec-server/src/remote.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 7819fa223a..cd3e93c9cb 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2701,6 +2701,7 @@ dependencies = [ "serde", "serde_json", "serial_test", + "sha2", "tempfile", "test-case", "thiserror 2.0.18", @@ -2709,6 +2710,7 @@ dependencies = [ "tokio-util", "tracing", "uuid", + "wiremock", ] [[package]] diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index 5b5ad7a2a3..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) } 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; From ee02cf26d684c3e00f20e19225f8e6e6c176cb62 Mon Sep 17 00:00:00 2001 From: Tom Date: Tue, 5 May 2026 15:25:19 -0700 Subject: [PATCH 4/7] codex: use ThreadStore history for core review forks (#20577) - fork loaded parent threads from `ThreadStore` history in core agent control paths - migrate guardian review fork history to loaded session history instead of rereading rollout files ## Verification - `cargo test -p codex-core spawn_agent_fork` --- codex-rs/core/src/agent/control.rs | 12 ++++-------- codex-rs/core/src/guardian/review_session.rs | 10 ++++------ codex-rs/core/src/rollout.rs | 1 + 3 files changed, 9 insertions(+), 14 deletions(-) diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 3917ad5e95..3c98fd15d7 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -354,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 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/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; } From 84fadf895946554f24f1a59e552941d52c6b25ff Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Tue, 5 May 2026 15:44:05 -0700 Subject: [PATCH 5/7] linux-sandbox: use standalone bundled bwrap --- .github/workflows/bazel.yml | 16 + codex-rs/Cargo.lock | 12 +- codex-rs/Cargo.toml | 1 + codex-rs/bwrap/BUILD.bazel | 35 ++ codex-rs/bwrap/Cargo.toml | 19 + codex-rs/bwrap/build.rs | 99 +++++ codex-rs/bwrap/config.h | 1 + codex-rs/bwrap/src/main.rs | 45 +++ codex-rs/core/README.md | 16 +- codex-rs/linux-sandbox/BUILD.bazel | 30 -- codex-rs/linux-sandbox/Cargo.toml | 5 +- codex-rs/linux-sandbox/README.md | 10 +- codex-rs/linux-sandbox/build.rs | 110 +----- codex-rs/linux-sandbox/config.h | 3 - codex-rs/linux-sandbox/src/launcher.rs | 338 ++++++++++++++++-- codex-rs/linux-sandbox/src/lib.rs | 2 - codex-rs/linux-sandbox/src/vendored_bwrap.rs | 78 ---- .../linux-sandbox/tests/suite/landlock.rs | 2 +- .../tests/suite/managed_proxy.rs | 4 +- codex-rs/sandboxing/src/bwrap.rs | 2 +- codex-rs/tui/src/resume_picker.rs | 3 - 21 files changed, 558 insertions(+), 273 deletions(-) create mode 100644 codex-rs/bwrap/BUILD.bazel create mode 100644 codex-rs/bwrap/Cargo.toml create mode 100644 codex-rs/bwrap/build.rs create mode 100644 codex-rs/bwrap/config.h create mode 100644 codex-rs/bwrap/src/main.rs delete mode 100644 codex-rs/linux-sandbox/config.h delete mode 100644 codex-rs/linux-sandbox/src/vendored_bwrap.rs diff --git a/.github/workflows/bazel.yml b/.github/workflows/bazel.yml index aed3f7a1b3..fa4b3f20d7 100644 --- a/.github/workflows/bazel.yml +++ b/.github/workflows/bazel.yml @@ -371,6 +371,22 @@ jobs: -- \ "${bazel_targets[@]}" + - name: Smoke-test Bazel-built bwrap + if: runner.os == 'Linux' + env: + BUILDBUDDY_API_KEY: ${{ secrets.BUILDBUDDY_API_KEY }} + shell: bash + run: | + ./.github/scripts/run-bazel-ci.sh \ + --remote-download-toplevel \ + --print-failed-action-summary \ + -- \ + run \ + --build_metadata=COMMIT_SHA=${GITHUB_SHA} \ + --build_metadata=TAG_job=verify-bwrap \ + -- \ + //codex-rs/bwrap:bwrap -- --version + - name: Upload Bazel execution logs if: always() && !cancelled() continue-on-error: true diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index cd3e93c9cb..1b92b5afd0 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2125,6 +2125,15 @@ dependencies = [ "serde_with", ] +[[package]] +name = "codex-bwrap" +version = "0.0.0" +dependencies = [ + "cc", + "libc", + "pkg-config", +] + [[package]] name = "codex-chatgpt" version = "0.0.0" @@ -2905,7 +2914,6 @@ dependencies = [ name = "codex-linux-sandbox" version = "0.0.0" dependencies = [ - "cc", "clap", "codex-core", "codex-process-hardening", @@ -2915,11 +2923,11 @@ dependencies = [ "globset", "landlock", "libc", - "pkg-config", "pretty_assertions", "seccompiler", "serde", "serde_json", + "sha2", "tempfile", "tokio", "url", diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index ce3e91626d..a58655ed20 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -5,6 +5,7 @@ members = [ "agent-graph-store", "agent-identity", "backend-client", + "bwrap", "ansi-escape", "async-utils", "app-server", diff --git a/codex-rs/bwrap/BUILD.bazel b/codex-rs/bwrap/BUILD.bazel new file mode 100644 index 0000000000..3d0b89b966 --- /dev/null +++ b/codex-rs/bwrap/BUILD.bazel @@ -0,0 +1,35 @@ +load("@rules_cc//cc:defs.bzl", "cc_library") +load("//:defs.bzl", "codex_rust_crate") + +codex_rust_crate( + name = "bwrap", + crate_name = "codex_bwrap", + # Bazel wires vendored bubblewrap + libcap via :bwrap-ffi below and sets + # bwrap_available explicitly, so we skip Cargo's build.rs in Bazel builds. + build_script_enabled = False, + deps_extra = select({ + "@platforms//os:linux": [":bwrap-ffi"], + "//conditions:default": [], + }), + rustc_flags_extra = select({ + "@platforms//os:linux": ["--cfg=bwrap_available"], + "//conditions:default": [], + }), +) + +cc_library( + name = "bwrap-ffi", + srcs = ["//codex-rs/vendor:bubblewrap_c_sources"], + hdrs = [ + "config.h", + "//codex-rs/vendor:bubblewrap_headers", + ], + copts = [ + "-D_GNU_SOURCE", + "-Dmain=bwrap_main", + ], + includes = ["."], + deps = ["@libcap//:libcap"], + target_compatible_with = ["@platforms//os:linux"], + visibility = ["//visibility:private"], +) diff --git a/codex-rs/bwrap/Cargo.toml b/codex-rs/bwrap/Cargo.toml new file mode 100644 index 0000000000..ed7010c8fd --- /dev/null +++ b/codex-rs/bwrap/Cargo.toml @@ -0,0 +1,19 @@ +[package] +name = "codex-bwrap" +version.workspace = true +edition.workspace = true +license.workspace = true + +[[bin]] +name = "bwrap" +path = "src/main.rs" + +[lints] +workspace = true + +[target.'cfg(target_os = "linux")'.dependencies] +libc = { workspace = true } + +[build-dependencies] +cc = "1" +pkg-config = "0.3" diff --git a/codex-rs/bwrap/build.rs b/codex-rs/bwrap/build.rs new file mode 100644 index 0000000000..40c2712839 --- /dev/null +++ b/codex-rs/bwrap/build.rs @@ -0,0 +1,99 @@ +use std::env; +use std::path::Path; +use std::path::PathBuf; + +fn main() { + println!("cargo:rustc-check-cfg=cfg(bwrap_available)"); + println!("cargo:rerun-if-env-changed=CODEX_BWRAP_SOURCE_DIR"); + println!("cargo:rerun-if-env-changed=PKG_CONFIG_ALLOW_CROSS"); + println!("cargo:rerun-if-env-changed=PKG_CONFIG_PATH"); + println!("cargo:rerun-if-env-changed=PKG_CONFIG_SYSROOT_DIR"); + println!("cargo:rerun-if-env-changed=CODEX_SKIP_BWRAP_BUILD"); + + let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap_or_default()); + let vendor_dir = manifest_dir.join("../vendor/bubblewrap"); + for source in ["bubblewrap.c", "bind-mount.c", "network.c", "utils.c"] { + println!( + "cargo:rerun-if-changed={}", + vendor_dir.join(source).display() + ); + } + + let target_os = env::var("CARGO_CFG_TARGET_OS").unwrap_or_default(); + if target_os != "linux" || env::var_os("CODEX_SKIP_BWRAP_BUILD").is_some() { + return; + } + + if let Err(err) = try_build_bwrap() { + panic!("failed to compile bubblewrap for Linux target: {err}"); + } +} + +fn try_build_bwrap() -> Result<(), String> { + let manifest_dir = + PathBuf::from(env::var("CARGO_MANIFEST_DIR").map_err(|err| err.to_string())?); + let out_dir = PathBuf::from(env::var("OUT_DIR").map_err(|err| err.to_string())?); + let src_dir = resolve_bwrap_source_dir(&manifest_dir)?; + let libcap = pkg_config::Config::new() + .probe("libcap") + .map_err(|err| format!("libcap not available via pkg-config: {err}"))?; + + let config_h = out_dir.join("config.h"); + std::fs::write( + &config_h, + r#"#pragma once +#define PACKAGE_STRING "bubblewrap built for Codex" +"#, + ) + .map_err(|err| format!("failed to write {}: {err}", config_h.display()))?; + + let mut build = cc::Build::new(); + build + .file(src_dir.join("bubblewrap.c")) + .file(src_dir.join("bind-mount.c")) + .file(src_dir.join("network.c")) + .file(src_dir.join("utils.c")) + .include(&out_dir) + .include(&src_dir) + .define("_GNU_SOURCE", None) + // Rename `main` so the Rust wrapper can expose the Cargo-built binary. + .define("main", Some("bwrap_main")); + for include_path in libcap.include_paths { + // Use -idirafter so target sysroot headers win (musl cross builds), + // while still allowing libcap headers from the host toolchain. + build.flag(format!("-idirafter{}", include_path.display())); + } + + build.compile("standalone_bwrap"); + println!("cargo:rustc-cfg=bwrap_available"); + Ok(()) +} + +/// Resolve the bubblewrap source directory used for build-time compilation. +/// +/// Priority: +/// 1. `CODEX_BWRAP_SOURCE_DIR` points at an existing bubblewrap checkout. +/// 2. The vendored bubblewrap tree under `codex-rs/vendor/bubblewrap`. +fn resolve_bwrap_source_dir(manifest_dir: &Path) -> Result { + if let Ok(path) = env::var("CODEX_BWRAP_SOURCE_DIR") { + let src_dir = PathBuf::from(path); + if src_dir.exists() { + return Ok(src_dir); + } + return Err(format!( + "CODEX_BWRAP_SOURCE_DIR was set but does not exist: {}", + src_dir.display() + )); + } + + let vendor_dir = manifest_dir.join("../vendor/bubblewrap"); + if vendor_dir.exists() { + return Ok(vendor_dir); + } + + Err(format!( + "expected vendored bubblewrap at {}, but it was not found.\n\ +Set CODEX_BWRAP_SOURCE_DIR to an existing checkout or vendor bubblewrap under codex-rs/vendor.", + vendor_dir.display() + )) +} diff --git a/codex-rs/bwrap/config.h b/codex-rs/bwrap/config.h new file mode 100644 index 0000000000..f73932a0f8 --- /dev/null +++ b/codex-rs/bwrap/config.h @@ -0,0 +1 @@ +#define PACKAGE_STRING "bubblewrap built for Codex" diff --git a/codex-rs/bwrap/src/main.rs b/codex-rs/bwrap/src/main.rs new file mode 100644 index 0000000000..09c624aa9e --- /dev/null +++ b/codex-rs/bwrap/src/main.rs @@ -0,0 +1,45 @@ +#[cfg(all(target_os = "linux", bwrap_available))] +fn main() { + use std::ffi::CStr; + use std::ffi::CString; + use std::os::raw::c_char; + use std::os::unix::ffi::OsStrExt; + + unsafe extern "C" { + fn bwrap_main(argc: libc::c_int, argv: *const *const c_char) -> libc::c_int; + } + + let cstrings = std::env::args_os() + .map(|arg| { + CString::new(arg.as_os_str().as_bytes()) + .unwrap_or_else(|err| panic!("failed to convert argv to CString: {err}")) + }) + .collect::>(); + let mut argv_ptrs = cstrings + .iter() + .map(CString::as_c_str) + .map(CStr::as_ptr) + .collect::>(); + argv_ptrs.push(std::ptr::null()); + + // SAFETY: We provide a null-terminated argv vector whose pointers remain + // valid for the duration of the call. + let exit_code = unsafe { bwrap_main(cstrings.len() as libc::c_int, argv_ptrs.as_ptr()) }; + std::process::exit(exit_code); +} + +#[cfg(all(target_os = "linux", not(bwrap_available)))] +fn main() { + panic!( + r#"bubblewrap is not available in this build. +Notes: +- ensure the target OS is Linux +- libcap headers must be available via pkg-config +- bubblewrap sources expected at codex-rs/vendor/bubblewrap (default)"# + ); +} + +#[cfg(not(target_os = "linux"))] +fn main() { + panic!("bwrap is only supported on Linux"); +} diff --git a/codex-rs/core/README.md b/codex-rs/core/README.md index be222a1673..3283ba2c3e 100644 --- a/codex-rs/core/README.md +++ b/codex-rs/core/README.md @@ -39,14 +39,14 @@ The Linux sandbox helper prefers the first `bwrap` found on `PATH` outside the current working directory whenever it is available. If `bwrap` is present but too old to support `--argv0`, the helper keeps using system bubblewrap and switches to a no-`--argv0` compatibility path for the inner re-exec. If -`bwrap` is missing, it falls back to the vendored bubblewrap path compiled into -the binary and Codex surfaces a startup warning through its normal notification -path instead of printing directly from the sandbox helper. Codex also surfaces -a startup warning when bubblewrap cannot create user namespaces. WSL2 uses the -normal Linux bubblewrap path. WSL1 is not supported for bubblewrap sandboxing -because it cannot create the required user namespaces, so Codex rejects -sandboxed shell commands that would enter the bubblewrap path before invoking -`bwrap`. +`bwrap` is missing, it falls back to the bundled `codex-resources/bwrap` +binary shipped with Codex and Codex surfaces a startup warning through its +normal notification path instead of printing directly from the sandbox helper. +Codex also surfaces a startup warning when bubblewrap cannot create user +namespaces. WSL2 uses the normal Linux bubblewrap path. WSL1 is not supported +for bubblewrap sandboxing because it cannot create the required user +namespaces, so Codex rejects sandboxed shell commands that would enter the +bubblewrap path before invoking `bwrap`. ### Windows diff --git a/codex-rs/linux-sandbox/BUILD.bazel b/codex-rs/linux-sandbox/BUILD.bazel index 87ca8ba066..aeef302ab3 100644 --- a/codex-rs/linux-sandbox/BUILD.bazel +++ b/codex-rs/linux-sandbox/BUILD.bazel @@ -1,36 +1,6 @@ -load("@rules_cc//cc:defs.bzl", "cc_library") load("//:defs.bzl", "codex_rust_crate") codex_rust_crate( name = "linux-sandbox", crate_name = "codex_linux_sandbox", - # Bazel wires vendored bubblewrap + libcap via :vendored-bwrap-ffi below - # and sets vendored_bwrap_available explicitly, so we skip Cargo's - # build.rs in Bazel builds. - build_script_enabled = False, - deps_extra = select({ - "@platforms//os:linux": [":vendored-bwrap-ffi"], - "//conditions:default": [], - }), - rustc_flags_extra = select({ - "@platforms//os:linux": ["--cfg=vendored_bwrap_available"], - "//conditions:default": [], - }), -) - -cc_library( - name = "vendored-bwrap-ffi", - srcs = ["//codex-rs/vendor:bubblewrap_c_sources"], - hdrs = [ - "config.h", - "//codex-rs/vendor:bubblewrap_headers", - ], - copts = [ - "-D_GNU_SOURCE", - "-Dmain=bwrap_main", - ], - includes = ["."], - deps = ["@libcap//:libcap"], - target_compatible_with = ["@platforms//os:linux"], - visibility = ["//visibility:private"], ) diff --git a/codex-rs/linux-sandbox/Cargo.toml b/codex-rs/linux-sandbox/Cargo.toml index 05967661e2..aaacee27fa 100644 --- a/codex-rs/linux-sandbox/Cargo.toml +++ b/codex-rs/linux-sandbox/Cargo.toml @@ -27,6 +27,7 @@ libc = { workspace = true } seccompiler = { workspace = true } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } +sha2 = { workspace = true } url = { workspace = true } [target.'cfg(target_os = "linux")'.dev-dependencies] @@ -40,7 +41,3 @@ tokio = { workspace = true, features = [ "rt-multi-thread", "signal", ] } - -[build-dependencies] -cc = "1" -pkg-config = "0.3" diff --git a/codex-rs/linux-sandbox/README.md b/codex-rs/linux-sandbox/README.md index 5745f4816c..1b4c1e5aa7 100644 --- a/codex-rs/linux-sandbox/README.md +++ b/codex-rs/linux-sandbox/README.md @@ -12,10 +12,10 @@ outside the current working directory whenever it is available. If `bwrap` is present but too old to support `--argv0`, the helper keeps using system bubblewrap and switches to a no-`--argv0` compatibility path for the inner re-exec. If `bwrap` is missing, -the helper falls back to the vendored bubblewrap path compiled into this -binary. +the helper falls back to the bundled `codex-resources/bwrap` binary shipped +with Codex. Codex also surfaces a startup warning when `bwrap` is missing so users know it -is falling back to the vendored helper. Codex surfaces the same startup warning +is falling back to the bundled helper. Codex surfaces the same startup warning path when bubblewrap cannot create user namespaces. WSL2 follows the normal Linux bubblewrap path. WSL1 is not supported for bubblewrap sandboxing because it cannot create the required user namespaces, so Codex rejects sandboxed shell @@ -28,8 +28,8 @@ commands that would enter the bubblewrap path. helper uses it. - If `bwrap` is present but too old to support `--argv0`, the helper uses a no-`--argv0` compatibility path for the inner re-exec. -- If `bwrap` is missing, the helper falls back to the vendored bubblewrap - path. +- If `bwrap` is missing, the helper falls back to the bundled + `codex-resources/bwrap` path. - If `bwrap` is missing, Codex also surfaces a startup warning instead of printing directly from the sandbox helper. - If bubblewrap cannot create user namespaces, Codex surfaces a startup warning diff --git a/codex-rs/linux-sandbox/build.rs b/codex-rs/linux-sandbox/build.rs index a2b4ca86ca..968cfc7e67 100644 --- a/codex-rs/linux-sandbox/build.rs +++ b/codex-rs/linux-sandbox/build.rs @@ -1,111 +1,3 @@ -use std::env; -use std::path::Path; -use std::path::PathBuf; - fn main() { - // Tell rustc/clippy that this is an expected cfg value. - println!("cargo:rustc-check-cfg=cfg(vendored_bwrap_available)"); - println!("cargo:rerun-if-env-changed=CODEX_BWRAP_SOURCE_DIR"); - println!("cargo:rerun-if-env-changed=PKG_CONFIG_ALLOW_CROSS"); - println!("cargo:rerun-if-env-changed=PKG_CONFIG_PATH"); - println!("cargo:rerun-if-env-changed=PKG_CONFIG_SYSROOT_DIR"); - println!("cargo:rerun-if-env-changed=CODEX_SKIP_VENDORED_BWRAP"); - - // Rebuild if the vendored bwrap sources change. - let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap_or_default()); - let vendor_dir = manifest_dir.join("../vendor/bubblewrap"); - println!( - "cargo:rerun-if-changed={}", - vendor_dir.join("bubblewrap.c").display() - ); - println!( - "cargo:rerun-if-changed={}", - vendor_dir.join("bind-mount.c").display() - ); - println!( - "cargo:rerun-if-changed={}", - vendor_dir.join("network.c").display() - ); - println!( - "cargo:rerun-if-changed={}", - vendor_dir.join("utils.c").display() - ); - - let target_os = env::var("CARGO_CFG_TARGET_OS").unwrap_or_default(); - if target_os != "linux" || env::var_os("CODEX_SKIP_VENDORED_BWRAP").is_some() { - return; - } - - if let Err(err) = try_build_vendored_bwrap() { - panic!("failed to compile vendored bubblewrap for Linux target: {err}"); - } -} - -fn try_build_vendored_bwrap() -> Result<(), String> { - let manifest_dir = - PathBuf::from(env::var("CARGO_MANIFEST_DIR").map_err(|err| err.to_string())?); - let out_dir = PathBuf::from(env::var("OUT_DIR").map_err(|err| err.to_string())?); - let src_dir = resolve_bwrap_source_dir(&manifest_dir)?; - let libcap = pkg_config::Config::new() - .probe("libcap") - .map_err(|err| format!("libcap not available via pkg-config: {err}"))?; - - let config_h = out_dir.join("config.h"); - std::fs::write( - &config_h, - r#"#pragma once -#define PACKAGE_STRING "bubblewrap built at codex build-time" -"#, - ) - .map_err(|err| format!("failed to write {}: {err}", config_h.display()))?; - - let mut build = cc::Build::new(); - build - .file(src_dir.join("bubblewrap.c")) - .file(src_dir.join("bind-mount.c")) - .file(src_dir.join("network.c")) - .file(src_dir.join("utils.c")) - .include(&out_dir) - .include(&src_dir) - .define("_GNU_SOURCE", None) - // Rename `main` so we can call it via FFI. - .define("main", Some("bwrap_main")); - for include_path in libcap.include_paths { - // Use -idirafter so target sysroot headers win (musl cross builds), - // while still allowing libcap headers from the host toolchain. - build.flag(format!("-idirafter{}", include_path.display())); - } - - build.compile("build_time_bwrap"); - println!("cargo:rustc-cfg=vendored_bwrap_available"); - Ok(()) -} - -/// Resolve the bubblewrap source directory used for build-time compilation. -/// -/// Priority: -/// 1. `CODEX_BWRAP_SOURCE_DIR` points at an existing bubblewrap checkout. -/// 2. The vendored bubblewrap tree under `codex-rs/vendor/bubblewrap`. -fn resolve_bwrap_source_dir(manifest_dir: &Path) -> Result { - if let Ok(path) = env::var("CODEX_BWRAP_SOURCE_DIR") { - let src_dir = PathBuf::from(path); - if src_dir.exists() { - return Ok(src_dir); - } - return Err(format!( - "CODEX_BWRAP_SOURCE_DIR was set but does not exist: {}", - src_dir.display() - )); - } - - let vendor_dir = manifest_dir.join("../vendor/bubblewrap"); - if vendor_dir.exists() { - return Ok(vendor_dir); - } - - Err(format!( - "expected vendored bubblewrap at {}, but it was not found.\n\ -Set CODEX_BWRAP_SOURCE_DIR to an existing checkout or vendor bubblewrap under codex-rs/vendor.", - vendor_dir.display() - )) + println!("cargo:rerun-if-env-changed=CODEX_BWRAP_SHA256"); } diff --git a/codex-rs/linux-sandbox/config.h b/codex-rs/linux-sandbox/config.h deleted file mode 100644 index f08aa6fcee..0000000000 --- a/codex-rs/linux-sandbox/config.h +++ /dev/null @@ -1,3 +0,0 @@ -#pragma once - -#define PACKAGE_STRING "bubblewrap built at codex build-time" diff --git a/codex-rs/linux-sandbox/src/launcher.rs b/codex-rs/linux-sandbox/src/launcher.rs index 577ef52405..e7320995fa 100644 --- a/codex-rs/linux-sandbox/src/launcher.rs +++ b/codex-rs/linux-sandbox/src/launcher.rs @@ -1,20 +1,29 @@ +use std::ffi::CStr; use std::ffi::CString; use std::fs::File; +use std::io::Read; use std::os::fd::AsRawFd; use std::os::raw::c_char; use std::os::unix::ffi::OsStrExt; +use std::os::unix::fs::PermissionsExt; use std::path::Path; +use std::path::PathBuf; use std::process::Command; use std::sync::OnceLock; -use crate::vendored_bwrap::exec_vendored_bwrap; use codex_sandboxing::find_system_bwrap_in_path; use codex_utils_absolute_path::AbsolutePathBuf; +use sha2::Digest as _; +use sha2::Sha256; + +const SHA256_HEX_LEN: usize = 64; +const NULL_SHA256_DIGEST: [u8; 32] = [0; 32]; #[derive(Debug, Clone, PartialEq, Eq)] enum BubblewrapLauncher { System(SystemBwrapLauncher), - Vendored, + Bundled(BundledBwrapLauncher), + Unavailable, } #[derive(Debug, Clone, PartialEq, Eq)] @@ -23,6 +32,11 @@ struct SystemBwrapLauncher { supports_argv0: bool, } +#[derive(Debug, Clone, PartialEq, Eq)] +struct BundledBwrapLauncher { + program: AbsolutePathBuf, +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] struct SystemBwrapCapabilities { supports_argv0: bool, @@ -34,30 +48,46 @@ pub(crate) fn exec_bwrap(argv: Vec, preserved_files: Vec) -> ! { BubblewrapLauncher::System(launcher) => { exec_system_bwrap(&launcher.program, argv, preserved_files) } - BubblewrapLauncher::Vendored => exec_vendored_bwrap(argv, preserved_files), + BubblewrapLauncher::Bundled(launcher) => { + exec_bundled_bwrap(&launcher.program, argv, preserved_files) + } + BubblewrapLauncher::Unavailable => { + panic!( + "bubblewrap is unavailable: no system bwrap was found on PATH and no bundled \ + codex-resources/bwrap binary was found next to the Codex executable" + ) + } } } fn preferred_bwrap_launcher() -> BubblewrapLauncher { static LAUNCHER: OnceLock = OnceLock::new(); LAUNCHER - .get_or_init(|| match find_system_bwrap_in_path() { - Some(path) => preferred_bwrap_launcher_for_path(&path), - None => BubblewrapLauncher::Vendored, + .get_or_init(|| { + if let Some(path) = find_system_bwrap_in_path() + && let Some(launcher) = system_bwrap_launcher_for_path(&path) + { + return BubblewrapLauncher::System(launcher); + } + + match bundled_bwrap_launcher() { + Some(launcher) => BubblewrapLauncher::Bundled(launcher), + None => BubblewrapLauncher::Unavailable, + } }) .clone() } -fn preferred_bwrap_launcher_for_path(system_bwrap_path: &Path) -> BubblewrapLauncher { - preferred_bwrap_launcher_for_path_with_probe(system_bwrap_path, system_bwrap_capabilities) +fn system_bwrap_launcher_for_path(system_bwrap_path: &Path) -> Option { + system_bwrap_launcher_for_path_with_probe(system_bwrap_path, system_bwrap_capabilities) } -fn preferred_bwrap_launcher_for_path_with_probe( +fn system_bwrap_launcher_for_path_with_probe( system_bwrap_path: &Path, system_bwrap_capabilities: impl FnOnce(&Path) -> Option, -) -> BubblewrapLauncher { +) -> Option { if !system_bwrap_path.is_file() { - return BubblewrapLauncher::Vendored; + return None; } let Some(SystemBwrapCapabilities { @@ -65,7 +95,7 @@ fn preferred_bwrap_launcher_for_path_with_probe( supports_perms: true, }) = system_bwrap_capabilities(system_bwrap_path) else { - return BubblewrapLauncher::Vendored; + return None; }; let system_bwrap_path = match AbsolutePathBuf::from_absolute_path(system_bwrap_path) { Ok(path) => path, @@ -74,7 +104,7 @@ fn preferred_bwrap_launcher_for_path_with_probe( system_bwrap_path.display() ), }; - BubblewrapLauncher::System(SystemBwrapLauncher { + Some(SystemBwrapLauncher { program: system_bwrap_path, supports_argv0, }) @@ -83,10 +113,50 @@ fn preferred_bwrap_launcher_for_path_with_probe( pub(crate) fn preferred_bwrap_supports_argv0() -> bool { match preferred_bwrap_launcher() { BubblewrapLauncher::System(launcher) => launcher.supports_argv0, - BubblewrapLauncher::Vendored => true, + BubblewrapLauncher::Bundled(_) | BubblewrapLauncher::Unavailable => true, } } +fn bundled_bwrap_launcher() -> Option { + let current_exe = std::env::current_exe().ok()?; + find_bundled_bwrap_for_exe(¤t_exe).map(|program| BundledBwrapLauncher { program }) +} + +fn find_bundled_bwrap_for_exe(exe: &Path) -> Option { + bundled_bwrap_candidates_for_exe(exe) + .into_iter() + .find(|candidate| is_executable_file(candidate)) + .map(|path| { + AbsolutePathBuf::from_absolute_path(&path).unwrap_or_else(|err| { + panic!( + "failed to normalize bundled bubblewrap path {}: {err}", + path.display() + ) + }) + }) +} + +fn bundled_bwrap_candidates_for_exe(exe: &Path) -> Vec { + let Some(exe_dir) = exe.parent() else { + return Vec::new(); + }; + + let mut candidates = Vec::new(); + candidates.push(exe_dir.join("codex-resources").join("bwrap")); + if let Some(package_target_dir) = exe_dir.parent() { + candidates.push(package_target_dir.join("codex-resources").join("bwrap")); + } + candidates.push(exe_dir.join("bwrap")); + candidates +} + +fn is_executable_file(path: &Path) -> bool { + let Ok(metadata) = path.metadata() else { + return false; + }; + metadata.is_file() && metadata.permissions().mode() & 0o111 != 0 +} + fn system_bwrap_capabilities(system_bwrap_path: &Path) -> Option { // bubblewrap added `--argv0` in v0.9.0: // https://github.com/containers/bubblewrap/releases/tag/v0.9.0 @@ -117,7 +187,11 @@ fn exec_system_bwrap( let program = CString::new(program.as_path().as_os_str().as_bytes()) .unwrap_or_else(|err| panic!("invalid system bubblewrap path: {err}")); let cstrings = argv_to_cstrings(&argv); - let mut argv_ptrs: Vec<*const c_char> = cstrings.iter().map(|arg| arg.as_ptr()).collect(); + let mut argv_ptrs: Vec<*const c_char> = cstrings + .iter() + .map(CString::as_c_str) + .map(CStr::as_ptr) + .collect(); argv_ptrs.push(std::ptr::null()); // SAFETY: `program` and every entry in `argv_ptrs` are valid C strings for @@ -129,6 +203,128 @@ fn exec_system_bwrap( panic!("failed to exec system bubblewrap {program_path}: {err}"); } +fn exec_bundled_bwrap( + program: &AbsolutePathBuf, + argv: Vec, + preserved_files: Vec, +) -> ! { + let bwrap_file = File::open(program.as_path()).unwrap_or_else(|err| { + panic!( + "failed to open bundled bubblewrap {}: {err}", + program.as_path().display() + ) + }); + verify_bundled_bwrap_digest( + &bwrap_file, + expected_bundled_bwrap_sha256(), + program.as_path(), + ) + .unwrap_or_else(|err| panic!("{err}")); + + make_files_inheritable(&preserved_files); + + let fd_path = format!("/proc/self/fd/{}", bwrap_file.as_raw_fd()); + let program_cstring = CString::new(fd_path.as_str()) + .unwrap_or_else(|err| panic!("invalid bundled bubblewrap fd path: {err}")); + let cstrings = argv_to_cstrings(&argv); + let mut argv_ptrs: Vec<*const c_char> = cstrings + .iter() + .map(CString::as_c_str) + .map(CStr::as_ptr) + .collect(); + argv_ptrs.push(std::ptr::null()); + + // SAFETY: `program_cstring` and every entry in `argv_ptrs` are valid C + // strings for the duration of the call. On success `execv` does not return. + unsafe { + libc::execv(program_cstring.as_ptr(), argv_ptrs.as_ptr()); + } + let err = std::io::Error::last_os_error(); + panic!( + "failed to exec bundled bubblewrap {} via {fd_path}: {err}", + program.as_path().display() + ); +} + +fn expected_bundled_bwrap_sha256() -> Option<[u8; 32]> { + static EXPECTED: OnceLock> = OnceLock::new(); + *EXPECTED.get_or_init(|| { + let Some(raw_digest) = option_env!("CODEX_BWRAP_SHA256") else { + return None; + }; + let digest = parse_sha256_hex(raw_digest) + .unwrap_or_else(|err| panic!("invalid CODEX_BWRAP_SHA256 value: {err}")); + (digest != NULL_SHA256_DIGEST).then_some(digest) + }) +} + +fn verify_bundled_bwrap_digest( + file: &File, + expected: Option<[u8; 32]>, + path: &Path, +) -> Result<(), String> { + let Some(expected) = expected else { + return Ok(()); + }; + + let mut file = file + .try_clone() + .map_err(|err| format!("failed to clone bundled bubblewrap fd: {err}"))?; + let mut hasher = Sha256::new(); + let mut buffer = [0_u8; 8192]; + loop { + let read = file.read(&mut buffer).map_err(|err| { + format!( + "failed to read bundled bubblewrap {} for digest verification: {err}", + path.display() + ) + })?; + if read == 0 { + break; + } + hasher.update(&buffer[..read]); + } + + let actual: [u8; 32] = hasher.finalize().into(); + if actual == expected { + return Ok(()); + } + + Err(format!( + "bundled bubblewrap digest mismatch for {}: expected sha256:{}, got sha256:{}", + path.display(), + bytes_to_hex(&expected), + bytes_to_hex(&actual), + )) +} + +fn parse_sha256_hex(raw: &str) -> Result<[u8; 32], String> { + if raw.len() != SHA256_HEX_LEN { + return Err(format!( + "expected {SHA256_HEX_LEN} hex characters, got {}", + raw.len() + )); + } + + let mut digest = [0_u8; 32]; + for (index, byte) in digest.iter_mut().enumerate() { + let start = index * 2; + *byte = u8::from_str_radix(&raw[start..start + 2], 16) + .map_err(|err| format!("invalid hex byte at offset {start}: {err}"))?; + } + Ok(digest) +} + +fn bytes_to_hex(bytes: &[u8; 32]) -> String { + const HEX: &[u8; 16] = b"0123456789abcdef"; + let mut hex = String::with_capacity(SHA256_HEX_LEN); + for byte in bytes { + hex.push(HEX[(byte >> 4) as usize] as char); + hex.push(HEX[(byte & 0x0f) as usize] as char); + } + hex +} + fn argv_to_cstrings(argv: &[String]) -> Vec { let mut cstrings: Vec = Vec::with_capacity(argv.len()); for arg in argv { @@ -170,7 +366,10 @@ fn clear_cloexec(fd: libc::c_int) { mod tests { use super::*; use pretty_assertions::assert_eq; + use std::fs; + use std::os::unix::fs::PermissionsExt; use tempfile::NamedTempFile; + use tempfile::tempdir; #[test] fn prefers_system_bwrap_when_help_lists_argv0() { @@ -179,13 +378,13 @@ mod tests { let expected = AbsolutePathBuf::from_absolute_path(fake_bwrap_path).expect("absolute"); assert_eq!( - preferred_bwrap_launcher_for_path_with_probe(fake_bwrap_path, |_| { + system_bwrap_launcher_for_path_with_probe(fake_bwrap_path, |_| { Some(SystemBwrapCapabilities { supports_argv0: true, supports_perms: true, }) }), - BubblewrapLauncher::System(SystemBwrapLauncher { + Some(SystemBwrapLauncher { program: expected, supports_argv0: true, }) @@ -198,13 +397,13 @@ mod tests { let fake_bwrap_path = fake_bwrap.path(); assert_eq!( - preferred_bwrap_launcher_for_path_with_probe(fake_bwrap_path, |_| { + system_bwrap_launcher_for_path_with_probe(fake_bwrap_path, |_| { Some(SystemBwrapCapabilities { supports_argv0: false, supports_perms: true, }) }), - BubblewrapLauncher::System(SystemBwrapLauncher { + Some(SystemBwrapLauncher { program: AbsolutePathBuf::from_absolute_path(fake_bwrap_path).expect("absolute"), supports_argv0: false, }) @@ -212,28 +411,108 @@ mod tests { } #[test] - fn falls_back_to_vendored_when_system_bwrap_lacks_perms() { + fn ignores_system_bwrap_when_system_bwrap_lacks_perms() { let fake_bwrap = NamedTempFile::new().expect("temp file"); assert_eq!( - preferred_bwrap_launcher_for_path_with_probe(fake_bwrap.path(), |_| { + system_bwrap_launcher_for_path_with_probe(fake_bwrap.path(), |_| { Some(SystemBwrapCapabilities { supports_argv0: false, supports_perms: false, }) }), - BubblewrapLauncher::Vendored + None ); } #[test] - fn falls_back_to_vendored_when_system_bwrap_is_missing() { + fn ignores_system_bwrap_when_system_bwrap_is_missing() { assert_eq!( - preferred_bwrap_launcher_for_path(Path::new("/definitely/not/a/bwrap")), - BubblewrapLauncher::Vendored + system_bwrap_launcher_for_path(Path::new("/definitely/not/a/bwrap")), + None ); } + #[test] + fn finds_standalone_bundled_bwrap_next_to_exe_resources() { + let temp_dir = tempdir().expect("temp dir"); + let exe = temp_dir.path().join("codex"); + let expected_bwrap = temp_dir.path().join("codex-resources").join("bwrap"); + write_executable(&exe); + write_executable(&expected_bwrap); + + assert_eq!( + find_bundled_bwrap_for_exe(&exe), + Some(AbsolutePathBuf::from_absolute_path(&expected_bwrap).expect("absolute")) + ); + } + + #[test] + fn finds_npm_bundled_bwrap_next_to_target_vendor_dir() { + let temp_dir = tempdir().expect("temp dir"); + let target_dir = temp_dir.path().join("vendor/x86_64-unknown-linux-musl"); + let exe = target_dir.join("codex").join("codex"); + let expected_bwrap = target_dir.join("codex-resources").join("bwrap"); + write_executable(&exe); + write_executable(&expected_bwrap); + + assert_eq!( + find_bundled_bwrap_for_exe(&exe), + Some(AbsolutePathBuf::from_absolute_path(&expected_bwrap).expect("absolute")) + ); + } + + #[test] + fn finds_adjacent_dev_bwrap() { + let temp_dir = tempdir().expect("temp dir"); + let exe = temp_dir.path().join("codex"); + let expected_bwrap = temp_dir.path().join("bwrap"); + write_executable(&exe); + write_executable(&expected_bwrap); + + assert_eq!( + find_bundled_bwrap_for_exe(&exe), + Some(AbsolutePathBuf::from_absolute_path(&expected_bwrap).expect("absolute")) + ); + } + + #[test] + fn bundled_digest_verification_skips_missing_expected_digest() { + let file = NamedTempFile::new().expect("temp file"); + fs::write(file.path(), b"contents").expect("write file"); + + verify_bundled_bwrap_digest(file.as_file(), /*expected*/ None, file.path()) + .expect("missing digest should skip verification"); + } + + #[test] + fn bundled_digest_verification_accepts_matching_digest() { + let file = NamedTempFile::new().expect("temp file"); + fs::write(file.path(), b"contents").expect("write file"); + let expected: [u8; 32] = Sha256::digest(b"contents").into(); + + verify_bundled_bwrap_digest(file.as_file(), Some(expected), file.path()) + .expect("matching digest should verify"); + } + + #[test] + fn bundled_digest_verification_rejects_mismatched_digest() { + let file = NamedTempFile::new().expect("temp file"); + fs::write(file.path(), b"contents").expect("write file"); + + let err = verify_bundled_bwrap_digest(file.as_file(), Some([0xab; 32]), file.path()) + .expect_err("mismatched digest should fail"); + assert!(err.contains("bundled bubblewrap digest mismatch")); + } + + #[test] + fn parses_sha256_hex_digest() { + assert_eq!(parse_sha256_hex(&"ab".repeat(32)), Ok([0xab; 32])); + assert_eq!(parse_sha256_hex(&"00".repeat(32)), Ok(NULL_SHA256_DIGEST)); + assert!(parse_sha256_hex("ab").is_err()); + assert!(parse_sha256_hex(&format!("{}xx", "00".repeat(31))).is_err()); + } + #[test] fn preserved_files_are_made_inheritable_for_system_exec() { let file = NamedTempFile::new().expect("temp file"); @@ -244,6 +523,15 @@ mod tests { assert_eq!(fd_flags(file.as_file().as_raw_fd()) & libc::FD_CLOEXEC, 0); } + fn write_executable(path: &Path) { + if let Some(parent) = path.parent() { + fs::create_dir_all(parent).expect("create parent dir"); + } + fs::write(path, b"").expect("write executable"); + fs::set_permissions(path, fs::Permissions::from_mode(0o755)) + .expect("set executable permissions"); + } + fn set_cloexec(fd: libc::c_int) { let flags = fd_flags(fd); // SAFETY: `fd` is valid for the duration of the test. diff --git a/codex-rs/linux-sandbox/src/lib.rs b/codex-rs/linux-sandbox/src/lib.rs index 900287c99d..29ca45812a 100644 --- a/codex-rs/linux-sandbox/src/lib.rs +++ b/codex-rs/linux-sandbox/src/lib.rs @@ -13,8 +13,6 @@ mod launcher; mod linux_run_main; #[cfg(target_os = "linux")] mod proxy_routing; -#[cfg(target_os = "linux")] -mod vendored_bwrap; #[cfg(target_os = "linux")] pub fn run_main() -> ! { diff --git a/codex-rs/linux-sandbox/src/vendored_bwrap.rs b/codex-rs/linux-sandbox/src/vendored_bwrap.rs deleted file mode 100644 index a2da14db05..0000000000 --- a/codex-rs/linux-sandbox/src/vendored_bwrap.rs +++ /dev/null @@ -1,78 +0,0 @@ -//! Build-time bubblewrap entrypoint. -//! -//! On Linux targets, the build script compiles bubblewrap's C sources and -//! exposes a `bwrap_main` symbol that we can call via FFI. - -#[cfg(vendored_bwrap_available)] -mod imp { - use std::ffi::CString; - use std::fs::File; - use std::os::raw::c_char; - - unsafe extern "C" { - fn bwrap_main(argc: libc::c_int, argv: *const *const c_char) -> libc::c_int; - } - - fn argv_to_cstrings(argv: &[String]) -> Vec { - let mut cstrings: Vec = Vec::with_capacity(argv.len()); - for arg in argv { - match CString::new(arg.as_str()) { - Ok(value) => cstrings.push(value), - Err(err) => panic!("failed to convert argv to CString: {err}"), - } - } - cstrings - } - - /// Run the build-time bubblewrap `main` function and return its exit code. - /// - /// On success, bubblewrap will `execve` into the target program and this - /// function will never return. A return value therefore implies failure. - pub(crate) fn run_vendored_bwrap_main( - argv: &[String], - _preserved_files: &[File], - ) -> libc::c_int { - let cstrings = argv_to_cstrings(argv); - - let mut argv_ptrs: Vec<*const c_char> = cstrings.iter().map(|arg| arg.as_ptr()).collect(); - argv_ptrs.push(std::ptr::null()); - - // SAFETY: We provide a null-terminated argv vector whose pointers - // remain valid for the duration of the call. - unsafe { bwrap_main(cstrings.len() as libc::c_int, argv_ptrs.as_ptr()) } - } - - /// Execute the build-time bubblewrap `main` function with the given argv. - pub(crate) fn exec_vendored_bwrap(argv: Vec, preserved_files: Vec) -> ! { - let exit_code = run_vendored_bwrap_main(&argv, &preserved_files); - std::process::exit(exit_code); - } -} - -#[cfg(not(vendored_bwrap_available))] -mod imp { - use std::fs::File; - - /// Panics with a clear error when the build-time bwrap path is not enabled. - pub(crate) fn run_vendored_bwrap_main( - _argv: &[String], - _preserved_files: &[File], - ) -> libc::c_int { - panic!( - r#"build-time bubblewrap is not available in this build. -codex-linux-sandbox should always compile vendored bubblewrap on Linux targets. -Notes: -- ensure the target OS is Linux -- libcap headers must be available via pkg-config -- bubblewrap sources expected at codex-rs/vendor/bubblewrap (default)"# - ); - } - - /// Panics with a clear error when the build-time bwrap path is not enabled. - pub(crate) fn exec_vendored_bwrap(_argv: Vec, _preserved_files: Vec) -> ! { - let _ = run_vendored_bwrap_main(&[], &[]); - unreachable!("run_vendored_bwrap_main should always panic in this configuration") - } -} - -pub(crate) use imp::exec_vendored_bwrap; diff --git a/codex-rs/linux-sandbox/tests/suite/landlock.rs b/codex-rs/linux-sandbox/tests/suite/landlock.rs index 31021fb9ce..87e4ce68ae 100644 --- a/codex-rs/linux-sandbox/tests/suite/landlock.rs +++ b/codex-rs/linux-sandbox/tests/suite/landlock.rs @@ -40,7 +40,7 @@ const NETWORK_TIMEOUT_MS: u64 = 10_000; #[cfg(target_arch = "aarch64")] const NETWORK_TIMEOUT_MS: u64 = 10_000; -const BWRAP_UNAVAILABLE_ERR: &str = "build-time bubblewrap is not available in this build."; +const BWRAP_UNAVAILABLE_ERR: &str = "bubblewrap is unavailable: no system bwrap was found"; fn create_env_from_core_vars() -> HashMap { let policy = ShellEnvironmentPolicy::default(); diff --git a/codex-rs/linux-sandbox/tests/suite/managed_proxy.rs b/codex-rs/linux-sandbox/tests/suite/managed_proxy.rs index 932b7981d3..d1aa6856c4 100644 --- a/codex-rs/linux-sandbox/tests/suite/managed_proxy.rs +++ b/codex-rs/linux-sandbox/tests/suite/managed_proxy.rs @@ -15,7 +15,7 @@ use std::process::Stdio; use std::time::Duration; use tokio::process::Command; -const BWRAP_UNAVAILABLE_ERR: &str = "build-time bubblewrap is not available in this build."; +const BWRAP_UNAVAILABLE_ERR: &str = "bubblewrap is unavailable: no system bwrap was found"; const NETWORK_TIMEOUT_MS: u64 = 4_000; const MANAGED_PROXY_PERMISSION_ERR_SNIPPETS: &[&str] = &[ "loopback: Failed RTM_NEWADDR", @@ -82,7 +82,7 @@ fn is_managed_proxy_permission_error(stderr: &str) -> bool { async fn managed_proxy_skip_reason() -> Option { if should_skip_bwrap_tests().await { - return Some("vendored bwrap was not built in this environment".to_string()); + return Some("bubblewrap is unavailable in this environment".to_string()); } let mut env = create_env_from_core_vars(); diff --git a/codex-rs/sandboxing/src/bwrap.rs b/codex-rs/sandboxing/src/bwrap.rs index 840bb3e62d..e0eee177fb 100644 --- a/codex-rs/sandboxing/src/bwrap.rs +++ b/codex-rs/sandboxing/src/bwrap.rs @@ -18,7 +18,7 @@ const MISSING_BWRAP_WARNING: &str = concat!( "Install bubblewrap with your OS package manager. ", "See the sandbox prerequisites: ", "https://developers.openai.com/codex/concepts/sandboxing#prerequisites. ", - "Codex will use the vendored bubblewrap in the meantime.", + "Codex will use the bundled bubblewrap in the meantime.", ); const USER_NAMESPACE_WARNING: &str = "Codex's Linux sandbox uses bubblewrap and needs access to create user namespaces."; diff --git a/codex-rs/tui/src/resume_picker.rs b/codex-rs/tui/src/resume_picker.rs index e06ccfd118..06ad0a61a7 100644 --- a/codex-rs/tui/src/resume_picker.rs +++ b/codex-rs/tui/src/resume_picker.rs @@ -5753,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, @@ -5805,7 +5804,6 @@ session_picker_view = "dense" 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, @@ -5861,7 +5859,6 @@ session_picker_view = "dense" 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, From 3ef40b7282201214ffd338c17caf6ee2fd9e4a02 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Tue, 5 May 2026 15:44:05 -0700 Subject: [PATCH 6/7] release: publish standalone bwrap artifacts --- .github/workflows/rust-release.yml | 26 +++++++++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/.github/workflows/rust-release.yml b/.github/workflows/rust-release.yml index faab87f97b..71aab45394 100644 --- a/.github/workflows/rust-release.yml +++ b/.github/workflows/rust-release.yml @@ -255,6 +255,24 @@ jobs: with: target: ${{ matrix.target }} + - if: ${{ contains(matrix.target, 'linux') }} + name: Build bwrap and export digest + shell: bash + run: | + set -euo pipefail + target="${{ matrix.target }}" + cargo build --target "$target" --release --timings --bin bwrap + + bwrap_path="target/${target}/release/bwrap" + if [[ ! -f "$bwrap_path" ]]; then + echo "bwrap binary ${bwrap_path} not found" + exit 1 + fi + + digest="$(sha256sum "$bwrap_path" | awk '{print $1}')" + echo "CODEX_BWRAP_SHA256=${digest}" >> "$GITHUB_ENV" + echo "Built bwrap ${bwrap_path} with sha256:${digest}" + - name: Cargo build shell: bash run: | @@ -278,7 +296,7 @@ jobs: with: target: ${{ matrix.target }} artifacts-dir: ${{ github.workspace }}/codex-rs/target/${{ matrix.target }}/release - binaries: ${{ matrix.binaries }} + binaries: ${{ matrix.binaries }} bwrap - if: ${{ runner.os == 'macOS' }} name: MacOS code signing (binaries) @@ -361,6 +379,12 @@ jobs: fi done + if [[ "${{ matrix.target }}" == *linux* && "${{ matrix.bundle }}" == "primary" ]]; then + cp "target/${{ matrix.target }}/release/bwrap" "$dest/bwrap-${{ matrix.target }}" + cp "target/${{ matrix.target }}/release/bwrap.sigstore" \ + "$dest/bwrap-${{ matrix.target }}.sigstore" + fi + if [[ "${{ matrix.build_dmg }}" == "true" ]]; then cp target/${{ matrix.target }}/release/codex-${{ matrix.target }}.dmg "$dest/codex-${{ matrix.target }}.dmg" fi From 9630e3613a6fc3b7059855f3cb7ea738fc96f0bc Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Tue, 5 May 2026 15:44:05 -0700 Subject: [PATCH 7/7] npm: bundle standalone bwrap on Linux --- .github/workflows/ci.yml | 2 ++ codex-cli/scripts/build_npm_package.py | 20 ++++++++++++++++++-- codex-cli/scripts/install_native_deps.py | 9 ++++++++- scripts/install/install.sh | 10 +++++++++- scripts/stage_npm_packages.py | 19 +++++++++++++++++-- 5 files changed, 54 insertions(+), 6 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index bb5ed5b5bf..699ab070c8 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -52,10 +52,12 @@ jobs: CODEX_VERSION=0.125.0 WORKFLOW_URL="https://github.com/openai/codex/actions/runs/24901475298" OUTPUT_DIR="${RUNNER_TEMP}" + # This reused workflow predates the standalone bwrap artifact. python3 ./scripts/stage_npm_packages.py \ --release-version "$CODEX_VERSION" \ --workflow-url "$WORKFLOW_URL" \ --package codex \ + --allow-missing-native-component bwrap \ --output-dir "$OUTPUT_DIR" PACK_OUTPUT="${OUTPUT_DIR}/codex-npm-${CODEX_VERSION}.tgz" echo "pack_output=$PACK_OUTPUT" >> "$GITHUB_OUTPUT" diff --git a/codex-cli/scripts/build_npm_package.py b/codex-cli/scripts/build_npm_package.py index eda6c26152..c38fed68d9 100755 --- a/codex-cli/scripts/build_npm_package.py +++ b/codex-cli/scripts/build_npm_package.py @@ -69,8 +69,8 @@ PACKAGE_EXPANSIONS: dict[str, list[str]] = { PACKAGE_NATIVE_COMPONENTS: dict[str, list[str]] = { "codex": [], - "codex-linux-x64": ["codex", "rg"], - "codex-linux-arm64": ["codex", "rg"], + "codex-linux-x64": ["bwrap", "codex", "rg"], + "codex-linux-arm64": ["bwrap", "codex", "rg"], "codex-darwin-x64": ["codex", "rg"], "codex-darwin-arm64": ["codex", "rg"], "codex-win32-x64": ["codex", "rg", "codex-windows-sandbox-setup", "codex-command-runner"], @@ -87,6 +87,7 @@ PACKAGE_TARGET_FILTERS: dict[str, str] = { PACKAGE_CHOICES = tuple(PACKAGE_NATIVE_COMPONENTS) COMPONENT_DEST_DIR: dict[str, str] = { + "bwrap": "codex-resources", "codex": "codex", "codex-responses-api-proxy": "codex-responses-api-proxy", "codex-windows-sandbox-setup": "codex", @@ -137,6 +138,16 @@ def parse_args() -> argparse.Namespace: type=Path, help="Directory containing pre-installed native binaries to bundle (vendor root).", ) + parser.add_argument( + "--allow-missing-native-component", + dest="allow_missing_native_components", + action="append", + default=[], + help=( + "Native component that may be absent from --vendor-src. Intended for CI " + "compatibility with older artifact workflows; releases should not use this." + ), + ) return parser.parse_args() @@ -177,6 +188,7 @@ def main() -> int: staging_dir, native_components, target_filter={target_filter} if target_filter else None, + allow_missing_components=set(args.allow_missing_native_components), ) if release_version: @@ -365,12 +377,14 @@ def copy_native_binaries( staging_dir: Path, components: list[str], target_filter: set[str] | None = None, + allow_missing_components: set[str] | None = None, ) -> None: vendor_src = vendor_src.resolve() if not vendor_src.exists(): raise RuntimeError(f"Vendor source directory not found: {vendor_src}") components_set = {component for component in components if component in COMPONENT_DEST_DIR} + allow_missing_components = allow_missing_components or set() if not components_set: return @@ -399,6 +413,8 @@ def copy_native_binaries( src_component_dir = target_dir / dest_dir_name if not src_component_dir.exists(): + if component in allow_missing_components: + continue raise RuntimeError( f"Missing native component '{component}' in vendor source: {src_component_dir}" ) diff --git a/codex-cli/scripts/install_native_deps.py b/codex-cli/scripts/install_native_deps.py index 58fbd370fc..19d0af5bae 100755 --- a/codex-cli/scripts/install_native_deps.py +++ b/codex-cli/scripts/install_native_deps.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -"""Install Codex native binaries (Rust CLI plus ripgrep helpers).""" +"""Install Codex native binaries (Rust CLI, bwrap, and ripgrep helpers).""" import argparse from contextlib import contextmanager @@ -42,8 +42,15 @@ class BinaryComponent: WINDOWS_TARGETS = tuple(target for target in BINARY_TARGETS if "windows" in target) +LINUX_TARGETS = tuple(target for target in BINARY_TARGETS if "linux" in target) BINARY_COMPONENTS = { + "bwrap": BinaryComponent( + artifact_prefix="bwrap", + dest_dir="codex-resources", + binary_basename="bwrap", + targets=LINUX_TARGETS, + ), "codex": BinaryComponent( artifact_prefix="codex", dest_dir="codex", diff --git a/scripts/install/install.sh b/scripts/install/install.sh index 8c225e4d3b..2fc585d7e9 100755 --- a/scripts/install/install.sh +++ b/scripts/install/install.sh @@ -596,6 +596,10 @@ install_release() { cp "$vendor_root/path/rg" "$stage_release/codex-resources/rg" chmod 0755 "$stage_release/codex" chmod 0755 "$stage_release/codex-resources/rg" + if [ -f "$vendor_root/codex-resources/bwrap" ]; then + cp "$vendor_root/codex-resources/bwrap" "$stage_release/codex-resources/bwrap" + chmod 0755 "$stage_release/codex-resources/bwrap" + fi if [ -e "$release_dir" ] || [ -L "$release_dir" ]; then rm -rf "$release_dir" @@ -611,7 +615,11 @@ release_dir_is_complete() { [ -d "$release_dir" ] && [ -x "$release_dir/codex" ] && [ -x "$release_dir/codex-resources/rg" ] && - [ "$(basename "$release_dir")" = "$expected_version-$expected_target" ] + [ "$(basename "$release_dir")" = "$expected_version-$expected_target" ] && + case "$expected_target" in + *linux*) [ -x "$release_dir/codex-resources/bwrap" ] ;; + *) true ;; + esac } update_current_link() { diff --git a/scripts/stage_npm_packages.py b/scripts/stage_npm_packages.py index 5bbee755e9..2a0cd0028d 100755 --- a/scripts/stage_npm_packages.py +++ b/scripts/stage_npm_packages.py @@ -58,6 +58,16 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Retain temporary staging directories instead of deleting them.", ) + parser.add_argument( + "--allow-missing-native-component", + dest="allow_missing_native_components", + action="append", + default=[], + help=( + "Native component that may be absent from reused workflow artifacts. " + "Intended for CI compatibility only; release staging should not use this." + ), + ) return parser.parse_args() @@ -147,6 +157,8 @@ def main() -> int: packages = expand_packages(list(args.packages)) native_components = collect_native_components(packages) + allow_missing_native_components = set(args.allow_missing_native_components) + native_components_to_install = native_components - allow_missing_native_components vendor_temp_root: Path | None = None vendor_src: Path | None = None @@ -155,12 +167,12 @@ def main() -> int: final_messages = [] try: - if native_components: + if native_components_to_install: workflow_url, resolved_head_sha = resolve_workflow_url( args.release_version, args.workflow_url ) vendor_temp_root = Path(tempfile.mkdtemp(prefix="npm-native-", dir=runner_temp)) - install_native_components(workflow_url, native_components, vendor_temp_root) + install_native_components(workflow_url, native_components_to_install, vendor_temp_root) vendor_src = vendor_temp_root / "vendor" if resolved_head_sha: @@ -185,6 +197,9 @@ def main() -> int: if vendor_src is not None: cmd.extend(["--vendor-src", str(vendor_src)]) + for component in sorted(allow_missing_native_components): + cmd.extend(["--allow-missing-native-component", component]) + try: run_command(cmd) finally: