🤖 fix: prefer ui_only diffs for compaction

Ensure edited file diff extraction prefers ui_only.file_edit.diff when present.

---

_Generated with `mux` • Model: `openai:gpt-5.2-codex` • Thinking: `xhigh` • Cost: `8.50`_

Author: ammar-agent

src/common/utils/messages/extractEditedFiles.test.ts

@@ -1,5 +1,6 @@
 import { describe, expect, it } from "bun:test";
 import { createPatch } from "diff";
+import { FILE_EDIT_DIFF_OMITTED_MESSAGE } from "@/common/types/tools";
 import type { MuxMessage } from "@/common/types/message";
 import { extractEditedFileDiffs, extractEditedFilePaths } from "./extractEditedFiles";
 
@@ -11,6 +12,7 @@ function createAssistantMessage(
     toolName: string;
     filePath: string;
     diff: string;
+    uiOnlyDiff?: string;
     success?: boolean;
   }>
 ): MuxMessage {
@@ -23,7 +25,19 @@ function createAssistantMessage(
       toolName: tc.toolName,
       state: "output-available" as const,
       input: { file_path: tc.filePath },
-      output: { success: tc.success ?? true, diff: tc.diff },
+      output: {
+        success: tc.success ?? true,
+        diff: tc.diff,
+        ...(tc.uiOnlyDiff
+          ? {
+              ui_only: {
+                file_edit: {
+                  diff: tc.uiOnlyDiff,
+                },
+              },
+            }
+          : {}),
+      },
     })),
   };
 }
@@ -129,6 +143,27 @@ describe("extractEditedFileDiffs", () => {
     expect(result[0].diff).toBe(diff);
   });
 
+  it("should prefer ui_only diffs when present", () => {
+    const originalContent = "line1\nline2\nline3";
+    const newContent = "line1\nmodified\nline3";
+    const diff = makeDiff("/path/to/file.ts", originalContent, newContent);
+
+    const messages: MuxMessage[] = [
+      createAssistantMessage([
+        {
+          toolName: "file_edit_replace_string",
+          filePath: "/path/to/file.ts",
+          diff: FILE_EDIT_DIFF_OMITTED_MESSAGE,
+          uiOnlyDiff: diff,
+        },
+      ]),
+    ];
+
+    const result = extractEditedFileDiffs(messages);
+    expect(result).toHaveLength(1);
+    expect(result[0].diff).toBe(diff);
+  });
+
   it("should combine multiple non-overlapping diffs for the same file", () => {
     // Edit 1: change line 2
     const original = "line1\nline2\nline3\nline4\nline5";

src/common/utils/messages/extractEditedFiles.ts

@@ -1,4 +1,5 @@
 import type { MuxMessage } from "@/common/types/message";
+import { getToolOutputUiOnly } from "@/common/utils/tools/toolOutputUiOnly";
 import { FILE_EDIT_TOOL_NAMES } from "@/common/types/tools";
 import { MAX_EDITED_FILES, MAX_FILE_CONTENT_SIZE } from "@/common/constants/attachments";
 import { applyPatch, createPatch, parsePatch } from "diff";
@@ -198,7 +199,11 @@ export function extractEditedFileDiffs(messages: MuxMessage[]): FileEditDiff[] {
       if (part.state !== "output-available") continue;
 
       const output = part.output as FileEditToolOutput | undefined;
-      if (!output?.success || !output.diff) continue;
+      if (!output?.success) continue;
+
+      const uiOnly = getToolOutputUiOnly(output);
+      const diff = uiOnly?.file_edit?.diff ?? output.diff;
+      if (!diff) continue;
 
       const input = part.input as FileEditToolInput | undefined;
       const filePath = input?.file_path;
@@ -208,7 +213,7 @@ export function extractEditedFileDiffs(messages: MuxMessage[]): FileEditDiff[] {
       if (!diffsByPath.has(filePath)) {
         diffsByPath.set(filePath, []);
       }
-      diffsByPath.get(filePath)!.push(output.diff);
+      diffsByPath.get(filePath)!.push(diff);
 
       // Update edit order (move to end if already exists)
       const idx = editOrder.indexOf(filePath);