From 38aaa44b95022f40fb830aeb440eb35028c46752 Mon Sep 17 00:00:00 2001
From: HugoCasa <hugo@casademont.ch>
Date: Mon, 2 Jun 2025 12:02:47 +0200
Subject: [PATCH] fix: reactivity

---
 .../src/lib/components/FlowBuilder.svelte     |   3 +-
 .../lib/components/copilot/chat/AIChat.svelte |  92 +++---
 .../copilot/chat/AIChatDisplay.svelte         |  11 +-
 .../components/copilot/chat/ChatMode.svelte   |   9 +-
 .../copilot/chat/ContextManager.svelte        | 300 ------------------
 .../copilot/chat/ContextManager.svelte.ts     | 282 ++++++++++++++++
 .../copilot/chat/flow/FlowAIChat.svelte       |  28 +-
 .../flows/content/FlowModuleComponent.svelte  |   9 +-
 frontend/src/lib/stores.ts                    |   1 +
 9 files changed, 364 insertions(+), 371 deletions(-)
 delete mode 100644 frontend/src/lib/components/copilot/chat/ContextManager.svelte
 create mode 100644 frontend/src/lib/components/copilot/chat/ContextManager.svelte.ts

diff --git a/frontend/src/lib/components/FlowBuilder.svelte b/frontend/src/lib/components/FlowBuilder.svelte
index 3fc5b2dcb4..78a24b7aaf 100644
--- a/frontend/src/lib/components/FlowBuilder.svelte
+++ b/frontend/src/lib/components/FlowBuilder.svelte
@@ -8,8 +8,7 @@
 		type OpenFlow,
 		type InputTransform,
 		type TriggersCount,
-		CaptureService,
-		type HubScriptKind
+		CaptureService
 	} from '$lib/gen'
 	import { initHistory, redo, undo } from '$lib/history'
 	import {
diff --git a/frontend/src/lib/components/copilot/chat/AIChat.svelte b/frontend/src/lib/components/copilot/chat/AIChat.svelte
index b2968ba845..c5f315e473 100644
--- a/frontend/src/lib/components/copilot/chat/AIChat.svelte
+++ b/frontend/src/lib/components/copilot/chat/AIChat.svelte
@@ -18,7 +18,7 @@
 	import { onDestroy, setContext, untrack, type Snippet } from 'svelte'
 	import { type OpenFlow, type ScriptLang } from '$lib/gen'
 	import { sendUserToast } from '$lib/toast'
-	import ContextManager from './ContextManager.svelte'
+	import ContextManager, { type ScriptOptions } from './ContextManager.svelte'
 	import HistoryManager from './HistoryManager.svelte'
 	import {
 		flowTools,
@@ -30,17 +30,9 @@
 		ChatCompletionMessageParam,
 		ChatCompletionSystemMessageParam
 	} from 'openai/resources/index.mjs'
+	import { chatMode, copilotSessionModel, dbSchemas, workspaceStore } from '$lib/stores'
 	interface Props {
-		scriptOptions?: {
-			lang: ScriptLang | 'bunnative'
-			code: string
-			error: string | undefined
-			args: Record<string, any>
-			path: string | undefined
-			lastSavedCode?: string | undefined
-			lastDeployedCode?: string | undefined
-			diffMode: boolean
-		}
+		scriptOptions?: ScriptOptions
 		flowHelpers?: FlowAIChatHelpers & {
 			getFlow: () => OpenFlow
 		}
@@ -60,16 +52,16 @@
 		script: scriptOptions !== undefined,
 		flow: flowHelpers !== undefined
 	})
-	let mode: 'script' | 'flow' = $state(flowHelpers ? 'flow' : 'script')
 
 	async function updateMode(currentMode: 'script' | 'flow') {
 		if (!allowedModes[currentMode]) {
-			mode = currentMode === 'script' ? 'flow' : 'script'
+			chatMode.set(currentMode === 'script' ? 'flow' : 'script')
 		}
 	}
 	$effect(() => {
-		updateMode(untrack(() => mode))
+		updateMode(untrack(() => $chatMode))
 	})
+
 	let displayMessages: DisplayMessage[] = $state([])
 	let abortController: AbortController | undefined = undefined
 	let messages: ChatCompletionMessageParam[] = $state([])
