codex/codex-rs/tools.md

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

• pub trait ToolHandler: Send + Sync {
      fn spec(&self) -> Option<&ToolSpec>;
      fn kind(&self) -> ToolKind;
      async fn handle(&self, invocation: ToolInvocation<'_>) -> Result<ToolOutput, FunctionCallError>;
  }
  

  • 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.

• 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.).

• pub enum ToolOutput {
      Function(FunctionCallOutputPayload),
      Custom(String),
  }
  

  • Helper constructors convert to ResponseInputItem in one place.

• pub struct ToolRegistry {
      handlers: HashMap<ToolName, Arc<dyn ToolHandler>>,
  }
  

  • Provides register_handler, iter_specs() for prompt assembly, and dispatch(call: ToolCall) -> Result<ResponseInputItem, FunctionCallError>.

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<HashMap<..>>) -> (Vec<ToolSpec>, 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<dyn ToolHandler>) 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.