diff --git a/packages/cli/src/ui/App.tsx b/packages/cli/src/ui/App.tsx index f07a5386..66396c36 100644 --- a/packages/cli/src/ui/App.tsx +++ b/packages/cli/src/ui/App.tsx @@ -505,6 +505,7 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => { performMemoryRefresh, modelSwitchedFromQuotaError, setModelSwitchedFromQuotaError, + refreshStatic, ); // Input handling diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index e53e77dc..63ba961f 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -93,6 +93,7 @@ export const useGeminiStream = ( performMemoryRefresh: () => Promise, modelSwitchedFromQuotaError: boolean, setModelSwitchedFromQuotaError: React.Dispatch>, + onEditorClose: () => void, ) => { const [initError, setInitError] = useState(null); const abortControllerRef = useRef(null); @@ -133,6 +134,7 @@ export const useGeminiStream = ( config, setPendingHistoryItem, getPreferredEditor, + onEditorClose, ); const pendingToolCallGroupDisplay = useMemo( diff --git a/packages/cli/src/ui/hooks/useReactToolScheduler.ts b/packages/cli/src/ui/hooks/useReactToolScheduler.ts index 307a90cf..01993650 100644 --- a/packages/cli/src/ui/hooks/useReactToolScheduler.ts +++ b/packages/cli/src/ui/hooks/useReactToolScheduler.ts @@ -70,6 +70,7 @@ export function useReactToolScheduler( React.SetStateAction >, 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, ], ); diff --git a/packages/core/src/core/coreToolScheduler.test.ts b/packages/core/src/core/coreToolScheduler.test.ts index 80651a14..623fb436 100644 --- a/packages/core/src/core/coreToolScheduler.test.ts +++ b/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(); diff --git a/packages/core/src/core/coreToolScheduler.ts b/packages/core/src/core/coreToolScheduler.ts index b4c10a64..5f2cc895 100644 --- a/packages/core/src/core/coreToolScheduler.ts +++ b/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, editorType, signal, + this.onEditorClose, ); this.setArgsInternal(callId, updatedParams); this.setStatusInternal(callId, 'awaiting_approval', { 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( modifyContext: ModifyContext, editorType: EditorType, _abortSignal: AbortSignal, + onEditorClose: () => void, ): Promise> { const currentContent = await modifyContext.getCurrentContent(originalParams); const proposedContent = @@ -150,7 +151,7 @@ export async function modifyWithEditor( ); try { - await openDiff(oldPath, newPath, editorType); + await openDiff(oldPath, newPath, editorType, onEditorClose); const result = getUpdatedParams( oldPath, newPath, diff --git a/packages/core/src/utils/editor.test.ts b/packages/core/src/utils/editor.test.ts index 203223ae..afdc2b24 100644 --- a/packages/core/src/utils/editor.test.ts +++ b/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', () => { diff --git a/packages/core/src/utils/editor.ts b/packages/core/src/utils/editor.ts index 704d1cbb..f22297df 100644 --- a/packages/core/src/utils/editor.ts +++ b/packages/core/src/utils/editor.ts @@ -164,6 +164,7 @@ export async function openDiff( oldPath: string, newPath: string, editor: EditorType, + onEditorClose: () => void, ): Promise { 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(' ')}`; - execSync(command, { - stdio: 'inherit', - encoding: 'utf8', - }); + try { + execSync(command, { + stdio: 'inherit', + encoding: 'utf8', + }); + } catch (e) { + console.error('Error in onEditorClose callback:', e); + } finally { + onEditorClose(); + } break; }