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

packages/cli/src/ui/App.tsx

@@ -505,6 +505,7 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => {
     performMemoryRefresh,
     modelSwitchedFromQuotaError,
     setModelSwitchedFromQuotaError,
+    refreshStatic,
   );
 
   // Input handling

packages/cli/src/ui/hooks/useGeminiStream.ts

@@ -93,6 +93,7 @@ export const useGeminiStream = (
   performMemoryRefresh: () => Promise<void>,
   modelSwitchedFromQuotaError: boolean,
   setModelSwitchedFromQuotaError: React.Dispatch<React.SetStateAction<boolean>>,
+  onEditorClose: () => void,
 ) => {
   const [initError, setInitError] = useState<string | null>(null);
   const abortControllerRef = useRef<AbortController | null>(null);
@@ -133,6 +134,7 @@ export const useGeminiStream = (
       config,
       setPendingHistoryItem,
       getPreferredEditor,
+      onEditorClose,
     );
 
   const pendingToolCallGroupDisplay = useMemo(

packages/cli/src/ui/hooks/useReactToolScheduler.ts

@@ -70,6 +70,7 @@ export function useReactToolScheduler(
     React.SetStateAction<HistoryItemWithoutId | null>
   >,
   getPreferredEditor: () => EditorType | undefined,
+  onEditorClose: () => void,
 ): [TrackedToolCall[], ScheduleFn, MarkToolsAsSubmittedFn] {
   const [toolCallsForDisplay, setToolCallsForDisplay] = useState<
     TrackedToolCall[]
@@ -140,6 +141,7 @@ export function useReactToolScheduler(
         onToolCallsUpdate: toolCallsUpdateHandler,
         getPreferredEditor,
         config,
+        onEditorClose,
       }),
     [
       config,
@@ -147,6 +149,7 @@ export function useReactToolScheduler(
       allToolCallsCompleteHandler,
       toolCallsUpdateHandler,
       getPreferredEditor,
+      onEditorClose,
     ],
   );
 

packages/core/src/core/coreToolScheduler.test.ts

@@ -136,6 +136,7 @@ describe('CoreToolScheduler', () => {
       onAllToolCallsComplete,
       onToolCallsUpdate,
       getPreferredEditor: () => 'vscode',
+      onEditorClose: vi.fn(),
     });
 
     const abortController = new AbortController();
@@ -205,6 +206,7 @@ describe('CoreToolScheduler with payload', () => {
       onAllToolCallsComplete,
       onToolCallsUpdate,
       getPreferredEditor: () => 'vscode',
+      onEditorClose: vi.fn(),
     });
 
     const abortController = new AbortController();
@@ -482,6 +484,7 @@ describe('CoreToolScheduler edit cancellation', () => {
       onAllToolCallsComplete,
       onToolCallsUpdate,
       getPreferredEditor: () => 'vscode',
+      onEditorClose: vi.fn(),
     });
 
     const abortController = new AbortController();
@@ -571,6 +574,7 @@ describe('CoreToolScheduler YOLO mode', () => {
       onAllToolCallsComplete,
       onToolCallsUpdate,
       getPreferredEditor: () => 'vscode',
+      onEditorClose: vi.fn(),
     });
 
     const abortController = new AbortController();

packages/core/src/core/coreToolScheduler.ts

@@ -224,6 +224,7 @@ interface CoreToolSchedulerOptions {
   onToolCallsUpdate?: ToolCallsUpdateHandler;
   getPreferredEditor: () => EditorType | undefined;
   config: Config;
+  onEditorClose: () => void;
 }
 
 export class CoreToolScheduler {
@@ -234,6 +235,7 @@ export class CoreToolScheduler {
   private onToolCallsUpdate?: ToolCallsUpdateHandler;
   private getPreferredEditor: () => EditorType | undefined;
   private config: Config;
+  private onEditorClose: () => void;
 
   constructor(options: CoreToolSchedulerOptions) {
     this.config = options.config;
@@ -242,6 +244,7 @@ export class CoreToolScheduler {
     this.onAllToolCallsComplete = options.onAllToolCallsComplete;
     this.onToolCallsUpdate = options.onToolCallsUpdate;
     this.getPreferredEditor = options.getPreferredEditor;
+    this.onEditorClose = options.onEditorClose;
   }
 
   private setStatusInternal(
@@ -563,6 +566,7 @@ export class CoreToolScheduler {
           modifyContext as ModifyContext<typeof waitingToolCall.request.args>,
           editorType,
           signal,
+          this.onEditorClose,
         );
         this.setArgsInternal(callId, updatedParams);
         this.setStatusInternal(callId, 'awaiting_approval', {

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

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,

packages/core/src/utils/editor.test.ts

@@ -331,7 +331,7 @@ describe('editor utils', () => {
           }),
         };
         (spawn as Mock).mockReturnValue(mockSpawn);
-        await openDiff('old.txt', 'new.txt', editor);
+        await openDiff('old.txt', 'new.txt', editor, () => {});
         const diffCommand = getDiffCommand('old.txt', 'new.txt', editor)!;
         expect(spawn).toHaveBeenCalledWith(
           diffCommand.command,
@@ -361,9 +361,9 @@ describe('editor utils', () => {
           }),
         };
         (spawn as Mock).mockReturnValue(mockSpawn);
-        await expect(openDiff('old.txt', 'new.txt', editor)).rejects.toThrow(
-          'spawn error',
-        );
+        await expect(
+          openDiff('old.txt', 'new.txt', editor, () => {}),
+        ).rejects.toThrow('spawn error');
       });
 
       it(`should reject if ${editor} exits with non-zero code`, async () => {
@@ -375,9 +375,9 @@ describe('editor utils', () => {
           }),
         };
         (spawn as Mock).mockReturnValue(mockSpawn);
-        await expect(openDiff('old.txt', 'new.txt', editor)).rejects.toThrow(
-          `${editor} exited with code 1`,
-        );
+        await expect(
+          openDiff('old.txt', 'new.txt', editor, () => {}),
+        ).rejects.toThrow(`${editor} exited with code 1`);
       });
     }
 
@@ -385,7 +385,7 @@ describe('editor utils', () => {
     for (const editor of execSyncEditors) {
       it(`should call execSync for ${editor} on non-windows`, async () => {
         Object.defineProperty(process, 'platform', { value: 'linux' });
-        await openDiff('old.txt', 'new.txt', editor);
+        await openDiff('old.txt', 'new.txt', editor, () => {});
         expect(execSync).toHaveBeenCalledTimes(1);
         const diffCommand = getDiffCommand('old.txt', 'new.txt', editor)!;
         const expectedCommand = `${
@@ -399,7 +399,7 @@ describe('editor utils', () => {
 
       it(`should call execSync for ${editor} on windows`, async () => {
         Object.defineProperty(process, 'platform', { value: 'win32' });
-        await openDiff('old.txt', 'new.txt', editor);
+        await openDiff('old.txt', 'new.txt', editor, () => {});
         expect(execSync).toHaveBeenCalledTimes(1);
         const diffCommand = getDiffCommand('old.txt', 'new.txt', editor)!;
         const expectedCommand = `${diffCommand.command} ${diffCommand.args.join(
@@ -417,11 +417,46 @@ describe('editor utils', () => {
         .spyOn(console, 'error')
         .mockImplementation(() => {});
       // @ts-expect-error Testing unsupported editor
-      await openDiff('old.txt', 'new.txt', 'foobar');
+      await openDiff('old.txt', 'new.txt', 'foobar', () => {});
       expect(consoleErrorSpy).toHaveBeenCalledWith(
         'No diff tool available. Install a supported editor.',
       );
     });
+
+    describe('onEditorClose callback', () => {
+      it('should call onEditorClose for execSync editors', async () => {
+        (execSync as Mock).mockReturnValue(Buffer.from(`/usr/bin/`));
+        const onEditorClose = vi.fn();
+        await openDiff('old.txt', 'new.txt', 'vim', onEditorClose);
+        expect(execSync).toHaveBeenCalledTimes(1);
+        expect(onEditorClose).toHaveBeenCalledTimes(1);
+      });
+
+      it('should call onEditorClose for execSync editors when an error is thrown', async () => {
+        (execSync as Mock).mockImplementation(() => {
+          throw new Error('test error');
+        });
+        const onEditorClose = vi.fn();
+        openDiff('old.txt', 'new.txt', 'vim', onEditorClose);
+        expect(execSync).toHaveBeenCalledTimes(1);
+        expect(onEditorClose).toHaveBeenCalledTimes(1);
+      });
+
+      it('should not call onEditorClose for spawn editors', async () => {
+        const onEditorClose = vi.fn();
+        const mockSpawn = {
+          on: vi.fn((event, cb) => {
+            if (event === 'close') {
+              cb(0);
+            }
+          }),
+        };
+        (spawn as Mock).mockReturnValue(mockSpawn);
+        await openDiff('old.txt', 'new.txt', 'vscode', onEditorClose);
+        expect(spawn).toHaveBeenCalledTimes(1);
+        expect(onEditorClose).not.toHaveBeenCalled();
+      });
+    });
   });
 
   describe('allowEditorTypeInSandbox', () => {

packages/core/src/utils/editor.ts

@@ -164,6 +164,7 @@ export async function openDiff(
   oldPath: string,
   newPath: string,
   editor: EditorType,
+  onEditorClose: () => void,
 ): Promise<void> {
   const diffCommand = getDiffCommand(oldPath, newPath, editor);
   if (!diffCommand) {
@@ -206,10 +207,16 @@ export async function openDiff(
           process.platform === 'win32'
             ? `${diffCommand.command} ${diffCommand.args.join(' ')}`
             : `${diffCommand.command} ${diffCommand.args.map((arg) => `"${arg}"`).join(' ')}`;
+        try {
           execSync(command, {
             stdio: 'inherit',
             encoding: 'utf8',
           });
+        } catch (e) {
+          console.error('Error in onEditorClose callback:', e);
+        } finally {
+          onEditorClose();
+        }
         break;
       }