diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index ef07a36daf..a0dd913374 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -751,6 +751,7 @@ dependencies = [ "codex-common", "codex-core", "codex-ollama", + "core_test_support", "libc", "owo-colors", "predicates", @@ -760,6 +761,7 @@ dependencies = [ "tokio", "tracing", "tracing-subscriber", + "wiremock", ] [[package]] diff --git a/codex-rs/core/src/apply_patch.rs b/codex-rs/core/src/apply_patch.rs index 21e80406e5..4f9292b6d7 100644 --- a/codex-rs/core/src/apply_patch.rs +++ b/codex-rs/core/src/apply_patch.rs @@ -1,4 +1,5 @@ use crate::codex::Session; +use crate::codex::TurnContext; use crate::models::FunctionCallOutputPayload; use crate::models::ResponseInputItem; use crate::protocol::FileChange; @@ -8,7 +9,6 @@ use crate::safety::assess_patch_safety; use codex_apply_patch::ApplyPatchAction; use codex_apply_patch::ApplyPatchFileChange; use std::collections::HashMap; -use std::path::Path; use std::path::PathBuf; pub const CODEX_APPLY_PATCH_ARG1: &str = "--codex-run-as-apply-patch"; @@ -41,17 +41,16 @@ impl From for InternalApplyPatchInvocation { pub(crate) async fn apply_patch( sess: &Session, + turn_context: &TurnContext, sub_id: &str, call_id: &str, action: ApplyPatchAction, ) -> InternalApplyPatchInvocation { - let writable_roots_snapshot = sess.get_writable_roots().to_vec(); - match assess_patch_safety( &action, - sess.get_approval_policy(), - &writable_roots_snapshot, - sess.get_cwd(), + turn_context.approval_policy, + &turn_context.sandbox_policy, + &turn_context.cwd, ) { SafetyCheck::AutoApprove { .. } => { InternalApplyPatchInvocation::DelegateToExec(ApplyPatchExec { @@ -124,30 +123,3 @@ pub(crate) fn convert_apply_patch_to_protocol( } result } - -pub(crate) fn get_writable_roots(cwd: &Path) -> Vec { - let mut writable_roots = Vec::new(); - if cfg!(target_os = "macos") { - // On macOS, $TMPDIR is private to the user. - writable_roots.push(std::env::temp_dir()); - - // Allow pyenv to update its shims directory. Without this, any tool - // that happens to be managed by `pyenv` will fail with an error like: - // - // pyenv: cannot rehash: $HOME/.pyenv/shims isn't writable - // - // which is emitted every time `pyenv` tries to run `rehash` (for - // example, after installing a new Python package that drops an entry - // point). Although the sandbox is intentionally read‑only by default, - // writing to the user's local `pyenv` directory is safe because it - // is already user‑writable and scoped to the current user account. - if let Ok(home_dir) = std::env::var("HOME") { - let pyenv_dir = PathBuf::from(home_dir).join(".pyenv"); - writable_roots.push(pyenv_dir); - } - } - - writable_roots.push(cwd.to_path_buf()); - - writable_roots -} diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 3d26bd0880..686ec79dcb 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -56,7 +56,7 @@ struct Error { message: Option, } -#[derive(Clone)] +#[derive(Debug, Clone)] pub struct ModelClient { config: Arc, auth: Option, diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index b7d1322c35..bbb192d03d 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -1,13 +1,10 @@ -// Poisoned mutex should fail the program -#![expect(clippy::unwrap_used)] - use std::borrow::Cow; use std::collections::HashMap; use std::collections::HashSet; -use std::path::Path; use std::path::PathBuf; use std::sync::Arc; use std::sync::Mutex; +use std::sync::MutexGuard; use std::sync::atomic::AtomicU64; use std::time::Duration; @@ -31,12 +28,11 @@ use tracing::warn; use uuid::Uuid; use crate::ModelProviderInfo; +use crate::apply_patch; use crate::apply_patch::ApplyPatchExec; use crate::apply_patch::CODEX_APPLY_PATCH_ARG1; use crate::apply_patch::InternalApplyPatchInvocation; use crate::apply_patch::convert_apply_patch_to_protocol; -use crate::apply_patch::get_writable_roots; -use crate::apply_patch::{self}; use crate::client::ModelClient; use crate::client_common::Prompt; use crate::client_common::ResponseEvent; @@ -68,6 +64,7 @@ use crate::models::ReasoningItemReasoningSummary; use crate::models::ResponseInputItem; use crate::models::ResponseItem; use crate::models::ShellToolCallParams; +use crate::openai_tools::ApplyPatchToolArgs; use crate::openai_tools::ToolsConfig; use crate::openai_tools::get_openai_tools; use crate::parse_command::parse_command; @@ -109,6 +106,21 @@ use crate::turn_diff_tracker::TurnDiffTracker; use crate::user_notification::UserNotification; use crate::util::backoff; +// A convenience extension trait for acquiring mutex locks where poisoning is +// unrecoverable and should abort the program. This avoids scattered `.unwrap()` +// calls on `lock()` while still surfacing a clear panic message when a lock is +// poisoned. +trait MutexExt { + fn lock_unchecked(&self) -> MutexGuard<'_, T>; +} + +impl MutexExt for Mutex { + fn lock_unchecked(&self) -> MutexGuard<'_, T> { + #[expect(clippy::expect_used)] + self.lock().expect("poisoned lock") + } +} + /// The high-level interface to the Codex system. /// It operates as a queue pair where you send submissions and receive events. pub struct Codex { @@ -154,16 +166,17 @@ impl Codex { }; // Generate a unique ID for the lifetime of this Codex session. - let session = Session::new(configure_session, config.clone(), auth, tx_event.clone()) - .await - .map_err(|e| { - error!("Failed to create session: {e:#}"); - CodexErr::InternalAgentDied - })?; + let (session, turn_context) = + Session::new(configure_session, config.clone(), auth, tx_event.clone()) + .await + .map_err(|e| { + error!("Failed to create session: {e:#}"); + CodexErr::InternalAgentDied + })?; let session_id = session.session_id; // This task will run until Op::Shutdown is received. - tokio::spawn(submission_loop(session, config, rx_sub)); + tokio::spawn(submission_loop(session, turn_context, config, rx_sub)); let codex = Codex { next_id: AtomicU64::new(0), tx_sub, @@ -219,22 +232,8 @@ struct State { /// A session has at most 1 running task at a time, and can be interrupted by user input. pub(crate) struct Session { session_id: Uuid, - client: ModelClient, tx_event: Sender, - /// The session's current working directory. All relative paths provided by - /// the model as well as sandbox policies are resolved against this path - /// instead of `std::env::current_dir()`. - cwd: PathBuf, - base_instructions: Option, - user_instructions: Option, - approval_policy: AskForApproval, - sandbox_policy: SandboxPolicy, - shell_environment_policy: ShellEnvironmentPolicy, - writable_roots: Vec, - disable_response_storage: bool, - tools_config: ToolsConfig, - /// Manager for external MCP servers/tools. mcp_connection_manager: McpConnectionManager, @@ -251,6 +250,31 @@ pub(crate) struct Session { show_raw_agent_reasoning: bool, } +/// The context needed for a single turn of the conversation. +#[derive(Debug)] +pub(crate) struct TurnContext { + pub(crate) client: ModelClient, + /// The session's current working directory. All relative paths provided by + /// the model as well as sandbox policies are resolved against this path + /// instead of `std::env::current_dir()`. + pub(crate) cwd: PathBuf, + pub(crate) base_instructions: Option, + pub(crate) user_instructions: Option, + pub(crate) approval_policy: AskForApproval, + pub(crate) sandbox_policy: SandboxPolicy, + pub(crate) shell_environment_policy: ShellEnvironmentPolicy, + pub(crate) disable_response_storage: bool, + pub(crate) tools_config: ToolsConfig, +} + +impl TurnContext { + fn resolve_path(&self, path: Option) -> PathBuf { + path.as_ref() + .map(PathBuf::from) + .map_or_else(|| self.cwd.clone(), |p| self.cwd.join(p)) + } +} + /// Configure the model session. struct ConfigureSession { /// Provider identifier ("openai", "openrouter", ...). @@ -298,7 +322,7 @@ impl Session { config: Arc, auth: Option, tx_event: Sender, - ) -> anyhow::Result> { + ) -> anyhow::Result<(Arc, TurnContext)> { let ConfigureSession { provider, model, @@ -410,8 +434,6 @@ impl Session { state.history.record_items(&restored_items); } - let writable_roots = get_writable_roots(&cwd); - // Handle MCP manager result and record any startup failures. let (mcp_connection_manager, failed_clients) = match mcp_res { Ok((mgr, failures)) => (mgr, failures), @@ -448,29 +470,31 @@ impl Session { model_reasoning_summary, session_id, ); - let sess = Arc::new(Session { - session_id, + let turn_context = TurnContext { client, tools_config: ToolsConfig::new( &config.model_family, approval_policy, sandbox_policy.clone(), config.include_plan_tool, + config.include_apply_patch_tool, ), - tx_event: tx_event.clone(), user_instructions, base_instructions, approval_policy, sandbox_policy, shell_environment_policy: config.shell_environment_policy.clone(), cwd, - writable_roots, + disable_response_storage, + }; + let sess = Arc::new(Session { + session_id, + tx_event: tx_event.clone(), mcp_connection_manager, notify, state: Mutex::new(state), rollout: Mutex::new(rollout_recorder), codex_linux_sandbox_exe: config.codex_linux_sandbox_exe.clone(), - disable_response_storage, user_shell: default_shell, show_raw_agent_reasoning: config.show_raw_agent_reasoning, }); @@ -478,13 +502,13 @@ impl Session { // record the initial user instructions and environment context, // regardless of whether we restored items. let mut conversation_items = Vec::::with_capacity(2); - if let Some(user_instructions) = sess.user_instructions.as_deref() { + if let Some(user_instructions) = turn_context.user_instructions.as_deref() { conversation_items.push(Prompt::format_user_instructions_message(user_instructions)); } conversation_items.push(ResponseItem::from(EnvironmentContext::new( - sess.get_cwd().to_path_buf(), - sess.get_approval_policy(), - sess.sandbox_policy.clone(), + turn_context.cwd.to_path_buf(), + turn_context.approval_policy, + turn_context.sandbox_policy.clone(), ))); sess.record_conversation_items(&conversation_items).await; @@ -505,29 +529,11 @@ impl Session { } } - Ok(sess) - } - - pub(crate) fn get_writable_roots(&self) -> &[PathBuf] { - &self.writable_roots - } - - pub(crate) fn get_approval_policy(&self) -> AskForApproval { - self.approval_policy - } - - pub(crate) fn get_cwd(&self) -> &Path { - &self.cwd - } - - fn resolve_path(&self, path: Option) -> PathBuf { - path.as_ref() - .map(PathBuf::from) - .map_or_else(|| self.cwd.clone(), |p| self.cwd.join(p)) + Ok((sess, turn_context)) } pub fn set_task(&self, task: AgentTask) { - let mut state = self.state.lock().unwrap(); + let mut state = self.state.lock_unchecked(); if let Some(current_task) = state.current_task.take() { current_task.abort(); } @@ -535,7 +541,7 @@ impl Session { } pub fn remove_task(&self, sub_id: &str) { - let mut state = self.state.lock().unwrap(); + let mut state = self.state.lock_unchecked(); if let Some(task) = &state.current_task { if task.sub_id == sub_id { state.current_task.take(); @@ -571,7 +577,7 @@ impl Session { }; let _ = self.tx_event.send(event).await; { - let mut state = self.state.lock().unwrap(); + let mut state = self.state.lock_unchecked(); state.pending_approvals.insert(sub_id, tx_approve); } rx_approve @@ -597,21 +603,21 @@ impl Session { }; let _ = self.tx_event.send(event).await; { - let mut state = self.state.lock().unwrap(); + let mut state = self.state.lock_unchecked(); state.pending_approvals.insert(sub_id, tx_approve); } rx_approve } pub fn notify_approval(&self, sub_id: &str, decision: ReviewDecision) { - let mut state = self.state.lock().unwrap(); + let mut state = self.state.lock_unchecked(); if let Some(tx_approve) = state.pending_approvals.remove(sub_id) { tx_approve.send(decision).ok(); } } pub fn add_approved_command(&self, cmd: Vec) { - let mut state = self.state.lock().unwrap(); + let mut state = self.state.lock_unchecked(); state.approved_commands.insert(cmd); } @@ -621,14 +627,14 @@ impl Session { debug!("Recording items for conversation: {items:?}"); self.record_state_snapshot(items).await; - self.state.lock().unwrap().history.record_items(items); + self.state.lock_unchecked().history.record_items(items); } async fn record_state_snapshot(&self, items: &[ResponseItem]) { let snapshot = { crate::rollout::SessionStateSnapshot {} }; let recorder = { - let guard = self.rollout.lock().unwrap(); + let guard = self.rollout.lock_unchecked(); guard.as_ref().cloned() }; @@ -806,12 +812,12 @@ impl Session { /// Build the full turn input by concatenating the current conversation /// history with additional items for this turn. pub fn turn_input_with_history(&self, extra: Vec) -> Vec { - [self.state.lock().unwrap().history.contents(), extra].concat() + [self.state.lock_unchecked().history.contents(), extra].concat() } /// Returns the input if there was no task running to inject into pub fn inject_input(&self, input: Vec) -> Result<(), Vec> { - let mut state = self.state.lock().unwrap(); + let mut state = self.state.lock_unchecked(); if state.current_task.is_some() { state.pending_input.push(input.into()); Ok(()) @@ -821,7 +827,7 @@ impl Session { } pub fn get_pending_input(&self) -> Vec { - let mut state = self.state.lock().unwrap(); + let mut state = self.state.lock_unchecked(); if state.pending_input.is_empty() { Vec::with_capacity(0) } else { @@ -845,7 +851,7 @@ impl Session { fn abort(&self) { info!("Aborting existing session"); - let mut state = self.state.lock().unwrap(); + let mut state = self.state.lock_unchecked(); state.pending_approvals.clear(); state.pending_input.clear(); if let Some(task) = state.current_task.take() { @@ -912,9 +918,19 @@ pub(crate) struct AgentTask { } impl AgentTask { - fn spawn(sess: Arc, sub_id: String, input: Vec) -> Self { - let handle = - tokio::spawn(run_task(Arc::clone(&sess), sub_id.clone(), input)).abort_handle(); + fn spawn( + sess: Arc, + turn_context: Arc, + sub_id: String, + input: Vec, + ) -> Self { + let handle = { + let sess = sess.clone(); + let sub_id = sub_id.clone(); + let tc = Arc::clone(&turn_context); + tokio::spawn(async move { run_task(sess, tc.as_ref(), sub_id, input).await }) + .abort_handle() + }; Self { sess, sub_id, @@ -924,17 +940,20 @@ impl AgentTask { fn compact( sess: Arc, + turn_context: Arc, sub_id: String, input: Vec, compact_instructions: String, ) -> Self { - let handle = tokio::spawn(run_compact_task( - Arc::clone(&sess), - sub_id.clone(), - input, - compact_instructions, - )) - .abort_handle(); + let handle = { + let sess = sess.clone(); + let sub_id = sub_id.clone(); + let tc = Arc::clone(&turn_context); + tokio::spawn(async move { + run_compact_task(sess, tc.as_ref(), sub_id, input, compact_instructions).await + }) + .abort_handle() + }; Self { sess, sub_id, @@ -959,7 +978,14 @@ impl AgentTask { } } -async fn submission_loop(sess: Arc, config: Arc, rx_sub: Receiver) { +async fn submission_loop( + sess: Arc, + turn_context: TurnContext, + config: Arc, + rx_sub: Receiver, +) { + // Wrap once to avoid cloning TurnContext for each task. + let turn_context = Arc::new(turn_context); // To break out of this loop, send Op::Shutdown. while let Ok(sub) = rx_sub.recv().await { debug!(?sub, "Submission"); @@ -971,7 +997,65 @@ async fn submission_loop(sess: Arc, config: Arc, rx_sub: Receiv // attempt to inject input into current task if let Err(items) = sess.inject_input(items) { // no current task, spawn a new one - let task = AgentTask::spawn(sess.clone(), sub.id, items); + let task = + AgentTask::spawn(sess.clone(), Arc::clone(&turn_context), sub.id, items); + sess.set_task(task); + } + } + Op::UserTurn { + items, + cwd, + approval_policy, + sandbox_policy, + model, + effort, + summary, + } => { + // attempt to inject input into current task + if let Err(items) = sess.inject_input(items) { + // Derive a fresh TurnContext for this turn using the provided overrides. + let provider = turn_context.client.get_provider(); + + // Derive a model family for the requested model; fall back to the session's. + let model_family = find_family_for_model(&model) + .unwrap_or_else(|| config.model_family.clone()); + + // Create a per‑turn Config clone with the requested model/family. + let mut per_turn_config = (*config).clone(); + per_turn_config.model = model.clone(); + per_turn_config.model_family = model_family.clone(); + + // Build a new client with per‑turn reasoning settings. + // Reuse the same provider and session id; auth defaults to env/API key. + let client = ModelClient::new( + Arc::new(per_turn_config), + None, + provider, + effort, + summary, + sess.session_id, + ); + + let fresh_turn_context = TurnContext { + client, + tools_config: ToolsConfig::new( + &model_family, + approval_policy, + sandbox_policy.clone(), + config.include_plan_tool, + ), + user_instructions: turn_context.user_instructions.clone(), + base_instructions: turn_context.base_instructions.clone(), + approval_policy, + sandbox_policy, + shell_environment_policy: turn_context.shell_environment_policy.clone(), + cwd, + disable_response_storage: turn_context.disable_response_storage, + }; + + // no current task, spawn a new one with the per‑turn context + let task = + AgentTask::spawn(sess.clone(), Arc::new(fresh_turn_context), sub.id, items); sess.set_task(task); } } @@ -1094,6 +1178,7 @@ async fn submission_loop(sess: Arc, config: Arc, rx_sub: Receiv }]) { let task = AgentTask::compact( sess.clone(), + Arc::clone(&turn_context), sub.id, items, SUMMARIZATION_PROMPT.to_string(), @@ -1106,7 +1191,7 @@ async fn submission_loop(sess: Arc, config: Arc, rx_sub: Receiv // Gracefully flush and shutdown rollout recorder on session end so tests // that inspect the rollout file do not race with the background writer. - let recorder_opt = sess.rollout.lock().unwrap().take(); + let recorder_opt = sess.rollout.lock_unchecked().take(); if let Some(rec) = recorder_opt { if let Err(e) = rec.shutdown().await { warn!("failed to shutdown rollout recorder: {e}"); @@ -1149,7 +1234,12 @@ async fn submission_loop(sess: Arc, config: Arc, rx_sub: Receiv /// back to the model in the next turn. /// - If the model sends only an assistant message, we record it in the /// conversation history and consider the task complete. -async fn run_task(sess: Arc, sub_id: String, input: Vec) { +async fn run_task( + sess: Arc, + turn_context: &TurnContext, + sub_id: String, + input: Vec, +) { if input.is_empty() { return; } @@ -1201,7 +1291,15 @@ async fn run_task(sess: Arc, sub_id: String, input: Vec) { }) }) .collect(); - match run_turn(&sess, &mut turn_diff_tracker, sub_id.clone(), turn_input).await { + match run_turn( + &sess, + turn_context, + &mut turn_diff_tracker, + sub_id.clone(), + turn_input, + ) + .await + { Ok(turn_output) => { let mut items_to_record_in_conversation_history = Vec::::new(); let mut responses = Vec::::new(); @@ -1330,25 +1428,26 @@ async fn run_task(sess: Arc, sub_id: String, input: Vec) { async fn run_turn( sess: &Session, + turn_context: &TurnContext, turn_diff_tracker: &mut TurnDiffTracker, sub_id: String, input: Vec, ) -> CodexResult> { let tools = get_openai_tools( - &sess.tools_config, + &turn_context.tools_config, Some(sess.mcp_connection_manager.list_all_tools()), ); let prompt = Prompt { input, - store: !sess.disable_response_storage, + store: !turn_context.disable_response_storage, tools, - base_instructions_override: sess.base_instructions.clone(), + base_instructions_override: turn_context.base_instructions.clone(), }; let mut retries = 0; loop { - match try_run_turn(sess, turn_diff_tracker, &sub_id, &prompt).await { + match try_run_turn(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)), @@ -1357,7 +1456,7 @@ async fn run_turn( } Err(e) => { // Use the configured provider-specific stream retry budget. - let max_retries = sess.client.get_provider().stream_max_retries(); + let max_retries = turn_context.client.get_provider().stream_max_retries(); if retries < max_retries { retries += 1; let delay = match e { @@ -1400,6 +1499,7 @@ struct ProcessedResponseItem { async fn try_run_turn( sess: &Session, + turn_context: &TurnContext, turn_diff_tracker: &mut TurnDiffTracker, sub_id: &str, prompt: &Prompt, @@ -1460,7 +1560,7 @@ async fn try_run_turn( }) }; - let mut stream = sess.client.clone().stream(&prompt).await?; + let mut stream = turn_context.client.clone().stream(&prompt).await?; let mut output = Vec::new(); loop { @@ -1489,9 +1589,14 @@ async fn try_run_turn( match event { ResponseEvent::Created => {} ResponseEvent::OutputItemDone(item) => { - let response = - handle_response_item(sess, turn_diff_tracker, sub_id, item.clone()).await?; - + let response = handle_response_item( + sess, + turn_context, + turn_diff_tracker, + sub_id, + item.clone(), + ) + .await?; output.push(ProcessedResponseItem { item, response }); } ResponseEvent::Completed { @@ -1522,7 +1627,7 @@ async fn try_run_turn( } ResponseEvent::OutputTextDelta(delta) => { { - let mut st = sess.state.lock().unwrap(); + let mut st = sess.state.lock_unchecked(); st.history.append_assistant_text(&delta); } @@ -1563,6 +1668,7 @@ async fn try_run_turn( async fn run_compact_task( sess: Arc, + turn_context: &TurnContext, sub_id: String, input: Vec, compact_instructions: String, @@ -1581,16 +1687,16 @@ async fn run_compact_task( let prompt = Prompt { input: turn_input, - store: !sess.disable_response_storage, + store: !turn_context.disable_response_storage, tools: Vec::new(), base_instructions_override: Some(compact_instructions.clone()), }; - let max_retries = sess.client.get_provider().stream_max_retries(); + let max_retries = turn_context.client.get_provider().stream_max_retries(); let mut retries = 0; loop { - let attempt_result = drain_to_completed(&sess, &sub_id, &prompt).await; + let attempt_result = drain_to_completed(&sess, turn_context, &sub_id, &prompt).await; match attempt_result { Ok(()) => break, @@ -1638,12 +1744,13 @@ async fn run_compact_task( }; sess.send_event(event).await; - let mut state = sess.state.lock().unwrap(); + let mut state = sess.state.lock_unchecked(); state.history.keep_last_messages(1); } async fn handle_response_item( sess: &Session, + turn_context: &TurnContext, turn_diff_tracker: &mut TurnDiffTracker, sub_id: &str, item: ResponseItem, @@ -1678,8 +1785,9 @@ async fn handle_response_item( }; sess.tx_event.send(event).await.ok(); } - if sess.show_raw_agent_reasoning && content.is_some() { - let content = content.unwrap(); + if sess.show_raw_agent_reasoning + && let Some(content) = content + { for item in content { let text = match item { ReasoningItemContent::ReasoningText { text } => text, @@ -1706,6 +1814,7 @@ async fn handle_response_item( Some( handle_function_call( sess, + turn_context, turn_diff_tracker, sub_id.to_string(), name, @@ -1745,11 +1854,12 @@ async fn handle_response_item( } }; - let exec_params = to_exec_params(params, sess); + let exec_params = to_exec_params(params, turn_context); Some( handle_container_exec_with_params( exec_params, sess, + turn_context, turn_diff_tracker, sub_id.to_string(), effective_call_id, @@ -1768,6 +1878,7 @@ async fn handle_response_item( async fn handle_function_call( sess: &Session, + turn_context: &TurnContext, turn_diff_tracker: &mut TurnDiffTracker, sub_id: String, name: String, @@ -1776,13 +1887,44 @@ async fn handle_function_call( ) -> ResponseInputItem { match name.as_str() { "container.exec" | "shell" => { - let params = match parse_container_exec_arguments(arguments, sess, &call_id) { + let params = match parse_container_exec_arguments(arguments, turn_context, &call_id) { Ok(params) => params, Err(output) => { return *output; } }; - handle_container_exec_with_params(params, sess, turn_diff_tracker, sub_id, call_id) + handle_container_exec_with_params( + params, + sess, + turn_context, + turn_diff_tracker, + sub_id, + call_id, + ) + .await + } + "apply_patch" => { + let args = match serde_json::from_str::(&arguments) { + Ok(a) => a, + Err(e) => { + return ResponseInputItem::FunctionCallOutput { + call_id, + output: FunctionCallOutputPayload { + content: format!("failed to parse function arguments: {e}"), + success: None, + }, + }; + } + }; + let exec_params = ExecParams { + command: vec!["apply_patch".to_string(), args.input.clone()], + cwd: sess.cwd.clone(), + timeout_ms: None, + env: HashMap::new(), + with_escalated_permissions: None, + justification: None, + }; + handle_container_exec_with_params(exec_params, sess, turn_diff_tracker, sub_id, call_id) .await } "update_plan" => handle_update_plan(sess, arguments, sub_id, call_id).await, @@ -1811,12 +1953,12 @@ async fn handle_function_call( } } -fn to_exec_params(params: ShellToolCallParams, sess: &Session) -> ExecParams { +fn to_exec_params(params: ShellToolCallParams, turn_context: &TurnContext) -> ExecParams { ExecParams { command: params.command, - cwd: sess.resolve_path(params.workdir.clone()), + cwd: turn_context.resolve_path(params.workdir.clone()), timeout_ms: params.timeout_ms, - env: create_env(&sess.shell_environment_policy), + env: create_env(&turn_context.shell_environment_policy), with_escalated_permissions: params.with_escalated_permissions, justification: params.justification, } @@ -1824,12 +1966,12 @@ fn to_exec_params(params: ShellToolCallParams, sess: &Session) -> ExecParams { fn parse_container_exec_arguments( arguments: String, - sess: &Session, + turn_context: &TurnContext, call_id: &str, ) -> Result> { // parse command match serde_json::from_str::(&arguments) { - Ok(shell_tool_call_params) => Ok(to_exec_params(shell_tool_call_params, sess)), + Ok(shell_tool_call_params) => Ok(to_exec_params(shell_tool_call_params, turn_context)), Err(e) => { // allow model to re-sample let output = ResponseInputItem::FunctionCallOutput { @@ -1852,8 +1994,12 @@ pub struct ExecInvokeArgs<'a> { pub stdout_stream: Option, } -fn maybe_run_with_user_profile(params: ExecParams, sess: &Session) -> ExecParams { - if sess.shell_environment_policy.use_profile { +fn maybe_run_with_user_profile( + params: ExecParams, + sess: &Session, + turn_context: &TurnContext, +) -> ExecParams { + if turn_context.shell_environment_policy.use_profile { let command = sess .user_shell .format_default_shell_invocation(params.command.clone()); @@ -1867,6 +2013,7 @@ fn maybe_run_with_user_profile(params: ExecParams, sess: &Session) -> ExecParams async fn handle_container_exec_with_params( params: ExecParams, sess: &Session, + turn_context: &TurnContext, turn_diff_tracker: &mut TurnDiffTracker, sub_id: String, call_id: String, @@ -1874,7 +2021,7 @@ async fn handle_container_exec_with_params( // check if this was a patch, and apply it if so let apply_patch_exec = match maybe_parse_apply_patch_verified(¶ms.command, ¶ms.cwd) { MaybeApplyPatchVerified::Body(changes) => { - match apply_patch::apply_patch(sess, &sub_id, &call_id, changes).await { + match apply_patch::apply_patch(sess, turn_context, &sub_id, &call_id, changes).await { InternalApplyPatchInvocation::Output(item) => return item, InternalApplyPatchInvocation::DelegateToExec(apply_patch_exec) => { Some(apply_patch_exec) @@ -1936,8 +2083,8 @@ async fn handle_container_exec_with_params( } } else { assess_safety_for_untrusted_command( - sess.approval_policy, - &sess.sandbox_policy, + turn_context.approval_policy, + &turn_context.sandbox_policy, params.with_escalated_permissions.unwrap_or(false), ) }; @@ -1949,11 +2096,11 @@ async fn handle_container_exec_with_params( } None => { let safety = { - let state = sess.state.lock().unwrap(); + let state = sess.state.lock_unchecked(); assess_command_safety( ¶ms.command, - sess.approval_policy, - &sess.sandbox_policy, + turn_context.approval_policy, + &turn_context.sandbox_policy, &state.approved_commands, params.with_escalated_permissions.unwrap_or(false), ) @@ -2023,7 +2170,7 @@ async fn handle_container_exec_with_params( ), }; - let params = maybe_run_with_user_profile(params, sess); + let params = maybe_run_with_user_profile(params, sess, turn_context); let output_result = sess .run_exec_with_events( turn_diff_tracker, @@ -2031,7 +2178,7 @@ async fn handle_container_exec_with_params( ExecInvokeArgs { params: params.clone(), sandbox_type, - sandbox_policy: &sess.sandbox_policy, + sandbox_policy: &turn_context.sandbox_policy, codex_linux_sandbox_exe: &sess.codex_linux_sandbox_exe, stdout_stream: Some(StdoutStream { sub_id: sub_id.clone(), @@ -2064,6 +2211,7 @@ async fn handle_container_exec_with_params( error, sandbox_type, sess, + turn_context, ) .await } @@ -2084,6 +2232,7 @@ async fn handle_sandbox_error( error: SandboxErr, sandbox_type: SandboxType, sess: &Session, + turn_context: &TurnContext, ) -> ResponseInputItem { let call_id = exec_command_context.call_id.clone(); let sub_id = exec_command_context.sub_id.clone(); @@ -2091,7 +2240,7 @@ async fn handle_sandbox_error( // Early out if either the user never wants to be asked for approval, or // we're letting the model manage escalation requests. Otherwise, continue - match sess.approval_policy { + match turn_context.approval_policy { AskForApproval::Never | AskForApproval::OnRequest => { return ResponseInputItem::FunctionCallOutput { call_id, @@ -2162,7 +2311,7 @@ async fn handle_sandbox_error( ExecInvokeArgs { params, sandbox_type: SandboxType::None, - sandbox_policy: &sess.sandbox_policy, + sandbox_policy: &turn_context.sandbox_policy, codex_linux_sandbox_exe: &sess.codex_linux_sandbox_exe, stdout_stream: Some(StdoutStream { sub_id: sub_id.clone(), @@ -2276,8 +2425,13 @@ fn get_last_assistant_message_from_turn(responses: &[ResponseItem]) -> Option CodexResult<()> { - let mut stream = sess.client.clone().stream(prompt).await?; +async fn drain_to_completed( + sess: &Session, + turn_context: &TurnContext, + sub_id: &str, + prompt: &Prompt, +) -> CodexResult<()> { + let mut stream = turn_context.client.clone().stream(prompt).await?; loop { let maybe_event = stream.next().await; let Some(event) = maybe_event else { @@ -2289,7 +2443,7 @@ async fn drain_to_completed(sess: &Session, sub_id: &str, prompt: &Prompt) -> Co match event { Ok(ResponseEvent::OutputItemDone(item)) => { // Record only to in-memory conversation history; avoid state snapshot. - let mut state = sess.state.lock().unwrap(); + let mut state = sess.state.lock_unchecked(); state.history.record_items(std::slice::from_ref(&item)); } Ok(ResponseEvent::Completed { diff --git a/codex-rs/core/src/config.rs b/codex-rs/core/src/config.rs index b17bb80815..e2a68d07dc 100644 --- a/codex-rs/core/src/config.rs +++ b/codex-rs/core/src/config.rs @@ -156,6 +156,11 @@ pub struct Config { /// Include an experimental plan tool that the model can use to update its current plan and status of each step. pub include_plan_tool: bool, + /// Include the `apply_patch` tool for models that benefit from invoking + /// file edits as a structured tool call. When unset, this falls back to the + /// model family's default preference. + pub include_apply_patch_tool: bool, + /// The value for the `originator` header included with Responses API requests. pub internal_originator: Option, } @@ -480,6 +485,7 @@ pub struct ConfigOverrides { pub codex_linux_sandbox_exe: Option, pub base_instructions: Option, pub include_plan_tool: Option, + pub include_apply_patch_tool: Option, pub disable_response_storage: Option, pub show_raw_agent_reasoning: Option, } @@ -505,6 +511,7 @@ impl Config { codex_linux_sandbox_exe, base_instructions, include_plan_tool, + include_apply_patch_tool, disable_response_storage, show_raw_agent_reasoning, } = overrides; @@ -581,6 +588,7 @@ impl Config { needs_special_apply_patch_instructions: false, supports_reasoning_summaries, uses_local_shell_tool: false, + uses_apply_patch_tool: false, } }); @@ -607,6 +615,9 @@ impl Config { Self::get_base_instructions(experimental_instructions_path, &resolved_cwd)?; let base_instructions = base_instructions.or(file_base_instructions); + let include_apply_patch_tool_val = + include_apply_patch_tool.unwrap_or(model_family.uses_apply_patch_tool); + let config = Self { model, model_family, @@ -659,6 +670,7 @@ impl Config { experimental_resume, include_plan_tool: include_plan_tool.unwrap_or(false), + include_apply_patch_tool: include_apply_patch_tool_val, internal_originator: cfg.internal_originator, }; Ok(config) @@ -1022,6 +1034,7 @@ disable_response_storage = true experimental_resume: None, base_instructions: None, include_plan_tool: false, + include_apply_patch_tool: false, internal_originator: None, }, o3_profile_config @@ -1073,6 +1086,7 @@ disable_response_storage = true experimental_resume: None, base_instructions: None, include_plan_tool: false, + include_apply_patch_tool: false, internal_originator: None, }; @@ -1139,6 +1153,7 @@ disable_response_storage = true experimental_resume: None, base_instructions: None, include_plan_tool: false, + include_apply_patch_tool: false, internal_originator: None, }; diff --git a/codex-rs/core/src/model_family.rs b/codex-rs/core/src/model_family.rs index fa4826d76f..6d1c2efcc1 100644 --- a/codex-rs/core/src/model_family.rs +++ b/codex-rs/core/src/model_family.rs @@ -23,6 +23,10 @@ pub struct ModelFamily { // the model such that its description can be omitted. // See https://platform.openai.com/docs/guides/tools-local-shell pub uses_local_shell_tool: bool, + + /// True if the model performs better when `apply_patch` is provided as + /// a tool call instead of just a bash command. + pub uses_apply_patch_tool: bool, } macro_rules! model_family { @@ -36,6 +40,7 @@ macro_rules! model_family { needs_special_apply_patch_instructions: false, supports_reasoning_summaries: false, uses_local_shell_tool: false, + uses_apply_patch_tool: false, }; // apply overrides $( @@ -55,6 +60,7 @@ macro_rules! simple_model_family { needs_special_apply_patch_instructions: false, supports_reasoning_summaries: false, uses_local_shell_tool: false, + uses_apply_patch_tool: false, }) }}; } @@ -88,10 +94,10 @@ pub fn find_family_for_model(slug: &str) -> Option { slug, "gpt-4.1", needs_special_apply_patch_instructions: true, ) + } else if slug.starts_with("gpt-oss") { + model_family!(slug, "gpt-oss", uses_apply_patch_tool: true) } else if slug.starts_with("gpt-4o") { simple_model_family!(slug, "gpt-4o") - } else if slug.starts_with("gpt-oss") { - simple_model_family!(slug, "gpt-oss") } else if slug.starts_with("gpt-3.5") { simple_model_family!(slug, "gpt-3.5") } else if slug.starts_with("gpt-5") { diff --git a/codex-rs/core/src/openai_tools.rs b/codex-rs/core/src/openai_tools.rs index 7c65880433..32ead20e7d 100644 --- a/codex-rs/core/src/openai_tools.rs +++ b/codex-rs/core/src/openai_tools.rs @@ -43,6 +43,7 @@ pub enum ConfigShellToolType { pub struct ToolsConfig { pub shell_type: ConfigShellToolType, pub plan_tool: bool, + pub apply_patch_tool: bool, } impl ToolsConfig { @@ -51,6 +52,7 @@ impl ToolsConfig { approval_policy: AskForApproval, sandbox_policy: SandboxPolicy, include_plan_tool: bool, + include_apply_patch_tool: bool, ) -> Self { let mut shell_type = if model_family.uses_local_shell_tool { ConfigShellToolType::LocalShell @@ -66,6 +68,7 @@ impl ToolsConfig { Self { shell_type, plan_tool: include_plan_tool, + apply_patch_tool: include_apply_patch_tool || model_family.uses_apply_patch_tool, } } } @@ -235,6 +238,87 @@ The shell tool is used to execute shell commands. }) } +#[derive(Serialize, Deserialize)] +pub(crate) struct ApplyPatchToolArgs { + pub(crate) input: String, +} + +fn create_apply_patch_tool() -> OpenAiTool { + // Minimal schema: one required string argument containing the patch body + let mut properties = BTreeMap::new(); + properties.insert( + "input".to_string(), + JsonSchema::String { + description: Some(r#"The entire contents of the apply_patch command"#.to_string()), + }, + ); + + OpenAiTool::Function(ResponsesApiTool { + name: "apply_patch".to_string(), + description: r#"Use this 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: + +**_ Begin Patch +[ one or more file sections ] +_** End Patch + +Within that envelope, you get a sequence of file operations. +You MUST include a header to specify the action you are taking. +Each operation starts with one of three headers: + +**_ Add File: - create a new file. Every following line is a + line (the initial contents). +_** Delete File: - remove an existing file. Nothing follows. +\*\*\* Update File: - patch an existing file in place (optionally with a rename). + +May be immediately followed by \*\*\* Move to: if you want to rename the file. +Then one or more “hunks”, each introduced by @@ (optionally followed by a hunk header). +Within a hunk each line starts with: + +- for inserted text, + +* for removed text, or + space ( ) for context. + At the end of a truncated hunk you can emit \*\*\* End of File. + +Patch := Begin { FileOp } End +Begin := "**_ Begin Patch" NEWLINE +End := "_** End Patch" NEWLINE +FileOp := AddFile | DeleteFile | UpdateFile +AddFile := "**_ Add File: " path NEWLINE { "+" line NEWLINE } +DeleteFile := "_** Delete File: " path NEWLINE +UpdateFile := "**_ Update File: " path NEWLINE [ MoveTo ] { Hunk } +MoveTo := "_** Move to: " newPath NEWLINE +Hunk := "@@" [ header ] NEWLINE { HunkLine } [ "*** End of File" NEWLINE ] +HunkLine := (" " | "-" | "+") text NEWLINE + +A full patch can combine several operations: + +**_ Begin Patch +_** Add File: hello.txt ++Hello world +**_ Update File: src/app.py +_** Move to: src/main.py +@@ def greet(): +-print("Hi") ++print("Hello, world!") +**_ Delete File: obsolete.txt +_** End Patch + +It is important to remember: + +- You must include a header with your intended action (Add/Delete/Update) +- You must prefix new lines with `+` even when creating a new file +"# + .to_string(), + strict: false, + parameters: JsonSchema::Object { + properties, + required: Some(vec!["input".to_string()]), + additional_properties: Some(false), + }, + }) +} + /// Returns JSON values that are compatible with Function Calling in the /// Responses API: /// https://platform.openai.com/docs/guides/function-calling?api-mode=responses @@ -455,6 +539,10 @@ pub(crate) fn get_openai_tools( tools.push(PLAN_TOOL.clone()); } + if config.apply_patch_tool { + tools.push(create_apply_patch_tool()); + } + if let Some(mcp_tools) = mcp_tools { for (name, tool) in mcp_tools { match mcp_tool_to_openai_tool(name.clone(), tool.clone()) { @@ -508,6 +596,7 @@ mod tests { AskForApproval::Never, SandboxPolicy::ReadOnly, true, + model_family.uses_apply_patch_tool, ); let tools = get_openai_tools(&config, Some(HashMap::new())); @@ -522,6 +611,7 @@ mod tests { AskForApproval::Never, SandboxPolicy::ReadOnly, true, + model_family.uses_apply_patch_tool, ); let tools = get_openai_tools(&config, Some(HashMap::new())); @@ -536,6 +626,7 @@ mod tests { AskForApproval::Never, SandboxPolicy::ReadOnly, false, + model_family.uses_apply_patch_tool, ); let tools = get_openai_tools( &config, @@ -629,6 +720,7 @@ mod tests { AskForApproval::Never, SandboxPolicy::ReadOnly, false, + model_family.uses_apply_patch_tool, ); let tools = get_openai_tools( @@ -684,6 +776,7 @@ mod tests { AskForApproval::Never, SandboxPolicy::ReadOnly, false, + model_family.uses_apply_patch_tool, ); let tools = get_openai_tools( @@ -734,6 +827,7 @@ mod tests { AskForApproval::Never, SandboxPolicy::ReadOnly, false, + model_family.uses_apply_patch_tool, ); let tools = get_openai_tools( @@ -787,6 +881,7 @@ mod tests { AskForApproval::Never, SandboxPolicy::ReadOnly, false, + model_family.uses_apply_patch_tool, ); let tools = get_openai_tools( diff --git a/codex-rs/core/src/protocol.rs b/codex-rs/core/src/protocol.rs index 4184f1a945..d334c2eb86 100644 --- a/codex-rs/core/src/protocol.rs +++ b/codex-rs/core/src/protocol.rs @@ -185,10 +185,31 @@ pub enum SandboxPolicy { /// not modified by the agent. #[derive(Debug, Clone, PartialEq, Eq)] pub struct WritableRoot { + /// Absolute path, by construction. pub root: PathBuf, + + /// Also absolute paths, by construction. pub read_only_subpaths: Vec, } +impl WritableRoot { + pub(crate) fn is_path_writable(&self, path: &Path) -> bool { + // Check if the path is under the root. + if !path.starts_with(&self.root) { + return false; + } + + // Check if the path is under any of the read-only subpaths. + for subpath in &self.read_only_subpaths { + if path.starts_with(subpath) { + return false; + } + } + + true + } +} + impl FromStr for SandboxPolicy { type Err = serde_json::Error; diff --git a/codex-rs/core/src/safety.rs b/codex-rs/core/src/safety.rs index 74872ddc4f..c878a71110 100644 --- a/codex-rs/core/src/safety.rs +++ b/codex-rs/core/src/safety.rs @@ -21,7 +21,7 @@ pub enum SafetyCheck { pub fn assess_patch_safety( action: &ApplyPatchAction, policy: AskForApproval, - writable_roots: &[PathBuf], + sandbox_policy: &SandboxPolicy, cwd: &Path, ) -> SafetyCheck { if action.is_empty() { @@ -45,7 +45,7 @@ pub fn assess_patch_safety( // is possible that paths in the patch are hard links to files outside the // writable roots, so we should still run `apply_patch` in a sandbox in that // case. - if is_write_patch_constrained_to_writable_paths(action, writable_roots, cwd) + if is_write_patch_constrained_to_writable_paths(action, sandbox_policy, cwd) || policy == AskForApproval::OnFailure { // Only auto‑approve when we can actually enforce a sandbox. Otherwise @@ -171,13 +171,19 @@ pub fn get_platform_sandbox() -> Option { fn is_write_patch_constrained_to_writable_paths( action: &ApplyPatchAction, - writable_roots: &[PathBuf], + sandbox_policy: &SandboxPolicy, cwd: &Path, ) -> bool { // Early‑exit if there are no declared writable roots. - if writable_roots.is_empty() { - return false; - } + let writable_roots = match sandbox_policy { + SandboxPolicy::ReadOnly => { + return false; + } + SandboxPolicy::DangerFullAccess => { + return true; + } + SandboxPolicy::WorkspaceWrite { .. } => sandbox_policy.get_writable_roots_with_cwd(cwd), + }; // Normalize a path by removing `.` and resolving `..` without touching the // filesystem (works even if the file does not exist). @@ -209,15 +215,9 @@ fn is_write_patch_constrained_to_writable_paths( None => return false, }; - writable_roots.iter().any(|root| { - let root_abs = if root.is_absolute() { - root.clone() - } else { - normalize(&cwd.join(root)).unwrap_or_else(|| cwd.join(root)) - }; - - abs.starts_with(&root_abs) - }) + writable_roots + .iter() + .any(|writable_root| writable_root.is_path_writable(&abs)) }; for (path, change) in action.changes() { @@ -246,38 +246,56 @@ fn is_write_patch_constrained_to_writable_paths( #[cfg(test)] mod tests { use super::*; + use tempfile::TempDir; #[test] fn test_writable_roots_constraint() { - let cwd = std::env::current_dir().unwrap(); + // Use a temporary directory as our workspace to avoid touching + // the real current working directory. + let tmp = TempDir::new().unwrap(); + let cwd = tmp.path().to_path_buf(); let parent = cwd.parent().unwrap().to_path_buf(); - // Helper to build a single‑entry map representing a patch that adds a - // file at `p`. + // Helper to build a single‑entry patch that adds a file at `p`. let make_add_change = |p: PathBuf| ApplyPatchAction::new_add_for_test(&p, "".to_string()); let add_inside = make_add_change(cwd.join("inner.txt")); let add_outside = make_add_change(parent.join("outside.txt")); + // Policy limited to the workspace only; exclude system temp roots so + // only `cwd` is writable by default. + let policy_workspace_only = SandboxPolicy::WorkspaceWrite { + writable_roots: vec![], + network_access: false, + exclude_tmpdir_env_var: true, + exclude_slash_tmp: true, + }; + assert!(is_write_patch_constrained_to_writable_paths( &add_inside, - &[PathBuf::from(".")], + &policy_workspace_only, &cwd, )); - let add_outside_2 = make_add_change(parent.join("outside.txt")); assert!(!is_write_patch_constrained_to_writable_paths( - &add_outside_2, - &[PathBuf::from(".")], + &add_outside, + &policy_workspace_only, &cwd, )); - // With parent dir added as writable root, it should pass. + // With the parent dir explicitly added as a writable root, the + // outside write should be permitted. + let policy_with_parent = SandboxPolicy::WorkspaceWrite { + writable_roots: vec![parent.clone()], + network_access: false, + exclude_tmpdir_env_var: true, + exclude_slash_tmp: true, + }; assert!(is_write_patch_constrained_to_writable_paths( &add_outside, - &[PathBuf::from("..")], + &policy_with_parent, &cwd, - )) + )); } #[test] diff --git a/codex-rs/core/tests/common/lib.rs b/codex-rs/core/tests/common/lib.rs index 0a9c8d5aa8..244d093e7d 100644 --- a/codex-rs/core/tests/common/lib.rs +++ b/codex-rs/core/tests/common/lib.rs @@ -47,6 +47,26 @@ pub fn load_sse_fixture(path: impl AsRef) -> String { .collect() } +pub fn load_sse_fixture_with_id_from_str(raw: &str, id: &str) -> String { + let replaced = raw.replace("__ID__", id); + let events: Vec = + serde_json::from_str(&replaced).expect("parse JSON fixture"); + events + .into_iter() + .map(|e| { + let kind = e + .get("type") + .and_then(|v| v.as_str()) + .expect("fixture event missing type"); + if e.as_object().map(|o| o.len() == 1).unwrap_or(false) { + format!("event: {kind}\n\n") + } else { + format!("event: {kind}\ndata: {e}\n\n") + } + }) + .collect() +} + /// Same as [`load_sse_fixture`], but replaces the placeholder `__ID__` in the /// fixture template with the supplied identifier before parsing. This lets a /// single JSON template be reused by multiple tests that each need a unique diff --git a/codex-rs/exec/Cargo.toml b/codex-rs/exec/Cargo.toml index b7c20df321..9847788d92 100644 --- a/codex-rs/exec/Cargo.toml +++ b/codex-rs/exec/Cargo.toml @@ -44,3 +44,5 @@ assert_cmd = "2" libc = "0.2" predicates = "3" tempfile = "3.13.0" +wiremock = "0.6" +core_test_support = { path = "../core/tests/common" } diff --git a/codex-rs/exec/src/lib.rs b/codex-rs/exec/src/lib.rs index ff6123d74b..e6b4d7fb0c 100644 --- a/codex-rs/exec/src/lib.rs +++ b/codex-rs/exec/src/lib.rs @@ -146,6 +146,7 @@ pub async fn run_main(cli: Cli, codex_linux_sandbox_exe: Option) -> any codex_linux_sandbox_exe, base_instructions: None, include_plan_tool: None, + include_apply_patch_tool: None, disable_response_storage: oss.then_some(true), show_raw_agent_reasoning: oss.then_some(true), }; diff --git a/codex-rs/exec/tests/apply_patch.rs b/codex-rs/exec/tests/apply_patch.rs index f65d32e1c8..ecce43d732 100644 --- a/codex-rs/exec/tests/apply_patch.rs +++ b/codex-rs/exec/tests/apply_patch.rs @@ -1,3 +1,5 @@ +#![allow(clippy::expect_used, clippy::unwrap_used)] + use anyhow::Context; use assert_cmd::prelude::*; use codex_core::CODEX_APPLY_PATCH_ARG1; @@ -37,3 +39,152 @@ fn test_standalone_exec_cli_can_use_apply_patch() -> anyhow::Result<()> { ); Ok(()) } + +#[cfg(not(target_os = "windows"))] +#[tokio::test] +async fn test_apply_patch_tool() -> anyhow::Result<()> { + use core_test_support::load_sse_fixture_with_id_from_str; + use tempfile::TempDir; + use wiremock::Mock; + use wiremock::MockServer; + use wiremock::ResponseTemplate; + use wiremock::matchers::method; + use wiremock::matchers::path; + + const SSE_TOOL_CALL_ADD: &str = r#"[ + { + "type": "response.output_item.done", + "item": { + "type": "function_call", + "name": "apply_patch", + "arguments": "{\n \"input\": \"*** Begin Patch\\n*** Add File: test.md\\n+Hello world\\n*** End Patch\"\n}", + "call_id": "__ID__" + } + }, + { + "type": "response.completed", + "response": { + "id": "__ID__", + "usage": { + "input_tokens": 0, + "input_tokens_details": null, + "output_tokens": 0, + "output_tokens_details": null, + "total_tokens": 0 + }, + "output": [] + } + } +]"#; + + const SSE_TOOL_CALL_UPDATE: &str = r#"[ + { + "type": "response.output_item.done", + "item": { + "type": "function_call", + "name": "apply_patch", + "arguments": "{\n \"input\": \"*** Begin Patch\\n*** Update File: test.md\\n@@\\n-Hello world\\n+Final text\\n*** End Patch\"\n}", + "call_id": "__ID__" + } + }, + { + "type": "response.completed", + "response": { + "id": "__ID__", + "usage": { + "input_tokens": 0, + "input_tokens_details": null, + "output_tokens": 0, + "output_tokens_details": null, + "total_tokens": 0 + }, + "output": [] + } + } +]"#; + + const SSE_TOOL_CALL_COMPLETED: &str = r#"[ + { + "type": "response.completed", + "response": { + "id": "__ID__", + "usage": { + "input_tokens": 0, + "input_tokens_details": null, + "output_tokens": 0, + "output_tokens_details": null, + "total_tokens": 0 + }, + "output": [] + } + } +]"#; + + // Start a mock model server + let server = MockServer::start().await; + + // First response: model calls apply_patch to create test.md + let first = ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_raw( + load_sse_fixture_with_id_from_str(SSE_TOOL_CALL_ADD, "call1"), + "text/event-stream", + ); + + Mock::given(method("POST")) + // .and(path("/v1/responses")) + .respond_with(first) + .up_to_n_times(1) + .mount(&server) + .await; + + // Second response: model calls apply_patch to update test.md + let second = ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_raw( + load_sse_fixture_with_id_from_str(SSE_TOOL_CALL_UPDATE, "call2"), + "text/event-stream", + ); + + Mock::given(method("POST")) + .and(path("/v1/responses")) + .respond_with(second) + .up_to_n_times(1) + .mount(&server) + .await; + + let final_completed = ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_raw( + load_sse_fixture_with_id_from_str(SSE_TOOL_CALL_COMPLETED, "resp3"), + "text/event-stream", + ); + + Mock::given(method("POST")) + // .and(path("/v1/responses")) + .respond_with(final_completed) + .expect(1) + .mount(&server) + .await; + + let tmp_cwd = TempDir::new().unwrap(); + Command::cargo_bin("codex-exec") + .context("should find binary for codex-exec")? + .current_dir(tmp_cwd.path()) + .env("CODEX_HOME", tmp_cwd.path()) + .env("OPENAI_API_KEY", "dummy") + .env("OPENAI_BASE_URL", format!("{}/v1", server.uri())) + .arg("--skip-git-repo-check") + .arg("-s") + .arg("workspace-write") + .arg("foo") + .assert() + .success(); + + // Verify final file contents + let final_path = tmp_cwd.path().join("test.md"); + let contents = std::fs::read_to_string(&final_path) + .unwrap_or_else(|e| panic!("failed reading {}: {e}", final_path.display())); + assert_eq!(contents, "Final text\n"); + Ok(()) +} diff --git a/codex-rs/mcp-server/src/codex_message_processor.rs b/codex-rs/mcp-server/src/codex_message_processor.rs index 54a418fd7a..d930c03b71 100644 --- a/codex-rs/mcp-server/src/codex_message_processor.rs +++ b/codex-rs/mcp-server/src/codex_message_processor.rs @@ -424,6 +424,7 @@ fn derive_config_from_params( config: cli_overrides, base_instructions, include_plan_tool, + include_apply_patch_tool, } = params; let overrides = ConfigOverrides { model, @@ -435,6 +436,7 @@ fn derive_config_from_params( codex_linux_sandbox_exe, base_instructions, include_plan_tool, + include_apply_patch_tool, disable_response_storage: None, show_raw_agent_reasoning: None, }; diff --git a/codex-rs/mcp-server/src/codex_tool_config.rs b/codex-rs/mcp-server/src/codex_tool_config.rs index 548f29334f..906921a030 100644 --- a/codex-rs/mcp-server/src/codex_tool_config.rs +++ b/codex-rs/mcp-server/src/codex_tool_config.rs @@ -160,6 +160,7 @@ impl CodexToolCallParam { codex_linux_sandbox_exe, base_instructions, include_plan_tool, + include_apply_patch_tool: None, disable_response_storage: None, show_raw_agent_reasoning: None, }; diff --git a/codex-rs/mcp-server/src/tool_handlers/create_conversation.rs b/codex-rs/mcp-server/src/tool_handlers/create_conversation.rs index 77a68f0128..eee2e1d5f4 100644 --- a/codex-rs/mcp-server/src/tool_handlers/create_conversation.rs +++ b/codex-rs/mcp-server/src/tool_handlers/create_conversation.rs @@ -52,6 +52,7 @@ pub(crate) async fn handle_create_conversation( codex_linux_sandbox_exe: None, base_instructions, include_plan_tool: None, + include_apply_patch_tool: None, disable_response_storage: None, show_raw_agent_reasoning: None, }; diff --git a/codex-rs/mcp-server/src/wire_format.rs b/codex-rs/mcp-server/src/wire_format.rs index 5df3344216..68d9aeb9eb 100644 --- a/codex-rs/mcp-server/src/wire_format.rs +++ b/codex-rs/mcp-server/src/wire_format.rs @@ -99,6 +99,10 @@ pub struct NewConversationParams { /// Whether to include the plan tool in the conversation. #[serde(skip_serializing_if = "Option::is_none")] pub include_plan_tool: Option, + + /// Whether to include the apply patch tool in the conversation. + #[serde(skip_serializing_if = "Option::is_none")] + pub include_apply_patch_tool: Option, } #[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] @@ -267,6 +271,7 @@ mod tests { config: None, base_instructions: None, include_plan_tool: None, + include_apply_patch_tool: None, }, }; assert_eq!( diff --git a/codex-rs/mcp-server/tests/codex_message_processor_flow.rs b/codex-rs/mcp-server/tests/codex_message_processor_flow.rs index 2cc55c6d62..e0c7a83209 100644 --- a/codex-rs/mcp-server/tests/codex_message_processor_flow.rs +++ b/codex-rs/mcp-server/tests/codex_message_processor_flow.rs @@ -1,14 +1,21 @@ use std::path::Path; +use codex_core::config_types::ReasoningEffort; +use codex_core::config_types::ReasoningSummary; +use codex_core::protocol::AskForApproval; +use codex_core::protocol::SandboxPolicy; use codex_core::spawn::CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR; use codex_mcp_server::wire_format::AddConversationListenerParams; use codex_mcp_server::wire_format::AddConversationSubscriptionResponse; +use codex_mcp_server::wire_format::EXEC_COMMAND_APPROVAL_METHOD; use codex_mcp_server::wire_format::NewConversationParams; use codex_mcp_server::wire_format::NewConversationResponse; use codex_mcp_server::wire_format::RemoveConversationListenerParams; use codex_mcp_server::wire_format::RemoveConversationSubscriptionResponse; use codex_mcp_server::wire_format::SendUserMessageParams; use codex_mcp_server::wire_format::SendUserMessageResponse; +use codex_mcp_server::wire_format::SendUserTurnParams; +use codex_mcp_server::wire_format::SendUserTurnResponse; use mcp_test_support::McpProcess; use mcp_test_support::create_final_assistant_message_sse_response; use mcp_test_support::create_mock_chat_completions_server; @@ -167,6 +174,184 @@ fn to_response(response: JSONRPCResponse) -> anyhow::Result Ok(codex_response) } +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_send_user_turn_changes_approval_policy_behavior() { + if env::var(CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR).is_ok() { + println!( + "Skipping test because it cannot execute when network is disabled in a Codex sandbox." + ); + return; + } + + let tmp = TempDir::new().expect("tmp dir"); + let codex_home = tmp.path().join("codex_home"); + std::fs::create_dir(&codex_home).expect("create codex home dir"); + let working_directory = tmp.path().join("workdir"); + std::fs::create_dir(&working_directory).expect("create working directory"); + + // Mock server will request a python shell call for the first and second turn, then finish. + let responses = vec![ + create_shell_sse_response( + vec![ + "python3".to_string(), + "-c".to_string(), + "print(42)".to_string(), + ], + Some(&working_directory), + Some(5000), + "call1", + ) + .expect("create first shell sse response"), + create_final_assistant_message_sse_response("done 1") + .expect("create final assistant message 1"), + create_shell_sse_response( + vec![ + "python3".to_string(), + "-c".to_string(), + "print(42)".to_string(), + ], + Some(&working_directory), + Some(5000), + "call2", + ) + .expect("create second shell sse response"), + create_final_assistant_message_sse_response("done 2") + .expect("create final assistant message 2"), + ]; + let server = create_mock_chat_completions_server(responses).await; + create_config_toml(&codex_home, &server.uri()).expect("write config"); + + // Start MCP server and initialize. + let mut mcp = McpProcess::new(&codex_home).await.expect("spawn mcp"); + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()) + .await + .expect("init timeout") + .expect("init error"); + + // 1) Start conversation with approval_policy=untrusted + let new_conv_id = mcp + .send_new_conversation_request(NewConversationParams { + cwd: Some(working_directory.to_string_lossy().into_owned()), + ..Default::default() + }) + .await + .expect("send newConversation"); + let new_conv_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(new_conv_id)), + ) + .await + .expect("newConversation timeout") + .expect("newConversation resp"); + let NewConversationResponse { + conversation_id, .. + } = to_response::(new_conv_resp) + .expect("deserialize newConversation response"); + + // 2) addConversationListener + let add_listener_id = mcp + .send_add_conversation_listener_request(AddConversationListenerParams { conversation_id }) + .await + .expect("send addConversationListener"); + let _: AddConversationSubscriptionResponse = + to_response::( + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(add_listener_id)), + ) + .await + .expect("addConversationListener timeout") + .expect("addConversationListener resp"), + ) + .expect("deserialize addConversationListener response"); + + // 3) sendUserMessage triggers a shell call; approval policy is Untrusted so we should get an elicitation + let send_user_id = mcp + .send_send_user_message_request(SendUserMessageParams { + conversation_id, + items: vec![codex_mcp_server::wire_format::InputItem::Text { + text: "run python".to_string(), + }], + }) + .await + .expect("send sendUserMessage"); + let _send_user_resp: SendUserMessageResponse = to_response::( + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(send_user_id)), + ) + .await + .expect("sendUserMessage timeout") + .expect("sendUserMessage resp"), + ) + .expect("deserialize sendUserMessage response"); + + // Expect an ExecCommandApproval request (elicitation) + let request = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_request_message(), + ) + .await + .expect("waiting for exec approval request timeout") + .expect("exec approval request"); + assert_eq!(request.method, EXEC_COMMAND_APPROVAL_METHOD); + + // Approve so the first turn can complete + mcp.send_response( + request.id, + serde_json::json!({ "decision": codex_core::protocol::ReviewDecision::Approved }), + ) + .await + .expect("send approval response"); + + // Wait for first TaskComplete + let _ = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("codex/event/task_complete"), + ) + .await + .expect("task_complete 1 timeout") + .expect("task_complete 1 notification"); + + // 4) sendUserTurn with approval_policy=never should run without elicitation + let send_turn_id = mcp + .send_send_user_turn_request(SendUserTurnParams { + conversation_id, + items: vec![codex_mcp_server::wire_format::InputItem::Text { + text: "run python again".to_string(), + }], + cwd: working_directory.clone(), + approval_policy: AskForApproval::Never, + sandbox_policy: SandboxPolicy::new_read_only_policy(), + model: "mock-model".to_string(), + effort: ReasoningEffort::Medium, + summary: ReasoningSummary::Auto, + }) + .await + .expect("send sendUserTurn"); + // Acknowledge sendUserTurn + let _send_turn_resp: SendUserTurnResponse = to_response::( + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(send_turn_id)), + ) + .await + .expect("sendUserTurn timeout") + .expect("sendUserTurn resp"), + ) + .expect("deserialize sendUserTurn response"); + + // Ensure we do NOT receive an ExecCommandApproval request before the task completes. + // If any Request is seen while waiting for task_complete, the helper will error and the test fails. + let _ = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("codex/event/task_complete"), + ) + .await + .expect("task_complete 2 timeout") + .expect("task_complete 2 notification"); +} + // Helper: minimal config.toml pointing at mock provider. fn create_config_toml(codex_home: &Path, server_uri: &str) -> std::io::Result<()> { let config_toml = codex_home.join("config.toml"); @@ -175,7 +360,7 @@ fn create_config_toml(codex_home: &Path, server_uri: &str) -> std::io::Result<() format!( r#" model = "mock-model" -approval_policy = "never" +approval_policy = "untrusted" model_provider = "mock_provider" diff --git a/codex-rs/mcp-server/tests/common/mcp_process.rs b/codex-rs/mcp-server/tests/common/mcp_process.rs index 35484264fa..dc7833441c 100644 --- a/codex-rs/mcp-server/tests/common/mcp_process.rs +++ b/codex-rs/mcp-server/tests/common/mcp_process.rs @@ -22,6 +22,7 @@ use codex_mcp_server::wire_format::AddConversationListenerParams; use codex_mcp_server::wire_format::NewConversationParams; use codex_mcp_server::wire_format::RemoveConversationListenerParams; use codex_mcp_server::wire_format::SendUserMessageParams; +use codex_mcp_server::wire_format::SendUserTurnParams; use mcp_types::CallToolRequestParams; use mcp_types::ClientCapabilities; @@ -281,6 +282,15 @@ impl McpProcess { .await } + /// Send a `sendUserTurn` JSON-RPC request. + pub async fn send_send_user_turn_request( + &mut self, + params: SendUserTurnParams, + ) -> anyhow::Result { + let params = Some(serde_json::to_value(params)?); + self.send_request("sendUserTurn", params).await + } + async fn send_request( &mut self, method: &str, diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index f1a7d99b79..7d605d683c 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -116,6 +116,7 @@ pub async fn run_main( codex_linux_sandbox_exe, base_instructions: None, include_plan_tool: Some(true), + include_apply_patch_tool: None, disable_response_storage: cli.oss.then_some(true), show_raw_agent_reasoning: cli.oss.then_some(true), };