diff --git a/src/downshift.test.ts b/src/downshift.test.ts new file mode 100644 index 0000000..2a110c4 --- /dev/null +++ b/src/downshift.test.ts @@ -0,0 +1,141 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import type { + ExtensionAPI, + ExtensionCommandContext, + ExtensionContext, +} from "@earendil-works/pi-coding-agent"; + +const fsMocks = vi.hoisted(() => ({ + config: { + enabled: true, + threshold: { percent: 50 }, + economy: { provider: "test", model: "economy", thinkingLevel: "off" }, + premiumSource: "current", + startOnPremium: false, + upshiftAfterCompaction: false, + handoffBeforeDownshift: true, + }, + readFile: vi.fn(), + writeFile: vi.fn(), +})); + +vi.mock("node:fs/promises", () => ({ + readFile: fsMocks.readFile, + writeFile: fsMocks.writeFile, +})); + +import downshift from "./downshift"; + +type EventHandler = ( + event: unknown, + ctx: ExtensionContext, +) => void | Promise; + +type CommandHandler = ( + args: string, + ctx: ExtensionCommandContext, +) => void | Promise; + +type TestExtension = { + handlers: Map; + commands: Map; + pi: ExtensionAPI; +}; + +function createExtension(): TestExtension { + const handlers = new Map(); + const commands = new Map(); + const pi = { + on: (event: string, handler: EventHandler) => { + handlers.set(event, handler); + }, + registerCommand: (name: string, options: { handler: CommandHandler }) => { + commands.set(name, options.handler); + }, + appendEntry: vi.fn(), + sendUserMessage: vi.fn(), + setModel: vi.fn(), + setThinkingLevel: vi.fn(), + } as unknown as ExtensionAPI; + downshift(pi); + return { handlers, commands, pi }; +} + +function createContext(usage: { + current: { tokens: number; percent: number }; +}) { + const status = vi.fn(); + const select = vi.fn(); + const input = vi.fn(); + const ctx = { + hasUI: true, + getContextUsage: () => usage.current, + ui: { + setStatus: status, + select, + input, + notify: vi.fn(), + }, + }; + return { + commandContext: ctx as unknown as ExtensionCommandContext, + context: ctx as unknown as ExtensionContext, + input, + select, + status, + }; +} + +describe("downshift lifecycle adapter", () => { + beforeEach(() => { + fsMocks.readFile.mockClear(); + fsMocks.writeFile.mockClear(); + fsMocks.readFile.mockImplementation(async () => + JSON.stringify(fsMocks.config), + ); + fsMocks.writeFile.mockResolvedValue(undefined); + }); + + it("refreshes status at turn_end and agent_settled without downshifting", async () => { + const usage = { current: { tokens: 100, percent: 10 } }; + const { handlers, pi } = createExtension(); + const { context, status } = createContext(usage); + + await handlers.get("turn_end")?.({}, context); + expect(status).toHaveBeenLastCalledWith( + "downshift", + "⇣ premium (40% left)", + ); + + usage.current = { tokens: 200, percent: 20 }; + await handlers.get("agent_settled")?.({}, context); + expect(status).toHaveBeenLastCalledWith( + "downshift", + "⇣ premium (30% left)", + ); + expect(pi.sendUserMessage).not.toHaveBeenCalled(); + expect(pi.setModel).not.toHaveBeenCalled(); + }); + + it("refreshes status immediately after saving configuration", async () => { + const usage = { current: { tokens: 100, percent: 10 } }; + const { commands } = createExtension(); + const { commandContext, select, input, status } = createContext(usage); + select + .mockResolvedValueOnce("threshold: 50%") + .mockResolvedValueOnce("percent") + .mockResolvedValueOnce(undefined); + input.mockResolvedValueOnce("60"); + + await commands.get("downshift")?.("", commandContext); + + expect(fsMocks.writeFile).toHaveBeenCalledOnce(); + expect(status).toHaveBeenLastCalledWith( + "downshift", + "⇣ premium (50% left)", + ); + expect(fsMocks.writeFile.mock.invocationCallOrder[0]).toBeLessThan( + status.mock.invocationCallOrder[0], + ); + }); +}); diff --git a/src/downshift.ts b/src/downshift.ts index 97974b9..1197a56 100644 --- a/src/downshift.ts +++ b/src/downshift.ts @@ -179,6 +179,14 @@ function updateStatus(ctx: ExtensionContext, config?: DownshiftConfig): void { ); } +async function refreshStatus( + ctx: ExtensionContext, +): Promise { + const config = await readConfig(); + updateStatus(ctx, config); + return config; +} + function saveState(pi: ExtensionAPI, patch?: Partial): void { if (patch) runtime.state = { ...runtime.state, ...patch }; pi.appendEntry(CUSTOM_TYPE, { version: 1, ...runtime.state }); @@ -526,7 +534,7 @@ async function configureInitial(ctx: ExtensionCommandContext): Promise { const threshold = { tokens: 100000, percent: 50 }; const economy = await selectTarget(ctx, "Select economy model"); if (!economy) return; - await writeConfig({ + const config: DownshiftConfig = { enabled, threshold, economy, @@ -535,7 +543,9 @@ async function configureInitial(ctx: ExtensionCommandContext): Promise { startOnPremium, upshiftAfterCompaction: false, handoffBeforeDownshift: true, - }); + }; + await writeConfig(config); + updateStatus(ctx, config); ctx.ui.notify("downshift config created", "info"); } @@ -556,6 +566,7 @@ async function configureMenu( if (!next) continue; config = next; await writeConfig(config); + updateStatus(ctx, config); ctx.ui.notify("downshift config saved", "info"); } } @@ -775,7 +786,7 @@ async function downshiftNow( runtime, ctx.isIdle() ? "immediate" : "steer", ); - updateStatus(ctx, await readConfig()); + await refreshStatus(ctx); } async function setSessionEnabled( @@ -795,7 +806,7 @@ async function disableSession( handoff: "idle", continueAfterHandoff: false, }); - updateStatus(ctx, await readConfig()); + await refreshStatus(ctx); ctx.ui.notify("downshift off for this session", "info"); } @@ -901,6 +912,13 @@ function hasExplicitStartPremium( return !!config?.enabled && config.startOnPremium && !!config.premium; } +type ExtensionAPIWithAgentSettled = ExtensionAPI & { + on( + event: "agent_settled", + handler: (event: unknown, ctx: ExtensionContext) => void | Promise, + ): void; +}; + export default function downshift(pi: ExtensionAPI): void { pi.on("session_start", async (event, ctx) => { await handleSessionStart(pi, event, ctx); @@ -915,6 +933,17 @@ export default function downshift(pi: ExtensionAPI): void { ); }); + pi.on("turn_end", async (_event, ctx) => { + await refreshStatus(ctx); + }); + + (pi as ExtensionAPIWithAgentSettled).on( + "agent_settled", + async (_event, ctx) => { + await refreshStatus(ctx); + }, + ); + pi.on("before_agent_start", async (event, ctx) => { await handleBeforeAgentStart(coreDeps(pi, ctx), runtime, event, ctx); });