import { readdirSync, readFileSync } from "node:fs"; import { join } from "node:path"; import { fileURLToPath } from "node:url"; import { fauxAssistantMessage, fauxToolCall } from "@earendil-works/pi-ai"; import { defineTool, type ExtensionAPI, type ExtensionFactory, } from "@earendil-works/pi-coding-agent"; import type { Component, TUI } from "@earendil-works/pi-tui"; import { Type } from "typebox"; import { afterEach, describe, expect, it, vi } from "vitest"; import { createPiSession, type PiSession, } from "../../../test/fixture/create-pi-session.ts"; let session: PiSession | undefined; const finalReport = "Completed work with evidence and no unresolved blockers."; function extensionPath(): string { return fileURLToPath(new URL("../index.ts", import.meta.url)); } function command(name: "afk" | "back") { const result = session?.extensions .flatMap((extension) => [...extension.commands.entries()]) .find(([commandName]) => commandName === name)?.[1]; if (!result) throw new Error(`missing /${name}`); return result; } async function runAfk(): Promise { await command("afk").handler( "", session?.session.extensionRunner.createContext() as never, ); await session?.session.waitForIdle(); } const humanTool: ExtensionFactory = (pi: ExtensionAPI) => { pi.registerTool( defineTool({ name: "ask_user_question", label: "ask user", description: "Ask a human a question.", parameters: Type.Object({}), async execute() { return { content: [{ type: "text", text: "human answered" }], details: undefined, }; }, }), ); }; function queueFollowUp(pi: ExtensionAPI): void { let queued = false; pi.on("agent_start", () => { if (queued) return; queued = true; pi.sendUserMessage("Finish queued follow-up work.", { deliverAs: "followUp", }); }); } function stopAfkWithoutAborting(pi: ExtensionAPI): void { let stopped = false; pi.on("agent_start", async (_event, ctx) => { if (stopped) return; stopped = true; await command("back").handler("", ctx as never); }); } afterEach(async () => { await session?.dispose(); session = undefined; }); describe("afk contract", () => { it("loads through real Pi and leaves completion inactive", async () => { session = await createPiSession({ extensionPaths: [extensionPath()] }); const extension = session.extensions[0]!; expect(session.errors).toEqual([]); expect([...extension.commands.keys()]).toEqual(["afk", "back"]); expect(extension.tools.get("afk_complete")?.definition.defaultActive).toBe( false, ); expect(session.session.getActiveToolNames()).not.toContain("afk_complete"); }); it("uses real tool follow-up before napping after completion", async () => { session = await createPiSession({ extensionPaths: [extensionPath()], responses: [ fauxAssistantMessage("Investigated first approach."), fauxAssistantMessage(fauxToolCall("afk_complete", {}), { stopReason: "toolUse", }), fauxAssistantMessage(finalReport), ], }); await runAfk(); expect(session.providerState.callCount).toBe(3); expect(session.session.getActiveToolNames()).toContain("afk_complete"); expect(session.session.sessionManager.getBranch()).toContainEqual( expect.objectContaining({ message: expect.objectContaining({ role: "assistant", content: expect.any(Array), }), }), ); }); it("requires a fresh completion signal when reactivated after napping", async () => { session = await createPiSession({ extensionPaths: [extensionPath()], responses: [ fauxAssistantMessage(fauxToolCall("afk_complete", {}), { stopReason: "toolUse", }), fauxAssistantMessage(finalReport), ], }); await runAfk(); session.setResponses([ fauxAssistantMessage("More work remains."), fauxAssistantMessage(fauxToolCall("afk_complete", {}), { stopReason: "toolUse", }), fauxAssistantMessage("Reactivated work complete."), ]); await runAfk(); const branch = session.session.sessionManager.getBranch(); expect( branch.filter( (entry) => entry.type === "custom_message" && entry.customType === "afk-start", ), ).toHaveLength(2); expect( branch.filter( (entry) => entry.type === "custom_message" && entry.customType === "afk-continue", ), ).toHaveLength(1); expect(session.errors).toEqual([]); }); it("ignores a completion call selected before /back and resets on reactivation", async () => { let returned = false; const returnBeforeCompletion: ExtensionFactory = (pi) => { pi.on("tool_call", async (event, ctx) => { if (event.toolName !== "afk_complete" || returned) return; returned = true; await command("back").handler("", ctx as never); }); }; session = await createPiSession({ extensionPaths: [extensionPath()], extensionFactories: [returnBeforeCompletion], responses: [ fauxAssistantMessage(fauxToolCall("afk_complete", {}), { stopReason: "toolUse", }), fauxAssistantMessage("Current work finished normally."), ], }); await runAfk(); expect(returned).toBe(true); expect(session.session.getActiveToolNames()).not.toContain("afk_complete"); expect(session.session.sessionManager.getBranch()).toContainEqual( expect.objectContaining({ message: expect.objectContaining({ role: "toolResult", toolName: "afk_complete", isError: false, }), }), ); session.setResponses([ fauxAssistantMessage("New work remains."), fauxAssistantMessage(fauxToolCall("afk_complete", {}), { stopReason: "toolUse", }), fauxAssistantMessage(finalReport), ]); await runAfk(); expect( session.session.sessionManager .getBranch() .filter( (entry) => entry.type === "custom_message" && entry.customType === "afk-continue", ), ).toHaveLength(1); expect(session.errors).toEqual([]); }); it("returns blocked human tool result to model and receives a useful follow-up", async () => { session = await createPiSession({ extensionPaths: [extensionPath()], extensionFactories: [humanTool], responses: [ fauxAssistantMessage(fauxToolCall("ask_user_question", {}), { stopReason: "toolUse", }), fauxAssistantMessage( "Question unavailable, so I used the documented default.", ), fauxAssistantMessage(fauxToolCall("afk_complete", {}), { stopReason: "toolUse", }), fauxAssistantMessage(finalReport), ], }); await runAfk(); expect(session.providerState.callCount).toBe(4); expect(session.session.sessionManager.getBranch()).toContainEqual( expect.objectContaining({ message: expect.objectContaining({ role: "toolResult", isError: true, content: expect.arrayContaining([ expect.objectContaining({ text: expect.stringContaining( "unavailable while the user is away", ), }), ]), }), }), ); }); it("drains queued follow-up work before accepting completion", async () => { session = await createPiSession({ extensionPaths: [extensionPath()], extensionFactories: [queueFollowUp], responses: [ fauxAssistantMessage(fauxToolCall("afk_complete", {}), { stopReason: "toolUse", }), fauxAssistantMessage("Initial task report."), fauxAssistantMessage("Queued follow-up is still unfinished."), fauxAssistantMessage(fauxToolCall("afk_complete", {}), { stopReason: "toolUse", }), fauxAssistantMessage("Queued follow-up complete."), ], }); await runAfk(); expect(session.providerState.callCount).toBe(5); expect(session.session.getActiveToolNames()).toContain("afk_complete"); }); it("activates during busy work without reordering already queued follow-ups", async () => { session = await createPiSession({ extensionPaths: [extensionPath()], responses: [ async () => { await session!.session.sendUserMessage("Follow-up A", { deliverAs: "followUp", }); await session!.session.sendUserMessage("Follow-up B", { deliverAs: "followUp", }); await command("afk").handler( "", session!.session.extensionRunner.createContext() as never, ); return fauxAssistantMessage("Current task finished."); }, fauxAssistantMessage("Follow-up A finished."), fauxAssistantMessage("Follow-up B finished."), fauxAssistantMessage(fauxToolCall("afk_complete", {}), { stopReason: "toolUse", }), fauxAssistantMessage(finalReport), ], }); await session.session.prompt("Start current task."); const branch = session.session.sessionManager.getBranch(); const users = branch.flatMap((entry) => entry.type === "message" && entry.message.role === "user" ? [entry.message.content] : [], ); expect( users.map((content) => typeof content === "string" ? content : content .filter((block) => block.type === "text") .map((block) => block.text) .join(""), ), ).toEqual(["Start current task.", "Follow-up A", "Follow-up B"]); expect(session.providerState.callCount).toBe(5); expect(session.errors).toEqual([]); }); it("does not abort active work when /back disables continuation", async () => { session = await createPiSession({ extensionPaths: [extensionPath()], extensionFactories: [stopAfkWithoutAborting], responses: [fauxAssistantMessage("Current work completed normally.")], }); await runAfk(); expect(session.providerState.callCount).toBe(1); expect(session.session.getActiveToolNames()).not.toContain("afk_complete"); }); it("stops after a bounded number of no-progress continuations", async () => { session = await createPiSession({ extensionPaths: [extensionPath()], responses: [ fauxAssistantMessage("Still considering it."), fauxAssistantMessage("Still considering it."), fauxAssistantMessage("Still considering it."), ], }); await runAfk(); expect(session.providerState.callCount).toBe(3); expect(session.session.getActiveToolNames()).not.toContain("afk_complete"); expect(session.session.sessionManager.getBranch()).toContainEqual( expect.objectContaining({ customType: "afk-incomplete" }), ); }); it.each(["aborted", "error"] as const)( "cleans up mode after %s without restarting", async (stopReason) => { session = await createPiSession({ extensionPaths: [extensionPath()], responses: [fauxAssistantMessage("Stopped.", { stopReason })], }); await runAfk(); expect(session.providerState.callCount).toBe(1); expect(session.session.getActiveToolNames()).not.toContain( "afk_complete", ); }, ); it("does not duplicate activation while AFK work is active", async () => { session = await createPiSession({ extensionPaths: [extensionPath()] }); const ctx = session.session.extensionRunner.createContext(); await command("afk").handler("", ctx as never); await command("afk").handler("", ctx as never); await session.session.waitForIdle(); expect(session.providerState.callCount).toBe(1); }); it("renders a napping creature through the real TUI binding and disposes it on /back", async () => { session = await createPiSession({ extensionPaths: [extensionPath()], responses: [ fauxAssistantMessage(fauxToolCall("afk_complete", {}), { stopReason: "toolUse", }), fauxAssistantMessage(finalReport), ], }); let creature: (Component & { dispose?(): void }) | undefined; const requestRender = vi.fn(); await session.session.bindExtensions({ mode: "tui", uiContext: { ...session.session.extensionRunner.getUIContext(), setWidget: (_key, content) => { if (typeof content === "function") { creature = content( { requestRender } as unknown as TUI, { fg: (_color: string, text: string) => text } as never, ); } else creature = undefined; }, }, }); await runAfk(); expect(creature?.render(80).join("\n")).toContain("zzz"); expect(requestRender).toHaveBeenCalled(); const dispose = vi.spyOn(creature!, "dispose"); await command("back").handler( "", session.session.extensionRunner.createContext() as never, ); expect(dispose).toHaveBeenCalledOnce(); expect(creature).toBeUndefined(); expect(session.errors).toEqual([]); }); it.each(["alone", "first", "last"])( "recovers idle compaction with AFK loaded %s", async (order) => { const paths = [extensionPath()]; if (order !== "alone") paths.push( fileURLToPath(new URL("../../continuity/index.ts", import.meta.url)), ); if (order === "last") paths.reverse(); session = await createPiSession({ extensionPaths: paths, responses: [ fauxAssistantMessage(fauxToolCall("afk_complete", {}), { stopReason: "toolUse", }), fauxAssistantMessage(finalReport), fauxAssistantMessage("Ordinary compaction recovery complete."), ], }); // Leave AFK enabled but idle to exercise recovery rather than the normal run loop. const start = vi .spyOn(session.session, "sendCustomMessage") .mockImplementationOnce(async () => {}); await command("afk").handler( "", session.session.extensionRunner.createContext() as never, ); start.mockRestore(); expect(session.providerState.callCount).toBe(0); const compact = () => session!.session.extensionRunner.emit({ type: "session_compact", compactionEntry: { id: "afk-compact" }, fromExtension: false, reason: "manual", willRetry: false, } as never); await compact(); await vi.waitFor(() => expect(session!.providerState.callCount).toBe(2)); await session.session.waitForIdle(); const continuations = (type: string) => session!.session.sessionManager .getBranch() .filter( (entry) => entry.type === "custom_message" && entry.customType === type, ); expect(continuations("afk-continue")).toHaveLength(1); expect(continuations("continuity-resume")).toHaveLength(0); await command("back").handler( "", session.session.extensionRunner.createContext() as never, ); await compact(); if (order !== "alone") { await vi.waitFor(() => expect(session!.providerState.callCount).toBe(3), ); await session.session.waitForIdle(); expect(continuations("continuity-resume")).toHaveLength(1); } else { await new Promise((resolve) => setTimeout(resolve, 10)); expect(session.providerState.callCount).toBe(2); } expect(continuations("afk-continue")).toHaveLength(1); expect(session.errors).toEqual([]); }, ); it.each(["first", "last"])( "does not restart napping AFK loaded %s, even without its completion tool active", async (order) => { const paths = [ extensionPath(), fileURLToPath(new URL("../../continuity/index.ts", import.meta.url)), ]; if (order === "last") paths.reverse(); session = await createPiSession({ extensionPaths: paths, responses: [ fauxAssistantMessage(fauxToolCall("afk_complete", {}), { stopReason: "toolUse", }), fauxAssistantMessage(finalReport), ], }); await runAfk(); session.session.setActiveToolsByName( session.session .getActiveToolNames() .filter((name) => name !== "afk_complete"), ); const send = vi.spyOn(session.session, "sendCustomMessage"); await session.session.extensionRunner.emit({ type: "session_compact", compactionEntry: { id: "afk-compact" }, fromExtension: false, reason: "manual", willRetry: false, } as never); await new Promise((resolve) => setTimeout(resolve, 10)); expect(send).not.toHaveBeenCalled(); expect(session.providerState.callCount).toBe(2); }, ); it("writes only schema-bound lifecycle diagnostics when explicitly enabled", async () => { session = await createPiSession({ extensionPaths: [extensionPath()], env: { PI_AFK_DEBUG: "1" }, }); await session.session.extensionRunner.emit({ type: "session_start", reason: "startup", }); await session.session.extensionRunner.emit({ type: "session_shutdown", reason: "quit", }); const directory = join( session.env["XDG_STATE_HOME"]!, "pi-ext", "debug", "afk", ); const events = await vi.waitFor(() => { const records = readdirSync(directory).flatMap((file) => readFileSync(join(directory, file), "utf8") .trim() .split("\n") .filter(Boolean) .map((line) => JSON.parse(line)), ); expect(records).toContainEqual( expect.objectContaining({ extension: "afk", event: "session.shutdown", }), ); return records.map((record) => record.event); }); expect(events).toContain("session.start"); }); it("cleans AFK mode on session-tree navigation and shutdown", async () => { session = await createPiSession({ extensionPaths: [extensionPath()] }); await command("afk").handler( "", session.session.extensionRunner.createContext() as never, ); await session.session.extensionRunner.emit({ type: "session_tree", newLeafId: null, oldLeafId: null, }); expect(session.session.getActiveToolNames()).not.toContain("afk_complete"); await session.session.extensionRunner.emit({ type: "session_shutdown", reason: "quit", }); expect(session.errors).toEqual([]); }); });