From 56b9fd4eb7d7e84cfd340d3ac7a0117719b82acf Mon Sep 17 00:00:00 2001
From: Ruben Fiszel <ruben@windmill.dev>
Date: Fri, 26 Dec 2025 14:15:08 +0000
Subject: [PATCH] feat: allow @ selection for raw apps

---
 .../lib/components/copilot/chat/AIChat.svelte |   4 +-
 .../copilot/chat/AIChatInput.svelte           | 308 ++++++++++++++++++
 .../copilot/chat/AIChatManager.svelte.ts      | 110 ++++++-
 .../chat/AppAvailableContextList.svelte       | 289 ++++++++++++++++
 .../copilot/chat/ContextElementBadge.svelte   |  44 +++
 .../chat/__tests__/app/appEvalHelpers.ts      |   8 +
 .../lib/components/copilot/chat/app/core.ts   |  99 +++++-
 .../lib/components/copilot/chat/context.ts    |  69 +++-
 .../components/raw_apps/RawAppEditor.svelte   |  48 ++-
 9 files changed, 968 insertions(+), 11 deletions(-)
 create mode 100644 frontend/src/lib/components/copilot/chat/AppAvailableContextList.svelte

diff --git a/frontend/src/lib/components/copilot/chat/AIChat.svelte b/frontend/src/lib/components/copilot/chat/AIChat.svelte
index 9e41034fbe..ff31146391 100644
--- a/frontend/src/lib/components/copilot/chat/AIChat.svelte
+++ b/frontend/src/lib/components/copilot/chat/AIChat.svelte
@@ -98,7 +98,9 @@
 		() => aiChatManager.contextManager.getSelectedContext(),
 		(sc) => aiChatManager.contextManager.setSelectedContext(sc)
 	}
-	availableContext={aiChatManager.contextManager.getAvailableContext()}
+	availableContext={aiChatManager.mode === AIMode.APP
+		? aiChatManager.getAppAvailableContext()
+		: aiChatManager.contextManager.getAvailableContext()}
 	messages={aiChatManager.currentReply
 		? [
 				...aiChatManager.displayMessages,
diff --git a/frontend/src/lib/components/copilot/chat/AIChatInput.svelte b/frontend/src/lib/components/copilot/chat/AIChatInput.svelte
index cf609a5848..457e315d94 100644
--- a/frontend/src/lib/components/copilot/chat/AIChatInput.svelte
+++ b/frontend/src/lib/components/copilot/chat/AIChatInput.svelte
@@ -1,6 +1,7 @@
 <script lang="ts">
 	import Popover from '$lib/components/meltComponents/Popover.svelte'
 	import AvailableContextList from './AvailableContextList.svelte'
+	import AppAvailableContextList from './AppAvailableContextList.svelte'
 	import ContextElementBadge from './ContextElementBadge.svelte'
 	import ContextTextarea from './ContextTextarea.svelte'
 	import autosize from '$lib/autosize'
@@ -8,6 +9,9 @@
 	import { aiChatManager, AIMode } from './AIChatManager.svelte'
 	import { twMerge } from 'tailwind-merge'
 	import type { Snippet } from 'svelte'
+	import Portal from '$lib/components/Portal.svelte'
+	import { zIndexes } from '$lib/zIndexes'
+	import { tick } from 'svelte'
 
 	interface Props {
 		availableContext: ContextElement[]
@@ -75,6 +79,13 @@
 	let instructionsTextareaComponent: HTMLTextAreaElement | undefined = $state()
 	let instructions = $state(initialInstructions)
 
+	// App mode @ mention state
+	let showAppContextTooltip = $state(false)
+	let appContextTooltipWord = $state('')
+	let appTooltipPosition = $state({ x: 0, y: 0 })
+	let appTooltipElement = $state<HTMLDivElement | undefined>(undefined)
+	let appTooltipCurrentViewNumber = $state(0)
+
 	export function focusInput() {
 		if (aiChatManager.mode === AIMode.SCRIPT || aiChatManager.mode === AIMode.FLOW) {
 			contextTextareaComponent?.focus()
@@ -110,6 +121,19 @@
 			)
 		) {
 			selectedContext = [...selectedContext, contextElement]
+
+			// If it's a datatable table, add it to the app's whitelisted tables
+			if (
+				contextElement.type === 'app_datatable' &&
+				aiChatManager.mode === AIMode.APP &&
+				aiChatManager.appAiChatHelpers
+			) {
+				aiChatManager.appAiChatHelpers.addTableToWhitelist(
+					contextElement.datatableName,
+					contextElement.schemaName,
+					contextElement.tableName
+				)
+			}
 		}
 	}
 
@@ -131,6 +155,198 @@
 			focusInput()
 		}
 	})
