From 0231fb388694c57ddce83565c7f2262f4db5aa75 Mon Sep 17 00:00:00 2001 From: no value <300561837+trevorleibert-mixpanel@users.noreply.github.com> Date: Tue, 18 Aug 2026 21:07:49 +0000 Subject: [PATCH] fix(pi-fff): share finders across in-process sessions --- packages/fff-bun/test/multi-session.test.ts | 15 +++ packages/pi-fff/src/aux-finders.ts | 13 +- packages/pi-fff/src/file-picker.ts | 137 +++++++++++++++----- packages/pi-fff/src/index.ts | 20 ++- packages/pi-fff/test/aux-dedup.test.ts | 35 +++++ packages/pi-fff/test/aux-pool.test.ts | 13 +- packages/pi-fff/test/extension.test.ts | 65 +++++++++- 7 files changed, 254 insertions(+), 44 deletions(-) diff --git a/packages/fff-bun/test/multi-session.test.ts b/packages/fff-bun/test/multi-session.test.ts index 35bd7f3d6..555d25ce1 100644 --- a/packages/fff-bun/test/multi-session.test.ts +++ b/packages/fff-bun/test/multi-session.test.ts @@ -198,6 +198,21 @@ function startSession( } describe("pi-fff: in-process double activation works (#760)", () => { + test("same-workspace sessions share a finder across one shutdown", async () => { + const dbs = makeDbPaths(); + const workspace = makeWorkspace("shared-session"); + const first = startSession(workspace, dbs); + const second = startSession(workspace, dbs); + + await Promise.all([first.start(), second.start()]); + expect(first.errors()).toEqual([]); + expect(second.errors()).toEqual([]); + + await first.shutdown(); + expect(await second.find("gamma")).toContain("gamma.ts"); + await second.shutdown(); + }, 40_000); + test("two sessions in one process both search against the same dbs", async () => { const dbs = makeDbPaths(); const first = startSession(makeWorkspace("session1"), dbs); diff --git a/packages/pi-fff/src/aux-finders.ts b/packages/pi-fff/src/aux-finders.ts index 45f4170b4..2ae9ec001 100644 --- a/packages/pi-fff/src/aux-finders.ts +++ b/packages/pi-fff/src/aux-finders.ts @@ -23,6 +23,7 @@ export interface AuxOpts { export class AuxFinderPool { private entries: AuxPicker[] = []; + private destroyed = false; // In-flight creations keyed by root. Concurrent acquire() calls for the same // (or a covering) root share one finder/scan instead of each starting a full // duplicate traversal — issue #746. Mirrors the main finder's finderPromise. @@ -30,8 +31,9 @@ export class AuxFinderPool { constructor(private opts: AuxOpts) {} destroy(): void { + this.destroyed = true; for (const e of this.entries) { - e.finder.destroy(); + this.opts.pickers.release(e.finder); } this.entries = []; @@ -42,7 +44,7 @@ export class AuxFinderPool { const kept: AuxPicker[] = []; for (const e of this.entries) { if (now - e.lastUsed > IDLE_TTL_MS) { - if (!e.finder.isDestroyed) e.finder.destroy(); + this.opts.pickers.release(e.finder); } else { kept.push(e); } @@ -54,6 +56,7 @@ export class AuxFinderPool { maybeRoot: string, opts?: { exact?: boolean }, ): Promise<{ finder: FileFinderApi; root: string }> { + if (this.destroyed) throw new Error("FFF auxiliary finder pool is destroyed"); this.sweepIdle(); let covering: AuxPicker | null = null; for (const e of this.entries) { @@ -89,7 +92,7 @@ export class AuxFinderPool { if (this.entries.length >= MAX_AUX) { let oldest = this.entries[0]; for (const e of this.entries) if (e.lastUsed < oldest.lastUsed) oldest = e; - if (!oldest.finder.isDestroyed) oldest.finder.destroy(); + this.opts.pickers.release(oldest.finder); this.entries = this.entries.filter((e) => e !== oldest); } @@ -105,6 +108,10 @@ export class AuxFinderPool { enableHomeDirScanning, enableFsRootScanning: this.opts.enableFsRootScanning, }); + if (this.destroyed) { + this.opts.pickers.release(finder); + throw new Error("FFF auxiliary finder pool was destroyed during initialization"); + } const entry: AuxPicker = { root, finder, lastUsed: Date.now() }; this.entries.push(entry); diff --git a/packages/pi-fff/src/file-picker.ts b/packages/pi-fff/src/file-picker.ts index 1be6681dd..1caf2148a 100644 --- a/packages/pi-fff/src/file-picker.ts +++ b/packages/pi-fff/src/file-picker.ts @@ -1,3 +1,4 @@ +import path from "node:path"; import type { FileFinderApi, InitOptions, Result } from "@ff-labs/fff-node"; import { type FileFinderStatic, loadSdk, SCAN_TIMEOUT_MS } from "./sdk"; @@ -7,13 +8,46 @@ export interface PickerOptions { enableFsRootScanning?: boolean; } +interface SharedFinder { + finder: FileFinderApi; + refs: number; +} + +const SHARED_FINDERS = Symbol.for("@ff-labs/pi-fff:shared-finders"); + +function sharedFinders(): Map { + const global = globalThis as typeof globalThis & { + [SHARED_FINDERS]?: Map; + }; + return (global[SHARED_FINDERS] ??= new Map()); +} + +function finderKey(options: InitOptions): string { + return JSON.stringify({ + basePath: path.resolve(options.basePath), + frecencyDbPath: + options.frecencyDbPath !== undefined + ? path.resolve(options.frecencyDbPath) + : undefined, + historyDbPath: + options.historyDbPath !== undefined + ? path.resolve(options.historyDbPath) + : undefined, + enableHomeDirScanning: options.enableHomeDirScanning ?? false, + enableFsRootScanning: options.enableFsRootScanning ?? false, + aiMode: true, + }); +} + /** Opens every picker in this pi process — the cwd picker and the aux pickers — - * on the same frecency/history databases. */ + * on the same frecency/history databases. Identical pickers are shared across + * in-process pi sessions, including retained subagent sessions. */ export class FilePickerFactory { private dbDisabled = false; private readonly frecencyDbPath: string; private readonly historyDbPath: string; private readonly onDbFailure?: (error: string) => void; + private readonly owned = new Map(); constructor(opts: { frecencyDbPath: string; @@ -33,41 +67,84 @@ export class FilePickerFactory { /** Opens a scanned, ready-to-use picker. Throws if it cannot be created. */ async create(options: PickerOptions): Promise { const { FileFinder } = await loadSdk(); - const result = this.openWithDbFallback(FileFinder, options); + const init: InitOptions = { ...options, aiMode: true }; + + if (!this.dbDisabled) { + const withDbs = { + ...init, + frecencyDbPath: this.frecencyDbPath, + historyDbPath: this.historyDbPath, + }; + const result = this.acquire(FileFinder, withDbs); + if (result.ok) return this.waitForScan(result.value); - if (!result.ok) { - throw new Error( - `Failed to create FFF file picker for ${options.basePath}: ${result.error}`, - ); + const dbLess = this.acquire(FileFinder, init); + if (!dbLess.ok) { + throw this.createError(options.basePath, result.error); + } + this.dbDisabled = true; + this.onDbFailure?.(result.error); + return this.waitForScan(dbLess.value); } - // waitForScan() also resolves on timeout, so this bounds startup rather - // than guaranteeing a complete index. - await result.value.waitForScan(SCAN_TIMEOUT_MS); - return result.value; + const result = this.acquire(FileFinder, init); + if (!result.ok) throw this.createError(options.basePath, result.error); + return this.waitForScan(result.value); + } + + /** Releases this factory's ownership without disrupting other sessions. */ + release(finder: FileFinderApi): void { + const ownership = this.owned.get(finder); + if (!ownership) return; + + if (--ownership.refs === 0) this.owned.delete(finder); + + const shared = sharedFinders().get(ownership.key); + if (!shared || shared.finder !== finder || --shared.refs > 0) return; + sharedFinders().delete(ownership.key); + if (!finder.isDestroyed) finder.destroy(); } - private openWithDbFallback( + private acquire( FileFinder: FileFinderStatic, - options: PickerOptions, + options: InitOptions, ): Result { - const init: InitOptions = { ...options, aiMode: true }; - if (this.dbDisabled) return FileFinder.create(init); - - const result = FileFinder.create({ - ...init, - frecencyDbPath: this.frecencyDbPath, - historyDbPath: this.historyDbPath, - }); - if (result.ok) return result; - - // A failure here is usually transient (broken lock, corruption) and self-heals - // on restart, so drop the databases instead of leaving pi without a picker - const dbLess = FileFinder.create(init); - if (!dbLess.ok) return result; // db error is the more useful one to report - - this.dbDisabled = true; - this.onDbFailure?.(result.error); - return dbLess; + const key = finderKey(options); + const existing = sharedFinders().get(key); + if (existing && !existing.finder.isDestroyed) { + existing.refs++; + this.own(existing.finder, key); + return { ok: true, value: existing.finder }; + } + if (existing) sharedFinders().delete(key); + + const result = FileFinder.create(options); + if (!result.ok) return result; + + sharedFinders().set(key, { finder: result.value, refs: 1 }); + this.own(result.value, key); + return result; + } + + private own(finder: FileFinderApi, key: string): void { + const existing = this.owned.get(finder); + if (existing) existing.refs++; + else this.owned.set(finder, { key, refs: 1 }); + } + + private async waitForScan(finder: FileFinderApi): Promise { + try { + // waitForScan() also resolves on timeout, so this bounds startup rather + // than guaranteeing a complete index. + await finder.waitForScan(SCAN_TIMEOUT_MS); + return finder; + } catch (error) { + this.release(finder); + throw error; + } + } + + private createError(basePath: string, error: string): Error { + return new Error(`Failed to create FFF file picker for ${basePath}: ${error}`); } } diff --git a/packages/pi-fff/src/index.ts b/packages/pi-fff/src/index.ts index 66578b564..ea4ae16c7 100644 --- a/packages/pi-fff/src/index.ts +++ b/packages/pi-fff/src/index.ts @@ -443,22 +443,28 @@ export default function fffExtension(pi: ExtensionAPI) { if (finderPromise) return finderPromise; finderPromise = (async () => { - if (mainFinder && !mainFinder.isDestroyed) { - mainFinder.destroy(); + if (mainFinder) { + pickers?.release(mainFinder); mainFinder = null; finderCwd = null; } // if the dbs can't be opened the factory falls back to a db-less picker, // e.g. when some other process corrupts the lock - if (!pickers) throw new Error("FFF picker factory is not initialized"); - mainFinder = await pickers.create({ + const factory = pickers; + if (!factory) throw new Error("FFF picker factory is not initialized"); + const finder = await factory.create({ basePath: cwd, enableHomeDirScanning, enableFsRootScanning, }); + if (pickers !== factory) { + factory.release(finder); + throw new Error("FFF session shut down during initialization"); + } + mainFinder = finder; finderCwd = cwd; - return mainFinder; + return finder; })().finally(() => { finderPromise = null; }); @@ -500,8 +506,8 @@ export default function fffExtension(pi: ExtensionAPI) { function destroyFinder() { stopHomeScanStatus(); - if (mainFinder && !mainFinder.isDestroyed) { - mainFinder.destroy(); + if (mainFinder) { + pickers?.release(mainFinder); mainFinder = null; finderCwd = null; } diff --git a/packages/pi-fff/test/aux-dedup.test.ts b/packages/pi-fff/test/aux-dedup.test.ts index 3a6a2b7ac..b6150f508 100644 --- a/packages/pi-fff/test/aux-dedup.test.ts +++ b/packages/pi-fff/test/aux-dedup.test.ts @@ -45,6 +45,25 @@ function makePickers() { }); } +describe("process-wide finder sharing", () => { + test("one factory releases each acquisition independently", async () => { + created.length = 0; + const pickers = makePickers(); + const [first, second] = await Promise.all([ + pickers.create({ basePath: "/workspace" }), + pickers.create({ basePath: "/workspace" }), + ]); + + expect(created).toHaveLength(1); + expect(first).toBe(second); + + pickers.release(first); + expect(created[0].isDestroyed).toBe(false); + pickers.release(second); + expect(created[0].isDestroyed).toBe(true); + }); +}); + describe("AuxFinderPool concurrent dedup (#746)", () => { test("two concurrent acquires for same root share one finder", async () => { created.length = 0; @@ -58,6 +77,21 @@ describe("AuxFinderPool concurrent dedup (#746)", () => { ]); expect(created.length).toBe(1); expect(a.finder).toBe(b.finder); + pool.destroy(); + }); + + test("destroying a pool releases a finder whose scan is still starting", async () => { + created.length = 0; + const pool = new AuxFinderPool({ + enableFsRootScanning: false, + pickers: makePickers(), + }); + const pending = pool.acquire("/Users/pending"); + + pool.destroy(); + expect(pending).rejects.toThrow("destroyed during initialization"); + await pending.catch(() => undefined); + expect(created[0].isDestroyed).toBe(true); }); test("sequential acquire after in-flight one resolves still reuses", async () => { @@ -72,5 +106,6 @@ describe("AuxFinderPool concurrent dedup (#746)", () => { const third = await pool.acquire("/Users/x"); expect(created.length).toBe(1); expect(third.root).toBe("/Users/x"); + pool.destroy(); }); }); diff --git a/packages/pi-fff/test/aux-pool.test.ts b/packages/pi-fff/test/aux-pool.test.ts index 86aa56130..37a3520fb 100644 --- a/packages/pi-fff/test/aux-pool.test.ts +++ b/packages/pi-fff/test/aux-pool.test.ts @@ -1,4 +1,4 @@ -import { describe, expect, mock, test } from "bun:test"; +import { afterEach, describe, expect, mock, test } from "bun:test"; import os from "node:os"; import path from "node:path"; @@ -52,16 +52,20 @@ mock.module("@ff-labs/fff-bun", () => finderModule); const { AuxFinderPool } = await import("../src/aux-finders"); const { FilePickerFactory } = await import("../src/file-picker"); +const activePools: Array> = []; + function makePool(opts: Record = {}) { created.length = 0; createOptions.length = 0; failDbCreates = false; failAllCreates = false; - return new AuxFinderPool({ + const pool = new AuxFinderPool({ enableFsRootScanning: false, pickers: makePickers(), ...opts, }); + activePools.push(pool); + return pool; } function makePickers(onDbFailure?: (error: string) => void) { @@ -72,6 +76,10 @@ function makePickers(onDbFailure?: (error: string) => void) { }); } +afterEach(() => { + for (const pool of activePools.splice(0)) pool.destroy(); +}); + describe("AuxFinderPool covering reuse", () => { test("reuses a picker rooted at an ancestor of the requested path", async () => { const pool = makePool(); @@ -192,6 +200,7 @@ describe("AuxFinderPool covering reuse", () => { expect(createOptions[0].frecencyDbPath).toBeUndefined(); expect(createOptions[0].historyDbPath).toBeUndefined(); expect(failures).toEqual(["db locked"]); + pickers.release(main as any); }); test("create throws when the picker cannot be opened at all", async () => { diff --git a/packages/pi-fff/test/extension.test.ts b/packages/pi-fff/test/extension.test.ts index 4021c3c9d..c4199d835 100644 --- a/packages/pi-fff/test/extension.test.ts +++ b/packages/pi-fff/test/extension.test.ts @@ -1,4 +1,4 @@ -import { afterAll, beforeEach, describe, expect, mock, test } from "bun:test"; +import { afterAll, afterEach, beforeEach, describe, expect, mock, test } from "bun:test"; import fs from "node:fs"; import os from "node:os"; import path from "node:path"; @@ -15,11 +15,12 @@ const createCalls: unknown[] = []; let finders: MockFinder[] = []; let mixedSearchImpl: ((query: string, options: unknown) => unknown) | undefined; let scanProgressImpl: (() => unknown) | undefined; +let waitForScanImpl: (() => Promise) | undefined; function createMockFinder(): MockFinder { return { isDestroyed: false, - waitForScan: mock(async () => undefined), + waitForScan: mock(async () => waitForScanImpl?.()), getScanProgress: mock(() => { if (scanProgressImpl) return scanProgressImpl(); return { @@ -153,6 +154,8 @@ function createContext(cwd = "/tmp/workspace") { }; } +const activeSetups: Array<{ events: Map }> = []; + async function start(mode?: string, cwd?: string, flags: Record = {}) { const setup = createPi(mode, flags); const ctx = createContext(cwd); @@ -161,6 +164,7 @@ async function start(mode?: string, cwd?: string, flags: Record const sessionStart = setup.events.get("session_start"); expect(sessionStart).toBeDefined(); await sessionStart?.({ reason: "startup" }, ctx); + activeSetups.push(setup); return { ...setup, ctx }; } @@ -203,12 +207,17 @@ beforeEach(() => { finders = []; mixedSearchImpl = undefined; scanProgressImpl = undefined; + waitForScanImpl = undefined; for (const key of CONFIG_ENV_KEYS) delete process.env[key]; process.env.PI_CODING_AGENT_DIR = agentDir; fs.rmSync(configPath, { force: true }); }); +afterEach(async () => { + await Promise.all(activeSetups.splice(0).map(shutdown)); +}); + afterAll(() => { for (const key of CONFIG_ENV_KEYS) { const value = savedEnv[key]; @@ -296,6 +305,57 @@ function writeConfig(config: Record): void { fs.writeFileSync(configPath, JSON.stringify(config)); } +describe("pi-fff shared finders", () => { + test("same-workspace subagent sessions share one finder", async () => { + const sessions = await Promise.all( + Array.from({ length: 42 }, () => start(undefined, "/tmp/shared-workspace")), + ); + + expect(createCalls).toHaveLength(1); + expect(finders).toHaveLength(1); + + await Promise.all(sessions.slice(0, -1).map(shutdown)); + expect(finders[0].isDestroyed).toBe(false); + + await shutdown(sessions.at(-1)!); + expect(finders[0].destroy).toHaveBeenCalledTimes(1); + }); + + test("shutdown releases a finder whose scan is still starting", async () => { + let finishScan!: () => void; + waitForScanImpl = () => + new Promise((resolve) => { + finishScan = resolve; + }); + + const setup = createPi(); + const ctx = createContext("/tmp/pending-workspace"); + fffExtension(setup.pi as any); + const starting = setup.events.get("session_start")?.({ reason: "startup" }, ctx); + + while (createCalls.length === 0) await Promise.resolve(); + await shutdown(setup); + finishScan(); + await starting; + + expect(finders[0].destroy).toHaveBeenCalledTimes(1); + }); + + test("different workspaces keep separate finders", async () => { + const first = await start(undefined, "/tmp/workspace-one"); + const second = await start(undefined, "/tmp/workspace-two"); + + expect(createCalls).toHaveLength(2); + expect(finders).toHaveLength(2); + + await shutdown(first); + expect(finders[0].isDestroyed).toBe(true); + expect(finders[1].isDestroyed).toBe(false); + + await shutdown(second); + }); +}); + describe("pi-fff session mode", () => { test("registers tools only after restoring the saved mode", async () => { const setup = createPi("tools-and-ui"); @@ -470,6 +530,7 @@ describe("pi-fff autocomplete registration", () => { expect(ctx.ui.notify).not.toHaveBeenCalled(); expect(createCalls).toHaveLength(1); + await shutdown(setup); }); test("delegates non-@ completions to the current provider", async () => {