roboomp
05668faa5f
The flat single-spawn task wire schema carries arktype `"+": "delete"`, so a
batch `{ context, tasks[] }` payload sent while `task.batch` is disabled has
those keys stripped and is then rejected as `task must be a string (was
missing)` in the agent loop. That preempts the tool's own actionable checks
(validateShapeParams / validateSpawnParams), so the model only ever saw the
misleading arktype error instead of "task.batch is disabled…".
Mark TaskTool with lenientArgValidation so the agent loop forwards the raw
args to execute() on any arktype failure, letting the tool's shape checks
surface the real reason. Valid calls still normalize through arktype; the
success path is unchanged. Mirrors the existing yield-tool pattern.
Fixes #6039
517 lines
19 KiB
TypeScript
517 lines
19 KiB
TypeScript
/**
|
|
* 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`, `outputSchema`, and
|
|
* `schemaMode` — live in the items), the flat form exposes those fields
|
|
* directly. The stale `schema` field is never accepted.
|
|
* 2. Shape validation rejects stale `schema`, `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<string, unknown>;
|
|
agentId?: string;
|
|
planMode?: boolean;
|
|
} = {},
|
|
): ToolSession {
|
|
return {
|
|
cwd: "/tmp",
|
|
hasUI: false,
|
|
settings: Settings.isolated(options.settings ?? {}),
|
|
getSessionFile: () => null,
|
|
getSessionSpawns: () => "*",
|
|
getAgentId: () => options.agentId ?? null,
|
|
getPlanModeState: options.planMode ? () => ({ enabled: true }) : undefined,
|
|
asyncJobManager: options.manager,
|
|
} as unknown as ToolSession;
|
|
}
|
|
|
|
function getSchemaProperties(tool: TaskTool): Record<string, unknown> {
|
|
const wire = toolWireSchema(tool) as { properties?: Record<string, unknown> };
|
|
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> = {}): 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(agent: AgentDefinition = taskAgent): void {
|
|
vi.spyOn(discoveryModule, "discoverAgents").mockResolvedValue({
|
|
agents: [agent],
|
|
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();
|
|
expect(offProperties.outputSchema).toBeDefined();
|
|
expect(typeof offProperties.outputSchema).toBe("object");
|
|
expect(offProperties.schemaMode).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();
|
|
expect(onProperties.outputSchema).toBeUndefined();
|
|
expect(onProperties.schemaMode).toBeUndefined();
|
|
const items = (onProperties.tasks as { items?: { properties?: Record<string, unknown> } }).items;
|
|
expect(items?.properties?.task).toBeDefined();
|
|
expect(items?.properties?.name).toBeDefined();
|
|
expect(items?.properties?.agent).toBeDefined();
|
|
expect(items?.properties?.outputSchema).toBeDefined();
|
|
expect(typeof items?.properties?.outputSchema).toBe("object");
|
|
expect(items?.properties?.schemaMode).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<string, unknown> } }).items;
|
|
expect(items?.properties?.isolated).toBeDefined();
|
|
});
|
|
|
|
it("hides isolation from the dynamic batch schema in plan mode", async () => {
|
|
mockDiscovery();
|
|
const tool = await TaskTool.create(
|
|
createSession({
|
|
planMode: true,
|
|
settings: { "task.batch": true, "task.isolation.mode": "auto" },
|
|
}),
|
|
);
|
|
const properties = getSchemaProperties(tool);
|
|
const items = (properties.tasks as { items?: { properties?: Record<string, unknown> } }).items;
|
|
expect(items?.properties?.isolated).toBeUndefined();
|
|
expect(tool.description).not.toContain("`isolated`");
|
|
});
|
|
|
|
it("exposes outputSchema but never the stale schema field", async () => {
|
|
mockDiscovery();
|
|
|
|
const flat = await TaskTool.create(createSession({ settings: { "task.batch": false } }));
|
|
expect(getSchemaProperties(flat).outputSchema).toBeDefined();
|
|
expect(getSchemaProperties(flat).schema).toBeUndefined();
|
|
|
|
const batch = await TaskTool.create(createSession({ settings: { "task.batch": true } }));
|
|
expect(getSchemaProperties(batch).schema).toBeUndefined();
|
|
});
|
|
});
|
|
|
|
describe("task.batch validation", () => {
|
|
afterEach(() => {
|
|
vi.restoreAllMocks();
|
|
});
|
|
|
|
async function executeText(params: unknown, settings: Record<string, unknown> = {}): Promise<string> {
|
|
mockDiscovery();
|
|
const tool = await TaskTool.create(createSession({ settings }));
|
|
const result = await tool.execute("tool-call", params);
|
|
return getFirstText(result);
|
|
}
|
|
|
|
it("rejects the stale 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("uses `outputSchema`");
|
|
}
|
|
});
|
|
|
|
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");
|
|
});
|
|
|
|
it("marks lenientArgValidation so execute() surfaces the actionable shape error", async () => {
|
|
// Regression (#6039): the flat single-spawn wire schema carries
|
|
// `"+": "delete"`, so a batch `{ context, tasks[] }` payload is stripped
|
|
// by arktype and rejected as `task must be a string (was missing)` in the
|
|
// agent loop — preempting the tool's own actionable message. The lenient
|
|
// flag makes the loop forward the raw args to execute() on that failure.
|
|
mockDiscovery();
|
|
const tool = await TaskTool.create(createSession({ settings: { "task.batch": false } }));
|
|
expect(tool.lenientArgValidation).toBe(true);
|
|
|
|
// The raw batch payload the loop would forward reaches execute() and
|
|
// yields the actionable reason, never arktype's misleading missing-`task`.
|
|
const text = await executeText(
|
|
{ context: "Background.", tasks: [{ name: "Alpha", task: "Work." }] },
|
|
{ "task.batch": false },
|
|
);
|
|
expect(text).toContain("task.batch is disabled");
|
|
expect(text).not.toContain("was missing");
|
|
});
|
|
});
|
|
|
|
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 independent schemas with shared context", async () => {
|
|
mockDiscovery({
|
|
...taskAgent,
|
|
output: { type: "object", properties: { staleAgentOutput: { type: "boolean" } } },
|
|
});
|
|
const seen: Array<{
|
|
id?: string;
|
|
context?: string;
|
|
assignment?: string;
|
|
parentAgentId?: string;
|
|
outputSchema?: unknown;
|
|
outputSchemaMode?: "permissive" | "strict";
|
|
outputSchemaSource?: "caller" | "agent" | "session" | "none";
|
|
outputSchemaOverridesAgent?: boolean;
|
|
}> = [];
|
|
vi.spyOn(executorModule, "runSubprocess").mockImplementation(async options => {
|
|
seen.push({
|
|
id: options.id,
|
|
context: options.context,
|
|
assignment: options.assignment,
|
|
parentAgentId: options.parentAgentId,
|
|
outputSchema: options.outputSchema,
|
|
outputSchemaMode: options.outputSchemaMode,
|
|
outputSchemaSource: options.outputSchemaSource,
|
|
outputSchemaOverridesAgent: options.outputSchemaOverridesAgent,
|
|
});
|
|
return makeResult(options.id ?? "?");
|
|
});
|
|
|
|
const manager = createManager();
|
|
const tool = await TaskTool.create(
|
|
createSession({ manager, agentId: "ParentA", settings: { "async.enabled": true, "task.batch": true } }),
|
|
);
|
|
const alphaSchema = { type: "object", properties: { alpha: { type: "string" } } };
|
|
const betaSchema = { type: "object", properties: { beta: { type: "number" } } };
|
|
const result = await tool.execute("tc-batch", {
|
|
context: "# Goal\nShared background.",
|
|
tasks: [
|
|
{ name: "Alpha", task: "Do A.", outputSchema: alphaSchema, schemaMode: "strict" },
|
|
{ name: "Beta", task: "Do B.", outputSchema: betaSchema, schemaMode: "permissive" },
|
|
],
|
|
} 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(seen).toHaveLength(2);
|
|
for (const spawn of seen) {
|
|
expect(spawn.context).toBe("# Goal\nShared background.");
|
|
expect(spawn.outputSchemaSource).toBe("caller");
|
|
expect(spawn.outputSchemaOverridesAgent).toBe(true);
|
|
}
|
|
const byId = new Map(seen.map(spawn => [spawn.id, spawn]));
|
|
expect(byId.get("Alpha")?.outputSchema).toEqual(alphaSchema);
|
|
expect(byId.get("Alpha")?.outputSchemaMode).toBe("strict");
|
|
expect(byId.get("Beta")?.outputSchema).toEqual(betaSchema);
|
|
expect(byId.get("Beta")?.outputSchemaMode).toBe("permissive");
|
|
expect(seen.map(spawn => spawn.assignment).sort()).toEqual(["Do A.", "Do B."]);
|
|
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({ ...taskAgent, output: { type: "object", properties: { agent: { type: "string" } } } });
|
|
let captured:
|
|
| {
|
|
outputSchema?: unknown;
|
|
outputSchemaMode?: "permissive" | "strict";
|
|
outputSchemaSource?: "caller" | "agent" | "session" | "none";
|
|
outputSchemaOverridesAgent?: boolean;
|
|
}
|
|
| undefined;
|
|
vi.spyOn(executorModule, "runSubprocess").mockImplementation(async options => {
|
|
captured = {
|
|
outputSchema: options.outputSchema,
|
|
outputSchemaMode: options.outputSchemaMode,
|
|
outputSchemaSource: options.outputSchemaSource,
|
|
outputSchemaOverridesAgent: options.outputSchemaOverridesAgent,
|
|
};
|
|
return makeResult(options.id ?? "?");
|
|
});
|
|
|
|
const manager = createManager();
|
|
const tool = await TaskTool.create(
|
|
createSession({ manager, settings: { "async.enabled": true, "task.batch": true } }),
|
|
);
|
|
|
|
const callerSchema = { type: "object", properties: { caller: { type: "number" } } };
|
|
const result = await tool.execute("tc-flat", {
|
|
agent: "task",
|
|
name: "Flat",
|
|
task: "Do the thing.",
|
|
outputSchema: callerSchema,
|
|
schemaMode: "strict",
|
|
} 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");
|
|
expect(captured?.outputSchema).toEqual(callerSchema);
|
|
expect(captured?.outputSchemaMode).toBe("strict");
|
|
expect(captured?.outputSchemaSource).toBe("caller");
|
|
expect(captured?.outputSchemaOverridesAgent).toBe(true);
|
|
});
|
|
|
|
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<string, { promise: Promise<void>; resolve: () => void }>();
|
|
vi.spyOn(executorModule, "runSubprocess").mockImplementation(async options => {
|
|
const id = options.id ?? "?";
|
|
started.push(id);
|
|
const { promise, resolve } = Promise.withResolvers<void>();
|
|
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");
|
|
});
|
|
});
|