@@ -92,7 +84,7 @@
 		} = {}
 	) {
 		if (options.mode) {
-			mode = options.mode
+			$chatMode = options.mode
 		}
 		if (options.instructions) {
 			instructions = options.instructions
@@ -102,7 +94,7 @@
 		}
 		try {
 			const oldSelectedContext = contextManager?.getSelectedContext() ?? []
-			if (mode === 'script') {
+			if ($chatMode === 'script') {
 				contextManager?.updateContextOnRequest(options)
 			}
 			loading.set(true)
@@ -114,20 +106,20 @@
 				{
 					role: 'user',
 					content: instructions,
-					contextElements: mode === 'script' ? oldSelectedContext : undefined
+					contextElements: $chatMode === 'script' ? oldSelectedContext : undefined
 				}
 			]
 			const oldInstructions = instructions
 			instructions = ''
 
 			const systemMessage =
-				mode === 'script' ? prepareScriptSystemMessage() : prepareFlowSystemMessage()
+				$chatMode === 'script' ? prepareScriptSystemMessage() : prepareFlowSystemMessage()
 
-			if (mode === 'flow' && !flowHelpers) {
+			if ($chatMode === 'flow' && !flowHelpers) {
 				throw new Error('No flow helpers passed')
 			}
 
-			if (mode === 'script' && !scriptOptions && !options.lang) {
+			if ($chatMode === 'script' && !scriptOptions && !options.lang) {
 				throw new Error('No script options passed')
 			}
 
@@ -135,7 +127,7 @@
 			const isPreprocessor = scriptOptions?.path === 'preprocessor' || options.isPreprocessor
 
 			const userMessage =
-				mode === 'flow'
+				$chatMode === 'flow'
 					? prepareFlowUserMessage(oldInstructions, flowHelpers!.getFlow())
 					: await prepareScriptUserMessage(oldInstructions, lang, oldSelectedContext, {
 							isPreprocessor
@@ -168,7 +160,7 @@
 									role: 'assistant',
 									content: $currentReply,
 									contextElements:
-										mode === 'script'
+										$chatMode === 'script'
 											? oldSelectedContext.filter((c) => c.type === 'code')
 											: undefined
 								}
@@ -193,7 +185,7 @@
 				}
 			}
 
-			if (mode === 'flow') {
+			if ($chatMode === 'flow') {
 				if (!flowHelpers) {
 					throw new Error('No flow helpers found')
 				}
@@ -227,7 +219,9 @@
 						role: 'assistant',
 						content: $currentReply,
 						contextElements:
-							mode === 'script' ? oldSelectedContext.filter((c) => c.type === 'code') : undefined
+							$chatMode === 'script'
+								? oldSelectedContext.filter((c) => c.type === 'code')
+								: undefined
 					}
 				]
 				currentReply.set('')
@@ -272,7 +266,7 @@
 			throw new Error('No script options passed')
 		}
 		instructions = prompt
-		contextManager?.setAskAiContext(options)
+		contextManager.setAskAiContext(options)
 		sendRequest({
 			removeDiff: options.withDiff,
 			addBackCode: options.withCode === false
@@ -295,42 +289,48 @@
 	})
 
 	let aiChatDisplay: AIChatDisplay | undefined = $state(undefined)
-	let contextManager: ContextManager | undefined = $state(undefined)
-</script>
+	// let contextManager: ContextManager | undefined = $state(undefined)
 
-{#if mode === 'script' && scriptOptions}
-	<ContextManager
-		bind:this={contextManager}
-		code={scriptOptions.code}
-		lang={scriptOptions.lang}
-		path={scriptOptions.path}
-		args={scriptOptions.args}
-		lastSavedCode={scriptOptions.lastSavedCode}
-		lastDeployedCode={scriptOptions.lastDeployedCode}
-		error={scriptOptions.error}
-		bind:displayMessages
-	/>
-{/if}
+	const contextManager = new ContextManager()
+
+	$effect(() => {
+		if (scriptOptions) {
+			contextManager.updateAvailableContext(
+				scriptOptions,
+				$dbSchemas,
+				$workspaceStore ?? '',
+				!$copilotSessionModel?.model.endsWith('/thinking'),
+				untrack(() => contextManager.getSelectedContext())
+			)
+		}
+	})
+
+	$effect(() => {
+		displayMessages = ContextManager.updateDisplayMessages(
+			untrack(() => displayMessages),
+			$dbSchemas
+		)
+	})
+</script>
 
 <AIChatDisplay
 	bind:this={aiChatDisplay}
-	bind:mode
 	{allowedModes}
 	pastChats={historyManager.getPastChats()}
 	bind:selectedContext={
-		() => contextManager?.getSelectedContext() ?? [],
+		() => contextManager.getSelectedContext(),
 		(sc) => {
-			contextManager?.setSelectedContext(sc)
+			scriptOptions && contextManager.setSelectedContext(sc)
 		}
 	}
-	availableContext={contextManager?.getAvailableContext() ?? []}
+	availableContext={contextManager.getAvailableContext()}
 	messages={$currentReply
 		? [
 				...displayMessages,
 				{
 					role: 'assistant',
 					content: $currentReply,
-					contextElements: contextManager?.getSelectedContext()?.filter((c) => c.type === 'code')
+					contextElements: contextManager.getSelectedContext().filter((c) => c.type === 'code')
 				}
 			]
 		: displayMessages}
@@ -345,7 +345,7 @@
 	loadPastChat={(id) => {
 		const chat = historyManager.loadPastChat(id)
 		if (chat) {
-			displayMessages = chat.displayMessages
+			displayMessages = ContextManager.updateDisplayMessages(chat.displayMessages, $dbSchemas)
 			messages = chat.actualMessages
 			aiChatDisplay?.enableAutomaticScroll()
 		}
diff --git a/frontend/src/lib/components/copilot/chat/AIChatDisplay.svelte b/frontend/src/lib/components/copilot/chat/AIChatDisplay.svelte
index 19f2155851..62122e6efa 100644
--- a/frontend/src/lib/components/copilot/chat/AIChatDisplay.svelte
+++ b/frontend/src/lib/components/copilot/chat/AIChatDisplay.svelte
@@ -14,9 +14,9 @@
 	import ChatQuickActions from './ChatQuickActions.svelte'
 	import ProviderModelSelector from './ProviderModelSelector.svelte'
 	import ChatMode from './ChatMode.svelte'
+	import { chatMode } from '$lib/stores'
 
 	let {
-		mode = $bindable(),
 		allowedModes,
 		messages,
 		instructions = $bindable(),
@@ -34,7 +34,6 @@
 		headerLeft,
 		headerRight
 	}: {
-		mode: 'script' | 'flow'
 		allowedModes: {
 			script: boolean
 			flow: boolean
@@ -249,7 +248,7 @@
 	{/if}
 
 	<div class:border-t={messages.length > 0}>
-		{#if mode === 'script'}
+		{#if $chatMode === 'script'}
 			<div class="flex flex-row gap-1 mb-1 overflow-scroll pt-2 px-2 no-scrollbar">
 				<Popover>
 					<svelte:fragment slot="trigger">
@@ -314,14 +313,14 @@
 		{/if}
 		<div
 			class={`flex flex-row ${
-				mode === 'script' && hasDiff ? 'justify-between' : 'justify-end'
+				$chatMode === 'script' && hasDiff ? 'justify-between' : 'justify-end'
 			} items-center px-0.5`}
 		>
-			{#if mode === 'script' && hasDiff}
+			{#if $chatMode === 'script' && hasDiff}
 				<ChatQuickActions {askAi} {diffMode} />
 			{/if}
 			<div class="flex flex-row gap-2 min-w-0">
-				<ChatMode bind:mode {allowedModes} />
+				<ChatMode {allowedModes} />
 				<ProviderModelSelector />
 			</div>
 		</div>
diff --git a/frontend/src/lib/components/copilot/chat/ChatMode.svelte b/frontend/src/lib/components/copilot/chat/ChatMode.svelte
index d03eac13aa..d366fd7011 100644
--- a/frontend/src/lib/components/copilot/chat/ChatMode.svelte
+++ b/frontend/src/lib/components/copilot/chat/ChatMode.svelte
@@ -2,12 +2,11 @@
 	import { ChevronDown } from 'lucide-svelte'
 	import Popover from '$lib/components/meltComponents/Popover.svelte'
 	import { twMerge } from 'tailwind-merge'
+	import { chatMode } from '$lib/stores'
 
 	let {
-		mode = $bindable(),
 		allowedModes
 	}: {
-		mode: 'script' | 'flow'
 		allowedModes: {
 			script: boolean
 			flow: boolean
@@ -22,7 +21,7 @@
 				class="text-tertiary text-xs flex flex-row items-center font-normal gap-0.5 border px-1 rounded-lg"
 			>
 				<span class={`truncate`}>
-					{mode} mode
+					{$chatMode} mode
 				</span>
 				{#if allowedModes.script && allowedModes.flow}
 					<div class="shrink-0">
@@ -37,10 +36,10 @@
 					<button
 						class={twMerge(
 							'text-left text-xs hover:bg-surface-hover rounded-md p-1 font-normal',
-							mode === possibleMode && 'bg-surface-hover'
+							$chatMode === possibleMode && 'bg-surface-hover'
 						)}
 						onclick={() => {
-							mode = possibleMode as 'script' | 'flow'
+							$chatMode = possibleMode as 'script' | 'flow'
 							close()
 						}}
 					>
diff --git a/frontend/src/lib/components/copilot/chat/ContextManager.svelte b/frontend/src/lib/components/copilot/chat/ContextManager.svelte
deleted file mode 100644
index ccd62fb930..0000000000
--- a/frontend/src/lib/components/copilot/chat/ContextManager.svelte
+++ /dev/null
@@ -1,300 +0,0 @@
-<script lang="ts">
-	import { langToExt } from '$lib/editorUtils'
-	import { ResourceService, type ListResourceResponse, type ScriptLang } from '$lib/gen'
-	import { scriptLangToEditorLang } from '$lib/scripts'
-	import {
-		copilotSessionModel,
-		dbSchemas,
-		SQLSchemaLanguages,
-		workspaceStore,
-		type DBSchema
-	} from '$lib/stores'
-	import { untrack } from 'svelte'
-	import type { ContextElement } from './context'
-	import type { DisplayMessage } from './shared'
-	import { diffLines } from 'diff'
-
-	let {
-		displayMessages = $bindable(),
-		code,
-		lang,
-		path,
-		args,
-		lastSavedCode,
-		lastDeployedCode,
-		error
-	}: {
-		displayMessages: DisplayMessage[]
-		code: string
-		lang: ScriptLang | 'bunnative'
-		path: string | undefined
-		args: Record<string, any>
-		lastSavedCode: string | undefined
-		lastDeployedCode: string | undefined
-		error: string | undefined
-	} = $props()
-
-	let contextCodePath = $derived(
-		path
-			? (path.split('/').pop() ?? 'script') + '.' + langToExt(scriptLangToEditorLang(lang))
-			: undefined
-	)
-
-	let selectedContext: ContextElement[] = $derived([
-		{
-			type: 'code',
-			title: contextCodePath ?? '',
-			content: code,
-			lang
-		}
-	])
-
-	export function getSelectedContext() {
-		return selectedContext
-	}
-
-	export function setSelectedContext(newSelectedContext: ContextElement[]) {
-		selectedContext = newSelectedContext
-	}
-
-	let db: { schema: DBSchema; resource: string } | undefined = $state(undefined)
-
-	function updateSchema() {
-		try {
-			const schemaRes = lang === 'graphql' ? args.api : args.database
-			if (typeof schemaRes === 'string') {
-				const schemaPath = schemaRes.replace('$res:', '')
-				const schema = $dbSchemas[schemaPath]
-				if (schema && schema.lang === lang) {
-					db = { schema, resource: schemaPath }
-				} else {
-					db = undefined
-				}
-			} else {
-				db = undefined
-			}
-		} catch (err) {
-			console.error('Could not update schema', err)
-		}
-	}
-
-	$effect(() => {
-		updateSchema()
-	})
-
-	let dbResources: ListResourceResponse = $state([])
-
-	async function updateDBResources() {
-		const workspace = $workspaceStore
-		if (workspace) {
-			dbResources = await ResourceService.listResource({
-				workspace: workspace,
-				resourceType: SQLSchemaLanguages.join(',')
-			})
-		}
-	}
-
-	$effect(() => {
-		updateDBResources()
-	})
-
-	function updateAvailableContext() {
-		if (!contextCodePath) {
-			return []
-		}
-		const providerModel = $copilotSessionModel
-		try {
-			let newAvailableContext: ContextElement[] = [
-				{
-					type: 'code',
-					title: contextCodePath,
-					content: code,
-					lang
-				}
-			]
-			if (!providerModel?.model.endsWith('/thinking')) {
-				for (const d of 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 } : {})
-					})
-				}
-			}
-
-			if (lastSavedCode && lastSavedCode !== code) {
-				newAvailableContext.push({
-					type: 'diff',
-					title: 'diff_with_last_saved_draft',
-					content: lastSavedCode ?? '',
-					diff: diffLines(lastSavedCode ?? '', code),
-					lang
-				})
-			}
-
-			if (lastDeployedCode && lastDeployedCode !== code) {
-				newAvailableContext.push({
-					type: 'diff',
-					title: 'diff_with_last_deployed_version',
-					content: lastDeployedCode ?? '',
-					diff: diffLines(lastDeployedCode ?? '', code),
-					lang
-				})
-			}
-
-			if (error) {
-				newAvailableContext = [
-					...newAvailableContext,
-					{
-						type: 'error',
-						title: 'error',
-						content: error
-					}
-				]
-			}
-
-			return newAvailableContext
-		} catch (err) {
-			console.error('Could not update available context', err)
-			return []
-		}
-	}
-
-	function updateDisplayMessages(currentDisplayMessages: DisplayMessage[]): DisplayMessage[] {
-		const schemas = $dbSchemas
-		return currentDisplayMessages.map((m) => ({
-			...m,
-			contextElements:
-				m.role !== 'tool' && m.contextElements
-					? m.contextElements.map((c) =>
-							c.type === 'db'
-								? {
-										type: 'db',
-										title: c.title,
-										schema: schemas[c.title]
-									}
-								: c
-						)
-					: undefined
-		}))
-	}
-
-	$effect(() => {
-		displayMessages = updateDisplayMessages(untrack(() => displayMessages))
-	})
-
-	let availableContext: ContextElement[] = $derived.by(updateAvailableContext)
-
-	export function getAvailableContext() {
-		return availableContext
-	}
-
-	function updateSelectedContext(currentSelectedContext: ContextElement[]) {
-		if (!contextCodePath) {
-			return currentSelectedContext
-		}
-		let newSelectedContext: ContextElement[] = [...currentSelectedContext]
-
-		// If the db is already fetched, add it to the selected context
-		if (
-			db &&
-			!newSelectedContext.find((c) => c.type === 'db' && db && c.title === db.resource) &&
-			!$copilotSessionModel?.model.endsWith('/thinking')
-		) {
-			newSelectedContext = [
-				...newSelectedContext,
-				{
-					type: 'db',
-					title: db.resource,
-					schema: db.schema
-				}
-			]
-		}
-
-		newSelectedContext = newSelectedContext
-			.map((c) =>
-				c.type === 'code_piece' && code.includes(c.content)
-					? c
-					: availableContext.find((ac) => ac.type === c.type && ac.title === c.title)
-			)
-			.filter((c) => c !== undefined) as ContextElement[]
-
-		selectedContext = newSelectedContext
-	}
-
-	$effect(() => {
-		updateSelectedContext(untrack(() => selectedContext))
-	})
-
-	export function addSelectedLinesToContext(lines: string, startLine: number, endLine: number) {
-		if (
-			selectedContext.find(
-				(c) => c.type === 'code_piece' && c.title === `L${startLine}-L${endLine}`
-			)
-		) {
-			return
-		}
-		selectedContext = [
-			...selectedContext,
-			{
-				type: 'code_piece',
-				title: `L${startLine}-L${endLine}`,
-				startLine,
-				endLine,
-				content: lines,
-				lang
-			}
-		]
-	}
-
-	export function setFixContext() {
-		const codeContext = contextCodePath
-			? availableContext.find((c) => c.type === 'code' && c.title === contextCodePath)
-			: undefined
-		const errorContext = availableContext.find((c) => c.type === 'error')
-
-		if (codeContext && errorContext) {
-			selectedContext = [codeContext, errorContext]
-		}
-	}
-
-	export function setAskAiContext(options: { withCode?: boolean; withDiff?: boolean }) {
-		const codeContext = availableContext.find(
-			(c) => c.type === 'code' && c.title === contextCodePath
-		)
-		if (!codeContext) {
-			return
-		}
-		selectedContext = [
-			...(options.withCode === false ? [] : [codeContext]),
-			...(options.withDiff
-				? [
-						{
-							type: 'diff' as const,
-							title: 'diff_with_last_deployed_version',
-							content: lastDeployedCode ?? '',
-							diff: diffLines(lastDeployedCode ?? '', code),
-							lang
-						}
-					]
-				: [])
-		]
-	}
-
-	export function updateContextOnRequest(options: { removeDiff?: boolean; addBackCode?: boolean }) {
-		selectedContext = selectedContext.filter((c) => c.type !== 'code_piece')
-		if (options.removeDiff) {
-			selectedContext = selectedContext.filter((c) => c.type !== 'diff')
-		}
-		if (options.addBackCode) {
-			const codeContext = availableContext.find(
-				(c) => c.type === 'code' && c.title === contextCodePath
-			)
-			if (codeContext) {
-				selectedContext = [...selectedContext, codeContext]
-			}
-		}
-	}
-</script>
diff --git a/frontend/src/lib/components/copilot/chat/ContextManager.svelte.ts b/frontend/src/lib/components/copilot/chat/ContextManager.svelte.ts
new file mode 100644
index 0000000000..e05e6fd622
--- /dev/null
+++ b/frontend/src/lib/components/copilot/chat/ContextManager.svelte.ts
@@ -0,0 +1,282 @@
+import { langToExt } from '$lib/editorUtils'
+import { ResourceService, 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 { DisplayMessage } from './shared'
+
+export interface ScriptOptions {
+	lang: ScriptLang | 'bunnative'
+	code: string
+	error: string | undefined
+	args: Record<string, any>
+	path: string | undefined
+	lastSavedCode?: string
+	lastDeployedCode?: string
+	diffMode: boolean
+}
+
+export default class ContextManager {
+	private selectedContext: ContextElement[] = $state([])
+	private availableContext: ContextElement[] = $state([])
+
+	private workspace: string | undefined = undefined
+	private dbResources: ListResourceResponse = []
+	private scriptOptions: ScriptOptions | undefined = undefined
+
+	private async refreshDbResources(workspace: string) {
+		this.dbResources = await ResourceService.listResource({
+			workspace: workspace,
+			resourceType: SQLSchemaLanguages.join(',')
+		})
+	}
+
+	private getSelectedDBSchema(scriptOptions: ScriptOptions, dbSchemas: DBSchemas) {
+		const schemaRes =
+			scriptOptions.lang === 'graphql' ? scriptOptions.args.api : scriptOptions.args.database
+		if (typeof schemaRes === 'string') {
+			const schemaPath = schemaRes.replace('$res:', '')
+			const schema = dbSchemas[schemaPath]
+			if (schema && schema.lang === scriptOptions.lang) {
+				return { schema, resource: schemaPath }
+			} else {
+				return { schema: undefined, resource: schemaPath }
+			}
+		}
+	}
+
+	private getContextCodePath(scriptOptions: ScriptOptions) {
+		return (
+			(scriptOptions.path?.split('/').pop() ?? 'script') +
+			'.' +
+			langToExt(scriptLangToEditorLang(scriptOptions.lang))
+		)
+	}
+
+	async updateAvailableContext(
+		scriptOptions: ScriptOptions,
+		dbSchemas: DBSchemas,
+		workspace: string,
+		toolSupport: boolean,
+		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',
+					title: this.getContextCodePath(scriptOptions) ?? '',
+					content: scriptOptions.code,
+					lang: scriptOptions.lang
+				}
+			]
+
+			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 } : {})
+					})
+				}
+			}
+
+			if (scriptOptions.lastSavedCode && scriptOptions.lastSavedCode !== scriptOptions.code) {
+				newAvailableContext.push({
+					type: 'diff',
+					title: 'diff_with_last_saved_draft',
+					content: scriptOptions.lastSavedCode ?? '',
+					diff: diffLines(scriptOptions.lastSavedCode ?? '', scriptOptions.code),
+					lang: scriptOptions.lang
+				})
+			}
+
+			if (scriptOptions.lastDeployedCode && scriptOptions.lastDeployedCode !== scriptOptions.code) {
+				newAvailableContext.push({
+					type: 'diff',
+					title: 'diff_with_last_deployed_version',
+					content: scriptOptions.lastDeployedCode ?? '',
+					diff: diffLines(scriptOptions.lastDeployedCode ?? '', scriptOptions.code),
+					lang: scriptOptions.lang
+				})
+			}
+
+			if (scriptOptions.error) {
+				newAvailableContext = [
+					...newAvailableContext,
+					{
+						type: 'error',
+						title: 'error',
+						content: scriptOptions.error
+					}
+				]
+			}
+
+			let newSelectedContext: ContextElement[] = [...currentlySelectedContext]
+
+			if (firstTime) {
+				newSelectedContext = [
+					{
+						type: 'code',
+						title: this.getContextCodePath(scriptOptions) ?? '',
+						content: scriptOptions.code,
+						lang: scriptOptions.lang
+					}
+				]
+			}
+
+			const db = this.getSelectedDBSchema(scriptOptions, dbSchemas)
+			if (
+				db &&
+				!newSelectedContext.find((c) => c.type === 'db' && db && c.title === db.resource) &&
+				toolSupport
+			) {
+				newSelectedContext = [
+					...newSelectedContext,
+					{
+						type: 'db',
+						title: db.resource,
+						schema: db.schema
+					}
+				]
+			}
+
+			newSelectedContext = newSelectedContext
+				.filter(
+					(c) =>
+						(c.type === 'code_piece' && scriptOptions.code.includes(c.content)) ||
+						c.type === 'code' ||
+						newAvailableContext.some((ac) => ac.type === c.type && ac.title === c.title)
+				)
+				.map((c) =>
+					c.type === 'code'
+						? {
+								...c,
+								content: scriptOptions.code,
+								title: this.getContextCodePath(scriptOptions)
+							}
+						: 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', err)
+		}
+	}
+
+	getSelectedContext() {
+		return this.selectedContext
+	}
+
+	setSelectedContext(newSelectedContext: ContextElement[]) {
+		this.selectedContext = newSelectedContext
+	}
+
+	getAvailableContext() {
+		return this.availableContext
+	}
+
+	addSelectedLinesToContext(lines: string, startLine: number, endLine: number) {
+		if (
+			!this.scriptOptions ||
+			this.selectedContext.find(
+				(c) => c.type === 'code_piece' && c.title === `L${startLine}-L${endLine}`
+			)
+		) {
+			return
+		}
+		this.selectedContext = [
+			...this.selectedContext,
+			{
+				type: 'code_piece',
+				title: `L${startLine}-L${endLine}`,
+				startLine,
+				endLine,
+				content: lines,
+				lang: this.scriptOptions.lang
+			}
+		]
+	}
+
+	setFixContext() {
+		const codeContext = this.availableContext.find((c) => c.type === 'code')
+		const errorContext = this.availableContext.find((c) => c.type === 'error')
+
+		if (codeContext && errorContext) {
+			this.selectedContext = [codeContext, errorContext]
+		}
+	}
+
+	setAskAiContext(options: { withCode?: boolean; withDiff?: boolean }) {
+		if (!this.scriptOptions) {
+			return
+		}
+		const codeContext = this.availableContext.find((c) => c.type === 'code')
+		if (!codeContext) {
+			return
+		}
+		this.selectedContext = [
+			...(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
+						}
+					]
+				: [])
+		]
+	}
+
+	updateContextOnRequest(options: { removeDiff?: boolean; addBackCode?: boolean }) {
+		this.selectedContext = this.selectedContext.filter((c) => c.type !== 'code_piece')
+		if (options.removeDiff) {
+			this.selectedContext = this.selectedContext.filter((c) => c.type !== 'diff')
+		}
+		if (options.addBackCode) {
+			const codeContext = this.availableContext.find((c) => c.type === 'code')
+			if (codeContext) {
+				this.selectedContext = [...this.selectedContext, codeContext]
+			}
+		}
+	}
+
+	static updateDisplayMessages(
+		displayMessages: DisplayMessage[],
+		dbSchemas: DBSchemas
+	): DisplayMessage[] {
+		return displayMessages.map((m) => ({
+			...m,
+			contextElements:
+				m.role !== 'tool' && m.contextElements
+					? m.contextElements.map((c) =>
+							c.type === 'db'
+								? {
+										type: 'db',
+										title: c.title,
+										schema: dbSchemas[c.title]
+									}
+								: c
+						)
+					: undefined
+		}))
+	}
+}
diff --git a/frontend/src/lib/components/copilot/chat/flow/FlowAIChat.svelte b/frontend/src/lib/components/copilot/chat/flow/FlowAIChat.svelte
index 40ae053f83..787ef3187b 100644
--- a/frontend/src/lib/components/copilot/chat/flow/FlowAIChat.svelte
+++ b/frontend/src/lib/components/copilot/chat/flow/FlowAIChat.svelte
@@ -15,6 +15,7 @@
 		insertNewFailureModule,
 		insertNewPreprocessorModule
 	} from '$lib/components/flows/flowStateUtils'
