fix(core): prevent UI shift after vim edit (#5315)

2 files changed, 11 insertions, 1 deletions

diff --git a/packages/core/src/tools/modifiable-tool.test.ts b/packages/core/src/tools/modifiable-tool.test.ts
index 47cf41fe..eb7e8dbf 100644
--- a/packages/core/src/tools/modifiable-tool.test.ts
+++ b/packages/core/src/tools/modifiable-tool.test.ts
@@ -94,6 +94,7 @@ describe('modifyWithEditor', () => {
         mockModifyContext,
         'vscode' as EditorType,
         abortSignal,
+        vi.fn(),
       );
 
       expect(mockModifyContext.getCurrentContent).toHaveBeenCalledWith(
@@ -148,6 +149,7 @@ describe('modifyWithEditor', () => {
         mockModifyContext,
         'vscode' as EditorType,
         abortSignal,
+        vi.fn(),
       );
 
       const stats = await fsp.stat(diffDir);
@@ -165,6 +167,7 @@ describe('modifyWithEditor', () => {
         mockModifyContext,
         'vscode' as EditorType,
         abortSignal,
+        vi.fn(),
       );
 
       expect(mkdirSpy).not.toHaveBeenCalled();
@@ -183,6 +186,7 @@ describe('modifyWithEditor', () => {
       mockModifyContext,
       'vscode' as EditorType,
       abortSignal,
+      vi.fn(),
     );
 
     expect(mockCreatePatch).toHaveBeenCalledWith(
@@ -211,6 +215,7 @@ describe('modifyWithEditor', () => {
       mockModifyContext,
       'vscode' as EditorType,
       abortSignal,
+      vi.fn(),
     );
 
     expect(mockCreatePatch).toHaveBeenCalledWith(
@@ -241,6 +246,7 @@ describe('modifyWithEditor', () => {
         mockModifyContext,
         'vscode' as EditorType,
         abortSignal,
+        vi.fn(),
       ),
     ).rejects.toThrow('Editor failed to open');
 
@@ -267,6 +273,7 @@ describe('modifyWithEditor', () => {
       mockModifyContext,
       'vscode' as EditorType,
       abortSignal,
+      vi.fn(),
     );
 
     expect(consoleErrorSpy).toHaveBeenCalledTimes(2);
@@ -290,6 +297,7 @@ describe('modifyWithEditor', () => {
       mockModifyContext,
       'vscode' as EditorType,
       abortSignal,
+      vi.fn(),
     );
 
     expect(mockOpenDiff).toHaveBeenCalledOnce();
@@ -311,6 +319,7 @@ describe('modifyWithEditor', () => {
       mockModifyContext,
       'vscode' as EditorType,
       abortSignal,
+      vi.fn(),
     );
 
     expect(mockOpenDiff).toHaveBeenCalledOnce();
diff --git a/packages/core/src/tools/modifiable-tool.ts b/packages/core/src/tools/modifiable-tool.ts
index 4f96a49c..42de3eb6 100644
--- a/packages/core/src/tools/modifiable-tool.ts
+++ b/packages/core/src/tools/modifiable-tool.ts
@@ -138,6 +138,7 @@ export async function modifyWithEditor<ToolParams>(
   modifyContext: ModifyContext<ToolParams>,
   editorType: EditorType,
   _abortSignal: AbortSignal,
+  onEditorClose: () => void,
 ): Promise<ModifyResult<ToolParams>> {
   const currentContent = await modifyContext.getCurrentContent(originalParams);
   const proposedContent =
@@ -150,7 +151,7 @@ export async function modifyWithEditor<ToolParams>(
   );
 
   try {
-    await openDiff(oldPath, newPath, editorType);
+    await openDiff(oldPath, newPath, editorType, onEditorClose);
     const result = getUpdatedParams(
       oldPath,
       newPath,