diff --git a/.env.production.example b/.env.production.example index 20791d1f1..c3b2c6b51 100644 --- a/.env.production.example +++ b/.env.production.example @@ -14,6 +14,11 @@ PREVIEW_PROXY_SUBDOMAIN_SUFFIX=preview # the separate develop soak env file. # 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. +# R_API_SHUTDOWN_DRAIN_MS=20000 + # Replace local defaults before exposing a shared deployment. DATABASE_URL=postgres://postgres:password@postgres:5432/roomote_development REDIS_URL=redis://redis:6379 diff --git a/apps/api/src/__tests__/graceful-shutdown.test.ts b/apps/api/src/__tests__/graceful-shutdown.test.ts index ac00384d9..37f0a7e78 100644 --- a/apps/api/src/__tests__/graceful-shutdown.test.ts +++ b/apps/api/src/__tests__/graceful-shutdown.test.ts @@ -1,9 +1,14 @@ 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}).`); @@ -17,46 +22,90 @@ import type { ServerType } from '@hono/node-server'; import { gracefullyShutdownApi, installApiGracefulShutdown, + resolveApiShutdownDrainMs, } from '../graceful-shutdown'; +describe('resolveApiShutdownDrainMs', () => { + it('defaults to a bounded drain window', () => { + expect(resolveApiShutdownDrainMs({})).toBe(20_000); + expect(resolveApiShutdownDrainMs({ R_API_SHUTDOWN_DRAIN_MS: '' })).toBe( + 20_000, + ); + expect( + resolveApiShutdownDrainMs({ R_API_SHUTDOWN_DRAIN_MS: 'not-a-number' }), + ).toBe(20_000); + expect(resolveApiShutdownDrainMs({ R_API_SHUTDOWN_DRAIN_MS: '-5' })).toBe( + 20_000, + ); + }); + + it('honors an explicit window, including the abort-immediately kill switch', () => { + expect(resolveApiShutdownDrainMs({ R_API_SHUTDOWN_DRAIN_MS: '5000' })).toBe( + 5_000, + ); + expect(resolveApiShutdownDrainMs({ R_API_SHUTDOWN_DRAIN_MS: '0' })).toBe(0); + }); +}); + describe('gracefullyShutdownApi', () => { beforeEach(() => { vi.clearAllMocks(); }); - it('aborts Fast turns before waiting for active requests to close', async () => { + it('drains active Fast turns before aborting the stragglers', async () => { let finishClose: ((error?: Error) => void) | undefined; const server = { close: vi.fn((callback: (error?: Error) => void) => { finishClose = callback; }), } as unknown as ServerType; - let finishAbort: (() => void) | undefined; - const abortTurns = vi.fn( + const beginDrain = vi.fn(); + let finishDrain: ((remaining: number) => void) | undefined; + const waitForTurns = vi.fn( () => new Promise((resolve) => { - finishAbort = () => resolve(1); + finishDrain = resolve; }), ); + const abortTurns = vi.fn().mockResolvedValue(1); const exitProcess = vi.fn() as unknown as (code?: number) => never; const flushSentry = vi.fn().mockResolvedValue(undefined); + const logWarn = vi.fn(); const shutdown = gracefullyShutdownApi(server, 'SIGTERM', { abortTurns, + beginDrain, + waitForTurns, + drainMs: 12_345, exitProcess, flushSentry, + logWarn, }); - expect(abortTurns).toHaveBeenCalledWith( + // Admissions close and the drain starts immediately; nothing is aborted + // while in-flight turns still have time to finish on their own. + expect(beginDrain).toHaveBeenCalledWith( expect.objectContaining({ name: 'FastAgentProcessShutdownError', signal: 'SIGTERM', }), ); expect(server.close).toHaveBeenCalledOnce(); - expect(exitProcess).not.toHaveBeenCalled(); + expect(waitForTurns).toHaveBeenCalledWith(12_345); + expect(abortTurns).not.toHaveBeenCalled(); + + finishDrain?.(1); + await vi.waitFor(() => expect(abortTurns).toHaveBeenCalledOnce()); + expect(abortTurns).toHaveBeenCalledWith( + expect.objectContaining({ + name: 'FastAgentProcessShutdownError', + signal: 'SIGTERM', + }), + ); + expect(logWarn).toHaveBeenCalledWith( + '[api] Aborting 1 Fast turn(s) still active after the 12345ms shutdown drain.', + ); - finishAbort?.(); finishClose?.(); await shutdown; @@ -64,6 +113,28 @@ describe('gracefullyShutdownApi', () => { expect(exitProcess).toHaveBeenCalledWith(0); }); + it('stays quiet when every turn settles inside the drain window', async () => { + const server = { + close: vi.fn((callback: (error?: Error) => void) => callback()), + } as unknown as ServerType; + const abortTurns = vi.fn().mockResolvedValue(0); + const exitProcess = vi.fn() as unknown as (code?: number) => never; + const logWarn = vi.fn(); + + await gracefullyShutdownApi(server, 'SIGTERM', { + abortTurns, + beginDrain: vi.fn(), + waitForTurns: vi.fn().mockResolvedValue(0), + drainMs: 20_000, + exitProcess, + logWarn, + }); + + expect(logWarn).not.toHaveBeenCalled(); + expect(abortTurns).toHaveBeenCalledOnce(); + expect(exitProcess).toHaveBeenCalledWith(0); + }); + it('exits unsuccessfully when the HTTP server cannot close', async () => { const closeError = new Error('close failed'); const server = { @@ -76,6 +147,9 @@ describe('gracefullyShutdownApi', () => { await gracefullyShutdownApi(server, 'SIGINT', { abortTurns, + beginDrain: vi.fn(), + waitForTurns: vi.fn().mockResolvedValue(0), + drainMs: 0, exitProcess, flushSentry, logError, @@ -105,6 +179,8 @@ describe('gracefullyShutdownApi', () => { try { const cleanup = installApiGracefulShutdown(server, { abortTurns: vi.fn(() => new Promise(() => undefined)), + beginDrain: vi.fn(), + waitForTurns: vi.fn(() => new Promise(() => undefined)), exitProcess, }); diff --git a/apps/api/src/graceful-shutdown.ts b/apps/api/src/graceful-shutdown.ts index 69bb0413c..59e9b0e50 100644 --- a/apps/api/src/graceful-shutdown.ts +++ b/apps/api/src/graceful-shutdown.ts @@ -1,14 +1,38 @@ import type { ServerType } from '@hono/node-server'; import { abortActiveFastAgentTurns, + beginFastAgentTurnDrain, FastAgentProcessShutdownError, + waitForActiveFastAgentTurnsToSettle, } 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; + +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; +} + type ApiShutdownOptions = { abortTurns?: typeof abortActiveFastAgentTurns; + beginDrain?: typeof beginFastAgentTurnDrain; + waitForTurns?: typeof waitForActiveFastAgentTurnsToSettle; + drainMs?: number; exitProcess?: (code?: number) => never; flushSentry?: () => Promise; logError?: (...args: Parameters) => void; + logWarn?: (...args: Parameters) => void; }; export async function gracefullyShutdownApi( @@ -16,15 +40,31 @@ export async function gracefullyShutdownApi( signal: NodeJS.Signals, { abortTurns = abortActiveFastAgentTurns, + beginDrain = beginFastAgentTurnDrain, + waitForTurns = waitForActiveFastAgentTurnsToSettle, + drainMs = resolveApiShutdownDrainMs(), exitProcess = process.exit, flushSentry = async () => undefined, logError = (...args) => console.error(...args), + logWarn = (...args) => console.warn(...args), }: ApiShutdownOptions = {}, ): Promise { - const abortPromise = abortTurns(new FastAgentProcessShutdownError(signal)); - const closeError = await new Promise((resolve) => { + const reason = new FastAgentProcessShutdownError(signal); + // 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); + const closeError = await closePromise; await abortPromise; if (closeError) { logError('[api] Graceful shutdown failed', closeError); diff --git a/packages/cloud-agents/src/server/fast-agent/__tests__/fast-agent-turn-drain.test.ts b/packages/cloud-agents/src/server/fast-agent/__tests__/fast-agent-turn-drain.test.ts new file mode 100644 index 000000000..411eabefe --- /dev/null +++ b/packages/cloud-agents/src/server/fast-agent/__tests__/fast-agent-turn-drain.test.ts @@ -0,0 +1,110 @@ +const { acquireRedisLockMock } = vi.hoisted(() => ({ + acquireRedisLockMock: vi.fn(), +})); + +vi.mock('@roomote/redis', () => ({ + acquireRedisLock: acquireRedisLockMock, +})); + +type TurnLockModule = typeof import('../fast-agent-turn-lock'); + +const conversation = { + surface: 'slack', + workspaceId: 'workspace-1', + conversationId: 'conversation-1', + replyTarget: { channelId: 'channel-1', threadId: 'conversation-1' }, +} as const; + +const otherConversation = { + surface: 'slack', + workspaceId: 'workspace-1', + conversationId: 'conversation-2', + replyTarget: { channelId: 'channel-1', threadId: 'conversation-2' }, +} as const; + +function buildRedisLock() { + return Object.assign(vi.fn().mockResolvedValue(undefined), { + renew: vi.fn().mockResolvedValue(true), + renewDetailed: vi.fn().mockResolvedValue('renewed'), + }); +} + +// The drain flag is module-global by design (a draining process never +// accepts turns again), so every test runs against a fresh module instance. +describe('Fast turn shutdown drain', () => { + let turnLock: TurnLockModule; + + beforeEach(async () => { + vi.resetModules(); + vi.clearAllMocks(); + acquireRedisLockMock.mockImplementation(async () => buildRedisLock()); + turnLock = await import('../fast-agent-turn-lock'); + }); + + it('refuses new admissions during drain without aborting active turns', async () => { + const activeLock = await turnLock.acquireFastAgentTurnLock({ + conversation, + }); + expect(activeLock).not.toBeNull(); + + turnLock.beginFastAgentTurnDrain( + new turnLock.FastAgentProcessShutdownError('SIGTERM'), + ); + + await expect( + turnLock.acquireFastAgentTurnLock({ conversation: otherConversation }), + ).resolves.toBeNull(); + expect(activeLock!.signal.aborted).toBe(false); + + await activeLock!(); + }); + + it('resolves the drain wait as soon as the last active turn settles', async () => { + await expect( + turnLock.waitForActiveFastAgentTurnsToSettle(5_000), + ).resolves.toBe(0); + + const activeLock = await turnLock.acquireFastAgentTurnLock({ + conversation, + }); + let settled = false; + const waiting = turnLock + .waitForActiveFastAgentTurnsToSettle(5_000) + .then((remaining) => { + settled = true; + return remaining; + }); + await new Promise((resolve) => setTimeout(resolve, 20)); + expect(settled).toBe(false); + + await activeLock!(); + await expect(waiting).resolves.toBe(0); + }); + + it('reports the stragglers still active at the drain deadline', async () => { + const straggler = await turnLock.acquireFastAgentTurnLock({ conversation }); + + await expect( + turnLock.waitForActiveFastAgentTurnsToSettle(20), + ).resolves.toBe(1); + expect(straggler!.signal.aborted).toBe(false); + + await straggler!(); + }); + + it('aborts stragglers with the reason the drain began with', async () => { + const straggler = await turnLock.acquireFastAgentTurnLock({ conversation }); + const drainReason = new turnLock.FastAgentProcessShutdownError('SIGTERM'); + + turnLock.beginFastAgentTurnDrain(drainReason); + await expect( + turnLock.abortActiveFastAgentTurns( + new turnLock.FastAgentProcessShutdownError('SIGINT'), + ), + ).resolves.toBe(1); + + expect(straggler!.signal.aborted).toBe(true); + expect(straggler!.signal.reason).toBe(drainReason); + await straggler!(); + }); +}); 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 275064bf6..974b2e579 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 @@ -63,6 +63,45 @@ export async function abortActiveFastAgentTurns( return activeLocks.length; } +const turnSettleWaiters = new Set<() => void>(); + +function notifyTurnSettleWaitersIfIdle() { + if (activeFastAgentTurnLocks.size > 0) return; + for (const waiter of [...turnSettleWaiters]) waiter(); +} + +/** + * Refuse new Fast turn admissions without aborting the active ones, so a + * shutdown can let in-flight turns finish before interrupting the remainder. + */ +export function beginFastAgentTurnDrain( + reason: FastAgentProcessShutdownError, +): void { + processShutdownReason ??= reason; +} + +/** + * Resolve once every active Fast turn has settled or the deadline passes. + * Returns the number of turns still active at that point. + */ +export async function waitForActiveFastAgentTurnsToSettle( + timeoutMs: number, +): Promise { + if (timeoutMs > 0 && activeFastAgentTurnLocks.size > 0) { + await new Promise((resolve) => { + const settle = () => { + clearTimeout(deadline); + turnSettleWaiters.delete(settle); + resolve(); + }; + turnSettleWaiters.add(settle); + const deadline = setTimeout(settle, timeoutMs); + deadline.unref(); + }); + } + return activeFastAgentTurnLocks.size; +} + /** Serialize every human and platform-generated Fast turn for one chat. */ export function buildFastAgentTurnLockKey( conversation: FastAgentConversation, @@ -140,6 +179,7 @@ export async function acquireFastAgentTurnLock(params: { if (turnSettled) return; turnSettled = true; activeFastAgentTurnLocks.delete(releaseTurnLock); + notifyTurnSettleWaitersIfIdle(); try { if ( ownership.signal.reason instanceof FastAgentProcessShutdownError