feat(aiagent): allow mcp as tools (#6790)

* draft mcp client * testing * fix * cleaning * mcp resource in inputtransforms * cleaning * big cleaning * cleaning * no arc * add utils file * refactor tools * add mcp actions * draft frontend * send arguments from backend * better frontend * cleaning * use token for auth * add logo * rm * fix * fix * chore: refactor mcp for ai agents (#6829) * Add Tool enum for AIAgent with backward compatibility - Created Tool enum that can be either Windmill (FlowModule) or Mcp (resource reference) - Created McpToolRef struct to hold MCP resource path - Implemented custom Deserialize for Tool with backward compatibility: - New format: {type: 'windmill'|'mcp', ...} - Old format: FlowModule objects (automatically wrapped in Tool::Windmill) - Updated AIAgent to use Vec<Tool> instead of Vec<FlowModule> - Updated FlowValue::traverse_leafs to handle Tool enum - Backward compatible: old flows with Vec<FlowModule> will deserialize correctly * Refactor AI executor to process Tool enum instead of extracting MCP from input_transforms - Separate Windmill tools and MCP resource paths from tools list - Process Windmill FlowModules into Tool definitions - Load MCP tools from resource paths in Tool::Mcp variants - Remove old logic that extracted mcp_resources from input_transforms - Import FlowModule, remove unused InputTransform - Fix type issues: use .as_str() for path and handle Option<bool> properly * handle in args * mcp as flowmodule * frontend * config for mcp * simplify logic * fix ai executor logic * cleaning * clean frontend * fix * better resource picker * fix and styling * add endpoint to fetch tools * apply tool filtering * fix name validation * better ui * use cache * fix * fix merge * refactor: Separate MCP tools from FlowModule in AIAgent - Add new AgentTool, ToolValue, and McpToolValue types - Update AIAgent to use Vec<AgentTool> instead of Vec<FlowModule> - Implement From traits for clean conversion between AgentTool and FlowModule - Add backward compatibility via custom deserializer for AgentTool - Simplify resolve_module logic by reusing existing resolve_modules function - Update traverse_leafs to handle AgentTool structure This refactoring separates MCP tools from FlowModule tools, making the type system clearer and eliminating the need to treat MCP servers as a special case of FlowModule. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> * refactor: Update ai_executor and worker_lockfiles for AgentTool - Update ai_executor.rs to handle new AgentTool structure - Separate MCP tools from FlowModule tools using ToolValue enum - Convert AgentTool to FlowModule for backward compatibility - Add imports for AgentTool and ToolValue types - Update worker_lockfiles.rs for lazy loading optimization - Convert AgentTool <-> FlowModule in insert_flow_modules - Preserve lazy loading for FlowModule tools via modules_node - Keep MCP tools inline (lightweight, no need for lazy loading) - Maintain backward compatibility with existing flows This enables the lazy loading optimization for FlowModule tools while keeping MCP tools inline, balancing performance and simplicity. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> * cleaning * adapt frontend * cleaning * cleaning * type fix * cleaning * fix back comp * move mcp button position * nit * cleaning * fix nested removal * cleaning * opti * fix chat markdown display * fix chat messages layout * fix back comp frontend * fix deserializer * nit * simpler serializer * use if else --------- Co-authored-by: Claude <noreply@anthropic.com>
2026-09-07 16:03:21 +00:00 · 2025-10-21 15:28:52 +00:00
parent ed3ac2d928
commit 1afa36ceeb
43 changed files with 2618 additions and 1044 deletions
@@ -11,6 +11,7 @@
 	import FlowLogViewerWrapper from './FlowLogViewerWrapper.svelte'
 	import { z } from 'zod'
 	import { onMount } from 'svelte'
+	import type { AgentTool } from './flows/agentToolUtils'

 	type AgentActionWithContent = NonNullable<FlowStatusModule['agent_actions']>[number] & {
 		content?: unknown
@@ -29,6 +30,13 @@
 							module_id: z.string(),
 							function_name: z.string()
 						}),
+						z.object({
+							type: z.literal('mcp_tool_call'),
+							call_id: z.string(),
+							function_name: z.string(),
+							resource_path: z.string(),
+							arguments: z.record(z.unknown()).optional()
+						}),
 						z.object({
 							type: z.literal('message')
 						})
@@ -39,7 +47,7 @@
 	})

 	interface Props {
-		tools: FlowModule[]
+		tools: AgentTool[]
 		agentJob: Partial<CompletedJob> & Pick<CompletedJob, 'id'> & { type: 'CompletedJob' }
 		workspaceId?: string | undefined
 		storedToolCallJobs?: Record<number, Job>
@@ -69,6 +77,13 @@
 					job_id: toolCall.job_id
 				}
 				onToolJobLoaded?.(job, idx)
+			} else if (toolCall.type === 'mcp_tool_call') {
+				fakeModuleStates[idx.toString()] = {
+					type: 'Success',
+					args: toolCall.arguments ?? {},
+					logs: '',
+					result: toolCall.content
+				}
 			} else {
 				fakeModuleStates[idx.toString()] = {
 					type: 'Success',
@@ -104,7 +119,15 @@
 									module_id: m.agent_action.module_id,
 									function_name: m.agent_action.function_name
 								}
-							: undefined) as AgentActionWithContent | undefined
+							: m.agent_action?.type === 'mcp_tool_call'
+								? {
+										type: 'mcp_tool_call',
+										content: m.content,
+										call_id: m.agent_action.call_id,
+										function_name: m.agent_action.function_name,
+										arguments: m.agent_action.arguments
+									}
+								: undefined) as AgentActionWithContent | undefined
 			)
 			.filter((m) => m !== undefined)

@@ -122,13 +145,22 @@
 									type: 'identity' as const
 								}
 							}
+						} else if (toolCall.type === 'mcp_tool_call') {
+							return {
+								id: idx.toString(),
+								value: {
+									type: 'identity' as const
+								},
+								summary: toolCall.function_name,
+								arguments: toolCall.arguments
+							}
 						} else {
 							const module = tools.find((m) => m.summary === toolCall.function_name)
 							return module
-								? {
+								? ({
 										...module,
 										id: idx.toString()
-									}
+									} as FlowModule)
 								: undefined
 						}
 					})
@@ -2,8 +2,9 @@
 	import { Loader2 } from 'lucide-svelte'
 	import DisplayResult from './DisplayResult.svelte'
 	import LogViewer from './LogViewer.svelte'
-	import type { CompletedJob, FlowModule, Job } from '$lib/gen'
+	import type { CompletedJob, Job } from '$lib/gen'
 	import AiAgentLogViewer from './AIAgentLogViewer.svelte'
