nits(aichat): add keyboard navigation in context list (#6443)

* better flow module code peice

* better keyboard nav for availablecontextlist

* cleaning

* escape to close + cleaning

* nit tab handling

* fix module extraction

* remove code category

* fixes

* comment

* fix

* fix tool params display

frontend/src/lib/components/Editor.svelte

@@ -1237,7 +1237,13 @@
 
 		try {
 			editor = meditor.create(divEl as HTMLDivElement, {
-				...editorConfig(code ?? '', lang, automaticLayout, fixedOverflowWidgets, $relativeLineNumbers),
+				...editorConfig(
+					code ?? '',
+					lang,
+					automaticLayout,
+					fixedOverflowWidgets,
+					$relativeLineNumbers
+				),
 				model,
 				fontSize: !small ? 14 : 12,
 				lineNumbersMinChars,
@@ -1657,7 +1663,7 @@
 		files && model && untrack(() => onFileChanges())
 	})
 	$effect(() => {
-		editor?.updateOptions({ 
+		editor?.updateOptions({
 			lineNumbers: $relativeLineNumbers ? 'relative' : 'on'
 		})
 	})

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

@@ -150,7 +150,6 @@
 								addContextToSelection(element)
 								close()
 							}}
-							categorize
 						/>
 					</svelte:fragment>
 				</Popover>

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

@@ -3,37 +3,40 @@
 	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'
+	import { ArrowLeft, Diff, Database, ChevronRight } from 'lucide-svelte'
 
 	interface Props {
 		availableContext: ContextElement[]
 		selectedContext: ContextElement[]
 		onSelect: (element: ContextElement) => void
+		setShowing?: (showing: boolean) => void
 		showAllAvailable?: boolean
 		stringSearch?: string
-		selectedIndex?: number
-		categorize?: boolean
+		onViewChange?: (newNumber: number) => void
 	}
 
 	const {
 		availableContext,
 		selectedContext,
 		onSelect,
+		setShowing,
 		showAllAvailable = false,
 		stringSearch = '',
-		selectedIndex = 0,
-		categorize = false
+		onViewChange
 	}: Props = $props()
 
 	// Current view state: 'categories' or specific category type
-	let currentView = $state<'categories' | 'diffs' | 'modules' | 'databases' | 'code'>('categories')
+	let currentView = $state<'categories' | 'diffs' | 'modules' | 'databases'>('categories')
+
+	// Selected index for keyboard navigation
+	let itemSelectedIndex = $state(0)
+	let categorySelectedIndex = $state(0)
 
 	// 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 }
+		{ id: 'databases', label: 'Databases', icon: Database }
 	]
 
 	const filteredAvailableContext = $derived(
@@ -52,15 +55,13 @@
 		const grouped: Record<string, ContextElement[]> = {
 			diffs: [],
 			modules: [],
-			databases: [],
-			code: []
+			databases: []
 		}
 
 		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
@@ -70,77 +71,140 @@
 		currentView !== 'categories' ? contextByCategory[currentView] : []
 	)
 
+	// Filter to only show categories with items
+	const availableCategories = $derived(
+		categories.filter((cat) => contextByCategory[cat.id].length > 0)
+	)
+
+	// Report view changes
+	$effect(() => {
+		if (onViewChange) {
+			if (currentView === 'categories') {
+				onViewChange(availableCategories.length)
+			} else {
+				onViewChange(currentCategoryItems.length + 1)
+			}
+		}
+	})
+
 	function handleCategoryClick(categoryId: string) {
 		currentView = categoryId as typeof currentView
 	}
 
 	function handleBackClick() {
 		currentView = 'categories'
+		itemSelectedIndex = 0
 	}
+
+	function handleKeyDown(e: KeyboardEvent) {
+		if (stringSearch.length > 0) {
+			// Navigation in search view (flat list)
+			if (e.key === 'ArrowDown') {
+				e.preventDefault()
+				e.stopPropagation()
+				if (filteredAvailableContext.length > 0) {
+					itemSelectedIndex = (itemSelectedIndex + 1) % filteredAvailableContext.length
+				}
+			} else if (e.key === 'ArrowUp') {
+				e.preventDefault()
+				e.stopPropagation()
+				if (filteredAvailableContext.length > 0) {
+					itemSelectedIndex =
+						(itemSelectedIndex - 1 + filteredAvailableContext.length) %
+						filteredAvailableContext.length
+				}
+			} else if (e.key === 'Enter' || e.key === 'Tab') {
+				if (e.key === 'Tab') e.preventDefault()
+				e.stopPropagation()
+				const selectedItem = filteredAvailableContext[itemSelectedIndex]
+				if (selectedItem) {
+					onSelect(selectedItem)
+				}
+			}
+		} else if (currentView === 'categories') {
+			// Navigation in categories view
+			if (e.key === 'ArrowDown') {
+				e.preventDefault()
+				e.stopPropagation()
+				categorySelectedIndex = (categorySelectedIndex + 1) % availableCategories.length
+			} else if (e.key === 'ArrowUp') {
+				e.preventDefault()
+				e.stopPropagation()
+				categorySelectedIndex =
+					(categorySelectedIndex - 1 + availableCategories.length) % availableCategories.length
+			} else if (e.key === 'Enter' || e.key === 'ArrowRight' || e.key === 'Tab') {
+				e.preventDefault()
+				e.stopPropagation()
+				const selectedCategory = availableCategories[categorySelectedIndex]
+				if (selectedCategory) {
+					handleCategoryClick(selectedCategory.id)
+				}
+			} else if (e.key === 'Escape' || e.key === 'ArrowLeft') {
+				e.preventDefault()
+				e.stopPropagation()
+				setShowing?.(false)
+			}
+		} else {
+			// Navigation in category items view
+			if (e.key === 'ArrowDown') {
+				e.preventDefault()
+				e.stopPropagation()
+				if (currentCategoryItems.length > 0) {
+					itemSelectedIndex = (itemSelectedIndex + 1) % currentCategoryItems.length
+				}
+			} else if (e.key === 'ArrowUp') {
+				e.preventDefault()
+				e.stopPropagation()
+				if (currentCategoryItems.length > 0) {
+					itemSelectedIndex =
+						(itemSelectedIndex - 1 + currentCategoryItems.length) % currentCategoryItems.length
+				}
+			} else if (e.key === 'Enter' || e.key === 'Tab') {
+				if (e.key === 'Tab') e.preventDefault()
+				e.stopPropagation()
+				const selectedItem = currentCategoryItems[itemSelectedIndex]
+				if (selectedItem) {
+					onSelect(selectedItem)
+					currentView = 'categories' // Go back to categories after selection
+				}
+			} else if (e.key === 'ArrowLeft' || e.key === 'Escape') {
+				e.preventDefault()
+				e.stopPropagation()
+				handleBackClick()
+			}
+		}
+	}
+
+	// Listen for keyboard events
+	$effect(() => {
+		document.addEventListener('keydown', handleKeyDown)
+		return () => {
+			document.removeEventListener('keydown', handleKeyDown)
+		}
+	})
+
+	$effect(() => {
+		if (stringSearch.length > 0) {
+			itemSelectedIndex = 0
+		}
+	})
 </script>
 
-<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}
+<div
+	class="flex flex-col gap-1 text-tertiary text-xs p-1 pr-0 min-w-24 max-h-48 overflow-y-scroll"
+	onmousedown={(e) =>
+		// avoids triggering onblur on the textinput and closing the tooltip
+		e.preventDefault()}
+	role="listbox"
+	tabindex={0}
+>
+	{#if stringSearch.length > 0}
+		<!-- Search view - show flat list -->
 		{#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 transition-colors {i ===
-				selectedIndex
+				itemSelectedIndex
 					? 'bg-surface-hover'
 					: ''}"
 				onclick={() => {
@@ -159,5 +223,63 @@
 				</span>
 			</button>
 		{/each}
+		{#if filteredAvailableContext.length === 0}
+			<div class="text-center text-tertiary text-xs py-2">No matching context</div>
+		{/if}
+	{:else if currentView === 'categories'}
+		<!-- Categories view -->
+		{#each availableCategories as category, i}
+			{@const Icon = category.icon}
+			<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 {i ===
+				categorySelectedIndex
+					? 'bg-surface-hover'
+					: ''}"
+				onclick={() => handleCategoryClick(category.id)}
+			>
+				<Icon size={16} />
+				<span class="flex-1">{category.label}</span>
+				<ChevronRight size={16} />
+			</button>
+		{/each}
+		{#if availableCategories.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, 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 transition-colors {i ===
+					itemSelectedIndex
+						? 'bg-surface-hover'
+						: ''}"
+					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.title.replace(/_/g, ' ') : element.title}
+					</span>
+				</button>
+			{/each}
+		{/if}
 	{/if}
 </div>

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

@@ -42,7 +42,7 @@
 			<button onclick={isDeletable ? onDelete : undefined} class:cursor-default={!isDeletable}>
 				{#if showDelete && isDeletable}
 					<X size={16} />
-				{:else if contextElement.type === 'flow_module'}
+				{:else if contextElement.type === 'flow_module' || contextElement.type === 'flow_module_code_piece'}
 					<FlowModuleIcon module={contextElement as FlowModule} size={16} />
 				{:else}
 					{@const SvelteComponent = icon}
@@ -50,7 +50,7 @@
 				{/if}
 			</button>
 			<span class="truncate">
-				{contextElement.type === 'diff' || contextElement.type === 'flow_module'
+				{contextElement.type === 'diff'
 					? contextElement.title.replace(/_/g, ' ')
 					: contextElement.title}
 			</span>
@@ -78,7 +78,7 @@
 					<div class="text-tertiary">Not loaded yet</div>
 				{/if}
 			</div>
-		{:else if contextElement.type === 'code' || contextElement.type === 'code_piece' || contextElement.type === 'diff'}
+		{:else if contextElement.type === 'code' || contextElement.type === 'code_piece' || contextElement.type === 'diff' || contextElement.type === 'flow_module_code_piece'}
 			<div class="max-w-96 max-h-[300px] text-xs overflow-auto">
 				<HighlightCode
 					language={contextElement.lang}

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

@@ -2,7 +2,7 @@ import { ResourceService, type Flow, type ListResourceResponse, type ScriptLang
 import { scriptLangToEditorLang } from '$lib/scripts'
 import { SQLSchemaLanguages, type DBSchemas } from '$lib/stores'
 import { diffLines } from 'diff'
-import type { ContextElement } from './context'
+import type { ContextElement, FlowModuleElement } from './context'
 import type { FlowModule } from '$lib/gen'
 
 import type { DisplayMessage } from './shared'
@@ -109,7 +109,7 @@ export default class ContextManager {
 				newAvailableContext.push({
 					type: 'flow_module',
 					id: module.id,
-					title: `module_[${module.id}]`,
+					title: `${module.id}`,
 					value: {
 						language: 'language' in module.value ? module.value.language : 'bunnative',
 						path: 'path' in module.value ? module.value.path : '',
@@ -290,24 +290,51 @@ export default class ContextManager {
 	}
 
 	addSelectedLinesToContext(lines: string, startLine: number, endLine: number, moduleId?: string) {
-		const title = moduleId ? `[${moduleId}] L${startLine}-L${endLine}` : `L${startLine}-L${endLine}`
+		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 === title)
+			this.selectedContext.find(
+				(c) =>
+					(c.type === 'code_piece' && c.title === title) ||
+					(c.type === 'flow_module_code_piece' && c.id === moduleId && c.title === title)
+			)
 		) {
 			return
 		}
-		this.selectedContext = [
-			...this.selectedContext,
-			{
-				type: 'code_piece',
-				title: title,
-				startLine,
-				endLine,
-				content: lines,
-				lang: this.scriptOptions.lang
+		if (moduleId) {
+			const module = [...this.availableContext, ...this.selectedContext].find(
+				(c) => c.type === 'flow_module' && c.id === moduleId
+			) as FlowModuleElement
+			if (!module) {
+				console.error('Module not found', moduleId)
+				return
 			}
-		]
+			this.selectedContext = [
+				...this.selectedContext,
+				{
+					type: 'flow_module_code_piece',
+					id: moduleId,
+					title: title,
+					startLine,
+					endLine,
+					content: lines,
+					lang: this.scriptOptions.lang,
+					value: module.value
+				}
+			]
+		} else {
+			this.selectedContext = [
+				...this.selectedContext,
+				{
+					type: 'code_piece',
+					title: title,
+					startLine,
+					endLine,
+					content: lines,
+					lang: this.scriptOptions.lang
+				}
+			]
+		}
 	}
 
 	setFixContext() {

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

@@ -38,7 +38,7 @@
 	let tooltipPosition = $state({ x: 0, y: 0 })
 	let textarea = $state<HTMLTextAreaElement | undefined>(undefined)
 	let tooltipElement = $state<HTMLDivElement | undefined>(undefined)
-	let selectedSuggestionIndex = $state(0)
+	let tooltipCurrentViewNumber = $state(0)
 
 	// Properties to copy for caret position calculation
 	const properties = [
@@ -182,27 +182,19 @@
 		showContextTooltip = false
 	}
 
-	async function updateTooltipPosition(
-		availableContext: ContextElement[],
-		showContextTooltip: boolean,
-		contextTooltipWord: string
-	) {
-		if (!textarea || !showContextTooltip) return
+	async function updateTooltipPosition(currentViewItemsNumber: number) {
+		if (!textarea) return
 
 		try {
 			const coords = getCaretCoordinates(textarea, textarea.selectionEnd)
 			const rect = textarea.getBoundingClientRect()
 
-			const filteredAvailableContext = availableContext.filter(
-				(c) => !contextTooltipWord || c.title.toLowerCase().includes(contextTooltipWord.slice(1))
-			)
-
 			const itemHeight = 28 // Estimated height of one item + gap (Button: p-1(8px) + text-xs(16px) = 24px; Parent: gap-1(4px) = 28px)
 			const containerPadding = 8 // p-1 top + p-1 bottom = 4px + 4px = 8px
 			const maxHeight = 192 + containerPadding // max-h-48 (192px) + containerPadding (8px)
 
 			// Calculate uncapped height, subtract gap from last item as it's not needed
-			const numItems = filteredAvailableContext.length
+			const numItems = currentViewItemsNumber
 			let uncappedHeight =
 				numItems > 0 ? numItems * itemHeight - 4 + containerPadding : containerPadding
 			// Ensure height is at least containerPadding even if no items
@@ -270,68 +262,33 @@
 		} else {
 			showContextTooltip = false
 			contextTooltipWord = ''
-			selectedSuggestionIndex = 0
-		}
-	}
-
-	function handleKeyPress(e: KeyboardEvent) {
-		if (e.key === 'Enter' && !e.shiftKey) {
-			e.preventDefault()
-			if (contextTooltipWord) {
-				const filteredContext = availableContext.filter(
-					(c) => !contextTooltipWord || c.title.toLowerCase().includes(contextTooltipWord.slice(1))
-				)
-				const contextElement = filteredContext[selectedSuggestionIndex]
-				if (contextElement) {
-					const isInSelectedContext = selectedContext.find(
-						(c) => c.title === contextElement.title && c.type === contextElement.type
-					)
-					// If the context element is already in the selected context and the last word in the instructions is the same as the context element title, send request
-					if (isInSelectedContext && value.split(' ').pop() === '@' + contextElement.title) {
-						onSendRequest()
-						return
-					}
-					handleContextSelection(contextElement)
-				} else if (contextTooltipWord === '@' && availableContext.length > 0) {
-					handleContextSelection(availableContext[0])
-				}
-			} else {
-				onSendRequest()
-			}
 		}
 	}
 
 	function handleKeyDown(e: KeyboardEvent) {
+		// Pass to parent first if provided
 		if (onKeyDown) {
 			onKeyDown(e)
 		}
 
-		if (!showContextTooltip) return
-
-		const filteredContext = availableContext.filter(
-			(c) => !contextTooltipWord || c.title.toLowerCase().includes(contextTooltipWord.slice(1))
-		)
-
-		if (e.key === 'Tab') {
-			e.preventDefault()
-			const contextElement = filteredContext[selectedSuggestionIndex]
-			if (contextElement) {
-				handleContextSelection(contextElement)
+		if (showContextTooltip) {
+			// avoid new line after Enter in the tooltip
+			if (e.key === 'Enter') {
+				e.preventDefault()
 			}
+			return
 		}
 
-		if (e.key === 'ArrowDown') {
+		if (e.key === 'Enter' && !e.shiftKey) {
 			e.preventDefault()
-			selectedSuggestionIndex = (selectedSuggestionIndex + 1) % filteredContext.length
-		} else if (e.key === 'ArrowUp') {
-			e.preventDefault()
-			selectedSuggestionIndex =
-				(selectedSuggestionIndex - 1 + filteredContext.length) % filteredContext.length
+			onSendRequest()
 		}
 	}
 
 	$effect(() => {
-		updateTooltipPosition(availableContext, showContextTooltip, contextTooltipWord)
+		if (showContextTooltip) {
+			updateTooltipPosition(tooltipCurrentViewNumber)
+		}
 	})
 
 	export function focus() {
@@ -352,7 +309,6 @@
 	</div>
 	<textarea
 		bind:this={textarea}
-		onkeypress={handleKeyPress}
 		onkeydown={handleKeyDown}
 		bind:value
 		use:autosize
@@ -388,7 +344,12 @@
 				}}
 				showAllAvailable={true}
 				stringSearch={contextTooltipWord.slice(1)}
-				selectedIndex={selectedSuggestionIndex}
+				onViewChange={(newNumber) => {
+					tooltipCurrentViewNumber = newNumber
+				}}
+				setShowing={(showing) => {
+					showContextTooltip = showing
+				}}
 			/>
 		</div>
 	</Portal>

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

@@ -25,6 +25,14 @@
 					return obj
 				}
 			}
+			for (const key in obj) {
+				try {
+					const parsed = JSON.parse(obj[key])
+					obj[key] = parsed
+				} catch (e) {
+					console.error('Failed to parse JSON:', e)
+				}
+			}
 			return JSON.stringify(obj, null, 2)
 		} catch {
 			return String(obj)
@@ -81,7 +89,9 @@
 			<div
 				class="bg-surface-secondary border border-gray-200 dark:border-gray-700 rounded p-3 overflow-x-auto max-h-64 overflow-y-auto"
 			>
-				<pre class="text-2xs text-primary whitespace-pre-wrap">{formatJson(content)}</pre>
+				<pre class="text-2xs text-primary whitespace-pre-wrap"
+					>{formatJson($state.snapshot(content))}</pre
+				>
 			</div>
 		{:else}
 			<div

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

@@ -25,8 +25,8 @@
 	<!-- Collapsible Header -->
 	<button
 		class={twMerge(
-			"w-full p-3 bg-surface-secondary hover:bg-surface-hover transition-colors flex items-center justify-between text-left border-b border-gray-200 dark:border-gray-700",
-			message.needsConfirmation ? "opacity-80" : ""
+			'w-full p-3 bg-surface-secondary hover:bg-surface-hover transition-colors flex items-center justify-between text-left border-b border-gray-200 dark:border-gray-700',
+			message.needsConfirmation ? 'opacity-80' : ''
 		)}
 		onclick={() => (isExpanded = !isExpanded)}
 		disabled={!message.showDetails}
@@ -57,7 +57,7 @@
 	{#if isExpanded}
 		<div class="p-3 bg-surface space-y-3">
 			<!-- Parameters Section -->
-			<div class={message.needsConfirmation ? "opacity-80" : ""}>
+			<div class={message.needsConfirmation ? 'opacity-80' : ''}>
 				<ToolContentDisplay title="Parameters" content={message.parameters} />
 			</div>
 

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

@@ -48,7 +48,7 @@ export interface CodePieceElement {
 	lang: ScriptLang | 'bunnative'
 }
 
-export interface FlowModule {
+export interface FlowModuleElement {
 	type: 'flow_module'
 	id: string
 	title: string
@@ -61,13 +61,20 @@ export interface FlowModule {
 	}
 }
 
+export interface FlowModuleCodePieceElement extends Omit<CodePieceElement, 'type'> {
+	type: 'flow_module_code_piece'
+	id: string
+	value: FlowModuleElement['value']
+}
+
 export type ContextElement = (
 	| CodeElement
 	| ErrorElement
 	| DBElement
 	| DiffElement
 	| CodePieceElement
-	| FlowModule
+	| FlowModuleElement
+	| FlowModuleCodePieceElement
 ) & {
 	deletable?: boolean
 }

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

@@ -860,7 +860,7 @@ 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 },
-	selectedContext?: ContextElement[]
+	selectedContext: ContextElement[] = []
 ): ChatCompletionUserMessageParam {
 	const flow = flowAndSelectedId?.flow
 	const selectedId = flowAndSelectedId?.selectedId
@@ -877,7 +877,7 @@ ${instructions}`
 		}
 	}
 
-	const codePieces = selectedContext?.filter((c) => c.type === 'code_piece') ?? []
+	const codePieces = selectedContext.filter((c) => c.type === 'flow_module_code_piece')
 	const flowModulesYaml = applyCodePiecesToFlowModules(codePieces, flow.value.modules)
 
 	let flowContent = `## FLOW:

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

@@ -4,7 +4,7 @@ import type {
 	ChatCompletionTool
 } from 'openai/resources/chat/completions.mjs'
 import { get } from 'svelte/store'
-import type { CodePieceElement, ContextElement } from './context'
+import type { CodePieceElement, ContextElement, FlowModuleCodePieceElement } from './context'
 import { copilotSessionModel, workspaceStore } from '$lib/stores'
 import type { ExtendedOpenFlow } from '$lib/components/flows/types'
 import type { FunctionParameters } from 'openai/resources/shared.mjs'
@@ -23,6 +23,24 @@ export interface ContextStringResult {
 	hasFlowModule: boolean
 }
 
+export const extractAllModules = (modules: FlowModule[]): FlowModule[] => {
+	return modules.flatMap((m) => {
+		if (m.value.type === 'forloopflow' || m.value.type === 'whileloopflow') {
+			return [m, ...extractAllModules(m.value.modules)]
+		}
+		if (m.value.type === 'branchall') {
+			return [m, ...extractAllModules(m.value.branches.flatMap((b) => b.modules))]
+		}
+		if (m.value.type === 'branchone') {
+			return [
+				m,
+				...extractAllModules([...m.value.branches.flatMap((b) => b.modules), ...m.value.default])
+			]
+		}
+		return [m]
+	})
+}
+
 export const findModuleById = (modules: FlowModule[], moduleId: string): FlowModule | undefined => {
 	for (const module of modules) {
 		if (module.id === moduleId) {
@@ -68,22 +86,16 @@ const applyCodePieceToCodeContext = (codePieces: CodePieceElement[], codeContext
 }
 
 export function applyCodePiecesToFlowModules(
-	codePieces: CodePieceElement[],
+	codePieces: FlowModuleCodePieceElement[],
 	flowModules: FlowModule[]
 ): string {
-	// Parse code piece titles to extract module IDs
-	// Format: "[id] L3-L5"
-	const moduleCodePieces = new Map<string, CodePieceElement[]>()
-
+	const moduleCodePieces = new Map<string, FlowModuleCodePieceElement[]>()
 	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)
+		const moduleId = codePiece.id
+		if (!moduleCodePieces.has(moduleId)) {
+			moduleCodePieces.set(moduleId, [])
 		}
+		moduleCodePieces.get(moduleId)!.push(codePiece)
 	}
 
 	// Clone modules to avoid mutation
@@ -93,7 +105,10 @@ export function applyCodePiecesToFlowModules(
 	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)
+			module.value.content = applyCodePieceToCodeContext(
+				pieces as unknown as CodePieceElement[],
+				module.value.content
+			)
 		}
 	}
 

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

@@ -19,6 +19,7 @@
 	import type { ModulesTestStates } from '../modulesTest.svelte'
 	import type { StateStore } from '$lib/utils'
 	import type { FlowOptions } from '../copilot/chat/ContextManager.svelte'
+	import { extractAllModules } from '../copilot/chat/shared'
 	const { flowStore } = getContext<FlowEditorContext>('FlowEditorContext')
 
 	interface Props {
@@ -110,7 +111,7 @@
 			lastDeployedFlow: savedFlow,
 			lastSavedFlow: savedFlow?.draft,
 			path: savedFlow?.path,
-			modules: flowStore.val.value.modules
+			modules: extractAllModules(flowStore.val.value.modules)
 		}
 		aiChatManager.flowOptions = options
 	})