Skip to content

Commit 404b46f

Browse files
icecrasher321icecrasher321
committed
fix tool refresh ux
1 parent 1c97aab commit 404b46f

File tree

9 files changed

+609
-112
lines changed

9 files changed

+609
-112
lines changed

‎apps/sim/app/api/mcp/servers/[id]/refresh/route.ts‎

Lines changed: 141 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,143 @@
11
import { db } from '@sim/db'
2-
import { mcpServers } from '@sim/db/schema'
2+
import { mcpServers, workflow, workflowBlocks } from '@sim/db/schema'
33
import { createLogger } from '@sim/logger'
44
import { and, eq, isNull } from 'drizzle-orm'
55
import type { NextRequest } from 'next/server'
66
import { withMcpAuth } from '@/lib/mcp/middleware'
77
import { mcpService } from '@/lib/mcp/service'
8-
import type { McpServerStatusConfig } from '@/lib/mcp/types'
8+
import type { McpServerStatusConfig, McpTool, McpToolSchema } from '@/lib/mcp/types'
99
import { createMcpErrorResponse, createMcpSuccessResponse } from '@/lib/mcp/utils'
1010

1111
const logger = createLogger('McpServerRefreshAPI')
1212

1313
export const dynamic = 'force-dynamic'
1414

15+
/** Schema stored in workflow blocks includes description from the tool. */
16+
type StoredToolSchema = McpToolSchema & { description?: string }
17+
18+
interface StoredTool {
19+
type: string
20+
title: string
21+
toolId: string
22+
params: {
23+
serverId: string
24+
serverUrl?: string
25+
toolName: string
26+
serverName?: string
27+
}
28+
schema?: StoredToolSchema
29+
[key: string]: unknown
30+
}
31+
32+
/** Core param keys that are metadata, not user-entered test values */
33+
const MCP_TOOL_CORE_PARAMS = new Set(['serverId', 'serverUrl', 'toolName', 'serverName'])
34+
35+
interface SyncResult {
36+
updatedCount: number
37+
updatedWorkflowIds: string[]
38+
}
39+
40+
/**
41+
* Syncs tool schemas from discovered MCP tools to all workflow blocks using those tools.
42+
* Returns the count and IDs of updated workflows.
43+
*/
44+
async function syncToolSchemasToWorkflows(
45+
workspaceId: string,
46+
serverId: string,
47+
tools: McpTool[],
48+
requestId: string
49+
): Promise<SyncResult> {
50+
const toolsByName = new Map(tools.map((t) => [t.name, t]))
51+
52+
const workspaceWorkflows = await db
53+
.select({ id: workflow.id })
54+
.from(workflow)
55+
.where(eq(workflow.workspaceId, workspaceId))
56+
57+
const workflowIds = workspaceWorkflows.map((w) => w.id)
58+
if (workflowIds.length === 0) return { updatedCount: 0, updatedWorkflowIds: [] }
59+
60+
const agentBlocks = await db
61+
.select({
62+
id: workflowBlocks.id,
63+
workflowId: workflowBlocks.workflowId,
64+
subBlocks: workflowBlocks.subBlocks,
65+
})
66+
.from(workflowBlocks)
67+
.where(eq(workflowBlocks.type, 'agent'))
68+
69+
const updatedWorkflowIds = new Set<string>()
70+
71+
for (const block of agentBlocks) {
72+
if (!workflowIds.includes(block.workflowId)) continue
73+
74+
const subBlocks = block.subBlocks as Record<string, unknown> | null
75+
if (!subBlocks) continue
76+
77+
const toolsSubBlock = subBlocks.tools as { value?: StoredTool[] } | undefined
78+
if (!toolsSubBlock?.value || !Array.isArray(toolsSubBlock.value)) continue
79+
80+
let hasUpdates = false
81+
const updatedTools = toolsSubBlock.value.map((tool) => {
82+
if (tool.type !== 'mcp' || tool.params?.serverId !== serverId) {
83+
return tool
84+
}
85+
86+
const freshTool = toolsByName.get(tool.params.toolName)
87+
if (!freshTool) return tool
88+
89+
const newSchema: StoredToolSchema = {
90+
...freshTool.inputSchema,
91+
description: freshTool.description,
92+
}
93+
94+
const schemasMatch = JSON.stringify(tool.schema) === JSON.stringify(newSchema)
95+
96+
if (!schemasMatch) {
97+
hasUpdates = true
98+
99+
const validParamKeys = new Set(Object.keys(newSchema.properties || {}))
100+
101+
const cleanedParams: Record<string, unknown> = {}
102+
for (const [key, value] of Object.entries(tool.params || {})) {
103+
if (MCP_TOOL_CORE_PARAMS.has(key) || validParamKeys.has(key)) {
104+
cleanedParams[key] = value
105+
}
106+
}
107+
108+
return { ...tool, schema: newSchema, params: cleanedParams }
109+
}
110+
111+
return tool
112+
})
113+
114+
if (hasUpdates) {
115+
const updatedSubBlocks = {
116+
...subBlocks,
117+
tools: { ...toolsSubBlock, value: updatedTools },
118+
}
119+
120+
await db
121+
.update(workflowBlocks)
122+
.set({ subBlocks: updatedSubBlocks, updatedAt: new Date() })
123+
.where(eq(workflowBlocks.id, block.id))
124+
125+
updatedWorkflowIds.add(block.workflowId)
126+
}
127+
}
128+
129+
if (updatedWorkflowIds.size > 0) {
130+
logger.info(
131+
`[${requestId}] Synced tool schemas to ${updatedWorkflowIds.size} workflow(s) for server ${serverId}`
132+
)
133+
}
134+
135+
return {
136+
updatedCount: updatedWorkflowIds.size,
137+
updatedWorkflowIds: Array.from(updatedWorkflowIds),
138+
}
139+
}
140+
15141
export const POST = withMcpAuth<{ id: string }>('read')(
16142
async (request: NextRequest, { userId, workspaceId, requestId }, { params }) => {
17143
const { id: serverId } = await params
@@ -42,6 +168,8 @@ export const POST = withMcpAuth<{ id: string }>('read')(
42168
let connectionStatus: 'connected' | 'disconnected' | 'error' = 'error'
43169
let toolCount = 0
44170
let lastError: string | null = null
171+
let syncResult: SyncResult = { updatedCount: 0, updatedWorkflowIds: [] }
172+
let discoveredTools: McpTool[] = []
45173

46174
const currentStatusConfig: McpServerStatusConfig =
47175
(server.statusConfig as McpServerStatusConfig | null) ?? {
@@ -50,10 +178,17 @@ export const POST = withMcpAuth<{ id: string }>('read')(
50178
}
51179

52180
try {
53-
const tools = await mcpService.discoverServerTools(userId, serverId, workspaceId)
181+
discoveredTools = await mcpService.discoverServerTools(userId, serverId, workspaceId)
54182
connectionStatus = 'connected'
55-
toolCount = tools.length
183+
toolCount = discoveredTools.length
56184
logger.info(`[${requestId}] Discovered ${toolCount} tools from server ${serverId}`)
185+
186+
syncResult = await syncToolSchemasToWorkflows(
187+
workspaceId,
188+
serverId,
189+
discoveredTools,
190+
requestId
191+
)
57192
} catch (error) {
58193
connectionStatus = 'error'
59194
lastError = error instanceof Error ? error.message : 'Connection test failed'
@@ -92,6 +227,8 @@ export const POST = withMcpAuth<{ id: string }>('read')(
92227
toolCount,
93228
lastConnected: refreshedServer?.lastConnected?.toISOString() || null,
94229
error: lastError,
230+
workflowsUpdated: syncResult.updatedCount,
231+
updatedWorkflowIds: syncResult.updatedWorkflowIds,
95232
})
96233
} catch (error) {
97234
logger.error(`[${requestId}] Error refreshing MCP server:`, error)

‎apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/editor/components/sub-block/components/tool-input/tool-input.tsx‎

Lines changed: 41 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,12 @@ import {
1717
import { McpIcon } from '@/components/icons'
1818
import { Switch } from '@/components/ui/switch'
1919
import { cn } from '@/lib/core/utils/cn'
20+
import {
21+
getIssueBadgeLabel,
22+
getIssueBadgeVariant,
23+
isToolUnavailable,
24+
getMcpToolIssue as validateMcpTool,
25+
} from '@/lib/mcp/tool-validation'
2026
import {
2127
getCanonicalScopesForProvider,
2228
getProviderIdFromServiceId,
@@ -49,7 +55,7 @@ import {
4955
type CustomTool as CustomToolDefinition,
5056
useCustomTools,
5157
} from '@/hooks/queries/custom-tools'
52-
import { useMcpServers } from '@/hooks/queries/mcp'
58+
import { useForceRefreshMcpTools, useMcpServers, useStoredMcpTools } from '@/hooks/queries/mcp'
5359
import { useWorkflows } from '@/hooks/queries/workflows'
5460
import { useMcpTools } from '@/hooks/use-mcp-tools'
5561
import { getProviderFromModel, supportsToolUsageControl } from '@/providers/utils'
@@ -826,8 +832,11 @@ export function ToolInput({
826832
} = useMcpTools(workspaceId)
827833

828834
const { data: mcpServers = [], isLoading: mcpServersLoading } = useMcpServers(workspaceId)
835+
const { data: storedMcpTools = [] } = useStoredMcpTools(workspaceId)
836+
const forceRefreshMcpTools = useForceRefreshMcpTools()
829837
const openSettingsModal = useSettingsModalStore((state) => state.openModal)
830838
const mcpDataLoading = mcpLoading || mcpServersLoading
839+
const hasRefreshedRef = useRef(false)
831840

832841
const value = isPreview ? previewValue : storeValue
833842

@@ -839,27 +848,49 @@ export function ToolInput({
839848
? (value as StoredTool[])
840849
: []
841850

851+
const hasMcpTools = selectedTools.some((tool) => tool.type === 'mcp')
852+
853+
useEffect(() => {
854+
if (hasMcpTools && !hasRefreshedRef.current) {
855+
hasRefreshedRef.current = true
856+
forceRefreshMcpTools(workspaceId)
857+
}
858+
}, [hasMcpTools, forceRefreshMcpTools, workspaceId])
859+
842860
/**
843-
* Returns issue info for an MCP tool using shared validation logic.
861+
* Returns issue info for an MCP tool.
862+
* Uses DB schema (storedMcpTools) when available for real-time updates after refresh,
863+
* otherwise falls back to Zustand schema (tool.schema) which is always available.
844864
*/
845865
const getMcpToolIssue = useCallback(
846866
(tool: StoredTool) => {
847867
if (tool.type !== 'mcp') return null
848868

849-
const { getMcpToolIssue: validateTool } = require('@/lib/mcp/tool-validation')
869+
const serverId = tool.params?.serverId as string
870+
const toolName = tool.params?.toolName as string
850871

851-
return validateTool(
872+
// Try to get fresh schema from DB (enables real-time updates after MCP refresh)
873+
const storedTool =
874+
storedMcpTools.find(
875+
(st) =>
876+
st.serverId === serverId && st.toolName === toolName && st.workflowId === workflowId
877+
) || storedMcpTools.find((st) => st.serverId === serverId && st.toolName === toolName)
878+
879+
// Use DB schema if available, otherwise use Zustand schema
880+
const schema = storedTool?.schema ?? tool.schema
881+
882+
return validateMcpTool(
852883
{
853-
serverId: tool.params?.serverId as string,
884+
serverId,
854885
serverUrl: tool.params?.serverUrl as string | undefined,
855-
toolName: tool.params?.toolName as string,
856-
schema: tool.schema,
886+
toolName,
887+
schema,
857888
},
858889
mcpServers.map((s) => ({
859890
id: s.id,
860891
url: s.url,
861892
connectionStatus: s.connectionStatus,
862-
lastError: s.lastError,
893+
lastError: s.lastError ?? undefined,
863894
})),
864895
mcpTools.map((t) => ({
865896
serverId: t.serverId,
@@ -868,59 +899,16 @@ export function ToolInput({
868899
}))
869900
)
870901
},
871-
[mcpTools, mcpServers]
902+
[mcpTools, mcpServers, storedMcpTools, workflowId]
872903
)
873904

874905
const isMcpToolUnavailable = useCallback(
875906
(tool: StoredTool): boolean => {
876-
const { isToolUnavailable } = require('@/lib/mcp/tool-validation')
877907
return isToolUnavailable(getMcpToolIssue(tool))
878908
},
879909
[getMcpToolIssue]
880910
)
881911

882-
const hasMcpToolIssue = useCallback(
883-
(tool: StoredTool): boolean => {
884-
return getMcpToolIssue(tool) !== null
885-
},
886-
[getMcpToolIssue]
887-
)
888-
889-
const handleUpdateMcpToolSchema = useCallback(
890-
(toolIndex: number) => {
891-
if (isPreview || disabled) return
892-
893-
const tool = selectedTools[toolIndex]
894-
if (tool.type !== 'mcp' || !tool.params?.toolName || !tool.params?.serverId) return
895-
896-
const discoveredTool = mcpTools.find(
897-
(t) => t.serverId === tool.params?.serverId && t.name === tool.params?.toolName
898-
)
899-
900-
if (!discoveredTool?.inputSchema) {
901-
logger.warn('No discovered schema found for MCP tool')
902-
return
903-
}
904-
905-
logger.info(`Updating schema for MCP tool: ${tool.params.toolName}`)
906-
907-
setStoreValue(
908-
selectedTools.map((t, index) =>
909-
index === toolIndex
910-
? {
911-
...t,
912-
schema: {
913-
...discoveredTool.inputSchema,
914-
description: discoveredTool.description,
915-
},
916-
}
917-
: t
918-
)
919-
)
920-
},
921-
[isPreview, disabled, selectedTools, mcpTools, setStoreValue]
922-
)
923-
924912
// Filter out MCP tools from unavailable servers for the dropdown
925913
const availableMcpTools = useMemo(() => {
926914
return mcpTools.filter((mcpTool) => {
@@ -2209,13 +2197,12 @@ export function ToolInput({
22092197
(() => {
22102198
const issue = getMcpToolIssue(tool)
22112199
if (!issue) return null
2212-
const { getIssueBadgeLabel } = require('@/lib/mcp/tool-validation')
22132200
const serverId = tool.params?.serverId
22142201
return (
22152202
<Tooltip.Root>
22162203
<Tooltip.Trigger asChild>
22172204
<Badge
2218-
variant='amber'
2205+
variant={getIssueBadgeVariant(issue)}
22192206
className='cursor-pointer'
22202207
size='sm'
22212208
dot

0 commit comments

Comments
 (0)