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); });