fix(isolated-vm): drain worker host deploys without failing in-flight executions (#9389)

* fix(isolated-vm): drain worker host deploys without failing in-flight executions

- Retire each remote worker connection at a max age (IVM_REMOTE_WORKER_MAX_AGE_MS,
  default 2 min) so a host the load balancer deregistered stops getting new
  executions and drains before its deregistration delay cuts the connections
- Re-dispatch an execution whose connection closed before it was sent; one that
  may have reached the host is never re-run and fails with a clear error
- Cap raw bytes a host connection may send before authenticating (16 KiB)
- Return a cancelled call before reading the worker-target flag

* fix(isolated-vm): address review on host drain

- Keep counting a refused connection's bytes until its socket closes
- Skip a redispatch the caller cancelled in between
- Look up the owner's state after the admission waits
- Parse IVM_REMOTE_WORKER_MAX_AGE_MS strictly, falling back below 1s
- Hold the inspected execution open until the worker check finishes

* test(isolated-vm): report an execution that fails before it is held

Author: waleedlatif1

apps/sim/lib/core/config/env.ts

@@ -459,6 +459,7 @@ export const env = createEnv({
     IVM_MAX_BROKERS_PER_EXECUTION:         z.string().optional().default('1000'),    // Max broker calls per sandbox task execution
     IVM_REMOTE_WORKER_URL:                 z.string().optional(),                    // ws(s)://<host>/v1/worker of a dedicated worker host; workers run there unless the isolated-vm-local-workers kill switch is on
     IVM_REMOTE_WORKER_TOKEN:               z.string().optional(),                    // Token the worker host expects (its IVM_HOST_TOKEN, >= 32 chars); required with IVM_REMOTE_WORKER_URL
+    IVM_REMOTE_WORKER_MAX_AGE_MS:          z.coerce.number().int().positive().optional(), // (default 120000) Age at which a worker host connection stops taking executions, so host deploys drain; keep below the host's deregistration delay
 
     // Knowledge Base Processing Configuration - Shared across all processing methods
     KB_CONFIG_MAX_DURATION:                z.number().optional().default(600),     // Max processing duration in seconds (10 minutes)

apps/sim/lib/execution/isolated-vm-host.cjs

@@ -14,7 +14,10 @@
  * are the objects the worker
  * exchanges over IPC. Host-originated frames are `{ type: 'host:heartbeat' }`
  * every 30s and `{ type: 'host:draining' }` on SIGTERM, after which the app
- * retires the connection once its executions finish. A breaking change to the
+ * retires the connection once its executions finish. A load balancer that
+ * deregisters this host sends nothing until SIGTERM, so the app also retires
+ * each connection at a maximum age, and a deregistered host drains as its
+ * connections age out. A breaking change to the
  * contract gets a new path served alongside this one until every client has
  * rolled.
  *
@@ -36,6 +39,12 @@ const { WebSocketServer } = require('ws')
 const WORKER_PATH = '/v1/worker'
 const HEARTBEAT_INTERVAL_MS = 30000
 const AUTH_TIMEOUT_MS = 5000
+/**
+ * Raw bytes a connection may send before it authenticates: the auth frame is
+ * well under 1 KiB. Without this, `maxPayload` lets every unauthenticated
+ * connection make the host buffer a full-size frame.
+ */
+const MAX_PRE_AUTH_BYTES = 16 * 1024
 const WORKER_START_TIMEOUT_MS = 10000
 const MAX_STARTUP_STDERR_CHARS = 64 * 1024
 const MIN_TOKEN_LENGTH = 32
@@ -179,13 +188,28 @@ server.on('upgrade', (req, socket, head) => {
     rejectUpgrade(socket, 503, 'Service Unavailable')
     return
   }
-  wss.handleUpgrade(req, socket, head, (ws) => authenticate(ws))
+  wss.handleUpgrade(req, socket, head, (ws) => authenticate(ws, socket))
 })
 
 /** Admits the connection to a worker slot only once its first frame proves it holds the token. */
-function authenticate(ws) {
+function authenticate(ws, socket) {
   pendingConnections++
   let settled = false
+  let authenticated = false
+  let preAuthBytes = 0
+  // Counts what arrives on the TCP socket, so an oversized frame is cut off as it streams in
+  // rather than after ws has buffered all of it.
+  const onPreAuthData = (chunk) => {
+    // ws parses each chunk before this listener runs, so the chunk that authenticated still arrives.
+    // A refused connection stays counted: its peer may ignore the close and keep sending.
+    if (authenticated) return
+    preAuthBytes += chunk.length
+    if (preAuthBytes <= MAX_PRE_AUTH_BYTES) return
+    if (settle()) {
+      log('warn', 'Rejected worker connection', { reason: 'too much data before authentication' })
+    }
+    socket.destroy()
+  }
   const settle = () => {
     if (settled) return false
     settled = true
@@ -199,6 +223,7 @@ function authenticate(ws) {
     ws.close(UNAUTHORIZED_CLOSE_CODE, 'unauthorized')
   }
   const authTimeout = setTimeout(() => refuse('authentication timed out'), AUTH_TIMEOUT_MS)
+  socket.on('data', onPreAuthData)
   const onPreAuthError = (error) => {
     log('warn', 'Worker connection error before authentication', { error: error.message })
   }
@@ -214,6 +239,8 @@ function authenticate(ws) {
       return
     }
     if (!settle()) return
+    authenticated = true
+    socket.off('data', onPreAuthData)
     if (draining) {
       ws.close(1012, 'draining')
       return

apps/sim/lib/execution/isolated-vm-remote.ts

@@ -9,6 +9,17 @@ const HOST_SILENCE_LIMIT_MS = 75_000
 const REMOTE_WORKER_URL_PATTERN = /^wss?:\/\/[^/]+\/v1\/worker$/
 const MIN_REMOTE_WORKER_TOKEN_LENGTH = 32
 
+/**
+ * Thrown by a remote worker's `send` once its connection is closing: nothing was
+ * delivered, so the message never reached the host.
+ */
+export class WorkerConnectionClosedError extends Error {
+  constructor() {
+    super('Remote isolated-vm worker connection is closed')
+    this.name = 'WorkerConnectionClosedError'
+  }
+}
+
 /**
  * The slice of `ChildProcess` the isolated-vm pool uses, so a local child and a
  * connection to a remote worker host are interchangeable.
@@ -113,9 +124,12 @@ export function connectRemoteWorker(options: RemoteWorkerOptions): IsolatedVMWor
   })
 
   return {
-    /** Returns false once the connection is closing; the pending `exit` settles this worker's executions. */
+    /**
+     * Throws {@link WorkerConnectionClosedError} once the connection is closing, so the pool can
+     * tell a message that never left this process from one the host may have acted on.
+     */
     send(message) {
-      if (socket.readyState !== WebSocket.OPEN) return false
+      if (socket.readyState !== WebSocket.OPEN) throw new WorkerConnectionClosedError()
       socket.send(JSON.stringify(message))
       return true
     },

apps/sim/lib/execution/isolated-vm.test.ts

@@ -275,6 +275,47 @@ class UnreachableWebSocket extends EventTarget {
   close() {}
 }
 
+/**
+ * A worker host that accepts every connection and answers each execution. A test closes one by
+ * setting `readyState` to CLOSING, the state a connection is in before its close event arrives.
+ */
+class ReachableWebSocket extends EventTarget {
+  static readonly OPEN = 1
+  static readonly instances: ReachableWebSocket[] = []
+  readyState = 0
+  readonly executionIds: number[] = []
+
+  constructor() {
+    super()
+    ReachableWebSocket.instances.push(this)
+    setImmediate(() => {
+      this.readyState = 1
+      this.dispatchEvent(new Event('open'))
+    })
+  }
+
+  send(data: string) {
+    const message = JSON.parse(data) as { type: string; executionId?: number }
+    if (message.type === 'auth') this.reply({ type: 'ready' })
+    if (message.type === 'execute' && message.executionId !== undefined) {
+      this.executionIds.push(message.executionId)
+      this.reply({
+        type: 'result',
+        executionId: message.executionId,
+        result: { result: 'ran remotely', stdout: '' },
+      })
+    }
+  }
+
+  close() {}
+
+  private reply(message: unknown) {
+    setImmediate(() =>
+      this.dispatchEvent(new MessageEvent('message', { data: JSON.stringify(message) }))
+    )
+  }
+}
+
 describe('isolated-vm scheduler', () => {
   it('recovers from an initial spawn failure and drains queued work', async () => {
     const { executeInIsolatedVM, spawnMock } = await loadExecutionModule({
@@ -416,6 +457,69 @@ describe('isolated-vm scheduler', () => {
     }
   })
 
+  it('returns a cancelled call without waiting on the worker target flag', async () => {
+    featureFlagsMockFns.mockIsFeatureEnabled.mockImplementation(() => new Promise(() => {}))
+    try {
+      const { executeInIsolatedVM } = await loadExecutionModule({
+        envOverrides: {
+          IVM_REMOTE_WORKER_URL: 'ws://worker-host.internal:3100/v1/worker',
+          IVM_REMOTE_WORKER_TOKEN: 'remote-worker-token-of-at-least-32-chars',
+        },
+        spawns: [],
+      })
+      const controller = new AbortController()
+      controller.abort()
+
+      const result = await executeInIsolatedVM(
+        {
+          code: 'return 1',
+          params: {},
+          envVars: {},
+          contextVariables: {},
+          timeoutMs: 100,
+          requestId: 'req-pre-aborted',
+        },
+        { signal: controller.signal }
+      )
+
+      expect(result.termination).toBe('cancelled')
+    } finally {
+      featureFlagsMockFns.mockIsFeatureEnabled.mockReset().mockResolvedValue(false)
+    }
+  })
+
+  it('moves an execution to a new worker when its connection closed before it was sent', async () => {
+    ReachableWebSocket.instances.length = 0
+    vi.stubGlobal('WebSocket', ReachableWebSocket)
+    const { executeInIsolatedVM } = await loadExecutionModule({
+      envOverrides: {
+        IVM_REMOTE_WORKER_URL: 'ws://worker-host.internal:3100/v1/worker',
+        IVM_REMOTE_WORKER_TOKEN: 'remote-worker-token-of-at-least-32-chars',
+      },
+      spawns: [],
+    })
+    const request = {
+      code: 'return 1',
+      params: {},
+      envVars: {},
+      contextVariables: {},
+      timeoutMs: 100,
+    }
+
+    const warm = await executeInIsolatedVM({ ...request, requestId: 'req-warm' })
+    expect(warm.result).toBe('ran remotely')
+    const [closing] = ReachableWebSocket.instances
+    if (!closing) throw new Error('Expected the first worker connection')
+    closing.readyState = 2
+
+    const result = await executeInIsolatedVM({ ...request, requestId: 'req-after-close' })
+
+    expect(result.error).toBeUndefined()
+    expect(result.result).toBe('ran remotely')
+    expect(closing.executionIds).toHaveLength(1)
+    expect(ReachableWebSocket.instances).toHaveLength(2)
+  })
+
   it('rejects new requests when the queue is full', async () => {
     const holder = createControllableReadyProc()
     const { executeInIsolatedVM } = await loadExecutionModule({