From 960d42ddae06242093d153da5b22f932665bc440 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Tue, 12 May 2026 16:34:37 -0700 Subject: [PATCH 1/5] code-mode: carry nested tool kind through runtime (#22377) ## Why Code mode only used nested spec lookup at execution time to rediscover whether a nested tool should be invoked as a function tool or a freeform tool. That information is already present in the enabled tool metadata that code mode builds to expose `tools.*` and `ALL_TOOLS`, so re-looking it up from the router was redundant and kept execution coupled to a separate spec lookup path. ## What Changed - thread `CodeModeToolKind` through the code-mode runtime `ToolCall` event and `CodeModeNestedToolCall` - emit the nested tool kind directly from the V8 callback using the already-enabled tool metadata - build nested tool payloads from the propagated kind instead of calling `find_spec` - remove the now-unused `find_spec` plumbing from the router and parallel runtime helpers - add unit coverage for function vs freeform payload shaping and update affected router tests ## Testing - `cargo test -p codex-code-mode` - `cargo test -p codex-core code_mode::tests` - `cargo test -p codex-core extension_tool_bundles_are_model_visible_and_dispatchable` - `cargo test -p codex-core model_visible_specs_filter_deferred_dynamic_tools` --- codex-rs/code-mode/src/runtime/callbacks.rs | 11 +-- codex-rs/code-mode/src/runtime/mod.rs | 3 + codex-rs/code-mode/src/service.rs | 8 +- codex-rs/core/src/tools/code_mode/mod.rs | 85 ++++++++++++------- codex-rs/core/src/tools/parallel.rs | 5 -- codex-rs/core/src/tools/router.rs | 25 ------ codex-rs/core/src/tools/router_tests.rs | 10 +-- .../src/chatwidget/tests/status_and_layout.rs | 1 + 8 files changed, 71 insertions(+), 77 deletions(-) diff --git a/codex-rs/code-mode/src/runtime/callbacks.rs b/codex-rs/code-mode/src/runtime/callbacks.rs index a9755f6eb0..c3a648ae32 100644 --- a/codex-rs/code-mode/src/runtime/callbacks.rs +++ b/codex-rs/code-mode/src/runtime/callbacks.rs @@ -42,20 +42,16 @@ pub(super) fn tool_callback( let promise = resolver.get_promise(scope); let resolver = v8::Global::new(scope, resolver); - let tool_name = { + let (tool_name, tool_kind) = { let Some(state) = scope.get_slot::() else { throw_type_error(scope, "runtime state unavailable"); return; }; - let Some(tool_name) = state - .enabled_tools - .get(tool_index) - .map(|tool| tool.tool_name.clone()) - else { + let Some(tool) = state.enabled_tools.get(tool_index) else { throw_type_error(scope, "tool callback data is out of range"); return; }; - tool_name + (tool.tool_name.clone(), tool.kind) }; let Some(state) = scope.get_slot_mut::() else { @@ -69,6 +65,7 @@ pub(super) fn tool_callback( let _ = event_tx.send(RuntimeEvent::ToolCall { id, name: tool_name, + kind: tool_kind, input, }); retval.set(promise.into()); diff --git a/codex-rs/code-mode/src/runtime/mod.rs b/codex-rs/code-mode/src/runtime/mod.rs index 200a47c989..f989312da1 100644 --- a/codex-rs/code-mode/src/runtime/mod.rs +++ b/codex-rs/code-mode/src/runtime/mod.rs @@ -14,6 +14,7 @@ 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; @@ -97,6 +98,7 @@ pub struct CodeModeNestedToolCall { pub cell_id: String, pub runtime_tool_call_id: String, pub tool_name: ToolName, + pub tool_kind: CodeModeToolKind, pub input: Option, } @@ -126,6 +128,7 @@ pub(crate) enum RuntimeEvent { ToolCall { id: String, name: ToolName, + kind: CodeModeToolKind, input: Option, }, Notify { diff --git a/codex-rs/code-mode/src/service.rs b/codex-rs/code-mode/src/service.rs index 7326c834e2..3ab8913f19 100644 --- a/codex-rs/code-mode/src/service.rs +++ b/codex-rs/code-mode/src/service.rs @@ -385,11 +385,17 @@ async fn run_session_control( text, }).await; } - RuntimeEvent::ToolCall { id, name, input } => { + RuntimeEvent::ToolCall { + id, + name, + kind, + input, + } => { let tool_call = CodeModeNestedToolCall { cell_id: cell_id.clone(), runtime_tool_call_id: id, tool_name: name, + tool_kind: kind, input, }; let _ = inner diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index 9ee8e0352a..e82faf6951 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -8,6 +8,7 @@ use std::sync::Arc; use std::time::Duration; use codex_code_mode::CodeModeNestedToolCall; +use codex_code_mode::CodeModeToolKind; use codex_code_mode::CodeModeTurnHost; use codex_code_mode::RuntimeResponse; use codex_protocol::models::FunctionCallOutputContentItem; @@ -33,7 +34,6 @@ use crate::tools::router::extension_tool_bundles; use crate::unified_exec::resolve_max_tokens; use codex_features::Feature; use codex_tools::ToolName; -use codex_tools::ToolSpec; use codex_tools::collect_code_mode_tool_definitions; use codex_utils_output_truncation::TruncationPolicy; use codex_utils_output_truncation::formatted_truncate_text_content_items_with_policy; @@ -300,6 +300,7 @@ async fn call_nested_tool( cell_id, runtime_tool_call_id, tool_name, + tool_kind, input, } = invocation; if is_exec_tool_name(&tool_name) { @@ -308,11 +309,10 @@ async fn call_nested_tool( ))); } - let payload = - match build_nested_tool_payload(tool_runtime.find_spec(&tool_name), &tool_name, input) { - Ok(payload) => payload, - Err(error) => return Err(FunctionCallError::RespondToModel(error)), - }; + let payload = match build_nested_tool_payload(tool_kind, &tool_name, input) { + Ok(payload) => payload, + Err(error) => return Err(FunctionCallError::RespondToModel(error)), + }; let call = ToolCall { tool_name, @@ -332,36 +332,14 @@ async fn call_nested_tool( Ok(result.code_mode_result()) } -fn tool_kind_for_spec(spec: &ToolSpec) -> codex_code_mode::CodeModeToolKind { - if matches!(spec, ToolSpec::Freeform(_)) { - codex_code_mode::CodeModeToolKind::Freeform - } else { - codex_code_mode::CodeModeToolKind::Function - } -} - -fn tool_kind_for_name( - spec: Option, - tool_name: &ToolName, -) -> Result { - spec.as_ref() - .map(tool_kind_for_spec) - .ok_or_else(|| format!("tool `{tool_name}` is not enabled in {PUBLIC_TOOL_NAME}")) -} - fn build_nested_tool_payload( - spec: Option, + tool_kind: CodeModeToolKind, tool_name: &ToolName, input: Option, ) -> Result { - let actual_kind = tool_kind_for_name(spec, tool_name)?; - match actual_kind { - codex_code_mode::CodeModeToolKind::Function => { - build_function_tool_payload(tool_name, input) - } - codex_code_mode::CodeModeToolKind::Freeform => { - build_freeform_tool_payload(tool_name, input) - } + match tool_kind { + CodeModeToolKind::Function => build_function_tool_payload(tool_name, input), + CodeModeToolKind::Freeform => build_freeform_tool_payload(tool_name, input), } } @@ -396,3 +374,46 @@ fn build_freeform_tool_payload( _ => Err(format!("tool `{tool_name}` expects a string input")), } } + +#[cfg(test)] +mod tests { + use super::build_nested_tool_payload; + use crate::tools::context::ToolPayload; + use codex_code_mode::CodeModeToolKind; + use codex_tools::ToolName; + use serde_json::json; + + #[test] + fn build_nested_tool_payload_uses_function_kind() { + let payload = build_nested_tool_payload( + CodeModeToolKind::Function, + &ToolName::plain("example"), + Some(json!({ "value": 1 })), + ) + .expect("function payload should serialize"); + + match payload { + ToolPayload::Function { arguments } => { + assert_eq!(arguments, r#"{"value":1}"#.to_string()); + } + other => panic!("expected function payload, got {other:?}"), + } + } + + #[test] + fn build_nested_tool_payload_uses_freeform_kind() { + let payload = build_nested_tool_payload( + CodeModeToolKind::Freeform, + &ToolName::plain("example"), + Some(json!("hello")), + ) + .expect("freeform payload should preserve string input"); + + match payload { + ToolPayload::Custom { input } => { + assert_eq!(input, "hello".to_string()); + } + other => panic!("expected freeform payload, got {other:?}"), + } + } +} diff --git a/codex-rs/core/src/tools/parallel.rs b/codex-rs/core/src/tools/parallel.rs index d7fe22e4f4..2c62e6c15d 100644 --- a/codex-rs/core/src/tools/parallel.rs +++ b/codex-rs/core/src/tools/parallel.rs @@ -22,7 +22,6 @@ use crate::tools::router::ToolCallSource; use crate::tools::router::ToolRouter; use codex_protocol::error::CodexErr; use codex_protocol::models::ResponseInputItem; -use codex_tools::ToolSpec; #[derive(Clone)] pub(crate) struct ToolCallRuntime { @@ -49,10 +48,6 @@ impl ToolCallRuntime { } } - pub(crate) fn find_spec(&self, tool_name: &codex_tools::ToolName) -> Option { - self.router.find_spec(tool_name) - } - pub(crate) fn create_diff_consumer( &self, tool_name: &codex_tools::ToolName, diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index ed3f1a720b..aed72186e0 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -103,31 +103,6 @@ impl ToolRouter { self.model_visible_specs.clone() } - pub fn find_spec(&self, tool_name: &ToolName) -> Option { - self.specs.iter().find_map(|spec| match spec { - ToolSpec::Function(tool) - if tool_name.namespace.is_none() && tool.name == tool_name.name => - { - Some(spec.clone()) - } - ToolSpec::Freeform(tool) - if tool_name.namespace.is_none() && tool.name == tool_name.name => - { - Some(spec.clone()) - } - ToolSpec::Namespace(namespace) => namespace.tools.iter().find_map(|tool| match tool { - ResponsesApiNamespaceTool::Function(tool) - if tool_name.namespace.as_deref() == Some(namespace.name.as_str()) - && tool.name == tool_name.name => - { - Some(ToolSpec::Function(tool.clone())) - } - _ => None, - }), - _ => None, - }) - } - pub(crate) fn create_diff_consumer( &self, tool_name: &ToolName, diff --git a/codex-rs/core/src/tools/router_tests.rs b/codex-rs/core/src/tools/router_tests.rs index e78bd9402f..76776b1add 100644 --- a/codex-rs/core/src/tools/router_tests.rs +++ b/codex-rs/core/src/tools/router_tests.rs @@ -273,11 +273,6 @@ async fn model_visible_specs_filter_deferred_dynamic_tools() -> anyhow::Result<( }, ); - assert!( - router - .find_spec(&ToolName::namespaced("codex_app", hidden_tool)) - .is_some() - ); assert_eq!( namespace_function_names(&router.specs(), "codex_app"), vec![hidden_tool.to_string(), visible_tool.to_string()] @@ -341,8 +336,9 @@ async fn extension_tool_bundles_are_model_visible_and_dispatchable() -> anyhow:: assert!( router - .find_spec(&ToolName::plain("extension_echo")) - .is_some(), + .specs() + .iter() + .any(|spec| spec.name() == "extension_echo"), "expected extension-provided tool spec to be registered" ); assert!( diff --git a/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs b/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs index 6cc9abc704..976a895479 100644 --- a/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs +++ b/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs @@ -401,6 +401,7 @@ async fn completed_plan_table_tail_skips_provisional_history_insert() { } #[tokio::test] +#[cfg_attr(target_os = "windows", ignore = "disabled on windows")] async fn configured_pet_load_is_deferred_until_after_construction() { let (tx_raw, mut rx) = unbounded_channel::(); let tx = AppEventSender::new(tx_raw); From 6dc3b3d7c84969e73270c8eef721d735f7e8bb95 Mon Sep 17 00:00:00 2001 From: Felipe Coury Date: Tue, 12 May 2026 20:50:35 -0300 Subject: [PATCH 2/5] test(tui): relax configured pet load timeout (#22392) ## Why Windows CI has been timing out in `configured_pet_load_is_deferred_until_after_construction` while waiting for the deferred configured-pet load event. The test still needs to prove construction returns before the pet image is available, but the background load slices the built-in pet spritesheet into frame cache files. That work can exceed the old 2 second deadline on slower or more contended CI machines. ## What Changed - Increased the test wait for `ConfiguredPetLoaded` from 2 seconds to 30 seconds. - Kept the post-construction assertion intact so the test still verifies that the pet is not loaded synchronously during `ChatWidget` construction. ## How to Test Targeted tests: - `cargo test -p codex-tui configured_pet_load_is_deferred_until_after_construction` - `just argument-comment-lint` Additional check: - `cargo test -p codex-tui` was run, but the broader crate suite did not complete successfully due to unrelated existing failures: - `status::tests::status_permissions_full_disk_managed_without_network_is_external_sandbox` - `status::tests::status_permissions_full_disk_managed_with_network_is_danger_full_access` - later abort in `tests::fork_last_filters_latest_session_by_cwd_unless_show_all` from stack overflow --- codex-rs/tui/src/chatwidget/tests/status_and_layout.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs b/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs index 976a895479..8b24baeafa 100644 --- a/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs +++ b/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs @@ -435,7 +435,7 @@ async fn configured_pet_load_is_deferred_until_after_construction() { let chat = ChatWidget::new_with_app_event(init); assert!(!chat.ambient_pet_image_enabled()); - let event = tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv()) + let event = tokio::time::timeout(std::time::Duration::from_secs(/*secs*/ 30), rx.recv()) .await .unwrap() .unwrap(); From ac466c0dbd3e7e0289c31e514344f3e24529b84e Mon Sep 17 00:00:00 2001 From: Anton Panasenko Date: Tue, 12 May 2026 16:50:45 -0700 Subject: [PATCH 3/5] feat(exec-server): use protobuf relay frames (#22343) ## Why Remote exec-server now needs one executor websocket to serve multiple harness JSON-RPC sessions. Rendezvous routes by `stream_id`, and the exec-server side needs to use the same stable relay frame contract instead of a hand-rolled JSON shape. The relay protocol also needs to make ownership boundaries clear: harness and executor endpoints own sequencing, acks, retries, duplicate suppression, segmentation, and reassembly; rendezvous only routes frames. ## What Changed - Add the checked-in `codex.exec_server.relay.v1.RelayMessageFrame` proto plus generated prost bindings for `codex-exec-server`. - Encode remote harness/executor relay traffic as binary protobuf websocket frames while keeping local websocket JSON-RPC unchanged. - Demux executor-side relay streams into independent `ConnectionProcessor` sessions keyed by `stream_id`. - Add a programmatic `RemoteExecutorConfig::with_bearer_token(...)` constructor for non-CLI callers and integration tests. - Add an integration test that starts the remote executor against a fake registry/rendezvous websocket and verifies two virtual streams share one executor websocket without cross-talk, including per-stream reset behavior. - Document the remote relay envelope, sequence ranges, `ack`/`ack_bits`, and endpoint responsibilities in `exec-server/README.md`. ## Verification - `cargo test -p codex-exec-server --test relay multiplexed_remote_executor_routes_independent_virtual_streams -- --exact` - `cargo test -p codex-exec-server --test relay` - `cargo test -p codex-exec-server` passed outside the sandbox. The sandboxed run hit macOS `sandbox-exec: sandbox_apply: Operation not permitted` in filesystem sandbox tests. --- codex-rs/Cargo.lock | 1 + codex-rs/exec-server/BUILD.bazel | 3 + codex-rs/exec-server/Cargo.toml | 1 + codex-rs/exec-server/README.md | 60 ++- codex-rs/exec-server/src/client_transport.rs | 26 +- codex-rs/exec-server/src/lib.rs | 2 + .../proto/codex.exec_server.relay.v1.proto | 37 ++ .../src/proto/codex.exec_server.relay.v1.rs | 54 +++ codex-rs/exec-server/src/relay.rs | 455 ++++++++++++++++++ codex-rs/exec-server/src/relay_proto.rs | 7 + codex-rs/exec-server/src/remote.rs | 11 +- codex-rs/exec-server/tests/relay.rs | 353 ++++++++++++++ 12 files changed, 993 insertions(+), 17 deletions(-) create mode 100644 codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.proto create mode 100644 codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.rs create mode 100644 codex-rs/exec-server/src/relay.rs create mode 100644 codex-rs/exec-server/src/relay_proto.rs create mode 100644 codex-rs/exec-server/tests/relay.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index afb3cbbed9..f1d78b7479 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2740,6 +2740,7 @@ dependencies = [ "ctor 0.6.3", "futures", "pretty_assertions", + "prost 0.14.3", "reqwest", "serde", "serde_json", diff --git a/codex-rs/exec-server/BUILD.bazel b/codex-rs/exec-server/BUILD.bazel index 224536da8e..e94a5c0043 100644 --- a/codex-rs/exec-server/BUILD.bazel +++ b/codex-rs/exec-server/BUILD.bazel @@ -10,6 +10,9 @@ codex_rust_crate( # they install process-global test-binary dispatch state, and the remote # exec-server cases already rely on serialization around the full CLI path. integration_test_args = ["--test-threads=1"], + integration_compile_data_extra = [ + "src/proto/codex.exec_server.relay.v1.rs", + ], extra_binaries = [ "//codex-rs/bwrap:bwrap", ], diff --git a/codex-rs/exec-server/Cargo.toml b/codex-rs/exec-server/Cargo.toml index 936fa412f1..09a9a71ea0 100644 --- a/codex-rs/exec-server/Cargo.toml +++ b/codex-rs/exec-server/Cargo.toml @@ -26,6 +26,7 @@ codex-utils-pty = { workspace = true } codex-utils-rustls-provider = { workspace = true } futures = { workspace = true } reqwest = { workspace = true, features = ["json", "rustls-tls", "stream"] } +prost = "0.14.3" serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } thiserror = { workspace = true } diff --git a/codex-rs/exec-server/README.md b/codex-rs/exec-server/README.md index 81664eaca0..1eaf6e69eb 100644 --- a/codex-rs/exec-server/README.md +++ b/codex-rs/exec-server/README.md @@ -30,7 +30,65 @@ It requires a bearer token in `CODEX_EXEC_SERVER_REMOTE_BEARER_TOKEN`. Wire framing: -- websocket: one JSON-RPC message per websocket text frame +- local websocket: one JSON-RPC message per websocket frame +- remote websocket: binary protobuf relay frames carrying JSON-RPC payloads + +## Remote Relay Message Format + +In remote mode, the harness and executor communicate through rendezvous using +`codex.exec_server.relay.v1.RelayMessageFrame`; the checked-in schema is in +`src/proto/codex.exec_server.relay.v1.proto`. The relay frame carries stream +identity plus endpoint-owned reliability metadata: + +```text +version +stream_id +body // data | ack_frame | resume | reset | heartbeat +ack // highest contiguous peer segment seq received +ack_bits // bitset for peer segment seqs after ack +seq // data only: segment sequence number +segment_index // data only: 0-based index within message +segment_count // data only: number of segments in message +payload // data only: JSON-RPC message bytes or segment bytes +next_seq // resume only: next sender seq +reason // reset only: reset reason +``` + +`stream_id` identifies one virtual harness/executor JSON-RPC session on the +executor websocket. The harness generates a UUIDv4 `stream_id`; the executor +demuxes frames by `stream_id` and runs an independent `ConnectionProcessor` per +stream. + +Use segment-level sequence numbers for reliability: + +```text +seq = 0, 1, 2, 3, ... +``` + +Use contiguous segment sequence ranges to identify and stitch a segmented +application message: + +```text +message_start_seq = seq - segment_index +segment_index = 0 +segment_count = 1 +``` + +`message_start_seq` is derived by the receiver, not sent on the wire. For +unsplit messages, `message_start_seq == seq`, `segment_index == 0`, and +`segment_count == 1`. + +Use cumulative `ack` plus fixed-size `ack_bits` instead of variable ack ranges: + +```text +ack = highest contiguous received segment seq +bit i in ack_bits acknowledges seq = ack + 1 + i +``` + +Send `ack` and `ack_bits` redundantly on every outbound frame. Acks are not +themselves acked. Acks, retries, duplicate suppression, segmentation, and +reassembly are endpoint responsibilities; rendezvous only routes relay frames +by `stream_id`. ## Lifecycle diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index 8ca1eb0280..4bdc09a80e 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -15,6 +15,7 @@ use crate::client_api::RemoteExecServerConnectArgs; use crate::client_api::StdioExecServerCommand; use crate::client_api::StdioExecServerConnectArgs; use crate::connection::JsonRpcConnection; +use crate::relay::harness_connection_from_websocket; const ENVIRONMENT_CLIENT_NAME: &str = "codex-environment"; @@ -69,14 +70,13 @@ impl ExecServerClient { source, })?; - Self::connect( - JsonRpcConnection::from_websocket( - stream, - format!("exec-server websocket {websocket_url}"), - ), - args.into(), - ) - .await + let connection_label = format!("exec-server websocket {websocket_url}"); + let connection = if is_rendezvous_harness_url(&websocket_url) { + harness_connection_from_websocket(stream, connection_label) + } else { + JsonRpcConnection::from_websocket(stream, connection_label) + }; + Self::connect(connection, args.into()).await } pub(crate) async fn connect_stdio_command( @@ -120,6 +120,16 @@ impl ExecServerClient { } } +fn is_rendezvous_harness_url(websocket_url: &str) -> bool { + let Some((_path, query)) = websocket_url.split_once('?') else { + return false; + }; + query + .split('&') + .filter_map(|pair| pair.split_once('=')) + .any(|(key, value)| key == "role" && value == "harness") +} + fn stdio_command_process(stdio_command: &StdioExecServerCommand) -> Command { let mut command = Command::new(&stdio_command.program); command.args(&stdio_command.args); diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index d8c147127c..872f16ce32 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -13,6 +13,8 @@ mod local_process; mod process; mod process_id; mod protocol; +mod relay; +mod relay_proto; mod remote; mod remote_file_system; mod remote_process; diff --git a/codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.proto b/codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.proto new file mode 100644 index 0000000000..46527d80cc --- /dev/null +++ b/codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.proto @@ -0,0 +1,37 @@ +syntax = "proto3"; + +package codex.exec_server.relay.v1; + +message RelayMessageFrame { + uint32 version = 1; + string stream_id = 2; + uint32 ack = 3; + uint32 ack_bits = 4; + + oneof body { + RelayData data = 5; + RelayAck ack_frame = 6; + RelayResume resume = 7; + RelayReset reset = 8; + RelayHeartbeat heartbeat = 9; + } +} + +message RelayData { + uint32 seq = 1; + uint32 segment_index = 2; + uint32 segment_count = 3; + bytes payload = 4; +} + +message RelayAck {} + +message RelayResume { + uint32 next_seq = 1; +} + +message RelayReset { + string reason = 1; +} + +message RelayHeartbeat {} diff --git a/codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.rs b/codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.rs new file mode 100644 index 0000000000..072a003dac --- /dev/null +++ b/codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.rs @@ -0,0 +1,54 @@ +// This file is @generated by prost-build. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct RelayMessageFrame { + #[prost(uint32, tag = "1")] + pub version: u32, + #[prost(string, tag = "2")] + pub stream_id: ::prost::alloc::string::String, + #[prost(uint32, tag = "3")] + pub ack: u32, + #[prost(uint32, tag = "4")] + pub ack_bits: u32, + #[prost(oneof = "relay_message_frame::Body", tags = "5, 6, 7, 8, 9")] + pub body: ::core::option::Option, +} +pub mod relay_message_frame { + #[derive(Clone, PartialEq, ::prost::Oneof)] + pub enum Body { + #[prost(message, tag = "5")] + Data(super::RelayData), + #[prost(message, tag = "6")] + AckFrame(super::RelayAck), + #[prost(message, tag = "7")] + Resume(super::RelayResume), + #[prost(message, tag = "8")] + Reset(super::RelayReset), + #[prost(message, tag = "9")] + Heartbeat(super::RelayHeartbeat), + } +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct RelayData { + #[prost(uint32, tag = "1")] + pub seq: u32, + #[prost(uint32, tag = "2")] + pub segment_index: u32, + #[prost(uint32, tag = "3")] + pub segment_count: u32, + #[prost(bytes = "vec", tag = "4")] + pub payload: ::prost::alloc::vec::Vec, +} +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct RelayAck {} +#[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] +pub struct RelayResume { + #[prost(uint32, tag = "1")] + pub next_seq: u32, +} +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct RelayReset { + #[prost(string, tag = "1")] + pub reason: ::prost::alloc::string::String, +} +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct RelayHeartbeat {} diff --git a/codex-rs/exec-server/src/relay.rs b/codex-rs/exec-server/src/relay.rs new file mode 100644 index 0000000000..bce787cfc2 --- /dev/null +++ b/codex-rs/exec-server/src/relay.rs @@ -0,0 +1,455 @@ +use std::collections::HashMap; + +use codex_app_server_protocol::JSONRPCMessage; +use futures::SinkExt; +use futures::StreamExt; +use prost::Message as ProstMessage; +use tokio::io::AsyncRead; +use tokio::io::AsyncWrite; +use tokio::sync::mpsc; +use tokio::sync::watch; +use tokio_tungstenite::WebSocketStream; +use tokio_tungstenite::tungstenite::Message; +use tracing::debug; +use tracing::warn; +use uuid::Uuid; + +use crate::ExecServerError; +use crate::connection::CHANNEL_CAPACITY; +use crate::connection::JsonRpcConnection; +use crate::connection::JsonRpcConnectionEvent; +use crate::connection::JsonRpcTransport; +use crate::relay_proto::RelayData; +use crate::relay_proto::RelayMessageFrame; +use crate::relay_proto::RelayResume; +use crate::relay_proto::relay_message_frame; +use crate::server::ConnectionProcessor; + +const RELAY_MESSAGE_FRAME_VERSION: u32 = 1; + +#[derive(Debug, Clone, Copy, Eq, PartialEq)] +enum RelayFrameBodyKind { + Data, + Ack, + Resume, + Reset, + Heartbeat, +} + +impl RelayMessageFrame { + fn data(stream_id: String, seq: u32, payload: Vec) -> Self { + Self { + version: RELAY_MESSAGE_FRAME_VERSION, + stream_id, + ack: 0, + ack_bits: 0, + body: Some(relay_message_frame::Body::Data(RelayData { + seq, + segment_index: 0, + segment_count: 1, + payload, + })), + } + } + + fn resume(stream_id: String) -> Self { + Self { + version: RELAY_MESSAGE_FRAME_VERSION, + stream_id, + ack: 0, + ack_bits: 0, + body: Some(relay_message_frame::Body::Resume(RelayResume { + next_seq: 0, + })), + } + } + + fn validate(&self) -> Result { + if self.version != RELAY_MESSAGE_FRAME_VERSION { + return Err(ExecServerError::Protocol(format!( + "unsupported relay message frame version {}", + self.version + ))); + } + if self.stream_id.trim().is_empty() { + return Err(ExecServerError::Protocol( + "relay message frame is missing stream_id".to_string(), + )); + } + match self.body.as_ref() { + Some(relay_message_frame::Body::Data(data)) => { + if data.segment_index != 0 || data.segment_count != 1 || data.payload.is_empty() { + return Err(ExecServerError::Protocol( + "relay data message frame is missing required fields".to_string(), + )); + } + Ok(RelayFrameBodyKind::Data) + } + Some(relay_message_frame::Body::AckFrame(_)) => Ok(RelayFrameBodyKind::Ack), + Some(relay_message_frame::Body::Resume(_)) => Ok(RelayFrameBodyKind::Resume), + Some(relay_message_frame::Body::Reset(reset)) => { + if reset.reason.is_empty() { + return Err(ExecServerError::Protocol( + "relay reset message frame is missing reason".to_string(), + )); + } + Ok(RelayFrameBodyKind::Reset) + } + Some(relay_message_frame::Body::Heartbeat(_)) => Ok(RelayFrameBodyKind::Heartbeat), + None => Err(ExecServerError::Protocol( + "relay message frame is missing body".to_string(), + )), + } + } + + fn into_jsonrpc_message(self) -> Result { + let kind = self.validate()?; + if kind != RelayFrameBodyKind::Data { + return Err(ExecServerError::Protocol( + "expected relay data message frame".to_string(), + )); + } + let payload = match self.body { + Some(relay_message_frame::Body::Data(data)) => data.payload, + _ => Vec::new(), + }; + serde_json::from_slice(&payload).map_err(ExecServerError::Json) + } + + fn into_reset_reason(self) -> Option { + match self.body { + Some(relay_message_frame::Body::Reset(reset)) if !reset.reason.is_empty() => { + Some(reset.reason) + } + _ => None, + } + } +} + +fn encode_relay_message_frame(frame: &RelayMessageFrame) -> Vec { + frame.encode_to_vec() +} + +fn decode_relay_message_frame(payload: &[u8]) -> Result { + RelayMessageFrame::decode(payload) + .map_err(|err| ExecServerError::Protocol(format!("invalid relay message frame: {err}"))) +} + +fn jsonrpc_payload(message: &JSONRPCMessage) -> Result, ExecServerError> { + serde_json::to_vec(message).map_err(ExecServerError::Json) +} + +pub(crate) fn harness_connection_from_websocket( + stream: WebSocketStream, + connection_label: String, +) -> JsonRpcConnection +where + S: AsyncRead + AsyncWrite + Unpin + Send + 'static, +{ + let stream_id = Uuid::new_v4().to_string(); + let (mut websocket_writer, mut websocket_reader) = stream.split(); + let (outgoing_tx, mut outgoing_rx) = mpsc::channel(CHANNEL_CAPACITY); + let (incoming_tx, incoming_rx) = mpsc::channel(CHANNEL_CAPACITY); + let (disconnected_tx, disconnected_rx) = watch::channel(false); + + let reader_label = connection_label; + let reader_stream_id = stream_id.clone(); + let incoming_tx_for_reader = incoming_tx; + let disconnected_tx_for_reader = disconnected_tx.clone(); + let reader_task = tokio::spawn(async move { + loop { + match websocket_reader.next().await { + Some(Ok(Message::Binary(payload))) => { + let frame = match decode_relay_message_frame(payload.as_ref()) { + Ok(frame) => frame, + Err(err) => { + let _ = incoming_tx_for_reader + .send(JsonRpcConnectionEvent::MalformedMessage { + reason: format!( + "failed to parse relay message frame from {reader_label}: {err}" + ), + }) + .await; + continue; + } + }; + if frame.stream_id != reader_stream_id { + continue; + } + let kind = match frame.validate() { + Ok(kind) => kind, + Err(err) => { + let _ = incoming_tx_for_reader + .send(JsonRpcConnectionEvent::MalformedMessage { + reason: err.to_string(), + }) + .await; + continue; + } + }; + match kind { + RelayFrameBodyKind::Data => match frame.into_jsonrpc_message() { + Ok(message) => { + if incoming_tx_for_reader + .send(JsonRpcConnectionEvent::Message(message)) + .await + .is_err() + { + break; + } + } + Err(err) => { + let _ = incoming_tx_for_reader + .send(JsonRpcConnectionEvent::MalformedMessage { + reason: err.to_string(), + }) + .await; + } + }, + RelayFrameBodyKind::Reset => { + let _ = disconnected_tx_for_reader.send(true); + let _ = incoming_tx_for_reader + .send(JsonRpcConnectionEvent::Disconnected { + reason: frame.into_reset_reason(), + }) + .await; + break; + } + RelayFrameBodyKind::Ack + | RelayFrameBodyKind::Resume + | RelayFrameBodyKind::Heartbeat => {} + } + } + Some(Ok(Message::Close(_))) | None => { + let _ = disconnected_tx_for_reader.send(true); + let _ = incoming_tx_for_reader + .send(JsonRpcConnectionEvent::Disconnected { reason: None }) + .await; + break; + } + Some(Ok(Message::Ping(_) | Message::Pong(_) | Message::Frame(_))) => {} + Some(Ok(Message::Text(_))) => { + let _ = incoming_tx_for_reader + .send(JsonRpcConnectionEvent::MalformedMessage { + reason: "relay exec-server transport expects binary protobuf frames" + .to_string(), + }) + .await; + } + Some(Err(err)) => { + let _ = disconnected_tx_for_reader.send(true); + let _ = incoming_tx_for_reader + .send(JsonRpcConnectionEvent::Disconnected { + reason: Some(format!( + "failed to read relay websocket frame from {reader_label}: {err}" + )), + }) + .await; + break; + } + } + } + }); + + let writer_task = tokio::spawn(async move { + let resume = RelayMessageFrame::resume(stream_id.clone()); + if websocket_writer + .send(Message::Binary(encode_relay_message_frame(&resume).into())) + .await + .is_err() + { + let _ = disconnected_tx.send(true); + return; + } + + let mut next_seq = 0u32; + while let Some(message) = outgoing_rx.recv().await { + let payload = match jsonrpc_payload(&message) { + Ok(payload) => payload, + Err(err) => { + warn!("failed to serialize JSON-RPC payload for relay transport: {err}"); + break; + } + }; + let frame = RelayMessageFrame::data(stream_id.clone(), next_seq, payload); + next_seq = next_seq.wrapping_add(1); + if websocket_writer + .send(Message::Binary(encode_relay_message_frame(&frame).into())) + .await + .is_err() + { + let _ = disconnected_tx.send(true); + break; + } + } + }); + + JsonRpcConnection { + outgoing_tx, + incoming_rx, + disconnected_rx, + task_handles: vec![reader_task, writer_task], + transport: JsonRpcTransport::Plain, + } +} + +pub(crate) async fn run_multiplexed_executor( + stream: WebSocketStream, + processor: ConnectionProcessor, +) where + S: AsyncRead + AsyncWrite + Unpin + Send + 'static, +{ + let (mut websocket_writer, mut websocket_reader) = stream.split(); + let (physical_outgoing_tx, mut physical_outgoing_rx) = + mpsc::channel::>(CHANNEL_CAPACITY); + let writer_task = tokio::spawn(async move { + while let Some(encoded) = physical_outgoing_rx.recv().await { + if websocket_writer + .send(Message::Binary(encoded.into())) + .await + .is_err() + { + break; + } + } + }); + + let mut streams: HashMap = HashMap::new(); + loop { + let frame = match websocket_reader.next().await { + Some(Ok(Message::Binary(payload))) => { + match decode_relay_message_frame(payload.as_ref()) { + Ok(frame) => frame, + Err(err) => { + warn!("dropping malformed relay message frame from harness: {err}"); + continue; + } + } + } + Some(Ok(Message::Close(_))) | None => break, + Some(Ok(Message::Ping(_) | Message::Pong(_) | Message::Frame(_))) => continue, + Some(Ok(Message::Text(_))) => { + warn!("dropping non-binary relay message frame from harness"); + continue; + } + Some(Err(err)) => { + debug!("multiplexed executor websocket read failed: {err}"); + break; + } + }; + + let kind = match frame.validate() { + Ok(kind) => kind, + Err(err) => { + warn!("dropping invalid relay message frame: {err}"); + continue; + } + }; + + match kind { + RelayFrameBodyKind::Data => { + let stream_id = frame.stream_id.clone(); + let message = match frame.into_jsonrpc_message() { + Ok(message) => message, + Err(err) => { + warn!("dropping malformed relay data message frame: {err}"); + continue; + } + }; + let stream = streams.entry(stream_id.clone()).or_insert_with(|| { + spawn_virtual_stream( + stream_id.clone(), + processor.clone(), + physical_outgoing_tx.clone(), + ) + }); + if stream + .incoming_tx + .send(JsonRpcConnectionEvent::Message(message)) + .await + .is_err() + { + streams.remove(&stream_id); + } + } + RelayFrameBodyKind::Reset => { + if let Some(stream) = streams.remove(&frame.stream_id) { + stream.disconnect(frame.into_reset_reason()).await; + } + } + RelayFrameBodyKind::Ack + | RelayFrameBodyKind::Resume + | RelayFrameBodyKind::Heartbeat => {} + } + } + + for (_stream_id, stream) in streams { + stream.disconnect(/*reason*/ None).await; + } + drop(physical_outgoing_tx); + let _ = writer_task.await; +} + +struct VirtualStream { + incoming_tx: mpsc::Sender, + disconnected_tx: watch::Sender, +} + +impl VirtualStream { + async fn disconnect(self, reason: Option) { + let _ = self.disconnected_tx.send(true); + let _ = self + .incoming_tx + .send(JsonRpcConnectionEvent::Disconnected { reason }) + .await; + } +} + +fn spawn_virtual_stream( + stream_id: String, + processor: ConnectionProcessor, + physical_outgoing_tx: mpsc::Sender>, +) -> VirtualStream { + let (json_outgoing_tx, mut json_outgoing_rx) = mpsc::channel(CHANNEL_CAPACITY); + let (incoming_tx, incoming_rx) = mpsc::channel(CHANNEL_CAPACITY); + let (disconnected_tx, disconnected_rx) = watch::channel(false); + + let writer_stream_id = stream_id; + let writer_task = tokio::spawn(async move { + let mut next_seq = 0u32; + while let Some(message) = json_outgoing_rx.recv().await { + let payload = match jsonrpc_payload(&message) { + Ok(payload) => payload, + Err(err) => { + warn!("failed to serialize virtual stream JSON-RPC payload: {err}"); + break; + } + }; + let frame = RelayMessageFrame::data(writer_stream_id.clone(), next_seq, payload); + next_seq = next_seq.wrapping_add(1); + if physical_outgoing_tx + .send(encode_relay_message_frame(&frame)) + .await + .is_err() + { + break; + } + } + }); + + let connection = JsonRpcConnection { + outgoing_tx: json_outgoing_tx, + incoming_rx, + disconnected_rx, + task_handles: vec![writer_task], + transport: JsonRpcTransport::Plain, + }; + tokio::spawn(async move { + processor.run_connection(connection).await; + }); + + VirtualStream { + incoming_tx, + disconnected_tx, + } +} diff --git a/codex-rs/exec-server/src/relay_proto.rs b/codex-rs/exec-server/src/relay_proto.rs new file mode 100644 index 0000000000..b8a938b8c7 --- /dev/null +++ b/codex-rs/exec-server/src/relay_proto.rs @@ -0,0 +1,7 @@ +#[path = "proto/codex.exec_server.relay.v1.rs"] +mod generated; + +pub(crate) use generated::RelayData; +pub(crate) use generated::RelayMessageFrame; +pub(crate) use generated::RelayResume; +pub(crate) use generated::relay_message_frame; diff --git a/codex-rs/exec-server/src/remote.rs b/codex-rs/exec-server/src/remote.rs index bb22105c19..32c0d5bc8e 100644 --- a/codex-rs/exec-server/src/remote.rs +++ b/codex-rs/exec-server/src/remote.rs @@ -11,7 +11,7 @@ use codex_utils_rustls_provider::ensure_rustls_crypto_provider; use crate::ExecServerError; use crate::ExecServerRuntimePaths; -use crate::connection::JsonRpcConnection; +use crate::relay::run_multiplexed_executor; use crate::server::ConnectionProcessor; pub const CODEX_EXEC_SERVER_REMOTE_BEARER_TOKEN_ENV_VAR: &str = @@ -113,7 +113,7 @@ impl RemoteExecutorConfig { Self::with_bearer_token(base_url, executor_id, read_remote_bearer_token_from_env()?) } - fn with_bearer_token( + pub fn with_bearer_token( base_url: String, executor_id: String, bearer_token: String, @@ -150,12 +150,7 @@ pub async fn run_remote_executor( match connect_async(response.url.as_str()).await { Ok((websocket, _)) => { backoff = Duration::from_secs(1); - processor - .run_connection(JsonRpcConnection::from_websocket( - websocket, - "remote exec-server websocket".to_string(), - )) - .await; + run_multiplexed_executor(websocket, processor.clone()).await; } Err(err) => { warn!("failed to connect remote exec-server websocket: {err}"); diff --git a/codex-rs/exec-server/tests/relay.rs b/codex-rs/exec-server/tests/relay.rs new file mode 100644 index 0000000000..d228db9ee5 --- /dev/null +++ b/codex-rs/exec-server/tests/relay.rs @@ -0,0 +1,353 @@ +mod common; + +#[path = "../src/proto/codex.exec_server.relay.v1.rs"] +mod relay_proto; + +use std::collections::HashMap; +use std::time::Duration; + +use anyhow::Context; +use anyhow::Result; +use anyhow::anyhow; +use anyhow::bail; +use codex_app_server_protocol::JSONRPCError; +use codex_app_server_protocol::JSONRPCMessage; +use codex_app_server_protocol::JSONRPCNotification; +use codex_app_server_protocol::JSONRPCRequest; +use codex_app_server_protocol::JSONRPCResponse; +use codex_app_server_protocol::RequestId; +use codex_exec_server::ExecServerRuntimePaths; +use codex_exec_server::InitializeParams; +use codex_exec_server::InitializeResponse; +use codex_exec_server::RemoteExecutorConfig; +use futures::SinkExt; +use futures::StreamExt; +use pretty_assertions::assert_eq; +use prost::Message as ProstMessage; +use relay_proto::RelayData; +use relay_proto::RelayMessageFrame; +use relay_proto::RelayReset; +use relay_proto::relay_message_frame; +use tokio::net::TcpListener; +use tokio::time::timeout; +use tokio_tungstenite::WebSocketStream; +use tokio_tungstenite::accept_async; +use tokio_tungstenite::tungstenite::Message; +use uuid::Uuid; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::header; +use wiremock::matchers::method; +use wiremock::matchers::path; + +const EXECUTOR_ID: &str = "exec-mux-test"; +const REGISTRY_TOKEN: &str = "registry-token"; +const RELAY_MESSAGE_FRAME_VERSION: u32 = 1; +const TEST_TIMEOUT: Duration = Duration::from_secs(5); + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn multiplexed_remote_executor_routes_independent_virtual_streams() -> Result<()> { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let rendezvous_url = format!("ws://{}", listener.local_addr()?); + let registry = MockServer::start().await; + Mock::given(method("POST")) + .and(path(format!("/cloud/executor/{EXECUTOR_ID}/register"))) + .and(header("authorization", format!("Bearer {REGISTRY_TOKEN}"))) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "executor_id": EXECUTOR_ID, + "url": rendezvous_url, + }))) + .mount(®istry) + .await; + + let (codex_exe, codex_linux_sandbox_exe) = common::current_test_binary_helper_paths()?; + let runtime_paths = ExecServerRuntimePaths::new(codex_exe, codex_linux_sandbox_exe)?; + let config = RemoteExecutorConfig::with_bearer_token( + registry.uri(), + EXECUTOR_ID.to_string(), + REGISTRY_TOKEN.to_string(), + )?; + let remote_executor = tokio::spawn(codex_exec_server::run_remote_executor( + config, + runtime_paths, + )); + + let (socket, _peer_addr) = timeout(TEST_TIMEOUT, listener.accept()) + .await + .context("remote executor should connect to fake rendezvous")??; + let mut websocket = timeout(TEST_TIMEOUT, accept_async(socket)) + .await + .context("fake rendezvous should accept executor websocket")??; + + let stream_a = "stream-a"; + let stream_b = "stream-b"; + send_relay_message( + &mut websocket, + stream_a, + /*seq*/ 0, + initialize_request(/*id*/ 1, "relay-test-a")?, + ) + .await?; + send_relay_message( + &mut websocket, + stream_b, + /*seq*/ 0, + initialize_request(/*id*/ 1, "relay-test-b")?, + ) + .await?; + + let initialize_responses = read_relay_messages_by_stream(&mut websocket, /*count*/ 2).await?; + let session_a = + assert_initialize_response(initialize_responses.get(stream_a), stream_a, /*id*/ 1)?; + let session_b = + assert_initialize_response(initialize_responses.get(stream_b), stream_b, /*id*/ 1)?; + assert_ne!(session_a, session_b); + + send_relay_message( + &mut websocket, + stream_a, + /*seq*/ 1, + notification("initialized", serde_json::json!({})), + ) + .await?; + send_relay_message( + &mut websocket, + stream_b, + /*seq*/ 1, + notification("initialized", serde_json::json!({})), + ) + .await?; + + send_relay_message( + &mut websocket, + stream_a, + /*seq*/ 2, + request(/*id*/ 2, "test/unknown-a", serde_json::json!({})), + ) + .await?; + send_relay_message( + &mut websocket, + stream_b, + /*seq*/ 2, + request(/*id*/ 2, "test/unknown-b", serde_json::json!({})), + ) + .await?; + + let unknown_method_responses = + read_relay_messages_by_stream(&mut websocket, /*count*/ 2).await?; + assert_error_response( + unknown_method_responses.get(stream_a), + stream_a, + /*id*/ 2, + "test/unknown-a", + )?; + assert_error_response( + unknown_method_responses.get(stream_b), + stream_b, + /*id*/ 2, + "test/unknown-b", + )?; + + send_relay_reset(&mut websocket, stream_a, "test_reset").await?; + send_relay_message( + &mut websocket, + stream_b, + /*seq*/ 3, + request( + /*id*/ 3, + "test/unknown-b-after-reset", + serde_json::json!({}), + ), + ) + .await?; + + let (stream_id, message) = read_relay_message(&mut websocket).await?; + assert_eq!(stream_id, stream_b); + assert_error_response( + Some(&message), + stream_b, + /*id*/ 3, + "test/unknown-b-after-reset", + )?; + + websocket.close(None).await?; + remote_executor.abort(); + let _ = remote_executor.await; + Ok(()) +} + +async fn send_relay_message( + websocket: &mut WebSocketStream, + stream_id: &str, + seq: u32, + message: JSONRPCMessage, +) -> Result<()> +where + S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin, +{ + let payload = serde_json::to_vec(&message)?; + let frame = RelayMessageFrame { + version: RELAY_MESSAGE_FRAME_VERSION, + stream_id: stream_id.to_string(), + ack: 0, + ack_bits: 0, + body: Some(relay_message_frame::Body::Data(RelayData { + seq, + segment_index: 0, + segment_count: 1, + payload, + })), + }; + send_relay_frame(websocket, frame).await +} + +async fn send_relay_reset( + websocket: &mut WebSocketStream, + stream_id: &str, + reason: &str, +) -> Result<()> +where + S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin, +{ + send_relay_frame( + websocket, + RelayMessageFrame { + version: RELAY_MESSAGE_FRAME_VERSION, + stream_id: stream_id.to_string(), + ack: 0, + ack_bits: 0, + body: Some(relay_message_frame::Body::Reset(RelayReset { + reason: reason.to_string(), + })), + }, + ) + .await +} + +async fn send_relay_frame( + websocket: &mut WebSocketStream, + frame: RelayMessageFrame, +) -> Result<()> +where + S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin, +{ + websocket + .send(Message::Binary(frame.encode_to_vec().into())) + .await?; + Ok(()) +} + +async fn read_relay_messages_by_stream( + websocket: &mut WebSocketStream, + count: usize, +) -> Result> +where + S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin, +{ + let mut messages = HashMap::new(); + for _ in 0..count { + let (stream_id, message) = read_relay_message(websocket).await?; + if messages.insert(stream_id.clone(), message).is_some() { + bail!("received duplicate response for stream {stream_id}"); + } + } + Ok(messages) +} + +async fn read_relay_message( + websocket: &mut WebSocketStream, +) -> Result<(String, JSONRPCMessage)> +where + S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin, +{ + loop { + let frame = timeout(TEST_TIMEOUT, websocket.next()) + .await + .context("timed out waiting for relay frame")? + .ok_or_else(|| anyhow!("executor websocket closed"))??; + match frame { + Message::Binary(bytes) => { + let frame = RelayMessageFrame::decode(bytes.as_ref())?; + let stream_id = frame.stream_id; + let Some(relay_message_frame::Body::Data(data)) = frame.body else { + continue; + }; + let message = serde_json::from_slice(&data.payload)?; + return Ok((stream_id, message)); + } + Message::Ping(_) | Message::Pong(_) => {} + Message::Close(_) => bail!("executor websocket closed"), + Message::Text(_) => bail!("executor sent text frame on relay websocket"), + Message::Frame(_) => {} + } + } +} + +fn initialize_request(id: i64, client_name: &str) -> Result { + Ok(request( + id, + "initialize", + serde_json::to_value(InitializeParams { + client_name: client_name.to_string(), + resume_session_id: None, + })?, + )) +} + +fn request(id: i64, method: &str, params: serde_json::Value) -> JSONRPCMessage { + JSONRPCMessage::Request(JSONRPCRequest { + id: RequestId::Integer(id), + method: method.to_string(), + params: Some(params), + trace: None, + }) +} + +fn notification(method: &str, params: serde_json::Value) -> JSONRPCMessage { + JSONRPCMessage::Notification(JSONRPCNotification { + method: method.to_string(), + params: Some(params), + }) +} + +fn assert_initialize_response( + message: Option<&JSONRPCMessage>, + stream_id: &str, + id: i64, +) -> Result { + let message = message.ok_or_else(|| anyhow!("missing initialize response for {stream_id}"))?; + let JSONRPCMessage::Response(JSONRPCResponse { + id: response_id, + result, + }) = message + else { + bail!("expected initialize response for {stream_id}, got {message:?}"); + }; + assert_eq!(response_id, &RequestId::Integer(id)); + let response: InitializeResponse = serde_json::from_value(result.clone())?; + Ok(Uuid::parse_str(&response.session_id)?) +} + +fn assert_error_response( + message: Option<&JSONRPCMessage>, + stream_id: &str, + id: i64, + expected_method: &str, +) -> Result<()> { + let message = message.ok_or_else(|| anyhow!("missing error response for {stream_id}"))?; + let JSONRPCMessage::Error(JSONRPCError { + id: response_id, + error, + }) = message + else { + bail!("expected error response for {stream_id}, got {message:?}"); + }; + assert_eq!(response_id, &RequestId::Integer(id)); + assert!( + error.message.contains(expected_method), + "expected error for {stream_id} to mention {expected_method}, got {}", + error.message + ); + Ok(()) +} From b6e718591b3d516172b6c812c4c3f4e513637604 Mon Sep 17 00:00:00 2001 From: richardopenai Date: Tue, 12 May 2026 17:07:33 -0700 Subject: [PATCH 4/5] [codex] Remove workspace owner usage nudge gate (#20509) ## Summary - make workspace owner nudge handling unconditional in the TUI now that it is fully rolled out - keep `workspace_owner_usage_nudge` as a removed no-op compatibility flag so old configs/app overrides remain accepted during rollout - remove flag-disabled test setup ## Companion PR - https://github.com/openai/openai/pull/876351 removes the Codex Apps Statsig rollout gate override after this change is available to the app/runtime path ## Validation - `just write-config-schema` - `just fmt` - `cargo test -p codex-features` - `cargo test -p codex-tui status_and_layout` --- codex-rs/features/src/lib.rs | 5 +- codex-rs/tui/src/chatwidget.rs | 22 +------ .../src/chatwidget/tests/status_and_layout.rs | 58 ------------------- 3 files changed, 4 insertions(+), 81 deletions(-) diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index 955ae44c00..3be2404a9a 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -231,7 +231,8 @@ pub enum Feature { TuiAppServer, /// Prevent idle system sleep while a turn is actively running. PreventIdleSleep, - /// Enable workspace-specific owner nudge copy and prompts in the TUI. + /// Removed compatibility flag retained as a no-op now that workspace owner + /// usage nudges are always enabled. WorkspaceOwnerUsageNudge, /// Legacy rollout flag for Responses API WebSocket transport experiments. ResponsesWebsockets, @@ -1153,7 +1154,7 @@ pub const FEATURES: &[FeatureSpec] = &[ FeatureSpec { id: Feature::WorkspaceOwnerUsageNudge, key: "workspace_owner_usage_nudge", - stage: Stage::UnderDevelopment, + stage: Stage::Removed, default_enabled: false, }, FeatureSpec { diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 9efc995313..38b6181221 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -2890,18 +2890,7 @@ impl ChatWidget { self.maybe_send_next_queued_input(); } - fn workspace_owner_usage_nudge_enabled(&self) -> bool { - self.config - .features - .enabled(Feature::WorkspaceOwnerUsageNudge) - } - fn on_rate_limit_error(&mut self, error_kind: RateLimitErrorKind, message: String) { - if !self.workspace_owner_usage_nudge_enabled() { - self.on_error(message); - return; - } - let rate_limit_reached_type = self.codex_rate_limit_reached_type.map(|kind| { if matches!(error_kind, RateLimitErrorKind::UsageLimit) { match kind { @@ -6919,9 +6908,7 @@ impl ChatWidget { } fn open_workspace_owner_nudge_prompt(&mut self, credit_type: AddCreditsNudgeCreditType) { - if !self.workspace_owner_usage_nudge_enabled() - || self.add_credits_nudge_email_in_flight.is_some() - { + if self.add_credits_nudge_email_in_flight.is_some() { return; } @@ -6969,10 +6956,6 @@ impl ChatWidget { &mut self, credit_type: AddCreditsNudgeCreditType, ) -> bool { - if !self.workspace_owner_usage_nudge_enabled() { - return false; - } - self.add_credits_nudge_email_in_flight = Some(credit_type); true } @@ -6985,9 +6968,6 @@ impl ChatWidget { .add_credits_nudge_email_in_flight .take() .unwrap_or(AddCreditsNudgeCreditType::Credits); - if !self.workspace_owner_usage_nudge_enabled() { - return; - } let message = match (credit_type, result) { (AddCreditsNudgeCreditType::Credits, Ok(AddCreditsNudgeEmailStatus::Sent)) => { "Workspace owner notified." diff --git a/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs b/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs index 8b24baeafa..4cc01d7591 100644 --- a/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs +++ b/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs @@ -797,57 +797,9 @@ async fn rate_limit_switch_prompt_popup_snapshot() { assert_chatwidget_snapshot!("rate_limit_switch_prompt_popup", popup); } -#[tokio::test] -async fn workspace_owner_usage_nudge_flag_disabled_keeps_generic_rate_limit_error() { - { - let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; - let mut limits = snapshot(/*percent*/ 100.0); - limits.rate_limit_reached_type = - Some(RateLimitReachedType::WorkspaceMemberUsageLimitReached); - chat.on_rate_limit_snapshot(Some(limits)); - - chat.on_rate_limit_error( - RateLimitErrorKind::UsageLimit, - "Usage limit reached.".to_string(), - ); - let rendered = drain_insert_history(&mut rx) - .into_iter() - .map(|lines| lines_to_single_string(&lines)) - .collect::(); - assert!( - rendered.contains("Usage limit reached."), - "rendered: {rendered}" - ); - } - - { - let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; - let mut limits = snapshot(/*percent*/ 100.0); - limits.rate_limit_reached_type = - Some(RateLimitReachedType::WorkspaceMemberUsageLimitReached); - chat.on_rate_limit_snapshot(Some(limits)); - - chat.on_rate_limit_error( - RateLimitErrorKind::UsageLimit, - "Usage limit reached.".to_string(), - ); - let popup = render_bottom_popup(&chat, /*width*/ 100); - assert!( - !popup.contains("Request a limit increase from your owner"), - "popup: {popup}" - ); - assert_no_owner_nudge_or_rate_limit_refresh(&mut rx); - } -} - -fn enable_workspace_owner_usage_nudge(chat: &mut ChatWidget) { - chat.set_feature_enabled(Feature::WorkspaceOwnerUsageNudge, /*enabled*/ true); -} - #[tokio::test] async fn workspace_member_credits_depleted_prompts_and_sends_credits() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; - enable_workspace_owner_usage_nudge(&mut chat); let mut limits = snapshot(/*percent*/ 100.0); limits.rate_limit_reached_type = Some(RateLimitReachedType::WorkspaceMemberCreditsDepleted); chat.on_rate_limit_snapshot(Some(limits)); @@ -867,7 +819,6 @@ async fn workspace_member_credits_depleted_prompts_and_sends_credits() { #[tokio::test] async fn workspace_member_usage_limit_prompts_and_sends_usage_limit() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; - enable_workspace_owner_usage_nudge(&mut chat); let mut limits = snapshot(/*percent*/ 100.0); limits.rate_limit_reached_type = Some(RateLimitReachedType::WorkspaceMemberUsageLimitReached); chat.on_rate_limit_snapshot(Some(limits)); @@ -887,7 +838,6 @@ async fn workspace_member_usage_limit_prompts_and_sends_usage_limit() { #[tokio::test] async fn header_rate_limit_snapshot_preserves_member_limit_type_for_error_prompt() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; - enable_workspace_owner_usage_nudge(&mut chat); let mut usage_limits = snapshot(/*percent*/ 100.0); usage_limits.rate_limit_reached_type = Some(RateLimitReachedType::WorkspaceMemberUsageLimitReached); @@ -917,7 +867,6 @@ async fn header_rate_limit_snapshot_preserves_member_limit_type_for_error_prompt #[tokio::test] async fn usage_limit_error_remaps_stale_member_credits_state_to_usage_limit_prompt() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; - enable_workspace_owner_usage_nudge(&mut chat); let mut limits = snapshot(/*percent*/ 100.0); limits.rate_limit_reached_type = Some(RateLimitReachedType::WorkspaceMemberCreditsDepleted); chat.on_rate_limit_snapshot(Some(limits)); @@ -954,7 +903,6 @@ async fn workspace_owner_limit_states_do_not_prompt_for_owner_nudge() { ), ] { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; - enable_workspace_owner_usage_nudge(&mut chat); let mut limits = snapshot(/*percent*/ 100.0); limits.rate_limit_reached_type = Some(limit_type); chat.on_rate_limit_snapshot(Some(limits)); @@ -984,7 +932,6 @@ async fn workspace_owner_limit_states_render_state_specific_messages() { let mut rendered_cases = Vec::new(); for (limit_type, error_kind, expected) in cases { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; - enable_workspace_owner_usage_nudge(&mut chat); let mut limits = snapshot(/*percent*/ 100.0); limits.rate_limit_reached_type = Some(limit_type); chat.on_rate_limit_snapshot(Some(limits)); @@ -1007,7 +954,6 @@ async fn workspace_owner_limit_states_render_state_specific_messages() { #[tokio::test] async fn missing_rate_limit_reached_type_does_not_prompt_or_refresh() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; - enable_workspace_owner_usage_nudge(&mut chat); chat.on_rate_limit_snapshot(Some(snapshot(/*percent*/ 100.0))); chat.on_rate_limit_error( @@ -1022,7 +968,6 @@ async fn missing_rate_limit_reached_type_does_not_prompt_or_refresh() { #[tokio::test] async fn workspace_owner_nudge_default_no_dismisses_without_sending() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; - enable_workspace_owner_usage_nudge(&mut chat); let mut limits = snapshot(/*percent*/ 100.0); limits.rate_limit_reached_type = Some(RateLimitReachedType::WorkspaceMemberCreditsDepleted); chat.on_rate_limit_snapshot(Some(limits)); @@ -1039,7 +984,6 @@ async fn workspace_owner_nudge_default_no_dismisses_without_sending() { #[tokio::test] async fn workspace_owner_nudge_reappears_after_dismissing_no() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; - enable_workspace_owner_usage_nudge(&mut chat); let mut limits = snapshot(/*percent*/ 100.0); limits.rate_limit_reached_type = Some(RateLimitReachedType::WorkspaceMemberUsageLimitReached); chat.on_rate_limit_snapshot(Some(limits)); @@ -1082,7 +1026,6 @@ async fn workspace_owner_credits_nudge_completion_renders_feedback() { let mut rendered_cases = Vec::new(); for (result, expected) in cases { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; - enable_workspace_owner_usage_nudge(&mut chat); chat.start_add_credits_nudge_email_request(AddCreditsNudgeCreditType::Credits); chat.finish_add_credits_nudge_email_request(result); let rendered = drain_insert_history(&mut rx) @@ -1119,7 +1062,6 @@ async fn workspace_owner_usage_limit_nudge_completion_renders_feedback() { let mut rendered_cases = Vec::new(); for (result, expected) in cases { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; - enable_workspace_owner_usage_nudge(&mut chat); chat.start_add_credits_nudge_email_request(AddCreditsNudgeCreditType::UsageLimit); chat.finish_add_credits_nudge_email_request(result); let rendered = drain_insert_history(&mut rx) From 0173f711437a0f34ca1cd76b0cbf3aa7c7f4e5bc Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Tue, 12 May 2026 17:09:14 -0700 Subject: [PATCH 5/5] Refactor namespaced tool spec registration (#22256) ## Summary This refactor makes tool handlers the owner of the specs they can publish, so registry construction can register handlers once and separately publish only the specs that should be model-visible. The main motivation is deferred tools: MCP and dynamic tools still need handlers registered up front, but deferred tools should be discoverable through `tool_search` rather than emitted in the initial tool spec list. ## What changed - `McpHandler` and `DynamicToolHandler` can return their own `ToolSpec`. - `build_tool_registry_builder` now collects handlers, registers them through the no-spec path, and publishes only non-deferred handler specs. - Deferred MCP and dynamic tool names are combined into one `all_deferred_tools` set that drives spec filtering, code-mode deferred-tool signaling, and `tool_search` registration. - `tool_search` registration now requires both deferred tools and `namespace_tools`. - Namespace specs are merged in `spec_plan`, preserving top-level spec order, sorting tools within each namespace, and backfilling empty namespace descriptions. - Hosted web search and image-generation specs are included in the collected spec vector before namespace merge/publication, and tool-name tests that should not care about hosted relative order now compare sets. ## Testing - `cargo test -p codex-core tools::spec::tests:: -- --nocapture` - `cargo test -p codex-core tools::spec_plan::tests:: -- --nocapture` - `cargo test -p codex-core tools::router::tests::specs_filter_deferred_dynamic_tools -- --nocapture` - `cargo test -p codex-core suite::prompt_caching::prompt_tools_are_consistent_across_requests -- --nocapture` - `just fmt` - `just fix -p codex-core` - `cargo test -p codex-core -- --skip tools::handlers::multi_agents::tests::tool_handlers_cascade_close_and_resume_and_keep_explicitly_closed_subtrees_closed` passed the library suite after skipping the known stack-overflowing unit test. Full `cargo test -p codex-core` currently hits a stack overflow in `tools::handlers::multi_agents::tests::tool_handlers_cascade_close_and_resume_and_keep_explicitly_closed_subtrees_closed`; the same focused test reproduces on `origin/main`. --- .../src/tools/code_mode/execute_handler.rs | 12 +- codex-rs/core/src/tools/code_mode/mod.rs | 33 -- codex-rs/core/src/tools/handlers/dynamic.rs | 17 +- codex-rs/core/src/tools/handlers/mcp.rs | 24 + codex-rs/core/src/tools/registry.rs | 42 +- codex-rs/core/src/tools/registry_tests.rs | 9 +- codex-rs/core/src/tools/router.rs | 6 - codex-rs/core/src/tools/router_tests.rs | 13 +- codex-rs/core/src/tools/spec.rs | 7 +- codex-rs/core/src/tools/spec_plan.rs | 416 +++++++++--------- codex-rs/core/src/tools/spec_plan_tests.rs | 45 +- codex-rs/core/src/tools/spec_tests.rs | 68 +-- codex-rs/core/src/tools/tool_search_entry.rs | 6 +- codex-rs/core/tests/suite/prompt_caching.rs | 2 +- 14 files changed, 342 insertions(+), 358 deletions(-) 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 80a72a88ec..8eb1697fc6 100644 --- a/codex-rs/core/src/tools/code_mode/execute_handler.rs +++ b/codex-rs/core/src/tools/code_mode/execute_handler.rs @@ -8,17 +8,20 @@ use codex_tools::ToolSpec; use super::ExecContext; use super::PUBLIC_TOOL_NAME; -use super::build_enabled_tools; use super::handle_runtime_response; use super::is_exec_tool_name; pub struct CodeModeExecuteHandler { spec: ToolSpec, + nested_tool_specs: Vec, } impl CodeModeExecuteHandler { - pub(crate) fn new(spec: ToolSpec) -> Self { - Self { spec } + pub(crate) fn new(spec: ToolSpec, nested_tool_specs: Vec) -> Self { + Self { + spec, + nested_tool_specs, + } } async fn execute( @@ -31,7 +34,8 @@ impl CodeModeExecuteHandler { let args = codex_code_mode::parse_exec_source(&code).map_err(FunctionCallError::RespondToModel)?; let exec = ExecContext { session, turn }; - let enabled_tools = build_enabled_tools(&exec).await; + let enabled_tools = + codex_tools::collect_code_mode_tool_definitions(&self.nested_tool_specs); let stored_values = exec .session .services diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index e82faf6951..ff9f8c8893 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -29,12 +29,9 @@ use crate::tools::context::ToolPayload; use crate::tools::parallel::ToolCallRuntime; use crate::tools::router::ToolCall; use crate::tools::router::ToolCallSource; -use crate::tools::router::ToolRouterParams; -use crate::tools::router::extension_tool_bundles; use crate::unified_exec::resolve_max_tokens; use codex_features::Feature; use codex_tools::ToolName; -use codex_tools::collect_code_mode_tool_definitions; use codex_utils_output_truncation::TruncationPolicy; use codex_utils_output_truncation::formatted_truncate_text_content_items_with_policy; use codex_utils_output_truncation::truncate_function_output_items_with_policy; @@ -260,36 +257,6 @@ fn truncate_code_mode_result( truncate_function_output_items_with_policy(&items, policy) } -pub(super) async fn build_enabled_tools( - exec: &ExecContext, -) -> Vec { - let router = build_nested_router(exec).await; - let specs = router.specs(); - collect_code_mode_tool_definitions(&specs) -} - -#[expect( - clippy::await_holding_invalid_type, - reason = "nested tool router construction reads through the session-owned manager guard" -)] -async fn build_nested_router(exec: &ExecContext) -> ToolRouter { - let nested_tools_config = exec.turn.tools_config.for_code_mode_nested_tools(); - let mcp_connection_manager = exec.session.services.mcp_connection_manager.read().await; - let listed_mcp_tools = mcp_connection_manager.list_all_tools().await; - - ToolRouter::from_config( - &nested_tools_config, - ToolRouterParams { - deferred_mcp_tools: None, - mcp_tools: Some(listed_mcp_tools), - unavailable_called_tools: Vec::new(), - discoverable_tools: None, - extension_tool_bundles: extension_tool_bundles(exec.session.as_ref()), - dynamic_tools: exec.turn.dynamic_tools.as_slice(), - }, - ) -} - async fn call_nested_tool( _exec: ExecContext, tool_runtime: ToolCallRuntime, diff --git a/codex-rs/core/src/tools/handlers/dynamic.rs b/codex-rs/core/src/tools/handlers/dynamic.rs index 7cc9e61e30..9ecdf2340d 100644 --- a/codex-rs/core/src/tools/handlers/dynamic.rs +++ b/codex-rs/core/src/tools/handlers/dynamic.rs @@ -9,10 +9,13 @@ use crate::tools::registry::ToolHandler; use crate::turn_timing::now_unix_timestamp_ms; use codex_protocol::dynamic_tools::DynamicToolCallRequest; use codex_protocol::dynamic_tools::DynamicToolResponse; +use codex_protocol::dynamic_tools::DynamicToolSpec; use codex_protocol::models::FunctionCallOutputContentItem; use codex_protocol::protocol::DynamicToolCallResponseEvent; use codex_protocol::protocol::EventMsg; use codex_tools::ToolName; +use codex_tools::ToolSpec; +use codex_tools::dynamic_tool_to_loadable_tool_spec; use serde_json::Value; use std::time::Instant; use tokio::sync::oneshot; @@ -20,11 +23,17 @@ use tracing::warn; pub struct DynamicToolHandler { tool_name: ToolName, + spec: Option, } impl DynamicToolHandler { - pub fn new(tool_name: ToolName) -> Self { - Self { tool_name } + pub fn new(tool: &DynamicToolSpec) -> Option { + let tool_name = ToolName::new(tool.namespace.clone(), tool.name.clone()); + let spec = dynamic_tool_to_loadable_tool_spec(tool).ok()?.into(); + Some(Self { + tool_name, + spec: Some(spec), + }) } } @@ -35,6 +44,10 @@ impl ToolHandler for DynamicToolHandler { self.tool_name.clone() } + fn spec(&self) -> Option { + self.spec.clone() + } + async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { session, diff --git a/codex-rs/core/src/tools/handlers/mcp.rs b/codex-rs/core/src/tools/handlers/mcp.rs index 8bb8337337..89666c80aa 100644 --- a/codex-rs/core/src/tools/handlers/mcp.rs +++ b/codex-rs/core/src/tools/handlers/mcp.rs @@ -14,7 +14,11 @@ use crate::tools::registry::PreToolUsePayload; use crate::tools::registry::ToolHandler; use crate::tools::registry::ToolTelemetryTags; use codex_mcp::ToolInfo; +use codex_tools::ResponsesApiNamespace; +use codex_tools::ResponsesApiNamespaceTool; use codex_tools::ToolName; +use codex_tools::ToolSpec; +use codex_tools::mcp_tool_to_responses_api_tool; use serde_json::Map; use serde_json::Value; @@ -35,6 +39,26 @@ impl ToolHandler for McpHandler { self.tool_info.canonical_tool_name() } + fn spec(&self) -> Option { + let tool_name = self.tool_name(); + let namespace_name = tool_name.namespace.as_ref()?; + let tool = mcp_tool_to_responses_api_tool(&tool_name, &self.tool_info.tool).ok()?; + let description = self + .tool_info + .namespace_description + .as_deref() + .map(str::trim) + .filter(|description| !description.is_empty()) + .map(str::to_string) + .unwrap_or_default(); + + Some(ToolSpec::Namespace(ResponsesApiNamespace { + name: namespace_name.clone(), + description, + tools: vec![ResponsesApiNamespaceTool::Function(tool)], + })) + } + fn supports_parallel_tool_calls(&self) -> bool { self.tool_info.supports_parallel_tool_calls } diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index 01b3caa34c..95e7d17ce5 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -168,7 +168,11 @@ pub(crate) struct PostToolUsePayload { pub(crate) tool_response: Value, } -trait AnyToolHandler: Send + Sync { +pub(crate) trait AnyToolHandler: Send + Sync { + fn tool_name(&self) -> ToolName; + + fn spec(&self) -> Option; + fn supports_parallel_tool_calls(&self) -> bool; fn matches_kind(&self, payload: &ToolPayload) -> bool; @@ -197,6 +201,14 @@ impl AnyToolHandler for T where T: ToolHandler, { + fn tool_name(&self) -> ToolName { + ToolHandler::tool_name(self) + } + + fn spec(&self) -> Option { + ToolHandler::spec(self) + } + fn supports_parallel_tool_calls(&self) -> bool { ToolHandler::supports_parallel_tool_calls(self) } @@ -531,24 +543,17 @@ impl ToolRegistry { pub struct ToolRegistryBuilder { handlers: HashMap>, specs: Vec, - code_mode_enabled: bool, } impl ToolRegistryBuilder { - pub fn new(code_mode_enabled: bool) -> Self { + pub fn new() -> Self { Self { handlers: HashMap::new(), specs: Vec::new(), - code_mode_enabled, } } pub(crate) fn push_spec(&mut self, spec: ToolSpec) { - let spec = if self.code_mode_enabled { - codex_tools::augment_tool_spec_for_code_mode(spec) - } else { - spec - }; self.specs.push(spec); } @@ -556,17 +561,32 @@ impl ToolRegistryBuilder { where H: ToolHandler + 'static, { + self.register_any_handler(handler); + } + + pub(crate) fn register_any_handler(&mut self, handler: Arc) { + self.register_any_handler_internal(handler, /*include_spec*/ true); + } + + pub(crate) fn register_any_handler_without_spec(&mut self, handler: Arc) { + self.register_any_handler_internal(handler, /*include_spec*/ false); + } + + fn register_any_handler_internal( + &mut self, + handler: Arc, + include_spec: bool, + ) { let name = handler.tool_name(); if self.handlers.contains_key(&name) { error_or_panic(format!("handler for tool {name} already registered")); return; } - if let Some(spec) = handler.spec() { + if include_spec && let Some(spec) = handler.spec() { self.push_spec(spec); } - let handler: Arc = handler; self.handlers.insert(name, handler); } diff --git a/codex-rs/core/src/tools/registry_tests.rs b/codex-rs/core/src/tools/registry_tests.rs index d6acb80fb4..dc744321d1 100644 --- a/codex-rs/core/src/tools/registry_tests.rs +++ b/codex-rs/core/src/tools/registry_tests.rs @@ -63,16 +63,13 @@ fn handler_looks_up_namespaced_aliases_explicitly() { } #[test] -fn register_handler_adds_handler_and_augments_specs_for_code_mode() { - let mut builder = ToolRegistryBuilder::new(/*code_mode_enabled*/ true); +fn register_handler_adds_handler_and_spec() { + let mut builder = ToolRegistryBuilder::new(); builder.register_handler(Arc::new(GetGoalHandler)); let (specs, registry) = builder.build(); assert_eq!(specs.len(), 1); - assert_eq!( - specs[0], - codex_tools::augment_tool_spec_for_code_mode(create_get_goal_tool()) - ); + assert_eq!(specs[0], create_get_goal_tool()); assert!(registry.has_handler(&codex_tools::ToolName::plain(GET_GOAL_TOOL_NAME))); } diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index aed72186e0..4bc366bf91 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -37,7 +37,6 @@ pub struct ToolCall { pub struct ToolRouter { registry: ToolRegistry, - specs: Vec, model_visible_specs: Vec, } @@ -90,15 +89,10 @@ impl ToolRouter { Self { registry, - specs, model_visible_specs, } } - pub fn specs(&self) -> Vec { - self.specs.clone() - } - pub fn model_visible_specs(&self) -> Vec { self.model_visible_specs.clone() } diff --git a/codex-rs/core/src/tools/router_tests.rs b/codex-rs/core/src/tools/router_tests.rs index 76776b1add..4d6119cd5c 100644 --- a/codex-rs/core/src/tools/router_tests.rs +++ b/codex-rs/core/src/tools/router_tests.rs @@ -232,7 +232,7 @@ async fn tools_without_handlers_do_not_support_parallel() -> anyhow::Result<()> } #[tokio::test] -async fn model_visible_specs_filter_deferred_dynamic_tools() -> anyhow::Result<()> { +async fn specs_filter_deferred_dynamic_tools() -> anyhow::Result<()> { let (_, turn) = make_session_and_context().await; let hidden_tool = "hidden_dynamic_tool"; let visible_tool = "visible_dynamic_tool"; @@ -273,10 +273,6 @@ async fn model_visible_specs_filter_deferred_dynamic_tools() -> anyhow::Result<( }, ); - assert_eq!( - namespace_function_names(&router.specs(), "codex_app"), - vec![hidden_tool.to_string(), visible_tool.to_string()] - ); assert_eq!( namespace_function_names(&router.model_visible_specs(), "codex_app"), vec![visible_tool.to_string()] @@ -334,13 +330,6 @@ async fn extension_tool_bundles_are_model_visible_and_dispatchable() -> anyhow:: }, ); - assert!( - router - .specs() - .iter() - .any(|spec| spec.name() == "extension_echo"), - "expected extension-provided tool spec to be registered" - ); assert!( router .model_visible_specs() diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index fc21ff3375..8c54c21479 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -70,7 +70,6 @@ pub(crate) fn build_specs_with_discoverable_tools( use crate::tools::tool_search_entry::build_tool_search_entries_for_config; let mcp_tool_plan_inputs = mcp_tools.as_deref().map(map_mcp_tools_for_plan); - let deferred_mcp_tool_sources = deferred_mcp_tools.as_deref(); let default_agent_type_description = crate::agent::role::spawn_tool_spec::build(&std::collections::BTreeMap::new()); let min_wait_timeout_ms = if config.multi_agent_v2 { @@ -85,12 +84,12 @@ pub(crate) fn build_specs_with_discoverable_tools( DEFAULT_WAIT_TIMEOUT_MS.clamp(min_wait_timeout_ms, MAX_WAIT_TIMEOUT_MS); let deferred_dynamic_tools = dynamic_tools .iter() - .filter(|tool| tool.defer_loading && (config.namespace_tools || tool.namespace.is_none())) + .filter(|tool| tool.defer_loading) .cloned() .collect::>(); let tool_search_entries = build_tool_search_entries_for_config( config, - deferred_mcp_tool_sources, + deferred_mcp_tools.as_deref(), &deferred_dynamic_tools, ); let mut builder = build_tool_registry_builder( @@ -99,7 +98,7 @@ pub(crate) fn build_specs_with_discoverable_tools( mcp_tools: mcp_tool_plan_inputs .as_ref() .map(|inputs| inputs.mcp_tools.as_slice()), - deferred_mcp_tools: deferred_mcp_tool_sources, + deferred_mcp_tools: deferred_mcp_tools.as_deref(), tool_namespaces: mcp_tool_plan_inputs .as_ref() .map(|inputs| &inputs.tool_namespaces), diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index 69456a36ef..7232dddd85 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -27,6 +27,7 @@ use crate::tools::handlers::ViewImageHandler; use crate::tools::handlers::WriteStdinHandler; use crate::tools::handlers::agent_jobs::ReportAgentJobResultHandler; use crate::tools::handlers::agent_jobs::SpawnAgentsOnCsvHandler; +use crate::tools::handlers::extension_tools::extension_tool_spec; use crate::tools::handlers::multi_agents::CloseAgentHandler; use crate::tools::handlers::multi_agents::ResumeAgentHandler; use crate::tools::handlers::multi_agents::SendInputHandler; @@ -44,12 +45,11 @@ use crate::tools::handlers::view_image_spec::ViewImageToolOptions; use crate::tools::hosted_spec::WebSearchToolOptions; use crate::tools::hosted_spec::create_image_generation_tool; use crate::tools::hosted_spec::create_web_search_tool; +use crate::tools::registry::AnyToolHandler; use crate::tools::registry::ToolRegistryBuilder; use crate::tools::spec_plan_types::ToolRegistryBuildParams; use crate::tools::spec_plan_types::agent_type_description; -use codex_mcp::ToolInfo; use codex_protocol::openai_models::ConfigShellToolType; -use codex_tools::ResponsesApiNamespace; use codex_tools::ResponsesApiNamespaceTool; use codex_tools::ToolEnvironmentMode; use codex_tools::ToolName; @@ -57,12 +57,9 @@ use codex_tools::ToolSearchSource; use codex_tools::ToolSearchSourceInfo; use codex_tools::ToolSpec; use codex_tools::ToolsConfig; -use codex_tools::coalesce_loadable_tool_specs; use codex_tools::collect_code_mode_exec_prompt_tool_definitions; use codex_tools::collect_tool_search_source_infos; use codex_tools::default_namespace_description; -use codex_tools::dynamic_tool_to_loadable_tool_spec; -use codex_tools::mcp_tool_to_responses_api_tool; use std::collections::BTreeMap; use std::collections::HashSet; use std::sync::Arc; @@ -71,8 +68,21 @@ pub fn build_tool_registry_builder( config: &ToolsConfig, params: ToolRegistryBuildParams<'_>, ) -> ToolRegistryBuilder { - let mut builder = ToolRegistryBuilder::new(config.code_mode_enabled); - let exec_permission_approvals_enabled = config.exec_permission_approvals_enabled; + let mut builder = ToolRegistryBuilder::new(); + let all_deferred_tools = params + .deferred_mcp_tools + .into_iter() + .flatten() + .map(codex_mcp::ToolInfo::canonical_tool_name) + .chain( + params + .dynamic_tools + .iter() + .filter(|tool| tool.defer_loading) + .map(|tool| ToolName::new(tool.namespace.clone(), tool.name.clone())), + ) + .collect::>(); + let handlers = collect_handler_tools(config, params); if config.code_mode_enabled { let namespace_descriptions = params @@ -89,16 +99,18 @@ pub fn build_tool_registry_builder( ) }) .collect::>(); - let nested_config = config.for_code_mode_nested_tools(); - let nested_builder = build_tool_registry_builder( - &nested_config, - ToolRegistryBuildParams { - discoverable_tools: None, - ..params - }, + let mut code_mode_nested_tool_specs = handlers + .iter() + .filter_map(|handler| handler.spec()) + .collect::>(); + code_mode_nested_tool_specs.extend( + params + .extension_tool_bundles + .iter() + .filter_map(|bundle| extension_tool_spec(bundle.spec()).ok()), ); let mut enabled_tools = - collect_code_mode_exec_prompt_tool_definitions(nested_builder.specs().iter()); + collect_code_mode_exec_prompt_tool_definitions(code_mode_nested_tool_specs.iter()); enabled_tools .sort_by(|left, right| compare_code_mode_tools(left, right, &namespace_descriptions)); builder.register_handler(Arc::new(CodeModeExecuteHandler::new( @@ -106,40 +118,162 @@ pub fn build_tool_registry_builder( &enabled_tools, &namespace_descriptions, config.code_mode_only_enabled, - config.search_tool - && params - .deferred_mcp_tools - .is_some_and(|tools| !tools.is_empty()), + config.search_tool && !all_deferred_tools.is_empty(), ), + code_mode_nested_tool_specs, ))); builder.register_handler(Arc::new(CodeModeWaitHandler)); } + let mut non_deferred_specs = Vec::new(); + for handler in &handlers { + let tool_name = handler.tool_name(); + if !all_deferred_tools.contains(&tool_name) + && let Some(spec) = handler.spec() + { + non_deferred_specs.push(spec); + } + } + + if let Some(web_search_tool) = create_web_search_tool(WebSearchToolOptions { + web_search_mode: config.web_search_mode, + web_search_config: config.web_search_config.as_ref(), + web_search_tool_type: config.web_search_tool_type, + }) { + non_deferred_specs.push(web_search_tool); + } + if config.image_gen_tool { + non_deferred_specs.push(create_image_generation_tool("png")); + } + + for spec in merge_into_namespaces(non_deferred_specs) { + if !config.namespace_tools && matches!(spec, ToolSpec::Namespace(_)) { + continue; + } + let spec = if config.code_mode_enabled { + codex_tools::augment_tool_spec_for_code_mode(spec) + } else { + spec + }; + builder.push_spec(spec); + } + + for handler in handlers { + builder.register_any_handler_without_spec(handler); + } + + if config.search_tool && config.namespace_tools && !all_deferred_tools.is_empty() { + let mut search_source_infos = params + .deferred_mcp_tools + .map(|mcp_tools| { + collect_tool_search_source_infos(mcp_tools.iter().map(|tool| ToolSearchSource { + server_name: tool.server_name.as_str(), + connector_name: tool.connector_name.as_deref(), + description: tool.namespace_description.as_deref(), + })) + }) + .unwrap_or_default(); + + if params.dynamic_tools.iter().any(|tool| { + all_deferred_tools.contains(&ToolName::new(tool.namespace.clone(), tool.name.clone())) + }) { + search_source_infos.push(ToolSearchSourceInfo { + name: "Dynamic tools".to_string(), + description: Some("Tools provided by the current Codex thread.".to_string()), + }); + } + + builder.register_handler(Arc::new(ToolSearchHandler::new( + params.tool_search_entries.to_vec(), + search_source_infos, + ))); + } + + for bundle in params.extension_tool_bundles.iter().cloned() { + builder.register_tool_bundle(bundle); + } + + builder +} + +fn merge_into_namespaces(specs: Vec) -> Vec { + let mut merged_specs = Vec::with_capacity(specs.len()); + let mut namespace_indices = BTreeMap::::new(); + for spec in specs { + match spec { + ToolSpec::Namespace(mut namespace) => { + if let Some(index) = namespace_indices.get(&namespace.name).copied() { + let ToolSpec::Namespace(existing_namespace) = &mut merged_specs[index] else { + unreachable!("namespace index must point to a namespace spec"); + }; + if existing_namespace.description.trim().is_empty() + && !namespace.description.trim().is_empty() + { + existing_namespace.description = namespace.description; + } + existing_namespace.tools.append(&mut namespace.tools); + continue; + } + + namespace_indices.insert(namespace.name.clone(), merged_specs.len()); + merged_specs.push(ToolSpec::Namespace(namespace)); + } + spec => merged_specs.push(spec), + } + } + + for spec in &mut merged_specs { + let ToolSpec::Namespace(namespace) = spec else { + continue; + }; + + namespace.tools.sort_by(|left, right| match (left, right) { + ( + ResponsesApiNamespaceTool::Function(left), + ResponsesApiNamespaceTool::Function(right), + ) => left.name.cmp(&right.name), + }); + + if namespace.description.trim().is_empty() { + namespace.description = default_namespace_description(&namespace.name); + } + } + + merged_specs +} + +fn collect_handler_tools( + config: &ToolsConfig, + params: ToolRegistryBuildParams<'_>, +) -> Vec> { + let exec_permission_approvals_enabled = config.exec_permission_approvals_enabled; + let mut handlers = Vec::>::new(); + if config.environment_mode.has_environment() { let include_environment_id = matches!(config.environment_mode, ToolEnvironmentMode::Multiple); match &config.shell_type { ConfigShellToolType::Default => { - builder.register_handler(Arc::new(ShellHandler::new(ShellToolOptions { + handlers.push(Arc::new(ShellHandler::new(ShellToolOptions { exec_permission_approvals_enabled, }))); } ConfigShellToolType::Local => { - builder.register_handler(Arc::new(LocalShellHandler::new())); + handlers.push(Arc::new(LocalShellHandler::new())); } ConfigShellToolType::UnifiedExec => { - builder.register_handler(Arc::new(ExecCommandHandler::new( + handlers.push(Arc::new(ExecCommandHandler::new( ExecCommandHandlerOptions { allow_login_shell: config.allow_login_shell, exec_permission_approvals_enabled, include_environment_id, }, ))); - builder.register_handler(Arc::new(WriteStdinHandler)); + handlers.push(Arc::new(WriteStdinHandler)); } ConfigShellToolType::Disabled => {} ConfigShellToolType::ShellCommand => { - builder.register_handler(Arc::new(ShellCommandHandler::new( + handlers.push(Arc::new(ShellCommandHandler::new( ShellCommandHandlerOptions { backend_config: config.shell_command_backend, allow_login_shell: config.allow_login_shell, @@ -155,101 +289,62 @@ pub fn build_tool_registry_builder( { match &config.shell_type { ConfigShellToolType::Default => { - builder.register_handler(Arc::new(ContainerExecHandler)); - builder.register_handler(Arc::new(LocalShellHandler::default())); - builder.register_handler(Arc::new(ShellCommandHandler::from( + handlers.push(Arc::new(ContainerExecHandler)); + handlers.push(Arc::new(LocalShellHandler::default())); + handlers.push(Arc::new(ShellCommandHandler::from( config.shell_command_backend, ))); } ConfigShellToolType::Local => { - builder.register_handler(Arc::new(ShellHandler::default())); - builder.register_handler(Arc::new(ContainerExecHandler)); - builder.register_handler(Arc::new(ShellCommandHandler::from( + handlers.push(Arc::new(ShellHandler::default())); + handlers.push(Arc::new(ContainerExecHandler)); + handlers.push(Arc::new(ShellCommandHandler::from( config.shell_command_backend, ))); } ConfigShellToolType::UnifiedExec => { - builder.register_handler(Arc::new(ShellHandler::default())); - builder.register_handler(Arc::new(ContainerExecHandler)); - builder.register_handler(Arc::new(LocalShellHandler::default())); - builder.register_handler(Arc::new(ShellCommandHandler::from( + handlers.push(Arc::new(ShellHandler::default())); + handlers.push(Arc::new(ContainerExecHandler)); + handlers.push(Arc::new(LocalShellHandler::default())); + handlers.push(Arc::new(ShellCommandHandler::from( config.shell_command_backend, ))); } ConfigShellToolType::ShellCommand => { - builder.register_handler(Arc::new(ShellHandler::default())); - builder.register_handler(Arc::new(ContainerExecHandler)); - builder.register_handler(Arc::new(LocalShellHandler::default())); + handlers.push(Arc::new(ShellHandler::default())); + handlers.push(Arc::new(ContainerExecHandler)); + handlers.push(Arc::new(LocalShellHandler::default())); } ConfigShellToolType::Disabled => {} } } if params.mcp_tools.is_some() { - builder.register_handler(Arc::new(ListMcpResourcesHandler)); - builder.register_handler(Arc::new(ListMcpResourceTemplatesHandler)); - builder.register_handler(Arc::new(ReadMcpResourceHandler)); + handlers.push(Arc::new(ListMcpResourcesHandler)); + handlers.push(Arc::new(ListMcpResourceTemplatesHandler)); + handlers.push(Arc::new(ReadMcpResourceHandler)); } - builder.register_handler(Arc::new(PlanHandler)); + handlers.push(Arc::new(PlanHandler)); if config.goal_tools { - builder.register_handler(Arc::new(GetGoalHandler)); - builder.register_handler(Arc::new(CreateGoalHandler)); - builder.register_handler(Arc::new(UpdateGoalHandler)); + handlers.push(Arc::new(GetGoalHandler)); + handlers.push(Arc::new(CreateGoalHandler)); + handlers.push(Arc::new(UpdateGoalHandler)); } - builder.register_handler(Arc::new(RequestUserInputHandler { + handlers.push(Arc::new(RequestUserInputHandler { available_modes: config.request_user_input_available_modes.clone(), })); if config.request_permissions_tool_enabled { - builder.register_handler(Arc::new(RequestPermissionsHandler)); - } - - let deferred_dynamic_tools = params - .dynamic_tools - .iter() - .filter(|tool| tool.defer_loading && (config.namespace_tools || tool.namespace.is_none())) - .collect::>(); - let deferred_mcp_tools_for_search = if config.namespace_tools { - params.deferred_mcp_tools - } else { - None - }; - - if config.search_tool - && (deferred_mcp_tools_for_search.is_some() || !deferred_dynamic_tools.is_empty()) - { - let mut search_source_infos = deferred_mcp_tools_for_search - .map(|deferred_mcp_tools| { - collect_tool_search_source_infos(deferred_mcp_tools.iter().map(|tool| { - ToolSearchSource { - server_name: tool.server_name.as_str(), - connector_name: tool.connector_name.as_deref(), - description: tool.namespace_description.as_deref(), - } - })) - }) - .unwrap_or_default(); - - if !deferred_dynamic_tools.is_empty() { - search_source_infos.push(ToolSearchSourceInfo { - name: "Dynamic tools".to_string(), - description: Some("Tools provided by the current Codex thread.".to_string()), - }); - } - - builder.register_handler(Arc::new(ToolSearchHandler::new( - params.tool_search_entries.to_vec(), - search_source_infos, - ))); + handlers.push(Arc::new(RequestPermissionsHandler)); } if config.tool_suggest && let Some(discoverable_tools) = params.discoverable_tools.filter(|tools| !tools.is_empty()) { - builder.register_handler(Arc::new(RequestPluginInstallHandler::new( + handlers.push(Arc::new(RequestPluginInstallHandler::new( discoverable_tools, ))); } @@ -257,7 +352,7 @@ pub fn build_tool_registry_builder( if config.environment_mode.has_environment() && config.apply_patch_tool_type.is_some() { let include_environment_id = matches!(config.environment_mode, ToolEnvironmentMode::Multiple); - builder.register_handler(Arc::new(ApplyPatchHandler::new(include_environment_id))); + handlers.push(Arc::new(ApplyPatchHandler::new(include_environment_id))); } if config @@ -265,25 +360,13 @@ pub fn build_tool_registry_builder( .iter() .any(|tool| tool == "test_sync_tool") { - builder.register_handler(Arc::new(TestSyncHandler)); - } - - if let Some(web_search_tool) = create_web_search_tool(WebSearchToolOptions { - web_search_mode: config.web_search_mode, - web_search_config: config.web_search_config.as_ref(), - web_search_tool_type: config.web_search_tool_type, - }) { - builder.push_spec(web_search_tool); - } - - if config.image_gen_tool { - builder.push_spec(create_image_generation_tool("png")); + handlers.push(Arc::new(TestSyncHandler)); } if config.environment_mode.has_environment() { let include_environment_id = matches!(config.environment_mode, ToolEnvironmentMode::Multiple); - builder.register_handler(Arc::new(ViewImageHandler::new(ViewImageToolOptions { + handlers.push(Arc::new(ViewImageHandler::new(ViewImageToolOptions { can_request_original_image_detail: config.can_request_original_image_detail, include_environment_id, }))); @@ -293,7 +376,7 @@ pub fn build_tool_registry_builder( if config.multi_agent_v2 { let agent_type_description = agent_type_description(config, params.default_agent_type_description); - builder.register_handler(Arc::new(SpawnAgentHandlerV2::new(SpawnAgentToolOptions { + handlers.push(Arc::new(SpawnAgentHandlerV2::new(SpawnAgentToolOptions { available_models: config.available_models.clone(), agent_type_description, hide_agent_type_model_reasoning: config.hide_spawn_agent_metadata, @@ -301,17 +384,17 @@ pub fn build_tool_registry_builder( usage_hint_text: config.spawn_agent_usage_hint_text.clone(), max_concurrent_threads_per_session: config.max_concurrent_threads_per_session, }))); - builder.register_handler(Arc::new(SendMessageHandlerV2)); - builder.register_handler(Arc::new(FollowupTaskHandlerV2)); - builder.register_handler(Arc::new(WaitAgentHandlerV2::new( + handlers.push(Arc::new(SendMessageHandlerV2)); + handlers.push(Arc::new(FollowupTaskHandlerV2)); + handlers.push(Arc::new(WaitAgentHandlerV2::new( params.wait_agent_timeouts, ))); - builder.register_handler(Arc::new(CloseAgentHandlerV2)); - builder.register_handler(Arc::new(ListAgentsHandlerV2)); + handlers.push(Arc::new(CloseAgentHandlerV2)); + handlers.push(Arc::new(ListAgentsHandlerV2)); } else { let agent_type_description = agent_type_description(config, params.default_agent_type_description); - builder.register_handler(Arc::new(SpawnAgentHandler::new(SpawnAgentToolOptions { + handlers.push(Arc::new(SpawnAgentHandler::new(SpawnAgentToolOptions { available_models: config.available_models.clone(), agent_type_description, hide_agent_type_model_reasoning: config.hide_spawn_agent_metadata, @@ -319,122 +402,45 @@ pub fn build_tool_registry_builder( usage_hint_text: config.spawn_agent_usage_hint_text.clone(), max_concurrent_threads_per_session: config.max_concurrent_threads_per_session, }))); - builder.register_handler(Arc::new(SendInputHandler)); - builder.register_handler(Arc::new(ResumeAgentHandler)); - builder.register_handler(Arc::new(WaitAgentHandler::new(params.wait_agent_timeouts))); - builder.register_handler(Arc::new(CloseAgentHandler)); + handlers.push(Arc::new(SendInputHandler)); + handlers.push(Arc::new(ResumeAgentHandler)); + handlers.push(Arc::new(WaitAgentHandler::new(params.wait_agent_timeouts))); + handlers.push(Arc::new(CloseAgentHandler)); } } if config.agent_jobs_tools { - builder.register_handler(Arc::new(SpawnAgentsOnCsvHandler)); + handlers.push(Arc::new(SpawnAgentsOnCsvHandler)); if config.agent_jobs_worker_tools { - builder.register_handler(Arc::new(ReportAgentJobResultHandler)); + handlers.push(Arc::new(ReportAgentJobResultHandler)); } } if let Some(mcp_tools) = params.mcp_tools { - let mut entries = mcp_tools - .iter() - .map(|tool| (tool.canonical_tool_name(), tool)) - .collect::>(); - entries.sort_by(|(left, _), (right, _)| left.cmp(right)); - let mut namespace_entries = BTreeMap::new(); - - for (tool_name, tool) in entries { - let Some(namespace) = tool_name.namespace.as_ref() else { - tracing::error!("Skipping MCP tool `{tool_name}`: MCP tools must be namespaced"); - continue; - }; - namespace_entries - .entry(namespace.clone()) - .or_insert_with(Vec::new) - .push((tool_name, tool)); - } - - for (namespace, mut entries) in namespace_entries { - entries.sort_by_key(|(tool_name, _)| tool_name.name.clone()); - let tool_namespace = params - .tool_namespaces - .and_then(|namespaces| namespaces.get(&namespace)); - let description = tool_namespace - .and_then(|namespace| namespace.description.as_deref()) - .map(str::trim) - .filter(|description| !description.is_empty()) - .map(str::to_string) - .unwrap_or_else(|| { - let namespace_name = tool_namespace - .map(|namespace| namespace.name.as_str()) - .unwrap_or(namespace.as_str()); - default_namespace_description(namespace_name) - }); - let mut tools = Vec::new(); - for (tool_name, tool) in entries { - match mcp_tool_to_responses_api_tool(&tool_name, &tool.tool) { - Ok(converted_tool) => { - tools.push(ResponsesApiNamespaceTool::Function(converted_tool)); - builder.register_handler(Arc::new(McpHandler::new(tool.clone()))); - } - Err(error) => { - tracing::error!( - "Failed to convert `{tool_name}` MCP tool to OpenAI tool: {error:?}" - ); - } - } - } - - if config.namespace_tools && !tools.is_empty() { - builder.push_spec(ToolSpec::Namespace(ResponsesApiNamespace { - name: namespace, - description, - tools, - })); - } - } - } - - let mut dynamic_tool_specs = Vec::new(); - for tool in params.dynamic_tools { - match dynamic_tool_to_loadable_tool_spec(tool) { - Ok(loadable_tool) => { - let handler_name = ToolName::new(tool.namespace.clone(), tool.name.clone()); - dynamic_tool_specs.push(loadable_tool); - builder.register_handler(Arc::new(DynamicToolHandler::new(handler_name))); - } - Err(error) => { - tracing::error!( - "Failed to convert dynamic tool {:?} to OpenAI tool: {error:?}", - tool.name - ); - } - } - } - for spec in coalesce_loadable_tool_specs(dynamic_tool_specs) { - let spec = spec.into(); - if config.namespace_tools || !matches!(spec, ToolSpec::Namespace(_)) { - builder.push_spec(spec); + for tool in mcp_tools { + handlers.push(Arc::new(McpHandler::new(tool.clone()))); } } if let Some(deferred_mcp_tools) = params.deferred_mcp_tools { - let directly_registered_mcp_tools = params - .mcp_tools - .into_iter() - .flatten() - .map(ToolInfo::canonical_tool_name) - .collect::>(); for tool in deferred_mcp_tools { - if !directly_registered_mcp_tools.contains(&tool.canonical_tool_name()) { - builder.register_handler(Arc::new(McpHandler::new(tool.clone()))); - } + handlers.push(Arc::new(McpHandler::new(tool.clone()))); } } - for bundle in params.extension_tool_bundles.iter().cloned() { - builder.register_tool_bundle(bundle); + for tool in params.dynamic_tools { + let Some(handler) = DynamicToolHandler::new(tool).map(Arc::new) else { + tracing::error!( + "Failed to convert dynamic tool {:?} to OpenAI tool", + tool.name + ); + continue; + }; + + handlers.push(handler); } - builder + handlers } fn compare_code_mode_tools( diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index a4132391a2..7c876d3738 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -1579,20 +1579,7 @@ fn search_tool_description_lists_each_mcp_source_once() { let (tools, registry) = build_specs( &tools_config, - Some(HashMap::from([ - ( - ToolName::namespaced("mcp__codex_apps__calendar", "_create_event"), - mcp_tool( - "calendar_create_event", - "Create calendar event", - serde_json::json!({"type": "object"}), - ), - ), - ( - ToolName::namespaced("mcp__rmcp__", "echo"), - mcp_tool("echo", "Echo", serde_json::json!({"type": "object"})), - ), - ])), + /*mcp_tools*/ None, Some(vec![ deferred_mcp_tool( "_create_event", @@ -1810,32 +1797,15 @@ fn search_tool_registers_for_deferred_dynamic_tools() { panic!("expected tool_search tool"); }; assert!(description.contains("- Dynamic tools: Tools provided by the current Codex thread.")); - assert_contains_tool_names(&tools, &[TOOL_SEARCH_TOOL_NAME, "codex_app"]); - assert_eq!( - tools - .iter() - .filter(|tool| tool.name() == "codex_app") - .count(), - 1 - ); - assert_eq!( - namespace_function_names(&tools, "codex_app"), - vec![ - "automation_update".to_string(), - "automation_list".to_string() - ] - ); - for tool_name in ["automation_update", "automation_list"] { - let dynamic_tool = find_namespace_function_tool(&tools, "codex_app", tool_name); - assert_eq!(dynamic_tool.defer_loading, Some(true)); - } + assert_contains_tool_names(&tools, &[TOOL_SEARCH_TOOL_NAME]); + assert_lacks_tool_name(&tools, "codex_app"); assert!(registry.has_handler(&ToolName::plain(TOOL_SEARCH_TOOL_NAME))); assert!(registry.has_handler(&ToolName::namespaced("codex_app", "automation_update"))); assert!(registry.has_handler(&ToolName::namespaced("codex_app", "automation_list"))); } #[test] -fn search_tool_keeps_plain_deferred_dynamic_tools_when_namespace_tools_are_disabled() { +fn search_tool_is_hidden_for_deferred_dynamic_tools_when_namespace_tools_are_disabled() { let model_info = search_capable_model_info(); let mut features = Features::with_defaults(); features.enable(Feature::ToolSearch); @@ -1875,9 +1845,12 @@ fn search_tool_keeps_plain_deferred_dynamic_tools_when_namespace_tools_are_disab &dynamic_tools, ); - assert_contains_tool_names(&tools, &[TOOL_SEARCH_TOOL_NAME, "plain_dynamic"]); + assert_lacks_tool_name(&tools, TOOL_SEARCH_TOOL_NAME); assert_lacks_tool_name(&tools, "codex_app"); - assert!(registry.has_handler(&ToolName::plain(TOOL_SEARCH_TOOL_NAME))); + assert_lacks_tool_name(&tools, "plain_dynamic"); + assert!(!registry.has_handler(&ToolName::plain(TOOL_SEARCH_TOOL_NAME))); + assert!(registry.has_handler(&ToolName::namespaced("codex_app", "automation_update"))); + assert!(registry.has_handler(&ToolName::plain("plain_dynamic"))); } #[test] diff --git a/codex-rs/core/src/tools/spec_tests.rs b/codex-rs/core/src/tools/spec_tests.rs index 0086542859..52f92b281d 100644 --- a/codex-rs/core/src/tools/spec_tests.rs +++ b/codex-rs/core/src/tools/spec_tests.rs @@ -365,7 +365,7 @@ async fn assert_model_tools( .iter() .map(ToolSpec::name) .collect::>(); - assert_eq!(&tool_names, &expected_tools,); + assert_eq!(&tool_names, &expected_tools); } async fn assert_default_model_tools( @@ -396,14 +396,14 @@ async fn test_build_specs_gpt5_codex_default() { "update_plan", "request_user_input", "apply_patch", - "web_search", - "image_generation", "view_image", "spawn_agent", "send_input", "resume_agent", "wait_agent", "close_agent", + "web_search", + "image_generation", ], ) .await; @@ -421,14 +421,14 @@ async fn test_build_specs_gpt51_codex_default() { "update_plan", "request_user_input", "apply_patch", - "web_search", - "image_generation", "view_image", "spawn_agent", "send_input", "resume_agent", "wait_agent", "close_agent", + "web_search", + "image_generation", ], ) .await; @@ -448,14 +448,14 @@ async fn test_build_specs_gpt5_codex_unified_exec_web_search() { "update_plan", "request_user_input", "apply_patch", - "web_search", - "image_generation", "view_image", "spawn_agent", "send_input", "resume_agent", "wait_agent", "close_agent", + "web_search", + "image_generation", ], ) .await; @@ -475,14 +475,14 @@ async fn test_build_specs_gpt51_codex_unified_exec_web_search() { "update_plan", "request_user_input", "apply_patch", - "web_search", - "image_generation", "view_image", "spawn_agent", "send_input", "resume_agent", "wait_agent", "close_agent", + "web_search", + "image_generation", ], ) .await; @@ -500,14 +500,14 @@ async fn test_gpt_5_1_codex_max_defaults() { "update_plan", "request_user_input", "apply_patch", - "web_search", - "image_generation", "view_image", "spawn_agent", "send_input", "resume_agent", "wait_agent", "close_agent", + "web_search", + "image_generation", ], ) .await; @@ -525,14 +525,14 @@ async fn test_codex_5_1_mini_defaults() { "update_plan", "request_user_input", "apply_patch", - "web_search", - "image_generation", "view_image", "spawn_agent", "send_input", "resume_agent", "wait_agent", "close_agent", + "web_search", + "image_generation", ], ) .await; @@ -550,14 +550,14 @@ async fn test_gpt_5_defaults() { "update_plan", "request_user_input", "apply_patch", - "web_search", - "image_generation", "view_image", "spawn_agent", "send_input", "resume_agent", "wait_agent", "close_agent", + "web_search", + "image_generation", ], ) .await; @@ -575,14 +575,14 @@ async fn test_gpt_5_1_defaults() { "update_plan", "request_user_input", "apply_patch", - "web_search", - "image_generation", "view_image", "spawn_agent", "send_input", "resume_agent", "wait_agent", "close_agent", + "web_search", + "image_generation", ], ) .await; @@ -602,14 +602,14 @@ async fn test_gpt_5_1_codex_max_unified_exec_web_search() { "update_plan", "request_user_input", "apply_patch", - "web_search", - "image_generation", "view_image", "spawn_agent", "send_input", "resume_agent", "wait_agent", "close_agent", + "web_search", + "image_generation", ], ) .await; @@ -841,7 +841,7 @@ async fn request_plugin_install_requires_apps_and_plugins_features() { } #[tokio::test] -async fn search_tool_description_handles_no_enabled_mcp_tools() { +async fn search_tool_is_hidden_without_deferred_tools() { let model_info = search_capable_model_info().await; let mut features = Features::with_defaults(); features.enable(Feature::Apps); @@ -865,13 +865,11 @@ async fn search_tool_description_handles_no_enabled_mcp_tools() { &[], ) .build(); - let search_tool = find_tool(&tools, TOOL_SEARCH_TOOL_NAME); - let ToolSpec::ToolSearch { description, .. } = search_tool else { - panic!("expected tool_search tool"); - }; - - assert!(description.contains("None currently enabled.")); - assert!(!description.contains("{{source_descriptions}}")); + assert!( + !tools + .iter() + .any(|tool| tool.name() == TOOL_SEARCH_TOOL_NAME) + ); } #[tokio::test] @@ -1003,7 +1001,7 @@ async fn search_tool_registers_namespaced_mcp_tool_aliases() { } #[tokio::test] -async fn tool_search_entries_skip_namespace_outputs_when_namespace_tools_are_disabled() { +async fn tool_search_entries_skip_mcp_outputs_when_namespace_tools_are_disabled() { let model_info = search_capable_model_info().await; let mut features = Features::with_defaults(); features.enable(Feature::ToolSearch); @@ -1048,11 +1046,15 @@ async fn tool_search_entries_skip_namespace_outputs_when_namespace_tools_are_dis .map(|entry| entry.output) .collect::>(); - assert_eq!(outputs.len(), 1); - match &outputs[0] { - LoadableToolSpec::Function(tool) => assert_eq!(tool.name, "plain_dynamic"), - LoadableToolSpec::Namespace(_) => panic!("namespace tool_search output should be hidden"), - } + assert_eq!(outputs.len(), 2); + assert!(outputs.iter().any(|output| match output { + LoadableToolSpec::Namespace(namespace) => namespace.name == "codex_app", + LoadableToolSpec::Function(_) => false, + })); + assert!(outputs.iter().any(|output| match output { + LoadableToolSpec::Function(tool) => tool.name == "plain_dynamic", + LoadableToolSpec::Namespace(_) => false, + })); } #[tokio::test] diff --git a/codex-rs/core/src/tools/tool_search_entry.rs b/codex-rs/core/src/tools/tool_search_entry.rs index 6a8ccd6402..cbaa4d87c9 100644 --- a/codex-rs/core/src/tools/tool_search_entry.rs +++ b/codex-rs/core/src/tools/tool_search_entry.rs @@ -63,11 +63,7 @@ pub(crate) fn build_tool_search_entries_for_config( } else { None }; - let dynamic_tools = dynamic_tools - .iter() - .filter(|tool| config.namespace_tools || tool.namespace.is_none()) - .cloned() - .collect::>(); + let dynamic_tools = dynamic_tools.to_vec(); build_tool_search_entries(mcp_tools, &dynamic_tools) } diff --git a/codex-rs/core/tests/suite/prompt_caching.rs b/codex-rs/core/tests/suite/prompt_caching.rs index b81bb06bb9..5b861ab219 100644 --- a/codex-rs/core/tests/suite/prompt_caching.rs +++ b/codex-rs/core/tests/suite/prompt_caching.rs @@ -179,13 +179,13 @@ async fn prompt_tools_are_consistent_across_requests() -> anyhow::Result<()> { "update_plan", "request_user_input", "apply_patch", - "web_search", "view_image", "spawn_agent", "send_input", "resume_agent", "wait_agent", "close_agent", + "web_search", ]); let body0 = req1.single_request().body_json();