diff --git a/codex-rs/core/src/client_common.rs b/codex-rs/core/src/client_common.rs index b695581deb..8e98bd6918 100644 --- a/codex-rs/core/src/client_common.rs +++ b/codex-rs/core/src/client_common.rs @@ -1,6 +1,6 @@ use crate::error::Result; use crate::model_family::ModelFamily; -use crate::openai_tools::OpenAiTool; +use crate::openai_tools::ToolSpec; use crate::protocol::RateLimitSnapshot; use crate::protocol::TokenUsage; use codex_apply_patch::APPLY_PATCH_TOOL_INSTRUCTIONS; @@ -29,7 +29,7 @@ pub struct Prompt { /// Tools available to the model, including additional tools sourced from /// external MCP servers. - pub(crate) tools: Vec, + pub(crate) tools: Vec, /// Optional override for the built-in BASE_INSTRUCTIONS. pub base_instructions_override: Option, @@ -49,8 +49,8 @@ impl Prompt { // AND // - there is no apply_patch tool present let is_apply_patch_tool_present = self.tools.iter().any(|tool| match tool { - OpenAiTool::Function(f) => f.name == "apply_patch", - OpenAiTool::Freeform(f) => f.name == "apply_patch", + ToolSpec::Function(f) => f.name == "apply_patch", + ToolSpec::Freeform(f) => f.name == "apply_patch", _ => false, }); if self.base_instructions_override.is_none() diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 3fc0da7829..974828f9ea 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -27,7 +27,6 @@ use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnContextItem; use futures::prelude::*; use mcp_types::CallToolResult; -use serde::Deserialize; use serde::Serialize; use serde_json; use serde_json::Value; @@ -59,26 +58,19 @@ use crate::exec::ExecToolCallOutput; use crate::exec::StdoutStream; #[cfg(test)] use crate::exec::StreamOutput; -use crate::exec_command::EXEC_COMMAND_TOOL_NAME; use crate::exec_command::ExecCommandParams; use crate::exec_command::ExecSessionManager; -use crate::exec_command::WRITE_STDIN_TOOL_NAME; use crate::exec_command::WriteStdinParams; -use crate::exec_env::create_env; use crate::executor::ExecutionMode; use crate::executor::Executor; use crate::executor::ExecutorConfig; use crate::executor::normalize_exec_result; use crate::mcp_connection_manager::McpConnectionManager; -use crate::mcp_tool_call::handle_mcp_tool_call; use crate::model_family::find_family_for_model; use crate::openai_model_info::get_model_info; -use crate::openai_tools::ApplyPatchToolArgs; use crate::openai_tools::ToolsConfig; use crate::openai_tools::ToolsConfigParams; -use crate::openai_tools::get_openai_tools; use crate::parse_command::parse_command; -use crate::plan_tool::handle_update_plan; use crate::project_doc::get_user_instructions; use crate::protocol::AgentMessageDeltaEvent; use crate::protocol::AgentReasoningDeltaEvent; @@ -118,6 +110,7 @@ use crate::state::SessionServices; use crate::tasks::CompactTask; use crate::tasks::RegularTask; use crate::tasks::ReviewTask; +use crate::tools::Router; use crate::turn_diff_tracker::TurnDiffTracker; use crate::unified_exec::UnifiedExecSessionManager; use crate::user_instructions::UserInstructions; @@ -129,10 +122,8 @@ use codex_protocol::config_types::ReasoningSummary as ReasoningSummaryConfig; use codex_protocol::custom_prompts::CustomPrompt; use codex_protocol::models::ContentItem; use codex_protocol::models::FunctionCallOutputPayload; -use codex_protocol::models::LocalShellAction; use codex_protocol::models::ResponseInputItem; use codex_protocol::models::ResponseItem; -use codex_protocol::models::ShellToolCallParams; use codex_protocol::protocol::InitialHistory; pub mod compact; @@ -285,7 +276,7 @@ pub(crate) struct TurnContext { } impl TurnContext { - fn resolve_path(&self, path: Option) -> PathBuf { + pub(crate) fn resolve_path(&self, path: Option) -> PathBuf { path.as_ref() .map(PathBuf::from) .map_or_else(|| self.cwd.clone(), |p| self.cwd.join(p)) @@ -1034,6 +1025,49 @@ impl Session { .await } + pub(crate) fn parse_mcp_tool_name(&self, tool_name: &str) -> Option<(String, String)> { + self.services + .mcp_connection_manager + .parse_tool_name(tool_name) + } + + pub(crate) async fn handle_exec_command_tool( + &self, + params: ExecCommandParams, + ) -> Result { + let result = self + .services + .session_manager + .handle_exec_command_request(params) + .await; + match result { + Ok(output) => Ok(output.to_text_output()), + Err(err) => Err(FunctionCallError::RespondToModel(err)), + } + } + + pub(crate) async fn handle_write_stdin_tool( + &self, + params: WriteStdinParams, + ) -> Result { + self.services + .session_manager + .handle_write_stdin_request(params) + .await + .map(|output| output.to_text_output()) + .map_err(FunctionCallError::RespondToModel) + } + + pub(crate) async fn run_unified_exec_request( + &self, + request: crate::unified_exec::UnifiedExecRequest<'_>, + ) -> Result { + self.services + .unified_exec_manager + .handle_request(request) + .await + } + pub async fn interrupt_task(self: &Arc) { info!("interrupt received: abort current task, if any"); self.abort_all_tasks(TurnAbortReason::Interrupted).await; @@ -1890,21 +1924,28 @@ async fn run_turn( sub_id: String, input: Vec, ) -> CodexResult { - let tools = get_openai_tools( - &turn_context.tools_config, - Some(sess.services.mcp_connection_manager.list_all_tools()), - ); + let mcp_tools = sess.services.mcp_connection_manager.list_all_tools(); + let router = Router::from_config(&turn_context.tools_config, Some(mcp_tools)); let prompt = Prompt { input, - tools, + tools: router.specs().to_vec(), base_instructions_override: turn_context.base_instructions.clone(), output_schema: turn_context.final_output_json_schema.clone(), }; let mut retries = 0; loop { - match try_run_turn(sess, turn_context, turn_diff_tracker, &sub_id, &prompt).await { + match try_run_turn( + &router, + sess, + turn_context, + turn_diff_tracker, + &sub_id, + &prompt, + ) + .await + { Ok(output) => return Ok(output), Err(CodexErr::Interrupted) => return Err(CodexErr::Interrupted), Err(CodexErr::EnvVar(var)) => return Err(CodexErr::EnvVar(var)), @@ -1966,6 +2007,7 @@ struct TurnRunResult { } async fn try_run_turn( + router: &crate::tools::Router, sess: &Session, turn_context: &TurnContext, turn_diff_tracker: &mut TurnDiffTracker, @@ -2067,6 +2109,7 @@ async fn try_run_turn( ResponseEvent::Created => {} ResponseEvent::OutputItemDone(item) => { let response = handle_response_item( + router, sess, turn_context, turn_diff_tracker, @@ -2157,6 +2200,7 @@ async fn try_run_turn( } async fn handle_response_item( + router: &crate::tools::Router, sess: &Session, turn_context: &TurnContext, turn_diff_tracker: &mut TurnDiffTracker, @@ -2164,441 +2208,69 @@ async fn handle_response_item( item: ResponseItem, ) -> CodexResult> { debug!(?item, "Output item"); - let output = match item { - ResponseItem::FunctionCall { - name, - arguments, - call_id, - .. - } => { - info!("FunctionCall: {name}({arguments})"); - if let Some((server, tool_name)) = - sess.services.mcp_connection_manager.parse_tool_name(&name) - { - let resp = handle_mcp_tool_call( - sess, - sub_id, - call_id.clone(), - server, - tool_name, - arguments, - ) + + match Router::build_tool_call(sess, item.clone()) { + Ok(Some(call)) => { + let payload_preview = call.payload.log_payload().into_owned(); + tracing::info!("ToolCall: {} {}", call.tool_name, payload_preview); + let response = router + .dispatch_tool_call(sess, turn_context, turn_diff_tracker, sub_id, call) .await; - Some(resp) - } else { - let result = turn_context + Ok(Some(response)) + } + Ok(None) => { + match &item { + ResponseItem::Message { .. } + | ResponseItem::Reasoning { .. } + | ResponseItem::WebSearchCall { .. } => { + let msgs = match &item { + ResponseItem::Message { .. } if turn_context.is_review_mode => { + trace!("suppressing assistant Message in review mode"); + Vec::new() + } + _ => map_response_item_to_event_messages( + &item, + sess.show_raw_agent_reasoning(), + ), + }; + for msg in msgs { + let event = Event { + id: sub_id.to_string(), + msg, + }; + sess.send_event(event).await; + } + } + ResponseItem::FunctionCallOutput { .. } + | ResponseItem::CustomToolCallOutput { .. } => { + debug!("unexpected tool output from stream"); + } + _ => {} + } + + Ok(None) + } + Err(FunctionCallError::RespondToModel(msg)) => { + if msg == "LocalShellCall without call_id or id" { + turn_context .client .get_otel_event_manager() - .log_tool_result(name.as_str(), call_id.as_str(), arguments.as_str(), || { - handle_function_call( - sess, - turn_context, - turn_diff_tracker, - sub_id.to_string(), - name.to_owned(), - arguments.to_owned(), - call_id.clone(), - ) - }) - .await; - - let output = match result { - Ok(content) => FunctionCallOutputPayload { - content, - success: Some(true), - }, - Err(FunctionCallError::RespondToModel(msg)) => FunctionCallOutputPayload { - content: msg, - success: Some(false), - }, - }; - Some(ResponseInputItem::FunctionCallOutput { call_id, output }) - } - } - ResponseItem::LocalShellCall { - id, - call_id, - status: _, - action, - } => { - let name = "local_shell"; - let LocalShellAction::Exec(action) = action; - tracing::info!("LocalShellCall: {action:?}"); - let params = ShellToolCallParams { - command: action.command, - workdir: action.working_directory, - timeout_ms: action.timeout_ms, - with_escalated_permissions: None, - justification: None, - }; - let effective_call_id = match (call_id, id) { - (Some(call_id), _) => call_id, - (None, Some(id)) => id, - (None, None) => { - let error_message = "LocalShellCall without call_id or id"; - - turn_context - .client - .get_otel_event_manager() - .log_tool_failed(name, error_message); - - error!(error_message); - return Ok(Some(ResponseInputItem::FunctionCallOutput { - call_id: "".to_string(), - output: FunctionCallOutputPayload { - content: error_message.to_string(), - success: None, - }, - })); - } - }; - - let exec_params = to_exec_params(params, turn_context); - { - let result = turn_context - .client - .get_otel_event_manager() - .log_tool_result( - name, - effective_call_id.as_str(), - exec_params.command.join(" ").as_str(), - || { - handle_container_exec_with_params( - name, - exec_params, - sess, - turn_context, - turn_diff_tracker, - sub_id.to_string(), - effective_call_id.clone(), - ) - }, - ) - .await; - - let output = match result { - Ok(content) => FunctionCallOutputPayload { - content, - success: Some(true), - }, - Err(FunctionCallError::RespondToModel(msg)) => FunctionCallOutputPayload { - content: msg, - success: Some(false), - }, - }; - Some(ResponseInputItem::FunctionCallOutput { - call_id: effective_call_id, - output, - }) - } - } - ResponseItem::CustomToolCall { - id: _, - call_id, - name, - input, - status: _, - } => { - let result = turn_context - .client - .get_otel_event_manager() - .log_tool_result(name.as_str(), call_id.as_str(), input.as_str(), || { - handle_custom_tool_call( - sess, - turn_context, - turn_diff_tracker, - sub_id.to_string(), - name.to_owned(), - input.to_owned(), - call_id.clone(), - ) - }) - .await; - - let output = match result { - Ok(content) => content, - Err(FunctionCallError::RespondToModel(msg)) => msg, - }; - Some(ResponseInputItem::CustomToolCallOutput { call_id, output }) - } - ResponseItem::FunctionCallOutput { .. } => { - debug!("unexpected FunctionCallOutput from stream"); - None - } - ResponseItem::CustomToolCallOutput { .. } => { - debug!("unexpected CustomToolCallOutput from stream"); - None - } - ResponseItem::Message { .. } - | ResponseItem::Reasoning { .. } - | ResponseItem::WebSearchCall { .. } => { - // In review child threads, suppress assistant message events but - // keep reasoning/web search. - let msgs = match &item { - ResponseItem::Message { .. } if turn_context.is_review_mode => { - trace!("suppressing assistant Message in review mode"); - Vec::new() - } - _ => map_response_item_to_event_messages(&item, sess.show_raw_agent_reasoning()), - }; - for msg in msgs { - let event = Event { - id: sub_id.to_string(), - msg, - }; - sess.send_event(event).await; - } - None - } - ResponseItem::Other => None, - }; - Ok(output) -} - -async fn handle_unified_exec_tool_call( - sess: &Session, - session_id: Option, - arguments: Vec, - timeout_ms: Option, -) -> Result { - let parsed_session_id = if let Some(session_id) = session_id { - match session_id.parse::() { - Ok(parsed) => Some(parsed), - Err(output) => { - return Err(FunctionCallError::RespondToModel(format!( - "invalid session_id: {session_id} due to error {output:?}" - ))); - } - } - } else { - None - }; - - let request = crate::unified_exec::UnifiedExecRequest { - session_id: parsed_session_id, - input_chunks: &arguments, - timeout_ms, - }; - - let value = sess - .services - .unified_exec_manager - .handle_request(request) - .await - .map_err(|err| { - FunctionCallError::RespondToModel(format!("unified exec failed: {err:?}")) - })?; - - #[derive(Serialize)] - struct SerializedUnifiedExecResult { - session_id: Option, - output: String, - } - - serde_json::to_string(&SerializedUnifiedExecResult { - session_id: value.session_id.map(|id| id.to_string()), - output: value.output, - }) - .map_err(|err| { - FunctionCallError::RespondToModel(format!( - "failed to serialize unified exec output: {err:?}" - )) - }) -} - -async fn handle_function_call( - sess: &Session, - turn_context: &TurnContext, - turn_diff_tracker: &mut TurnDiffTracker, - sub_id: String, - name: String, - arguments: String, - call_id: String, -) -> Result { - match name.as_str() { - "container.exec" | "shell" => { - let params = parse_container_exec_arguments(arguments, turn_context, &call_id)?; - handle_container_exec_with_params( - name.as_str(), - params, - sess, - turn_context, - turn_diff_tracker, - sub_id, - call_id, - ) - .await - } - "unified_exec" => { - #[derive(Deserialize)] - struct UnifiedExecArgs { - input: Vec, - #[serde(default)] - session_id: Option, - #[serde(default)] - timeout_ms: Option, + .log_tool_failed("local_shell", &msg); + error!(msg); } - let args: UnifiedExecArgs = serde_json::from_str(&arguments).map_err(|err| { - FunctionCallError::RespondToModel(format!( - "failed to parse function arguments: {err:?}" - )) - })?; - - handle_unified_exec_tool_call(sess, args.session_id, args.input, args.timeout_ms).await - } - "view_image" => { - #[derive(serde::Deserialize)] - struct SeeImageArgs { - path: String, - } - let args: SeeImageArgs = serde_json::from_str(&arguments).map_err(|e| { - FunctionCallError::RespondToModel(format!( - "failed to parse function arguments: {e:?}" - )) - })?; - let abs = turn_context.resolve_path(Some(args.path)); - sess.inject_input(vec![InputItem::LocalImage { path: abs }]) - .await - .map_err(|_| { - FunctionCallError::RespondToModel( - "unable to attach image (no active task)".to_string(), - ) - })?; - - Ok("attached local image path".to_string()) - } - "apply_patch" => { - let args: ApplyPatchToolArgs = serde_json::from_str(&arguments).map_err(|e| { - FunctionCallError::RespondToModel(format!( - "failed to parse function arguments: {e:?}" - )) - })?; - let exec_params = ExecParams { - command: vec!["apply_patch".to_string(), args.input.clone()], - cwd: turn_context.cwd.clone(), - timeout_ms: None, - env: HashMap::new(), - with_escalated_permissions: None, - justification: None, - }; - handle_container_exec_with_params( - name.as_str(), - exec_params, - sess, - turn_context, - turn_diff_tracker, - sub_id, - call_id, - ) - .await - } - "update_plan" => handle_update_plan(sess, arguments, sub_id, call_id).await, - EXEC_COMMAND_TOOL_NAME => { - // TODO(mbolin): Sandbox check. - let exec_params: ExecCommandParams = serde_json::from_str(&arguments).map_err(|e| { - FunctionCallError::RespondToModel(format!( - "failed to parse function arguments: {e:?}" - )) - })?; - let result = sess - .services - .session_manager - .handle_exec_command_request(exec_params) - .await; - match result { - Ok(output) => Ok(output.to_text_output()), - Err(err) => Err(FunctionCallError::RespondToModel(err)), - } - } - WRITE_STDIN_TOOL_NAME => { - let write_stdin_params = - serde_json::from_str::(&arguments).map_err(|e| { - FunctionCallError::RespondToModel(format!( - "failed to parse function arguments: {e:?}" - )) - })?; - - let result = sess - .services - .session_manager - .handle_write_stdin_request(write_stdin_params) - .await - .map_err(FunctionCallError::RespondToModel)?; - - Ok(result.to_text_output()) - } - _ => Err(FunctionCallError::RespondToModel(format!( - "unsupported call: {name}" - ))), - } -} - -async fn handle_custom_tool_call( - sess: &Session, - turn_context: &TurnContext, - turn_diff_tracker: &mut TurnDiffTracker, - sub_id: String, - name: String, - input: String, - call_id: String, -) -> Result { - info!("CustomToolCall: {name} {input}"); - match name.as_str() { - "apply_patch" => { - let exec_params = ExecParams { - command: vec!["apply_patch".to_string(), input.clone()], - cwd: turn_context.cwd.clone(), - timeout_ms: None, - env: HashMap::new(), - with_escalated_permissions: None, - justification: None, - }; - - handle_container_exec_with_params( - name.as_str(), - exec_params, - sess, - turn_context, - turn_diff_tracker, - sub_id, - call_id, - ) - .await - } - _ => { - debug!("unexpected CustomToolCall from stream"); - Err(FunctionCallError::RespondToModel(format!( - "unsupported custom tool call: {name}" - ))) + Ok(Some(ResponseInputItem::FunctionCallOutput { + call_id: String::new(), + output: FunctionCallOutputPayload { + content: msg, + success: None, + }, + })) } } } -fn to_exec_params(params: ShellToolCallParams, turn_context: &TurnContext) -> ExecParams { - ExecParams { - command: params.command, - cwd: turn_context.resolve_path(params.workdir.clone()), - timeout_ms: params.timeout_ms, - env: create_env(&turn_context.shell_environment_policy), - with_escalated_permissions: params.with_escalated_permissions, - justification: params.justification, - } -} - -fn parse_container_exec_arguments( - arguments: String, - turn_context: &TurnContext, - _call_id: &str, -) -> Result { - serde_json::from_str::(&arguments) - .map(|p| to_exec_params(p, turn_context)) - .map_err(|e| { - FunctionCallError::RespondToModel(format!("failed to parse function arguments: {e:?}")) - }) -} - -async fn handle_container_exec_with_params( +pub(crate) async fn handle_container_exec_with_params( tool_name: &str, params: ExecParams, sess: &Session, diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 4c7dfdcb3c..32f12adee4 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -66,6 +66,7 @@ pub mod shell; pub mod spawn; pub mod terminal; mod tool_apply_patch; +mod tools; pub mod turn_diff_tracker; pub use rollout::ARCHIVED_SESSIONS_SUBDIR; pub use rollout::RolloutRecorder; diff --git a/codex-rs/core/src/openai_tools.rs b/codex-rs/core/src/openai_tools.rs index 48dca79671..0e10f909ec 100644 --- a/codex-rs/core/src/openai_tools.rs +++ b/codex-rs/core/src/openai_tools.rs @@ -1,1046 +1 @@ -use serde::Deserialize; -use serde::Serialize; -use serde_json::Value as JsonValue; -use serde_json::json; -use std::collections::BTreeMap; -use std::collections::HashMap; - -use crate::model_family::ModelFamily; -use crate::plan_tool::PLAN_TOOL; -use crate::tool_apply_patch::ApplyPatchToolType; -use crate::tool_apply_patch::create_apply_patch_freeform_tool; -use crate::tool_apply_patch::create_apply_patch_json_tool; - -#[derive(Debug, Clone, Serialize, PartialEq)] -pub struct ResponsesApiTool { - pub(crate) name: String, - pub(crate) description: String, - /// TODO: Validation. When strict is set to true, the JSON schema, - /// `required` and `additional_properties` must be present. All fields in - /// `properties` must be present in `required`. - pub(crate) strict: bool, - pub(crate) parameters: JsonSchema, -} - -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] -pub struct FreeformTool { - pub(crate) name: String, - pub(crate) description: String, - pub(crate) format: FreeformToolFormat, -} - -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] -pub struct FreeformToolFormat { - pub(crate) r#type: String, - pub(crate) syntax: String, - pub(crate) definition: String, -} - -/// When serialized as JSON, this produces a valid "Tool" in the OpenAI -/// Responses API. -#[derive(Debug, Clone, Serialize, PartialEq)] -#[serde(tag = "type")] -pub(crate) enum OpenAiTool { - #[serde(rename = "function")] - Function(ResponsesApiTool), - #[serde(rename = "local_shell")] - LocalShell {}, - // TODO: Understand why we get an error on web_search although the API docs say it's supported. - // https://platform.openai.com/docs/guides/tools-web-search?api-mode=responses#:~:text=%7B%20type%3A%20%22web_search%22%20%7D%2C - #[serde(rename = "web_search")] - WebSearch {}, - #[serde(rename = "custom")] - Freeform(FreeformTool), -} - -#[derive(Debug, Clone)] -pub enum ConfigShellToolType { - Default, - Local, - Streamable, -} - -#[derive(Debug, Clone)] -pub(crate) struct ToolsConfig { - pub shell_type: ConfigShellToolType, - pub plan_tool: bool, - pub apply_patch_tool_type: Option, - pub web_search_request: bool, - pub include_view_image_tool: bool, - pub experimental_unified_exec_tool: bool, -} - -pub(crate) struct ToolsConfigParams<'a> { - pub(crate) model_family: &'a ModelFamily, - pub(crate) include_plan_tool: bool, - pub(crate) include_apply_patch_tool: bool, - pub(crate) include_web_search_request: bool, - pub(crate) use_streamable_shell_tool: bool, - pub(crate) include_view_image_tool: bool, - pub(crate) experimental_unified_exec_tool: bool, -} - -impl ToolsConfig { - pub fn new(params: &ToolsConfigParams) -> Self { - let ToolsConfigParams { - model_family, - include_plan_tool, - include_apply_patch_tool, - include_web_search_request, - use_streamable_shell_tool, - include_view_image_tool, - experimental_unified_exec_tool, - } = params; - let shell_type = if *use_streamable_shell_tool { - ConfigShellToolType::Streamable - } else if model_family.uses_local_shell_tool { - ConfigShellToolType::Local - } else { - ConfigShellToolType::Default - }; - - let apply_patch_tool_type = match model_family.apply_patch_tool_type { - Some(ApplyPatchToolType::Freeform) => Some(ApplyPatchToolType::Freeform), - Some(ApplyPatchToolType::Function) => Some(ApplyPatchToolType::Function), - None => { - if *include_apply_patch_tool { - Some(ApplyPatchToolType::Freeform) - } else { - None - } - } - }; - - Self { - shell_type, - plan_tool: *include_plan_tool, - apply_patch_tool_type, - web_search_request: *include_web_search_request, - include_view_image_tool: *include_view_image_tool, - experimental_unified_exec_tool: *experimental_unified_exec_tool, - } - } -} - -/// Generic JSON‑Schema subset needed for our tool definitions -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] -#[serde(tag = "type", rename_all = "lowercase")] -pub(crate) enum JsonSchema { - Boolean { - #[serde(skip_serializing_if = "Option::is_none")] - description: Option, - }, - String { - #[serde(skip_serializing_if = "Option::is_none")] - description: Option, - }, - /// MCP schema allows "number" | "integer" for Number - #[serde(alias = "integer")] - Number { - #[serde(skip_serializing_if = "Option::is_none")] - description: Option, - }, - Array { - items: Box, - - #[serde(skip_serializing_if = "Option::is_none")] - description: Option, - }, - Object { - properties: BTreeMap, - #[serde(skip_serializing_if = "Option::is_none")] - required: Option>, - #[serde( - rename = "additionalProperties", - skip_serializing_if = "Option::is_none" - )] - additional_properties: Option, - }, -} - -fn create_unified_exec_tool() -> OpenAiTool { - let mut properties = BTreeMap::new(); - properties.insert( - "input".to_string(), - JsonSchema::Array { - items: Box::new(JsonSchema::String { description: None }), - description: Some( - "When no session_id is provided, treat the array as the command and arguments \ - to launch. When session_id is set, concatenate the strings (in order) and write \ - them to the session's stdin." - .to_string(), - ), - }, - ); - properties.insert( - "session_id".to_string(), - JsonSchema::String { - description: Some( - "Identifier for an existing interactive session. If omitted, a new command \ - is spawned." - .to_string(), - ), - }, - ); - properties.insert( - "timeout_ms".to_string(), - JsonSchema::Number { - description: Some( - "Maximum time in milliseconds to wait for output after writing the input." - .to_string(), - ), - }, - ); - - OpenAiTool::Function(ResponsesApiTool { - name: "unified_exec".to_string(), - description: - "Runs a command in a PTY. Provide a session_id to reuse an existing interactive session.".to_string(), - strict: false, - parameters: JsonSchema::Object { - properties, - required: Some(vec!["input".to_string()]), - additional_properties: Some(false), - }, - }) -} - -fn create_shell_tool() -> OpenAiTool { - let mut properties = BTreeMap::new(); - properties.insert( - "command".to_string(), - JsonSchema::Array { - items: Box::new(JsonSchema::String { description: None }), - description: Some("The command to execute".to_string()), - }, - ); - properties.insert( - "workdir".to_string(), - JsonSchema::String { - description: Some("The working directory to execute the command in".to_string()), - }, - ); - properties.insert( - "timeout_ms".to_string(), - JsonSchema::Number { - description: Some("The timeout for the command in milliseconds".to_string()), - }, - ); - - properties.insert( - "with_escalated_permissions".to_string(), - JsonSchema::Boolean { - description: Some("Whether to request escalated permissions. Set to true if command needs to be run without sandbox restrictions".to_string()), - }, - ); - properties.insert( - "justification".to_string(), - JsonSchema::String { - description: Some("Only set if with_escalated_permissions is true. 1-sentence explanation of why we want to run this command.".to_string()), - }, - ); - - OpenAiTool::Function(ResponsesApiTool { - name: "shell".to_string(), - description: "Runs a shell command and returns its output.".to_string(), - strict: false, - parameters: JsonSchema::Object { - properties, - required: Some(vec!["command".to_string()]), - additional_properties: Some(false), - }, - }) -} - -fn create_view_image_tool() -> OpenAiTool { - // Support only local filesystem path. - let mut properties = BTreeMap::new(); - properties.insert( - "path".to_string(), - JsonSchema::String { - description: Some("Local filesystem path to an image file".to_string()), - }, - ); - - OpenAiTool::Function(ResponsesApiTool { - name: "view_image".to_string(), - description: - "Attach a local image (by filesystem path) to the conversation context for this turn." - .to_string(), - strict: false, - parameters: JsonSchema::Object { - properties, - required: Some(vec!["path".to_string()]), - additional_properties: Some(false), - }, - }) -} -/// TODO(dylan): deprecate once we get rid of json tool -#[derive(Serialize, Deserialize)] -pub(crate) struct ApplyPatchToolArgs { - pub(crate) input: String, -} - -/// Returns JSON values that are compatible with Function Calling in the -/// Responses API: -/// https://platform.openai.com/docs/guides/function-calling?api-mode=responses -pub fn create_tools_json_for_responses_api( - tools: &[OpenAiTool], -) -> crate::error::Result> { - let mut tools_json = Vec::new(); - - for tool in tools { - let json = serde_json::to_value(tool)?; - tools_json.push(json); - } - - Ok(tools_json) -} -/// Returns JSON values that are compatible with Function Calling in the -/// Chat Completions API: -/// https://platform.openai.com/docs/guides/function-calling?api-mode=chat -pub(crate) fn create_tools_json_for_chat_completions_api( - tools: &[OpenAiTool], -) -> crate::error::Result> { - // We start with the JSON for the Responses API and than rewrite it to match - // the chat completions tool call format. - let responses_api_tools_json = create_tools_json_for_responses_api(tools)?; - let tools_json = responses_api_tools_json - .into_iter() - .filter_map(|mut tool| { - if tool.get("type") != Some(&serde_json::Value::String("function".to_string())) { - return None; - } - - if let Some(map) = tool.as_object_mut() { - // Remove "type" field as it is not needed in chat completions. - map.remove("type"); - Some(json!({ - "type": "function", - "function": map, - })) - } else { - None - } - }) - .collect::>(); - Ok(tools_json) -} - -pub(crate) fn mcp_tool_to_openai_tool( - fully_qualified_name: String, - tool: mcp_types::Tool, -) -> Result { - let mcp_types::Tool { - description, - mut input_schema, - .. - } = tool; - - // OpenAI models mandate the "properties" field in the schema. The Agents - // SDK fixed this by inserting an empty object for "properties" if it is not - // already present https://github.com/openai/openai-agents-python/issues/449 - // so here we do the same. - if input_schema.properties.is_none() { - input_schema.properties = Some(serde_json::Value::Object(serde_json::Map::new())); - } - - // Serialize to a raw JSON value so we can sanitize schemas coming from MCP - // servers. Some servers omit the top-level or nested `type` in JSON - // Schemas (e.g. using enum/anyOf), or use unsupported variants like - // `integer`. Our internal JsonSchema is a small subset and requires - // `type`, so we coerce/sanitize here for compatibility. - let mut serialized_input_schema = serde_json::to_value(input_schema)?; - sanitize_json_schema(&mut serialized_input_schema); - let input_schema = serde_json::from_value::(serialized_input_schema)?; - - Ok(ResponsesApiTool { - name: fully_qualified_name, - description: description.unwrap_or_default(), - strict: false, - parameters: input_schema, - }) -} - -/// Sanitize a JSON Schema (as serde_json::Value) so it can fit our limited -/// JsonSchema enum. This function: -/// - Ensures every schema object has a "type". If missing, infers it from -/// common keywords (properties => object, items => array, enum/const/format => string) -/// and otherwise defaults to "string". -/// - Fills required child fields (e.g. array items, object properties) with -/// permissive defaults when absent. -fn sanitize_json_schema(value: &mut JsonValue) { - match value { - JsonValue::Bool(_) => { - // JSON Schema boolean form: true/false. Coerce to an accept-all string. - *value = json!({ "type": "string" }); - } - JsonValue::Array(arr) => { - for v in arr.iter_mut() { - sanitize_json_schema(v); - } - } - JsonValue::Object(map) => { - // First, recursively sanitize known nested schema holders - if let Some(props) = map.get_mut("properties") - && let Some(props_map) = props.as_object_mut() - { - for (_k, v) in props_map.iter_mut() { - sanitize_json_schema(v); - } - } - if let Some(items) = map.get_mut("items") { - sanitize_json_schema(items); - } - // Some schemas use oneOf/anyOf/allOf - sanitize their entries - for combiner in ["oneOf", "anyOf", "allOf", "prefixItems"] { - if let Some(v) = map.get_mut(combiner) { - sanitize_json_schema(v); - } - } - - // Normalize/ensure type - let mut ty = map.get("type").and_then(|v| v.as_str()).map(str::to_string); - - // If type is an array (union), pick first supported; else leave to inference - if ty.is_none() - && let Some(JsonValue::Array(types)) = map.get("type") - { - for t in types { - if let Some(tt) = t.as_str() - && matches!( - tt, - "object" | "array" | "string" | "number" | "integer" | "boolean" - ) - { - ty = Some(tt.to_string()); - break; - } - } - } - - // Infer type if still missing - if ty.is_none() { - if map.contains_key("properties") - || map.contains_key("required") - || map.contains_key("additionalProperties") - { - ty = Some("object".to_string()); - } else if map.contains_key("items") || map.contains_key("prefixItems") { - ty = Some("array".to_string()); - } else if map.contains_key("enum") - || map.contains_key("const") - || map.contains_key("format") - { - ty = Some("string".to_string()); - } else if map.contains_key("minimum") - || map.contains_key("maximum") - || map.contains_key("exclusiveMinimum") - || map.contains_key("exclusiveMaximum") - || map.contains_key("multipleOf") - { - ty = Some("number".to_string()); - } - } - // If we still couldn't infer, default to string - let ty = ty.unwrap_or_else(|| "string".to_string()); - map.insert("type".to_string(), JsonValue::String(ty.to_string())); - - // Ensure object schemas have properties map - if ty == "object" { - if !map.contains_key("properties") { - map.insert( - "properties".to_string(), - JsonValue::Object(serde_json::Map::new()), - ); - } - // If additionalProperties is an object schema, sanitize it too. - // Leave booleans as-is, since JSON Schema allows boolean here. - if let Some(ap) = map.get_mut("additionalProperties") { - let is_bool = matches!(ap, JsonValue::Bool(_)); - if !is_bool { - sanitize_json_schema(ap); - } - } - } - - // Ensure array schemas have items - if ty == "array" && !map.contains_key("items") { - map.insert("items".to_string(), json!({ "type": "string" })); - } - } - _ => {} - } -} - -/// Returns a list of OpenAiTools based on the provided config and MCP tools. -/// Note that the keys of mcp_tools should be fully qualified names. See -/// [`McpConnectionManager`] for more details. -pub(crate) fn get_openai_tools( - config: &ToolsConfig, - mcp_tools: Option>, -) -> Vec { - let mut tools: Vec = Vec::new(); - - if config.experimental_unified_exec_tool { - tools.push(create_unified_exec_tool()); - } else { - match &config.shell_type { - ConfigShellToolType::Default => { - tools.push(create_shell_tool()); - } - ConfigShellToolType::Local => { - tools.push(OpenAiTool::LocalShell {}); - } - ConfigShellToolType::Streamable => { - tools.push(OpenAiTool::Function( - crate::exec_command::create_exec_command_tool_for_responses_api(), - )); - tools.push(OpenAiTool::Function( - crate::exec_command::create_write_stdin_tool_for_responses_api(), - )); - } - } - } - - if config.plan_tool { - tools.push(PLAN_TOOL.clone()); - } - - if let Some(apply_patch_tool_type) = &config.apply_patch_tool_type { - match apply_patch_tool_type { - ApplyPatchToolType::Freeform => { - tools.push(create_apply_patch_freeform_tool()); - } - ApplyPatchToolType::Function => { - tools.push(create_apply_patch_json_tool()); - } - } - } - - if config.web_search_request { - tools.push(OpenAiTool::WebSearch {}); - } - - // Include the view_image tool so the agent can attach images to context. - if config.include_view_image_tool { - tools.push(create_view_image_tool()); - } - if let Some(mcp_tools) = mcp_tools { - // Ensure deterministic ordering to maximize prompt cache hits. - let mut entries: Vec<(String, mcp_types::Tool)> = mcp_tools.into_iter().collect(); - entries.sort_by(|a, b| a.0.cmp(&b.0)); - - for (name, tool) in entries.into_iter() { - match mcp_tool_to_openai_tool(name.clone(), tool.clone()) { - Ok(converted_tool) => tools.push(OpenAiTool::Function(converted_tool)), - Err(e) => { - tracing::error!("Failed to convert {name:?} MCP tool to OpenAI tool: {e:?}"); - } - } - } - } - - tools -} - -#[cfg(test)] -mod tests { - use crate::model_family::find_family_for_model; - use mcp_types::ToolInputSchema; - use pretty_assertions::assert_eq; - - use super::*; - - fn assert_eq_tool_names(tools: &[OpenAiTool], expected_names: &[&str]) { - let tool_names = tools - .iter() - .map(|tool| match tool { - OpenAiTool::Function(ResponsesApiTool { name, .. }) => name, - OpenAiTool::LocalShell {} => "local_shell", - OpenAiTool::WebSearch {} => "web_search", - OpenAiTool::Freeform(FreeformTool { name, .. }) => name, - }) - .collect::>(); - - assert_eq!( - tool_names.len(), - expected_names.len(), - "tool_name mismatch, {tool_names:?}, {expected_names:?}", - ); - for (name, expected_name) in tool_names.iter().zip(expected_names.iter()) { - assert_eq!( - name, expected_name, - "tool_name mismatch, {name:?}, {expected_name:?}" - ); - } - } - - #[test] - fn test_get_openai_tools() { - let model_family = find_family_for_model("codex-mini-latest") - .expect("codex-mini-latest should be a valid model family"); - let config = ToolsConfig::new(&ToolsConfigParams { - model_family: &model_family, - include_plan_tool: true, - include_apply_patch_tool: false, - include_web_search_request: true, - use_streamable_shell_tool: false, - include_view_image_tool: true, - experimental_unified_exec_tool: true, - }); - let tools = get_openai_tools(&config, Some(HashMap::new())); - - assert_eq_tool_names( - &tools, - &["unified_exec", "update_plan", "web_search", "view_image"], - ); - } - - #[test] - fn test_get_openai_tools_default_shell() { - let model_family = find_family_for_model("o3").expect("o3 should be a valid model family"); - let config = ToolsConfig::new(&ToolsConfigParams { - model_family: &model_family, - include_plan_tool: true, - include_apply_patch_tool: false, - include_web_search_request: true, - use_streamable_shell_tool: false, - include_view_image_tool: true, - experimental_unified_exec_tool: true, - }); - let tools = get_openai_tools(&config, Some(HashMap::new())); - - assert_eq_tool_names( - &tools, - &["unified_exec", "update_plan", "web_search", "view_image"], - ); - } - - #[test] - fn test_get_openai_tools_mcp_tools() { - let model_family = find_family_for_model("o3").expect("o3 should be a valid model family"); - let config = ToolsConfig::new(&ToolsConfigParams { - model_family: &model_family, - include_plan_tool: false, - include_apply_patch_tool: false, - include_web_search_request: true, - use_streamable_shell_tool: false, - include_view_image_tool: true, - experimental_unified_exec_tool: true, - }); - let tools = get_openai_tools( - &config, - Some(HashMap::from([( - "test_server/do_something_cool".to_string(), - mcp_types::Tool { - name: "do_something_cool".to_string(), - input_schema: ToolInputSchema { - properties: Some(serde_json::json!({ - "string_argument": { - "type": "string", - }, - "number_argument": { - "type": "number", - }, - "object_argument": { - "type": "object", - "properties": { - "string_property": { "type": "string" }, - "number_property": { "type": "number" }, - }, - "required": [ - "string_property", - "number_property", - ], - "additionalProperties": Some(false), - }, - })), - required: None, - r#type: "object".to_string(), - }, - output_schema: None, - title: None, - annotations: None, - description: Some("Do something cool".to_string()), - }, - )])), - ); - - assert_eq_tool_names( - &tools, - &[ - "unified_exec", - "web_search", - "view_image", - "test_server/do_something_cool", - ], - ); - - assert_eq!( - tools[3], - OpenAiTool::Function(ResponsesApiTool { - name: "test_server/do_something_cool".to_string(), - parameters: JsonSchema::Object { - properties: BTreeMap::from([ - ( - "string_argument".to_string(), - JsonSchema::String { description: None } - ), - ( - "number_argument".to_string(), - JsonSchema::Number { description: None } - ), - ( - "object_argument".to_string(), - JsonSchema::Object { - properties: BTreeMap::from([ - ( - "string_property".to_string(), - JsonSchema::String { description: None } - ), - ( - "number_property".to_string(), - JsonSchema::Number { description: None } - ), - ]), - required: Some(vec![ - "string_property".to_string(), - "number_property".to_string(), - ]), - additional_properties: Some(false), - }, - ), - ]), - required: None, - additional_properties: None, - }, - description: "Do something cool".to_string(), - strict: false, - }) - ); - } - - #[test] - fn test_get_openai_tools_mcp_tools_sorted_by_name() { - let model_family = find_family_for_model("o3").expect("o3 should be a valid model family"); - let config = ToolsConfig::new(&ToolsConfigParams { - model_family: &model_family, - include_plan_tool: false, - include_apply_patch_tool: false, - include_web_search_request: false, - use_streamable_shell_tool: false, - include_view_image_tool: true, - experimental_unified_exec_tool: true, - }); - - // Intentionally construct a map with keys that would sort alphabetically. - let tools_map: HashMap = HashMap::from([ - ( - "test_server/do".to_string(), - mcp_types::Tool { - name: "a".to_string(), - input_schema: ToolInputSchema { - properties: Some(serde_json::json!({})), - required: None, - r#type: "object".to_string(), - }, - output_schema: None, - title: None, - annotations: None, - description: Some("a".to_string()), - }, - ), - ( - "test_server/something".to_string(), - mcp_types::Tool { - name: "b".to_string(), - input_schema: ToolInputSchema { - properties: Some(serde_json::json!({})), - required: None, - r#type: "object".to_string(), - }, - output_schema: None, - title: None, - annotations: None, - description: Some("b".to_string()), - }, - ), - ( - "test_server/cool".to_string(), - mcp_types::Tool { - name: "c".to_string(), - input_schema: ToolInputSchema { - properties: Some(serde_json::json!({})), - required: None, - r#type: "object".to_string(), - }, - output_schema: None, - title: None, - annotations: None, - description: Some("c".to_string()), - }, - ), - ]); - - let tools = get_openai_tools(&config, Some(tools_map)); - // Expect unified_exec first, followed by MCP tools sorted by fully-qualified name. - assert_eq_tool_names( - &tools, - &[ - "unified_exec", - "view_image", - "test_server/cool", - "test_server/do", - "test_server/something", - ], - ); - } - - #[test] - fn test_mcp_tool_property_missing_type_defaults_to_string() { - let model_family = find_family_for_model("o3").expect("o3 should be a valid model family"); - let config = ToolsConfig::new(&ToolsConfigParams { - model_family: &model_family, - include_plan_tool: false, - include_apply_patch_tool: false, - include_web_search_request: true, - use_streamable_shell_tool: false, - include_view_image_tool: true, - experimental_unified_exec_tool: true, - }); - - let tools = get_openai_tools( - &config, - Some(HashMap::from([( - "dash/search".to_string(), - mcp_types::Tool { - name: "search".to_string(), - input_schema: ToolInputSchema { - properties: Some(serde_json::json!({ - "query": { - "description": "search query" - } - })), - required: None, - r#type: "object".to_string(), - }, - output_schema: None, - title: None, - annotations: None, - description: Some("Search docs".to_string()), - }, - )])), - ); - - assert_eq_tool_names( - &tools, - &["unified_exec", "web_search", "view_image", "dash/search"], - ); - - assert_eq!( - tools[3], - OpenAiTool::Function(ResponsesApiTool { - name: "dash/search".to_string(), - parameters: JsonSchema::Object { - properties: BTreeMap::from([( - "query".to_string(), - JsonSchema::String { - description: Some("search query".to_string()) - } - )]), - required: None, - additional_properties: None, - }, - description: "Search docs".to_string(), - strict: false, - }) - ); - } - - #[test] - fn test_mcp_tool_integer_normalized_to_number() { - let model_family = find_family_for_model("o3").expect("o3 should be a valid model family"); - let config = ToolsConfig::new(&ToolsConfigParams { - model_family: &model_family, - include_plan_tool: false, - include_apply_patch_tool: false, - include_web_search_request: true, - use_streamable_shell_tool: false, - include_view_image_tool: true, - experimental_unified_exec_tool: true, - }); - - let tools = get_openai_tools( - &config, - Some(HashMap::from([( - "dash/paginate".to_string(), - mcp_types::Tool { - name: "paginate".to_string(), - input_schema: ToolInputSchema { - properties: Some(serde_json::json!({ - "page": { "type": "integer" } - })), - required: None, - r#type: "object".to_string(), - }, - output_schema: None, - title: None, - annotations: None, - description: Some("Pagination".to_string()), - }, - )])), - ); - - assert_eq_tool_names( - &tools, - &["unified_exec", "web_search", "view_image", "dash/paginate"], - ); - assert_eq!( - tools[3], - OpenAiTool::Function(ResponsesApiTool { - name: "dash/paginate".to_string(), - parameters: JsonSchema::Object { - properties: BTreeMap::from([( - "page".to_string(), - JsonSchema::Number { description: None } - )]), - required: None, - additional_properties: None, - }, - description: "Pagination".to_string(), - strict: false, - }) - ); - } - - #[test] - fn test_mcp_tool_array_without_items_gets_default_string_items() { - let model_family = find_family_for_model("o3").expect("o3 should be a valid model family"); - let config = ToolsConfig::new(&ToolsConfigParams { - model_family: &model_family, - include_plan_tool: false, - include_apply_patch_tool: false, - include_web_search_request: true, - use_streamable_shell_tool: false, - include_view_image_tool: true, - experimental_unified_exec_tool: true, - }); - - let tools = get_openai_tools( - &config, - Some(HashMap::from([( - "dash/tags".to_string(), - mcp_types::Tool { - name: "tags".to_string(), - input_schema: ToolInputSchema { - properties: Some(serde_json::json!({ - "tags": { "type": "array" } - })), - required: None, - r#type: "object".to_string(), - }, - output_schema: None, - title: None, - annotations: None, - description: Some("Tags".to_string()), - }, - )])), - ); - - assert_eq_tool_names( - &tools, - &["unified_exec", "web_search", "view_image", "dash/tags"], - ); - assert_eq!( - tools[3], - OpenAiTool::Function(ResponsesApiTool { - name: "dash/tags".to_string(), - parameters: JsonSchema::Object { - properties: BTreeMap::from([( - "tags".to_string(), - JsonSchema::Array { - items: Box::new(JsonSchema::String { description: None }), - description: None - } - )]), - required: None, - additional_properties: None, - }, - description: "Tags".to_string(), - strict: false, - }) - ); - } - - #[test] - fn test_mcp_tool_anyof_defaults_to_string() { - let model_family = find_family_for_model("o3").expect("o3 should be a valid model family"); - let config = ToolsConfig::new(&ToolsConfigParams { - model_family: &model_family, - include_plan_tool: false, - include_apply_patch_tool: false, - include_web_search_request: true, - use_streamable_shell_tool: false, - include_view_image_tool: true, - experimental_unified_exec_tool: true, - }); - - let tools = get_openai_tools( - &config, - Some(HashMap::from([( - "dash/value".to_string(), - mcp_types::Tool { - name: "value".to_string(), - input_schema: ToolInputSchema { - properties: Some(serde_json::json!({ - "value": { "anyOf": [ { "type": "string" }, { "type": "number" } ] } - })), - required: None, - r#type: "object".to_string(), - }, - output_schema: None, - title: None, - annotations: None, - description: Some("AnyOf Value".to_string()), - }, - )])), - ); - - assert_eq_tool_names( - &tools, - &["unified_exec", "web_search", "view_image", "dash/value"], - ); - assert_eq!( - tools[3], - OpenAiTool::Function(ResponsesApiTool { - name: "dash/value".to_string(), - parameters: JsonSchema::Object { - properties: BTreeMap::from([( - "value".to_string(), - JsonSchema::String { description: None } - )]), - required: None, - additional_properties: None, - }, - description: "AnyOf Value".to_string(), - strict: false, - }) - ); - } - - #[test] - fn test_shell_tool() { - let tool = super::create_shell_tool(); - let OpenAiTool::Function(ResponsesApiTool { - description, name, .. - }) = &tool - else { - panic!("expected function tool"); - }; - assert_eq!(name, "shell"); - - let expected = "Runs a shell command and returns its output."; - assert_eq!(description, expected); - } -} +pub use crate::tools::spec::*; diff --git a/codex-rs/core/src/plan_tool.rs b/codex-rs/core/src/plan_tool.rs index c81675698f..b4be53fea3 100644 --- a/codex-rs/core/src/plan_tool.rs +++ b/codex-rs/core/src/plan_tool.rs @@ -4,8 +4,8 @@ use std::sync::LazyLock; use crate::codex::Session; use crate::function_tool::FunctionCallError; use crate::openai_tools::JsonSchema; -use crate::openai_tools::OpenAiTool; use crate::openai_tools::ResponsesApiTool; +use crate::openai_tools::ToolSpec; use crate::protocol::Event; use crate::protocol::EventMsg; @@ -17,7 +17,7 @@ pub use codex_protocol::plan_tool::UpdatePlanArgs; // Types for the TODO tool arguments matching codex-vscode/todo-mcp/src/main.rs -pub(crate) static PLAN_TOOL: LazyLock = LazyLock::new(|| { +pub(crate) static PLAN_TOOL: LazyLock = LazyLock::new(|| { let mut plan_item_props = BTreeMap::new(); plan_item_props.insert("step".to_string(), JsonSchema::String { description: None }); plan_item_props.insert( @@ -43,7 +43,7 @@ pub(crate) static PLAN_TOOL: LazyLock = LazyLock::new(|| { ); properties.insert("plan".to_string(), plan_items_schema); - OpenAiTool::Function(ResponsesApiTool { + ToolSpec::Function(ResponsesApiTool { name: "update_plan".to_string(), description: r#"Updates the task plan. Provide an optional explanation and a list of plan items, each with a step and status. diff --git a/codex-rs/core/src/tool_apply_patch.rs b/codex-rs/core/src/tool_apply_patch.rs index 164a1e127f..7d1ea2a528 100644 --- a/codex-rs/core/src/tool_apply_patch.rs +++ b/codex-rs/core/src/tool_apply_patch.rs @@ -5,8 +5,8 @@ use std::collections::BTreeMap; use crate::openai_tools::FreeformTool; use crate::openai_tools::FreeformToolFormat; use crate::openai_tools::JsonSchema; -use crate::openai_tools::OpenAiTool; use crate::openai_tools::ResponsesApiTool; +use crate::openai_tools::ToolSpec; const APPLY_PATCH_LARK_GRAMMAR: &str = include_str!("tool_apply_patch.lark"); @@ -19,8 +19,8 @@ pub enum ApplyPatchToolType { /// Returns a custom tool that can be used to edit files. Well-suited for GPT-5 models /// https://platform.openai.com/docs/guides/function-calling#custom-tools -pub(crate) fn create_apply_patch_freeform_tool() -> OpenAiTool { - OpenAiTool::Freeform(FreeformTool { +pub(crate) fn create_apply_patch_freeform_tool() -> ToolSpec { + ToolSpec::Freeform(FreeformTool { name: "apply_patch".to_string(), description: "Use the `apply_patch` tool to edit files".to_string(), format: FreeformToolFormat { @@ -32,7 +32,7 @@ pub(crate) fn create_apply_patch_freeform_tool() -> OpenAiTool { } /// Returns a json tool that can be used to edit files. Should only be used with gpt-oss models -pub(crate) fn create_apply_patch_json_tool() -> OpenAiTool { +pub(crate) fn create_apply_patch_json_tool() -> ToolSpec { let mut properties = BTreeMap::new(); properties.insert( "input".to_string(), @@ -41,7 +41,7 @@ pub(crate) fn create_apply_patch_json_tool() -> OpenAiTool { }, ); - OpenAiTool::Function(ResponsesApiTool { + ToolSpec::Function(ResponsesApiTool { name: "apply_patch".to_string(), description: r#"Use the `apply_patch` tool to edit files. Your patch language is a stripped‑down, file‑oriented diff format designed to be easy to parse and safe to apply. You can think of it as a high‑level envelope: diff --git a/codex-rs/core/src/tools/context.rs b/codex-rs/core/src/tools/context.rs new file mode 100644 index 0000000000..dd4aef3d60 --- /dev/null +++ b/codex-rs/core/src/tools/context.rs @@ -0,0 +1,111 @@ +use std::borrow::Cow; +use std::collections::HashMap; +use std::path::PathBuf; +use codex_otel::otel_event_manager::OtelEventManager; +use codex_protocol::models::FunctionCallOutputPayload; +use codex_protocol::models::ResponseInputItem; +use codex_protocol::models::ShellToolCallParams; +use codex_protocol::protocol::FileChange; +use mcp_types::CallToolResult; + +use crate::codex::Session; +use crate::codex::TurnContext; +use crate::turn_diff_tracker::TurnDiffTracker; + +pub struct ToolInvocation<'a> { + pub session: &'a Session, + pub turn: &'a TurnContext, + pub tracker: &'a mut TurnDiffTracker, + pub sub_id: &'a str, + pub call_id: String, + pub tool_name: String, + pub payload: ToolPayload, +} + +#[derive(Clone)] +pub enum ToolPayload { + Function { + arguments: String, + }, + Custom { + input: String, + }, + LocalShell { + params: ShellToolCallParams, + }, + UnifiedExec { + arguments: String, + }, + Mcp { + server: String, + tool: String, + arguments: Option, + raw_arguments: String, + }, +} + +impl ToolPayload { + pub fn log_payload(&self) -> Cow<'_, str> { + match self { + ToolPayload::Function { arguments } => Cow::Borrowed(arguments), + ToolPayload::Custom { input } => Cow::Borrowed(input), + ToolPayload::LocalShell { params } => Cow::Owned(params.command.join(" ")), + ToolPayload::UnifiedExec { arguments } => Cow::Borrowed(arguments), + ToolPayload::Mcp { raw_arguments, .. } => Cow::Borrowed(raw_arguments), + } + } +} + +#[derive(Clone)] +pub enum ToolOutput { + Function { + content: String, + success: bool, + }, + Mcp { + result: Result, + }, +} + +impl ToolOutput { + pub fn log_preview(&self) -> String { + match self { + ToolOutput::Function { content, .. } => content.clone(), + ToolOutput::Mcp { result } => format!("{result:?}"), + } + } + + pub fn into_response(self, call_id: &str) -> ResponseInputItem { + match self { + ToolOutput::Function { content, success } => ResponseInputItem::FunctionCallOutput { + call_id: call_id.to_string(), + output: FunctionCallOutputPayload { + content, + success: Some(success), + }, + }, + ToolOutput::Mcp { result } => ResponseInputItem::McpToolCallOutput { + call_id: call_id.to_string(), + result, + }, + } + } +} + + +#[derive(Clone, Debug)] +pub(crate) struct ExecCommandContext { + pub(crate) sub_id: String, + pub(crate) call_id: String, + pub(crate) command_for_display: Vec, + pub(crate) cwd: PathBuf, + pub(crate) apply_patch: Option, + pub(crate) tool_name: String, + pub(crate) otel_event_manager: OtelEventManager, +} + +#[derive(Clone, Debug)] +pub(crate) struct ApplyPatchCommandContext { + pub(crate) user_explicitly_approved_this_action: bool, + pub(crate) changes: HashMap, +} \ No newline at end of file diff --git a/codex-rs/core/src/tools/handlers/apply_patch.rs b/codex-rs/core/src/tools/handlers/apply_patch.rs new file mode 100644 index 0000000000..4a38de4ef1 --- /dev/null +++ b/codex-rs/core/src/tools/handlers/apply_patch.rs @@ -0,0 +1,85 @@ +use std::collections::HashMap; + +use async_trait::async_trait; + +use crate::exec::ExecParams; +use crate::function_tool::FunctionCallError; +use crate::tools::context::ToolInvocation; +use crate::tools::context::ToolOutput; +use crate::tools::context::ToolPayload; +use crate::tools::registry::ToolHandler; +use crate::tools::registry::ToolKind; +use crate::tools::spec::ApplyPatchToolArgs; + +pub struct ApplyPatchHandler; + +#[async_trait] +impl ToolHandler for ApplyPatchHandler { + fn kind(&self) -> ToolKind { + ToolKind::Function + } + + fn matches_kind(&self, payload: &ToolPayload) -> bool { + matches!( + payload, + ToolPayload::Function { .. } | ToolPayload::Custom { .. } + ) + } + + async fn handle( + &self, + invocation: ToolInvocation<'_>, + ) -> Result { + let ToolInvocation { + session, + turn, + tracker, + sub_id, + call_id, + tool_name, + payload, + } = invocation; + + let patch_input = match payload { + ToolPayload::Function { arguments } => { + let args: ApplyPatchToolArgs = serde_json::from_str(&arguments).map_err(|e| { + FunctionCallError::RespondToModel(format!( + "failed to parse function arguments: {e:?}" + )) + })?; + args.input + } + ToolPayload::Custom { input } => input, + _ => { + return Err(FunctionCallError::RespondToModel( + "apply_patch handler received unsupported payload".to_string(), + )); + } + }; + + let exec_params = ExecParams { + command: vec!["apply_patch".to_string(), patch_input.clone()], + cwd: turn.cwd.clone(), + timeout_ms: None, + env: HashMap::new(), + with_escalated_permissions: None, + justification: None, + }; + + let content = crate::codex::handle_container_exec_with_params( + tool_name.as_str(), + exec_params, + session, + turn, + tracker, + sub_id.to_string(), + call_id.clone(), + ) + .await?; + + Ok(ToolOutput::Function { + content, + success: true, + }) + } +} diff --git a/codex-rs/core/src/tools/handlers/exec_stream.rs b/codex-rs/core/src/tools/handlers/exec_stream.rs new file mode 100644 index 0000000000..d931d3ddbd --- /dev/null +++ b/codex-rs/core/src/tools/handlers/exec_stream.rs @@ -0,0 +1,71 @@ +use async_trait::async_trait; + +use crate::exec_command::EXEC_COMMAND_TOOL_NAME; +use crate::exec_command::ExecCommandParams; +use crate::exec_command::WRITE_STDIN_TOOL_NAME; +use crate::exec_command::WriteStdinParams; +use crate::function_tool::FunctionCallError; +use crate::tools::context::ToolInvocation; +use crate::tools::context::ToolOutput; +use crate::tools::context::ToolPayload; +use crate::tools::registry::ToolHandler; +use crate::tools::registry::ToolKind; + +pub struct ExecStreamHandler; + +#[async_trait] +impl ToolHandler for ExecStreamHandler { + fn kind(&self) -> ToolKind { + ToolKind::Function + } + + async fn handle( + &self, + invocation: ToolInvocation<'_>, + ) -> Result { + let ToolInvocation { + session, + tool_name, + payload, + .. + } = invocation; + + let arguments = match payload { + ToolPayload::Function { arguments } => arguments, + _ => { + return Err(FunctionCallError::RespondToModel( + "exec_stream handler received unsupported payload".to_string(), + )); + } + }; + + let content = match tool_name.as_str() { + EXEC_COMMAND_TOOL_NAME => { + let params: ExecCommandParams = serde_json::from_str(&arguments).map_err(|e| { + FunctionCallError::RespondToModel(format!( + "failed to parse function arguments: {e:?}" + )) + })?; + session.handle_exec_command_tool(params).await? + } + WRITE_STDIN_TOOL_NAME => { + let params: WriteStdinParams = serde_json::from_str(&arguments).map_err(|e| { + FunctionCallError::RespondToModel(format!( + "failed to parse function arguments: {e:?}" + )) + })?; + session.handle_write_stdin_tool(params).await? + } + _ => { + return Err(FunctionCallError::RespondToModel(format!( + "exec_stream handler does not support tool {tool_name}" + ))); + } + }; + + Ok(ToolOutput::Function { + content, + success: true, + }) + } +} diff --git a/codex-rs/core/src/tools/handlers/mcp.rs b/codex-rs/core/src/tools/handlers/mcp.rs new file mode 100644 index 0000000000..8379915aaa --- /dev/null +++ b/codex-rs/core/src/tools/handlers/mcp.rs @@ -0,0 +1,74 @@ +use async_trait::async_trait; + +use crate::function_tool::FunctionCallError; +use crate::mcp_tool_call::handle_mcp_tool_call; +use crate::tools::context::ToolInvocation; +use crate::tools::context::ToolOutput; +use crate::tools::context::ToolPayload; +use crate::tools::registry::ToolHandler; +use crate::tools::registry::ToolKind; + +pub struct McpHandler; + +#[async_trait] +impl ToolHandler for McpHandler { + fn kind(&self) -> ToolKind { + ToolKind::Mcp + } + + async fn handle( + &self, + invocation: ToolInvocation<'_>, + ) -> Result { + let ToolInvocation { + session, + sub_id, + call_id, + payload, + .. + } = invocation; + + let payload = match payload { + ToolPayload::Mcp { + server, + tool, + arguments, + raw_arguments, + } => (server, tool, arguments, raw_arguments), + _ => { + return Err(FunctionCallError::RespondToModel( + "mcp handler received unsupported payload".to_string(), + )); + } + }; + + let (server, tool, _arguments, raw_arguments) = payload; + let arguments_str = raw_arguments; + + let response = handle_mcp_tool_call( + session, + sub_id, + call_id.clone(), + server, + tool, + arguments_str, + ) + .await; + + match response { + codex_protocol::models::ResponseInputItem::McpToolCallOutput { result, .. } => { + Ok(ToolOutput::Mcp { result }) + } + codex_protocol::models::ResponseInputItem::FunctionCallOutput { output, .. } => { + let success = output.success.unwrap_or(false); + Ok(ToolOutput::Function { + content: output.content, + success, + }) + } + _ => Err(FunctionCallError::RespondToModel( + "mcp handler received unexpected response variant".to_string(), + )), + } + } +} diff --git a/codex-rs/core/src/tools/handlers/mod.rs b/codex-rs/core/src/tools/handlers/mod.rs new file mode 100644 index 0000000000..93249f8873 --- /dev/null +++ b/codex-rs/core/src/tools/handlers/mod.rs @@ -0,0 +1,15 @@ +mod apply_patch; +mod exec_stream; +mod mcp; +mod plan; +mod shell; +mod unified_exec; +mod view_image; + +pub use apply_patch::ApplyPatchHandler; +pub use exec_stream::ExecStreamHandler; +pub use mcp::McpHandler; +pub use plan::PlanHandler; +pub use shell::ShellHandler; +pub use unified_exec::UnifiedExecHandler; +pub use view_image::ViewImageHandler; diff --git a/codex-rs/core/src/tools/handlers/plan.rs b/codex-rs/core/src/tools/handlers/plan.rs new file mode 100644 index 0000000000..f1dc584177 --- /dev/null +++ b/codex-rs/core/src/tools/handlers/plan.rs @@ -0,0 +1,47 @@ +use async_trait::async_trait; + +use crate::function_tool::FunctionCallError; +use crate::plan_tool::handle_update_plan; +use crate::tools::context::ToolInvocation; +use crate::tools::context::ToolOutput; +use crate::tools::context::ToolPayload; +use crate::tools::registry::ToolHandler; +use crate::tools::registry::ToolKind; + +pub struct PlanHandler; + +#[async_trait] +impl ToolHandler for PlanHandler { + fn kind(&self) -> ToolKind { + ToolKind::Function + } + + async fn handle( + &self, + invocation: ToolInvocation<'_>, + ) -> Result { + let ToolInvocation { + session, + sub_id, + call_id, + payload, + .. + } = invocation; + + let arguments = match payload { + ToolPayload::Function { arguments } => arguments, + _ => { + return Err(FunctionCallError::RespondToModel( + "update_plan handler received unsupported payload".to_string(), + )); + } + }; + + let content = handle_update_plan(session, arguments, sub_id.to_string(), call_id).await?; + + Ok(ToolOutput::Function { + content, + success: true, + }) + } +} diff --git a/codex-rs/core/src/tools/handlers/shell.rs b/codex-rs/core/src/tools/handlers/shell.rs new file mode 100644 index 0000000000..5e98a79add --- /dev/null +++ b/codex-rs/core/src/tools/handlers/shell.rs @@ -0,0 +1,102 @@ +use async_trait::async_trait; +use codex_protocol::models::ShellToolCallParams; + +use crate::codex::TurnContext; +use crate::exec::ExecParams; +use crate::exec_env::create_env; +use crate::function_tool::FunctionCallError; +use crate::tools::context::ToolInvocation; +use crate::tools::context::ToolOutput; +use crate::tools::context::ToolPayload; +use crate::tools::registry::ToolHandler; +use crate::tools::registry::ToolKind; + +pub struct ShellHandler; + +impl ShellHandler { + fn to_exec_params(params: ShellToolCallParams, turn_context: &TurnContext) -> ExecParams { + ExecParams { + command: params.command, + cwd: turn_context.resolve_path(params.workdir.clone()), + timeout_ms: params.timeout_ms, + env: create_env(&turn_context.shell_environment_policy), + with_escalated_permissions: params.with_escalated_permissions, + justification: params.justification, + } + } +} + +#[async_trait] +impl ToolHandler for ShellHandler { + fn kind(&self) -> ToolKind { + ToolKind::Function + } + + fn matches_kind(&self, payload: &ToolPayload) -> bool { + matches!( + payload, + ToolPayload::Function { .. } | ToolPayload::LocalShell { .. } + ) + } + + async fn handle( + &self, + invocation: ToolInvocation<'_>, + ) -> Result { + let ToolInvocation { + session, + turn, + tracker, + sub_id, + call_id, + tool_name, + payload, + } = invocation; + + match payload { + ToolPayload::Function { arguments } => { + let params: ShellToolCallParams = + serde_json::from_str(&arguments).map_err(|e| { + FunctionCallError::RespondToModel(format!( + "failed to parse function arguments: {e:?}" + )) + })?; + let exec_params = Self::to_exec_params(params, turn); + let content = crate::codex::handle_container_exec_with_params( + tool_name.as_str(), + exec_params, + session, + turn, + tracker, + sub_id.to_string(), + call_id.clone(), + ) + .await?; + Ok(ToolOutput::Function { + content, + success: true, + }) + } + ToolPayload::LocalShell { params } => { + let exec_params = Self::to_exec_params(params, turn); + let content = crate::codex::handle_container_exec_with_params( + tool_name.as_str(), + exec_params, + session, + turn, + tracker, + sub_id.to_string(), + call_id.clone(), + ) + .await?; + Ok(ToolOutput::Function { + content, + success: true, + }) + } + _ => Err(FunctionCallError::RespondToModel(format!( + "unsupported payload for shell handler: {tool_name}" + ))), + } + } +} diff --git a/codex-rs/core/src/tools/handlers/unified_exec.rs b/codex-rs/core/src/tools/handlers/unified_exec.rs new file mode 100644 index 0000000000..4afd7d3b83 --- /dev/null +++ b/codex-rs/core/src/tools/handlers/unified_exec.rs @@ -0,0 +1,112 @@ +use async_trait::async_trait; +use serde::Deserialize; + +use crate::function_tool::FunctionCallError; +use crate::tools::context::ToolInvocation; +use crate::tools::context::ToolOutput; +use crate::tools::context::ToolPayload; +use crate::tools::registry::ToolHandler; +use crate::tools::registry::ToolKind; +use crate::unified_exec::UnifiedExecRequest; + +pub struct UnifiedExecHandler; + +#[derive(Deserialize)] +struct UnifiedExecArgs { + input: Vec, + #[serde(default)] + session_id: Option, + #[serde(default)] + timeout_ms: Option, +} + +#[async_trait] +impl ToolHandler for UnifiedExecHandler { + fn kind(&self) -> ToolKind { + ToolKind::UnifiedExec + } + + fn matches_kind(&self, payload: &ToolPayload) -> bool { + matches!( + payload, + ToolPayload::UnifiedExec { .. } | ToolPayload::Function { .. } + ) + } + + async fn handle( + &self, + invocation: ToolInvocation<'_>, + ) -> Result { + let ToolInvocation { + session, payload, .. + } = invocation; + + let args = match payload { + ToolPayload::UnifiedExec { arguments } | ToolPayload::Function { arguments } => { + serde_json::from_str::(&arguments).map_err(|err| { + FunctionCallError::RespondToModel(format!( + "failed to parse function arguments: {err:?}" + )) + })? + } + _ => { + return Err(FunctionCallError::RespondToModel( + "unified_exec handler received unsupported payload".to_string(), + )); + } + }; + + let UnifiedExecArgs { + input, + session_id, + timeout_ms, + } = args; + + let parsed_session_id = if let Some(session_id) = session_id { + match session_id.parse::() { + Ok(parsed) => Some(parsed), + Err(output) => { + return Err(FunctionCallError::RespondToModel(format!( + "invalid session_id: {session_id} due to error {output:?}" + ))); + } + } + } else { + None + }; + + let request = UnifiedExecRequest { + session_id: parsed_session_id, + input_chunks: &input, + timeout_ms, + }; + + let value = session + .run_unified_exec_request(request) + .await + .map_err(|err| { + FunctionCallError::RespondToModel(format!("unified exec failed: {err:?}")) + })?; + + #[derive(serde::Serialize)] + struct SerializedUnifiedExecResult { + session_id: Option, + output: String, + } + + let content = serde_json::to_string(&SerializedUnifiedExecResult { + session_id: value.session_id.map(|id| id.to_string()), + output: value.output, + }) + .map_err(|err| { + FunctionCallError::RespondToModel(format!( + "failed to serialize unified exec output: {err:?}" + )) + })?; + + Ok(ToolOutput::Function { + content, + success: true, + }) + } +} diff --git a/codex-rs/core/src/tools/handlers/view_image.rs b/codex-rs/core/src/tools/handlers/view_image.rs new file mode 100644 index 0000000000..b525004a1f --- /dev/null +++ b/codex-rs/core/src/tools/handlers/view_image.rs @@ -0,0 +1,64 @@ +use async_trait::async_trait; +use serde::Deserialize; + +use crate::function_tool::FunctionCallError; +use crate::protocol::InputItem; +use crate::tools::context::ToolInvocation; +use crate::tools::context::ToolOutput; +use crate::tools::context::ToolPayload; +use crate::tools::registry::ToolHandler; +use crate::tools::registry::ToolKind; + +pub struct ViewImageHandler; + +#[derive(Deserialize)] +struct ViewImageArgs { + path: String, +} + +#[async_trait] +impl ToolHandler for ViewImageHandler { + fn kind(&self) -> ToolKind { + ToolKind::Function + } + + async fn handle( + &self, + invocation: ToolInvocation<'_>, + ) -> Result { + let ToolInvocation { + session, + turn, + payload, + .. + } = invocation; + + let arguments = match payload { + ToolPayload::Function { arguments } => arguments, + _ => { + return Err(FunctionCallError::RespondToModel( + "view_image handler received unsupported payload".to_string(), + )); + } + }; + + let args: ViewImageArgs = serde_json::from_str(&arguments).map_err(|e| { + FunctionCallError::RespondToModel(format!("failed to parse function arguments: {e:?}")) + })?; + + let abs_path = turn.resolve_path(Some(args.path)); + session + .inject_input(vec![InputItem::LocalImage { path: abs_path }]) + .await + .map_err(|_| { + FunctionCallError::RespondToModel( + "unable to attach image (no active task)".to_string(), + ) + })?; + + Ok(ToolOutput::Function { + content: "attached local image path".to_string(), + success: true, + }) + } +} diff --git a/codex-rs/core/src/tools/mod.rs b/codex-rs/core/src/tools/mod.rs new file mode 100644 index 0000000000..ee78cafdfa --- /dev/null +++ b/codex-rs/core/src/tools/mod.rs @@ -0,0 +1,7 @@ +pub mod context; +pub(crate) mod handlers; +pub mod registry; +pub mod router; +pub mod spec; + +pub use router::Router; diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs new file mode 100644 index 0000000000..a0e0cf5fd7 --- /dev/null +++ b/codex-rs/core/src/tools/registry.rs @@ -0,0 +1,174 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use async_trait::async_trait; +use codex_protocol::models::ResponseInputItem; +use tracing::warn; + +use crate::function_tool::FunctionCallError; +use crate::tools::context::ToolInvocation; +use crate::tools::context::ToolOutput; +use crate::tools::context::ToolPayload; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub enum ToolKind { + Function, + Custom, + LocalShell, + UnifiedExec, + Mcp, +} + +#[async_trait] +pub trait ToolHandler: Send + Sync { + fn spec(&self) -> Option<&crate::tools::spec::ResponsesApiTool> { + None + } + + fn kind(&self) -> ToolKind; + + fn matches_kind(&self, payload: &ToolPayload) -> bool { + matches!( + (self.kind(), payload), + (ToolKind::Function, ToolPayload::Function { .. }) + | (ToolKind::Function, ToolPayload::UnifiedExec { .. }) + | (ToolKind::Custom, ToolPayload::Custom { .. }) + | (ToolKind::LocalShell, ToolPayload::LocalShell { .. }) + | (ToolKind::UnifiedExec, ToolPayload::UnifiedExec { .. }) + | (ToolKind::Mcp, ToolPayload::Mcp { .. }) + ) + } + + async fn handle(&self, invocation: ToolInvocation<'_>) + -> Result; +} + +pub struct ToolRegistry { + handlers: HashMap>, +} + +impl ToolRegistry { + pub fn new(handlers: HashMap>) -> Self { + Self { handlers } + } + + pub fn handler(&self, name: &str) -> Option> { + self.handlers.get(name).map(Arc::clone) + } + + // TODO(jif) for dynamic tools. + // pub fn register(&mut self, name: impl Into, handler: Arc) { + // let name = name.into(); + // if self.handlers.insert(name.clone(), handler).is_some() { + // warn!("overwriting handler for tool {name}"); + // } + // } + + pub fn iter_handlers(&self) -> impl Iterator)> { + self.handlers.iter() + } + + pub async fn dispatch<'a>( + &self, + invocation: ToolInvocation<'a>, + ) -> Result { + let tool_name = invocation.tool_name.clone(); + let handler = self.handler(tool_name.as_ref()).ok_or_else(|| { + FunctionCallError::RespondToModel(format!("unsupported call: {tool_name}")) + })?; + + if !handler.matches_kind(&invocation.payload) { + return Err(FunctionCallError::RespondToModel(format!( + "tool {tool_name} invoked with incompatible payload" + ))); + } + + let call_id_owned = invocation.call_id.clone(); + let otel = invocation.turn.client.get_otel_event_manager(); + let log_payload = invocation.payload.log_payload().into_owned(); + let output_cell = std::sync::Mutex::new(None); + + let result = otel + .log_tool_result(tool_name.as_ref(), &call_id_owned, &log_payload, || { + let handler = handler.clone(); + let output_cell = &output_cell; + let invocation = invocation; + async move { + match handler.handle(invocation).await { + Ok(output) => { + let preview = output.log_preview(); + let mut guard = output_cell.lock().expect("mutex poisoned"); + *guard = Some(output); + Ok(preview) + } + Err(err) => Err(err), + } + } + }) + .await; + + match result { + Ok(_) => { + let mut guard = output_cell.lock().expect("mutex poisoned"); + let output = guard.take().ok_or_else(|| { + FunctionCallError::RespondToModel("tool produced no output".to_string()) + })?; + Ok(output.into_response(&call_id_owned)) + } + Err(err) => Err(err), + } + } +} + +pub struct ToolRegistryBuilder { + handlers: HashMap>, +} + +impl ToolRegistryBuilder { + pub fn new() -> Self { + Self { + handlers: HashMap::new(), + } + } + + pub fn register_handler(&mut self, name: impl Into, handler: Arc) { + let name = name.into(); + if self + .handlers + .insert(name.clone(), handler.clone()) + .is_some() + { + warn!("overwriting handler for tool {name}"); + } + } + + // TODO(jif) for dynamic tools. + // pub fn register_many(&mut self, names: I, handler: Arc) + // where + // I: IntoIterator, + // I::Item: Into, + // { + // for name in names { + // let name = name.into(); + // if self + // .handlers + // .insert(name.clone(), handler.clone()) + // .is_some() + // { + // warn!("overwriting handler for tool {name}"); + // } + // } + // } + + pub fn extend(&mut self, other: ToolRegistryBuilder) { + for (name, handler) in other.handlers { + if self.handlers.insert(name.clone(), handler).is_some() { + warn!("overwriting handler for tool {name}"); + } + } + } + + pub fn build(self) -> ToolRegistry { + ToolRegistry::new(self.handlers) + } +} diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs new file mode 100644 index 0000000000..372b652605 --- /dev/null +++ b/codex-rs/core/src/tools/router.rs @@ -0,0 +1,187 @@ +use std::collections::HashMap; + +use crate::codex::Session; +use crate::codex::TurnContext; +use crate::function_tool::FunctionCallError; +use crate::tools::context::ToolInvocation; +use crate::tools::context::ToolPayload; +use crate::tools::registry::ToolRegistry; +use crate::tools::spec::ToolSpec; +use crate::tools::spec::ToolsConfig; +use crate::tools::spec::build_specs; +use crate::turn_diff_tracker::TurnDiffTracker; +use codex_protocol::models::FunctionCallOutputPayload; +use codex_protocol::models::LocalShellAction; +use codex_protocol::models::ResponseInputItem; +use codex_protocol::models::ResponseItem; +use codex_protocol::models::ShellToolCallParams; + +#[derive(Clone)] +pub struct ToolCall { + pub tool_name: String, + pub call_id: String, + pub payload: ToolPayload, +} + +pub struct Router { + registry: ToolRegistry, + specs: Vec, +} + +impl Router { + pub fn from_config( + config: &ToolsConfig, + mcp_tools: Option>, + ) -> Self { + let (specs, builder) = build_specs(config, mcp_tools); + let registry = builder.build(); + Self { registry, specs } + } + + pub fn new(registry: ToolRegistry, specs: Vec) -> Self { + Self { registry, specs } + } + + pub fn specs(&self) -> &[ToolSpec] { + &self.specs + } + + pub fn build_tool_call( + session: &Session, + item: ResponseItem, + ) -> Result, FunctionCallError> { + match item { + ResponseItem::FunctionCall { + name, + arguments, + call_id, + .. + } => { + if let Some((server, tool)) = session.parse_mcp_tool_name(&name) { + let parsed_arguments = if arguments.trim().is_empty() { + None + } else { + Some(serde_json::from_str(&arguments).map_err(|err| { + FunctionCallError::RespondToModel(format!( + "failed to parse tool call arguments: {err}" + )) + })?) + }; + Ok(Some(ToolCall { + tool_name: name, + call_id, + payload: ToolPayload::Mcp { + server, + tool, + arguments: parsed_arguments, + raw_arguments: arguments, + }, + })) + } else { + let payload = if name == "unified_exec" { + ToolPayload::UnifiedExec { arguments } + } else { + ToolPayload::Function { arguments } + }; + Ok(Some(ToolCall { + tool_name: name, + call_id, + payload, + })) + } + } + ResponseItem::CustomToolCall { + name, + input, + call_id, + .. + } => Ok(Some(ToolCall { + tool_name: name, + call_id, + payload: ToolPayload::Custom { input }, + })), + ResponseItem::LocalShellCall { + id, + call_id, + action, + .. + } => { + let call_id = call_id.or(id).ok_or_else(|| { + FunctionCallError::RespondToModel( + "LocalShellCall without call_id or id".to_string(), + ) + })?; + + match action { + LocalShellAction::Exec(exec) => { + let params = ShellToolCallParams { + command: exec.command, + workdir: exec.working_directory, + timeout_ms: exec.timeout_ms, + with_escalated_permissions: None, + justification: None, + }; + Ok(Some(ToolCall { + tool_name: "local_shell".to_string(), + call_id, + payload: ToolPayload::LocalShell { params }, + })) + } + } + } + _ => Ok(None), + } + } + + pub async fn dispatch_tool_call( + &self, + session: &Session, + turn: &TurnContext, + tracker: &mut TurnDiffTracker, + sub_id: &str, + call: ToolCall, + ) -> ResponseInputItem { + let payload_clone = call.payload.clone(); + let ToolCall { + tool_name, + call_id, + payload, + } = call; + + let invocation = ToolInvocation { + session, + turn, + tracker, + sub_id, + call_id: call_id.clone(), + tool_name: tool_name.clone(), + payload, + }; + + match self.registry.dispatch(invocation).await { + Ok(response) => response, + Err(err) => Self::failure_response(call_id, payload_clone, err), + } + } + + fn failure_response( + call_id: String, + payload: ToolPayload, + err: FunctionCallError, + ) -> ResponseInputItem { + let message = err.to_string(); + match payload { + ToolPayload::Custom { .. } => ResponseInputItem::CustomToolCallOutput { + call_id, + output: message, + }, + _ => ResponseInputItem::FunctionCallOutput { + call_id, + output: codex_protocol::models::FunctionCallOutputPayload { + content: message, + success: Some(false), + }, + }, + } + } +} diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs new file mode 100644 index 0000000000..6a2351f727 --- /dev/null +++ b/codex-rs/core/src/tools/spec.rs @@ -0,0 +1,1085 @@ +use crate::tools::registry::ToolRegistryBuilder; +use serde::Deserialize; +use serde::Serialize; +use serde_json::Value as JsonValue; +use serde_json::json; +use std::collections::BTreeMap; +use std::collections::HashMap; + +use crate::model_family::ModelFamily; +use crate::plan_tool::PLAN_TOOL; +use crate::tool_apply_patch::ApplyPatchToolType; +use crate::tool_apply_patch::create_apply_patch_freeform_tool; +use crate::tool_apply_patch::create_apply_patch_json_tool; + +#[derive(Debug, Clone, Serialize, PartialEq)] +pub struct ResponsesApiTool { + pub(crate) name: String, + pub(crate) description: String, + /// TODO: Validation. When strict is set to true, the JSON schema, + /// `required` and `additional_properties` must be present. All fields in + /// `properties` must be present in `required`. + pub(crate) strict: bool, + pub(crate) parameters: JsonSchema, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct FreeformTool { + pub(crate) name: String, + pub(crate) description: String, + pub(crate) format: FreeformToolFormat, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct FreeformToolFormat { + pub(crate) r#type: String, + pub(crate) syntax: String, + pub(crate) definition: String, +} + +/// When serialized as JSON, this produces a valid "Tool" in the OpenAI +/// Responses API. +#[derive(Debug, Clone, Serialize, PartialEq)] +#[serde(tag = "type")] +pub(crate) enum ToolSpec { + #[serde(rename = "function")] + Function(ResponsesApiTool), + #[serde(rename = "local_shell")] + LocalShell {}, + // TODO: Understand why we get an error on web_search although the API docs say it's supported. + // https://platform.openai.com/docs/guides/tools-web-search?api-mode=responses#:~:text=%7B%20type%3A%20%22web_search%22%20%7D%2C + #[serde(rename = "web_search")] + WebSearch {}, + #[serde(rename = "custom")] + Freeform(FreeformTool), +} + +#[derive(Debug, Clone)] +pub enum ConfigShellToolType { + Default, + Local, + Streamable, +} + +#[derive(Debug, Clone)] +pub(crate) struct ToolsConfig { + pub shell_type: ConfigShellToolType, + pub plan_tool: bool, + pub apply_patch_tool_type: Option, + pub web_search_request: bool, + pub include_view_image_tool: bool, + pub experimental_unified_exec_tool: bool, +} + +pub(crate) struct ToolsConfigParams<'a> { + pub(crate) model_family: &'a ModelFamily, + pub(crate) include_plan_tool: bool, + pub(crate) include_apply_patch_tool: bool, + pub(crate) include_web_search_request: bool, + pub(crate) use_streamable_shell_tool: bool, + pub(crate) include_view_image_tool: bool, + pub(crate) experimental_unified_exec_tool: bool, +} + +impl ToolsConfig { + pub fn new(params: &ToolsConfigParams) -> Self { + let ToolsConfigParams { + model_family, + include_plan_tool, + include_apply_patch_tool, + include_web_search_request, + use_streamable_shell_tool, + include_view_image_tool, + experimental_unified_exec_tool, + } = params; + let shell_type = if *use_streamable_shell_tool { + ConfigShellToolType::Streamable + } else if model_family.uses_local_shell_tool { + ConfigShellToolType::Local + } else { + ConfigShellToolType::Default + }; + + let apply_patch_tool_type = match model_family.apply_patch_tool_type { + Some(ApplyPatchToolType::Freeform) => Some(ApplyPatchToolType::Freeform), + Some(ApplyPatchToolType::Function) => Some(ApplyPatchToolType::Function), + None => { + if *include_apply_patch_tool { + Some(ApplyPatchToolType::Freeform) + } else { + None + } + } + }; + + Self { + shell_type, + plan_tool: *include_plan_tool, + apply_patch_tool_type, + web_search_request: *include_web_search_request, + include_view_image_tool: *include_view_image_tool, + experimental_unified_exec_tool: *experimental_unified_exec_tool, + } + } +} + +/// Generic JSON‑Schema subset needed for our tool definitions +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(tag = "type", rename_all = "lowercase")] +pub(crate) enum JsonSchema { + Boolean { + #[serde(skip_serializing_if = "Option::is_none")] + description: Option, + }, + String { + #[serde(skip_serializing_if = "Option::is_none")] + description: Option, + }, + /// MCP schema allows "number" | "integer" for Number + #[serde(alias = "integer")] + Number { + #[serde(skip_serializing_if = "Option::is_none")] + description: Option, + }, + Array { + items: Box, + + #[serde(skip_serializing_if = "Option::is_none")] + description: Option, + }, + Object { + properties: BTreeMap, + #[serde(skip_serializing_if = "Option::is_none")] + required: Option>, + #[serde( + rename = "additionalProperties", + skip_serializing_if = "Option::is_none" + )] + additional_properties: Option, + }, +} + +fn create_unified_exec_tool() -> ToolSpec { + let mut properties = BTreeMap::new(); + properties.insert( + "input".to_string(), + JsonSchema::Array { + items: Box::new(JsonSchema::String { description: None }), + description: Some( + "When no session_id is provided, treat the array as the command and arguments \ + to launch. When session_id is set, concatenate the strings (in order) and write \ + them to the session's stdin." + .to_string(), + ), + }, + ); + properties.insert( + "session_id".to_string(), + JsonSchema::String { + description: Some( + "Identifier for an existing interactive session. If omitted, a new command \ + is spawned." + .to_string(), + ), + }, + ); + properties.insert( + "timeout_ms".to_string(), + JsonSchema::Number { + description: Some( + "Maximum time in milliseconds to wait for output after writing the input." + .to_string(), + ), + }, + ); + + ToolSpec::Function(ResponsesApiTool { + name: "unified_exec".to_string(), + description: + "Runs a command in a PTY. Provide a session_id to reuse an existing interactive session.".to_string(), + strict: false, + parameters: JsonSchema::Object { + properties, + required: Some(vec!["input".to_string()]), + additional_properties: Some(false), + }, + }) +} + +fn create_shell_tool() -> ToolSpec { + let mut properties = BTreeMap::new(); + properties.insert( + "command".to_string(), + JsonSchema::Array { + items: Box::new(JsonSchema::String { description: None }), + description: Some("The command to execute".to_string()), + }, + ); + properties.insert( + "workdir".to_string(), + JsonSchema::String { + description: Some("The working directory to execute the command in".to_string()), + }, + ); + properties.insert( + "timeout_ms".to_string(), + JsonSchema::Number { + description: Some("The timeout for the command in milliseconds".to_string()), + }, + ); + + properties.insert( + "with_escalated_permissions".to_string(), + JsonSchema::Boolean { + description: Some("Whether to request escalated permissions. Set to true if command needs to be run without sandbox restrictions".to_string()), + }, + ); + properties.insert( + "justification".to_string(), + JsonSchema::String { + description: Some("Only set if with_escalated_permissions is true. 1-sentence explanation of why we want to run this command.".to_string()), + }, + ); + + ToolSpec::Function(ResponsesApiTool { + name: "shell".to_string(), + description: "Runs a shell command and returns its output.".to_string(), + strict: false, + parameters: JsonSchema::Object { + properties, + required: Some(vec!["command".to_string()]), + additional_properties: Some(false), + }, + }) +} + +fn create_view_image_tool() -> ToolSpec { + // Support only local filesystem path. + let mut properties = BTreeMap::new(); + properties.insert( + "path".to_string(), + JsonSchema::String { + description: Some("Local filesystem path to an image file".to_string()), + }, + ); + + ToolSpec::Function(ResponsesApiTool { + name: "view_image".to_string(), + description: + "Attach a local image (by filesystem path) to the conversation context for this turn." + .to_string(), + strict: false, + parameters: JsonSchema::Object { + properties, + required: Some(vec!["path".to_string()]), + additional_properties: Some(false), + }, + }) +} +/// TODO(dylan): deprecate once we get rid of json tool +#[derive(Serialize, Deserialize)] +pub(crate) struct ApplyPatchToolArgs { + pub(crate) input: String, +} + +/// Returns JSON values that are compatible with Function Calling in the +/// Responses API: +/// https://platform.openai.com/docs/guides/function-calling?api-mode=responses +pub fn create_tools_json_for_responses_api( + tools: &[ToolSpec], +) -> crate::error::Result> { + let mut tools_json = Vec::new(); + + for tool in tools { + let json = serde_json::to_value(tool)?; + tools_json.push(json); + } + + Ok(tools_json) +} +/// Returns JSON values that are compatible with Function Calling in the +/// Chat Completions API: +/// https://platform.openai.com/docs/guides/function-calling?api-mode=chat +pub(crate) fn create_tools_json_for_chat_completions_api( + tools: &[ToolSpec], +) -> crate::error::Result> { + // We start with the JSON for the Responses API and than rewrite it to match + // the chat completions tool call format. + let responses_api_tools_json = create_tools_json_for_responses_api(tools)?; + let tools_json = responses_api_tools_json + .into_iter() + .filter_map(|mut tool| { + if tool.get("type") != Some(&serde_json::Value::String("function".to_string())) { + return None; + } + + if let Some(map) = tool.as_object_mut() { + // Remove "type" field as it is not needed in chat completions. + map.remove("type"); + Some(json!({ + "type": "function", + "function": map, + })) + } else { + None + } + }) + .collect::>(); + Ok(tools_json) +} + +pub(crate) fn mcp_tool_to_openai_tool( + fully_qualified_name: String, + tool: mcp_types::Tool, +) -> Result { + let mcp_types::Tool { + description, + mut input_schema, + .. + } = tool; + + // OpenAI models mandate the "properties" field in the schema. The Agents + // SDK fixed this by inserting an empty object for "properties" if it is not + // already present https://github.com/openai/openai-agents-python/issues/449 + // so here we do the same. + if input_schema.properties.is_none() { + input_schema.properties = Some(serde_json::Value::Object(serde_json::Map::new())); + } + + // Serialize to a raw JSON value so we can sanitize schemas coming from MCP + // servers. Some servers omit the top-level or nested `type` in JSON + // Schemas (e.g. using enum/anyOf), or use unsupported variants like + // `integer`. Our internal JsonSchema is a small subset and requires + // `type`, so we coerce/sanitize here for compatibility. + let mut serialized_input_schema = serde_json::to_value(input_schema)?; + sanitize_json_schema(&mut serialized_input_schema); + let input_schema = serde_json::from_value::(serialized_input_schema)?; + + Ok(ResponsesApiTool { + name: fully_qualified_name, + description: description.unwrap_or_default(), + strict: false, + parameters: input_schema, + }) +} + +/// Sanitize a JSON Schema (as serde_json::Value) so it can fit our limited +/// JsonSchema enum. This function: +/// - Ensures every schema object has a "type". If missing, infers it from +/// common keywords (properties => object, items => array, enum/const/format => string) +/// and otherwise defaults to "string". +/// - Fills required child fields (e.g. array items, object properties) with +/// permissive defaults when absent. +fn sanitize_json_schema(value: &mut JsonValue) { + match value { + JsonValue::Bool(_) => { + // JSON Schema boolean form: true/false. Coerce to an accept-all string. + *value = json!({ "type": "string" }); + } + JsonValue::Array(arr) => { + for v in arr.iter_mut() { + sanitize_json_schema(v); + } + } + JsonValue::Object(map) => { + // First, recursively sanitize known nested schema holders + if let Some(props) = map.get_mut("properties") + && let Some(props_map) = props.as_object_mut() + { + for (_k, v) in props_map.iter_mut() { + sanitize_json_schema(v); + } + } + if let Some(items) = map.get_mut("items") { + sanitize_json_schema(items); + } + // Some schemas use oneOf/anyOf/allOf - sanitize their entries + for combiner in ["oneOf", "anyOf", "allOf", "prefixItems"] { + if let Some(v) = map.get_mut(combiner) { + sanitize_json_schema(v); + } + } + + // Normalize/ensure type + let mut ty = map.get("type").and_then(|v| v.as_str()).map(str::to_string); + + // If type is an array (union), pick first supported; else leave to inference + if ty.is_none() + && let Some(JsonValue::Array(types)) = map.get("type") + { + for t in types { + if let Some(tt) = t.as_str() + && matches!( + tt, + "object" | "array" | "string" | "number" | "integer" | "boolean" + ) + { + ty = Some(tt.to_string()); + break; + } + } + } + + // Infer type if still missing + if ty.is_none() { + if map.contains_key("properties") + || map.contains_key("required") + || map.contains_key("additionalProperties") + { + ty = Some("object".to_string()); + } else if map.contains_key("items") || map.contains_key("prefixItems") { + ty = Some("array".to_string()); + } else if map.contains_key("enum") + || map.contains_key("const") + || map.contains_key("format") + { + ty = Some("string".to_string()); + } else if map.contains_key("minimum") + || map.contains_key("maximum") + || map.contains_key("exclusiveMinimum") + || map.contains_key("exclusiveMaximum") + || map.contains_key("multipleOf") + { + ty = Some("number".to_string()); + } + } + // If we still couldn't infer, default to string + let ty = ty.unwrap_or_else(|| "string".to_string()); + map.insert("type".to_string(), JsonValue::String(ty.to_string())); + + // Ensure object schemas have properties map + if ty == "object" { + if !map.contains_key("properties") { + map.insert( + "properties".to_string(), + JsonValue::Object(serde_json::Map::new()), + ); + } + // If additionalProperties is an object schema, sanitize it too. + // Leave booleans as-is, since JSON Schema allows boolean here. + if let Some(ap) = map.get_mut("additionalProperties") { + let is_bool = matches!(ap, JsonValue::Bool(_)); + if !is_bool { + sanitize_json_schema(ap); + } + } + } + + // Ensure array schemas have items + if ty == "array" && !map.contains_key("items") { + map.insert("items".to_string(), json!({ "type": "string" })); + } + } + _ => {} + } +} + +/// Builds the tool specs along with the registry builder containing matching handlers. +pub(crate) fn build_specs( + config: &ToolsConfig, + mcp_tools: Option>, +) -> (Vec, ToolRegistryBuilder) { + use crate::exec_command::EXEC_COMMAND_TOOL_NAME; + use crate::exec_command::WRITE_STDIN_TOOL_NAME; + use crate::exec_command::create_exec_command_tool_for_responses_api; + use crate::exec_command::create_write_stdin_tool_for_responses_api; + use crate::tools::handlers::ApplyPatchHandler; + use crate::tools::handlers::ExecStreamHandler; + use crate::tools::handlers::McpHandler; + use crate::tools::handlers::PlanHandler; + use crate::tools::handlers::ShellHandler; + use crate::tools::handlers::UnifiedExecHandler; + use crate::tools::handlers::ViewImageHandler; + use std::sync::Arc; + + let mut specs: Vec = Vec::new(); + let mut builder = ToolRegistryBuilder::new(); + + let shell_handler = Arc::new(ShellHandler); + let exec_stream_handler = Arc::new(ExecStreamHandler); + let unified_exec_handler = Arc::new(UnifiedExecHandler); + let plan_handler = Arc::new(PlanHandler); + let apply_patch_handler = Arc::new(ApplyPatchHandler); + let view_image_handler = Arc::new(ViewImageHandler); + let mcp_handler = Arc::new(McpHandler); + + if config.experimental_unified_exec_tool { + specs.push(create_unified_exec_tool()); + builder.register_handler("unified_exec", unified_exec_handler.clone()); + } else { + match &config.shell_type { + ConfigShellToolType::Default => { + specs.push(create_shell_tool()); + } + ConfigShellToolType::Local => { + specs.push(ToolSpec::LocalShell {}); + } + ConfigShellToolType::Streamable => { + specs.push(ToolSpec::Function( + create_exec_command_tool_for_responses_api(), + )); + specs.push(ToolSpec::Function( + create_write_stdin_tool_for_responses_api(), + )); + builder.register_handler(EXEC_COMMAND_TOOL_NAME, exec_stream_handler.clone()); + builder.register_handler(WRITE_STDIN_TOOL_NAME, exec_stream_handler.clone()); + } + } + } + + // 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.clone()); + + if config.plan_tool { + specs.push(PLAN_TOOL.clone()); + builder.register_handler("update_plan", plan_handler.clone()); + } + + if let Some(apply_patch_tool_type) = &config.apply_patch_tool_type { + match apply_patch_tool_type { + ApplyPatchToolType::Freeform => { + specs.push(create_apply_patch_freeform_tool()); + } + ApplyPatchToolType::Function => { + specs.push(create_apply_patch_json_tool()); + } + } + builder.register_handler("apply_patch", apply_patch_handler.clone()); + } + + if config.web_search_request { + specs.push(ToolSpec::WebSearch {}); + } + + if config.include_view_image_tool { + specs.push(create_view_image_tool()); + builder.register_handler("view_image", view_image_handler.clone()); + } + + if let Some(mcp_tools) = mcp_tools { + let mut entries: Vec<(String, mcp_types::Tool)> = mcp_tools.into_iter().collect(); + entries.sort_by(|a, b| a.0.cmp(&b.0)); + + for (name, tool) in entries.into_iter() { + match mcp_tool_to_openai_tool(name.clone(), tool.clone()) { + Ok(converted_tool) => { + specs.push(ToolSpec::Function(converted_tool)); + builder.register_handler(name, mcp_handler.clone()); + } + Err(e) => { + tracing::error!("Failed to convert {name:?} MCP tool to OpenAI tool: {e:?}"); + } + } + } + } + + (specs, builder) +} + +#[cfg(test)] +mod tests { + use crate::model_family::find_family_for_model; + use mcp_types::ToolInputSchema; + use pretty_assertions::assert_eq; + + use super::*; + + fn assert_eq_tool_names(tools: &[ToolSpec], expected_names: &[&str]) { + let tool_names = tools + .iter() + .map(|tool| match tool { + ToolSpec::Function(ResponsesApiTool { name, .. }) => name, + ToolSpec::LocalShell {} => "local_shell", + ToolSpec::WebSearch {} => "web_search", + ToolSpec::Freeform(FreeformTool { name, .. }) => name, + }) + .collect::>(); + + assert_eq!( + tool_names.len(), + expected_names.len(), + "tool_name mismatch, {tool_names:?}, {expected_names:?}", + ); + for (name, expected_name) in tool_names.iter().zip(expected_names.iter()) { + assert_eq!( + name, expected_name, + "tool_name mismatch, {name:?}, {expected_name:?}" + ); + } + } + + #[test] + fn test_build_specs() { + let model_family = find_family_for_model("codex-mini-latest") + .expect("codex-mini-latest should be a valid model family"); + let config = ToolsConfig::new(&ToolsConfigParams { + model_family: &model_family, + include_plan_tool: true, + include_apply_patch_tool: false, + include_web_search_request: true, + use_streamable_shell_tool: false, + include_view_image_tool: true, + experimental_unified_exec_tool: true, + }); + let tools = build_specs(&config, Some(HashMap::new())).0; + + assert_eq_tool_names( + &tools, + &["unified_exec", "update_plan", "web_search", "view_image"], + ); + } + + #[test] + fn test_build_specs_default_shell() { + let model_family = find_family_for_model("o3").expect("o3 should be a valid model family"); + let config = ToolsConfig::new(&ToolsConfigParams { + model_family: &model_family, + include_plan_tool: true, + include_apply_patch_tool: false, + include_web_search_request: true, + use_streamable_shell_tool: false, + include_view_image_tool: true, + experimental_unified_exec_tool: true, + }); + let tools = build_specs(&config, Some(HashMap::new())).0; + + assert_eq_tool_names( + &tools, + &["unified_exec", "update_plan", "web_search", "view_image"], + ); + } + + #[test] + fn test_build_specs_mcp_tools() { + let model_family = find_family_for_model("o3").expect("o3 should be a valid model family"); + let config = ToolsConfig::new(&ToolsConfigParams { + model_family: &model_family, + include_plan_tool: false, + include_apply_patch_tool: false, + include_web_search_request: true, + use_streamable_shell_tool: false, + include_view_image_tool: true, + experimental_unified_exec_tool: true, + }); + let tools = build_specs( + &config, + Some(HashMap::from([( + "test_server/do_something_cool".to_string(), + mcp_types::Tool { + name: "do_something_cool".to_string(), + input_schema: ToolInputSchema { + properties: Some(serde_json::json!({ + "string_argument": { + "type": "string", + }, + "number_argument": { + "type": "number", + }, + "object_argument": { + "type": "object", + "properties": { + "string_property": { "type": "string" }, + "number_property": { "type": "number" }, + }, + "required": [ + "string_property", + "number_property", + ], + "additionalProperties": Some(false), + }, + })), + required: None, + r#type: "object".to_string(), + }, + output_schema: None, + title: None, + annotations: None, + description: Some("Do something cool".to_string()), + }, + )])), + ) + .0; + + assert_eq_tool_names( + &tools, + &[ + "unified_exec", + "web_search", + "view_image", + "test_server/do_something_cool", + ], + ); + + assert_eq!( + tools[3], + ToolSpec::Function(ResponsesApiTool { + name: "test_server/do_something_cool".to_string(), + parameters: JsonSchema::Object { + properties: BTreeMap::from([ + ( + "string_argument".to_string(), + JsonSchema::String { description: None } + ), + ( + "number_argument".to_string(), + JsonSchema::Number { description: None } + ), + ( + "object_argument".to_string(), + JsonSchema::Object { + properties: BTreeMap::from([ + ( + "string_property".to_string(), + JsonSchema::String { description: None } + ), + ( + "number_property".to_string(), + JsonSchema::Number { description: None } + ), + ]), + required: Some(vec![ + "string_property".to_string(), + "number_property".to_string(), + ]), + additional_properties: Some(false), + }, + ), + ]), + required: None, + additional_properties: None, + }, + description: "Do something cool".to_string(), + strict: false, + }) + ); + } + + #[test] + fn test_build_specs_mcp_tools_sorted_by_name() { + let model_family = find_family_for_model("o3").expect("o3 should be a valid model family"); + let config = ToolsConfig::new(&ToolsConfigParams { + model_family: &model_family, + include_plan_tool: false, + include_apply_patch_tool: false, + include_web_search_request: false, + use_streamable_shell_tool: false, + include_view_image_tool: true, + experimental_unified_exec_tool: true, + }); + + // Intentionally construct a map with keys that would sort alphabetically. + let tools_map: HashMap = HashMap::from([ + ( + "test_server/do".to_string(), + mcp_types::Tool { + name: "a".to_string(), + input_schema: ToolInputSchema { + properties: Some(serde_json::json!({})), + required: None, + r#type: "object".to_string(), + }, + output_schema: None, + title: None, + annotations: None, + description: Some("a".to_string()), + }, + ), + ( + "test_server/something".to_string(), + mcp_types::Tool { + name: "b".to_string(), + input_schema: ToolInputSchema { + properties: Some(serde_json::json!({})), + required: None, + r#type: "object".to_string(), + }, + output_schema: None, + title: None, + annotations: None, + description: Some("b".to_string()), + }, + ), + ( + "test_server/cool".to_string(), + mcp_types::Tool { + name: "c".to_string(), + input_schema: ToolInputSchema { + properties: Some(serde_json::json!({})), + required: None, + r#type: "object".to_string(), + }, + output_schema: None, + title: None, + annotations: None, + description: Some("c".to_string()), + }, + ), + ]); + + let tools = build_specs(&config, Some(tools_map)).0; + // Expect unified_exec first, followed by MCP tools sorted by fully-qualified name. + assert_eq_tool_names( + &tools, + &[ + "unified_exec", + "view_image", + "test_server/cool", + "test_server/do", + "test_server/something", + ], + ); + } + + #[test] + fn test_mcp_tool_property_missing_type_defaults_to_string() { + let model_family = find_family_for_model("o3").expect("o3 should be a valid model family"); + let config = ToolsConfig::new(&ToolsConfigParams { + model_family: &model_family, + include_plan_tool: false, + include_apply_patch_tool: false, + include_web_search_request: true, + use_streamable_shell_tool: false, + include_view_image_tool: true, + experimental_unified_exec_tool: true, + }); + + let tools = build_specs( + &config, + Some(HashMap::from([( + "dash/search".to_string(), + mcp_types::Tool { + name: "search".to_string(), + input_schema: ToolInputSchema { + properties: Some(serde_json::json!({ + "query": { + "description": "search query" + } + })), + required: None, + r#type: "object".to_string(), + }, + output_schema: None, + title: None, + annotations: None, + description: Some("Search docs".to_string()), + }, + )])), + ) + .0; + + assert_eq_tool_names( + &tools, + &["unified_exec", "web_search", "view_image", "dash/search"], + ); + + assert_eq!( + tools[3], + ToolSpec::Function(ResponsesApiTool { + name: "dash/search".to_string(), + parameters: JsonSchema::Object { + properties: BTreeMap::from([( + "query".to_string(), + JsonSchema::String { + description: Some("search query".to_string()) + } + )]), + required: None, + additional_properties: None, + }, + description: "Search docs".to_string(), + strict: false, + }) + ); + } + + #[test] + fn test_mcp_tool_integer_normalized_to_number() { + let model_family = find_family_for_model("o3").expect("o3 should be a valid model family"); + let config = ToolsConfig::new(&ToolsConfigParams { + model_family: &model_family, + include_plan_tool: false, + include_apply_patch_tool: false, + include_web_search_request: true, + use_streamable_shell_tool: false, + include_view_image_tool: true, + experimental_unified_exec_tool: true, + }); + + let tools = build_specs( + &config, + Some(HashMap::from([( + "dash/paginate".to_string(), + mcp_types::Tool { + name: "paginate".to_string(), + input_schema: ToolInputSchema { + properties: Some(serde_json::json!({ + "page": { "type": "integer" } + })), + required: None, + r#type: "object".to_string(), + }, + output_schema: None, + title: None, + annotations: None, + description: Some("Pagination".to_string()), + }, + )])), + ) + .0; + + assert_eq_tool_names( + &tools, + &["unified_exec", "web_search", "view_image", "dash/paginate"], + ); + assert_eq!( + tools[3], + ToolSpec::Function(ResponsesApiTool { + name: "dash/paginate".to_string(), + parameters: JsonSchema::Object { + properties: BTreeMap::from([( + "page".to_string(), + JsonSchema::Number { description: None } + )]), + required: None, + additional_properties: None, + }, + description: "Pagination".to_string(), + strict: false, + }) + ); + } + + #[test] + fn test_mcp_tool_array_without_items_gets_default_string_items() { + let model_family = find_family_for_model("o3").expect("o3 should be a valid model family"); + let config = ToolsConfig::new(&ToolsConfigParams { + model_family: &model_family, + include_plan_tool: false, + include_apply_patch_tool: false, + include_web_search_request: true, + use_streamable_shell_tool: false, + include_view_image_tool: true, + experimental_unified_exec_tool: true, + }); + + let tools = build_specs( + &config, + Some(HashMap::from([( + "dash/tags".to_string(), + mcp_types::Tool { + name: "tags".to_string(), + input_schema: ToolInputSchema { + properties: Some(serde_json::json!({ + "tags": { "type": "array" } + })), + required: None, + r#type: "object".to_string(), + }, + output_schema: None, + title: None, + annotations: None, + description: Some("Tags".to_string()), + }, + )])), + ) + .0; + + assert_eq_tool_names( + &tools, + &["unified_exec", "web_search", "view_image", "dash/tags"], + ); + assert_eq!( + tools[3], + ToolSpec::Function(ResponsesApiTool { + name: "dash/tags".to_string(), + parameters: JsonSchema::Object { + properties: BTreeMap::from([( + "tags".to_string(), + JsonSchema::Array { + items: Box::new(JsonSchema::String { description: None }), + description: None + } + )]), + required: None, + additional_properties: None, + }, + description: "Tags".to_string(), + strict: false, + }) + ); + } + + #[test] + fn test_mcp_tool_anyof_defaults_to_string() { + let model_family = find_family_for_model("o3").expect("o3 should be a valid model family"); + let config = ToolsConfig::new(&ToolsConfigParams { + model_family: &model_family, + include_plan_tool: false, + include_apply_patch_tool: false, + include_web_search_request: true, + use_streamable_shell_tool: false, + include_view_image_tool: true, + experimental_unified_exec_tool: true, + }); + + let tools = build_specs( + &config, + Some(HashMap::from([( + "dash/value".to_string(), + mcp_types::Tool { + name: "value".to_string(), + input_schema: ToolInputSchema { + properties: Some(serde_json::json!({ + "value": { "anyOf": [ { "type": "string" }, { "type": "number" } ] } + })), + required: None, + r#type: "object".to_string(), + }, + output_schema: None, + title: None, + annotations: None, + description: Some("AnyOf Value".to_string()), + }, + )])), + ) + .0; + + assert_eq_tool_names( + &tools, + &["unified_exec", "web_search", "view_image", "dash/value"], + ); + assert_eq!( + tools[3], + ToolSpec::Function(ResponsesApiTool { + name: "dash/value".to_string(), + parameters: JsonSchema::Object { + properties: BTreeMap::from([( + "value".to_string(), + JsonSchema::String { description: None } + )]), + required: None, + additional_properties: None, + }, + description: "AnyOf Value".to_string(), + strict: false, + }) + ); + } + + #[test] + fn test_shell_tool() { + let tool = super::create_shell_tool(); + let ToolSpec::Function(ResponsesApiTool { + description, name, .. + }) = &tool + else { + panic!("expected function tool"); + }; + assert_eq!(name, "shell"); + + let expected = "Runs a shell command and returns its output."; + assert_eq!(description, expected); + } +} diff --git a/codex-rs/tools.md b/codex-rs/tools.md new file mode 100644 index 0000000000..cdc8eba32b --- /dev/null +++ b/codex-rs/tools.md @@ -0,0 +1,98 @@ +# Tool Handling Refactor Plan + +## Current Pain Points +- Tool specs and dispatch live in multiple files (`openai_tools.rs`, `codex.rs`, `shell.rs`, `apply_patch.rs`, `exec_command`, `plan_tool`), so every tool is stitched together by hand with duplicated string matches. +- `codex.rs` mixes concerns: argument parsing, apply-patch verification, approval policy enforcement, telemetry, and response construction for Function/Custom/LocalShell calls are all intertwined. +- Adding or adjusting a tool requires touching 5–7 sites and remembering subtle invariants (e.g., apply_patch needs extra verification, view_image requires path resolution), which increases review overhead and bug risk. +- Spec configuration (`ToolsConfig`, `get_openai_tools`) is disconnected from runtime dispatch, so there is no single registry that knows which tools are active or how to invoke them. + +## Goals +- Centralize tool definitions, specs, and dispatching inside a dedicated `tools` module with clearly defined traits and types. +- Unify handling for function calls, custom calls, and local shell calls, sharing approval checks, apply_patch guards, and telemetry. +- Slim down `codex.rs` so it only converts `ResponseItem` into a `ToolCall` and delegates the rest. +- Make new tool onboarding require (at most) one handler + registration entry + optional tests. + +## Core Traits & Types +- ```rust + pub trait ToolHandler: Send + Sync { + fn spec(&self) -> Option<&ToolSpec>; + fn kind(&self) -> ToolKind; + async fn handle(&self, invocation: ToolInvocation<'_>) -> Result; + } + ``` + - `spec`: supplies the OpenAI tool definition for function/custom tools (returns `None` for implicit tools like local shell). + - `kind`: declares whether the handler services `ToolCall::Function`, `ToolCall::Custom`, `ToolCall::LocalShell`, etc., for dispatch filtering. + - `handle`: runs the tool logic, returning structured output + optional streaming side effects. +- ```rust + pub struct ToolInvocation<'a> { + pub session: &'a Session, + pub turn: &'a TurnContext, + pub tracker: &'a mut TurnDiffTracker, + pub sub_id: &'a str, + pub call_id: &'a str, + pub payload: ToolPayload<'a>, + } + ``` + - `payload` captures parsed arguments (`serde_json::Value` for functions/custom, `ShellToolCallParams` for local shell, etc.). +- ```rust + pub enum ToolOutput { + Function(FunctionCallOutputPayload), + Custom(String), + } + ``` + - Helper constructors convert to `ResponseInputItem` in one place. +- ```rust + pub struct ToolRegistry { + handlers: HashMap>, + } + ``` + - Provides `register_handler`, `iter_specs()` for prompt assembly, and `dispatch(call: ToolCall) -> Result`. + +## Proposed Refactor +1. **Scaffold `core/src/tools` Module** + - `mod.rs` re-exports `context`, `registry`, `router`, `spec`, and `handlers::*`. + - `context.rs` defines `ToolInvocation`, `ToolPayload`, `ToolOutput`, helper constructors for `ResponseInputItem`. + - `registry.rs` implements `ToolRegistry` + `ToolKind` enum and houses the instrumentation wrapper (`otel.log_tool_result`, approval policy checks). + - `router.rs` provides `ToolCall` enum (variants: `Function`, `Custom`, `LocalShell`, `UnifiedExec`, `Mcp`), plus `Router::from_turn_context(&TurnContext)` to build the registry + spec list together. + +2. **Move Spec Logic Into `tools::spec`** + - Relocate `ToolsConfig`, `ToolsConfigParams`, `get_openai_tools`, `ConfigShellToolType`, and tool builders (shell/unified exec/apply_patch/etc.) into `spec.rs`. + - Add `ToolSpec` struct mirroring current `OpenAiTool`, keeping existing tests but updated paths. + - Expose `pub fn build_specs(config: &ToolsConfig, mcp: Option>) -> (Vec, ToolRegistryBuilder)` to keep spec generation and handler registration in sync. + +3. **Implement Concrete Handlers** + - Split existing logic into focused files: + 1. `handlers/shell.rs`: wraps `handle_container_exec_with_params`, approval policy enforcement, apply_patch verification, and streaming support. + 2. `handlers/apply_patch.rs`: retains apply_patch-specific parsing when invoked as a function/custom tool. + 3. `handlers/unified_exec.rs`: hosts `handle_unified_exec_tool_call` logic. + 4. `handlers/plan.rs`: wraps `handle_update_plan`. + 5. `handlers/view_image.rs`: handles local path resolution + `inject_input` call. + 6. `handlers/exec_stream.rs`: handles `exec_command` and `write_stdin` (for streamable exec shell variant). + 7. `handlers/mcp.rs`: optional adapter that calls through to `handle_mcp_tool_call` while emitting consistent telemetry. + - Each handler implements `ToolHandler` with explicit payload parsing (`serde_json::from_value` helpers) and reuses shared utilities (`create_env`, `maybe_parse_apply_patch_verified`). + +4. **Registry Wiring** + - Add `ToolRegistryBuilder` helper collecting `(ToolName, Arc)` pairs. + - Builder methods (e.g., `with_shell_handler`, `with_apply_patch_handler`) are invoked from `spec::build_specs` based on `ToolsConfig` decisions so that enabling a spec automatically registers the handler. + - Integrate MCP tools by registering a generic `McpHandler` per tool when MCP discovery runs. + +5. **Update `codex.rs`** + - Replace `handle_function_call`, `handle_custom_tool_call`, `to_exec_params`, and `handle_container_exec_with_params` with: + - `let router = tools::router::Router::new(sess, turn_context, turn_diff_tracker);` + - Convert `ResponseItem` into `ToolCall` and call `router.dispatch(...)`. + - Remove duplicate instrumentation; `ToolRegistry` handles `log_tool_result` and error wrapping. + - Keep only MCP discovery + fallback logic, delegating to `ToolRegistry` for execution. + +6. **Tests & Validation** + - Port existing unit tests in `openai_tools.rs` to `tools::spec` and add new tests for `ToolRegistry` covering success, error, and approval-policy rejection paths. + - Add smoke tests for converting `ResponseItem` to `ToolCall` in `tools::router`. + - Ensure apply_patch verification works both for function and local shell flows via dedicated tests (mock `apply_patch::apply_patch`). + +## Testing & Rollout +- Run `cargo test -p codex-core tools::spec` (once added) plus focused handler tests. +- Execute existing integration suites that cover exec/apply_patch/unified_exec flows. +- No behavior change expected, but snapshot suites (esp. TUI if any output differs) should be re-run as a sanity check. + +## Follow-ups +- With registry centralization, expose structured telemetry (counts, durations) and hook into metrics. +- Evaluate deletion of legacy local shell pathway once unified_exec proves stable.