feat(mothership): format generated function code

apps/sim/lib/copilot/tools/server/workflow/edit-workflow/engine.ts

@@ -164,6 +164,8 @@ export function applyOperationsToWorkflowState(
   // Collect skipped items across all operations
   const skippedItems: SkippedItem[] = []
 
+  const functionCodeBlockIds = new Set<string>()
+
   // Normalize block IDs to UUIDs before processing
   const { normalizedOperations } = normalizeBlockIdsInOperations(operations)
 
@@ -184,6 +186,7 @@ export function applyOperationsToWorkflowState(
     skippedItems,
     validationErrors,
     permissionConfig,
+    functionCodeBlockIds,
     deferredConnections: [],
   }
 
@@ -288,7 +291,7 @@ export function applyOperationsToWorkflowState(
     )
   }
 
-  return { state: modifiedState, validationErrors, skippedItems }
+  return { state: modifiedState, validationErrors, skippedItems, functionCodeBlockIds }
 }
 
 /**

apps/sim/lib/copilot/tools/server/workflow/edit-workflow/function-code-formatting.test.ts

@@ -0,0 +1,190 @@
+/**
+ * @vitest-environment node
+ */
+import { describe, expect, it } from 'vitest'
+import { formatChangedFunctionCode } from '@/lib/copilot/tools/server/workflow/edit-workflow/function-code-formatting'
+import { CodeLanguage } from '@/lib/execution/languages'
+
+function functionBlock(
+  code: string,
+  language: CodeLanguage = CodeLanguage.JavaScript
+): Record<string, unknown> {
+  return {
+    type: 'function',
+    subBlocks: {
+      code: { value: code },
+      language: { value: language },
+    },
+  }
+}
+
+describe('formatChangedFunctionCode', () => {
+  it('formats changed JavaScript and Python Function blocks without mutating previous state', async () => {
+    const previousWorkflowState = {
+      blocks: {
+        'existing-function': functionBlock('return 0;'),
+        'python-function': functionBlock('value={"foo": 1}\nreturn value', CodeLanguage.Python),
+      },
+    }
+    const modifiedWorkflowState = {
+      blocks: {
+        'existing-function': functionBlock('const value=1;return value;'),
+        'new-function': functionBlock('const value=<start.value>;return {value};'),
+        'python-function': functionBlock('value={"foo":1};return value', CodeLanguage.Python),
+        agent: {
+          type: 'agent',
+          subBlocks: { code: { value: 'const value=1;return value;' } },
+        },
+      },
+    }
+
+    const result = await formatChangedFunctionCode(
+      new Set(['existing-function', 'new-function', 'python-function']),
+      modifiedWorkflowState
+    )
+
+    expect(modifiedWorkflowState.blocks['existing-function'].subBlocks.code.value).toBe(
+      'const value = 1;\nreturn value;'
+    )
+    expect(modifiedWorkflowState.blocks['new-function'].subBlocks.code.value).toBe(
+      'const value = <start.value>;\nreturn { value };'
+    )
+    expect(modifiedWorkflowState.blocks['python-function'].subBlocks.code.value).toBe(
+      'value = {"foo": 1}\nreturn value'
+    )
+    expect(modifiedWorkflowState.blocks.agent.subBlocks.code.value).toBe(
+      'const value=1;return value;'
+    )
+    expect(previousWorkflowState.blocks['existing-function'].subBlocks.code.value).toBe('return 0;')
+    expect(result).toEqual({
+      changedBlockIds: ['existing-function', 'new-function', 'python-function'],
+      errors: [],
+    })
+  })
+
+  it('formats using the final block language after all operations are applied', async () => {
+    const modifiedWorkflowState = {
+      blocks: {
+        'function-1': functionBlock('value={"foo":1};return value', CodeLanguage.Python),
+      },
+    }
+
+    await formatChangedFunctionCode(new Set(['function-1']), modifiedWorkflowState)
+
+    expect(modifiedWorkflowState.blocks['function-1'].subBlocks.code.value).toBe(
+      'value = {"foo": 1}\nreturn value'
+    )
+  })
+
+  it('formats a replacement Function block from the final state', async () => {
+    const modifiedWorkflowState = {
+      blocks: {
+        replacement: functionBlock('const value=1;return value;', CodeLanguage.JavaScript),
+      },
+    }
+
+    await formatChangedFunctionCode(new Set(['replacement']), modifiedWorkflowState)
+
+    expect(modifiedWorkflowState.blocks.replacement.subBlocks.code.value).toBe(
+      'const value = 1;\nreturn value;'
+    )
+  })
+
+  it('formats same-source nested edits after the engine resolves child identity', async () => {
+    const code = 'const value=1;return value;'
+    const modifiedWorkflowState = {
+      blocks: { 'existing-child': functionBlock(code, CodeLanguage.JavaScript) },
+    }
+
+    await formatChangedFunctionCode(new Set(['existing-child']), modifiedWorkflowState)
+
+    expect(modifiedWorkflowState.blocks['existing-child'].subBlocks.code.value).toBe(
+      'const value = 1;\nreturn value;'
+    )
+  })
+
+  it('does not format untouched JavaScript Function blocks', async () => {
+    const untouchedCode = 'const untouched=1;return untouched;'
+    const modifiedWorkflowState = {
+      blocks: {
+        'changed-function': functionBlock('const changed=1;return changed;'),
+        'untouched-function': functionBlock(untouchedCode),
+      },
+    }
+
+    await formatChangedFunctionCode(new Set(['changed-function']), modifiedWorkflowState)
+
+    expect(modifiedWorkflowState.blocks['changed-function'].subBlocks.code.value).toBe(
+      'const changed = 1;\nreturn changed;'
+    )
+    expect(modifiedWorkflowState.blocks['untouched-function'].subBlocks.code.value).toBe(
+      untouchedCode
+    )
+  })
+
+  it('does not format an untouched Function block with identical submitted code', async () => {
+    const code = 'const value=1;return value;'
+    const modifiedWorkflowState = {
+      blocks: {
+        'changed-function': functionBlock(code),
+        'untouched-function': functionBlock(code),
+      },
+    }
+
+    await formatChangedFunctionCode(new Set(['changed-function']), modifiedWorkflowState)
+
+    expect(modifiedWorkflowState.blocks['changed-function'].subBlocks.code.value).toBe(
+      'const value = 1;\nreturn value;'
+    )
+    expect(modifiedWorkflowState.blocks['untouched-function'].subBlocks.code.value).toBe(code)
+  })
+
+  it('keeps invalid submitted JavaScript unchanged', async () => {
+    const code = 'if ('
+    const modifiedWorkflowState = {
+      blocks: { 'function-1': functionBlock(code) },
+    }
+
+    const result = await formatChangedFunctionCode(new Set(['function-1']), modifiedWorkflowState)
+
+    expect(modifiedWorkflowState.blocks['function-1'].subBlocks.code.value).toBe(code)
+    expect(result).toEqual({
+      changedBlockIds: [],
+      errors: [
+        {
+          blockId: 'function-1',
+          language: CodeLanguage.JavaScript,
+          error: expect.any(String),
+        },
+      ],
+    })
+  })
+
+  it('keeps invalid submitted Python unchanged', async () => {
+    const code = 'if ('
+    const modifiedWorkflowState = {
+      blocks: { 'function-1': functionBlock(code, CodeLanguage.Python) },
+    }
+
+    const result = await formatChangedFunctionCode(new Set(['function-1']), modifiedWorkflowState)
+
+    expect(modifiedWorkflowState.blocks['function-1'].subBlocks.code.value).toBe(code)
+    expect(result).toEqual({
+      changedBlockIds: [],
+      errors: [
+        {
+          blockId: 'function-1',
+          language: CodeLanguage.Python,
+          error: expect.any(String),
+        },
+      ],
+    })
+  })
+
+  it('returns before inspecting workflow state when no code blocks changed', async () => {
+    await expect(formatChangedFunctionCode(new Set(), null)).resolves.toEqual({
+      changedBlockIds: [],
+      errors: [],
+    })
+  })
+})

