From 0465de303aa1ddb798aa0b11bc4ebe62bcb704a4 Mon Sep 17 00:00:00 2001 From: Krushna Korade Date: Thu, 29 Jan 2026 00:37:58 +0530 Subject: [PATCH] fix: persist and restore workspace directories on session resume (#17454) Co-authored-by: Jacob Richman --- .../src/ui/commands/directoryCommand.test.tsx | 3 + .../cli/src/ui/commands/directoryCommand.tsx | 7 ++ .../cli/src/ui/hooks/useSessionResume.test.ts | 78 +++++++++++++++++++ packages/cli/src/ui/hooks/useSessionResume.ts | 11 +++ .../src/services/chatRecordingService.test.ts | 71 +++++++++++++++++ .../core/src/services/chatRecordingService.ts | 19 +++++ 6 files changed, 189 insertions(+) diff --git a/packages/cli/src/ui/commands/directoryCommand.test.tsx b/packages/cli/src/ui/commands/directoryCommand.test.tsx index 26e9bc727c..91ace7fca5 100644 --- a/packages/cli/src/ui/commands/directoryCommand.test.tsx +++ b/packages/cli/src/ui/commands/directoryCommand.test.tsx @@ -67,6 +67,9 @@ describe('directoryCommand', () => { isRestrictiveSandbox: vi.fn().mockReturnValue(false), getGeminiClient: vi.fn().mockReturnValue({ addDirectoryContext: vi.fn(), + getChatRecordingService: vi.fn().mockReturnValue({ + recordDirectories: vi.fn(), + }), }), getWorkingDir: () => path.resolve('/test/dir'), shouldLoadMemoryFromIncludeDirectories: () => false, diff --git a/packages/cli/src/ui/commands/directoryCommand.tsx b/packages/cli/src/ui/commands/directoryCommand.tsx index 53ec7acb7f..2da2f107df 100644 --- a/packages/cli/src/ui/commands/directoryCommand.tsx +++ b/packages/cli/src/ui/commands/directoryCommand.tsx @@ -57,6 +57,13 @@ async function finishAddingDirectories( const gemini = config.getGeminiClient(); if (gemini) { await gemini.addDirectoryContext(); + + // Persist directories to session file for resume support + const chatRecordingService = gemini.getChatRecordingService(); + const workspaceContext = config.getWorkspaceContext(); + chatRecordingService?.recordDirectories( + workspaceContext.getDirectories(), + ); } addItem({ type: MessageType.INFO, diff --git a/packages/cli/src/ui/hooks/useSessionResume.test.ts b/packages/cli/src/ui/hooks/useSessionResume.test.ts index 071fe5878b..9350cc167a 100644 --- a/packages/cli/src/ui/hooks/useSessionResume.test.ts +++ b/packages/cli/src/ui/hooks/useSessionResume.test.ts @@ -177,6 +177,84 @@ describe('useSessionResume', () => { expect(mockRefreshStatic).toHaveBeenCalledTimes(1); expect(mockGeminiClient.resumeChat).toHaveBeenCalledWith([], resumedData); }); + + it('should restore directories from resumed session data', async () => { + const mockAddDirectories = vi + .fn() + .mockReturnValue({ added: [], failed: [] }); + const mockWorkspaceContext = { + addDirectories: mockAddDirectories, + }; + const configWithWorkspace = { + ...mockConfig, + getWorkspaceContext: vi.fn().mockReturnValue(mockWorkspaceContext), + }; + + const { result } = renderHook(() => + useSessionResume({ + ...getDefaultProps(), + config: configWithWorkspace as unknown as Config, + }), + ); + + const resumedData: ResumedSessionData = { + conversation: { + sessionId: 'test-123', + projectHash: 'project-123', + startTime: '2025-01-01T00:00:00Z', + lastUpdated: '2025-01-01T01:00:00Z', + messages: [] as MessageRecord[], + directories: ['/restored/dir1', '/restored/dir2'], + }, + filePath: '/path/to/session.json', + }; + + await act(async () => { + await result.current.loadHistoryForResume([], [], resumedData); + }); + + expect(configWithWorkspace.getWorkspaceContext).toHaveBeenCalled(); + expect(mockAddDirectories).toHaveBeenCalledWith([ + '/restored/dir1', + '/restored/dir2', + ]); + }); + + it('should not call addDirectories when no directories in resumed session', async () => { + const mockAddDirectories = vi.fn(); + const mockWorkspaceContext = { + addDirectories: mockAddDirectories, + }; + const configWithWorkspace = { + ...mockConfig, + getWorkspaceContext: vi.fn().mockReturnValue(mockWorkspaceContext), + }; + + const { result } = renderHook(() => + useSessionResume({ + ...getDefaultProps(), + config: configWithWorkspace as unknown as Config, + }), + ); + + const resumedData: ResumedSessionData = { + conversation: { + sessionId: 'test-123', + projectHash: 'project-123', + startTime: '2025-01-01T00:00:00Z', + lastUpdated: '2025-01-01T01:00:00Z', + messages: [] as MessageRecord[], + // No directories field + }, + filePath: '/path/to/session.json', + }; + + await act(async () => { + await result.current.loadHistoryForResume([], [], resumedData); + }); + + expect(mockAddDirectories).not.toHaveBeenCalled(); + }); }); describe('callback stability', () => { diff --git a/packages/cli/src/ui/hooks/useSessionResume.ts b/packages/cli/src/ui/hooks/useSessionResume.ts index 21b9d0884f..9889c4bd12 100644 --- a/packages/cli/src/ui/hooks/useSessionResume.ts +++ b/packages/cli/src/ui/hooks/useSessionResume.ts @@ -71,6 +71,17 @@ export function useSessionResume({ }); refreshStaticRef.current(); // Force Static component to re-render with the updated history. + // Restore directories from the resumed session + if ( + resumedData.conversation.directories && + resumedData.conversation.directories.length > 0 + ) { + const workspaceContext = config.getWorkspaceContext(); + // Add back any directories that were saved in the session + // but filter out ones that no longer exist + workspaceContext.addDirectories(resumedData.conversation.directories); + } + // Give the history to the Gemini client. await config.getGeminiClient()?.resumeChat(clientHistory, resumedData); } catch (error) { diff --git a/packages/core/src/services/chatRecordingService.test.ts b/packages/core/src/services/chatRecordingService.test.ts index ff4fe51879..cba0a2e977 100644 --- a/packages/core/src/services/chatRecordingService.test.ts +++ b/packages/core/src/services/chatRecordingService.test.ts @@ -402,6 +402,77 @@ describe('ChatRecordingService', () => { }); }); + describe('recordDirectories', () => { + beforeEach(() => { + chatRecordingService.initialize(); + }); + + it('should save directories to the conversation', () => { + const writeFileSyncSpy = vi + .spyOn(fs, 'writeFileSync') + .mockImplementation(() => undefined); + const initialConversation = { + sessionId: 'test-session-id', + projectHash: 'test-project-hash', + messages: [ + { + id: '1', + type: 'user', + content: 'Hello', + timestamp: new Date().toISOString(), + }, + ], + }; + vi.spyOn(fs, 'readFileSync').mockReturnValue( + JSON.stringify(initialConversation), + ); + + chatRecordingService.recordDirectories([ + '/path/to/dir1', + '/path/to/dir2', + ]); + + expect(writeFileSyncSpy).toHaveBeenCalled(); + const conversation = JSON.parse( + writeFileSyncSpy.mock.calls[0][1] as string, + ) as ConversationRecord; + expect(conversation.directories).toEqual([ + '/path/to/dir1', + '/path/to/dir2', + ]); + }); + + it('should overwrite existing directories', () => { + const writeFileSyncSpy = vi + .spyOn(fs, 'writeFileSync') + .mockImplementation(() => undefined); + const initialConversation = { + sessionId: 'test-session-id', + projectHash: 'test-project-hash', + messages: [ + { + id: '1', + type: 'user', + content: 'Hello', + timestamp: new Date().toISOString(), + }, + ], + directories: ['/old/dir'], + }; + vi.spyOn(fs, 'readFileSync').mockReturnValue( + JSON.stringify(initialConversation), + ); + + chatRecordingService.recordDirectories(['/new/dir1', '/new/dir2']); + + expect(writeFileSyncSpy).toHaveBeenCalled(); + const conversation = JSON.parse( + writeFileSyncSpy.mock.calls[0][1] as string, + ) as ConversationRecord; + expect(conversation.directories).toEqual(['/new/dir1', '/new/dir2']); + }); + }); + describe('rewindTo', () => { it('should rewind the conversation to a specific message ID', () => { chatRecordingService.initialize(); diff --git a/packages/core/src/services/chatRecordingService.ts b/packages/core/src/services/chatRecordingService.ts index 2a920df8b7..a0a8034ce8 100644 --- a/packages/core/src/services/chatRecordingService.ts +++ b/packages/core/src/services/chatRecordingService.ts @@ -96,6 +96,8 @@ export interface ConversationRecord { lastUpdated: string; messages: MessageRecord[]; summary?: string; + /** Workspace directories added during the session via /dir add */ + directories?: string[]; } /** @@ -486,6 +488,23 @@ export class ChatRecordingService { } } + /** + * Records workspace directories to the session file. + * Called when directories are added via /dir add. + */ + recordDirectories(directories: readonly string[]): void { + if (!this.conversationFile) return; + + try { + this.updateConversation((conversation) => { + conversation.directories = [...directories]; + }); + } catch (error) { + debugLogger.error('Error saving directories to chat history.', error); + // Don't throw - we want graceful degradation + } + } + /** * Gets the current conversation data (for summary generation). */