feat: connect fix btn in flow editor to ai chat (#5863)

* feat: connect fix btn in flow editor to ai chat

* adapt to unified chat

* cleaning

* cleanup

frontend/src/lib/components/FlowBuilder.svelte

@@ -749,12 +749,6 @@
 	}
 
 	let flowPreviewButtons: FlowPreviewButtons
-
-	let flowEditor: FlowEditor | undefined = undefined
-	$: if (flowEditor) {
-		flowCopilotContext.toggleAiPanel = flowEditor.toggleAiPanel
-		flowCopilotContext.addSelectedLinesToAiChat = flowEditor.addSelectedLinesToAiChat
-	}
 </script>
 
 <svelte:window on:keydown={onKeyDown} />
@@ -971,7 +965,6 @@
 			<!-- metadata -->
 			{#if $flowStateStore}
 				<FlowEditor
-					bind:this={flowEditor}
 					{disabledFlowInputs}
 					disableAi={disableAi || customUi?.stepInputs?.ai == false}
 					disableSettings={customUi?.settingsPanel === false}

frontend/src/lib/components/ModulePreviewResultViewer.svelte

@@ -9,6 +9,7 @@
 	import { Pane, Splitpanes } from 'svelte-splitpanes'
 	import type { FlowEditorContext } from './flows/types'
 	import { getContext } from 'svelte'
+	import { getStringError } from './copilot/chat/utils'
 
 	export let lang: Script['language']
 	export let editor: Editor | undefined
@@ -70,14 +71,8 @@
 			{disableHistory}
 		>
 			<svelte:fragment slot="copilot-fix">
-				{#if lang && editor && diffEditor && $testStepStore[mod.id] && selectedJob && 'result' in selectedJob && selectedJob.result && typeof selectedJob.result == 'object' && `error` in selectedJob.result && selectedJob.result.error}
-					<ScriptFix
-						error={JSON.stringify(selectedJob.result.error)}
-						{lang}
-						{editor}
-						{diffEditor}
-						args={$testStepStore[mod.id]}
-					/>
+				{#if lang && editor && diffEditor && $testStepStore[mod.id] && selectedJob?.type === 'CompletedJob' && !selectedJob.success && getStringError(selectedJob.result)}
+					<ScriptFix {lang} />
 				{/if}
 			</svelte:fragment>
 		</OutputPickerInner>

frontend/src/lib/components/ScriptEditor.svelte

@@ -312,16 +312,6 @@
 	let testPanelSize = 30
 	let storedTestPanelSize = testPanelSize
 
-	function addSelectedLinesToAiChat(
-		e: CustomEvent<{ lines: string; startLine: number; endLine: number }>
-	) {
-		if (!aiChatManager.open) {
-			aiChatManager.toggleOpen()
-		}
-		aiChatManager.addSelectedLinesToContext(e.detail.lines, e.detail.startLine, e.detail.endLine)
-		// aiChatManager.focusTextArea() TODO: Add this back
-	}
-
 	$: !SUPPORTED_CHAT_SCRIPT_LANGUAGES.includes(lang ?? '') &&
 		!aiChatManager.open &&
 		aiChatManager.toggleOpen()
@@ -530,7 +520,10 @@
 						}}
 						on:saveDraft
 						on:toggleAiPanel={() => aiChatManager.toggleOpen()}
-						on:addSelectedLinesToAiChat={addSelectedLinesToAiChat}
+						on:addSelectedLinesToAiChat={(e) => {
+							const { lines, startLine, endLine } = e.detail
+							aiChatManager.addSelectedLinesToContext(lines, startLine, endLine)
+						}}
 						on:toggleTestPanel={toggleTestPanel}
 						cmdEnterAction={async () => {
 							await inferSchema(code)
@@ -678,10 +671,6 @@
 					<Pane size={67} class="relative">
 						<LogPanel
 							bind:setFocusToLogs
-							on:fix={() => {
-								aiChatManager.fix()
-							}}
-							fixChatMode
 							{lang}
 							previewJob={testJob}
 							{pastPreviews}

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

@@ -2,172 +2,23 @@
 	import { base } from '$lib/base'
 	import { Button } from '../common'
 
-	import { SUPPORTED_LANGUAGES, copilot } from './lib'
+	import { SUPPORTED_LANGUAGES } from './lib'
 	import type { SupportedLanguage } from '$lib/common'
-	import { sendUserToast } from '$lib/toast'
-	import type Editor from '../Editor.svelte'
-	import { dbSchemas, copilotInfo, type DBSchema, workspaceStore } from '$lib/stores'
-	import type DiffEditor from '../DiffEditor.svelte'
-	import { scriptLangToEditorLang } from '$lib/scripts'
 	import Popover from '$lib/components/meltComponents/Popover.svelte'
-	import { writable } from 'svelte/store'
-	import { WindmillIcon } from '../icons'
-	import HighlightCode from '../HighlightCode.svelte'
-	import LoadingIcon from '../apps/svelte-select/lib/LoadingIcon.svelte'
 	import { autoPlacement } from '@floating-ui/core'
-	import { Check, Wand2, X, RotateCw } from 'lucide-svelte'
-	import { createEventDispatcher } from 'svelte'
+	import { WandSparkles } from 'lucide-svelte'
+	import { aiChatManager } from './chat/AIChatManager.svelte'
+	import { copilotInfo } from '$lib/stores'
 
-	// props
-	export let lang: SupportedLanguage
-	export let editor: Editor | undefined
-	export let diffEditor: DiffEditor | undefined
-	export let error: string
-	export let args: Record<string, any>
-	export let chatMode: boolean = false
-
-	// state
-	let genLoading: boolean = false
-	let generatedCode = writable<string>('')
-	let generatedExplanation = writable<string>('')
-	let dbSchema: DBSchema | undefined = undefined
-	let abortController: AbortController | undefined = undefined
-
-	const dispatch = createEventDispatcher<{
-		fix: null
-	}>()
-
-	async function onFix() {
-		if (!error) {
-			return
-		}
-		try {
-			genLoading = true
-			abortController = new AbortController()
-			await copilot(
-				{
-					language: lang,
-					code: editor?.getCode() || '',
-					error,
-					dbSchema: dbSchema,
-					type: 'fix',
-					workspace: $workspaceStore!
-				},
-				generatedCode,
-				abortController,
-				generatedExplanation
-			)
-			setupDiff()
-			diffEditor?.setModified($generatedCode)
-			showDiff()
-		} catch (err) {
-			if (!abortController?.signal.aborted) {
-				if (err?.message) {
-					sendUserToast('Failed to generate code: ' + err.message, true)
-				} else {
-					sendUserToast('Failed to generate code', true)
-					console.error(err)
-				}
-			}
-		} finally {
-			genLoading = false
-		}
-	}
-
-	function acceptDiff() {
-		editor?.setCode(diffEditor?.getModified() || '')
-		editor?.format()
-		clear()
-	}
-
-	function rejectDiff() {
-		clear()
-	}
-
-	function setupDiff() {
-		diffEditor?.setupModel(scriptLangToEditorLang(lang))
-		diffEditor?.setOriginal(editor?.getCode() || '')
-	}
-
-	function showDiff() {
-		diffEditor?.show()
-		editor?.hide()
-	}
-
-	function hideDiff() {
-		editor?.show()
-		diffEditor?.hide()
-	}
-
-	function clear() {
-		$generatedCode = ''
-		$generatedExplanation = ''
-	}
-
-	$: lang && clear()
-
-	$: !$generatedCode && hideDiff()
-
-	function updateSchema(lang, args) {
-		const schemaRes = lang === 'graphql' ? args.api : args.database
-		if (typeof schemaRes === 'string') {
-			dbSchema = $dbSchemas[schemaRes.replace('$res:', '')]
-		}
-	}
-	$: updateSchema(lang, args)
-
-	let popover: Popover | undefined = undefined
+	let {
+		lang
+	}: {
+		lang: SupportedLanguage
+	} = $props()
 </script>
 
 {#if SUPPORTED_LANGUAGES.has(lang)}
-	{#if !genLoading && $generatedCode.length > 0}
-		<div class="flex gap-1">
-			<Button
-				title="Discard generated code"
-				size="xs"
-				color="red"
-				spacingSize="xs2"
-				on:click={() => {
-					popover?.close()
-					rejectDiff()
-				}}
-				variant="contained"
-				startIcon={{ icon: X }}
-				propagateEvent={true}
-			>
-				Discard
-			</Button>
-			<Button
-				title="Accept generated code"
-				size="xs"
-				color="green"
-				spacingSize="xs2"
-				on:click={() => {
-					popover?.close()
-					acceptDiff()
-				}}
-				startIcon={{ icon: Check }}
-				propagateEvent={true}
-			>
-				Accept
-			</Button>
-			<Button
-				size="xs"
-				color="light"
-				spacingSize="xs2"
-				on:click={() => {
-					$generatedExplanation = ''
-					popover?.open()
-					onFix()
-				}}
-				startIcon={{ icon: RotateCw }}
-				title="Retry"
-				btnClasses="text-violet-800 dark:text-violet-400"
-			/>
-		</div>
-	{/if}
 	<Popover
-		bind:this={popover}
 		floatingConfig={{
 			middleware: [
 				autoPlacement({
@@ -175,75 +26,39 @@
 				})
 			]
 		}}
-		closeOnOutsideClick={!genLoading}
-		closeButton={!genLoading}
 		displayArrow={true}
 	>
-		<svelte:fragment slot="trigger" let:isOpen>
+		<svelte:fragment slot="trigger">
 			<div class="flex flex-row">
 				<Button
 					title="Fix code"
 					size="xs"
-					color={genLoading ? 'red' : 'light'}
+					color="light"
 					spacingSize="xs2"
-					startIcon={genLoading ? undefined : { icon: Wand2 }}
-					propagateEvent={!chatMode}
-					on:click={chatMode
-						? () => dispatch('fix')
-						: genLoading
-						? () => abortController?.abort()
-						: $generatedCode.length > 0
-							? () => {}
-							: () => onFix()}
-					btnClasses={genLoading
-						? ''
-						: 'text-violet-800 dark:text-violet-400 bg-violet-100 dark:bg-gray-700 min-w-[84px]'}
+					startIcon={{ icon: WandSparkles }}
+					on:click={() => {
+						if ($copilotInfo.enabled) {
+							aiChatManager.fix()
+						}
+					}}
+					btnClasses="text-violet-800 dark:text-violet-400 bg-violet-100 dark:bg-gray-700 min-w-[84px]"
+					propagateEvent={!$copilotInfo.enabled}
 				>
-					{#if genLoading}
-						<WindmillIcon
-							white
-							class="mr-1 text-white"
-							height="16px"
-							width="20px"
-							spin="veryfast"
-						/>
-						Stop
-					{:else if $generatedCode.length > 0}
-						<span class="text-xs">{isOpen ? 'Hide' : 'Show'}</span>
-					{:else}
-						AI Fix
-					{/if}
+					AI Fix
 				</Button>
 			</div>
 		</svelte:fragment>
 		<svelte:fragment slot="content">
 			<div class="p-4">
-				{#if $copilotInfo.enabled}
-					<div class="w-[42rem] min-h-[3rem] max-h-[34rem] overflow-y-auto">
-						{#if $generatedCode.length > 0 && genLoading}
-							<div class="overflow-x-scroll">
-								<HighlightCode language={lang} code={$generatedCode} />
-							</div>
-						{:else if genLoading}
-							<LoadingIcon />
-						{/if}
-						{#if $generatedExplanation.length > 0}
-							<p class="text-sm mt-2"
-								><span class="font-bold">Explanation: test</span> {$generatedExplanation}</p
-							>
-						{/if}
-					</div>
-				{:else}
-					<div class="w-80">
-						<p class="text-sm"
-							>Enable Windmill AI in the <a
-								class="inline-flex flex-row items-center gap-1"
-								href="{base}/workspace_settings?tab=ai"
-								target="_blank">workspace settings</a
-							></p
-						></div
-					>
-				{/if}
+				<div class="w-80">
+					<p class="text-sm"
+						>Enable Windmill AI in the <a
+							class="inline-flex flex-row items-center gap-1"
+							href="{base}/workspace_settings?tab=ai"
+							target="_blank">workspace settings</a
+						></p
+					></div
+				>
 			</div>
 		</svelte:fragment>
 	</Popover>

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

@@ -474,12 +474,20 @@ class AIChatManager {
 	}
 
 	fix = () => {
+		if (!this.open) {
+			this.toggleOpen()
+		}
+		this.changeMode('script')
 		this.instructions = 'Fix the error'
 		this.contextManager?.setFixContext()
 		this.sendRequest()
 	}
 
 	addSelectedLinesToContext = (lines: string, startLine: number, endLine: number) => {
+		if (!this.open) {
+			this.toggleOpen()
+		}
+		this.changeMode('script')
 		this.contextManager?.addSelectedLinesToContext(lines, startLine, endLine)
 	}
 

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

@@ -12,6 +12,4 @@ export type FlowCopilotContext = {
 	exprsToSet: Writable<{
 		[key: string]: InputTransform | undefined
 	}>
-	toggleAiPanel?: () => void
-	addSelectedLinesToAiChat?: (lines: string, startLine: number, endLine: number) => void
 }

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

@@ -39,14 +39,6 @@
 		pickablePropertiesFiltered: writable<PickableProperties | undefined>(undefined)
 	})
 
-	export function addSelectedLinesToAiChat(lines: string, startLine: number, endLine: number) {
-		aiChatManager.addSelectedLinesToContext(lines, startLine, endLine)
-		if (!aiChatManager.open) {
-			aiChatManager.openChat()
-			aiChatManager.changeMode('script')
-		}
-	}
-
 	onMount(() => {
 		aiChatManager.changeMode('flow')
 	})

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

@@ -51,7 +51,7 @@
 	import { workspaceStore } from '$lib/stores'
 	import { checkIfParentLoop } from '../utils'
 	import ModulePreviewResultViewer from '$lib/components/ModulePreviewResultViewer.svelte'
-	import type { FlowCopilotContext } from '$lib/components/copilot/flow'
+	import { aiChatManager } from '$lib/components/copilot/chat/AIChatManager.svelte'
 
 	const {
 		selectedId,
@@ -66,9 +66,6 @@
 		executionCount
 	} = getContext<FlowEditorContext>('FlowEditorContext')
 
-	const { toggleAiPanel, addSelectedLinesToAiChat } =
-		getContext<FlowCopilotContext | undefined>('FlowCopilotContext') ?? {}
-
 	export let flowModule: FlowModule
 	export let failureModule: boolean = false
 	export let preprocessorModule: boolean = false
@@ -392,10 +389,10 @@
 										<Editor
 											on:addSelectedLinesToAiChat={(e) => {
 												const { lines, startLine, endLine } = e.detail
-												addSelectedLinesToAiChat?.(lines, startLine, endLine)
+												aiChatManager.addSelectedLinesToContext(lines, startLine, endLine)
 											}}
 											on:toggleAiPanel={() => {
-												toggleAiPanel?.()
+												aiChatManager.toggleOpen()
 											}}
 											loadAsync
 											folding

frontend/src/lib/components/scriptEditor/LogPanel.svelte

@@ -29,6 +29,7 @@
 	import WorkflowTimeline from '../WorkflowTimeline.svelte'
 	import Tooltip from '$lib/components/Tooltip.svelte'
 	import type { PreviewPanelUi } from '../custom_ui'
+	import { getStringError } from '../copilot/chat/utils'
 
 	export let lang: Preview['language'] | undefined
 	export let previewIsLoading = false
@@ -40,7 +41,6 @@
 	export let workspace: string | undefined = undefined
 	export let showCaptures: boolean = false
 	export let customUi: PreviewPanelUi | undefined = undefined
-	export let fixChatMode: boolean = false
 
 	type DrawerContent = {
 		mode: 'json' | Preview['language'] | 'plain'
@@ -141,16 +141,8 @@
 												language={lang}
 											>
 												<svelte:fragment slot="copilot-fix">
-													{#if lang && editor && diffEditor && args && previewJob?.result && typeof previewJob?.result == 'object' && `error` in previewJob?.result && previewJob?.result.error}
-														<ScriptFix
-															on:fix
-															chatMode={fixChatMode}
-															error={JSON.stringify(previewJob.result.error)}
-															{lang}
-															{editor}
-															{diffEditor}
-															{args}
-														/>
+													{#if lang && editor && diffEditor && args && previewJob && !previewJob.success && getStringError(previewJob.result)}
+														<ScriptFix {lang} />
 													{/if}
 												</svelte:fragment>
 											</DisplayResult>