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.

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<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,

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<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,

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");

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<Box<dyn ToolArgumentDiffConsumer>> {
+        self.router.create_diff_consumer(tool_name)
+    }
+
     #[instrument(level = "trace", skip_all)]
     pub(crate) fn handle_tool_call(
         self,

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<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();

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<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;

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(_)