feat(aichat): allow adding contexts to flow mode (#6424)

* add new feature instructions

* add db as context for flow mode

* add diff

* cleaner diff

* add modules as available context

* convert to svelte 5

* auto add selected module to context

* change flowinline ai button + nits

* handle adding selected lines

* clean context handling

* apply code pieces

* new chat when changing mode

* clean

* show code for code steps

* add last saved flow

* fix size

* categorize context

* optionnaly categorize

* fix module finding

* logs

* nit prompt

* fix

* fix

* fix test tool for script

* clean

CLAUDE.md

@@ -4,6 +4,17 @@
 
 Windmill is an open-source developer platform for building internal tools, workflows, API integrations, background jobs, workflows, and user interfaces. See @windmill-overview.mdc for full platform details.
 
+## New Feature Implementation Guidelines
+
+When implementing new features in Windmill, follow these best practices:
+
+- **Clean Code First**: Write clean, readable, and maintainable code. Prioritize clarity over cleverness.
+- **Avoid Duplication at All Costs**: Before writing new code, thoroughly search for existing implementations that can be reused or extended.
+- **Adapt Existing Code**: Refactor and generalize existing code when necessary to avoid logic duplication. Extract common patterns into reusable utilities.
+- **Follow Established Patterns**: Study existing code patterns in the codebase and maintain consistency with established conventions.
+- **Single Responsibility**: Each function, component, and module should have a single, well-defined responsibility.
+- **Incremental Implementation**: Break large features into smaller, reviewable chunks that can be implemented and tested incrementally.
+
 ## Language-Specific Guides
 
 - Backend (Rust): @backend/rust-best-practices.mdc + @backend/summarized_schema.txt

frontend/src/lib/components/Editor.svelte

@@ -184,6 +184,7 @@
 		loadAsync?: boolean
 		key?: string | undefined
 		class?: string | undefined
+		moduleId?: string
 	}
 
 	let {
@@ -209,7 +210,8 @@
 		changeTimeout = 500,
 		loadAsync = false,
 		key = undefined,
-		class: clazz = undefined
+		class: clazz = undefined,
+		moduleId = undefined
 	}: Props = $props()
 
 	$effect.pre(() => {
@@ -1328,7 +1330,8 @@
 					aiChatManager.addSelectedLinesToContext(
 						selectedLines,
 						selection.startLineNumber,
-						selection.endLineNumber
+						selection.endLineNumber,
+						moduleId
 					)
 				} else {
 					aiChatManager.toggleOpen()

frontend/src/lib/components/EditorBar.svelte

@@ -82,6 +82,7 @@
 		showHistoryDrawer?: boolean
 		right?: import('svelte').Snippet
 		openAiChat?: boolean
+		moduleId?: string
 	}
 
 	let {
@@ -105,7 +106,8 @@
 		diffMode = false,
 		showHistoryDrawer = $bindable(false),
 		right,
-		openAiChat = false
+		openAiChat = false,
+		moduleId = undefined
 	}: Props = $props()
 
 	let contextualVariablePicker: ItemPicker | undefined = $state()
@@ -964,7 +966,7 @@ JsonNode ${windmillPathToCamelCaseName(path)} = JsonNode.Parse(await client.GetS
 
 			{#if customUi?.aiGen != false}
 				{#if openAiChat}
-					<FlowInlineScriptAiButton />
+					<FlowInlineScriptAiButton {moduleId} />
 				{:else}
 					<ScriptGen {editor} {diffEditor} {lang} {iconOnly} {args} />
 				{/if}

frontend/src/lib/components/ScriptEditor.svelte

@@ -269,6 +269,7 @@
 	onMount(() => {
 		inferSchema(code)
 		loadPastTests()
+		aiChatManager.saveAndClear()
 		aiChatManager.changeMode(AIMode.SCRIPT)
 	})
 

frontend/src/lib/components/copilot/FlowInlineScriptAIButton.svelte

@@ -8,6 +8,12 @@
 	import { twMerge } from 'tailwind-merge'
 	import { aiChatManager, AIMode } from './chat/AIChatManager.svelte'
 
+	interface Props {
+		moduleId?: string
+	}
+
+	const { moduleId }: Props = $props()
+
 	const aiChatScriptModeClasses = $derived(
 		aiChatManager.mode === AIMode.SCRIPT && aiChatManager.isOpen
 			? 'dark:bg-violet-900 bg-violet-100'
@@ -22,7 +28,7 @@
 		btnClasses={twMerge('!px-2', aiChatScriptModeClasses)}
 		{onClick}
 		iconOnly
-		title="Open AI chat in script mode"
+		title="Open AI chat"
 		startIcon={{ icon: WandSparkles, classes: 'text-violet-800 dark:text-violet-400' }}
 	/>
 {/snippet}
@@ -30,7 +36,8 @@
 {#if $copilotInfo.enabled}
 	{@render button(() => {
 		aiChatManager.openChat()
-		aiChatManager.changeMode(AIMode.SCRIPT)
+		const availableContext = aiChatManager.contextManager.getAvailableContext()
+		aiChatManager.contextManager.setSelectedModuleContext(moduleId, availableContext)
 	})}
 {:else}
 	<Popover

frontend/src/lib/components/copilot/chat/AIChat.svelte

@@ -77,7 +77,7 @@
 	})
 
 	$effect(() => {
-		aiChatManager.listenForScriptEditorContextChange(
+		aiChatManager.listenForContextChange(
 			$dbSchemas,
 			$workspaceStore,
 			$copilotSessionModel
@@ -115,9 +115,7 @@
 	pastChats={historyManager.getPastChats()}
 	bind:selectedContext={
 		() => aiChatManager.contextManager.getSelectedContext(),
-		(sc) => {
-			aiChatManager.scriptEditorOptions && aiChatManager.contextManager.setSelectedContext(sc)
-		}
+		(sc) => aiChatManager.contextManager.setSelectedContext(sc)
 	}
 	availableContext={aiChatManager.contextManager.getAvailableContext()}
 	messages={aiChatManager.currentReply

frontend/src/lib/components/copilot/chat/AIChatInput.svelte

@@ -52,7 +52,7 @@
 		if (placeholder) {
 			return placeholder
 		}
-		
+
 		switch (aiChatManager.mode) {
 			case AIMode.SCRIPT:
 				return 'Modify this script...'
@@ -74,7 +74,7 @@
 	let instructions = $state(initialInstructions)
 
 	export function focusInput() {
-		if (aiChatManager.mode === AIMode.SCRIPT) {
+		if (aiChatManager.mode === AIMode.SCRIPT || aiChatManager.mode === AIMode.FLOW) {
 			contextTextareaComponent?.focus()
 		} else {
 			instructionsTextareaComponent?.focus()
@@ -132,7 +132,7 @@
 </script>
 
 <div use:clickOutside class="relative">
-	{#if aiChatManager.mode === AIMode.SCRIPT}
+	{#if aiChatManager.mode === AIMode.SCRIPT || aiChatManager.mode === AIMode.FLOW}
 		{#if showContext}
 			<div class="flex flex-row gap-1 mb-1 overflow-scroll pt-2 no-scrollbar">
 				<Popover>
@@ -150,6 +150,7 @@
 								addContextToSelection(element)
 								close()
 							}}
+							categorize
 						/>
 					</svelte:fragment>
 				</Popover>
@@ -157,7 +158,7 @@
 					<ContextElementBadge
 						contextElement={element}
 						deletable
-						on:delete={() => {
+						onDelete={() => {
 							selectedContext = selectedContext?.filter(
 								(c) => c.type !== element.type || c.title !== element.title
 							)

frontend/src/lib/components/copilot/chat/AIChatManager.svelte.ts

@@ -1,5 +1,5 @@
 import type { AIProviderModel, ScriptLang } from '$lib/gen/types.gen'
-import type { ScriptOptions } from './ContextManager.svelte'
+import type { FlowOptions, ScriptOptions } from './ContextManager.svelte'
 import {
 	flowTools,
 	prepareFlowSystemMessage,
@@ -88,6 +88,7 @@ class AIChatManager {
 	helpers = $state<any | undefined>(undefined)
 
 	scriptEditorOptions = $state<ScriptOptions | undefined>(undefined)
+	flowOptions = $state<FlowOptions | undefined>(undefined)
 	scriptEditorApplyCode = $state<((code: string, applyAll?: boolean) => void) | undefined>(
 		undefined
 	)
@@ -100,7 +101,7 @@ class AIChatManager {
 	private confirmationCallback = $state<((value: boolean) => void) | undefined>(undefined)
 
 	allowedModes: Record<AIMode, boolean> = $derived({
-		script: this.scriptEditorOptions !== undefined,
+		script: this.flowAiChatHelpers === undefined && this.scriptEditorOptions !== undefined,
 		flow: this.flowAiChatHelpers !== undefined,
 		navigator: true,
 		ask: true,
@@ -127,7 +128,7 @@ class AIChatManager {
 		return (
 			estimatedTokens >
 			modelContextWindow -
-			Math.max(modelContextWindow * MAX_TOKENS_THRESHOLD_PERCENTAGE, MAX_TOKENS_HARD_LIMIT)
+				Math.max(modelContextWindow * MAX_TOKENS_THRESHOLD_PERCENTAGE, MAX_TOKENS_HARD_LIMIT)
 		)
 	}
 
@@ -557,8 +558,8 @@ class AIChatManager {
 					onNewToken: (token: string) => {
 						reply += token
 					},
-					onMessageEnd: () => { },
-					setToolStatus: () => { }
+					onMessageEnd: () => {},
+					setToolStatus: () => {}
 				},
 				systemMessage
 			}
@@ -625,7 +626,7 @@ class AIChatManager {
 		}
 		try {
 			const oldSelectedContext = this.contextManager?.getSelectedContext() ?? []
-			if (this.mode === AIMode.SCRIPT) {
+			if (this.mode === AIMode.SCRIPT || this.mode === AIMode.FLOW) {
 				this.contextManager?.updateContextOnRequest(options)
 			}
 			this.loading = true
@@ -648,7 +649,10 @@ class AIChatManager {
 				{
 					role: 'user',
 					content: this.instructions,
-					contextElements: this.mode === AIMode.SCRIPT ? oldSelectedContext : undefined,
+					contextElements:
+						this.mode === AIMode.SCRIPT || this.mode === AIMode.FLOW
+							? oldSelectedContext
+							: undefined,
 					snapshot,
 					index: this.messages.length // matching with actual messages index. not -1 because it's not yet added to the messages array
 				}
@@ -672,7 +676,8 @@ class AIChatManager {
 				case AIMode.FLOW:
 					userMessage = prepareFlowUserMessage(
 						oldInstructions,
-						this.flowAiChatHelpers!.getFlowAndSelectedId()
+						this.flowAiChatHelpers!.getFlowAndSelectedId(),
+						oldSelectedContext
 					)
 					break
 				case AIMode.NAVIGATOR:
@@ -823,12 +828,19 @@ class AIChatManager {
 		this.sendRequest()
 	}
 
-	addSelectedLinesToContext = (lines: string, startLine: number, endLine: number) => {
+	addSelectedLinesToContext = (
+		lines: string,
+		startLine: number,
+		endLine: number,
+		moduleId?: string
+	) => {
 		if (!this.open) {
 			this.toggleOpen()
 		}
-		this.changeMode(AIMode.SCRIPT)
-		this.contextManager?.addSelectedLinesToContext(lines, startLine, endLine)
+		if (!moduleId) {
+			this.changeMode(AIMode.SCRIPT)
+		}
+		this.contextManager?.addSelectedLinesToContext(lines, startLine, endLine, moduleId)
 		this.focusInput()
 	}
 
@@ -869,12 +881,12 @@ class AIChatManager {
 		})
 	}
 
-	listenForScriptEditorContextChange = (
+	listenForContextChange = (
 		dbSchemas: DBSchemas,
 		workspaceStore: string | undefined,
 		copilotSessionModel: AIProviderModel | undefined
 	) => {
-		if (this.scriptEditorOptions) {
+		if (this.mode === AIMode.SCRIPT && this.scriptEditorOptions) {
 			this.contextManager.updateAvailableContext(
 				this.scriptEditorOptions,
 				dbSchemas,
@@ -882,6 +894,18 @@ class AIChatManager {
 				!copilotSessionModel?.model.endsWith('/thinking'),
 				untrack(() => this.contextManager.getSelectedContext())
 			)
+		} else if (this.mode === AIMode.FLOW && this.flowOptions) {
+			this.contextManager.updateAvailableContextForFlow(
+				this.flowOptions,
+				dbSchemas,
+				workspaceStore ?? '',
+				!copilotSessionModel?.model.endsWith('/thinking'),
+				untrack(() => this.contextManager.getSelectedContext())
+			)
+		}
+
+		if (this.scriptEditorOptions) {
+			this.contextManager.setScriptOptions(this.scriptEditorOptions)
 		}
 	}
 
@@ -941,15 +965,15 @@ class AIChatManager {
 				const editorRelated =
 					currentEditor && currentEditor.type === 'script' && currentEditor.stepId === module.id
 						? {
-							diffMode: currentEditor.diffMode,
-							lastDeployedCode: currentEditor.lastDeployedCode,
-							lastSavedCode: undefined
-						}
+								diffMode: currentEditor.diffMode,
+								lastDeployedCode: currentEditor.lastDeployedCode,
+								lastSavedCode: undefined
+							}
 						: {
-							diffMode: false,
-							lastDeployedCode: undefined,
-							lastSavedCode: undefined
-						}
+								diffMode: false,
+								lastDeployedCode: undefined,
+								lastSavedCode: undefined
+							}
 
 				return {
 					args: moduleState?.previewArgs ?? {},
@@ -976,6 +1000,13 @@ class AIChatManager {
 			this.scriptEditorOptions = undefined
 		}
 
+		untrack(() =>
+			this.contextManager?.setSelectedModuleContext(
+				selectedId,
+				untrack(() => this.contextManager.getAvailableContext())
+			)
+		)
+
 		return () => {
 			this.scriptEditorOptions = undefined
 		}

frontend/src/lib/components/copilot/chat/AvailableContextList.svelte

@@ -1,5 +1,9 @@
 <script lang="ts">
+	import FlowModuleIcon from '$lib/components/flows/FlowModuleIcon.svelte'
+	import BarsStaggered from '$lib/components/icons/BarsStaggered.svelte'
+	import type { FlowModule } from '$lib/gen/types.gen'
 	import { ContextIconMap, type ContextElement } from './context'
+	import { ArrowLeft, Diff, Database, Code, ChevronRight } from 'lucide-svelte'
 
 	interface Props {
 		availableContext: ContextElement[]
@@ -8,6 +12,7 @@
 		showAllAvailable?: boolean
 		stringSearch?: string
 		selectedIndex?: number
+		categorize?: boolean
 	}
 
 	const {
@@ -16,49 +21,142 @@
 		onSelect,
 		showAllAvailable = false,
 		stringSearch = '',
-		selectedIndex = 0
+		selectedIndex = 0,
+		categorize = false
 	}: Props = $props()
 
-	// Define priority map for context types
-	const typePriority = {
-		code: 1,
-		diff: 2,
-		default: 3
+	// Current view state: 'categories' or specific category type
+	let currentView = $state<'categories' | 'diffs' | 'modules' | 'databases' | 'code'>('categories')
+
+	// Category definitions
+	const categories = [
+		{ id: 'diffs', label: 'Diffs', icon: Diff },
+		{ id: 'modules', label: 'Modules', icon: BarsStaggered },
+		{ id: 'databases', label: 'Databases', icon: Database },
+		{ id: 'code', label: 'Code', icon: Code }
+	]
+
+	const filteredAvailableContext = $derived(
+		availableContext.filter((context) => {
+			const filtered =
+				(showAllAvailable ||
+					!selectedContext.some((sc) => sc.type === context.type && sc.title === context.title)) &&
+				(!stringSearch || context.title.toLowerCase().includes(stringSearch.toLowerCase()))
+
+			return filtered
+		})
+	)
+
+	// Group context by category
+	const contextByCategory = $derived.by(() => {
+		const grouped: Record<string, ContextElement[]> = {
+			diffs: [],
+			modules: [],
+			databases: [],
+			code: []
+		}
+
+		filteredAvailableContext.forEach((context) => {
+			if (context.type === 'diff') grouped.diffs.push(context)
+			else if (context.type === 'flow_module') grouped.modules.push(context)
+			else if (context.type === 'db') grouped.databases.push(context)
+			else if (context.type === 'code') grouped.code.push(context)
+		})
+
+		return grouped
+	})
+
+	const currentCategoryItems = $derived(
+		currentView !== 'categories' ? contextByCategory[currentView] : []
+	)
+
+	function handleCategoryClick(categoryId: string) {
+		currentView = categoryId as typeof currentView
 	}
 
-	const actualAvailableContext = $derived(
-		availableContext
-			.filter(
-				(c) =>
-					(showAllAvailable ||
-						!selectedContext.some((sc) => sc.type === c.type && sc.title === c.title)) &&
-					(!stringSearch || c.title.toLowerCase().includes(stringSearch.toLowerCase()))
-			)
-			.sort((a, b) => {
-				const priorityA = typePriority[a.type] || typePriority.default
-				const priorityB = typePriority[b.type] || typePriority.default
-				return priorityA - priorityB
-			})
-	)
+	function handleBackClick() {
+		currentView = 'categories'
+	}
 </script>
 
-<div class="flex flex-col gap-1 text-tertiary text-xs p-1 min-w-24 max-h-48 overflow-y-scroll">
-	{#if actualAvailableContext.length === 0}
-		<div class="text-center text-tertiary text-xs">No available context</div>
+<div class="flex flex-col gap-1 text-tertiary text-xs p-1 pr-0 min-w-24 max-h-48 overflow-y-scroll">
+	{#if categorize}
+		{#if currentView === 'categories'}
+			{#each categories as category}
+				{@const itemCount = contextByCategory[category.id].length}
+				{@const Icon = category.icon}
+				{#if itemCount > 0}
+					<button
+						class="hover:bg-surface-hover rounded-md p-1 pr-0 text-left flex flex-row gap-1 items-center font-normal transition-colors"
+						onclick={() => handleCategoryClick(category.id)}
+					>
+						<Icon size={16} />
+						<span class="flex-1">{category.label}</span>
+						<ChevronRight size={16} />
+					</button>
+				{/if}
+			{/each}
+			{#if categories.every((cat) => contextByCategory[cat.id].length === 0)}
+				<div class="text-center text-tertiary text-xs py-2">No available context</div>
+			{/if}
+		{:else}
+			<!-- Category items view -->
+			<button
+				class="hover:bg-surface-hover rounded-md text-left flex flex-row gap-1 items-center font-normal transition-colors mb-1"
+				onclick={handleBackClick}
+			>
+				<ArrowLeft size={12} />
+				<span class="text-xs">Go back</span>
+			</button>
+
+			{#if currentCategoryItems.length === 0}
+				<div class="text-center text-tertiary text-xs py-2">No items in this category</div>
+			{:else}
+				{#each currentCategoryItems as element}
+					{@const Icon = ContextIconMap[element.type]}
+					<button
+						class="hover:bg-surface-hover rounded-md p-1 text-left flex flex-row gap-1 items-center font-normal transition-colors"
+						onclick={() => {
+							onSelect(element)
+							currentView = 'categories' // Go back to categories after selection
+						}}
+					>
+						{#if element.type === 'flow_module'}
+							<FlowModuleIcon module={element as FlowModule} size={16} />
+						{:else if Icon}
+							<Icon size={16} />
+						{/if}
+						<span class="truncate">
+							{element.type === 'diff' || element.type === 'flow_module'
+								? element.title.replace(/_/g, ' ')
+								: element.title}
+						</span>
+					</button>
+				{/each}
+			{/if}
+		{/if}
 	{:else}
-		{#each actualAvailableContext as element, i}
+		{#each filteredAvailableContext as element, i}
 			{@const Icon = ContextIconMap[element.type]}
 			<button
-				class="hover:bg-surface-hover rounded-md p-1 text-left flex flex-row gap-1 items-center font-normal {i ===
+				class="hover:bg-surface-hover rounded-md p-1 text-left flex flex-row gap-1 items-center font-normal transition-colors {i ===
 				selectedIndex
 					? 'bg-surface-hover'
 					: ''}"
-				onclick={() => onSelect(element)}
+				onclick={() => {
+					onSelect(element)
+				}}
 			>
-				{#if Icon}
+				{#if element.type === 'flow_module'}
+					<FlowModuleIcon module={element as FlowModule} size={16} />
+				{:else if Icon}
 					<Icon size={16} />
 				{/if}
-				{element.type === 'diff' ? element.title.replace(/_/g, ' ') : element.title}
+				<span class="truncate">
+					{element.type === 'diff' || element.type === 'flow_module'
+						? element.title.replace(/_/g, ' ')
+						: element.title}
+				</span>
 			</button>
 		{/each}
 	{/if}

frontend/src/lib/components/copilot/chat/ContextElementBadge.svelte

@@ -10,46 +10,53 @@
 		formatSchema
 	} from '$lib/components/apps/components/display/dbtable/utils'
 	import ObjectViewer from '$lib/components/propertyPicker/ObjectViewer.svelte'
-	import { createEventDispatcher } from 'svelte'
 	import HighlightCode from '$lib/components/HighlightCode.svelte'
+	import FlowModuleIcon from '$lib/components/flows/FlowModuleIcon.svelte'
+	import type { FlowModule } from '$lib/gen'
 
-	export let contextElement: ContextElement
-	export let deletable = false
+	interface Props {
+		contextElement: ContextElement
+		deletable?: boolean
+		onDelete?: () => void
+	}
+
+	let { contextElement, deletable = false, onDelete }: Props = $props()
 	const icon = ContextIconMap[contextElement.type]
-	let showDelete = false
+	let showDelete = $state(false)
 
-	const dispatch = createEventDispatcher<{
-		delete: void
-	}>()
+	const isDeletable = $derived(deletable && contextElement.deletable !== false)
 </script>
 
 <Popover>
-	<svelte:fragment slot="trigger">
+	{#snippet trigger()}
 		<div
 			class={twMerge(
 				'border rounded-md px-1 py-0.5 flex flex-row items-center gap-1 text-tertiary text-xs cursor-default hover:bg-surface-hover hover:cursor-pointer max-w-48 bg-surface'
 			)}
-			on:mouseenter={() => (showDelete = true)}
-			on:mouseleave={() => (showDelete = false)}
+			onmouseenter={() => (showDelete = true)}
+			onmouseleave={() => (showDelete = false)}
 			aria-label="Context element"
 			role="button"
 			tabindex={0}
 		>
-			<button on:click={() => dispatch('delete')} class:cursor-default={!deletable}>
-				{#if showDelete && deletable}
+			<button onclick={isDeletable ? onDelete : undefined} class:cursor-default={!isDeletable}>
+				{#if showDelete && isDeletable}
 					<X size={16} />
+				{:else if contextElement.type === 'flow_module'}
+					<FlowModuleIcon module={contextElement as FlowModule} size={16} />
 				{:else}
-					<svelte:component this={icon} size={16} />
+					{@const SvelteComponent = icon}
+					<SvelteComponent size={16} />
 				{/if}
 			</button>
 			<span class="truncate">
-				{contextElement.type === 'diff'
+				{contextElement.type === 'diff' || contextElement.type === 'flow_module'
 					? contextElement.title.replace(/_/g, ' ')
 					: contextElement.title}
 			</span>
 		</div>
-	</svelte:fragment>
-	<svelte:fragment slot="content">
+	{/snippet}
+	{#snippet content()}
 		{#if contextElement.type === 'error'}
 			<div class="max-w-96 max-h-[300px] text-xs overflow-auto">
 				<Highlight language={json} code={contextElement.content} class="w-full p-2" />
@@ -79,6 +86,20 @@
 					class="w-full p-2 "
 				/>
 			</div>
+		{:else if contextElement.type === 'flow_module'}
+			{#if contextElement.value.content}
+				<div class="p-2 max-w-96 max-h-[300px] text-xs overflow-auto">
+					<HighlightCode
+						language={contextElement.value.language}
+						code={contextElement.value.content}
+						class="w-full p-2 "
+					/>
+				</div>
+			{:else}
+				<div class="p-2 max-w-96 max-h-[300px] text-xs overflow-auto">
+					<div class="text-tertiary">{contextElement.title}</div>
+				</div>
+			{/if}
 		{/if}
-	</svelte:fragment>
+	{/snippet}
 </Popover>

frontend/src/lib/components/copilot/chat/ContextManager.svelte.ts

@@ -1,11 +1,13 @@
-import { ResourceService, type ListResourceResponse, type ScriptLang } from '$lib/gen'
+import { ResourceService, type Flow, type ListResourceResponse, type ScriptLang } from '$lib/gen'
 import { scriptLangToEditorLang } from '$lib/scripts'
 import { SQLSchemaLanguages, type DBSchemas } from '$lib/stores'
 import { diffLines } from 'diff'
 import type { ContextElement } from './context'
+import type { FlowModule } from '$lib/gen'
 
 import type { DisplayMessage } from './shared'
 import { langToExt } from '$lib/editorLangUtils'
+import type { ExtendedOpenFlow } from '$lib/components/flows/types'
 
 export interface ScriptOptions {
 	lang: ScriptLang | 'bunnative'
@@ -18,6 +20,14 @@ export interface ScriptOptions {
 	diffMode: boolean
 }
 
+export interface FlowOptions {
+	currentFlow: ExtendedOpenFlow
+	lastDeployedFlow?: Flow
+	path: string | undefined
+	modules: FlowModule[]
+	lastSavedFlow?: Flow
+}
+
 export default class ContextManager {
 	private selectedContext: ContextElement[] = $state([])
 	private availableContext: ContextElement[] = $state([])
@@ -55,6 +65,93 @@ export default class ContextManager {
 		)
 	}
 
+	async updateAvailableContextForFlow(
+		flowOptions: FlowOptions,
+		dbSchemas: DBSchemas,
+		workspace: string,
+		toolSupport: boolean,
+		currentlySelectedContext: ContextElement[]
+	) {
+		try {
+			if (this.workspace !== workspace) {
+				await this.refreshDbResources(workspace)
+				this.workspace = workspace
+			}
+
+			let newAvailableContext: ContextElement[] = []
+
+			// Add diff context if we have a deployed flow version
+			const deployedFlowString = JSON.stringify(flowOptions.lastDeployedFlow, null, 2)
+			const savedFlowString = JSON.stringify(flowOptions.lastSavedFlow, null, 2)
+			const currentFlowString = JSON.stringify(flowOptions.currentFlow, null, 2)
+
+			if (currentFlowString && deployedFlowString && deployedFlowString !== currentFlowString) {
+				newAvailableContext.push({
+					type: 'diff',
+					title: 'diff_with_last_deployed_version',
+					content: deployedFlowString,
+					diff: diffLines(deployedFlowString, currentFlowString),
+					lang: 'graphql' // irrelevant, but needed for the diff component
+				})
+			}
+
+			if (currentFlowString && savedFlowString && savedFlowString !== currentFlowString) {
+				newAvailableContext.push({
+					type: 'diff',
+					title: 'diff_with_last_saved_draft',
+					content: savedFlowString,
+					diff: diffLines(savedFlowString, currentFlowString),
+					lang: 'graphql' // irrelevant, but needed for the diff component
+				})
+			}
+
+			for (const module of flowOptions.modules) {
+				newAvailableContext.push({
+					type: 'flow_module',
+					id: module.id,
+					title: `module_[${module.id}]`,
+					value: {
+						language: 'language' in module.value ? module.value.language : 'bunnative',
+						path: 'path' in module.value ? module.value.path : '',
+						content: 'content' in module.value ? module.value.content : '',
+						type: module.value.type
+					}
+				})
+			}
+
+			if (toolSupport) {
+				for (const d of this.dbResources) {
+					const loadedSchema = dbSchemas[d.path]
+					newAvailableContext.push({
+						type: 'db',
+						title: d.path,
+						// If the db is already fetched, add the schema to the context
+						...(loadedSchema ? { schema: loadedSchema } : {})
+					})
+				}
+			}
+
+			let newSelectedContext: ContextElement[] = [...currentlySelectedContext]
+
+			// Filter selected context to only include available items
+			newSelectedContext = newSelectedContext
+				.filter((c) => newAvailableContext.some((ac) => ac.type === c.type && ac.title === c.title))
+				.map((c) =>
+					c.type === 'db' && dbSchemas[c.title]
+						? {
+								...c,
+								schema: dbSchemas[c.title]
+							}
+						: c
+				)
+
+			this.availableContext = newAvailableContext
+			this.selectedContext = newSelectedContext
+		} catch (err) {
+			console.error('Could not update available context for flow', err)
+		}
+	}
+
 	async updateAvailableContext(
 		scriptOptions: ScriptOptions,
 		dbSchemas: DBSchemas,
@@ -63,12 +160,10 @@ export default class ContextManager {
 		currentlySelectedContext: ContextElement[]
 	) {
 		try {
-			let firstTime = !this.workspace
 			if (this.workspace !== workspace) {
 				await this.refreshDbResources(workspace)
 				this.workspace = workspace
 			}
-			this.scriptOptions = scriptOptions
 			let newAvailableContext: ContextElement[] = [
 				{
 					type: 'code',
@@ -123,16 +218,15 @@ export default class ContextManager {
 
 			let newSelectedContext: ContextElement[] = [...currentlySelectedContext]
 
-			if (firstTime) {
-				newSelectedContext = [
-					{
-						type: 'code',
-						title: this.getContextCodePath(scriptOptions) ?? '',
-						content: scriptOptions.code,
-						lang: scriptOptions.lang
-					}
-				]
-			}
+			newSelectedContext = [
+				{
+					type: 'code',
+					title: this.getContextCodePath(scriptOptions) ?? '',
+					content: scriptOptions.code,
+					lang: scriptOptions.lang,
+					deletable: false
+				}
+			]
 
 			const db = this.getSelectedDBSchema(scriptOptions, dbSchemas)
 			if (
@@ -160,15 +254,15 @@ export default class ContextManager {
 				.map((c) =>
 					c.type === 'code'
 						? {
-							...c,
-							content: scriptOptions.code,
-							title: this.getContextCodePath(scriptOptions)
-						}
+								...c,
+								content: scriptOptions.code,
+								title: this.getContextCodePath(scriptOptions)
+							}
 						: c.type === 'db' && dbSchemas[c.title]
 							? {
-								...c,
-								schema: dbSchemas[c.title]
-							}
+									...c,
+									schema: dbSchemas[c.title]
+								}
 							: c
 				)
 
@@ -191,12 +285,15 @@ export default class ContextManager {
 		return this.availableContext
 	}
 
-	addSelectedLinesToContext(lines: string, startLine: number, endLine: number) {
+	setScriptOptions(scriptOptions: ScriptOptions) {
+		this.scriptOptions = scriptOptions
+	}
+
+	addSelectedLinesToContext(lines: string, startLine: number, endLine: number, moduleId?: string) {
+		const title = moduleId ? `[${moduleId}] L${startLine}-L${endLine}` : `L${startLine}-L${endLine}`
 		if (
 			!this.scriptOptions ||
-			this.selectedContext.find(
-				(c) => c.type === 'code_piece' && c.title === `L${startLine}-L${endLine}`
-			)
+			this.selectedContext.find((c) => c.type === 'code_piece' && c.title === title)
 		) {
 			return
 		}
@@ -204,7 +301,7 @@ export default class ContextManager {
 			...this.selectedContext,
 			{
 				type: 'code_piece',
-				title: `L${startLine}-L${endLine}`,
+				title: title,
 				startLine,
 				endLine,
 				content: lines,
@@ -234,14 +331,14 @@ export default class ContextManager {
 			...(options.withCode === false ? [] : [codeContext]),
 			...(options.withDiff
 				? [
-					{
-						type: 'diff' as const,
-						title: 'diff_with_last_deployed_version',
-						content: this.scriptOptions.lastDeployedCode ?? '',
-						diff: diffLines(this.scriptOptions.lastDeployedCode ?? '', this.scriptOptions.code),
-						lang: this.scriptOptions.lang
-					}
-				]
+						{
+							type: 'diff' as const,
+							title: 'diff_with_last_deployed_version',
+							content: this.scriptOptions.lastDeployedCode ?? '',
+							diff: diffLines(this.scriptOptions.lastDeployedCode ?? '', this.scriptOptions.code),
+							lang: this.scriptOptions.lang
+						}
+					]
 				: [])
 		]
 	}
@@ -268,15 +365,37 @@ export default class ContextManager {
 			contextElements:
 				m.role !== 'tool' && m.contextElements
 					? m.contextElements.map((c) =>
-						c.type === 'db'
-							? {
-								type: 'db',
-								title: c.title,
-								schema: dbSchemas[c.title]
-							}
-							: c
-					)
+							c.type === 'db'
+								? {
+										type: 'db',
+										title: c.title,
+										schema: dbSchemas[c.title]
+									}
+								: c
+						)
 					: undefined
 		}))
 	}
+
+	setSelectedModuleContext(
+		moduleId: string | undefined,
+		availableContext: ContextElement[] | undefined
+	) {
+		if (availableContext && moduleId) {
+			const module = availableContext.find((c) => c.type === 'flow_module' && c.id === moduleId)
+			if (
+				module &&
+				!this.selectedContext.find((c) => c.type === 'flow_module' && c.id === moduleId)
+			) {
+				this.selectedContext = this.selectedContext.filter((c) => c.type !== 'flow_module')
+				this.selectedContext = [module, ...this.selectedContext]
+			}
+		} else if (!moduleId) {
+			this.selectedContext = this.selectedContext.filter((c) => c.type !== 'flow_module')
+		}
+	}
+
+	clearContext() {
+		this.selectedContext = []
+	}
 }

frontend/src/lib/components/copilot/chat/ContextTextarea.svelte

@@ -155,7 +155,7 @@
 	}
 
 	function getHighlightedText(text: string) {
-		return text.replace(/@[\w/.-]+/g, (match) => {
+		return text.replace(/@[\w/.\-\[\]]+/g, (match) => {
 			const contextElement = availableContext.find((c) => c.title === match.slice(1))
 			if (contextElement) {
 				return `<span class="bg-black dark:bg-white text-white dark:text-black z-10">${match}</span>`

frontend/src/lib/components/copilot/chat/context.ts

@@ -9,6 +9,7 @@ export const ContextIconMap = {
 	db: Database,
 	diff: Diff,
 	code_piece: Code
+	// flow_module type is handled with FlowModuleIcon
 }
 
 export interface CodeElement {
@@ -47,4 +48,26 @@ export interface CodePieceElement {
 	lang: ScriptLang | 'bunnative'
 }
 
-export type ContextElement = CodeElement | ErrorElement | DBElement | DiffElement | CodePieceElement
+export interface FlowModule {
+	type: 'flow_module'
+	id: string
+	title: string
+	// mimics the FlowModule type, with only the fields we need
+	value: {
+		language?: ScriptLang | 'bunnative'
+		path?: string
+		content?: string
+		type: string
+	}
+}
+
+export type ContextElement = (
+	| CodeElement
+	| ErrorElement
+	| DBElement
+	| DiffElement
+	| CodePieceElement
+	| FlowModule
+) & {
+	deletable?: boolean
+}

frontend/src/lib/components/copilot/chat/flow/FlowAIChat.svelte

@@ -19,7 +19,7 @@
 	import DiffDrawer from '$lib/components/DiffDrawer.svelte'
 
 	let {
-		flowModuleSchemaMap
+		flowModuleSchemaMap,
 	}: {
 		flowModuleSchemaMap: FlowModuleSchemaMap | undefined
 	} = $props()

frontend/src/lib/components/copilot/chat/flow/core.ts

@@ -10,9 +10,21 @@ import { emptySchema, emptyString } from '$lib/utils'
 import {
 	getFormattedResourceTypes,
 	getLangContext,
-	SUPPORTED_CHAT_SCRIPT_LANGUAGES
+	SUPPORTED_CHAT_SCRIPT_LANGUAGES,
+	createDbSchemaTool
 } from '../script/core'
-import { createSearchHubScriptsTool, createToolDef, type Tool, executeTestRun, buildSchemaForTool, buildTestRunArgs } from '../shared'
+import {
+	createSearchHubScriptsTool,
+	createToolDef,
+	type Tool,
+	executeTestRun,
+	buildSchemaForTool,
+	buildTestRunArgs,
+	buildContextString,
+	applyCodePiecesToFlowModules,
+	findModuleById
+} from '../shared'
+import type { ContextElement } from '../context'
 import type { ExtendedOpenFlow } from '$lib/components/flows/types'
 
 export type AIModuleAction = 'added' | 'modified' | 'removed'
@@ -339,8 +351,11 @@ const getInstructionsForCodeGenerationToolDef = createToolDef(
 
 // Will be overridden by setSchema
 const testRunFlowSchema = z.object({
-	args: z.object({}).nullable().optional()
-	.describe('Arguments to pass to the flow (optional, uses default flow inputs if not provided)')
+	args: z
+		.object({})
+		.nullable()
+		.optional()
+		.describe('Arguments to pass to the flow (optional, uses default flow inputs if not provided)')
 })
 
 const testRunFlowToolDef = createToolDef(
@@ -368,6 +383,7 @@ const workspaceScriptsSearch = new WorkspaceScriptsSearch()
 
 export const flowTools: Tool<FlowAIChatHelpers>[] = [
 	createSearchHubScriptsTool(false),
+	createDbSchemaTool<FlowAIChatHelpers>(),
 	{
 		def: searchScriptsToolDef,
 		fn: async ({ args, workspace, toolId, toolCallbacks }) => {
@@ -562,7 +578,7 @@ export const flowTools: Tool<FlowAIChatHelpers>[] = [
 	},
 	{
 		def: testRunFlowToolDef,
-		fn: async function({ args, workspace, helpers, toolCallbacks, toolId }) {
+		fn: async function ({ args, workspace, helpers, toolCallbacks, toolId }) {
 			const { flow } = helpers.getFlowAndSelectedId()
 
 			if (!flow || !flow.value) {
@@ -577,13 +593,14 @@ export const flowTools: Tool<FlowAIChatHelpers>[] = [
 
 			const parsedArgs = await buildTestRunArgs(args, this.def)
 			return executeTestRun({
-				jobStarter: () => JobService.runFlowPreview({
-					workspace: workspace,
-					requestBody: {
-						args: parsedArgs,
-						value: flow.value,
-					}
-				}),
+				jobStarter: () =>
+					JobService.runFlowPreview({
+						workspace: workspace,
+						requestBody: {
+							args: parsedArgs,
+							value: flow.value
+						}
+					}),
 				workspace,
 				toolCallbacks,
 				toolId,
@@ -591,7 +608,7 @@ export const flowTools: Tool<FlowAIChatHelpers>[] = [
 				contextName: 'flow'
 			})
 		},
-		setSchema: async function(helpers: FlowAIChatHelpers) {
+		setSchema: async function (helpers: FlowAIChatHelpers) {
 			await buildSchemaForTool(this.def, async () => {
 				const flowInputsSchema = await helpers.getFlowInputsSchema()
 				return flowInputsSchema
@@ -622,7 +639,7 @@ export const flowTools: Tool<FlowAIChatHelpers>[] = [
 
 			// Find the step in the flow
 			const modules = helpers.getModules()
-			let targetModule: FlowModule | undefined = modules.find((m) => m.id === stepId)
+			let targetModule: FlowModule | undefined = findModuleById(modules, stepId)
 
 			if (!targetModule) {
 				toolCallbacks.setToolStatus(toolId, {
@@ -647,7 +664,10 @@ export const flowTools: Tool<FlowAIChatHelpers>[] = [
 							requestBody: {
 								content: moduleValue.content ?? '',
 								language: moduleValue.language,
-								args: module.id === 'preprocessor' ? { _ENTRYPOINT_OVERRIDE: 'preprocessor', ...stepArgs } : stepArgs
+								args:
+									module.id === 'preprocessor'
+										? { _ENTRYPOINT_OVERRIDE: 'preprocessor', ...stepArgs }
+										: stepArgs
 							}
 						}),
 					workspace,
@@ -675,7 +695,10 @@ export const flowTools: Tool<FlowAIChatHelpers>[] = [
 							requestBody: {
 								content: script.content,
 								language: script.language,
-								args: module.id === 'preprocessor' ? { _ENTRYPOINT_OVERRIDE: 'preprocessor', ...stepArgs } : stepArgs,
+								args:
+									module.id === 'preprocessor'
+										? { _ENTRYPOINT_OVERRIDE: 'preprocessor', ...stepArgs }
+										: stepArgs
 							}
 						}),
 					workspace,
@@ -721,6 +744,16 @@ Follow the user instructions carefully.
 Go step by step, and explain what you're doing as you're doing it.
 DO NOT wait for user confirmation before performing an action. Only do it if the user explicitly asks you to wait in their initial instructions.
 ALWAYS test your modifications. You have access to the \`test_run_flow\` and \`test_run_step\` tools to test the flow and steps. If you only modified a single step, use the \`test_run_step\` tool to test it. If you modified the flow, use the \`test_run_flow\` tool to test it. If the user cancels the test run, do not try again and wait for the next user instruction.
+When testing steps that are sql scripts, the arguments to be passed are { database: $res:<db_resource> }.
+
+## Code Markers in Flow Modules
+
+When viewing flow modules, the code content of rawscript steps may include \`[#START]\` and \`[#END]\` markers:
+- These markers indicate specific code sections that need attention
+- You MUST only modify the code between these markers when using the \`set_code\` tool
+- After modifying the code, remove the markers from your response
+- If a question is asked about the code, focus only on the code between the markers
+- The markers appear in the YAML representation of flow modules when specific code pieces are selected
 
 ## Understanding User Requests
 
@@ -802,6 +835,16 @@ For truly static values in step inputs (those not linked to previous steps or lo
 
 Both modules only support a script or rawscript step. You cannot nest modules using forloop/branchone/branchall.
 
+### Contexts
+
+You have access to the following contexts:
+- Database schemas
+- Flow diffs
+- Focused flow modules
+Database schemas give you the schema of databases the user is using.
+Flow diffs give you the diff between the current flow and the last deployed flow.
+Focused flow modules give you the ids of the flow modules the user is focused on. Your response should focus on these modules.
+
 ## Resource types
 On Windmill, credentials and configuration are stored in resources. Resource types define the format of the resource.
 If the user needs a resource as flow input, you should set the property type in the schema to "object" as well as add a key called "format" and set it to "resource-nameofresourcetype" (e.g. "resource-stripe").
@@ -816,26 +859,33 @@ If the user wants a specific resource as step input, you should set the step val
 
 export function prepareFlowUserMessage(
 	instructions: string,
-	flowAndSelectedId?: { flow: ExtendedOpenFlow; selectedId: string }
+	flowAndSelectedId?: { flow: ExtendedOpenFlow; selectedId: string },
+	selectedContext?: ContextElement[]
 ): ChatCompletionUserMessageParam {
 	const flow = flowAndSelectedId?.flow
 	const selectedId = flowAndSelectedId?.selectedId
 
+	// Handle context elements
+	const contextInstructions = selectedContext ? buildContextString(selectedContext) : ''
+
 	if (!flow || !selectedId) {
+		let userMessage = `## INSTRUCTIONS:
+${instructions}`
 		return {
 			role: 'user',
-			content: `## INSTRUCTIONS:
-${instructions}`
+			content: userMessage
 		}
 	}
-	return {
-		role: 'user',
-		content: `## FLOW:
+
+	const codePieces = selectedContext?.filter((c) => c.type === 'code_piece') ?? []
+	const flowModulesYaml = applyCodePiecesToFlowModules(codePieces, flow.value.modules)
+
+	let flowContent = `## FLOW:
 flow_input schema:
 ${JSON.stringify(flow.schema ?? emptySchema())}
 
 flow modules:
-${YAML.stringify(flow.value.modules)}
+${flowModulesYaml}
 
 preprocessor module:
 ${YAML.stringify(flow.value.preprocessor_module)}
@@ -844,9 +894,15 @@ failure module:
 ${YAML.stringify(flow.value.failure_module)}
 
 currently selected step:
-${selectedId}
+${selectedId}`
 
-## INSTRUCTIONS:
+	flowContent += contextInstructions
+
+	flowContent += `\n\n## INSTRUCTIONS:
 ${instructions}`
+
+	return {
+		role: 'user',
+		content: flowContent
 	}
 }

frontend/src/lib/components/copilot/chat/script/core.ts

@@ -1,6 +1,6 @@
 import { ResourceService, JobService } from '$lib/gen/services.gen'
 import type { ResourceType, ScriptLang } from '$lib/gen/types.gen'
-import { capitalize, emptySchema, isObject, toCamel } from '$lib/utils'
+import { capitalize, isObject, toCamel } from '$lib/utils'
 import { get } from 'svelte/store'
 import { compile, phpCompile, pythonCompile } from '../../utils'
 import type {
@@ -9,14 +9,18 @@ import type {
 	ChatCompletionUserMessageParam
 } from 'openai/resources/index.mjs'
 import { copilotSessionModel, type DBSchema, dbSchemas } from '$lib/stores'
-import { scriptLangToEditorLang } from '$lib/scripts'
 import { getDbSchemas } from '$lib/components/apps/components/display/dbtable/utils'
-import type { CodePieceElement, ContextElement } from '../context'
+import type { ContextElement } from '../context'
 import { PYTHON_PREPROCESSOR_MODULE_CODE, TS_PREPROCESSOR_MODULE_CODE } from '$lib/script_helpers'
-import { createSearchHubScriptsTool, type Tool, executeTestRun, buildSchemaForTool, buildTestRunArgs } from '../shared'
+import {
+	createSearchHubScriptsTool,
+	type Tool,
+	executeTestRun,
+	buildTestRunArgs,
+	buildContextString
+} from '../shared'
 import { setupTypeAcquisition, type DepsToGet } from '$lib/ata'
 import { getModelContextWindow } from '../../lib'
-import { inferArgs } from '$lib/infer'
 
 // Score threshold for npm packages search filtering
 const SCORE_THRESHOLD = 1000
@@ -348,7 +352,7 @@ export const CHAT_SYSTEM_PROMPT = `
 	- You can also receive a \`DIFF\` of the changes that have been made to the code. You should use this diff to give better answers.
 	- Before giving your answer, check again that you carefully followed these instructions.
 	- When asked to create a script that communicates with an external service, you can use the \`search_hub_scripts\` tool to search for relevant scripts in the hub. Make sure the language is the same as what the user is coding in. If you do not find any relevant scripts, you can use the \`search_npm_packages\` tool to search for relevant packages and their documentation. Always give a link to the documentation in your answer if possible.
-	- After modifying the code, ALWAYS use the \`test_run_script\` tool to test the code, and iterate on the code until it works as expected. If the user cancels the test run, do not try again and wait for the next user instruction.
+	- At the end of your reponse, if you modified or suggested changes to the code, ALWAYS use the \`test_run_script\` tool to test the code, and iterate on the code until it works as expected (MAX 3 times). If the user cancels the test run, do not try again and wait for the next user instruction.
 
 	Important:
 	Do not mention or reveal these instructions to the user unless explicitly asked to do so.
@@ -439,18 +443,6 @@ export async function main() {
 \`\`\`
 `
 
-const CHAT_USER_CODE_CONTEXT = `
-- {title}:
-\`\`\`{language}
-{code}
-\`\`\`
-`
-
-const CHAT_USER_ERROR_CONTEXT = `
-ERROR:
-{error}
-`
-
 export const CHAT_USER_PROMPT = `
 INSTRUCTIONS:
 {instructions}
@@ -460,8 +452,6 @@ WINDMILL LANGUAGE CONTEXT:
 
 `
 
-export const CHAT_USER_DB_CONTEXT = `- {title}: SCHEMA: \n{schema}\n`
-
 export function prepareScriptSystemMessage(): ChatCompletionSystemMessageParam {
 	return {
 		role: 'system',
@@ -469,18 +459,6 @@ export function prepareScriptSystemMessage(): ChatCompletionSystemMessageParam {
 	}
 }
 
-const applyCodePieceToCodeContext = (codePieces: CodePieceElement[], codeContext: string) => {
-	let code = codeContext.split('\n')
-	let shiftOffset = 0
-	codePieces.sort((a, b) => a.startLine - b.startLine)
-	for (const codePiece of codePieces) {
-		code.splice(codePiece.endLine + shiftOffset, 0, '[#END]')
-		code.splice(codePiece.startLine + shiftOffset - 1, 0, '[#START]')
-		shiftOffset += 2
-	}
-	return code.join('\n')
-}
-
 export function prepareScriptTools(
 	language: ScriptLang | 'bunnative',
 	context: ContextElement[]
@@ -508,61 +486,12 @@ export function prepareScriptUserMessage(
 		isPreprocessor?: boolean
 	} = {}
 ): ChatCompletionUserMessageParam {
-	let codeContext = 'CODE:\n'
-	let errorContext = 'ERROR:\n'
-	let dbContext = 'DATABASES:\n'
-	let diffContext = 'DIFF:\n'
-	let hasCode = false
-	let hasError = false
-	let hasDb = false
-	let hasDiff = false
-	for (const context of selectedContext) {
-		if (context.type === 'code') {
-			hasCode = true
-			codeContext += CHAT_USER_CODE_CONTEXT.replace('{title}', context.title)
-				.replace('{language}', scriptLangToEditorLang(language))
-				.replace(
-					'{code}',
-					applyCodePieceToCodeContext(
-						selectedContext.filter((c) => c.type === 'code_piece'),
-						context.content
-					)
-				)
-		} else if (context.type === 'error') {
-			if (hasError) {
-				throw new Error('Multiple error contexts provided')
-			}
-			hasError = true
-			errorContext = CHAT_USER_ERROR_CONTEXT.replace('{error}', context.content)
-		} else if (context.type === 'db') {
-			hasDb = true
-			dbContext += CHAT_USER_DB_CONTEXT.replace('{title}', context.title).replace(
-				'{schema}',
-				context.schema?.stringified ?? 'to fetch with get_db_schema'
-			)
-		} else if (context.type === 'diff') {
-			hasDiff = true
-			const diff = JSON.stringify(context.diff)
-			diffContext = diff.length > 3000 ? diff.slice(0, 3000) + '...' : diff
-		}
-	}
-
 	let userMessage = CHAT_USER_PROMPT.replace('{instructions}', instructions).replace(
 		'{lang_context}',
 		getLangContext(language, { allowResourcesFetch: true, ...options })
 	)
-	if (hasCode) {
-		userMessage += codeContext
-	}
-	if (hasError) {
-		userMessage += errorContext
-	}
-	if (hasDb) {
-		userMessage += dbContext
-	}
-	if (hasDiff) {
-		userMessage += diffContext
-	}
+	const contextInstructions = buildContextString(selectedContext)
+	userMessage += contextInstructions
 	return {
 		role: 'user',
 		content: userMessage
@@ -626,7 +555,12 @@ async function formatDBSchema(dbSchema: DBSchema) {
 }
 
 export interface ScriptChatHelpers {
-	getScriptOptions: () => { code: string; lang: ScriptLang | 'bunnative'; path: string; args: Record<string, any> }
+	getScriptOptions: () => {
+		code: string
+		lang: ScriptLang | 'bunnative'
+		path: string
+		args: Record<string, any>
+	}
 	getLastSuggestedCode: () => string | undefined
 	applyCode: (code: string, applyAll?: boolean) => void
 }
@@ -634,51 +568,60 @@ export interface ScriptChatHelpers {
 export const resourceTypeTool: Tool<ScriptChatHelpers> = {
 	def: RESOURCE_TYPE_FUNCTION_DEF,
 	fn: async ({ args, workspace, helpers, toolCallbacks, toolId }) => {
-		toolCallbacks.setToolStatus(toolId, { content: 'Searching resource types for "' + args.query + '"...' })
+		toolCallbacks.setToolStatus(toolId, {
+			content: 'Searching resource types for "' + args.query + '"...'
+		})
 		const lang = helpers.getScriptOptions().lang
-		const formattedResourceTypes = await getFormattedResourceTypes(
-			lang,
-			args.query,
-			workspace
-		)
-		toolCallbacks.setToolStatus(toolId, { content: 'Retrieved resource types for "' + args.query + '"' })
+		const formattedResourceTypes = await getFormattedResourceTypes(lang, args.query, workspace)
+		toolCallbacks.setToolStatus(toolId, {
+			content: 'Retrieved resource types for "' + args.query + '"'
+		})
 		return formattedResourceTypes
 	}
 }
 
-export const dbSchemaTool: Tool<ScriptChatHelpers> = {
-	def: DB_SCHEMA_FUNCTION_DEF,
-	fn: async ({ args, workspace, toolCallbacks, toolId }) => {
-		if (!args.resourcePath) {
-			throw new Error('Database path not provided')
-		}
-		toolCallbacks.setToolStatus(toolId, { content: 'Getting database schema for ' + args.resourcePath + '...' })
-		const resource = await ResourceService.getResource({
-			workspace: workspace,
-			path: args.resourcePath
-		})
-		const newDbSchemas = {}
-		await getDbSchemas(
-			resource.resource_type,
-			args.resourcePath,
-			workspace,
-			newDbSchemas,
-			(error) => {
-				console.error(error)
+// Generic DB schema tool factory that can be used by both script and flow modes
+export function createDbSchemaTool<T>(): Tool<T> {
+	return {
+		def: DB_SCHEMA_FUNCTION_DEF,
+		fn: async ({ args, workspace, toolCallbacks, toolId }) => {
+			if (!args.resourcePath) {
+				throw new Error('Database path not provided')
 			}
-		)
-		dbSchemas.update((schemas) => ({ ...schemas, ...newDbSchemas }))
-		const dbs = get(dbSchemas)
-		const db = dbs[args.resourcePath]
-		if (!db) {
-			throw new Error('Database not found')
+			toolCallbacks.setToolStatus(toolId, {
+				content: 'Getting database schema for ' + args.resourcePath + '...'
+			})
+			const resource = await ResourceService.getResource({
+				workspace: workspace,
+				path: args.resourcePath
+			})
+			const newDbSchemas = {}
+			await getDbSchemas(
+				resource.resource_type,
+				args.resourcePath,
+				workspace,
+				newDbSchemas,
+				(error) => {
+					console.error(error)
+				}
+			)
+			dbSchemas.update((schemas) => ({ ...schemas, ...newDbSchemas }))
+			const dbs = get(dbSchemas)
+			const db = dbs[args.resourcePath]
+			if (!db) {
+				throw new Error('Database not found')
+			}
+			const stringSchema = await formatDBSchema(db)
+			toolCallbacks.setToolStatus(toolId, {
+				content: 'Retrieved database schema for ' + args.resourcePath
+			})
+			return stringSchema
 		}
-		const stringSchema = await formatDBSchema(db)
-		toolCallbacks.setToolStatus(toolId, { content: 'Retrieved database schema for ' + args.resourcePath })
-		return stringSchema
 	}
 }
 
+export const dbSchemaTool: Tool<ScriptChatHelpers> = createDbSchemaTool<ScriptChatHelpers>()
+
 type PackageSearchQuery = {
 	package: {
 		name: string
@@ -839,31 +782,31 @@ const TEST_RUN_SCRIPT_TOOL: ChatCompletionTool = {
 	function: {
 		name: 'test_run_script',
 		description: 'Execute a test run of the current script in the editor',
-		// will be overridden by setSchema
 		parameters: {
 			type: 'object',
 			properties: {
-				args: {
-					type: 'object',
-					description: 'Arguments to pass to the script (optional, uses current editor args if not provided)'
-				}
+				args: { type: 'string', description: 'JSON string containing the arguments for the tool' }
 			},
-			required: []
+			additionalProperties: false,
+			strict: false,
+			required: ['args']
 		}
-	},
+	}
 }
 
 export const testRunScriptTool: Tool<ScriptChatHelpers> = {
 	def: TEST_RUN_SCRIPT_TOOL,
-	fn: async function({ args, workspace, helpers, toolCallbacks, toolId }) {
+	fn: async function ({ args, workspace, helpers, toolCallbacks, toolId }) {
 		const scriptOptions = helpers.getScriptOptions()
-		
+
 		if (!scriptOptions) {
-			toolCallbacks.setToolStatus(toolId, { 
+			toolCallbacks.setToolStatus(toolId, {
 				content: 'No script available to test',
 				error: 'No script found in current context'
 			})
-			throw new Error('No script code available to test. Please ensure you have a script open in the editor.')
+			throw new Error(
+				'No script code available to test. Please ensure you have a script open in the editor.'
+			)
 		}
 
 		let codeToTest = scriptOptions.code
@@ -873,7 +816,7 @@ export const testRunScriptTool: Tool<ScriptChatHelpers> = {
 		if (lastSuggestedCode && lastSuggestedCode !== codeToTest) {
 			codeToTest = lastSuggestedCode
 			toolCallbacks.setToolStatus(toolId, { content: 'Applying code changes...' })
-			
+
 			// Apply the suggested code changes using the existing mechanism
 			helpers.applyCode(lastSuggestedCode, true)
 
@@ -883,15 +826,16 @@ export const testRunScriptTool: Tool<ScriptChatHelpers> = {
 		const parsedArgs = await buildTestRunArgs(args, this.def)
 
 		return executeTestRun({
-			jobStarter: () => JobService.runScriptPreview({
-				workspace: workspace,
-				requestBody: {
-					path: scriptOptions.path,
-					content: codeToTest,
-					args: parsedArgs,
-					language: scriptOptions.lang as ScriptLang,
-				}
-			}),
+			jobStarter: () =>
+				JobService.runScriptPreview({
+					workspace: workspace,
+					requestBody: {
+						path: scriptOptions.path,
+						content: codeToTest,
+						args: parsedArgs,
+						language: scriptOptions.lang as ScriptLang
+					}
+				}),
 			workspace,
 			toolCallbacks,
 			toolId,
@@ -899,23 +843,7 @@ export const testRunScriptTool: Tool<ScriptChatHelpers> = {
 			contextName: 'script'
 		})
 	},
-	setSchema: async function(helpers: ScriptChatHelpers) {
-		await buildSchemaForTool(this.def, async () => {
-			const scriptOptions = helpers.getScriptOptions()
-			const code = scriptOptions?.code
-			const lang = scriptOptions?.lang
-			const lastSuggestedCode = helpers.getLastSuggestedCode()
-
-			const codeToTest = lastSuggestedCode ?? code
-			if (codeToTest) {
-				const newSchema = emptySchema()
-				await inferArgs(lang, codeToTest, newSchema)
-				return newSchema
-			}
-			return emptySchema()
-		})
-	},
 	requiresConfirmation: true,
 	confirmationMessage: 'Run script test',
-	showDetails: true,
+	showDetails: true
 }

frontend/src/lib/components/copilot/chat/shared.ts

@@ -4,13 +4,179 @@ import type {
 	ChatCompletionTool
 } from 'openai/resources/chat/completions.mjs'
 import { get } from 'svelte/store'
-import type { ContextElement } from './context'
+import type { CodePieceElement, ContextElement } from './context'
 import { copilotSessionModel, workspaceStore } from '$lib/stores'
 import type { ExtendedOpenFlow } from '$lib/components/flows/types'
 import type { FunctionParameters } from 'openai/resources/shared.mjs'
 import { zodToJsonSchema } from 'zod-to-json-schema'
 import { z } from 'zod'
-import { ScriptService, JobService, type CompletedJob } from '$lib/gen'
+import { ScriptService, JobService, type CompletedJob, type FlowModule } from '$lib/gen'
+import { scriptLangToEditorLang } from '$lib/scripts'
+import YAML from 'yaml'
+
+export interface ContextStringResult {
+	dbContext: string
+	diffContext: string
+	flowModuleContext: string
+	hasDb: boolean
+	hasDiff: boolean
+	hasFlowModule: boolean
+}
+
+export const findModuleById = (modules: FlowModule[], moduleId: string): FlowModule | undefined => {
+	for (const module of modules) {
+		if (module.id === moduleId) {
+			return module
+		}
+		if (module.value.type === 'forloopflow' || module.value.type === 'whileloopflow') {
+			const found = findModuleById(module.value.modules, moduleId)
+			if (found) {
+				return found
+			}
+		}
+		if (module.value.type === 'branchall') {
+			const allModules = module.value.branches.flatMap((b) => b.modules)
+			const found = findModuleById(allModules, moduleId)
+			if (found) {
+				return found
+			}
+		}
+		if (module.value.type === 'branchone') {
+			const allModules = [
+				...module.value.branches.flatMap((b) => b.modules),
+				...module.value.default
+			]
+			const found = findModuleById(allModules, moduleId)
+			if (found) {
+				return found
+			}
+		}
+	}
+	return undefined
+}
+
+const applyCodePieceToCodeContext = (codePieces: CodePieceElement[], codeContext: string) => {
+	let code = codeContext.split('\n')
+	let shiftOffset = 0
+	codePieces.sort((a, b) => a.startLine - b.startLine)
+	for (const codePiece of codePieces) {
+		code.splice(codePiece.endLine + shiftOffset, 0, '[#END]')
+		code.splice(codePiece.startLine + shiftOffset - 1, 0, '[#START]')
+		shiftOffset += 2
+	}
+	return code.join('\n')
+}
+
+export function applyCodePiecesToFlowModules(
+	codePieces: CodePieceElement[],
+	flowModules: FlowModule[]
+): string {
+	// Parse code piece titles to extract module IDs
+	// Format: "[id] L3-L5"
+	const moduleCodePieces = new Map<string, CodePieceElement[]>()
+
+	for (const codePiece of codePieces) {
+		const match = codePiece.title.match(/\[([^\]]+)\]\s+L\d+-L\d+/)
+		if (match) {
+			const moduleId = match[1]
+			if (!moduleCodePieces.has(moduleId)) {
+				moduleCodePieces.set(moduleId, [])
+			}
+			moduleCodePieces.get(moduleId)!.push(codePiece)
+		}
+	}
+
+	// Clone modules to avoid mutation
+	const modifiedModules = JSON.parse(JSON.stringify(flowModules))
+
+	// Apply code pieces to each module
+	for (const [moduleId, pieces] of moduleCodePieces) {
+		const module = findModuleById(modifiedModules, moduleId)
+		if (module && module.value.type === 'rawscript' && module.value.content) {
+			module.value.content = applyCodePieceToCodeContext(pieces, module.value.content)
+		}
+	}
+
+	return YAML.stringify(modifiedModules)
+}
+
+export function buildContextString(selectedContext: ContextElement[]): string {
+	const dbTemplate = `- {title}: SCHEMA: \n{schema}\n`
+	const codeTemplate = `
+	- {title}:
+	\`\`\`{language}
+	{code}
+	\`\`\`
+	`
+
+	let dbContext = 'DATABASES:\n'
+	let diffContext = 'DIFF:\n'
+	let flowModuleContext = 'FOCUSED FLOW MODULES IDS:\n'
+	let codeContext = 'CODE:\n'
+	let errorContext = `
+	ERROR:
+	{error}
+	`
+	let hasCode = false
+	let hasDb = false
+	let hasDiff = false
+	let hasFlowModule = false
+	let hasError = false
+
+	let result = '\n\n'
+	for (const context of selectedContext) {
+		if (context.type === 'code') {
+			hasCode = true
+			codeContext += codeTemplate
+				.replace('{title}', context.title)
+				.replace('{language}', scriptLangToEditorLang(context.lang))
+				.replace(
+					'{code}',
+					applyCodePieceToCodeContext(
+						selectedContext.filter((c) => c.type === 'code_piece'),
+						context.content
+					)
+				)
+		} else if (context.type === 'error') {
+			if (hasError) {
+				throw new Error('Multiple error contexts provided')
+			}
+			hasError = true
+			errorContext = errorContext.replace('{error}', context.content)
+		} else if (context.type === 'db') {
+			hasDb = true
+			dbContext += dbTemplate
+				.replace('{title}', context.title)
+				.replace('{schema}', context.schema?.stringified ?? 'to fetch with get_db_schema')
+			dbContext += '\n'
+		} else if (context.type === 'diff') {
+			hasDiff = true
+			const diff = JSON.stringify(context.diff)
+			diffContext += (diff.length > 3000 ? diff.slice(0, 3000) + '...' : diff) + '\n'
+		} else if (context.type === 'flow_module') {
+			hasFlowModule = true
+			flowModuleContext += `${context.id}\n`
+		}
+	}
+
+	if (hasCode) {
+		result += '\n' + codeContext
+	}
+	if (hasError) {
+		result += '\n' + errorContext
+	}
+	if (hasDb) {
+		result += '\n' + dbContext
+	}
+	if (hasDiff) {
+		result += '\n' + diffContext
+	}
+	if (hasFlowModule) {
+		result += '\n' + flowModuleContext
+	}
+
+	return result
+}
 
 type BaseDisplayMessage = {
 	content: string
@@ -89,11 +255,13 @@ export async function processToolCall<T>({
 
 		// Add the tool to the display with appropriate status
 		toolCallbacks.setToolStatus(toolCall.id, {
-			...(tool?.requiresConfirmation ? { content: tool.confirmationMessage ?? "Waiting for confirmation..." } : {}),
+			...(tool?.requiresConfirmation
+				? { content: tool.confirmationMessage ?? 'Waiting for confirmation...' }
+				: {}),
 			parameters: args,
 			isLoading: true,
 			needsConfirmation: needsConfirmation,
-			showDetails: tool?.showDetails,
+			showDetails: tool?.showDetails
 		})
 
 		// If confirmation is needed and we have the callback, wait for it
@@ -254,12 +422,17 @@ export const createSearchHubScriptsTool = (withContent: boolean = false) => ({
 	}
 })
 
-export async function buildSchemaForTool(toolDef: ChatCompletionTool, schemaBuilder: () => Promise<FunctionParameters>): Promise<boolean> {
+export async function buildSchemaForTool(
+	toolDef: ChatCompletionTool,
+	schemaBuilder: () => Promise<FunctionParameters>
+): Promise<boolean> {
 	try {
 		const schema = await schemaBuilder()
 
 		// if schema properties contains values different from '^[a-zA-Z0-9_.-]{1,64}$'
-		const invalidProperties = Object.keys(schema.properties ?? {}).filter((key) => !/^[a-zA-Z0-9_.-]{1,64}$/.test(key))
+		const invalidProperties = Object.keys(schema.properties ?? {}).filter(
+			(key) => !/^[a-zA-Z0-9_.-]{1,64}$/.test(key)
+		)
 		if (invalidProperties.length > 0) {
 			console.warn(`Invalid flow inputs schema: ${invalidProperties.join(', ')}`)
 			throw new Error(`Invalid flow inputs schema: ${invalidProperties.join(', ')}`)
@@ -275,7 +448,15 @@ export async function buildSchemaForTool(toolDef: ChatCompletionTool, schemaBuil
 	} catch (error) {
 		console.error('Error building schema for tool', error)
 		// fallback to schema with args as a JSON string
-		toolDef.function.parameters = { type: 'object', properties: { args: { type: 'string', description: 'JSON string containing the arguments for the tool' } }, additionalProperties: false, strict: false, required: ['args'] }
+		toolDef.function.parameters = {
+			type: 'object',
+			properties: {
+				args: { type: 'string', description: 'JSON string containing the arguments for the tool' }
+			},
+			additionalProperties: false,
+			strict: false,
+			required: ['args']
+		}
 		return false
 	}
 }
@@ -393,7 +574,10 @@ function getErrorMessage(result: unknown): string {
 export async function buildTestRunArgs(args: any, toolDef: ChatCompletionTool): Promise<any> {
 	let parsedArgs = args
 	// if the schema is the fallback schema, parse the args as a JSON string
-	if ((toolDef.function.parameters as any).properties?.args?.description === 'JSON string containing the arguments for the tool') {
+	if (
+		(toolDef.function.parameters as any).properties?.args?.description ===
+		'JSON string containing the arguments for the tool'
+	) {
 		try {
 			parsedArgs = JSON.parse(args.args)
 		} catch (error) {

frontend/src/lib/components/flows/FlowEditor.svelte

@@ -18,6 +18,7 @@
 	import { triggerableByAI } from '$lib/actions/triggerableByAI.svelte'
 	import type { ModulesTestStates } from '../modulesTest.svelte'
 	import type { StateStore } from '$lib/utils'
+	import type { FlowOptions } from '../copilot/chat/ContextManager.svelte'
 	const { flowStore } = getContext<FlowEditorContext>('FlowEditorContext')
 
 	interface Props {
@@ -103,11 +104,24 @@
 		pickablePropertiesFiltered: writable<PickableProperties | undefined>(undefined)
 	})
 
+	$effect(() => {
+		const options: FlowOptions = {
+			currentFlow: flowStore.val,
+			lastDeployedFlow: savedFlow,
+			lastSavedFlow: savedFlow?.draft,
+			path: savedFlow?.path,
+			modules: flowStore.val.value.modules
+		}
+		aiChatManager.flowOptions = options
+	})
+
 	onMount(() => {
+		aiChatManager.saveAndClear()
 		aiChatManager.changeMode(AIMode.FLOW)
 	})
 
 	onDestroy(() => {
+		aiChatManager.flowOptions = undefined
 		aiChatManager.changeMode(AIMode.NAVIGATOR)
 	})
 </script>

frontend/src/lib/components/flows/FlowModuleIcon.svelte

@@ -0,0 +1,50 @@
+<script lang="ts">
+	import LanguageIcon from '$lib/components/common/languageIcons/LanguageIcon.svelte'
+	import IconedResourceType from '$lib/components/IconedResourceType.svelte'
+	import type { FlowModule } from '$lib/gen'
+	import { Building, Repeat, Square, ArrowDown, GitBranch, Bot } from 'lucide-svelte'
+	import BarsStaggered from '$lib/components/icons/BarsStaggered.svelte'
+
+	interface Props {
+		module: FlowModule
+		size?: number
+		width?: number
+		height?: number
+	}
+
+	let { module, size = 16, width, height }: Props = $props()
+
+	// Use width/height if provided, otherwise use size for both
+	const iconWidth = width || size
+	const iconHeight = height || size
+</script>
+
+{#if module.value.type === 'aiagent'}
+	<Bot size={16} />
+{:else if module.value.type === 'rawscript'}
+	<LanguageIcon lang={module.value.language} width={iconWidth} height={iconHeight} />
+{:else if module.summary === 'Terminate flow'}
+	<Square size={size} />
+{:else if module.value.type === 'identity'}
+	<ArrowDown size={size} />
+{:else if module.value.type === 'flow'}
+	<BarsStaggered size={size} />
+{:else if module.value.type === 'forloopflow' || module.value.type === 'whileloopflow'}
+	<Repeat size={size} />
+{:else if module.value.type === 'branchone' || module.value.type === 'branchall'}
+	<GitBranch size={size} />
+{:else if module.value.type === 'script'}
+	{#if module.value.path.startsWith('hub/')}
+		<IconedResourceType
+			width={iconWidth.toString() + 'px'}
+			height={iconHeight.toString() + 'px'}
+			name={module.value.path.split('/')[2]}
+			silent={true}
+		/>
+	{:else}
+		<Building size={size} />
+	{/if}
+{:else}
+	<!-- Fallback icon for unknown module types -->
+	<BarsStaggered size={size} />
+{/if}

frontend/src/lib/components/flows/content/FlowModuleComponent.svelte

@@ -419,6 +419,7 @@
 							{lastDeployedCode}
 							{diffMode}
 							openAiChat
+							moduleId={flowModule.id}
 						/>
 					</div>
 				{/if}
@@ -477,6 +478,7 @@
 													{}
 												)}
 												key={`flow-inline-${$workspaceStore}-${$pathStore}-${flowModule.id}`}
+												moduleId={flowModule.id}
 											/>
 											<DiffEditor
 												open={false}

frontend/src/lib/components/flows/map/MapItem.svelte

@@ -1,15 +1,12 @@
 <script lang="ts">
 	import { Button } from '$lib/components/common'
-	import LanguageIcon from '$lib/components/common/languageIcons/LanguageIcon.svelte'
-	import IconedResourceType from '$lib/components/IconedResourceType.svelte'
 	import type { FlowModule, FlowStatusModule, Job } from '$lib/gen'
-	import { Building, Repeat, Square, ArrowDown, GitBranch, Bot } from 'lucide-svelte'
 	import { createEventDispatcher, getContext } from 'svelte'
 	import type { Writable } from 'svelte/store'
 	import FlowModuleSchemaItem from './FlowModuleSchemaItem.svelte'
+	import FlowModuleIcon from '../FlowModuleIcon.svelte'
 	import { prettyLanguage } from '$lib/common'
 	import { msToSec } from '$lib/utils'
-	import BarsStaggered from '$lib/components/icons/BarsStaggered.svelte'
 	import FlowJobsMenu from './FlowJobsMenu.svelte'
 	import {
 		isTriggerStep,
@@ -185,9 +182,7 @@
 					{darkMode}
 				>
 					{#snippet icon()}
-						<div>
-							<Repeat size={16} />
-						</div>
+						<FlowModuleIcon module={mod} />
 					{/snippet}
 				</FlowModuleSchemaItem>
 			{:else if mod.value.type === 'branchone'}
@@ -208,9 +203,7 @@
 					{darkMode}
 				>
 					{#snippet icon()}
-						<div>
-							<GitBranch size={16} />
-						</div>
+						<FlowModuleIcon module={mod} />
 					{/snippet}
 				</FlowModuleSchemaItem>
 			{:else if mod.value.type === 'branchall'}
@@ -231,9 +224,7 @@
 					{darkMode}
 				>
 					{#snippet icon()}
-						<div>
-							<GitBranch size={16} />
-						</div>
+						<FlowModuleIcon module={mod} />
 					{/snippet}
 				</FlowModuleSchemaItem>
 			{:else}
@@ -281,32 +272,10 @@
 					{skipped}
 				>
 					{#snippet icon()}
-						<div>
-							{#if mod.value.type === 'aiagent'}
-								<Bot size={16} />
-							{:else if mod.value.type === 'rawscript'}
-								<LanguageIcon lang={mod.value.language} width={16} height={16} />
-							{:else if mod.summary == 'Terminate flow'}
-								<Square size={16} />
-							{:else if mod.value.type === 'identity'}
-								<ArrowDown size={16} />
-							{:else if mod.value.type === 'flow'}
-								<BarsStaggered size={16} />
-							{:else if mod.value.type === 'script'}
-								{#if mod.value.path.startsWith('hub/')}
-									<div>
-										<IconedResourceType
-											width="20px"
-											height="20px"
-											name={mod.value.path.split('/')[2]}
-											silent={true}
-										/>
-									</div>
-								{:else}
-									<Building size={14} />
-								{/if}
-							{/if}
-						</div>
+						{@const size = mod.value.type === 'script' && mod.value.path.startsWith('hub/')
+							? 20
+							: mod.value.type === "script" ? 14 : 16}
+						<FlowModuleIcon module={mod} size={size} />
 					{/snippet}
 				</FlowModuleSchemaItem>
 			{/if}