+	import type { AgentTool } from './flows/agentToolUtils'

 	interface Props {
 		waitingForExecutor?: boolean
@@ -21,7 +22,7 @@
 		downloadLogs?: boolean
 		tagLabel?: string | undefined
 		aiAgentStatus?: {
-			tools: FlowModule[]
+			tools: AgentTool[]
 			agentJob: Partial<CompletedJob> & Pick<CompletedJob, 'id'> & { type: 'CompletedJob' }
 			storedToolCallJobs?: Record<number, Job>
 			onToolJobLoaded?: (job: Job, idx: number) => void
@@ -693,12 +693,16 @@

 												<div class="flex items-center gap-2">
 													<span class="text-xs font-mono text-left">
-														<b>
-															{mode === 'aiagent'
-																? module.summary
-																	? 'Tool call'
-																	: 'Message'
-																: module.id}
+														<b class="flex items-center gap-1">
+															{#if mode === 'aiagent'}
+																{#if module.summary}
+																	Tool call: {module.summary}
+																{:else}
+																	Message
+																{/if}
+															{:else}
+																{module.id}
+															{/if}
 														</b>
 														{#if mode === 'flow'}
 															{#if module.value.type === 'forloopflow'}
@@ -715,7 +719,7 @@
 																Step
 															{/if}
 														{/if}
-														{#if module.summary}
+														{#if module.summary && mode !== 'aiagent'}
 															: {module.summary}
 														{/if}
 														{#if hasEmptySubflowValue}
@@ -43,9 +43,11 @@
 	import {
 		AI_TOOL_CALL_PREFIX,
 		AI_TOOL_MESSAGE_PREFIX,
+		AI_MCP_TOOL_CALL_PREFIX,
 		getToolCallId
 	} from './graph/renderers/nodes/AIToolNode.svelte'
 	import JobAssetsViewer from './assets/JobAssetsViewer.svelte'
+	import McpToolCallDetails from './McpToolCallDetails.svelte'

 	let {
 		flowState: flowStateStore,
@@ -678,6 +680,12 @@
 									job_id: action.job_id,
 									type: success != undefined ? (success ? 'Success' : 'Failure') : 'InProgress'
 								})
+							} else if (action.type == 'mcp_tool_call') {
+								const mcpToolCallId = AI_MCP_TOOL_CALL_PREFIX + '-' + mod.id + '-' + idx
+								const success = mod.agent_actions_success?.[idx]
+								setModuleState(mcpToolCallId, {
+									type: success != undefined ? (success ? 'Success' : 'Failure') : 'InProgress'
+								})
 							} else if (action.type == 'message') {
 								const toolCallId = getToolCallId(idx, mod.id)
 								setModuleState(toolCallId, {
@@ -1806,6 +1814,27 @@
 									<div class="pt-2 px-4 pb-4">
 										<Alert type="info" title="Message output is available on the AI agent node" />
 									</div>
+								{:else if selectedNode?.startsWith(AI_MCP_TOOL_CALL_PREFIX)}
+									{@const [, agentModuleId, toolCallIndex] = selectedNode.split('-')}
+									{@const agentNode = localModuleStates?.[agentModuleId]}
+									{@const agentActions = agentNode?.agent_actions}
+									{@const mcpActionIndex = parseInt(toolCallIndex)}
+									{@const mcpAction =
+										agentActions && mcpActionIndex >= 0 && mcpActionIndex < agentActions.length
+											? agentActions[mcpActionIndex]
+											: undefined}
+									{#if mcpAction?.type === 'mcp_tool_call' && agentNode?.result?.messages}
+										{@const message = agentNode.result.messages.find(
+											(m) => m.agent_action?.call_id === mcpAction.call_id
+										)}
+										<McpToolCallDetails
+											functionName={mcpAction.function_name}
+											args={mcpAction.arguments ?? {}}
+											result={message?.content}
+											type="Success"
+											workspaceId={job?.workspace_id}
+										/>
+									{/if}
 								{:else if selectedNode}
 									{@const node = localModuleStates[selectedNode]}
 									{#if selectedNode == 'end'}
@@ -0,0 +1,44 @@
+<script lang="ts">
+	import { Badge } from './common'
+	import JobArgs from './JobArgs.svelte'
+	import DisplayResult from './DisplayResult.svelte'
+	import type { FlowStatusModule } from '$lib/gen'
+
+	interface Props {
+		functionName: string
+		args: any
+		result: any
+		type: FlowStatusModule['type']
+		workspaceId?: string | undefined
+	}
+
+	let { functionName, args, result, type, workspaceId = undefined }: Props = $props()
+</script>
+
+<div class="p-2 flex flex-col gap-4">
+	<!-- Header -->
+	<div class="flex items-center gap-2">
+		<span class="font-semibold text-sm">{functionName}</span>
+		<Badge color={type === 'Success' ? 'green' : type === 'Failure' ? 'red' : 'gray'}>
+			{type}
+		</Badge>
+	</div>
+
+	<!-- Arguments Section -->
+	<div>
+		<h3 class="text-sm font-semibold mb-2 text-secondary">Arguments</h3>
+		{#if args && typeof args === 'object' && Object.keys(args).length > 0}
+			<JobArgs {args} argLabel="Parameter" />
+		{:else}
+			<p class="text-xs text-tertiary italic">No arguments</p>
+		{/if}
+	</div>
+
+	<!-- Result Section -->
+	<div>
+		<h3 class="text-sm font-semibold mb-2 text-secondary">Result</h3>
+		<div class="border rounded">
+			<DisplayResult {result} {workspaceId} />
+		</div>
+	</div>
+</div>
@@ -17,6 +17,7 @@
 	import { aiChatManager } from '../AIChatManager.svelte'
 	import { refreshStateStore } from '$lib/svelte5Utils.svelte'
 	import DiffDrawer from '$lib/components/DiffDrawer.svelte'
+	import type { AgentTool } from '$lib/components/flows/agentToolUtils'

 	let {
 		flowModuleSchemaMap
@@ -267,7 +268,7 @@

 			const indexToInsertAt = index + 1

-			let newModules: FlowModule[] | undefined = undefined
+			let newModules: FlowModule[] | AgentTool[] | undefined = undefined
 			switch (step.type) {
 				case 'rawscript': {
 					const inlineScript = {
@@ -532,7 +533,7 @@
 					module.value.parallelism = undefined
 				} else if (module.value.parallel || opts.parallel === true) {
 					// Only set parallelism if parallel is enabled
-					const n = Math.max(1, Math.floor(Math.abs(opts.parallelism)));
+					const n = Math.max(1, Math.floor(Math.abs(opts.parallelism)))
 					module.value.parallelism = {
 						type: 'static',
 						value: n
@@ -0,0 +1,58 @@
+import type { AiAgent, FlowModule, FlowModuleValue } from '$lib/gen'
+
+// Type aliases for better readability
+export type AgentTool = AiAgent['tools'][number]
+export type FlowModuleTool = AgentTool & { value: { tool_type: 'flowmodule' } & FlowModuleValue }
+export type McpTool = AgentTool & {
+	value: {
+		tool_type: 'mcp'
+		resource_path: string
+		include_tools?: string[]
+		exclude_tools?: string[]
+	}
+}
+
+/**
+ * Type guard to check if a tool is a FlowModule tool
+ */
+export function isFlowModuleTool(tool: AgentTool): tool is FlowModuleTool {
+	return tool.value.tool_type === undefined || tool.value.tool_type === 'flowmodule'
+}
+
+/**
+ * Type guard to check if a tool is an MCP tool
+ */
+export function isMcpTool(tool: AgentTool): tool is McpTool {
+	return tool.value.tool_type === 'mcp'
+}
+
+/**
+ * Create an MCP tool from resource path
+ */
+export function createMcpTool(id: string): McpTool {
+	return {
+		id,
+		summary: '',
+		value: {
+			tool_type: 'mcp',
+			resource_path: '',
+			include_tools: [],
+			exclude_tools: []
+		}
+	}
+}
+
+/**
+ * Convert a FlowModule back to an AgentTool
+ * Used when saving changes back to the AI Agent tools array
+ */
+export function flowModuleToAgentTool(flowModule: FlowModule): AgentTool {
+	return {
+		id: flowModule.id,
+		summary: flowModule.summary,
+		value: {
+			tool_type: 'flowmodule',
+			...flowModule.value
+		} as FlowModuleTool['value']
+	}
+}
@@ -0,0 +1,49 @@
+<script lang="ts">
+	import type { AgentTool } from '../agentToolUtils'
+	import { isFlowModuleTool, isMcpTool, type McpTool } from '../agentToolUtils'
+	import type { FlowModule } from '$lib/gen'
+	import FlowModuleComponent from './FlowModuleComponent.svelte'
+	import McpToolEditor from './McpToolEditor.svelte'
+
+	interface Props {
+		tool: AgentTool
+		noEditor?: boolean
+		enableAi?: boolean
+		parentModule?: FlowModule | undefined
+		previousModule?: FlowModule | undefined
+		forceTestTab?: Record<string, boolean>
+		highlightArg?: Record<string, string | undefined>
+	}
+
+	let {
+		tool = $bindable(),
+		noEditor = false,
+		enableAi = false,
+		parentModule = undefined,
+		previousModule = undefined,
+		forceTestTab,
+		highlightArg
+	}: Props = $props()
+</script>
+
+{#if isFlowModuleTool(tool)}
+	<!-- FlowModule tool - use existing FlowModuleComponent -->
+	<FlowModuleComponent
+		{noEditor}
+		flowModule={tool as FlowModule}
+		{parentModule}
+		{previousModule}
+		failureModule={false}
+		preprocessorModule={false}
+		scriptKind="script"
+		scriptTemplate="script"
+		{enableAi}
+		savedModule={undefined}
+		forceTestTab={forceTestTab?.[tool.id]}
+		highlightArg={highlightArg?.[tool.id]}
+		isAgentTool={true}
+	/>
+{:else if isMcpTool(tool)}
+	<!-- MCP tool - use McpToolEditor -->
+	<McpToolEditor bind:tool={tool as McpTool} {noEditor} />
+{/if}
@@ -172,7 +172,10 @@
 	]

 	let topLevelNodes: [string, string][] = $state([])
-	function computeToplevelNodeChoices(funcDesc: string, preFilter: 'all' | 'workspace' | 'hub') {
+	function computeToplevelNodeChoices(
+		funcDesc: string,
+		preFilter: 'all' | 'workspace' | 'hub'
+	) {
 		if (funcDesc.length > 0 && preFilter == 'all' && kind == 'script') {
 			topLevelNodes = allToplevelNodes.filter((node) =>
 				node[0].toLowerCase().startsWith(funcDesc.toLowerCase())
@@ -21,6 +21,7 @@
 	import FlowWhileLoop from './FlowWhileLoop.svelte'
 	import type { TriggerContext } from '$lib/components/triggers'
 	import { formatCron } from '$lib/utils'
+	import AgentToolWrapper from './AgentToolWrapper.svelte'

 	const { selectedId, flowStateStore } = getContext<FlowEditorContext>('FlowEditorContext')

@@ -293,12 +294,16 @@
 		{/if}
 	{/each}
 {:else if flowModule.value.type === 'aiagent'}
-	{#each flowModule.value.tools as _, index (index)}
-		<FlowModuleWrapper
+	{@const toolIndex = flowModule.value.tools.findIndex((t) => t.id === $selectedId)}
+	{#if toolIndex !== -1}
+		<AgentToolWrapper
 			{noEditor}
-			bind:flowModule={flowModule.value.tools[index]}
-			bind:parentModule={flowModule}
-			isAgentTool
+			bind:tool={flowModule.value.tools[toolIndex]}
+			parentModule={flowModule}
+			{previousModule}
+			{enableAi}
+			{forceTestTab}
+			{highlightArg}
 		/>
-	{/each}
+	{/if}
 {/if}
@@ -0,0 +1,194 @@
+<script module lang="ts">
+	import { get } from 'svelte/store'
+	import { workspaceStore, userStore } from '$lib/stores'
+	import { ResourceService } from '$lib/gen'
+	import { createCache } from '$lib/utils'
+
+	let loadToolsCached = createCache(
+		({ workspace, path }: { workspace?: string; path?: string; refreshCount?: number }) =>
+			workspace && path && get(userStore)
+				? ResourceService.getMcpTools({ workspace, path })
+				: undefined,
+		{
+			initial: { workspace: get(workspaceStore), path: undefined, refreshCount: 0 },
+			invalidateMs: 1000 * 60
+		} // Cache for 60 seconds
+	)
+</script>
+
+<script lang="ts">
+	import type { McpTool } from '../agentToolUtils'
+	import Section from '$lib/components/Section.svelte'
+	import Label from '$lib/components/Label.svelte'
+	import { Button } from '$lib/components/common'
+	import { RefreshCw } from 'lucide-svelte'
+	import MultiSelect from '$lib/components/select/MultiSelect.svelte'
+	import { safeSelectItems } from '$lib/components/select/utils.svelte'
+	import ResourcePicker from '$lib/components/ResourcePicker.svelte'
+	import { usePromise } from '$lib/svelte5Utils.svelte'
+	import { untrack } from 'svelte'
+	import Alert from '$lib/components/common/alert/Alert.svelte'
+
+	interface Props {
+		tool: McpTool
+		noEditor: boolean
+	}
+
+	let { tool = $bindable() }: Props = $props()
+
+	let refreshCount = $state(0)
+
+	let tools = usePromise(
+		async () =>
+			await loadToolsCached({
+				workspace: $workspaceStore!,
+				path: tool.value.resource_path,
+				refreshCount
+			}),
+		{ loadInit: false, clearValueOnRefresh: false }
+	)
+
+	// Options for the multiselect
+	let toolOptions = $derived(safeSelectItems((tools.value ?? []).map((t) => t.name)))
+
+	// Watch for resource_path changes and refresh tools
+	$effect(() => {
+		// Track reactive dependencies
+		tool.value.resource_path
+		$workspaceStore
+		refreshCount
+		// Trigger refresh when resource_path or workspace changes
+		untrack(() => {
+			if (tool.value.resource_path?.length > 0) {
+				tools.refresh()
+			}
+		})
+	})
+
+	$effect(() => {
+		if (!tool.value.include_tools) {
+			tool.value.include_tools = []
+		}
+		if (!tool.value.exclude_tools) {
+			tool.value.exclude_tools = []
+		}
+	})
+
+	$effect(() => {
+		if (tool.value.resource_path?.length > 0 && tool.summary?.length === 0) {
+			tool.summary = `MCP: ${tool.value.resource_path}`
+		}
+	})
+</script>
+
+<div class="flex flex-col gap-4 p-4">
+	<!-- Explanatory Section -->
+	<Alert type="info" title="MCP Client Configuration">
+		{#snippet children()}
+			<p class="mb-2 text-sm">
+				MCP clients allow AI agents to access and execute a list of tools made available by an MCP
+				server.
+				<br />
+				Choose an MCP resource to make its tools available to the agent.
+				<br />
+				<br />
+				<strong>Note:</strong> Only HTTP streamable MCP servers are supported.
+			</p>
+		{/snippet}
+	</Alert>
+
+	<!-- Resource Path Section -->
+	<div class="w-full">
+		<Label label="MCP Resource">
+			<ResourcePicker resourceType="mcp" bind:value={tool.value.resource_path} />
+		</Label>
+	</div>
+
+	{#if tool.value.resource_path?.length > 0}
+		<!-- Summary Section -->
+		<div class="w-full">
+			<Label label="Summary">
+				<input
+					type="text"
+					bind:value={tool.summary}
+					placeholder="e.g., GitHub MCP"
+					class="text-sm w-full"
+				/>
+			</Label>
+		</div>
+
+		<!-- Available Tools Section -->
+		<Section label="Available Tools">
+			{#snippet action()}
+				<Button
+					size="xs"
+					color="light"
+					on:click={() => (refreshCount += 1)}
+					startIcon={{ icon: RefreshCw }}
+					disabled={tools.status === 'loading'}
+				>
+					{tools.status === 'loading' ? 'Loading...' : 'Refresh Tools'}
+				</Button>
+			{/snippet}
+			<div class="w-full flex flex-col gap-2">
+				{#if tools.error}
+					<div class="text-xs text-red-600 p-2 border border-red-300 rounded bg-red-50">
+						{tools.error?.body?.message ||
+							tools.error?.message ||
+							'Failed to load tools from MCP server'}
+					</div>
+				{/if}
+				<div class="max-h-48 overflow-y-auto border rounded p-2 bg-surface-secondary">
+					{#if tools.status === 'loading'}
+						<div class="text-xs text-secondary italic">Loading tools...</div>
+					{:else if (tools.value ?? []).length === 0}
+						<div class="text-xs text-secondary italic">
+							{tools.error
+								? 'Failed to load tools. Please check the resource path and try again.'
+								: 'No tools loaded yet. Click "Refresh Tools" to fetch tools from the MCP server.'}
+						</div>
+					{:else}
+						<div class="flex flex-col gap-1">
+							{#each tools.value ?? [] as tool}
+								<div class="text-xs">
+									<span class="font-semibold">{tool.name}</span>
+									{#if tool.description}
+										<span class="text-secondary">— {tool.description}</span>
+									{/if}
+								</div>
+							{/each}
+						</div>
+					{/if}
+				</div>
+			</div>
+		</Section>
+
+		<!-- Tool Filtering Section -->
+		{#if tool.value.include_tools && tool.value.exclude_tools}
+			<Section label="Tool Filtering">
+				<div class="w-full flex flex-col gap-3">
+					<div class="flex flex-col gap-2">
+						<Label label="Only include specified tools">
+							<MultiSelect
+								bind:value={tool.value.include_tools}
+								items={toolOptions}
+								placeholder="Choose tools to include..."
+								disablePortal
+							/>
+						</Label>
+					</div>
+					<div class="flex flex-col gap-2">
+						<Label label="Exclude specified tools">
+							<MultiSelect
+								bind:value={tool.value.exclude_tools}
+								items={toolOptions}
+								placeholder="Choose tools to exclude..."
+								disablePortal
+							/>
+						</Label>
+					</div>
+				</div>
+			</Section>
+		{/if}
+	{/if}
+</div>
@@ -72,69 +72,67 @@
 	}
 </script>

-<div class="flex flex-col h-full w-full">
-	<div class="flex-1 flex flex-col min-h-0 w-full">
-		<!-- Messages Container -->
-		<div
-			bind:this={manager.messagesContainer}
-			class="flex-1 overflow-y-auto p-4 bg-background"
-			onscroll={manager.handleScroll}
-		>
-			{#if deploymentInProgress}
-				<Alert type="warning" title="Deployment in progress" size="xs" />
-			{/if}
-			{#if manager.isLoadingMessages}
-				<div class="flex items-center justify-center h-full">
-					<Loader2 size={32} class="animate-spin" />
-				</div>
-			{:else if manager.messages.length === 0}
-				<div class="text-center text-tertiary flex items-center justify-center flex-col h-full">
-					<MessageCircle size={48} class="mx-auto mb-4 opacity-50" />
-					<p class="text-lg font-medium">Start a conversation</p>
-					<p class="text-sm">Send a message to run the flow and see the results</p>
-				</div>
-			{:else}
-				<div class="max-w-7xl mx-auto space-y-4">
-					{#each manager.messages as message (message.id)}
-						<FlowChatMessage {message} />
-					{/each}
-					{#if manager.isWaitingForResponse}
-						<div class="flex items-center gap-2 text-tertiary">
-							<Loader2 size={16} class="animate-spin" />
-							<span class="text-sm">Processing...</span>
-						</div>
-					{/if}
-				</div>
-			{/if}
-		</div>
+<div class="flex flex-col h-full flex-1 min-w-0">
+	<!-- Messages Container -->
+	<div
+		bind:this={manager.messagesContainer}
+		class="flex-1 min-h-0 overflow-y-auto p-4 bg-background"
+		onscroll={manager.handleScroll}
+	>
+		{#if deploymentInProgress}
+			<Alert type="warning" title="Deployment in progress" size="xs" />
+		{/if}
+		{#if manager.isLoadingMessages}
+			<div class="flex items-center justify-center h-full">
+				<Loader2 size={32} class="animate-spin" />
+			</div>
+		{:else if manager.messages.length === 0}
+			<div class="text-center text-tertiary flex items-center justify-center flex-col h-full">
+				<MessageCircle size={48} class="mx-auto mb-4 opacity-50" />
+				<p class="text-lg font-medium">Start a conversation</p>
+				<p class="text-sm">Send a message to run the flow and see the results</p>
+			</div>
+		{:else}
+			<div class="w-full xl:max-w-7xl mx-auto space-y-4">
+				{#each manager.messages as message (message.id)}
+					<FlowChatMessage {message} />
+				{/each}
+				{#if manager.isWaitingForResponse}
+					<div class="flex items-center gap-2 text-tertiary">
+						<Loader2 size={16} class="animate-spin" />
+						<span class="text-sm">Processing...</span>
+					</div>
+				{/if}
+			</div>
+		{/if}
+	</div>

-		<!-- Chat Input -->
-		<div class="p-2 bg-surface">
-			<div
-				class="flex items-center gap-2 rounded-lg border border-gray-200 dark:border-gray-600 bg-surface"
-				class:opacity-50={deploymentInProgress}
-			>
-				<textarea
-					bind:this={manager.inputElement}
-					bind:value={manager.inputMessage}
-					use:autosize
-					onkeydown={manager.handleKeyDown}
-					placeholder="Type your message here..."
-					class="flex-1 min-h-[24px] max-h-32 resize-none !border-0 !bg-transparent text-sm placeholder-gray-400 !outline-none !ring-0 p-0 !shadow-none focus:!border-0 focus:!outline-none focus:!ring-0 focus:!shadow-none"
-					rows={3}
-				></textarea>
-				<div class="flex-shrink-0 pr-2">
-					<Button
-						color="blue"
-						size="xs2"
-						btnClasses="!rounded-full !p-1.5"
-						startIcon={{ icon: ArrowUp }}
-						disabled={!manager.inputMessage?.trim() || manager.isLoading || deploymentInProgress}
-						on:click={() => manager.sendMessage()}
-						iconOnly
-						title={deploymentInProgress ? 'Deployment in progress' : 'Send message (Enter)'}
-					/>
-				</div>
+	<!-- Chat Input -->
+	<div class="p-2 bg-surface">
+		<div
+			class="flex items-center gap-2 rounded-lg border border-gray-200 dark:border-gray-600 bg-surface"
+			class:opacity-50={deploymentInProgress}
+		>
+			<textarea
+				bind:this={manager.inputElement}
+				bind:value={manager.inputMessage}
+				use:autosize
+				onkeydown={manager.handleKeyDown}
+				placeholder="Type your message here..."
+				class="flex-1 min-h-[24px] max-h-32 resize-none !border-0 !bg-transparent text-sm placeholder-gray-400 !outline-none !ring-0 p-0 !shadow-none focus:!border-0 focus:!outline-none focus:!ring-0 focus:!shadow-none"
+				rows={3}
+			></textarea>
+			<div class="flex-shrink-0 pr-2">
+				<Button
+					color="blue"
+					size="xs2"
+					btnClasses="!rounded-full !p-1.5"
+					startIcon={{ icon: ArrowUp }}
+					disabled={!manager.inputMessage?.trim() || manager.isLoading || deploymentInProgress}
+					on:click={() => manager.sendMessage()}
+					iconOnly
+					title={deploymentInProgress ? 'Deployment in progress' : 'Send message (Enter)'}
+				/>
 			</div>
 		</div>
 	</div>
@@ -380,9 +380,6 @@ class FlowChatManager {
 								success
 							} = this.parseStreamDeltas(data.new_result_stream)
 							accumulatedContent += newContent
-							if (accumulatedContent.length > 0 || type === 'tool_result') {
-								this.isWaitingForResponse = false
-							}

 							// Create tool message if type is tool_result
 							if (type === 'tool_result') {
@@ -11,64 +11,59 @@
 	}

 	let { message }: Props = $props()
+
+	const messageClass = $derived.by(() => {
+		const base = 'max-w-[90%] min-w-0 rounded-lg w-fit'
+		if (message.message_type === 'user') {
+			return `${base} ml-auto bg-surface-secondary p-3`
+		}
+		return `${base} mr-auto bg-surface border ${message.success !== false ? 'border-gray-200 dark:border-gray-600' : '!border-red-500'}`
+	})
 </script>

-<div
-	class={`flex ${message.message_type === 'user' ? 'justify-end' : 'justify-start'} ${message.loading || message.streaming ? 'min-h-[200px] items-start' : ''}`}
-	data-message-id={message.id}
->
-	<div
-		class="max-w-[90%] min-w-0 rounded-lg
-			{message.message_type === 'user'
-			? 'bg-surface-secondary p-3'
-			: `bg-surface border ${message.success !== false ? 'border-gray-200 dark:border-gray-600' : '!border-red-500'}`}"
-	>
-		{#if message.step_name}
-			<div
-				class="bg-surface-secondary text-2xs text-tertiary mb-2 font-medium py-1 px-2 rounded-t-lg"
-				>{message.step_name}</div
-			>
-		{/if}
+<div class={messageClass} data-message-id={message.id}>
+	{#if message.step_name}
+		<div class="bg-surface-secondary text-2xs text-tertiary mb-2 font-medium py-1 px-2 rounded-t-lg"
+			>{message.step_name}</div
+		>
+	{/if}

-		{#if message.message_type === 'user'}
-			<p class="whitespace-pre-wrap text-sm break-words">{message.content}</p>
-		{:else if message.loading}
-			<div class="flex items-center gap-2 text-tertiary">
-				<Loader2 size={16} class="animate-spin" />
-				<span>Processing...</span>
-			</div>
-		{:else if message.content}
-			<div
-				class="flex flex-row items-center gap-2 px-3 pb-3 text-sm {!message.step_name
-					? 'pt-3'
-					: ''}"
-			>
-				{#if message.message_type === 'tool'}
-					{#if message.success !== false}
-						<CheckCircle2 class="w-3.5 h-3.5 text-green-500" />
-					{:else}
-						<AlertTriangle class="w-3.5 h-3.5 text-red-500" />
-					{/if}
+	{#if message.message_type === 'user'}
+		<p class="whitespace-pre-wrap text-sm break-words text-right">{message.content}</p>
+	{:else if message.loading}
+		<div class="flex items-center gap-2 text-tertiary">
+			<Loader2 size={16} class="animate-spin" />
+			<span>Processing...</span>
+		</div>
+	{:else if message.content}
+		<div
+			class="flex flex-row items-center gap-2 px-3 pb-3 text-sm {!message.step_name
+				? 'pt-3'
+				: ''} overflow-x-auto"
+		>
+			{#if message.message_type === 'tool'}
+				{#if message.success !== false}
+					<CheckCircle2 class="w-3.5 h-3.5 text-green-500" />
+				{:else}
+					<AlertTriangle class="w-3.5 h-3.5 text-red-500" />
 				{/if}
-				<div
-					class="prose prose-sm dark:prose-invert break-words whitespace-pre-wrap prose-headings:!text-base"
-				>
-					<Markdown
-						md={message.content}
-						plugins={[
-							gfmPlugin(),
-							{
-								renderer: {
-									pre: CodeDisplay,
-									a: LinkRenderer
-								}
+			{/if}
+			<div class="prose prose-sm dark:prose-invert break-words prose-headings:!text-base">
+				<Markdown
+					md={message.content}
+					plugins={[
+						gfmPlugin(),
+						{
+							renderer: {
+								pre: CodeDisplay,
+								a: LinkRenderer
 							}
-						]}
-					/>
-				</div>
+						}
+					]}
+				/>
 			</div>
-		{:else}
-			<p class="text-tertiary text-sm">No result</p>
-		{/if}
-	</div>
+		</div>
+	{:else}
+		<p class="text-tertiary text-sm">No result</p>
+	{/if}
 </div>
@@ -24,8 +24,8 @@ export function dfs<T>(
 				result = result.concat(dfs(branch, f, opts))
 			}
 		} else if (module.value.type == 'aiagent' && !opts.skipToolNodes) {
-			result = result.concat(f(module, modules, [module.value.tools]))
-			result = result.concat(dfs(module.value.tools, f, opts))
+			result = result.concat(f(module, modules, [module.value.tools as FlowModule[]]))
+			result = result.concat(dfs(module.value.tools as FlowModule[], f, opts))
 		} else {
 			result.push(f(module, modules, []))
 		}
@@ -41,6 +41,7 @@
 	import { getStepHistoryLoaderContext } from '$lib/components/stepHistoryLoader.svelte'
 	import { ModulesTestStates } from '$lib/components/modulesTest.svelte'
 	import type { StateStore } from '$lib/utils'
+	import { type AgentTool, flowModuleToAgentTool, createMcpTool } from '../agentToolUtils'

 	interface Props {
 		sidebarSize?: number | undefined
@@ -111,13 +112,14 @@
 	const { flowPropPickerConfig } = getContext<PropPickerContext>('PropPickerContext')

 	export async function insertNewModuleAtIndex(
-		modules: FlowModule[],
+		modules: FlowModule[] | AgentTool[],
 		index: number,
 		kind: InsertKind,
 		wsScript?: { path: string; summary: string; hash: string | undefined },
 		wsFlow?: { path: string; summary: string },
-		inlineScript?: InlineScript
-	): Promise<FlowModule[]> {
+		inlineScript?: InlineScript,
+		toolKind?: 'mcpTool' | 'flowmoduleTool'
+	): Promise<FlowModule[] | AgentTool[]> {
 		push(history, flowStore.val)
 		let module = emptyModule(flowStateStore.val, flowStore.val, kind == 'flow')
 		let state = emptyFlowModuleState()
@@ -167,8 +169,35 @@
 		}

 		if (!modules) return [module]
-		modules.splice(index, 0, module)
-		return modules
+
+		if (toolKind === 'mcpTool') {
+			// Create MCP AgentTool
+			const mcpTool = createMcpTool(module.id)
+			;(modules as AgentTool[]).splice(index, 0, mcpTool)
+			return modules as AgentTool[]
+		} else if (toolKind === 'flowmoduleTool') {
+			// Create AgentTool from FlowModule
+			const agentTool = flowModuleToAgentTool(module)
+			;(modules as AgentTool[]).splice(index, 0, agentTool)
+			return modules as AgentTool[]
+		} else {
+			// Standard FlowModule insertion (existing behavior)
+			modules.splice(index, 0, module)
+			return modules
+		}
+	}
+
+	/**
+	 * Helper function to remove an AgentTool by id from the tools array
+	 * Tools are always leaf nodes, so we just need to delete their state directly
+	 */
+	function removeAgentToolById(tools: AgentTool[], id: string): AgentTool[] {
+		const index = tools.findIndex((tool) => tool.id == id)
+		if (index != -1) {
+			const [removed] = tools.splice(index, 1)
+			deleteFlowStateById(removed.id, flowStateStore)
+		}
+		return tools
 	}

 	export function removeAtId(modules: FlowModule[], id: string): FlowModule[] {
@@ -194,7 +223,7 @@
 				})
 				mod.value.default = removeAtId(mod.value.default, id)
 			} else if (mod.value.type == 'aiagent') {
-				mod.value.tools = removeAtId(mod.value.tools, id)
+				mod.value.tools = removeAgentToolById(mod.value.tools, id)
 			}
 			return mod
 		})
@@ -489,6 +518,11 @@
 								}
 							} else {
 								const index = (detail.agentId ? targetModules?.length : detail.index) ?? 0
+								const toolKind = detail.agentId
+									? detail.kind === 'mcpTool'
+										? 'mcpTool'
+										: 'flowmoduleTool'
+									: undefined

 								await insertNewModuleAtIndex(
 									targetModules,
@@ -496,7 +530,8 @@
 									detail.kind,
 									detail.script,
 									detail.flow,
-									detail.inlineScript
+									detail.inlineScript,
+									toolKind
 								)
 								const id = targetModules[index].id
 								$selectedId = id
@@ -18,7 +18,7 @@
 		disableAi?: boolean
 		kind?: 'script' | 'trigger' | 'preprocessor' | 'failure'
 		allowTrigger?: boolean
-		scriptOnly?: boolean
+		toolMode?: boolean
 	}

 	let {
@@ -27,7 +27,7 @@
 		disableAi = false,
 		kind = 'script',
 		allowTrigger = true,
-		scriptOnly = false
+		toolMode = false
 	}: Props = $props()

 	let customUi: undefined | FlowBuilderWhitelabelCustomUi = getContext('customUi')
@@ -73,86 +73,96 @@
 	</div>

 	<div class="flex flex-row grow min-h-0">
-		{#if kind === 'script' && !scriptOnly}
+		{#if kind === 'script'}
 			<div class="flex-none flex flex-col text-xs text-primary">
 				<TopLevelNode
 					label="Action"
 					selected={selectedKind === 'script'}
-					on:select={() => {
+					onSelect={() => {
 						selectedKind = 'script'
 					}}
 				/>
-				{#if customUi?.triggers != false && allowTrigger}
+				{#if toolMode}
 					<TopLevelNode
-						label="Trigger"
-						selected={selectedKind === 'trigger'}
-						on:select={() => {
-							selectedKind = 'trigger'
-						}}
-					/>
-				{/if}
-				<TopLevelNode
-					label="Approval/Prompt"
-					selected={selectedKind === 'approval'}
-					on:select={() => {
-						selectedKind = 'approval'
-					}}
-				/>
-				{#if customUi?.flowNode != false}
-					<TopLevelNode
-						label="Flow"
-						selected={selectedKind === 'flow'}
-						on:select={() => {
-							selectedKind = 'flow'
-						}}
-					/>
-				{/if}
-				{#if stop}
-					<TopLevelNode
-						label="End flow"
-						selected={selectedKind === 'script'}
-						on:select={() => {
-							selectedKind = 'script'
-						}}
-					/>
-				{/if}
-
-				<TopLevelNode
-					label="For loop"
-					on:select={() => {
-						dispatch('close')
-						dispatch('new', { kind: 'forloop' })
-					}}
-				/>
-				<TopLevelNode
-					label="While loop"
-					on:select={() => {
-						dispatch('close')
-						dispatch('new', { kind: 'whileloop' })
-					}}
-				/>
-				<TopLevelNode
-					label="Branch to one"
-					on:select={() => {
-						dispatch('close')
-						dispatch('new', { kind: 'branchone' })
-					}}
-				/>
-				<TopLevelNode
-					label="Branch to all"
-					on:select={() => {
-						dispatch('close')
-						dispatch('new', { kind: 'branchall' })
-					}}
-				/>
-				{#if customUi?.aiAgent != false}
-					<TopLevelNode
-						label="AI Agent"
-						on:select={() => {
+						label="MCP"
+						onSelect={() => {
+							dispatch('pickMcpTool')
 							dispatch('close')
-							dispatch('new', { kind: 'aiagent' })
 						}}
 					/>
+				{:else}
+					{#if customUi?.triggers != false && allowTrigger}
+						<TopLevelNode
+							label="Trigger"
+							selected={selectedKind === 'trigger'}
+							onSelect={() => {
+								selectedKind = 'trigger'
+							}}
+						/>
+					{/if}
+					<TopLevelNode
+						label="Approval/Prompt"
+						selected={selectedKind === 'approval'}
+						onSelect={() => {
+							selectedKind = 'approval'
+						}}
+					/>
+					{#if customUi?.flowNode != false}
+						<TopLevelNode
+							label="Flow"
+							selected={selectedKind === 'flow'}
+							onSelect={() => {
+								selectedKind = 'flow'
+							}}
+						/>
+					{/if}
+					{#if stop}
+						<TopLevelNode
+							label="End flow"
+							selected={selectedKind === 'script'}
+							onSelect={() => {
+								selectedKind = 'script'
+							}}
+						/>
+					{/if}
+
+					<TopLevelNode
+						label="For loop"
+						onSelect={() => {
+							dispatch('close')
+							dispatch('new', { kind: 'forloop' })
+						}}
+					/>
+					<TopLevelNode
+						label="While loop"
+						onSelect={() => {
+							dispatch('close')
+							dispatch('new', { kind: 'whileloop' })
+						}}
+					/>
+					<TopLevelNode
+						label="Branch to one"
+						onSelect={() => {
+							dispatch('close')
+							dispatch('new', { kind: 'branchone' })
+						}}
+					/>
+					<TopLevelNode
+						label="Branch to all"
+						onSelect={() => {
+							dispatch('close')
+							dispatch('new', { kind: 'branchall' })
+						}}
+					/>
+					{#if customUi?.aiAgent != false}
+						<TopLevelNode
+							label="AI Agent"
+							onSelect={() => {
+								dispatch('close')
+								dispatch('new', { kind: 'aiagent' })
+							}}
+						/>
+					{/if}
 				{/if}
 			</div>
 		{/if}
@@ -19,4 +19,4 @@

 <svelte:window on:keydown={handleKeydown} />

-<TopLevelNode class="px-3" {label} {selected} returnIcon on:select={click} />
+<TopLevelNode class="px-3" {label} {selected} returnIcon onSelect={click} />
@@ -6,65 +6,70 @@
 		ChevronRight,
 		Code,
 		GitBranch,
+		Plug,
 		Repeat,
 		Square,
 		Zap
 	} from 'lucide-svelte'
-	import { createEventDispatcher } from 'svelte'
 	import { twMerge } from 'tailwind-merge'
+	import type { ComponentType } from 'svelte'

-	export let label: string
-	export let selected = false
-	export let returnIcon = false
-	const dispatch = createEventDispatcher()
+	interface Props {
+		label: string
+		selected?: boolean
+		returnIcon?: boolean
+		onSelect: () => void
+		class?: string
+	}
+
+	let { label, selected, returnIcon, onSelect, class: className }: Props = $props()
+
+	interface IconConfig {
+		icon: ComponentType
+		showChevron?: boolean
+		iconClass?: string
+	}
+
+	const iconMap: Record<string, IconConfig> = {
+		Action: { icon: Code, showChevron: true },
+		Trigger: { icon: Zap, showChevron: true },
+		'Approval/Prompt': { icon: CheckCircle2, showChevron: true },
+		Flow: { icon: BarsStaggered as unknown as ComponentType, showChevron: true },
+		'End Flow': { icon: Square },
+		'For loop': { icon: Repeat },
+		'While loop': { icon: Repeat },
+		'Branch to one': { icon: GitBranch },
+		'Branch to all': { icon: GitBranch },
+		'AI Agent': { icon: BotIcon, iconClass: 'text-violet-800 dark:text-violet-400' },
+		MCP: { icon: Plug, showChevron: true }
+	}
+
+	const config = $derived(iconMap[label])
 </script>

+{#snippet iconWithText(icon: ComponentType, showChevron = false, iconClass = '')}
+	{@const Icon = icon}
+	<Icon size={14} class={iconClass} />
+	{label}
+	{#if showChevron}
+		<ChevronRight size={12} class="ml-auto" color="#4c566a" />
+	{/if}
+{/snippet}
+
 <button
 	id={`flow-editor-flow-kind-${label.replaceAll(' ', '-').toLowerCase()}`}
 	class={twMerge(
 		'w-full text-left py-2 px-1.5 hover:bg-surface-hover text-xs font-medium transition-all whitespace-nowrap flex flex-row gap-2 items-center rounded-md',
 		selected ? 'bg-surface-hover' : '',
-		$$props.class
+		className
 	)}
-	on:pointerdown={() => dispatch('select', label)}
+	onpointerdown={onSelect}
 	role="menuitem"
 	tabindex="-1"
 >
 	<span class="grow flex items-center gap-2">
-		{#if label === 'Action'}
-			<Code size={14} />
-			Action
-			<ChevronRight size={12} class="ml-auto" color="#4c566a" />
-		{:else if label === 'Trigger'}
-			<Zap size={14} />
-			Trigger
-			<ChevronRight size={12} class="ml-auto" color="#4c566a" />
-		{:else if label === 'Approval/Prompt'}
-			<CheckCircle2 size={14} />
-			Approval/Prompt
-			<ChevronRight size={12} class="ml-auto" color="#4c566a" />
-		{:else if label === 'Flow'}
-			<BarsStaggered size={14} />
-			Flow
-			<ChevronRight size={12} class="ml-auto" color="#4c566a" />
-		{:else if label === 'End Flow'}
-			<Square size={14} />
-			End Flow
-		{:else if label === 'For loop'}
-			<Repeat size={14} />
-			For Loop
-		{:else if label === 'While loop'}
-			<Repeat size={14} />
-			While Loop
-		{:else if label === 'Branch to one'}
-			<GitBranch size={14} />
-			Branch to one
-		{:else if label === 'Branch to all'}
-			<GitBranch size={14} />
-			Branch to all
-		{:else if label === 'AI Agent'}
-			<BotIcon size={14} class="text-violet-800 dark:text-violet-400" />
-			AI Agent
+		{#if config}
+			{@render iconWithText(config.icon, config.showChevron, config.iconClass ?? '')}
 		{/if}
 	</span>
 	{#if returnIcon && selected}
@@ -19,6 +19,7 @@ export type InsertKind =
 	| 'approval'
 	| 'end'
 	| 'aiagent'
+	| 'mcpTool'

 export type InlineScript = {
 	language: RawScript['language']
@@ -305,6 +306,7 @@ export type AiToolN = {
 	type: 'aiTool'
 	data: {
 		tool: string
+		type?: string
 		eventHandlers: GraphEventHandlers
 		moduleId: string
 		insertable: boolean
@@ -1,5 +1,8 @@
 <script module lang="ts">
-	export function validateToolName(name: string) {
+	export function validateToolName(name: string, type?: string) {
+		if (type === 'mcp') {
+			return name.length > 0
+		}
 		return /^[a-zA-Z0-9_]+$/.test(name)
 	}

@@ -8,6 +11,7 @@
 	export const BELOW_ADDITIONAL_OFFSET = 19

 	export const AI_TOOL_CALL_PREFIX = '_wm_ai_agent_tool_call'
+	export const AI_MCP_TOOL_CALL_PREFIX = '_wm_ai_mcp_tool_call'
 	export const AI_TOOL_MESSAGE_PREFIX = '_wm_ai_agent_message'

 	const ROW_WIDTH = 275
@@ -79,11 +83,22 @@
 			let tools: {
 				id: string
 				name: string
+				type?: string
 				stateType?: GraphModuleState['type']
-			}[] = node.data.module.value.tools.map((t) => ({
-				id: t.id,
-				name: t.summary ?? ''
-			}))
+			}[] = node.data.module.value.tools.map((t, idx) => {
+				// Handle both FlowModule tools and MCP tools
+				const toolType =
+					t.value.tool_type === 'mcp'
+						? 'mcp'
+						: t.value.tool_type === 'flowmodule'
+							? t.value.type
+							: undefined
+				return {
+					id: t.id,
+					name: t.summary ?? '',
+					type: toolType
+				}
+			})

 			const agentActions = !insertable && flowModuleStates?.[node.id]?.agent_actions
 			if (agentActions) {
@@ -91,8 +106,11 @@
 				baseOffset = BELOW_ADDITIONAL_OFFSET + AI_TOOL_BASE_OFFSET
 				rowOffset = AI_TOOL_ROW_OFFSET
 				tools = agentActions.map((a, idx) => {
-					if (a.type === 'tool_call') {
-						const id = getToolCallId(idx, node.id, a.module_id)
+					if (a.type === 'tool_call' || a.type === 'mcp_tool_call') {
+						const id =
+							a.type === 'tool_call'
+								? getToolCallId(idx, node.id, a.module_id)
+								: AI_MCP_TOOL_CALL_PREFIX + '-' + node.id + '-' + idx
 						return {
 							id,
 							name: a.function_name
@@ -131,6 +149,7 @@
 					parentId: node.id,
 					data: {
 						tool: tool.name,
+						type: tool.type,
 						eventHandlers,
 						moduleId: tool.id,
 						insertable,
@@ -232,7 +251,7 @@
 		NewAiToolN,
 		NodeLayout
 	} from '../../graphBuilder.svelte'
-	import { MessageCircle, Play, Wrench, X } from 'lucide-svelte'
+	import { MessageCircle, Play, Plug, Wrench, X } from 'lucide-svelte'
 	import { twMerge } from 'tailwind-merge'
 	import { getContext } from 'svelte'
 	import type { Edge, Node } from '@xyflow/svelte'
@@ -274,8 +293,10 @@
 			>
 				{#if data.moduleId.startsWith(AI_TOOL_MESSAGE_PREFIX)}
 					<MessageCircle size={16} class="ml-1 shrink-0" />
-				{:else if data.moduleId.startsWith(AI_TOOL_CALL_PREFIX)}
+				{:else if data.moduleId.startsWith(AI_TOOL_CALL_PREFIX) || data.moduleId.startsWith(AI_MCP_TOOL_CALL_PREFIX)}
 					<Play size={16} class="ml-1 shrink-0" />
+				{:else if data.type === 'mcp'}
+					<Plug size={16} class="ml-1 shrink-0" />
 				{:else}
 					<Wrench size={16} class="ml-1 shrink-0" />
 				{/if}
@@ -283,10 +304,10 @@
 				<span
 					class={twMerge(
 						'text-3xs truncate flex-1',
-						!validateToolName(data.tool) && 'text-red-400'
+						!validateToolName(data.tool, data.type) && 'text-red-400'
 					)}
 				>
-					{data.tool || 'No tool name'}
+					{data.tool || 'Missing name'}
 				</span>
 			</button>
 			{#if data.insertable}
@@ -45,7 +45,7 @@
 	{#snippet children({ close })}
 		<InsertModuleInner
 			bind:funcDesc
-			scriptOnly
+			toolMode
 			on:close={() => {
 				close()
 			}}
@@ -79,6 +79,14 @@
 				})
 				close()
 			}}
+			on:pickMcpTool={(e) => {
+				data.eventHandlers.insert({
+					index: -1,
+					agentId: data.agentModuleId,
+					kind: 'mcpTool'
+				})
+				close()
+			}}
 		/>
 	{/snippet}
 </PopupV2>
@@ -0,0 +1,23 @@
+<script lang="ts">
+	interface Props {
+		height?: number
+		width?: number
+	}
+
+	let { height = 24, width = 24 }: Props = $props()
+</script>
+
+<svg
+	fill="currentColor"
+	fill-rule="evenodd"
+	{height}
+	style="flex:none;line-height:1"
+	viewBox="0 0 24 24"
+	{width}
+	xmlns="http://www.w3.org/2000/svg"
+	><title>ModelContextProtocol</title><path
+		d="M15.688 2.343a2.588 2.588 0 00-3.61 0l-9.626 9.44a.863.863 0 01-1.203 0 .823.823 0 010-1.18l9.626-9.44a4.313 4.313 0 016.016 0 4.116 4.116 0 011.204 3.54 4.3 4.3 0 013.609 1.18l.05.05a4.115 4.115 0 010 5.9l-8.706 8.537a.274.274 0 000 .393l1.788 1.754a.823.823 0 010 1.18.863.863 0 01-1.203 0l-1.788-1.753a1.92 1.92 0 010-2.754l8.706-8.538a2.47 2.47 0 000-3.54l-.05-.049a2.588 2.588 0 00-3.607-.003l-7.172 7.034-.002.002-.098.097a.863.863 0 01-1.204 0 .823.823 0 010-1.18l7.273-7.133a2.47 2.47 0 00-.003-3.537z"
+	></path><path
+		d="M14.485 4.703a.823.823 0 000-1.18.863.863 0 00-1.204 0l-7.119 6.982a4.115 4.115 0 000 5.9 4.314 4.314 0 006.016 0l7.12-6.982a.823.823 0 000-1.18.863.863 0 00-1.204 0l-7.119 6.982a2.588 2.588 0 01-3.61 0 2.47 2.47 0 010-3.54l7.12-6.982z"
+	></path></svg
+>
@@ -99,6 +99,7 @@ import XeroIcon from './XeroIcon.svelte'
 import KafkaIcon from './KafkaIcon.svelte'
 import NatsIcon from './NatsIcon.svelte'
 import MqttIcon from './MqttIcon.svelte'
+import McpIcon from './McpIcon.svelte'
 import SageIcon from './SageIcon.svelte'
 import ZohoIcon from './ZohoIcon.svelte'
 export const APP_TO_ICON_COMPONENT = {
@@ -206,6 +207,7 @@ export const APP_TO_ICON_COMPONENT = {
 	kafka: KafkaIcon,
 	nats: NatsIcon,
 	mqtt: MqttIcon,
+	mcp: McpIcon,
 	zoho: ZohoIcon
 } as const

@@ -305,5 +307,6 @@ export {
 	KafkaIcon,
 	NatsIcon,
 	MqttIcon,
+	McpIcon,
 	ZohoIcon
 }
@@ -91,9 +91,8 @@ export function updateFlowModuleById(
 				module.value.branches.forEach((branch) => dfs(branch.modules))
 			} else if (module.value.type === 'branchall') {
 				module.value.branches.forEach((branch) => dfs(branch.modules))
-			} else if (module.value.type === 'aiagent') {
-				dfs(module.value.tools)
 			}
+			// AI agent tools are leaf nodes - no traversal needed
 		}
 	}

@@ -619,28 +619,24 @@
 						<div
 							class="flex border border-gray-200 dark:border-gray-700 rounded-lg overflow-hidden flex-1"
 						>
-							<div class="flex-shrink-0">
-								<FlowConversationsSidebar
-									bind:this={flowConversationsSidebar}
-									flowPath={flow?.path ?? ''}
-									{selectedConversationId}
-									onNewConversation={handleNewConversation}
-									onSelectConversation={handleSelectConversation}
-									onDeleteConversation={handleDeleteConversation}
-								/>
-							</div>
-							<div class="flex-1">
-								<FlowChatInterface
-									bind:this={flowChatInterface}
-									onRunFlow={runFlowForChat}
-									useStreaming={shouldUseStreaming}
-									{refreshConversations}
-									conversationId={selectedConversationId}
-									{deploymentInProgress}
-									createConversation={handleNewConversation}
-									{path}
-								/>
-							</div>
+							<FlowConversationsSidebar
+								bind:this={flowConversationsSidebar}
+								flowPath={flow?.path ?? ''}
+								{selectedConversationId}
+								onNewConversation={handleNewConversation}
+								onSelectConversation={handleSelectConversation}
+								onDeleteConversation={handleDeleteConversation}
+							/>
+							<FlowChatInterface
+								bind:this={flowChatInterface}
+								onRunFlow={runFlowForChat}
+								useStreaming={shouldUseStreaming}
+								{refreshConversations}
+								conversationId={selectedConversationId}
+								{deploymentInProgress}
+								createConversation={handleNewConversation}
+								{path}
+							/>
 						</div>
 					{:else}
 						<!-- Normal Mode: Form Layout -->