From cb4d57cb625fdaa4dd7d9d07e46fd87a35fcc145 Mon Sep 17 00:00:00 2001
From: "claude[bot]" <209825114+claude[bot]@users.noreply.github.com>
Date: Fri, 4 Jul 2025 13:11:29 +0200
Subject: [PATCH] feat: allow editing messages in AI chat (#6117)

* feat: add message editing functionality to AI chat

- Add editMessage() method to AIChatManager for editing user messages
- Add edit button UI that appears on hover for all user messages
- Transform message display into textarea when editing
- Remove subsequent messages and resend on Enter
- Add keyboard shortcuts: Enter to save, Escape to cancel
- Maintain existing restart generation functionality

Co-authored-by: Ruben Fiszel <rubenfiszel@users.noreply.github.com>

* fix

* big cleaning

* use aichatinput when editing message

* cleaning

* fix

* remove logs

* better retry button

* fix

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

Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>

* cleaning

* fix retry logic

---------

Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com>
Co-authored-by: Ruben Fiszel <rubenfiszel@users.noreply.github.com>
Co-authored-by: centdix <farhadg110@gmail.com>
Co-authored-by: centdix <40307056+centdix@users.noreply.github.com>
Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>
---
 .../copilot/chat/AIChatDisplay.svelte         | 206 ++----------------
 .../copilot/chat/AIChatInput.svelte           | 162 ++++++++++++++
 .../copilot/chat/AIChatManager.svelte.ts      |  69 ++++--
 .../copilot/chat/AIChatMessage.svelte         | 112 ++++++++++
 .../copilot/chat/ContextTextarea.svelte       |  42 ++--
 .../src/lib/components/copilot/chat/shared.ts |  35 ++-
 6 files changed, 386 insertions(+), 240 deletions(-)
 create mode 100644 frontend/src/lib/components/copilot/chat/AIChatInput.svelte
 create mode 100644 frontend/src/lib/components/copilot/chat/AIChatMessage.svelte

diff --git a/frontend/src/lib/components/copilot/chat/AIChatDisplay.svelte b/frontend/src/lib/components/copilot/chat/AIChatDisplay.svelte
index aea971c0f7..11f42445de 100644
--- a/frontend/src/lib/components/copilot/chat/AIChatDisplay.svelte
+++ b/frontend/src/lib/components/copilot/chat/AIChatDisplay.svelte
@@ -1,31 +1,17 @@
 <script lang="ts">
-	import { twMerge } from 'tailwind-merge'
-	import AssistantMessage from './AssistantMessage.svelte'
+	import AIChatMessage from './AIChatMessage.svelte'
 	import { type Snippet } from 'svelte'
-	import {
-		CheckIcon,
-		HistoryIcon,
-		Loader2,
-		Plus,
-		RefreshCwIcon,
-		StopCircleIcon,
-		Undo2Icon,
-		X,
-		XIcon
-	} from 'lucide-svelte'
-	import autosize from '$lib/autosize'
+	import { CheckIcon, HistoryIcon, Loader2, Plus, StopCircleIcon, X, XIcon } from 'lucide-svelte'
 	import Button from '$lib/components/common/button/Button.svelte'
 	import Popover from '$lib/components/meltComponents/Popover.svelte'
 	import { type DisplayMessage } from './shared'
 	import type { ContextElement } from './context'
-	import ContextElementBadge from './ContextElementBadge.svelte'
-	import ContextTextarea from './ContextTextarea.svelte'
-	import AvailableContextList from './AvailableContextList.svelte'
 	import ChatQuickActions from './ChatQuickActions.svelte'
 	import ProviderModelSelector from './ProviderModelSelector.svelte'
 	import ChatMode from './ChatMode.svelte'
 	import Markdown from 'svelte-exmarkdown'
 	import { aiChatManager, AIMode } from './AIChatManager.svelte'
+	import AIChatInput from './AIChatInput.svelte'
 
 	let {
 		messages,
@@ -63,16 +49,8 @@
 		suggestions?: string[]
 	} = $props()
 
-	let contextTextareaComponent: ContextTextarea | undefined = $state()
-	let instructionsTextarea: HTMLTextAreaElement | undefined = $state()
-
-	export function focusInput() {
-		if (aiChatManager.mode === 'script') {
-			contextTextareaComponent?.focus()
-		} else {
-			instructionsTextarea?.focus()
-		}
-	}
+	let aiChatInput: AIChatInput | undefined = $state()
+	let editingMessageIndex = $state<number | null>(null)
 
 	let scrollEl: HTMLDivElement | undefined = $state()
 	async function scrollDown() {
@@ -87,38 +65,12 @@
 		aiChatManager.automaticScroll && height && scrollDown()
 	})
 
