Stream apply_patch changes (#17862)

Adds new events for streaming apply_patch changes from responses api.
This is to enable clients to show progress during file writes.

Caveat: This does not work with apply_patch in function call mode, since
that required adding streaming json parsing.
This commit is contained in:
Akshay Nathan
2026-04-16 18:12:19 -07:00
committed by GitHub
Unverified
parent 9effa0509f
commit 7995c66032
20 changed files with 729 additions and 29 deletions
+2 -1
View File
@@ -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;
+233 -26
View File
@@ -132,6 +132,14 @@ pub fn parse_patch(patch: &str) -> Result<ApplyPatchArgs, ParseError> {
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<ApplyPatchArgs, ParseError> {
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<ApplyPatchArgs, ParseError> {
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<Hunk> = 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<ApplyPatchArgs, Pars
})
}
fn check_patch_boundaries_streaming<'a>(
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 `<<EOF`
@@ -222,8 +253,12 @@ fn check_patch_boundaries_strict(lines: &[&str]) -> 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 == &"<<EOF" || first == &"<<'EOF'" || first == &"<<\"EOF\"")
@@ -231,10 +266,7 @@ fn check_patch_boundaries_lenient<'a>(
&& 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<_>>(),
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(),
+30
View File
@@ -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",
})
);
}
}
+5
View File
@@ -80,6 +80,11 @@ pub enum ResponseEvent {
token_usage: Option<TokenUsage>,
},
OutputTextDelta(String),
ToolCallInputDelta {
item_id: String,
call_id: Option<String>,
delta: String,
},
ReasoningSummaryDelta {
delta: String,
summary_index: i64,
+45
View File
@@ -167,6 +167,8 @@ pub struct ResponsesStreamEvent {
headers: Option<Value>,
response: Option<Value>,
item: Option<Value>,
item_id: Option<String>,
call_id: Option<String>,
delta: Option<String>,
summary_index: Option<i64>,
content_index: Option<i64>,
@@ -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!({
+6
View File
@@ -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"
},
+34
View File
@@ -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<String> {
| 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<String> = None;
let mut active_item: Option<TurnItem> = None;
let mut active_tool_argument_diff_consumer: Option<(
String,
Box<dyn ToolArgumentDiffConsumer>,
)> = 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,
+120 -2
View File
@@ -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<Vec<Hunk>>,
}
impl ToolArgumentDiffConsumer for ApplyPatchArgumentDiffConsumer {
fn consume_diff(
&mut self,
turn: &TurnContext,
call_id: String,
diff: &str,
) -> Option<EventMsg> {
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<PatchApplyUpdatedEvent> {
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<PathBuf, FileChange> {
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<AbsolutePathBuf> {
let mut keys = Vec::new();
let cwd = &action.cwd;
@@ -142,6 +256,10 @@ impl ToolHandler for ApplyPatchHandler {
true
}
fn create_diff_consumer(&self) -> Option<Box<dyn ToolArgumentDiffConsumer>> {
Some(Box::<ApplyPatchArgumentDiffConsumer>::default())
}
async fn handle(&self, invocation: ToolInvocation) -> Result<Self::Output, FunctionCallError> {
let ToolInvocation {
session,
@@ -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");
+8
View File
@@ -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<Box<dyn ToolArgumentDiffConsumer>> {
self.router.create_diff_consumer(tool_name)
}
#[instrument(level = "trace", skip_all)]
pub(crate) fn handle_tool_call(
self,
+28
View File
@@ -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<Box<dyn ToolArgumentDiffConsumer>> {
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<Output = Result<Self::Output, FunctionCallError>> + 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<EventMsg>;
}
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<PostToolUsePayload>;
fn create_diff_consumer(&self) -> Option<Box<dyn ToolArgumentDiffConsumer>>;
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<Box<dyn ToolArgumentDiffConsumer>> {
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<Box<dyn ToolArgumentDiffConsumer>> {
self.handler(name)?.create_diff_consumer()
}
// TODO(jif) for dynamic tools.
// pub fn register(&mut self, name: impl Into<String>, handler: Arc<dyn ToolHandler>) {
// let name = name.into();
+8
View File
@@ -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<Box<dyn ToolArgumentDiffConsumer>> {
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;
+1
View File
@@ -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(_)
@@ -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<_>>(),
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(()));
+8
View File
@@ -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",
@@ -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(_)
@@ -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 { .. } => {
+11
View File
@@ -1546,6 +1546,9 @@ pub enum EventMsg {
/// `ExecCommandBegin` so frontends 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<PathBuf, FileChange>,
}
#[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<PathBuf, FileChange>,
}
#[derive(Debug, Clone, Deserialize, Serialize, JsonSchema, TS)]
pub struct PatchApplyEndEvent {
/// Identifier for the PatchApplyBegin that finished.
+1
View File
@@ -156,6 +156,7 @@ fn event_msg_persistence_mode(ev: &EventMsg) -> Option<EventPersistenceMode> {
| EventMsg::BackgroundEvent(_)
| EventMsg::StreamError(_)
| EventMsg::PatchApplyBegin(_)
| EventMsg::PatchApplyUpdated(_)
| EventMsg::TurnDiff(_)
| EventMsg::GetHistoryEntryResponse(_)
| EventMsg::UndoStarted(_)
+2
View File
@@ -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(_)