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

@@ -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",

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\":{}}",

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",
+)

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"

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"],
 )

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 }

codex-rs/code-mode-protocol/build.rs

@@ -0,0 +1,17 @@
+use std::path::PathBuf;
+
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    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::<Result<Vec<_>, _>>()?;
+
+    tonic_prost_build::configure()
+        .build_client(/*enable*/ true)
+        .build_server(/*enable*/ true)
+        .compile_with_config(config, &proto_files, &[PathBuf::from("src/grpc")])?;
+
+    Ok(())
+}

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;
+}

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");

codex-rs/code-mode-protocol/src/lib.rs

@@ -1,4 +1,5 @@
 mod description;
+pub mod grpc;
 pub mod host;
 mod response;
 mod runtime;