From 549e7ade2d848906e505c9ae972c9f3533d80132 Mon Sep 17 00:00:00 2001 From: Lily Shen <115414357+lilyshen0722@users.noreply.github.com> Date: Sun, 25 Oct 2026 01:44:59 -0700 Subject: [PATCH 1/3] fix(runtime): drain hosted turns before backend shutdown Co-authored-by: rivet Agent-Model: codex/gpt-6.1-sol/medium --- .../service/hosted-turn-drain.test.js | 106 ++++++++++++++++ .../unit/routes/health.ready.test.js | 9 ++ .../services/hostedShutdownService.test.js | 36 ++++++ backend/routes/health.ts | 5 + backend/server.ts | 1 + backend/services/agentEventService.ts | 120 +++++++++++++++++- backend/services/hostedShutdownService.ts | 37 ++++++ docs/runbooks/hosted-turn-shutdown.md | 19 +++ .../templates/core/backend-deployment.yaml | 2 +- 9 files changed, 327 insertions(+), 8 deletions(-) create mode 100644 backend/__tests__/service/hosted-turn-drain.test.js create mode 100644 backend/__tests__/unit/services/hostedShutdownService.test.js create mode 100644 backend/services/hostedShutdownService.ts create mode 100644 docs/runbooks/hosted-turn-shutdown.md diff --git a/backend/__tests__/service/hosted-turn-drain.test.js b/backend/__tests__/service/hosted-turn-drain.test.js new file mode 100644 index 0000000000..11df4c5126 --- /dev/null +++ b/backend/__tests__/service/hosted-turn-drain.test.js @@ -0,0 +1,106 @@ +process.env.PG_HOST = ''; +jest.mock('../../services/nativeRuntimeService', () => ({ runAgent: jest.fn() })); +jest.mock('../../services/agentWebSocketService', () => ({ pushEvent: jest.fn() })); +jest.mock('../../services/agentDigestFoldService', () => ({ readFoldedDigest: jest.fn().mockResolvedValue(null), settleFoldedDigest: jest.fn().mockResolvedValue(undefined), foldExpiredDigests: jest.fn().mockResolvedValue(0) })); +const mongoose = require('mongoose'); +const AgentEvent = require('../../models/AgentEvent'); +const Service = require('../../services/agentEventService'); +const { runAgent } = require('../../services/nativeRuntimeService'); +const { setupMongoDb, clearMongoDb, closeMongoDb, waitFor } = require('../utils/testUtils'); +const podId = new mongoose.Types.ObjectId(); +const install = { agentName: 'clerk', instanceId: 'default', runtimeType: 'internal', podId }; +const input = { agentName: 'clerk', podId, type: 'test.long-turn', payload: { content: 'work' } }; +const deferred = () => { let resolve; const promise = new Promise((done) => { resolve = done; }); return { promise, resolve }; }; +const success = { status: 'succeeded', runId: 'run-long', totalTurns: 1 }; +beforeAll(setupMongoDb); +beforeEach(async () => { + // Simulate a fresh process for each rollout, without a production reset API. + Service.inProcessRequeued = 0; Service.inProcessDraining = false; Service.inProcessDrain = undefined; + Service.inProcessWork.clear(); Service.inProcessRunning.clear(); + jest.clearAllMocks(); await clearMongoDb(); + jest.spyOn(Service, 'resolveInProcessInstallation').mockResolvedValue(install); +}); +afterEach(() => jest.restoreAllMocks()); +afterAll(closeMongoDb); + +test('fake long turn is acked before drain resolves; new native events stay pending', async () => { + const turn = deferred(); runAgent.mockReturnValue(turn.promise); + const event = await Service.enqueue(input); + await waitFor(() => runAgent.mock.calls.length === 1); + let finished = false; + const drain = Service.drainInProcess(2000).then((result) => { finished = true; return result; }); + expect(Service.isInProcessDraining()).toBe(true); + const queued = await Service.enqueue(input); + expect((await AgentEvent.findById(queued._id)).status).toBe('pending'); + expect(runAgent).toHaveBeenCalledTimes(1); expect(finished).toBe(false); + turn.resolve(success); + await expect(drain).resolves.toEqual({ timedOut: false, requeued: 0 }); + expect((await AgentEvent.findById(event._id)).status).toBe('acked'); + expect(await Service.rescueStrandedPendingInProcess(new Date(), 5)).toEqual({ rescued: 0, closed: 0 }); + expect(await Service.redispatchStuckInProcess(new Date(), 5)).toEqual({ redispatched: 0, acked: 0 }); +}); +test('deadline requeues only local turns; late success cannot ack a replacement', async () => { + const turn = deferred(); runAgent.mockReturnValue(turn.promise); + const event = await Service.enqueue(input); await waitFor(() => runAgent.mock.calls.length === 1); + const sibling = await AgentEvent.create({ ...input, status: 'delivered', deliveryNonce: 'sibling-nonce', attempts: 1 }); + await expect(Service.drainInProcess(20)).resolves.toEqual({ timedOut: true, requeued: 1 }); + expect(await AgentEvent.findById(event._id).lean()).toMatchObject({ status: 'pending', deliveryNonce: null, attempts: 1 }); + expect((await AgentEvent.findById(sibling._id)).status).toBe('delivered'); + await AgentEvent.updateOne({ _id: event._id }, { $set: { status: 'delivered', deliveryNonce: 'replacement-nonce' } }); + turn.resolve(success); await waitFor(() => Service.inProcessWork.size === 0); + expect(await AgentEvent.findById(event._id).lean()).toMatchObject({ status: 'delivered', deliveryNonce: 'replacement-nonce' }); +}); +test('a create claim crossing SIGTERM is requeued without starting a model', async () => { + const created = deferred(), release = deferred(), original = AgentEvent.create.bind(AgentEvent); + jest.spyOn(AgentEvent, 'create').mockImplementationOnce(async (...args) => { + const row = await original(...args); created.resolve(row); await release.promise; return row; + }); + const enqueue = Service.enqueue(input), claimed = await created.promise; + const drain = Service.drainInProcess(2000); release.resolve(); await enqueue; + await expect(drain).resolves.toEqual({ timedOut: false, requeued: 1 }); + expect((await AgentEvent.findById(claimed._id)).status).toBe('pending'); expect(runAgent).not.toHaveBeenCalled(); +}); +test('a failed in-flight turn settles failure before drain completes', async () => { + const turn = deferred(); runAgent.mockReturnValue(turn.promise); + const event = await Service.enqueue(input); await waitFor(() => runAgent.mock.calls.length === 1); + const drain = Service.drainInProcess(2000); turn.resolve({ status: 'failed', error: 'provider refused' }); await drain; + expect((await AgentEvent.findById(event._id)).status).toBe('failed'); +}); +test('deadline cleanup cannot steal a delivery reclaimed by a sibling', async () => { + const turn = deferred(); runAgent.mockReturnValue(turn.promise); + const event = await Service.enqueue(input); await waitFor(() => runAgent.mock.calls.length === 1); + await AgentEvent.updateOne({ _id: event._id }, { $set: { deliveryNonce: 'newer-owner' } }); + await expect(Service.drainInProcess(20)).resolves.toEqual({ timedOut: true, requeued: 0 }); + expect(await AgentEvent.findById(event._id).lean()).toMatchObject({ status: 'delivered', deliveryNonce: 'newer-owner' }); + turn.resolve(success); await waitFor(() => Service.inProcessWork.size === 0); +}); + + +test('drain waits for the successful terminal write, not just the model promise', async () => { + const ack = deferred(), entered = deferred(), original = Service.acknowledge.bind(Service); + jest.spyOn(Service, 'acknowledge').mockImplementation(async (...args) => { + entered.resolve(); await ack.promise; return original(...args); + }); + runAgent.mockResolvedValue(success); + const event = await Service.enqueue(input); await entered.promise; + let complete = false; + const drain = Service.drainInProcess(2000).then(() => { complete = true; }); + await Promise.resolve(); expect(complete).toBe(false); + ack.resolve(); await drain; + expect((await AgentEvent.findById(event._id)).status).toBe('acked'); +}); + +test.each(['pending', 'delivered'])('SIGTERM during a %s recovery lookup prevents the claim', async (status) => { + const entered = deferred(), release = deferred(); + Service.resolveInProcessInstallation.mockImplementation(async () => { entered.resolve(); await release.promise; return install; }); + jest.spyOn(Service, 'inProcessEventTargets').mockResolvedValue([{ agentName: 'clerk', instanceId: 'default', podId }]); + const old = new Date(Date.now() - 5 * 60000); + const event = await AgentEvent.create({ ...input, status, createdAt: old, deliveredAt: old, deliveryNonce: 'old-nonce', attempts: 1 }); + const sweep = status === 'pending' + ? Service.rescueStrandedPendingInProcess(new Date(), 5) + : Service.redispatchStuckInProcess(new Date(), 5); + await entered.promise; + const drain = Service.drainInProcess(2000); release.resolve(); await sweep; await drain; + expect(await AgentEvent.findById(event._id).lean()).toMatchObject({ status, deliveryNonce: 'old-nonce', attempts: 1 }); + expect(runAgent).not.toHaveBeenCalled(); +}); diff --git a/backend/__tests__/unit/routes/health.ready.test.js b/backend/__tests__/unit/routes/health.ready.test.js index ab080597ff..624b71aa37 100644 --- a/backend/__tests__/unit/routes/health.ready.test.js +++ b/backend/__tests__/unit/routes/health.ready.test.js @@ -10,6 +10,8 @@ const mockPool = { query: jest.fn().mockResolvedValue({ rows: [{ ok: 1 }] }), }; const mockCreateClient = jest.fn(); +const mockIsDraining = jest.fn().mockReturnValue(false); +jest.mock('../../../services/agentEventService', () => ({ isInProcessDraining: () => mockIsDraining() })); jest.mock('mongoose', () => mockMongoose); jest.mock('redis', () => ({ createClient: mockCreateClient })); @@ -41,6 +43,7 @@ const buildApp = () => { describe('GET /api/health/ready', () => { beforeEach(() => { + mockIsDraining.mockReturnValue(false); mockMongoose.connection.readyState = 1; mockPool.query.mockResolvedValue({ rows: [{ ok: 1 }] }); mockPgRoutesAreMounted.mockReturnValue(true); @@ -54,6 +57,12 @@ describe('GET /api/health/ready', () => { else process.env.AGENT_PROVISIONER_K8S = originalK8sMode; }); + it('refuses readiness while hosted turns drain', async () => { + mockIsDraining.mockReturnValue(true); + const res = await request(buildApp()).get('/api/health/ready').expect(503); + expect(res.body.reason).toBe('hosted runtime draining'); + }); + it('returns ready when MongoDB and PostgreSQL are available without probing Redis', async () => { const res = await request(buildApp()).get('/api/health/ready').expect(200); diff --git a/backend/__tests__/unit/services/hostedShutdownService.test.js b/backend/__tests__/unit/services/hostedShutdownService.test.js new file mode 100644 index 0000000000..12f6b642f9 --- /dev/null +++ b/backend/__tests__/unit/services/hostedShutdownService.test.js @@ -0,0 +1,36 @@ +jest.mock('../../../services/agentEventService', () => ({ drainInProcess: jest.fn() })); +jest.mock('../../../services/schedulerService', () => ({ stop: jest.fn() })); +const fs = require('fs'), path = require('path'); +const service = require('../../../services/agentEventService'); +const scheduler = require('../../../services/schedulerService'); +const { installHostedShutdown, HOSTED_SHUTDOWN_WATCHDOG_MS } = require('../../../services/hostedShutdownService'); +let handlers, server, exit; +beforeEach(() => { + jest.useFakeTimers(); jest.clearAllMocks(); handlers = {}; + jest.spyOn(process, 'once').mockImplementation((name, handler) => { handlers[name] = handler; return process; }); + exit = jest.spyOn(process, 'exit').mockImplementation(() => {}); + server = { close: jest.fn(), closeIdleConnections: jest.fn() }; +}); +afterEach(() => { jest.restoreAllMocks(); jest.useRealTimers(); }); +const flush = async () => { await Promise.resolve(); await Promise.resolve(); }; +test('SIGTERM gates admission before closing HTTP; exits after drain; repeated signals are idempotent', async () => { + let resolve; service.drainInProcess.mockImplementation(() => new Promise((done) => { resolve = done; })); + installHostedShutdown(server); handlers.SIGTERM(); handlers.SIGINT(); + expect(service.drainInProcess).toHaveBeenCalledTimes(1); expect(scheduler.stop).toHaveBeenCalledTimes(1); + expect(service.drainInProcess.mock.invocationCallOrder[0]).toBeLessThan(server.close.mock.invocationCallOrder[0]); + expect(server.close).toHaveBeenCalledTimes(1); expect(exit).not.toHaveBeenCalled(); + resolve({ timedOut: false, requeued: 0 }); await flush(); expect(exit).toHaveBeenCalledWith(0); + jest.advanceTimersByTime(HOSTED_SHUTDOWN_WATCHDOG_MS); expect(exit).toHaveBeenCalledTimes(1); +}); +test('watchdog bounds stuck cleanup inside the Kubernetes budget', () => { + service.drainInProcess.mockReturnValue(new Promise(() => {})); + installHostedShutdown(server); handlers.SIGTERM(); + jest.advanceTimersByTime(HOSTED_SHUTDOWN_WATCHDOG_MS - 1); expect(exit).not.toHaveBeenCalled(); + jest.advanceTimersByTime(1); expect(exit).toHaveBeenCalledWith(1); + const chart = fs.readFileSync(path.resolve(__dirname, '../../../../k8s/helm/commonly/templates/core/backend-deployment.yaml'), 'utf8'); + expect(Number(/terminationGracePeriodSeconds: (\d+)/.exec(chart)[1]) * 1000).toBeGreaterThan(10000 + HOSTED_SHUTDOWN_WATCHDOG_MS); +}); +test('cleanup failure exits nonzero', async () => { + service.drainInProcess.mockRejectedValue(new Error('database unavailable')); + installHostedShutdown(server); handlers.SIGTERM(); await flush(); expect(exit).toHaveBeenCalledWith(1); +}); diff --git a/backend/routes/health.ts b/backend/routes/health.ts index ece324cd37..46acc1d7d7 100644 --- a/backend/routes/health.ts +++ b/backend/routes/health.ts @@ -168,6 +168,11 @@ router.get('/db', (_req: unknown, res: Res) => { }); router.get('/ready', async (_req: unknown, res: Res) => { + // Draining is not a liveness failure; only readiness stops new traffic. + const AgentEventService = require('../services/agentEventService'); + if (AgentEventService.isInProcessDraining()) { + return res.status(503).json({ status: 'not_ready', reason: 'hosted runtime draining' }); + } try { if (mongoose.connection.readyState !== 1) { return res.status(503).json({ status: 'not_ready', reason: 'MongoDB not connected' }); diff --git a/backend/server.ts b/backend/server.ts index 13a0e2de4b..83d9408c34 100644 --- a/backend/server.ts +++ b/backend/server.ts @@ -834,6 +834,7 @@ io.on('connection', (socket: any) => { // Start the server only when executed directly if (require.main === module) { + require('./services/hostedShutdownService').installHostedShutdown(server); server.listen(PORT, () => console.log(`Server running on port ${PORT}`)); } diff --git a/backend/services/agentEventService.ts b/backend/services/agentEventService.ts index 8476c93b41..1fcc87050b 100644 --- a/backend/services/agentEventService.ts +++ b/backend/services/agentEventService.ts @@ -492,6 +492,66 @@ const signalAgentTyping = async (event: EventDoc): Promise => { }; class AgentEventService { + private static inProcessDraining = false; + private static inProcessWork = new Set>(); + private static inProcessRequeued = 0; + private static inProcessRunning = new Map(); + private static inProcessDrain: Promise<{ timedOut: boolean; requeued: number }> | undefined; + + static isInProcessDraining(): boolean { return this.inProcessDraining; } + + // Track admission DB work as well as turns: SIGTERM can race a claim whose + // database reply has not arrived yet. Dispatch will requeue that late claim. + private static trackInProcessWork(work: () => Promise): Promise { + const pending = Promise.resolve().then(work); + this.inProcessWork.add(pending); + void pending.finally(() => this.inProcessWork.delete(pending)).catch(() => undefined); + return pending; + } + + private static async requeueInProcess(event: EventDoc): Promise { + const result = await AgentEvent.updateOne( + { _id: event._id, status: 'delivered', deliveryNonce: event.deliveryNonce ?? null }, + { $set: { status: 'pending', deliveredAt: null, deliveryNonce: null } }, + ) as { modifiedCount?: number }; + if (result.modifiedCount) { + this.inProcessRequeued += result.modifiedCount; + event.status = 'pending'; + event.deliveryNonce = null; + } + return result.modifiedCount || 0; + } + + /** Close native admission synchronously, then wait through terminal event writes. + * The caller exits after this returns. Deadline recovery is at-least-once; + * provider effects already performed by an unfinished turn cannot be undone. + */ + static drainInProcess(timeoutMs = 150_000): Promise<{ timedOut: boolean; requeued: number }> { + this.inProcessDraining = true; + if (this.inProcessDrain) return this.inProcessDrain; + this.inProcessDrain = (async () => { + let timer: ReturnType | undefined; + const settled = async () => { + while (this.inProcessWork.size) { + await Promise.allSettled([...this.inProcessWork]); + } + return false; + }; + const timedOut = await Promise.race([ + settled(), + new Promise((resolve) => { timer = setTimeout(() => resolve(true), timeoutMs); }), + ]); + if (timer) clearTimeout(timer); + if (!timedOut) return { timedOut: false, requeued: this.inProcessRequeued }; + // Requeue only this process's unfinished deliveries, never a pod-wide + // query that could steal a sibling replica's turn. The nonce fences late + // success/failure settlement from a replacement's claim. + await Promise.all([...this.inProcessRunning.values()].map((event) => this.requeueInProcess(event))); + return { timedOut: true, requeued: this.inProcessRequeued }; + })(); + return this.inProcessDrain; + } + static getContextOverflowRetryLimit(): number { const parsed = Number.parseInt(process.env.AGENT_CONTEXT_OVERFLOW_RETRY_LIMIT || '', 10); if (!Number.isFinite(parsed) || parsed < 0) return 1; @@ -998,7 +1058,12 @@ class AgentEventService { * acked, not answered twice. Either way the dead turn's AgentRun, if it was * left 'queued' or 'running', is marked 'interrupted'. */ - static async redispatchStuckInProcess( + static redispatchStuckInProcess(threshold: Date, attemptCap: number): Promise<{ redispatched: number; acked: number }> { + if (this.inProcessDraining) return Promise.resolve({ redispatched: 0, acked: 0 }); + return this.trackInProcessWork(() => this.redispatchStuckInProcessWhileServing(threshold, attemptCap)); + } + + private static async redispatchStuckInProcessWhileServing( threshold: Date, attemptCap: number, ): Promise<{ redispatched: number; acked: number }> { @@ -1023,6 +1088,7 @@ class AgentEventService { const runKey = { podId: stale.podId, agentName, instanceId, triggerEventId: stale._id }; // eslint-disable-next-line no-await-in-loop const finished = await AgentRun.exists({ ...runKey, status: 'succeeded' }); + if (this.inProcessDraining) break; const deliveryNonce = mintDeliveryNonce(); // eslint-disable-next-line no-await-in-loop const claimed = await AgentEvent.findOneAndUpdate( @@ -1103,7 +1169,12 @@ class AgentEventService { * attempt cap bounds a poison event. Only targets that would run are * selected, and each is re-checked by enqueue's own routing before it runs. */ - static async rescueStrandedPendingInProcess( + static rescueStrandedPendingInProcess(now: Date, attemptCap: number): Promise<{ rescued: number; closed: number }> { + if (this.inProcessDraining) return Promise.resolve({ rescued: 0, closed: 0 }); + return this.trackInProcessWork(() => this.rescueStrandedPendingInProcessWhileServing(now, attemptCap)); + } + + private static async rescueStrandedPendingInProcessWhileServing( now: Date, attemptCap: number, ): Promise<{ rescued: number; closed: number }> { @@ -1142,6 +1213,7 @@ class AgentEventService { .catch(() => null); // Changed since the targets were read (a deprovision, an install removed). if (!installation) continue; + if (this.inProcessDraining) break; const deliveryNonce = mintDeliveryNonce(); // eslint-disable-next-line no-await-in-loop const claimed = await AgentEvent.findOneAndUpdate( @@ -1502,6 +1574,24 @@ class AgentEventService { instanceId: string, type: string, ): void { + void this.trackInProcessWork(async () => { + if (this.inProcessDraining) { await this.requeueInProcess(event); return; } + const deliveryKey = `${String(event._id)}:${event.deliveryNonce || ''}`; + this.inProcessRunning.set(deliveryKey, event); + try { + await this.runInProcess(event, nativeInstallation, agentName, instanceId, type, () => { + this.inProcessRunning.delete(deliveryKey); + }); + } finally { + this.inProcessRunning.delete(deliveryKey); + } + }).catch((error: Error) => console.warn('[native-runtime] dispatch error:', error.message)); + } + + private static async runInProcess( + event: EventDoc, nativeInstallation: InstallationDoc, agentName: string, + instanceId: string, type: string, turnFinished: () => void, + ): Promise { try { // eslint-disable-next-line global-require, @typescript-eslint/no-require-imports const { runAgent } = require('./nativeRuntimeService'); @@ -1541,7 +1631,14 @@ class AgentEventService { ...event.payload, content: `${baseContent}\n\n${summary}`, } : event.payload; + if (this.inProcessDraining) { + await this.requeueInProcess(event); + return { status: 'drained' }; + } const result = await runAgent(nativeInstallation, { type, eventId: eventIdStr, payload }); + // The model is done even if a digest/ack write stalls. Leave that + // delivery for succeeded-run recovery instead of replaying its output. + turnFinished(); // Leave digest unread when the native turn fails. On a successful turn, // mark only the rows actually included; older backlog remains pending. const modelTurnRan = result?.status === 'succeeded' @@ -1556,14 +1653,17 @@ class AgentEventService { } return result; }; - Promise.resolve(runWithDigest()).then( + await Promise.resolve(runWithDigest()).then( // Two-argument .then, deliberately, NOT .then().catch(): a chained // .catch would also catch a rejection thrown by the success handler, // so a failing acknowledge() would fall through and mark a run that // actually succeeded as 'failed'. The rejection handler here is scoped // to runAgent alone. Both terminal calls swallow their own errors so // this stays fire-and-forget for the caller. - (runResult: { status?: string; error?: string; runId?: string; totalTurns?: number }) => (runResult?.status === 'failed' + (runResult: { status?: string; error?: string; runId?: string; totalTurns?: number }) => { + turnFinished(); + if (runResult?.status === 'drained') return undefined; + return (runResult?.status === 'failed' ? this.recordFailure( eventIdStr, agentName, @@ -1582,8 +1682,10 @@ class AgentEventService { 'native', ).catch((ackErr: Error) => { console.warn('[native-runtime] ack after successful run failed:', ackErr?.message || ackErr); - })), + })); + }, (err: Error) => { + turnFinished(); console.error('[native-runtime] runAgent failed:', err?.message || err); // Terminal on failure too — a native event that errored must not // be left for the requeue to hand to an external poller that @@ -1605,7 +1707,11 @@ class AgentEventService { } } - static async enqueue({ + static enqueue(options: EnqueueOptions): Promise { + return this.trackInProcessWork(() => this.enqueueEvent(options)); + } + + private static async enqueueEvent({ agentName, podId, type, payload = {}, instanceId = 'default', deliveryMode = 'wake', }: EnqueueOptions): Promise { if (!agentName || !podId || !type) { @@ -1646,7 +1752,7 @@ class AgentEventService { } // A digest is data for a later turn, not a turn on its own. Native // installations must leave it pending just like external installations. - const routedToNative = Boolean(nativeInstallation) && deliveryMode === 'wake'; + const routedToNative = Boolean(nativeInstallation) && deliveryMode === 'wake' && !this.inProcessDraining; const event = await AgentEvent.create({ agentName: agentName.toLowerCase(), diff --git a/backend/services/hostedShutdownService.ts b/backend/services/hostedShutdownService.ts new file mode 100644 index 0000000000..08bd6727e6 --- /dev/null +++ b/backend/services/hostedShutdownService.ts @@ -0,0 +1,37 @@ +import type { Server } from 'http'; +import AgentEventService from './agentEventService'; +import schedulerService from './schedulerService'; + +// Kubernetes grants 180s, including the existing 10s preStop delay. Native +// admission closes before HTTP stops; 150s covers the ~2 minute turn bound, +// leaving 15s for fenced requeues and 5s before the kubelet's hard kill. +export const HOSTED_SHUTDOWN_WATCHDOG_MS = 165_000; + +export const installHostedShutdown = (server: Server): void => { + let stopping = false; + const shutdown = () => { + if (stopping) return; + stopping = true; + const watchdog = setTimeout(() => { + console.error('[shutdown] hosted drain exceeded shutdown budget'); + process.exit(1); + }, HOSTED_SHUTDOWN_WATCHDOG_MS); + const drained = AgentEventService.drainInProcess(); + schedulerService.stop(); + // Do not await WebSocket/keepalive clients: their lifetime does not bound + // agent execution. The drain waits through its own terminal DB writes. + server.close(); + server.closeIdleConnections(); + void drained.then((result) => { + console.log('[shutdown] hosted drain complete', result); + clearTimeout(watchdog); + process.exit(0); + }, (error: Error) => { + console.error('[shutdown] hosted drain failed:', error.message); + clearTimeout(watchdog); + process.exit(1); + }); + }; + process.once('SIGTERM', shutdown); + process.once('SIGINT', shutdown); +}; diff --git a/docs/runbooks/hosted-turn-shutdown.md b/docs/runbooks/hosted-turn-shutdown.md new file mode 100644 index 0000000000..25f411ba01 --- /dev/null +++ b/docs/runbooks/hosted-turn-shutdown.md @@ -0,0 +1,19 @@ +# Hosted turn shutdown + +TASK-500 adds graceful shutdown to the backend server entry point. It pairs with the background-worker split (#2475), whose worker deployment renders the same backend template and starts the same server entry point. + +On SIGTERM or SIGINT: + +1. Close in-process admission synchronously. Native enqueue persists new wakes as pending; both recovery sweeps stop claiming. Async claims already in flight are tracked and dispatch returns them to pending instead of starting another model turn. +2. Stop the local scheduler and accepting new HTTP connections. Readiness returns 503 while draining. Liveness is unchanged. +3. Await all tracked admission, native turns and terminal event writes. A successful model promise alone is not completion: its event acknowledgement must finish too. +4. After 150 seconds, return this process's unfinished deliveries to pending with a status + delivery-nonce compare-and-swap. Keep attempts unchanged. A sibling's delivery is never selected by a broad pod query; late settlement from the old nonce cannot acknowledge its replacement. +5. Exit when cleanup finishes. A 165-second watchdog exits nonzero if cleanup hangs. The chart allows 180 seconds total, including the existing 10-second preStop delay, leaving 15 seconds after the normal drain deadline for DB cleanup and 5 seconds after the watchdog before SIGKILL. + +This is at-least-once recovery, not exactly-once provider execution. A provider side effect from a timed-out turn may already have happened. Existing approval/deduplication behavior still applies. SIGKILL, process crashes and unavailable DB writes still use the existing stale-delivery recovery sweeps; this change does not claim lossless recovery from those failures. Completed model turns whose acknowledgement failed remain delivered for the existing succeeded-run recovery path, rather than being blindly requeued and answered twice. + +## Verify a rollout + +Use a non-production test pod and a controlled long-turn provider. Start one native turn, then send SIGTERM. Confirm readiness refuses new traffic, a second wake remains pending, the original turn and its event acknowledgement finish, and `[shutdown] hosted drain complete` appears before process exit. Check the replacement process handles the pending wake. A forced deadline case should show the unfinished event pending with its old nonce invalidated; confirm any provider write independently before retrying it. + +Automated tests use real Mongo AgentEvent rows and a fake long `runAgent`, plus signal-handler tests. They cover graceful success/failure, claim/lookup races, waiting for the acknowledgement, timeout requeue, sibling ownership and late settlement. An actual Kubernetes rollout remains an operator deployment check. diff --git a/k8s/helm/commonly/templates/core/backend-deployment.yaml b/k8s/helm/commonly/templates/core/backend-deployment.yaml index 8f4fbdd1db..58c292cfd8 100644 --- a/k8s/helm/commonly/templates/core/backend-deployment.yaml +++ b/k8s/helm/commonly/templates/core/backend-deployment.yaml @@ -25,7 +25,7 @@ spec: {{- include "commonly.backend.selectorLabels" . | nindent 8 }} spec: serviceAccountName: {{ .Values.rbac.serviceAccount.name }} - terminationGracePeriodSeconds: 30 + terminationGracePeriodSeconds: 180 {{- if and .Values.skillsCatalogStorage.enabled .Values.skillsCatalogStorage.bootstrapFromImage }} initContainers: - name: skills-catalog-bootstrap From 367db01e499f3b50c382f85e1a4dee4353d0f4bb Mon Sep 17 00:00:00 2001 From: Lily Shen <115414357+lilyshen0722@users.noreply.github.com> Date: Sun, 25 Oct 2026 01:48:50 -0700 Subject: [PATCH 2/3] fix(ci): align hosted drain grace assertion and recovery runbook Co-authored-by: rivet Agent-Model: codex/gpt-6.1-sol/medium --- .github/workflows/tests.yml | 2 +- docs/runbooks/hosted-turn-shutdown.md | 4 ++++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 9b82cd9f61..a32c39cec4 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -395,7 +395,7 @@ jobs: test "$(yq 'select(.kind == "Deployment" and .metadata.name == "redis") | .spec.template.spec.tolerations | map(select(.key == "workload-tier")) | length' /tmp/dev-rendered.yaml)" = "0" test "$(yq 'select(.kind == "Deployment" and .metadata.name == "redis") | .spec.template.spec.tolerations | map(select(.key == "pool" and .value == "dev" and .effect == "NoSchedule")) | length' /tmp/dev-rendered.yaml)" = "1" test "$(yq 'select(.kind == "Deployment" and .metadata.name == "backend") | .spec.minReadySeconds' /tmp/dev-rendered.yaml)" = "10" - test "$(yq 'select(.kind == "Deployment" and .metadata.name == "backend") | .spec.template.spec.terminationGracePeriodSeconds' /tmp/dev-rendered.yaml)" = "30" + test "$(yq 'select(.kind == "Deployment" and .metadata.name == "backend") | .spec.template.spec.terminationGracePeriodSeconds' /tmp/dev-rendered.yaml)" = "180" test "$(yq 'select(.kind == "Deployment" and .metadata.name == "backend") | .spec.template.spec.containers[] | select(.name == "backend") | .readinessProbe.httpGet.path' /tmp/dev-rendered.yaml)" = "/api/health/ready" test "$(yq 'select(.kind == "Deployment" and .metadata.name == "backend") | .spec.template.spec.containers[] | select(.name == "backend") | .lifecycle.preStop.exec.command[2]' /tmp/dev-rendered.yaml)" = "sleep 10" # Two frontend replicas never share a node, so a spot reclaim leaves one serving; with the spot diff --git a/docs/runbooks/hosted-turn-shutdown.md b/docs/runbooks/hosted-turn-shutdown.md index 25f411ba01..2aa79bc010 100644 --- a/docs/runbooks/hosted-turn-shutdown.md +++ b/docs/runbooks/hosted-turn-shutdown.md @@ -16,4 +16,8 @@ This is at-least-once recovery, not exactly-once provider execution. A provider Use a non-production test pod and a controlled long-turn provider. Start one native turn, then send SIGTERM. Confirm readiness refuses new traffic, a second wake remains pending, the original turn and its event acknowledgement finish, and `[shutdown] hosted drain complete` appears before process exit. Check the replacement process handles the pending wake. A forced deadline case should show the unfinished event pending with its old nonce invalidated; confirm any provider write independently before retrying it. +Also force-kill a test process during a controlled turn. SIGKILL runs no shutdown handler, so confirm the delivery stays delivered until the replacement's stuck-redispatch sweep recovers it. A spot reclaim can provide less notice than the drain budget and needs the same recovery path. Record the recovery delay and verify that an old nonce cannot settle the replacement's claim. + +A deadline requeue reruns the tool-using turn. Any provider write or posted message that happened before interruption can repeat; inspect those effects independently before retrying. + Automated tests use real Mongo AgentEvent rows and a fake long `runAgent`, plus signal-handler tests. They cover graceful success/failure, claim/lookup races, waiting for the acknowledgement, timeout requeue, sibling ownership and late settlement. An actual Kubernetes rollout remains an operator deployment check. From 793e3871bbf3e0adb7124b45ac698d7e7d3ed50d Mon Sep 17 00:00:00 2001 From: Lily Shen <115414357+lilyshen0722@users.noreply.github.com> Date: Sun, 25 Oct 2026 02:01:33 -0700 Subject: [PATCH 3/3] fix(runtime): finish sibling requeues and absorb repeated signals Co-authored-by: rivet Agent-Model: codex/gpt-6.1-sol/medium --- .../service/hosted-turn-drain.test.js | 32 +++++++++++++++++++ .../services/hostedShutdownService.test.js | 4 +-- backend/services/agentEventService.ts | 6 +++- backend/services/hostedShutdownService.ts | 4 +-- docs/runbooks/hosted-turn-shutdown.md | 2 ++ 5 files changed, 43 insertions(+), 5 deletions(-) diff --git a/backend/__tests__/service/hosted-turn-drain.test.js b/backend/__tests__/service/hosted-turn-drain.test.js index 11df4c5126..f13263cab3 100644 --- a/backend/__tests__/service/hosted-turn-drain.test.js +++ b/backend/__tests__/service/hosted-turn-drain.test.js @@ -66,6 +66,38 @@ test('a failed in-flight turn settles failure before drain completes', async () const drain = Service.drainInProcess(2000); turn.resolve({ status: 'failed', error: 'provider refused' }); await drain; expect((await AgentEvent.findById(event._id)).status).toBe('failed'); }); +test('deadline waits for sibling requeues before propagating a cleanup failure', async () => { + const turn = deferred(), cleanup = deferred(), entered = deferred(); + runAgent.mockReturnValue(turn.promise); + await Service.enqueue(input); const second = await Service.enqueue(input); + await waitFor(() => runAgent.mock.calls.length === 2); + const original = AgentEvent.updateOne.bind(AgentEvent); + jest.spyOn(AgentEvent, 'updateOne').mockImplementation((filter, update, ...rest) => { + if (update.$set?.status !== 'pending') return original(filter, update, ...rest); + if (String(filter._id) !== String(second._id)) return Promise.reject(new Error('first requeue failed')); + entered.resolve(); + return cleanup.promise.then(() => original(filter, update, ...rest)); + }); + let settled = false; + const drain = Service.drainInProcess(20).finally(() => { settled = true; }); + const rejection = expect(drain).rejects.toThrow('first requeue failed'); + await entered.promise; await new Promise((resolve) => setImmediate(resolve)); + expect(settled).toBe(false); + cleanup.resolve(); await rejection; + expect((await AgentEvent.findById(second._id)).status).toBe('pending'); + turn.resolve(success); await waitFor(() => Service.inProcessWork.size === 0); +}); +test('SIGTERM during digest lookup requeues before runAgent starts', async () => { + const entered = deferred(), release = deferred(); + const digest = require('../../services/agentDigestFoldService'); + digest.readFoldedDigest.mockImplementationOnce(async () => { entered.resolve(); await release.promise; return null; }); + const event = await Service.enqueue(input); await entered.promise; + expect(Service.inProcessRunning.size).toBe(1); + const drain = Service.drainInProcess(2000); release.resolve(); + await expect(drain).resolves.toEqual({ timedOut: false, requeued: 1 }); + expect((await AgentEvent.findById(event._id)).status).toBe('pending'); + expect(runAgent).not.toHaveBeenCalled(); +}); test('deadline cleanup cannot steal a delivery reclaimed by a sibling', async () => { const turn = deferred(); runAgent.mockReturnValue(turn.promise); const event = await Service.enqueue(input); await waitFor(() => runAgent.mock.calls.length === 1); diff --git a/backend/__tests__/unit/services/hostedShutdownService.test.js b/backend/__tests__/unit/services/hostedShutdownService.test.js index 12f6b642f9..7a4212a11d 100644 --- a/backend/__tests__/unit/services/hostedShutdownService.test.js +++ b/backend/__tests__/unit/services/hostedShutdownService.test.js @@ -7,7 +7,7 @@ const { installHostedShutdown, HOSTED_SHUTDOWN_WATCHDOG_MS } = require('../../.. let handlers, server, exit; beforeEach(() => { jest.useFakeTimers(); jest.clearAllMocks(); handlers = {}; - jest.spyOn(process, 'once').mockImplementation((name, handler) => { handlers[name] = handler; return process; }); + jest.spyOn(process, 'on').mockImplementation((name, handler) => { handlers[name] = handler; return process; }); exit = jest.spyOn(process, 'exit').mockImplementation(() => {}); server = { close: jest.fn(), closeIdleConnections: jest.fn() }; }); @@ -15,7 +15,7 @@ afterEach(() => { jest.restoreAllMocks(); jest.useRealTimers(); }); const flush = async () => { await Promise.resolve(); await Promise.resolve(); }; test('SIGTERM gates admission before closing HTTP; exits after drain; repeated signals are idempotent', async () => { let resolve; service.drainInProcess.mockImplementation(() => new Promise((done) => { resolve = done; })); - installHostedShutdown(server); handlers.SIGTERM(); handlers.SIGINT(); + installHostedShutdown(server); handlers.SIGTERM(); handlers.SIGTERM(); handlers.SIGINT(); expect(service.drainInProcess).toHaveBeenCalledTimes(1); expect(scheduler.stop).toHaveBeenCalledTimes(1); expect(service.drainInProcess.mock.invocationCallOrder[0]).toBeLessThan(server.close.mock.invocationCallOrder[0]); expect(server.close).toHaveBeenCalledTimes(1); expect(exit).not.toHaveBeenCalled(); diff --git a/backend/services/agentEventService.ts b/backend/services/agentEventService.ts index 1fcc87050b..ef5fbaef44 100644 --- a/backend/services/agentEventService.ts +++ b/backend/services/agentEventService.ts @@ -546,7 +546,11 @@ class AgentEventService { // Requeue only this process's unfinished deliveries, never a pod-wide // query that could steal a sibling replica's turn. The nonce fences late // success/failure settlement from a replacement's claim. - await Promise.all([...this.inProcessRunning.values()].map((event) => this.requeueInProcess(event))); + const requeues = await Promise.allSettled( + [...this.inProcessRunning.values()].map((event) => this.requeueInProcess(event)), + ); + const failed = requeues.find((result) => result.status === 'rejected'); + if (failed?.status === 'rejected') throw failed.reason; return { timedOut: true, requeued: this.inProcessRequeued }; })(); return this.inProcessDrain; diff --git a/backend/services/hostedShutdownService.ts b/backend/services/hostedShutdownService.ts index 08bd6727e6..1c83cf43a9 100644 --- a/backend/services/hostedShutdownService.ts +++ b/backend/services/hostedShutdownService.ts @@ -32,6 +32,6 @@ export const installHostedShutdown = (server: Server): void => { process.exit(1); }); }; - process.once('SIGTERM', shutdown); - process.once('SIGINT', shutdown); + process.on('SIGTERM', shutdown); + process.on('SIGINT', shutdown); }; diff --git a/docs/runbooks/hosted-turn-shutdown.md b/docs/runbooks/hosted-turn-shutdown.md index 2aa79bc010..8190922728 100644 --- a/docs/runbooks/hosted-turn-shutdown.md +++ b/docs/runbooks/hosted-turn-shutdown.md @@ -14,6 +14,8 @@ This is at-least-once recovery, not exactly-once provider execution. A provider ## Verify a rollout +Verify that the replacement process responsible for scheduling actually starts the scheduler, especially in worker-split deployments. Startup garbage collection can recover eligible wakes immediately; otherwise recovery waits for the ten-minute scheduled sweep. Its one-minute minimum event age can skip a fresh wake at one sweep, so allow roughly 11 minutes for pickup plus turn execution before a human-visible reply. This is scheduled recovery, not immediate delivery. + Use a non-production test pod and a controlled long-turn provider. Start one native turn, then send SIGTERM. Confirm readiness refuses new traffic, a second wake remains pending, the original turn and its event acknowledgement finish, and `[shutdown] hosted drain complete` appears before process exit. Check the replacement process handles the pending wake. A forced deadline case should show the unfinished event pending with its old nonce invalidated; confirm any provider write independently before retrying it. Also force-kill a test process during a controlled turn. SIGKILL runs no shutdown handler, so confirm the delivery stays delivered until the replacement's stuck-redispatch sweep recovers it. A spot reclaim can provide less notice than the drain budget and needs the same recovery path. Record the recovery delay and verify that an old nonce cannot settle the replacement's claim.