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, }; } 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), } 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, }); 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 = { prompt: vi.fn(async () => { promptStarted.resolve(); await abortObserved.promise; }), waitForIdle: vi.fn(async () => {}), getLastAssistantMessage: vi.fn(() => assistantText("finished")), abort: vi.fn(() => { abortObserved.resolve(); }), dispose: vi.fn(async () => {}), }; 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 = { prompt: vi.fn(async () => {}), waitForIdle: vi.fn(async () => {}), getLastAssistantMessage: vi.fn(() => assistantText("done")), abort: vi.fn(), dispose: vi.fn(async () => {}), }; 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 clone = { prompt: vi.fn(async () => {}), waitForIdle: vi.fn(async () => {}), getLastAssistantMessage: vi.fn(() => assistantText("done")), abort: vi.fn(), dispose: vi.fn(async () => {}), }; 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"); // 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(); }); });