From 6df8e353146c4cc2f2b0841c465dd85c35a5cae0 Mon Sep 17 00:00:00 2001 From: Dylan Date: Fri, 15 Aug 2025 11:55:53 -0400 Subject: [PATCH 1/7] [tools] Add apply_patch tool (#2303) ## Summary We've been seeing a number of issues and reports with our synthetic `apply_patch` tool, e.g. #802. Let's make this a real tool - in my anecdotal testing, it's critical for GPT-OSS models, but I'd like to make it the standard across GPT-5 and codex models as well. ## Testing - [x] Tested locally - [x] Integration test --- codex-rs/Cargo.lock | 2 + codex-rs/core/src/codex.rs | 26 +++ codex-rs/core/src/config.rs | 15 ++ codex-rs/core/src/model_family.rs | 10 +- codex-rs/core/src/openai_tools.rs | 95 +++++++++++ codex-rs/core/tests/common/lib.rs | 20 +++ codex-rs/exec/Cargo.toml | 2 + codex-rs/exec/src/lib.rs | 1 + codex-rs/exec/tests/apply_patch.rs | 151 ++++++++++++++++++ .../mcp-server/src/codex_message_processor.rs | 2 + codex-rs/mcp-server/src/codex_tool_config.rs | 1 + .../src/tool_handlers/create_conversation.rs | 1 + codex-rs/mcp-server/src/wire_format.rs | 5 + codex-rs/tui/src/lib.rs | 1 + 14 files changed, 330 insertions(+), 2 deletions(-) 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/codex.rs b/codex-rs/core/src/codex.rs index 482ac2f1ab..edc034ee43 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -67,6 +67,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; @@ -455,6 +456,7 @@ impl Session { approval_policy, sandbox_policy.clone(), config.include_plan_tool, + config.include_apply_patch_tool, ), tx_event: tx_event.clone(), user_instructions, @@ -1727,6 +1729,30 @@ async fn handle_function_call( handle_container_exec_with_params(params, sess, 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, _ => { match sess.mcp_connection_manager.parse_tool_name(&name) { 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/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 2495feb9e0..cdf24d9214 100644 --- a/codex-rs/mcp-server/src/codex_message_processor.rs +++ b/codex-rs/mcp-server/src/codex_message_processor.rs @@ -367,6 +367,7 @@ fn derive_config_from_params( config: cli_overrides, base_instructions, include_plan_tool, + include_apply_patch_tool, } = params; let overrides = ConfigOverrides { model, @@ -378,6 +379,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 034c72a1c2..e2ba729eb5 100644 --- a/codex-rs/mcp-server/src/wire_format.rs +++ b/codex-rs/mcp-server/src/wire_format.rs @@ -90,6 +90,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)] @@ -241,6 +245,7 @@ mod tests { config: None, base_instructions: None, include_plan_tool: None, + include_apply_patch_tool: None, }, }; assert_eq!( 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), }; From 26c8373821607d41dce50ec4c6a4717053d3c062 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Fri, 15 Aug 2025 09:06:15 -0700 Subject: [PATCH 2/7] fix: tighten up checks against writable folders for SandboxPolicy (#2338) I was looking at the implementation of `Session::get_writable_roots()`, which did not seem right, as it was a copy of writable roots, which is not guaranteed to be in sync with the `sandbox_policy` field. I looked at who was calling `get_writable_roots()` and its only call site was `apply_patch()` in `codex-rs/core/src/apply_patch.rs`, which took the roots and forwarded them to `assess_patch_safety()` in `safety.rs`. I updated `assess_patch_safety()` to take `sandbox_policy: &SandboxPolicy` instead of `writable_roots: &[PathBuf]` (and replaced `Session::get_writable_roots()` with `Session::get_sandbox_policy()`). Within `safety.rs`, it was fairly easy to update `is_write_patch_constrained_to_writable_paths()` to work with `SandboxPolicy`, and in particular, it is far more accurate because, for better or worse, `SandboxPolicy::get_writable_roots_with_cwd()` _returns an empty vec_ for `SandboxPolicy::DangerFullAccess`, suggesting that _nothing_ is writable when in reality _everything_ is writable. With this PR, `is_write_patch_constrained_to_writable_paths()` now does the right thing for each variant of `SandboxPolicy`. I thought this would be the end of the story, but it turned out that `test_writable_roots_constraint()` in `safety.rs` needed to be updated, as well. In particular, the test was writing to `std::env::current_dir()` instead of a `TempDir`, which I suspect was a holdover from earlier when `SandboxPolicy::WorkspaceWrite` would always make `TMPDIR` writable on macOS, which made it hard to write tests to verify `SandboxPolicy` in `TMPDIR`. Fortunately, we now have `exclude_tmpdir_env_var` as an option on `SandboxPolicy::WorkspaceWrite`, so I was able to update the test to preserve the existing behavior, but to no longer write to `std::env::current_dir()`. --- [//]: # (BEGIN SAPLING FOOTER) Stack created with [Sapling](https://sapling-scm.com). Best reviewed with [ReviewStack](https://reviewstack.dev/openai/codex/pull/2338). * #2345 * #2329 * #2343 * #2340 * __->__ #2338 --- codex-rs/core/src/apply_patch.rs | 32 +-------------- codex-rs/core/src/codex.rs | 15 +++---- codex-rs/core/src/protocol.rs | 21 ++++++++++ codex-rs/core/src/safety.rs | 68 ++++++++++++++++++++------------ 4 files changed, 70 insertions(+), 66 deletions(-) diff --git a/codex-rs/core/src/apply_patch.rs b/codex-rs/core/src/apply_patch.rs index 21e80406e5..fcccb40f8c 100644 --- a/codex-rs/core/src/apply_patch.rs +++ b/codex-rs/core/src/apply_patch.rs @@ -8,7 +8,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"; @@ -45,12 +44,10 @@ pub(crate) async fn apply_patch( 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_sandbox_policy(), sess.get_cwd(), ) { SafetyCheck::AutoApprove { .. } => { @@ -124,30 +121,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/codex.rs b/codex-rs/core/src/codex.rs index edc034ee43..c66ed72132 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -31,12 +31,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; @@ -231,7 +230,6 @@ pub(crate) struct Session { approval_policy: AskForApproval, sandbox_policy: SandboxPolicy, shell_environment_policy: ShellEnvironmentPolicy, - writable_roots: Vec, disable_response_storage: bool, tools_config: ToolsConfig, @@ -410,8 +408,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), @@ -465,7 +461,6 @@ impl Session { sandbox_policy, shell_environment_policy: config.shell_environment_policy.clone(), cwd, - writable_roots, mcp_connection_manager, notify, state: Mutex::new(state), @@ -509,14 +504,14 @@ 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_sandbox_policy(&self) -> &SandboxPolicy { + &self.sandbox_policy + } + pub(crate) fn get_cwd(&self) -> &Path { &self.cwd } diff --git a/codex-rs/core/src/protocol.rs b/codex-rs/core/src/protocol.rs index ac95b6a20a..1d264d3ed1 100644 --- a/codex-rs/core/src/protocol.rs +++ b/codex-rs/core/src/protocol.rs @@ -156,10 +156,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] From 6730592433e8bdb549b4168e6fa5579e5bafe489 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Fri, 15 Aug 2025 09:14:44 -0700 Subject: [PATCH 3/7] fix: introduce MutexExt::lock_unchecked() so we stop ignoring unwrap() throughout codex.rs (#2340) This way we are sure a dangerous `unwrap()` does not sneak in! --- [//]: # (BEGIN SAPLING FOOTER) Stack created with [Sapling](https://sapling-scm.com). Best reviewed with [ReviewStack](https://reviewstack.dev/openai/codex/pull/2340). * #2345 * #2329 * #2343 * __->__ #2340 * #2338 --- codex-rs/core/src/codex.rs | 58 +++++++++++++++++++++++--------------- 1 file changed, 36 insertions(+), 22 deletions(-) diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index c66ed72132..f66abe0072 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -1,6 +1,3 @@ -// Poisoned mutex should fail the program -#![expect(clippy::unwrap_used)] - use std::borrow::Cow; use std::collections::HashMap; use std::collections::HashSet; @@ -8,6 +5,7 @@ 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; @@ -108,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 { @@ -523,7 +536,7 @@ impl Session { } 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(); } @@ -531,7 +544,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(); @@ -567,7 +580,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 @@ -593,21 +606,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); } @@ -617,14 +630,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() }; @@ -802,12 +815,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(()) @@ -817,7 +830,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 { @@ -841,7 +854,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() { @@ -1045,7 +1058,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}"); @@ -1461,7 +1474,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); } @@ -1577,7 +1590,7 @@ 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); } @@ -1617,8 +1630,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, @@ -1912,7 +1926,7 @@ 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, @@ -2252,7 +2266,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 { From 265fd89e3141a6a0cce5f69c8f0e5edefd240c6e Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Fri, 15 Aug 2025 09:17:20 -0700 Subject: [PATCH 4/7] fix: try to fix flakiness in test_shell_command_approval_triggers_elicitation (#2344) I still see flakiness in `test_shell_command_approval_triggers_elicitation()` on occasion where `MockServer` claims it has not received all of its expected requests. I recently introduced a similar type of test in #2264, `test_codex_jsonrpc_conversation_flow()`, which I have not seen flake (yet!), so this PR pulls over two things I did in that test: - increased `worker_threads` from `2` to `4` - added an assertion to make sure the `task_complete` notification is received Honestly, I'm still not sure why `MockServer` claims it sometimes does not receive all its expected requests given that we assert that the final `JSONRPCResponse` is read on the stream, but let's give this a shot. Assuming this fixes things, my hypothesis is that the increase in `worker_threads` helps because perhaps there are async tasks in `MockServer` that do not reliably complete fully when there are not enough threads available? If that is correct, it seems like the test would still be flaky, though perhaps with lower frequency? --- codex-rs/mcp-server/tests/codex_tool.rs | 12 +++++- .../mcp-server/tests/common/mcp_process.rs | 42 +++++++++++++++++++ 2 files changed, 53 insertions(+), 1 deletion(-) diff --git a/codex-rs/mcp-server/tests/codex_tool.rs b/codex-rs/mcp-server/tests/codex_tool.rs index 92f11eaa4f..1ebd10a77a 100644 --- a/codex-rs/mcp-server/tests/codex_tool.rs +++ b/codex-rs/mcp-server/tests/codex_tool.rs @@ -35,7 +35,7 @@ const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs /// Test that a shell command that is not on the "trusted" list triggers an /// elicitation request to the MCP and that sending the approval runs the /// command, as expected. -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_shell_command_approval_triggers_elicitation() { if env::var(CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR).is_ok() { println!( @@ -114,6 +114,16 @@ async fn shell_command_approval_triggers_elicitation() -> anyhow::Result<()> { ) .await?; + // Verify task_complete notification arrives before the tool call completes. + #[expect(clippy::expect_used)] + let _task_complete = timeout( + DEFAULT_READ_TIMEOUT, + mcp_process.read_stream_until_legacy_task_complete_notification(), + ) + .await + .expect("task_complete_notification timeout") + .expect("task_complete_notification resp"); + // Verify the original `codex` tool call completes and that `git init` ran // successfully. let codex_response = timeout( diff --git a/codex-rs/mcp-server/tests/common/mcp_process.rs b/codex-rs/mcp-server/tests/common/mcp_process.rs index a659b1d950..35484264fa 100644 --- a/codex-rs/mcp-server/tests/common/mcp_process.rs +++ b/codex-rs/mcp-server/tests/common/mcp_process.rs @@ -474,4 +474,46 @@ impl McpProcess { })) .await } + + /// Reads notifications until a legacy TaskComplete event is observed: + /// Method "codex/event" with params.msg.type == "task_complete". + pub async fn read_stream_until_legacy_task_complete_notification( + &mut self, + ) -> anyhow::Result { + loop { + let message = self.read_jsonrpc_message().await?; + eprint!("message: {message:?}"); + + match message { + JSONRPCMessage::Notification(notification) => { + let is_match = if notification.method == "codex/event" { + if let Some(params) = ¬ification.params { + params + .get("msg") + .and_then(|m| m.get("type")) + .and_then(|t| t.as_str()) + == Some("task_complete") + } else { + false + } + } else { + false + }; + + if is_match { + return Ok(notification); + } + } + JSONRPCMessage::Request(_) => { + anyhow::bail!("unexpected JSONRPCMessage::Request: {message:?}"); + } + JSONRPCMessage::Error(_) => { + anyhow::bail!("unexpected JSONRPCMessage::Error: {message:?}"); + } + JSONRPCMessage::Response(_) => { + anyhow::bail!("unexpected JSONRPCMessage::Response: {message:?}"); + } + } + } + } } From c8ee33807c5817fa0b2c7780d8299e02d36041db Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Fri, 15 Aug 2025 09:20:24 -0700 Subject: [PATCH 5/7] feat: introduce TurnContext --- codex-rs/core/src/apply_patch.rs | 8 +- codex-rs/core/src/client.rs | 2 +- codex-rs/core/src/codex.rs | 256 +++++++++++++++++++------------ 3 files changed, 165 insertions(+), 101 deletions(-) diff --git a/codex-rs/core/src/apply_patch.rs b/codex-rs/core/src/apply_patch.rs index fcccb40f8c..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; @@ -40,15 +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 { match assess_patch_safety( &action, - sess.get_approval_policy(), - sess.get_sandbox_policy(), - sess.get_cwd(), + turn_context.approval_policy, + &turn_context.sandbox_policy, + &turn_context.cwd, ) { SafetyCheck::AutoApprove { .. } => { InternalApplyPatchInvocation::DelegateToExec(ApplyPatchExec { 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 f66abe0072..41e37640fa 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -1,7 +1,6 @@ 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; @@ -166,16 +165,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, @@ -231,21 +231,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, - disable_response_storage: bool, - tools_config: ToolsConfig, - /// Manager for external MCP servers/tools. mcp_connection_manager: McpConnectionManager, @@ -262,6 +249,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", ...). @@ -309,7 +321,7 @@ impl Session { config: Arc, auth: Option, tx_event: Sender, - ) -> anyhow::Result> { + ) -> anyhow::Result<(Arc, TurnContext)> { let ConfigureSession { provider, model, @@ -457,8 +469,7 @@ 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, @@ -467,19 +478,22 @@ impl Session { 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, + 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, }); @@ -487,13 +501,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; @@ -514,25 +528,7 @@ impl Session { } } - Ok(sess) - } - - pub(crate) fn get_approval_policy(&self) -> AskForApproval { - self.approval_policy - } - - pub(crate) fn get_sandbox_policy(&self) -> &SandboxPolicy { - &self.sandbox_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) { @@ -921,9 +917,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, @@ -933,17 +939,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, @@ -968,7 +977,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"); @@ -980,7 +996,8 @@ 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); } } @@ -1046,6 +1063,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(), @@ -1101,7 +1119,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; } @@ -1153,7 +1176,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(); @@ -1282,25 +1313,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)), @@ -1309,7 +1341,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 { @@ -1352,6 +1384,7 @@ struct ProcessedResponseItem { async fn try_run_turn( sess: &Session, + turn_context: &TurnContext, turn_diff_tracker: &mut TurnDiffTracker, sub_id: &str, prompt: &Prompt, @@ -1412,7 +1445,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 { @@ -1441,9 +1474,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 { @@ -1515,6 +1553,7 @@ async fn try_run_turn( async fn run_compact_task( sess: Arc, + turn_context: &TurnContext, sub_id: String, input: Vec, compact_instructions: String, @@ -1533,16 +1572,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, @@ -1596,6 +1635,7 @@ async fn run_compact_task( async fn handle_response_item( sess: &Session, + turn_context: &TurnContext, turn_diff_tracker: &mut TurnDiffTracker, sub_id: &str, item: ResponseItem, @@ -1659,6 +1699,7 @@ async fn handle_response_item( Some( handle_function_call( sess, + turn_context, turn_diff_tracker, sub_id.to_string(), name, @@ -1698,11 +1739,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, @@ -1721,6 +1763,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, @@ -1729,14 +1772,21 @@ 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) - .await + 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) { @@ -1788,12 +1838,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, } @@ -1801,12 +1851,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 { @@ -1829,8 +1879,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()); @@ -1844,6 +1898,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, @@ -1851,7 +1906,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) @@ -1913,8 +1968,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), ) }; @@ -1929,8 +1984,8 @@ async fn handle_container_exec_with_params( 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), ) @@ -2000,7 +2055,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, @@ -2008,7 +2063,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(), @@ -2041,6 +2096,7 @@ async fn handle_container_exec_with_params( error, sandbox_type, sess, + turn_context, ) .await } @@ -2061,6 +2117,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(); @@ -2068,7 +2125,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, @@ -2139,7 +2196,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(), @@ -2253,8 +2310,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 { From ed1d34542bf99844a07146cdd67b540a5d6c34c2 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Fri, 15 Aug 2025 09:20:24 -0700 Subject: [PATCH 6/7] feat: introduce Op:UserTurn --- codex-rs/core/src/codex.rs | 58 +++++++++++++++++++++++++++++++++++ codex-rs/core/src/protocol.rs | 29 ++++++++++++++++++ 2 files changed, 87 insertions(+) diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 41e37640fa..215a2e41a1 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -55,6 +55,7 @@ use crate::exec::process_exec_tool_call; use crate::exec_env::create_env; 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::models::ContentItem; use crate::models::FunctionCallOutputPayload; use crate::models::LocalShellAction; @@ -1001,6 +1002,63 @@ async fn submission_loop( 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); + } + } Op::ExecApproval { id, decision } => match decision { ReviewDecision::Abort => { sess.abort(); diff --git a/codex-rs/core/src/protocol.rs b/codex-rs/core/src/protocol.rs index 1d264d3ed1..d334c2eb86 100644 --- a/codex-rs/core/src/protocol.rs +++ b/codex-rs/core/src/protocol.rs @@ -17,6 +17,8 @@ use serde_bytes::ByteBuf; use strum_macros::Display; use uuid::Uuid; +use crate::config_types::ReasoningEffort as ReasoningEffortConfig; +use crate::config_types::ReasoningSummary as ReasoningSummaryConfig; use crate::message_history::HistoryEntry; use crate::parse_command::ParsedCommand; use crate::plan_tool::UpdatePlanArgs; @@ -46,6 +48,33 @@ pub enum Op { items: Vec, }, + /// Similar to [`Op::UserInput`], but contains additional context required + /// for a turn of a [`crate::codex_conversation::CodexConversation`]. + UserTurn { + /// User input items, see `InputItem` + items: Vec, + + /// `cwd` to use with the [`SandboxPolicy`] and potentially tool calls + /// such as `local_shell`. + cwd: PathBuf, + + /// Policy to use for command approval. + approval_policy: AskForApproval, + + /// Policy to use for tool calls such as `local_shell`. + sandbox_policy: SandboxPolicy, + + /// Must be a valid model slug for the [`crate::client::ModelClient`] + /// associated with this conversation. + model: String, + + /// Will only be honored if the model is configured to use reasoning. + effort: ReasoningEffortConfig, + + /// Will only be honored if the model is configured to use reasoning. + summary: ReasoningSummaryConfig, + }, + /// Approve a command execution ExecApproval { /// The id of the submission we are approving From fb653dc147c1f605de8bbe1b1d518b458ccd46f0 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Fri, 15 Aug 2025 09:20:24 -0700 Subject: [PATCH 7/7] feat: introduce ClientRequest::SendUserTurn --- .../mcp-server/src/codex_message_processor.rs | 57 ++++++ codex-rs/mcp-server/src/wire_format.rs | 26 +++ .../tests/codex_message_processor_flow.rs | 187 +++++++++++++++++- .../mcp-server/tests/common/mcp_process.rs | 10 + 4 files changed, 279 insertions(+), 1 deletion(-) diff --git a/codex-rs/mcp-server/src/codex_message_processor.rs b/codex-rs/mcp-server/src/codex_message_processor.rs index cdf24d9214..d930c03b71 100644 --- a/codex-rs/mcp-server/src/codex_message_processor.rs +++ b/codex-rs/mcp-server/src/codex_message_processor.rs @@ -42,6 +42,8 @@ use crate::wire_format::RemoveConversationListenerParams; use crate::wire_format::RemoveConversationSubscriptionResponse; use crate::wire_format::SendUserMessageParams; use crate::wire_format::SendUserMessageResponse; +use crate::wire_format::SendUserTurnParams; +use crate::wire_format::SendUserTurnResponse; use codex_core::protocol::InputItem as CoreInputItem; use codex_core::protocol::Op; @@ -78,6 +80,9 @@ impl CodexMessageProcessor { ClientRequest::SendUserMessage { request_id, params } => { self.send_user_message(request_id, params).await; } + ClientRequest::SendUserTurn { request_id, params } => { + self.send_user_turn(request_id, params).await; + } ClientRequest::InterruptConversation { request_id, params } => { self.interrupt_conversation(request_id, params).await; } @@ -169,6 +174,58 @@ impl CodexMessageProcessor { .await; } + async fn send_user_turn(&self, request_id: RequestId, params: SendUserTurnParams) { + let SendUserTurnParams { + conversation_id, + items, + cwd, + approval_policy, + sandbox_policy, + model, + effort, + summary, + } = params; + + let Ok(conversation) = self + .conversation_manager + .get_conversation(conversation_id.0) + .await + else { + let error = JSONRPCErrorError { + code: INVALID_REQUEST_ERROR_CODE, + message: format!("conversation not found: {conversation_id}"), + data: None, + }; + self.outgoing.send_error(request_id, error).await; + return; + }; + + let mapped_items: Vec = items + .into_iter() + .map(|item| match item { + WireInputItem::Text { text } => CoreInputItem::Text { text }, + WireInputItem::Image { image_url } => CoreInputItem::Image { image_url }, + WireInputItem::LocalImage { path } => CoreInputItem::LocalImage { path }, + }) + .collect(); + + let _ = conversation + .submit(Op::UserTurn { + items: mapped_items, + cwd, + approval_policy, + sandbox_policy, + model, + effort, + summary, + }) + .await; + + self.outgoing + .send_response(request_id, SendUserTurnResponse {}) + .await; + } + async fn interrupt_conversation( &mut self, request_id: RequestId, diff --git a/codex-rs/mcp-server/src/wire_format.rs b/codex-rs/mcp-server/src/wire_format.rs index e2ba729eb5..68d9aeb9eb 100644 --- a/codex-rs/mcp-server/src/wire_format.rs +++ b/codex-rs/mcp-server/src/wire_format.rs @@ -2,8 +2,12 @@ use std::collections::HashMap; use std::fmt::Display; use std::path::PathBuf; +use codex_core::config_types::ReasoningEffort; +use codex_core::config_types::ReasoningSummary; +use codex_core::protocol::AskForApproval; use codex_core::protocol::FileChange; use codex_core::protocol::ReviewDecision; +use codex_core::protocol::SandboxPolicy; use mcp_types::RequestId; use serde::Deserialize; use serde::Serialize; @@ -36,6 +40,11 @@ pub enum ClientRequest { request_id: RequestId, params: SendUserMessageParams, }, + SendUserTurn { + #[serde(rename = "id")] + request_id: RequestId, + params: SendUserTurnParams, + }, InterruptConversation { #[serde(rename = "id")] request_id: RequestId, @@ -120,6 +129,23 @@ pub struct SendUserMessageParams { pub items: Vec, } +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] +#[serde(rename_all = "camelCase")] +pub struct SendUserTurnParams { + pub conversation_id: ConversationId, + pub items: Vec, + pub cwd: PathBuf, + pub approval_policy: AskForApproval, + pub sandbox_policy: SandboxPolicy, + pub model: String, + pub effort: ReasoningEffort, + pub summary: ReasoningSummary, +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] +#[serde(rename_all = "camelCase")] +pub struct SendUserTurnResponse {} + #[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] #[serde(rename_all = "camelCase")] pub struct InterruptConversationParams { 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,