fix(mothership): back off the transport poll instead of retrying every second

Author: waleedlatif1

apps/sim/lib/mothership/transport/receiver.test.ts

@@ -29,7 +29,7 @@ import { receiveSimControls, startSimReceivers } from '@/lib/mothership/transpor
 const mocks = {
   fetch: mothershipGoFetchMockFns.mockFetchGo,
   execute,
-  sleep: utilsHelpersMockFns.mockSleep,
+  sleep: utilsHelpersMockFns.mockInterruptibleSleep,
 }
 
 function request() {
@@ -232,3 +232,69 @@ describe('outbound receiver lifecycle', () => {
     }
   )
 })
+
+describe('outbound receiver backoff', () => {
+  const logger = () => getMockLogger('MothershipTransportReceiver')
+
+  beforeEach(() => {
+    vi.resetAllMocks()
+    mocks.sleep.mockResolvedValue(undefined)
+  })
+
+  /** Runs the receiver over scripted poll outcomes and returns each wait it took. */
+  async function delaysFor(outcomes: (() => Response)[]): Promise<number[]> {
+    const controller = new AbortController()
+    let polls = 0
+    mocks.fetch.mockImplementation(async () => {
+      const next = outcomes[polls++]
+      if (next) return next()
+      controller.abort()
+      throw new Error('shutdown')
+    })
+    await receiveSimControls('https://worker.test', 'a'.repeat(64), controller.signal)
+    for (const [, signal] of mocks.sleep.mock.calls) expect(signal).toBe(controller.signal)
+    return mocks.sleep.mock.calls.map(([ms]) => ms)
+  }
+
+  const refused = (status: number) => () => new Response(null, { status })
+  const lost = () => {
+    throw new Error('connection lost')
+  }
+  const empty = () => Response.json({ requests: [] })
+
+  it.each([401, 403])(
+    'waits minutes, not seconds, between polls the worker refuses (%i)',
+    async (status) => {
+      const delays = await delaysFor(Array.from({ length: 8 }, () => refused(status)))
+      expect(delays).toHaveLength(8)
+      expect(Math.min(...delays)).toBeGreaterThanOrEqual(24_000)
+      expect(delays[7]).toBeGreaterThanOrEqual(480_000)
+      expect(Math.max(...delays)).toBeLessThanOrEqual(720_000)
+      expect(logger().warn).toHaveBeenCalledOnce()
+    }
+  )
+
+  it('backs off transient failures to a ceiling under the 5 s run-control deadline, warning once', async () => {
+    const delays = await delaysFor(Array.from({ length: 8 }, () => lost))
+    expect(delays[0]).toBeLessThanOrEqual(1200)
+    expect(delays[7]).toBeGreaterThanOrEqual(3200)
+    expect(Math.max(...delays)).toBeLessThan(5000)
+    expect(logger().warn).toHaveBeenCalledOnce()
+  })
+
+  it('starts over after a poll succeeds and says it reconnected', async () => {
+    const delays = await delaysFor([lost, lost, lost, empty, lost])
+    expect(delays).toHaveLength(4)
+    expect(delays[2]).toBeGreaterThanOrEqual(3200)
+    expect(delays[3]).toBeLessThanOrEqual(1200)
+    expect(logger().warn).toHaveBeenCalledTimes(2)
+    expect(logger().info).toHaveBeenCalledOnce()
+  })
+
+  it('warns again, and backs off from the start, when an outage turns into a refused key', async () => {
+    const delays = await delaysFor([lost, lost, lost, refused(401), refused(401)])
+    expect(delays[3]).toBeLessThanOrEqual(36_000)
+    expect(delays[4]).toBeGreaterThanOrEqual(48_000)
+    expect(logger().warn).toHaveBeenCalledTimes(2)
+  })
+})

apps/sim/lib/mothership/transport/receiver.ts

@@ -1,6 +1,7 @@
 import { createLogger } from '@sim/logger'
 import { getErrorMessage } from '@sim/utils/errors'
-import { sleep } from '@sim/utils/helpers'
+import { interruptibleSleep } from '@sim/utils/helpers'
+import { type BackoffOptions, backoffWithJitter } from '@sim/utils/retry'
 import { truncate } from '@sim/utils/string'
 import { env } from '@/lib/core/config/env'
 import { isMothershipBenchmarkEnabled } from '@/lib/core/config/env-flags'
@@ -22,6 +23,25 @@ import { executeSimControl } from '@/lib/mothership/transport/control'
 const logger = createLogger('MothershipTransportReceiver')
 const receivers = new Map<string, AbortController>()
 
+/**
+ * A lost connection or a worker deploy: retry within a second, slowing to 4 s, so with jitter a
+ * healthy deployment is back inside the worker's 5 s run-control deadline.
+ */
+const TRANSIENT_BACKOFF: BackoffOptions = { baseMs: 1000, maxMs: 4000 }
+/**
+ * The worker rejected this deployment's key, which only an operator can fix: start at 30 s and
+ * slow to 10 minutes, so a misconfigured deployment polls a few hundred times a day, not every
+ * second, and still reconnects on its own once the key is accepted.
+ */
+const REFUSED_BACKOFF: BackoffOptions = { baseMs: 30_000, maxMs: 600_000 }
+
+/** The worker refused this deployment's key (401/403), as opposed to failing to answer. */
+class TransportKeyRefused extends Error {
+  constructor(readonly status: number) {
+    super(`Transport key refused (HTTP ${status})`)
+  }
+}
+
 interface ReceivedControl {
   id: string
   answer: () => Promise<SimControlResult>
@@ -61,6 +81,8 @@ export async function receiveSimControls(
   channelId: string,
   signal: AbortSignal
 ): Promise<void> {
+  let failures = 0
+  let outage: 'disconnected' | 'key_refused' | undefined
   while (!signal.aborted) {
     let phase: 'poll' | 'parse' | 'reply' = 'poll'
     try {
@@ -75,8 +97,15 @@ export async function receiveSimControls(
       })
       if (!response.ok) {
         await response.body?.cancel()
+        if (response.status === 401 || response.status === 403) {
+          throw new TransportKeyRefused(response.status)
+        }
         throw new Error(`Transport poll refused (HTTP ${response.status})`)
       }
+      if (outage) {
+        logger.info('Sim transport reconnected')
+        outage = undefined
+      }
       phase = 'parse'
       const batch = SimChannelBatchEnvelope.parse(await response.json())
       const controls = batch.requests.flatMap((raw) => receiveControl(raw) ?? [])
@@ -103,11 +132,28 @@ export async function receiveSimControls(
       if (failed?.status === 'rejected') throw failed.reason
     } catch (error) {
       if (signal.aborted) return
-      logger.warn('Sim transport disconnected; reconnecting', {
-        phase,
-        error: getErrorMessage(error),
-      })
-      await sleep(1000)
+      const keyRefused = error instanceof TransportKeyRefused
+      const kind = keyRefused ? 'key_refused' : 'disconnected'
+      // One line per outage, and again only when it changes kind, not one per retry.
+      if (outage !== kind) {
+        outage = kind
+        failures = 0
+        if (keyRefused) {
+          logger.warn('Sim transport key refused by the worker; check COPILOT_API_KEY', {
+            status: error.status,
+          })
+        } else {
+          logger.warn('Sim transport disconnected; reconnecting', {
+            phase,
+            error: getErrorMessage(error),
+          })
+        }
+      }
+      failures++
+      await interruptibleSleep(
+        backoffWithJitter(failures, null, keyRefused ? REFUSED_BACKOFF : TRANSIENT_BACKOFF),
+        signal
+      )
     }
   }
 }