@aaroncql/pim-agent 0.2.0 → 0.4.0

package/src/extensions/file-picker/filePickerWorkerMessages.ts

@@ -0,0 +1,39 @@
+import type { AutocompleteItem } from "@earendil-works/pi-tui";
+
+export type FilePickerWorkerRequest =
+  | {
+      readonly id: number;
+      readonly type: "refreshRelative";
+      readonly root: string;
+    }
+  | {
+      readonly id: number;
+      readonly type: "rank";
+      readonly query: string;
+      readonly limit?: number;
+    };
+
+export type FilePickerWorkerResponse =
+  | {
+      readonly id: number;
+      readonly type: "refreshRelative";
+      readonly ok: true;
+    }
+  | {
+      readonly id: number;
+      readonly type: "refreshRelative";
+      readonly ok: false;
+      readonly error: string;
+    }
+  | {
+      readonly id: number;
+      readonly type: "rank";
+      readonly ok: true;
+      readonly items: readonly AutocompleteItem[] | undefined;
+    }
+  | {
+      readonly id: number;
+      readonly type: "rank";
+      readonly ok: false;
+      readonly error: string;
+    };

package/src/extensions/file-picker/index.ts

@@ -3,116 +3,171 @@ import type {
   ExtensionAPI,
 } from "@earendil-works/pi-coding-agent";
 import type { AutocompleteProvider } from "@earendil-works/pi-tui";
-import { type FileCandidate, loadRelative } from "./catalog";
-import { rank } from "./ranker";
+import type { FilePickerSuggestionEngine } from "./FilePickerSuggestionEngine";
+import { WorkerFilePickerSuggestionEngine } from "./WorkerFilePickerSuggestionEngine";
 
 const MAX_VISIBLE_ROWS = 50;
 const AT_PREFIX = /(?:^|\s)@(\S*)$/;
 
