diff --git a/server/src/index.ts b/server/src/index.ts index 6645c7ef067..7bfaf346c27 100644 --- a/server/src/index.ts +++ b/server/src/index.ts @@ -1,4 +1,5 @@ import type { IdleDeadlineOwners } from "./services/idle-deadline-types.js"; +import { closeProcessRunAdmission } from "./services/heartbeat/run-admission.js"; import { registerOwnershipReconciliationWork } from "./services/native-runtime/native-finalization-reconciler.js"; import { bindIdleDeadlineOwner } from "./services/idle-deadline-owner.js"; import { readIdleRetrySchedule } from "./services/heartbeat/idle-retry-schedule.js"; @@ -2044,7 +2045,7 @@ async function startServerWithDatabaseTeardown( liveEventsWebSocketServer.close(); // Finish admitted HTTP requests while their services, workers and database // are still available. Railway must switch routing before signaling A. - await closeHttpListenerForShutdown({ server, signal, timeoutMs: 30_000, log: logger, + await closeHttpListenerForShutdown({ server, signal, timeoutMs: 30_000, log: logger, closeRunAdmission: closeProcessRunAdmission, drainRequests: abort => drainHttpRequestWork(app, abort) }); const preDrain = createPreDrainShutdownBudget({ signal, log: logger }); await preDrain.run("systemd_notify", () => diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index 2997874085a..6fbd9ad75eb 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -1,5 +1,6 @@ import { registerHeldPreparationResumer } from "./interrupted-preparation-hold-settlement.js"; import { resolveProviderApiPlacement } from "./native-runtime/provider-api-placement.js"; +import { isProcessRunAdmissionClosed } from "./heartbeat/run-admission.js"; import type { StopInvocationOptions } from "./heartbeat/run-control.js"; import { isContextCompactionRun } from "./context-compaction.js"; import { isUnstartedTaskQuestionAnswer } from "./unstarted-task-question-answer.js"; @@ -1321,6 +1322,7 @@ export function heartbeatService( isHeartbeatRunTerminalStatus, liveRunExecutions, getSchedulingSuppression, + isRunAdmissionClosed: () => shutdownInProgress || isProcessRunAdmissionClosed(), resumeExecutionWaitComments, getWorktreeExecutionCutoff, resumeQueuedCommentInterrupt, diff --git a/server/src/services/heartbeat/queue.test.ts b/server/src/services/heartbeat/queue.test.ts index 9d504c65874..3f288527682 100644 --- a/server/src/services/heartbeat/queue.test.ts +++ b/server/src/services/heartbeat/queue.test.ts @@ -9,6 +9,7 @@ import { resolveExternalTestDatabaseUrl } from "../../__tests__/helpers/external import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "../../__tests__/helpers/embedded-postgres.js"; import { createPostgresWakeQueueAdapter, createWakeAdmissionWriter } from "../../modules/wake-queue/adapters/postgres.js"; import { createReleaseIssueExecution } from "../../modules/wake-queue/application/use-cases.js"; +import { closeProcessRunAdmission, resetProcessRunAdmissionForTests } from "./run-admission.js"; import { createHeartbeatQueue, DEFERRED_WAKE_CONTEXT_KEY, shouldDeferFollowupWakeForSameIssue, shouldQueueFollowupForRunningIssueWake, type HeartbeatQueueDependencies } from "./queue.js"; import { createWakeQueue } from "../../modules/wake-queue/index.js"; import { TransactionScope } from "../../modules/wake-queue/application/ports.js"; @@ -50,6 +51,30 @@ describe("heartbeat queue module boundary", () => { expect(deps.executeRun).not.toHaveBeenCalled(); }); + it.each(["startNextQueuedRunForAgent", "resumeQueuedRuns"] as const)("starts no run once shutdown closed run admission (%s)", async (operation) => { + const database = guardedDatabase(); + const deps = callbacks(database.db); + deps.isRunAdmissionClosed = () => true; + const queue = createHeartbeatQueue(database.db, deps); + await queue[operation]("agent"); + expect(database.access).not.toHaveBeenCalled(); + expect(deps.getAgent).not.toHaveBeenCalled(); + expect(deps.executeRun).not.toHaveBeenCalled(); + }); + + it("honors the process-wide run admission gate in every heartbeat instance", async () => { + const database = guardedDatabase(); + const deps = callbacks(database.db); + closeProcessRunAdmission(); + try { + await createHeartbeatQueue(database.db, deps).startNextQueuedRunForAgent("agent"); + } finally { + resetProcessRunAdmissionForTests(); + } + expect(database.access).not.toHaveBeenCalled(); + expect(deps.getAgent).not.toHaveBeenCalled(); + }); + it("tracks pending wakes across instances and clears each rejected wake", async () => { const database = guardedDatabase(); const deps = callbacks(database.db); diff --git a/server/src/services/heartbeat/queue.ts b/server/src/services/heartbeat/queue.ts index 0f0ebdf8f7d..5f3c2fe48d2 100644 --- a/server/src/services/heartbeat/queue.ts +++ b/server/src/services/heartbeat/queue.ts @@ -1,4 +1,5 @@ import { CONTEXT_COMPACTION_WAKE_REASON, ContextCompactionError, isContextCompactionRun, readContextCompactionRequest } from "../context-compaction.js"; +import { isProcessRunAdmissionClosed } from "./run-admission.js"; import { deliveredQueuedCommentIds } from "../queued-comment-delivery.js"; import { deferredWorkIntentWait } from "../work-intent-outbox.js"; import { createHeartbeatTimerPolicy } from "./timer-policy.js"; @@ -227,6 +228,7 @@ export interface HeartbeatQueueDependencies extends Pick string | null; filterZombieCoalesceTarget: (run: T | null, live: { has(id: string): boolean }) => T | null; getSchedulingSuppression: () => Promise<{ suppressed: boolean; reason: "worktree_instance" | "database_restore_in_progress" | "task_drain" | null }>; + isRunAdmissionClosed?: () => boolean; canStartNewExternalBotWork: (agent: typeof agents.$inferSelect) => Promise; executeRun: (runId: string) => Promise; releaseIssueExecutionAndPromote: (run: Pick, options?: { suppressImmediateRecovery?: boolean; resumeStrandedQueue?: boolean; deferredPostCommitEffects?: WakeQueuePostCommitEffect[] }) => Promise; @@ -411,6 +413,7 @@ export function createHeartbeatQueue(db: Db, dependencies: HeartbeatQueueDepende isHeartbeatRunTerminalStatus, liveRunExecutions, getSchedulingSuppression, + isRunAdmissionClosed = isProcessRunAdmissionClosed, resumeExecutionWaitComments, getWorktreeExecutionCutoff, resumeQueuedCommentInterrupt, @@ -1947,6 +1950,7 @@ export function createHeartbeatQueue(db: Db, dependencies: HeartbeatQueueDepende } async function resumeQueuedRuns() { + if (isRunAdmissionClosed()) return; await resumeDeferredExecutionReceipts(); await resumeQueuedRunsAfterDrain(); } @@ -1954,7 +1958,7 @@ export function createHeartbeatQueue(db: Db, dependencies: HeartbeatQueueDepende // Release admission promptly without first scanning unrelated recovery work. // Every run still passes the ordinary per-agent claim and dispatch gates. async function resumeQueuedRunsAfterDrain(scope?: { companyId: string }) { - if ((await getSchedulingSuppression()).suppressed) return; + if (isRunAdmissionClosed() || (await getSchedulingSuppression()).suppressed) return; const cutoff = await getWorktreeExecutionCutoff(); const queuedRuns = await db .select({ agentId: heartbeatRuns.agentId }) @@ -2048,7 +2052,7 @@ export function createHeartbeatQueue(db: Db, dependencies: HeartbeatQueueDepende } async function startNextQueuedRunForAgent(agentId: string) { - if ((await getSchedulingSuppression()).suppressed) return []; + if (isRunAdmissionClosed() || (await getSchedulingSuppression()).suppressed) return []; const cutoff = await getWorktreeExecutionCutoff(); // Cancelled after the start lock is released: cancelRunInternal promotes the // agent's next queued run, which takes this same lock. diff --git a/server/src/services/heartbeat/run-admission.ts b/server/src/services/heartbeat/run-admission.ts new file mode 100644 index 00000000000..d61a0b48bd0 --- /dev/null +++ b/server/src/services/heartbeat/run-admission.ts @@ -0,0 +1,13 @@ +let runAdmissionClosed = false; + +export function closeProcessRunAdmission(): void { + runAdmissionClosed = true; +} + +export function isProcessRunAdmissionClosed(): boolean { + return runAdmissionClosed; +} + +export function resetProcessRunAdmissionForTests(): void { + runAdmissionClosed = false; +} diff --git a/server/src/shutdown.test.ts b/server/src/shutdown.test.ts index 83176e35210..728896f0535 100644 --- a/server/src/shutdown.test.ts +++ b/server/src/shutdown.test.ts @@ -281,6 +281,19 @@ describe("closeHttpListenerForShutdown", () => { expect(server.closeAllConnections).not.toHaveBeenCalled(); }); + it("closes run admission before it stops accepting requests", async () => { + const order: string[] = []; + const server = fakeServer({ listening: true, closeDelayMs: 0 }); + server.close.mockImplementation((callback?: (err?: Error) => void) => { order.push("listener"); setTimeout(() => callback?.(), 0); }); + await closeHttpListenerForShutdown({ server, signal: "SIGTERM", timeoutMs: 1_000, log: stubLogger(), + closeRunAdmission: () => order.push("run_admission") }); + expect(order).toEqual(["run_admission", "listener"]); + const stopped = vi.fn(); + await expect(closeHttpListenerForShutdown({ server: fakeServer({ listening: false }), signal: "SIGTERM", log: stubLogger(), + closeRunAdmission: stopped })).resolves.toBe("not_listening"); + expect(stopped).toHaveBeenCalledOnce(); + }); + it("closes the remaining connections when the drain outlives the grace period", async () => { vi.useFakeTimers(); try { diff --git a/server/src/shutdown.ts b/server/src/shutdown.ts index ebc9e590b8f..475df7c095b 100644 --- a/server/src/shutdown.ts +++ b/server/src/shutdown.ts @@ -127,8 +127,10 @@ export async function closeHttpListenerForShutdown(input: { signal: "SIGINT" | "SIGTERM"; timeoutMs?: number; drainRequests?: (signal: AbortSignal) => Promise; + closeRunAdmission?: () => void; log: ShutdownLogger; }): Promise<"closed" | "timed_out" | "not_listening"> { + input.closeRunAdmission?.(); if (!input.server.listening) return "not_listening"; const timeoutMs = input.timeoutMs ?? 5_000; let timer: NodeJS.Timeout | null = null;