+
+	// Properties to copy for caret position calculation (app mode)
+	const caretProperties = [
+		'direction',
+		'boxSizing',
+		'width',
+		'height',
+		'overflowX',
+		'overflowY',
+		'borderTopWidth',
+		'borderRightWidth',
+		'borderBottomWidth',
+		'borderLeftWidth',
+		'borderStyle',
+		'paddingTop',
+		'paddingRight',
+		'paddingBottom',
+		'paddingLeft',
+		'fontStyle',
+		'fontVariant',
+		'fontWeight',
+		'fontStretch',
+		'fontSize',
+		'fontSizeAdjust',
+		'lineHeight',
+		'fontFamily',
+		'textAlign',
+		'textTransform',
+		'textIndent',
+		'textDecoration',
+		'letterSpacing',
+		'wordSpacing',
+		'tabSize',
+		'MozTabSize'
+	]
+
+	function getCaretCoordinates(element: HTMLTextAreaElement, position: number) {
+		const div = document.createElement('div')
+		div.id = 'input-textarea-caret-position-mirror-div'
+		document.body.appendChild(div)
+
+		const style = div.style
+		const computed = window.getComputedStyle(element)
+		const isInput = element.nodeName === 'INPUT'
+
+		style.whiteSpace = 'pre-wrap'
+		if (!isInput) style.wordWrap = 'break-word'
+
+		style.position = 'absolute'
+		style.visibility = 'hidden'
+
+		caretProperties.forEach(function (prop) {
+			if (isInput && prop === 'lineHeight') {
+				if (computed.boxSizing === 'border-box') {
+					const height = parseInt(computed.height)
+					const outerHeight =
+						parseInt(computed.paddingTop) +
+						parseInt(computed.paddingBottom) +
+						parseInt(computed.borderTopWidth) +
+						parseInt(computed.borderBottomWidth)
+					const targetHeight = outerHeight + parseInt(computed.lineHeight)
+					if (height > targetHeight) {
+						style.lineHeight = height - outerHeight + 'px'
+					} else if (height === targetHeight) {
+						style.lineHeight = computed.lineHeight
+					} else {
+						style.lineHeight = '0'
+					}
+				} else {
+					style.lineHeight = computed.height
+				}
+			} else {
+				style[prop] = computed[prop]
+			}
+		})
+
+		const isFirefox =
+			(window as typeof window & { mozInnerScreenX: number }).mozInnerScreenX != null
+		if (isFirefox) {
+			if (element.scrollHeight > parseInt(computed.height)) style.overflowY = 'scroll'
+		} else {
+			style.overflow = 'hidden'
+		}
+
+		div.textContent = element.value.substring(0, position)
+
+		if (isInput) div.textContent = div.textContent.replace(/\s/g, '\u00a0')
+
+		const span = document.createElement('span')
+		span.textContent = element.value.substring(position) || '.'
+		div.appendChild(span)
+
+		const coordinates = {
+			top: span.offsetTop + parseInt(computed['borderTopWidth']),
+			left: span.offsetLeft + parseInt(computed['borderLeftWidth']),
+			height: parseInt(computed['lineHeight'])
+		}
+
+		document.body.removeChild(div)
+
+		return coordinates
+	}
+
+	async function updateAppTooltipPosition(currentViewItemsNumber: number) {
+		if (!instructionsTextareaComponent) return
+
+		try {
+			const coords = getCaretCoordinates(
+				instructionsTextareaComponent,
+				instructionsTextareaComponent.selectionEnd
+			)
+			const rect = instructionsTextareaComponent.getBoundingClientRect()
+
+			const itemHeight = 28
+			const containerPadding = 8
+			const maxHeight = 192 + containerPadding
+
+			const numItems = currentViewItemsNumber
+			let uncappedHeight =
+				numItems > 0 ? numItems * itemHeight - 4 + containerPadding : containerPadding
+			uncappedHeight = Math.max(uncappedHeight, containerPadding)
+
+			const estimatedTooltipHeight = Math.min(uncappedHeight, maxHeight)
+			const margin = 6
+
+			let finalX = rect.left + coords.left - 70
+			let finalY: number
+
+			if (isFirstMessage) {
+				finalY = rect.top + coords.top + coords.height - 3
+			} else {
+				finalY = rect.top + coords.top - estimatedTooltipHeight - margin
+			}
+
+			appTooltipPosition = {
+				x: finalX,
+				y: finalY
+			}
+
+			await tick()
+
+			if (appTooltipElement) {
+				const tooltipRect = appTooltipElement.getBoundingClientRect()
+				const tooltipWidth = tooltipRect.width
+
+				if (finalX + tooltipWidth > window.innerWidth) {
+					finalX = Math.max(10, window.innerWidth - tooltipWidth - 10)
+
+					appTooltipPosition = {
+						x: finalX,
+						y: finalY
+					}
+				}
+			}
+		} catch (error) {
+			console.error('Error updating tooltip position', error)
+			showAppContextTooltip = false
+		}
+	}
+
+	function handleAppInput(_e: Event) {
+		const words = instructions.split(/\s+/)
+		const lastWord = words[words.length - 1]
+
+		if (
+			lastWord.startsWith('@') &&
+			(!availableContext.find((c) => c.title === lastWord.slice(1)) ||
+				!selectedContext.find((c) => c.title === lastWord.slice(1)))
+		) {
+			showAppContextTooltip = true
+			appContextTooltipWord = lastWord
+		} else {
+			showAppContextTooltip = false
+			appContextTooltipWord = ''
+		}
+	}
+
+	function handleAppContextSelection(contextElement: ContextElement) {
+		addContextToSelection(contextElement)
+		// Update instructions with the selected context title
+		const index = instructions.lastIndexOf('@')
+		if (index !== -1) {
+			instructions = instructions.substring(0, index) + `@${contextElement.title}`
+		}
+		showAppContextTooltip = false
+	}
+
+	$effect(() => {
+		if (showAppContextTooltip) {
+			updateAppTooltipPosition(appTooltipCurrentViewNumber)
+		}
+	})
 </script>
 
 <div use:clickOutside class="relative">
