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/backend/__tests__/service/hosted-turn-drain.test.js b/backend/__tests__/service/hosted-turn-drain.test.js new file mode 100644 index 0000000000..f13263cab3 --- /dev/null +++ b/backend/__tests__/service/hosted-turn-drain.test.js @@ -0,0 +1,138 @@ +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 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); + 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..7a4212a11d --- /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, 'on').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.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..ef5fbaef44 100644 --- a/backend/services/agentEventService.ts +++ b/backend/services/agentEventService.ts @@ -492,6 +492,70 @@ 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. + 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; + } + 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 +1062,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 +1092,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 +1173,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 +1217,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 +1578,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 +1635,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 +1657,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 +1686,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 +1711,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 +1756,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..1c83cf43a9 --- /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.on('SIGTERM', shutdown); + process.on('SIGINT', shutdown); +}; diff --git a/docs/runbooks/hosted-turn-shutdown.md b/docs/runbooks/hosted-turn-shutdown.md new file mode 100644 index 0000000000..8190922728 --- /dev/null +++ b/docs/runbooks/hosted-turn-shutdown.md @@ -0,0 +1,25 @@ +# 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 + +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. + +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. 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