+// Pi cancels autocomplete after Tab; for directories we want to keep
+// drilling, so re-enter Tab on the next tick.
+function keepDrilling(): void {
+  setTimeout(() => {
+    try {
+      process.stdin.emit("data", "\t");
+    } catch {}
+  }, 0);
+}
+
+function activeAtTokenFromMatch(
+  match: RegExpMatchArray,
+  cursorLine: number
+): ActiveAtToken {
+  const matchedText = match[0] ?? "";
+  const matchCol = match.index ?? 0;
+  const atCol = matchCol + (matchedText.startsWith("@") ? 0 : 1);
+  return { cursorLine, atCol };
+}
+
+function sameActiveAtToken(
+  a: ActiveAtToken | undefined,
+  b: ActiveAtToken
+): boolean {
+  return a?.cursorLine === b.cursorLine && a.atCol === b.atCol;
+}
+
 export type FilePickerProviderFactoryOptions = {
-  readonly loadRelativeCatalog: () => Promise<readonly FileCandidate[]>;
+  readonly engine: FilePickerSuggestionEngine;
+};
+
+type ActiveAtToken = {
+  readonly cursorLine: number;
+  readonly atCol: number;
 };
 
 export function createFilePickerProviderFactory(
   options: FilePickerProviderFactoryOptions
 ): AutocompleteProviderFactory {
-  let cachedRelative: readonly FileCandidate[] | undefined;
-  let relativeRefresh: Promise<void> | undefined;
-
   const refreshRelative = (): void => {
-    relativeRefresh ??= options
-      .loadRelativeCatalog()
-      .then((catalog) => {
-        cachedRelative = catalog;
-      })
-      .catch(() => {
-        if (cachedRelative === undefined) {
-          cachedRelative = [];
-        }
-      })
-      .finally(() => {
-        relativeRefresh = undefined;
-      });
+    void options.engine.refreshRelative();
   };
 
-  refreshRelative();
+  return (current: AutocompleteProvider): AutocompleteProvider => {
+    let activeAtToken: ActiveAtToken | undefined;
 
-  return (current: AutocompleteProvider): AutocompleteProvider => ({
-    async getSuggestions(lines, cursorLine, cursorCol, autocompleteOptions) {
-      const line = lines[cursorLine] ?? "";
-      const beforeCursor = line.slice(0, cursorCol);
+    return {
+      async getSuggestions(lines, cursorLine, cursorCol, autocompleteOptions) {
+        const line = lines[cursorLine] ?? "";
+        const beforeCursor = line.slice(0, cursorCol);
 
-      const atMatch = beforeCursor.match(AT_PREFIX);
-      if (!atMatch) {
-        return current.getSuggestions(
-          lines,
-          cursorLine,
-          cursorCol,
-          autocompleteOptions
-        );
-      }
+        const atMatch = beforeCursor.match(AT_PREFIX);
+        if (!atMatch) {
+          activeAtToken = undefined;
+          return current.getSuggestions(
+            lines,
+            cursorLine,
+            cursorCol,
+            autocompleteOptions
+          );
+        }
+
+        const query = atMatch[1] ?? "";
+        const atToken = activeAtTokenFromMatch(atMatch, cursorLine);
+        if (!sameActiveAtToken(activeAtToken, atToken)) {
+          activeAtToken = atToken;
+          refreshRelative();
+        }
+
+        const items = await options.engine
+          .rank(query, {
+            limit: MAX_VISIBLE_ROWS,
+            signal: autocompleteOptions.signal,
+          })
+          .catch(() => undefined);
+        if (items === undefined) {
+          return current.getSuggestions(
+            lines,
+            cursorLine,
+            cursorCol,
+            autocompleteOptions
+          );
+        }
+        if (items.length === 0) {
+          return null;
+        }
+        return {
+          items: items.map((item) => ({
+            ...item,
+            value: `@${item.value}`,
+          })),
+          prefix: `@${query}`,
+        };
+      },
+
+      applyCompletion(lines, cursorLine, cursorCol, item, prefix) {
+        // Pi appends a trailing space after file completions; apply @ items
+        // ourselves so Tab inserts the bare path.
+        if (prefix.startsWith("@")) {
+          const line = lines[cursorLine] ?? "";
+          const beforePrefix = line.slice(0, cursorCol - prefix.length);
+          const afterCursor = line.slice(cursorCol);
+          const hasTrailingQuote = item.value.endsWith('"');
+          const adjustedAfterCursor =
+            prefix.startsWith('@"') &&
+            hasTrailingQuote &&
+            afterCursor.startsWith('"')
+              ? afterCursor.slice(1)
+              : afterCursor;
+
+          const newLines = [...lines];
+          newLines[cursorLine] =
+            `${beforePrefix}${item.value}${adjustedAfterCursor}`;
 
-      const query = atMatch[1] ?? "";
-      refreshRelative();
+          const isDirectory = item.label.endsWith("/");
+          const cursorOffset =
+            isDirectory && hasTrailingQuote
+              ? item.value.length - 1
+              : item.value.length;
 
-      const items = await rank(query, {
-        cachedRelative,
-        limit: MAX_VISIBLE_ROWS,
-      });
-      if (items === undefined) {
-        return current.getSuggestions(
+          if (isDirectory) {
+            keepDrilling();
+          }
+
+          return {
+            lines: newLines,
+            cursorLine,
+            cursorCol: beforePrefix.length + cursorOffset,
+          };
+        }
+
+        const result = current.applyCompletion(
           lines,
           cursorLine,
           cursorCol,
-          autocompleteOptions
+          item,
+          prefix
         );
-      }
-      if (items.length === 0) {
-        return null;
-      }
-      return {
-        items: items.map((item) => ({
-          ...item,
-          value: `@${item.value}`,
-        })),
-        prefix: `@${query}`,
-      };
-    },
-
-    applyCompletion(lines, cursorLine, cursorCol, item, prefix) {
-      const result = current.applyCompletion(
-        lines,
-        cursorLine,
-        cursorCol,
-        item,
-        prefix
-      );
-      // Pi cancels autocomplete after Tab; for directories we want to keep
-      // drilling, so re-enter Tab on the next tick.
-      if (typeof item.value === "string" && item.value.endsWith("/")) {
-        setTimeout(() => {
-          try {
-            process.stdin.emit("data", "\t");
-          } catch {}
-        }, 0);
-      }
-      return result;
-    },
-
-    shouldTriggerFileCompletion(lines, cursorLine, cursorCol) {
-      return (
-        current.shouldTriggerFileCompletion?.(lines, cursorLine, cursorCol) ??
-        true
-      );
-    },
-  });
+        if (item.value.endsWith("/")) {
+          keepDrilling();
+        }
+        return result;
+      },
+
+      shouldTriggerFileCompletion(lines, cursorLine, cursorCol) {
+        return (
+          current.shouldTriggerFileCompletion?.(lines, cursorLine, cursorCol) ??
+          true
+        );
+      },
+    };
+  };
 }
 
 export default function (pi: ExtensionAPI): void {
   pi.on("session_start", (_event, ctx) => {
     ctx.ui.addAutocompleteProvider(
       createFilePickerProviderFactory({
-        loadRelativeCatalog: () => loadRelative({ root: ctx.cwd }),
+        engine: new WorkerFilePickerSuggestionEngine(ctx.cwd),
       })
     );
   });

package/src/extensions/file-picker/ranker.ts

@@ -14,31 +14,175 @@ export type FileRankOptions = {
 
 const isAbsoluteQuery = (query: string): boolean =>
   query.startsWith("/") || query.startsWith("~");
+const GLOBAL_FUZZY_QUERY_MIN_LENGTH = 3;
 
 let cachedIndex:
   | {
       readonly source: readonly FileCandidate[];
-      readonly index: FuzzyIndex<FileCandidate>;
+      readonly index: RelativeRankingIndex;
     }
   | undefined;
 
 const indexFor = (
   candidates: readonly FileCandidate[]
-): FuzzyIndex<FileCandidate> => {
+): RelativeRankingIndex => {
   if (cachedIndex?.source === candidates) {
     return cachedIndex.index;
   }
-  const fuzzy: FuzzyCandidate<FileCandidate>[] = candidates.map(
-    (candidate) => ({
-      item: candidate,
-      haystacks: [candidate.matchHaystack],
-    })
-  );
-  const index = FuzzyMatcher.prepare(fuzzy);
+
+  const index = new RelativeRankingIndex(candidates);
   cachedIndex = { source: candidates, index };
   return index;
 };
 
+type LoweredCandidate = {
+  readonly candidate: FileCandidate;
+  readonly nameLower: string;
+  readonly haystackLower: string;
+};
+
+class RelativeRankingIndex {
+  private readonly childrenByDirectory = new Map<string, FileCandidate[]>();
+  private readonly loweredSource: readonly LoweredCandidate[];
+
+  private readonly scopedIndexes = new Map<string, FuzzyIndex<FileCandidate>>();
+  private globalIndex: FuzzyIndex<FileCandidate> | undefined;
+
+  public constructor(private readonly source: readonly FileCandidate[]) {
+    const lowered: LoweredCandidate[] = [];
+
+    for (const candidate of source) {
+      const slash = candidate.insertPath.lastIndexOf("/");
+      const directory =
+        slash === -1 ? "" : candidate.insertPath.slice(0, slash);
+      const children = this.childrenByDirectory.get(directory) ?? [];
+      children.push(candidate);
+      this.childrenByDirectory.set(directory, children);
+
+      lowered.push({
+        candidate,
+        nameLower: basename(candidate.insertPath).toLocaleLowerCase(),
+        haystackLower: candidate.matchHaystack.toLocaleLowerCase(),
+      });
+    }
+
+    this.loweredSource = lowered;
+  }
+
+  public rank(query: string, limit: number | undefined): AutocompleteItem[] {
+    const scoped = this.scopedCandidates(query);
+    if (scoped !== undefined) {
+      return rankCandidates(scoped.candidates, scoped.residualQuery, limit, {
+        index: () => this.indexForScope(scoped.directory, scoped.candidates),
+      });
+    }
+
+    const literalHits = this.literalRank(query, limit);
+    if (literalHits !== undefined) {
+      return literalHits;
+    }
+
+    return rankCandidates(this.source, query, limit, {
+      index: () => this.indexForGlobal(),
+    });
+  }
+
+  private scopedCandidates(query: string):
+    | {
+        readonly directory: string;
+        readonly residualQuery: string;
+        readonly candidates: readonly FileCandidate[];
+      }
+    | undefined {
+    const slash = query.lastIndexOf("/");
+    if (slash === -1) {
+      return undefined;
+    }
+
+    const directory = query.slice(0, slash);
+    const candidates = this.childrenByDirectory.get(directory);
+    if (candidates === undefined) {
+      return undefined;
+    }
+
+    return {
+      directory,
+      residualQuery: query.slice(slash + 1),
+      candidates,
+    };
+  }
+
+  private indexForScope(
+    directory: string,
+    candidates: readonly FileCandidate[]
+  ): FuzzyIndex<FileCandidate> {
+    const cached = this.scopedIndexes.get(directory);
+    if (cached !== undefined) {
+      return cached;
+    }
+
+    const index = prepareIndex(candidates, (candidate) =>
+      basename(candidate.insertPath)
+    );
+    this.scopedIndexes.set(directory, index);
+    return index;
+  }
+
+  private indexForGlobal(): FuzzyIndex<FileCandidate> {
+    this.globalIndex ??= prepareIndex(
+      this.source,
+      (candidate) => candidate.matchHaystack
+    );
+    return this.globalIndex;
+  }
+
+  private literalRank(
+    query: string,
+    limit: number | undefined
+  ): AutocompleteItem[] | undefined {
+    const needle = query.trim().toLocaleLowerCase();
+    if (needle.length === 0 || query.includes("/")) {
+      return undefined;
+    }
+
+    const prefixHits: FileCandidate[] = [];
+    const substringHits: FileCandidate[] = [];
+    const limitSize = limit ?? Infinity;
+
+    for (const { candidate, nameLower, haystackLower } of this.loweredSource) {
+      if (nameLower.startsWith(needle) || haystackLower.startsWith(needle)) {
+        prefixHits.push(candidate);
+      } else if (nameLower.includes(needle) || haystackLower.includes(needle)) {
+        substringHits.push(candidate);
+      }
+
+      // Prefix hits always outrank substring hits, so we can stop only once we
+      // have enough prefixes to fill the limit; otherwise a late-sorting prefix
+      // could be dropped for an earlier substring match.
+      if (prefixHits.length >= limitSize) {
+        break;
+      }
+    }
+
+    const hitCount = prefixHits.length + substringHits.length;
+    if (hitCount === 0) {
+      return undefined;
+    }
+
+    if (
+      prefixHits.length === 0 &&
+      needle.length >= GLOBAL_FUZZY_QUERY_MIN_LENGTH &&
+      hitCount < limitSize
+    ) {
+      return undefined;
+    }
+
+    return [...prefixHits, ...substringHits]
+      .slice(0, limit)
+      .map((candidate) => toItem(candidate));
+  }
+}
+
 export async function rank(
   query: string,
   options: FileRankOptions
@@ -52,18 +196,42 @@ export async function rank(
     return undefined;
   }
 
-  return rankCandidates(options.cachedRelative, query, options.limit);
+  return indexFor(options.cachedRelative).rank(query, options.limit);
 }
 
+type RankCandidatesOptions = {
+  readonly index?: () => FuzzyIndex<FileCandidate>;
+};
+
 const rankCandidates = (
   candidates: readonly FileCandidate[],
   query: string,
-  limit: number | undefined
+  limit: number | undefined,
+  options: RankCandidatesOptions = {}
 ): AutocompleteItem[] => {
-  const hits = indexFor(candidates).find(query, { limit });
+  if (query.trim().length === 0) {
+    return candidates.slice(0, limit).map((candidate) => toItem(candidate));
+  }
+
+  const index = options.index ?? (() => prepareIndex(candidates));
+  const hits = index().find(query, { limit });
   return hits.map((hit) => toItem(hit.item));
 };
 
+const prepareIndex = (
+  candidates: readonly FileCandidate[],
+  haystack: (candidate: FileCandidate) => string = (candidate) =>
+    candidate.matchHaystack
+): FuzzyIndex<FileCandidate> => {
+  const fuzzy: FuzzyCandidate<FileCandidate>[] = candidates.map(
+    (candidate) => ({
+      item: candidate,
+      haystacks: [haystack(candidate)],
+    })
+  );
+  return FuzzyMatcher.prepare(fuzzy);
+};
+
 const toItem = (candidate: FileCandidate): AutocompleteItem => {
   const suffix = candidate.isDirectory ? "/" : "";
   const value = `${candidate.insertPath}${suffix}`;

package/src/extensions/grep/grep.ts

@@ -2,7 +2,7 @@ import { FileScanner, type FileScanOptions } from "../../shared/FileScanner";
 import { FsErrors } from "../../shared/FsErrors";
 import { Lines } from "../../shared/Lines";
 
-const MATCH_CONCURRENCY = 32;
+const MATCH_CONCURRENCY = 16;
 
 export type GrepLine = {
   readonly lineNumber: number;
@@ -25,10 +25,33 @@ export type GrepMatch = {
 export type GrepMatcher = {
   readonly regex: RegExp;
   readonly matchAcrossLines: boolean;
+  /**
+   * Raw-byte needle for the literal fast path: present only when the pattern is
+   * a pure literal, case-sensitive, and single-line, so `matchFile` can reject a
+   * non-matching file with `Buffer.indexOf` before decoding it. Undefined
+   * otherwise, in which case the regex path runs unchanged.
+   */
+  readonly literal: Buffer | undefined;
 };
 
 export type GrepScanOptions = FileScanOptions;
 
+// Characters that stand for themselves in both a default-flag regex and raw
+// UTF-8 bytes (all ASCII, so they never alias a multibyte sequence). A pattern
+// made only of these is a literal we can scan on bytes.
+const PURE_LITERAL = /^[A-Za-z0-9_ \-/]+$/;
+
+function literalNeedle(
+  pattern: string,
+  caseInsensitive: boolean,
+  matchAcrossLines: boolean
+): Buffer | undefined {
+  if (caseInsensitive || matchAcrossLines || !PURE_LITERAL.test(pattern)) {
+    return undefined;
+  }
+  return Buffer.from(pattern, "utf8");
+}
+
 export function buildMatcher(options: {
   readonly pattern: string;
   readonly caseInsensitive: boolean;
@@ -42,6 +65,11 @@ export function buildMatcher(options: {
     return {
       regex: new RegExp(options.pattern, flags),
       matchAcrossLines: options.matchAcrossLines,
+      literal: literalNeedle(
+        options.pattern,
+        options.caseInsensitive,
+        options.matchAcrossLines
+      ),
     };
   } catch (error) {
     const message = error instanceof Error ? error.message : String(error);
@@ -87,11 +115,26 @@ async function matchFile(
 ): Promise<GrepMatch | undefined> {
   const file = Bun.file(filePath);
 
+  // Binary skip reads only the first 8KB, so a binary file is never fully read.
   if (await Lines.isBinary(file)) {
     return undefined;
   }
 
-  const content = Lines.normalize(await file.text());
+  let text: string;
+  if (matcher.literal !== undefined) {
+    // Literal fast path: scan raw bytes and bail on a miss without decoding. An
+    // ASCII literal can't match across a normalized newline or alias a
+    // multibyte char, so a raw-byte hit/miss matches the decoded result.
+    const bytes = Buffer.from(await file.arrayBuffer());
+    if (bytes.indexOf(matcher.literal) < 0) {
+      return undefined;
+    }
+    text = bytes.toString("utf8");
+  } else {
+    text = await file.text();
+  }
+
+  const content = Lines.normalize(text);
   const fileLines = Lines.split(content);
   const ranges = matcher.matchAcrossLines
     ? regexRanges(content, matcher.regex)

package/src/extensions/web-search/ExaMcpClient.ts

@@ -1,9 +1,11 @@
 import { McpClient, type McpFetch } from "../../shared/McpClient";
+import { RateLimiter } from "../../shared/RateLimiter";
 
 type ExaMcpClientOptions = {
   readonly endpoint?: string;
   readonly apiKey?: string;
   readonly fetch?: McpFetch;
+  readonly rateLimiter?: RateLimiter;
 };
 
 type ExaSearchInput = {
@@ -28,15 +30,30 @@ class ExaSearchError extends Error {
 export class ExaMcpClient {
   private static readonly defaultEndpoint = "https://mcp.exa.ai/mcp";
   private static readonly toolName = "web_search_exa";
+  private static readonly maxRequestsPerWindow = 3;
+  private static readonly windowMs = 1000;
 
   private readonly client: McpClient;
 
   public constructor(options: ExaMcpClientOptions = {}) {
+    const apiKey =
+      options.apiKey === undefined || options.apiKey.length === 0
+        ? undefined
+        : options.apiKey;
+    // Throttle only on the free tier; an API key lifts the request rate limit.
+    const rateLimiter =
+      apiKey !== undefined
+        ? undefined
+        : (options.rateLimiter ??
+          new RateLimiter({
+            maxRequests: ExaMcpClient.maxRequestsPerWindow,
+            windowMs: ExaMcpClient.windowMs,
+          }));
+
     this.client = new McpClient({
       endpoint: options.endpoint ?? ExaMcpClient.defaultEndpoint,
-      ...(options.apiKey === undefined || options.apiKey.length === 0
-        ? {}
-        : { headers: { "x-api-key": options.apiKey } }),
+      ...(apiKey === undefined ? {} : { headers: { "x-api-key": apiKey } }),
+      ...(rateLimiter === undefined ? {} : { rateLimiter }),
       ...(options.fetch === undefined ? {} : { fetch: options.fetch }),
     });
   }