diff --git a/codex-rs/apply-patch/src/lib.rs b/codex-rs/apply-patch/src/lib.rs index 34fb4e95c..09d777ef6 100644 --- a/codex-rs/apply-patch/src/lib.rs +++ b/codex-rs/apply-patch/src/lib.rs @@ -18,8 +18,9 @@ use codex_utils_absolute_path::AbsolutePathBuf; pub use parser::Hunk; pub use parser::ParseError; use parser::ParseError::*; -use parser::UpdateFileChunk; +pub use parser::UpdateFileChunk; pub use parser::parse_patch; +pub use parser::parse_patch_streaming; use similar::TextDiff; use thiserror::Error; diff --git a/codex-rs/apply-patch/src/parser.rs b/codex-rs/apply-patch/src/parser.rs index 24ebcb146..64d3685d5 100644 --- a/codex-rs/apply-patch/src/parser.rs +++ b/codex-rs/apply-patch/src/parser.rs @@ -132,6 +132,14 @@ pub fn parse_patch(patch: &str) -> Result { parse_patch_text(patch, mode) } +/// Parses streamed patch text that may not have reached `*** End Patch` yet. +/// +/// This entry point is for progress reporting only; callers must not use its +/// output to apply a patch. +pub fn parse_patch_streaming(patch: &str) -> Result { + parse_patch_text(patch, ParseMode::Streaming) +} + enum ParseMode { /// Parse the patch text argument as is. Strict, @@ -169,32 +177,33 @@ enum ParseMode { /// `<<'EOF'` and ends with `EOF\n`. If so, we strip off these markers, /// trim() the result, and treat what is left as the patch text. Lenient, + + /// Parse partial patch text for progress reporting while the model is + /// still streaming tool input. This mode requires a begin marker but does + /// not require an end marker, and its output must not be used to apply a + /// patch. + Streaming, } fn parse_patch_text(patch: &str, mode: ParseMode) -> Result { let lines: Vec<&str> = patch.trim().lines().collect(); - let lines: &[&str] = match check_patch_boundaries_strict(&lines) { - Ok(()) => &lines, - Err(e) => match mode { - ParseMode::Strict => { - return Err(e); - } - ParseMode::Lenient => check_patch_boundaries_lenient(&lines, e)?, - }, + let (patch_lines, hunk_lines) = match mode { + ParseMode::Strict => check_patch_boundaries_strict(&lines)?, + ParseMode::Lenient => check_patch_boundaries_lenient(&lines)?, + ParseMode::Streaming => check_patch_boundaries_streaming(&lines)?, }; let mut hunks: Vec = Vec::new(); - // The above checks ensure that lines.len() >= 2. - let last_line_index = lines.len().saturating_sub(1); - let mut remaining_lines = &lines[1..last_line_index]; + let mut remaining_lines = hunk_lines; let mut line_number = 2; + let allow_incomplete = matches!(mode, ParseMode::Streaming); while !remaining_lines.is_empty() { - let (hunk, hunk_lines) = parse_one_hunk(remaining_lines, line_number)?; + let (hunk, hunk_lines) = parse_one_hunk(remaining_lines, line_number, allow_incomplete)?; hunks.push(hunk); line_number += hunk_lines; remaining_lines = &remaining_lines[hunk_lines..] } - let patch = lines.join("\n"); + let patch = patch_lines.join("\n"); Ok(ApplyPatchArgs { hunks, patch, @@ -202,15 +211,37 @@ fn parse_patch_text(patch: &str, mode: ParseMode) -> Result( + original_lines: &'a [&'a str], +) -> Result<(&'a [&'a str], &'a [&'a str]), ParseError> { + match original_lines { + [first, ..] if first.trim() == BEGIN_PATCH_MARKER => { + let body_lines = if original_lines + .last() + .is_some_and(|line| line.trim() == END_PATCH_MARKER) + { + &original_lines[1..original_lines.len() - 1] + } else { + &original_lines[1..] + }; + Ok((original_lines, body_lines)) + } + _ => check_patch_boundaries_strict(original_lines), + } +} + /// Checks the start and end lines of the patch text for `apply_patch`, /// returning an error if they do not match the expected markers. -fn check_patch_boundaries_strict(lines: &[&str]) -> Result<(), ParseError> { +fn check_patch_boundaries_strict<'a>( + lines: &'a [&'a str], +) -> Result<(&'a [&'a str], &'a [&'a str]), ParseError> { let (first_line, last_line) = match lines { [] => (None, None), [first] => (Some(first), Some(first)), [first, .., last] => (Some(first), Some(last)), }; - check_start_and_end_lines_strict(first_line, last_line) + check_start_and_end_lines_strict(first_line, last_line)?; + Ok((lines, &lines[1..lines.len() - 1])) } /// If we are in lenient mode, we check if the first line starts with `< Result<(), ParseError> { /// contents, excluding the heredoc markers. fn check_patch_boundaries_lenient<'a>( original_lines: &'a [&'a str], - original_parse_error: ParseError, -) -> Result<&'a [&'a str], ParseError> { +) -> Result<(&'a [&'a str], &'a [&'a str]), ParseError> { + let original_parse_error = match check_patch_boundaries_strict(original_lines) { + Ok(lines) => return Ok(lines), + Err(e) => e, + }; + match original_lines { [first, .., last] => { if (first == &"<( && original_lines.len() >= 4 { let inner_lines = &original_lines[1..original_lines.len() - 1]; - match check_patch_boundaries_strict(inner_lines) { - Ok(()) => Ok(inner_lines), - Err(e) => Err(e), - } + check_patch_boundaries_strict(inner_lines) } else { Err(original_parse_error) } @@ -265,7 +297,11 @@ fn check_start_and_end_lines_strict( /// Attempts to parse a single hunk from the start of lines. /// Returns the parsed hunk and the number of lines parsed (or a ParseError). -fn parse_one_hunk(lines: &[&str], line_number: usize) -> Result<(Hunk, usize), ParseError> { +fn parse_one_hunk( + lines: &[&str], + line_number: usize, + allow_incomplete: bool, +) -> Result<(Hunk, usize), ParseError> { // Be tolerant of case mismatches and extra padding around marker strings. let first_line = lines[0].trim(); if let Some(path) = first_line.strip_prefix(ADD_FILE_MARKER) { @@ -321,15 +357,26 @@ fn parse_one_hunk(lines: &[&str], line_number: usize) -> Result<(Hunk, usize), P continue; } - if remaining_lines[0].starts_with("***") { + if remaining_lines[0].starts_with('*') { break; } - let (chunk, chunk_lines) = parse_update_file_chunk( + if allow_incomplete && remaining_lines[0] == "@" { + break; + } + + let parsed_chunk = parse_update_file_chunk( remaining_lines, line_number + parsed_lines, chunks.is_empty(), - )?; + ); + let (chunk, chunk_lines) = match parsed_chunk { + Ok(parsed) => parsed, + Err(InvalidHunkError { .. }) if allow_incomplete && !chunks.is_empty() => { + break; + } + Err(err) => return Err(err), + }; chunks.push(chunk); parsed_lines += chunk_lines; remaining_lines = &remaining_lines[chunk_lines..] @@ -453,6 +500,166 @@ fn parse_update_file_chunk( Ok((chunk, parsed_lines + start_index)) } +#[test] +fn test_parse_patch_streaming() { + assert_eq!( + parse_patch_streaming("*** Begin Patch\n*** Add File: src/hello.txt\n+hello\n+wor"), + Ok(ApplyPatchArgs { + hunks: vec![AddFile { + path: PathBuf::from("src/hello.txt"), + contents: "hello\nwor\n".to_string(), + }], + patch: "*** Begin Patch\n*** Add File: src/hello.txt\n+hello\n+wor".to_string(), + workdir: None, + }) + ); + + assert_eq!( + parse_patch_streaming( + "*** Begin Patch\n*** Update File: src/old.rs\n*** Move to: src/new.rs\n@@\n-old\n+new", + ), + Ok(ApplyPatchArgs { + hunks: vec![UpdateFile { + path: PathBuf::from("src/old.rs"), + move_path: Some(PathBuf::from("src/new.rs")), + chunks: vec![UpdateFileChunk { + change_context: None, + old_lines: vec!["old".to_string()], + new_lines: vec!["new".to_string()], + is_end_of_file: false, + }], + }], + patch: "*** Begin Patch\n*** Update File: src/old.rs\n*** Move to: src/new.rs\n@@\n-old\n+new".to_string(), + workdir: None, + }) + ); + + assert!( + parse_patch_text( + "*** Begin Patch\n*** Delete File: gone.txt", + ParseMode::Streaming + ) + .is_ok() + ); + assert!( + parse_patch_text( + "*** Begin Patch\n*** Delete File: gone.txt", + ParseMode::Strict + ) + .is_err() + ); + + assert_eq!( + parse_patch_streaming( + "*** Begin Patch\n*** Add File: src/one.txt\n+one\n*** Delete File: src/two.txt\n", + ), + Ok(ApplyPatchArgs { + hunks: vec![ + AddFile { + path: PathBuf::from("src/one.txt"), + contents: "one\n".to_string(), + }, + DeleteFile { + path: PathBuf::from("src/two.txt"), + }, + ], + patch: "*** Begin Patch\n*** Add File: src/one.txt\n+one\n*** Delete File: src/two.txt" + .to_string(), + workdir: None, + }) + ); +} + +#[test] +fn test_parse_patch_streaming_large_patch_by_character() { + let patch = "\ +*** Begin Patch +*** Add File: docs/release-notes.md ++# Release notes ++ ++## CLI ++- Surface apply_patch progress while arguments stream. ++- Keep final patch application gated on the completed tool call. ++- Include file summaries in the progress event payload. +*** Update File: src/config.rs +@@ impl Config +- pub apply_patch_progress: bool, ++ pub stream_apply_patch_progress: bool, + pub include_diagnostics: bool, +@@ fn default_progress_interval() +- Duration::from_millis(500) ++ Duration::from_millis(250) +*** Delete File: src/legacy_patch_progress.rs +*** Update File: crates/cli/src/main.rs +*** Move to: crates/cli/src/bin/codex.rs +@@ fn run() +- let args = Args::parse(); +- dispatch(args) ++ let cli = Cli::parse(); ++ dispatch(cli) +*** Add File: tests/fixtures/apply_patch_progress.json ++{ ++ \"type\": \"apply_patch_progress\", ++ \"hunks\": [ ++ { \"operation\": \"add\", \"path\": \"docs/release-notes.md\" }, ++ { \"operation\": \"update\", \"path\": \"src/config.rs\" } ++ ] ++} +*** Update File: README.md +@@ Development workflow + Build the Rust workspace before opening a pull request. ++When touching streamed tool calls, include parser coverage for partial input. ++Prefer tests that exercise the exact event payload shape. +*** Delete File: docs/old-apply-patch-progress.md +*** End Patch"; + + let mut max_hunk_count = 0; + let mut saw_hunk_counts = Vec::new(); + for i in 1..=patch.len() { + let partial = &patch[..i]; + if let Ok(parsed) = parse_patch_streaming(partial) { + let hunk_count = parsed.hunks.len(); + assert!( + hunk_count >= max_hunk_count, + "hunk count should never decrease while streaming: {hunk_count} < {max_hunk_count} for {partial:?}", + ); + if hunk_count > max_hunk_count { + saw_hunk_counts.push(hunk_count); + max_hunk_count = hunk_count; + } + } + } + + assert_eq!(saw_hunk_counts, vec![1, 2, 3, 4, 5, 6, 7]); + let parsed = parse_patch_streaming(patch).unwrap(); + assert_eq!(parsed.hunks.len(), 7); + assert_eq!( + parsed + .hunks + .iter() + .map(|hunk| match hunk { + AddFile { .. } => "add", + DeleteFile { .. } => "delete", + UpdateFile { + move_path: Some(_), .. + } => "move-update", + UpdateFile { + move_path: None, .. + } => "update", + }) + .collect::>(), + vec![ + "add", + "update", + "delete", + "move-update", + "add", + "update", + "delete" + ] + ); +} + #[test] fn test_parse_patch() { assert_eq!( @@ -794,7 +1001,7 @@ fn test_parse_patch_lenient() { #[test] fn test_parse_one_hunk() { assert_eq!( - parse_one_hunk(&["bad"], /*line_number*/ 234), + parse_one_hunk(&["bad"], /*line_number*/ 234, /*allow_incomplete*/ false), Err(InvalidHunkError { message: "'bad' is not a valid hunk header. \ Valid hunk headers: '*** Add File: {path}', '*** Delete File: {path}', '*** Update File: {path}'".to_string(), diff --git a/codex-rs/cli/src/responses_cmd.rs b/codex-rs/cli/src/responses_cmd.rs index 044c7e0b2..2dc7a6714 100644 --- a/codex-rs/cli/src/responses_cmd.rs +++ b/codex-rs/cli/src/responses_cmd.rs @@ -104,6 +104,18 @@ fn response_event_to_json(event: codex_api::ResponseEvent) -> serde_json::Value codex_api::ResponseEvent::OutputTextDelta(delta) => { json!({ "type": "response.output_text.delta", "delta": delta }) } + codex_api::ResponseEvent::ToolCallInputDelta { + item_id, + call_id, + delta, + } => { + json!({ + "type": "response.tool_call_input.delta", + "item_id": item_id, + "call_id": call_id, + "delta": delta, + }) + } codex_api::ResponseEvent::ReasoningSummaryDelta { delta, summary_index, @@ -216,4 +228,22 @@ mod tests { }) ); } + + #[test] + fn tool_call_input_delta_uses_responses_event_name() { + let delta = response_event_to_json(codex_api::ResponseEvent::ToolCallInputDelta { + item_id: "item-1".to_string(), + call_id: Some("call-1".to_string()), + delta: "patch".to_string(), + }); + assert_eq!( + delta, + json!({ + "type": "response.tool_call_input.delta", + "item_id": "item-1", + "call_id": "call-1", + "delta": "patch", + }) + ); + } } diff --git a/codex-rs/codex-api/src/common.rs b/codex-rs/codex-api/src/common.rs index ef4dd5fb8..d08820649 100644 --- a/codex-rs/codex-api/src/common.rs +++ b/codex-rs/codex-api/src/common.rs @@ -80,6 +80,11 @@ pub enum ResponseEvent { token_usage: Option, }, OutputTextDelta(String), + ToolCallInputDelta { + item_id: String, + call_id: Option, + delta: String, + }, ReasoningSummaryDelta { delta: String, summary_index: i64, diff --git a/codex-rs/codex-api/src/sse/responses.rs b/codex-rs/codex-api/src/sse/responses.rs index bf6d75506..da7e87461 100644 --- a/codex-rs/codex-api/src/sse/responses.rs +++ b/codex-rs/codex-api/src/sse/responses.rs @@ -167,6 +167,8 @@ pub struct ResponsesStreamEvent { headers: Option, response: Option, item: Option, + item_id: Option, + call_id: Option, delta: Option, summary_index: Option, content_index: Option, @@ -250,6 +252,17 @@ pub fn process_responses_event( return Ok(Some(ResponseEvent::OutputTextDelta(delta))); } } + "response.custom_tool_call_input.delta" => { + if let (Some(delta), Some(item_id)) = + (event.delta, event.item_id.clone().or(event.call_id.clone())) + { + return Ok(Some(ResponseEvent::ToolCallInputDelta { + item_id, + call_id: event.call_id, + delta, + })); + } + } "response.reasoning_summary_text.delta" => { if let (Some(delta), Some(summary_index)) = (event.delta, event.summary_index) { return Ok(Some(ResponseEvent::ReasoningSummaryDelta { @@ -692,6 +705,38 @@ mod tests { ); } + #[tokio::test] + async fn parses_tool_call_input_deltas() { + let events = run_sse(vec![ + json!({ + "type": "response.custom_tool_call_input.delta", + "item_id": "ctc_1", + "call_id": "call_1", + "delta": "*** Begin", + }), + json!({ + "type": "response.function_call_arguments.delta", + "item_id": "fc_1", + "delta": "{\"input\":\"", + }), + json!({ + "type": "response.completed", + "response": { "id": "resp1" } + }), + ]) + .await; + + assert_matches!( + &events[0], + ResponseEvent::ToolCallInputDelta { + item_id, + call_id: Some(call_id), + delta, + } if item_id == "ctc_1" && call_id == "call_1" && delta == "*** Begin" + ); + assert_matches!(&events[1], ResponseEvent::Completed { .. }); + } + #[tokio::test] async fn emits_completed_without_stream_end() { let completed = json!({ diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 67d051a69..eb44f5c8f 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -335,6 +335,9 @@ "apply_patch_freeform": { "type": "boolean" }, + "apply_patch_streaming_events": { + "type": "boolean" + }, "apps": { "type": "boolean" }, @@ -2225,6 +2228,9 @@ "apply_patch_freeform": { "type": "boolean" }, + "apply_patch_streaming_events": { + "type": "boolean" + }, "apps": { "type": "boolean" }, diff --git a/codex-rs/core/src/codex/turn.rs b/codex-rs/core/src/codex/turn.rs index fa24bd692..84610804a 100644 --- a/codex-rs/core/src/codex/turn.rs +++ b/codex-rs/core/src/codex/turn.rs @@ -50,6 +50,7 @@ use crate::stream_events_utils::record_completed_response_item; use crate::tools::ToolRouter; use crate::tools::context::SharedTurnDiffTracker; use crate::tools::parallel::ToolCallRuntime; +use crate::tools::registry::ToolArgumentDiffConsumer; use crate::tools::router::ToolRouterParams; use crate::turn_diff_tracker::TurnDiffTracker; use crate::turn_timing::record_turn_ttft_metric; @@ -91,6 +92,7 @@ use codex_protocol::protocol::ReasoningRawContentDeltaEvent; use codex_protocol::protocol::TurnDiffEvent; use codex_protocol::protocol::WarningEvent; use codex_protocol::user_input::UserInput; +use codex_tools::ToolName; use codex_tools::filter_tool_suggest_discoverable_tools_for_client; use codex_utils_stream_parser::AssistantTextChunk; use codex_utils_stream_parser::AssistantTextStreamParser; @@ -1478,6 +1480,7 @@ pub(super) fn realtime_text_for_event(msg: &EventMsg) -> Option { | EventMsg::TerminalInteraction(_) | EventMsg::ExecCommandEnd(_) | EventMsg::PatchApplyBegin(_) + | EventMsg::PatchApplyUpdated(_) | EventMsg::PatchApplyEnd(_) | EventMsg::ViewImageToolCall(_) | EventMsg::ImageGenerationBegin(_) @@ -1858,6 +1861,10 @@ async fn try_run_sampling_request( let mut needs_follow_up = false; let mut last_agent_message: Option = None; let mut active_item: Option = None; + let mut active_tool_argument_diff_consumer: Option<( + String, + Box, + )> = None; let mut should_emit_turn_diff = false; let plan_mode = turn_context.collaboration_mode.mode == ModeKind::Plan; let mut assistant_message_stream_parsers = AssistantMessageStreamParsers::new(plan_mode); @@ -1901,6 +1908,7 @@ async fn try_run_sampling_request( match event { ResponseEvent::Created => {} ResponseEvent::OutputItemDone(item) => { + active_tool_argument_diff_consumer = None; let previously_active_item = active_item.take(); if let Some(previous) = previously_active_item.as_ref() && matches!(previous, TurnItem::AgentMessage(_)) @@ -1953,6 +1961,14 @@ async fn try_run_sampling_request( needs_follow_up |= output_result.needs_follow_up; } ResponseEvent::OutputItemAdded(item) => { + if let ResponseItem::CustomToolCall { call_id, name, .. } = &item { + let tool_name = ToolName::plain(name.as_str()); + active_tool_argument_diff_consumer = tool_runtime + .create_diff_consumer(&tool_name) + .map(|consumer| (call_id.clone(), consumer)); + } else if matches!(&item, ResponseItem::FunctionCall { .. }) { + active_tool_argument_diff_consumer = None; + } if let Some(turn_item) = handle_non_tool_response_item( sess.as_ref(), turn_context.as_ref(), @@ -2080,6 +2096,24 @@ async fn try_run_sampling_request( error_or_panic("OutputTextDelta without active item".to_string()); } } + ResponseEvent::ToolCallInputDelta { + item_id: _, + call_id, + delta, + } => { + let Some((active_call_id, consumer)) = active_tool_argument_diff_consumer.as_mut() + else { + continue; + }; + let call_id = match call_id { + Some(call_id) if call_id.as_str() != active_call_id.as_str() => continue, + Some(call_id) => call_id, + None => active_call_id.clone(), + }; + if let Some(event) = consumer.consume_diff(turn_context.as_ref(), call_id, &delta) { + sess.send_event(&turn_context, event).await; + } + } ResponseEvent::ReasoningSummaryDelta { delta, summary_index, diff --git a/codex-rs/core/src/tools/handlers/apply_patch.rs b/codex-rs/core/src/tools/handlers/apply_patch.rs index 8e932c3ad..0452b4578 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch.rs @@ -1,4 +1,8 @@ +use std::collections::BTreeSet; +use std::collections::HashMap; use std::path::Path; +use std::path::PathBuf; +use std::sync::Arc; use crate::apply_patch; use crate::apply_patch::InternalApplyPatchInvocation; @@ -16,26 +20,136 @@ use crate::tools::events::ToolEventCtx; use crate::tools::handlers::apply_granted_turn_permissions; use crate::tools::handlers::parse_arguments; use crate::tools::orchestrator::ToolOrchestrator; +use crate::tools::registry::ToolArgumentDiffConsumer; use crate::tools::registry::ToolHandler; use crate::tools::registry::ToolKind; use crate::tools::runtimes::apply_patch::ApplyPatchRequest; use crate::tools::runtimes::apply_patch::ApplyPatchRuntime; use crate::tools::sandboxing::ToolCtx; use codex_apply_patch::ApplyPatchAction; +use codex_apply_patch::ApplyPatchArgs; use codex_apply_patch::ApplyPatchFileChange; +use codex_apply_patch::Hunk; +use codex_apply_patch::parse_patch_streaming; use codex_exec_server::ExecutorFileSystem; +use codex_features::Feature; use codex_protocol::models::FileSystemPermissions; use codex_protocol::models::PermissionProfile; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::FileChange; +use codex_protocol::protocol::PatchApplyUpdatedEvent; use codex_sandboxing::policy_transforms::effective_file_system_sandbox_policy; use codex_sandboxing::policy_transforms::merge_permission_profiles; use codex_sandboxing::policy_transforms::normalize_additional_permissions; use codex_tools::ApplyPatchToolArgs; use codex_utils_absolute_path::AbsolutePathBuf; -use std::collections::BTreeSet; -use std::sync::Arc; pub struct ApplyPatchHandler; +#[derive(Default)] +struct ApplyPatchArgumentDiffConsumer { + input: String, + last_progress: Option>, +} + +impl ToolArgumentDiffConsumer for ApplyPatchArgumentDiffConsumer { + fn consume_diff( + &mut self, + turn: &TurnContext, + call_id: String, + diff: &str, + ) -> Option { + if !turn.features.enabled(Feature::ApplyPatchStreamingEvents) { + return None; + } + + self.push_delta(call_id, diff) + .map(EventMsg::PatchApplyUpdated) + } +} + +impl ApplyPatchArgumentDiffConsumer { + fn push_delta(&mut self, call_id: String, delta: &str) -> Option { + self.input.push_str(delta); + + let ApplyPatchArgs { hunks, .. } = parse_patch_streaming(&self.input).ok()?; + if hunks.is_empty() { + return None; + } + if self.last_progress.as_ref() == Some(&hunks) { + return None; + } + + let changes = convert_apply_patch_hunks_to_protocol(&hunks); + self.last_progress = Some(hunks); + Some(PatchApplyUpdatedEvent { call_id, changes }) + } +} + +fn convert_apply_patch_hunks_to_protocol(hunks: &[Hunk]) -> HashMap { + hunks + .iter() + .map(|hunk| { + let path = hunk_source_path(hunk).to_path_buf(); + let change = match hunk { + Hunk::AddFile { contents, .. } => FileChange::Add { + content: contents.clone(), + }, + Hunk::DeleteFile { .. } => FileChange::Delete { + content: String::new(), + }, + Hunk::UpdateFile { + chunks, move_path, .. + } => FileChange::Update { + unified_diff: format_update_chunks_for_progress(chunks), + move_path: move_path.clone(), + }, + }; + (path, change) + }) + .collect() +} + +fn hunk_source_path(hunk: &Hunk) -> &Path { + match hunk { + Hunk::AddFile { path, .. } | Hunk::DeleteFile { path } | Hunk::UpdateFile { path, .. } => { + path + } + } +} + +fn format_update_chunks_for_progress(chunks: &[codex_apply_patch::UpdateFileChunk]) -> String { + let mut unified_diff = String::new(); + for chunk in chunks { + match &chunk.change_context { + Some(context) => { + unified_diff.push_str("@@ "); + unified_diff.push_str(context); + unified_diff.push('\n'); + } + None => { + unified_diff.push_str("@@"); + unified_diff.push('\n'); + } + } + for line in &chunk.old_lines { + unified_diff.push('-'); + unified_diff.push_str(line); + unified_diff.push('\n'); + } + for line in &chunk.new_lines { + unified_diff.push('+'); + unified_diff.push_str(line); + unified_diff.push('\n'); + } + if chunk.is_end_of_file { + unified_diff.push_str("*** End of File"); + unified_diff.push('\n'); + } + } + unified_diff +} + fn file_paths_for_action(action: &ApplyPatchAction) -> Vec { let mut keys = Vec::new(); let cwd = &action.cwd; @@ -142,6 +256,10 @@ impl ToolHandler for ApplyPatchHandler { true } + fn create_diff_consumer(&self) -> Option> { + Some(Box::::default()) + } + async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { session, diff --git a/codex-rs/core/src/tools/handlers/apply_patch_tests.rs b/codex-rs/core/src/tools/handlers/apply_patch_tests.rs index 86e05fb8a..a2c9c23af 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch_tests.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch_tests.rs @@ -2,12 +2,75 @@ use super::*; use codex_apply_patch::MaybeApplyPatchVerified; use codex_exec_server::LOCAL_FS; use codex_protocol::permissions::FileSystemSandboxPolicy; +use codex_protocol::protocol::FileChange; use codex_protocol::protocol::SandboxPolicy; use core_test_support::PathBufExt; use core_test_support::PathExt; use pretty_assertions::assert_eq; +use std::collections::HashMap; +use std::path::PathBuf; use tempfile::TempDir; +#[test] +fn diff_consumer_does_not_stream_json_tool_call_arguments() { + let mut consumer = ApplyPatchArgumentDiffConsumer::default(); + assert!( + consumer + .push_delta("call-1".to_string(), r#"{"input":"*** Begin Patch\n"#) + .is_none() + ); + assert!( + consumer + .push_delta( + "call-1".to_string(), + r#"*** Add File: hello.txt\n+hello\n*** End Patch\n"}"# + ) + .is_none() + ); +} + +#[test] +fn diff_consumer_streams_apply_patch_changes() { + let mut consumer = ApplyPatchArgumentDiffConsumer::default(); + assert!( + consumer + .push_delta("call-1".to_string(), "*** Begin Patch\n") + .is_none() + ); + + let event = consumer + .push_delta("call-1".to_string(), "*** Add File: hello.txt\n+hello") + .expect("progress event"); + assert_eq!( + (event.call_id, event.changes), + ( + "call-1".to_string(), + HashMap::from([( + PathBuf::from("hello.txt"), + FileChange::Add { + content: "hello\n".to_string(), + }, + )]), + ) + ); + + let event = consumer + .push_delta("call-1".to_string(), "\n+world") + .expect("progress event"); + assert_eq!( + (event.call_id, event.changes), + ( + "call-1".to_string(), + HashMap::from([( + PathBuf::from("hello.txt"), + FileChange::Add { + content: "hello\nworld\n".to_string(), + }, + )]), + ) + ); +} + #[tokio::test] async fn approval_keys_include_move_destination() { let tmp = TempDir::new().expect("tmp"); diff --git a/codex-rs/core/src/tools/parallel.rs b/codex-rs/core/src/tools/parallel.rs index b4dbfaa6c..2789c9c2f 100644 --- a/codex-rs/core/src/tools/parallel.rs +++ b/codex-rs/core/src/tools/parallel.rs @@ -16,6 +16,7 @@ use crate::tools::context::AbortedToolOutput; use crate::tools::context::SharedTurnDiffTracker; use crate::tools::context::ToolPayload; use crate::tools::registry::AnyToolResult; +use crate::tools::registry::ToolArgumentDiffConsumer; use crate::tools::router::ToolCall; use crate::tools::router::ToolCallSource; use crate::tools::router::ToolRouter; @@ -52,6 +53,13 @@ impl ToolCallRuntime { self.router.find_spec(tool_name) } + pub(crate) fn create_diff_consumer( + &self, + tool_name: &codex_tools::ToolName, + ) -> Option> { + self.router.create_diff_consumer(tool_name) + } + #[instrument(level = "trace", skip_all)] pub(crate) fn handle_tool_call( self, diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index 7e3713dcb..9c0291940 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use std::time::Duration; use std::time::Instant; +use crate::codex::TurnContext; use crate::function_tool::FunctionCallError; use crate::hook_runtime::record_additional_contexts; use crate::hook_runtime::run_post_tool_use_hooks; @@ -21,6 +22,7 @@ use codex_hooks::HookToolInput; use codex_hooks::HookToolInputLocalShell; use codex_hooks::HookToolKind; use codex_protocol::models::ResponseInputItem; +use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::SandboxPolicy; use codex_tools::ConfiguredToolSpec; use codex_tools::ToolName; @@ -74,6 +76,11 @@ pub trait ToolHandler: Send + Sync { None } + /// Creates an optional consumer for streamed tool argument diffs. + fn create_diff_consumer(&self) -> Option> { + None + } + /// Perform the actual [ToolInvocation] and returns a [ToolOutput] containing /// the final output to return to the model. fn handle( @@ -82,6 +89,14 @@ pub trait ToolHandler: Send + Sync { ) -> impl std::future::Future> + Send; } +/// Consumes streamed argument diffs for a tool call and emits protocol events +/// derived from partial tool input. +pub(crate) trait ToolArgumentDiffConsumer: Send { + /// Consume the next argument diff for a tool call. + fn consume_diff(&mut self, turn: &TurnContext, call_id: String, diff: &str) + -> Option; +} + pub(crate) struct AnyToolResult { pub(crate) call_id: String, pub(crate) payload: ToolPayload, @@ -132,6 +147,8 @@ trait AnyToolHandler: Send + Sync { result: &dyn ToolOutput, ) -> Option; + fn create_diff_consumer(&self) -> Option>; + fn handle_any<'a>( &'a self, invocation: ToolInvocation, @@ -163,6 +180,10 @@ where ToolHandler::post_tool_use_payload(self, call_id, payload, result) } + fn create_diff_consumer(&self) -> Option> { + ToolHandler::create_diff_consumer(self) + } + fn handle_any<'a>( &'a self, invocation: ToolInvocation, @@ -198,6 +219,13 @@ impl ToolRegistry { self.handler(name).is_some() } + pub(crate) fn create_diff_consumer( + &self, + name: &ToolName, + ) -> Option> { + self.handler(name)?.create_diff_consumer() + } + // TODO(jif) for dynamic tools. // pub fn register(&mut self, name: impl Into, handler: Arc) { // let name = name.into(); diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index ef7ae32d4..d57ca5473 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -6,6 +6,7 @@ use crate::tools::context::SharedTurnDiffTracker; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; use crate::tools::registry::AnyToolResult; +use crate::tools::registry::ToolArgumentDiffConsumer; use crate::tools::registry::ToolRegistry; use crate::tools::spec::build_specs_with_discoverable_tools; use codex_mcp::ToolInfo; @@ -131,6 +132,13 @@ impl ToolRouter { }) } + pub(crate) fn create_diff_consumer( + &self, + tool_name: &ToolName, + ) -> Option> { + self.registry.create_diff_consumer(tool_name) + } + fn configured_tool_supports_parallel(&self, tool_name: &ToolName) -> bool { if tool_name.namespace.is_some() { return false; diff --git a/codex-rs/core/src/turn_timing.rs b/codex-rs/core/src/turn_timing.rs index 6c47d3b52..a3653560c 100644 --- a/codex-rs/core/src/turn_timing.rs +++ b/codex-rs/core/src/turn_timing.rs @@ -134,6 +134,7 @@ fn response_event_records_turn_ttft(event: &ResponseEvent) -> bool { ResponseEvent::Created | ResponseEvent::ServerModel(_) | ResponseEvent::ServerReasoningIncluded(_) + | ResponseEvent::ToolCallInputDelta { .. } | ResponseEvent::Completed { .. } | ResponseEvent::ReasoningSummaryPartAdded { .. } | ResponseEvent::RateLimits(_) diff --git a/codex-rs/core/tests/suite/apply_patch_cli.rs b/codex-rs/core/tests/suite/apply_patch_cli.rs index cb264a0b2..df8654696 100644 --- a/codex-rs/core/tests/suite/apply_patch_cli.rs +++ b/codex-rs/core/tests/suite/apply_patch_cli.rs @@ -935,6 +935,128 @@ async fn apply_patch_cli_can_use_shell_command_output_as_patch_input() -> Result Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn apply_patch_custom_tool_streaming_emits_updated_changes() -> Result<()> { + skip_if_no_network!(Ok(())); + + let harness = apply_patch_harness_with(|builder| { + builder.with_config(|config| { + config + .features + .enable(Feature::ApplyPatchStreamingEvents) + .expect("enable apply_patch streaming events"); + }) + }) + .await?; + let test = harness.test(); + let codex = test.codex.clone(); + let call_id = "apply-patch-streaming"; + let patch = "*** Begin Patch\n*** Add File: streamed.txt\n+hello\n+world\n*** End Patch"; + mount_sse_sequence( + harness.server(), + vec![ + sse(vec![ + ev_response_created("resp-1"), + json!({ + "type": "response.output_item.added", + "item": { + "type": "custom_tool_call", + "call_id": call_id, + "name": "apply_patch", + "input": "", + } + }), + json!({ + "type": "response.custom_tool_call_input.delta", + "call_id": call_id, + "delta": "*** Begin Patch\n", + }), + json!({ + "type": "response.custom_tool_call_input.delta", + "call_id": call_id, + "delta": "*** Add File: streamed.txt\n+hello", + }), + json!({ + "type": "response.custom_tool_call_input.delta", + "call_id": call_id, + "delta": "\n+world\n*** End Patch", + }), + ev_apply_patch_custom_tool_call(call_id, patch), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_assistant_message("msg-1", "done"), + ev_completed("resp-2"), + ]), + ], + ) + .await; + + codex + .submit(Op::UserTurn { + items: vec![UserInput::Text { + text: "create streamed file".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + cwd: harness.cwd().to_path_buf(), + approval_policy: AskForApproval::Never, + approvals_reviewer: None, + sandbox_policy: SandboxPolicy::DangerFullAccess, + model: test.session_configured.model.clone(), + effort: None, + summary: None, + service_tier: None, + collaboration_mode: None, + personality: None, + }) + .await?; + + let mut updates = Vec::new(); + wait_for_event(&codex, |event| match event { + EventMsg::PatchApplyUpdated(update) => { + updates.push(update.clone()); + false + } + EventMsg::TurnComplete(_) => true, + _ => false, + }) + .await; + + assert_eq!( + updates + .iter() + .map(|update| update.call_id.as_str()) + .collect::>(), + vec![call_id, call_id] + ); + assert_eq!( + updates + .first() + .expect("first update") + .changes + .get(&std::path::PathBuf::from("streamed.txt")), + Some(&codex_protocol::protocol::FileChange::Add { + content: "hello\n".to_string(), + }) + ); + assert_eq!( + updates + .last() + .expect("last update") + .changes + .get(&std::path::PathBuf::from("streamed.txt")), + Some(&codex_protocol::protocol::FileChange::Add { + content: "hello\nworld\n".to_string(), + }) + ); + assert_eq!( + harness.read_file_text("streamed.txt").await?, + "hello\nworld\n" + ); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn apply_patch_shell_command_heredoc_with_cd_emits_turn_diff() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index 586c36685..888f0a409 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -91,6 +91,8 @@ pub enum Feature { ShellZshFork, /// Include the freeform apply_patch tool. ApplyPatchFreeform, + /// Stream structured progress while apply_patch input is being generated. + ApplyPatchStreamingEvents, /// Allow exec tools to request additional permissions while staying sandboxed. ExecPermissionApprovals, /// Enable Claude-style lifecycle hooks loaded from hooks.json files. @@ -720,6 +722,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::UnderDevelopment, default_enabled: false, }, + FeatureSpec { + id: Feature::ApplyPatchStreamingEvents, + key: "apply_patch_streaming_events", + stage: Stage::UnderDevelopment, + default_enabled: false, + }, FeatureSpec { id: Feature::ExecPermissionApprovals, key: "exec_permission_approvals", diff --git a/codex-rs/mcp-server/src/codex_tool_runner.rs b/codex-rs/mcp-server/src/codex_tool_runner.rs index eee234720..9c9680017 100644 --- a/codex-rs/mcp-server/src/codex_tool_runner.rs +++ b/codex-rs/mcp-server/src/codex_tool_runner.rs @@ -347,6 +347,7 @@ async fn run_codex_tool_session_inner( | EventMsg::BackgroundEvent(_) | EventMsg::StreamError(_) | EventMsg::PatchApplyBegin(_) + | EventMsg::PatchApplyUpdated(_) | EventMsg::PatchApplyEnd(_) | EventMsg::TurnDiff(_) | EventMsg::WebSearchBegin(_) diff --git a/codex-rs/otel/src/events/session_telemetry.rs b/codex-rs/otel/src/events/session_telemetry.rs index b1bf9b094..14ae1cb3e 100644 --- a/codex-rs/otel/src/events/session_telemetry.rs +++ b/codex-rs/otel/src/events/session_telemetry.rs @@ -1051,6 +1051,7 @@ impl SessionTelemetry { } ResponseEvent::Completed { .. } => "completed".into(), ResponseEvent::OutputTextDelta(_) => "text_delta".into(), + ResponseEvent::ToolCallInputDelta { .. } => "tool_input_delta".into(), ResponseEvent::ReasoningSummaryDelta { .. } => "reasoning_summary_delta".into(), ResponseEvent::ReasoningContentDelta { .. } => "reasoning_content_delta".into(), ResponseEvent::ReasoningSummaryPartAdded { .. } => { diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 05eefc29b..b777b4d32 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -1546,6 +1546,9 @@ pub enum EventMsg { /// `ExecCommandBegin` so front‑ends can show progress indicators. PatchApplyBegin(PatchApplyBeginEvent), + /// Latest model-generated structured changes for an `apply_patch` call. + PatchApplyUpdated(PatchApplyUpdatedEvent), + /// Notification that a patch application has finished. PatchApplyEnd(PatchApplyEndEvent), @@ -3240,6 +3243,14 @@ pub struct PatchApplyBeginEvent { pub changes: HashMap, } +#[derive(Debug, Clone, Deserialize, Serialize, JsonSchema, TS)] +pub struct PatchApplyUpdatedEvent { + /// Identifier for the originating `apply_patch` tool call. + pub call_id: String, + /// Structured file changes parsed from the model-generated patch input so far. + pub changes: HashMap, +} + #[derive(Debug, Clone, Deserialize, Serialize, JsonSchema, TS)] pub struct PatchApplyEndEvent { /// Identifier for the PatchApplyBegin that finished. diff --git a/codex-rs/rollout/src/policy.rs b/codex-rs/rollout/src/policy.rs index 4b50781e7..df41d2c76 100644 --- a/codex-rs/rollout/src/policy.rs +++ b/codex-rs/rollout/src/policy.rs @@ -156,6 +156,7 @@ fn event_msg_persistence_mode(ev: &EventMsg) -> Option { | EventMsg::BackgroundEvent(_) | EventMsg::StreamError(_) | EventMsg::PatchApplyBegin(_) + | EventMsg::PatchApplyUpdated(_) | EventMsg::TurnDiff(_) | EventMsg::GetHistoryEntryResponse(_) | EventMsg::UndoStarted(_) diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 7cfcd56e1..e0812e488 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -6649,6 +6649,7 @@ impl ChatWidget { | EventMsg::PlanDelta(_) | EventMsg::AgentReasoningDelta(_) | EventMsg::TerminalInteraction(_) + | EventMsg::PatchApplyUpdated(_) | EventMsg::ExecCommandOutputDelta(_) => {} _ => { tracing::trace!("handle_codex_event: {:?}", msg); @@ -6864,6 +6865,7 @@ impl ChatWidget { EventMsg::RawResponseItem(_) | EventMsg::ItemStarted(_) | EventMsg::AgentMessageContentDelta(_) + | EventMsg::PatchApplyUpdated(_) | EventMsg::ReasoningContentDelta(_) | EventMsg::ReasoningRawContentDelta(_) | EventMsg::DynamicToolCallRequest(_)