apps/sim/lib/copilot/tools/server/workflow/edit-workflow/function-code-formatting.ts

@@ -0,0 +1,79 @@
+import { isRecordLike } from '@sim/utils/object'
+import { DEFAULT_CODE_LANGUAGE } from '@/lib/execution/languages'
+import {
+  type FormattableCodeLanguage,
+  formatFunctionCode,
+  isFormattableCodeLanguage,
+} from '@/lib/workflows/blocks/format-function-code'
+
+interface FunctionCodeField {
+  code: string
+  codeSubBlock: Record<string, unknown>
+  language: string
+}
+
+export interface ChangedFunctionCodeFormatResult {
+  changedBlockIds: string[]
+  errors: Array<{
+    blockId: string
+    language: FormattableCodeLanguage
+    error: string
+  }>
+}
+
+function createEmptyFormatResult(): ChangedFunctionCodeFormatResult {
+  return { changedBlockIds: [], errors: [] }
+}
+
+function readFunctionCodeField(block: unknown): FunctionCodeField | undefined {
+  if (!isRecordLike(block) || block.type !== 'function' || !isRecordLike(block.subBlocks)) {
+    return undefined
+  }
+
+  const codeSubBlock = block.subBlocks.code
+  if (!isRecordLike(codeSubBlock) || typeof codeSubBlock.value !== 'string') return undefined
+
+  const languageSubBlock = block.subBlocks.language
+  const language =
+    isRecordLike(languageSubBlock) && typeof languageSubBlock.value === 'string'
+      ? languageSubBlock.value
+      : DEFAULT_CODE_LANGUAGE
+
+  return { code: codeSubBlock.value, codeSubBlock, language }
+}
+
+/**
+ * Formats changed JavaScript and Python Function code after the workflow engine resolves operation
+ * ordering, nested block identity, and the final code language.
+ */
+export async function formatChangedFunctionCode(
+  functionCodeBlockIds: ReadonlySet<string>,
+  modifiedWorkflowState: unknown
+): Promise<ChangedFunctionCodeFormatResult> {
+  if (functionCodeBlockIds.size === 0) return createEmptyFormatResult()
+  if (!isRecordLike(modifiedWorkflowState) || !isRecordLike(modifiedWorkflowState.blocks)) {
+    return createEmptyFormatResult()
+  }
+
+  const blocks = modifiedWorkflowState.blocks
+  const targets = [...functionCodeBlockIds].flatMap((blockId) => {
+    const currentField = readFunctionCodeField(blocks[blockId])
+    if (!currentField || !isFormattableCodeLanguage(currentField.language)) return []
+    return [{ blockId, ...currentField, language: currentField.language }]
+  })
+
+  const results = await Promise.all(
+    targets.map(async ({ blockId, code, codeSubBlock, language }) => {
+      const result = await formatFunctionCode(code, language)
+      if (result.changed) codeSubBlock.value = result.code
+      return { blockId, language, result }
+    })
+  )
+
+  return {
+    changedBlockIds: results.filter(({ result }) => result.changed).map(({ blockId }) => blockId),
+    errors: results.flatMap(({ blockId, language, result }) =>
+      result.error ? [{ blockId, language, error: result.error }] : []
+    ),
+  }
+}

