diff --git a/backend/windmill-worker/src/ai_executor.rs b/backend/windmill-worker/src/ai_executor.rs
index 25638cbae9..806737ed19 100644
--- a/backend/windmill-worker/src/ai_executor.rs
+++ b/backend/windmill-worker/src/ai_executor.rs
@@ -118,7 +118,7 @@ struct Tool {
 #[derive(Deserialize, Debug)]
 struct AIAgentArgs {
     provider: Provider,
-    system_prompt: String,
+    system_prompt: Option<String>,
     user_message: String,
     temperature: Option<f32>,
     max_completion_tokens: Option<u32>,
@@ -608,18 +608,21 @@ async fn run_agent(
     hostname: &str,
     killpill_rx: &mut tokio::sync::broadcast::Receiver<()>,
 ) -> error::Result<Box<RawValue>> {
-    let mut messages = vec![
-        OpenAIMessage {
+    let mut messages = if let Some(system_prompt) = args.system_prompt.filter(|s| !s.is_empty()) {
+        vec![OpenAIMessage {
             role: "system".to_string(),
-            content: Some(args.system_prompt),
+            content: Some(system_prompt),
             ..Default::default()
-        },
-        OpenAIMessage {
-            role: "user".to_string(),
-            content: Some(args.user_message),
-            ..Default::default()
-        },
-    ];
+        }]
+    } else {
+        vec![]
+    };
+
+    messages.push(OpenAIMessage {
+        role: "user".to_string(),
+        content: Some(args.user_message),
+        ..Default::default()
+    });
 
     let mut actions = vec![];
 
diff --git a/frontend/src/lib/components/FlowStatusViewer.svelte b/frontend/src/lib/components/FlowStatusViewer.svelte
index 9b11554149..6687a04655 100644
--- a/frontend/src/lib/components/FlowStatusViewer.svelte
+++ b/frontend/src/lib/components/FlowStatusViewer.svelte
@@ -110,6 +110,9 @@
 	let updateGlobalRefresh = (moduleId: string, updateFn: (clear, root) => Promise<void>) => {
 		globalRefreshes[moduleId] = [...(globalRefreshes[moduleId] ?? []), updateFn]
 	}
+
+	let storedToolCallJobs: Record<string, Job> = $state({})
+	let toolCallIndicesToLoad: string[] = $state([])
 </script>
 
 <FlowStatusViewerInner
@@ -141,4 +144,30 @@
 	isNodeSelected={true}
 	{refreshGlobal}
 	{updateGlobalRefresh}
+	toolCallStore={{
+		getStoredToolCallJob: (storeKey: string) => storedToolCallJobs[storeKey],
+		setStoredToolCallJob: (storeKey: string, job: Job) => {
+			storedToolCallJobs[storeKey] = job
+		},
+		getLocalToolCallJobs: (prefix: string) => {
+			// we return a map from tool call index to job
+			// to do so, we filter the storedToolCallJobs object by the prefix and we make sure what's left in the key is a tool call index: 2 part of format agentModuleId-toolCallIndex
+			// and not a further nested tool call index
+			return Object.fromEntries(
+				Object.entries(storedToolCallJobs)
+					.filter(
+						([key]) => key.startsWith(prefix) && key.replace(prefix, '').split('-').length === 2
+					)
+					.map(([key, job]) => [Number(key.replace(prefix, '').split('-').pop()), job])
+			)
+		},
+		isToolCallToBeLoaded: (storeKey: string) => {
+			return toolCallIndicesToLoad.includes(storeKey)
+		},
+		addToolCallToLoad: (storeKey: string) => {
+			if (!toolCallIndicesToLoad.includes(storeKey)) {
+				toolCallIndicesToLoad.push(storeKey)
+			}
+		}
+	}}
 />
diff --git a/frontend/src/lib/components/FlowStatusViewerInner.svelte b/frontend/src/lib/components/FlowStatusViewerInner.svelte
index 5d00608245..a2012f73c3 100644
--- a/frontend/src/lib/components/FlowStatusViewerInner.svelte
+++ b/frontend/src/lib/components/FlowStatusViewerInner.svelte
@@ -33,6 +33,7 @@
 	import { deepEqual } from 'fast-equals'
 	import FlowTimeline from './FlowTimeline.svelte'
 	import { dfs } from './flows/dfs'
+	import { dfs as dfsPreviousResults } from '$lib/components/flows/previousResults'
 	import Alert from './common/alert/Alert.svelte'
 	import FlowGraphViewerStep from './FlowGraphViewerStep.svelte'
 	import FlowGraphV2 from './graph/FlowGraphV2.svelte'
@@ -116,6 +117,13 @@
 		onStart?: () => void
 		onJobsLoaded?: ({ job, force }: { job: Job; force: boolean }) => void
 		onDone?: ({ job }: { job: CompletedJob }) => void
+		toolCallStore?: {
+			getStoredToolCallJob: (storeKey: string) => Job | undefined
+			setStoredToolCallJob: (storeKey: string, job: Job) => void
+			getLocalToolCallJobs: (prefix: string) => Record<number, Job>
+			isToolCallToBeLoaded: (storeKey: string) => boolean
+			addToolCallToLoad: (storeKey: string) => void
+		}
 	}
 
 	let {
@@ -155,7 +163,8 @@
 		loadExtraLogs = undefined,
 		onStart = undefined,
 		onJobsLoaded = undefined,
-		onDone = undefined
+		onDone = undefined,
+		toolCallStore
 	}: Props = $props()
 
 	let getTopModuleStates = $derived(topModuleStates ?? localModuleStates)
@@ -913,9 +922,7 @@
 
 	let storedListJobs: Record<number, Job> = $state({})
 
-	let storedToolCallJobs: Record<number, Job> = $state({})
-	let selectedToolCall: number | undefined = $state(undefined)
-	let toolCallIndicesToLoad: number[] = $state([])
+	let selectedToolCall: string | undefined = $state(undefined)
 
 	let wrapperHeight: number = $state(0)
 
@@ -950,8 +957,10 @@
 					let nprefix = buildPrefix(prefix, oid)
 					return fms
 						? rec(
-								dfs(fms, (x) =>
-									x.id.startsWith('subflow:') ? x.id : buildSubflowKey(x.id, nprefix)
+								dfs(
+									fms,
+									(x) => (x.id.startsWith('subflow:') ? x.id : buildSubflowKey(x.id, nprefix)),
+									{ skipToolNodes: true }
 								),
 								nprefix
 							)
@@ -1009,6 +1018,11 @@
 				selectedForLoopSetManually: false
 			})
 		}
