From 960d42ddae06242093d153da5b22f932665bc440 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Tue, 12 May 2026 16:34:37 -0700 Subject: [PATCH] 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);