refactor(coding-agent): centralize agent-message admission in AgentSession (#1861)
* refactor(coding-agent): query agent peers on demand instead of broadcasting mirrors * refactor(coding-agent): centralize agent-message admission in AgentSession * fix(coding-agent): respond to unknown worker commands instead of dropping them Legacy worker_sync_agent_peers commands from an older supervisor fell through handleWorkerCommand without a response, forcing the requester to wait out its full 5s timeout on every peer sync. A default case now writes an immediate failure for any unhandled worker command. * fix(daemon): re-check cron job runnability after prompt admission fence runCronJob snapshotted the runnable job before promptHeartbeat/ promptUntilAccepted, which can wait on the session admission fence. A job cancelled or completed during that wait was still delivered. Re-check via the admissionCommitted hook and skip when the job is no longer runnable. * fix(daemon): skip cron delivery when the job content changed during admission updateRlmHeartbeat mutates prompt/deliveryMode in place under the same job id, so the existence-only admissionCommitted re-check still delivered the stale snapshot captured before the admission fence wait. Compare the refreshed job's prompt and deliveryMode and skip on mismatch; the skip reschedules nextRunAt so the updated instruction is delivered next cycle. * fix(coding-agent): re-check prompt admission after async input normalization _prompt ran admissionCommitted before _normalizeSubmission, but async extension input handlers are awaited between that check and _admitSessionInput, so a cron job cancelled or updated during the await still delivered its stale prompt. Re-run the callback after the awaited normalization so invalidated content is not admitted.
1Agent messages now use core session admission to choose immediate or queued delivery.1- Made cross-worker agent lists current without broadcasting duplicate peer rosters.10 unmodified lines11121314151617181920212223383 unmodified linesExpand all407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716386 unmodified linesExpand all10 unmodified lines type TextContent,} from "@earendil-works/pi-ai";import { Type } from "typebox";import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";import { createAgentSessionMessage } from "../src/core/agent-messages.js";import { AgentSession } from "../src/core/agent-session.js";import { AuthStorage } from "../src/core/auth-storage.js";import type { CompactionResult } from "../src/core/compaction/index.js";import type { AgentCronJob } from "../src/core/cron-jobs.js";import { ModelRegistry } from "../src/core/model-registry.js";import { SessionManager } from "../src/core/session-manager.js";import { SettingsManager } from "../src/core/settings-manager.js";import type { BuildSystemPromptOptions } from "../src/core/system-prompt.js";383 unmodified linesExpand all expect(inputCalls).toBe(0); expect(receivedUserText).toBe("agent-to-agent payload"); });
it("queues an agent message behind an active turn", async () => { createSession(); const activeTurn = session.prompt("active turn"); await vi.waitFor(() => expect(session.isStreaming).toBe(true)); let queued: boolean | undefined;
await session.acceptAgentMessagePrompt("message during turn", { queueIfBusy: true, streamingBehavior: "steer", preflightResult: (accepted, didQueue) => { expect(accepted).toBe(true); queued = didQueue; }, });
expect(queued).toBe(true); expect(session.getSessionActionSnapshot().steering).toContain("message during turn"); await session.abort(); await activeTurn.catch(() => undefined); });
it("serializes an agent message behind cron admission", async () => { createSession(); const now = new Date().toISOString(); const job: AgentCronJob = { id: "cron-overlap", status: "active", source: "cron", activeSessionId: "active-1", sessionId: session.sessionId, sessionFile: "/tmp/session.jsonl", cwd: tempDir, prompt: "cron prompt", schedule: { kind: "interval", expression: "every 5m", intervalMs: 300_000 }, createdAt: now, updatedAt: now, runCount: 0, }; const cronTurn = session.promptHeartbeat(job, { streamingBehavior: "followUp", source: "rpc" }); let queued: boolean | undefined;
await session.acceptAgentMessagePrompt("message during cron admission", { queueIfBusy: true, streamingBehavior: "steer", preflightResult: (accepted, didQueue) => { expect(accepted).toBe(true); queued = didQueue; }, });
expect(queued).toBe(true); expect(session.getSessionActionSnapshot().steering).toContain("message during cron admission"); await session.abort(); await cronTurn.catch(() => undefined); });
it("does not admit a cron prompt invalidated while async input handlers ran", async () => { const model = getModel("anthropic", "claude-sonnet-4-5")!; let deliveredUserText: string | undefined; const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: "Test", tools: [], }, streamFn: (_model, context) => { const stream = new MockAssistantStream(); queueMicrotask(() => { const user = context.messages.filter((message) => message.role === "user").at(-1); if (user && typeof user.content !== "string") { deliveredUserText = user.content .filter( (part): part is TextContent => typeof part === "object" && part !== null && part.type === "text", ) .map((part) => part.text) .join("\n"); } stream.push({ type: "start", partial: createAssistantMessage("") }); stream.push({ type: "done", reason: "stop", message: createAssistantMessage("Delivered") }); }); return stream; }, });
const sessionManager = SessionManager.inMemory(); const settingsManager = SettingsManager.create(tempDir, tempDir); const authStorage = AuthStorage.create(join(tempDir, "auth.json")); const modelRegistry = ModelRegistry.create(authStorage, tempDir); authStorage.setRuntimeApiKey("anthropic", "test-key"); let jobInvalidated = false; const extensionsResult = await createTestExtensionsResult([ (pi) => { pi.on("input", async () => { jobInvalidated = true; return undefined; }); }, ]);
session = new AgentSession({ agent, sessionManager, settingsManager, cwd: tempDir, modelRegistry, resourceLoader: createTestResourceLoader({ extensionsResult }), });
// Mirrors the daemon cron path: admissionCommitted re-checks the job and // throws when it was cancelled or updated before the prompt is admitted. await expect( session.promptUntilAccepted("stale cron prompt", { streamingBehavior: "followUp", source: "rpc", admissionCommitted: () => { if (jobInvalidated) throw new Error("Cron job became unrunnable before admission"); }, }), ).rejects.toThrow("Cron job became unrunnable before admission");
await session.agent.waitForIdle(); expect(deliveredUserText).toBeUndefined(); expect(session.isStreaming).toBe(false); });
it("serializes concurrent agent messages", async () => { createSession(); const dispositions = new Map<string, boolean | undefined>(); const accept = (message: string) => session.acceptAgentMessagePrompt(message, { queueIfBusy: true, streamingBehavior: "steer", preflightResult: (accepted, queued) => { expect(accepted).toBe(true); dispositions.set(message, queued); }, });
await Promise.all([accept("first concurrent message"), accept("second concurrent message")]);
expect(dispositions.get("first concurrent message")).toBe(false); expect(dispositions.get("second concurrent message")).toBe(true); expect(session.getSessionActionSnapshot().steering).toContain("second concurrent message"); await session.abort(); });
it("enforces the agent message queue cap inside core admission", async () => { createSession(); const accept = (index: number) => { const message = createAgentSessionMessage({ id: `agentmsg-${index}`, source: "agent_message", message: `message ${index}`, from: { clientId: "test" }, target: { activeSessionId: "target", sessionId: session.sessionId, runtimeKind: "subagent", }, }); return session.queueAgentMessagePrompt(message.content, "steer", message); };
for (let index = 0; index < 20; index++) { await accept(index); } await expect(accept(20)).rejects.toThrow("Target session has too many pending messages"); expect(session.unfinishedActionCount).toBe(20); const snapshot = session.getSessionActionRecoverySnapshot(); session.dispose();
createSession(); await expect(session.restoreSessionActions(snapshot)).resolves.toBe(20); expect(session.unfinishedActionCount).toBe(20); });
it("rejects an agent message when clear wins core admission", async () => { createSession(); const internals = session as unknown as { _acquireSessionActionCommitFence(): Promise<{ release(): void }>; }; const fence = await internals._acquireSessionActionCommitFence(); const message = createAgentSessionMessage({ id: "agentmsg-clear-race", source: "agent_message", message: "clear race", from: { clientId: "test" }, target: { activeSessionId: "target", sessionId: session.sessionId, runtimeKind: "subagent", }, }); const accepted = session.acceptAgentMessagePrompt(message.content, { customMessage: message, queueIfBusy: true, streamingBehavior: "steer", }); await Promise.resolve();
session.clearQueuedAgentMessages(); fence.release();
await expect(accepted).rejects.toThrow("Agent message was cleared before admission"); expect(session.unfinishedActionCount).toBe(0); });
it("accepts agent messages during and after direct compaction", async () => { createSession(); let releaseCompaction: (result: CompactionResult) => void = () => {}; const compactionGate = new Promise<CompactionResult>((resolve) => { releaseCompaction = resolve; }); const internals = session as unknown as { _performCompaction: () => Promise<CompactionResult>; }; vi.spyOn(internals, "_performCompaction").mockImplementation(() => compactionGate); const compaction = session.compact(); await vi.waitFor(() => expect(session.isCompacting).toBe(true)); let queued: boolean | undefined;
await session.acceptAgentMessagePrompt("message during direct compaction", { queueIfBusy: true, streamingBehavior: "steer", preflightResult: (accepted, didQueue) => { expect(accepted).toBe(true); queued = didQueue; }, });
expect(queued).toBe(true); expect(session.getSessionActionSnapshot().steering).toContain("message during direct compaction"); releaseCompaction({ summary: "compacted", firstKeptEntryId: "entry-1", tokensBefore: 1 }); await compaction; let postCompactionQueued: boolean | undefined; await session.acceptAgentMessagePrompt("message after direct compaction", { queueIfBusy: true, streamingBehavior: "steer", preflightResult: (accepted, didQueue) => { expect(accepted).toBe(true); postCompactionQueued = didQueue; }, }); expect(postCompactionQueued).toBe(true); const deliveryCount = (text: string) => session.messages.filter( (message) => message.role === "user" && typeof message.content !== "string" && message.content.some((part) => part.type === "text" && part.text === text), ).length; expect(deliveryCount("message during direct compaction")).toBe(0); expect(deliveryCount("message after direct compaction")).toBe(0);
session.resumeQueuedWork(); await vi.waitFor(() => expect(deliveryCount("message during direct compaction")).toBe(1)); expect(session.getSessionActionSnapshot().steering).toContain("message after direct compaction"); });
it("delivers an accepted agent message once after scheduler-owned compaction", async () => { createSession(); let releaseCompaction: (result: CompactionResult) => void = () => {}; const compactionGate = new Promise<CompactionResult>((resolve) => { releaseCompaction = resolve; }); const internals = session as unknown as { _performCompaction: () => Promise<CompactionResult>; }; vi.spyOn(internals, "_performCompaction").mockImplementation(() => compactionGate); const compaction = session.compact(undefined, { skipAbort: true }); await vi.waitFor(() => expect(session.isCompacting).toBe(true)); let queued: boolean | undefined;
await session.acceptAgentMessagePrompt("message during scheduled compaction", { queueIfBusy: true, streamingBehavior: "steer", preflightResult: (accepted, didQueue) => { expect(accepted).toBe(true); queued = didQueue; }, });
expect(queued).toBe(true); expect(session.getSessionActionSnapshot().steering).toContain("message during scheduled compaction"); releaseCompaction({ summary: "compacted", firstKeptEntryId: "entry-1", tokensBefore: 1 }); await compaction; const deliveryCount = () => session.messages.filter( (message) => message.role === "user" && typeof message.content !== "string" && message.content.some( (part) => part.type === "text" && part.text === "message during scheduled compaction", ), ).length;
await vi.waitFor(() => expect(deliveryCount()).toBe(1)); expect(deliveryCount()).toBe(1); });
it("delivers internal prompts without extension input interception", async () => { const model = getModel("anthropic", "claude-sonnet-4-5")!; let inputCalls = 0; let receivedUserText: string | undefined;386 unmodified linesExpand all