diff --git a/src/main/worker/entry/index.ts b/src/main/worker/entry/index.ts index 908df9c9..0cfe3301 100644 --- a/src/main/worker/entry/index.ts +++ b/src/main/worker/entry/index.ts @@ -90,11 +90,14 @@ async function startRun(job: RunJob) { await runner.run(job); } finally { cancelledRuns.delete(job.runId); - active.delete(job.runId); - tickQueue(); } } +function acknowledgeRunCompletion(runId: string) { + const released = active.delete(runId); + if (released) tickQueue(); +} + function cancelRun(runId: string) { const queuedIndex = queue.findIndex((entry) => entry.runId === runId); if (queuedIndex >= 0) { @@ -129,6 +132,10 @@ parentPort.on("message", (event) => { cancelRun(message.runId); return; } + if (message.type === "run-completion-ack") { + acknowledgeRunCompletion(message.runId); + return; + } if (message.type === "steer-run") { const accepted = runner.steer(message.runId, message.message); post({ diff --git a/src/main/worker/entry/message-handler.ts b/src/main/worker/entry/message-handler.ts index 42059428..f5612cc6 100644 --- a/src/main/worker/entry/message-handler.ts +++ b/src/main/worker/entry/message-handler.ts @@ -3,8 +3,10 @@ import type { workerMessageSchema } from "./types"; import { handleRunLog } from "../handlers/run-log-handler"; import { handleRunStatus } from "../handlers/run-status-handler"; import { handleRunFinished } from "../handlers/run-finished-handler"; +import { releaseCompletedRunWorkspaceLock } from "../handlers/run-finished-helpers"; import { handleContextCompaction } from "../handlers/context-compaction-handler"; import type { HandlerContext, HandlerDeps } from "../handlers/handler-context"; +import type { RunFinishedContext } from "../handlers/run-finished-helpers"; export async function handleWorkerMessage( message: unknown, @@ -47,11 +49,17 @@ export async function handleWorkerMessage( handleRunStatus(event, deps, context); return; case "run-finished": { + let completion: RunFinishedContext | null = null; try { - handleRunFinished(event, deps, context); + completion = handleRunFinished(event, deps, context); } finally { deps.getRuntimeBridge?.()?.releaseRunResources(event.runId); } + releaseCompletedRunWorkspaceLock(completion, deps); + if (!worker) { + throw new Error(`run-completion-ack worker missing for run ${event.runId}`); + } + worker.postMessage({ type: "run-completion-ack", runId: event.runId }); return; } case "agent-event": { diff --git a/src/main/worker/handlers/run-finished-handler.ts b/src/main/worker/handlers/run-finished-handler.ts index afbb53fb..97f44a38 100644 --- a/src/main/worker/handlers/run-finished-handler.ts +++ b/src/main/worker/handlers/run-finished-handler.ts @@ -6,17 +6,19 @@ import { triggerSelfWake, } from "./run-finished-helpers"; import type { HandlerContext, HandlerDeps, WorkerMessageEvent } from "./handler-context"; +import type { RunFinishedContext } from "./run-finished-helpers"; export function handleRunFinished( event: Extract, deps: HandlerDeps, context: HandlerContext, -): void { +): RunFinishedContext | null { const rfCtx = finalizeRun(event, deps, context); - if (!rfCtx) return; + if (!rfCtx) return null; accumulateAndRotateSession(event, deps, rfCtx); emitAndNotifyCompletion(event, deps, rfCtx); const { hadDeferredWakes } = runPostCompletionAutomation(event, deps, context, rfCtx); triggerSelfWake(event, deps, context, rfCtx, hadDeferredWakes); + return rfCtx; } diff --git a/src/main/worker/handlers/run-finished-helpers.ts b/src/main/worker/handlers/run-finished-helpers.ts index 736f8cbc..4b1e9cda 100644 --- a/src/main/worker/handlers/run-finished-helpers.ts +++ b/src/main/worker/handlers/run-finished-helpers.ts @@ -71,9 +71,6 @@ export function finalizeRun( leaseRenewalMap.delete(event.runId); const runOutcome = db.core.runs.finishRun({ ...event, unattributedCost }); - if (run.workspaceId) { - db.core.workspaces.releaseWorkspaceLock(run.workspaceId, run.id); - } db.core.activity.addActivity({ companyId: run.companyId, actor: run.agentId, @@ -86,6 +83,15 @@ export function finalizeRun( return { run, task, runOutcome }; } +export function releaseCompletedRunWorkspaceLock( + rfCtx: RunFinishedContext | null, + deps: Pick, +): void { + const workspaceId = rfCtx?.run.workspaceId; + if (!workspaceId) return; + deps.db.core.workspaces.releaseWorkspaceLock(workspaceId, rfCtx.run.id); +} + export function accumulateAndRotateSession( event: RunFinishedEvent, deps: HandlerDeps, diff --git a/src/main/worker/runtime/helpers.ts b/src/main/worker/runtime/helpers.ts index f73f015f..cce4ca0b 100644 --- a/src/main/worker/runtime/helpers.ts +++ b/src/main/worker/runtime/helpers.ts @@ -47,6 +47,7 @@ export type RunAgentLifecycleState = "agent_active" | "closing"; export type ParentMessage = | { type: "enqueue-run"; payload: RunJob } | { type: "cancel-run"; runId: string } + | { type: "run-completion-ack"; runId: string } | { type: "shutdown" } | { type: "steer-run"; runId: string; correlationId: string; message: AgentMessage } | { type: "rpc-response"; rpcId: string; response: { ok: true; result: unknown } | { ok: false; error: string } }; diff --git a/tests/unit/worker-completion-boundary.test.ts b/tests/unit/worker-completion-boundary.test.ts new file mode 100644 index 00000000..043e8f38 --- /dev/null +++ b/tests/unit/worker-completion-boundary.test.ts @@ -0,0 +1,236 @@ +import { EventEmitter } from "node:events"; +import { describe, it, expect, vi, afterEach } from "vitest"; +import { handleWorkerMessage } from "@main/worker/entry/message-handler"; +import { workerMessageSchema } from "@main/worker/entry/types"; +import type { RunJob } from "@main/worker/runtime/helpers"; + +const RUN_ID = "00000000-0000-4000-8000-000000000001"; +const NEXT_RUN_ID = "00000000-0000-4000-8000-000000000002"; + +function runFinishedMessage(runId: string) { + return { + type: "run-finished" as const, + runId, + status: "succeeded" as const, + summary: null, + errorMessage: null, + exitCode: null, + signal: null, + model: "test-model", + usageStatus: "ok" as const, + inputTokens: 100, + outputTokens: 50, + }; +} + +function createHandlerDeps(calls: string[]) { + const releaseWorkspaceLock = vi.fn(() => calls.push("releaseWorkspaceLock")); + const finishRun = vi.fn(() => { + calls.push("finishRun"); + return { budgetHardStopped: false }; + }); + const rotateSession = vi.fn(() => calls.push("rotateSession")); + const processRunCompletion = vi.fn(() => calls.push("processRunCompletion")); + const releaseRunResources = vi.fn(() => calls.push("releaseRunResources")); + + const db = { + core: { + runs: { + getRun: vi.fn(() => ({ + id: RUN_ID, + companyId: "company-1", + agentId: "agent-1", + taskId: null, + workspaceId: "workspace-1", + })), + finishRun, + }, + workspaces: { releaseWorkspaceLock }, + activity: { addActivity: vi.fn(() => calls.push("addActivity")) }, + entity: { + getAgent: vi.fn(() => ({ + id: "agent-1", + companyId: "company-1", + reportsTo: "manager-1", + sessionId: "session-1", + })), + countActiveGoalsForCompany: vi.fn(() => 0), + }, + tasks: { + getSessionById: vi.fn(() => ({ id: "session-1" })), + listAgentTasks: vi.fn(() => []), + }, + agents: { + getDirectReports: vi.fn(() => []), + getChainOfCommand: vi.fn(() => []), + }, + approvals: { + listCompanyApprovalsByType: vi.fn(() => []), + }, + }, + automation: { + wakeup: { + promoteWakeupRequests: vi.fn(() => 0), + getPreferredWakeTrigger: vi.fn(() => null), + }, + }, + }; + + const sessionService = { + accumulateTokenUsage: vi.fn(() => calls.push("accumulateTokenUsage")), + evaluateRotation: vi.fn(() => ({ shouldRotate: true, reason: "unit-test" })), + rotateSession, + }; + + return { + deps: { + db, + logger: { + debug: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + }, + isQuitting: vi.fn(() => false), + emitEvent: vi.fn(() => calls.push("emitEvent")), + wakeAgentIfPossible: vi.fn(() => null), + processRunCompletion, + notify: vi.fn(), + sessionService, + getRuntimeBridge: vi.fn(() => ({ releaseRunResources })), + getRuntimeMessageInjector: vi.fn(() => null), + }, + releaseWorkspaceLock, + rotateSession, + processRunCompletion, + releaseRunResources, + }; +} + +function runJob(runId: string): RunJob { + return { + runId, + companyId: "company-1", + taskId: null, + agentId: "agent-1", + agentName: "Agent", + connectorId: "openai-codex", + workspaceId: "workspace-1", + workspacePath: "/tmp/agentcompany-worker-boundary", + workspaceMode: "exclusive", + model: null, + thinkingEffort: null, + env: {}, + systemPrompt: null, + trigger: "unit-test", + taskSessionId: `session-${runId}`, + executionMode: "ai_ceo", + initialUserTurn: "Run a unit test.", + internalTools: [], + }; +} + +describe("worker completion boundary", () => { + afterEach(() => { + vi.restoreAllMocks(); + vi.resetModules(); + Reflect.deleteProperty(process, "parentPort"); + }); + + it("releases the workspace lock after session, automation, and runtime resource release", async () => { + const calls: string[] = []; + const { + deps, + releaseWorkspaceLock, + rotateSession, + processRunCompletion, + releaseRunResources, + } = createHandlerDeps(calls); + const worker = { postMessage: vi.fn(() => calls.push("completionAck")) }; + + await handleWorkerMessage( + runFinishedMessage(RUN_ID), + worker as never, + deps as never, + workerMessageSchema, + false, + 2, + 30_000, + new Map(), + 3, + new Set(), + new Map(), + ); + + expect(releaseWorkspaceLock).toHaveBeenCalledWith("workspace-1", RUN_ID); + expect(worker.postMessage).toHaveBeenCalledWith({ type: "run-completion-ack", runId: RUN_ID }); + expect(releaseWorkspaceLock.mock.invocationCallOrder[0]).toBeGreaterThan( + rotateSession.mock.invocationCallOrder[0]!, + ); + expect(releaseWorkspaceLock.mock.invocationCallOrder[0]).toBeGreaterThan( + processRunCompletion.mock.invocationCallOrder[0]!, + ); + expect(releaseWorkspaceLock.mock.invocationCallOrder[0]).toBeGreaterThan( + releaseRunResources.mock.invocationCallOrder[0]!, + ); + }); + + it("keeps same-workspace runs queued until parent acknowledges completion", async () => { + const parentPort = new EventEmitter() as EventEmitter & { + postMessage: ReturnType; + }; + parentPort.postMessage = vi.fn(); + Object.defineProperty(process, "parentPort", { + configurable: true, + value: parentPort, + }); + + vi.doMock("@main/worker/runtime", () => ({ + WorkerRunner: class TestWorkerRunner { + constructor( + private readonly deps: { + post: (message: unknown) => void; + isCancelled: (runId: string) => boolean; + }, + ) {} + + async run(job: RunJob): Promise { + this.deps.post(runFinishedMessage(job.runId)); + } + + cancel(): void {} + + steer(): boolean { + return false; + } + + async cancelAll(): Promise {} + }, + })); + + await import("@main/worker/entry/index"); + + parentPort.emit("message", { data: { type: "enqueue-run", payload: runJob(RUN_ID) } }); + await Promise.resolve(); + await Promise.resolve(); + + parentPort.emit("message", { data: { type: "enqueue-run", payload: runJob(NEXT_RUN_ID) } }); + await Promise.resolve(); + await Promise.resolve(); + + expect(parentPort.postMessage).toHaveBeenCalledWith( + expect.objectContaining({ type: "run-finished", runId: RUN_ID }), + ); + expect(parentPort.postMessage).not.toHaveBeenCalledWith( + expect.objectContaining({ type: "run-finished", runId: NEXT_RUN_ID }), + ); + + parentPort.emit("message", { data: { type: "run-completion-ack", runId: RUN_ID } }); + await Promise.resolve(); + await Promise.resolve(); + + expect(parentPort.postMessage).toHaveBeenCalledWith( + expect.objectContaining({ type: "run-finished", runId: NEXT_RUN_ID }), + ); + }); +});