fix(fs): cap file list and abort untracked scan to prevent session init timeout on large repos (#2819)

Co-authored-by: Charles Vien <charles.v@posthog.com>

Author: ricardo-leiva

packages/git/src/queries.test.ts

@@ -13,6 +13,7 @@ import {
   getChangedFilesDetailed,
   getGitBusyState,
   getLinkedWorktreeMainPath,
+  listAllFiles,
   remoteBranchExists,
   splitUnifiedDiffByFile,
 } from "./queries";
@@ -597,3 +598,53 @@ describe("getLinkedWorktreeMainPath", () => {
     expect(getLinkedWorktreeMainPath(worktreeDir)).toBeNull();
   });
 });
+
+describe("listAllFiles", () => {
+  let repoDir: string;
+
+  afterEach(async () => {
+    if (repoDir) {
+      await rm(repoDir, { recursive: true, force: true });
+    }
+  });
+
+  it("combines tracked and untracked files uncapped by default", async () => {
+    repoDir = await setupRepo();
+    await writeFile(path.join(repoDir, "untracked.txt"), "content");
+
+    const files = await listAllFiles(repoDir);
+
+    expect(files.sort()).toEqual(["file.txt", "untracked.txt"]);
+  });
+
+  it("truncates to maxFiles", async () => {
+    repoDir = await setupRepo();
+    const git = createGitClient(repoDir);
+    await writeFile(path.join(repoDir, "b.txt"), "content");
+    await writeFile(path.join(repoDir, "c.txt"), "content");
+    await git.add(["b.txt", "c.txt"]);
+    await git.commit("add more files");
+
+    const files = await listAllFiles(repoDir, { maxFiles: 2 });
+
+    expect(files.length).toBe(2);
+  });
+
+  it("keeps tracked files over untracked ones when truncating", async () => {
+    repoDir = await setupRepo();
+    await writeFile(path.join(repoDir, "untracked.txt"), "content");
+
+    const files = await listAllFiles(repoDir, { maxFiles: 1 });
+
+    expect(files).toEqual(["file.txt"]);
+  });
+
+  it("returns tracked files when the untracked scan times out", async () => {
+    repoDir = await setupRepo();
+    await writeFile(path.join(repoDir, "untracked.txt"), "content");
+
+    const files = await listAllFiles(repoDir, { timeoutMs: 0 });
+
+    expect(files).toContain("file.txt");
+  });
+});

packages/git/src/queries.ts

@@ -1181,15 +1181,37 @@ export async function listUntrackedFiles(
   );
 }
 
+export interface ListAllFilesOptions {
+  maxFiles?: number;
+  timeoutMs?: number;
+}
+
 export async function listAllFiles(
   baseDir: string,
-  options?: CreateGitClientOptions,
+  options?: ListAllFilesOptions,
 ): Promise<string[]> {
-  const [tracked, untracked] = await Promise.all([
-    listFiles(baseDir, options),
-    listUntrackedFiles(baseDir, options),
-  ]);
-  return [...tracked, ...untracked];
+  const { maxFiles, timeoutMs } = options ?? {};
+  const controller =
+    timeoutMs !== undefined ? new AbortController() : undefined;
+  const timer =
+    controller && timeoutMs !== undefined
+      ? setTimeout(() => controller.abort(), timeoutMs)
+      : undefined;
+  try {
+    const [tracked, untracked] = await Promise.all([
+      listFiles(baseDir).catch((): string[] => []),
+      listUntrackedFiles(baseDir, { abortSignal: controller?.signal }).catch(
+        (): string[] => [],
+      ),
+    ]);
+    const combined = tracked.concat(untracked);
+    if (maxFiles !== undefined && combined.length > maxFiles) {
+      combined.splice(maxFiles);
+    }
+    return combined;
+  } finally {
+    if (timer) clearTimeout(timer);
+  }
 }
 
 // Tracked + untracked files containing `pattern` (literal, case-insensitive).

packages/workspace-server/src/services/fs/service.test.ts

@@ -48,6 +48,35 @@ describe("FsService.listRepoFiles", () => {
       { path: "src/sub/c.ts", kind: "file" },
     ]);
   });
+
+  it("passes the file cap and timeout through to listAllFiles", async () => {
+    vi.mocked(getChangedFiles).mockResolvedValue(new Set());
+    vi.mocked(listAllFiles).mockResolvedValue([]);
+
+    const service = new FsService();
+    await service.listRepoFiles("/repo");
+
+    expect(listAllFiles).toHaveBeenCalledWith("/repo", {
+      maxFiles: 50_000,
+      timeoutMs: 8_000,
+    });
+  });
+
+  it("total entries can exceed the file cap when derived directories are included", async () => {
+    vi.mocked(getChangedFiles).mockResolvedValue(new Set());
+    const cappedList = Array.from(
+      { length: 50_000 },
+      (_, i) => `src/sub${i}/file.ts`,
+    );
+    vi.mocked(listAllFiles).mockResolvedValue(cappedList);
+
+    const service = new FsService();
+    const entries = await service.listRepoFiles("/repo");
+
+    const fileEntries = entries.filter((e) => e.kind === "file");
+    expect(fileEntries.length).toBe(50_000);
+    expect(entries.length).toBeGreaterThan(50_000);
+  });
 });
 
 describe("FsService repo file IO", () => {

packages/workspace-server/src/services/fs/service.ts

@@ -8,6 +8,8 @@ import type { BoundedReadResult, DirectoryEntry, FileEntry } from "./schemas";
 export class FsService {
   private static readonly CACHE_TTL = 30000;
   private static readonly READ_REPO_FILES_CONCURRENCY = 24;
+  private static readonly MAX_REPO_FILES = 50_000;
+  private static readonly LIST_FILES_TIMEOUT_MS = 8_000;
   private cache = new Map<string, { files: FileEntry[]; timestamp: number }>();
 
   async listDirectory(dirPath: string): Promise<DirectoryEntry[]> {
@@ -43,7 +45,7 @@ export class FsService {
       const changedFiles = await getChangedFiles(repoPath);
 
       if (query?.trim()) {
-        const allFiles = await listAllFiles(repoPath);
+        const allFiles = await this.listAllFilesBounded(repoPath);
         const directories = this.deriveDirectories(allFiles);
         const lowerQuery = query.toLowerCase();
         const matchingDirs = directories.filter((d) =>
@@ -64,7 +66,7 @@ export class FsService {
         return limit ? cached.files.slice(0, limit) : cached.files;
       }
 
-      const files = await listAllFiles(repoPath);
+      const files = await this.listAllFilesBounded(repoPath);
       const directories = this.deriveDirectories(files);
       const entries = [
         ...this.toDirectoryEntries(directories),
@@ -221,6 +223,13 @@ export class FsService {
     }));
   }
 
+  private listAllFilesBounded(repoPath: string): Promise<string[]> {
+    return listAllFiles(repoPath, {
+      maxFiles: FsService.MAX_REPO_FILES,
+      timeoutMs: FsService.LIST_FILES_TIMEOUT_MS,
+    });
+  }
+
   private deriveDirectories(files: string[]): string[] {
     const dirs = new Set<string>();
     for (const file of files) {