From a6f5dc6910dbff22663fea49a9a54abbfbde8648 Mon Sep 17 00:00:00 2001
From: Aleksander Grygier <aleksander.grygier@gmail.com>
Date: Wed, 12 Aug 2026 10:00:19 +0200
Subject: [PATCH] refactor: Contexts & types

---
 tools/ui/src/app.d.ts                         |   2 +-
 .../ChatFormActionAddDropdown.svelte          |  60 ++++------
 .../ChatFormActionAddSheet.svelte             |  51 +++-----
 .../ChatFormActionsAdd.svelte                 |  55 +--------
 .../ChatFormActions/ChatFormActions.svelte    |  51 ++++++--
 .../ChatMessage/ChatMessage.svelte            |  19 ++-
 .../ChatMessageAssistant.svelte               |   4 +-
 .../ChatMessageMcpPrompt.svelte               |   4 +-
 .../ChatMessageSystem.svelte                  |   4 +-
 .../ChatMessageToolCallBlock.svelte           |   3 +-
 .../ChatMessageToolCallBlockDefault.svelte    |   3 +-
 ...essageToolCallBlockExecShellCommand.svelte |   3 +-
 .../ChatMessageUser/ChatMessageUser.svelte    |   4 +-
 .../ChatMessageUserPending.svelte             |   4 +-
 .../ChatMessageActionIcons.svelte             |   4 +-
 .../ChatMessageAgenticContent.svelte          |   2 +-
 .../ChatMessages/ChatMessageEditForm.svelte   |   4 +-
 .../ChatMessageReasoningBlock.svelte          |   3 +-
 .../app/chat/ChatMessages/ChatMessages.svelte |   7 +-
 .../lib/constants/context-keys.constants.ts   |   4 +-
 .../src/lib/contexts/chat-actions.context.ts  |  39 ------
 .../lib/contexts/chat-form-actions.context.ts |  19 +++
 .../contexts/chat-message-actions.context.ts  |  20 +---
 .../lib/contexts/chat-message-edit.context.ts |  19 +++
 tools/ui/src/lib/contexts/index.ts            |  20 +---
 .../src/lib/contexts/message-edit.context.ts  |  51 --------
 ...> use-chat-message-edit-context.svelte.ts} |   8 +-
 tools/ui/src/lib/types/chat.d.ts              | 111 +++++++++++++++++-
 tools/ui/src/lib/types/index.ts               |   8 ++
 .../tests/stories/ChatMessage.stories.svelte  |  15 ++-
 30 files changed, 294 insertions(+), 307 deletions(-)
 delete mode 100644 tools/ui/src/lib/contexts/chat-actions.context.ts
 create mode 100644 tools/ui/src/lib/contexts/chat-form-actions.context.ts
 create mode 100644 tools/ui/src/lib/contexts/chat-message-edit.context.ts
 delete mode 100644 tools/ui/src/lib/contexts/message-edit.context.ts
 rename tools/ui/src/lib/hooks/{use-message-edit-context.svelte.ts => use-chat-message-edit-context.svelte.ts} (91%)

diff --git a/tools/ui/src/app.d.ts b/tools/ui/src/app.d.ts
index 6f26214e1c..5309dce8f4 100644
--- a/tools/ui/src/app.d.ts
+++ b/tools/ui/src/app.d.ts
@@ -32,8 +32,8 @@ import type {
 	ApiRouterModelsStatusResponse,
 	ApiRouterModelsUnloadRequest,
 	ApiRouterModelsUnloadResponse,
-	// Chat types
 	ChatAttachmentDisplayItem,
+	// Chat types
 	ChatMessagePromptProgress,
 	ChatMessageSiblingInfo,
 	ChatMessageTimings,
diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddDropdown.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddDropdown.svelte
index 99aaee7845..02bfadb7e4 100644
--- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddDropdown.svelte
+++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddDropdown.svelte
@@ -15,37 +15,16 @@
 		ICON_CLASS_DEFAULT,
 		TOOLTIP_DELAY_DURATION
 	} from '$lib/constants';
+	import { getChatFormActionsContext } from '$lib/contexts';
 	import { useAttachmentMenu } from '$lib/hooks/use-attachment-menu.svelte';
 
 	interface Props {
 		class?: string;
-		disabled?: boolean;
-		hasAudioModality?: boolean;
-		hasVideoModality?: boolean;
-		hasVisionModality?: boolean;
-		hasMcpPromptsSupport?: boolean;
-		hasMcpResourcesSupport?: boolean;
-		onFileUpload?: () => void;
-		onSystemPromptClick?: () => void;
-		onMcpPromptClick?: () => void;
-		onMcpSettingsClick?: () => void;
-		onMcpResourcesClick?: () => void;
 	}
 
-	let {
-		class: className = '',
-		disabled = false,
-		hasAudioModality = false,
-		hasMcpPromptsSupport = false,
-		hasMcpResourcesSupport = false,
-		hasVideoModality = false,
-		hasVisionModality = false,
-		onFileUpload,
-		onMcpPromptClick,
-		onMcpResourcesClick,
-		onMcpSettingsClick,
-		onSystemPromptClick
-	}: Props = $props();
+	let { class: className = '' }: Props = $props();
+
+	const chatFormActions = getChatFormActionsContext();
 
 	let dropdownOpen = $state(false);
 	// The system message action moves focus to the message editor, so the menu
@@ -54,18 +33,23 @@
 
 	function handleMcpSettingsClick() {
 		dropdownOpen = false;
-		onMcpSettingsClick?.();
+		chatFormActions.onMcpSettingsClick?.();
 	}
 
 	const attachmentMenu = useAttachmentMenu(
 		() => ({
-			hasAudioModality,
-			hasMcpPromptsSupport,
-			hasMcpResourcesSupport,
-			hasVideoModality,
-			hasVisionModality
+			hasAudioModality: chatFormActions.hasAudioModality,
+			hasMcpPromptsSupport: chatFormActions.hasMcpPromptsSupport,
+			hasMcpResourcesSupport: chatFormActions.hasMcpResourcesSupport,
+			hasVideoModality: chatFormActions.hasVideoModality,
+			hasVisionModality: chatFormActions.hasVisionModality
+		}),
+		() => ({
+			onFileUpload: chatFormActions.onFileUpload,
+			onMcpPromptClick: chatFormActions.onMcpPromptClick,
+			onMcpResourcesClick: chatFormActions.onMcpResourcesClick,
+			onSystemPromptClick: chatFormActions.onSystemPromptClick
 		}),
-		() => ({ onFileUpload, onMcpPromptClick, onMcpResourcesClick, onSystemPromptClick }),
 		() => {
 			dropdownOpen = false;
 		}
@@ -85,7 +69,7 @@
 							buttonVariants({ variant: 'secondary' }),
 							'file-upload-button h-8 w-8 cursor-pointer rounded-full p-0'
 						)}
