From ef500f4d2389afd09a2a7877570d375addc8bec3 Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Fri, 5 Jun 2026 06:23:58 +0000 Subject: [PATCH] Move code mode behind an IPC host Split the code-mode protocol and client from the V8-backed runtime so core and app-server no longer link codex-code-mode in production. ThreadManager now provisions durable code-mode sessions through a shared external host process, while tests can still inject the in-process provider. The IPC protocol uses a persistent stdin/stdout transport. Each frame is a 4-byte big-endian length followed by JSON, with a 16 MiB frame limit. Client requests carry u64 request IDs so create, execute, wait, terminate, and shutdown operations can be multiplexed over one process. Session IDs isolate durable stored values. Execute returns an ExecutionStarted response immediately and an asynchronous InitialResponse when the initial yield or completion is available. Nested tool calls and notifications travel from the host back to the client as delegate requests with their own IDs. Delegate responses, cancellation, and cell-closed lifecycle messages use the same framed channel. Wire operations encode errors as Result values. A dead connection fails pending operations, cancels outstanding delegates, and lets the provider spawn a new host for later sessions. Build codex-code-mode-host with V8 pointer-compression sandbox support and add it to canonical primary and app-server packages, legacy Linux and Windows bundles, signing verification, installers, Python runtime packages, and release CI for macOS, Linux, and Windows. The host is discovered next to the current executable, through CODEX_CODE_MODE_HOST_PATH, or on PATH. OS-level seccomp or seatbelt restrictions remain a follow-up to this cross-platform process split. Benchmarks were run from release builds on Linux x86_64 with the V8 sandbox profile and a text('ok') workload. Cold measurements used 30 samples, warm session provisioning used 200, and warm command execution used 500. Values are mean/p50/p95 in milliseconds: - session startup: in-process 0.002/0.002/0.005, IPC 2.623/2.599/2.894 - fresh-session command: in-process 1.831/1.758/1.915, IPC 7.428/7.252/8.306 - warm session provisioning: in-process 0.002/0.002/0.003, IPC 0.471/0.463/0.581 - warm command: in-process 1.759/1.757/1.940, IPC 2.005/2.001/2.166 The steady-state median command overhead is approximately 0.244 ms. The median fresh host plus first command cost is 7.252 ms. Validation: - 62/62 core code-mode integration tests passed against the external host - focused protocol, client, runtime, host, tools, and trace tests passed - Cargo and Bazel real-process host IPC tests passed - 11/11 package builder tests passed - Bazel lock verification, scoped Clippy fixes, and repository formatting passed --- .github/actions/setup-rusty-v8/action.yml | 21 +- .../scripts/build-codex-package-archive.sh | 17 + .github/workflows/rust-ci-full.yml | 1 + .github/workflows/rust-release-windows.yml | 14 +- .github/workflows/rust-release.yml | 52 ++- codex-rs/Cargo.lock | 42 +- codex-rs/Cargo.toml | 4 + codex-rs/code-mode-client/BUILD.bazel | 6 + codex-rs/code-mode-client/Cargo.toml | 23 + codex-rs/code-mode-client/src/connection.rs | 415 ++++++++++++++++++ codex-rs/code-mode-client/src/lib.rs | 231 ++++++++++ codex-rs/code-mode-client/src/tests.rs | 53 +++ codex-rs/code-mode-host/Cargo.toml | 11 + codex-rs/code-mode-host/src/main.rs | 333 +++++++++++++- codex-rs/code-mode-host/tests/host.rs | 123 ++++++ codex-rs/code-mode-protocol/BUILD.bazel | 6 + codex-rs/code-mode-protocol/Cargo.toml | 23 + .../src/description.rs | 2 +- codex-rs/code-mode-protocol/src/lib.rs | 46 ++ .../src/response.rs | 0 codex-rs/code-mode-protocol/src/runtime.rs | 89 ++++ codex-rs/code-mode-protocol/src/session.rs | 138 ++++++ .../code-mode-protocol/src/session_tests.rs | 19 + codex-rs/code-mode-protocol/src/wire.rs | 153 +++++++ codex-rs/code-mode-protocol/src/wire_tests.rs | 28 ++ codex-rs/code-mode/Cargo.toml | 1 + codex-rs/code-mode/src/lib.rs | 41 +- codex-rs/code-mode/src/runtime/callbacks.rs | 2 +- codex-rs/code-mode/src/runtime/mod.rs | 127 +----- codex-rs/code-mode/src/runtime/value.rs | 6 +- codex-rs/code-mode/src/service.rs | 138 +----- codex-rs/core/Cargo.toml | 3 +- codex-rs/core/src/codex_delegate.rs | 1 + codex-rs/core/src/session/mod.rs | 4 + codex-rs/core/src/session/session.rs | 6 +- codex-rs/core/src/session/tests.rs | 17 +- .../core/src/session/tests/guardian_tests.rs | 3 + codex-rs/core/src/state/service.rs | 2 + codex-rs/core/src/thread_manager.rs | 18 + codex-rs/core/src/tools/code_mode/delegate.rs | 10 +- .../src/tools/code_mode/execute_handler.rs | 11 +- .../core/src/tools/code_mode/execute_spec.rs | 16 +- codex-rs/core/src/tools/code_mode/mod.rs | 61 +-- .../src/tools/code_mode/response_adapter.rs | 25 +- .../core/src/tools/code_mode/wait_handler.rs | 21 +- .../core/src/tools/code_mode/wait_spec.rs | 12 +- codex-rs/core/src/tools/spec_plan.rs | 29 +- codex-rs/core/src/tools/spec_plan_tests.rs | 19 +- .../src/tools/tool_dispatch_trace_tests.rs | 60 +++ codex-rs/core/tests/common/Cargo.toml | 1 + codex-rs/core/tests/common/test_codex.rs | 7 + .../tests/suite/model_runtime_selectors.rs | 8 +- codex-rs/rollout-trace/Cargo.toml | 2 +- codex-rs/rollout-trace/src/code_cell.rs | 2 +- codex-rs/rollout-trace/src/tool_dispatch.rs | 6 +- codex-rs/tools/Cargo.toml | 2 +- codex-rs/tools/src/code_mode.rs | 19 +- codex-rs/tools/src/code_mode_tests.rs | 8 +- scripts/codex_package/README.md | 4 +- scripts/codex_package/cargo.py | 13 +- scripts/codex_package/cli.py | 14 + scripts/codex_package/layout.py | 8 + scripts/codex_package/targets.py | 1 + scripts/codex_package/test_cargo.py | 21 + scripts/codex_package/test_layout.py | 79 ++++ scripts/codex_package/test_v8.py | 57 +++ scripts/codex_package/v8.py | 7 +- scripts/install/install.ps1 | 1 + scripts/install/install.sh | 6 +- 69 files changed, 2304 insertions(+), 445 deletions(-) create mode 100644 codex-rs/code-mode-client/BUILD.bazel create mode 100644 codex-rs/code-mode-client/Cargo.toml create mode 100644 codex-rs/code-mode-client/src/connection.rs create mode 100644 codex-rs/code-mode-client/src/lib.rs create mode 100644 codex-rs/code-mode-client/src/tests.rs create mode 100644 codex-rs/code-mode-host/tests/host.rs create mode 100644 codex-rs/code-mode-protocol/BUILD.bazel create mode 100644 codex-rs/code-mode-protocol/Cargo.toml rename codex-rs/{code-mode => code-mode-protocol}/src/description.rs (99%) create mode 100644 codex-rs/code-mode-protocol/src/lib.rs rename codex-rs/{code-mode => code-mode-protocol}/src/response.rs (100%) create mode 100644 codex-rs/code-mode-protocol/src/runtime.rs create mode 100644 codex-rs/code-mode-protocol/src/session.rs create mode 100644 codex-rs/code-mode-protocol/src/session_tests.rs create mode 100644 codex-rs/code-mode-protocol/src/wire.rs create mode 100644 codex-rs/code-mode-protocol/src/wire_tests.rs create mode 100644 scripts/codex_package/test_layout.py create mode 100644 scripts/codex_package/test_v8.py diff --git a/.github/actions/setup-rusty-v8/action.yml b/.github/actions/setup-rusty-v8/action.yml index d9c4484657..bfe4a1444e 100644 --- a/.github/actions/setup-rusty-v8/action.yml +++ b/.github/actions/setup-rusty-v8/action.yml @@ -4,6 +4,10 @@ inputs: target: description: Rust target triple with Codex-built V8 release artifacts. required: true + sandbox: + description: Use the V8 pointer-compression sandbox artifact profile. + required: false + default: "false" runs: using: composite @@ -11,6 +15,7 @@ runs: - name: Configure rusty_v8 artifact overrides and verify checksums shell: bash env: + SANDBOX: ${{ inputs.sandbox }} TARGET: ${{ inputs.target }} run: | set -euo pipefail @@ -19,14 +24,18 @@ runs: release_tag="rusty-v8-v${version}" base_url="https://github.com/openai/codex/releases/download/${release_tag}" binding_dir="${RUNNER_TEMP}/rusty_v8" - archive_path="${binding_dir}/librusty_v8_release_${TARGET}.a.gz" - binding_path="${binding_dir}/src_binding_release_${TARGET}.rs" - checksums_path="${binding_dir}/rusty_v8_release_${TARGET}.sha256" + artifact_profile="release" + if [[ "${SANDBOX}" == "true" ]]; then + artifact_profile="ptrcomp_sandbox_release" + fi + archive_path="${binding_dir}/librusty_v8_${artifact_profile}_${TARGET}.a.gz" + binding_path="${binding_dir}/src_binding_${artifact_profile}_${TARGET}.rs" + checksums_path="${binding_dir}/rusty_v8_${artifact_profile}_${TARGET}.sha256" mkdir -p "${binding_dir}" - curl -fsSL "${base_url}/librusty_v8_release_${TARGET}.a.gz" -o "${archive_path}" - curl -fsSL "${base_url}/src_binding_release_${TARGET}.rs" -o "${binding_path}" - curl -fsSL "${base_url}/rusty_v8_release_${TARGET}.sha256" -o "${checksums_path}" + curl -fsSL "${base_url}/$(basename "${archive_path}")" -o "${archive_path}" + curl -fsSL "${base_url}/$(basename "${binding_path}")" -o "${binding_path}" + curl -fsSL "${base_url}/$(basename "${checksums_path}")" -o "${checksums_path}" if [[ "$(wc -l < "${checksums_path}")" -ne 2 ]]; then echo "Expected exactly two checksums for ${TARGET} in ${checksums_path}" >&2 diff --git a/.github/scripts/build-codex-package-archive.sh b/.github/scripts/build-codex-package-archive.sh index 80da4cf20c..5559562566 100644 --- a/.github/scripts/build-codex-package-archive.sh +++ b/.github/scripts/build-codex-package-archive.sh @@ -8,6 +8,7 @@ Usage: build-codex-package-archive.sh \ --bundle \ --entrypoint-dir \ --archive-dir \ + [--code-mode-host-bin ] \ [--bwrap-bin ] \ [--codex-command-runner-bin ] \ [--codex-windows-sandbox-setup-bin ] \ @@ -19,6 +20,7 @@ target="" bundle="" entrypoint_dir="" archive_dir="" +code_mode_host_bin="" target_suffixed_entrypoint="false" resource_args=() bwrap_bin_provided="false" @@ -43,6 +45,10 @@ while [[ $# -gt 0 ]]; do archive_dir="${2:?--archive-dir requires a value}" shift 2 ;; + --code-mode-host-bin) + code_mode_host_bin="${2:?--code-mode-host-bin requires a value}" + shift 2 + ;; --bwrap-bin) resource_args+=(--bwrap-bin "${2:?--bwrap-bin requires a value}") bwrap_bin_provided="true" @@ -110,8 +116,18 @@ case "$target" in esac entrypoint_name="$entrypoint" +code_mode_host_name="codex-code-mode-host" if [[ "$target_suffixed_entrypoint" == "true" ]]; then entrypoint_name="${entrypoint_name}-${target}" + code_mode_host_name="${code_mode_host_name}-${target}" +fi + +if [[ -z "$code_mode_host_bin" ]]; then + code_mode_host_bin="${entrypoint_dir%/}/${code_mode_host_name}${exe_suffix}" +fi +if [[ ! -f "$code_mode_host_bin" ]]; then + echo "Code-mode host binary ${code_mode_host_bin} not found" >&2 + exit 1 fi case "$target" in @@ -159,6 +175,7 @@ python_args=( --target "$target" --variant "$variant" --entrypoint-bin "${entrypoint_dir%/}/${entrypoint_name}${exe_suffix}" + --code-mode-host-bin "$code_mode_host_bin" --cargo-profile release --package-dir "$package_dir" --archive-output "$gzip_archive_path" diff --git a/.github/workflows/rust-ci-full.yml b/.github/workflows/rust-ci-full.yml index 7ad1d4b3ad..0b8b2ed031 100644 --- a/.github/workflows/rust-ci-full.yml +++ b/.github/workflows/rust-ci-full.yml @@ -382,6 +382,7 @@ jobs: name: Configure rusty_v8 artifact overrides and verify checksums uses: ./.github/actions/setup-rusty-v8 with: + sandbox: "true" target: ${{ matrix.target }} - name: Install cargo-chef diff --git a/.github/workflows/rust-release-windows.yml b/.github/workflows/rust-release-windows.yml index 49e627c1be..fd9549e006 100644 --- a/.github/workflows/rust-release-windows.yml +++ b/.github/workflows/rust-release-windows.yml @@ -41,14 +41,14 @@ jobs: - runner: windows-x64 target: x86_64-pc-windows-msvc bundle: helpers - binaries: "codex-windows-sandbox-setup codex-command-runner" + binaries: "codex-windows-sandbox-setup codex-command-runner codex-code-mode-host" runs_on: group: ${{ github.event.repository.name }}-runners labels: ${{ github.event.repository.name }}-windows-x64 - runner: windows-arm64 target: aarch64-pc-windows-msvc bundle: helpers - binaries: "codex-windows-sandbox-setup codex-command-runner" + binaries: "codex-windows-sandbox-setup codex-command-runner codex-code-mode-host" runs_on: group: ${{ github.event.repository.name }}-runners labels: ${{ github.event.repository.name }}-windows-arm64 @@ -158,7 +158,7 @@ jobs: run: working-directory: codex-rs env: - WINDOWS_BINARIES: "codex codex-responses-api-proxy codex-windows-sandbox-setup codex-command-runner codex-app-server" + WINDOWS_BINARIES: "codex codex-responses-api-proxy codex-windows-sandbox-setup codex-command-runner codex-code-mode-host codex-app-server" strategy: fail-fast: false @@ -340,16 +340,18 @@ jobs: bundle_dir="$(mktemp -d)" runner_src="$dest/codex-command-runner-${{ matrix.target }}.exe" setup_src="$dest/codex-windows-sandbox-setup-${{ matrix.target }}.exe" - if [[ -f "$runner_src" && -f "$setup_src" ]]; then + code_mode_host_src="$dest/codex-code-mode-host-${{ matrix.target }}.exe" + if [[ -f "$runner_src" && -f "$setup_src" && -f "$code_mode_host_src" ]]; then cp "$dest/$base" "$bundle_dir/$base" cp "$runner_src" "$bundle_dir/codex-command-runner.exe" cp "$setup_src" "$bundle_dir/codex-windows-sandbox-setup.exe" + cp "$code_mode_host_src" "$bundle_dir/codex-code-mode-host.exe" # Use an absolute path so bundle zips land in the real dist # dir even when 7z runs from a temp directory. (cd "$bundle_dir" && 7z a "$repo_root/$dest/${base}.zip" .) else - echo "warning: missing sandbox binaries; falling back to single-binary zip" - echo "warning: expected $runner_src and $setup_src" + echo "warning: missing bundled binaries; falling back to single-binary zip" + echo "warning: expected $runner_src, $setup_src, and $code_mode_host_src" (cd "$dest" && 7z a "${base}.zip" "$base") fi rm -rf "$bundle_dir" diff --git a/.github/workflows/rust-release.yml b/.github/workflows/rust-release.yml index 22b873f7d7..c5a4c834fa 100644 --- a/.github/workflows/rust-release.yml +++ b/.github/workflows/rust-release.yml @@ -80,50 +80,50 @@ jobs: target: aarch64-apple-darwin bundle: primary artifact_name: aarch64-apple-darwin - binaries: "codex codex-responses-api-proxy" + binaries: "codex codex-responses-api-proxy codex-code-mode-host" build_dmg: "true" - runner: macos-15-xlarge target: aarch64-apple-darwin bundle: app-server artifact_name: aarch64-apple-darwin-app-server - binaries: "codex-app-server" + binaries: "codex-app-server codex-code-mode-host" build_dmg: "false" - runner: macos-15-xlarge target: x86_64-apple-darwin bundle: primary artifact_name: x86_64-apple-darwin - binaries: "codex codex-responses-api-proxy" + binaries: "codex codex-responses-api-proxy codex-code-mode-host" build_dmg: "true" - runner: macos-15-xlarge target: x86_64-apple-darwin bundle: app-server artifact_name: x86_64-apple-darwin-app-server - binaries: "codex-app-server" + binaries: "codex-app-server codex-code-mode-host" build_dmg: "false" # Release artifacts intentionally ship MUSL-linked Linux binaries. - runner: ${{ github.event.repository.name }}-linux-x64-xl target: x86_64-unknown-linux-musl bundle: primary artifact_name: x86_64-unknown-linux-musl - binaries: "codex codex-responses-api-proxy bwrap" + binaries: "codex codex-responses-api-proxy codex-code-mode-host bwrap" build_dmg: "false" - runner: ${{ github.event.repository.name }}-linux-x64-xl target: x86_64-unknown-linux-musl bundle: app-server artifact_name: x86_64-unknown-linux-musl-app-server - binaries: "codex-app-server" + binaries: "codex-app-server codex-code-mode-host" build_dmg: "false" - runner: ${{ github.event.repository.name }}-linux-arm64 target: aarch64-unknown-linux-musl bundle: primary artifact_name: aarch64-unknown-linux-musl - binaries: "codex codex-responses-api-proxy bwrap" + binaries: "codex codex-responses-api-proxy codex-code-mode-host bwrap" build_dmg: "false" - runner: ${{ github.event.repository.name }}-linux-arm64 target: aarch64-unknown-linux-musl bundle: app-server artifact_name: aarch64-unknown-linux-musl-app-server - binaries: "codex-app-server" + binaries: "codex-app-server codex-code-mode-host" build_dmg: "false" steps: @@ -208,6 +208,7 @@ jobs: - name: Configure rusty_v8 artifact overrides and verify checksums uses: ./.github/actions/setup-rusty-v8 with: + sandbox: "true" target: ${{ matrix.target }} - if: ${{ contains(matrix.target, 'linux') }} @@ -344,9 +345,15 @@ jobs: rm -rf "$bundle_root" mkdir -p "$bundle_root/codex-resources" cp "$dest/codex-${{ matrix.target }}" "$bundle_root/codex" + cp "$dest/codex-code-mode-host-${{ matrix.target }}" \ + "$bundle_root/codex-code-mode-host" cp "$dest/bwrap-${{ matrix.target }}" "$bundle_root/codex-resources/bwrap" - chmod 0755 "$bundle_root/codex" "$bundle_root/codex-resources/bwrap" - tar -C "$bundle_root" -cf - codex codex-resources/bwrap | + chmod 0755 \ + "$bundle_root/codex" \ + "$bundle_root/codex-code-mode-host" \ + "$bundle_root/codex-resources/bwrap" + tar -C "$bundle_root" -cf - \ + codex codex-code-mode-host codex-resources/bwrap | zstd -T0 -19 -o "$dest/codex-${{ matrix.target }}-bundle.tar.zst" fi @@ -484,19 +491,19 @@ jobs: - target: aarch64-apple-darwin bundle: primary artifact_name: aarch64-apple-darwin - binaries: "codex codex-responses-api-proxy" + binaries: "codex codex-responses-api-proxy codex-code-mode-host" - target: aarch64-apple-darwin bundle: app-server artifact_name: aarch64-apple-darwin-app-server - binaries: "codex-app-server" + binaries: "codex-app-server codex-code-mode-host" - target: x86_64-apple-darwin bundle: primary artifact_name: x86_64-apple-darwin - binaries: "codex codex-responses-api-proxy" + binaries: "codex codex-responses-api-proxy codex-code-mode-host" - target: x86_64-apple-darwin bundle: app-server artifact_name: x86_64-apple-darwin-app-server - binaries: "codex-app-server" + binaries: "codex-app-server codex-code-mode-host" steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 @@ -602,22 +609,22 @@ jobs: - target: aarch64-apple-darwin bundle: primary artifact_name: aarch64-apple-darwin - binaries: "codex codex-responses-api-proxy" + binaries: "codex codex-responses-api-proxy codex-code-mode-host" build_dmg: "true" - target: aarch64-apple-darwin bundle: app-server artifact_name: aarch64-apple-darwin-app-server - binaries: "codex-app-server" + binaries: "codex-app-server codex-code-mode-host" build_dmg: "false" - target: x86_64-apple-darwin bundle: primary artifact_name: x86_64-apple-darwin - binaries: "codex codex-responses-api-proxy" + binaries: "codex codex-responses-api-proxy codex-code-mode-host" build_dmg: "true" - target: x86_64-apple-darwin bundle: app-server artifact_name: x86_64-apple-darwin-app-server - binaries: "codex-app-server" + binaries: "codex-app-server codex-code-mode-host" build_dmg: "false" steps: @@ -889,22 +896,22 @@ jobs: - target: aarch64-apple-darwin bundle: primary artifact_name: aarch64-apple-darwin - binaries: "codex codex-responses-api-proxy" + binaries: "codex codex-responses-api-proxy codex-code-mode-host" verify_dmg: "true" - target: aarch64-apple-darwin bundle: app-server artifact_name: aarch64-apple-darwin-app-server - binaries: "codex-app-server" + binaries: "codex-app-server codex-code-mode-host" verify_dmg: "false" - target: x86_64-apple-darwin bundle: primary artifact_name: x86_64-apple-darwin - binaries: "codex codex-responses-api-proxy" + binaries: "codex codex-responses-api-proxy codex-code-mode-host" verify_dmg: "true" - target: x86_64-apple-darwin bundle: app-server artifact_name: x86_64-apple-darwin-app-server - binaries: "codex-app-server" + binaries: "codex-app-server codex-code-mode-host" verify_dmg: "false" steps: @@ -992,6 +999,7 @@ jobs: mkdir -p "$package_dir" tar -xzf "${packaged_dir}/${package_stem}-${target}.tar.gz" -C "$package_dir" verify_signed_binary "${package_dir}/bin/${package_entrypoint}" + verify_signed_binary "${package_dir}/bin/codex-code-mode-host" if [[ "${{ matrix.verify_dmg }}" != "true" ]]; then exit 0 diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index d13238cb04..0950ca1235 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2483,6 +2483,7 @@ dependencies = [ name = "codex-code-mode" version = "0.0.0" dependencies = [ + "codex-code-mode-protocol", "codex-protocol", "deno_core_icudata", "pretty_assertions", @@ -2494,9 +2495,42 @@ dependencies = [ "v8", ] +[[package]] +name = "codex-code-mode-client" +version = "0.0.0" +dependencies = [ + "codex-code-mode-protocol", + "pretty_assertions", + "serde_json", + "tokio", + "tokio-util", + "tracing", +] + [[package]] name = "codex-code-mode-host" version = "0.0.0" +dependencies = [ + "codex-code-mode", + "codex-code-mode-protocol", + "codex-utils-cargo-bin", + "pretty_assertions", + "serde_json", + "tokio", + "tokio-util", +] + +[[package]] +name = "codex-code-mode-protocol" +version = "0.0.0" +dependencies = [ + "codex-protocol", + "pretty_assertions", + "serde", + "serde_json", + "tokio", + "tokio-util", +] [[package]] name = "codex-collaboration-mode-templates" @@ -2593,7 +2627,8 @@ dependencies = [ "codex-app-server-protocol", "codex-apply-patch", "codex-async-utils", - "codex-code-mode", + "codex-code-mode-client", + "codex-code-mode-protocol", "codex-config", "codex-connectors", "codex-context-fragments", @@ -3709,7 +3744,7 @@ name = "codex-rollout-trace" version = "0.0.0" dependencies = [ "anyhow", - "codex-code-mode", + "codex-code-mode-protocol", "codex-protocol", "http 1.4.0", "pretty_assertions", @@ -3918,7 +3953,7 @@ name = "codex-tools" version = "0.0.0" dependencies = [ "codex-app-server-protocol", - "codex-code-mode", + "codex-code-mode-protocol", "codex-features", "codex-protocol", "codex-utils-absolute-path", @@ -4577,6 +4612,7 @@ dependencies = [ "assert_cmd", "base64 0.22.1", "codex-arg0", + "codex-code-mode", "codex-config", "codex-core", "codex-exec-server", diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index 3b86ecf415..30e3c9ae48 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -21,7 +21,9 @@ members = [ "install-context", "codex-backend-openapi-models", "code-mode", + "code-mode-client", "code-mode-host", + "code-mode-protocol", "cloud-config", "cloud-tasks", "cloud-tasks-client", @@ -158,6 +160,8 @@ codex-cloud-config = { path = "cloud-config" } codex-cloud-tasks-client = { path = "cloud-tasks-client" } codex-cloud-tasks-mock-client = { path = "cloud-tasks-mock-client" } codex-code-mode = { path = "code-mode" } +codex-code-mode-client = { path = "code-mode-client" } +codex-code-mode-protocol = { path = "code-mode-protocol" } codex-config = { path = "config" } codex-connectors = { path = "connectors" } codex-context-fragments = { path = "context-fragments" } diff --git a/codex-rs/code-mode-client/BUILD.bazel b/codex-rs/code-mode-client/BUILD.bazel new file mode 100644 index 0000000000..071598cc2c --- /dev/null +++ b/codex-rs/code-mode-client/BUILD.bazel @@ -0,0 +1,6 @@ +load("//:defs.bzl", "codex_rust_crate") + +codex_rust_crate( + name = "code-mode-client", + crate_name = "codex_code_mode_client", +) diff --git a/codex-rs/code-mode-client/Cargo.toml b/codex-rs/code-mode-client/Cargo.toml new file mode 100644 index 0000000000..f659ab9a6b --- /dev/null +++ b/codex-rs/code-mode-client/Cargo.toml @@ -0,0 +1,23 @@ +[package] +edition.workspace = true +license.workspace = true +name = "codex-code-mode-client" +version.workspace = true + +[lib] +doctest = false +name = "codex_code_mode_client" +path = "src/lib.rs" + +[lints] +workspace = true + +[dependencies] +codex-code-mode-protocol = { workspace = true } +serde_json = { workspace = true } +tokio = { workspace = true, features = ["io-std", "io-util", "process", "rt", "sync"] } +tokio-util = { workspace = true, features = ["rt"] } +tracing = { workspace = true } + +[dev-dependencies] +pretty_assertions = { workspace = true } diff --git a/codex-rs/code-mode-client/src/connection.rs b/codex-rs/code-mode-client/src/connection.rs new file mode 100644 index 0000000000..85454a8c3d --- /dev/null +++ b/codex-rs/code-mode-client/src/connection.rs @@ -0,0 +1,415 @@ +use std::collections::HashMap; +use std::process::Stdio; +use std::sync::Arc; +use std::sync::Weak; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; + +use codex_code_mode_protocol::CodeModeSessionDelegate; +use codex_code_mode_protocol::ExecuteRequest; +use codex_code_mode_protocol::StartedCell; +use codex_code_mode_protocol::wire::ClientMessage; +use codex_code_mode_protocol::wire::DelegateRequest; +use codex_code_mode_protocol::wire::DelegateRequestId; +use codex_code_mode_protocol::wire::DelegateResponse; +use codex_code_mode_protocol::wire::HostMessage; +use codex_code_mode_protocol::wire::HostRequest; +use codex_code_mode_protocol::wire::HostResponse; +use codex_code_mode_protocol::wire::RequestId; +use codex_code_mode_protocol::wire::SessionId; +use codex_code_mode_protocol::wire::read_frame; +use codex_code_mode_protocol::wire::write_frame; +use tokio::io::AsyncBufReadExt; +use tokio::io::BufReader; +use tokio::process::Command; +use tokio::sync::Mutex; +use tokio::sync::mpsc; +use tokio::sync::oneshot; +use tokio_util::sync::CancellationToken; +use tracing::debug; +use tracing::warn; + +use crate::CodeModeHostCommand; + +const IPC_CHANNEL_CAPACITY: usize = 128; + +pub(super) struct Connection { + state: Arc, + cancellation: CancellationToken, +} + +impl Connection { + pub(super) async fn spawn(command: &CodeModeHostCommand) -> Result { + let mut child = host_process(command) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true) + .spawn() + .map_err(|err| { + format!( + "failed to spawn code-mode host {}: {err}", + command.program.display() + ) + })?; + let stdin = child + .stdin + .take() + .ok_or_else(|| "spawned code-mode host has no stdin".to_string())?; + let stdout = child + .stdout + .take() + .ok_or_else(|| "spawned code-mode host has no stdout".to_string())?; + let stderr = child.stderr.take(); + let (outgoing_tx, mut outgoing_rx) = mpsc::channel(IPC_CHANNEL_CAPACITY); + let cancellation = CancellationToken::new(); + let state = Arc::new(ConnectionState::new(outgoing_tx)); + + let writer_state = Arc::downgrade(&state); + let writer_cancellation = cancellation.clone(); + tokio::spawn(async move { + let mut stdin = stdin; + loop { + tokio::select! { + _ = writer_cancellation.cancelled() => break, + message = outgoing_rx.recv() => { + let Some(message) = message else { + break; + }; + if let Err(err) = write_frame(&mut stdin, &message).await { + warn!("failed to write code-mode host message: {err}"); + if let Some(state) = writer_state.upgrade() { + state + .fail(format!("failed to write code-mode host message: {err}")) + .await; + } + break; + } + } + } + } + }); + + let reader_state = Arc::downgrade(&state); + let reader_cancellation = cancellation.clone(); + tokio::spawn(async move { + drive_reader(stdout, reader_state, reader_cancellation).await; + }); + + if let Some(stderr) = stderr { + tokio::spawn(async move { + let mut lines = BufReader::new(stderr).lines(); + loop { + match lines.next_line().await { + Ok(Some(line)) => debug!("code-mode host stderr: {line}"), + Ok(None) => break, + Err(err) => { + warn!("failed to read code-mode host stderr: {err}"); + break; + } + } + } + }); + } + + let supervisor_state = Arc::downgrade(&state); + let supervisor_cancellation = cancellation.clone(); + tokio::spawn(async move { + tokio::select! { + result = child.wait() => { + let reason = match result { + Ok(status) => format!("code-mode host exited with status {status}"), + Err(err) => format!("failed waiting for code-mode host: {err}"), + }; + if let Some(state) = supervisor_state.upgrade() { + state.fail(reason).await; + } + } + _ = supervisor_cancellation.cancelled() => { + let _ = child.start_kill(); + let _ = child.wait().await; + } + } + }); + + Ok(Self { + state, + cancellation, + }) + } + + pub(super) fn is_alive(&self) -> bool { + self.state.alive.load(Ordering::Acquire) + } + + pub(super) async fn request(&self, request: HostRequest) -> Result { + let id = self.state.next_request_id.fetch_add(1, Ordering::Relaxed); + let (response_tx, response_rx) = oneshot::channel(); + self.state.pending.lock().await.insert(id, response_tx); + if let Err(err) = self + .state + .send(ClientMessage::Request { id, request }) + .await + { + self.state.pending.lock().await.remove(&id); + return Err(err); + } + response_rx + .await + .map_err(|_| self.state.failure_message())? + } + + pub(super) async fn execute( + &self, + session_id: SessionId, + request: ExecuteRequest, + ) -> Result { + let id = self.state.next_request_id.fetch_add(1, Ordering::Relaxed); + let (response_tx, response_rx) = oneshot::channel(); + let (initial_tx, initial_rx) = oneshot::channel(); + self.state.pending.lock().await.insert(id, response_tx); + self.state + .initial_responses + .lock() + .await + .insert(id, initial_tx); + if let Err(err) = self + .state + .send(ClientMessage::Request { + id, + request: HostRequest::Execute { + session_id, + request, + }, + }) + .await + { + self.state.pending.lock().await.remove(&id); + self.state.initial_responses.lock().await.remove(&id); + return Err(err); + } + let response = match response_rx.await { + Ok(Ok(response)) => response, + Ok(Err(err)) => { + self.state.initial_responses.lock().await.remove(&id); + return Err(err); + } + Err(_) => { + self.state.initial_responses.lock().await.remove(&id); + return Err(self.state.failure_message()); + } + }; + match response { + HostResponse::ExecutionStarted { cell_id } => { + Ok(StartedCell::from_result_receiver(cell_id, initial_rx)) + } + _ => { + self.state.initial_responses.lock().await.remove(&id); + Err("code-mode host returned an invalid execute response".to_string()) + } + } + } + + pub(super) async fn register_delegate( + &self, + session_id: SessionId, + delegate: Arc, + ) { + self.state + .delegates + .lock() + .await + .insert(session_id, delegate); + } + + pub(super) async fn remove_delegate(&self, session_id: SessionId) { + self.state.delegates.lock().await.remove(&session_id); + } +} + +impl Drop for Connection { + fn drop(&mut self) { + self.cancellation.cancel(); + } +} + +struct ConnectionState { + outgoing_tx: mpsc::Sender, + pending: Mutex>>>, + initial_responses: Mutex< + HashMap< + RequestId, + oneshot::Sender>, + >, + >, + delegates: Mutex>>, + delegate_cancellations: Mutex>, + next_request_id: AtomicU64, + alive: AtomicBool, + failure: std::sync::Mutex>, +} + +impl ConnectionState { + fn new(outgoing_tx: mpsc::Sender) -> Self { + Self { + outgoing_tx, + pending: Mutex::new(HashMap::new()), + initial_responses: Mutex::new(HashMap::new()), + delegates: Mutex::new(HashMap::new()), + delegate_cancellations: Mutex::new(HashMap::new()), + next_request_id: AtomicU64::new(1), + alive: AtomicBool::new(true), + failure: std::sync::Mutex::new(None), + } + } + + async fn send(&self, message: ClientMessage) -> Result<(), String> { + if !self.alive.load(Ordering::Acquire) { + return Err(self.failure_message()); + } + self.outgoing_tx + .send(message) + .await + .map_err(|_| self.failure_message()) + } + + fn failure_message(&self) -> String { + self.failure + .lock() + .ok() + .and_then(|failure| failure.clone()) + .unwrap_or_else(|| "code-mode host connection closed".to_string()) + } + + async fn fail(&self, reason: String) { + if !self.alive.swap(false, Ordering::AcqRel) { + return; + } + if let Ok(mut failure) = self.failure.lock() { + *failure = Some(reason.clone()); + } + for (_, sender) in self.pending.lock().await.drain() { + let _ = sender.send(Err(reason.clone())); + } + self.initial_responses.lock().await.clear(); + for (_, cancellation) in self.delegate_cancellations.lock().await.drain() { + cancellation.cancel(); + } + } +} + +async fn drive_reader( + mut stdout: tokio::process::ChildStdout, + state: Weak, + cancellation: CancellationToken, +) { + loop { + let message = tokio::select! { + _ = cancellation.cancelled() => return, + result = read_frame::<_, HostMessage>(&mut stdout) => result, + }; + match message { + Ok(Some(message)) => { + let Some(state) = state.upgrade() else { + return; + }; + handle_host_message(state, message).await; + } + Ok(None) => { + if let Some(state) = state.upgrade() { + state + .fail("code-mode host closed its stdout".to_string()) + .await; + } + return; + } + Err(err) => { + if let Some(state) = state.upgrade() { + state + .fail(format!("failed to read code-mode host message: {err}")) + .await; + } + return; + } + } + } +} + +async fn handle_host_message(state: Arc, message: HostMessage) { + match message { + HostMessage::Response { id, response } => { + if let Some(sender) = state.pending.lock().await.remove(&id) { + let _ = sender.send(response); + } + } + HostMessage::InitialResponse { id, response } => { + if let Some(sender) = state.initial_responses.lock().await.remove(&id) { + let _ = sender.send(response); + } + } + HostMessage::DelegateRequest { + id, + session_id, + request, + } => { + let delegate = state.delegates.lock().await.get(&session_id).cloned(); + let Some(delegate) = delegate else { + let _ = state + .send(ClientMessage::DelegateResponse { + id, + response: Err(format!("unknown code-mode session {session_id}")), + }) + .await; + return; + }; + let cancellation = CancellationToken::new(); + state + .delegate_cancellations + .lock() + .await + .insert(id, cancellation.clone()); + tokio::spawn(async move { + let response = match request { + DelegateRequest::InvokeTool(invocation) => delegate + .invoke_tool(invocation, cancellation) + .await + .map(DelegateResponse::ToolResult), + DelegateRequest::Notify { + call_id, + cell_id, + text, + } => delegate + .notify(call_id, cell_id, text, cancellation) + .await + .map(|()| DelegateResponse::NotificationDelivered), + }; + state.delegate_cancellations.lock().await.remove(&id); + let _ = state + .send(ClientMessage::DelegateResponse { id, response }) + .await; + }); + } + HostMessage::CancelDelegateRequest { id } => { + if let Some(cancellation) = state.delegate_cancellations.lock().await.remove(&id) { + cancellation.cancel(); + } + } + HostMessage::CellClosed { + session_id, + cell_id, + } => { + if let Some(delegate) = state.delegates.lock().await.get(&session_id).cloned() { + delegate.cell_closed(&cell_id); + } + } + } +} + +fn host_process(command: &CodeModeHostCommand) -> Command { + let mut process = Command::new(&command.program); + process.args(&command.args); + #[cfg(unix)] + process.process_group(0); + process +} diff --git a/codex-rs/code-mode-client/src/lib.rs b/codex-rs/code-mode-client/src/lib.rs new file mode 100644 index 0000000000..fe53aaed3a --- /dev/null +++ b/codex-rs/code-mode-client/src/lib.rs @@ -0,0 +1,231 @@ +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; + +use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::CodeModeSession; +use codex_code_mode_protocol::CodeModeSessionDelegate; +use codex_code_mode_protocol::CodeModeSessionProvider; +use codex_code_mode_protocol::CodeModeSessionProviderFuture; +use codex_code_mode_protocol::CodeModeSessionResultFuture; +use codex_code_mode_protocol::ExecuteRequest; +use codex_code_mode_protocol::StartedCell; +use codex_code_mode_protocol::WaitOutcome; +use codex_code_mode_protocol::WaitRequest; +use codex_code_mode_protocol::wire::HostRequest; +use codex_code_mode_protocol::wire::HostResponse; +use codex_code_mode_protocol::wire::SessionId; +use tokio::sync::Semaphore; + +const CODE_MODE_HOST_PATH_ENV: &str = "CODEX_CODE_MODE_HOST_PATH"; + +mod connection; + +use connection::Connection; + +#[derive(Clone, Debug)] +pub struct CodeModeHostCommand { + pub program: PathBuf, + pub args: Vec, +} + +impl Default for CodeModeHostCommand { + fn default() -> Self { + Self { + program: default_host_program(), + args: Vec::new(), + } + } +} + +pub struct IpcCodeModeSessionProvider { + command: CodeModeHostCommand, + connection: std::sync::Mutex>>, + spawn_permit: Semaphore, +} + +impl Default for IpcCodeModeSessionProvider { + fn default() -> Self { + Self::new(CodeModeHostCommand::default()) + } +} + +impl IpcCodeModeSessionProvider { + pub fn new(command: CodeModeHostCommand) -> Self { + Self { + command, + connection: std::sync::Mutex::new(None), + spawn_permit: Semaphore::new(/*permits*/ 1), + } + } + + async fn connection(&self) -> Result, String> { + if let Some(connection) = { + let current = self + .connection + .lock() + .map_err(|_| "code-mode host connection lock poisoned".to_string())?; + current + .as_ref() + .filter(|connection| connection.is_alive()) + .cloned() + } { + return Ok(connection); + } + + let _spawn_permit = self + .spawn_permit + .acquire() + .await + .map_err(|_| "code-mode host spawn coordinator closed".to_string())?; + if let Some(connection) = { + let current = self + .connection + .lock() + .map_err(|_| "code-mode host connection lock poisoned".to_string())?; + current + .as_ref() + .filter(|connection| connection.is_alive()) + .cloned() + } { + return Ok(connection); + } + let connection = Arc::new(Connection::spawn(&self.command).await?); + *self + .connection + .lock() + .map_err(|_| "code-mode host connection lock poisoned".to_string())? = + Some(Arc::clone(&connection)); + Ok(connection) + } +} + +impl CodeModeSessionProvider for IpcCodeModeSessionProvider { + fn create_session<'a>( + &'a self, + delegate: Arc, + ) -> CodeModeSessionProviderFuture<'a> { + Box::pin(async move { + let connection = self.connection().await?; + let response = connection.request(HostRequest::CreateSession).await?; + let HostResponse::SessionCreated { session_id } = response else { + return Err( + "code-mode host returned an invalid create-session response".to_string() + ); + }; + connection.register_delegate(session_id, delegate).await; + let session: Arc = Arc::new(IpcCodeModeSession { + connection, + session_id, + shutdown: AtomicBool::new(false), + }); + Ok(session) + }) + } +} + +struct IpcCodeModeSession { + connection: Arc, + session_id: SessionId, + shutdown: AtomicBool, +} + +impl CodeModeSession for IpcCodeModeSession { + fn execute<'a>( + &'a self, + request: ExecuteRequest, + ) -> CodeModeSessionResultFuture<'a, StartedCell> { + Box::pin(async move { + self.ensure_active()?; + self.connection.execute(self.session_id, request).await + }) + } + + fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome> { + Box::pin(async move { + self.ensure_active()?; + let response = self + .connection + .request(HostRequest::Wait { + session_id: self.session_id, + request, + }) + .await?; + match response { + HostResponse::WaitCompleted { outcome } => Ok(outcome), + _ => Err("code-mode host returned an invalid wait response".to_string()), + } + }) + } + + fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome> { + Box::pin(async move { + self.ensure_active()?; + let response = self + .connection + .request(HostRequest::Terminate { + session_id: self.session_id, + cell_id, + }) + .await?; + match response { + HostResponse::WaitCompleted { outcome } => Ok(outcome), + _ => Err("code-mode host returned an invalid terminate response".to_string()), + } + }) + } + + fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()> { + Box::pin(async move { + if self.shutdown.swap(true, Ordering::AcqRel) { + return Ok(()); + } + let result = self + .connection + .request(HostRequest::ShutdownSession { + session_id: self.session_id, + }) + .await; + self.connection.remove_delegate(self.session_id).await; + match result? { + HostResponse::SessionShutdown => Ok(()), + _ => Err("code-mode host returned an invalid shutdown response".to_string()), + } + }) + } +} + +impl IpcCodeModeSession { + fn ensure_active(&self) -> Result<(), String> { + if self.shutdown.load(Ordering::Acquire) { + Err("code mode session is shutting down".to_string()) + } else { + Ok(()) + } + } +} + +fn default_host_program() -> PathBuf { + if let Some(path) = std::env::var_os(CODE_MODE_HOST_PATH_ENV) { + return PathBuf::from(path); + } + let executable_name = if cfg!(windows) { + "codex-code-mode-host.exe" + } else { + "codex-code-mode-host" + }; + if let Ok(current_exe) = std::env::current_exe() + && let Some(parent) = current_exe.parent() + { + let sibling = parent.join(executable_name); + if sibling.is_file() { + return sibling; + } + } + PathBuf::from(executable_name) +} + +#[cfg(test)] +#[path = "tests.rs"] +mod tests; diff --git a/codex-rs/code-mode-client/src/tests.rs b/codex-rs/code-mode-client/src/tests.rs new file mode 100644 index 0000000000..3b38f03fc3 --- /dev/null +++ b/codex-rs/code-mode-client/src/tests.rs @@ -0,0 +1,53 @@ +use std::path::PathBuf; +use std::sync::Arc; + +use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::CodeModeNestedToolCall; +use codex_code_mode_protocol::CodeModeSessionDelegate; +use codex_code_mode_protocol::CodeModeSessionProvider; +use codex_code_mode_protocol::NotificationFuture; +use codex_code_mode_protocol::ToolInvocationFuture; +use tokio_util::sync::CancellationToken; + +use super::CodeModeHostCommand; +use super::IpcCodeModeSessionProvider; + +struct TestDelegate; + +impl CodeModeSessionDelegate for TestDelegate { + fn invoke_tool<'a>( + &'a self, + _invocation: CodeModeNestedToolCall, + _cancellation_token: CancellationToken, + ) -> ToolInvocationFuture<'a> { + Box::pin(async { Err("unexpected tool invocation".to_string()) }) + } + + fn notify<'a>( + &'a self, + _call_id: String, + _cell_id: CellId, + _text: String, + _cancellation_token: CancellationToken, + ) -> NotificationFuture<'a> { + Box::pin(async { Ok(()) }) + } + + fn cell_closed(&self, _cell_id: &CellId) {} +} + +#[tokio::test] +async fn create_session_reports_host_spawn_errors() { + let provider = IpcCodeModeSessionProvider::new(CodeModeHostCommand { + program: PathBuf::from("codex-code-mode-host-does-not-exist"), + args: Vec::new(), + }); + + let error = provider + .create_session(Arc::new(TestDelegate)) + .await + .err() + .expect("session creation should fail"); + + assert!(error.contains("failed to spawn code-mode host")); +} diff --git a/codex-rs/code-mode-host/Cargo.toml b/codex-rs/code-mode-host/Cargo.toml index a2c384d010..87e2eba092 100644 --- a/codex-rs/code-mode-host/Cargo.toml +++ b/codex-rs/code-mode-host/Cargo.toml @@ -10,3 +10,14 @@ path = "src/main.rs" [lints] workspace = true + +[dependencies] +codex-code-mode = { workspace = true, features = ["sandbox"] } +codex-code-mode-protocol = { workspace = true } +serde_json = { workspace = true } +tokio = { workspace = true, features = ["io-std", "io-util", "macros", "rt-multi-thread", "sync"] } +tokio-util = { workspace = true, features = ["rt"] } + +[dev-dependencies] +codex-utils-cargo-bin = { workspace = true } +pretty_assertions = { workspace = true } diff --git a/codex-rs/code-mode-host/src/main.rs b/codex-rs/code-mode-host/src/main.rs index f328e4d9d0..6b35d92021 100644 --- a/codex-rs/code-mode-host/src/main.rs +++ b/codex-rs/code-mode-host/src/main.rs @@ -1 +1,332 @@ -fn main() {} +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; + +use codex_code_mode::CodeModeService; +use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::CodeModeNestedToolCall; +use codex_code_mode_protocol::CodeModeSessionDelegate; +use codex_code_mode_protocol::NotificationFuture; +use codex_code_mode_protocol::ToolInvocationFuture; +use codex_code_mode_protocol::wire::ClientMessage; +use codex_code_mode_protocol::wire::DelegateRequest; +use codex_code_mode_protocol::wire::DelegateRequestId; +use codex_code_mode_protocol::wire::DelegateResponse; +use codex_code_mode_protocol::wire::HostMessage; +use codex_code_mode_protocol::wire::HostRequest; +use codex_code_mode_protocol::wire::HostResponse; +use codex_code_mode_protocol::wire::SessionId; +use codex_code_mode_protocol::wire::read_frame; +use codex_code_mode_protocol::wire::write_frame; +use tokio::sync::Mutex; +use tokio::sync::mpsc; +use tokio::sync::oneshot; +use tokio_util::sync::CancellationToken; + +const IPC_CHANNEL_CAPACITY: usize = 128; + +fn main() { + let runtime = build_runtime() + .unwrap_or_else(|err| panic!("failed to build code-mode host runtime: {err}")); + if let Err(err) = runtime.block_on(run()) { + eprintln!("codex-code-mode-host failed: {err}"); + std::process::exit(1); + } +} + +fn build_runtime() -> std::io::Result { + tokio::runtime::Builder::new_current_thread() + .enable_io() + .enable_time() + .build() +} + +async fn run() -> Result<(), String> { + let (outgoing_tx, mut outgoing_rx) = mpsc::channel(IPC_CHANNEL_CAPACITY); + let peer = Arc::new(HostPeer::new(outgoing_tx)); + let state = Arc::new(HostState { + sessions: Mutex::new(HashMap::new()), + next_session_id: AtomicU64::new(1), + peer, + }); + + let writer = tokio::spawn(async move { + let mut stdout = tokio::io::stdout(); + while let Some(message) = outgoing_rx.recv().await { + write_frame(&mut stdout, &message) + .await + .map_err(|err| err.to_string())?; + } + Ok::<(), String>(()) + }); + + let mut stdin = tokio::io::stdin(); + while let Some(message) = read_frame::<_, ClientMessage>(&mut stdin) + .await + .map_err(|err| err.to_string())? + { + match message { + ClientMessage::Request { id, request } => { + let state = Arc::clone(&state); + tokio::spawn(async move { + state.handle_request(id, request).await; + }); + } + ClientMessage::DelegateResponse { id, response } => { + state.peer.complete(id, response).await; + } + } + } + + let sessions = state + .sessions + .lock() + .await + .drain() + .map(|(_, session)| session) + .collect::>(); + for session in sessions { + let _ = session.shutdown().await; + } + drop(state); + writer.await.map_err(|err| err.to_string())? +} + +struct HostState { + sessions: Mutex>>, + next_session_id: AtomicU64, + peer: Arc, +} + +impl HostState { + async fn handle_request(self: Arc, request_id: u64, request: HostRequest) { + match request { + HostRequest::CreateSession => { + let session_id = self.next_session_id.fetch_add(1, Ordering::Relaxed); + let delegate = Arc::new(RemoteDelegate { + session_id, + peer: Arc::clone(&self.peer), + }); + self.sessions.lock().await.insert( + session_id, + Arc::new(CodeModeService::with_delegate(delegate)), + ); + self.respond(request_id, Ok(HostResponse::SessionCreated { session_id })) + .await; + } + HostRequest::Execute { + session_id, + request, + } => { + let error = match self.session(session_id).await { + Ok(session) => match session.execute(request).await { + Ok(started) => { + let cell_id = started.cell_id.clone(); + self.respond( + request_id, + Ok(HostResponse::ExecutionStarted { cell_id }), + ) + .await; + let peer = Arc::clone(&self.peer); + tokio::spawn(async move { + let response = started.initial_response().await; + peer.send(HostMessage::InitialResponse { + id: request_id, + response, + }) + .await; + }); + return; + } + Err(err) => err, + }, + Err(err) => err, + }; + self.respond(request_id, Err(error)).await; + } + HostRequest::Wait { + session_id, + request, + } => { + let result = match self.session(session_id).await { + Ok(session) => session + .wait(request) + .await + .map(|outcome| HostResponse::WaitCompleted { outcome }), + Err(err) => Err(err), + }; + self.respond(request_id, result).await; + } + HostRequest::Terminate { + session_id, + cell_id, + } => { + let result = match self.session(session_id).await { + Ok(session) => session + .terminate(cell_id) + .await + .map(|outcome| HostResponse::WaitCompleted { outcome }), + Err(err) => Err(err), + }; + self.respond(request_id, result).await; + } + HostRequest::ShutdownSession { session_id } => { + let session = self.sessions.lock().await.remove(&session_id); + let result = match session { + Some(session) => session + .shutdown() + .await + .map(|()| HostResponse::SessionShutdown), + None => Ok(HostResponse::SessionShutdown), + }; + self.respond(request_id, result).await; + } + } + } + + async fn session(&self, session_id: SessionId) -> Result, String> { + self.sessions + .lock() + .await + .get(&session_id) + .cloned() + .ok_or_else(|| format!("unknown code-mode session {session_id}")) + } + + async fn respond(&self, id: u64, response: Result) { + self.peer.send(HostMessage::Response { id, response }).await; + } +} + +struct HostPeer { + outgoing_tx: mpsc::Sender, + pending: Mutex>>>, + next_request_id: AtomicU64, +} + +impl HostPeer { + fn new(outgoing_tx: mpsc::Sender) -> Self { + Self { + outgoing_tx, + pending: Mutex::new(HashMap::new()), + next_request_id: AtomicU64::new(1), + } + } + + async fn send(&self, message: HostMessage) { + let _ = self.outgoing_tx.send(message).await; + } + + async fn call( + &self, + session_id: SessionId, + request: DelegateRequest, + cancellation_token: CancellationToken, + ) -> Result { + let id = self.next_request_id.fetch_add(1, Ordering::Relaxed); + let (response_tx, response_rx) = oneshot::channel(); + self.pending.lock().await.insert(id, response_tx); + if self + .outgoing_tx + .send(HostMessage::DelegateRequest { + id, + session_id, + request, + }) + .await + .is_err() + { + self.pending.lock().await.remove(&id); + return Err("code-mode client connection closed".to_string()); + } + tokio::select! { + response = response_rx => { + response.map_err(|_| "code-mode client closed before returning delegate output".to_string())? + } + _ = cancellation_token.cancelled() => { + self.pending.lock().await.remove(&id); + self.send(HostMessage::CancelDelegateRequest { id }).await; + Err("code mode delegate request cancelled".to_string()) + } + } + } + + async fn complete(&self, id: DelegateRequestId, response: Result) { + if let Some(sender) = self.pending.lock().await.remove(&id) { + let _ = sender.send(response); + } + } +} + +struct RemoteDelegate { + session_id: SessionId, + peer: Arc, +} + +impl CodeModeSessionDelegate for RemoteDelegate { + fn invoke_tool<'a>( + &'a self, + invocation: CodeModeNestedToolCall, + cancellation_token: CancellationToken, + ) -> ToolInvocationFuture<'a> { + Box::pin(async move { + match self + .peer + .call( + self.session_id, + DelegateRequest::InvokeTool(invocation), + cancellation_token, + ) + .await? + { + DelegateResponse::ToolResult(result) => Ok(result), + DelegateResponse::NotificationDelivered => { + Err("code-mode client returned an invalid tool result".to_string()) + } + } + }) + } + + fn notify<'a>( + &'a self, + call_id: String, + cell_id: CellId, + text: String, + cancellation_token: CancellationToken, + ) -> NotificationFuture<'a> { + Box::pin(async move { + match self + .peer + .call( + self.session_id, + DelegateRequest::Notify { + call_id, + cell_id, + text, + }, + cancellation_token, + ) + .await? + { + DelegateResponse::NotificationDelivered => Ok(()), + DelegateResponse::ToolResult(_) => { + Err("code-mode client returned an invalid notification result".to_string()) + } + } + }) + } + + fn cell_closed(&self, cell_id: &CellId) { + let peer = Arc::clone(&self.peer); + let cell_id = cell_id.clone(); + let session_id = self.session_id; + tokio::spawn(async move { + peer.send(HostMessage::CellClosed { + session_id, + cell_id, + }) + .await; + }); + } +} diff --git a/codex-rs/code-mode-host/tests/host.rs b/codex-rs/code-mode-host/tests/host.rs new file mode 100644 index 0000000000..aa32957cd0 --- /dev/null +++ b/codex-rs/code-mode-host/tests/host.rs @@ -0,0 +1,123 @@ +use std::process::Stdio; +use std::time::Duration; + +use codex_code_mode_protocol::ExecuteRequest; +use codex_code_mode_protocol::FunctionCallOutputContentItem; +use codex_code_mode_protocol::RuntimeResponse; +use codex_code_mode_protocol::wire::ClientMessage; +use codex_code_mode_protocol::wire::HostMessage; +use codex_code_mode_protocol::wire::HostRequest; +use codex_code_mode_protocol::wire::HostResponse; +use codex_code_mode_protocol::wire::read_frame; +use codex_code_mode_protocol::wire::write_frame; +use pretty_assertions::assert_eq; +use tokio::process::Command; + +#[tokio::test] +async fn serves_code_mode_sessions_over_stdio() { + let mut child = Command::new( + codex_utils_cargo_bin::cargo_bin("codex-code-mode-host") + .expect("resolve codex-code-mode-host binary"), + ) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::inherit()) + .kill_on_drop(true) + .spawn() + .expect("spawn codex-code-mode-host"); + let mut stdin = child.stdin.take().expect("host stdin"); + let mut stdout = child.stdout.take().expect("host stdout"); + + write_frame( + &mut stdin, + &ClientMessage::Request { + id: 1, + request: HostRequest::CreateSession, + }, + ) + .await + .expect("create session request"); + let session_id = match read_frame(&mut stdout).await.expect("create response") { + Some(HostMessage::Response { + id: 1, + response: Ok(HostResponse::SessionCreated { session_id }), + }) => session_id, + message => panic!("unexpected create-session response: {message:?}"), + }; + + write_frame( + &mut stdin, + &ClientMessage::Request { + id: 2, + request: HostRequest::Execute { + session_id, + request: ExecuteRequest { + tool_call_id: "call-1".to_string(), + enabled_tools: Vec::new(), + source: "text('hello')".to_string(), + yield_time_ms: None, + max_output_tokens: None, + }, + }, + }, + ) + .await + .expect("execute request"); + let cell_id = match read_frame(&mut stdout).await.expect("execute response") { + Some(HostMessage::Response { + id: 2, + response: Ok(HostResponse::ExecutionStarted { cell_id }), + }) => cell_id, + message => panic!("unexpected execute response: {message:?}"), + }; + let response = match read_frame(&mut stdout) + .await + .expect("initial response frame") + { + Some(HostMessage::InitialResponse { + id: 2, + response: Ok(response), + }) => response, + message => panic!("unexpected initial response: {message:?}"), + }; + assert_eq!( + response, + RuntimeResponse::Result { + cell_id, + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "hello".to_string(), + }], + error_text: None, + } + ); + + write_frame( + &mut stdin, + &ClientMessage::Request { + id: 3, + request: HostRequest::ShutdownSession { session_id }, + }, + ) + .await + .expect("shutdown request"); + loop { + match read_frame(&mut stdout).await.expect("shutdown response") { + Some(HostMessage::CellClosed { + session_id: closed_session_id, + .. + }) if closed_session_id == session_id => {} + Some(HostMessage::Response { + id: 3, + response: Ok(HostResponse::SessionShutdown), + }) => break, + message => panic!("unexpected shutdown response: {message:?}"), + } + } + + drop(stdin); + let status = tokio::time::timeout(Duration::from_secs(5), child.wait()) + .await + .expect("host exit timeout") + .expect("wait for host"); + assert!(status.success(), "host exited with {status}"); +} diff --git a/codex-rs/code-mode-protocol/BUILD.bazel b/codex-rs/code-mode-protocol/BUILD.bazel new file mode 100644 index 0000000000..ebf43b9bba --- /dev/null +++ b/codex-rs/code-mode-protocol/BUILD.bazel @@ -0,0 +1,6 @@ +load("//:defs.bzl", "codex_rust_crate") + +codex_rust_crate( + name = "code-mode-protocol", + crate_name = "codex_code_mode_protocol", +) diff --git a/codex-rs/code-mode-protocol/Cargo.toml b/codex-rs/code-mode-protocol/Cargo.toml new file mode 100644 index 0000000000..808d4c3775 --- /dev/null +++ b/codex-rs/code-mode-protocol/Cargo.toml @@ -0,0 +1,23 @@ +[package] +edition.workspace = true +license.workspace = true +name = "codex-code-mode-protocol" +version.workspace = true + +[lib] +doctest = false +name = "codex_code_mode_protocol" +path = "src/lib.rs" + +[lints] +workspace = true + +[dependencies] +codex-protocol = { workspace = true } +serde = { workspace = true, features = ["derive"] } +serde_json = { workspace = true } +tokio = { workspace = true, features = ["io-util", "sync"] } +tokio-util = { workspace = true, features = ["rt"] } + +[dev-dependencies] +pretty_assertions = { workspace = true } diff --git a/codex-rs/code-mode/src/description.rs b/codex-rs/code-mode-protocol/src/description.rs similarity index 99% rename from codex-rs/code-mode/src/description.rs rename to codex-rs/code-mode-protocol/src/description.rs index 7e0801ed25..b4a54b8acc 100644 --- a/codex-rs/code-mode/src/description.rs +++ b/codex-rs/code-mode-protocol/src/description.rs @@ -128,7 +128,7 @@ pub enum CodeModeToolKind { Freeform, } -#[derive(Clone, Debug, PartialEq)] +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] pub struct ToolDefinition { pub name: String, pub tool_name: ToolName, diff --git a/codex-rs/code-mode-protocol/src/lib.rs b/codex-rs/code-mode-protocol/src/lib.rs new file mode 100644 index 0000000000..3fdb7a0373 --- /dev/null +++ b/codex-rs/code-mode-protocol/src/lib.rs @@ -0,0 +1,46 @@ +mod description; +mod response; +mod runtime; +mod session; +pub mod wire; + +pub use description::CODE_MODE_PRAGMA_PREFIX; +pub use description::CodeModeToolKind; +pub use description::EnabledToolMetadata; +pub use description::ToolDefinition; +pub use description::ToolNamespaceDescription; +pub use description::augment_tool_definition; +pub use description::build_exec_tool_description; +pub use description::build_wait_tool_description; +pub use description::enabled_tool_metadata; +pub use description::is_code_mode_nested_tool; +pub use description::normalize_code_mode_identifier; +pub use description::parse_exec_source; +pub use description::render_code_mode_sample; +pub use description::render_json_schema_to_typescript; +pub use response::DEFAULT_IMAGE_DETAIL; +pub use response::FunctionCallOutputContentItem; +pub use response::ImageDetail; +pub use runtime::CodeModeNestedToolCall; +pub use runtime::DEFAULT_EXEC_YIELD_TIME_MS; +pub use runtime::DEFAULT_MAX_OUTPUT_TOKENS_PER_EXEC_CALL; +pub use runtime::DEFAULT_WAIT_YIELD_TIME_MS; +pub use runtime::ExecuteRequest; +pub use runtime::ExecuteToPendingOutcome; +pub use runtime::RuntimeResponse; +pub use runtime::WaitOutcome; +pub use runtime::WaitRequest; +pub use runtime::WaitToPendingOutcome; +pub use runtime::WaitToPendingRequest; +pub use session::CellId; +pub use session::CodeModeSession; +pub use session::CodeModeSessionDelegate; +pub use session::CodeModeSessionProvider; +pub use session::CodeModeSessionProviderFuture; +pub use session::CodeModeSessionResultFuture; +pub use session::NotificationFuture; +pub use session::StartedCell; +pub use session::ToolInvocationFuture; + +pub const PUBLIC_TOOL_NAME: &str = "exec"; +pub const WAIT_TOOL_NAME: &str = "wait"; diff --git a/codex-rs/code-mode/src/response.rs b/codex-rs/code-mode-protocol/src/response.rs similarity index 100% rename from codex-rs/code-mode/src/response.rs rename to codex-rs/code-mode-protocol/src/response.rs diff --git a/codex-rs/code-mode-protocol/src/runtime.rs b/codex-rs/code-mode-protocol/src/runtime.rs new file mode 100644 index 0000000000..147822063f --- /dev/null +++ b/codex-rs/code-mode-protocol/src/runtime.rs @@ -0,0 +1,89 @@ +use codex_protocol::ToolName; +use serde::Deserialize; +use serde::Serialize; +use serde_json::Value as JsonValue; + +use crate::CellId; +use crate::CodeModeToolKind; +use crate::FunctionCallOutputContentItem; +use crate::ToolDefinition; + +pub const DEFAULT_EXEC_YIELD_TIME_MS: u64 = 10_000; +pub const DEFAULT_WAIT_YIELD_TIME_MS: u64 = 10_000; +pub const DEFAULT_MAX_OUTPUT_TOKENS_PER_EXEC_CALL: usize = 10_000; + +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +pub struct ExecuteRequest { + pub tool_call_id: String, + pub enabled_tools: Vec, + pub source: String, + pub yield_time_ms: Option, + pub max_output_tokens: Option, +} + +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +pub struct WaitRequest { + pub cell_id: CellId, + pub yield_time_ms: u64, +} + +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +pub struct WaitToPendingRequest { + pub cell_id: CellId, +} + +#[derive(Debug, Deserialize, PartialEq, Serialize)] +pub enum WaitOutcome { + LiveCell(RuntimeResponse), + MissingCell(RuntimeResponse), +} + +#[derive(Debug, Deserialize, PartialEq, Serialize)] +pub enum ExecuteToPendingOutcome { + Pending { + cell_id: CellId, + content_items: Vec, + pending_tool_call_ids: Vec, + }, + Completed(RuntimeResponse), +} + +#[derive(Debug, Deserialize, PartialEq, Serialize)] +pub enum WaitToPendingOutcome { + LiveCell(ExecuteToPendingOutcome), + MissingCell(RuntimeResponse), +} + +impl From for RuntimeResponse { + fn from(outcome: WaitOutcome) -> Self { + match outcome { + WaitOutcome::LiveCell(response) | WaitOutcome::MissingCell(response) => response, + } + } +} + +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +pub enum RuntimeResponse { + Yielded { + cell_id: CellId, + content_items: Vec, + }, + Terminated { + cell_id: CellId, + content_items: Vec, + }, + Result { + cell_id: CellId, + content_items: Vec, + error_text: Option, + }, +} + +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +pub struct CodeModeNestedToolCall { + pub cell_id: CellId, + pub runtime_tool_call_id: String, + pub tool_name: ToolName, + pub tool_kind: CodeModeToolKind, + pub input: Option, +} diff --git a/codex-rs/code-mode-protocol/src/session.rs b/codex-rs/code-mode-protocol/src/session.rs new file mode 100644 index 0000000000..57669c4314 --- /dev/null +++ b/codex-rs/code-mode-protocol/src/session.rs @@ -0,0 +1,138 @@ +use std::fmt; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; + +use serde::Deserialize; +use serde::Serialize; +use serde_json::Value as JsonValue; +use tokio::sync::oneshot; +use tokio_util::sync::CancellationToken; + +use crate::CodeModeNestedToolCall; +use crate::ExecuteRequest; +use crate::RuntimeResponse; +use crate::WaitOutcome; +use crate::WaitRequest; + +pub type CodeModeSessionResultFuture<'a, T> = + Pin> + Send + 'a>>; +pub type CodeModeSessionProviderFuture<'a> = + CodeModeSessionResultFuture<'a, Arc>; +pub type ToolInvocationFuture<'a> = + Pin> + Send + 'a>>; +pub type NotificationFuture<'a> = Pin> + Send + 'a>>; + +#[derive(Clone, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)] +pub struct CellId(String); + +impl CellId { + pub fn new(value: String) -> Self { + Self(value) + } + + pub fn as_str(&self) -> &str { + &self.0 + } +} + +impl AsRef for CellId { + fn as_ref(&self) -> &str { + self.as_str() + } +} + +impl fmt::Display for CellId { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(self.as_str()) + } +} + +pub struct StartedCell { + pub cell_id: CellId, + initial_response: CodeModeSessionResultFuture<'static, RuntimeResponse>, +} + +impl StartedCell { + pub fn new(cell_id: CellId, initial_response_rx: oneshot::Receiver) -> Self { + Self { + cell_id, + initial_response: Box::pin(async move { + initial_response_rx + .await + .map_err(|_| "exec runtime ended unexpectedly".to_string()) + }), + } + } + + pub fn from_result_receiver( + cell_id: CellId, + initial_response_rx: oneshot::Receiver>, + ) -> Self { + Self { + cell_id, + initial_response: Box::pin(async move { + initial_response_rx + .await + .map_err(|_| "exec runtime ended unexpectedly".to_string())? + }), + } + } + + pub async fn initial_response(self) -> Result { + self.initial_response.await + } +} + +/// Host callbacks used by a code-mode session while cells are executing. +pub trait CodeModeSessionDelegate: Send + Sync { + fn invoke_tool<'a>( + &'a self, + invocation: CodeModeNestedToolCall, + cancellation_token: CancellationToken, + ) -> ToolInvocationFuture<'a>; + + fn notify<'a>( + &'a self, + call_id: String, + cell_id: CellId, + text: String, + cancellation_token: CancellationToken, + ) -> NotificationFuture<'a>; + + /// Releases delegate state associated with a cell after it reaches a terminal state. + fn cell_closed(&self, cell_id: &CellId); +} + +/// A durable code-mode session owned by one Codex thread. +/// +/// Cells executed in the same session share stored values. Separate sessions +/// must keep those values isolated. Implementations may execute cells +/// in-process or remotely. +pub trait CodeModeSession: Send + Sync { + fn execute<'a>( + &'a self, + request: ExecuteRequest, + ) -> CodeModeSessionResultFuture<'a, StartedCell>; + + fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome>; + + fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome>; + + fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()>; +} + +/// Creates code-mode sessions for Codex threads. +/// +/// Implementations may share a remote host process across all sessions created +/// by one provider. +pub trait CodeModeSessionProvider: Send + Sync { + fn create_session<'a>( + &'a self, + delegate: Arc, + ) -> CodeModeSessionProviderFuture<'a>; +} + +#[cfg(test)] +#[path = "session_tests.rs"] +mod tests; diff --git a/codex-rs/code-mode-protocol/src/session_tests.rs b/codex-rs/code-mode-protocol/src/session_tests.rs new file mode 100644 index 0000000000..d0f6491105 --- /dev/null +++ b/codex-rs/code-mode-protocol/src/session_tests.rs @@ -0,0 +1,19 @@ +use pretty_assertions::assert_eq; +use tokio::sync::oneshot; + +use super::CellId; +use super::StartedCell; + +#[tokio::test] +async fn started_cell_preserves_remote_initial_response_errors() { + let (response_tx, response_rx) = oneshot::channel(); + response_tx + .send(Err("remote runtime failed".to_string())) + .expect("initial response receiver should be open"); + let started = StartedCell::from_result_receiver(CellId::new("1".to_string()), response_rx); + + assert_eq!( + started.initial_response().await, + Err("remote runtime failed".to_string()) + ); +} diff --git a/codex-rs/code-mode-protocol/src/wire.rs b/codex-rs/code-mode-protocol/src/wire.rs new file mode 100644 index 0000000000..7a637cb767 --- /dev/null +++ b/codex-rs/code-mode-protocol/src/wire.rs @@ -0,0 +1,153 @@ +use std::io; + +use serde::Deserialize; +use serde::Serialize; +use serde::de::DeserializeOwned; +use serde_json::Value as JsonValue; +use tokio::io::AsyncRead; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWrite; +use tokio::io::AsyncWriteExt; + +use crate::CellId; +use crate::CodeModeNestedToolCall; +use crate::ExecuteRequest; +use crate::RuntimeResponse; +use crate::WaitOutcome; +use crate::WaitRequest; + +const MAX_FRAME_BYTES: usize = 16 * 1024 * 1024; + +pub type RequestId = u64; +pub type SessionId = u64; +pub type DelegateRequestId = u64; + +#[derive(Debug, Deserialize, Serialize)] +pub enum ClientMessage { + Request { + id: RequestId, + request: HostRequest, + }, + DelegateResponse { + id: DelegateRequestId, + response: Result, + }, +} + +#[derive(Debug, Deserialize, Serialize)] +pub enum HostMessage { + Response { + id: RequestId, + response: Result, + }, + InitialResponse { + id: RequestId, + response: Result, + }, + DelegateRequest { + id: DelegateRequestId, + session_id: SessionId, + request: DelegateRequest, + }, + CancelDelegateRequest { + id: DelegateRequestId, + }, + CellClosed { + session_id: SessionId, + cell_id: CellId, + }, +} + +#[derive(Debug, Deserialize, Serialize)] +pub enum HostRequest { + CreateSession, + Execute { + session_id: SessionId, + request: ExecuteRequest, + }, + Wait { + session_id: SessionId, + request: WaitRequest, + }, + Terminate { + session_id: SessionId, + cell_id: CellId, + }, + ShutdownSession { + session_id: SessionId, + }, +} + +#[derive(Debug, Deserialize, Serialize)] +pub enum HostResponse { + SessionCreated { session_id: SessionId }, + ExecutionStarted { cell_id: CellId }, + WaitCompleted { outcome: WaitOutcome }, + SessionShutdown, +} + +#[derive(Debug, Deserialize, Serialize)] +pub enum DelegateRequest { + InvokeTool(CodeModeNestedToolCall), + Notify { + call_id: String, + cell_id: CellId, + text: String, + }, +} + +#[derive(Debug, Deserialize, Serialize)] +pub enum DelegateResponse { + ToolResult(JsonValue), + NotificationDelivered, +} + +pub async fn read_frame(reader: &mut R) -> io::Result> +where + R: AsyncRead + Unpin, + T: DeserializeOwned, +{ + let length = match reader.read_u32().await { + Ok(length) => length as usize, + Err(err) if err.kind() == io::ErrorKind::UnexpectedEof => return Ok(None), + Err(err) => return Err(err), + }; + if length > MAX_FRAME_BYTES { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("code-mode IPC frame exceeds {MAX_FRAME_BYTES} bytes"), + )); + } + let mut payload = vec![0; length]; + reader.read_exact(&mut payload).await?; + serde_json::from_slice(&payload) + .map(Some) + .map_err(io::Error::other) +} + +pub async fn write_frame(writer: &mut W, message: &T) -> io::Result<()> +where + W: AsyncWrite + Unpin, + T: Serialize, +{ + let payload = serde_json::to_vec(message).map_err(io::Error::other)?; + let length = u32::try_from(payload.len()).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidInput, + "code-mode IPC frame length exceeds u32", + ) + })?; + if payload.len() > MAX_FRAME_BYTES { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!("code-mode IPC frame exceeds {MAX_FRAME_BYTES} bytes"), + )); + } + writer.write_u32(length).await?; + writer.write_all(&payload).await?; + writer.flush().await +} + +#[cfg(test)] +#[path = "wire_tests.rs"] +mod tests; diff --git a/codex-rs/code-mode-protocol/src/wire_tests.rs b/codex-rs/code-mode-protocol/src/wire_tests.rs new file mode 100644 index 0000000000..61b672063b --- /dev/null +++ b/codex-rs/code-mode-protocol/src/wire_tests.rs @@ -0,0 +1,28 @@ +use tokio::io::duplex; + +use super::ClientMessage; +use super::HostRequest; +use super::read_frame; +use super::write_frame; + +#[tokio::test] +async fn frame_round_trip_preserves_message() { + let (mut client, mut server) = duplex(1024); + let message = ClientMessage::Request { + id: 7, + request: HostRequest::CreateSession, + }; + + write_frame(&mut client, &message) + .await + .expect("write frame"); + let decoded = read_frame(&mut server).await.expect("read frame"); + + assert!(matches!( + decoded, + Some(ClientMessage::Request { + id: 7, + request: HostRequest::CreateSession, + }) + )); +} diff --git a/codex-rs/code-mode/Cargo.toml b/codex-rs/code-mode/Cargo.toml index 879404d8aa..c2e9a68338 100644 --- a/codex-rs/code-mode/Cargo.toml +++ b/codex-rs/code-mode/Cargo.toml @@ -16,6 +16,7 @@ sandbox = ["v8/v8_enable_sandbox"] workspace = true [dependencies] +codex-code-mode-protocol = { workspace = true } codex-protocol = { workspace = true } deno_core_icudata = { workspace = true } serde = { workspace = true, features = ["derive"] } diff --git a/codex-rs/code-mode/src/lib.rs b/codex-rs/code-mode/src/lib.rs index 37ba1d2a91..2699258271 100644 --- a/codex-rs/code-mode/src/lib.rs +++ b/codex-rs/code-mode/src/lib.rs @@ -1,46 +1,7 @@ -mod description; -mod response; mod runtime; mod service; -pub use description::CODE_MODE_PRAGMA_PREFIX; -pub use description::CodeModeToolKind; -pub use description::ToolDefinition; -pub use description::ToolNamespaceDescription; -pub use description::augment_tool_definition; -pub use description::build_exec_tool_description; -pub use description::build_wait_tool_description; -pub use description::is_code_mode_nested_tool; -pub use description::normalize_code_mode_identifier; -pub use description::parse_exec_source; -pub use description::render_code_mode_sample; -pub use description::render_json_schema_to_typescript; -pub use response::DEFAULT_IMAGE_DETAIL; -pub use response::FunctionCallOutputContentItem; -pub use response::ImageDetail; -pub use runtime::CodeModeNestedToolCall; -pub use runtime::DEFAULT_EXEC_YIELD_TIME_MS; -pub use runtime::DEFAULT_MAX_OUTPUT_TOKENS_PER_EXEC_CALL; -pub use runtime::DEFAULT_WAIT_YIELD_TIME_MS; -pub use runtime::ExecuteRequest; -pub use runtime::ExecuteToPendingOutcome; -pub use runtime::RuntimeResponse; -pub use runtime::WaitOutcome; -pub use runtime::WaitRequest; -pub use runtime::WaitToPendingOutcome; -pub use runtime::WaitToPendingRequest; -pub use service::CellId; +pub use codex_code_mode_protocol::*; pub use service::CodeModeService; -pub use service::CodeModeSession; -pub use service::CodeModeSessionDelegate; -pub use service::CodeModeSessionProvider; -pub use service::CodeModeSessionProviderFuture; -pub use service::CodeModeSessionResultFuture; pub use service::InProcessCodeModeSessionProvider; pub use service::NoopCodeModeSessionDelegate; -pub use service::NotificationFuture; -pub use service::StartedCell; -pub use service::ToolInvocationFuture; - -pub const PUBLIC_TOOL_NAME: &str = "exec"; -pub const WAIT_TOOL_NAME: &str = "wait"; diff --git a/codex-rs/code-mode/src/runtime/callbacks.rs b/codex-rs/code-mode/src/runtime/callbacks.rs index dde63617ec..7b09ba5ffd 100644 --- a/codex-rs/code-mode/src/runtime/callbacks.rs +++ b/codex-rs/code-mode/src/runtime/callbacks.rs @@ -1,4 +1,4 @@ -use crate::response::FunctionCallOutputContentItem; +use codex_code_mode_protocol::FunctionCallOutputContentItem; use super::EXIT_SENTINEL; use super::RuntimeEvent; diff --git a/codex-rs/code-mode/src/runtime/mod.rs b/codex-rs/code-mode/src/runtime/mod.rs index 21757328ee..36ffd926cc 100644 --- a/codex-rs/code-mode/src/runtime/mod.rs +++ b/codex-rs/code-mode/src/runtime/mod.rs @@ -9,125 +9,17 @@ use std::sync::OnceLock; use std::sync::mpsc as std_mpsc; use std::thread; +use codex_code_mode_protocol::CodeModeToolKind; +use codex_code_mode_protocol::EnabledToolMetadata; +use codex_code_mode_protocol::ExecuteRequest; +use codex_code_mode_protocol::FunctionCallOutputContentItem; +use codex_code_mode_protocol::enabled_tool_metadata; use codex_protocol::ToolName; -use serde::Serialize; use serde_json::Value as JsonValue; use tokio::sync::mpsc; -use crate::description::CodeModeToolKind; -use crate::description::EnabledToolMetadata; -use crate::description::ToolDefinition; -use crate::description::enabled_tool_metadata; -use crate::response::FunctionCallOutputContentItem; -use crate::service::CellId; - -pub const DEFAULT_EXEC_YIELD_TIME_MS: u64 = 10_000; -pub const DEFAULT_WAIT_YIELD_TIME_MS: u64 = 10_000; -pub const DEFAULT_MAX_OUTPUT_TOKENS_PER_EXEC_CALL: usize = 10_000; const EXIT_SENTINEL: &str = "__codex_code_mode_exit__"; -#[derive(Clone, Debug)] -pub struct ExecuteRequest { - pub tool_call_id: String, - pub enabled_tools: Vec, - pub source: String, - pub yield_time_ms: Option, - pub max_output_tokens: Option, -} - -#[derive(Clone, Debug)] -pub struct WaitRequest { - pub cell_id: CellId, - pub yield_time_ms: u64, -} - -#[derive(Clone, Debug)] -pub struct WaitToPendingRequest { - pub cell_id: CellId, -} - -/// Result of waiting on a code-mode cell. -/// -/// The wrapped `RuntimeResponse` is the model-facing wait result. The enum -/// variant carries the extra lifecycle provenance that `RuntimeResponse` cannot: -/// a failed real cell and a missing-cell wait both use -/// `RuntimeResponse::Result { error_text: Some(..), .. }`, but only the former -/// should be treated as a code-cell lifecycle event. -#[derive(Debug, PartialEq)] -pub enum WaitOutcome { - /// The requested code cell was live when the wait command was accepted. - LiveCell(RuntimeResponse), - /// The requested code cell was not live. - MissingCell(RuntimeResponse), -} - -/// Result of executing a code-mode cell until it either completes or reaches a -/// quiescent pending state. -#[derive(Debug, PartialEq)] -pub enum ExecuteToPendingOutcome { - /// The cell is waiting for more runtime input after draining the runtime - /// input queue that was ready at the pending boundary. - Pending { - cell_id: CellId, - content_items: Vec, - /// Runtime tool-call ids emitted before this paused execution frontier - /// sealed. Hosts can use these ids to drain their tool-call transport - /// before surfacing the pending boundary to callers. - pending_tool_call_ids: Vec, - }, - /// The cell reached a terminal runtime response before going pending. - Completed(RuntimeResponse), -} - -/// Result of resuming a live code-mode cell until it completes or becomes -/// quiescent again. -#[derive(Debug, PartialEq)] -pub enum WaitToPendingOutcome { - /// The requested code cell was live when the wait command was accepted. - LiveCell(ExecuteToPendingOutcome), - /// The requested code cell was not live. - MissingCell(RuntimeResponse), -} - -impl From for RuntimeResponse { - fn from(outcome: WaitOutcome) -> Self { - match outcome { - WaitOutcome::LiveCell(response) | WaitOutcome::MissingCell(response) => response, - } - } -} - -#[derive(Debug, PartialEq, Serialize)] -pub enum RuntimeResponse { - Yielded { - cell_id: CellId, - content_items: Vec, - }, - Terminated { - cell_id: CellId, - content_items: Vec, - }, - Result { - cell_id: CellId, - content_items: Vec, - error_text: Option, - }, -} - -/// Nested tool request emitted by one code-mode cell. -/// -/// Code mode owns the per-cell runtime id. Hosts should preserve it for -/// provenance/debugging, but should still assign their own runtime tool call id -/// if their tool-call graph requires globally unique ids. -#[derive(Debug)] -pub struct CodeModeNestedToolCall { - pub cell_id: CellId, - pub runtime_tool_call_id: String, - pub tool_name: ToolName, - pub tool_kind: CodeModeToolKind, - pub input: Option, -} - #[derive(Debug)] pub(crate) enum RuntimeCommand { ToolResponse { id: String, result: JsonValue }, @@ -326,12 +218,9 @@ fn run_runtime( } let mut pending_promise = pending_promise; - loop { - let Some(command) = next_runtime_command(&event_tx, &command_rx, &control_rx, pending_mode) - else { - break; - }; - + while let Some(command) = + next_runtime_command(&event_tx, &command_rx, &control_rx, pending_mode) + { match command { RuntimeCommand::Terminate => break, RuntimeCommand::ToolResponse { id, result } => { diff --git a/codex-rs/code-mode/src/runtime/value.rs b/codex-rs/code-mode/src/runtime/value.rs index 8d76a832d3..cd6ff6e938 100644 --- a/codex-rs/code-mode/src/runtime/value.rs +++ b/codex-rs/code-mode/src/runtime/value.rs @@ -1,8 +1,8 @@ use serde_json::Value as JsonValue; -use crate::response::DEFAULT_IMAGE_DETAIL; -use crate::response::FunctionCallOutputContentItem; -use crate::response::ImageDetail; +use codex_code_mode_protocol::DEFAULT_IMAGE_DETAIL; +use codex_code_mode_protocol::FunctionCallOutputContentItem; +use codex_code_mode_protocol::ImageDetail; const IMAGE_HELPER_EXPECTS_MESSAGE: &str = "image expects a non-empty image URL string, an object with image_url and optional detail, or a raw MCP image block"; const CODEX_IMAGE_DETAIL_META_KEY: &str = "codex/imageDetail"; diff --git a/codex-rs/code-mode/src/service.rs b/codex-rs/code-mode/src/service.rs index b348d1d35f..ca53edbe54 100644 --- a/codex-rs/code-mode/src/service.rs +++ b/codex-rs/code-mode/src/service.rs @@ -1,15 +1,29 @@ use std::collections::HashMap; -use std::fmt; -use std::future::Future; -use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; use std::time::Duration; -use serde::Deserialize; -use serde::Serialize; +use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::CodeModeNestedToolCall; +use codex_code_mode_protocol::CodeModeSession; +use codex_code_mode_protocol::CodeModeSessionDelegate; +use codex_code_mode_protocol::CodeModeSessionProvider; +use codex_code_mode_protocol::CodeModeSessionProviderFuture; +use codex_code_mode_protocol::CodeModeSessionResultFuture; +use codex_code_mode_protocol::DEFAULT_EXEC_YIELD_TIME_MS; +use codex_code_mode_protocol::ExecuteRequest; +use codex_code_mode_protocol::ExecuteToPendingOutcome; +use codex_code_mode_protocol::FunctionCallOutputContentItem; +use codex_code_mode_protocol::NotificationFuture; +use codex_code_mode_protocol::RuntimeResponse; +use codex_code_mode_protocol::StartedCell; +use codex_code_mode_protocol::ToolInvocationFuture; +use codex_code_mode_protocol::WaitOutcome; +use codex_code_mode_protocol::WaitRequest; +use codex_code_mode_protocol::WaitToPendingOutcome; +use codex_code_mode_protocol::WaitToPendingRequest; use serde_json::Value as JsonValue; use tokio::sync::Mutex; use tokio::sync::mpsc; @@ -18,88 +32,12 @@ use tokio::task::JoinSet; use tokio_util::sync::CancellationToken; use tracing::warn; -use crate::FunctionCallOutputContentItem; -use crate::runtime::CodeModeNestedToolCall; -use crate::runtime::DEFAULT_EXEC_YIELD_TIME_MS; -use crate::runtime::ExecuteRequest; -use crate::runtime::ExecuteToPendingOutcome; use crate::runtime::PendingRuntimeMode; use crate::runtime::RuntimeCommand; use crate::runtime::RuntimeControlCommand; use crate::runtime::RuntimeEvent; -use crate::runtime::RuntimeResponse; -use crate::runtime::WaitOutcome; -use crate::runtime::WaitRequest; -use crate::runtime::WaitToPendingOutcome; -use crate::runtime::WaitToPendingRequest; use crate::runtime::spawn_runtime; -pub type CodeModeSessionResultFuture<'a, T> = - Pin> + Send + 'a>>; -pub type CodeModeSessionProviderFuture<'a> = - CodeModeSessionResultFuture<'a, Arc>; -pub type ToolInvocationFuture<'a> = - Pin> + Send + 'a>>; -pub type NotificationFuture<'a> = Pin> + Send + 'a>>; - -#[derive(Clone, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)] -pub struct CellId(String); - -impl CellId { - pub fn new(value: String) -> Self { - Self(value) - } - - pub fn as_str(&self) -> &str { - &self.0 - } -} - -impl AsRef for CellId { - fn as_ref(&self) -> &str { - self.as_str() - } -} - -impl fmt::Display for CellId { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter.write_str(self.as_str()) - } -} - -pub struct StartedCell { - pub cell_id: CellId, - initial_response_rx: oneshot::Receiver, -} - -impl StartedCell { - pub async fn initial_response(self) -> Result { - self.initial_response_rx - .await - .map_err(|_| "exec runtime ended unexpectedly".to_string()) - } -} - -/// Host callbacks used by a code-mode session while cells are executing. -pub trait CodeModeSessionDelegate: Send + Sync { - fn invoke_tool<'a>( - &'a self, - invocation: CodeModeNestedToolCall, - cancellation_token: CancellationToken, - ) -> ToolInvocationFuture<'a>; - - fn notify<'a>( - &'a self, - call_id: String, - cell_id: CellId, - text: String, - cancellation_token: CancellationToken, - ) -> NotificationFuture<'a>; - - /// Releases delegate state associated with a cell after it reaches a terminal state. - fn cell_closed(&self, cell_id: &CellId); -} - pub struct NoopCodeModeSessionDelegate; impl CodeModeSessionDelegate for NoopCodeModeSessionDelegate { @@ -127,35 +65,6 @@ impl CodeModeSessionDelegate for NoopCodeModeSessionDelegate { fn cell_closed(&self, _cell_id: &CellId) {} } -/// A durable code-mode session owned by one Codex thread. -/// -/// Cells executed in the same session share stored values. Separate sessions -/// must keep those values isolated. Implementations may execute cells -/// in-process or remotely. -pub trait CodeModeSession: Send + Sync { - fn execute<'a>( - &'a self, - request: ExecuteRequest, - ) -> CodeModeSessionResultFuture<'a, StartedCell>; - - fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome>; - - fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome>; - - fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()>; -} - -/// Creates code-mode sessions for one Codex thread. -/// -/// Providers choose where a session executes and receive the host delegate that -/// the session should use for nested tool calls and notifications. -pub trait CodeModeSessionProvider: Send + Sync { - fn create_session<'a>( - &'a self, - delegate: Arc, - ) -> CodeModeSessionProviderFuture<'a>; -} - #[derive(Default)] pub struct InProcessCodeModeSessionProvider; @@ -233,10 +142,7 @@ impl CodeModeService { ) .await?; - Ok(StartedCell { - cell_id, - initial_response_rx: response_rx, - }) + Ok(StartedCell::new(cell_id, response_rx)) } pub async fn execute_to_pending( @@ -876,10 +782,10 @@ mod tests { use super::WaitToPendingRequest; use super::run_cell_control; use crate::CodeModeToolKind; + use crate::ExecuteRequest; + use crate::ExecuteToPendingOutcome; use crate::FunctionCallOutputContentItem; use crate::ToolDefinition; - use crate::runtime::ExecuteRequest; - use crate::runtime::ExecuteToPendingOutcome; use crate::runtime::RuntimeEvent; use crate::runtime::spawn_runtime; diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index 352e8051e4..079119fd03 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -29,7 +29,8 @@ codex-api = { workspace = true } codex-app-server-protocol = { workspace = true } codex-apply-patch = { workspace = true } codex-async-utils = { workspace = true } -codex-code-mode = { workspace = true } +codex-code-mode-client = { workspace = true } +codex-code-mode-protocol = { workspace = true } codex-connectors = { workspace = true } codex-context-fragments = { workspace = true } codex-config = { workspace = true } diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 5b68ecf917..4da86da3fe 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -107,6 +107,7 @@ pub(crate) async fn run_codex_thread_interactive( analytics_events_client: Some(parent_session.services.analytics_events_client.clone()), thread_store: Arc::clone(&parent_session.services.thread_store), attestation_provider: parent_session.services.attestation_provider.clone(), + code_mode_session_provider: Arc::clone(&parent_session.services.code_mode_session_provider), inherited_multi_agent_version: Some(MultiAgentVersion::Disabled), })) .or_cancel(&cancel_token) diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 8424ce216f..b1485aa82e 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -422,6 +422,8 @@ pub(crate) struct CodexSpawnArgs { pub(crate) analytics_events_client: Option, pub(crate) thread_store: Arc, pub(crate) attestation_provider: Option>, + pub(crate) code_mode_session_provider: + Arc, pub(crate) inherited_multi_agent_version: Option, } @@ -503,6 +505,7 @@ impl Codex { analytics_events_client, thread_store, attestation_provider, + code_mode_session_provider, inherited_multi_agent_version, } = args; let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY); @@ -651,6 +654,7 @@ impl Codex { thread_store, parent_rollout_thread_trace, attestation_provider, + code_mode_session_provider, multi_agent_version, )) .await diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index e42cf202e8..b7f9ff51dd 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -491,6 +491,7 @@ impl Session { thread_store: Arc, parent_rollout_thread_trace: ThreadTraceContext, attestation_provider: Option>, + code_mode_session_provider: Arc, multi_agent_version: Option, ) -> anyhow::Result> { debug!( @@ -1033,7 +1034,10 @@ impl Session { session_configuration.parent_thread_id, ), ), - code_mode_service: crate::tools::code_mode::CodeModeService::new(), + code_mode_service: crate::tools::code_mode::CodeModeService::new( + Arc::clone(&code_mode_session_provider), + ), + code_mode_session_provider, environment_manager, }; let (out_of_band_elicitation_paused, _out_of_band_elicitation_paused_rx) = diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 29de7ca662..25e4f12678 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -4723,6 +4723,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() { )), codex_rollout_trace::ThreadTraceContext::disabled(), /*attestation_provider*/ None, + Arc::new(codex_code_mode_client::IpcCodeModeSessionProvider::default()), Some(config.multi_agent_version_from_features()), ) .await; @@ -4892,7 +4893,12 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { Session::build_model_client_beta_features_header(config.as_ref()), /*attestation_provider*/ None, ), - code_mode_service: crate::tools::code_mode::CodeModeService::new(), + code_mode_service: crate::tools::code_mode::CodeModeService::new(Arc::new( + codex_code_mode_client::IpcCodeModeSessionProvider::default(), + )), + code_mode_session_provider: Arc::new( + codex_code_mode_client::IpcCodeModeSessionProvider::default(), + ), environment_manager: Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), }; @@ -5064,6 +5070,7 @@ async fn make_session_with_config_and_rx( )), codex_rollout_trace::ThreadTraceContext::disabled(), /*attestation_provider*/ None, + Arc::new(codex_code_mode_client::IpcCodeModeSessionProvider::default()), Some(config.multi_agent_version_from_features()), ) .await?; @@ -5173,6 +5180,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx( )), codex_rollout_trace::ThreadTraceContext::disabled(), /*attestation_provider*/ None, + Arc::new(codex_code_mode_client::IpcCodeModeSessionProvider::default()), Some(config.multi_agent_version_from_features()), ) .await?; @@ -6970,7 +6978,12 @@ where Session::build_model_client_beta_features_header(config.as_ref()), /*attestation_provider*/ None, ), - code_mode_service: crate::tools::code_mode::CodeModeService::new(), + code_mode_service: crate::tools::code_mode::CodeModeService::new(Arc::new( + codex_code_mode_client::IpcCodeModeSessionProvider::default(), + )), + code_mode_session_provider: Arc::new( + codex_code_mode_client::IpcCodeModeSessionProvider::default(), + ), environment_manager: Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), }; diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index c73add2920..facfd07cc5 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -735,6 +735,9 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { analytics_events_client: None, thread_store, attestation_provider: None, + code_mode_session_provider: Arc::new( + codex_code_mode_client::IpcCodeModeSessionProvider::default(), + ), inherited_multi_agent_version: None, }) .await diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index b4c499118d..afd68bc3f6 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -19,6 +19,7 @@ use anyhow::Result; use arc_swap::ArcSwap; use arc_swap::ArcSwapOption; use codex_analytics::AnalyticsEventsClient; +use codex_code_mode_protocol::CodeModeSessionProvider; use codex_core_plugins::PluginsManager; use codex_exec_server::EnvironmentManager; use codex_extension_api::ExtensionData; @@ -79,6 +80,7 @@ pub(crate) struct SessionServices { /// Session-scoped model client shared across turns. pub(crate) model_client: ModelClient, pub(crate) code_mode_service: CodeModeService, + pub(crate) code_mode_session_provider: Arc, /// Shared process-level environment registry. Sessions carry an `Arc` handle so they can pass /// the same manager through child-thread spawn paths without reconstructing it. pub(crate) environment_manager: Arc, diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 5758c0cadf..031714a4b9 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -19,6 +19,8 @@ use crate::tasks::interrupted_turn_history_marker; use codex_analytics::AnalyticsEventsClient; use codex_app_server_protocol::ThreadHistoryBuilder; use codex_app_server_protocol::TurnStatus; +use codex_code_mode_client::IpcCodeModeSessionProvider; +use codex_code_mode_protocol::CodeModeSessionProvider; use codex_core_plugins::PluginsManager; use codex_exec_server::EnvironmentManager; use codex_extension_api::ExtensionDataInit; @@ -215,6 +217,7 @@ pub(crate) struct ThreadManagerState { installation_id: String, analytics_events_client: Option, state_db: Option, + code_mode_session_provider: Arc, // Captures submitted ops for testing purpose when test mode is enabled. ops_log: Option, } @@ -298,6 +301,7 @@ impl ThreadManager { installation_id, analytics_events_client, state_db, + code_mode_session_provider: Arc::new(IpcCodeModeSessionProvider::default()), ops_log: should_use_test_thread_manager_behavior() .then(|| Arc::new(std::sync::Mutex::new(Vec::new()))), }), @@ -399,6 +403,7 @@ impl ThreadManager { installation_id, analytics_events_client: None, state_db, + code_mode_session_provider: Arc::new(IpcCodeModeSessionProvider::default()), ops_log: should_use_test_thread_manager_behavior() .then(|| Arc::new(std::sync::Mutex::new(Vec::new()))), }), @@ -410,6 +415,18 @@ impl ThreadManager { self.state.session_source.clone() } + #[doc(hidden)] + pub fn with_code_mode_session_provider( + mut self, + provider: Arc, + ) -> Self { + let Some(state) = Arc::get_mut(&mut self.state) else { + panic!("code-mode provider must be configured before sharing the thread manager"); + }; + state.code_mode_session_provider = provider; + self + } + pub fn auth_manager(&self) -> Arc { self.state.auth_manager.clone() } @@ -1348,6 +1365,7 @@ impl ThreadManagerState { analytics_events_client: self.analytics_events_client.clone(), thread_store: Arc::clone(&self.thread_store), attestation_provider: self.attestation_provider.clone(), + code_mode_session_provider: Arc::clone(&self.code_mode_session_provider), inherited_multi_agent_version: multi_agent_version, })) .await?; diff --git a/codex-rs/core/src/tools/code_mode/delegate.rs b/codex-rs/core/src/tools/code_mode/delegate.rs index 06bb71a280..cc56caeab9 100644 --- a/codex-rs/core/src/tools/code_mode/delegate.rs +++ b/codex-rs/core/src/tools/code_mode/delegate.rs @@ -2,11 +2,11 @@ use std::collections::HashMap; use std::sync::Arc; use std::sync::Mutex; -use codex_code_mode::CellId; -use codex_code_mode::CodeModeNestedToolCall; -use codex_code_mode::CodeModeSessionDelegate; -use codex_code_mode::NotificationFuture; -use codex_code_mode::ToolInvocationFuture; +use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::CodeModeNestedToolCall; +use codex_code_mode_protocol::CodeModeSessionDelegate; +use codex_code_mode_protocol::NotificationFuture; +use codex_code_mode_protocol::ToolInvocationFuture; use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::models::ResponseItem; use serde_json::Value as JsonValue; diff --git a/codex-rs/core/src/tools/code_mode/execute_handler.rs b/codex-rs/core/src/tools/code_mode/execute_handler.rs index 2e46c1bd6b..ea24be4e15 100644 --- a/codex-rs/core/src/tools/code_mode/execute_handler.rs +++ b/codex-rs/core/src/tools/code_mode/execute_handler.rs @@ -33,8 +33,8 @@ impl CodeModeExecuteHandler { call_id: String, code: String, ) -> Result { - let args = - codex_code_mode::parse_exec_source(&code).map_err(FunctionCallError::RespondToModel)?; + let args = codex_code_mode_protocol::parse_exec_source(&code) + .map_err(FunctionCallError::RespondToModel)?; let exec = ExecContext { session, turn }; let enabled_tools = codex_tools::collect_code_mode_tool_definitions(&self.nested_tool_specs); @@ -43,7 +43,7 @@ impl CodeModeExecuteHandler { .session .services .code_mode_service - .execute(codex_code_mode::ExecuteRequest { + .execute(codex_code_mode_protocol::ExecuteRequest { tool_call_id: call_id.clone(), enabled_tools, source: args.code.clone(), @@ -78,7 +78,10 @@ impl CodeModeExecuteHandler { code_cell_trace.record_initial_response(&response); // Yielded cells keep running, so terminal lifecycle is only emitted // here when the first response also ended the runtime. - if !matches!(response, codex_code_mode::RuntimeResponse::Yielded { .. }) { + if !matches!( + response, + codex_code_mode_protocol::RuntimeResponse::Yielded { .. } + ) { code_cell_trace.record_ended(&response); exec.session .services diff --git a/codex-rs/core/src/tools/code_mode/execute_spec.rs b/codex-rs/core/src/tools/code_mode/execute_spec.rs index 0a858bd206..404525bea6 100644 --- a/codex-rs/core/src/tools/code_mode/execute_spec.rs +++ b/codex-rs/core/src/tools/code_mode/execute_spec.rs @@ -1,4 +1,4 @@ -use codex_code_mode::ToolDefinition as CodeModeToolDefinition; +use codex_code_mode_protocol::ToolDefinition as CodeModeToolDefinition; use codex_tools::FreeformTool; use codex_tools::FreeformToolFormat; use codex_tools::ToolSpec; @@ -6,7 +6,7 @@ use std::collections::BTreeMap; pub(crate) fn create_code_mode_tool( enabled_tools: &[CodeModeToolDefinition], - namespace_descriptions: &BTreeMap, + namespace_descriptions: &BTreeMap, code_mode_only: bool, deferred_tools_available: bool, ) -> ToolSpec { @@ -21,8 +21,8 @@ SOURCE: /[\s\S]+/ "#; ToolSpec::Freeform(FreeformTool { - name: codex_code_mode::PUBLIC_TOOL_NAME.to_string(), - description: codex_code_mode::build_exec_tool_description( + name: codex_code_mode_protocol::PUBLIC_TOOL_NAME.to_string(), + description: codex_code_mode_protocol::build_exec_tool_description( enabled_tools, namespace_descriptions, code_mode_only, @@ -44,11 +44,11 @@ mod tests { #[test] fn create_code_mode_tool_matches_expected_spec() { - let enabled_tools = vec![codex_code_mode::ToolDefinition { + let enabled_tools = vec![codex_code_mode_protocol::ToolDefinition { name: "update_plan".to_string(), tool_name: ToolName::plain("update_plan"), description: "Update the plan".to_string(), - kind: codex_code_mode::CodeModeToolKind::Function, + kind: codex_code_mode_protocol::CodeModeToolKind::Function, input_schema: None, output_schema: None, }]; @@ -61,8 +61,8 @@ mod tests { /*deferred_tools_available*/ false, ), ToolSpec::Freeform(FreeformTool { - name: codex_code_mode::PUBLIC_TOOL_NAME.to_string(), - description: codex_code_mode::build_exec_tool_description( + name: codex_code_mode_protocol::PUBLIC_TOOL_NAME.to_string(), + description: codex_code_mode_protocol::build_exec_tool_description( &enabled_tools, &BTreeMap::new(), /*code_mode_only*/ true, diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index ade23cf3b7..88086326fe 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -8,13 +8,15 @@ pub(crate) mod wait_spec; use std::sync::Arc; use std::time::Duration; -use codex_code_mode::CellId; -use codex_code_mode::CodeModeNestedToolCall; -use codex_code_mode::CodeModeSession; -use codex_code_mode::CodeModeToolKind; -use codex_code_mode::RuntimeResponse; +use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::CodeModeNestedToolCall; +use codex_code_mode_protocol::CodeModeSession; +use codex_code_mode_protocol::CodeModeSessionProvider; +use codex_code_mode_protocol::CodeModeToolKind; +use codex_code_mode_protocol::RuntimeResponse; use codex_protocol::models::FunctionCallOutputContentItem; use serde_json::Value as JsonValue; +use tokio::sync::OnceCell; use tokio_util::sync::CancellationToken; use crate::function_tool::FunctionCallError; @@ -42,9 +44,10 @@ pub(crate) use execute_handler::CodeModeExecuteHandler; use response_adapter::into_function_call_output_content_items; pub(crate) use wait_handler::CodeModeWaitHandler; -pub(crate) const PUBLIC_TOOL_NAME: &str = codex_code_mode::PUBLIC_TOOL_NAME; -pub(crate) const WAIT_TOOL_NAME: &str = codex_code_mode::WAIT_TOOL_NAME; -pub(crate) const DEFAULT_WAIT_YIELD_TIME_MS: u64 = codex_code_mode::DEFAULT_WAIT_YIELD_TIME_MS; +pub(crate) const PUBLIC_TOOL_NAME: &str = codex_code_mode_protocol::PUBLIC_TOOL_NAME; +pub(crate) const WAIT_TOOL_NAME: &str = codex_code_mode_protocol::WAIT_TOOL_NAME; +pub(crate) const DEFAULT_WAIT_YIELD_TIME_MS: u64 = + codex_code_mode_protocol::DEFAULT_WAIT_YIELD_TIME_MS; /// Returns true for the un-namespaced code-mode `exec` tool. pub(crate) fn is_exec_tool_name(tool_name: &ToolName) -> bool { @@ -58,50 +61,50 @@ pub(crate) struct ExecContext { } pub(crate) struct CodeModeService { - session: Option>, + session: OnceCell>, + provider: Arc, dispatch_broker: Arc, } impl CodeModeService { - pub(crate) fn new() -> Self { + pub(crate) fn new(provider: Arc) -> Self { let dispatch_broker = Arc::new(CodeModeDispatchBroker::new()); Self { - session: Some(Arc::new(codex_code_mode::CodeModeService::with_delegate( - dispatch_broker.clone(), - ))), + session: OnceCell::new(), + provider, dispatch_broker, } } pub(crate) async fn execute( &self, - request: codex_code_mode::ExecuteRequest, - ) -> Result { - self.session()?.execute(request).await + request: codex_code_mode_protocol::ExecuteRequest, + ) -> Result { + self.session().await?.execute(request).await } pub(crate) async fn wait( &self, - request: codex_code_mode::WaitRequest, - ) -> Result { - self.session()?.wait(request).await + request: codex_code_mode_protocol::WaitRequest, + ) -> Result { + self.session().await?.wait(request).await } pub(crate) async fn terminate( &self, cell_id: CellId, - ) -> Result { - self.session()?.terminate(cell_id).await + ) -> Result { + self.session().await?.terminate(cell_id).await } pub(crate) async fn shutdown(&self) -> Result<(), String> { - match &self.session { + match self.session.get() { Some(session) => session.shutdown().await, None => Ok(()), } } - pub(crate) fn mark_cell_ready_for_dispatch(&self, cell_id: &codex_code_mode::CellId) { + pub(crate) fn mark_cell_ready_for_dispatch(&self, cell_id: &codex_code_mode_protocol::CellId) { self.dispatch_broker.mark_cell_ready_for_dispatch(cell_id); } @@ -116,9 +119,7 @@ impl CodeModeService { router: Arc, tracker: SharedTurnDiffTracker, ) -> Option { - if !matches!(turn.tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly) - || self.session.is_none() - { + if !matches!(turn.tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly) { return None; } @@ -132,10 +133,10 @@ impl CodeModeService { ) } - fn session(&self) -> Result<&Arc, String> { + async fn session(&self) -> Result<&Arc, String> { self.session - .as_ref() - .ok_or_else(|| "code mode is unavailable".to_string()) + .get_or_try_init(|| self.provider.create_session(self.dispatch_broker.clone())) + .await } } @@ -322,7 +323,7 @@ fn build_freeform_tool_payload( mod tests { use super::build_nested_tool_payload; use crate::tools::context::ToolPayload; - use codex_code_mode::CodeModeToolKind; + use codex_code_mode_protocol::CodeModeToolKind; use codex_tools::ToolName; use serde_json::json; diff --git a/codex-rs/core/src/tools/code_mode/response_adapter.rs b/codex-rs/core/src/tools/code_mode/response_adapter.rs index e20cf6a071..37782f6f74 100644 --- a/codex-rs/core/src/tools/code_mode/response_adapter.rs +++ b/codex-rs/core/src/tools/code_mode/response_adapter.rs @@ -1,4 +1,4 @@ -use codex_code_mode::ImageDetail as CodeModeImageDetail; +use codex_code_mode_protocol::ImageDetail as CodeModeImageDetail; use codex_protocol::models::DEFAULT_IMAGE_DETAIL; use codex_protocol::models::FunctionCallOutputContentItem; use codex_protocol::models::ImageDetail; @@ -8,7 +8,7 @@ trait IntoProtocol { } pub(super) fn into_function_call_output_content_items( - items: Vec, + items: Vec, ) -> Vec { items.into_iter().map(IntoProtocol::into_protocol).collect() } @@ -26,22 +26,23 @@ impl IntoProtocol for CodeModeImageDetail { } impl IntoProtocol - for codex_code_mode::FunctionCallOutputContentItem + for codex_code_mode_protocol::FunctionCallOutputContentItem { fn into_protocol(self) -> FunctionCallOutputContentItem { let value = self; match value { - codex_code_mode::FunctionCallOutputContentItem::InputText { text } => { + codex_code_mode_protocol::FunctionCallOutputContentItem::InputText { text } => { FunctionCallOutputContentItem::InputText { text } } - codex_code_mode::FunctionCallOutputContentItem::InputImage { image_url, detail } => { - FunctionCallOutputContentItem::InputImage { - image_url, - detail: detail - .map(IntoProtocol::into_protocol) - .or(Some(DEFAULT_IMAGE_DETAIL)), - } - } + codex_code_mode_protocol::FunctionCallOutputContentItem::InputImage { + image_url, + detail, + } => FunctionCallOutputContentItem::InputImage { + image_url, + detail: detail + .map(IntoProtocol::into_protocol) + .or(Some(DEFAULT_IMAGE_DETAIL)), + }, } } } diff --git a/codex-rs/core/src/tools/code_mode/wait_handler.rs b/codex-rs/core/src/tools/code_mode/wait_handler.rs index 953391a573..0db26296ec 100644 --- a/codex-rs/core/src/tools/code_mode/wait_handler.rs +++ b/codex-rs/core/src/tools/code_mode/wait_handler.rs @@ -78,7 +78,7 @@ impl CodeModeWaitHandler { let args: ExecWaitArgs = parse_arguments(&arguments)?; let exec = ExecContext { session, turn }; let started_at = std::time::Instant::now(); - let cell_id = codex_code_mode::CellId::new(args.cell_id); + let cell_id = codex_code_mode_protocol::CellId::new(args.cell_id); let wait_response = if args.terminate { exec.session .services @@ -89,23 +89,30 @@ impl CodeModeWaitHandler { exec.session .services .code_mode_service - .wait(codex_code_mode::WaitRequest { + .wait(codex_code_mode_protocol::WaitRequest { cell_id, yield_time_ms: args.yield_time_ms, }) .await } .map_err(FunctionCallError::RespondToModel)?; - if let codex_code_mode::WaitOutcome::LiveCell(response) = &wait_response - && !matches!(response, codex_code_mode::RuntimeResponse::Yielded { .. }) + if let codex_code_mode_protocol::WaitOutcome::LiveCell(response) = &wait_response + && !matches!( + response, + codex_code_mode_protocol::RuntimeResponse::Yielded { .. } + ) { // Only a live-cell wait can close a CodeCell. A missing // cell is still an ordinary `wait` tool result, but there // is no runtime object for the reducer to complete. let runtime_cell_id = match response { - codex_code_mode::RuntimeResponse::Yielded { cell_id, .. } - | codex_code_mode::RuntimeResponse::Terminated { cell_id, .. } - | codex_code_mode::RuntimeResponse::Result { cell_id, .. } => cell_id, + codex_code_mode_protocol::RuntimeResponse::Yielded { cell_id, .. } + | codex_code_mode_protocol::RuntimeResponse::Terminated { + cell_id, .. + } + | codex_code_mode_protocol::RuntimeResponse::Result { cell_id, .. } => { + cell_id + } }; exec.session .services diff --git a/codex-rs/core/src/tools/code_mode/wait_spec.rs b/codex-rs/core/src/tools/code_mode/wait_spec.rs index 72bb09084c..8232705137 100644 --- a/codex-rs/core/src/tools/code_mode/wait_spec.rs +++ b/codex-rs/core/src/tools/code_mode/wait_spec.rs @@ -30,11 +30,11 @@ pub(crate) fn create_wait_tool() -> ToolSpec { ]); ToolSpec::Function(ResponsesApiTool { - name: codex_code_mode::WAIT_TOOL_NAME.to_string(), + name: codex_code_mode_protocol::WAIT_TOOL_NAME.to_string(), description: format!( "Waits on a yielded `{}` cell and returns new output or completion.\n{}", - codex_code_mode::PUBLIC_TOOL_NAME, - codex_code_mode::build_wait_tool_description().trim() + codex_code_mode_protocol::PUBLIC_TOOL_NAME, + codex_code_mode_protocol::build_wait_tool_description().trim() ), strict: false, parameters: JsonSchema::object( @@ -57,11 +57,11 @@ mod tests { assert_eq!( create_wait_tool(), ToolSpec::Function(ResponsesApiTool { - name: codex_code_mode::WAIT_TOOL_NAME.to_string(), + name: codex_code_mode_protocol::WAIT_TOOL_NAME.to_string(), description: format!( "Waits on a yielded `{}` cell and returns new output or completion.\n{}", - codex_code_mode::PUBLIC_TOOL_NAME, - codex_code_mode::build_wait_tool_description().trim() + codex_code_mode_protocol::PUBLIC_TOOL_NAME, + codex_code_mode_protocol::build_wait_tool_description().trim() ), strict: false, defer_loading: None, diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index 697eba1d4a..1544c20b7d 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -240,7 +240,7 @@ fn spec_for_model_request( ToolMode::CodeMode | ToolMode::CodeModeOnly ) && exposure != ToolExposure::DirectModelOnly && !is_excluded_from_code_mode(turn_context, tool_name) - && codex_code_mode::is_code_mode_nested_tool(spec.name()) + && codex_code_mode_protocol::is_code_mode_nested_tool(spec.name()) { codex_tools::augment_tool_spec_for_code_mode(spec) } else { @@ -415,9 +415,9 @@ fn is_hidden_by_code_mode_only( ) -> bool { turn_context.tool_mode == ToolMode::CodeModeOnly && exposure != ToolExposure::DirectModelOnly - && codex_code_mode::is_code_mode_nested_tool(&codex_tools::code_mode_name_for_tool_name( - tool_name, - )) + && codex_code_mode_protocol::is_code_mode_nested_tool( + &codex_tools::code_mode_name_for_tool_name(tool_name), + ) } fn is_excluded_from_code_mode(turn_context: &TurnContext, tool_name: &ToolName) -> bool { @@ -540,7 +540,7 @@ fn merge_into_namespaces(specs: Vec) -> Vec { fn code_mode_namespace_descriptions( specs: &[ToolSpec], -) -> BTreeMap { +) -> BTreeMap { let mut namespace_descriptions = BTreeMap::new(); for spec in specs { let ToolSpec::Namespace(namespace) = spec else { @@ -549,7 +549,7 @@ fn code_mode_namespace_descriptions( let entry = namespace_descriptions .entry(namespace.name.clone()) - .or_insert_with(|| codex_code_mode::ToolNamespaceDescription { + .or_insert_with(|| codex_code_mode_protocol::ToolNamespaceDescription { name: namespace.name.clone(), description: namespace.description.clone(), }); @@ -891,8 +891,8 @@ fn append_extension_tool_executors( turn_context.tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly ) { - reserved_tool_names.insert(ToolName::plain(codex_code_mode::PUBLIC_TOOL_NAME)); - reserved_tool_names.insert(ToolName::plain(codex_code_mode::WAIT_TOOL_NAME)); + reserved_tool_names.insert(ToolName::plain(codex_code_mode_protocol::PUBLIC_TOOL_NAME)); + reserved_tool_names.insert(ToolName::plain(codex_code_mode_protocol::WAIT_TOOL_NAME)); } if search_tool_enabled(turn_context) && namespace_tools_enabled(turn_context) @@ -991,9 +991,9 @@ impl CoreToolRuntime for MultiAgentV2NamespaceOverride { } fn compare_code_mode_tools( - left: &codex_code_mode::ToolDefinition, - right: &codex_code_mode::ToolDefinition, - namespace_descriptions: &BTreeMap, + left: &codex_code_mode_protocol::ToolDefinition, + right: &codex_code_mode_protocol::ToolDefinition, + namespace_descriptions: &BTreeMap, ) -> std::cmp::Ordering { let left_namespace = code_mode_namespace_name(left, namespace_descriptions); let right_namespace = code_mode_namespace_name(right, namespace_descriptions); @@ -1005,8 +1005,11 @@ fn compare_code_mode_tools( } fn code_mode_namespace_name<'a>( - tool: &codex_code_mode::ToolDefinition, - namespace_descriptions: &'a BTreeMap, + tool: &codex_code_mode_protocol::ToolDefinition, + namespace_descriptions: &'a BTreeMap< + String, + codex_code_mode_protocol::ToolNamespaceDescription, + >, ) -> Option<&'a str> { tool.tool_name .namespace diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index dce1082bb8..ea57654afe 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -876,8 +876,8 @@ async fn code_mode_only_exposes_code_executor_and_hides_nested_tools() { &["lookup".to_string()] ); plain.assert_visible_lacks(&[ - codex_code_mode::PUBLIC_TOOL_NAME, - codex_code_mode::WAIT_TOOL_NAME, + codex_code_mode_protocol::PUBLIC_TOOL_NAME, + codex_code_mode_protocol::WAIT_TOOL_NAME, ]); let code_mode_only = probe_with( @@ -895,8 +895,8 @@ async fn code_mode_only_exposes_code_executor_and_hides_nested_tools() { ) .await; code_mode_only.assert_visible_contains(&[ - codex_code_mode::PUBLIC_TOOL_NAME, - codex_code_mode::WAIT_TOOL_NAME, + codex_code_mode_protocol::PUBLIC_TOOL_NAME, + codex_code_mode_protocol::WAIT_TOOL_NAME, ]); assert_eq!( code_mode_only.namespace_function_names("codex_app"), @@ -926,7 +926,8 @@ async fn excluded_deferred_namespaces_do_not_enable_nested_tool_guidance() { ) .await; - let ToolSpec::Freeform(exec) = plan.visible_spec(codex_code_mode::PUBLIC_TOOL_NAME) else { + let ToolSpec::Freeform(exec) = plan.visible_spec(codex_code_mode_protocol::PUBLIC_TOOL_NAME) + else { panic!("expected code mode exec tool"); }; assert!( @@ -1071,8 +1072,8 @@ async fn tool_mode_selector_overrides_feature_flags() { }) .await; direct.assert_visible_lacks(&[ - codex_code_mode::PUBLIC_TOOL_NAME, - codex_code_mode::WAIT_TOOL_NAME, + codex_code_mode_protocol::PUBLIC_TOOL_NAME, + codex_code_mode_protocol::WAIT_TOOL_NAME, ]); } @@ -1307,8 +1308,8 @@ async fn hosted_tools_follow_provider_auth_model_and_config_gates() { code_mode_only.visible_names, vec![ // Code-mode entrypoints. - codex_code_mode::PUBLIC_TOOL_NAME, - codex_code_mode::WAIT_TOOL_NAME, + codex_code_mode_protocol::PUBLIC_TOOL_NAME, + codex_code_mode_protocol::WAIT_TOOL_NAME, // Multi-agent v2 tools. "spawn_agent", "send_message", diff --git a/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs b/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs index b09996872e..b025fa4367 100644 --- a/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs +++ b/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs @@ -3,6 +3,17 @@ use std::path::Path; use std::path::PathBuf; use std::sync::Arc; +use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::CodeModeSession; +use codex_code_mode_protocol::CodeModeSessionDelegate; +use codex_code_mode_protocol::CodeModeSessionProvider; +use codex_code_mode_protocol::CodeModeSessionProviderFuture; +use codex_code_mode_protocol::CodeModeSessionResultFuture; +use codex_code_mode_protocol::ExecuteRequest; +use codex_code_mode_protocol::RuntimeResponse; +use codex_code_mode_protocol::StartedCell; +use codex_code_mode_protocol::WaitOutcome; +use codex_code_mode_protocol::WaitRequest; use codex_protocol::protocol::SessionSource; use codex_rollout_trace::ExecutionStatus; use codex_rollout_trace::ThreadStartedTraceMetadata; @@ -58,6 +69,51 @@ impl ToolExecutor for TestHandler { impl CoreToolRuntime for TestHandler {} +struct MissingCellSessionProvider; + +impl CodeModeSessionProvider for MissingCellSessionProvider { + fn create_session<'a>( + &'a self, + _delegate: Arc, + ) -> CodeModeSessionProviderFuture<'a> { + Box::pin(async { + let session: Arc = Arc::new(MissingCellSession); + Ok(session) + }) + } +} + +struct MissingCellSession; + +impl CodeModeSession for MissingCellSession { + fn execute<'a>( + &'a self, + _request: ExecuteRequest, + ) -> CodeModeSessionResultFuture<'a, StartedCell> { + Box::pin(async { Err("execute is not supported by this test session".to_string()) }) + } + + fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome> { + Box::pin(async move { Ok(missing_cell_outcome(request.cell_id)) }) + } + + fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome> { + Box::pin(async move { Ok(missing_cell_outcome(cell_id)) }) + } + + fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()> { + Box::pin(async { Ok(()) }) + } +} + +fn missing_cell_outcome(cell_id: CellId) -> WaitOutcome { + WaitOutcome::MissingCell(RuntimeResponse::Result { + error_text: Some(format!("exec cell {cell_id} not found")), + cell_id, + content_items: Vec::new(), + }) +} + #[tokio::test] async fn dispatch_lifecycle_trace_records_direct_and_code_mode_requesters() -> anyhow::Result<()> { let temp = TempDir::new()?; @@ -213,6 +269,10 @@ async fn dispatch_lifecycle_trace_records_incompatible_payload_failures() -> any async fn missing_code_mode_wait_traces_only_the_wait_tool_call() -> anyhow::Result<()> { let temp = TempDir::new()?; let (mut session, turn) = make_session_and_context().await; + let provider: Arc = Arc::new(MissingCellSessionProvider); + session.services.code_mode_service = + crate::tools::code_mode::CodeModeService::new(Arc::clone(&provider)); + session.services.code_mode_session_provider = provider; attach_test_trace(&mut session, &turn, temp.path())?; let registry = ToolRegistry::with_handler_for_test(Arc::new(CodeModeWaitHandler)); diff --git a/codex-rs/core/tests/common/Cargo.toml b/codex-rs/core/tests/common/Cargo.toml index fb6221cfce..c91903646a 100644 --- a/codex-rs/core/tests/common/Cargo.toml +++ b/codex-rs/core/tests/common/Cargo.toml @@ -17,6 +17,7 @@ assert_cmd = { workspace = true } base64 = { workspace = true } codex-arg0 = { workspace = true } codex-config = { workspace = true } +codex-code-mode = { workspace = true } codex-core = { workspace = true } codex-extension-api = { workspace = true } codex-exec-server = { workspace = true } diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index 62f422ba1c..01c3d90415 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -520,6 +520,13 @@ impl TestCodexBuilder { installation_id, /*attestation_provider*/ None, ); + let thread_manager = if std::env::var_os("CODEX_CODE_MODE_HOST_PATH").is_some() { + thread_manager + } else { + thread_manager.with_code_mode_session_provider(Arc::new( + codex_code_mode::InProcessCodeModeSessionProvider, + )) + }; let thread_manager = Arc::new(thread_manager); let user_shell_override = self.user_shell_override.clone(); diff --git a/codex-rs/core/tests/suite/model_runtime_selectors.rs b/codex-rs/core/tests/suite/model_runtime_selectors.rs index 38c95394ea..f14f998c05 100644 --- a/codex-rs/core/tests/suite/model_runtime_selectors.rs +++ b/codex-rs/core/tests/suite/model_runtime_selectors.rs @@ -158,8 +158,8 @@ async fn remote_tool_mode_selector_overrides_feature_flags() -> Result<()> { assert!( direct_tools .iter() - .all(|name| name != codex_code_mode::PUBLIC_TOOL_NAME - && name != codex_code_mode::WAIT_TOOL_NAME), + .all(|name| name != codex_code_mode_protocol::PUBLIC_TOOL_NAME + && name != codex_code_mode_protocol::WAIT_TOOL_NAME), "direct mode should override enabled code mode flags: {direct_tools:?}" ); @@ -171,8 +171,8 @@ async fn remote_tool_mode_selector_overrides_feature_flags() -> Result<()> { tool_names(&code_mode_only_body), vec![ // Code-mode entrypoints. - codex_code_mode::PUBLIC_TOOL_NAME.to_string(), - codex_code_mode::WAIT_TOOL_NAME.to_string(), + codex_code_mode_protocol::PUBLIC_TOOL_NAME.to_string(), + codex_code_mode_protocol::WAIT_TOOL_NAME.to_string(), // Hosted Responses tools. "web_search".to_string(), "image_generation".to_string(), diff --git a/codex-rs/rollout-trace/Cargo.toml b/codex-rs/rollout-trace/Cargo.toml index b368c9acc5..1335fecc20 100644 --- a/codex-rs/rollout-trace/Cargo.toml +++ b/codex-rs/rollout-trace/Cargo.toml @@ -14,7 +14,7 @@ workspace = true [dependencies] anyhow = { workspace = true } -codex-code-mode = { workspace = true } +codex-code-mode-protocol = { workspace = true } codex-protocol = { workspace = true } http = { workspace = true } serde = { workspace = true, features = ["derive"] } diff --git a/codex-rs/rollout-trace/src/code_cell.rs b/codex-rs/rollout-trace/src/code_cell.rs index 5f2603b70a..e215ad41e6 100644 --- a/codex-rs/rollout-trace/src/code_cell.rs +++ b/codex-rs/rollout-trace/src/code_cell.rs @@ -7,7 +7,7 @@ use std::sync::Arc; -use codex_code_mode::RuntimeResponse; +use codex_code_mode_protocol::RuntimeResponse; use serde::Serialize; use tracing::warn; diff --git a/codex-rs/rollout-trace/src/tool_dispatch.rs b/codex-rs/rollout-trace/src/tool_dispatch.rs index 36c6651b10..28a5c003e0 100644 --- a/codex-rs/rollout-trace/src/tool_dispatch.rs +++ b/codex-rs/rollout-trace/src/tool_dispatch.rs @@ -194,7 +194,7 @@ impl ToolDispatchTraceContext { fn suppresses_tool_dispatch_trace(invocation: &ToolDispatchInvocation) -> bool { matches!(invocation.payload, ToolDispatchPayload::Custom { .. }) && invocation.tool_namespace.is_none() - && invocation.tool_name == codex_code_mode::PUBLIC_TOOL_NAME + && invocation.tool_name == codex_code_mode_protocol::PUBLIC_TOOL_NAME } fn record_started(context: &EnabledToolDispatchTraceContext, invocation: ToolDispatchInvocation) { @@ -394,7 +394,7 @@ mod tests { #[test] fn suppresses_only_noncanonical_dispatch_boundaries() { assert!(suppresses_tool_dispatch_trace(&invocation( - codex_code_mode::PUBLIC_TOOL_NAME, + codex_code_mode_protocol::PUBLIC_TOOL_NAME, /*tool_namespace*/ None, ToolDispatchRequester::Model { model_visible_call_id: "call-exec".to_string(), @@ -414,7 +414,7 @@ mod tests { }, ))); assert!(!suppresses_tool_dispatch_trace(&invocation( - codex_code_mode::PUBLIC_TOOL_NAME, + codex_code_mode_protocol::PUBLIC_TOOL_NAME, Some("mcp__server".to_string()), ToolDispatchRequester::Model { model_visible_call_id: "call-namespaced".to_string(), diff --git a/codex-rs/tools/Cargo.toml b/codex-rs/tools/Cargo.toml index 4b2a3257f3..f75145f3cb 100644 --- a/codex-rs/tools/Cargo.toml +++ b/codex-rs/tools/Cargo.toml @@ -9,7 +9,7 @@ workspace = true [dependencies] codex-app-server-protocol = { workspace = true } -codex-code-mode = { workspace = true } +codex-code-mode-protocol = { workspace = true } codex-features = { workspace = true } codex-protocol = { workspace = true } codex-utils-absolute-path = { workspace = true } diff --git a/codex-rs/tools/src/code_mode.rs b/codex-rs/tools/src/code_mode.rs index 4876486f61..be94ac69eb 100644 --- a/codex-rs/tools/src/code_mode.rs +++ b/codex-rs/tools/src/code_mode.rs @@ -1,8 +1,8 @@ use crate::ResponsesApiNamespaceTool; use crate::ToolName; use crate::ToolSpec; -use codex_code_mode::CodeModeToolKind; -use codex_code_mode::ToolDefinition as CodeModeToolDefinition; +use codex_code_mode_protocol::CodeModeToolKind; +use codex_code_mode_protocol::ToolDefinition as CodeModeToolDefinition; /// Augment tool descriptions with code-mode-specific exec samples. pub fn augment_tool_spec_for_code_mode(spec: ToolSpec) -> ToolSpec { @@ -40,7 +40,8 @@ pub fn augment_tool_spec_for_code_mode(spec: ToolSpec) -> ToolSpec { output_schema: tool.output_schema.clone(), }; tool.description = - codex_code_mode::augment_tool_definition(definition).description; + codex_code_mode_protocol::augment_tool_definition(definition) + .description; } } } @@ -54,8 +55,8 @@ pub fn augment_tool_spec_for_code_mode(spec: ToolSpec) -> ToolSpec { /// including the code-mode-specific description sample. pub fn tool_spec_to_code_mode_tool_definition(spec: &ToolSpec) -> Option { let definition = code_mode_tool_definition_for_spec(spec)?; - codex_code_mode::is_code_mode_nested_tool(&definition.name) - .then(|| codex_code_mode::augment_tool_definition(definition)) + codex_code_mode_protocol::is_code_mode_nested_tool(&definition.name) + .then(|| codex_code_mode_protocol::augment_tool_definition(definition)) } pub fn collect_code_mode_tool_definitions<'a>( @@ -64,8 +65,8 @@ pub fn collect_code_mode_tool_definitions<'a>( let mut tool_definitions = specs .into_iter() .flat_map(code_mode_tool_definitions_for_spec) - .filter(|definition| codex_code_mode::is_code_mode_nested_tool(&definition.name)) - .map(codex_code_mode::augment_tool_definition) + .filter(|definition| codex_code_mode_protocol::is_code_mode_nested_tool(&definition.name)) + .map(codex_code_mode_protocol::augment_tool_definition) .collect::>(); tool_definitions.sort_by(|left, right| left.name.cmp(&right.name)); tool_definitions.dedup_by(|left, right| left.name == right.name); @@ -78,7 +79,7 @@ pub fn collect_code_mode_exec_prompt_tool_definitions<'a>( let mut tool_definitions = specs .into_iter() .flat_map(code_mode_tool_definitions_for_spec) - .filter(|definition| codex_code_mode::is_code_mode_nested_tool(&definition.name)) + .filter(|definition| codex_code_mode_protocol::is_code_mode_nested_tool(&definition.name)) .collect::>(); tool_definitions.sort_by(|left, right| left.name.cmp(&right.name)); tool_definitions.dedup_by(|left, right| left.name == right.name); @@ -87,7 +88,7 @@ pub fn collect_code_mode_exec_prompt_tool_definitions<'a>( fn augmented_description_for_spec(spec: &ToolSpec) -> Option { code_mode_tool_definition_for_spec(spec) - .map(codex_code_mode::augment_tool_definition) + .map(codex_code_mode_protocol::augment_tool_definition) .map(|definition| definition.description) } diff --git a/codex-rs/tools/src/code_mode_tests.rs b/codex-rs/tools/src/code_mode_tests.rs index c4c4c7ce26..c5c7b7b3f1 100644 --- a/codex-rs/tools/src/code_mode_tests.rs +++ b/codex-rs/tools/src/code_mode_tests.rs @@ -69,7 +69,7 @@ declare const tools: { lookup_order(args: { order_id: string; }): Promise<{ ok: fn augment_tool_spec_for_code_mode_preserves_exec_tool_description() { assert_eq!( augment_tool_spec_for_code_mode(ToolSpec::Freeform(FreeformTool { - name: codex_code_mode::PUBLIC_TOOL_NAME.to_string(), + name: codex_code_mode_protocol::PUBLIC_TOOL_NAME.to_string(), description: "Run code".to_string(), format: FreeformToolFormat { r#type: "grammar".to_string(), @@ -78,7 +78,7 @@ fn augment_tool_spec_for_code_mode_preserves_exec_tool_description() { }, })), ToolSpec::Freeform(FreeformTool { - name: codex_code_mode::PUBLIC_TOOL_NAME.to_string(), + name: codex_code_mode_protocol::PUBLIC_TOOL_NAME.to_string(), description: "Run code".to_string(), format: FreeformToolFormat { r#type: "grammar".to_string(), @@ -103,7 +103,7 @@ fn tool_spec_to_code_mode_tool_definition_returns_augmented_nested_tools() { assert_eq!( tool_spec_to_code_mode_tool_definition(&spec), - Some(codex_code_mode::ToolDefinition { + Some(codex_code_mode_protocol::ToolDefinition { name: "apply_patch".to_string(), tool_name: ToolName::plain("apply_patch"), description: r#"Apply a patch @@ -113,7 +113,7 @@ exec tool declaration: declare const tools: { apply_patch(input: string): Promise; }; ```"# .to_string(), - kind: codex_code_mode::CodeModeToolKind::Freeform, + kind: codex_code_mode_protocol::CodeModeToolKind::Freeform, input_schema: None, output_schema: None, }) diff --git a/scripts/codex_package/README.md b/scripts/codex_package/README.md index 323a3ce5ba..3fe93441d4 100644 --- a/scripts/codex_package/README.md +++ b/scripts/codex_package/README.md @@ -10,7 +10,8 @@ The builder creates a canonical Codex package directory: . ├── codex-package.json ├── bin -│ └── [.exe] +│ ├── [.exe] +│ └── codex-code-mode-host[.exe] ├── codex-resources │ ├── bwrap # Linux only │ ├── zsh/bin/zsh # supported Unix targets only @@ -40,6 +41,7 @@ grouped `cargo build` command per package when they are needed and no prebuilt override was provided: - all targets: the selected entrypoint, unless `--entrypoint-bin` is provided +- all targets: `codex-code-mode-host`, unless `--code-mode-host-bin` is provided - Linux targets: `bwrap`, unless `--bwrap-bin` is provided - Windows targets: `codex-command-runner` and `codex-windows-sandbox-setup`, unless the corresponding prebuilt helper flags are provided diff --git a/scripts/codex_package/cargo.py b/scripts/codex_package/cargo.py index 208d85d174..f2dbe6c774 100644 --- a/scripts/codex_package/cargo.py +++ b/scripts/codex_package/cargo.py @@ -17,6 +17,7 @@ CODEX_RS_ROOT = REPO_ROOT / "codex-rs" @dataclass(frozen=True) class SourceBuildOutputs: entrypoint_bin: Path + code_mode_host_bin: Path bwrap_bin: Path | None codex_command_runner_bin: Path | None codex_windows_sandbox_setup_bin: Path | None @@ -29,6 +30,7 @@ def build_source_binaries( cargo: str, profile: str, entrypoint_bin: Path | None, + code_mode_host_bin: Path | None, bwrap_bin: Path | None, codex_command_runner_bin: Path | None, codex_windows_sandbox_setup_bin: Path | None, @@ -43,6 +45,7 @@ def build_source_binaries( spec, variant, build_entrypoint=entrypoint_bin is None, + build_code_mode_host=code_mode_host_bin is None, build_bwrap=spec.is_linux and bwrap_bin is None, build_codex_command_runner=spec.is_windows and codex_command_runner_bin is None, build_codex_windows_sandbox_setup=spec.is_windows @@ -61,7 +64,7 @@ def build_source_binaries( cmd.extend(["--bin", binary]) cargo_env = None - if entrypoint_bin is None: + if entrypoint_bin is None or code_mode_host_bin is None: codex_v8_env = resolve_codex_v8_cargo_env(spec) if codex_v8_env: cargo_env = {**os.environ, **codex_v8_env} @@ -80,6 +83,10 @@ def build_source_binaries( entrypoint_bin, output_dir / variant.entrypoint_name(spec), ), + code_mode_host_bin=resolve_output_path( + code_mode_host_bin, + output_dir / f"codex-code-mode-host{spec.exe_suffix}", + ), bwrap_bin=resolve_output_path( bwrap_bin, output_dir / "bwrap" if spec.is_linux else None, @@ -102,6 +109,7 @@ def source_binaries_for_target( variant: PackageVariant, *, build_entrypoint: bool, + build_code_mode_host: bool, build_bwrap: bool, build_codex_command_runner: bool, build_codex_windows_sandbox_setup: bool, @@ -109,6 +117,8 @@ def source_binaries_for_target( binaries = [] if build_entrypoint: binaries.append(variant.cargo_bin) + if build_code_mode_host: + binaries.append("codex-code-mode-host") if build_bwrap: binaries.append("bwrap") if build_codex_command_runner: @@ -174,6 +184,7 @@ def cargo_profile_dirname(profile: str) -> str: def validate_source_outputs(outputs: SourceBuildOutputs) -> None: for path in [ outputs.entrypoint_bin, + outputs.code_mode_host_bin, outputs.bwrap_bin, outputs.codex_command_runner_bin, outputs.codex_windows_sandbox_setup_bin, diff --git a/scripts/codex_package/cli.py b/scripts/codex_package/cli.py index b7d919e4ff..d3c6fafe33 100644 --- a/scripts/codex_package/cli.py +++ b/scripts/codex_package/cli.py @@ -82,6 +82,14 @@ def parse_args() -> argparse.Namespace: "variant. If omitted, the entrypoint is built with Cargo." ), ) + parser.add_argument( + "--code-mode-host-bin", + type=Path, + help=( + "Optional prebuilt codex-code-mode-host executable. If omitted, " + "codex-code-mode-host is built with Cargo." + ), + ) parser.add_argument( "--bwrap-bin", type=Path, @@ -140,6 +148,11 @@ def main() -> int: "prebuilt entrypoint executable", "--entrypoint-bin", ), + code_mode_host_bin=resolve_optional_input_path( + args.code_mode_host_bin, + "prebuilt code-mode host executable", + "--code-mode-host-bin", + ), bwrap_bin=resolve_optional_input_path( args.bwrap_bin, "prebuilt Linux bwrap executable", @@ -159,6 +172,7 @@ def main() -> int: version = read_workspace_version() inputs = PackageInputs( entrypoint_bin=source_outputs.entrypoint_bin, + code_mode_host_bin=source_outputs.code_mode_host_bin, rg_bin=resolve_rg_bin(spec, args.rg_bin), zsh_bin=resolve_zsh_bin(spec), bwrap_bin=source_outputs.bwrap_bin, diff --git a/scripts/codex_package/layout.py b/scripts/codex_package/layout.py index 63598672ea..131de5f5bc 100644 --- a/scripts/codex_package/layout.py +++ b/scripts/codex_package/layout.py @@ -51,6 +51,11 @@ def build_package_dir( bin_dir / entrypoint_name, is_windows=spec.is_windows, ) + copy_executable( + inputs.code_mode_host_bin, + bin_dir / f"codex-code-mode-host{spec.exe_suffix}", + is_windows=spec.is_windows, + ) copy_executable(inputs.rg_bin, path_dir / spec.rg_name, is_windows=spec.is_windows) if inputs.zsh_bin is not None: @@ -83,6 +88,7 @@ def build_package_dir( "target": spec.target, "variant": variant.name, "entrypoint": f"bin/{entrypoint_name}", + "codeModeHost": f"bin/codex-code-mode-host{spec.exe_suffix}", "resourcesDir": "codex-resources", "pathDir": "codex-path", } @@ -118,6 +124,7 @@ def validate_package_dir( "target": spec.target, "variant": variant.name, "entrypoint": f"bin/{variant.entrypoint_name(spec)}", + "codeModeHost": f"bin/codex-code-mode-host{spec.exe_suffix}", "resourcesDir": "codex-resources", "pathDir": "codex-path", } @@ -130,6 +137,7 @@ def validate_package_dir( required_files = [ Path("bin") / variant.entrypoint_name(spec), + Path("bin") / f"codex-code-mode-host{spec.exe_suffix}", Path("codex-path") / spec.rg_name, ] executable_files = list(required_files) diff --git a/scripts/codex_package/targets.py b/scripts/codex_package/targets.py index 8307a3e630..016202cf5f 100644 --- a/scripts/codex_package/targets.py +++ b/scripts/codex_package/targets.py @@ -39,6 +39,7 @@ class PackageVariant: @dataclass(frozen=True) class PackageInputs: entrypoint_bin: Path + code_mode_host_bin: Path rg_bin: Path zsh_bin: Path | None bwrap_bin: Path | None diff --git a/scripts/codex_package/test_cargo.py b/scripts/codex_package/test_cargo.py index 6185f8b8b3..6f2df7a3ff 100644 --- a/scripts/codex_package/test_cargo.py +++ b/scripts/codex_package/test_cargo.py @@ -20,6 +20,7 @@ class SourceBinariesForTargetTest(unittest.TestCase): TARGET_SPECS["aarch64-apple-darwin"], PACKAGE_VARIANTS["codex"], build_entrypoint=False, + build_code_mode_host=False, build_bwrap=False, build_codex_command_runner=False, build_codex_windows_sandbox_setup=False, @@ -35,6 +36,7 @@ class SourceBinariesForTargetTest(unittest.TestCase): TARGET_SPECS["x86_64-unknown-linux-musl"], PACKAGE_VARIANTS["codex"], build_entrypoint=False, + build_code_mode_host=False, build_bwrap=False, build_codex_command_runner=False, build_codex_windows_sandbox_setup=False, @@ -50,6 +52,7 @@ class SourceBinariesForTargetTest(unittest.TestCase): TARGET_SPECS["x86_64-pc-windows-msvc"], PACKAGE_VARIANTS["codex"], build_entrypoint=False, + build_code_mode_host=False, build_bwrap=False, build_codex_command_runner=False, build_codex_windows_sandbox_setup=False, @@ -63,6 +66,7 @@ class SourceBinariesForTargetTest(unittest.TestCase): TARGET_SPECS["x86_64-pc-windows-msvc"], PACKAGE_VARIANTS["codex"], build_entrypoint=False, + build_code_mode_host=False, build_bwrap=False, build_codex_command_runner=True, build_codex_windows_sandbox_setup=True, @@ -70,10 +74,25 @@ class SourceBinariesForTargetTest(unittest.TestCase): ["codex-command-runner", "codex-windows-sandbox-setup"], ) + def test_missing_code_mode_host_is_built(self) -> None: + self.assertEqual( + source_binaries_for_target( + TARGET_SPECS["x86_64-unknown-linux-musl"], + PACKAGE_VARIANTS["codex"], + build_entrypoint=False, + build_code_mode_host=True, + build_bwrap=False, + build_codex_command_runner=False, + build_codex_windows_sandbox_setup=False, + ), + ["codex-code-mode-host"], + ) + def test_build_uses_prebuilt_windows_helpers_without_running_cargo(self) -> None: with tempfile.TemporaryDirectory() as temp_dir: root = Path(temp_dir) entrypoint = touch_file(root / "codex.exe") + code_mode_host = touch_file(root / "codex-code-mode-host.exe") command_runner = touch_file(root / "codex-command-runner.exe") sandbox_setup = touch_file(root / "codex-windows-sandbox-setup.exe") @@ -83,12 +102,14 @@ class SourceBinariesForTargetTest(unittest.TestCase): cargo=str(root / "cargo-that-should-not-run"), profile="release", entrypoint_bin=entrypoint, + code_mode_host_bin=code_mode_host, bwrap_bin=None, codex_command_runner_bin=command_runner, codex_windows_sandbox_setup_bin=sandbox_setup, ) self.assertEqual(outputs.entrypoint_bin, entrypoint) + self.assertEqual(outputs.code_mode_host_bin, code_mode_host) self.assertEqual(outputs.codex_command_runner_bin, command_runner) self.assertEqual(outputs.codex_windows_sandbox_setup_bin, sandbox_setup) diff --git a/scripts/codex_package/test_layout.py b/scripts/codex_package/test_layout.py new file mode 100644 index 0000000000..bd3e7bfa9e --- /dev/null +++ b/scripts/codex_package/test_layout.py @@ -0,0 +1,79 @@ +#!/usr/bin/env python3 + +import json +from pathlib import Path +import sys +import tempfile +import unittest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from codex_package.layout import build_package_dir +from codex_package.layout import validate_package_dir +from codex_package.targets import PACKAGE_VARIANTS +from codex_package.targets import PackageInputs +from codex_package.targets import TARGET_SPECS + + +class PackageLayoutTest(unittest.TestCase): + def test_code_mode_host_is_packaged_and_required(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + package_dir = root / "package" + package_dir.mkdir() + entrypoint = touch_file(root / "codex") + code_mode_host = touch_file(root / "codex-code-mode-host") + rg = touch_file(root / "rg") + spec = TARGET_SPECS["aarch64-apple-darwin"] + variant = PACKAGE_VARIANTS["codex"] + + build_package_dir( + package_dir, + "0.0.0-test", + variant, + spec, + PackageInputs( + entrypoint_bin=entrypoint, + code_mode_host_bin=code_mode_host, + rg_bin=rg, + zsh_bin=None, + bwrap_bin=None, + codex_command_runner_bin=None, + codex_windows_sandbox_setup_bin=None, + ), + ) + + packaged_host = package_dir / "bin" / "codex-code-mode-host" + self.assertTrue(packaged_host.is_file()) + self.assertTrue(packaged_host.stat().st_mode & 0o100) + metadata = json.loads( + (package_dir / "codex-package.json").read_text(encoding="utf-8") + ) + self.assertEqual(metadata["codeModeHost"], "bin/codex-code-mode-host") + validate_package_dir( + package_dir, + variant, + spec, + include_zsh=False, + ) + + packaged_host.unlink() + with self.assertRaisesRegex( + RuntimeError, + "Missing package file: bin/codex-code-mode-host", + ): + validate_package_dir( + package_dir, + variant, + spec, + include_zsh=False, + ) + + +def touch_file(path: Path) -> Path: + path.touch() + return path + + +if __name__ == "__main__": + unittest.main() diff --git a/scripts/codex_package/test_v8.py b/scripts/codex_package/test_v8.py new file mode 100644 index 0000000000..1ae9069eb7 --- /dev/null +++ b/scripts/codex_package/test_v8.py @@ -0,0 +1,57 @@ +#!/usr/bin/env python3 + +import hashlib +from pathlib import Path +import sys +import tempfile +import unittest +from unittest.mock import patch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from codex_package.targets import TARGET_SPECS +from codex_package.v8 import fetch_codex_v8_artifacts + + +class RustyV8ArtifactsTest(unittest.TestCase): + def test_fetches_pointer_compression_sandbox_artifacts(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + payload = b"artifact" + digest = hashlib.sha256(payload).hexdigest() + + def download_file(url: str, dest: Path) -> None: + if dest.suffix == ".sha256": + target = "x86_64-unknown-linux-musl" + dest.parent.mkdir(parents=True, exist_ok=True) + dest.write_text( + "\n".join( + [ + f"{digest} librusty_v8_ptrcomp_sandbox_release_{target}.a.gz", + f"{digest} src_binding_ptrcomp_sandbox_release_{target}.rs", + ] + ) + + "\n", + encoding="utf-8", + ) + else: + dest.write_bytes(payload) + + with patch("codex_package.v8.download_file", side_effect=download_file): + artifacts = fetch_codex_v8_artifacts( + TARGET_SPECS["x86_64-unknown-linux-musl"], + version="147.4.0", + cache_root=Path(temp_dir), + ) + + self.assertEqual( + artifacts.archive.name, + "librusty_v8_ptrcomp_sandbox_release_x86_64-unknown-linux-musl.a.gz", + ) + self.assertEqual( + artifacts.binding.name, + "src_binding_ptrcomp_sandbox_release_x86_64-unknown-linux-musl.rs", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/scripts/codex_package/v8.py b/scripts/codex_package/v8.py index 4033e6822f..16c9a740a9 100644 --- a/scripts/codex_package/v8.py +++ b/scripts/codex_package/v8.py @@ -16,6 +16,7 @@ from .targets import TargetSpec DOWNLOAD_TIMEOUT_SECS = 120 +ARTIFACT_PROFILE = "ptrcomp_sandbox_release" @dataclass(frozen=True) @@ -70,9 +71,9 @@ def fetch_codex_v8_artifacts( ) target = spec.target cache_dir = (cache_root or default_cache_root()) / f"rusty-v8-{version}-{target}" - archive = cache_dir / f"librusty_v8_release_{target}.a.gz" - binding = cache_dir / f"src_binding_release_{target}.rs" - checksums = cache_dir / f"rusty_v8_release_{target}.sha256" + archive = cache_dir / f"librusty_v8_{ARTIFACT_PROFILE}_{target}.a.gz" + binding = cache_dir / f"src_binding_{ARTIFACT_PROFILE}_{target}.rs" + checksums = cache_dir / f"rusty_v8_{ARTIFACT_PROFILE}_{target}.sha256" download_file(f"{release_url}/{checksums.name}", checksums) expected_checksums = load_checksums(checksums, {archive.name, binding.name}) diff --git a/scripts/install/install.ps1 b/scripts/install/install.ps1 index 6973d482e2..31aef584f3 100644 --- a/scripts/install/install.ps1 +++ b/scripts/install/install.ps1 @@ -537,6 +537,7 @@ function Test-PackageContentsAreComplete { $expectedFiles = @( "codex-package.json", "bin\codex.exe", + "bin\codex-code-mode-host.exe", "codex-path\rg.exe", "codex-resources\codex-command-runner.exe", "codex-resources\codex-windows-sandbox-setup.exe" diff --git a/scripts/install/install.sh b/scripts/install/install.sh index 23980b0c5e..f5304090b6 100755 --- a/scripts/install/install.sh +++ b/scripts/install/install.sh @@ -677,7 +677,10 @@ install_package_release() { rm -rf "$stage_release" mkdir -p "$stage_release" tar -xzf "$archive_path" -C "$stage_release" - chmod 0755 "$stage_release/bin/codex" "$stage_release/codex-path/rg" + chmod 0755 \ + "$stage_release/bin/codex" \ + "$stage_release/bin/codex-code-mode-host" \ + "$stage_release/codex-path/rg" if [ -f "$stage_release/codex-resources/bwrap" ]; then chmod 0755 "$stage_release/codex-resources/bwrap" fi @@ -730,6 +733,7 @@ release_dir_is_complete() { package) [ -f "$release_dir/codex-package.json" ] && [ -x "$release_dir/bin/codex" ] && + [ -x "$release_dir/bin/codex-code-mode-host" ] && [ -x "$release_dir/codex" ] && [ -x "$release_dir/codex-path/rg" ] || return 1