From c54da23b408e6210343a65a8e3ee2ccf0b41fcc8 Mon Sep 17 00:00:00 2001 From: ashwinnathan-openai Date: Wed, 11 Mar 2026 00:00:22 -0700 Subject: [PATCH] refactor --- codex-rs/Cargo.lock | 2 + codex-rs/core/Cargo.toml | 2 + codex-rs/core/src/tools/registry.rs | 112 ++++++++++++++++++++++++++-- codex-rs/core/src/tools/spec.rs | 57 +++++--------- 4 files changed, 127 insertions(+), 46 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 64abfaac3c..9ac4cc945a 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1877,6 +1877,8 @@ dependencies = [ "sha2", "shlex", "similar", + "strum 0.27.2", + "strum_macros 0.28.0", "tempfile", "test-case", "test-log", diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index e406586524..e0e0c43fc8 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -92,6 +92,8 @@ similar = { workspace = true } tempfile = { workspace = true } test-log = { workspace = true } thiserror = { workspace = true } +strum = { workspace = true } +strum_macros = { workspace = true } time = { workspace = true, features = [ "formatting", "parsing", diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index eed13c7465..460a89d930 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -12,6 +12,7 @@ use crate::sandbox_tags::sandbox_tag; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolOutput; use crate::tools::context::ToolPayload; +use crate::tools::handlers::SEARCH_TOOL_BM25_TOOL_NAME; use async_trait::async_trait; use codex_hooks::HookEvent; use codex_hooks::HookEventAfterToolUse; @@ -21,9 +22,93 @@ use codex_hooks::HookToolInput; use codex_hooks::HookToolInputLocalShell; use codex_hooks::HookToolKind; use codex_protocol::models::ResponseInputItem; +use codex_protocol::models::VIEW_IMAGE_TOOL_NAME; use codex_utils_readiness::Readiness; +use strum::IntoEnumIterator; +use strum_macros::EnumIter; use tracing::warn; +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, EnumIter)] +pub(crate) enum BuiltinToolKey { + CodeMode, + ExecCommand, + WriteStdin, + ListMcpResources, + ListMcpResourceTemplates, + ReadMcpResource, + UpdatePlan, + JsRepl, + JsReplReset, + RequestUserInput, + RequestPermissions, + SearchToolBm25, + ApplyPatch, + GrepFiles, + ReadFile, + ListDir, + TestSyncTool, + WebSearch, + ImageGeneration, + ViewImage, + Artifacts, + SpawnAgent, + SendInput, + ResumeAgent, + Wait, + CloseAgent, + SpawnAgentsOnCsv, + ReportAgentJobResult, +} + +impl BuiltinToolKey { + pub(crate) const fn invocation_names(self) -> &'static [&'static str] { + match self { + Self::CodeMode => &["code_mode"], + Self::ExecCommand => &[ + "exec_command", + "shell", + "container.exec", + "local_shell", + "shell_command", + ], + Self::WriteStdin => &["write_stdin"], + Self::ListMcpResources => &["list_mcp_resources"], + Self::ListMcpResourceTemplates => &["list_mcp_resource_templates"], + Self::ReadMcpResource => &["read_mcp_resource"], + Self::UpdatePlan => &["update_plan"], + Self::JsRepl => &["js_repl"], + Self::JsReplReset => &["js_repl_reset"], + Self::RequestUserInput => &["request_user_input"], + Self::RequestPermissions => &["request_permissions"], + Self::SearchToolBm25 => &[SEARCH_TOOL_BM25_TOOL_NAME], + Self::ApplyPatch => &["apply_patch"], + Self::GrepFiles => &["grep_files"], + Self::ReadFile => &["read_file"], + Self::ListDir => &["list_dir"], + Self::TestSyncTool => &["test_sync_tool"], + Self::WebSearch => &["web_search"], + Self::ImageGeneration => &["image_generation"], + Self::ViewImage => &[VIEW_IMAGE_TOOL_NAME], + Self::Artifacts => &["artifacts"], + Self::SpawnAgent => &["spawn_agent"], + Self::SendInput => &["send_input"], + Self::ResumeAgent => &["resume_agent"], + Self::Wait => &["wait"], + Self::CloseAgent => &["close_agent"], + Self::SpawnAgentsOnCsv => &["spawn_agents_on_csv"], + Self::ReportAgentJobResult => &["report_agent_job_result"], + } + } +} + +pub(crate) fn builtin_tool_key(name: &str) -> Option { + BuiltinToolKey::iter().find(|key| key.invocation_names().contains(&name)) +} + +pub(crate) fn builtin_tool_invocation_names() -> impl Iterator { + BuiltinToolKey::iter().flat_map(|key| key.invocation_names().iter().copied()) +} + #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] pub enum ToolKind { Function, @@ -346,15 +431,17 @@ impl ToolRegistryBuilder { } } - pub fn push_spec(&mut self, spec: ToolSpec) { - self.push_spec_with_parallel_support(spec, false); + pub fn push_spec(&mut self, key: BuiltinToolKey, spec: ToolSpec) { + self.push_spec_with_parallel_support(key, spec, false); } pub fn push_spec_with_parallel_support( &mut self, + key: BuiltinToolKey, spec: ToolSpec, supports_parallel_tool_calls: bool, ) { + debug_assert_eq!(builtin_tool_key(spec.name()), Some(key)); self.specs.push(ConfiguredToolSpec::new( spec, supports_parallel_tool_calls, @@ -378,13 +465,26 @@ impl ToolRegistryBuilder { )); } - pub fn register_handler(&mut self, name: impl Into, handler: Arc) + pub fn register_builtin_handler(&mut self, key: BuiltinToolKey, handler: Arc) where H: ToolHandler + 'static, { - let name = name.into(); - self.builtin_handler_names.insert(name.clone()); - self.register_external_handler(name, handler); + self.register_builtin_handler_for_names(key, key.invocation_names(), handler); + } + + pub fn register_builtin_handler_for_names( + &mut self, + key: BuiltinToolKey, + names: &[&str], + handler: Arc, + ) where + H: ToolHandler + 'static, + { + for name in names { + debug_assert_eq!(builtin_tool_key(name), Some(key)); + self.builtin_handler_names.insert((*name).to_string()); + self.register_external_handler(*name, handler.clone()); + } } pub fn register_external_handler(&mut self, name: impl Into, handler: Arc) diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index 8801526c71..19475457d8 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -20,7 +20,9 @@ use crate::tools::handlers::multi_agents::MAX_WAIT_TIMEOUT_MS; use crate::tools::handlers::multi_agents::MIN_WAIT_TIMEOUT_MS; use crate::tools::handlers::request_permissions_tool_description; use crate::tools::handlers::request_user_input_tool_description; +use crate::tools::registry::BuiltinToolKey; use crate::tools::registry::ToolRegistryBuilder; +use crate::tools::registry::builtin_tool_invocation_names; use codex_protocol::config_types::WebSearchConfig; use codex_protocol::config_types::WebSearchMode; use codex_protocol::dynamic_tools::DynamicToolSpec; @@ -273,41 +275,10 @@ impl ToolsConfig { } } -fn builtin_tool_names() -> BTreeSet<&'static str> { - BTreeSet::from([ - "apply_patch", - "artifacts", - "close_agent", - "code_mode", - "container.exec", - "exec_command", - "grep_files", - "image_generation", - "js_repl", - "js_repl_reset", - "list_dir", - "list_mcp_resource_templates", - "list_mcp_resources", - "local_shell", - "read_file", - "read_mcp_resource", - "report_agent_job_result", - "request_permissions", - "request_user_input", - "resume_agent", - "search_tool_bm25", - "send_input", - "shell", - "shell_command", - "spawn_agent", - "spawn_agents_on_csv", - "test_sync_tool", - "update_plan", - "view_image", - "wait", - "web_search", - "write_stdin", - ]) +fn builtin_tool_names() -> BTreeSet { + builtin_tool_invocation_names() + .map(str::to_string) + .collect() } pub fn validate_builtin_tools_request(requested_builtin_tools: &[String]) -> Result<(), String> { @@ -2224,10 +2195,16 @@ pub(crate) fn build_specs( if config.shell_type != ConfigShellToolType::Disabled { // Always register shell aliases so older prompts remain compatible. - builder.register_handler("shell", shell_handler.clone()); - builder.register_handler("container.exec", shell_handler.clone()); - builder.register_handler("local_shell", shell_handler); - builder.register_handler("shell_command", shell_command_handler); + builder.register_builtin_handler_for_names( + BuiltinToolKey::ExecCommand, + &["shell", "container.exec", "local_shell"], + shell_handler, + ); + builder.register_builtin_handler_for_names( + BuiltinToolKey::ExecCommand, + &["shell_command"], + shell_command_handler, + ); } if mcp_tools.is_some() { @@ -2331,7 +2308,7 @@ pub(crate) fn build_specs( ); } } - builder.register_handler("apply_patch", apply_patch_handler); + builder.register_builtin_handler(BuiltinToolKey::ApplyPatch, apply_patch_handler); } if config