chore: import upstream snapshot with attribution

apps/sim/executor/handlers/mothership/mothership-handler.test.ts

@@ -0,0 +1,599 @@
+import '@sim/testing/mocks/executor'
+
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+import { BlockType } from '@/executor/constants'
+import { MothershipBlockHandler } from '@/executor/handlers/mothership/mothership-handler'
+import type { ExecutionContext, StreamingExecution } from '@/executor/types'
+import type { SerializedBlock } from '@/serializer/types'
+
+const {
+  mockBuildAuthHeaders,
+  mockBuildAPIUrl,
+  mockExtractAPIErrorMessage,
+  mockGenerateId,
+  mockIsExecutionCancelled,
+  mockIsRedisCancellationEnabled,
+  mockReadUserFileContent,
+} = vi.hoisted(() => ({
+  mockBuildAuthHeaders: vi.fn(),
+  mockBuildAPIUrl: vi.fn(),
+  mockExtractAPIErrorMessage: vi.fn(),
+  mockGenerateId: vi.fn(),
+  mockIsExecutionCancelled: vi.fn(),
+  mockIsRedisCancellationEnabled: vi.fn(),
+  mockReadUserFileContent: vi.fn(),
+}))
+
+vi.mock('@/executor/utils/http', () => ({
+  buildAuthHeaders: mockBuildAuthHeaders,
+  buildAPIUrl: mockBuildAPIUrl,
+  extractAPIErrorMessage: mockExtractAPIErrorMessage,
+}))
+
+vi.mock('@sim/utils/id', () => ({
+  generateId: mockGenerateId,
+}))
+
+vi.mock('@/lib/execution/cancellation', () => ({
+  isExecutionCancelled: mockIsExecutionCancelled,
+  isRedisCancellationEnabled: mockIsRedisCancellationEnabled,
+}))
+
+vi.mock('@/lib/execution/payloads/materialization.server', () => ({
+  readUserFileContent: mockReadUserFileContent,
+}))
+
+function createAbortError(): Error {
+  const error = new Error('The operation was aborted')
+  error.name = 'AbortError'
+  return error
+}
+
+function createAbortableFetchPromise(signal?: AbortSignal): Promise<Response> {
+  return new Promise((_resolve, reject) => {
+    if (signal?.aborted) {
+      reject(createAbortError())
+      return
+    }
+
+    signal?.addEventListener(
+      'abort',
+      () => {
+        reject(createAbortError())
+      },
+      { once: true }
+    )
+  })
+}
+
+async function readStreamText(stream: ReadableStream): Promise<string> {
+  const reader = stream.getReader()
+  const decoder = new TextDecoder()
+  let text = ''
+
+  while (true) {
+    const { done, value } = await reader.read()
+    if (done) break
+    text += decoder.decode(value, { stream: true })
+  }
+
+  text += decoder.decode()
+  reader.releaseLock()
+  return text
+}
+
+describe('MothershipBlockHandler', () => {
+  let handler: MothershipBlockHandler
+  let block: SerializedBlock
+  let context: ExecutionContext
+  let fetchMock: ReturnType<typeof vi.fn>
+
+  beforeEach(() => {
+    handler = new MothershipBlockHandler()
+    fetchMock = vi.fn()
+    vi.stubGlobal('fetch', fetchMock)
+
+    mockBuildAuthHeaders.mockResolvedValue({ Authorization: 'Bearer internal' })
+    mockBuildAPIUrl.mockReturnValue(new URL('/api/mothership/execute', 'http://localhost:3000'))
+    mockExtractAPIErrorMessage.mockResolvedValue('boom')
+    mockGenerateId.mockReset()
+    mockIsExecutionCancelled.mockReset()
+    mockIsRedisCancellationEnabled.mockReset()
+    mockIsRedisCancellationEnabled.mockReturnValue(false)
+    mockReadUserFileContent.mockReset()
+
+    block = {
+      id: 'mothership-block-1',
+      metadata: { id: BlockType.MOTHERSHIP, name: 'Mothership' },
+      position: { x: 0, y: 0 },
+      config: { tool: BlockType.MOTHERSHIP, params: {} },
+      inputs: { prompt: 'string', conversationId: 'string', files: 'file[]' },
+      outputs: {},
+      enabled: true,
+    } as SerializedBlock
+
+    context = {
+      workflowId: 'workflow-1',
+      executionId: 'execution-1',
+      workspaceId: 'workspace-1',
+      userId: 'user-1',
+      blockStates: new Map(),
+      blockLogs: [],
+      metadata: { duration: 0 },
+      environmentVariables: {},
+      decisions: { router: new Map(), condition: new Map() },
+      loopExecutions: new Map(),
+      completedLoops: new Set(),
+      executedBlocks: new Set(),
+      activeExecutionPath: new Set(),
+    } as ExecutionContext
+  })
+
+  afterEach(() => {
+    vi.useRealTimers()
+    vi.clearAllMocks()
+    vi.unstubAllGlobals()
+  })
+
+  function createNdjsonResponse(events: unknown[]): Response {
+    const encoder = new TextEncoder()
+    return new Response(
+      new ReadableStream({
+        start(controller) {
+          for (const event of events) {
+            controller.enqueue(encoder.encode(`${JSON.stringify(event)}\n`))
+          }
+          controller.close()
+        },
+      }),
+      {
+        status: 200,
+        headers: { 'Content-Type': 'application/x-ndjson; charset=utf-8' },
+      }
+    )
+  }
+
+  it('forwards workflow and execution metadata with generated UUID ids', async () => {
+    mockGenerateId.mockReturnValueOnce('chat-uuid')
+    mockGenerateId.mockReturnValueOnce('message-uuid')
+    mockGenerateId.mockReturnValueOnce('request-uuid')
+
+    fetchMock.mockResolvedValue(
+      new Response(
+        JSON.stringify({
+          content: 'done',
+          model: 'mothership',
+          conversationId: 'chat-uuid',
+          tokens: { total: 5 },
+          toolCalls: [],
+        }),
+        {
+          status: 200,
+          headers: { 'Content-Type': 'application/json' },
+        }
+      )
+    )
+
+    const result = await handler.execute(context, block, { prompt: 'Hello from workflow' })
+
+    expect(result).toEqual({
+      content: 'done',
+      model: 'mothership',
+      conversationId: 'chat-uuid',
+      tokens: { total: 5 },
+      toolCalls: { list: [], count: 0 },
+      cost: undefined,
+    })
+
+    expect(fetchMock).toHaveBeenCalledTimes(1)
+    const [url, options] = fetchMock.mock.calls[0] as [string, RequestInit]
+    expect(url).toBe('http://localhost:3000/api/mothership/execute')
+    expect(options.method).toBe('POST')
+    expect(options.signal).toBeInstanceOf(AbortSignal)
+    expect(options.headers).toMatchObject({
+      Accept: 'application/x-ndjson',
+      'X-Mothership-Execute-Stream': 'ndjson',
+    })
+
+    const body = JSON.parse(String(options.body))
+    expect(body).toEqual({
+      messages: [{ role: 'user', content: 'Hello from workflow' }],
+      workspaceId: 'workspace-1',
+      userId: 'user-1',
+      chatId: 'chat-uuid',
+      messageId: 'message-uuid',
+      requestId: 'request-uuid',
+      workflowId: 'workflow-1',
+      executionId: 'execution-1',
+    })
+  })
+
+  it('uses a provided conversation ID as the mothership chat ID', async () => {
+    mockGenerateId.mockReturnValueOnce('message-uuid')
+    mockGenerateId.mockReturnValueOnce('request-uuid')
+
+    fetchMock.mockResolvedValue(
+      new Response(
+        JSON.stringify({
+          content: 'continued',
+          model: 'mothership',
+          conversationId: 'existing-chat-id',
+          tokens: {},
+          toolCalls: [],
+        }),
+        {
+          status: 200,
+          headers: { 'Content-Type': 'application/json' },
+        }
+      )
+    )
+
+    const result = await handler.execute(context, block, {
+      prompt: 'Continue this thread',
+      conversationId: ' existing-chat-id ',
+    })
+
+    expect(result).toEqual({
+      content: 'continued',
+      model: 'mothership',
+      conversationId: 'existing-chat-id',
+      tokens: {},
+      toolCalls: { list: [], count: 0 },
+      cost: undefined,
+    })
+
+    const [, options] = fetchMock.mock.calls[0] as [string, RequestInit]
+    const body = JSON.parse(String(options.body))
+    expect(body).toEqual({
+      messages: [{ role: 'user', content: 'Continue this thread' }],
+      workspaceId: 'workspace-1',
+      userId: 'user-1',
+      chatId: 'existing-chat-id',
+      messageId: 'message-uuid',
+      requestId: 'request-uuid',
+      workflowId: 'workflow-1',
+      executionId: 'execution-1',
+    })
+    expect(mockGenerateId).toHaveBeenCalledTimes(2)
+  })
+
+  it('consumes mothership execute heartbeat streams until the final result', async () => {
+    mockGenerateId.mockReturnValueOnce('chat-uuid')
+    mockGenerateId.mockReturnValueOnce('message-uuid')
+    mockGenerateId.mockReturnValueOnce('request-uuid')
+
+    fetchMock.mockResolvedValue(
+      createNdjsonResponse([
+        { type: 'heartbeat', timestamp: '2026-05-15T18:13:48.000Z' },
+        {
+          type: 'final',
+          data: {
+            content: 'streamed done',
+            model: 'mothership',
+            conversationId: 'chat-uuid',
+            tokens: { total: 7 },
+            toolCalls: [{ name: 'tool_a', params: { a: 1 }, result: 'ok', durationMs: 42 }],
+            cost: { total: 0.1 },
+          },
+        },
+      ])
+    )
+
+    const result = await handler.execute(context, block, { prompt: 'Hello from workflow' })
+
+    expect(result).toEqual({
+      content: 'streamed done',
+      model: 'mothership',
+      conversationId: 'chat-uuid',
+      tokens: { total: 7 },
+      toolCalls: {
+        list: [
+          {
+            name: 'tool_a',
+            arguments: { a: 1 },
+            result: 'ok',
+            error: undefined,
+            duration: 42,
+          },
+        ],
+        count: 1,
+      },
+      cost: { total: 0.1 },
+    })
+  })
+
+  it('preserves failed tool calls as output metadata without throwing', async () => {
+    mockGenerateId.mockReturnValueOnce('chat-uuid')
+    mockGenerateId.mockReturnValueOnce('message-uuid')
+    mockGenerateId.mockReturnValueOnce('request-uuid')
+
+    fetchMock.mockResolvedValue(
+      createNdjsonResponse([
+        {
+          type: 'final',
+          data: {
+            content: 'The lookup failed, so I could not use that result.',
+            model: 'mothership',
+            conversationId: 'chat-uuid',
+            tokens: { total: 7 },
+            toolCalls: [
+              {
+                name: 'lookup_customer',
+                status: 'error',
+                params: { email: 'missing@example.com' },
+                result: { success: false, error: 'Customer not found' },
+                error: 'Customer not found',
+                durationMs: 42,
+              },
+            ],
+          },
+        },
+      ])
+    )
+
+    const result = await handler.execute(context, block, { prompt: 'Hello from workflow' })
+
+    expect(result).toEqual({
+      content: 'The lookup failed, so I could not use that result.',
+      model: 'mothership',
+      conversationId: 'chat-uuid',
+      tokens: { total: 7 },
+      toolCalls: {
+        list: [
+          expect.objectContaining({
+            name: 'lookup_customer',
+            status: 'error',
+            arguments: { email: 'missing@example.com' },
+            result: { success: false, error: 'Customer not found' },
+            error: 'Customer not found',
+            duration: 42,
+          }),
+        ],
+        count: 1,
+      },
+      cost: undefined,
+    })
+  })
+
+  it('surfaces mothership execute stream errors', async () => {
+    mockGenerateId.mockReturnValueOnce('chat-uuid')
+    mockGenerateId.mockReturnValueOnce('message-uuid')
+    mockGenerateId.mockReturnValueOnce('request-uuid')
+
+    fetchMock.mockResolvedValue(
+      createNdjsonResponse([
+        { type: 'heartbeat', timestamp: '2026-05-15T18:13:48.000Z' },
+        { type: 'error', error: 'Mothership execution aborted' },
+      ])
+    )
+
+    await expect(
+      handler.execute(context, block, { prompt: 'Hello from workflow' })
+    ).rejects.toThrow('Sim execution failed: Mothership execution aborted')
+  })
+
+  it('streams mothership assistant chunks and preserves final metadata', async () => {
+    context.stream = true
+    context.selectedOutputs = [`${block.id}_content`]
+    mockGenerateId.mockReturnValueOnce('chat-uuid')
+    mockGenerateId.mockReturnValueOnce('message-uuid')
+    mockGenerateId.mockReturnValueOnce('request-uuid')
+
+    fetchMock.mockResolvedValue(
+      createNdjsonResponse([
+        { type: 'heartbeat', timestamp: '2026-05-15T18:13:48.000Z' },
+        { type: 'chunk', content: 'Hello' },
+        { type: 'heartbeat', timestamp: '2026-05-15T18:14:03.000Z' },
+        { type: 'chunk', content: ' world' },
+        {
+          type: 'final',
+          data: {
+            content: 'Hello world',
+            model: 'mothership',
+            conversationId: 'chat-uuid',
+            tokens: { total: 7 },
+            toolCalls: [{ name: 'tool_a', params: { a: 1 }, result: 'ok', durationMs: 42 }],
+            cost: { total: 0.1 },
+          },
+        },
+      ])
+    )
+
+    const result = await handler.execute(context, block, { prompt: 'Hello from workflow' })
+    expect(result).toHaveProperty('stream')
+
+    const streamingExecution = result as StreamingExecution
+    await expect(readStreamText(streamingExecution.stream)).resolves.toBe('Hello world')
+    expect(streamingExecution.execution.output).toEqual({
+      content: 'Hello world',
+      model: 'mothership',
+      conversationId: 'chat-uuid',
+      tokens: { total: 7 },
+      toolCalls: {
+        list: [
+          {
+            name: 'tool_a',
+            arguments: { a: 1 },
+            result: 'ok',
+            error: undefined,
+            duration: 42,
+          },
+        ],
+        count: 1,
+      },
+      cost: { total: 0.1 },
+    })
+  })
+
+  it('surfaces mothership streaming errors while streaming selected content', async () => {
+    context.stream = true
+    context.selectedOutputs = [`${block.id}_content`]
+    mockGenerateId.mockReturnValueOnce('chat-uuid')
+    mockGenerateId.mockReturnValueOnce('message-uuid')
+    mockGenerateId.mockReturnValueOnce('request-uuid')
+
+    fetchMock.mockResolvedValue(
+      createNdjsonResponse([
+        { type: 'chunk', content: 'partial' },
+        { type: 'error', error: 'Mothership execution aborted' },
+      ])
+    )
+
+    const result = (await handler.execute(context, block, {
+      prompt: 'Hello from workflow',
+    })) as StreamingExecution
+
+    await expect(readStreamText(result.stream)).rejects.toThrow(
+      'Sim execution failed: Mothership execution aborted'
+    )
+  })
+
+  it('embeds attached files for the mothership execute request', async () => {
+    const fileContent = Buffer.from('hello mothership', 'utf8').toString('base64')
+    mockGenerateId.mockReturnValueOnce('chat-uuid')
+    mockGenerateId.mockReturnValueOnce('message-uuid')
+    mockGenerateId.mockReturnValueOnce('request-uuid')
+    mockReadUserFileContent.mockResolvedValueOnce(fileContent)
+
+    fetchMock.mockResolvedValue(
+      new Response(
+        JSON.stringify({
+          content: 'analyzed',
+          model: 'mothership',
+          conversationId: 'chat-uuid',
+          tokens: {},
+          toolCalls: [],
+        }),
+        {
+          status: 200,
+          headers: { 'Content-Type': 'application/json' },
+        }
+      )
+    )
+
+    const result = await handler.execute(context, block, {
+      prompt: 'Analyze this file',
+      files: [
+        {
+          name: 'notes.txt',
+          key: 'workspace/workspace-1/notes.txt',
+          size: 16,
+          type: 'text/plain',
+        },
+      ],
+    })
+
+    expect(result).toMatchObject({
+      content: 'analyzed',
+      model: 'mothership',
+      conversationId: 'chat-uuid',
+    })
+    expect(mockReadUserFileContent).toHaveBeenCalledWith(
+      expect.objectContaining({
+        id: expect.stringMatching(/^file-/),
+        key: 'workspace/workspace-1/notes.txt',
+        name: 'notes.txt',
+        url: '',
+        size: 16,
+        type: 'text/plain',
+      }),
+      expect.objectContaining({
+        encoding: 'base64',
+        userId: 'user-1',
+        workspaceId: 'workspace-1',
+        workflowId: 'workflow-1',
+        executionId: 'execution-1',
+        requestId: 'request-uuid',
+      })
+    )
+
+    const [, options] = fetchMock.mock.calls[0] as [string, RequestInit]
+    const body = JSON.parse(String(options.body))
+    expect(body.fileAttachments).toEqual([
+      {
+        type: 'document',
+        source: {
+          type: 'base64',
+          media_type: 'text/plain',
+          data: fileContent,
+        },
+        filename: 'notes.txt',
+      },
+    ])
+  })
+
+  it('propagates local aborts to the mothership request', async () => {
+    const abortController = new AbortController()
+    context.abortSignal = abortController.signal
+
+    mockGenerateId.mockReturnValueOnce('chat-uuid')
+    mockGenerateId.mockReturnValueOnce('message-uuid')
+    mockGenerateId.mockReturnValueOnce('request-uuid')
+
+    fetchMock.mockImplementation((_url: string, options?: RequestInit) =>
+      createAbortableFetchPromise(options?.signal as AbortSignal | undefined)
+    )
+
+    const executionPromise = handler.execute(context, block, { prompt: 'Abort me' })
+    const abortedExecution = executionPromise.catch((error) => error)
+
+    abortController.abort()
+
+    await expect(abortedExecution).resolves.toMatchObject({ name: 'AbortError' })
+  })
+
+  it('propagates durable workflow cancellation to the mothership request', async () => {
+    vi.useFakeTimers()
+
+    mockGenerateId.mockReturnValueOnce('chat-uuid')
+    mockGenerateId.mockReturnValueOnce('message-uuid')
+    mockGenerateId.mockReturnValueOnce('request-uuid')
+    mockIsRedisCancellationEnabled.mockReturnValue(true)
+    mockIsExecutionCancelled.mockResolvedValueOnce(false).mockResolvedValueOnce(true)
+
+    fetchMock.mockImplementation((_url: string, options?: RequestInit) =>
+      createAbortableFetchPromise(options?.signal as AbortSignal | undefined)
+    )
+
+    const executionPromise = handler.execute(context, block, { prompt: 'Cancel me durably' })
+    const abortedExecution = executionPromise.catch((error) => error)
+
+    await vi.advanceTimersByTimeAsync(1000)
+
+    await expect(abortedExecution).resolves.toMatchObject({ name: 'AbortError' })
+    expect(mockIsExecutionCancelled).toHaveBeenCalledWith('execution-1')
+  })
+
+  it('aborts the mothership request when selected-output streaming is cancelled', async () => {
+    context.stream = true
+    context.selectedOutputs = [`${block.id}_content`]
+
+    mockGenerateId.mockReturnValueOnce('chat-uuid')
+    mockGenerateId.mockReturnValueOnce('message-uuid')
+    mockGenerateId.mockReturnValueOnce('request-uuid')
+
+    let fetchSignal: AbortSignal | undefined
+    fetchMock.mockImplementation((_url: string, options?: RequestInit) => {
+      fetchSignal = options?.signal as AbortSignal | undefined
+      return Promise.resolve(
+        new Response(
+          new ReadableStream({
+            start() {},
+          }),
+          {
+            status: 200,
+            headers: { 'Content-Type': 'application/x-ndjson; charset=utf-8' },
+          }
+        )
+      )
+    })
+
+    const result = (await handler.execute(context, block, { prompt: 'Cancel stream' })) as
+      | StreamingExecution
+      | undefined
+
+    await result?.stream.cancel('client_cancelled')
+
+    expect(fetchSignal?.aborted).toBe(true)
+  })
+})