diff --git a/server/src/__tests__/heartbeat-process-recovery.test.ts b/server/src/__tests__/heartbeat-process-recovery.test.ts index e1065ae4f18..6e47ecc3c68 100644 --- a/server/src/__tests__/heartbeat-process-recovery.test.ts +++ b/server/src/__tests__/heartbeat-process-recovery.test.ts @@ -43,6 +43,7 @@ import { import { runningProcesses } from "../adapters/index.ts"; const mockTelemetryClient = vi.hoisted(() => ({ track: vi.fn() })); const mockTrackAgentFirstHeartbeat = vi.hoisted(() => vi.fn()); +const mockTerminateLocalService = vi.hoisted(() => vi.fn()); const mockAdapterExecute = vi.hoisted(() => vi.fn(async () => ({ exitCode: 0, @@ -59,6 +60,17 @@ vi.mock("../telemetry.ts", () => ({ getTelemetryClient: () => mockTelemetryClient, })); +vi.mock("../services/local-service-supervisor.js", async () => { + const actual = await vi.importActual( + "../services/local-service-supervisor.js", + ); + mockTerminateLocalService.mockImplementation(actual.terminateLocalService); + return { + ...actual, + terminateLocalService: mockTerminateLocalService, + }; +}); + vi.mock("@paperclipai/shared/telemetry", async () => { const actual = await vi.importActual( "@paperclipai/shared/telemetry", @@ -277,6 +289,10 @@ describeEmbeddedPostgres("heartbeat orphaned process recovery", () => { afterEach(async () => { vi.clearAllMocks(); + const localServiceSupervisor = await vi.importActual( + "../services/local-service-supervisor.js", + ); + mockTerminateLocalService.mockImplementation(localServiceSupervisor.terminateLocalService); mockAdapterExecute.mockImplementation(async () => ({ exitCode: 0, signal: null, @@ -1893,6 +1909,35 @@ describeEmbeddedPostgres("heartbeat orphaned process recovery", () => { ); }); + it("terminates the in-memory process before persisting cancellation status", async () => { + const { runId } = await seedRunFixture({ + agentStatus: "running", + includeIssue: false, + }); + const heartbeat = heartbeatService(db); + runningProcesses.set(runId, { + child: { pid: 12345 } as ChildProcess, + graceSec: 1, + processGroupId: null, + }); + mockTerminateLocalService.mockResolvedValueOnce(undefined); + const updateSpy = vi.spyOn(db, "update"); + updateSpy.mockImplementationOnce((() => { + throw new Error("db update unavailable"); + }) as typeof db.update); + + try { + await expect(heartbeat.cancelRun(runId)).rejects.toThrow("db update unavailable"); + expect(mockTerminateLocalService).toHaveBeenCalledWith( + expect.objectContaining({ pid: 12345, processGroupId: null }), + { forceAfterMs: 1000 }, + ); + expect(runningProcesses.has(runId)).toBe(false); + } finally { + updateSpy.mockRestore(); + } + }); + it("records manual cancellation stop metadata", async () => { const { runId } = await seedRunFixture({ agentStatus: "running", diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index 79a321e422d..910d395a4cd 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -10636,32 +10636,37 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) } : options.resultJson; + const running = runningProcesses.get(run.id); + try { + if (running) { + await terminateHeartbeatRunProcess({ + pid: running.child.pid ?? run.processPid, + processGroupId: running.processGroupId ?? run.processGroupId, + graceMs: Math.max(1, running.graceSec) * 1000, + }); + } else if (run.processPid || run.processGroupId) { + await terminateHeartbeatRunProcess({ + pid: run.processPid, + processGroupId: run.processGroupId, + }); + } + } finally { + runningProcesses.delete(run.id); + } + + const finishedAt = new Date(); const cancelled = await setRunStatus(run.id, "cancelled", { - finishedAt: new Date(), + finishedAt, error: reason, errorCode, ...(resultJson ? { resultJson } : {}), }); await setWakeupStatus(run.wakeupRequestId, "cancelled", { - finishedAt: new Date(), + finishedAt, error: reason, }); - const running = runningProcesses.get(run.id); - if (running) { - await terminateHeartbeatRunProcess({ - pid: running.child.pid ?? run.processPid, - processGroupId: running.processGroupId ?? run.processGroupId, - graceMs: Math.max(1, running.graceSec) * 1000, - }); - } else if (run.processPid || run.processGroupId) { - await terminateHeartbeatRunProcess({ - pid: run.processPid, - processGroupId: run.processGroupId, - }); - } - if (cancelled) { await appendRunEvent(cancelled, 1, { eventType: "lifecycle", @@ -10673,7 +10678,6 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) await releaseIssueExecutionAndPromote(cancelled); } - runningProcesses.delete(run.id); await finalizeAgentStatus(run.agentId, "cancelled"); await startNextQueuedRunForAgent(run.agentId); return cancelled;