@@ -185,6 +401,98 @@
 			{disabled}
 			{onKeyDown}
 		/>
+	{:else if aiChatManager.mode === AIMode.APP}
+		{#if showContext}
+			<div class="flex flex-row gap-1 mb-1 overflow-scroll pt-2 no-scrollbar">
+				<Popover>
+					<svelte:fragment slot="trigger">
+						<div
+							class="border rounded-md px-1 py-0.5 font-normal text-primary text-xs hover:bg-surface-hover bg-surface"
+							>@</div
+						>
+					</svelte:fragment>
+					<svelte:fragment slot="content" let:close>
+						<AppAvailableContextList
+							{availableContext}
+							{selectedContext}
+							onSelect={(element) => {
+								addContextToSelection(element)
+								close()
+							}}
+						/>
+					</svelte:fragment>
+				</Popover>
+				{#each selectedContext as element (element.type + '-' + element.title)}
+					<ContextElementBadge
+						contextElement={element}
+						deletable
+						onDelete={() => {
+							selectedContext = selectedContext?.filter(
+								(c) => c.type !== element.type || c.title !== element.title
+							)
+						}}
+					/>
+				{/each}
+			</div>
+		{/if}
+		<div class={twMerge('relative w-full scroll-pb-2', className)}>
+			<textarea
+				bind:this={instructionsTextareaComponent}
+				bind:value={instructions}
+				use:autosize
+				oninput={handleAppInput}
+				onblur={() => {
+					setTimeout(() => {
+						showAppContextTooltip = false
+					}, 200)
+				}}
+				onkeydown={(e) => {
+					if (onKeyDown) {
+						onKeyDown(e)
+					}
+					if (showAppContextTooltip) {
+						// avoid new line after Enter in the tooltip
+						if (e.key === 'Enter') {
+							e.preventDefault()
+						}
+						return
+					}
+					if (e.key === 'Enter' && !e.shiftKey) {
+						e.preventDefault()
+						sendRequest()
+					}
+				}}
+				rows={3}
+				placeholder={modePlaceholder}
+				class="resize-none"
+				{disabled}
+			></textarea>
+		</div>
+		{#if showAppContextTooltip}
+			<Portal target="body">
+				<div
+					bind:this={appTooltipElement}
+					class="absolute bg-white dark:bg-gray-800 border border-gray-200 dark:border-gray-700 rounded-md shadow-lg"
+					style="left: {appTooltipPosition.x}px; top: {appTooltipPosition.y}px; z-index: {zIndexes.tooltip};"
+				>
+					<AppAvailableContextList
+						{availableContext}
+						{selectedContext}
+						onSelect={(element) => {
+							handleAppContextSelection(element)
+						}}
+						showAllAvailable={true}
+						stringSearch={appContextTooltipWord.slice(1)}
+						onViewChange={(newNumber) => {
+							appTooltipCurrentViewNumber = newNumber
+						}}
+						setShowing={(showing) => {
+							showAppContextTooltip = showing
+						}}
+					/>
+				</div>
+			</Portal>
+		{/if}
 	{:else}
 		<div class={twMerge('relative w-full scroll-pb-2 pt-2', className)}>
 			<textarea
diff --git a/frontend/src/lib/components/copilot/chat/AIChatManager.svelte.ts b/frontend/src/lib/components/copilot/chat/AIChatManager.svelte.ts
index 1d5110fc94..2e0d05ee0f 100644
--- a/frontend/src/lib/components/copilot/chat/AIChatManager.svelte.ts
+++ b/frontend/src/lib/components/copilot/chat/AIChatManager.svelte.ts
@@ -45,7 +45,12 @@ import { untrack } from 'svelte'
 import { type DBSchemas } from '$lib/stores'
 import { askTools, prepareAskSystemMessage, prepareAskUserMessage } from './ask/core'
 import { chatState, DEFAULT_SIZE, triggerablesByAi } from './sharedChatState.svelte'
-import type { ContextElement } from './context'
+import type {
+	ContextElement,
+	AppFrontendFileElement,
+	AppBackendRunnableElement,
+	AppDatatableElement
+} from './context'
 import type { Selection } from 'monaco-editor'
 import type AIChatInput from './AIChatInput.svelte'
 import { prepareApiSystemMessage, prepareApiUserMessage } from './api/core'
@@ -110,6 +115,8 @@ class AIChatManager {
 	pendingNewCode = $state<string | undefined>(undefined)
 	apiTools = $state<Tool<any>[]>([])
 	aiChatInput = $state<AIChatInput | null>(null)
+	/** Cached datatables for app context (fetched asynchronously) */
+	cachedDatatables = $state<AppDatatableElement[]>([])
 
 	private confirmationCallback = $state<((value: boolean) => void) | undefined>(undefined)
 
@@ -713,7 +720,8 @@ class AIChatManager {
 				case AIMode.APP:
 					userMessage = prepareAppUserMessage(
 						oldInstructions,
-						this.appAiChatHelpers?.getSelectedContext()
+						this.appAiChatHelpers?.getSelectedContext(),
+						oldSelectedContext
 					)
 					break
 			}
@@ -1089,11 +1097,109 @@ class AIChatManager {
 		}
 	}
 
+	/**
+	 * Refresh cached datatables from the app helpers (async)
+	 * Creates one context element per table (not per datatable)
+	 */
+	refreshDatatables = async (): Promise<void> => {
+		if (!this.appAiChatHelpers) {
+			this.cachedDatatables = []
+			return
+		}
+
+		try {
+			const datatables = await this.appAiChatHelpers.getDatatables()
+			console.log('Refreshed datatables:', datatables)
+
+			// Flatten to individual tables
+			const tableElements: AppDatatableElement[] = []
+			for (const dt of datatables) {
+				if (dt.error) {
+					// Skip datatables with errors
+					continue
+				}
+				for (const [schemaName, tables] of Object.entries(dt.schemas)) {
+					for (const [tableName, columns] of Object.entries(tables)) {
+						// Format title as "datatable/schema:table" or "datatable/table" if schema is public
+						const title =
+							schemaName === 'public'
+								? `${dt.datatable_name}/${tableName}`
+								: `${dt.datatable_name}/${schemaName}:${tableName}`
+						tableElements.push({
+							type: 'app_datatable',
+							datatableName: dt.datatable_name,
+							schemaName,
+							tableName,
+							title,
+							columns
+						})
+					}
+				}
+			}
+			this.cachedDatatables = tableElements
+		} catch (err) {
+			console.error('Failed to refresh datatables:', err)
+			this.cachedDatatables = []
+		}
+	}
+
+	/**
+	 * Get available context elements for app mode (frontend files + backend runnables + datatables)
+	 */
+	getAppAvailableContext = (): ContextElement[] => {
+		if (!this.appAiChatHelpers) {
+			return []
+		}
+
+		const context: ContextElement[] = []
+
+		// Add frontend files
+		const frontendFiles = this.appAiChatHelpers.listFrontendFiles()
+		for (const path of frontendFiles) {
+			const content = this.appAiChatHelpers.getFrontendFile(path)
+			if (content !== undefined) {
+				const element: AppFrontendFileElement = {
+					type: 'app_frontend_file',
+					path,
+					title: path,
+					content
+				}
+				context.push(element)
+			}
+		}
+
+		// Add backend runnables
+		const runnables = this.appAiChatHelpers.listBackendRunnables()
+		for (const { key } of runnables) {
+			const runnable = this.appAiChatHelpers.getBackendRunnable(key)
+			if (runnable) {
+				const element: AppBackendRunnableElement = {
+					type: 'app_backend_runnable',
+					key,
+					title: key,
+					runnable
+				}
+				context.push(element)
+			}
+		}
+
+		// Add cached datatables
+		context.push(...this.cachedDatatables)
+
+		return context
+	}
+
 	setAppHelpers = (appHelpers: AppAIChatHelpers) => {
 		this.appAiChatHelpers = appHelpers
+		// Refresh datatables when app helpers are set (deferred to avoid loop)
+		// Use setTimeout to ensure this runs after the effect completes
+		setTimeout(() => {
+			this.refreshDatatables()
+		}, 50)
 
 		return () => {
 			this.appAiChatHelpers = undefined
+			this.cachedDatatables = []
 		}
 	}
 }
diff --git a/frontend/src/lib/components/copilot/chat/AppAvailableContextList.svelte b/frontend/src/lib/components/copilot/chat/AppAvailableContextList.svelte
new file mode 100644
index 0000000000..f120c89221
--- /dev/null
+++ b/frontend/src/lib/components/copilot/chat/AppAvailableContextList.svelte
@@ -0,0 +1,289 @@
+<script lang="ts">
+	import { ContextIconMap, type ContextElement } from './context'
+	import { ArrowLeft, ChevronRight, FileCode, Code2, Table2 } from 'lucide-svelte'
+
+	interface Props {
+		availableContext: ContextElement[]
+		selectedContext: ContextElement[]
+		onSelect: (element: ContextElement) => void
+		setShowing?: (showing: boolean) => void
+		showAllAvailable?: boolean
+		stringSearch?: string
+		onViewChange?: (newNumber: number) => void
+	}
+
+	const {
+		availableContext,
+		selectedContext,
+		onSelect,
+		setShowing,
+		showAllAvailable = false,
+		stringSearch = '',
+		onViewChange
+	}: Props = $props()
+
+	// Current view state: 'categories' or specific category type
+	let currentView = $state<'categories' | 'files' | 'runnables' | 'datatables'>('categories')
+
+	// Selected index for keyboard navigation
+	let itemSelectedIndex = $state(0)
+	let categorySelectedIndex = $state(0)
+
+	// Category definitions for app mode
+	const categories = [
+		{ id: 'files', label: 'Frontend Files', icon: FileCode },
+		{ id: 'runnables', label: 'Backend Runnables', icon: Code2 },
+		{ id: 'datatables', label: 'Datatables', icon: Table2 }
+	]
+
+	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[]> = {
+			files: [],
+			runnables: [],
+			datatables: []
+		}
+
+		filteredAvailableContext.forEach((context) => {
+			if (context.type === 'app_frontend_file') grouped.files.push(context)
+			else if (context.type === 'app_backend_runnable') grouped.runnables.push(context)
+			else if (context.type === 'app_datatable') grouped.datatables.push(context)
+		})
+
+		return grouped
+	})
+
+	const currentCategoryItems = $derived(
+		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
+		itemSelectedIndex = 0
+	}
+
+	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
+		}
+	})
+
+	function getDisplayTitle(element: ContextElement): string {
+		if (element.type === 'app_frontend_file') {
+			return element.path
+		} else if (element.type === 'app_backend_runnable') {
+			return element.key
+		} else if (element.type === 'app_datatable') {
+			// Show as datatable/table or datatable/schema:table
+			return element.title
+		}
+		return element.title
+	}
+</script>
+
+<div
+	class="flex flex-col gap-1 text-primary 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 (element.type + '-' + element.title)}
+			{@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)
+				}}
+			>
+				{#if Icon}
+					<Icon size={16} />
+				{/if}
+				<span class="truncate">
+					{getDisplayTitle(element)}
+				</span>
+			</button>
+		{/each}
+		{#if filteredAvailableContext.length === 0}
+			<div class="text-center text-primary text-xs py-2">No matching context</div>
+		{/if}
+	{:else if currentView === 'categories'}
+		<!-- Categories view -->
+		{#each availableCategories as category, i (category.id)}
+			{@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-primary 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-primary text-xs py-2">No items in this category</div>
+		{:else}
+			{#each currentCategoryItems as element, i (element.type + '-' + element.title)}
+				{@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 Icon}
+						<Icon size={16} />
+					{/if}
+					<span class="truncate">
+						{getDisplayTitle(element)}
+					</span>
+				</button>
+			{/each}
+		{/if}
+	{/if}
+</div>
diff --git a/frontend/src/lib/components/copilot/chat/ContextElementBadge.svelte b/frontend/src/lib/components/copilot/chat/ContextElementBadge.svelte
index 65ae5108b1..186414c8dd 100644
--- a/frontend/src/lib/components/copilot/chat/ContextElementBadge.svelte
+++ b/frontend/src/lib/components/copilot/chat/ContextElementBadge.svelte
@@ -100,6 +100,50 @@
 					<div class="text-primary">{contextElement.title}</div>
 				</div>
 			{/if}
+		{:else if contextElement.type === 'app_frontend_file'}
+			<div class="max-w-96 max-h-[300px] text-xs overflow-auto">
+				<HighlightCode
+					language={contextElement.path.endsWith('.tsx') || contextElement.path.endsWith('.ts')
+						? 'bun'
+						: contextElement.path.endsWith('.css')
+							? 'bash'
+							: 'bun'}
+					code={contextElement.content}
+					className="w-full p-2"
+				/>
+			</div>
+		{:else if contextElement.type === 'app_backend_runnable'}
+			<div class="p-2 max-w-96 max-h-[300px] text-xs overflow-auto">
+				{#if contextElement.runnable.inlineScript}
+					<HighlightCode
+						language={contextElement.runnable.inlineScript.language}
+						code={contextElement.runnable.inlineScript.content}
+						className="w-full p-2"
+					/>
+				{:else}
+					<ObjectViewer json={contextElement.runnable} pureViewer collapseLevel={2} />
+				{/if}
+			</div>
+		{:else if contextElement.type === 'app_code_selection'}
+			<div class="max-w-96 max-h-[300px] text-xs overflow-auto">
+				<div class="text-tertiary text-xs mb-1 px-2 pt-1">
+					{contextElement.source} (L{contextElement.startLine}-L{contextElement.endLine})
+				</div>
+				<HighlightCode
+					language="bun"
+					code={contextElement.content}
+					className="w-full p-2"
+				/>
+			</div>
+		{:else if contextElement.type === 'app_datatable'}
+			<div class="p-2 max-w-96 max-h-[300px] text-xs overflow-auto">
+				<div class="text-tertiary text-xs mb-1">
+					{contextElement.datatableName}/{contextElement.schemaName === 'public'
+						? ''
+						: contextElement.schemaName + ':'}{contextElement.tableName}
+				</div>
+				<ObjectViewer json={contextElement.columns} pureViewer collapseLevel={1} />
+			</div>
 		{/if}
 	{/snippet}
 </Popover>
diff --git a/frontend/src/lib/components/copilot/chat/__tests__/app/appEvalHelpers.ts b/frontend/src/lib/components/copilot/chat/__tests__/app/appEvalHelpers.ts
index daf934ef90..4dbc8d58a0 100644
--- a/frontend/src/lib/components/copilot/chat/__tests__/app/appEvalHelpers.ts
+++ b/frontend/src/lib/components/copilot/chat/__tests__/app/appEvalHelpers.ts
@@ -124,6 +124,14 @@ export function createAppEvalHelpers(
 		) => {
 			// Return success with empty result for eval testing
 			return { success: true, result: [] }
+		},
+
+		addTableToWhitelist: (
+			_datatableName: string,
+			_schemaName: string,
+			_tableName: string
+		) => {
+			// No-op for eval testing - tables are not tracked in test context
 		}
 	}
 
diff --git a/frontend/src/lib/components/copilot/chat/app/core.ts b/frontend/src/lib/components/copilot/chat/app/core.ts
index 6ff59f7708..6b6f4854e8 100644
--- a/frontend/src/lib/components/copilot/chat/app/core.ts
+++ b/frontend/src/lib/components/copilot/chat/app/core.ts
@@ -8,6 +8,13 @@ import { FlowService, ScriptService, type Flow, type Script } from '$lib/gen'
 import uFuzzy from '@leeoniya/ufuzzy'
 import { emptyString } from '$lib/utils'
 import { aiChatManager } from '../AIChatManager.svelte'
+import type {
+	ContextElement,
+	AppFrontendFileElement,
+	AppBackendRunnableElement,
+	AppCodeSelectionElement,
+	AppDatatableElement
+} from '../context'
 
 // Backend runnable types
 export type BackendRunnableType = 'script' | 'flow' | 'hubscript' | 'inline'
@@ -93,8 +100,10 @@ export interface SelectedContext {
 	clearInspector?: () => void
 	/** Function to clear the runnable selection (go back to frontend view) */
 	clearRunnable?: () => void
-	// Future: text selection within the file
-	// textSelection?: { startLine: number; endLine: number; startColumn: number; endColumn: number }
+	/** Code selection from the editor (either frontend or backend) */
+	codeSelection?: AppCodeSelectionElement
+	/** Function to clear the code selection */
+	clearCodeSelection?: () => void
 }
 
 /**
@@ -142,6 +151,8 @@ export interface AppAIChatHelpers {
 		sql: string,
 		newTable?: { schema: string; name: string }
 	) => Promise<{ success: boolean; result?: Record<string, any>[]; error?: string }>
+	/** Add a table to the app's whitelisted tables (called when user selects a table via @) */
+	addTableToWhitelist: (datatableName: string, schemaName: string, tableName: string) => void
 }
 
 // ============= Utility =============
@@ -1085,15 +1096,22 @@ const MAX_CONTEXT_CONTENT_LENGTH = 3000
 
 export function prepareAppUserMessage(
 	instructions: string,
-	selectedContext?: SelectedContext
+	selectedContext?: SelectedContext,
+	additionalContext?: ContextElement[]
 ): ChatCompletionUserMessageParam {
 	let content = ''
 
-	if (selectedContext && (selectedContext.type !== 'none' || selectedContext.inspectorElement)) {
+	// Check if we have any context to add
+	const hasSelectedContext =
+		selectedContext && (selectedContext.type !== 'none' || selectedContext.inspectorElement)
+	const hasAdditionalContext = additionalContext && additionalContext.length > 0
+
+	if (hasSelectedContext || hasAdditionalContext) {
 		content += `## SELECTED CONTEXT:\n`
 
 		// Add frontend file context with content (unless excluded)
 		if (
+			selectedContext &&
 			selectedContext.type === 'frontend' &&
 			selectedContext.frontendPath &&
 			!selectedContext.selectionExcluded
@@ -1111,6 +1129,7 @@ export function prepareAppUserMessage(
 
 		// Add backend runnable context with content (unless excluded)
 		if (
+			selectedContext &&
 			selectedContext.type === 'backend' &&
 			selectedContext.backendKey &&
 			!selectedContext.selectionExcluded
@@ -1139,7 +1158,7 @@ export function prepareAppUserMessage(
 		}
 
 		// Add inspector element context if available
-		if (selectedContext.inspectorElement) {
+		if (selectedContext?.inspectorElement) {
 			const el = selectedContext.inspectorElement
 			content += `\nThe user has selected an element in the app preview using the inspector tool:\n`
 			content += `- **Element**: ${el.tagName}${el.id ? `#${el.id}` : ''}${el.className ? `.${el.className.split(' ').join('.')}` : ''}\n`
@@ -1154,6 +1173,76 @@ export function prepareAppUserMessage(
 			const truncatedHtml = el.html.length > 500 ? el.html.slice(0, 500) + '...' : el.html
 			content += `- **HTML**:\n\`\`\`html\n${truncatedHtml}\n\`\`\`\n`
 		}
+
+		// Add code selection context if available
+		if (selectedContext?.codeSelection) {
+			const selection = selectedContext.codeSelection
+			content += `\n### CODE SELECTION:\n`
+			content += `The user has selected code in the ${selection.sourceType} editor:\n`
+			content += `- **File/Source**: ${selection.source}\n`
+			content += `- **Lines**: ${selection.startLine}-${selection.endLine}\n`
+			const truncatedCode =
+				selection.content.length > MAX_CONTEXT_CONTENT_LENGTH
+					? selection.content.slice(0, MAX_CONTEXT_CONTENT_LENGTH) + '\n... [TRUNCATED]'
+					: selection.content
+			content += `\`\`\`\n${truncatedCode}\n\`\`\`\n`
+		}
+
+		// Add additional context from @ mentions
+		if (additionalContext && additionalContext.length > 0) {
+			content += `\n### ADDITIONAL CONTEXT (mentioned by user):\n`
+
+			for (const ctx of additionalContext) {
+				if (ctx.type === 'app_frontend_file') {
+					const fileCtx = ctx as AppFrontendFileElement
+					content += `\n**Frontend File: ${fileCtx.path}**\n`
+					const truncatedContent =
+						fileCtx.content.length > MAX_CONTEXT_CONTENT_LENGTH
+							? fileCtx.content.slice(0, MAX_CONTEXT_CONTENT_LENGTH) + '\n... [TRUNCATED]'
+							: fileCtx.content
+					content += `\`\`\`\n${truncatedContent}\n\`\`\`\n`
+				} else if (ctx.type === 'app_backend_runnable') {
+					const runnableCtx = ctx as AppBackendRunnableElement
+					const runnable = runnableCtx.runnable
+					content += `\n**Backend Runnable: ${runnableCtx.key}**\n`
+					content += `- **Name**: ${runnable.name}\n`
+					content += `- **Type**: ${runnable.type}\n`
+					if (runnable.path) {
+						content += `- **Path**: ${runnable.path}\n`
+					}
+					if (runnable.inlineScript) {
+						const truncatedCode =
+							runnable.inlineScript.content.length > MAX_CONTEXT_CONTENT_LENGTH
+								? runnable.inlineScript.content.slice(0, MAX_CONTEXT_CONTENT_LENGTH) +
+									'\n... [TRUNCATED]'
+								: runnable.inlineScript.content
+						content += `- **Language**: ${runnable.inlineScript.language}\n`
+						content += `- **Code**:\n\`\`\`${runnable.inlineScript.language === 'bun' ? 'typescript' : 'python'}\n${truncatedCode}\n\`\`\`\n`
+					}
+					if (runnable.staticInputs && Object.keys(runnable.staticInputs).length > 0) {
+						content += `- **Static inputs**: ${JSON.stringify(runnable.staticInputs)}\n`
+					}
+				} else if (ctx.type === 'app_datatable') {
+					const datatableCtx = ctx as AppDatatableElement
+					const tableRef =
+						datatableCtx.schemaName === 'public'
+							? `${datatableCtx.datatableName}/${datatableCtx.tableName}`
+							: `${datatableCtx.datatableName}/${datatableCtx.schemaName}:${datatableCtx.tableName}`
+					content += `\n**Table: ${tableRef}**\n`
+					content += `- **Datatable**: ${datatableCtx.datatableName}\n`
+					content += `- **Schema**: ${datatableCtx.schemaName}\n`
+					content += `- **Table**: ${datatableCtx.tableName}\n`
+					// Format columns as column_name: type
+					const columnsStr = JSON.stringify(datatableCtx.columns, null, 2)
+					const truncatedColumns =
+						columnsStr.length > MAX_CONTEXT_CONTENT_LENGTH
+							? columnsStr.slice(0, MAX_CONTEXT_CONTENT_LENGTH) + '\n... [TRUNCATED]'
+							: columnsStr
+					content += `- **Columns** (column_name -> type):\n\`\`\`json\n${truncatedColumns}\n\`\`\`\n`
+				}
+			}
+		}
+
 		content += '\n'
 	}
 
diff --git a/frontend/src/lib/components/copilot/chat/context.ts b/frontend/src/lib/components/copilot/chat/context.ts
index 5c02a84835..f0f9af2310 100644
--- a/frontend/src/lib/components/copilot/chat/context.ts
+++ b/frontend/src/lib/components/copilot/chat/context.ts
@@ -1,14 +1,19 @@
-import { Code, Database, TriangleAlert, Diff } from 'lucide-svelte'
+import { Code, Database, TriangleAlert, Diff, FileCode, Code2, TextSelect, Table2 } from 'lucide-svelte'
 import type { ScriptLang } from '$lib/gen/types.gen'
 import { type DBSchema } from '$lib/stores'
 import { type Change } from 'diff'
+import type { BackendRunnable } from './app/core'
 
 export const ContextIconMap = {
 	code: Code,
 	error: TriangleAlert,
 	db: Database,
 	diff: Diff,
-	code_piece: Code
+	code_piece: Code,
+	app_frontend_file: FileCode,
+	app_backend_runnable: Code2,
+	app_code_selection: TextSelect,
+	app_datatable: Table2
 	// flow_module type is handled with FlowModuleIcon
 }
 
@@ -67,6 +72,62 @@ export interface FlowModuleCodePieceElement extends Omit<CodePieceElement, 'type
 	value: FlowModuleElement['value']
 }
 
+/** App frontend file context element */
+export interface AppFrontendFileElement {
+	type: 'app_frontend_file'
+	/** The file path (e.g., /index.tsx, /styles.css) */
+	path: string
+	/** Title for display (the path) */
+	title: string
+	/** The file content */
+	content: string
+}
+
+/** App backend runnable context element */
+export interface AppBackendRunnableElement {
+	type: 'app_backend_runnable'
+	/** The runnable key */
+	key: string
+	/** Title for display (the key) */
+	title: string
+	/** The runnable configuration */
+	runnable: BackendRunnable
+}
+
+/** App code selection context element (from frontend or backend editor) */
+export interface AppCodeSelectionElement {
+	type: 'app_code_selection'
+	/** Source: frontend file path or backend runnable key */
+	source: string
+	/** Whether this is from frontend or backend */
+	sourceType: 'frontend' | 'backend'
+	/** Title for display */
+	title: string
+	/** The selected code content */
+	content: string
+	/** Line range (1-indexed) */
+	startLine: number
+	endLine: number
+	/** Column range (1-indexed) */
+	startColumn: number
+	endColumn: number
+}
+
+/** App datatable table context element (represents a single table within a datatable) */
+export interface AppDatatableElement {
+	type: 'app_datatable'
+	/** The datatable name (e.g., "main") */
+	datatableName: string
+	/** The schema name (e.g., "public") */
+	schemaName: string
+	/** The table name (e.g., "users") */
+	tableName: string
+	/** Title for display (e.g., "main/public:users" or "main/users") */
+	title: string
+	/** The table columns: column_name -> compact_type */
+	columns: Record<string, string>
+}
+
 export type ContextElement = (
 	| CodeElement
 	| ErrorElement
@@ -75,6 +136,10 @@ export type ContextElement = (
 	| CodePieceElement
 	| FlowModuleElement
 	| FlowModuleCodePieceElement
+	| AppFrontendFileElement
+	| AppBackendRunnableElement
+	| AppCodeSelectionElement
+	| AppDatatableElement
 ) & {
 	deletable?: boolean
 }
diff --git a/frontend/src/lib/components/raw_apps/RawAppEditor.svelte b/frontend/src/lib/components/raw_apps/RawAppEditor.svelte
index 2f2b006fb7..0b124492b5 100644
--- a/frontend/src/lib/components/raw_apps/RawAppEditor.svelte
+++ b/frontend/src/lib/components/raw_apps/RawAppEditor.svelte
@@ -19,6 +19,7 @@
 	import { aiChatManager, AIMode } from '../copilot/chat/AIChatManager.svelte'
 	import { onMount } from 'svelte'
 	import type { LintResult, DataTableSchema, InspectorElementInfo } from '../copilot/chat/app/core'
+	import type { AppCodeSelectionElement } from '../copilot/chat/context'
 	import { rawAppLintStore } from './lintStore'
 	import { dbSchemas } from '$lib/stores'
 	import { runScriptAndPollResult } from '../jobs/utils'
@@ -392,7 +393,11 @@
 					selectionExcluded: selectionExcludedFromPrompt,
 					toggleSelectionExcluded: toggleSelectionExcluded,
 					clearInspector: clearInspectorSelection,
-					clearRunnable: handleClearRunnable
+					clearRunnable: handleClearRunnable,
+					codeSelection: codeSelection,
+					clearCodeSelection: () => {
+						codeSelection = undefined
+					}
 				}
 				if (selectedRunnable) {
 					const runnable = convertToBackendRunnable(selectedRunnable, runnables[selectedRunnable])
@@ -440,6 +445,12 @@
 				// Get unique datatable names from dataTableRefs (the whitelisted tables)
 				const whitelistedDatatables = new Set(dataTableRefsObjects.map((ref) => ref.datatable))
 
+				// If no datatables are configured, return all available datatables
+				// This allows users to see all datatables in the @ context menu
+				if (whitelistedDatatables.size === 0) {
+					return allSchemas
+				}
+
 				// Build a map of whitelisted tables per datatable: datatable -> schema -> Set<table>
 				const whitelistedTables = new Map<string, Map<string, Set<string>>>()
 				for (const ref of dataTableRefsObjects) {
@@ -550,6 +561,19 @@
 					const errorMsg = e instanceof Error ? e.message : String(e)
 					return { success: false, error: errorMsg }
 				}
+			},
+			addTableToWhitelist: (datatableName: string, schemaName: string, tableName: string) => {
+				// Format the table reference
+				const newRef = formatDataTableRef({
+					datatable: datatableName,
+					schema: schemaName === 'public' ? undefined : schemaName,
+					table: tableName
+				})
+				// Only add if not already present
+				if (!data.tables.includes(newRef)) {
+					data.tables = [...data.tables, newRef]
+					saveFrontendDraft()
+				}
 			}
 		})
 	})
@@ -557,6 +581,7 @@
 	let selectedDocument: string | undefined = $state(undefined)
 	let inspectorElement: InspectorElementInfo | undefined = $state(undefined)
 	let selectionExcludedFromPrompt: boolean = $state(false)
+	let codeSelection: AppCodeSelectionElement | undefined = $state(undefined)
 
 	function toggleSelectionExcluded() {
 		selectionExcludedFromPrompt = !selectionExcludedFromPrompt
@@ -596,6 +621,27 @@
 		} else if (e.data.type === 'inspectorClear') {
 			// Clear the inspector element when user dismisses the selection
 			inspectorElement = undefined
+		} else if (e.data.type === 'editorSelection') {
+			// Handle code selection from the iframe editor
+			const selection = e.data.selection
+			if (selection === null) {
+				// Selection cleared
+				codeSelection = undefined
+			} else {
+				// Normalize path
+				const normalizedPath = selection.path?.replace(/\\/g, '/')
+				codeSelection = {
+					type: 'app_code_selection',
+					source: normalizedPath,
+					sourceType: 'frontend',
+					title: `${normalizedPath}:L${selection.range.startLine}-L${selection.range.endLine}`,
+					content: selection.content,
+					startLine: selection.range.startLine,
+					endLine: selection.range.endLine,
+					startColumn: selection.range.startColumn,
+					endColumn: selection.range.endColumn
+				}
+			}
 		}
 	}