diff --git a/.changeset/bullmq-fast-turn-drain.md b/.changeset/bullmq-fast-turn-drain.md new file mode 100644 index 0000000000..306684a7ee --- /dev/null +++ b/.changeset/bullmq-fast-turn-drain.md @@ -0,0 +1,5 @@ +--- +'@roomote/web': patch +--- + +The bullmq service now drains and aborts the Fast turns it is executing before it shuts down, the same way the API does for the turns it admits. Interrupted turns are resumed by the parent-event queue inside the bullmq process, so a deploy that restarted both services could kill a resumed turn a second time with its durable claim and conversation lock still held, leaving the row to wait out both leases (up to 15 minutes) before it ran again. Shutdown now closes admissions, stops fetching queue wakeups, gives in-flight turns the drain window to finish, and aborts the stragglers so each hands its row back to the queue immediately. The window is `R_BULLMQ_SHUTDOWN_DRAIN_MS`, falling back to `R_API_SHUTDOWN_DRAIN_MS` and then 20 seconds; the shutdown sequence itself is shared between the two services. diff --git a/.env.production.example b/.env.production.example index 223f141667..4c6f59ec3c 100644 --- a/.env.production.example +++ b/.env.production.example @@ -15,9 +15,12 @@ PREVIEW_PROXY_SUBDOMAIN_SUFFIX=preview # R_APP_ENV=production # How long an API shutdown lets in-flight Fast turns finish before aborting -# the remainder. Keep it under the platform's SIGTERM-to-SIGKILL grace -# window; 0 aborts active turns immediately. +# the remainder so they resume on the next process. Keep it under the +# platform's SIGTERM-to-SIGKILL grace window; 0 aborts active turns +# immediately. The bullmq service, which runs the turns the queue resumes, +# uses the same window unless R_BULLMQ_SHUTDOWN_DRAIN_MS overrides it. # R_API_SHUTDOWN_DRAIN_MS=20000 +# R_BULLMQ_SHUTDOWN_DRAIN_MS=20000 # Set to true to keep inference retry backoff inside the process that owns # the Fast turn instead of scheduling it durably (kill switch for durable # retry scheduling). diff --git a/apps/api/src/__tests__/graceful-shutdown.test.ts b/apps/api/src/__tests__/graceful-shutdown.test.ts index 37f0a7e78b..fac763d74c 100644 --- a/apps/api/src/__tests__/graceful-shutdown.test.ts +++ b/apps/api/src/__tests__/graceful-shutdown.test.ts @@ -1,22 +1,3 @@ -const mocks = vi.hoisted(() => ({ - abortActiveFastAgentTurns: vi.fn(), - beginFastAgentTurnDrain: vi.fn(), - waitForActiveFastAgentTurnsToSettle: vi.fn(), -})); - -vi.mock('@roomote/cloud-agents/server', () => ({ - abortActiveFastAgentTurns: mocks.abortActiveFastAgentTurns, - beginFastAgentTurnDrain: mocks.beginFastAgentTurnDrain, - waitForActiveFastAgentTurnsToSettle: - mocks.waitForActiveFastAgentTurnsToSettle, - FastAgentProcessShutdownError: class extends Error { - constructor(public readonly signal: NodeJS.Signals) { - super(`Fast turn interrupted by API shutdown (${signal}).`); - this.name = 'FastAgentProcessShutdownError'; - } - }, -})); - import type { ServerType } from '@hono/node-server'; import { diff --git a/apps/api/src/graceful-shutdown.ts b/apps/api/src/graceful-shutdown.ts index 59e9b0e504..c2d81954dc 100644 --- a/apps/api/src/graceful-shutdown.ts +++ b/apps/api/src/graceful-shutdown.ts @@ -1,33 +1,19 @@ import type { ServerType } from '@hono/node-server'; import { - abortActiveFastAgentTurns, - beginFastAgentTurnDrain, + drainAndAbortFastAgentTurns, FastAgentProcessShutdownError, - waitForActiveFastAgentTurnsToSettle, + resolveFastAgentShutdownDrainMs, + type FastAgentShutdownDrainDeps, } from '@roomote/cloud-agents/server'; -// Most Fast turns finish within seconds, so letting them settle turns a -// deploy-time interruption into a completed answer. The default leaves room -// for the straggler abort, closeout delivery, and Sentry flush inside a -// typical 30s SIGTERM-to-SIGKILL grace window. R_API_SHUTDOWN_DRAIN_MS -// overrides it; 0 restores the previous abort-immediately behavior. -const DEFAULT_API_SHUTDOWN_DRAIN_MS = 20_000; - +/** `R_API_SHUTDOWN_DRAIN_MS` overrides the shared default; 0 aborts at once. */ export function resolveApiShutdownDrainMs( env: NodeJS.ProcessEnv = process.env, ): number { - const raw = env.R_API_SHUTDOWN_DRAIN_MS?.trim(); - if (!raw) return DEFAULT_API_SHUTDOWN_DRAIN_MS; - const parsed = Number(raw); - return Number.isFinite(parsed) && parsed >= 0 - ? parsed - : DEFAULT_API_SHUTDOWN_DRAIN_MS; + return resolveFastAgentShutdownDrainMs(['R_API_SHUTDOWN_DRAIN_MS'], env); } -type ApiShutdownOptions = { - abortTurns?: typeof abortActiveFastAgentTurns; - beginDrain?: typeof beginFastAgentTurnDrain; - waitForTurns?: typeof waitForActiveFastAgentTurnsToSettle; +type ApiShutdownOptions = FastAgentShutdownDrainDeps & { drainMs?: number; exitProcess?: (code?: number) => never; flushSentry?: () => Promise; @@ -39,9 +25,9 @@ export async function gracefullyShutdownApi( server: ServerType, signal: NodeJS.Signals, { - abortTurns = abortActiveFastAgentTurns, - beginDrain = beginFastAgentTurnDrain, - waitForTurns = waitForActiveFastAgentTurnsToSettle, + abortTurns, + beginDrain, + waitForTurns, drainMs = resolveApiShutdownDrainMs(), exitProcess = process.exit, flushSentry = async () => undefined, @@ -53,19 +39,22 @@ export async function gracefullyShutdownApi( // Refuse new turn admissions and stop accepting connections first, then // give in-flight turns a bounded window to finish on their own. Only the // stragglers still active at the deadline are aborted. - beginDrain(reason); - const closePromise = new Promise((resolve) => { - server.close((error) => resolve(error ?? null)); - }); - const remaining = await waitForTurns(drainMs); - if (remaining > 0) { - logWarn( - `[api] Aborting ${remaining} Fast turn(s) still active after the ${drainMs}ms shutdown drain.`, - ); - } - const abortPromise = abortTurns(reason); + let closePromise: Promise = Promise.resolve(null); + await drainAndAbortFastAgentTurns( + { + reason, + drainMs, + service: 'api', + logWarn, + onDrainStarted: () => { + closePromise = new Promise((resolve) => { + server.close((error) => resolve(error ?? null)); + }); + }, + }, + { abortTurns, beginDrain, waitForTurns }, + ); const closeError = await closePromise; - await abortPromise; if (closeError) { logError('[api] Graceful shutdown failed', closeError); } diff --git a/apps/bullmq/src/graceful-shutdown.test.ts b/apps/bullmq/src/graceful-shutdown.test.ts new file mode 100644 index 0000000000..dbbf956cc3 --- /dev/null +++ b/apps/bullmq/src/graceful-shutdown.test.ts @@ -0,0 +1,187 @@ +import { + gracefullyShutdownBullMq, + installBullMqGracefulShutdown, + resolveBullMqShutdownDrainMs, +} from './graceful-shutdown'; + +describe('resolveBullMqShutdownDrainMs', () => { + it('prefers its own window and falls back to the API window, then the default', () => { + expect(resolveBullMqShutdownDrainMs({})).toBe(20_000); + expect( + resolveBullMqShutdownDrainMs({ R_API_SHUTDOWN_DRAIN_MS: '15000' }), + ).toBe(15_000); + expect( + resolveBullMqShutdownDrainMs({ + R_API_SHUTDOWN_DRAIN_MS: '15000', + R_BULLMQ_SHUTDOWN_DRAIN_MS: '45000', + }), + ).toBe(45_000); + expect( + resolveBullMqShutdownDrainMs({ R_BULLMQ_SHUTDOWN_DRAIN_MS: '0' }), + ).toBe(0); + }); +}); + +describe('gracefullyShutdownBullMq', () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it('stops the Fast worker as the drain starts and closes the rest only after the abort', async () => { + const order: string[] = []; + let finishWorkerClose: (() => void) | undefined; + const fastAgentWorker = { + close: vi.fn( + () => + new Promise((resolve) => { + order.push('worker.close'); + finishWorkerClose = resolve; + }), + ), + }; + const closeRemaining = vi.fn(async () => { + order.push('closeRemaining'); + }); + let finishDrain: ((remaining: number) => void) | undefined; + const waitForTurns = vi.fn( + () => + new Promise((resolve) => { + order.push('wait'); + finishDrain = resolve; + }), + ); + const abortTurns = vi.fn(async () => { + order.push('abort'); + // The aborted turns reject their jobs, which lets the worker close. + finishWorkerClose?.(); + return 1; + }); + const exitProcess = vi.fn() as unknown as (code?: number) => never; + const logWarn = vi.fn(); + + const shutdown = gracefullyShutdownBullMq('SIGTERM', { + fastAgentWorker, + closeRemaining, + beginDrain: vi.fn(() => order.push('begin')), + waitForTurns, + abortTurns, + drainMs: 4_321, + exitProcess, + logWarn, + logInfo: vi.fn(), + }); + + await vi.waitFor(() => expect(waitForTurns).toHaveBeenCalledWith(4_321)); + expect(fastAgentWorker.close).toHaveBeenCalledOnce(); + expect(abortTurns).not.toHaveBeenCalled(); + expect(closeRemaining).not.toHaveBeenCalled(); + + finishDrain?.(1); + await shutdown; + + expect(order).toEqual([ + 'begin', + 'worker.close', + 'wait', + 'abort', + 'closeRemaining', + ]); + expect(abortTurns).toHaveBeenCalledWith( + expect.objectContaining({ + name: 'FastAgentProcessShutdownError', + signal: 'SIGTERM', + }), + ); + expect(logWarn).toHaveBeenCalledWith( + '[bullmq] Aborting 1 Fast turn(s) still active after the 4321ms shutdown drain.', + ); + expect(exitProcess).toHaveBeenCalledWith(0); + }); + + it('does not let a worker that ignores its abort hold shutdown open', async () => { + const fastAgentWorker = { + close: vi.fn(() => new Promise(() => undefined)), + }; + const closeRemaining = vi.fn().mockResolvedValue(undefined); + const exitProcess = vi.fn() as unknown as (code?: number) => never; + const logWarn = vi.fn(); + + await gracefullyShutdownBullMq('SIGTERM', { + fastAgentWorker, + closeRemaining, + beginDrain: vi.fn(), + waitForTurns: vi.fn().mockResolvedValue(0), + abortTurns: vi.fn().mockResolvedValue(0), + drainMs: 0, + workerCloseTimeoutMs: 10, + exitProcess, + logWarn, + logInfo: vi.fn(), + }); + + expect(logWarn).toHaveBeenCalledWith( + '[bullmq] Fast parent event worker did not close within 10ms of the abort; continuing shutdown.', + ); + expect(closeRemaining).toHaveBeenCalledOnce(); + expect(exitProcess).toHaveBeenCalledWith(0); + }); + + it('still exits when closing the remaining queues fails', async () => { + const error = new Error('redis gone'); + const exitProcess = vi.fn() as unknown as (code?: number) => never; + const logError = vi.fn(); + + await gracefullyShutdownBullMq('SIGINT', { + fastAgentWorker: { close: vi.fn().mockResolvedValue(undefined) }, + closeRemaining: vi.fn().mockRejectedValue(error), + beginDrain: vi.fn(), + waitForTurns: vi.fn().mockResolvedValue(0), + abortTurns: vi.fn().mockResolvedValue(0), + drainMs: 0, + exitProcess, + logError, + logInfo: vi.fn(), + }); + + expect(logError).toHaveBeenCalledWith( + '[Shutdown] Error during shutdown:', + error, + ); + expect(exitProcess).toHaveBeenCalledWith(0); + }); + + it.each(['SIGTERM', 'SIGINT'] as const)( + 'forces exit when %s arrives again during shutdown', + (signal) => { + const exitProcess = vi.fn() as unknown as (code?: number) => never; + const signalHandlers = new Map void>(); + const on = vi.spyOn(process, 'on').mockImplementation((( + registeredSignal: NodeJS.Signals, + handler: () => void, + ) => { + signalHandlers.set(registeredSignal, handler); + return process; + }) as typeof process.on); + + try { + const cleanup = installBullMqGracefulShutdown({ + fastAgentWorker: { close: vi.fn().mockResolvedValue(undefined) }, + closeRemaining: vi.fn().mockResolvedValue(undefined), + beginDrain: vi.fn(), + waitForTurns: vi.fn(() => new Promise(() => undefined)), + abortTurns: vi.fn(() => new Promise(() => undefined)), + exitProcess, + logInfo: vi.fn(), + }); + + signalHandlers.get(signal)?.(); + expect(exitProcess).not.toHaveBeenCalled(); + signalHandlers.get(signal)?.(); + expect(exitProcess).toHaveBeenCalledWith(1); + cleanup(); + } finally { + on.mockRestore(); + } + }, + ); +}); diff --git a/apps/bullmq/src/graceful-shutdown.ts b/apps/bullmq/src/graceful-shutdown.ts new file mode 100644 index 0000000000..18e4fc32a6 --- /dev/null +++ b/apps/bullmq/src/graceful-shutdown.ts @@ -0,0 +1,142 @@ +import { + drainAndAbortFastAgentTurns, + FastAgentProcessShutdownError, + resolveFastAgentShutdownDrainMs, + type FastAgentShutdownDrainDeps, +} from '@roomote/cloud-agents/server'; + +/** + * Once the stragglers are aborted, each resumed turn's job rejects and the + * Fast worker's close resolves. A turn that ignores its abort would hold the + * close open; SIGKILL is coming anyway, so the wait is bounded. + */ +const DEFAULT_FAST_AGENT_WORKER_CLOSE_TIMEOUT_MS = 5_000; + +/** + * `R_BULLMQ_SHUTDOWN_DRAIN_MS` sizes the window for the turns this process + * resumes; it falls back to the API's window so one setting covers both. + */ +export function resolveBullMqShutdownDrainMs( + env: NodeJS.ProcessEnv = process.env, +): number { + return resolveFastAgentShutdownDrainMs( + ['R_BULLMQ_SHUTDOWN_DRAIN_MS', 'R_API_SHUTDOWN_DRAIN_MS'], + env, + ); +} + +type Closeable = { close: () => Promise }; + +type BullMqShutdownOptions = FastAgentShutdownDrainDeps & { + /** + * The worker that executes resumed Fast turns. Closed as the drain starts + * so it stops fetching new wakeups; its active jobs are the turns the + * drain waits for. + */ + fastAgentWorker: Closeable; + /** Every other queue, worker, and connection; closed once the Fast turns + * have finished or been handed back. */ + closeRemaining: () => Promise; + drainMs?: number; + workerCloseTimeoutMs?: number; + exitProcess?: (code?: number) => never; + logError?: (...args: Parameters) => void; + logWarn?: (...args: Parameters) => void; + logInfo?: (...args: Parameters) => void; +}; + +function waitWithTimeout(promise: Promise, timeoutMs: number) { + return new Promise<'closed' | 'timeout'>((resolve) => { + const timer = setTimeout(() => resolve('timeout'), timeoutMs); + timer.unref(); + void promise.then(() => { + clearTimeout(timer); + resolve('closed'); + }); + }); +} + +/** + * Resumed Fast turns run inside this process, so it owes them the same + * hand-off the API gives the turns it admits: close admissions, let in-flight + * turns finish inside the drain window, abort the stragglers so each releases + * its durable claim and wakes the queue, and only then close everything else. + * Without this the turn dies with SIGKILL holding its claim and the + * conversation lock, and the row waits out both leases before it resumes. + */ +export async function gracefullyShutdownBullMq( + signal: NodeJS.Signals, + { + abortTurns, + beginDrain, + waitForTurns, + fastAgentWorker, + closeRemaining, + drainMs = resolveBullMqShutdownDrainMs(), + workerCloseTimeoutMs = DEFAULT_FAST_AGENT_WORKER_CLOSE_TIMEOUT_MS, + exitProcess = process.exit, + logError = (...args) => console.error(...args), + logWarn = (...args) => console.warn(...args), + logInfo = (...args) => console.log(...args), + }: BullMqShutdownOptions, +): Promise { + logInfo('[Shutdown] Starting graceful shutdown...'); + const reason = new FastAgentProcessShutdownError(signal); + let workerClosed: Promise = Promise.resolve(); + await drainAndAbortFastAgentTurns( + { + reason, + drainMs, + service: 'bullmq', + logWarn, + onDrainStarted: () => { + workerClosed = fastAgentWorker.close().catch((error) => { + logError( + '[bullmq] Failed to close the Fast parent event worker', + error, + ); + }); + }, + }, + { abortTurns, beginDrain, waitForTurns }, + ); + if ( + (await waitWithTimeout(workerClosed, workerCloseTimeoutMs)) === 'timeout' + ) { + logWarn( + `[bullmq] Fast parent event worker did not close within ${workerCloseTimeoutMs}ms of the abort; continuing shutdown.`, + ); + } + try { + await closeRemaining(); + } catch (error) { + logError('[Shutdown] Error during shutdown:', error); + } + exitProcess(0); +} + +export function installBullMqGracefulShutdown( + options: BullMqShutdownOptions, +): () => void { + let shuttingDown = false; + const handlers = new Map void>(); + + for (const signal of ['SIGTERM', 'SIGINT'] as const) { + const handler = () => { + if (shuttingDown) { + (options.exitProcess ?? process.exit)(1); + return; + } + shuttingDown = true; + void gracefullyShutdownBullMq(signal, options); + }; + handlers.set(signal, handler); + process.on(signal, handler); + } + + return () => { + for (const [signal, handler] of handlers) { + process.off(signal, handler); + } + }; +} diff --git a/apps/bullmq/src/index.ts b/apps/bullmq/src/index.ts index 2a3a83e4ac..c53d8eac64 100644 --- a/apps/bullmq/src/index.ts +++ b/apps/bullmq/src/index.ts @@ -53,6 +53,7 @@ import { startPullRequestMergeabilityCheckQueue } from './pull-request-mergeabil import { startTaskSleepQueue } from './task-sleep-queue'; import { startAutomationRecommendationsQueue } from './automation-recommendations-queue'; import { startFastAgentParentEventQueue } from './fast-agent-parent-event-queue'; +import { installBullMqGracefulShutdown } from './graceful-shutdown'; // Deployments roll every service at once while migrations run only ahead // of the api service. A boot that reads a column the pending migration adds @@ -373,10 +374,11 @@ app.route('/admin/queues', serverAdapter.registerPlugin()); app.get('/', (c) => c.redirect('/admin/queues')); -async function gracefulShutdown() { - console.log('[Shutdown] Starting graceful shutdown...'); - - try { +// Resumed Fast turns execute inside this process, so shutdown drains and +// aborts them before anything else closes; see graceful-shutdown.ts. +installBullMqGracefulShutdown({ + fastAgentWorker: fastAgentParentEventWorker, + closeRemaining: async () => { await schedulerWorker.close(); await schedulerQueueEvents.close(); await schedulerQueue.close(); @@ -427,20 +429,12 @@ async function gracefulShutdown() { await pullRequestMergeabilityCheckWorker.close(); await pullRequestMergeabilityCheckQueueEvents.close(); await pullRequestMergeabilityCheckQueue.close(); - await fastAgentParentEventWorker.close(); await fastAgentParentEventQueueEvents.close(); await fastAgentParentEventQueue.close(); await discordGatewaySupervisor.stop(); await closeRedis(); - } catch (error) { - console.error('[Shutdown] Error during shutdown:', error); - } - - process.exit(0); -} - -process.on('SIGTERM', gracefulShutdown); -process.on('SIGINT', gracefulShutdown); + }, +}); const port = Number(process.env.PORT || 13002); diff --git a/apps/docs/environment-variables.mdx b/apps/docs/environment-variables.mdx index f34453241d..f28fcf71bb 100644 --- a/apps/docs/environment-variables.mdx +++ b/apps/docs/environment-variables.mdx @@ -128,7 +128,8 @@ as per-task auth tokens or workspace paths. | `R_INSTANCE_ID` | Optional | Stable anonymous deployment identifier sent with telemetry and version checks. Use a random, non-identifying value when overriding it. | | `R_STATUSPAGE_INCIDENTS_URL` | Optional | URL of a Statuspage-compatible unresolved-incidents JSON feed. Setting it enables incident banners and Slack warnings; leaving it unset disables Statuspage checks. | | `ROOMOTE_FORCE_TELEMETRY` | Development only | Force-enables telemetry in development or preview environments when a Ping endpoint is explicitly configured. | -| `R_API_SHUTDOWN_DRAIN_MS` | Optional | Milliseconds the API allows in-flight Fast turns to finish during shutdown before aborting the remainder. Defaults to `20000`; set to `0` to restore immediate aborts. Keep it below the hosting platform's SIGTERM-to-SIGKILL grace period. | +| `R_API_SHUTDOWN_DRAIN_MS` | Optional | Milliseconds the API allows in-flight Fast turns to finish during shutdown before aborting the remainder so they resume on the next process. Defaults to `20000`; set to `0` to abort immediately. Keep it below the hosting platform's SIGTERM-to-SIGKILL grace period. Also the fallback for `R_BULLMQ_SHUTDOWN_DRAIN_MS`. | +| `R_BULLMQ_SHUTDOWN_DRAIN_MS` | Optional | Same window for the bullmq service, which executes the Fast turns the queue resumes. Defaults to `R_API_SHUTDOWN_DRAIN_MS`, then `20000`. | | `R_FAST_DURABLE_RETRY_DISABLED` | Optional | Set to `true` to keep inference retry waits in the current process instead of parking in-flight Fast turns durably. Durable admission remains enabled. | | `R_CLOUD_ENABLED` | Roomote Cloud only | Deployment-managed switch for Roomote Cloud behavior, including required anonymous analytics and Cloud support integrations. Do not set this for self-hosted deployments. | | `R_CURATED_INTEGRATIONS_DISABLED` | Optional | Operator policy for the curated **Settings > Integrations** catalog, which is enabled by default. Set to `true` and restart Roomote to prevent those integrations from being configured or used. Existing connections remain stored while disabled and become available again once the value is unset. Communications, source-control, inference, sandbox providers, and environment-defined MCP servers are unaffected. | diff --git a/packages/cloud-agents/src/server/fast-agent/__tests__/fast-agent-turn-shutdown.test.ts b/packages/cloud-agents/src/server/fast-agent/__tests__/fast-agent-turn-shutdown.test.ts new file mode 100644 index 0000000000..119f822583 --- /dev/null +++ b/packages/cloud-agents/src/server/fast-agent/__tests__/fast-agent-turn-shutdown.test.ts @@ -0,0 +1,95 @@ +import { FastAgentProcessShutdownError } from '../fast-agent-turn-lock'; +import { + DEFAULT_FAST_AGENT_SHUTDOWN_DRAIN_MS, + drainAndAbortFastAgentTurns, + resolveFastAgentShutdownDrainMs, +} from '../fast-agent-turn-shutdown'; + +describe('resolveFastAgentShutdownDrainMs', () => { + it('defaults to a bounded drain window', () => { + expect(resolveFastAgentShutdownDrainMs(['A'], {})).toBe( + DEFAULT_FAST_AGENT_SHUTDOWN_DRAIN_MS, + ); + expect(resolveFastAgentShutdownDrainMs(['A'], { A: '' })).toBe(20_000); + expect(resolveFastAgentShutdownDrainMs(['A'], { A: 'nope' })).toBe(20_000); + expect(resolveFastAgentShutdownDrainMs(['A'], { A: '-5' })).toBe(20_000); + }); + + it('honors an explicit window, including the abort-immediately kill switch', () => { + expect(resolveFastAgentShutdownDrainMs(['A'], { A: '5000' })).toBe(5_000); + expect(resolveFastAgentShutdownDrainMs(['A'], { A: '0' })).toBe(0); + }); + + it('reads the first key that is set so a service can fall back to a shared one', () => { + expect(resolveFastAgentShutdownDrainMs(['A', 'B'], { B: '7000' })).toBe( + 7_000, + ); + expect( + resolveFastAgentShutdownDrainMs(['A', 'B'], { A: '3000', B: '7000' }), + ).toBe(3_000); + expect(resolveFastAgentShutdownDrainMs(['A', 'B'], { A: '', B: '0' })).toBe( + 0, + ); + }); +}); + +describe('drainAndAbortFastAgentTurns', () => { + it('closes admissions, waits out the window, then aborts the stragglers', async () => { + const reason = new FastAgentProcessShutdownError('SIGTERM'); + const order: string[] = []; + const beginDrain = vi.fn(() => order.push('begin')); + const onDrainStarted = vi.fn(() => order.push('started')); + let finishDrain: ((remaining: number) => void) | undefined; + const waitForTurns = vi.fn( + () => + new Promise((resolve) => { + order.push('wait'); + finishDrain = resolve; + }), + ); + const abortTurns = vi.fn(async () => { + order.push('abort'); + return 2; + }); + const logWarn = vi.fn(); + + const result = drainAndAbortFastAgentTurns( + { reason, drainMs: 1_500, service: 'bullmq', onDrainStarted, logWarn }, + { beginDrain, waitForTurns, abortTurns }, + ); + + expect(beginDrain).toHaveBeenCalledWith(reason); + expect(waitForTurns).toHaveBeenCalledWith(1_500); + expect(abortTurns).not.toHaveBeenCalled(); + + finishDrain?.(2); + await expect(result).resolves.toBe(2); + expect(order).toEqual(['begin', 'started', 'wait', 'abort']); + expect(abortTurns).toHaveBeenCalledWith(reason); + expect(logWarn).toHaveBeenCalledWith( + '[bullmq] Aborting 2 Fast turn(s) still active after the 1500ms shutdown drain.', + ); + }); + + it('stays quiet when every turn settles inside the window', async () => { + const logWarn = vi.fn(); + const abortTurns = vi.fn().mockResolvedValue(0); + + await drainAndAbortFastAgentTurns( + { + reason: new FastAgentProcessShutdownError('SIGINT'), + drainMs: 0, + service: 'api', + logWarn, + }, + { + beginDrain: vi.fn(), + waitForTurns: vi.fn().mockResolvedValue(0), + abortTurns, + }, + ); + + expect(logWarn).not.toHaveBeenCalled(); + expect(abortTurns).toHaveBeenCalledOnce(); + }); +}); diff --git a/packages/cloud-agents/src/server/fast-agent/fast-agent-turn-lock.ts b/packages/cloud-agents/src/server/fast-agent/fast-agent-turn-lock.ts index 117f994f67..2d8bb13891 100644 --- a/packages/cloud-agents/src/server/fast-agent/fast-agent-turn-lock.ts +++ b/packages/cloud-agents/src/server/fast-agent/fast-agent-turn-lock.ts @@ -19,9 +19,14 @@ export class FastAgentTurnLockLostError extends Error { } } +/** + * Abort reason for a Fast turn whose process is shutting down. Raised by + * every process that executes turns (the API for the turns it admits, the + * bullmq service for the turns the queue resumes). + */ export class FastAgentProcessShutdownError extends Error { - constructor(signal: NodeJS.Signals) { - super(`Fast turn interrupted by API shutdown (${signal}).`); + constructor(public readonly signal: NodeJS.Signals) { + super(`Fast turn interrupted by process shutdown (${signal}).`); this.name = 'FastAgentProcessShutdownError'; } } diff --git a/packages/cloud-agents/src/server/fast-agent/fast-agent-turn-shutdown.ts b/packages/cloud-agents/src/server/fast-agent/fast-agent-turn-shutdown.ts new file mode 100644 index 0000000000..d214cf94ba --- /dev/null +++ b/packages/cloud-agents/src/server/fast-agent/fast-agent-turn-shutdown.ts @@ -0,0 +1,83 @@ +import { + abortActiveFastAgentTurns, + beginFastAgentTurnDrain, + type FastAgentProcessShutdownError, + waitForActiveFastAgentTurnsToSettle, +} from './fast-agent-turn-lock'; + +/** + * Most Fast turns finish within seconds, so letting them settle turns a + * deploy-time interruption into a completed answer. The default leaves room + * for the straggler abort, the durable hand-back, and a Sentry flush inside + * a typical 30s SIGTERM-to-SIGKILL grace window. + */ +export const DEFAULT_FAST_AGENT_SHUTDOWN_DRAIN_MS = 20_000; + +/** + * Resolve the shutdown drain window from the first of `keys` that is set. A + * non-numeric or negative value falls back to the default; 0 is the + * abort-immediately kill switch. + */ +export function resolveFastAgentShutdownDrainMs( + keys: readonly string[], + env: NodeJS.ProcessEnv = process.env, +): number { + for (const key of keys) { + const raw = env[key]?.trim(); + if (!raw) continue; + const parsed = Number(raw); + return Number.isFinite(parsed) && parsed >= 0 + ? parsed + : DEFAULT_FAST_AGENT_SHUTDOWN_DRAIN_MS; + } + return DEFAULT_FAST_AGENT_SHUTDOWN_DRAIN_MS; +} + +export type FastAgentShutdownDrainDeps = { + beginDrain?: typeof beginFastAgentTurnDrain; + waitForTurns?: typeof waitForActiveFastAgentTurnsToSettle; + abortTurns?: typeof abortActiveFastAgentTurns; +}; + +/** + * The shutdown sequence for every process that executes Fast turns: the API + * runs the turns it admits, and the bullmq service runs the turns the queue + * resumes. Admissions close first, in-flight turns get the window to finish + * on their own, and only the stragglers are aborted, so each one hands its + * durable row back to the queue instead of dying with the claim and the + * conversation lock held. Returns how many turns were aborted. + */ +export async function drainAndAbortFastAgentTurns( + params: { + reason: FastAgentProcessShutdownError; + drainMs: number; + /** Process name for the straggler log line. */ + service: string; + /** + * Runs once admissions are closed, alongside the drain: stop accepting + * connections or fetching queue jobs so nothing new arrives while the + * turns already here finish. + */ + onDrainStarted?: () => void; + logWarn?: (message: string) => void; + }, + deps: FastAgentShutdownDrainDeps = {}, +): Promise { + const { + beginDrain = beginFastAgentTurnDrain, + waitForTurns = waitForActiveFastAgentTurnsToSettle, + abortTurns = abortActiveFastAgentTurns, + } = deps; + const logWarn = + params.logWarn ?? ((message: string) => console.warn(message)); + + beginDrain(params.reason); + params.onDrainStarted?.(); + const remaining = await waitForTurns(params.drainMs); + if (remaining > 0) { + logWarn( + `[${params.service}] Aborting ${remaining} Fast turn(s) still active after the ${params.drainMs}ms shutdown drain.`, + ); + } + return abortTurns(params.reason); +} diff --git a/packages/cloud-agents/src/server/fast-agent/index.ts b/packages/cloud-agents/src/server/fast-agent/index.ts index 45168378dd..d747be61e2 100644 --- a/packages/cloud-agents/src/server/fast-agent/index.ts +++ b/packages/cloud-agents/src/server/fast-agent/index.ts @@ -6,6 +6,7 @@ export * from './fast-agent-reply-stream'; export * from './fast-agent-surface-reply-stream'; export * from './fast-agent-service'; export * from './fast-agent-turn-lock'; +export * from './fast-agent-turn-shutdown'; export * from './fast-agent-session'; export * from './fast-agent-task-launcher'; export * from './fast-agent-title'; diff --git a/packages/sdk/src/server/lib/fast-agent-parent-event-queue.test.ts b/packages/sdk/src/server/lib/fast-agent-parent-event-queue.test.ts index 3aa002b3e4..947c1bdcd8 100644 --- a/packages/sdk/src/server/lib/fast-agent-parent-event-queue.test.ts +++ b/packages/sdk/src/server/lib/fast-agent-parent-event-queue.test.ts @@ -531,7 +531,20 @@ describe('Fast parent event durable queue', () => { .mockResolvedValueOnce({ deliveredAt: new Date(), discardedAt: null }) .mockResolvedValueOnce(undefined); mocks.acquireLock.mockResolvedValueOnce(mocks.releaseLock); - mocks.deliver.mockResolvedValueOnce('delivered'); + let boundDuringDelivery: + | { durableRowId?: string; durableResume?: () => Promise } + | undefined; + mocks.deliver.mockImplementationOnce(async () => { + const lock = mocks.releaseLock as typeof mocks.releaseLock & { + durableRowId?: string; + durableResume?: () => Promise; + }; + boundDuringDelivery = { + durableRowId: lock.durableRowId, + durableResume: lock.durableResume, + }; + return 'delivered'; + }); await drainFastAgentParentEvents({ conversationId: parent.sessionId, @@ -546,6 +559,19 @@ describe('Fast parent event durable queue', () => { }), mocks.releaseLock, ); + // While the resumed turn runs, its row is bound to the lock so a process + // shutdown that aborts it during setup can still release the claim and + // wake the queue; the binding does not outlive the delivery. + expect(boundDuringDelivery?.durableRowId).toBe('inline-1'); + await boundDuringDelivery?.durableResume?.(); + expect(mocks.queueAdd).toHaveBeenCalledWith( + 'deliver', + { conversationId: parent.sessionId, eventKey: inlineRow.eventKey }, + { jobId: inlineRow.eventKey }, + ); + expect( + (mocks.releaseLock as { durableRowId?: string }).durableRowId, + ).toBeUndefined(); // The resumed run settled its own row; the drain must not overwrite // that settlement (a replay-withdrawn row would otherwise also read as // delivered, losing its recorded reason). diff --git a/packages/sdk/src/server/lib/fast-agent-parent-event-queue.ts b/packages/sdk/src/server/lib/fast-agent-parent-event-queue.ts index 8927a6a514..22d917531b 100644 --- a/packages/sdk/src/server/lib/fast-agent-parent-event-queue.ts +++ b/packages/sdk/src/server/lib/fast-agent-parent-event-queue.ts @@ -374,6 +374,16 @@ export async function drainFastAgentParentEvents( conversationId: request.conversationId, eventKey: row.eventKey, }; + if (row.admission === 'inline') { + // Bind the row to the lock the way the inline surfaces do, so a + // process shutdown that aborts this turn before it reaches its own + // abort handling (still in setup, no inference yet) can release + // the claim and wake the queue instead of leaving the row held + // until its lease expires. + turnLock.durableRowId = row.id; + turnLock.durableResume = () => + wakeFastAgentParentEventNow(wakeRequest); + } await deliverFastAgentParentEventWithLock( { parent: row.parent, @@ -454,6 +464,11 @@ export async function drainFastAgentParentEvents( }) .where(eq(fastAgentParentEvents.id, row.id)); throw error; + } finally { + // The same lock carries every row this drain delivers; a settled or + // handed-back row must not stay bound to it. + delete turnLock.durableRowId; + delete turnLock.durableResume; } } } finally {