fix(copilot): fix custom tools (#2278)

* Fix title custom tool

* Checkpoitn (broken)

* Fix custom tool flash

* Edit workflow returns null fix

* Works

* Fix lint

Author: Sg312

apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.ts

@@ -16,6 +16,7 @@ import {
 } from '@/lib/workflows/triggers/triggers'
 import { useCurrentWorkflow } from '@/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-current-workflow'
 import type { BlockLog, ExecutionResult, StreamingExecution } from '@/executor/types'
+import { coerceValue } from '@/executor/utils/start-block'
 import { subscriptionKeys } from '@/hooks/queries/subscription'
 import { useExecutionStream } from '@/hooks/use-execution-stream'
 import { WorkflowValidationError } from '@/serializer'
@@ -757,7 +758,7 @@ export function useWorkflowExecution() {
       if (Array.isArray(inputFormatValue)) {
         inputFormatValue.forEach((field: any) => {
           if (field && typeof field === 'object' && field.name && field.value !== undefined) {
-            testInput[field.name] = field.value
+            testInput[field.name] = coerceValue(field.type, field.value)
           }
         })
       }

apps/sim/executor/utils/start-block.ts

@@ -132,7 +132,7 @@ function extractInputFormat(block: SerializedBlock): InputFormatField[] {
     .map((field) => field)
 }
 
-function coerceValue(type: string | null | undefined, value: unknown): unknown {
+export function coerceValue(type: string | null | undefined, value: unknown): unknown {
   if (value === undefined || value === null) {
     return value
   }

apps/sim/lib/copilot/registry.ts

@@ -32,6 +32,7 @@ export const ToolIds = z.enum([
   'navigate_ui',
   'knowledge_base',
   'manage_custom_tool',
+  'manage_mcp_tool',
 ])
 export type ToolId = z.infer<typeof ToolIds>
 
@@ -199,12 +200,6 @@ export const ToolArgSchemas = {
       .describe(
         'Required for edit and delete operations. The database ID of the custom tool (e.g., "0robnW7_JUVwZrDkq1mqj"). Use get_workflow_data with data_type "custom_tools" to get the list of tools and their IDs. Do NOT use the function name - use the actual "id" field from the tool.'
       ),
-    title: z
-      .string()
-      .optional()
-      .describe(
-        'The display title of the custom tool. Required for add. Should always be provided for edit/delete so the user knows which tool is being modified.'
-      ),
     schema: z
       .object({
         type: z.literal('function'),
@@ -227,6 +222,36 @@ export const ToolArgSchemas = {
         'Required for add. The JavaScript function body code. Use {{ENV_VAR}} for environment variables and reference parameters directly by name.'
       ),
   }),
+
+  manage_mcp_tool: z.object({
+    operation: z
+      .enum(['add', 'edit', 'delete'])
+      .describe('The operation to perform: add (create new), edit (update existing), or delete'),
+    serverId: z
+      .string()
+      .optional()
+      .describe(
+        'Required for edit and delete operations. The database ID of the MCP server. Use the MCP settings panel or API to get server IDs.'
+      ),
+    config: z
+      .object({
+        name: z.string().describe('The display name for the MCP server'),
+        transport: z
+          .enum(['streamable-http'])
+          .optional()
+          .default('streamable-http')
+          .describe('Transport protocol (currently only streamable-http is supported)'),
+        url: z.string().optional().describe('The MCP server endpoint URL (required for add)'),
+        headers: z
+          .record(z.string())
+          .optional()
+          .describe('Optional HTTP headers to send with requests'),
+        timeout: z.number().optional().describe('Request timeout in milliseconds (default: 30000)'),
+        enabled: z.boolean().optional().describe('Whether the server is enabled (default: true)'),
+      })
+      .optional()
+      .describe('Required for add and edit operations. The MCP server configuration.'),
+  }),
 } as const
 export type ToolArgSchemaMap = typeof ToolArgSchemas
 
@@ -292,6 +317,7 @@ export const ToolSSESchemas = {
   navigate_ui: toolCallSSEFor('navigate_ui', ToolArgSchemas.navigate_ui),
   knowledge_base: toolCallSSEFor('knowledge_base', ToolArgSchemas.knowledge_base),
   manage_custom_tool: toolCallSSEFor('manage_custom_tool', ToolArgSchemas.manage_custom_tool),
+  manage_mcp_tool: toolCallSSEFor('manage_mcp_tool', ToolArgSchemas.manage_mcp_tool),
 } as const
 export type ToolSSESchemaMap = typeof ToolSSESchemas
 
@@ -519,6 +545,13 @@ export const ToolResultSchemas = {
     title: z.string().optional(),
     message: z.string().optional(),
   }),
+  manage_mcp_tool: z.object({
+    success: z.boolean(),
+    operation: z.enum(['add', 'edit', 'delete']),
+    serverId: z.string().optional(),
+    serverName: z.string().optional(),
+    message: z.string().optional(),
+  }),
 } as const
 export type ToolResultSchemaMap = typeof ToolResultSchemas
 

apps/sim/lib/copilot/tools/client/workflow/edit-workflow.ts

@@ -93,6 +93,26 @@ export class EditWorkflowClientTool extends BaseClientTool {
     }
   }
 
+  /**
+   * Safely get the current workflow JSON sanitized for copilot without throwing.
+   * Used to ensure we always include workflow state in markComplete.
+   */
+  private getCurrentWorkflowJsonSafe(logger: ReturnType<typeof createLogger>): string | undefined {
+    try {
+      const currentState = useWorkflowStore.getState().getWorkflowState()
+      if (!currentState) {
+        logger.warn('No current workflow state available')
+        return undefined
+      }
+      return this.getSanitizedWorkflowJson(currentState)
+    } catch (error) {
+      logger.warn('Failed to get current workflow JSON safely', {
+        error: error instanceof Error ? error.message : String(error),
+      })
+      return undefined
+    }
+  }
+
   static readonly metadata: BaseClientToolMetadata = {
     displayNames: {
       [ClientToolCallState.generating]: { text: 'Editing your workflow', icon: Loader2 },
@@ -133,66 +153,16 @@ export class EditWorkflowClientTool extends BaseClientTool {
 
   async handleAccept(): Promise<void> {
     const logger = createLogger('EditWorkflowClientTool')
-    logger.info('handleAccept called', {
-      toolCallId: this.toolCallId,
-      state: this.getState(),
-      hasResult: this.lastResult !== undefined,
-    })
-    this.setState(ClientToolCallState.success)
-
-    // Read from the workflow store to get the actual state with diff applied
-    const workflowStore = useWorkflowStore.getState()
-    const currentState = workflowStore.getWorkflowState()
-
-    // Get the workflow state that was applied, merge subblocks, and sanitize
-    // This matches what get_user_workflow would return
-    const workflowJson = this.getSanitizedWorkflowJson(currentState)
-
-    // Build sanitized data including workflow JSON and any skipped/validation info from the result
-    const sanitizedData: Record<string, any> = {}
-    if (workflowJson) {
-      sanitizedData.userWorkflow = workflowJson
-    }
-
-    // Include skipped items and validation errors in the accept response for LLM feedback
-    if (this.lastResult?.skippedItems?.length > 0) {
-      sanitizedData.skippedItems = this.lastResult.skippedItems
-      sanitizedData.skippedItemsMessage = this.lastResult.skippedItemsMessage
-    }
-    if (this.lastResult?.inputValidationErrors?.length > 0) {
-      sanitizedData.inputValidationErrors = this.lastResult.inputValidationErrors
-      sanitizedData.inputValidationMessage = this.lastResult.inputValidationMessage
-    }
-
-    // Build a message that includes info about skipped items
-    let acceptMessage = 'Workflow edits accepted'
-    if (
-      this.lastResult?.skippedItems?.length > 0 ||
-      this.lastResult?.inputValidationErrors?.length > 0
-    ) {
-      const parts: string[] = []
-      if (this.lastResult?.skippedItems?.length > 0) {
-        parts.push(`${this.lastResult.skippedItems.length} operation(s) were skipped`)
-      }
-      if (this.lastResult?.inputValidationErrors?.length > 0) {
-        parts.push(`${this.lastResult.inputValidationErrors.length} input(s) were rejected`)
-      }
-      acceptMessage = `Workflow edits accepted. Note: ${parts.join(', ')}.`
-    }
-
-    await this.markToolComplete(
-      200,
-      acceptMessage,
-      Object.keys(sanitizedData).length > 0 ? sanitizedData : undefined
-    )
+    logger.info('handleAccept called', { toolCallId: this.toolCallId, state: this.getState() })
+    // Tool was already marked complete in execute() - this is just for UI state
     this.setState(ClientToolCallState.success)
   }
 
   async handleReject(): Promise<void> {
     const logger = createLogger('EditWorkflowClientTool')
     logger.info('handleReject called', { toolCallId: this.toolCallId, state: this.getState() })
+    // Tool was already marked complete in execute() - this is just for UI state
     this.setState(ClientToolCallState.rejected)
-    await this.markToolComplete(200, 'Workflow changes rejected')
   }
 
   async execute(args?: EditWorkflowArgs): Promise<void> {
@@ -202,9 +172,14 @@ export class EditWorkflowClientTool extends BaseClientTool {
     await this.executeWithTimeout(async () => {
       if (this.hasExecuted) {
         logger.info('execute skipped (already executed)', { toolCallId: this.toolCallId })
-        // Even if skipped, ensure we mark complete
+        // Even if skipped, ensure we mark complete with current workflow state
         if (!this.hasBeenMarkedComplete()) {
-          await this.markToolComplete(200, 'Tool already executed')
+          const currentWorkflowJson = this.getCurrentWorkflowJsonSafe(logger)
+          await this.markToolComplete(
+            200,
+            'Tool already executed',
+            currentWorkflowJson ? { userWorkflow: currentWorkflowJson } : undefined
+          )
         }
         return
       }
@@ -231,7 +206,12 @@ export class EditWorkflowClientTool extends BaseClientTool {
       const operations = args?.operations || []
       if (!operations.length) {
         this.setState(ClientToolCallState.error)
-        await this.markToolComplete(400, 'No operations provided for edit_workflow')
+        const currentWorkflowJson = this.getCurrentWorkflowJsonSafe(logger)
+        await this.markToolComplete(
+          400,
+          'No operations provided for edit_workflow',
+          currentWorkflowJson ? { userWorkflow: currentWorkflowJson } : undefined
+        )
         return
       }
 
@@ -281,12 +261,22 @@ export class EditWorkflowClientTool extends BaseClientTool {
 
         if (!res.ok) {
           const errorText = await res.text().catch(() => '')
+          let errorMessage: string
           try {
             const errorJson = JSON.parse(errorText)
-            throw new Error(errorJson.error || errorText || `Server error (${res.status})`)
+            errorMessage = errorJson.error || errorText || `Server error (${res.status})`
           } catch {
-            throw new Error(errorText || `Server error (${res.status})`)
+            errorMessage = errorText || `Server error (${res.status})`
           }
+          // Mark complete with error but include current workflow state
+          this.setState(ClientToolCallState.error)
+          const currentWorkflowJson = this.getCurrentWorkflowJsonSafe(logger)
+          await this.markToolComplete(
+            res.status,
+            errorMessage,
+            currentWorkflowJson ? { userWorkflow: currentWorkflowJson } : undefined
+          )
+          return
         }
 
         const json = await res.json()
@@ -318,7 +308,14 @@ export class EditWorkflowClientTool extends BaseClientTool {
 
         // Update diff directly with workflow state - no YAML conversion needed!
         if (!result.workflowState) {
-          throw new Error('No workflow state returned from server')
+          this.setState(ClientToolCallState.error)
+          const currentWorkflowJson = this.getCurrentWorkflowJsonSafe(logger)
+          await this.markToolComplete(
+            500,
+            'No workflow state returned from server',
+            currentWorkflowJson ? { userWorkflow: currentWorkflowJson } : undefined
+          )
+          return
         }
 
         let actualDiffWorkflow: WorkflowState | null = null
@@ -336,17 +333,37 @@ export class EditWorkflowClientTool extends BaseClientTool {
         actualDiffWorkflow = workflowStore.getWorkflowState()
 
         if (!actualDiffWorkflow) {
-          throw new Error('Failed to retrieve workflow state after applying changes')
+          this.setState(ClientToolCallState.error)
+          const currentWorkflowJson = this.getCurrentWorkflowJsonSafe(logger)
+          await this.markToolComplete(
+            500,
+            'Failed to retrieve workflow state after applying changes',
+            currentWorkflowJson ? { userWorkflow: currentWorkflowJson } : undefined
+          )
+          return
         }
 
         // Get the workflow state that was just applied, merge subblocks, and sanitize
         // This matches what get_user_workflow would return (the true state after edits were applied)
-        const workflowJson = this.getSanitizedWorkflowJson(actualDiffWorkflow)
+        let workflowJson = this.getSanitizedWorkflowJson(actualDiffWorkflow)
+
+        // Fallback: try to get current workflow state if sanitization failed
+        if (!workflowJson) {
+          workflowJson = this.getCurrentWorkflowJsonSafe(logger)
+        }
+
+        // userWorkflow must always be present on success - log error if missing
+        if (!workflowJson) {
+          logger.error('Failed to get workflow JSON on success path - this should not happen', {
+            toolCallId: this.toolCallId,
+            workflowId: this.workflowId,
+          })
+        }
 
         // Build sanitized data including workflow JSON and any skipped/validation info
-        const sanitizedData: Record<string, any> = {}
-        if (workflowJson) {
-          sanitizedData.userWorkflow = workflowJson
+        // Always include userWorkflow on success paths
+        const sanitizedData: Record<string, any> = {
+          userWorkflow: workflowJson ?? '{}', // Fallback to empty object JSON if all else fails
         }
 
         // Include skipped items and validation errors in the response for LLM feedback
@@ -372,21 +389,25 @@ export class EditWorkflowClientTool extends BaseClientTool {
           completeMessage = `Workflow diff ready for review. Note: ${parts.join(', ')}.`
         }
 
-        // Mark complete early to unblock LLM stream
-        await this.markToolComplete(
-          200,
-          completeMessage,
-          Object.keys(sanitizedData).length > 0 ? sanitizedData : undefined
-        )
+        // Mark complete early to unblock LLM stream - sanitizedData always has userWorkflow
+        await this.markToolComplete(200, completeMessage, sanitizedData)
 
         // Move into review state
         this.setState(ClientToolCallState.review, { result })
       } catch (fetchError: any) {
         clearTimeout(fetchTimeout)
-        if (fetchError.name === 'AbortError') {
-          throw new Error('Server request timed out')
-        }
-        throw fetchError
+        // Handle error with current workflow state
+        this.setState(ClientToolCallState.error)
+        const currentWorkflowJson = this.getCurrentWorkflowJsonSafe(logger)
+        const errorMessage =
+          fetchError.name === 'AbortError'
+            ? 'Server request timed out'
+            : fetchError.message || String(fetchError)
+        await this.markToolComplete(
+          500,
+          errorMessage,
+          currentWorkflowJson ? { userWorkflow: currentWorkflowJson } : undefined
+        )
       }
     })
   }