/** * Contracts: task.batch gating (batch spawning + shared context). * * 1. The wire schema is shape-swapped by `task.batch`: `{ context, tasks[] }` * when on (per-spawn fields — including `isolated` — live in the items), * the flat `{ name?, agent?, task, isolated? }` when off. Neither * shape exposes a per-call `schema` input (structured output comes from * agent frontmatter / inherited session schema / eval agent()). * 2. Shape validation rejects `schema` always, `tasks`/`context` while batch * is disabled, top-level `task` in batch calls, empty/invalid items, * duplicate names, and a missing shared `context`. * 3. With `async.enabled=true`, a batch call registers one background job per * item; with `async.enabled=false`, it blocks and returns merged results. * Both modes forward the shared `context`; the flat form stays accepted at * runtime for internal callers. */ import { afterEach, beforeEach, describe, expect, it, vi } from "bun:test"; import { toolWireSchema } from "@oh-my-pi/pi-ai/utils/schema"; import { AsyncJobManager } from "@oh-my-pi/pi-coding-agent/async/job-manager"; import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; import { AgentLifecycleManager } from "@oh-my-pi/pi-coding-agent/registry/agent-lifecycle"; import { AgentRegistry } from "@oh-my-pi/pi-coding-agent/registry/agent-registry"; import { TaskTool } from "@oh-my-pi/pi-coding-agent/task"; import * as discoveryModule from "@oh-my-pi/pi-coding-agent/task/discovery"; import * as executorModule from "@oh-my-pi/pi-coding-agent/task/executor"; import type { AgentDefinition, SingleResult, TaskParams } from "@oh-my-pi/pi-coding-agent/task/types"; import type { ToolSession } from "@oh-my-pi/pi-coding-agent/tools"; const taskAgent: AgentDefinition = { name: "task", description: "General-purpose task agent", systemPrompt: "You are a task agent.", source: "bundled", }; function createSession( options: { manager?: AsyncJobManager; settings?: Record; agentId?: string } = {}, ): ToolSession { return { cwd: "/tmp", hasUI: false, settings: Settings.isolated(options.settings ?? {}), getSessionFile: () => null, getSessionSpawns: () => "*", getAgentId: () => options.agentId ?? null, asyncJobManager: options.manager, } as unknown as ToolSession; } function getSchemaProperties(tool: TaskTool): Record { const wire = toolWireSchema(tool) as { properties?: Record }; return wire.properties ?? {}; } function getFirstText(result: { content: Array<{ type: string; text?: string }> }): string { const content = result.content.find(part => part.type === "text"); return content?.type === "text" ? (content.text ?? "") : ""; } function makeResult(id: string, overrides: Partial = {}): SingleResult { return { index: 0, id, agent: "task", agentSource: "bundled", task: "task prompt", assignment: "Do the thing.", exitCode: 0, output: "All done.", stderr: "", truncated: false, durationMs: 5, tokens: 0, requests: 1, ...overrides, }; } function mockDiscovery(): void { vi.spyOn(discoveryModule, "discoverAgents").mockResolvedValue({ agents: [taskAgent], projectAgentsDir: null, }); } describe("task.batch schema gating", () => { afterEach(() => { vi.restoreAllMocks(); }); it("swaps between the flat and batch wire shapes", async () => { mockDiscovery(); const off = await TaskTool.create(createSession({ settings: { "task.batch": false } })); const offProperties = getSchemaProperties(off); expect(offProperties.tasks).toBeUndefined(); expect(offProperties.context).toBeUndefined(); expect(offProperties.task).toBeDefined(); expect(offProperties.name).toBeDefined(); const on = await TaskTool.create(createSession({ settings: { "task.batch": true } })); const onProperties = getSchemaProperties(on); expect(onProperties.tasks).toBeDefined(); expect(onProperties.context).toBeDefined(); // The batch shape is { context, tasks[] } — the per-spawn fields live // only inside the task items. expect(onProperties.task).toBeUndefined(); expect(onProperties.name).toBeUndefined(); expect(onProperties.agent).toBeUndefined(); const items = (onProperties.tasks as { items?: { properties?: Record } }).items; expect(items?.properties?.task).toBeDefined(); expect(items?.properties?.name).toBeDefined(); expect(items?.properties?.agent).toBeDefined(); }); it("places isolated per item in the batch shape when isolation is enabled", async () => { mockDiscovery(); const tool = await TaskTool.create( createSession({ settings: { "task.batch": true, "task.isolation.mode": "auto" } }), ); const properties = getSchemaProperties(tool); expect(properties.isolated).toBeUndefined(); const items = (properties.tasks as { items?: { properties?: Record } }).items; expect(items?.properties?.isolated).toBeDefined(); }); it("never exposes a per-call schema input", async () => { mockDiscovery(); for (const settings of [{ "task.batch": false }, { "task.batch": true }]) { const tool = await TaskTool.create(createSession({ settings })); expect(getSchemaProperties(tool).schema).toBeUndefined(); } }); }); describe("task.batch validation", () => { afterEach(() => { vi.restoreAllMocks(); }); async function executeText(params: unknown, settings: Record = {}): Promise { mockDiscovery(); const tool = await TaskTool.create(createSession({ settings })); const result = await tool.execute("tool-call", params); return getFirstText(result); } it("rejects a schema argument regardless of batch mode", async () => { for (const batch of [false, true]) { const text = await executeText( { agent: "task", task: "Work.", schema: '{"properties":{}}' }, { "task.batch": batch }, ); expect(text).toContain("does not accept `schema`"); } }); it("rejects tasks and context while task.batch is disabled", async () => { const disabled = { "task.batch": false }; const text = await executeText({ agent: "task", tasks: [{ task: "Work." }] }, disabled); expect(text).toContain("task.batch is disabled"); const contextText = await executeText({ agent: "task", task: "Work.", context: "Background." }, disabled); expect(contextText).toContain("task.batch is disabled"); }); it("rejects top-level task in the batch shape", async () => { const text = await executeText({ task: "Work.", tasks: [{ task: "Other." }] }, { "task.batch": true }); expect(text).toContain("not part of the batch shape"); }); it("rejects empty task arrays and items without tasks", async () => { const empty = await executeText({ tasks: [] }, { "task.batch": true }); expect(empty).toContain("Missing `tasks`"); const missing = await executeText({ tasks: [{ task: "Work." }, { name: "Beta" }] }, { "task.batch": true }); expect(missing).toContain("Task 2 (`Beta`) is missing `task`"); }); it("requires a shared context for batch calls", async () => { const text = await executeText({ tasks: [{ task: "Work." }] }, { "task.batch": true }); expect(text).toContain("Missing `context`"); }); it("rejects duplicate provided names case-insensitively", async () => { const text = await executeText( { tasks: [ { name: "Anna", task: "A." }, { name: "anna", task: "B." }, ], }, { "task.batch": true }, ); expect(text).toContain("Duplicate task name"); }); }); describe("task.batch spawning", () => { const managers: AsyncJobManager[] = []; function createManager(): AsyncJobManager { const manager = new AsyncJobManager({ onJobComplete: () => {} }); managers.push(manager); return manager; } beforeEach(() => { AgentRegistry.resetGlobalForTests(); AgentLifecycleManager.resetGlobalForTests(); }); afterEach(async () => { vi.restoreAllMocks(); for (const manager of managers.splice(0)) { await manager.dispose({ timeoutMs: 1000 }); } AgentLifecycleManager.resetGlobalForTests(); AgentRegistry.resetGlobalForTests(); }); it("spawns one background job per task item and forwards the shared context", async () => { mockDiscovery(); const seen: Array<{ id?: string; context?: string; assignment?: string; parentAgentId?: string }> = []; vi.spyOn(executorModule, "runSubprocess").mockImplementation(async options => { seen.push({ id: options.id, context: options.context, assignment: options.assignment, parentAgentId: options.parentAgentId, }); return makeResult(options.id ?? "?"); }); const manager = createManager(); const tool = await TaskTool.create( createSession({ manager, agentId: "ParentA", settings: { "async.enabled": true, "task.batch": true } }), ); const result = await tool.execute("tc-batch", { context: "# Goal\nShared background.", tasks: [ { name: "Alpha", task: "Do A." }, { name: "Beta", task: "Do B." }, ], } as TaskParams); const text = getFirstText(result); expect(text).toContain("Spawned 2 background agents"); expect(text).toContain("- `Alpha`"); expect(text).toContain("- `Beta`"); expect(result.details?.progress?.map(progress => progress.id)).toEqual(["Alpha", "Beta"]); expect(result.details?.async?.state).toBe("running"); const alphaJob = manager.getJob("Alpha"); const betaJob = manager.getJob("Beta"); expect(alphaJob).toBeDefined(); expect(betaJob).toBeDefined(); await alphaJob!.promise; await betaJob!.promise; expect(alphaJob!.status).toBe("completed"); expect(betaJob!.status).toBe("completed"); expect(alphaJob!.resultText).toContain("Alpha is now idle"); expect(betaJob!.resultText).toContain("history://Beta"); expect(seen).toHaveLength(2); for (const spawn of seen) { expect(spawn.context).toBe("# Goal\nShared background."); } expect(seen.map(spawn => spawn.assignment).sort()).toEqual(["Do A.", "Do B."]); // Every spawn is parented to the spawning agent (not to itself): the // registry "of " link must be the caller, never the child's id. for (const spawn of seen) expect(spawn.parentAgentId).toBe("ParentA"); }); it("treats a one-item batch as a single spawn and forwards context", async () => { mockDiscovery(); let capturedContext: string | undefined; vi.spyOn(executorModule, "runSubprocess").mockImplementation(async options => { capturedContext = options.context; return makeResult(options.id ?? "?"); }); const manager = createManager(); const tool = await TaskTool.create( createSession({ manager, settings: { "async.enabled": true, "task.batch": true } }), ); const result = await tool.execute("tc-single", { context: "Shared notes.", tasks: [{ name: "Solo", task: "Do the thing." }], } as TaskParams); expect(getFirstText(result)).toContain("Spawned agent `Solo`"); const job = manager.getJob(result.details!.async!.jobId)!; await job.promise; expect(job.status).toBe("completed"); expect(capturedContext).toBe("Shared notes."); }); it("accepts the flat single-spawn form at runtime under batch mode", async () => { // Internal callers (e.g. the commit flow) and stale transcripts use the // flat shape directly; the wire schema is batch-only but runtime is not. mockDiscovery(); vi.spyOn(executorModule, "runSubprocess").mockImplementation(async options => makeResult(options.id ?? "?")); const manager = createManager(); const tool = await TaskTool.create( createSession({ manager, settings: { "async.enabled": true, "task.batch": true } }), ); const result = await tool.execute("tc-flat", { agent: "task", name: "Flat", task: "Do the thing.", } as TaskParams); expect(getFirstText(result)).toContain("Spawned agent `Flat`"); const job = manager.getJob(result.details!.async!.jobId)!; await job.promise; expect(job.status).toBe("completed"); }); it("blocks batch execution when async.enabled is false even with a job manager", async () => { mockDiscovery(); const seen: Array<{ id?: string; context?: string; assignment?: string }> = []; vi.spyOn(executorModule, "runSubprocess").mockImplementation(async options => { seen.push({ id: options.id, context: options.context, assignment: options.assignment }); return makeResult(options.id ?? "?"); }); const manager = createManager(); const tool = await TaskTool.create( createSession({ manager, settings: { "async.enabled": false, "task.batch": true } }), ); const result = await tool.execute("tc-sync-batch", { context: "# Goal\nShared synchronous context.", tasks: [ { name: "Alpha", task: "Do A." }, { name: "Beta", task: "Do B." }, ], } as TaskParams); expect(getFirstText(result)).toContain("All done."); expect(result.details?.async).toBeUndefined(); expect(result.details?.results.map(item => item.id).sort()).toEqual(["Alpha", "Beta"]); expect(manager.getJob("Alpha")).toBeUndefined(); expect(manager.getJob("Beta")).toBeUndefined(); expect(seen.map(spawn => spawn.context)).toEqual([ "# Goal\nShared synchronous context.", "# Goal\nShared synchronous context.", ]); }); it("settles the batch async aggregate when a queued spawn is cancelled mid-flight", async () => { mockDiscovery(); const started: string[] = []; const gates = new Map; resolve: () => void }>(); vi.spyOn(executorModule, "runSubprocess").mockImplementation(async options => { const id = options.id ?? "?"; started.push(id); const { promise, resolve } = Promise.withResolvers(); gates.set(id, { promise, resolve }); await promise; return makeResult(id); }); const manager = createManager(); const tool = await TaskTool.create( createSession({ manager, settings: { "async.enabled": true, "task.batch": true, "task.maxConcurrency": 1 }, }), ); const updates: Array<{ async?: { state?: string }; progress?: Array<{ id: string; status: string }> }> = []; const result = await tool.execute( "tc-batch-cancel", { context: "ctx", tasks: [ { name: "First", task: "Do A." }, { name: "Second", task: "Do B." }, ], } as TaskParams, undefined, update => { if (update.details) { updates.push({ async: update.details.async, progress: update.details.progress?.map(p => ({ id: p.id, status: p.status })), }); } }, ); expect(result.details?.async?.state).toBe("running"); const firstJob = manager.getJob("First")!; const secondJob = manager.getJob("Second")!; const deadline = Date.now() + 1_000; while (started.length === 0) { if (Date.now() > deadline) throw new Error("First spawn never reached the executor"); await Bun.sleep(5); } expect(started).toEqual(["First"]); expect(secondJob.queued).toBe(true); expect(manager.cancel(secondJob.id)).toBe(true); await secondJob.promise; gates.get("First")!.resolve(); await firstJob.promise; expect(secondJob.status).toBe("cancelled"); const last = updates.at(-1); // The acquire-time abort path has to flow through the same `onSettled` // the post-acquire abort path uses, otherwise the batch aggregate sticks // at "running" forever after the surviving spawn completes. expect(last?.async?.state).toBe("failed"); expect(last?.progress?.find(p => p.id === "Second")?.status).toBe("aborted"); expect(last?.progress?.find(p => p.id === "First")?.status).toBe("completed"); }); });