feat: Multi-Directory Workspace Support (part1: add `--include-directories` option) (#4605)

Co-authored-by: Allen Hutchison <[email protected]>

3 files changed, 417 insertions, 0 deletions

diff --git a/packages/core/src/utils/flashFallback.integration.test.ts b/packages/core/src/utils/flashFallback.integration.test.ts
index f5e354a0..9211ad2f 100644
--- a/packages/core/src/utils/flashFallback.integration.test.ts
+++ b/packages/core/src/utils/flashFallback.integration.test.ts
@@ -6,6 +6,7 @@
 
 import { describe, it, expect, beforeEach, vi } from 'vitest';
 import { Config } from '../config/config.js';
+import fs from 'node:fs';
 import {
   setSimulate429,
   disableSimulationAfterFallback,
@@ -17,10 +18,16 @@ import { DEFAULT_GEMINI_FLASH_MODEL } from '../config/models.js';
 import { retryWithBackoff } from './retry.js';
 import { AuthType } from '../core/contentGenerator.js';
 
+vi.mock('node:fs');
+
 describe('Flash Fallback Integration', () => {
   let config: Config;
 
   beforeEach(() => {
+    vi.mocked(fs.existsSync).mockReturnValue(true);
+    vi.mocked(fs.statSync).mockReturnValue({
+      isDirectory: () => true,
+    } as fs.Stats);
     config = new Config({
       sessionId: 'test-session',
       targetDir: '/test',
diff --git a/packages/core/src/utils/workspaceContext.test.ts b/packages/core/src/utils/workspaceContext.test.ts
new file mode 100644
index 00000000..67d06b62
--- /dev/null
+++ b/packages/core/src/utils/workspaceContext.test.ts
@@ -0,0 +1,283 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { describe, it, expect, beforeEach, vi } from 'vitest';
+import * as fs from 'fs';
+import * as path from 'path';
+import { WorkspaceContext } from './workspaceContext.js';
+
+vi.mock('fs');
+
+describe('WorkspaceContext', () => {
+  let workspaceContext: WorkspaceContext;
+  // Use path module to create platform-agnostic paths
+  const mockCwd = path.resolve(path.sep, 'home', 'user', 'project');
+  const mockExistingDir = path.resolve(
+    path.sep,
+    'home',
+    'user',
+    'other-project',
+  );
+  const mockNonExistentDir = path.resolve(
+    path.sep,
+    'home',
+    'user',
+    'does-not-exist',
+  );
+  const mockSymlinkDir = path.resolve(path.sep, 'home', 'user', 'symlink');
+  const mockRealPath = path.resolve(path.sep, 'home', 'user', 'real-directory');
+
+  beforeEach(() => {
+    vi.resetAllMocks();
+
+    // Mock fs.existsSync
+    vi.mocked(fs.existsSync).mockImplementation((path) => {
+      const pathStr = path.toString();
+      return (
+        pathStr === mockCwd ||
+        pathStr === mockExistingDir ||
+        pathStr === mockSymlinkDir ||
+        pathStr === mockRealPath
+      );
+    });
+
+    // Mock fs.statSync
+    vi.mocked(fs.statSync).mockImplementation((path) => {
+      const pathStr = path.toString();
+      if (pathStr === mockNonExistentDir) {
+        throw new Error('ENOENT');
+      }
+      return {
+        isDirectory: () => true,
+      } as fs.Stats;
+    });
+
+    // Mock fs.realpathSync
+    vi.mocked(fs.realpathSync).mockImplementation((path) => {
+      const pathStr = path.toString();
+      if (pathStr === mockSymlinkDir) {
+        return mockRealPath;
+      }
+      return pathStr;
+    });
+  });
+
+  describe('initialization', () => {
+    it('should initialize with a single directory (cwd)', () => {
+      workspaceContext = new WorkspaceContext(mockCwd);
+      const directories = workspaceContext.getDirectories();
+      expect(directories).toHaveLength(1);
+      expect(directories[0]).toBe(mockCwd);
+    });
+
+    it('should validate and resolve directories to absolute paths', () => {
+      const absolutePath = path.join(mockCwd, 'subdir');
+      vi.mocked(fs.existsSync).mockImplementation(
+        (p) => p === mockCwd || p === absolutePath,
+      );
+      vi.mocked(fs.realpathSync).mockImplementation((p) => p.toString());
+
+      workspaceContext = new WorkspaceContext(mockCwd, [absolutePath]);
+      const directories = workspaceContext.getDirectories();
+      expect(directories).toContain(absolutePath);
+    });
+
+    it('should reject non-existent directories', () => {
+      expect(() => {
+        new WorkspaceContext(mockCwd, [mockNonExistentDir]);
+      }).toThrow('Directory does not exist');
+    });
+
+    it('should handle empty initialization', () => {
+      workspaceContext = new WorkspaceContext(mockCwd, []);
+      const directories = workspaceContext.getDirectories();
+      expect(directories).toHaveLength(1);
+      expect(directories[0]).toBe(mockCwd);
+    });
+  });
+
+  describe('adding directories', () => {
+    beforeEach(() => {
+      workspaceContext = new WorkspaceContext(mockCwd);
+    });
+
+    it('should add valid directories', () => {
+      workspaceContext.addDirectory(mockExistingDir);
+      const directories = workspaceContext.getDirectories();
+      expect(directories).toHaveLength(2);
+      expect(directories).toContain(mockExistingDir);
+    });
+
+    it('should resolve relative paths to absolute', () => {
+      // Since we can't mock path.resolve, we'll test with absolute paths
+      workspaceContext.addDirectory(mockExistingDir);
+      const directories = workspaceContext.getDirectories();
+      expect(directories).toContain(mockExistingDir);
+    });
+
+    it('should reject non-existent directories', () => {
+      expect(() => {
+        workspaceContext.addDirectory(mockNonExistentDir);
+      }).toThrow('Directory does not exist');
+    });
+
+    it('should prevent duplicate directories', () => {
+      workspaceContext.addDirectory(mockExistingDir);
+      workspaceContext.addDirectory(mockExistingDir);
+      const directories = workspaceContext.getDirectories();
+      expect(directories.filter((d) => d === mockExistingDir)).toHaveLength(1);
+    });
+
+    it('should handle symbolic links correctly', () => {
+      workspaceContext.addDirectory(mockSymlinkDir);
+      const directories = workspaceContext.getDirectories();
+      expect(directories).toContain(mockRealPath);
+      expect(directories).not.toContain(mockSymlinkDir);
+    });
+  });
+
+  describe('path validation', () => {
+    beforeEach(() => {
+      workspaceContext = new WorkspaceContext(mockCwd, [mockExistingDir]);
+    });
+
+    it('should accept paths within workspace directories', () => {
+      const validPath1 = path.join(mockCwd, 'src', 'file.ts');
+      const validPath2 = path.join(mockExistingDir, 'lib', 'module.js');
+
+      expect(workspaceContext.isPathWithinWorkspace(validPath1)).toBe(true);
+      expect(workspaceContext.isPathWithinWorkspace(validPath2)).toBe(true);
+    });
+
+    it('should reject paths outside workspace', () => {
+      const invalidPath = path.resolve(
+        path.dirname(mockCwd),
+        'outside-workspace',
+        'file.txt',
+      );
+      expect(workspaceContext.isPathWithinWorkspace(invalidPath)).toBe(false);
+    });
+
+    it('should resolve symbolic links before validation', () => {
+      const symlinkPath = path.join(mockCwd, 'symlink-file');
+      const realPath = path.join(mockCwd, 'real-file');
+
+      vi.mocked(fs.existsSync).mockReturnValue(true);
+      vi.mocked(fs.realpathSync).mockImplementation((p) => {
+        if (p === symlinkPath) {
+          return realPath;
+        }
+        return p.toString();
+      });
+
+      expect(workspaceContext.isPathWithinWorkspace(symlinkPath)).toBe(true);
+    });
+
+    it('should handle nested directories correctly', () => {
+      const nestedPath = path.join(
+        mockCwd,
+        'deeply',
+        'nested',
+        'path',
+        'file.txt',
+      );
+      expect(workspaceContext.isPathWithinWorkspace(nestedPath)).toBe(true);
+    });
+
+    it('should handle edge cases (root, parent references)', () => {
+      const rootPath = '/';
+      const parentPath = path.dirname(mockCwd);
+
+      expect(workspaceContext.isPathWithinWorkspace(rootPath)).toBe(false);
+      expect(workspaceContext.isPathWithinWorkspace(parentPath)).toBe(false);
+    });
+
+    it('should handle non-existent paths correctly', () => {
+      const nonExistentPath = path.join(mockCwd, 'does-not-exist.txt');
+      vi.mocked(fs.existsSync).mockImplementation((p) => p !== nonExistentPath);
+
+      // Should still validate based on path structure
+      expect(workspaceContext.isPathWithinWorkspace(nonExistentPath)).toBe(
+        true,
+      );
+    });
+  });
+
+  describe('getDirectories', () => {
+    it('should return a copy of directories array', () => {
+      workspaceContext = new WorkspaceContext(mockCwd);
+      const dirs1 = workspaceContext.getDirectories();
+      const dirs2 = workspaceContext.getDirectories();
+
+      expect(dirs1).not.toBe(dirs2); // Different array instances
+      expect(dirs1).toEqual(dirs2); // Same content
+    });
+  });
+
+  describe('symbolic link security', () => {
+    beforeEach(() => {
+      workspaceContext = new WorkspaceContext(mockCwd);
+    });
+
+    it('should follow symlinks but validate resolved path', () => {
+      const symlinkInsideWorkspace = path.join(mockCwd, 'link-to-subdir');
+      const resolvedInsideWorkspace = path.join(mockCwd, 'subdir');
+
+      vi.mocked(fs.existsSync).mockReturnValue(true);
+      vi.mocked(fs.realpathSync).mockImplementation((p) => {
+        if (p === symlinkInsideWorkspace) {
+          return resolvedInsideWorkspace;
+        }
+        return p.toString();
+      });
+
+      expect(
+        workspaceContext.isPathWithinWorkspace(symlinkInsideWorkspace),
+      ).toBe(true);
+    });
+
+    it('should prevent sandbox escape via symlinks', () => {
+      const symlinkEscape = path.join(mockCwd, 'escape-link');
+      const resolvedOutside = path.resolve(mockCwd, '..', 'outside-file');
+
+      vi.mocked(fs.existsSync).mockImplementation((p) => {
+        const pathStr = p.toString();
+        return (
+          pathStr === symlinkEscape ||
+          pathStr === resolvedOutside ||
+          pathStr === mockCwd
+        );
+      });
+      vi.mocked(fs.realpathSync).mockImplementation((p) => {
+        if (p.toString() === symlinkEscape) {
+          return resolvedOutside;
+        }
+        return p.toString();
+      });
+      vi.mocked(fs.statSync).mockImplementation(
+        (p) =>
+          ({
+            isDirectory: () => p.toString() !== resolvedOutside,
+          }) as fs.Stats,
+      );
+
+      workspaceContext = new WorkspaceContext(mockCwd);
+      expect(workspaceContext.isPathWithinWorkspace(symlinkEscape)).toBe(false);
+    });
+
+    it('should handle circular symlinks', () => {
+      const circularLink = path.join(mockCwd, 'circular');
+
+      vi.mocked(fs.existsSync).mockReturnValue(true);
+      vi.mocked(fs.realpathSync).mockImplementation(() => {
+        throw new Error('ELOOP: too many symbolic links encountered');
+      });
+
+      // Should handle the error gracefully
+      expect(workspaceContext.isPathWithinWorkspace(circularLink)).toBe(false);
+    });
+  });
+});
diff --git a/packages/core/src/utils/workspaceContext.ts b/packages/core/src/utils/workspaceContext.ts
new file mode 100644
index 00000000..16d1b4c9
--- /dev/null
+++ b/packages/core/src/utils/workspaceContext.ts
@@ -0,0 +1,127 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as fs from 'fs';
+import * as path from 'path';
+
+/**
+ * WorkspaceContext manages multiple workspace directories and validates paths
+ * against them. This allows the CLI to operate on files from multiple directories
+ * in a single session.
+ */
+export class WorkspaceContext {
+  private directories: Set<string>;
+
+  /**
+   * Creates a new WorkspaceContext with the given initial directory and optional additional directories.
+   * @param initialDirectory The initial working directory (usually cwd)
+   * @param additionalDirectories Optional array of additional directories to include
+   */
+  constructor(initialDirectory: string, additionalDirectories: string[] = []) {
+    this.directories = new Set<string>();
+
+    this.addDirectoryInternal(initialDirectory);
+
+    for (const dir of additionalDirectories) {
+      this.addDirectoryInternal(dir);
+    }
+  }
+
+  /**
+   * Adds a directory to the workspace.
+   * @param directory The directory path to add (can be relative or absolute)
+   * @param basePath Optional base path for resolving relative paths (defaults to cwd)
+   */
+  addDirectory(directory: string, basePath: string = process.cwd()): void {
+    this.addDirectoryInternal(directory, basePath);
+  }
+
+  /**
+   * Internal method to add a directory with validation.
+   */
+  private addDirectoryInternal(
+    directory: string,
+    basePath: string = process.cwd(),
+  ): void {
+    const absolutePath = path.isAbsolute(directory)
+      ? directory
+      : path.resolve(basePath, directory);
+
+    if (!fs.existsSync(absolutePath)) {
+      throw new Error(`Directory does not exist: ${absolutePath}`);
+    }
+
+    const stats = fs.statSync(absolutePath);
+    if (!stats.isDirectory()) {
+      throw new Error(`Path is not a directory: ${absolutePath}`);
+    }
+
+    let realPath: string;
+    try {
+      realPath = fs.realpathSync(absolutePath);
+    } catch (_error) {
+      throw new Error(`Failed to resolve path: ${absolutePath}`);
+    }
+
+    this.directories.add(realPath);
+  }
+
+  /**
+   * Gets a copy of all workspace directories.
+   * @returns Array of absolute directory paths
+   */
+  getDirectories(): readonly string[] {
+    return Array.from(this.directories);
+  }
+
+  /**
+   * Checks if a given path is within any of the workspace directories.
+   * @param pathToCheck The path to validate
+   * @returns True if the path is within the workspace, false otherwise
+   */
+  isPathWithinWorkspace(pathToCheck: string): boolean {
+    try {
+      const absolutePath = path.resolve(pathToCheck);
+
+      let resolvedPath = absolutePath;
+      if (fs.existsSync(absolutePath)) {
+        try {
+          resolvedPath = fs.realpathSync(absolutePath);
+        } catch (_error) {
+          return false;
+        }
+      }
+
+      for (const dir of this.directories) {
+        if (this.isPathWithinRoot(resolvedPath, dir)) {
+          return true;
+        }
+      }
+
+      return false;
+    } catch (_error) {
+      return false;
+    }
+  }
+
+  /**
+   * Checks if a path is within a given root directory.
+   * @param pathToCheck The absolute path to check
+   * @param rootDirectory The absolute root directory
+   * @returns True if the path is within the root directory, false otherwise
+   */
+  private isPathWithinRoot(
+    pathToCheck: string,
+    rootDirectory: string,
+  ): boolean {
+    const relative = path.relative(rootDirectory, pathToCheck);
+    return (
+      !relative.startsWith(`..${path.sep}`) &&
+      relative !== '..' &&
+      !path.isAbsolute(relative)
+    );
+  }
+}