fix(mothership): parse each checkpoint control request on its own and declare accepted kinds (#9086)

The outbound receiver parsed a poll batch strictly, so one control kind
this build did not know threw for the whole batch: nothing in it was
answered and the worker waited out each request's timeout. Parse each
request alone, answer an unknown kind 501 sim_control_unsupported and a
malformed one 400, and keep running its batch-mates. Every poll now
declares the kinds this build's generated contract parses, so the worker
can stop sending kinds a receiver cannot handle.

Author: waleedlatif1

apps/sim/lib/mothership/generated/sim-transport.ts

@@ -7,9 +7,11 @@ import { MemoryScopeRequest } from "./memory-scope";
 import { RunControlRequest } from "./run-control";
 import { TaskWakeRequest, WorkflowWatchRequest } from "./tasks";
 
+const SimChannelId = z.string().regex(/^[a-f0-9]{64}$/);
+
 export const SimConnection = z.discriminatedUnion("mode", [
   z.object({ mode: z.literal("direct") }),
-  z.object({ mode: z.literal("checkpoint"), channelId: z.string().regex(/^[a-f0-9]{64}$/) }),
+  z.object({ mode: z.literal("checkpoint"), channelId: SimChannelId }),
 ]);
 export type SimConnection = z.infer<typeof SimConnection>;
 
@@ -51,9 +53,43 @@ export const SimControlResult = z.object({
 });
 export type SimControlResult = z.infer<typeof SimControlResult>;
 
-export const SimChannelPoll = z.object({ channelId: z.string().regex(/^[a-f0-9]{64}$/) });
+export type SimControlKind = SimControlOperation["kind"];
+/** Every kind this build's union parses: what a receiver declares in `accepts`. */
+export const SIM_CONTROL_KINDS: SimControlKind[] = SimControlOperation.options.map(
+  (option) => option.shape.kind.value,
+);
+
+/** The v1.0.0 kinds every receiver parses without declaring them. Frozen: a poll without
+ * `accepts` comes from a receiver that knows only these, so every later kind is negotiated. */
+export const SIM_CONTROL_BASELINE_KINDS = [
+  "run_control",
+  "workflow_status",
+  "wake",
+  "workspace_context",
+  "integration_catalog",
+] as const satisfies readonly SimControlKind[];
+export type BaselineSimControlKind = (typeof SIM_CONTROL_BASELINE_KINDS)[number];
+export type NegotiatedSimControlKind = Exclude<SimControlKind, BaselineSimControlKind>;
+
+/** Error code of the 501 answer for a kind the receiving Sim did not declare. */
+export const SIM_CONTROL_UNSUPPORTED = "sim_control_unsupported";
+
+/** `accepts` is free-form so a newer Sim's kinds never fail an older worker's poll. */
+export const SimChannelPoll = z.object({
+  channelId: SimChannelId,
+  accepts: z.array(z.string().min(1).max(64)).max(64).optional(),
+});
+/** The worker's typed batch; receivers parse `SimChannelBatchEnvelope` instead. */
 export const SimChannelBatch = z.object({ requests: z.array(SimControlRequest).max(16) });
-export const SimChannelReply = SimChannelPoll.extend({
+/** Receivers parse each request alone, so one unknown kind never poisons its batch-mates. */
+export const SimChannelBatchEnvelope = z.object({ requests: z.array(z.unknown()).max(16) });
+/** Enough of a request to answer it when its operation does not parse. */
+export const SimControlRequestHeader = z.object({
+  id: z.uuid(),
+  operation: z.object({ kind: z.string() }).optional().catch(undefined),
+});
+export const SimChannelReply = z.object({
+  channelId: SimChannelId,
   id: z.uuid(),
   result: SimControlResult,
 });

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

@@ -1,6 +1,7 @@
 import { createDeferred } from '@sim/testing/helpers/deferred'
 import { setEnv } from '@sim/testing/mocks/env.mock'
 import { setEnvFlags } from '@sim/testing/mocks/env-flags.mock'
+import { getMockLogger } from '@sim/testing/mocks/logger.mock'
 import {
   mothershipAgentUrlMock,
   mothershipAgentUrlMockFns,
@@ -22,6 +23,7 @@ vi.mock('@/lib/mothership/transport/control', () => ({ executeSimControl: execut
 vi.mock('@/lib/mothership/server/agent-url', () => mothershipAgentUrlMock)
 vi.mock('@sim/utils/helpers', () => utilsHelpersMock)
 
+import { SIM_CONTROL_KINDS } from '@/lib/mothership/generated/sim-transport'
 import { receiveSimControls, startSimReceivers } from '@/lib/mothership/transport/receiver'
 
 const mocks = {
@@ -87,18 +89,92 @@ describe('outbound receiver lifecycle', () => {
     expect(mocks.sleep).toHaveBeenCalledOnce()
   })
 
-  it('does not execute malformed or redirected responses as Sim operations', async () => {
+  it('does not execute redirected responses or requests it cannot read as Sim operations', async () => {
     const controller = new AbortController()
+    const replies: string[] = []
     mocks.fetch
       .mockResolvedValueOnce(new Response(null, { status: 302 }))
       .mockResolvedValueOnce(Response.json({ requests: [{ operation: { kind: 'unknown' } }] }))
-      .mockImplementationOnce(async () => {
+      .mockImplementation(async (url: string, init: RequestInit) => {
+        if (url.endsWith('/reply')) {
+          replies.push(String(init.body))
+          return Response.json({ accepted: true })
+        }
         controller.abort()
         throw new Error('shutdown')
       })
     await receiveSimControls('https://worker.test', 'a'.repeat(64), controller.signal)
     expect(mocks.execute).not.toHaveBeenCalled()
-    expect(mocks.sleep).toHaveBeenCalledTimes(2)
+    expect(replies).toEqual([])
+    expect(mocks.sleep).toHaveBeenCalledOnce()
+    expect(getMockLogger('MothershipTransportReceiver').warn).toHaveBeenCalledWith(
+      'Unreadable Sim control request dropped'
+    )
+  })
+
+  it('answers an unknown kind 501 and a malformed input 400 while their batch-mates still run', async () => {
+    const controller = new AbortController()
+    const valid = request()
+    const unknown = { ...request(), operation: { kind: 'a_future_kind', input: {} } }
+    const malformed = { ...request(), operation: { kind: 'run_control', input: { chatId: 7 } } }
+    const kindless = { ...request(), operation: { input: {} } }
+    const replies: { id: string; result: { status: number; body: string } }[] = []
+    let polls = 0
+    mocks.fetch.mockImplementation(async (url: string, init: RequestInit) => {
+      if (url.endsWith('/reply')) {
+        replies.push(JSON.parse(String(init.body)))
+        return Response.json({ accepted: true })
+      }
+      polls++
+      if (polls === 1) return Response.json({ requests: [unknown, valid, malformed, kindless] })
+      controller.abort()
+      throw new Error('shutdown')
+    })
+    await receiveSimControls('https://worker.test', 'a'.repeat(64), controller.signal)
+    expect(mocks.execute).toHaveBeenCalledOnce()
+    expect(mocks.execute).toHaveBeenCalledWith(valid)
+    expect(replies).toHaveLength(4)
+    expect(replies).toEqual(
+      expect.arrayContaining([
+        expect.objectContaining({
+          id: valid.id,
+          result: { status: 200, body: '{"stopped":false}' },
+        }),
+        expect.objectContaining({
+          id: unknown.id,
+          result: {
+            status: 501,
+            body: JSON.stringify({ error: 'sim_control_unsupported', kind: 'a_future_kind' }),
+          },
+        }),
+        expect.objectContaining({
+          id: malformed.id,
+          result: { status: 400, body: '{"error":"Invalid Sim control request"}' },
+        }),
+        expect.objectContaining({
+          id: kindless.id,
+          result: { status: 400, body: '{"error":"Invalid Sim control request"}' },
+        }),
+      ])
+    )
+    expect(mocks.sleep).not.toHaveBeenCalled()
+  })
+
+  it('declares on every poll exactly the kinds its generated contract parses', async () => {
+    const controller = new AbortController()
+    const polls: unknown[] = []
+    mocks.fetch.mockImplementation(async (_url: string, init: RequestInit) => {
+      polls.push(JSON.parse(String(init.body)))
+      if (polls.length === 1) throw new Error('connection lost')
+      controller.abort()
+      throw new Error('shutdown')
+    })
+    await receiveSimControls('https://worker.test', 'a'.repeat(64), controller.signal)
+    expect(SIM_CONTROL_KINDS).toContain('memory_scope')
+    expect(polls).toEqual([
+      { channelId: 'a'.repeat(64), accepts: SIM_CONTROL_KINDS },
+      { channelId: 'a'.repeat(64), accepts: SIM_CONTROL_KINDS },
+    ])
   })
 
   it('does not repeat an operation when its reply acknowledgement is lost', async () => {

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

@@ -1,9 +1,18 @@
 import { createLogger } from '@sim/logger'
 import { getErrorMessage } from '@sim/utils/errors'
 import { sleep } from '@sim/utils/helpers'
+import { truncate } from '@sim/utils/string'
 import { env } from '@/lib/core/config/env'
 import { isMothershipBenchmarkEnabled } from '@/lib/core/config/env-flags'
-import { SimChannelBatch, type SimConnection } from '@/lib/mothership/generated/sim-transport'
+import {
+  SIM_CONTROL_KINDS,
+  SIM_CONTROL_UNSUPPORTED,
+  SimChannelBatchEnvelope,
+  type SimConnection,
+  SimControlRequest,
+  SimControlRequestHeader,
+  type SimControlResult,
+} from '@/lib/mothership/generated/sim-transport'
 import { fetchGo } from '@/lib/mothership/request/go/fetch'
 import { mothershipRequestHeaders } from '@/lib/mothership/request/headers'
 import { getMothershipBaseURL } from '@/lib/mothership/server/agent-url'
@@ -13,17 +22,52 @@ import { executeSimControl } from '@/lib/mothership/transport/control'
 const logger = createLogger('MothershipTransportReceiver')
 const receivers = new Map<string, AbortController>()
 
+interface ReceivedControl {
+  id: string
+  answer: () => Promise<SimControlResult>
+}
+
+/** Each request parses alone, so one unknown or malformed request never costs its batch-mates. */
+function receiveControl(raw: unknown): ReceivedControl | null {
+  const request = SimControlRequest.safeParse(raw)
+  if (request.success) return { id: request.data.id, answer: () => executeSimControl(request.data) }
+  const header = SimControlRequestHeader.safeParse(raw)
+  if (!header.success) {
+    logger.warn('Unreadable Sim control request dropped')
+    return null
+  }
+  const { id, operation } = header.data
+  const kind = operation && truncate(operation.kind, 64)
+  if (operation && !SIM_CONTROL_KINDS.some((known) => known === operation.kind)) {
+    logger.warn('Unsupported Sim control kind', { kind })
+    const body = JSON.stringify({ error: SIM_CONTROL_UNSUPPORTED, kind })
+    return { id, answer: async () => ({ status: 501, body }) }
+  }
+  logger.warn('Invalid Sim control request', {
+    kind,
+    issues: request.error.issues.map((issue) => issue.path.join('.')),
+  })
+  return {
+    id,
+    answer: async () => ({
+      status: 400,
+      body: JSON.stringify({ error: 'Invalid Sim control request' }),
+    }),
+  }
+}
+
 export async function receiveSimControls(
   baseURL: string,
   channelId: string,
   signal: AbortSignal
 ): Promise<void> {
   while (!signal.aborted) {
+    let phase: 'poll' | 'parse' | 'reply' = 'poll'
     try {
       const response = await fetchGo(`${baseURL}/api/sim-transport/poll`, {
         method: 'POST',
         headers: mothershipRequestHeaders(),
-        body: JSON.stringify({ channelId }),
+        body: JSON.stringify({ channelId, accepts: SIM_CONTROL_KINDS }),
         signal: AbortSignal.any([signal, AbortSignal.timeout(30_000)]),
         redirect: 'error',
         spanName: 'sim → worker transport poll',
@@ -33,15 +77,18 @@ export async function receiveSimControls(
         await response.body?.cancel()
         throw new Error(`Transport poll refused (HTTP ${response.status})`)
       }
-      const batch = SimChannelBatch.parse(await response.json())
+      phase = 'parse'
+      const batch = SimChannelBatchEnvelope.parse(await response.json())
+      const controls = batch.requests.flatMap((raw) => receiveControl(raw) ?? [])
+      phase = 'reply'
       const replies = await Promise.allSettled(
-        batch.requests.map(async (request) => {
+        controls.map(async (control) => {
           signal.throwIfAborted()
-          const result = await executeSimControl(request)
+          const result = await control.answer()
           const reply = await fetchGo(`${baseURL}/api/sim-transport/reply`, {
             method: 'POST',
             headers: mothershipRequestHeaders(),
-            body: JSON.stringify({ channelId, id: request.id, result }),
+            body: JSON.stringify({ channelId, id: control.id, result }),
             redirect: 'error',
             signal: AbortSignal.any([signal, AbortSignal.timeout(5000)]),
             spanName: 'sim → worker transport reply',
@@ -56,7 +103,10 @@ export async function receiveSimControls(
       if (failed?.status === 'rejected') throw failed.reason
     } catch (error) {
       if (signal.aborted) return
-      logger.warn('Sim transport disconnected; reconnecting', { error: getErrorMessage(error) })
+      logger.warn('Sim transport disconnected; reconnecting', {
+        phase,
+        error: getErrorMessage(error),
+      })
       await sleep(1000)
     }
   }