fix: allow workers to drain executions before shutdown

.github/workflows/ci.yml

@@ -37,6 +37,9 @@ jobs:
       - name: KVM guest resolver handoff
         run: tests/kvm_guest_dns.sh
 
+      - name: Worker drains before sandbox shutdown
+        run: tests/worker_sandbox_shutdown.sh
+
       - name: Sandbox-runner liveness checks
         run: tests/sandbox_runner_healthcheck.sh
 

docker-compose.local-dev.yml

@@ -103,6 +103,7 @@ services:
   # Local Service API (combined API + Workers, no auth required)
   # For production, split into api + worker-sandbox pods
   service:
+    stop_grace_period: 400s
     build:
       context: .
       dockerfile: service/Dockerfile.local
@@ -111,6 +112,7 @@ services:
     ports:
       - 3112:3112
     environment:
+      - CODEAPI_SHUTDOWN_TIMEOUT_MS=${CODEAPI_SHUTDOWN_TIMEOUT_MS:-}
       - LOCAL_MODE=${LOCAL_MODE:-true}
       - SANDBOX_ENDPOINT=http://sandbox:2000/api/v2
       - FILE_SERVER_URL=http://file_server:3000

docker-compose.scalable.yml

@@ -166,6 +166,7 @@ services:
   # Each pod runs its own microVM — PYTHON_CONCURRENCY is per-sandbox
   # =============================================================================
   worker-sandbox:
+    stop_grace_period: 400s
     build:
       context: .
       dockerfile: docker/Dockerfile.worker-sandbox
@@ -174,6 +175,7 @@ services:
     devices:
       - ${KVM_DEVICE_PATH:-/dev/kvm}:/dev/kvm
     environment:
+      - CODEAPI_SHUTDOWN_TIMEOUT_MS=${CODEAPI_SHUTDOWN_TIMEOUT_MS:-}
       - KVM_ENABLED=${KVM_ENABLED:-true}
       - LOCAL_MODE=${LOCAL_MODE:-true}
       # Worker config

docker-compose.yaml

@@ -58,12 +58,14 @@ services:
             - tool_call_server
 
     service-worker:
+        stop_grace_period: 400s
         build:
             context: .
             dockerfile: service/Dockerfile
             target: worker
         container_name: service-worker
         environment:
+            - CODEAPI_SHUTDOWN_TIMEOUT_MS=${CODEAPI_SHUTDOWN_TIMEOUT_MS:-}
             - LOCAL_MODE=${LOCAL_MODE:-true}
             - CODEAPI_HARDENED_SANDBOX_MODE=${CODEAPI_HARDENED_SANDBOX_MODE:-true}
             - CODEAPI_BRIDGE_TOKEN=${CODEAPI_BRIDGE_TOKEN:-}

docker/supervisor.sh