apps/sim/lib/copilot/tools/server/workflow/edit-workflow/index.ts

@@ -13,6 +13,7 @@ import {
   type BaseServerTool,
   type ServerToolContext,
 } from '@/lib/copilot/tools/server/base-tool'
+import { formatChangedFunctionCode } from '@/lib/copilot/tools/server/workflow/edit-workflow/function-code-formatting'
 import { env } from '@/lib/core/config/env'
 import { getSocketServerUrl } from '@/lib/core/utils/urls'
 import {
@@ -161,8 +162,25 @@ export const editWorkflowServerTool: BaseServerTool<EditWorkflowParams, unknown>
       state: modifiedWorkflowState,
       validationErrors,
       skippedItems,
+      functionCodeBlockIds,
     } = applyOperationsToWorkflowState(workflowState, operationsToApply, permissionConfig)
 
+    const codeFormatting = await formatChangedFunctionCode(
+      functionCodeBlockIds,
+      modifiedWorkflowState
+    )
+    const codeFormattingFailures = codeFormatting.errors.map(({ blockId, language }) => ({
+      blockId,
+      language,
+    }))
+    if (codeFormattingFailures.length > 0) {
+      logger.warn('Function code formatting failed open', {
+        workflowId,
+        failureCount: codeFormattingFailures.length,
+        blocks: codeFormattingFailures,
+      })
+    }
+
     // Add credential validation errors
     validationErrors.push(...credentialErrors)
 
@@ -396,6 +414,15 @@ export const editWorkflowServerTool: BaseServerTool<EditWorkflowParams, unknown>
       workflowState: { ...finalWorkflowState, blocks: layoutedBlocks },
       workflowLint,
       ...(workflowLintMessage && { workflowLintMessage }),
+      ...(functionCodeBlockIds.size > 0 && {
+        codeFormatting: {
+          changedBlockIds: codeFormatting.changedBlockIds,
+          failures: codeFormattingFailures,
+        },
+        ...(codeFormattingFailures.length > 0 && {
+          codeFormattingMessage: `${codeFormattingFailures.length} Function block(s) could not be formatted, so their original code was kept unchanged.`,
+        }),
+      }),
       ...(inputErrors && {
         inputValidationErrors: inputErrors,
         inputValidationMessage: `${inputErrors.length} input(s) were rejected due to validation errors. The workflow was still updated with valid inputs only. Errors: ${inputErrors.join('; ')}`,