fix(runtime): drain hosted turns on SIGTERM (TASK-500)

.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

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();
+});

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);
 

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);
+});

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' });

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}`));
 }