+	import type { ScriptOptions } from '../ContextManager.svelte'
 
 	let {
 		flowModuleSchemaMap,
@@ -39,18 +40,25 @@
 		}
 	}
 
-	function getScriptOptions(id: string) {
+	function getScriptOptions(id: string): ScriptOptions | undefined {
 		const module = getModule(id)
 
-		if (
-			module &&
-			module.value.type === 'rawscript' &&
-			$currentEditor &&
-			$currentEditor.type === 'script' &&
-			$currentEditor.stepId === module.id
-		) {
+		if (module && module.value.type === 'rawscript') {
 			const moduleState: FlowModuleState | undefined = $flowStateStore[module.id]
 
+			const editorRelated =
+				$currentEditor && $currentEditor.type === 'script' && $currentEditor.stepId === module.id
+					? {
+							diffMode: $currentEditor.diffMode,
+							lastDeployedCode: $currentEditor.lastDeployedCode,
+							lastSavedCode: undefined
+						}
+					: {
+							diffMode: false,
+							lastDeployedCode: undefined,
+							lastSavedCode: undefined
+						}
+
 			return {
 				args: moduleState?.previewArgs ?? {},
 				error:
@@ -60,9 +68,7 @@
 				code: module.value.content,
 				lang: module.value.language,
 				path: module.id,
-				diffMode: $currentEditor.diffMode,
-				lastDeployedCode: $currentEditor.lastDeployedCode,
-				lastSavedCode: undefined
+				...editorRelated
 			}
 		}
 
diff --git a/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte b/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte
index 5f90e280bd..068e508501 100644
--- a/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte
+++ b/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte
@@ -13,7 +13,7 @@
 	import FlowModuleHeader from './FlowModuleHeader.svelte'
 	import { getLatestHashForScript, scriptLangToEditorLang } from '$lib/scripts'
 	import PropPickerWrapper from '../propPicker/PropPickerWrapper.svelte'
-	import { getContext, tick } from 'svelte'
+	import { getContext, onDestroy, tick } from 'svelte'
 	import type { FlowEditorContext } from '../types'
 	import FlowModuleScript from './FlowModuleScript.svelte'
 	import FlowModuleEarlyStop from './FlowModuleEarlyStop.svelte'
@@ -276,6 +276,10 @@
 			diffMode,
 			lastDeployedCode
 		})
+
+	onDestroy(() => {
+		$currentEditor = undefined
+	})
 </script>
 
 <svelte:window on:keydown={onKeyDown} />
@@ -385,6 +389,9 @@
 											on:addSelectedLinesToAiChat={(e) => {
 												// TODO
 											}}
+											on:toggleAiPanel={() => {
+												// TODO
+											}}
 											loadAsync
 											folding
 											path={$pathStore + '/' + flowModule.id}
diff --git a/frontend/src/lib/stores.ts b/frontend/src/lib/stores.ts
index f6d2e42685..3eddc0876d 100644
--- a/frontend/src/lib/stores.ts
+++ b/frontend/src/lib/stores.ts
@@ -105,6 +105,7 @@ export const copilotInfo = writable<{
 	defaultModel: undefined,
 	aiModels: []
 })
+export const chatMode = writable<'script' | 'flow'>('script')
 
 export function setCopilotInfo(aiConfig: AIConfig) {
 	if (Object.keys(aiConfig.providers ?? {}).length > 0) {