From 8073dbb20bbd57f2acdc18708a95e8fbbfc8b91f Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Fri, 7 Aug 2026 23:22:47 +0000 Subject: [PATCH] Define the code-mode host gRPC protocol (#37510) ## What changed - Add the `codex.code_mode.v1` protobuf API for managing code-mode sessions, executions, waits, tool callbacks, notifications, and content results. - Generate and export the Rust client/server bindings with `tonic` under Cargo. - Add Bazel protobuf targets and a `prost` toolchain that uses the workspace's Rust runtime versions. GitOrigin-RevId: 51d9d044e01dab2f4d25ab79aa7a52585485eee3 --- MODULE.bazel | 8 + MODULE.bazel.lock | 9 + bazel/toolchains/prost/BUILD.bazel | 35 +++ codex-rs/Cargo.lock | 70 +++++ codex-rs/code-mode-protocol/BUILD.bazel | 18 ++ codex-rs/code-mode-protocol/Cargo.toml | 12 + codex-rs/code-mode-protocol/build.rs | 17 ++ .../src/grpc/codex.code_mode.v1.proto | 257 ++++++++++++++++++ codex-rs/code-mode-protocol/src/grpc/mod.rs | 5 + codex-rs/code-mode-protocol/src/lib.rs | 1 + 10 files changed, 432 insertions(+) create mode 100644 bazel/toolchains/prost/BUILD.bazel create mode 100644 codex-rs/code-mode-protocol/build.rs create mode 100644 codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto create mode 100644 codex-rs/code-mode-protocol/src/grpc/mod.rs diff --git a/MODULE.bazel b/MODULE.bazel index 5fa0329a7f..0e0d1d3ed0 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -2,6 +2,10 @@ module(name = "codex") bazel_dep(name = "bazel_skylib", version = "1.9.0") bazel_dep(name = "platforms", version = "1.0.0") + +# Expose the public protobuf rules already used transitively by rules_rs. +bazel_dep(name = "protobuf", version = "34.0.bcr.1", repo_name = "com_google_protobuf") +bazel_dep(name = "rules_proto", version = "7.1.0") bazel_dep(name = "llvm", version = "0.8.11") # Patch hermetic LLVM for Codex's custom libc++ and Windows gnullvm runtime @@ -206,6 +210,10 @@ use_repo(rules_rust_bindgen, "rules_rust_bindgen") register_toolchains("@rules_rust_bindgen//:all") +# Generate Rust protobuf/gRPC libraries against the workspace's prost/tonic +# versions instead of the versions bundled with rules_rs's default toolchain. +register_toolchains("//bazel/toolchains/prost:toolchain") + crate = use_extension("@rules_rs//rs:extensions.bzl", "crate") crate.from_cargo( cargo_lock = "//codex-rs:Cargo.lock", diff --git a/MODULE.bazel.lock b/MODULE.bazel.lock index 6f8a694fb5..017d0e03c7 100644 --- a/MODULE.bazel.lock +++ b/MODULE.bazel.lock @@ -1375,6 +1375,15 @@ "prost-derive_0.14.3": "{\"dependencies\":[{\"name\":\"anyhow\",\"req\":\"^1.0.1\"},{\"name\":\"itertools\",\"req\":\">=0.10.1, <=0.14\"},{\"name\":\"proc-macro2\",\"req\":\"^1.0.60\"},{\"name\":\"quote\",\"req\":\"^1\"},{\"features\":[\"extra-traits\"],\"name\":\"syn\",\"req\":\"^2\"}],\"features\":{}}", "prost-types_0.14.3": "{\"dependencies\":[{\"features\":[\"derive\"],\"name\":\"arbitrary\",\"optional\":true,\"req\":\"^1.4\"},{\"default_features\":false,\"name\":\"chrono\",\"optional\":true,\"req\":\"^0.4.34\"},{\"kind\":\"dev\",\"name\":\"proptest\",\"req\":\"^1\"},{\"default_features\":false,\"features\":[\"derive\"],\"name\":\"prost\",\"req\":\"^0.14.3\"}],\"features\":{\"arbitrary\":[\"dep:arbitrary\"],\"default\":[\"std\"],\"std\":[\"prost/std\"]}}", "prost_0.14.3": "{\"dependencies\":[{\"default_features\":false,\"name\":\"bytes\",\"req\":\"^1\"},{\"default_features\":false,\"kind\":\"dev\",\"name\":\"criterion\",\"req\":\"^0.7\"},{\"kind\":\"dev\",\"name\":\"proptest\",\"req\":\"^1\"},{\"name\":\"prost-derive\",\"optional\":true,\"req\":\"^0.14.3\"},{\"kind\":\"dev\",\"name\":\"rand\",\"req\":\"^0.9\"}],\"features\":{\"default\":[\"derive\",\"std\"],\"derive\":[\"dep:prost-derive\"],\"no-recursion-limit\":[],\"std\":[]}}", + "protoc-bin-vendored-linux-aarch_64_3.2.0": "{\"dependencies\":[],\"features\":{}}", + "protoc-bin-vendored-linux-ppcle_64_3.2.0": "{\"dependencies\":[],\"features\":{}}", + "protoc-bin-vendored-linux-s390_64_3.2.0": "{\"dependencies\":[],\"features\":{}}", + "protoc-bin-vendored-linux-x86_32_3.2.0": "{\"dependencies\":[],\"features\":{}}", + "protoc-bin-vendored-linux-x86_64_3.2.0": "{\"dependencies\":[],\"features\":{}}", + "protoc-bin-vendored-macos-aarch_64_3.2.0": "{\"dependencies\":[],\"features\":{}}", + "protoc-bin-vendored-macos-x86_64_3.2.0": "{\"dependencies\":[],\"features\":{}}", + "protoc-bin-vendored-win32_3.2.0": "{\"dependencies\":[],\"features\":{}}", + "protoc-bin-vendored_3.2.0": "{\"dependencies\":[{\"name\":\"protoc-bin-vendored-linux-aarch_64\",\"req\":\"=3.2.0\"},{\"name\":\"protoc-bin-vendored-linux-ppcle_64\",\"req\":\"=3.2.0\"},{\"name\":\"protoc-bin-vendored-linux-s390_64\",\"req\":\"=3.2.0\"},{\"name\":\"protoc-bin-vendored-linux-x86_32\",\"req\":\"=3.2.0\"},{\"name\":\"protoc-bin-vendored-linux-x86_64\",\"req\":\"=3.2.0\"},{\"name\":\"protoc-bin-vendored-macos-aarch_64\",\"req\":\"=3.2.0\"},{\"name\":\"protoc-bin-vendored-macos-x86_64\",\"req\":\"=3.2.0\"},{\"name\":\"protoc-bin-vendored-win32\",\"req\":\"=3.2.0\"}],\"features\":{}}", "protoc-gen-prost_0.5.0": "{\"dependencies\":[{\"name\":\"once_cell\",\"req\":\"^1.21.3\"},{\"default_features\":false,\"features\":[\"std\"],\"name\":\"prost\",\"req\":\"^0.14.1\"},{\"default_features\":false,\"name\":\"prost-build\",\"req\":\"^0.14.1\"},{\"default_features\":false,\"name\":\"prost-types\",\"req\":\"^0.14.1\"},{\"default_features\":false,\"name\":\"regex\",\"req\":\"^1.11.1\"}],\"features\":{}}", "protoc-gen-tonic_0.5.0": "{\"dependencies\":[{\"name\":\"heck\",\"req\":\"^0.5.0\"},{\"name\":\"prettyplease\",\"req\":\"^0.2.37\"},{\"name\":\"proc-macro2\",\"req\":\"^1.0.103\"},{\"default_features\":false,\"features\":[\"std\"],\"name\":\"prost\",\"req\":\"^0.14.1\"},{\"default_features\":false,\"name\":\"prost-build\",\"req\":\"^0.14.1\"},{\"default_features\":false,\"name\":\"prost-types\",\"req\":\"^0.14.1\"},{\"name\":\"protoc-gen-prost\",\"req\":\"^0.5.0\"},{\"name\":\"quote\",\"req\":\"^1.0.42\"},{\"default_features\":false,\"name\":\"regex\",\"req\":\"^1.11.1\"},{\"features\":[\"parsing\",\"full\"],\"name\":\"syn\",\"req\":\"^2.0.109\"},{\"name\":\"tonic-build\",\"req\":\"^0.14.1\"}],\"features\":{}}", "psl-types_2.0.11": "{\"dependencies\":[],\"features\":{}}", diff --git a/bazel/toolchains/prost/BUILD.bazel b/bazel/toolchains/prost/BUILD.bazel new file mode 100644 index 0000000000..2b87000bb2 --- /dev/null +++ b/bazel/toolchains/prost/BUILD.bazel @@ -0,0 +1,35 @@ +load("@rules_rust//extensions/prost:defs.bzl", "rust_prost_toolchain") +load("@rules_rust//rust:defs.bzl", "rust_library_group") + +package(default_visibility = ["//visibility:public"]) + +rust_library_group( + name = "prost_runtime", + deps = ["@crates//:prost"], +) + +rust_library_group( + name = "tonic_runtime", + deps = [ + ":prost_runtime", + "@crates//:tonic", + "@crates//:tonic-prost", + ], +) + +rust_prost_toolchain( + name = "toolchain_impl", + compile_well_known_types = False, + prost_plugin = "@rules_rs//rs/private/prost:protoc-gen-prost__protoc-gen-prost", + prost_runtime = ":prost_runtime", + prost_types = "@crates//:prost-types-0.14.3", + tonic_plugin = "@rules_rs//rs/private/prost:protoc-gen-tonic__protoc-gen-tonic", + tonic_plugin_flag = "--plugin=protoc-gen-tonic=%s", + tonic_runtime = ":tonic_runtime", +) + +toolchain( + name = "toolchain", + toolchain = ":toolchain_impl", + toolchain_type = "@rules_rust//extensions/prost:toolchain_type", +) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 88f29ad0f2..024b7255a2 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2584,11 +2584,17 @@ name = "codex-code-mode-protocol" version = "0.0.0" dependencies = [ "codex-protocol", + "glob", "pretty_assertions", + "prost", + "protoc-bin-vendored", "serde", "serde_json", "tokio", "tokio-util", + "tonic", + "tonic-prost", + "tonic-prost-build", ] [[package]] @@ -10687,6 +10693,70 @@ dependencies = [ "prost", ] +[[package]] +name = "protoc-bin-vendored" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1c381df33c98266b5f08186583660090a4ffa0889e76c7e9a5e175f645a67fa" +dependencies = [ + "protoc-bin-vendored-linux-aarch_64", + "protoc-bin-vendored-linux-ppcle_64", + "protoc-bin-vendored-linux-s390_64", + "protoc-bin-vendored-linux-x86_32", + "protoc-bin-vendored-linux-x86_64", + "protoc-bin-vendored-macos-aarch_64", + "protoc-bin-vendored-macos-x86_64", + "protoc-bin-vendored-win32", +] + +[[package]] +name = "protoc-bin-vendored-linux-aarch_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c350df4d49b5b9e3ca79f7e646fde2377b199e13cfa87320308397e1f37e1a4c" + +[[package]] +name = "protoc-bin-vendored-linux-ppcle_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a55a63e6c7244f19b5c6393f025017eb5d793fd5467823a099740a7a4222440c" + +[[package]] +name = "protoc-bin-vendored-linux-s390_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1dba5565db4288e935d5330a07c264a4ee8e4a5b4a4e6f4e83fad824cc32f3b0" + +[[package]] +name = "protoc-bin-vendored-linux-x86_32" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8854774b24ee28b7868cd71dccaae8e02a2365e67a4a87a6cd11ee6cdbdf9cf5" + +[[package]] +name = "protoc-bin-vendored-linux-x86_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b38b07546580df720fa464ce124c4b03630a6fb83e05c336fea2a241df7e5d78" + +[[package]] +name = "protoc-bin-vendored-macos-aarch_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "89278a9926ce312e51f1d999fee8825d324d603213344a9a706daa009f1d8092" + +[[package]] +name = "protoc-bin-vendored-macos-x86_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "81745feda7ccfb9471d7a4de888f0652e806d5795b61480605d4943176299756" + +[[package]] +name = "protoc-bin-vendored-win32" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "95067976aca6421a523e491fce939a3e65249bac4b977adee0ee9771568e8aa3" + [[package]] name = "psl" version = "2.1.184" diff --git a/codex-rs/code-mode-protocol/BUILD.bazel b/codex-rs/code-mode-protocol/BUILD.bazel index ebf43b9bba..deff3c2e3a 100644 --- a/codex-rs/code-mode-protocol/BUILD.bazel +++ b/codex-rs/code-mode-protocol/BUILD.bazel @@ -1,6 +1,24 @@ +load("@rules_proto//proto:defs.bzl", "proto_library") +load("@rules_rust//extensions/prost:defs.bzl", "rust_prost_library") load("//:defs.bzl", "codex_rust_crate") +proto_library( + name = "code-mode-proto", + srcs = glob(["src/grpc/*.proto"]), + strip_import_prefix = "src/grpc", + visibility = ["//visibility:public"], +) + +rust_prost_library( + name = "code-mode-rust-proto", + proto = ":code-mode-proto", + visibility = ["//visibility:public"], +) + codex_rust_crate( name = "code-mode-protocol", + build_script_enabled = False, crate_name = "codex_code_mode_protocol", + deps_extra = [":code-mode-rust-proto"], + rustc_flags_extra = ["--cfg=codex_bazel"], ) diff --git a/codex-rs/code-mode-protocol/Cargo.toml b/codex-rs/code-mode-protocol/Cargo.toml index f6280fd091..cad10b47eb 100644 --- a/codex-rs/code-mode-protocol/Cargo.toml +++ b/codex-rs/code-mode-protocol/Cargo.toml @@ -14,10 +14,22 @@ workspace = true [dependencies] codex-protocol = { workspace = true } +prost = "0.14.3" serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } tokio = { workspace = true, features = ["io-util", "sync"] } tokio-util = { workspace = true, features = ["rt"] } +tonic = { workspace = true } +tonic-prost = { workspace = true } + +[build-dependencies] +glob = { workspace = true } +protoc-bin-vendored = "3.2.0" +tonic-prost-build = { version = "=0.14.3", default-features = false, features = ["transport"] } + +# Cargo-shear cannot inspect the prost and tonic-prost references in generated gRPC bindings. +[package.metadata.cargo-shear] +ignored = ["prost", "tonic-prost"] [dev-dependencies] pretty_assertions = { workspace = true } diff --git a/codex-rs/code-mode-protocol/build.rs b/codex-rs/code-mode-protocol/build.rs new file mode 100644 index 0000000000..d8d8bbff32 --- /dev/null +++ b/codex-rs/code-mode-protocol/build.rs @@ -0,0 +1,17 @@ +use std::path::PathBuf; + +fn main() -> Result<(), Box> { + println!("cargo:rustc-check-cfg=cfg(codex_bazel)"); + println!("cargo:rerun-if-changed=src/grpc"); + + let mut config = tonic_prost_build::Config::new(); + config.protoc_executable(protoc_bin_vendored::protoc_bin_path()?); + let proto_files = glob::glob("src/grpc/*.proto")?.collect::, _>>()?; + + tonic_prost_build::configure() + .build_client(/*enable*/ true) + .build_server(/*enable*/ true) + .compile_with_config(config, &proto_files, &[PathBuf::from("src/grpc")])?; + + Ok(()) +} diff --git a/codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto b/codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto new file mode 100644 index 0000000000..25db26de3b --- /dev/null +++ b/codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto @@ -0,0 +1,257 @@ +syntax = "proto3"; + +package codex.code_mode.v1; + +// Hosts stateful JavaScript execution and delegates nested tool calls to the +// session owner. Large tool inputs and outputs stay off the session event +// stream so independent HTTP/2 streams can make progress concurrently. +service CodeModeHost { + // Opens a session lease. The first event is always SessionOpened; dropping + // this stream closes the session and terminates its active cells. + rpc OpenSession(OpenSessionRequest) returns (stream SessionEvent); + rpc CloseSession(CloseSessionRequest) returns (CloseSessionResponse); + + // Each subscription owns an independent stream of matching invocations. An + // empty tool_names filter matches every tool. Each invocation is routed to + // exactly one matching subscription, even when filters overlap. + rpc SubscribeToToolCalls(SubscribeToToolCallsRequest) + returns (stream ToolCall); + + // Each result receives its own HTTP/2 stream, preventing a large response + // from blocking unrelated tool completions or session control events. + rpc CompleteToolCall(CompleteToolCallRequest) + returns (CompleteToolCallResponse); + rpc AcknowledgeNotification(AcknowledgeNotificationRequest) + returns (AcknowledgeNotificationResponse); + + // Emits ExecutionStarted immediately, followed by one ExecutionOutcome when + // the execution yields, completes, or is terminated. + rpc Execute(ExecuteRequest) returns (stream ExecuteEvent); + rpc Wait(WaitRequest) returns (WaitResponse); + + // Acknowledges that a canceled wait has retired before another wait starts. + rpc CancelWait(CancelWaitRequest) returns (CancelWaitResponse); + rpc Terminate(TerminateRequest) returns (WaitResponse); +} + +message OpenSessionRequest { + optional SessionCellExecutionLimits cell_execution_limits = 1; +} + +message SessionCellExecutionLimits { + optional uint64 max_yield_time_ms = 1; + optional uint64 max_heap_size_bytes = 2; +} + +message SessionEvent { + oneof event { + SessionOpened opened = 1; + ToolCallCancelled tool_call_cancelled = 2; + Notification notification = 3; + NotificationCancelled notification_cancelled = 4; + CellClosed cell_closed = 5; + } +} + +message SessionOpened { + string session_id = 1; +} + +message CloseSessionRequest { + string session_id = 1; +} + +message CloseSessionResponse {} + +message SubscribeToToolCallsRequest { + string session_id = 1; + repeated ToolName tool_names = 2; +} + +message ToolCall { + string session_id = 1; + + // Correlates callbacks with Execute before ExecutionStarted is received. + string execution_id = 2; + string cell_id = 3; + string invocation_id = 4; + string runtime_tool_call_id = 5; + ToolName tool_name = 6; + ToolKind tool_kind = 7; + optional bytes input_json = 8; + + // Starts at one and increases independently for each execution. + uint64 sequence = 9; +} + +message CompleteToolCallRequest { + string session_id = 1; + string invocation_id = 2; + + oneof outcome { + ToolCallSucceeded succeeded = 3; + ToolCallFailed failed = 4; + } +} + +message ToolCallSucceeded { + bytes output_json = 1; +} + +message ToolCallFailed { + string message = 1; +} + +message CompleteToolCallResponse {} + +message ToolCallCancelled { + string invocation_id = 1; + + // Cancellation can arrive before the corresponding ToolCall because session + // control events and tool subscriptions use independent HTTP/2 streams. +} + +message Notification { + string notification_id = 1; + string execution_id = 2; + string cell_id = 3; + string call_id = 4; + string text = 5; +} + +message NotificationCancelled { + string notification_id = 1; +} + +message AcknowledgeNotificationRequest { + string session_id = 1; + string notification_id = 2; +} + +message AcknowledgeNotificationResponse {} + +message CellClosed { + string execution_id = 1; + string cell_id = 2; + + // All lower sequences must be observed before the client retires the cell. + uint64 final_tool_call_sequence = 3; +} + +message ExecuteRequest { + string session_id = 1; + + // Chosen by the client so callbacks can be correlated before cell admission. + string execution_id = 2; + string tool_call_id = 3; + string source = 4; + repeated ToolDefinition enabled_tools = 5; + optional uint64 yield_time_ms = 6; + optional uint64 max_output_tokens = 7; +} + +message ExecuteEvent { + oneof event { + ExecutionStarted started = 1; + ExecutionOutcome outcome = 2; + } +} + +message ExecutionStarted { + string execution_id = 1; + string cell_id = 2; +} + +message WaitRequest { + string session_id = 1; + string cell_id = 2; + string wait_id = 3; + uint64 yield_time_ms = 4; +} + +message WaitResponse { + oneof state { + ExecutionOutcome live_cell = 1; + ExecutionOutcome missing_cell = 2; + } +} + +message CancelWaitRequest { + string session_id = 1; + string wait_id = 2; +} + +message CancelWaitResponse {} + +message TerminateRequest { + string session_id = 1; + string cell_id = 2; +} + +message ExecutionOutcome { + string cell_id = 1; + repeated ContentItem content_items = 2; + + oneof outcome { + ExecutionYielded yielded = 3; + ExecutionTerminated terminated = 4; + ExecutionCompleted completed = 5; + } +} + +message ExecutionYielded {} + +message ExecutionTerminated {} + +message ExecutionCompleted { + optional string error_text = 1; +} + +message ToolDefinition { + string name = 1; + ToolName tool_name = 2; + string description = 3; + ToolKind kind = 4; + optional bytes input_schema_json = 5; + optional bytes output_schema_json = 6; +} + +message ToolName { + string name = 1; + optional string namespace = 2; +} + +enum ToolKind { + TOOL_KIND_UNSPECIFIED = 0; + TOOL_KIND_FUNCTION = 1; + TOOL_KIND_FREEFORM = 2; +} + +message ContentItem { + oneof item { + TextContent text = 1; + ImageContent image = 2; + AudioContent audio = 3; + } +} + +message TextContent { + string text = 1; +} + +message ImageContent { + string image_url = 1; + optional ImageDetail detail = 2; +} + +message AudioContent { + string audio_url = 1; +} + +enum ImageDetail { + IMAGE_DETAIL_UNSPECIFIED = 0; + IMAGE_DETAIL_AUTO = 1; + IMAGE_DETAIL_LOW = 2; + IMAGE_DETAIL_HIGH = 3; + IMAGE_DETAIL_ORIGINAL = 4; +} diff --git a/codex-rs/code-mode-protocol/src/grpc/mod.rs b/codex-rs/code-mode-protocol/src/grpc/mod.rs new file mode 100644 index 0000000000..566c5cf969 --- /dev/null +++ b/codex-rs/code-mode-protocol/src/grpc/mod.rs @@ -0,0 +1,5 @@ +#[cfg(codex_bazel)] +pub use code_mode_proto::codex::code_mode::v1::*; + +#[cfg(not(codex_bazel))] +tonic::include_proto!("codex.code_mode.v1"); diff --git a/codex-rs/code-mode-protocol/src/lib.rs b/codex-rs/code-mode-protocol/src/lib.rs index 3ec8c9ef1a..6af942a4e8 100644 --- a/codex-rs/code-mode-protocol/src/lib.rs +++ b/codex-rs/code-mode-protocol/src/lib.rs @@ -1,4 +1,5 @@ mod description; +pub mod grpc; pub mod host; mod response; mod runtime;