Md Tahsin Rahman
41dfc2cd9e
SessionManager.open writes title+session before createAgentSession claims the id. Agent Hub parked that stub, so the spawn CAS failed with "already owned by another session generation" and the row could not be revived (no session_init).
508 lines
19 KiB
TypeScript
508 lines
19 KiB
TypeScript
import { afterEach, beforeEach, describe, expect, it, vi } from "bun:test";
|
|
import { ASYNC_JOB_MANAGER_SHUTDOWN_REASON, AsyncJobManager } from "@oh-my-pi/pi-coding-agent/async";
|
|
import type { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry";
|
|
import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
|
|
import type { LoadExtensionsResult } from "@oh-my-pi/pi-coding-agent/extensibility/extensions/types";
|
|
import { IrcBus } from "@oh-my-pi/pi-coding-agent/irc/bus";
|
|
import { RpcSubagentRegistry } from "@oh-my-pi/pi-coding-agent/modes/rpc/rpc-subagents";
|
|
import type { RpcSubagentFrame } from "@oh-my-pi/pi-coding-agent/modes/rpc/rpc-types";
|
|
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 { registerPersistedSubagents } from "@oh-my-pi/pi-coding-agent/registry/persisted-agents";
|
|
import type { CreateAgentSessionResult } from "@oh-my-pi/pi-coding-agent/sdk";
|
|
import * as sdkModule from "@oh-my-pi/pi-coding-agent/sdk";
|
|
import type { AgentSession, AgentSessionEvent, PromptOptions } from "@oh-my-pi/pi-coding-agent/session/agent-session";
|
|
import type { CustomMessage } from "@oh-my-pi/pi-coding-agent/session/messages";
|
|
import { resolveSoftRequestBudget, runSubprocess } from "@oh-my-pi/pi-coding-agent/task/executor";
|
|
import type { AgentDefinition } from "@oh-my-pi/pi-coding-agent/task/types";
|
|
import { EventBus } from "@oh-my-pi/pi-coding-agent/utils/event-bus";
|
|
import { TempDir } from "@oh-my-pi/pi-utils";
|
|
|
|
/**
|
|
* Contracts under test — the soft request budget must degrade gracefully
|
|
* instead of killing scouts into an unreachable state:
|
|
*
|
|
* 1. Crossing 1.5x the budget stops the free-running turn and drives ONE
|
|
* forced final `yield`, so the run finishes as a normal completion with a
|
|
* partial report — not as an abort with no output.
|
|
* 2. If the agent still refuses to yield (grace exhausted → hard abort), a
|
|
* kept-alive agent stays adopted (`idle`), so `irc` can message/resume it
|
|
* with the same RPC lifecycle/progress frames as the original run.
|
|
* 3. Caller-signal aborts remain terminal, and the irc bus names the aborted
|
|
* agent precisely instead of claiming it is unknown.
|
|
*/
|
|
|
|
interface MockSessionHandle {
|
|
session: AgentSession;
|
|
prompts: Array<{ text: string; options?: PromptOptions }>;
|
|
abortCalls: () => number;
|
|
disposeCalls: () => number;
|
|
}
|
|
|
|
function assistantText(text: string, stopReason: "stop" | "aborted" = "stop") {
|
|
return { role: "assistant" as const, content: [{ type: "text" as const, text }], stopReason };
|
|
}
|
|
|
|
function createMockSession(
|
|
onPrompt: (params: {
|
|
promptIndex: number;
|
|
emit: (event: AgentSessionEvent) => void;
|
|
pushMessage: (message: unknown) => void;
|
|
}) => void | Promise<void>,
|
|
onAbort?: () => void | Promise<void>,
|
|
): MockSessionHandle {
|
|
const listeners: Array<(event: AgentSessionEvent) => void> = [];
|
|
const messages: unknown[] = [];
|
|
const prompts: Array<{ text: string; options?: PromptOptions }> = [];
|
|
let abortCount = 0;
|
|
let disposeCount = 0;
|
|
let promptIndex = 0;
|
|
let ircWakeTurnObserver:
|
|
| ((records: CustomMessage[]) => ((error?: unknown) => void | Promise<void>) | undefined)
|
|
| undefined;
|
|
|
|
const emit = (event: AgentSessionEvent) => {
|
|
for (const listener of [...listeners]) listener(event);
|
|
};
|
|
|
|
const session: Partial<AgentSession> = {
|
|
state: { messages: [] } as never,
|
|
agent: { state: { systemPrompt: ["test"] } } as never,
|
|
model: { api: "anthropic-messages" } as never,
|
|
extensionRunner: undefined as never,
|
|
sessionManager: { appendSessionInit: () => {} } as never,
|
|
getActiveToolNames: () => ["read", "yield"],
|
|
getEnabledToolNames: () => ["read", "yield"],
|
|
setActiveToolsByName: async () => {},
|
|
subscribe: (listener: (event: AgentSessionEvent) => void) => {
|
|
listeners.push(listener);
|
|
return () => {
|
|
const index = listeners.indexOf(listener);
|
|
if (index >= 0) listeners.splice(index, 1);
|
|
};
|
|
},
|
|
prompt: async (text: string, options?: PromptOptions) => {
|
|
promptIndex += 1;
|
|
prompts.push({ text, options });
|
|
await onPrompt({ promptIndex, emit, pushMessage: message => messages.push(message) });
|
|
return true;
|
|
},
|
|
waitForIdle: async () => {},
|
|
getLastAssistantMessage: () => messages[messages.length - 1] as never,
|
|
sendUserMessage: async () => {},
|
|
setIrcWakeTurnObserver: observer => {
|
|
ircWakeTurnObserver = observer;
|
|
},
|
|
deliverIrcMessage: async msg => {
|
|
const record: CustomMessage = {
|
|
role: "custom",
|
|
customType: "irc:incoming",
|
|
content: msg.body,
|
|
display: true,
|
|
details: { id: msg.id, from: msg.from, message: msg.body },
|
|
attribution: "agent",
|
|
timestamp: msg.ts,
|
|
};
|
|
const finishObservation = ircWakeTurnObserver?.([record]);
|
|
const yieldMessage = {
|
|
role: "assistant" as const,
|
|
content: [
|
|
{
|
|
type: "toolCall" as const,
|
|
id: "tool-irc-yield",
|
|
name: "yield",
|
|
arguments: { result: { data: { report: "resumed findings" } } },
|
|
},
|
|
],
|
|
stopReason: "toolUse" as const,
|
|
};
|
|
messages.push(yieldMessage);
|
|
emit({ type: "agent_start" } as AgentSessionEvent);
|
|
emit({ type: "message_end", message: yieldMessage } as unknown as AgentSessionEvent);
|
|
emit({
|
|
type: "tool_execution_end",
|
|
toolCallId: "tool-irc-yield",
|
|
toolName: "yield",
|
|
result: {
|
|
content: [{ type: "text", text: "Result submitted." }],
|
|
details: { status: "success", data: { report: "resumed findings" } },
|
|
},
|
|
isError: false,
|
|
} as AgentSessionEvent);
|
|
emit({ type: "agent_end", messages: [yieldMessage] } as unknown as AgentSessionEvent);
|
|
await finishObservation?.();
|
|
return "woken";
|
|
},
|
|
abort: async () => {
|
|
abortCount += 1;
|
|
await onAbort?.();
|
|
},
|
|
dispose: async () => {
|
|
disposeCount += 1;
|
|
},
|
|
};
|
|
|
|
return {
|
|
session: session as AgentSession,
|
|
prompts,
|
|
abortCalls: () => abortCount,
|
|
disposeCalls: () => disposeCount,
|
|
};
|
|
}
|
|
|
|
function mockCreateAgentSession(session: AgentSession) {
|
|
return vi.spyOn(sdkModule, "createAgentSession").mockResolvedValue({
|
|
session,
|
|
extensionsResult: {} as unknown as LoadExtensionsResult,
|
|
setToolUIContext: () => {},
|
|
eventBus: new EventBus(),
|
|
} satisfies CreateAgentSessionResult);
|
|
}
|
|
// Use a bundled scout so these runSubprocess tests exercise the built-in
|
|
// ceiling together with a lower task.softRequestBudget setting.
|
|
const baseAgent: AgentDefinition = {
|
|
name: "scout",
|
|
description: "test",
|
|
systemPrompt: "test",
|
|
source: "bundled",
|
|
};
|
|
|
|
describe("runSubprocess soft request budget", () => {
|
|
let tempDir: TempDir;
|
|
|
|
beforeEach(() => {
|
|
AgentRegistry.resetGlobalForTests();
|
|
AgentLifecycleManager.resetGlobalForTests();
|
|
AsyncJobManager.resetForTests();
|
|
tempDir = TempDir.createSync("@pi-soft-budget-");
|
|
});
|
|
afterEach(() => {
|
|
vi.restoreAllMocks();
|
|
AgentLifecycleManager.resetGlobalForTests();
|
|
AgentRegistry.resetGlobalForTests();
|
|
AsyncJobManager.resetForTests();
|
|
tempDir[Symbol.dispose]();
|
|
});
|
|
|
|
function baseOptions(id: string, eventBus?: EventBus) {
|
|
return {
|
|
cwd: "/tmp",
|
|
agent: baseAgent,
|
|
task: "inventory the api surface",
|
|
index: 0,
|
|
id,
|
|
settings: Settings.isolated({ "task.softRequestBudget": 2 }),
|
|
modelRegistry: { refresh: async () => {} } as unknown as ModelRegistry,
|
|
enableLsp: false,
|
|
artifactsDir: tempDir.path(),
|
|
eventBus,
|
|
};
|
|
}
|
|
|
|
function registerRunning(id: string, session: AgentSession, sessionFile: string | null = null) {
|
|
AgentRegistry.global().register({
|
|
id,
|
|
displayName: id,
|
|
kind: "sub",
|
|
session,
|
|
sessionFile,
|
|
status: "running",
|
|
});
|
|
}
|
|
|
|
it("a budget stop drives one forced final yield and finishes as a normal completion", async () => {
|
|
const id = "BudgetScout";
|
|
let abortCallsAtReminder: number | undefined;
|
|
const handle = createMockSession(({ promptIndex, emit, pushMessage }) => {
|
|
if (promptIndex === 1) {
|
|
// Free-running exploration: budget 2 → stop threshold 3.
|
|
for (let i = 1; i <= 3; i++) {
|
|
const message = assistantText(`exploring ${i}`, i === 3 ? "aborted" : "stop");
|
|
pushMessage(message);
|
|
emit({ type: "message_end", message } as unknown as AgentSessionEvent);
|
|
}
|
|
return;
|
|
}
|
|
// The forced wrap-up reminder: answer it with a terminal yield.
|
|
abortCallsAtReminder = handle.abortCalls();
|
|
const yieldMessage = {
|
|
role: "assistant" as const,
|
|
content: [
|
|
{
|
|
type: "toolCall" as const,
|
|
id: "tool-forced-yield",
|
|
name: "yield",
|
|
arguments: { result: { data: { report: "partial findings" } } },
|
|
},
|
|
],
|
|
stopReason: "toolUse" as const,
|
|
};
|
|
pushMessage(yieldMessage);
|
|
emit({ type: "message_end", message: yieldMessage } as unknown as AgentSessionEvent);
|
|
emit({
|
|
type: "tool_execution_end",
|
|
toolCallId: "tool-forced-yield",
|
|
toolName: "yield",
|
|
result: {
|
|
content: [{ type: "text", text: "Result submitted." }],
|
|
details: { status: "success", data: { report: "partial findings" } },
|
|
},
|
|
isError: false,
|
|
} as AgentSessionEvent);
|
|
});
|
|
mockCreateAgentSession(handle.session);
|
|
registerRunning(id, handle.session);
|
|
|
|
const result = await runSubprocess(baseOptions(id));
|
|
|
|
// The budget stop aborted the free-running turn exactly once before the
|
|
// wrap-up reminder; the second abort (after the terminal yield) is the
|
|
// normal post-yield terminate.
|
|
expect(abortCallsAtReminder).toBe(1);
|
|
// The budget stop forces a synthetic terminal yield.
|
|
expect(handle.prompts).toHaveLength(2);
|
|
expect(handle.prompts[1]?.options?.synthetic).toBe(true);
|
|
expect(handle.prompts[1]?.options?.toolChoice).toEqual({ type: "tool", name: "yield" });
|
|
// The forced yield finalizes as a normal completion, not an abort.
|
|
expect(result.aborted).toBe(false);
|
|
expect(result.exitCode).toBe(0);
|
|
expect(result.abortReason).toBeUndefined();
|
|
expect(JSON.parse(result.output)).toEqual({ report: "partial findings" });
|
|
// The agent stays a live, adopted peer.
|
|
expect(AgentRegistry.global().get(id)?.status).toBe("idle");
|
|
expect(AgentLifecycleManager.global().has(id)).toBe(true);
|
|
expect(handle.disposeCalls()).toBe(0);
|
|
});
|
|
|
|
it("a budget hard-abort keeps the kept-alive agent adopted and messageable via irc", async () => {
|
|
const id = "StubbornScout";
|
|
const eventBus = new EventBus();
|
|
const frames: RpcSubagentFrame[] = [];
|
|
let resolveFollowUpTerminal: (() => void) | undefined;
|
|
const waitForFollowUpTerminal = (): Promise<void> => {
|
|
const deferred = Promise.withResolvers<void>();
|
|
resolveFollowUpTerminal = deferred.resolve;
|
|
return deferred.promise;
|
|
};
|
|
const rpcRegistry = new RpcSubagentRegistry(eventBus, frame => {
|
|
frames.push(frame);
|
|
if (frame.type !== "subagent_lifecycle" || frame.payload.status === "started") return;
|
|
const resolve = resolveFollowUpTerminal;
|
|
resolveFollowUpTerminal = undefined;
|
|
resolve?.();
|
|
});
|
|
rpcRegistry.setSubscriptionLevel("progress");
|
|
const handle = createMockSession(({ promptIndex, emit, pushMessage }) => {
|
|
if (promptIndex !== 1) return;
|
|
// Never yields: budget 2 → stop at 3, grace exhausted at 3 + 5 = 8.
|
|
for (let i = 1; i <= 8; i++) {
|
|
const message = assistantText(`burning request ${i}`);
|
|
pushMessage(message);
|
|
emit({ type: "message_end", message } as unknown as AgentSessionEvent);
|
|
}
|
|
});
|
|
mockCreateAgentSession(handle.session);
|
|
registerRunning(id, handle.session);
|
|
|
|
const result = await runSubprocess(baseOptions(id, eventBus));
|
|
|
|
expect(result.aborted).toBe(true);
|
|
expect(result.abortReason).toMatch(/Soft request budget exceeded/);
|
|
// Resumable stop, not a terminal kill: the ref stays adopted and live.
|
|
expect(AgentRegistry.global().get(id)?.status).toBe("idle");
|
|
expect(AgentLifecycleManager.global().has(id)).toBe(true);
|
|
expect(handle.disposeCalls()).toBe(0);
|
|
|
|
const expectRpcTurn = (): void => {
|
|
expect(frames[0]).toMatchObject({
|
|
type: "subagent_lifecycle",
|
|
payload: { id, status: "started" },
|
|
});
|
|
expect(frames.some(frame => frame.type === "subagent_progress")).toBe(true);
|
|
expect(frames.at(-1)).toMatchObject({
|
|
type: "subagent_lifecycle",
|
|
payload: { id, status: "completed" },
|
|
});
|
|
};
|
|
|
|
frames.length = 0;
|
|
const idleTerminal = waitForFollowUpTerminal();
|
|
const idleReceipt = await new IrcBus().send({ from: "Main", to: id, body: "resume your inventory" });
|
|
expect(idleReceipt.outcome).toBe("woken");
|
|
await idleTerminal;
|
|
expectRpcTurn();
|
|
|
|
await AgentLifecycleManager.global().park(id);
|
|
expect(AgentRegistry.global().get(id)?.status).toBe("parked");
|
|
frames.length = 0;
|
|
const revivedTerminal = waitForFollowUpTerminal();
|
|
const revivedReceipt = await new IrcBus().send({ from: "Main", to: id, body: "resume after parking" });
|
|
expect(revivedReceipt.outcome).toBe("revived");
|
|
await revivedTerminal;
|
|
expectRpcTurn();
|
|
rpcRegistry.dispose();
|
|
});
|
|
|
|
it("a shutdown racing a budget hard-abort follows the shutdown release path", async () => {
|
|
// Regression: a process shutdown that lands right after the soft-budget
|
|
// grace hard-aborts must supersede the budget reason, so the subagent is
|
|
// released (disposed + unregistered, restorable as parked) instead of
|
|
// being left adopted and alive past AgentLifecycleManager.dispose().
|
|
const id = "RacedScout";
|
|
const rootSessionFile = `${tempDir.path()}/main.jsonl`;
|
|
const workerSessionFile = `${tempDir.path()}/main/${id}.jsonl`;
|
|
await Bun.write(rootSessionFile, "");
|
|
await Bun.write(
|
|
workerSessionFile,
|
|
[
|
|
JSON.stringify({ type: "session", version: 3, id, timestamp: "2026-08-13T17:14:48.000Z", cwd: "/tmp" }),
|
|
JSON.stringify({
|
|
type: "session_init",
|
|
id: "si",
|
|
parentId: null,
|
|
timestamp: "2026-08-13T17:14:48.000Z",
|
|
systemPrompt: "system",
|
|
task: "work",
|
|
tools: ["read"],
|
|
}),
|
|
].join("\n"),
|
|
);
|
|
const controller = new AbortController();
|
|
// abort #1 = budget soft-stop (abortSent still false); abort #2 =
|
|
// budget hard-abort's abortActiveSession (abortReason already "budget").
|
|
// Fire the shutdown only on #2 so it must supersede the budget reason.
|
|
let abortInvocations = 0;
|
|
const handle = createMockSession(
|
|
({ promptIndex, emit, pushMessage }) => {
|
|
if (promptIndex !== 1) return;
|
|
// Never yields: budget 2 → stop at 3, grace exhausted at 8.
|
|
for (let i = 1; i <= 8; i++) {
|
|
const message = assistantText(`burning request ${i}`);
|
|
pushMessage(message);
|
|
emit({ type: "message_end", message } as unknown as AgentSessionEvent);
|
|
}
|
|
},
|
|
() => {
|
|
abortInvocations += 1;
|
|
if (abortInvocations >= 2 && !controller.signal.aborted) {
|
|
controller.abort(ASYNC_JOB_MANAGER_SHUTDOWN_REASON);
|
|
}
|
|
},
|
|
);
|
|
mockCreateAgentSession(handle.session);
|
|
registerRunning(id, handle.session, workerSessionFile);
|
|
|
|
const result = await runSubprocess({ ...baseOptions(id), signal: controller.signal });
|
|
|
|
expect(result.aborted).toBe(true);
|
|
expect(AgentRegistry.global().get(id)).toBeUndefined();
|
|
expect(handle.disposeCalls()).toBeGreaterThanOrEqual(1);
|
|
expect(await Bun.file(`${workerSessionFile}.tombstone`).exists()).toBe(false);
|
|
const restored = new AgentRegistry();
|
|
await registerPersistedSubagents(restored, rootSessionFile);
|
|
expect(restored.get(id)?.status).toBe("parked");
|
|
});
|
|
|
|
it("manager shutdown restores a running kept-alive agent as parked without a tombstone", async () => {
|
|
const id = "ShutdownScout";
|
|
const rootSessionFile = `${tempDir.path()}/main.jsonl`;
|
|
const workerSessionFile = `${tempDir.path()}/main/${id}.jsonl`;
|
|
await Bun.write(rootSessionFile, "");
|
|
await Bun.write(
|
|
workerSessionFile,
|
|
[
|
|
JSON.stringify({ type: "session", version: 3, id, timestamp: "2026-08-13T17:14:48.000Z", cwd: "/tmp" }),
|
|
JSON.stringify({
|
|
type: "session_init",
|
|
id: "si",
|
|
parentId: null,
|
|
timestamp: "2026-08-13T17:14:48.000Z",
|
|
systemPrompt: "system",
|
|
task: "work",
|
|
tools: ["read"],
|
|
}),
|
|
].join("\n"),
|
|
);
|
|
const promptStarted = Promise.withResolvers<void>();
|
|
const promptStopped = Promise.withResolvers<void>();
|
|
const handle = createMockSession(
|
|
async ({ promptIndex }) => {
|
|
if (promptIndex !== 1) return;
|
|
promptStarted.resolve();
|
|
await promptStopped.promise;
|
|
},
|
|
() => promptStopped.resolve(),
|
|
);
|
|
mockCreateAgentSession(handle.session);
|
|
registerRunning(id, handle.session, workerSessionFile);
|
|
const manager = new AsyncJobManager({ maxRunningJobs: 1 });
|
|
AsyncJobManager.setInstance(manager);
|
|
manager.register(
|
|
"task",
|
|
"shutdown regression",
|
|
async ({ signal }) => {
|
|
const result = await runSubprocess({ ...baseOptions(id), signal });
|
|
return result.output;
|
|
},
|
|
{ ownerId: "Main", agentId: id },
|
|
);
|
|
|
|
await promptStarted.promise;
|
|
await manager.dispose({ timeoutMs: 1_000 });
|
|
AsyncJobManager.setInstance(undefined);
|
|
|
|
expect(await Bun.file(`${workerSessionFile}.tombstone`).exists()).toBe(false);
|
|
expect(AgentRegistry.global().get(id)).toBeUndefined();
|
|
const restoredRegistry = new AgentRegistry();
|
|
await registerPersistedSubagents(restoredRegistry, rootSessionFile);
|
|
expect(restoredRegistry.get(id)?.status).toBe("parked");
|
|
});
|
|
|
|
it("a caller-signal abort stays terminal and irc names the aborted agent precisely", async () => {
|
|
const id = "CancelledScout";
|
|
const controller = new AbortController();
|
|
const handle = createMockSession(({ promptIndex, emit, pushMessage }) => {
|
|
if (promptIndex !== 1) return;
|
|
const message = assistantText("working");
|
|
pushMessage(message);
|
|
emit({ type: "message_end", message } as unknown as AgentSessionEvent);
|
|
controller.abort();
|
|
});
|
|
mockCreateAgentSession(handle.session);
|
|
registerRunning(id, handle.session);
|
|
|
|
const result = await runSubprocess({ ...baseOptions(id), signal: controller.signal });
|
|
|
|
expect(result.aborted).toBe(true);
|
|
expect(AgentRegistry.global().get(id)?.status).toBe("aborted");
|
|
expect(handle.disposeCalls()).toBeGreaterThanOrEqual(1);
|
|
|
|
const receipt = await new IrcBus().send({ from: "Main", to: id, body: "resume" });
|
|
expect(receipt.outcome).toBe("failed");
|
|
expect(receipt.error).toMatch(/hard-aborted/);
|
|
expect(receipt.error).toMatch(new RegExp(`history://${id}`));
|
|
});
|
|
});
|
|
|
|
describe("resolveSoftRequestBudget", () => {
|
|
it("lets a configured budget lower a bundled agent's ceiling", () => {
|
|
expect(resolveSoftRequestBudget("scout", 20)).toBe(20);
|
|
expect(resolveSoftRequestBudget("sonic", 20)).toBe(20);
|
|
});
|
|
|
|
it("keeps the bundled ceiling when the configured budget is higher", () => {
|
|
expect(resolveSoftRequestBudget("scout", 200)).toBe(100);
|
|
expect(resolveSoftRequestBudget("sonic", 200)).toBe(100);
|
|
});
|
|
|
|
it("uses the configured budget for agents without a bundled entry", () => {
|
|
expect(resolveSoftRequestBudget("task", 20)).toBe(20);
|
|
});
|
|
|
|
it("keeps 0 disabled and normalizes negative or fractional budgets", () => {
|
|
expect(resolveSoftRequestBudget("scout", 0)).toBe(0);
|
|
expect(resolveSoftRequestBudget("scout", -5)).toBe(0);
|
|
expect(resolveSoftRequestBudget("scout", 20.9)).toBe(20);
|
|
});
|
|
});
|