Skip to content

Commit 560f539

Browse files
fix(mcp): preserve operation selection and discover managed connections
1 parent 0dda86a commit 560f539

6 files changed

Lines changed: 166 additions & 10 deletions

File tree

apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/editor/components/sub-block/hooks/use-sub-block-value.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -112,7 +112,6 @@ export function useSubBlockValue<T = any>(
112112
// Emit the value to socket/DB and update local store
113113
const emitValue = useCallback(
114114
(value: T) => {
115-
collaborativeSetSubblockValue(blockId, subBlockId, value)
116115
if (
117116
blockType === 'mcp' &&
118117
(subBlockId === 'serverSelector' || subBlockId === 'serverReference')
@@ -127,6 +126,7 @@ export function useSubBlockValue<T = any>(
127126
}
128127
}
129128
}
129+
collaborativeSetSubblockValue(blockId, subBlockId, value)
130130
lastEmittedValueRef.current = value
131131
},
132132
[

apps/sim/blocks/blocks/mcp.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@ export const McpBlock: BlockConfig<McpResponse> = {
9090
title: 'Operation',
9191
type: 'mcp-tool-selector',
9292
selectorKey: 'mcp.tools',
93+
dependsOn: ['server'],
9394
required: { field: 'operation', value: 'list', not: true },
9495
placeholder: 'Select an operation',
9596
description: 'Available tools from the selected MCP server',
Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
/** @vitest-environment node */
2+
import { renderToStaticMarkup } from 'react-dom/server'
3+
import { beforeEach, describe, expect, it, vi } from 'vitest'
4+
5+
const mocks = vi.hoisted(() => {
6+
const values: Record<string, unknown> = {}
7+
const modes: Record<string, 'basic' | 'advanced'> = {}
8+
const workflowState = { blocks: { mcp: { type: 'mcp', data: { canonicalModes: modes } } } }
9+
const registryState = { activeWorkflowId: 'workflow' }
10+
const subBlockState = {
11+
workflowValues: { workflow: { mcp: values } },
12+
getValue: (_blockId: string, field: string) => values[field],
13+
}
14+
return {
15+
values,
16+
modes,
17+
workflowState,
18+
registryState,
19+
subBlockState,
20+
setValue: vi.fn(),
21+
setMode: vi.fn(),
22+
}
23+
})
24+
25+
vi.mock('@/hooks/use-collaborative-workflow', () => ({
26+
useCollaborativeWorkflow: () => ({
27+
collaborativeSetSubblockValue: mocks.setValue,
28+
collaborativeSetBlockCanonicalMode: mocks.setMode,
29+
}),
30+
}))
31+
vi.mock('@/providers/utils', () => ({ getProviderFromModel: vi.fn() }))
32+
vi.mock('@/stores/workflows/workflow/store', () => ({
33+
useWorkflowStore: Object.assign(
34+
(selector: (state: typeof mocks.workflowState) => unknown) => selector(mocks.workflowState),
35+
{ getState: () => mocks.workflowState }
36+
),
37+
}))
38+
vi.mock('@/stores/workflows/registry/store', () => ({
39+
useWorkflowRegistry: Object.assign(
40+
(selector: (state: typeof mocks.registryState) => unknown) => selector(mocks.registryState),
41+
{ getState: () => mocks.registryState }
42+
),
43+
}))
44+
vi.mock('@/stores/workflows/subblock/store', () => ({
45+
useSubBlockStore: { getState: () => mocks.subBlockState },
46+
}))
47+
vi.mock('zustand/traditional', () => ({
48+
useStoreWithEqualityFn: (
49+
store: { getState: () => typeof mocks.subBlockState },
50+
selector: (state: typeof mocks.subBlockState) => unknown
51+
) => selector(store.getState()),
52+
}))
53+
vi.mock('@/stores/workflow-diff/store', () => ({
54+
useWorkflowDiffStore: () => ({ hasActiveDiff: false, isShowingDiff: false }),
55+
}))
56+
57+
import { buildSelectorRawContext } from '@/lib/selectors/context'
58+
import { getSubBlocksDependingOnChange } from '@/lib/workflows/subblocks/dependencies'
59+
import { useSubBlockValue } from '@/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/editor/components/sub-block/hooks/use-sub-block-value'
60+
import { McpBlock } from '@/blocks/blocks/mcp'
61+
import { getBlock } from '@/blocks/registry'
62+
63+
describe('MCP server mode switching', () => {
64+
beforeEach(() => {
65+
vi.clearAllMocks()
66+
for (const field of Object.keys(mocks.values)) delete mocks.values[field]
67+
mocks.values.toolSelector = 'read'
68+
mocks.modes.server = 'advanced'
69+
mocks.modes.tool = 'basic'
70+
mocks.setMode.mockImplementation((_blockId, field, value) => {
71+
mocks.modes[field] = value
72+
})
73+
mocks.setValue.mockImplementation((blockId: string, field: string, value: unknown) => {
74+
mocks.values[field] = value
75+
for (const dependent of getSubBlocksDependingOnChange(McpBlock.subBlocks, field)) {
76+
if (mocks.values[dependent.id]) mocks.setValue(blockId, dependent.id, '')
77+
}
78+
})
79+
vi.mocked(getBlock).mockReturnValue(McpBlock)
80+
})
81+
82+
it('preserves the exact selected operation before server dependency clearing', () => {
83+
let changeServer: ((value: string) => void) | undefined
84+
function Harness() {
85+
const [, setValue] = useSubBlockValue<string>('mcp', 'serverReference')
86+
changeServer = setValue
87+
return null
88+
}
89+
renderToStaticMarkup(<Harness />)
90+
changeServer?.('<connection.id>')
91+
expect(mocks.values.serverReference).toBe('<connection.id>')
92+
expect(mocks.values.toolSelector).toBe('')
93+
expect(mocks.values.toolReference).toBe('read')
94+
expect(mocks.modes.tool).toBe('advanced')
95+
})
96+
97+
it('projects only the active canonical server into operation discovery', () => {
98+
const subBlocks = {
99+
serverSelector: { value: 'inactive' },
100+
serverReference: { value: 'mcp-cg-abcdefghijklmnopqrstu' },
101+
}
102+
const input = {
103+
selectorKey: 'mcp.tools' as const,
104+
blockType: 'mcp',
105+
subBlocks,
106+
canonicalModes: mocks.modes,
107+
dependsOn: ['server'],
108+
}
109+
expect(buildSelectorRawContext(input)).toEqual({ mcpServerId: 'mcp-cg-abcdefghijklmnopqrstu' })
110+
subBlocks.serverReference.value = '<connection.id>'
111+
expect(buildSelectorRawContext(input)).toEqual({})
112+
})
113+
})

apps/sim/lib/selectors/server/providers/mcp.test.ts

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,13 @@
11
/** @vitest-environment node */
22
import { beforeEach, describe, expect, it, vi } from 'vitest'
33

4-
const { discover } = vi.hoisted(() => ({ discover: vi.fn() }))
4+
const { discover, discoverManaged } = vi.hoisted(() => ({
5+
discover: vi.fn(),
6+
discoverManaged: vi.fn(),
7+
}))
8+
vi.mock('@/lib/credentials/application/discover-managed-mcp-tools', () => ({
9+
discoverManagedMcpToolsUseCase: { execute: discoverManaged },
10+
}))
511
vi.mock('@/lib/mcp/application/use-cases', () => ({
612
discoverMcpServerToolsUseCase: { execute: discover },
713
}))
@@ -68,4 +74,25 @@ describe('MCP tools selector', () => {
6874
discover.mockRejectedValueOnce(new Error('Destination access denied'))
6975
await expect(execute(args({ kind: 'list' }))).rejects.toThrow('Destination access denied')
7076
})
77+
it('discovers a managed connection through its authorized use case with the acting principal', async () => {
78+
discoverManaged.mockResolvedValue({ tools: [{ name: 'read', canonicalServerId: 'parent' }] })
79+
const input = args({ kind: 'list' })
80+
input.context.mcpServerId = 'mcp-cg-abcdefghijklmnopqrstu'
81+
expect(await execute(input)).toEqual({
82+
kind: 'list',
83+
items: [{ id: 'read', label: 'read' }],
84+
})
85+
expect(discoverManaged).toHaveBeenCalledWith({
86+
principal: input.principal,
87+
input: {
88+
workspaceId: 'destination',
89+
credentialId: 'mcp-cg-abcdefghijklmnopqrstu',
90+
signal: input.signal,
91+
},
92+
})
93+
expect(discover).not.toHaveBeenCalled()
94+
discoverManaged.mockRejectedValueOnce(new Error('Credential revoked'))
95+
await expect(execute(input)).rejects.toThrow('Credential revoked')
96+
expect(discover).not.toHaveBeenCalled()
97+
})
7198
})

apps/sim/lib/selectors/server/providers/mcp.ts

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
import { OrchestrationError } from '@/lib/core/orchestration/types'
2+
import { discoverManagedMcpToolsUseCase } from '@/lib/credentials/application/discover-managed-mcp-tools'
23
import { discoverMcpServerToolsUseCase } from '@/lib/mcp/application/use-cases'
4+
import { isManagedMcpConnectionId, MANAGED_MCP_CONNECTION_PREFIX } from '@/lib/mcp/utils'
35
import {
46
definePreparedSelectorAttachment,
57
detailSelectorResult,
@@ -18,11 +20,24 @@ export const mcpSelectorAttachments = {
1820
'validation',
1921
'MCP tool discovery requires a destination server'
2022
)
23+
const serverId = args.context.mcpServerId
24+
if (serverId.startsWith(MANAGED_MCP_CONNECTION_PREFIX)) {
25+
if (!isManagedMcpConnectionId(serverId))
26+
throw new OrchestrationError('validation', 'Invalid managed MCP connection ID')
27+
return discoverManagedMcpToolsUseCase.execute({
28+
principal: args.principal,
29+
input: {
30+
workspaceId: args.workspaceId,
31+
credentialId: serverId,
32+
signal: args.signal,
33+
},
34+
})
35+
}
2136
return discoverMcpServerToolsUseCase.execute({
2237
principal: args.principal,
2338
input: {
2439
workspaceId: args.workspaceId,
25-
serverId: args.context.mcpServerId,
40+
serverId,
2641
signal: args.signal,
2742
requireComplete: true,
2843
},

scripts/check-tool-registry-boundary.baseline.json

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -298,16 +298,16 @@
298298
}
299299
},
300300
"app/workspace/[workspaceId]/knowledge/page.tsx": {
301-
"modules": 2271,
301+
"modules": 2318,
302302
"gateways": {
303303
"apps/sim/triggers/registry.ts": 485,
304-
"apps/sim/app/workspace/[workspaceId]/knowledge/prefetch.ts": 379,
304+
"apps/sim/app/workspace/[workspaceId]/knowledge/prefetch.ts": 389,
305305
"apps/sim/blocks/registry.ts": 354,
306-
"apps/sim/lib/knowledge/application/knowledge-bases.ts": 319,
307-
"apps/sim/lib/auth/index.ts": 222,
308-
"apps/sim/lib/knowledge/orchestration/index.ts": 208,
309-
"apps/sim/lib/knowledge/orchestration/connectors.ts": 204,
310-
"apps/sim/app/workspace/[workspaceId]/knowledge/knowledge.tsx": 175
306+
"apps/sim/lib/knowledge/application/knowledge-bases.ts": 330,
307+
"apps/sim/lib/auth/index.ts": 224,
308+
"apps/sim/lib/knowledge/orchestration/index.ts": 201,
309+
"apps/sim/lib/knowledge/orchestration/connectors.ts": 197,
310+
"apps/sim/app/workspace/[workspaceId]/knowledge/knowledge.tsx": 188
311311
}
312312
},
313313
"app/workspace/[workspaceId]/layout.tsx": {

0 commit comments

Comments
 (0)