From 8d56ae172e1cfc246253941d28d3144ce0384964 Mon Sep 17 00:00:00 2001 From: RunMintOn <1639562902@qq.com> Date: Mon, 7 Sep 2026 22:49:10 +0800 Subject: [PATCH] fix(pi-fff): preserve FFF tool renderers across reload Register FFF-named tool definitions while the extension loads so restored ffgrep/fffind rows can resolve their renderers (Pi rebuilds history rows before session_start). session_start still restores mode and enables the final tool names; it also prunes pre-registered names the final mode did not select, since Pi activates tools on every registerTool call. Renderer titles bind to the registered name so historical ffgrep rows keep that name under override mode. --- packages/pi-fff/src/index.ts | 74 +++++++++++++++----- packages/pi-fff/test/extension.test.ts | 93 +++++++++++++++++++++----- 2 files changed, 135 insertions(+), 32 deletions(-) diff --git a/packages/pi-fff/src/index.ts b/packages/pi-fff/src/index.ts index 75d93cfe..c6b60c01 100644 --- a/packages/pi-fff/src/index.ts +++ b/packages/pi-fff/src/index.ts @@ -637,30 +637,71 @@ export default function fffExtension(pi: ExtensionAPI) { }; const pendingTools: (() => string)[] = []; + const registeredToolNames = new Set(); + // A renderer is attached to a concrete registered name. Keep that name outside + // row state so old fffgrep rows retain their title after mode changes to override. + const renderToolNames = new WeakMap(); let toolsRegistered = false; + function getRenderToolName(context: object, fallback: string): string { + return renderToolNames.get(context) ?? fallback; + } + + function registerTool( + resolveName: () => string, + definition: PendingToolDefinition, + ): string { + const resolvedName = resolveName(); + if (registeredToolNames.has(resolvedName)) return resolvedName; + + const { promptGuidelines, renderCall, ...tool } = definition; + pi.registerTool({ + ...tool, + name: resolvedName, + label: resolvedName, + promptGuidelines: promptGuidelines?.(toolNames), + renderCall: renderCall + ? (args, theme, context) => { + renderToolNames.set(context, resolvedName); + return renderCall(args, theme, context); + } + : undefined, + }); + registeredToolNames.add(resolvedName); + return resolvedName; + } + function queueTool( resolveName: () => string, definition: PendingToolDefinition, ): void { - pendingTools.push(() => { - const { promptGuidelines, ...tool } = definition; - const resolvedName = resolveName(); - pi.registerTool({ - ...tool, - name: resolvedName, - label: resolvedName, - promptGuidelines: promptGuidelines?.(toolNames), - }); - return resolvedName; - }); + pendingTools.push(() => registerTool(resolveName, definition)); + + // Pi restores historical tool rows before session_start. Register the + // FFF-named tools now so their renderers resolve. Pi activates every + // newly registered tool, so registerPendingTools prunes the names the + // final mode did not select. + registerTool(resolveName, definition); } function registerPendingTools(): void { if (toolsRegistered) return; const registeredNames = pendingTools.map((register) => register()); - pi.setActiveTools([...new Set([...pi.getActiveTools(), ...registeredNames])]); + const finalNames = new Set(registeredNames); + // Pi activates tools on every registerTool call, so the early FFF + // registrations above are active by now. Drop the ones the final mode + // did not select (override: ffgrep/fffind). Only names this extension + // registered are candidates, so builtin tools sharing a name are safe. + const staleNames = new Set( + [...registeredToolNames].filter((name) => !finalNames.has(name)), + ); + pi.setActiveTools([ + ...new Set([ + ...pi.getActiveTools().filter((name) => !staleNames.has(name)), + ...registeredNames, + ]), + ]); toolsRegistered = true; } @@ -997,7 +1038,7 @@ export default function fffExtension(pi: ExtensionAPI) { const pattern = args?.pattern ?? ""; const path = args?.path ?? "."; let content = - theme.fg("toolTitle", theme.bold(toolNames.grep)) + + theme.fg("toolTitle", theme.bold(getRenderToolName(context, toolNames.grep))) + " " + theme.fg("accent", `/${pattern}/`) + theme.fg("toolOutput", ` in ${path}`); @@ -1140,7 +1181,7 @@ export default function fffExtension(pi: ExtensionAPI) { const pattern = args?.pattern ?? ""; const path = args?.path ?? "."; let content = - theme.fg("toolTitle", theme.bold(toolNames.find)) + + theme.fg("toolTitle", theme.bold(getRenderToolName(context, toolNames.find))) + " " + theme.fg("accent", pattern) + theme.fg("toolOutput", ` in ${path}`); @@ -1244,7 +1285,10 @@ export default function fffExtension(pi: ExtensionAPI) { const patterns = args?.patterns ?? []; const constraints = args?.constraints; let content = - theme.fg("toolTitle", theme.bold(toolNames.multiGrep)) + + theme.fg( + "toolTitle", + theme.bold(getRenderToolName(context, toolNames.multiGrep)), + ) + " " + theme.fg("accent", patterns.map((p: string) => `"${p}"`).join(", ")); if (constraints) content += theme.fg("toolOutput", ` (${constraints})`); diff --git a/packages/pi-fff/test/extension.test.ts b/packages/pi-fff/test/extension.test.ts index e9d7fd58..24eb4cc5 100644 --- a/packages/pi-fff/test/extension.test.ts +++ b/packages/pi-fff/test/extension.test.ts @@ -251,9 +251,12 @@ describe("pi-fff global config", () => { const setup = await start(); const toolNames = setup.pi.registerTool.mock.calls.map(([tool]) => tool.name); - expect(toolNames).toContain("grep"); - expect(toolNames).toContain("find"); - expect(toolNames).not.toContain("ffgrep"); + expect(toolNames).toEqual( + expect.arrayContaining(["ffgrep", "fffind", "grep", "find"]), + ); + expect(setup.pi.setActiveTools).toHaveBeenLastCalledWith( + expect.not.arrayContaining(["ffgrep", "fffind"]), + ); expect(createCalls[0]).toEqual({ basePath: "/tmp/workspace", frecencyDbPath: "/config/frecency", @@ -320,9 +323,12 @@ describe("pi-fff global config", () => { const setup = await start("invalid-flag-mode"); const toolNames = setup.pi.registerTool.mock.calls.map(([tool]) => tool.name); - expect(toolNames).toContain("grep"); - expect(toolNames).toContain("find"); - expect(toolNames).not.toContain("ffgrep"); + expect(toolNames).toEqual( + expect.arrayContaining(["ffgrep", "fffind", "grep", "find"]), + ); + expect(setup.pi.setActiveTools).toHaveBeenLastCalledWith( + expect.not.arrayContaining(["ffgrep", "fffind"]), + ); await shutdown(setup); }); }); @@ -332,7 +338,7 @@ function writeConfig(config: Record): void { } describe("pi-fff session mode", () => { - test("registers tools only after restoring the saved mode", async () => { + test("pre-registers FFF renderers before restoring the saved mode", async () => { const setup = createPi("tools-and-ui"); const ctx = createContext(); ctx.sessionManager.getEntries.mockReturnValue([ @@ -340,20 +346,43 @@ describe("pi-fff session mode", () => { ]); fffExtension(setup.pi as any); - expect(setup.pi.registerTool).not.toHaveBeenCalled(); + expect(setup.pi.registerTool.mock.calls.map(([tool]) => tool.name)).toEqual([ + "ffgrep", + "fffind", + ]); + expect(setup.pi.setActiveTools).not.toHaveBeenCalled(); await setup.events.get("session_start")?.({ reason: "startup" }, ctx); const tools = setup.pi.registerTool.mock.calls.map(([tool]) => tool); const toolNames = tools.map((tool) => tool.name); - expect(toolNames).toContain("grep"); - expect(toolNames).toContain("find"); - expect(toolNames).not.toContain("ffgrep"); - expect(toolNames).not.toContain("fffind"); - const grepTool = tools.find((tool) => tool.name === "grep"); - expect(grepTool.promptGuidelines[0].startsWith("grep:")).toBe(true); + expect(toolNames).toEqual( + expect.arrayContaining(["ffgrep", "fffind", "grep", "find"]), + ); + const historicalGrep = tools.find((tool) => tool.name === "ffgrep"); + const activeGrep = tools.find((tool) => tool.name === "grep"); + const theme = { + bold: (text: string) => text, + fg: (_color: string, text: string) => text, + }; + const historicalCall = historicalGrep.renderCall( + { pattern: "TODO", path: "." }, + theme, + { state: {}, invalidate: mock(() => undefined), isError: false }, + ); + const activeCall = activeGrep.renderCall({ pattern: "TODO", path: "." }, theme, { + state: {}, + invalidate: mock(() => undefined), + isError: false, + }); + expect(historicalCall.text).toBe("ffgrep /TODO/ in ."); + expect(activeCall.text).toBe("grep /TODO/ in ."); + expect(activeGrep.promptGuidelines[0].startsWith("grep:")).toBe(true); expect(setup.pi.setActiveTools).toHaveBeenCalledWith( expect.arrayContaining(["read", "grep", "find"]), ); + expect(setup.pi.setActiveTools).toHaveBeenLastCalledWith( + expect.not.arrayContaining(["ffgrep", "fffind"]), + ); await setup.commands.get("fff-mode").handler("", ctx); expect(ctx.ui.notify).toHaveBeenLastCalledWith( @@ -363,17 +392,47 @@ describe("pi-fff session mode", () => { await shutdown(setup); }); + test("prunes auto-activated FFF names when override is the final mode", async () => { + const setup = createPi("tools-and-ui"); + const ctx = createContext(); + ctx.sessionManager.getEntries.mockReturnValue([ + { type: "custom", customType: "fff-mode", data: { mode: "override" } }, + ]); + fffExtension(setup.pi as any); + + // Pi core activates every newly registered tool, so the early FFF + // registrations are active by the time session_start runs. + setup.pi.getActiveTools.mockReturnValue(["read", "ffgrep", "fffind"]); + await setup.events.get("session_start")?.({ reason: "startup" }, ctx); + + expect(setup.pi.setActiveTools).toHaveBeenLastCalledWith( + expect.arrayContaining(["read", "grep", "find"]), + ); + expect(setup.pi.setActiveTools).toHaveBeenLastCalledWith( + expect.not.arrayContaining(["ffgrep", "fffind"]), + ); + await shutdown(setup); + }); + test("registers tools before an unbound SDK session's first agent turn", async () => { const setup = createPi("override"); const ctx = createContext(); fffExtension(setup.pi as any); - expect(setup.pi.registerTool).not.toHaveBeenCalled(); + expect(setup.pi.registerTool.mock.calls.map(([tool]) => tool.name)).toEqual([ + "ffgrep", + "fffind", + ]); + expect(setup.pi.setActiveTools).not.toHaveBeenCalled(); await setup.events.get("before_agent_start")?.({}, ctx); const toolNames = setup.pi.registerTool.mock.calls.map(([tool]) => tool.name); - expect(toolNames).toContain("grep"); - expect(toolNames).toContain("find"); + expect(toolNames).toEqual( + expect.arrayContaining(["ffgrep", "fffind", "grep", "find"]), + ); + expect(setup.pi.setActiveTools).toHaveBeenLastCalledWith( + expect.not.arrayContaining(["ffgrep", "fffind"]), + ); expect(createCalls).toHaveLength(0); await shutdown(setup); });