-						{disabled}
+						disabled={chatFormActions.disabled}
 					>
 						<span class="sr-only">{ATTACHMENT_TOOLTIP_TEXT}</span>
 
@@ -162,7 +146,7 @@
 				class="flex cursor-pointer items-center gap-2"
 				onclick={() => {
 					suppressCloseAutoFocus = true;
-					onSystemPromptClick?.();
+					chatFormActions.onSystemPromptClick?.();
 				}}
 			>
 				<MessageSquare class={ICON_CLASS_DEFAULT} />
@@ -174,12 +158,12 @@
 
 			<ChatFormActionAddMcpServersSubmenu onMcpSettingsClick={handleMcpSettingsClick} />
 
-			{#if hasMcpPromptsSupport}
+			{#if chatFormActions.hasMcpPromptsSupport}
 				<DropdownMenu.Separator />
 
 				<DropdownMenu.Item
 					class="flex cursor-pointer items-center gap-2"
-					onclick={onMcpPromptClick}
+					onclick={chatFormActions.onMcpPromptClick}
 				>
 					<Zap class={ICON_CLASS_DEFAULT} />
 
@@ -187,10 +171,10 @@
 				</DropdownMenu.Item>
 			{/if}
 
-			{#if hasMcpResourcesSupport}
+			{#if chatFormActions.hasMcpResourcesSupport}
 				<DropdownMenu.Item
 					class="flex cursor-pointer items-center gap-2"
-					onclick={onMcpResourcesClick}
+					onclick={chatFormActions.onMcpResourcesClick}
 				>
 					<FolderOpen class={ICON_CLASS_DEFAULT} />
 
diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddSheet.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddSheet.svelte
index aa5397ca0b..63a8c267d8 100644
--- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddSheet.svelte
+++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddSheet.svelte
@@ -19,6 +19,7 @@
 		ICON_CLASS_DEFAULT,
 		TOOLTIP_DELAY_DURATION
 	} from '$lib/constants';
+	import { getChatFormActionsContext } from '$lib/contexts';
 	import { HealthCheckStatus } from '$lib/enums';
 	import { AttachmentAction } from '$lib/enums/attachment.enums';
 	import { useAttachmentMenu } from '$lib/hooks/use-attachment-menu.svelte';
@@ -29,33 +30,12 @@
 
 	interface Props {
 		class?: string;
-		disabled?: boolean;
-		hasAudioModality?: boolean;
-		hasVideoModality?: boolean;
-		hasVisionModality?: boolean;
-		hasMcpPromptsSupport?: boolean;
-		hasMcpResourcesSupport?: boolean;
-		onFileUpload?: () => void;
-		onSystemPromptClick?: () => void;
-		onMcpPromptClick?: () => void;
-		onMcpResourcesClick?: () => void;
 		trigger: Snippet<[{ disabled: boolean; onclick?: () => void }]>;
 	}
 
-	let {
-		class: className = '',
-		disabled = false,
-		hasAudioModality = false,
-		hasMcpPromptsSupport = false,
-		hasMcpResourcesSupport = false,
-		hasVideoModality = false,
-		hasVisionModality = false,
-		onFileUpload,
-		onMcpPromptClick,
-		onMcpResourcesClick,
-		onSystemPromptClick,
-		trigger
-	}: Props = $props();
+	let { class: className = '', trigger }: Props = $props();
+
+	const chatFormActions = getChatFormActionsContext();
 
 	let sheetOpen = $state(false);
 	let reasoningExpanded = $state(false);
@@ -65,13 +45,18 @@
 
 	const attachmentMenu = useAttachmentMenu(
 		() => ({
-			hasAudioModality,
-			hasMcpPromptsSupport,
-			hasMcpResourcesSupport,
-			hasVideoModality,
-			hasVisionModality
+			hasAudioModality: chatFormActions.hasAudioModality,
+			hasMcpPromptsSupport: chatFormActions.hasMcpPromptsSupport,
+			hasMcpResourcesSupport: chatFormActions.hasMcpResourcesSupport,
+			hasVideoModality: chatFormActions.hasVideoModality,
+			hasVisionModality: chatFormActions.hasVisionModality
+		}),
+		() => ({
+			onFileUpload: chatFormActions.onFileUpload,
+			onMcpPromptClick: chatFormActions.onMcpPromptClick,
+			onMcpResourcesClick: chatFormActions.onMcpResourcesClick,
+			onSystemPromptClick: chatFormActions.onSystemPromptClick
 		}),
-		() => ({ onFileUpload, onMcpPromptClick, onMcpResourcesClick, onSystemPromptClick }),
 		() => {
 			sheetOpen = false;
 		}
@@ -91,7 +76,7 @@
 
 <div class="flex items-center gap-1 {className}">
 	<Sheet.Root bind:open={sheetOpen}>
-		{@render trigger({ disabled, onclick: () => (sheetOpen = true) })}
+		{@render trigger({ disabled: chatFormActions.disabled, onclick: () => (sheetOpen = true) })}
 
 		<Sheet.Content side="bottom" class="max-h-[85vh] gap-0 overflow-y-auto">
 			<Sheet.Header>
@@ -350,7 +335,7 @@
 					<span>System Message</span>
 				</button>
 
-				{#if hasMcpPromptsSupport}
+				{#if chatFormActions.hasMcpPromptsSupport}
 					<button
 						type="button"
 						class={sheetItemClass}
@@ -362,7 +347,7 @@
 					</button>
 				{/if}
 
-				{#if hasMcpResourcesSupport}
+				{#if chatFormActions.hasMcpResourcesSupport}
 					<button
 						type="button"
 						class={sheetItemClass}
diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionsAdd.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionsAdd.svelte
index a31ffb394d..47bdb47a47 100644
--- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionsAdd.svelte
+++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionsAdd.svelte
@@ -3,65 +3,14 @@
 	import ChatFormActionAddDropdown from './ChatFormActionAddDropdown.svelte';
 	import ChatFormActionAddSheet from './ChatFormActionAddSheet.svelte';
 	import { isMobile } from '$lib/stores';
-
-	interface Props {
-		disabled?: boolean;
-		hasAudioModality?: boolean;
-		hasVideoModality?: boolean;
-		hasMcpPromptsSupport?: boolean;
-		hasMcpResourcesSupport?: boolean;
-		hasVisionModality?: boolean;
-		onFileUpload?: () => void;
-		onMcpPromptClick?: () => void;
-		onMcpResourcesClick?: () => void;
-		onMcpSettingsClick?: () => void;
-		onSystemPromptClick?: () => void;
-	}
-
-	let {
-		disabled = false,
-		hasAudioModality = false,
-		hasMcpPromptsSupport = false,
-		hasMcpResourcesSupport = false,
-		hasVideoModality = false,
-		hasVisionModality = false,
-		onFileUpload,
-		onMcpPromptClick,
-		onMcpResourcesClick,
-		onMcpSettingsClick,
-		onSystemPromptClick
-	}: Props = $props();
 </script>
 
 {#if isMobile.current}
-	<ChatFormActionAddSheet
-		{disabled}
-		{hasAudioModality}
-		{hasVideoModality}
-		{hasVisionModality}
-		{hasMcpPromptsSupport}
-		{hasMcpResourcesSupport}
-		{onFileUpload}
-		{onSystemPromptClick}
-		{onMcpPromptClick}
-		{onMcpResourcesClick}
-	>
+	<ChatFormActionAddSheet>
 		{#snippet trigger({ disabled, onclick })}
 			<ChatFormActionAddButton {disabled} {onclick} />
 		{/snippet}
 	</ChatFormActionAddSheet>
 {:else}
-	<ChatFormActionAddDropdown
-		{disabled}
-		{hasAudioModality}
-		{hasVideoModality}
-		{hasVisionModality}
-		{hasMcpPromptsSupport}
-		{hasMcpResourcesSupport}
-		{onFileUpload}
-		{onMcpPromptClick}
-		{onMcpResourcesClick}
-		{onMcpSettingsClick}
-		{onSystemPromptClick}
-	/>
+	<ChatFormActionAddDropdown />
 {/if}
diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActions.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActions.svelte
index 13e97cc73c..d8fad772dd 100644
--- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActions.svelte
+++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActions.svelte
@@ -11,6 +11,7 @@
 	} from '$lib/components/app';
 	import { Button } from '$lib/components/ui/button';
 	import { ICON_CLASS_DEFAULT, ROUTES } from '$lib/constants';
+	import { setChatFormActionsContext } from '$lib/contexts';
 	import { FileTypeCategory, MessageRole } from '$lib/enums';
 	import { ChatService } from '$lib/services';
 	import { chatStore, conversationsStore, mcpStore, settingsStore } from '$lib/stores';
@@ -132,6 +133,42 @@
 
 		return livePromptTokens > 0 || liveOutputTokens > 0;
 	});
+
+	setChatFormActionsContext({
+		get disabled() {
+			return disabled;
+		},
+		get hasAudioModality() {
+			return hasAudioModality;
+		},
+		get hasMcpPromptsSupport() {
+			return hasMcpPromptsSupport;
+		},
+		get hasMcpResourcesSupport() {
+			return hasMcpResourcesSupport;
+		},
+		get hasVideoModality() {
+			return hasVideoModality;
+		},
+		get hasVisionModality() {
+			return hasVisionModality;
+		},
+		get onFileUpload() {
+			return onFileUpload;
+		},
+		get onMcpPromptClick() {
+			return onMcpPromptClick;
+		},
+		get onMcpResourcesClick() {
+			return onMcpResourcesClick;
+		},
+		get onMcpSettingsClick() {
+			return () => goto(ROUTES.MCP_SERVERS);
+		},
+		get onSystemPromptClick() {
+			return onSystemPromptClick;
+		}
+	});
 </script>
 
 <div
@@ -140,19 +177,7 @@
 >
 	{#if showAddButton}
 		<div class="mr-auto flex items-center gap-2">
-			<ChatFormActionsAdd
-				{disabled}
-				{hasAudioModality}
-				{hasVideoModality}
-				{hasVisionModality}
-				{hasMcpPromptsSupport}
-				{hasMcpResourcesSupport}
-				{onFileUpload}
-				{onSystemPromptClick}
-				{onMcpPromptClick}
-				{onMcpResourcesClick}
-				onMcpSettingsClick={() => goto(ROUTES.MCP_SERVERS)}
-			/>
+			<ChatFormActionsAdd />
 		</div>
 	{/if}
 
diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessage.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessage.svelte
index ebb49df6b5..78cb887217 100644
--- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessage.svelte
+++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessage.svelte
@@ -8,21 +8,21 @@
 		ChatMessageUser
 	} from '$lib/components/app/chat';
 	import { REASONING_TAGS, ROUTES, SYSTEM_MESSAGE_PLACEHOLDER } from '$lib/constants';
-	import {
-		type ChatMessageDeletionInfo,
-		getChatActionsContext,
-		setChatMessageActionsContext,
-		setMessageEditContext
-	} from '$lib/contexts';
+	import { setChatMessageActionsContext, setChatMessageEditContext } from '$lib/contexts';
 	import { AgenticSectionType, AttachmentType, MessageRole } from '$lib/enums';
 	import { DatabaseService } from '$lib/services/database.service';
 	import { chatStore, conversationsStore, isMobile } from '$lib/stores';
-	import type { DatabaseMessageExtraMcpPrompt } from '$lib/types';
+	import type {
+		ChatMessageActions,
+		ChatMessageDeletionInfo,
+		DatabaseMessageExtraMcpPrompt
+	} from '$lib/types';
 	import { deriveAgenticSections } from '$lib/utils';
 	import { parseFilesToMessageExtras } from '$lib/utils/browser-only';
 
 	interface Props {
 		class?: string;
+		chatActions: ChatMessageActions;
 		message: DatabaseMessage;
 		toolMessages?: DatabaseMessage[];
 		isLastAssistantMessage?: boolean;
@@ -32,6 +32,7 @@
 	}
 
 	let {
+		chatActions,
 		class: className = '',
 		isLastAssistantMessage = false,
 		isLastUserMessage = false,
@@ -41,8 +42,6 @@
 		toolMessages = []
 	}: Props = $props();
 
-	const chatActions = getChatActionsContext();
-
 	let deletionInfo = $state<ChatMessageDeletionInfo | null>(null);
 	// The system message placeholder must never surface as editable content; keeping
 	// it in the derived (not just in handleEdit) guards against prop invalidation
@@ -112,7 +111,7 @@
 	let showSaveOnlyOption = $derived(message.role === MessageRole.USER);
 	let showBranchAfterEditOption = $derived(message.role === MessageRole.ASSISTANT);
 
-	setMessageEditContext({
+	setChatMessageEditContext({
 		cancel: handleCancelEdit,
 		get editedContent() {
 			return editedContent;
diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageAssistant/ChatMessageAssistant.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageAssistant/ChatMessageAssistant.svelte
index 95e99a961c..b92be9fbd6 100644
--- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageAssistant/ChatMessageAssistant.svelte
+++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageAssistant/ChatMessageAssistant.svelte
@@ -8,7 +8,7 @@
 		ChatMessageAssistantStatistics,
 		ChatMessageEditForm
 	} from '$lib/components/app';
-	import { getMessageEditContext } from '$lib/contexts';
+	import { getChatMessageEditContext } from '$lib/contexts';
 	import { MessageRole } from '$lib/enums';
 	import { useProcessingState } from '$lib/hooks/use-processing-state.svelte';
 	import { chatStore, modelsStore, serverStore, settingsStore } from '$lib/stores';
@@ -36,7 +36,7 @@
 	}: Props = $props();
 
 	// Get edit context
-	const editCtx = getMessageEditContext();
+	const editCtx = getChatMessageEditContext();
 
 	const isAgentic = $derived(hasAgenticContent(message, toolMessages));
 	const processingState = useProcessingState();
diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageMcpPrompt/ChatMessageMcpPrompt.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageMcpPrompt/ChatMessageMcpPrompt.svelte
index 6b2c964ddd..4563b1fa86 100644
--- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageMcpPrompt/ChatMessageMcpPrompt.svelte
+++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageMcpPrompt/ChatMessageMcpPrompt.svelte
@@ -4,7 +4,7 @@
 		ChatMessageEditForm,
 		ChatMessageMcpPromptContent
 	} from '$lib/components/app';
-	import { getMessageEditContext } from '$lib/contexts';
+	import { getChatMessageEditContext } from '$lib/contexts';
 	import { McpPromptVariant, MessageRole } from '$lib/enums';
 	import type { DatabaseMessageExtraMcpPrompt } from '$lib/types';
 
@@ -17,7 +17,7 @@
 	let { class: className = '', mcpPrompt, message }: Props = $props();
 
 	// Get edit context
-	const editCtx = getMessageEditContext();
+	const editCtx = getChatMessageEditContext();
 </script>
 
 <div
diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageSystem/ChatMessageSystem.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageSystem/ChatMessageSystem.svelte
index 5f16b0f679..c6222f568c 100644
--- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageSystem/ChatMessageSystem.svelte
+++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageSystem/ChatMessageSystem.svelte
@@ -4,7 +4,7 @@
 	import { Button } from '$lib/components/ui/button';
 	import { Card } from '$lib/components/ui/card';
 	import { INPUT_CLASSES } from '$lib/constants';
-	import { getMessageEditContext } from '$lib/contexts';
+	import { getChatMessageEditContext } from '$lib/contexts';
 	import { KeyboardKey, MessageRole } from '$lib/enums';
 	import { settingsStore } from '$lib/stores';
 	import { autoResizeTextarea, isIMEComposing } from '$lib/utils';
@@ -17,7 +17,7 @@
 
 	let { class: className = '', message, textareaElement = $bindable() }: Props = $props();
 
-	const editCtx = getMessageEditContext();
+	const editCtx = getChatMessageEditContext();
 
 	function handleEditKeydown(event: KeyboardEvent) {
 		if (event.key === KeyboardKey.ENTER && !event.shiftKey && !isIMEComposing(event)) {
diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlock.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlock.svelte
index d0d8684e17..4de634264f 100644
--- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlock.svelte
+++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlock.svelte
@@ -11,8 +11,7 @@
 	import ChatMessageToolCallBlockSearchResults from './ChatMessageToolCallBlockSearchResults.svelte';
 	import ChatMessageToolCallBlockWriteFile from './ChatMessageToolCallBlockWriteFile.svelte';
 	import { BuiltInTool } from '$lib/enums';
-	import type { AgenticSection } from '$lib/types';
-	import type { DatabaseMessageExtra } from '$lib/types';
+	import type { AgenticSection, DatabaseMessageExtra } from '$lib/types';
 	import { extractSearchQuery, extractSearchResults, isWebSearchToolName } from '$lib/utils';
 
 	interface Props {
diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockDefault.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockDefault.svelte
index 9009eae3d2..0010440376 100644
--- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockDefault.svelte
+++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockDefault.svelte
@@ -8,8 +8,7 @@
 	import { MarkdownContent, SyntaxHighlightedCode } from '$lib/components/app';
 	import { MAX_HEIGHT_CODE_BLOCK } from '$lib/constants';
 	import { FileTypeText, ToolResultKind } from '$lib/enums';
-	import type { AgenticSection } from '$lib/types';
-	import type { DatabaseMessageExtra } from '$lib/types';
+	import type { AgenticSection, DatabaseMessageExtra } from '$lib/types';
 	import {
 		classifyToolResult,
 		formatJsonPretty,
diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockExecShellCommand.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockExecShellCommand.svelte
index 395938b517..bd3a5f2630 100644
--- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockExecShellCommand.svelte
+++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockExecShellCommand.svelte
@@ -12,8 +12,7 @@
 	import { CollapsibleTerminalBlock } from '$lib/components/app';
 	import { SETTINGS_KEYS, TOOL_RUNTIME_SCROLL_AT_BOTTOM_THRESHOLD_PX } from '$lib/constants';
 	import { settingsStore, toolsStore } from '$lib/stores';
-	import type { AgenticSection, ToolResultLine } from '$lib/types';
-	import type { DatabaseMessageExtra } from '$lib/types';
+	import type { AgenticSection, DatabaseMessageExtra, ToolResultLine } from '$lib/types';
 	import {
 		abbreviateHome,
 		type ExecShellExitStatus,
diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUser.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUser.svelte
index 54f6400d5d..d217a9c382 100644
--- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUser.svelte
+++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUser.svelte
@@ -5,7 +5,7 @@
 		ChatMessageStatistics,
 		ChatMessageUserBubble
 	} from '$lib/components/app/chat';
-	import { getMessageEditContext } from '$lib/contexts';
+	import { getChatMessageEditContext } from '$lib/contexts';
 	import { ChatMessageStatisticsMode, MessageRole } from '$lib/enums';
 	import { useProcessingState } from '$lib/hooks/use-processing-state.svelte';
 	import { chatStore, settingsStore } from '$lib/stores';
@@ -25,7 +25,7 @@
 	}: Props = $props();
 
 	// Get contexts
-	const editCtx = getMessageEditContext();
+	const editCtx = getChatMessageEditContext();
 	const processingState = useProcessingState();
 
 	const currentConfig = $derived(settingsStore.config);
diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUserPending.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUserPending.svelte
index f3e8c4e91d..a072f2e84d 100644
--- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUserPending.svelte
+++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUserPending.svelte
@@ -1,7 +1,7 @@
 <script lang="ts">
 	import { ArrowUp, Edit, Trash2 } from '@lucide/svelte';
 	import { ActionIcon, ChatMessageEditForm, ChatMessageUserBubble } from '$lib/components/app';
-	import { useMessageEditContext } from '$lib/hooks/use-message-edit-context.svelte';
+	import { useChatMessageEditContext } from '$lib/hooks/use-chat-message-edit-context.svelte';
 
 	interface Props {
 		class?: string;
@@ -21,7 +21,7 @@
 		onSendImmediately
 	}: Props = $props();
 
-	const editCtx = useMessageEditContext({
+	const editCtx = useChatMessageEditContext({
 		getContent: () => content,
 		getExtras: () => extras,
 		onSave: (content, extras) => onEdit(content, extras)
diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageActions/ChatMessageActionIcons/ChatMessageActionIcons.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageActions/ChatMessageActionIcons/ChatMessageActionIcons.svelte
index e379bf2502..895f391255 100644
--- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageActions/ChatMessageActionIcons/ChatMessageActionIcons.svelte
+++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageActions/ChatMessageActionIcons/ChatMessageActionIcons.svelte
@@ -9,7 +9,7 @@
 	import Input from '$lib/components/ui/input/input.svelte';
 	import Label from '$lib/components/ui/label/label.svelte';
 	import { Switch } from '$lib/components/ui/switch';
-	import { getChatMessageActionsContext, getMessageEditContext } from '$lib/contexts';
+	import { getChatMessageActionsContext, getChatMessageEditContext } from '$lib/contexts';
 	import { MessageRole } from '$lib/enums';
 	import { conversationsStore } from '$lib/stores';
 
@@ -36,7 +36,7 @@
 	}: Props = $props();
 
 	const messageActions = getChatMessageActionsContext();
-	const editCtx = getMessageEditContext();
+	const editCtx = getChatMessageEditContext();
 
 	let showForkDialog = $state(false);
 	let forkName = $state('');
diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageAgenticContent.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageAgenticContent.svelte
index 8ce1ead01e..5849799794 100644
--- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageAgenticContent.svelte
+++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageAgenticContent.svelte
@@ -9,8 +9,8 @@
 	} from '$lib/components/app';
 	import { AgenticSectionType, ChatMessageStatsView, ToolPermissionDecision } from '$lib/enums';
 	import { agenticStore, settingsStore } from '$lib/stores';
-	import type { AgenticSection } from '$lib/types';
 	import type {
+		AgenticSection,
 		ChatMessageAgenticTimings,
 		ChatMessageAgenticTurnStats,
 		DatabaseMessage
diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageEditForm.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageEditForm.svelte
index feacd0ddde..369b6137b4 100644
--- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageEditForm.svelte
+++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageEditForm.svelte
@@ -3,12 +3,12 @@
 	import { ChatForm, DialogConfirmation } from '$lib/components/app';
 	import { Button } from '$lib/components/ui/button';
 	import { Switch } from '$lib/components/ui/switch';
-	import { getMessageEditContext } from '$lib/contexts';
+	import { getChatMessageEditContext } from '$lib/contexts';
 	import { KeyboardKey, MessageRole } from '$lib/enums';
 	import { chatStore } from '$lib/stores';
 	import { processFilesToChatUploaded } from '$lib/utils/browser-only';
 
-	const editCtx = getMessageEditContext();
+	const editCtx = getChatMessageEditContext();
 
 	let saveWithoutRegenerate = $state(false);
 	let showDiscardDialog = $state(false);
diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageReasoningBlock.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageReasoningBlock.svelte
index 78c9dcb25a..a0861fb961 100644
--- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageReasoningBlock.svelte
+++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageReasoningBlock.svelte
@@ -4,8 +4,7 @@
 	import { REASONING_SCROLL_AT_BOTTOM_THRESHOLD_PX } from '$lib/constants';
 	import { AgenticSectionType } from '$lib/enums';
 	import { settingsStore } from '$lib/stores';
-	import type { DatabaseMessageExtra } from '$lib/types';
-	import type { AgenticSection } from '$lib/types';
+	import type { AgenticSection, DatabaseMessageExtra } from '$lib/types';
 
 	interface Props {
 		section: AgenticSection;
diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessages.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessages.svelte
index 33d971f931..2a8f45ba53 100644
--- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessages.svelte
+++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessages.svelte
@@ -1,8 +1,8 @@
 <script lang="ts">
 	import { ChatMessage, ChatMessageUserPending } from '$lib/components/app';
-	import { setChatActionsContext } from '$lib/contexts';
 	import { MessageRole } from '$lib/enums';
 	import { agenticStore, chatStore, conversationsStore, settingsStore } from '$lib/stores';
+	import type { ChatMessageActions } from '$lib/types';
 	import {
 		buildSiblingInfoMap,
 		copyToClipboard,
@@ -22,7 +22,7 @@
 
 	const currentConfig = settingsStore.config;
 
-	setChatActionsContext({
+	const chatActions: ChatMessageActions = {
 		continueAssistantMessage: async (message: DatabaseMessage) => {
 			onUserAction?.();
 			await chatStore.continueAssistantMessage(message.id);
@@ -91,7 +91,7 @@
 			await chatStore.regenerateMessageWithBranching(message.id, modelOverride);
 			refreshAllMessages();
 		}
-	});
+	};
 
 	function refreshAllMessages() {
 		const conversation = conversationsStore.activeConversation;
@@ -228,6 +228,7 @@
 	{#each displayMessages as { isLastAssistantMessage, isLastUserMessage, message, nextAssistantMessage, siblingInfo, toolMessages } (message.id)}
 		<ChatMessage
 			class="mx-auto mt-12 w-full max-w-3xl"
+			{chatActions}
 			{message}
 			{toolMessages}
 			{isLastAssistantMessage}
diff --git a/tools/ui/src/lib/constants/context-keys.constants.ts b/tools/ui/src/lib/constants/context-keys.constants.ts
index 81966bbcac..62ff5413e1 100644
--- a/tools/ui/src/lib/constants/context-keys.constants.ts
+++ b/tools/ui/src/lib/constants/context-keys.constants.ts
@@ -1,3 +1,3 @@
-export const CONTEXT_KEY_MESSAGE_EDIT = 'chat-message-edit';
-export const CONTEXT_KEY_CHAT_ACTIONS = 'chat-actions';
+export const CONTEXT_KEY_CHAT_MESSAGE_EDIT = 'chat-message-edit';
 export const CONTEXT_KEY_CHAT_MESSAGE_ACTIONS = 'chat-message-actions';
+export const CONTEXT_KEY_CHAT_FORM_ACTIONS = 'chat-form-actions';
diff --git a/tools/ui/src/lib/contexts/chat-actions.context.ts b/tools/ui/src/lib/contexts/chat-actions.context.ts
deleted file mode 100644
index dffb5f3a36..0000000000
--- a/tools/ui/src/lib/contexts/chat-actions.context.ts
+++ /dev/null
@@ -1,39 +0,0 @@
-import { CONTEXT_KEY_CHAT_ACTIONS } from '$lib/constants';
-import { getContext, setContext } from 'svelte';
-
-export interface ChatActionsContext {
-	copy: (message: DatabaseMessage) => void;
-	delete: (message: DatabaseMessage) => void;
-	navigateToSibling: (siblingId: string) => void;
-	editWithBranching: (
-		message: DatabaseMessage,
-		newContent: string,
-		newExtras?: DatabaseMessageExtra[]
-	) => void;
-	editWithReplacement: (
-		message: DatabaseMessage,
-		newContent: string,
-		shouldBranch: boolean
-	) => void;
-	editUserMessagePreserveResponses: (
-		message: DatabaseMessage,
-		newContent: string,
-		newExtras?: DatabaseMessageExtra[]
-	) => void;
-	regenerateWithBranching: (message: DatabaseMessage, modelOverride?: string) => void;
-	continueAssistantMessage: (message: DatabaseMessage) => void;
-	forkConversation: (
-		message: DatabaseMessage,
-		options: { name: string; includeAttachments: boolean }
-	) => void;
-}
-
-const CHAT_ACTIONS_KEY = Symbol.for(CONTEXT_KEY_CHAT_ACTIONS);
-
-export function setChatActionsContext(ctx: ChatActionsContext): ChatActionsContext {
-	return setContext(CHAT_ACTIONS_KEY, ctx);
-}
-
-export function getChatActionsContext(): ChatActionsContext {
-	return getContext(CHAT_ACTIONS_KEY);
-}
diff --git a/tools/ui/src/lib/contexts/chat-form-actions.context.ts b/tools/ui/src/lib/contexts/chat-form-actions.context.ts
new file mode 100644
index 0000000000..a49f17447e
--- /dev/null
+++ b/tools/ui/src/lib/contexts/chat-form-actions.context.ts
@@ -0,0 +1,19 @@
+import { CONTEXT_KEY_CHAT_FORM_ACTIONS } from '$lib/constants';
+import type { ChatFormActionsContext } from '$lib/types';
+import { getContext, setContext } from 'svelte';
+
+const CHAT_FORM_ACTIONS_KEY = Symbol.for(CONTEXT_KEY_CHAT_FORM_ACTIONS);
+
+/**
+ * Sets the chat form actions context. Call in the parent component (ChatFormActions.svelte).
+ */
+export function setChatFormActionsContext(ctx: ChatFormActionsContext): ChatFormActionsContext {
+	return setContext(CHAT_FORM_ACTIONS_KEY, ctx);
+}
+
+/**
+ * Gets the chat form actions context. Call in child components.
+ */
+export function getChatFormActionsContext(): ChatFormActionsContext {
+	return getContext(CHAT_FORM_ACTIONS_KEY);
+}
diff --git a/tools/ui/src/lib/contexts/chat-message-actions.context.ts b/tools/ui/src/lib/contexts/chat-message-actions.context.ts
index cd4213a05a..fb075b3b02 100644
--- a/tools/ui/src/lib/contexts/chat-message-actions.context.ts
+++ b/tools/ui/src/lib/contexts/chat-message-actions.context.ts
@@ -1,25 +1,7 @@
 import { CONTEXT_KEY_CHAT_MESSAGE_ACTIONS } from '$lib/constants';
+import type { ChatMessageActionsContext } from '$lib/types';
 import { getContext, setContext } from 'svelte';
 
-export interface ChatMessageDeletionInfo {
-	totalCount: number;
-	userMessages: number;
-	assistantMessages: number;
-	messageTypes: string[];
-}
-
-export interface ChatMessageActionsContext {
-	readonly siblingInfo: ChatMessageSiblingInfo | null;
-	readonly deletionInfo: ChatMessageDeletionInfo | null;
-	readonly showDeleteDialog: boolean;
-	copy: () => void;
-	requestDelete: () => void;
-	confirmDelete: () => void;
-	setShowDeleteDialog: (show: boolean) => void;
-	navigateToSibling: (siblingId: string) => void;
-	forkConversation?: (options: { name: string; includeAttachments: boolean }) => void;
-}
-
 const CHAT_MESSAGE_ACTIONS_KEY = Symbol.for(CONTEXT_KEY_CHAT_MESSAGE_ACTIONS);
 
 /**
diff --git a/tools/ui/src/lib/contexts/chat-message-edit.context.ts b/tools/ui/src/lib/contexts/chat-message-edit.context.ts
new file mode 100644
index 0000000000..e9c053036e
--- /dev/null
+++ b/tools/ui/src/lib/contexts/chat-message-edit.context.ts
@@ -0,0 +1,19 @@
+import { CONTEXT_KEY_CHAT_MESSAGE_EDIT } from '$lib/constants';
+import type { ChatMessageEditContext } from '$lib/types';
+import { getContext, setContext } from 'svelte';
+
+const CHAT_MESSAGE_EDIT_KEY = Symbol.for(CONTEXT_KEY_CHAT_MESSAGE_EDIT);
+
+/**
+ * Sets the message edit context. Call this in the parent component (ChatMessage.svelte).
+ */
+export function setChatMessageEditContext(ctx: ChatMessageEditContext): ChatMessageEditContext {
+	return setContext(CHAT_MESSAGE_EDIT_KEY, ctx);
+}
+
+/**
+ * Gets the message edit context. Call this in child components.
+ */
+export function getChatMessageEditContext(): ChatMessageEditContext {
+	return getContext(CHAT_MESSAGE_EDIT_KEY);
+}
diff --git a/tools/ui/src/lib/contexts/index.ts b/tools/ui/src/lib/contexts/index.ts
index 4390d397bc..4aaec8148c 100644
--- a/tools/ui/src/lib/contexts/index.ts
+++ b/tools/ui/src/lib/contexts/index.ts
@@ -1,20 +1,8 @@
-export {
-	getMessageEditContext,
-	setMessageEditContext,
-	type MessageEditContext,
-	type MessageEditState,
-	type MessageEditActions
-} from './message-edit.context';
-
-export {
-	getChatActionsContext,
-	setChatActionsContext,
-	type ChatActionsContext
-} from './chat-actions.context';
+export { getChatMessageEditContext, setChatMessageEditContext } from './chat-message-edit.context';
 
 export {
 	getChatMessageActionsContext,
-	setChatMessageActionsContext,
-	type ChatMessageActionsContext,
-	type ChatMessageDeletionInfo
+	setChatMessageActionsContext
 } from './chat-message-actions.context';
+
+export { getChatFormActionsContext, setChatFormActionsContext } from './chat-form-actions.context';
diff --git a/tools/ui/src/lib/contexts/message-edit.context.ts b/tools/ui/src/lib/contexts/message-edit.context.ts
deleted file mode 100644
index 80f3e6eee2..0000000000
--- a/tools/ui/src/lib/contexts/message-edit.context.ts
+++ /dev/null
@@ -1,51 +0,0 @@
-import { CONTEXT_KEY_MESSAGE_EDIT } from '$lib/constants';
-import { MessageRole } from '$lib/enums';
-import { getContext, setContext } from 'svelte';
-
-export interface MessageEditState {
-	readonly isEditing: boolean;
-	readonly editedContent: string;
-	readonly editedExtras: DatabaseMessageExtra[];
-	readonly editedUploadedFiles: ChatUploadedFile[];
-	readonly originalContent: string;
-	readonly originalExtras: DatabaseMessageExtra[];
-	readonly showSaveOnlyOption: boolean;
-	readonly showBranchAfterEditOption: boolean;
-	readonly shouldBranchAfterEdit: boolean;
-	readonly messageRole: MessageRole;
-	readonly rawEditContent?: string;
-}
-
-export interface MessageEditActions {
-	setContent: (content: string) => void;
-	setExtras: (extras: DatabaseMessageExtra[]) => void;
-	setUploadedFiles: (files: ChatUploadedFile[]) => void;
-	save: () => void;
-	saveOnly: () => void;
-	cancel: () => void;
-	startEdit: () => void;
-}
-
-export interface AssistantEditActions {
-	setShouldBranchAfterEdit: (value: boolean) => void;
-}
-
-export type MessageEditContext = MessageEditState &
-	MessageEditActions &
-	Partial<AssistantEditActions>;
-
-const MESSAGE_EDIT_KEY = Symbol.for(CONTEXT_KEY_MESSAGE_EDIT);
-
-/**
- * Sets the message edit context. Call this in the parent component (ChatMessage.svelte).
- */
-export function setMessageEditContext(ctx: MessageEditContext): MessageEditContext {
-	return setContext(MESSAGE_EDIT_KEY, ctx);
-}
-
-/**
- * Gets the message edit context. Call this in child components.
- */
-export function getMessageEditContext(): MessageEditContext {
-	return getContext(MESSAGE_EDIT_KEY);
-}
diff --git a/tools/ui/src/lib/hooks/use-message-edit-context.svelte.ts b/tools/ui/src/lib/hooks/use-chat-message-edit-context.svelte.ts
similarity index 91%
rename from tools/ui/src/lib/hooks/use-message-edit-context.svelte.ts
rename to tools/ui/src/lib/hooks/use-chat-message-edit-context.svelte.ts
index 2716754045..de2994e739 100644
--- a/tools/ui/src/lib/hooks/use-message-edit-context.svelte.ts
+++ b/tools/ui/src/lib/hooks/use-chat-message-edit-context.svelte.ts
@@ -1,15 +1,15 @@
-import { setMessageEditContext } from '$lib/contexts';
+import { setChatMessageEditContext } from '$lib/contexts';
 import { MessageRole } from '$lib/enums';
 import { parseFilesToMessageExtras } from '$lib/utils/convert-files-to-extra';
 
-interface UseMessageEditContextOptions {
+interface UseChatMessageEditContextOptions {
 	getContent: () => string;
 	getExtras: () => DatabaseMessageExtra[];
 	showSaveOnlyOption?: boolean;
 	onSave: (content: string, extras?: DatabaseMessageExtra[]) => void;
 }
 
-export function useMessageEditContext(options: UseMessageEditContextOptions) {
+export function useChatMessageEditContext(options: UseChatMessageEditContextOptions) {
 	let isEditing = $state(false);
 	let editedContent = $state('');
 	let editedExtras = $state<DatabaseMessageExtra[]>([]);
@@ -45,7 +45,7 @@ export function useMessageEditContext(options: UseMessageEditContextOptions) {
 		isEditing = false;
 	}
 
-	setMessageEditContext({
+	setChatMessageEditContext({
 		cancel: handleCancelEdit,
 		get editedContent() {
 			return editedContent;
diff --git a/tools/ui/src/lib/types/chat.d.ts b/tools/ui/src/lib/types/chat.d.ts
index cb93e77951..f0f3a297e8 100644
--- a/tools/ui/src/lib/types/chat.d.ts
+++ b/tools/ui/src/lib/types/chat.d.ts
@@ -7,7 +7,8 @@ import type {
 	AttachmentMenuItemId,
 	ChatFormCommandAction,
 	ErrorDialogType,
-	FileMentionEntryType
+	FileMentionEntryType,
+	MessageRole
 } from '$lib/enums';
 import type { Component } from 'svelte';
 
@@ -235,3 +236,111 @@ export interface ChatFormCommand {
 	action: ChatFormCommandAction;
 	disabled: boolean;
 }
+
+/**
+ * Data shown in the message delete confirmation dialog.
+ */
+export interface ChatMessageDeletionInfo {
+	totalCount: number;
+	userMessages: number;
+	assistantMessages: number;
+	messageTypes: string[];
+}
+
+/**
+ * Conversation-level message operations owned by ChatMessages (store calls + list
+ * refresh + user-action notification), passed to each ChatMessage as a prop.
+ */
+export interface ChatMessageActions {
+	copy: (message: DatabaseMessage) => void;
+	delete: (message: DatabaseMessage) => void;
+	navigateToSibling: (siblingId: string) => void;
+	editWithBranching: (
+		message: DatabaseMessage,
+		newContent: string,
+		newExtras?: DatabaseMessageExtra[]
+	) => void;
+	editWithReplacement: (
+		message: DatabaseMessage,
+		newContent: string,
+		shouldBranch: boolean
+	) => void;
+	editUserMessagePreserveResponses: (
+		message: DatabaseMessage,
+		newContent: string,
+		newExtras?: DatabaseMessageExtra[]
+	) => void;
+	regenerateWithBranching: (message: DatabaseMessage, modelOverride?: string) => void;
+	continueAssistantMessage: (message: DatabaseMessage) => void;
+	forkConversation: (
+		message: DatabaseMessage,
+		options: { name: string; includeAttachments: boolean }
+	) => void;
+}
+
+/**
+ * Per-message actions and state. Set once per message in ChatMessage.svelte and
+ * consumed by its descendants (action icons, branching controls).
+ */
+export interface ChatMessageActionsContext {
+	readonly siblingInfo: ChatMessageSiblingInfo | null;
+	readonly deletionInfo: ChatMessageDeletionInfo | null;
+	readonly showDeleteDialog: boolean;
+	copy: () => void;
+	requestDelete: () => void;
+	confirmDelete: () => void;
+	setShowDeleteDialog: (show: boolean) => void;
+	navigateToSibling: (siblingId: string) => void;
+	forkConversation?: (options: { name: string; includeAttachments: boolean }) => void;
+}
+
+export interface ChatMessageEditState {
+	readonly isEditing: boolean;
+	readonly editedContent: string;
+	readonly editedExtras: DatabaseMessageExtra[];
+	readonly editedUploadedFiles: ChatUploadedFile[];
+	readonly originalContent: string;
+	readonly originalExtras: DatabaseMessageExtra[];
+	readonly showSaveOnlyOption: boolean;
+	readonly showBranchAfterEditOption: boolean;
+	readonly shouldBranchAfterEdit: boolean;
+	readonly messageRole: MessageRole;
+	readonly rawEditContent?: string;
+}
+
+export interface ChatMessageEditActions {
+	setContent: (content: string) => void;
+	setExtras: (extras: DatabaseMessageExtra[]) => void;
+	setUploadedFiles: (files: ChatUploadedFile[]) => void;
+	save: () => void;
+	saveOnly: () => void;
+	cancel: () => void;
+	startEdit: () => void;
+}
+
+export interface ChatMessageAssistantEditActions {
+	setShouldBranchAfterEdit: (value: boolean) => void;
+}
+
+export type ChatMessageEditContext = ChatMessageEditState &
+	ChatMessageEditActions &
+	Partial<ChatMessageAssistantEditActions>;
+
+/**
+ * Actions and capability flags for the ChatForm add-menu. Set once in
+ * ChatFormActions.svelte and consumed by its deep descendants (the add sheet,
+ * dropdown and MCP servers submenu) to avoid relaying them through props.
+ */
+export interface ChatFormActionsContext {
+	readonly disabled: boolean;
+	readonly hasAudioModality: boolean;
+	readonly hasVideoModality: boolean;
+	readonly hasVisionModality: boolean;
+	readonly hasMcpPromptsSupport: boolean;
+	readonly hasMcpResourcesSupport: boolean;
+	onFileUpload?: () => void;
+	onSystemPromptClick?: () => void;
+	onMcpPromptClick?: () => void;
+	onMcpResourcesClick?: () => void;
+	onMcpSettingsClick?: () => void;
+}
diff --git a/tools/ui/src/lib/types/index.ts b/tools/ui/src/lib/types/index.ts
index 917c9516b5..34767e2ae7 100644
--- a/tools/ui/src/lib/types/index.ts
+++ b/tools/ui/src/lib/types/index.ts
@@ -44,6 +44,14 @@ export type {
 	ChatUploadedFile,
 	ChatAttachmentDisplayItem,
 	ChatMessageSiblingInfo,
+	ChatMessageActions,
+	ChatMessageActionsContext,
+	ChatMessageDeletionInfo,
+	ChatMessageEditContext,
+	ChatMessageEditState,
+	ChatMessageEditActions,
+	ChatMessageAssistantEditActions,
+	ChatFormActionsContext,
 	ChatMessagePromptProgress,
 	ChatMessageTimings,
 	ChatMessageAgenticTimings,
diff --git a/tools/ui/tests/stories/ChatMessage.stories.svelte b/tools/ui/tests/stories/ChatMessage.stories.svelte
index 023ba1ac4f..84fee2ea1c 100644
--- a/tools/ui/tests/stories/ChatMessage.stories.svelte
+++ b/tools/ui/tests/stories/ChatMessage.stories.svelte
@@ -1,6 +1,7 @@
 <script module lang="ts">
 	import { defineMeta } from '@storybook/addon-svelte-csf';
 	import ChatMessage from '$lib/components/app/chat/ChatMessages/ChatMessage/ChatMessage.svelte';
+	import type { ChatMessageActions } from '$lib/types';
 
 	const { Story } = defineMeta({
 		component: ChatMessage,
@@ -10,6 +11,18 @@
 		title: 'Components/ChatScreen/ChatMessage'
 	});
 
+	const chatActions: ChatMessageActions = {
+		continueAssistantMessage: () => {},
+		copy: () => {},
+		delete: () => {},
+		editUserMessagePreserveResponses: () => {},
+		editWithBranching: () => {},
+		editWithReplacement: () => {},
+		forkConversation: () => {},
+		navigateToSibling: () => {},
+		regenerateWithBranching: () => {}
+	};
+
 	// Mock messages for different scenarios
 	const userMessage: DatabaseMessage = {
 		children: [],
@@ -190,7 +203,7 @@
 	}}
 >
 	<div class="w-[56rem]">
-		<ChatMessage message={streamingMessage} />
+		<ChatMessage message={streamingMessage} {chatActions} />
 	</div>
 </Story>