diff --git a/codex-rs/codex-mcp/src/mcp_connection_manager.rs b/codex-rs/codex-mcp/src/mcp_connection_manager.rs index 5c0ede3cba0..9021f7d7a9f 100644 --- a/codex-rs/codex-mcp/src/mcp_connection_manager.rs +++ b/codex-rs/codex-mcp/src/mcp_connection_manager.rs @@ -1191,11 +1191,14 @@ impl McpConnectionManager { .with_context(|| format!("resources/read failed for `{server}` ({uri})")) } - pub async fn parse_tool_name(&self, tool_name: &str) -> Option<(String, String)> { - self.list_all_tools() - .await - .get(tool_name) - .map(|tool| (tool.server_name.clone(), tool.tool.name.to_string())) + pub async fn resolve_tool_info(&self, name: &str, namespace: Option<&str>) -> Option { + let qualified_name = match namespace { + Some(namespace) if name.starts_with(namespace) => name.to_string(), + Some(namespace) => format!("{namespace}{name}"), + None => name.to_string(), + }; + + self.list_all_tools().await.get(&qualified_name).cloned() } pub async fn notify_sandbox_state_change(&self, sandbox_state: &SandboxState) -> Result<()> { diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index cd6f1e84fdd..f44a79528f7 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -77,6 +77,7 @@ use codex_login::auth_env_telemetry::collect_auth_env_telemetry; use codex_login::default_client::originator; use codex_mcp::McpConnectionManager; use codex_mcp::SandboxState; +use codex_mcp::ToolInfo; use codex_mcp::codex_apps_tools_cache_key; #[cfg(test)] use codex_models_manager::collaboration_mode_presets::CollaborationModesConfig; @@ -4412,25 +4413,16 @@ impl Session { .await } - pub(crate) async fn parse_mcp_tool_name( + pub(crate) async fn resolve_mcp_tool_info( &self, name: &str, - namespace: &Option, - ) -> Option<(String, String)> { - let tool_name = if let Some(namespace) = namespace { - if name.starts_with(namespace.as_str()) { - name - } else { - &format!("{namespace}{name}") - } - } else { - name - }; + namespace: Option<&str>, + ) -> Option { self.services .mcp_connection_manager .read() .await - .parse_tool_name(tool_name) + .resolve_tool_info(name, namespace) .await } diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index 6e901a03467..3660f436844 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -5618,8 +5618,7 @@ async fn rejects_escalated_permissions_when_policy_not_on_request() { turn: Arc::clone(&turn_context), tracker: Arc::clone(&turn_diff_tracker), call_id, - tool_name: tool_name.to_string(), - tool_namespace: None, + tool_name: codex_tools::ToolName::plain(tool_name), payload: ToolPayload::Function { arguments: serde_json::json!({ "command": params.command.clone(), @@ -5697,8 +5696,7 @@ async fn unified_exec_rejects_escalated_permissions_when_policy_not_on_request() turn: Arc::clone(&turn_context), tracker: Arc::clone(&tracker), call_id: "exec-call".to_string(), - tool_name: "exec_command".to_string(), - tool_namespace: None, + tool_name: codex_tools::ToolName::plain("exec_command"), payload: ToolPayload::Function { arguments: serde_json::json!({ "cmd": "echo hi", diff --git a/codex-rs/core/src/codex_tests_guardian.rs b/codex-rs/core/src/codex_tests_guardian.rs index 70a70a367c2..cc84ee1d870 100644 --- a/codex-rs/core/src/codex_tests_guardian.rs +++ b/codex-rs/core/src/codex_tests_guardian.rs @@ -144,8 +144,7 @@ async fn guardian_allows_shell_additional_permissions_requests_past_policy_valid turn: Arc::clone(&turn_context), tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())), call_id: "test-call".to_string(), - tool_name: "shell".to_string(), - tool_namespace: None, + tool_name: codex_tools::ToolName::plain("shell"), payload: ToolPayload::Function { arguments: serde_json::json!({ "command": params.command.clone(), @@ -211,8 +210,7 @@ async fn guardian_allows_unified_exec_additional_permissions_requests_past_polic turn: Arc::clone(&turn_context), tracker: Arc::clone(&tracker), call_id: "exec-call".to_string(), - tool_name: "exec_command".to_string(), - tool_namespace: None, + tool_name: codex_tools::ToolName::plain("exec_command"), payload: ToolPayload::Function { arguments: serde_json::json!({ "cmd": "echo hi", @@ -325,8 +323,7 @@ async fn shell_handler_allows_sticky_turn_permissions_without_inline_request_per turn: Arc::clone(&turn_context), tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())), call_id: "sticky-turn-grant".to_string(), - tool_name: "shell".to_string(), - tool_namespace: None, + tool_name: codex_tools::ToolName::plain("shell"), payload: ToolPayload::Function { arguments: serde_json::json!({ "command": [ diff --git a/codex-rs/core/src/memories/usage.rs b/codex-rs/core/src/memories/usage.rs index 574567d99e4..480ef5dd9b9 100644 --- a/codex-rs/core/src/memories/usage.rs +++ b/codex-rs/core/src/memories/usage.rs @@ -38,13 +38,14 @@ pub(crate) async fn emit_metric_for_tool_read(invocation: &ToolInvocation, succe } let success = if success { "true" } else { "false" }; + let tool_name = invocation.tool_name.display(); for kind in kinds { invocation.turn.session_telemetry.counter( MEMORIES_USAGE_METRIC, /*inc*/ 1, &[ ("kind", kind.as_tag()), - ("tool", invocation.tool_name.as_str()), + ("tool", &tool_name), ("success", success), ], ); @@ -77,8 +78,11 @@ fn shell_command_for_invocation(invocation: &ToolInvocation) -> Option<(Vec serde_json::from_str::(arguments) + match ( + invocation.tool_name.namespace.as_deref(), + invocation.tool_name.name.as_str(), + ) { + (None, "shell") => serde_json::from_str::(arguments) .ok() .map(|params| { ( @@ -86,7 +90,7 @@ fn shell_command_for_invocation(invocation: &ToolInvocation) -> Option<(Vec serde_json::from_str::(arguments) + (None, "shell_command") => serde_json::from_str::(arguments) .ok() .map(|params| { if !invocation.turn.tools_config.allow_login_shell && params.login == Some(true) { @@ -107,7 +111,7 @@ fn shell_command_for_invocation(invocation: &ToolInvocation) -> Option<(Vec serde_json::from_str::(arguments) + (None, "exec_command") => serde_json::from_str::(arguments) .ok() .and_then(|params| { let command = crate::tools::handlers::unified_exec::get_command( @@ -122,7 +126,7 @@ fn shell_command_for_invocation(invocation: &ToolInvocation) -> Option<(Vec None, + (Some(_), _) | (None, _) => None, } } diff --git a/codex-rs/core/src/stream_events_utils.rs b/codex-rs/core/src/stream_events_utils.rs index ac4334d6830..8c63616e7ed 100644 --- a/codex-rs/core/src/stream_events_utils.rs +++ b/codex-rs/core/src/stream_events_utils.rs @@ -221,7 +221,7 @@ pub(crate) async fn handle_output_item_done( tracing::info!( thread_id = %ctx.sess.conversation_id, "ToolCall: {} {}", - call.tool_name, + call.tool_name.display(), payload_preview ); 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 c46c3e04585..8eaa85b5648 100644 --- a/codex-rs/core/src/tools/code_mode/execute_handler.rs +++ b/codex-rs/core/src/tools/code_mode/execute_handler.rs @@ -73,7 +73,9 @@ impl ToolHandler for CodeModeExecuteHandler { } = invocation; match payload { - ToolPayload::Custom { input } if tool_name == PUBLIC_TOOL_NAME => { + ToolPayload::Custom { input } + if tool_name.namespace.is_none() && tool_name.name.as_str() == PUBLIC_TOOL_NAME => + { self.execute(session, turn, call_id, input).await } _ => Err(FunctionCallError::RespondToModel(format!( diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index c0320387f4c..99c8ec0e0d8 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -27,6 +27,7 @@ use crate::tools::router::ToolCallSource; use crate::tools::router::ToolRouterParams; 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; @@ -283,27 +284,29 @@ async fn call_nested_tool( ))); } - let payload = - if let Some((server, tool)) = exec.session.parse_mcp_tool_name(&tool_name, &None).await { - match serialize_function_tool_arguments(&tool_name, input) { - Ok(raw_arguments) => ToolPayload::Mcp { - server, - tool, - raw_arguments, - }, - Err(error) => return Err(FunctionCallError::RespondToModel(error)), - } - } else { - 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 = if let Some(tool_info) = exec + .session + .resolve_mcp_tool_info(&tool_name, /*namespace*/ None) + .await + { + match serialize_function_tool_arguments(&tool_name, input) { + Ok(raw_arguments) => ToolPayload::Mcp { + server: tool_info.server_name, + tool: tool_info.tool.name.to_string(), + raw_arguments, + }, + Err(error) => return Err(FunctionCallError::RespondToModel(error)), + } + } else { + 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 call = ToolCall { - tool_name: tool_name.clone(), + tool_name: ToolName::plain(tool_name.clone()), call_id: format!("{PUBLIC_TOOL_NAME}-{}", uuid::Uuid::new_v4()), - tool_namespace: None, payload, }; let result = tool_runtime diff --git a/codex-rs/core/src/tools/code_mode/wait_handler.rs b/codex-rs/core/src/tools/code_mode/wait_handler.rs index 8a9d24c1ffc..dcc416b8329 100644 --- a/codex-rs/core/src/tools/code_mode/wait_handler.rs +++ b/codex-rs/core/src/tools/code_mode/wait_handler.rs @@ -55,7 +55,9 @@ impl ToolHandler for CodeModeWaitHandler { } = invocation; match payload { - ToolPayload::Function { arguments } if tool_name == WAIT_TOOL_NAME => { + ToolPayload::Function { arguments } + if tool_name.namespace.is_none() && tool_name.name.as_str() == WAIT_TOOL_NAME => + { let args: ExecWaitArgs = parse_arguments(&arguments)?; let exec = ExecContext { session, turn }; let started_at = std::time::Instant::now(); diff --git a/codex-rs/core/src/tools/context.rs b/codex-rs/core/src/tools/context.rs index 01dd0548b4c..933e973f115 100644 --- a/codex-rs/core/src/tools/context.rs +++ b/codex-rs/core/src/tools/context.rs @@ -13,6 +13,7 @@ use codex_protocol::models::ResponseInputItem; use codex_protocol::models::SearchToolCallParams; use codex_protocol::models::ShellToolCallParams; use codex_protocol::models::function_call_output_content_items_to_text; +use codex_tools::ToolName; use codex_tools::ToolSearchOutputTool; use codex_utils_output_truncation::TruncationPolicy; use codex_utils_output_truncation::formatted_truncate_text; @@ -39,8 +40,7 @@ pub struct ToolInvocation { pub turn: Arc, pub tracker: SharedTurnDiffTracker, pub call_id: String, - pub tool_name: String, - pub tool_namespace: Option, + pub tool_name: ToolName, pub payload: ToolPayload, } diff --git a/codex-rs/core/src/tools/handlers/agent_jobs.rs b/codex-rs/core/src/tools/handlers/agent_jobs.rs index c9475eeb390..e6119a3fc9d 100644 --- a/codex-rs/core/src/tools/handlers/agent_jobs.rs +++ b/codex-rs/core/src/tools/handlers/agent_jobs.rs @@ -206,7 +206,7 @@ impl ToolHandler for BatchJobHandler { } }; - match tool_name.as_str() { + match tool_name.name.as_str() { "spawn_agents_on_csv" => spawn_agents_on_csv::handle(session, turn, arguments).await, "report_agent_job_result" => report_agent_job_result::handle(session, arguments).await, other => Err(FunctionCallError::RespondToModel(format!( diff --git a/codex-rs/core/src/tools/handlers/apply_patch.rs b/codex-rs/core/src/tools/handlers/apply_patch.rs index 0be920feb31..281c3aa6ece 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch.rs @@ -218,7 +218,7 @@ impl ToolHandler for ApplyPatchHandler { session: session.clone(), turn: turn.clone(), call_id: call_id.clone(), - tool_name: tool_name.to_string(), + tool_name: tool_name.display(), }; let out = orchestrator .run( diff --git a/codex-rs/core/src/tools/handlers/dynamic.rs b/codex-rs/core/src/tools/handlers/dynamic.rs index 5f73f096e28..742c69a4827 100644 --- a/codex-rs/core/src/tools/handlers/dynamic.rs +++ b/codex-rs/core/src/tools/handlers/dynamic.rs @@ -50,13 +50,14 @@ impl ToolHandler for DynamicToolHandler { }; let args: Value = parse_arguments(&arguments)?; - let response = request_dynamic_tool(&session, turn.as_ref(), call_id, tool_name, args) - .await - .ok_or_else(|| { - FunctionCallError::RespondToModel( - "dynamic tool call was cancelled before receiving a response".to_string(), - ) - })?; + let response = + request_dynamic_tool(&session, turn.as_ref(), call_id, tool_name.display(), args) + .await + .ok_or_else(|| { + FunctionCallError::RespondToModel( + "dynamic tool call was cancelled before receiving a response".to_string(), + ) + })?; let DynamicToolResponse { content_items, diff --git a/codex-rs/core/src/tools/handlers/mcp_resource.rs b/codex-rs/core/src/tools/handlers/mcp_resource.rs index d0b7ddc5a81..69058a62267 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource.rs @@ -205,7 +205,7 @@ impl ToolHandler for McpResourceHandler { let arguments_value = parse_arguments(arguments.as_str())?; - match tool_name.as_str() { + match tool_name.name.as_str() { "list_mcp_resources" => { handle_list_resources( Arc::clone(&session), diff --git a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs index ca74b29d9b9..733d9985361 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -67,8 +67,7 @@ fn invocation( turn, tracker: Arc::new(Mutex::new(TurnDiffTracker::default())), call_id: "call-1".to_string(), - tool_name: tool_name.to_string(), - tool_namespace: None, + tool_name: codex_tools::ToolName::plain(tool_name), payload, } } diff --git a/codex-rs/core/src/tools/handlers/request_user_input_tests.rs b/codex-rs/core/src/tools/handlers/request_user_input_tests.rs index 0275a12860c..1dac095f1b3 100644 --- a/codex-rs/core/src/tools/handlers/request_user_input_tests.rs +++ b/codex-rs/core/src/tools/handlers/request_user_input_tests.rs @@ -29,8 +29,7 @@ async fn multi_agent_v2_request_user_input_rejects_subagent_threads() { turn: Arc::new(turn), tracker: Arc::new(Mutex::new(TurnDiffTracker::default())), call_id: "call-1".to_string(), - tool_name: REQUEST_USER_INPUT_TOOL_NAME.to_string(), - tool_namespace: None, + tool_name: codex_tools::ToolName::plain(REQUEST_USER_INPUT_TOOL_NAME), payload: ToolPayload::Function { arguments: json!({ "questions": [{ diff --git a/codex-rs/core/src/tools/handlers/shell.rs b/codex-rs/core/src/tools/handlers/shell.rs index 970afd6f338..3ed21bd2ed2 100644 --- a/codex-rs/core/src/tools/handlers/shell.rs +++ b/codex-rs/core/src/tools/handlers/shell.rs @@ -239,7 +239,7 @@ impl ToolHandler for ShellHandler { let exec_params = Self::to_exec_params(¶ms, turn.as_ref(), session.conversation_id); Self::run_exec_like(RunExecLikeArgs { - tool_name: tool_name.clone(), + tool_name: tool_name.display(), exec_params, additional_permissions: params.additional_permissions.clone(), prefix_rule, @@ -256,7 +256,7 @@ impl ToolHandler for ShellHandler { let exec_params = Self::to_exec_params(¶ms, turn.as_ref(), session.conversation_id); Self::run_exec_like(RunExecLikeArgs { - tool_name: tool_name.clone(), + tool_name: tool_name.display(), exec_params, additional_permissions: None, prefix_rule: None, @@ -270,7 +270,8 @@ impl ToolHandler for ShellHandler { .await } _ => Err(FunctionCallError::RespondToModel(format!( - "unsupported payload for shell handler: {tool_name}" + "unsupported payload for shell handler: {}", + tool_name.display() ))), } } @@ -339,7 +340,8 @@ impl ToolHandler for ShellCommandHandler { let ToolPayload::Function { arguments } = payload else { return Err(FunctionCallError::RespondToModel(format!( - "unsupported payload for shell_command handler: {tool_name}" + "unsupported payload for shell_command handler: {}", + tool_name.display() ))); }; @@ -362,7 +364,7 @@ impl ToolHandler for ShellCommandHandler { turn.tools_config.allow_login_shell, )?; ShellHandler::run_exec_like(RunExecLikeArgs { - tool_name, + tool_name: tool_name.display(), exec_params, additional_permissions: params.additional_permissions.clone(), prefix_rule, diff --git a/codex-rs/core/src/tools/handlers/shell_tests.rs b/codex-rs/core/src/tools/handlers/shell_tests.rs index fee47c21473..5c92a988ced 100644 --- a/codex-rs/core/src/tools/handlers/shell_tests.rs +++ b/codex-rs/core/src/tools/handlers/shell_tests.rs @@ -225,8 +225,7 @@ async fn shell_pre_tool_use_payload_uses_joined_command() { turn: turn.into(), tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), call_id: "call-41".to_string(), - tool_name: "shell".to_string(), - tool_namespace: None, + tool_name: codex_tools::ToolName::plain("shell"), payload, }), Some(crate::tools::registry::PreToolUsePayload { @@ -251,8 +250,7 @@ async fn shell_command_pre_tool_use_payload_uses_raw_command() { turn: turn.into(), tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), call_id: "call-42".to_string(), - tool_name: "shell_command".to_string(), - tool_namespace: None, + tool_name: codex_tools::ToolName::plain("shell_command"), payload, }), Some(crate::tools::registry::PreToolUsePayload { diff --git a/codex-rs/core/src/tools/handlers/unified_exec.rs b/codex-rs/core/src/tools/handlers/unified_exec.rs index fc1ad94b4d7..a604e9762cf 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec.rs @@ -121,7 +121,9 @@ impl ToolHandler for UnifiedExecHandler { } fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option { - if invocation.tool_name != "exec_command" { + if invocation.tool_name.namespace.is_some() + || invocation.tool_name.name.as_str() != "exec_command" + { return None; } @@ -186,7 +188,7 @@ impl ToolHandler for UnifiedExecHandler { let manager: &UnifiedExecProcessManager = &session.services.unified_exec_manager; let context = UnifiedExecContext::new(session.clone(), turn.clone(), call_id.clone()); - let response = match tool_name.as_str() { + let response = match tool_name.name.as_str() { "exec_command" => { let cwd = resolve_workdir_base_path(&arguments, &context.turn.cwd)?; let args: ExecCommandArgs = parse_arguments_with_base_path(&arguments, &cwd)?; @@ -289,7 +291,7 @@ impl ToolHandler for UnifiedExecHandler { context.turn.clone(), Some(&tracker), &context.call_id, - tool_name.as_str(), + &tool_name.name, ) .await? { diff --git a/codex-rs/core/src/tools/handlers/unified_exec_tests.rs b/codex-rs/core/src/tools/handlers/unified_exec_tests.rs index 1bdcee05e87..b6418145166 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec_tests.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec_tests.rs @@ -212,8 +212,7 @@ async fn exec_command_pre_tool_use_payload_uses_raw_command() { turn: turn.into(), tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), call_id: "call-43".to_string(), - tool_name: "exec_command".to_string(), - tool_namespace: None, + tool_name: codex_tools::ToolName::plain("exec_command"), payload, }), Some(crate::tools::registry::PreToolUsePayload { @@ -236,8 +235,7 @@ async fn exec_command_pre_tool_use_payload_skips_write_stdin() { turn: turn.into(), tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), call_id: "call-44".to_string(), - tool_name: "write_stdin".to_string(), - tool_namespace: None, + tool_name: codex_tools::ToolName::plain("write_stdin"), payload, }), None diff --git a/codex-rs/core/src/tools/js_repl/mod.rs b/codex-rs/core/src/tools/js_repl/mod.rs index f3acf13d533..e307f9718cd 100644 --- a/codex-rs/core/src/tools/js_repl/mod.rs +++ b/codex-rs/core/src/tools/js_repl/mod.rs @@ -1572,14 +1572,14 @@ impl JsReplManager { }, ); - let payload = if let Some((server, tool)) = exec + let payload = if let Some(tool_info) = exec .session - .parse_mcp_tool_name(&req.tool_name, &None) + .resolve_mcp_tool_info(&req.tool_name, /*namespace*/ None) .await { crate::tools::context::ToolPayload::Mcp { - server, - tool, + server: tool_info.server_name, + tool: tool_info.tool.name.to_string(), raw_arguments: req.arguments.clone(), } } else if is_freeform_tool(&router.specs(), &req.tool_name) { @@ -1594,8 +1594,7 @@ impl JsReplManager { let tool_name = req.tool_name.clone(); let call = crate::tools::router::ToolCall { - tool_name: tool_name.clone(), - tool_namespace: None, + tool_name: codex_tools::ToolName::plain(tool_name.clone()), call_id: req.id.clone(), payload, }; diff --git a/codex-rs/core/src/tools/parallel.rs b/codex-rs/core/src/tools/parallel.rs index 0e649333df4..a0818dfc830 100644 --- a/codex-rs/core/src/tools/parallel.rs +++ b/codex-rs/core/src/tools/parallel.rs @@ -85,11 +85,12 @@ impl ToolCallRuntime { let tracker = Arc::clone(&self.tracker); let lock = Arc::clone(&self.parallel_execution); let started = Instant::now(); + let display_name = call.tool_name.display(); let dispatch_span = trace_span!( "dispatch_tool_call_with_code_mode_result", - otel.name = call.tool_name.as_str(), - tool_name = call.tool_name.as_str(), + otel.name = display_name.as_str(), + tool_name = display_name.as_str(), call_id = call.call_id.as_str(), aborted = false, ); @@ -171,11 +172,15 @@ impl ToolCallRuntime { } fn abort_message(call: &ToolCall, secs: f32) -> String { - match call.tool_name.as_str() { - "shell" | "container.exec" | "local_shell" | "shell_command" | "unified_exec" => { - format!("Wall time: {secs:.1} seconds\naborted by user") - } - _ => format!("aborted by user after {secs:.1}s"), + if call.tool_name.namespace.is_none() + && matches!( + call.tool_name.name.as_str(), + "shell" | "container.exec" | "local_shell" | "shell_command" | "unified_exec" + ) + { + format!("Wall time: {secs:.1} seconds\naborted by user") + } else { + format!("aborted by user after {secs:.1}s") } } } diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index 2723d8a9a44..6231c7ea1da 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -23,6 +23,7 @@ use codex_hooks::HookToolKind; use codex_protocol::models::ResponseInputItem; use codex_protocol::protocol::SandboxPolicy; use codex_tools::ConfiguredToolSpec; +use codex_tools::ToolName; use codex_tools::ToolSpec; use codex_utils_readiness::Readiness; use futures::future::BoxFuture; @@ -179,32 +180,22 @@ where } } -pub(crate) fn tool_handler_key(tool_name: &str, namespace: Option<&str>) -> String { - if let Some(namespace) = namespace { - format!("{namespace}:{tool_name}") - } else { - tool_name.to_string() - } -} - pub struct ToolRegistry { - handlers: HashMap>, + handlers: HashMap>, } impl ToolRegistry { - fn new(handlers: HashMap>) -> Self { + fn new(handlers: HashMap>) -> Self { Self { handlers } } - fn handler(&self, name: &str, namespace: Option<&str>) -> Option> { - self.handlers - .get(&tool_handler_key(name, namespace)) - .map(Arc::clone) + fn handler(&self, name: &ToolName) -> Option> { + self.handlers.get(name).map(Arc::clone) } #[cfg(test)] - pub(crate) fn has_handler(&self, name: &str, namespace: Option<&str>) -> bool { - self.handler(name, namespace).is_some() + pub(crate) fn has_handler(&self, name: &ToolName) -> bool { + self.handler(name).is_some() } // TODO(jif) for dynamic tools. @@ -220,7 +211,7 @@ impl ToolRegistry { invocation: ToolInvocation, ) -> Result { let tool_name = invocation.tool_name.clone(); - let tool_namespace = invocation.tool_namespace.clone(); + let display_name = tool_name.display(); let call_id_owned = invocation.call_id.clone(); let otel = invocation.turn.session_telemetry.clone(); let payload_for_response = invocation.payload.clone(); @@ -262,16 +253,12 @@ impl ToolRegistry { } } - let handler = match self.handler(tool_name.as_ref(), tool_namespace.as_deref()) { + let handler = match self.handler(&tool_name) { Some(handler) => handler, None => { - let message = unsupported_tool_call_message( - &invocation.payload, - tool_name.as_ref(), - tool_namespace.as_deref(), - ); + let message = unsupported_tool_call_message(&invocation.payload, &tool_name); otel.tool_result_with_tags( - tool_name.as_ref(), + &display_name, &call_id_owned, log_payload.as_ref(), Duration::ZERO, @@ -286,9 +273,9 @@ impl ToolRegistry { }; if !handler.matches_kind(&invocation.payload) { - let message = format!("tool {tool_name} invoked with incompatible payload"); + let message = format!("tool {display_name} invoked with incompatible payload"); otel.tool_result_with_tags( - tool_name.as_ref(), + &display_name, &call_id_owned, log_payload.as_ref(), Duration::ZERO, @@ -323,7 +310,7 @@ impl ToolRegistry { let started = Instant::now(); let result = otel .log_tool_result_with_tags( - tool_name.as_ref(), + &display_name, &call_id_owned, log_payload.as_ref(), &metric_tags, @@ -443,7 +430,7 @@ impl ToolRegistry { } pub struct ToolRegistryBuilder { - handlers: HashMap>, + handlers: HashMap>, specs: Vec, } @@ -468,18 +455,15 @@ impl ToolRegistryBuilder { .push(ConfiguredToolSpec::new(spec, supports_parallel_tool_calls)); } - pub fn register_handler(&mut self, name: impl Into, handler: Arc) + pub fn register_handler(&mut self, name: impl Into, handler: Arc) where H: ToolHandler + 'static, { let name = name.into(); + let display_name = name.display(); let handler: Arc = handler; - if self - .handlers - .insert(name.clone(), handler.clone()) - .is_some() - { - warn!("overwriting handler for tool {name}"); + if self.handlers.insert(name, handler).is_some() { + warn!("overwriting handler for tool {display_name}"); } } @@ -507,12 +491,8 @@ impl ToolRegistryBuilder { } } -fn unsupported_tool_call_message( - payload: &ToolPayload, - tool_name: &str, - namespace: Option<&str>, -) -> String { - let tool_name = tool_handler_key(tool_name, namespace); +fn unsupported_tool_call_message(payload: &ToolPayload, tool_name: &ToolName) -> String { + let tool_name = tool_name.display(); match payload { ToolPayload::Custom { .. } => format!("unsupported custom tool call: {tool_name}"), _ => format!("unsupported call: {tool_name}"), @@ -605,7 +585,7 @@ async fn dispatch_after_tool_use_hook( event: HookEventAfterToolUse { turn_id: turn.sub_id.clone(), call_id: invocation.call_id.clone(), - tool_name: invocation.tool_name.clone(), + tool_name: invocation.tool_name.display(), tool_kind: hook_tool_kind(&tool_input), tool_input, executed: dispatch.executed, @@ -628,7 +608,7 @@ async fn dispatch_after_tool_use_hook( HookResult::FailedContinue(error) => { warn!( call_id = %invocation.call_id, - tool_name = %invocation.tool_name, + tool_name = %invocation.tool_name.display(), hook_name = %hook_name, error = %error, "after_tool_use hook failed; continuing" @@ -637,7 +617,7 @@ async fn dispatch_after_tool_use_hook( HookResult::FailedAbort(error) => { warn!( call_id = %invocation.call_id, - tool_name = %invocation.tool_name, + tool_name = %invocation.tool_name.display(), hook_name = %hook_name, error = %error, "after_tool_use hook failed; aborting operation" diff --git a/codex-rs/core/src/tools/registry_tests.rs b/codex-rs/core/src/tools/registry_tests.rs index 287996d7d87..f4fb72e791c 100644 --- a/codex-rs/core/src/tools/registry_tests.rs +++ b/codex-rs/core/src/tools/registry_tests.rs @@ -21,15 +21,19 @@ fn handler_looks_up_namespaced_aliases_explicitly() { let namespaced_handler = Arc::new(TestHandler) as Arc; let namespace = "mcp__codex_apps__gmail"; let tool_name = "gmail_get_recent_emails"; - let namespaced_name = tool_handler_key(tool_name, Some(namespace)); + let plain_name = codex_tools::ToolName::plain(tool_name); + let namespaced_name = codex_tools::ToolName::namespaced(namespace, tool_name); let registry = ToolRegistry::new(HashMap::from([ - (tool_name.to_string(), Arc::clone(&plain_handler)), - (namespaced_name, Arc::clone(&namespaced_handler)), + (plain_name.clone(), Arc::clone(&plain_handler)), + (namespaced_name.clone(), Arc::clone(&namespaced_handler)), ])); - let plain = registry.handler(tool_name, /*namespace*/ None); - let namespaced = registry.handler(tool_name, Some(namespace)); - let missing_namespaced = registry.handler(tool_name, Some("mcp__codex_apps__calendar")); + let plain = registry.handler(&plain_name); + let namespaced = registry.handler(&namespaced_name); + let missing_namespaced = registry.handler(&codex_tools::ToolName::namespaced( + "mcp__codex_apps__calendar", + tool_name, + )); assert_eq!(plain.is_some(), true); assert_eq!(namespaced.is_some(), true); diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index ce59eebe019..af9d2785855 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -16,6 +16,7 @@ use codex_protocol::models::SearchToolCallParams; use codex_protocol::models::ShellToolCallParams; use codex_tools::ConfiguredToolSpec; use codex_tools::DiscoverableTool; +use codex_tools::ToolName; use codex_tools::ToolSpec; use codex_tools::ToolsConfig; use std::collections::HashMap; @@ -26,8 +27,7 @@ pub use crate::tools::context::ToolCallSource; #[derive(Clone, Debug)] pub struct ToolCall { - pub tool_name: String, - pub tool_namespace: Option, + pub tool_name: ToolName, pub call_id: String, pub payload: ToolPayload, } @@ -104,11 +104,13 @@ impl ToolRouter { .map(|config| config.spec.clone()) } - pub fn tool_supports_parallel(&self, tool_name: &str) -> bool { - self.specs - .iter() - .filter(|config| config.supports_parallel_tool_calls) - .any(|config| config.name() == tool_name) + pub fn tool_supports_parallel(&self, tool_name: &ToolName) -> bool { + tool_name.namespace.is_none() + && self + .specs + .iter() + .filter(|config| config.supports_parallel_tool_calls) + .any(|config| config.name() == tool_name.name.as_str()) } #[instrument(level = "trace", skip_all, err)] @@ -124,21 +126,26 @@ impl ToolRouter { call_id, .. } => { - if let Some((server, tool)) = session.parse_mcp_tool_name(&name, &namespace).await { + let mcp_tool = session + .resolve_mcp_tool_info(&name, namespace.as_deref()) + .await; + let tool_name = match namespace { + Some(namespace) => ToolName::namespaced(namespace, name), + None => ToolName::plain(name), + }; + if let Some(tool_info) = mcp_tool { Ok(Some(ToolCall { - tool_name: name, - tool_namespace: namespace, + tool_name, call_id, payload: ToolPayload::Mcp { - server, - tool, + server: tool_info.server_name, + tool: tool_info.tool.name.to_string(), raw_arguments: arguments, }, })) } else { Ok(Some(ToolCall { - tool_name: name, - tool_namespace: namespace, + tool_name, call_id, payload: ToolPayload::Function { arguments }, })) @@ -157,8 +164,7 @@ impl ToolRouter { )) })?; Ok(Some(ToolCall { - tool_name: "tool_search".to_string(), - tool_namespace: None, + tool_name: ToolName::plain("tool_search"), call_id, payload: ToolPayload::ToolSearch { arguments }, })) @@ -170,8 +176,7 @@ impl ToolRouter { call_id, .. } => Ok(Some(ToolCall { - tool_name: name, - tool_namespace: None, + tool_name: ToolName::plain(name), call_id, payload: ToolPayload::Custom { input }, })), @@ -197,8 +202,7 @@ impl ToolRouter { justification: None, }; Ok(Some(ToolCall { - tool_name: "local_shell".to_string(), - tool_namespace: None, + tool_name: ToolName::plain("local_shell"), call_id, payload: ToolPayload::LocalShell { params }, })) @@ -220,14 +224,15 @@ impl ToolRouter { ) -> Result { let ToolCall { tool_name, - tool_namespace, call_id, payload, } = call; + let direct_js_repl_call = tool_name.namespace.is_none() + && matches!(tool_name.name.as_str(), "js_repl" | "js_repl_reset"); if source == ToolCallSource::Direct && turn.tools_config.js_repl_tools_only - && !matches!(tool_name.as_str(), "js_repl" | "js_repl_reset") + && !direct_js_repl_call { return Err(FunctionCallError::RespondToModel( "direct tool calls are disabled; use js_repl and codex.tool(...) instead" @@ -241,7 +246,6 @@ impl ToolRouter { tracker, call_id, tool_name, - tool_namespace, payload, }; diff --git a/codex-rs/core/src/tools/router_tests.rs b/codex-rs/core/src/tools/router_tests.rs index e54981b3b1c..0478ebca643 100644 --- a/codex-rs/core/src/tools/router_tests.rs +++ b/codex-rs/core/src/tools/router_tests.rs @@ -5,6 +5,7 @@ use crate::function_tool::FunctionCallError; use crate::tools::context::ToolPayload; use crate::turn_diff_tracker::TurnDiffTracker; use codex_protocol::models::ResponseItem; +use codex_tools::ToolName; use super::ToolCall; use super::ToolCallSource; @@ -37,8 +38,7 @@ async fn js_repl_tools_only_blocks_direct_tool_calls() -> anyhow::Result<()> { ); let call = ToolCall { - tool_name: "shell".to_string(), - tool_namespace: None, + tool_name: ToolName::plain("shell"), call_id: "call-1".to_string(), payload: ToolPayload::Function { arguments: "{}".to_string(), @@ -90,8 +90,7 @@ async fn js_repl_tools_only_allows_js_repl_source_calls() -> anyhow::Result<()> ); let call = ToolCall { - tool_name: "shell".to_string(), - tool_namespace: None, + tool_name: ToolName::plain("shell"), call_id: "call-2".to_string(), payload: ToolPayload::Function { arguments: "{}".to_string(), @@ -118,6 +117,84 @@ async fn js_repl_tools_only_allows_js_repl_source_calls() -> anyhow::Result<()> Ok(()) } +#[tokio::test] +async fn js_repl_tools_only_blocks_namespaced_js_repl_tool() -> anyhow::Result<()> { + let (session, mut turn) = make_session_and_context().await; + turn.tools_config.js_repl_tools_only = true; + + let session = Arc::new(session); + let turn = Arc::new(turn); + let router = ToolRouter::from_config( + &turn.tools_config, + ToolRouterParams { + deferred_mcp_tools: None, + mcp_tools: None, + discoverable_tools: None, + dynamic_tools: turn.dynamic_tools.as_slice(), + }, + ); + + let call = ToolCall { + tool_name: ToolName::namespaced("mcp__server__", "js_repl"), + call_id: "call-namespaced-js-repl".to_string(), + payload: ToolPayload::Mcp { + server: "server".to_string(), + tool: "js_repl".to_string(), + raw_arguments: "{}".to_string(), + }, + }; + let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); + let err = router + .dispatch_tool_call_with_code_mode_result( + session, + turn, + tracker, + call, + ToolCallSource::Direct, + ) + .await + .err() + .expect("namespaced js_repl calls should be blocked"); + let FunctionCallError::RespondToModel(message) = err else { + panic!("expected RespondToModel, got {err:?}"); + }; + assert!(message.contains("direct tool calls are disabled")); + + Ok(()) +} + +#[tokio::test] +async fn parallel_support_does_not_match_namespaced_local_tool_names() -> anyhow::Result<()> { + let (session, turn) = make_session_and_context().await; + let mcp_tools = session + .services + .mcp_connection_manager + .read() + .await + .list_all_tools() + .await; + let router = ToolRouter::from_config( + &turn.tools_config, + ToolRouterParams { + deferred_mcp_tools: None, + mcp_tools: Some(mcp_tools), + discoverable_tools: None, + dynamic_tools: turn.dynamic_tools.as_slice(), + }, + ); + + let parallel_tool_name = ["shell", "local_shell", "exec_command", "shell_command"] + .into_iter() + .find(|name| router.tool_supports_parallel(&ToolName::plain(*name))) + .expect("test session should expose a parallel shell-like tool"); + + assert!( + !router.tool_supports_parallel(&ToolName::namespaced("mcp__server__", parallel_tool_name)) + ); + + Ok(()) +} + #[tokio::test] async fn build_tool_call_uses_namespace_for_registry_name() -> anyhow::Result<()> { let (session, _) = make_session_and_context().await; @@ -137,10 +214,9 @@ async fn build_tool_call_uses_namespace_for_registry_name() -> anyhow::Result<() .await? .expect("function_call should produce a tool call"); - assert_eq!(call.tool_name, tool_name); assert_eq!( - call.tool_namespace, - Some("mcp__codex_apps__calendar".to_string()) + call.tool_name, + ToolName::namespaced("mcp__codex_apps__calendar", tool_name) ); assert_eq!(call.call_id, "call-namespace"); match call.payload { diff --git a/codex-rs/core/src/tools/spec_tests.rs b/codex-rs/core/src/tools/spec_tests.rs index 14cef85b1f4..da76a1bb28f 100644 --- a/codex-rs/core/src/tools/spec_tests.rs +++ b/codex-rs/core/src/tools/spec_tests.rs @@ -3,7 +3,6 @@ use crate::shell::Shell; use crate::shell::ShellType; use crate::test_support::construct_model_info_offline; use crate::tools::ToolRouter; -use crate::tools::registry::tool_handler_key; use crate::tools::router::ToolRouterParams; use codex_app_server_protocol::AppInfo; use codex_features::Feature; @@ -24,6 +23,7 @@ use codex_tools::ResponsesApiTool; use codex_tools::ShellCommandBackendConfig; use codex_tools::TOOL_SEARCH_TOOL_NAME; use codex_tools::TOOL_SUGGEST_TOOL_NAME; +use codex_tools::ToolName; use codex_tools::ToolSpec; use codex_tools::ToolsConfig; use codex_tools::ToolsConfigParams; @@ -903,12 +903,12 @@ fn search_tool_registers_namespaced_mcp_tool_aliases() { ) .build(); - let app_alias = tool_handler_key("_create_event", Some("mcp__codex_apps__calendar")); - let mcp_alias = tool_handler_key("echo", Some("mcp__rmcp__")); + let app_alias = ToolName::namespaced("mcp__codex_apps__calendar", "_create_event"); + let mcp_alias = ToolName::namespaced("mcp__rmcp__", "echo"); - assert!(registry.has_handler(TOOL_SEARCH_TOOL_NAME, /*namespace*/ None)); - assert!(registry.has_handler(app_alias.as_str(), /*namespace*/ None)); - assert!(registry.has_handler(mcp_alias.as_str(), /*namespace*/ None)); + assert!(registry.has_handler(&ToolName::plain(TOOL_SEARCH_TOOL_NAME))); + assert!(registry.has_handler(&app_alias)); + assert!(registry.has_handler(&mcp_alias)); } #[test] diff --git a/codex-rs/tools/src/lib.rs b/codex-rs/tools/src/lib.rs index 9a3107dd666..c310222a0fd 100644 --- a/codex-rs/tools/src/lib.rs +++ b/codex-rs/tools/src/lib.rs @@ -18,6 +18,7 @@ mod responses_api; mod tool_config; mod tool_definition; mod tool_discovery; +mod tool_name; mod tool_registry_plan; mod tool_registry_plan_types; mod tool_spec; @@ -112,6 +113,7 @@ pub use tool_discovery::collect_tool_suggest_entries; pub use tool_discovery::create_tool_search_tool; pub use tool_discovery::create_tool_suggest_tool; pub use tool_discovery::filter_tool_suggest_discoverable_tools_for_client; +pub use tool_name::ToolName; pub use tool_registry_plan::build_tool_registry_plan; pub use tool_registry_plan_types::ToolHandlerKind; pub use tool_registry_plan_types::ToolHandlerSpec; diff --git a/codex-rs/tools/src/tool_name.rs b/codex-rs/tools/src/tool_name.rs new file mode 100644 index 00000000000..da18233abb5 --- /dev/null +++ b/codex-rs/tools/src/tool_name.rs @@ -0,0 +1,42 @@ +/// Identifies a callable tool, preserving the namespace split when the model +/// provides one. +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +pub struct ToolName { + pub name: String, + pub namespace: Option, +} + +impl ToolName { + pub fn plain(name: impl Into) -> Self { + Self { + name: name.into(), + namespace: None, + } + } + + pub fn namespaced(namespace: impl Into, name: impl Into) -> Self { + Self { + name: name.into(), + namespace: Some(namespace.into()), + } + } + + pub fn display(&self) -> String { + match &self.namespace { + Some(namespace) => format!("{namespace}{}", self.name), + None => self.name.clone(), + } + } +} + +impl From for ToolName { + fn from(name: String) -> Self { + Self::plain(name) + } +} + +impl From<&str> for ToolName { + fn from(name: &str) -> Self { + Self::plain(name) + } +} diff --git a/codex-rs/tools/src/tool_registry_plan.rs b/codex-rs/tools/src/tool_registry_plan.rs index 72be7843cce..82b4c3a1426 100644 --- a/codex-rs/tools/src/tool_registry_plan.rs +++ b/codex-rs/tools/src/tool_registry_plan.rs @@ -6,6 +6,7 @@ use crate::TOOL_SEARCH_DEFAULT_LIMIT; use crate::TOOL_SEARCH_TOOL_NAME; use crate::TOOL_SUGGEST_TOOL_NAME; use crate::ToolHandlerKind; +use crate::ToolName; use crate::ToolRegistryPlan; use crate::ToolRegistryPlanParams; use crate::ToolSearchSource; @@ -266,7 +267,7 @@ pub fn build_tool_registry_plan( for tool in deferred_mcp_tools { plan.register_handler( - format!("{}:{}", tool.tool_namespace, tool.tool_name), + ToolName::namespaced(tool.tool_namespace, tool.tool_name), ToolHandlerKind::Mcp, ); } diff --git a/codex-rs/tools/src/tool_registry_plan_tests.rs b/codex-rs/tools/src/tool_registry_plan_tests.rs index 3454a229966..9da58849c1b 100644 --- a/codex-rs/tools/src/tool_registry_plan_tests.rs +++ b/codex-rs/tools/src/tool_registry_plan_tests.rs @@ -1285,11 +1285,11 @@ fn search_tool_description_lists_each_mcp_source_once() { assert!(!description.contains("mcp__rmcp__echo")); assert!(handlers.contains(&ToolHandlerSpec { - name: "mcp__codex_apps__calendar:_create_event".to_string(), + name: ToolName::namespaced("mcp__codex_apps__calendar", "_create_event"), kind: ToolHandlerKind::Mcp, })); assert!(handlers.contains(&ToolHandlerSpec { - name: "mcp__rmcp__:echo".to_string(), + name: ToolName::namespaced("mcp__rmcp__", "echo"), kind: ToolHandlerKind::Mcp, })); } diff --git a/codex-rs/tools/src/tool_registry_plan_types.rs b/codex-rs/tools/src/tool_registry_plan_types.rs index fb628671711..7459954dcdd 100644 --- a/codex-rs/tools/src/tool_registry_plan_types.rs +++ b/codex-rs/tools/src/tool_registry_plan_types.rs @@ -1,5 +1,6 @@ use crate::ConfiguredToolSpec; use crate::DiscoverableTool; +use crate::ToolName; use crate::ToolSpec; use crate::ToolsConfig; use crate::WaitAgentTimeoutOptions; @@ -45,7 +46,7 @@ pub enum ToolHandlerKind { #[derive(Debug, Clone, PartialEq, Eq)] pub struct ToolHandlerSpec { - pub name: String, + pub name: ToolName, pub kind: ToolHandlerKind, } @@ -104,7 +105,7 @@ impl ToolRegistryPlan { .push(ConfiguredToolSpec::new(spec, supports_parallel_tool_calls)); } - pub(crate) fn register_handler(&mut self, name: impl Into, kind: ToolHandlerKind) { + pub(crate) fn register_handler(&mut self, name: impl Into, kind: ToolHandlerKind) { self.handlers.push(ToolHandlerSpec { name: name.into(), kind,