Stateful streaming apply_patch parser

codex-rs/core/src/session/turn.rs

@@ -1902,7 +1902,7 @@ async fn try_run_sampling_request(
             ResponseEvent::Created => {}
             ResponseEvent::OutputItemDone(item) => {
                 if let Some((_, mut consumer)) = active_tool_argument_diff_consumer.take()
-                    && let Some(event) = consumer.flush_on_complete()
+                    && let Ok(Some(event)) = consumer.finish()
                 {
                     sess.send_event(&turn_context, event).await;
                 }

codex-rs/core/src/tools/handlers/apply_patch.rs

@@ -33,10 +33,9 @@ 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_apply_patch::StreamingPatchParser;
 use codex_exec_server::ExecutorFileSystem;
 use codex_features::Feature;
 use codex_protocol::models::AdditionalPermissionProfile;
@@ -56,8 +55,7 @@ pub struct ApplyPatchHandler;
 
 #[derive(Default)]
 struct ApplyPatchArgumentDiffConsumer {
-    input: String,
-    last_progress: Option<Vec<Hunk>>,
+    parser: StreamingPatchParser,
     last_sent_at: Option<Instant>,
     pending: Option<PatchApplyUpdatedEvent>,
 }
@@ -77,26 +75,19 @@ impl ToolArgumentDiffConsumer for ApplyPatchArgumentDiffConsumer {
             .map(EventMsg::PatchApplyUpdated)
     }
 
-    fn flush_on_complete(&mut self) -> Option<EventMsg> {
-        self.flush_update_on_complete()
-            .map(EventMsg::PatchApplyUpdated)
+    fn finish(&mut self) -> Result<Option<EventMsg>, FunctionCallError> {
+        self.finish_update_on_complete()
+            .map(|event| event.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()?;
+        let hunks = self.parser.push_delta(delta).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);
         let event = PatchApplyUpdatedEvent { call_id, changes };
         let now = Instant::now();
         match self.last_sent_at {
@@ -114,12 +105,18 @@ impl ApplyPatchArgumentDiffConsumer {
         }
     }
 
-    fn flush_update_on_complete(&mut self) -> Option<PatchApplyUpdatedEvent> {
+    fn finish_update_on_complete(
+        &mut self,
+    ) -> Result<Option<PatchApplyUpdatedEvent>, FunctionCallError> {
+        self.parser.finish().map_err(|err| {
+            FunctionCallError::RespondToModel(format!("failed to parse apply_patch: {err}"))
+        })?;
+
         let event = self.pending.take();
         if event.is_some() {
             self.last_sent_at = Some(Instant::now());
         }
-        event
+        Ok(event)
     }
 }
 

codex-rs/core/src/tools/handlers/apply_patch_tests.rs

@@ -136,7 +136,7 @@ fn diff_consumer_streams_apply_patch_changes() {
             HashMap::from([(
                 PathBuf::from("hello.txt"),
                 FileChange::Add {
-                    content: "hello\n".to_string(),
+                    content: String::new(),
                 },
             )]),
         )
@@ -147,8 +147,16 @@ fn diff_consumer_streams_apply_patch_changes() {
             .push_delta("call-1".to_string(), "\n+world")
             .is_none()
     );
+    assert!(
+        consumer
+            .push_delta("call-1".to_string(), "\n*** End Patch")
+            .is_none()
+    );
 
-    let event = consumer.flush_update_on_complete().expect("progress event");
+    let event = consumer
+        .finish_update_on_complete()
+        .expect("finish parser")
+        .expect("progress event");
     assert_eq!(
         (event.call_id, event.changes),
         (
@@ -175,7 +183,7 @@ fn diff_consumer_sends_next_update_after_buffer_interval() {
         HashMap::from([(
             PathBuf::from("hello.txt"),
             FileChange::Add {
-                content: "hello\n".to_string(),
+                content: String::new(),
             },
         )])
     );
@@ -190,7 +198,7 @@ fn diff_consumer_sends_next_update_after_buffer_interval() {
         HashMap::from([(
             PathBuf::from("hello.txt"),
             FileChange::Add {
-                content: "hello\nworld\n".to_string(),
+                content: "hello\n".to_string(),
             },
         )])
     );

codex-rs/core/src/tools/registry.rs

@@ -98,9 +98,9 @@ pub(crate) trait ToolArgumentDiffConsumer: Send {
     fn consume_diff(&mut self, turn: &TurnContext, call_id: String, diff: &str)
     -> Option<EventMsg>;
 
-    /// Flush any buffered event before the tool call completes.
-    fn flush_on_complete(&mut self) -> Option<EventMsg> {
-        None
+    /// Finish consuming argument diffs before the tool call completes.
+    fn finish(&mut self) -> Result<Option<EventMsg>, FunctionCallError> {
+        Ok(None)
     }
 }