-	function addContextToSelection(contextElement: ContextElement) {
-		if (
-			selectedContext &&
-			availableContext &&
-			!selectedContext.find(
-				(c) => c.type === contextElement.type && c.title === contextElement.title
-			) &&
-			availableContext.find(
-				(c) => c.type === contextElement.type && c.title === contextElement.title
-			)
-		) {
-			selectedContext = [...selectedContext, contextElement]
-		}
-	}
-
 	function submitSuggestion(suggestion: string) {
-		aiChatManager.instructions = suggestion
-		aiChatManager.sendRequest()
+		aiChatManager.sendRequest({ instructions: suggestion })
 	}
 
-	function isLastUserMessage(messageIndex: number): boolean {
-		// Find the last user message index
-		for (let i = messages.length - 1; i >= 0; i--) {
-			if (messages[i].role === 'user') {
-				return i === messageIndex
-			}
-		}
-		return false
-	}
-
-	function restartGeneration(messageIndex: number) {
-		aiChatManager.restartLastGeneration(messageIndex)
+	export function focusInput() {
+		aiChatInput?.focusInput()
 	}
 </script>
 
@@ -208,68 +160,13 @@
 		>
 			<div class="flex flex-col" bind:clientHeight={height}>
 				{#each messages as message, messageIndex}
-					<div class={twMerge(message.role === 'user' && messageIndex > 0 && 'mt-6', 'mb-2')}>
-						{#if message.role === 'user' && message.contextElements}
-							<div class="flex flex-row gap-1 mb-1 overflow-scroll no-scrollbar px-2">
-								{#each message.contextElements as element}
-									<ContextElementBadge contextElement={element} />
-								{/each}
-							</div>
-						{/if}
-						<div
-							class={twMerge(
-								'text-sm py-1 mx-2',
-								message.role === 'user' &&
-									'px-2 border border-gray-300 dark:border-gray-600 bg-gray-50 dark:bg-gray-900 rounded-lg relative group',
-								(message.role === 'assistant' || message.role === 'tool') && 'px-[1px]',
-								message.role === 'tool' && 'text-tertiary'
-							)}
-						>
-							{#if message.role === 'assistant'}
-								<AssistantMessage {message} />
-							{:else}
-								{message.content}
-							{/if}
-
-							{#if message.role === 'user' && isLastUserMessage(messageIndex) && !aiChatManager.loading}
-								<div
-									class="absolute top-1 right-1 opacity-0 group-hover:opacity-100 transition-opacity"
-								>
-									<Button
-										size="xs2"
-										variant="border"
-										color="light"
-										iconOnly
-										title="Restart generation"
-										startIcon={{ icon: RefreshCwIcon }}
-										btnClasses="!p-1 !h-6 !w-6"
-										on:click={() => restartGeneration(messageIndex)}
-									/>
-								</div>
-							{/if}
-						</div>
-						{#if message.role === 'user' && message.snapshot}
-							<div
-								class="mx-2 text-sm text-tertiary flex flex-row items-center justify-between gap-2 mt-2"
-							>
-								Saved a flow snapshot
-								<Button
-									size="xs2"
-									variant="border"
-									color="light"
-									on:click={() => {
-										if (message.snapshot) {
-											aiChatManager.flowAiChatHelpers?.revertToSnapshot(message.snapshot)
-										}
-									}}
-									title="Revert to snapshot"
-									startIcon={{ icon: Undo2Icon }}
-								>
-									Revert
-								</Button>
-							</div>
-						{/if}
-					</div>
+					<AIChatMessage
+						{message}
+						{messageIndex}
+						{availableContext}
+						bind:selectedContext
+						bind:editingMessageIndex
+					/>
 				{/each}
 				{#if aiChatManager.loading && !aiChatManager.currentReply}
 					<div class="mb-6 py-1 px-2">
@@ -323,71 +220,14 @@
 				</Button>
 			</div>
 		{/if}
-		{#if aiChatManager.mode === 'script'}
-			<div class="flex flex-row gap-1 mb-1 overflow-scroll pt-2 px-2 no-scrollbar">
-				<Popover>
-					<svelte:fragment slot="trigger">
-						<div
-							class="border rounded-md px-1 py-0.5 font-normal text-tertiary text-xs hover:bg-surface-hover"
-							>@</div
-						>
-					</svelte:fragment>
-					<svelte:fragment slot="content" let:close>
-						<AvailableContextList
-							{availableContext}
-							{selectedContext}
-							onSelect={(element) => {
-								addContextToSelection(element)
-								close()
-							}}
-						/>
-					</svelte:fragment>
-				</Popover>
-				{#each selectedContext as element}
-					<ContextElementBadge
-						contextElement={element}
-						deletable
-						on:delete={() => {
-							selectedContext = selectedContext?.filter(
-								(c) => c.type !== element.type || c.title !== element.title
-							)
-						}}
-					/>
-				{/each}
-			</div>
-			<ContextTextarea
-				bind:this={contextTextareaComponent}
-				{availableContext}
-				{selectedContext}
-				isFirstMessage={messages.length === 0}
-				onAddContext={(contextElement) => addContextToSelection(contextElement)}
-				onSendRequest={() => {
-					if (!aiChatManager.loading) {
-						aiChatManager.sendRequest()
-					}
-				}}
-				onUpdateInstructions={(value) => (aiChatManager.instructions = value)}
-				{disabled}
-			/>
-		{:else}
-			<div class="relative w-full px-2 scroll-pb-2 pt-2">
-				<textarea
-					bind:this={instructionsTextarea}
-					bind:value={aiChatManager.instructions}
-					use:autosize
-					onkeydown={(e) => {
-						if (e.key === 'Enter' && !e.shiftKey && !aiChatManager.loading) {
-							e.preventDefault()
-							aiChatManager.sendRequest()
-						}
-					}}
-					rows={3}
-					placeholder={messages.length === 0 ? 'Ask anything' : 'Ask followup'}
-					class="resize-none"
-					{disabled}
-				></textarea>
-			</div>
-		{/if}
+		<AIChatInput
+			bind:this={aiChatInput}
+			bind:selectedContext
+			{availableContext}
+			{disabled}
+			isFirstMessage={messages.length === 0}
+			placeholder={messages.length === 0 ? 'Ask anything' : 'Ask followup'}
+		/>
 		<div
 			class={`flex flex-row ${
 				aiChatManager.mode === 'script' && hasDiff ? 'justify-between' : 'justify-end'
diff --git a/frontend/src/lib/components/copilot/chat/AIChatInput.svelte b/frontend/src/lib/components/copilot/chat/AIChatInput.svelte
new file mode 100644
index 0000000000..926c55be80
--- /dev/null
+++ b/frontend/src/lib/components/copilot/chat/AIChatInput.svelte
@@ -0,0 +1,162 @@
+<script lang="ts">
+	import Popover from '$lib/components/meltComponents/Popover.svelte'
+	import AvailableContextList from './AvailableContextList.svelte'
+	import ContextElementBadge from './ContextElementBadge.svelte'
+	import ContextTextarea from './ContextTextarea.svelte'
+	import autosize from '$lib/autosize'
+	import type { ContextElement } from './context'
+	import { aiChatManager } from './AIChatManager.svelte'
+
+	interface Props {
+		availableContext: ContextElement[]
+		selectedContext: ContextElement[]
+		isFirstMessage?: boolean
+		disabled?: boolean
+		placeholder?: string
+		initialInstructions?: string
+		editingMessageIndex?: number | null
+		onEditEnd?: () => void
+	}
+
+	let {
+		availableContext,
+		selectedContext = $bindable([]),
+		disabled = false,
+		isFirstMessage = false,
+		placeholder = 'Ask anything',
+		initialInstructions = '',
+		editingMessageIndex = null,
+		onEditEnd = () => {}
+	}: Props = $props()
+
+	let contextTextareaComponent: ContextTextarea | undefined = $state()
+	let instructionsTextareaComponent: HTMLTextAreaElement | undefined = $state()
+	let instructions = $state(initialInstructions)
+
+	export function focusInput() {
+		if (aiChatManager.mode === 'script') {
+			contextTextareaComponent?.focus()
+		} else {
+			instructionsTextareaComponent?.focus()
+		}
+	}
+
+	function clickOutside(node: HTMLElement) {
+		function handleClick(event: MouseEvent) {
+			if (node && !node.contains(event.target as Node) && editingMessageIndex !== null) {
+				onEditEnd()
+			}
+		}
+
+		document.addEventListener('click', handleClick, true)
+		return {
+			destroy() {
+				document.removeEventListener('click', handleClick, true)
+			}
+		}
+	}
+
+	function addContextToSelection(contextElement: ContextElement) {
+		if (
+			selectedContext &&
+			availableContext &&
+			!selectedContext.find(
+				(c) => c.type === contextElement.type && c.title === contextElement.title
+			) &&
+			availableContext.find(
+				(c) => c.type === contextElement.type && c.title === contextElement.title
+			)
+		) {
+			selectedContext = [...selectedContext, contextElement]
+		}
+	}
+
+	function sendRequest() {
+		if (aiChatManager.loading) {
+			return
+		}
+		if (editingMessageIndex !== null) {
+			aiChatManager.restartGeneration(editingMessageIndex, instructions)
+			onEditEnd()
+		} else {
+			aiChatManager.sendRequest({ instructions })
+			instructions = ''
+		}
+	}
+
+	$effect(() => {
+		if (editingMessageIndex !== null) {
+			focusInput()
+		}
+	})
+</script>
+
+<div use:clickOutside>
+	{#if aiChatManager.mode === 'script'}
+		<div class="flex flex-row gap-1 mb-1 overflow-scroll pt-2 px-2 no-scrollbar">
+			<Popover>
+				<svelte:fragment slot="trigger">
+					<div
+						class="border rounded-md px-1 py-0.5 font-normal text-tertiary text-xs hover:bg-surface-hover"
+						>@</div
+					>
+				</svelte:fragment>
+				<svelte:fragment slot="content" let:close>
+					<AvailableContextList
+						{availableContext}
+						{selectedContext}
+						onSelect={(element) => {
+							addContextToSelection(element)
+							close()
+						}}
+					/>
+				</svelte:fragment>
+			</Popover>
+			{#each selectedContext as element}
+				<ContextElementBadge
+					contextElement={element}
+					deletable
+					on:delete={() => {
+						selectedContext = selectedContext?.filter(
+							(c) => c.type !== element.type || c.title !== element.title
+						)
+					}}
+				/>
+			{/each}
+		</div>
+		<ContextTextarea
+			bind:this={contextTextareaComponent}
+			bind:value={instructions}
+			{availableContext}
+			{selectedContext}
+			{isFirstMessage}
+			{placeholder}
+			onAddContext={(contextElement) => addContextToSelection(contextElement)}
+			onSendRequest={() => {
+				sendRequest()
+			}}
+			{disabled}
+			onEscape={onEditEnd}
+		/>
+	{:else}
+		<div class="relative w-full px-2 scroll-pb-2 pt-2">
+			<textarea
+				bind:this={instructionsTextareaComponent}
+				bind:value={instructions}
+				use:autosize
+				onkeydown={(e) => {
+					if (e.key === 'Enter' && !e.shiftKey) {
+						e.preventDefault()
+						sendRequest()
+					} else if (e.key === 'Escape') {
+						onEditEnd()
+					}
+				}}
+				rows={3}
+				{placeholder}
+				class="resize-none"
+				{disabled}
+			></textarea>
+		</div>
+	{/if}
+</div>
diff --git a/frontend/src/lib/components/copilot/chat/AIChatManager.svelte.ts b/frontend/src/lib/components/copilot/chat/AIChatManager.svelte.ts
index 2480304522..ebff218afe 100644
--- a/frontend/src/lib/components/copilot/chat/AIChatManager.svelte.ts
+++ b/frontend/src/lib/components/copilot/chat/AIChatManager.svelte.ts
@@ -32,7 +32,6 @@ import type { DBSchemas } from '$lib/stores'
 import { askTools, prepareAskSystemMessage } from './ask/core'
 import { chatState, DEFAULT_SIZE, triggerablesByAi } from './sharedChatState.svelte'
 
-
 export enum AIMode {
 	SCRIPT = 'script',
 	FLOW = 'flow',
@@ -217,6 +216,32 @@ class AIChatManager {
 		}
 	}
 
+	retryRequest = (messageIndex: number) => {
+		const message = this.displayMessages[messageIndex]
+		if (message && message.role === 'user') {
+			this.restartGeneration(messageIndex)
+			message.error = false
+		} else {
+			throw new Error('No user message found at the specified index')
+		}
+	}
+
+	private getLastUserMessage = () => {
+		for (let i = this.displayMessages.length - 1; i >= 0; i--) {
+			const message = this.displayMessages[i]
+			if (message.role === 'user') {
+				return message
+			}
+		}
+	}
+
+	private flagLastMessageAsError = () => {
+		const lastUserMessage = this.getLastUserMessage()
+		if (lastUserMessage) {
+			lastUserMessage.error = true
+		}
+	}
+
 	private chatRequest = async ({
 		messages,
 		abortController,
@@ -403,7 +428,8 @@ class AIChatManager {
 					role: 'user',
 					content: this.instructions,
 					contextElements: this.mode === AIMode.SCRIPT ? oldSelectedContext : undefined,
-					snapshot
+					snapshot,
+					index: this.messages.length // matching with actual messages index. not -1 because it's not yet added to the messages array
 				}
 			]
 			const oldInstructions = this.instructions
@@ -423,8 +449,8 @@ class AIChatManager {
 					: this.mode === AIMode.NAVIGATOR
 						? prepareNavigatorUserMessage(oldInstructions)
 						: await prepareScriptUserMessage(oldInstructions, lang, oldSelectedContext, {
-							isPreprocessor
-						})
+								isPreprocessor
+							})
 
 			this.messages.push(userMessage)
 			await this.historyManager.saveChat(this.displayMessages, this.messages)
@@ -478,6 +504,7 @@ class AIChatManager {
 			await this.historyManager.saveChat(this.displayMessages, this.messages)
 		} catch (err) {
 			console.error(err)
+			this.flagLastMessageAsError()
 			if (err instanceof Error) {
 				sendUserToast('Failed to send request: ' + err.message, true)
 			} else {
@@ -492,7 +519,7 @@ class AIChatManager {
 		this.abortController?.abort()
 	}
 
-	restartLastGeneration = (displayMessageIndex: number) => {
+	restartGeneration = (displayMessageIndex: number, newContent?: string) => {
 		const userMessage = this.displayMessages[displayMessageIndex]
 
 		if (!userMessage || userMessage.role !== 'user') {
@@ -502,23 +529,17 @@ class AIChatManager {
 		// Remove all messages including and after the specified user message
 		this.displayMessages = this.displayMessages.slice(0, displayMessageIndex)
 
-		// Find the last user message in actual messages and remove it and everything after it
-		let lastActualUserMessageIndex = -1
-		for (let i = this.messages.length - 1; i >= 0; i--) {
-			if (this.messages[i].role === 'user') {
-				lastActualUserMessageIndex = i
-				break
-			}
-		}
+		// Find corresponding message in actual messages and remove it and everything after it
+		let actualMessageIndex = this.messages.findIndex((_, i) => i === userMessage.index)
 
-		if (lastActualUserMessageIndex === -1) {
+		if (actualMessageIndex === -1) {
 			throw new Error('No actual user message found to restart from')
 		}
 
-		this.messages = this.messages.slice(0, lastActualUserMessageIndex)
+		this.messages = this.messages.slice(0, actualMessageIndex)
 
 		// Resend the request with the same instructions
-		this.instructions = userMessage.content
+		this.instructions = newContent ?? userMessage.content
 		this.sendRequest()
 	}
 
@@ -648,15 +669,15 @@ class AIChatManager {
 				const editorRelated =
 					currentEditor && currentEditor.type === 'script' && currentEditor.stepId === module.id
 						? {
-							diffMode: currentEditor.diffMode,
-							lastDeployedCode: currentEditor.lastDeployedCode,
-							lastSavedCode: undefined
-						}
+								diffMode: currentEditor.diffMode,
+								lastDeployedCode: currentEditor.lastDeployedCode,
+								lastSavedCode: undefined
+							}
 						: {
-							diffMode: false,
-							lastDeployedCode: undefined,
-							lastSavedCode: undefined
-						}
+								diffMode: false,
+								lastDeployedCode: undefined,
+								lastSavedCode: undefined
+							}
 
 				return {
 					args: moduleState?.previewArgs ?? {},
diff --git a/frontend/src/lib/components/copilot/chat/AIChatMessage.svelte b/frontend/src/lib/components/copilot/chat/AIChatMessage.svelte
new file mode 100644
index 0000000000..66f6585a40
--- /dev/null
+++ b/frontend/src/lib/components/copilot/chat/AIChatMessage.svelte
@@ -0,0 +1,112 @@
+<script lang="ts">
+	import { twMerge } from 'tailwind-merge'
+	import type { DisplayMessage } from './shared'
+	import ContextElementBadge from './ContextElementBadge.svelte'
+	import AssistantMessage from './AssistantMessage.svelte'
+	import { aiChatManager } from './AIChatManager.svelte'
+	import { Button } from '$lib/components/common'
+	import { RefreshCwIcon, Undo2Icon } from 'lucide-svelte'
+	import AIChatInput from './AIChatInput.svelte'
+	import type { ContextElement } from './context'
+
+	interface Props {
+		availableContext: ContextElement[]
+		selectedContext: ContextElement[]
+		message: DisplayMessage
+		messageIndex: number
+		editingMessageIndex: number | null
+	}
+
+	let {
+		message,
+		messageIndex,
+		availableContext,
+		selectedContext = $bindable(),
+		editingMessageIndex = $bindable(null)
+	}: Props = $props()
+
+	function editMessage() {
+		if (message.role !== 'user' || editingMessageIndex !== null || aiChatManager.loading) {
+			return
+		}
+		editingMessageIndex = messageIndex
+	}
+</script>
+
+<div
+	class={twMerge(
+		message.role === 'user' && messageIndex > 0 && 'mt-6',
+		'mb-2',
+		message.role !== 'user' ? 'cursor-default' : 'cursor-pointer'
+	)}
+	role="button"
+	tabindex="0"
+	onclick={() => editMessage()}
+	onkeydown={() => {}}
+>
+	{#if message.role === 'user' && message.contextElements && editingMessageIndex !== messageIndex}
+		<div class="flex flex-row gap-1 mb-1 overflow-scroll no-scrollbar px-2">
+			{#each message.contextElements as element}
+				<ContextElementBadge contextElement={element} />
+			{/each}
+		</div>
+	{/if}
+	{#if message.role === 'user' && editingMessageIndex === messageIndex}
+		<AIChatInput
+			{availableContext}
+			bind:selectedContext
+			initialInstructions={message.content}
+			{editingMessageIndex}
+			onEditEnd={() => (editingMessageIndex = null)}
+		/>
+	{:else}
+		<div
+			class={twMerge(
+				'text-sm py-1 mx-2',
+				message.role === 'user' &&
+					'px-2 border border-gray-300 dark:border-gray-600 bg-gray-50 dark:bg-gray-900 rounded-lg relative group',
+				(message.role === 'assistant' || message.role === 'tool') && 'px-[1px]',
+				message.role === 'tool' && 'text-tertiary'
+			)}
+		>
+			{#if message.role === 'assistant'}
+				<AssistantMessage {message} />
+			{:else}
+				{message.content}
+			{/if}
+		</div>
+	{/if}
+	{#if message.role === 'user' && message.snapshot}
+		<div class="mx-2 text-sm text-tertiary flex flex-row items-center justify-between gap-2 mt-2">
+			Saved a flow snapshot
+			<Button
+				size="xs2"
+				variant="border"
+				color="light"
+				on:click={() => {
+					if (message.snapshot) {
+						aiChatManager.flowAiChatHelpers?.revertToSnapshot(message.snapshot)
+					}
+				}}
+				title="Revert to snapshot"
+				startIcon={{ icon: Undo2Icon }}
+			>
+				Revert
+			</Button>
+		</div>
+	{/if}
+</div>
+{#if message.role === 'user' && message.error}
+	<div class="flex justify-end px-2 -mt-1">
+		<Button
+			size="xs2"
+			variant="border"
+			title="Retry generation"
+			color="light"
+			startIcon={{ icon: RefreshCwIcon }}
+			onclick={() => aiChatManager.retryRequest(messageIndex)}
+		>
+			Retry
+		</Button>
+	</div>
+{/if}
diff --git a/frontend/src/lib/components/copilot/chat/ContextTextarea.svelte b/frontend/src/lib/components/copilot/chat/ContextTextarea.svelte
index cfec386ba3..b348862048 100644
--- a/frontend/src/lib/components/copilot/chat/ContextTextarea.svelte
+++ b/frontend/src/lib/components/copilot/chat/ContextTextarea.svelte
@@ -3,28 +3,31 @@
 	import { tick } from 'svelte'
 	import type { ContextElement } from './context'
 	import AvailableContextList from './AvailableContextList.svelte'
-	import { aiChatManager } from './AIChatManager.svelte'
 	import Portal from '$lib/components/Portal.svelte'
 	import { zIndexes } from '$lib/zIndexes'
 
 	interface Props {
+		value: string
 		availableContext: ContextElement[]
 		selectedContext: ContextElement[]
 		isFirstMessage: boolean
+		placeholder: string
 		disabled: boolean
-		onUpdateInstructions: (value: string) => void
 		onSendRequest: () => void
 		onAddContext: (contextElement: ContextElement) => void
+		onEscape: () => void
 	}
 
-	const {
+	let {
+		value = $bindable(''),
 		availableContext,
 		selectedContext,
 		isFirstMessage,
+		placeholder,
 		disabled,
-		onUpdateInstructions,
 		onSendRequest,
-		onAddContext
+		onAddContext,
+		onEscape
 	}: Props = $props()
 
 	let showContextTooltip = $state(false)
@@ -163,11 +166,10 @@
 	}
 
 	function updateInstructionsWithContext(contextElement: ContextElement) {
-		const index = aiChatManager.instructions.lastIndexOf('@')
+		const index = value.lastIndexOf('@')
 		if (index !== -1) {
-			const newInstructions =
-				aiChatManager.instructions.substring(0, index) + `@${contextElement.title}`
-			onUpdateInstructions(newInstructions)
+			const newInstructions = value.substring(0, index) + `@${contextElement.title}`
+			value = newInstructions
 		}
 	}
 
@@ -252,7 +254,7 @@
 
 	function handleInput(e: Event) {
 		textarea = e.target as HTMLTextAreaElement
-		const words = aiChatManager.instructions.split(/\s+/)
+		const words = value.split(/\s+/)
 		const lastWord = words[words.length - 1]
 
 		if (
@@ -267,7 +269,6 @@
 			contextTooltipWord = ''
 			selectedSuggestionIndex = 0
 		}
-		onUpdateInstructions(aiChatManager.instructions)
 	}
 
 	function handleKeyPress(e: KeyboardEvent) {
@@ -283,10 +284,7 @@
 						(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 &&
-						aiChatManager.instructions.split(' ').pop() === '@' + contextElement.title
-					) {
+					if (isInSelectedContext && value.split(' ').pop() === '@' + contextElement.title) {
 						onSendRequest()
 						return
 					}
@@ -301,6 +299,10 @@
 	}
 
 	function handleKeyDown(e: KeyboardEvent) {
+		if (e.key === 'Escape') {
+			onEscape()
+		}
+
 		if (!showContextTooltip) return
 
 		const filteredContext = availableContext.filter(
@@ -337,14 +339,14 @@
 <div class="relative w-full px-2 scroll-pb-2">
 	<div class="textarea-input absolute top-0 left-0 pointer-events-none">
 		<span class="break-words">
-			{@html getHighlightedText(aiChatManager.instructions)}
+			{@html getHighlightedText(value)}
 		</span>
 	</div>
 	<textarea
 		bind:this={textarea}
 		onkeypress={handleKeyPress}
 		onkeydown={handleKeyDown}
-		bind:value={aiChatManager.instructions}
+		bind:value
 		use:autosize
 		rows={3}
 		oninput={handleInput}
@@ -353,11 +355,9 @@
 				showContextTooltip = false
 			}, 200)
 		}}
-		placeholder={isFirstMessage ? 'Ask anything' : 'Ask followup'}
+		{placeholder}
 		class="textarea-input resize-none bg-transparent caret-black dark:caret-white"
-		style={aiChatManager.instructions.length > 0
-			? 'color: transparent; -webkit-text-fill-color: transparent;'
-			: ''}
+		style={value.length > 0 ? 'color: transparent; -webkit-text-fill-color: transparent;' : ''}
 		{disabled}
 	></textarea>
 </div>
diff --git a/frontend/src/lib/components/copilot/chat/shared.ts b/frontend/src/lib/components/copilot/chat/shared.ts
index f1fa174f38..e0ef74bb40 100644
--- a/frontend/src/lib/components/copilot/chat/shared.ts
+++ b/frontend/src/lib/components/copilot/chat/shared.ts
@@ -8,18 +8,29 @@ import type { ContextElement } from './context'
 import { workspaceStore } from '$lib/stores'
 import type { ExtendedOpenFlow } from '$lib/components/flows/types'
 
-export type DisplayMessage =
-	| {
-			role: 'user' | 'assistant'
-			content: string
-			contextElements?: ContextElement[]
-			snapshot?: ExtendedOpenFlow
-	  }
-	| {
-			role: 'tool'
-			tool_call_id: string
-			content: string
-	  }
+type BaseDisplayMessage = {
+	content: string
+	contextElements?: ContextElement[]
+	snapshot?: ExtendedOpenFlow
+}
+
+export type UserDisplayMessage = BaseDisplayMessage & {
+	role: 'user'
+	index: number // Used to match index with actual chat messages
+	error?: boolean
+}
+
+export type ToolDisplayMessage = {
+	role: 'tool'
+	tool_call_id: string
+	content: string
+}
+
+export type AssistantDisplayMessage = BaseDisplayMessage & {
+	role: 'assistant'
+}
+
+export type DisplayMessage = UserDisplayMessage | ToolDisplayMessage | AssistantDisplayMessage
 
 async function callTool<T>({
 	tools,