address comments

Author: icecrasher321

apps/sim/lib/mcp/orchestration/workflow-mcp-lifecycle.test.ts

@@ -4,6 +4,17 @@
 import { dbChainMock, dbChainMockFns, resetDbChainMock, schemaMock } from '@sim/testing'
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 
+vi.mock('@sim/audit', () => ({
+  AuditAction: {
+    MCP_SERVER_UPDATED: 'mcp_server_updated',
+    MCP_TOOL_UPDATED: 'mcp_tool_updated',
+  },
+  AuditResourceType: {
+    MCP_SERVER: 'mcp_server',
+    MCP_TOOL: 'mcp_tool',
+  },
+  recordAudit: vi.fn(),
+}))
 vi.mock('@sim/db', () => ({
   ...dbChainMock,
   workflow: schemaMock.workflow,
@@ -28,8 +39,11 @@ vi.mock('@/lib/mcp/workflow-mcp-sync', () => ({
   generateParameterSchemaForWorkflow: vi.fn().mockResolvedValue({ type: 'object', properties: {} }),
 }))
 
-import { MAX_MCP_TOOLS_PER_SERVER } from '@/lib/mcp/constants'
-import { performCreateWorkflowMcpServer } from '@/lib/mcp/orchestration/workflow-mcp-lifecycle'
+import { MAX_MCP_PARAMETER_SCHEMA_BYTES, MAX_MCP_TOOLS_PER_SERVER } from '@/lib/mcp/constants'
+import {
+  performCreateWorkflowMcpServer,
+  performUpdateWorkflowMcpTool,
+} from '@/lib/mcp/orchestration/workflow-mcp-lifecycle'
 import { hasValidStartBlock } from '@/lib/workflows/triggers/trigger-utils.server'
 
 describe('workflow MCP lifecycle orchestration', () => {
@@ -151,4 +165,39 @@ describe('workflow MCP lifecycle orchestration', () => {
     })
     expect(dbChainMockFns.insert).not.toHaveBeenCalled()
   })
+
+  it('allows updating tool metadata when an unchanged stored schema exceeds the new cap', async () => {
+    dbChainMockFns.limit.mockResolvedValueOnce([{ id: 'server-1' }]).mockResolvedValueOnce([
+      {
+        id: 'tool-1',
+        toolName: 'tool_a',
+        toolDescription: null,
+        parameterSchemaBytes: MAX_MCP_PARAMETER_SCHEMA_BYTES + 1,
+      },
+    ])
+    dbChainMockFns.returning.mockResolvedValueOnce([
+      {
+        id: 'tool-1',
+        serverId: 'server-1',
+        toolName: 'tool_a',
+        toolDescription: 'Updated description',
+      },
+    ])
+
+    const result = await performUpdateWorkflowMcpTool({
+      workspaceId: 'workspace-1',
+      userId: 'user-1',
+      serverId: 'server-1',
+      toolId: 'tool-1',
+      toolDescription: 'Updated description',
+    })
+
+    expect(result).toMatchObject({
+      success: true,
+      tool: {
+        toolDescription: 'Updated description',
+      },
+    })
+    expect(dbChainMockFns.update).toHaveBeenCalled()
+  })
 })

apps/sim/lib/mcp/orchestration/workflow-mcp-lifecycle.ts

@@ -4,7 +4,12 @@ import { createLogger } from '@sim/logger'
 import { generateId } from '@sim/utils/id'
 import { and, asc, eq, inArray, isNull, ne, sql } from 'drizzle-orm'
 import type { DbOrTx } from '@/lib/db/types'
-import { MAX_MCP_SERVERS_PER_WORKFLOW, MAX_MCP_TOOLS_PER_SERVER } from '@/lib/mcp/constants'
+import {
+  MAX_MCP_SERVER_PARAMETER_SCHEMAS_BYTES,
+  MAX_MCP_SERVER_TOOLS_METADATA_BYTES,
+  MAX_MCP_SERVERS_PER_WORKFLOW,
+  MAX_MCP_TOOLS_PER_SERVER,
+} from '@/lib/mcp/constants'
 import { mcpPubSub } from '@/lib/mcp/pubsub'
 import {
   acquireWorkflowMcpServerLock,
@@ -13,9 +18,13 @@ import {
 } from '@/lib/mcp/server-locks'
 import {
   addMcpToolMetadataUsage,
+  addMcpToolMetadataUsageRow,
+  createMcpToolMetadataUsageRow,
   getMcpServerToolMetadataUsageRows,
   getMcpToolDescriptionForStorage,
+  getMcpToolMetadataSizes,
   getMcpToolMetadataUsageFromRows,
+  type McpToolMetadataUsage,
   validateMcpServerToolMetadataBudget,
   validateMcpToolMetadataForStorage,
 } from '@/lib/mcp/tool-limits'
@@ -181,6 +190,25 @@ async function validateServerToolMetadataBudget(
   return validateMcpServerToolMetadataBudget(usage)
 }
 
+function validateServerToolMetadataBudgetForUpdate(
+  currentUsage: McpToolMetadataUsage,
+  proposedUsage: McpToolMetadataUsage
+): string | null {
+  if (
+    proposedUsage.schemaBytes > MAX_MCP_SERVER_PARAMETER_SCHEMAS_BYTES &&
+    proposedUsage.schemaBytes > currentUsage.schemaBytes
+  ) {
+    return `MCP server tool schemas exceed maximum size of ${MAX_MCP_SERVER_PARAMETER_SCHEMAS_BYTES} bytes`
+  }
+  if (
+    proposedUsage.metadataBytes > MAX_MCP_SERVER_TOOLS_METADATA_BYTES &&
+    proposedUsage.metadataBytes > currentUsage.metadataBytes
+  ) {
+    return `MCP server tool metadata exceeds maximum size of ${MAX_MCP_SERVER_TOOLS_METADATA_BYTES} bytes`
+  }
+  return null
+}
+
 async function prepareWorkflowMcpTool(params: {
   workflowRecord: WorkflowMcpToolWorkflowRecord
   toolName?: string
@@ -863,7 +891,7 @@ export async function performUpdateWorkflowMcpTool(
           id: workflowMcpTool.id,
           toolName: workflowMcpTool.toolName,
           toolDescription: workflowMcpTool.toolDescription,
-          parameterSchema: workflowMcpTool.parameterSchema,
+          parameterSchemaBytes: sql<number>`octet_length(${workflowMcpTool.parameterSchema}::text)`,
         })
         .from(workflowMcpTool)
         .where(
@@ -885,13 +913,13 @@ export async function performUpdateWorkflowMcpTool(
           ? updateData.toolDescription
           : currentTool.toolDescription
       const effectiveParameterSchema =
-        updateData.parameterSchema !== undefined
-          ? updateData.parameterSchema
-          : currentTool.parameterSchema
+        updateData.parameterSchema !== undefined ? updateData.parameterSchema : undefined
       const metadataLimitError = validateMcpToolMetadataForStorage({
         toolName: effectiveToolName,
         toolDescription: effectiveToolDescription,
-        parameterSchema: effectiveParameterSchema,
+        ...(effectiveParameterSchema !== undefined && {
+          parameterSchema: effectiveParameterSchema,
+        }),
       })
       if (metadataLimitError) {
         throw new WorkflowMcpExpectedError(metadataLimitError, 'validation')
@@ -919,18 +947,36 @@ export async function performUpdateWorkflowMcpTool(
         }
       }
 
-      const budgetError = await validateServerToolMetadataBudget(
-        params.serverId,
-        [
-          {
-            toolName: effectiveToolName,
-            toolDescription: effectiveToolDescription,
-            parameterSchema: effectiveParameterSchema,
-          },
-        ],
-        tx,
-        params.toolId
+      const baseUsage = getMcpToolMetadataUsageFromRows(
+        await getMcpServerToolMetadataUsageRows(tx, params.serverId, params.toolId)
+      )
+      const currentUsage = addMcpToolMetadataUsageRow(baseUsage, {
+        id: currentTool.id,
+        ...getMcpToolMetadataSizes({
+          toolName: currentTool.toolName,
+          toolDescription: currentTool.toolDescription,
+        }),
+        parameterSchemaBytes: Number(currentTool.parameterSchemaBytes) || 0,
+      })
+      const proposedUsage = addMcpToolMetadataUsageRow(
+        baseUsage,
+        effectiveParameterSchema !== undefined
+          ? createMcpToolMetadataUsageRow({
+              id: currentTool.id,
+              toolName: effectiveToolName,
+              toolDescription: effectiveToolDescription,
+              parameterSchema: effectiveParameterSchema,
+            })
+          : {
+              id: currentTool.id,
+              ...getMcpToolMetadataSizes({
+                toolName: effectiveToolName,
+                toolDescription: effectiveToolDescription,
+              }),
+              parameterSchemaBytes: Number(currentTool.parameterSchemaBytes) || 0,
+            }
       )
+      const budgetError = validateServerToolMetadataBudgetForUpdate(currentUsage, proposedUsage)
       if (budgetError) {
         throw new WorkflowMcpExpectedError(budgetError, 'validation')
       }