fix: persist and restore workspace directories on session resume (#17454)

Co-authored-by: Jacob Richman <jacob314@gmail.com>

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,

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,

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', () => {

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) {

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

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).
    */