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`
This commit is contained in:
pakrym-oai authored and GitHub committed 2026-05-12 23:34:37 +00:00
1 parent 8123bddb16
commit 960d42ddae
8 files changed
+71 -77

No files matched your search

+4 -7
View File
@@ -42,20 +42,16 @@ pub(super) fn tool_callback(
let promise = resolver.get_promise(scope); let promise = resolver.get_promise(scope);
let resolver = v8::Global::new(scope, resolver); let resolver = v8::Global::new(scope, resolver);
let tool_name = { let (tool_name, tool_kind) = {
let Some(state) = scope.get_slot::<RuntimeState>() else { let Some(state) = scope.get_slot::<RuntimeState>() else {
throw_type_error(scope, "runtime state unavailable"); throw_type_error(scope, "runtime state unavailable");
return; return;
}; };
let Some(tool_name) = state let Some(tool) = state.enabled_tools.get(tool_index) else {
.enabled_tools
.get(tool_index)
.map(|tool| tool.tool_name.clone())
else {
throw_type_error(scope, "tool callback data is out of range"); throw_type_error(scope, "tool callback data is out of range");
return; return;
}; };
tool_name (tool.tool_name.clone(), tool.kind)
}; };
let Some(state) = scope.get_slot_mut::<RuntimeState>() else { let Some(state) = scope.get_slot_mut::<RuntimeState>() else {
@@ -69,6 +65,7 @@ pub(super) fn tool_callback(
let _ = event_tx.send(RuntimeEvent::ToolCall { let _ = event_tx.send(RuntimeEvent::ToolCall {
id, id,
name: tool_name, name: tool_name,
kind: tool_kind,
input, input,
}); });
retval.set(promise.into()); retval.set(promise.into());
+3
View File
@@ -14,6 +14,7 @@ use serde::Serialize;
use serde_json::Value as JsonValue; use serde_json::Value as JsonValue;
use tokio::sync::mpsc; use tokio::sync::mpsc;
use crate::description::CodeModeToolKind;
use crate::description::EnabledToolMetadata; use crate::description::EnabledToolMetadata;
use crate::description::ToolDefinition; use crate::description::ToolDefinition;
use crate::description::enabled_tool_metadata; use crate::description::enabled_tool_metadata;
@@ -97,6 +98,7 @@ pub struct CodeModeNestedToolCall {
pub cell_id: String, pub cell_id: String,
pub runtime_tool_call_id: String, pub runtime_tool_call_id: String,
pub tool_name: ToolName, pub tool_name: ToolName,
pub tool_kind: CodeModeToolKind,
pub input: Option<JsonValue>, pub input: Option<JsonValue>,
} }
@@ -126,6 +128,7 @@ pub(crate) enum RuntimeEvent {
ToolCall { ToolCall {
id: String, id: String,
name: ToolName, name: ToolName,
kind: CodeModeToolKind,
input: Option<JsonValue>, input: Option<JsonValue>,
}, },
Notify { Notify {
+7 -1
View File
@@ -385,11 +385,17 @@ async fn run_session_control(
text, text,
}).await; }).await;
} }
RuntimeEvent::ToolCall { id, name, input } => { RuntimeEvent::ToolCall {
id,
name,
kind,
input,
} => {
let tool_call = CodeModeNestedToolCall { let tool_call = CodeModeNestedToolCall {
cell_id: cell_id.clone(), cell_id: cell_id.clone(),
runtime_tool_call_id: id, runtime_tool_call_id: id,
tool_name: name, tool_name: name,
tool_kind: kind,
input, input,
}; };
let _ = inner let _ = inner
+53 -32
View File
@@ -8,6 +8,7 @@ use std::sync::Arc;
use std::time::Duration; use std::time::Duration;
use codex_code_mode::CodeModeNestedToolCall; use codex_code_mode::CodeModeNestedToolCall;
use codex_code_mode::CodeModeToolKind;
use codex_code_mode::CodeModeTurnHost; use codex_code_mode::CodeModeTurnHost;
use codex_code_mode::RuntimeResponse; use codex_code_mode::RuntimeResponse;
use codex_protocol::models::FunctionCallOutputContentItem; use codex_protocol::models::FunctionCallOutputContentItem;
@@ -33,7 +34,6 @@ use crate::tools::router::extension_tool_bundles;
use crate::unified_exec::resolve_max_tokens; use crate::unified_exec::resolve_max_tokens;
use codex_features::Feature; use codex_features::Feature;
use codex_tools::ToolName; use codex_tools::ToolName;
use codex_tools::ToolSpec;
use codex_tools::collect_code_mode_tool_definitions; use codex_tools::collect_code_mode_tool_definitions;
use codex_utils_output_truncation::TruncationPolicy; use codex_utils_output_truncation::TruncationPolicy;
use codex_utils_output_truncation::formatted_truncate_text_content_items_with_policy; use codex_utils_output_truncation::formatted_truncate_text_content_items_with_policy;
@@ -300,6 +300,7 @@ async fn call_nested_tool(
cell_id, cell_id,
runtime_tool_call_id, runtime_tool_call_id,
tool_name, tool_name,
tool_kind,
input, input,
} = invocation; } = invocation;
if is_exec_tool_name(&tool_name) { if is_exec_tool_name(&tool_name) {
@@ -308,11 +309,10 @@ async fn call_nested_tool(
))); )));
} }
let payload = let payload = match build_nested_tool_payload(tool_kind, &tool_name, input) {
match build_nested_tool_payload(tool_runtime.find_spec(&tool_name), &tool_name, input) { Ok(payload) => payload,
Ok(payload) => payload, Err(error) => return Err(FunctionCallError::RespondToModel(error)),
Err(error) => return Err(FunctionCallError::RespondToModel(error)), };
};
let call = ToolCall { let call = ToolCall {
tool_name, tool_name,
@@ -332,36 +332,14 @@ async fn call_nested_tool(
Ok(result.code_mode_result()) 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<ToolSpec>,
tool_name: &ToolName,
) -> Result<codex_code_mode::CodeModeToolKind, String> {
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( fn build_nested_tool_payload(
spec: Option<ToolSpec>, tool_kind: CodeModeToolKind,
tool_name: &ToolName, tool_name: &ToolName,
input: Option<JsonValue>, input: Option<JsonValue>,
) -> Result<ToolPayload, String> { ) -> Result<ToolPayload, String> {
let actual_kind = tool_kind_for_name(spec, tool_name)?; match tool_kind {
match actual_kind { CodeModeToolKind::Function => build_function_tool_payload(tool_name, input),
codex_code_mode::CodeModeToolKind::Function => { CodeModeToolKind::Freeform => build_freeform_tool_payload(tool_name, input),
build_function_tool_payload(tool_name, input)
}
codex_code_mode::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")), _ => 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:?}"),
}
}
}
-5
View File
@@ -22,7 +22,6 @@ use crate::tools::router::ToolCallSource;
use crate::tools::router::ToolRouter; use crate::tools::router::ToolRouter;
use codex_protocol::error::CodexErr; use codex_protocol::error::CodexErr;
use codex_protocol::models::ResponseInputItem; use codex_protocol::models::ResponseInputItem;
use codex_tools::ToolSpec;
#[derive(Clone)] #[derive(Clone)]
pub(crate) struct ToolCallRuntime { pub(crate) struct ToolCallRuntime {
@@ -49,10 +48,6 @@ impl ToolCallRuntime {
} }
} }
pub(crate) fn find_spec(&self, tool_name: &codex_tools::ToolName) -> Option<ToolSpec> {
self.router.find_spec(tool_name)
}
pub(crate) fn create_diff_consumer( pub(crate) fn create_diff_consumer(
&self, &self,
tool_name: &codex_tools::ToolName, tool_name: &codex_tools::ToolName,
-25
View File
@@ -103,31 +103,6 @@ impl ToolRouter {
self.model_visible_specs.clone() self.model_visible_specs.clone()
} }
pub fn find_spec(&self, tool_name: &ToolName) -> Option<ToolSpec> {
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( pub(crate) fn create_diff_consumer(
&self, &self,
tool_name: &ToolName, tool_name: &ToolName,
+3 -7
View File
@@ -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!( assert_eq!(
namespace_function_names(&router.specs(), "codex_app"), namespace_function_names(&router.specs(), "codex_app"),
vec![hidden_tool.to_string(), visible_tool.to_string()] 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!( assert!(
router router
.find_spec(&ToolName::plain("extension_echo")) .specs()
.is_some(), .iter()
.any(|spec| spec.name() == "extension_echo"),
"expected extension-provided tool spec to be registered" "expected extension-provided tool spec to be registered"
); );
assert!( assert!(
@@ -401,6 +401,7 @@ async fn completed_plan_table_tail_skips_provisional_history_insert() {
} }
#[tokio::test] #[tokio::test]
#[cfg_attr(target_os = "windows", ignore = "disabled on windows")]
async fn configured_pet_load_is_deferred_until_after_construction() { async fn configured_pet_load_is_deferred_until_after_construction() {
let (tx_raw, mut rx) = unbounded_channel::<AppEvent>(); let (tx_raw, mut rx) = unbounded_channel::<AppEvent>();
let tx = AppEventSender::new(tx_raw); let tx = AppEventSender::new(tx_raw);