import { afterEach, describe, expect, it, vi } from "bun:test"; import * as path from "node:path"; import type { AssistantMessage, Model } from "@oh-my-pi/pi-ai"; import type { AsyncJobRegisterOptions } from "@oh-my-pi/pi-coding-agent/async/job-manager"; import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; import { TanCommandController } from "@oh-my-pi/pi-coding-agent/modes/controllers/tan-command-controller"; import type { InteractiveModeContext } from "@oh-my-pi/pi-coding-agent/modes/types"; import { AgentRegistry, MAIN_AGENT_ID } from "@oh-my-pi/pi-coding-agent/registry/agent-registry"; import type { CreateAgentSessionResult } from "@oh-my-pi/pi-coding-agent/sdk"; import * as sdkModule from "@oh-my-pi/pi-coding-agent/sdk"; import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager"; import { TempDir } from "@oh-my-pi/pi-utils"; interface CapturedJobRunContext { jobId: string; signal: AbortSignal; reportProgress: (text: string, details?: Record) => Promise; } type CapturedJobRun = (ctx: CapturedJobRunContext) => Promise; const model = { provider: "anthropic", id: "claude-sonnet-4-5" } as Model; function assistantText(text: string): AssistantMessage { return { role: "assistant", content: [{ type: "text", text }], api: "anthropic-messages", provider: model.provider, model: model.id, usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, stopReason: "stop", timestamp: 0, }; } interface TanSessionEvent { type: string; result?: unknown; aborted?: boolean; } /** Minimal tan clone session stub covering the surface `TanCommandController` drives. */ function createCloneStub(overrides?: { prompt?: () => Promise; abort?: () => void; sessionManager?: { appendSessionInit: (init: unknown) => void }; lastAssistantText?: string; }) { const appendMessage = vi.fn(); let listener: ((event: TanSessionEvent) => void) | undefined; const clone = { agent: { appendMessage }, sessionManager: overrides?.sessionManager, setTodoPhases: vi.fn(), subscribe: vi.fn((l: (event: TanSessionEvent) => void) => { listener = l; return () => { listener = undefined; }; }), prompt: vi.fn(overrides?.prompt ?? (async () => {})), waitForIdle: vi.fn(async () => {}), getLastAssistantMessage: vi.fn(() => assistantText(overrides?.lastAssistantText ?? "done")), abort: vi.fn(overrides?.abort ?? (() => {})), dispose: vi.fn(async () => {}), }; return { clone, appendMessage, get compactionListener() { return listener; }, }; } function createContext(overrides?: { isStreaming?: boolean; model?: Model; agentId?: string; register?: (run: CapturedJobRun, options?: AsyncJobRegisterOptions) => string; }) { const tempDir = TempDir.createSync("@omp-tan-controller-"); const parentFile = path.join(tempDir.path(), "parent.jsonl"); // The clone nests inside the parent's artifact directory, like a subagent. const cloneFile = path.join(parentFile.slice(0, -6), "clone.jsonl"); let capturedRun: CapturedJobRun | undefined; let capturedOptions: AsyncJobRegisterOptions | undefined; const sequence: string[] = []; const register = vi.fn( (_type: "bash" | "task", _label: string, run: CapturedJobRun, options?: AsyncJobRegisterOptions): string => { sequence.push("register"); capturedRun = run; capturedOptions = options; return overrides?.register ? overrides.register(run, options) : "job-123"; }, ); const session = { isStreaming: overrides?.isStreaming ?? false, model: overrides?.model ?? model, asyncJobManager: { register }, sessionId: "parent-session", configuredThinkingLevel: vi.fn(() => undefined), systemPrompt: ["system prompt"], getActiveToolNames: vi.fn(() => ["read", "bash"]), modelRegistry: { authStorage: { marker: "auth" } }, getAgentId: vi.fn(() => overrides?.agentId), sendCustomMessage: vi.fn(async () => { sequence.push("sendCustomMessage"); }), } as unknown as InteractiveModeContext["session"]; const sessionManager = { getSessionFile: vi.fn(() => parentFile), getCwd: vi.fn(() => tempDir.path()), getSessionDir: vi.fn(() => tempDir.path()), ensureOnDisk: vi.fn(async () => {}), flush: vi.fn(async () => {}), } as unknown as InteractiveModeContext["sessionManager"]; const cloneManager = { getSessionFile: vi.fn(() => cloneFile), appendCustomEntry: vi.fn(), } as unknown as SessionManager; const ctx = { session, sessionManager, settings: Settings.isolated({ "task.enableLsp": true }), showStatus: vi.fn(), showWarning: vi.fn(), showError: vi.fn(), rebuildChatFromMessages: vi.fn(), } as unknown as InteractiveModeContext; return { tempDir, parentFile, cloneFile, cloneManager, ctx, register, sequence, get capturedRun() { return capturedRun; }, get capturedOptions() { return capturedOptions; }, }; } describe("TanCommandController", () => { afterEach(() => { vi.restoreAllMocks(); }); it("rejects empty work before forking", async () => { const harness = createContext(); const forkSpy = vi.spyOn(SessionManager, "forkFrom").mockResolvedValue(harness.cloneManager); const controller = new TanCommandController(harness.ctx); await controller.start(" "); expect(forkSpy).not.toHaveBeenCalled(); expect(harness.ctx.showStatus).toHaveBeenCalledWith("Usage: /tan "); }); it("dispatches without disturbing an in-flight turn while streaming", async () => { const harness = createContext({ isStreaming: true }); const forkSpy = vi.spyOn(SessionManager, "forkFrom").mockResolvedValue(harness.cloneManager); const controller = new TanCommandController(harness.ctx); await controller.start("check something"); expect(forkSpy).toHaveBeenCalled(); expect(harness.ctx.showWarning).not.toHaveBeenCalled(); // The breadcrumb is queued for the next turn, not steered into the live one, // and the live chat is left to the streaming renderer (no synchronous rebuild). expect(harness.ctx.session.sendCustomMessage).toHaveBeenCalledWith( expect.objectContaining({ customType: "background-tan-dispatch" }), { triggerTurn: false, deliverAs: "nextTurn" }, ); expect(harness.ctx.rebuildChatFromMessages).not.toHaveBeenCalled(); expect(harness.ctx.showStatus).toHaveBeenCalledWith("Dispatched background tan job-123"); }); it("forks with breadcrumb suppression, registers under Main, and dispatches after receiving the job id", async () => { const harness = createContext(); const forkSpy = vi.spyOn(SessionManager, "forkFrom").mockResolvedValue(harness.cloneManager); const controller = new TanCommandController(harness.ctx); await controller.start("write the release note"); expect(forkSpy).toHaveBeenCalledWith( harness.parentFile, harness.tempDir.path(), harness.parentFile.slice(0, -6), undefined, { suppressBreadcrumb: true, sessionFile: expect.stringMatching(/Tan-.+\.jsonl$/) }, ); expect(harness.register).toHaveBeenCalledWith("task", "/tan write the release note", expect.any(Function), { ownerId: MAIN_AGENT_ID, agentId: expect.stringMatching(/^Tan-/) as unknown as string, }); expect(harness.capturedOptions?.ownerId).toBe(MAIN_AGENT_ID); expect(harness.sequence).toEqual(["register", "sendCustomMessage"]); expect(harness.ctx.session.sendCustomMessage).toHaveBeenCalledWith( expect.objectContaining({ customType: "background-tan-dispatch", details: { jobId: "job-123", work: "write the release note", sessionFile: expect.stringMatching(/Tan-.+\.jsonl$/), }, }), { triggerTurn: false, deliverAs: "nextTurn" }, ); expect(harness.ctx.rebuildChatFromMessages).toHaveBeenCalled(); expect(harness.ctx.showStatus).toHaveBeenCalledWith("Dispatched background tan job-123"); }); it("aborts the cloned agent when the background job signal aborts", async () => { const harness = createContext({ agentId: MAIN_AGENT_ID }); vi.spyOn(SessionManager, "forkFrom").mockResolvedValue(harness.cloneManager); const promptStarted = Promise.withResolvers(); const abortObserved = Promise.withResolvers(); const { clone } = createCloneStub({ prompt: async () => { promptStarted.resolve(); await abortObserved.promise; }, abort: () => { abortObserved.resolve(); }, lastAssistantText: "finished", }); const createAgentSessionSpy = vi .spyOn(sdkModule, "createAgentSession") .mockResolvedValue({ session: clone } as unknown as CreateAgentSessionResult); const controller = new TanCommandController(harness.ctx); await controller.start("follow the tangent"); const capturedRun = harness.capturedRun; expect(capturedRun).toBeDefined(); if (!capturedRun) throw new Error("run function was not captured"); const abortController = new AbortController(); const resultPromise = capturedRun({ jobId: "job-123", signal: abortController.signal, reportProgress: async () => {}, }); await promptStarted.promise; abortController.abort(); const result = await resultPromise; expect(result).toBe("finished"); expect(clone.abort).toHaveBeenCalled(); expect(clone.dispose).toHaveBeenCalled(); expect(createAgentSessionSpy.mock.calls[0]?.[0]).toEqual( expect.objectContaining({ providerPromptCacheKey: "parent-session", parentTaskPrefix: expect.stringMatching(/^Tan-/) as unknown as string, agentDisplayName: "tan", }), ); }); it("parents the tan clone to the spawning agent, not to the clone itself", async () => { const harness = createContext({ agentId: "FocusedParent" }); vi.spyOn(SessionManager, "forkFrom").mockResolvedValue(harness.cloneManager); const { clone } = createCloneStub(); const createAgentSessionSpy = vi .spyOn(sdkModule, "createAgentSession") .mockResolvedValue({ session: clone } as unknown as CreateAgentSessionResult); const controller = new TanCommandController(harness.ctx); await controller.start("follow the tangent"); const capturedRun = harness.capturedRun; if (!capturedRun) throw new Error("run function was not captured"); await capturedRun({ jobId: "job-1", signal: new AbortController().signal, reportProgress: async () => {} }); const opts = createAgentSessionSpy.mock.calls[0]?.[0]; // The clone's registry parent is the spawning (focused) agent. Its own // `Tan-` artifact prefix must never double as the parent link, or the // hub would render the tan parented to itself. expect(opts?.parentAgentId).toBe("FocusedParent"); expect(opts?.parentTaskPrefix).toMatch(/^Tan-/); expect(opts?.parentTaskPrefix).not.toBe("FocusedParent"); }); it("parks the finished tan in the registry so it stays visible in the Agent Hub", async () => { const harness = createContext(); vi.spyOn(SessionManager, "forkFrom").mockResolvedValue(harness.cloneManager); const appendSessionInit = vi.fn(); const { clone } = createCloneStub({ sessionManager: { appendSessionInit } }); vi.spyOn(sdkModule, "createAgentSession").mockResolvedValue({ session: clone, } as unknown as CreateAgentSessionResult); const registry = AgentRegistry.global(); const setStatus = vi.spyOn(registry, "setStatus"); const detachSession = vi.spyOn(registry, "detachSession"); const unregister = vi.spyOn(registry, "unregister"); const controller = new TanCommandController(harness.ctx); await controller.start("park me"); const run = harness.capturedRun; if (!run) throw new Error("run function was not captured"); const result = await run({ jobId: "job-123", signal: new AbortController().signal, reportProgress: async () => {}, }); expect(result).toBe("done"); expect(appendSessionInit).toHaveBeenCalledWith({ systemPrompt: "system prompt", task: "park me", tools: ["read", "bash"], }); // Parked (not unregistered) before dispose, then the disposed session is nulled // out — the hub keeps the ref and reads its transcript from the session file. expect(setStatus).toHaveBeenCalledWith(expect.stringMatching(/^Tan-/), "parked"); expect(detachSession).toHaveBeenCalledWith(expect.stringMatching(/^Tan-/)); expect(clone.dispose).toHaveBeenCalled(); expect(unregister).not.toHaveBeenCalled(); }); it("isolates the fork: clears inherited todos, injects the fork notice, and re-injects after compaction", async () => { const harness = createContext(); vi.spyOn(SessionManager, "forkFrom").mockResolvedValue(harness.cloneManager); const compacted = Promise.withResolvers(); const stub = createCloneStub({ prompt: async () => { // Simulate the clone's history compacting mid-run: the summarizer // erases the fork notice, so the controller must append it again. stub.compactionListener?.({ type: "auto_compaction_end", result: {}, aborted: false }); compacted.resolve(); }, }); vi.spyOn(sdkModule, "createAgentSession").mockResolvedValue({ session: stub.clone, } as unknown as CreateAgentSessionResult); const controller = new TanCommandController(harness.ctx); await controller.start("follow the tangent"); const run = harness.capturedRun; if (!run) throw new Error("run function was not captured"); await run({ jobId: "job-123", signal: new AbortController().signal, reportProgress: async () => {} }); await compacted.promise; // Inherited parent todos are wiped both in-memory and in the persisted // session so reloads agree; otherwise todo reminders drag the tan back // onto the parent's task. expect(stub.clone.setTodoPhases).toHaveBeenCalledWith([]); expect(harness.cloneManager.appendCustomEntry).toHaveBeenCalledWith("user_todo_edit", { phases: [] }); // Fork notice injected before the prompt and again after compaction. expect(stub.appendMessage).toHaveBeenCalledTimes(2); for (const call of stub.appendMessage.mock.calls) { expect(call[0]).toEqual( expect.objectContaining({ role: "developer", content: expect.stringContaining(''), }), ); } // The compaction listener is released once the tan finishes. expect(stub.compactionListener).toBeUndefined(); }); });