perf: batch remote filesystem walks

.changeset/batch-remote-filesystem-walks.md

@@ -0,0 +1,5 @@
+---
+"@moonshot-ai/kimi-code": patch
+---
+
+Speed up remote file search and completion without per-file metadata requests.

.github/workflows/ci.yml

@@ -62,6 +62,7 @@ jobs:
           cache: pnpm
 
       - run: pnpm install --frozen-lockfile
+      - run: sudo apt-get update && sudo apt-get install -y ripgrep
       - run: pnpm run test --shard=${{ matrix.shard }}/5
 
   # pi-tui's suite runs on node:test (not vitest), so the root `pnpm run test`

packages/acp-server/src/acp-fs/acpFsService.ts

@@ -20,6 +20,8 @@ import {
   HostFileSystem,
   type HostDirEntry,
   type HostFileStat,
+  type HostWalkOptions,
+  type HostWalkResult,
   IHostFileSystem,
   ISessionContext,
   LifecycleScope,
@@ -151,6 +153,10 @@ export class AcpHostFileSystem implements IHostFileSystem {
     return this.inner.readdir(path);
   }
 
+  walk(path: string, options: HostWalkOptions): Promise<HostWalkResult> {
+    return this.inner.walk(path, options);
+  }
+
   mkdir(path: string, options?: { readonly recursive?: boolean }): Promise<void> {
     return this.inner.mkdir(path, options);
   }

packages/agent-core-v2/src/agent/tools/os/grep/grepTool.ts

@@ -44,7 +44,6 @@ import { type GrepInput, GrepInputSchema, IGrepTool } from './grep';
 
 const RG_MAX_COLUMNS = 500;
 const DEFAULT_HEAD_LIMIT = 250;
-const MTIME_STAT_CONCURRENCY = 32;
 
 const VCS_DIRECTORIES_TO_EXCLUDE = ['.git', '.svn', '.hg', '.bzr', '.jj', '.sl'] as const;
 const SENSITIVE_KEY_BASENAMES = ['id_rsa', 'id_ed25519', 'id_ecdsa'] as const;
@@ -222,18 +221,7 @@ export class GrepTool implements IGrepTool {
 
     const filteredSensitive = new Set<string>();
     const keptLines = filterSensitiveLines(rawLines, mode, filteredSensitive, pathClass);
-    let orderedLines: ParsedGrepLine[];
-    try {
-      orderedLines =
-        mode === 'files_with_matches' && !timedOut
-          ? await this.sortFilesWithMatchesByMtime(fs, keptLines, signal)
-          : keptLines;
-    } catch (error) {
-      if (error instanceof GrepAbortedError) {
-        return { isError: true, output: 'Grep aborted' };
-      }
-      throw error;
-    }
+    const orderedLines = keptLines;
 
     const offset = args.offset ?? 0;
     const headLimit = args.head_limit ?? DEFAULT_HEAD_LIMIT;
@@ -329,32 +317,6 @@ export class GrepTool implements IGrepTool {
     };
   }
 
-  private async sortFilesWithMatchesByMtime(
-    fs: IHostFileSystem,
-    lines: readonly ParsedGrepLine[],
-    signal: AbortSignal,
-  ): Promise<ParsedGrepLine[]> {
-    const entries = await mapWithConcurrency(
-      lines,
-      MTIME_STAT_CONCURRENCY,
-      signal,
-      async (line, index) => {
-        const path =
-          line.kind === 'record' ? line.filePath : line.kind === 'legacy' ? line.text : undefined;
-        let mtime = 0;
-        if (path !== undefined) {
-          try {
-            const mtimeMs = (await fs.stat(path)).mtimeMs ?? 0;
-            mtime = Math.trunc(mtimeMs / 1000);
-          } catch {
-          }
-        }
-        return { line, mtime, index };
-      },
-    );
-    entries.sort((a, b) => b.mtime - a.mtime || a.index - b.index);
-    return entries.map((entry) => entry.line);
-  }
 }
 
 registerAgentToolService(IGrepTool, GrepTool, {
@@ -387,41 +349,6 @@ type ParsedGrepLine =
       readonly text: string;
     };
 
-class GrepAbortedError extends Error {
-  constructor() {
-    super('Grep aborted');
-    this.name = 'GrepAbortedError';
-  }
-}
-
-async function mapWithConcurrency<T, U>(
-  items: readonly T[],
-  concurrency: number,
-  signal: AbortSignal,
-  mapper: (item: T, index: number) => Promise<U>,
-): Promise<U[]> {
-  if (signal.aborted) throw new GrepAbortedError();
-  if (items.length === 0) return [];
-
-  const results: U[] = [];
-  results.length = items.length;
-  let nextIndex = 0;
-  const workerCount = Math.min(Math.max(1, concurrency), items.length);
-  await Promise.all(
-    Array.from({ length: workerCount }, async () => {
-      while (true) {
-        if (signal.aborted) return;
-        const index = nextIndex;
-        nextIndex += 1;
-        if (index >= items.length) return;
-        results[index] = await mapper(items[index] as T, index);
-      }
-    }),
-  );
-  if (signal.aborted) throw new GrepAbortedError();
-  return results;
-}
-
 function buildRgArgs(
   rgPath: string,
   args: GrepInput,
@@ -440,7 +367,7 @@ function buildRgArgs(
     cmd.push('--glob', `!${dir}`);
   }
 
-  if (mode === 'files_with_matches') cmd.push('-l');
+  if (mode === 'files_with_matches') cmd.push('-l', '--sortr=modified');
   else if (mode === 'count_matches') {
     cmd.push('--count-matches', '--with-filename');
   }

packages/agent-core-v2/src/os/backends/node-local/hostFsService.ts

@@ -14,8 +14,15 @@ import {
 } from 'node:fs/promises';
 import { decodeTextWithErrors, readUtf8Lines, type TextDecodeErrors } from '#/_base/execEnv/decodeText';
 
-import { type HostDirEntry, type HostFileStat, IHostFileSystem } from '#/os/interface/hostFileSystem';
+import {
+  type HostDirEntry,
+  type HostFileStat,
+  type HostWalkOptions,
+  type HostWalkResult,
+  IHostFileSystem,
+} from '#/os/interface/hostFileSystem';
 import { toHostFsError } from '#/os/interface/hostFsErrors';
+import { walkHostDirectory } from '#/os/interface/hostFsWalk';
 
 const READ_CHUNK_SIZE = 64 * 1024;
 
@@ -192,6 +199,10 @@ export class HostFileSystem implements IHostFileSystem {
     }
   }
 
+  async walk(path: string, options: HostWalkOptions): Promise<HostWalkResult> {
+    return walkHostDirectory(this, path, options);
+  }
+
   async mkdir(path: string, options?: { readonly recursive?: boolean }): Promise<void> {
     try {
       await mkdir(path, { recursive: options?.recursive ?? false });

packages/agent-core-v2/src/os/backends/remote/remoteFileSystem.ts

@@ -1,3 +1,5 @@
+import * as posixPath from 'node:path/posix';
+
 import {
   decodeTextWithErrors,
   readUtf8Lines,
@@ -7,8 +9,15 @@ import { splitLinesKeepingTerminator } from '#/_base/text/line-endings';
 import type {
   HostDirEntry,
   HostFileStat,
+  HostWalkOptions,
+  HostWalkResult,
   IHostFileSystem,
 } from '#/os/interface/hostFileSystem';
+import {
+  HOST_WALK_DEFAULT_TIMEOUT_MS,
+  HOST_WALK_MAX_IGNORE_BYTES,
+  walkHostDirectory,
+} from '#/os/interface/hostFsWalk';
 import {
   HostFsError,
   OsFsErrors,
@@ -28,16 +37,23 @@ import {
   FS_READ_FILE_METHOD,
   FS_REMOVE_METHOD,
   FS_RENAME_METHOD,
+  FS_WALK_METHOD,
   FS_WRITE_FILE_CHUNK_BYTES,
   FS_WRITE_FILE_METHOD,
   type FsGetMetadataResult,
   type FsReadDirectoryResult,
   type FsReadFileResult,
+  type FsWalkResult,
   type FsWriteMode,
 } from '#/remote/protocol/methods';
 import { rpcDomainError } from './rpcDomainError';
 
 const HOST_FS_CODES: ReadonlySet<string> = new Set(Object.values(OsFsErrors.codes));
+const REMOTE_WALK_PATH = {
+  join: (...paths: readonly string[]) => posixPath.join(...paths),
+  relative: (from: string, to: string) => posixPath.relative(from, to),
+  separator: posixPath.sep,
+};
 
 export function toRemoteFsError(error: unknown): Error {
   if (!(error instanceof RpcError)) {
@@ -232,6 +248,46 @@ export class RemoteFileSystem implements IHostFileSystem {
     }));
   }
 
+  async walk(path: string, options: HostWalkOptions): Promise<HostWalkResult> {
+    const remainingMs = options.deadlineMs === undefined ? undefined : options.deadlineMs - Date.now();
+    if (remainingMs !== undefined && remainingMs <= 0) {
+      return { entries: [], errors: [], truncated: true, responseBytes: 0 };
+    }
+    const timeoutMs = Math.max(1, Math.min(
+      remainingMs ?? HOST_WALK_DEFAULT_TIMEOUT_MS,
+      HOST_WALK_DEFAULT_TIMEOUT_MS,
+    ));
+    if (options.ignore !== undefined && Buffer.byteLength(JSON.stringify(options.ignore)) > HOST_WALK_MAX_IGNORE_BYTES) {
+      return this.walkFallback(path, options);
+    }
+    try {
+      const result = await this.connection.call(FS_WALK_METHOD, {
+        path,
+        maxDepth: options.maxDepth,
+        maxDirectories: options.maxDirectories,
+        maxEntries: options.maxEntries,
+        maxResponseBytes: options.maxResponseBytes,
+        ignore: options.ignore,
+        timeoutMs,
+      }) as FsWalkResult;
+      return {
+        entries: result.entries,
+        errors: result.errors.map(({ path: errorPath, error }) => ({
+          path: errorPath,
+          error: toRemoteFsError(RpcError.fromErrorBody(error)),
+        })),
+        truncated: result.truncated,
+        responseBytes: result.responseBytes,
+      };
+    } catch (error) {
+      throw toRemoteFsError(error);
+    }
+  }
+
+  private async walkFallback(path: string, options: HostWalkOptions): Promise<HostWalkResult> {
+    return walkHostDirectory(this, path, options, { path: REMOTE_WALK_PATH });
+  }
+
   async mkdir(path: string, options?: { readonly recursive?: boolean }): Promise<void> {
     await this.call(FS_CREATE_DIRECTORY_METHOD, {
       path,

packages/agent-core-v2/src/os/interface/hostFileSystem.ts

@@ -18,6 +18,27 @@ export interface HostDirEntry {
   readonly isSymbolicLink?: boolean;
 }
 
+export interface HostWalkOptions {
+  readonly maxDepth: number;
+  readonly maxDirectories: number;
+  readonly maxEntries: number;
+  readonly maxResponseBytes: number;
+  readonly ignore?: readonly string[];
+  readonly deadlineMs?: number;
+}
+
+export interface HostWalkError {
+  readonly path: string;
+  readonly error: Error;
+}
+
+export interface HostWalkResult {
+  readonly entries: readonly HostDirEntry[];
+  readonly errors: readonly HostWalkError[];
+  readonly truncated: boolean;
+  readonly responseBytes: number;
+}
+
 export interface IHostFileSystem {
   readonly _serviceBrand: undefined;
 
@@ -37,6 +58,7 @@ export interface IHostFileSystem {
   stat(path: string): Promise<HostFileStat>;
   lstat(path: string): Promise<HostFileStat>;
   readdir(path: string): Promise<readonly HostDirEntry[]>;
+  walk(path: string, options: HostWalkOptions): Promise<HostWalkResult>;
   mkdir(path: string, options?: { readonly recursive?: boolean }): Promise<void>;
   remove(path: string): Promise<void>;
   rename?(from: string, to: string): Promise<void>;