@@ -38,11 +38,13 @@ cleanup() {
 
     echo "Received shutdown signal, stopping processes..."
 
-    if [ -n "$SANDBOX_PID" ]; then
-        kill -TERM "$SANDBOX_PID" 2>/dev/null || true
-    fi
     if [ -n "$WORKER_PID" ]; then
         kill -TERM "$WORKER_PID" 2>/dev/null || true
+        # Active jobs still need the sandbox while the worker drains.
+        wait "$WORKER_PID" 2>/dev/null || true
+    fi
+    if [ -n "$SANDBOX_PID" ]; then
+        kill -TERM "$SANDBOX_PID" 2>/dev/null || true
     fi
 
     wait

helm/codeapi/Chart.yaml

@@ -3,7 +3,7 @@ apiVersion: v2
 name: codeapi
 description: A Helm chart for Code Interpreter API - scalable code execution service
 type: application
-version: 0.3.2       # Chart version (bump this when you change the chart)
+version: 0.3.3       # Chart version (bump this when you change the chart)
 appVersion: "2.0.0"  # App version (bump this when you change the app)
 
 # Keywords for searching

helm/codeapi/templates/worker-sandbox-deployment.yaml

@@ -121,6 +121,7 @@ spec:
       labels:
         {{- include "codeapi.serviceWorker.selectorLabels" . | nindent 8 }}
     spec:
+      terminationGracePeriodSeconds: {{ .Values.workerSandbox.serviceWorker.terminationGracePeriodSeconds | default (add 100 (div (add (int .Values.workerSandbox.config.jobTimeout) 999) 1000)) }}
       {{- if or .Values.workerSandbox.serviceAccount.create .Values.workerSandbox.serviceAccount.name }}
       serviceAccountName: {{ $serviceWorkerServiceAccount }}
       {{- end }}

helm/codeapi/values.yaml

@@ -207,6 +207,8 @@ workerSandbox:
   # nodeSelector/tolerations/affinity below. Use sandboxRunner to place only the
   # untrusted-code tier on KVM-capable nodes.
   serviceWorker:
+    # Defaults to jobTimeout (rounded up to seconds) + 100 seconds.
+    terminationGracePeriodSeconds: null
     # Defaults to workerSandbox.replicaCount when unset.
     # Keep this as fixed HA capacity; scale sandboxRunner for execution throughput.
     replicaCount: null

service/src/config.ts

@@ -250,6 +250,19 @@ export function hostedAppOperationTimeoutMs(): number {
     + env.HOSTED_APP_START_TIMEOUT_MS + 30_000;
 }
 
+export function gracefulShutdownTimeoutMs(hasWorkers: boolean): number {
+  const defaultTimeout = hasWorkers
+    ? Math.max(
+        jobCompletionWaitTimeoutMs(env.JOB_TIMEOUT, env.LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS, env.EGRESS_GATEWAY_REVOKE_TIMEOUT_MS),
+        env.HOSTED_APPS_ENABLED
+          ? hostedAppOperationTimeoutMs() + env.LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS + 30_000
+          : 0,
+      )
+    : 30_000;
+  // Node/Bun timers overflow to an immediate timeout above this limit.
+  return Math.min(resolvePositiveIntEnv(process.env.CODEAPI_SHUTDOWN_TIMEOUT_MS, defaultTimeout), 2_147_483_647);
+}
+
 function configuredNumber(raw: string | undefined, fallback: number): number {
   return raw == null || raw.trim() === '' ? fallback : Number(raw);
 }

service/src/lifecycle-timeout.test.ts

@@ -0,0 +1,60 @@
+import { afterEach, beforeEach, expect, test } from 'bun:test';
+import { env, gracefulShutdownTimeoutMs } from './config';
+
+const original = {
+  JOB_TIMEOUT: env.JOB_TIMEOUT,
+  LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS: env.LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS,
+  EGRESS_GATEWAY_REVOKE_TIMEOUT_MS: env.EGRESS_GATEWAY_REVOKE_TIMEOUT_MS,
+  HOSTED_APPS_ENABLED: env.HOSTED_APPS_ENABLED,
+  CHECKPOINT_TIMEOUT_MS: env.CHECKPOINT_TIMEOUT_MS,
+};
+const originalOverride = process.env.CODEAPI_SHUTDOWN_TIMEOUT_MS;
+
+beforeEach(() => {
+  Object.assign(env, {
+    JOB_TIMEOUT: 300_000,
+    LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS: 60_000,
+    EGRESS_GATEWAY_REVOKE_TIMEOUT_MS: 5_000,
+    HOSTED_APPS_ENABLED: false,
+    CHECKPOINT_TIMEOUT_MS: 60_000,
+  });
+  delete process.env.CODEAPI_SHUTDOWN_TIMEOUT_MS;
+});
+
+afterEach(() => {
+  Object.assign(env, original);
+  if (originalOverride === undefined) delete process.env.CODEAPI_SHUTDOWN_TIMEOUT_MS;
+  else process.env.CODEAPI_SHUTDOWN_TIMEOUT_MS = originalOverride;
+});
+
+test('workers can finish a five-minute job and cleanup while API-only shutdown stays bounded', () => {
+  expect(gracefulShutdownTimeoutMs(true)).toBe(370_000);
+  expect(gracefulShutdownTimeoutMs(false)).toBe(30_000);
+});
+
+test('worker shutdown follows configured execution and cleanup deadlines', () => {
+  env.JOB_TIMEOUT = 1_200_000;
+  env.LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS = 90_000;
+  env.EGRESS_GATEWAY_REVOKE_TIMEOUT_MS = 20_000;
+  expect(gracefulShutdownTimeoutMs(true)).toBe(1_315_000);
+});
+
+test('hosted-app workers retain their longer operation budget', () => {
+  env.HOSTED_APPS_ENABLED = true;
+  expect(gracefulShutdownTimeoutMs(true)).toBe(1_110_000);
+  env.JOB_TIMEOUT = 1_200_000;
+  expect(gracefulShutdownTimeoutMs(true)).toBe(1_270_000);
+  expect(gracefulShutdownTimeoutMs(false)).toBe(30_000);
+});
+
+test('explicit shutdown overrides remain finite and invalid values use the default', () => {
+  process.env.CODEAPI_SHUTDOWN_TIMEOUT_MS = '420000';
+  expect(gracefulShutdownTimeoutMs(true)).toBe(420_000);
+  expect(gracefulShutdownTimeoutMs(false)).toBe(420_000);
+  for (const invalid of ['', ' ', '0', '-1', 'NaN', 'Infinity']) {
+    process.env.CODEAPI_SHUTDOWN_TIMEOUT_MS = invalid;
+    expect(gracefulShutdownTimeoutMs(true)).toBe(370_000);
+  }
+  process.env.CODEAPI_SHUTDOWN_TIMEOUT_MS = '2147483648';
+  expect(gracefulShutdownTimeoutMs(true)).toBe(2_147_483_647);
+});

service/src/lifecycle.ts

@@ -7,7 +7,7 @@ import {
   closeQueueConnections,
 } from './queue';
 import { validateStartupAuthConfig } from './auth/startup';
-import { env, hostedAppOperationTimeoutMs } from './config';
+import { env, gracefulShutdownTimeoutMs } from './config';
 import {
   validateApiBridgePolicy,
   validateApiHardenedConfig,
@@ -241,9 +241,7 @@ export async function gracefulShutdown(): Promise<void> {
   const shutdownTimeout = setTimeout(() => {
     logger.error('Shutdown timeout reached, forcing exit');
     process.exit(1);
-  }, hasWorkers && env.HOSTED_APPS_ENABLED
-    ? hostedAppOperationTimeoutMs() + env.LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS + 30_000
-    : 30_000);
+  }, gracefulShutdownTimeoutMs(hasWorkers));
 
   try {
     if (hasWorkers) {

tests/worker_sandbox_shutdown.sh

@@ -0,0 +1,50 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+fixture_dir="$(mktemp -d)"
+trap 'rm -rf "$fixture_dir"' EXIT
+
+# Exercise the production cleanup function without booting a MicroVM.
+sed -n '/^cleanup() {/,/^trap cleanup SIGTERM SIGINT/p' "$repo_root/docker/supervisor.sh" \
+    | sed '$d' > "$fixture_dir/cleanup.sh"
+cat > "$fixture_dir/fixture.sh" <<'FIXTURE'
+#!/usr/bin/env bash
+set -u
+CLEANUP_IN_PROGRESS=false
+source "$1/cleanup.sh"
+fixture_dir="$1"
+
+sleep 30 &
+SANDBOX_PID=$!
+(
+    trap '
+        sleep 0.1
+        if kill -0 "$SANDBOX_PID" 2>/dev/null \
+            && [ "$(awk '\''{print $3}'\'' "/proc/$SANDBOX_PID/stat")" != Z ]; then
+            touch "$fixture_dir/drained"
+        fi
+        exit 0
+    ' TERM
+    touch "$fixture_dir/ready"
+    while true; do sleep 0.01; done
+) &
+WORKER_PID=$!
+while [ ! -f "$fixture_dir/ready" ]; do sleep 0.01; done
+cleanup
+FIXTURE
+
+timeout 5 bash "$fixture_dir/fixture.sh" "$fixture_dir"
+test -f "$fixture_dir/drained"
+echo 'Worker drained before sandbox shutdown.'
+
+for entry in docker-compose.yaml:service-worker docker-compose.scalable.yml:worker-sandbox docker-compose.local-dev.yml:service; do
+    docker compose -f "$repo_root/${entry%:*}" config --format json \
+        | python3 -c '
+import json, sys
+worker = json.load(sys.stdin)["services"][sys.argv[1]]
+assert worker["stop_grace_period"] == "6m40s", worker["stop_grace_period"]
+assert "CODEAPI_SHUTDOWN_TIMEOUT_MS" in worker["environment"]
+' "${entry#*:}"
+done
+echo 'Compose worker stop grace permits the default execution and cleanup budget.'