+		if (selectedNode?.startsWith(AI_TOOL_CALL_PREFIX)) {
+			const [, agentModuleId, toolCallIndex, _] = selectedNode.split('-')
+			const parentLoopsPrefix = getParentLoopsPrefix(agentModuleId)
+			toolCallStore?.addToolCallToLoad(parentLoopsPrefix + agentModuleId + '-' + toolCallIndex)
+		}
 	}
 
 	$effect(() => {
@@ -1039,6 +1053,29 @@
 	let animateLogsTab = $state(false)
 
 	let noLogs = $derived(graphTabOpen && !isNodeSelected)
+
+	/**
+	 * Returns a string like "forloopmodid1-{iter1}-forloopmodid2-{iter2}-forloopmodid3-{iter3}-"
+	 * that can be used to prefix tool call store keys for nested tool calls.
+	 */
+	function getParentLoopsPrefix(modId: string) {
+		if (job?.raw_flow) {
+			const indices: string[] = []
+			const parents = dfsPreviousResults(modId, { value: job?.raw_flow, summary: '' }, true)
+			for (const parent of parents) {
+				if (parent.value.type === 'forloopflow' || parent.value.type === 'whileloopflow') {
+					const state = localModuleStates[parent.id]
+					if (state?.selectedForloopIndex !== undefined) {
+						indices.push(parent.id + '-' + state.selectedForloopIndex.toString())
+					}
+				}
+			}
+			indices.reverse()
+			return indices.length > 0 ? indices.join('-') + '-' : ''
+		}
+
+		return ''
+	}
 </script>
 
 <JobLoader workspaceOverride={workspaceId} {noLogs} noCode bind:this={jobLoader} />
@@ -1173,6 +1210,10 @@
 							{@const forloopIsSelected =
 								forloop_selected == loopJobId ||
 								(innerModule?.type != 'forloopflow' && innerModule?.type != 'whileloopflow')}
+							{@const forLoopStoreKeyPrefix =
+								innerModule?.type == 'forloopflow' || innerModule?.type == 'whileloopflow'
+									? (flowJobIds?.moduleId ?? '') + '-' + j + '-'
+									: ''}
 							<!-- <LogId id={loopJobId} /> -->
 							<div class="border p-6" class:hidden={forloop_selected != loopJobId}>
 								<FlowStatusViewerInner
@@ -1207,6 +1248,18 @@
 									graphTabOpen={selected == 'graph' && graphTabOpen}
 									isNodeSelected={forloop_selected == loopJobId}
 									{globalIterationBounds}
+									toolCallStore={{
+										getStoredToolCallJob: (storeKey: string) =>
+											toolCallStore?.getStoredToolCallJob(forLoopStoreKeyPrefix + storeKey),
+										setStoredToolCallJob: (storeKey: string, job: Job) =>
+											toolCallStore?.setStoredToolCallJob(forLoopStoreKeyPrefix + storeKey, job),
+										getLocalToolCallJobs: (prefix: string) =>
+											toolCallStore?.getLocalToolCallJobs(forLoopStoreKeyPrefix + prefix) ?? {},
+										addToolCallToLoad: (storeKey: string) =>
+											toolCallStore?.addToolCallToLoad(forLoopStoreKeyPrefix + storeKey),
+										isToolCallToBeLoaded: (storeKey: string) =>
+											toolCallStore?.isToolCallToBeLoaded(forLoopStoreKeyPrefix + storeKey) ?? false
+									}}
 								/>
 							</div>
 						{/if}
@@ -1366,12 +1419,17 @@
 										graphTabOpen={selected == 'graph' && graphTabOpen}
 										isNodeSelected={localModuleStates?.[selectedNode ?? '']?.job_id == mod.job}
 										{globalIterationBounds}
+										{toolCallStore}
 									/>
-									{#if mod.agent_actions && mod.agent_actions.length > 0}
+									{#if mod.agent_actions && mod.agent_actions.length > 0 && mod.id}
+										{@const storeKeyPrefix = getParentLoopsPrefix(mod.id)}
 										{#each mod.agent_actions as agentAction, j}
-											{#if agentAction.type === 'tool_call' && mod.id}
+											{#if agentAction.type === 'tool_call'}
 												{@const toolCallId = getToolCallId(j, mod.id, agentAction.module_id)}
-												{@const isSelected = selectedToolCall === j}
+												{@const localToolCallKey = mod.id + '-' + j}
+												{@const storeKey = storeKeyPrefix + localToolCallKey}
+												{@const storedToolCallJob = toolCallStore?.getStoredToolCallJob(storeKey)}
+												{@const isSelected = localToolCallKey === selectedToolCall}
 												<Button
 													variant={isSelected ? 'contained' : 'border'}
 													color={mod.agent_actions_success?.[j] === false
@@ -1381,10 +1439,10 @@
 															: 'light'}
 													btnClasses="w-full flex justify-start"
 													on:click={async () => {
-														if (selectedToolCall == j) {
+														if (isSelected) {
 															selectedToolCall = undefined
 														} else {
-															selectedToolCall = j
+															selectedToolCall = localToolCallKey
 														}
 													}}
 													endIcon={{
@@ -1396,7 +1454,7 @@
 														Tool call: {agentAction.function_name}
 													</span>
 												</Button>
-												{#if isSelected || storedToolCallJobs[j] || toolCallIndicesToLoad.includes(j)}
+												{#if isSelected || storedToolCallJob || toolCallStore?.isToolCallToBeLoaded(storeKey)}
 													<FlowStatusViewerInner
 														topModuleStates={getTopModuleStates}
 														{refreshGlobal}
@@ -1414,11 +1472,11 @@
 														{subflowParentsDurationStatuses}
 														{isSelectedBranch}
 														jobId={agentAction.job_id}
-														job={storedToolCallJobs[j]}
-														initialJob={storedToolCallJobs[j]}
+														job={storedToolCallJob}
+														initialJob={storedToolCallJob}
 														{reducedPolling}
 														onJobsLoaded={({ job, force }) => {
-															storedToolCallJobs[j] = job
+															toolCallStore?.setStoredToolCallJob(storeKey, job)
 															onJobsLoadedInner({ id: toolCallId } as FlowStatusModule, job, force)
 														}}
 														loadExtraLogs={(logs) => {
@@ -1509,11 +1567,11 @@
 										stepDetail = mod
 										selectedNode = e
 										if (e.startsWith(AI_TOOL_CALL_PREFIX)) {
-											const [_prefix, _agentModuleId, j, _toolModuleId] = e.split('-')
+											const [_prefix, agentModuleId, j, _toolModuleId] = e.split('-')
+											const parentLoopsPrefix = getParentLoopsPrefix(agentModuleId)
 											const jIdx = Number(j)
-											if (!toolCallIndicesToLoad.includes(jIdx)) {
-												toolCallIndicesToLoad.push(jIdx)
-											}
+											const storeKey = parentLoopsPrefix + agentModuleId + '-' + jIdx
+											toolCallStore?.addToolCallToLoad(storeKey)
 										}
 									}
 								} else {
@@ -1603,6 +1661,7 @@
 											stepDetail && typeof stepDetail !== 'string' ? stepDetail : undefined}
 										{@const agentTools =
 											module && module.value.type === 'aiagent' ? module.value.tools : undefined}
+										{@const parentLoopsPrefix = getParentLoopsPrefix(module?.id ?? '')}
 										{#if node.flow_jobs_results}
 											<span class="pl-1 text-tertiary"
 												>Result of step as collection of all subflows</span
@@ -1667,7 +1726,7 @@
 											logs={node.logs}
 											downloadLogs={!hideDownloadLogs}
 											aiAgentStatus={agentTools &&
-											node.job_id &&
+											node?.job_id &&
 											(node.type === 'Success' || node.type === 'Failure')
 												? {
 														tools: agentTools,
@@ -1679,9 +1738,14 @@
 															success: node.type === 'Success',
 															type: 'CompletedJob'
 														},
-														storedToolCallJobs,
+														storedToolCallJobs: module
+															? toolCallStore?.getLocalToolCallJobs(parentLoopsPrefix)
+															: undefined,
 														onToolJobLoaded: (job, idx) => {
-															storedToolCallJobs[idx] = job
+															if (module) {
+																const storeKey = parentLoopsPrefix + module.id + '-' + idx
+																toolCallStore?.setStoredToolCallJob(storeKey, job)
+															}
 														}
 													}
 												: undefined}
diff --git a/frontend/src/lib/components/flows/FlowModuleIcon.svelte b/frontend/src/lib/components/flows/FlowModuleIcon.svelte
index ecadb9ea4f..02fe81047a 100644
--- a/frontend/src/lib/components/flows/FlowModuleIcon.svelte
+++ b/frontend/src/lib/components/flows/FlowModuleIcon.svelte
@@ -20,19 +20,19 @@
 </script>
 
 {#if module.value.type === 'aiagent'}
-	<Bot size={16} />
+	<Bot size={16} class="text-violet-800 dark:text-violet-400" />
 {:else if module.value.type === 'rawscript'}
 	<LanguageIcon lang={module.value.language} width={iconWidth} height={iconHeight} />
 {:else if module.summary === 'Terminate flow'}
-	<Square size={size} />
+	<Square {size} />
 {:else if module.value.type === 'identity'}
-	<ArrowDown size={size} />
+	<ArrowDown {size} />
 {:else if module.value.type === 'flow'}
-	<BarsStaggered size={size} />
+	<BarsStaggered {size} />
 {:else if module.value.type === 'forloopflow' || module.value.type === 'whileloopflow'}
-	<Repeat size={size} />
+	<Repeat {size} />
 {:else if module.value.type === 'branchone' || module.value.type === 'branchall'}
-	<GitBranch size={size} />
+	<GitBranch {size} />
 {:else if module.value.type === 'script'}
 	{#if module.value.path.startsWith('hub/')}
 		<IconedResourceType
@@ -42,9 +42,9 @@
 			silent={true}
 		/>
 	{:else}
-		<Building size={size} />
+		<Building {size} />
 	{/if}
 {:else}
 	<!-- Fallback icon for unknown module types -->
-	<BarsStaggered size={size} />
+	<BarsStaggered {size} />
 {/if}
diff --git a/frontend/src/lib/components/flows/common/FlowCardHeader.svelte b/frontend/src/lib/components/flows/common/FlowCardHeader.svelte
index 07d1a0eb37..39e3f615fd 100644
--- a/frontend/src/lib/components/flows/common/FlowCardHeader.svelte
+++ b/frontend/src/lib/components/flows/common/FlowCardHeader.svelte
@@ -154,6 +154,9 @@
 				{:else if flowModuleValue.type === 'flow'}
 					<Badge color="indigo" capitalize>flow</Badge>
 					<input bind:value={summary} placeholder="Summary" class="w-full grow" />
+				{:else if flowModuleValue.type === 'aiagent'}
+					<Badge color="indigo">AI Agent</Badge>
+					<input bind:value={summary} placeholder="Summary" class="w-full grow" />
 				{/if}
 			</div>
 		</span>
diff --git a/frontend/src/lib/components/flows/dfs.ts b/frontend/src/lib/components/flows/dfs.ts
index 20b487ac72..3e8709c206 100644
--- a/frontend/src/lib/components/flows/dfs.ts
+++ b/frontend/src/lib/components/flows/dfs.ts
@@ -3,29 +3,29 @@ import type { FlowModule } from '$lib/gen'
 export function dfs<T>(
 	modules: FlowModule[],
 	f: (x: FlowModule, modules: FlowModule[], branches: FlowModule[][]) => T,
-	{ skipToolNodes = false }: { skipToolNodes?: boolean } = {}
+	opts: { skipToolNodes?: boolean } = {}
 ): T[] {
 	let result: T[] = []
 	for (const module of modules) {
 		if (module.value.type == 'forloopflow' || module.value.type == 'whileloopflow') {
 			result = result.concat(f(module, modules, [module.value.modules]))
-			result = result.concat(dfs(module.value.modules, f))
+			result = result.concat(dfs(module.value.modules, f, opts))
 		} else if (module.value.type == 'branchone') {
 			const allBranches = [module.value.default, ...module.value.branches.map((b) => b.modules)]
 			result = result.concat(f(module, modules, allBranches))
 
 			for (const branch of allBranches) {
-				result = result.concat(dfs(branch, f))
+				result = result.concat(dfs(branch, f, opts))
 			}
 		} else if (module.value.type == 'branchall') {
 			const allBranches = module.value.branches.map((b) => b.modules)
 			result = result.concat(f(module, modules, allBranches))
 			for (const branch of allBranches) {
-				result = result.concat(dfs(branch, f))
+				result = result.concat(dfs(branch, f, opts))
 			}
-		} else if (module.value.type == 'aiagent' && !skipToolNodes) {
+		} 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))
+			result = result.concat(dfs(module.value.tools, f, opts))
 		} else {
 			result.push(f(module, modules, []))
 		}
diff --git a/frontend/src/lib/components/flows/flowInfers.ts b/frontend/src/lib/components/flows/flowInfers.ts
index 00b6c608c1..887fe6e418 100644
--- a/frontend/src/lib/components/flows/flowInfers.ts
+++ b/frontend/src/lib/components/flows/flowInfers.ts
@@ -96,13 +96,12 @@ export async function loadSchemaFromModule(module: FlowModule): Promise<{
 						}
 					]
 				},
-				system_prompt: {
-					type: 'string',
-					default: 'You are a helpful assistant'
-				},
 				user_message: {
 					type: 'string'
 				},
+				system_prompt: {
+					type: 'string'
+				},
 				max_completion_tokens: {
 					type: 'number'
 				},
@@ -110,13 +109,13 @@ export async function loadSchemaFromModule(module: FlowModule): Promise<{
 					type: 'number'
 				}
 			},
-			required: ['provider', 'model', 'system_prompt', 'user_message'],
+			required: ['provider', 'model', 'user_message'],
 			type: 'object',
 			order: [
 				'provider',
 				'model',
-				'system_prompt',
 				'user_message',
+				'system_prompt',
 				'max_completion_tokens',
 				'temperature'
 			]
diff --git a/frontend/src/lib/components/flows/flowStateUtils.svelte.ts b/frontend/src/lib/components/flows/flowStateUtils.svelte.ts
index 5b2e91320d..c27bae97f3 100644
--- a/frontend/src/lib/components/flows/flowStateUtils.svelte.ts
+++ b/frontend/src/lib/components/flows/flowStateUtils.svelte.ts
@@ -165,8 +165,7 @@ export async function createBranchAll(id: string): Promise<[FlowModule, FlowModu
 export async function createAiAgent(id: string): Promise<[FlowModule, FlowModuleState]> {
 	const aiAgentFlowModules: FlowModule = {
 		id,
-		value: { type: 'aiagent', tools: [], input_transforms: {} },
-		summary: 'AI Agent'
+		value: { type: 'aiagent', tools: [], input_transforms: {} }
 	}
 
 	const flowModuleState = await loadFlowModuleState(aiAgentFlowModules)
diff --git a/frontend/src/lib/components/flows/map/MapItem.svelte b/frontend/src/lib/components/flows/map/MapItem.svelte
index d19eaaaedf..c6483e1fae 100644
--- a/frontend/src/lib/components/flows/map/MapItem.svelte
+++ b/frontend/src/lib/components/flows/map/MapItem.svelte
@@ -248,6 +248,7 @@
 					{bgColor}
 					{bgHoverColor}
 					label={mod.summary ||
+						(mod.value.type === 'aiagent' ? 'AI Agent' : undefined) ||
 						(mod.id === 'preprocessor'
 							? 'Preprocessor'
 							: mod.id.startsWith('failure')
@@ -272,10 +273,13 @@
 					{skipped}
 				>
 					{#snippet icon()}
-						{@const size = mod.value.type === 'script' && mod.value.path.startsWith('hub/')
-							? 20
-							: mod.value.type === "script" ? 14 : 16}
-						<FlowModuleIcon module={mod} size={size} />
+						{@const size =
+							mod.value.type === 'script' && mod.value.path.startsWith('hub/')
+								? 20
+								: mod.value.type === 'script'
+									? 14
+									: 16}
+						<FlowModuleIcon module={mod} {size} />
 					{/snippet}
 				</FlowModuleSchemaItem>
 			{/if}
diff --git a/frontend/src/lib/components/flows/pickers/TopLevelNode.svelte b/frontend/src/lib/components/flows/pickers/TopLevelNode.svelte
index c00ed14d5c..ab6746515d 100644
--- a/frontend/src/lib/components/flows/pickers/TopLevelNode.svelte
+++ b/frontend/src/lib/components/flows/pickers/TopLevelNode.svelte
@@ -63,7 +63,7 @@
 			<GitBranch size={14} />
 			Branch to all
 		{:else if label === 'AI Agent'}
-			<BotIcon size={14} />
+			<BotIcon size={14} class="text-violet-800 dark:text-violet-400" />
 			AI Agent
 		{/if}
 	</span>
diff --git a/frontend/src/lib/components/graph/renderers/nodes/AIToolNode.svelte b/frontend/src/lib/components/graph/renderers/nodes/AIToolNode.svelte
index 8900e471f1..36235add24 100644
--- a/frontend/src/lib/components/graph/renderers/nodes/AIToolNode.svelte
+++ b/frontend/src/lib/components/graph/renderers/nodes/AIToolNode.svelte
@@ -11,7 +11,7 @@
 	export const AI_TOOL_MESSAGE_PREFIX = '_wm_ai_agent_message'
 
 	const ROW_WIDTH = 275
-	const NEW_TOOL_NODE_WIDTH = 40
+	const NEW_TOOL_NODE_WIDTH = 50
 	const MAX_TOOLS_PER_ROW = 2
 
 	let computeAIToolNodesCache:
@@ -140,13 +140,13 @@
 					width: inputToolWidth,
 					position: {
 						x:
-							tools.length === 1
+							(tools.length === 1
 								? (ROW_WIDTH - inputToolWidth) / 2
 								: (i + 1) % 2 === 0
 									? inputToolWidth + inputToolXGap
 									: isLastRow && tools.length % 2 === 1
 										? (ROW_WIDTH - inputToolWidth) / 2
-										: 0,
+										: 0) + node.data.offset,
 						y:
 							baseOffset +
 							rowOffset *
@@ -176,7 +176,7 @@
 					parentId: node.id,
 					width: NEW_TOOL_NODE_WIDTH,
 					position: {
-						x: (ROW_WIDTH - NEW_TOOL_NODE_WIDTH) / 2,
+						x: (ROW_WIDTH - NEW_TOOL_NODE_WIDTH) / 2 + node.data.offset,
 						y: baseOffset + rowOffset
 					}
 				} satisfies Node & NewAiToolN)
diff --git a/frontend/src/lib/components/graph/renderers/nodes/NewAIToolNode.svelte b/frontend/src/lib/components/graph/renderers/nodes/NewAIToolNode.svelte
index b9f3a900d2..cb30827dcb 100644
--- a/frontend/src/lib/components/graph/renderers/nodes/NewAIToolNode.svelte
+++ b/frontend/src/lib/components/graph/renderers/nodes/NewAIToolNode.svelte
@@ -1,69 +1,89 @@
 <script lang="ts">
+	import { preventDefault, stopPropagation } from 'svelte/legacy'
+
 	import { twMerge } from 'tailwind-merge'
 	import { type NewAiToolN } from '../../graphBuilder.svelte'
-	import NodeWrapper from './NodeWrapper.svelte'
-	import Popover from '$lib/components/meltComponents/Popover.svelte'
 	import InsertModuleInner from '$lib/components/flows/map/InsertModuleInner.svelte'
+	import { Cross } from 'lucide-svelte'
+	import PopupV2 from '$lib/components/common/popup/PopupV2.svelte'
+	import { flip, offset } from 'svelte-floating-ui/dom'
+	import type { ComputeConfig } from 'svelte-floating-ui'
 
 	let funcDesc = $state('')
 	interface Props {
 		data: NewAiToolN['data']
 	}
 	let { data }: Props = $props()
+
+	let floatingConfig: ComputeConfig = {
+		strategy: 'fixed',
+		// @ts-ignore
+		placement: 'bottom-center',
+		middleware: [offset(8), flip()],
+		autoUpdate: true
+	}
 </script>
 
-<NodeWrapper>
-	{#snippet children({ darkMode })}
-		<!-- svelte-ignore a11y_no_static_element_interactions -->
-		<Popover
-			portal={null}
-			usePointerDownOutside
+<!-- svelte-ignore a11y_no_static_element_interactions -->
+<PopupV2 {floatingConfig} target="#flow-editor">
+	{#snippet button({ pointerdown, pointerup })}
+		<button
+			title={`Add 'tool'
+			}`}
+			type="button"
 			class={twMerge(
-				'!w-full text-2xs font-normal bg-surface h-6 pr-0.5 flex justify-center items-center rounded-sm text-tertiary border',
-				'hover:bg-surface-hover'
+				'!w-full h-6 flex items-center justify-center !outline-[1px] outline dark:outline-gray-500 outline-gray-300 text-secondary bg-surface focus:outline-none hover:bg-surface-hover rounded'
 			)}
-			placement="top"
+			onpointerdown={stopPropagation(
+				preventDefault(() => {
+					pointerdown()
+				})
+			)}
+			onpointerup={pointerup}
 		>
-			<svelte:fragment slot="trigger">+tool</svelte:fragment>
-			<svelte:fragment slot="content" let:close>
-				<InsertModuleInner
-					bind:funcDesc
-					scriptOnly
-					on:close={() => {
-						close()
-					}}
-					on:new={(e) => {
-						data.eventHandlers.insert({
-							index: -1, // ignored when agentId is set
-							agentId: data.agentModuleId,
-							...e.detail
-						})
-						close()
-					}}
-					on:insert={(e) => {
-						data.eventHandlers.insert({
-							index: -1, // ignored when agentId is set
-							agentId: data.agentModuleId,
-							...e.detail
-						})
-						close()
-					}}
-					on:pickScript={(e) => {
-						data.eventHandlers.insert({
-							index: -1, // ignored when agentId is set
-							agentId: data.agentModuleId,
-							kind: e.detail.kind,
-							script: {
-								...e.detail,
-								summary: e.detail.summary
-									? e.detail.summary.replace(/\s/, '_').replace(/[^a-zA-Z0-9_]/g, '')
-									: e.detail.path.split('/').pop()
-							}
-						})
-						close()
-					}}
-				/>
-			</svelte:fragment>
-		</Popover>
+			<div class="flex flex-row items-center gap-1 font-medium text-2xs">
+				<Cross size={12} />
+				tool
+			</div>
+		</button>
 	{/snippet}
-</NodeWrapper>
+	{#snippet children({ close })}
+		<InsertModuleInner
+			bind:funcDesc
+			scriptOnly
+			on:close={() => {
+				close()
+			}}
+			on:new={(e) => {
+				data.eventHandlers.insert({
+					index: -1, // ignored when agentId is set
+					agentId: data.agentModuleId,
+					...e.detail
+				})
+				close()
+			}}
+			on:insert={(e) => {
+				data.eventHandlers.insert({
+					index: -1, // ignored when agentId is set
+					agentId: data.agentModuleId,
+					...e.detail
+				})
+				close()
+			}}
+			on:pickScript={(e) => {
+				data.eventHandlers.insert({
+					index: -1, // ignored when agentId is set
+					agentId: data.agentModuleId,
+					kind: e.detail.kind,
+					script: {
+						...e.detail,
+						summary: e.detail.summary
+							? e.detail.summary.replace(/\s/, '_').replace(/[^a-zA-Z0-9_]/g, '')
+							: e.detail.path.split('/').pop()
+					}
+				})
+				close()
+			}}
+		/>
+	{/snippet}
+</PopupV2>