chore: rework tools execution workflow (#5278)

Re-work the tool execution flow. Read `orchestrator.rs` to understand
the structure

codex-rs/core/src/tools/handlers/apply_patch.rs

@@ -8,7 +8,6 @@ use crate::client_common::tools::ResponsesApiTool;
 use crate::client_common::tools::ToolSpec;
 use crate::exec::ExecParams;
 use crate::function_tool::FunctionCallError;
-use crate::openai_tools::JsonSchema;
 use crate::tools::context::ToolInvocation;
 use crate::tools::context::ToolOutput;
 use crate::tools::context::ToolPayload;
@@ -16,6 +15,7 @@ use crate::tools::handle_container_exec_with_params;
 use crate::tools::registry::ToolHandler;
 use crate::tools::registry::ToolKind;
 use crate::tools::spec::ApplyPatchToolArgs;
+use crate::tools::spec::JsonSchema;
 use async_trait::async_trait;
 use serde::Deserialize;
 use serde::Serialize;
@@ -72,6 +72,7 @@ impl ToolHandler for ApplyPatchHandler {
             env: HashMap::new(),
             with_escalated_permissions: None,
             justification: None,
+            arg0: None,
         };
 
         let content = handle_container_exec_with_params(

codex-rs/core/src/tools/handlers/exec_stream.rs

@@ -1,68 +0,0 @@
-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<ToolOutput, FunctionCallError> {
-        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: Some(true),
-        })
-    }
-}

codex-rs/core/src/tools/handlers/mod.rs

@@ -1,5 +1,4 @@
 pub mod apply_patch;
-mod exec_stream;
 mod grep_files;
 mod list_dir;
 mod mcp;
@@ -14,7 +13,6 @@ mod view_image;
 pub use plan::PLAN_TOOL;
 
 pub use apply_patch::ApplyPatchHandler;
-pub use exec_stream::ExecStreamHandler;
 pub use grep_files::GrepFilesHandler;
 pub use list_dir::ListDirHandler;
 pub use mcp::McpHandler;

codex-rs/core/src/tools/handlers/plan.rs

@@ -2,12 +2,12 @@ use crate::client_common::tools::ResponsesApiTool;
 use crate::client_common::tools::ToolSpec;
 use crate::codex::Session;
 use crate::function_tool::FunctionCallError;
-use crate::openai_tools::JsonSchema;
 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::JsonSchema;
 use async_trait::async_trait;
 use codex_protocol::plan_tool::UpdatePlanArgs;
 use codex_protocol::protocol::Event;

codex-rs/core/src/tools/handlers/shell.rs

@@ -24,6 +24,7 @@ impl ShellHandler {
             env: create_env(&turn_context.shell_environment_policy),
             with_escalated_permissions: params.with_escalated_permissions,
             justification: params.justification,
+            arg0: None,
         }
     }
 }

codex-rs/core/src/tools/handlers/unified_exec.rs

@@ -35,7 +35,13 @@ impl ToolHandler for UnifiedExecHandler {
 
     async fn handle(&self, invocation: ToolInvocation) -> Result<ToolOutput, FunctionCallError> {
         let ToolInvocation {
-            session, payload, ..
+            session,
+            turn,
+            sub_id,
+            call_id,
+            tool_name: _tool_name,
+            payload,
+            ..
         } = invocation;
 
         let args = match payload {
@@ -73,13 +79,23 @@ impl ToolHandler for UnifiedExecHandler {
         };
 
         let request = UnifiedExecRequest {
-            session_id: parsed_session_id,
             input_chunks: &input,
             timeout_ms,
         };
 
         let value = session
-            .run_unified_exec_request(request)
+            .services
+            .unified_exec_manager
+            .handle_request(
+                request,
+                crate::unified_exec::UnifiedExecContext {
+                    session: &session,
+                    turn: turn.as_ref(),
+                    sub_id: &sub_id,
+                    call_id: &call_id,
+                    session_id: parsed_session_id,
+                },
+            )
             .await
             .map_err(|err| {
                 FunctionCallError::RespondToModel(format!("unified exec failed: {err:?}"))