wip

app/ui/app/codegen/gotypes.gen.ts

@@ -213,6 +213,44 @@ export class ChatResponse {
 	    return a;
 	}
 }
+export class MessageUpdateRequest {
+    content: string;
+
+    constructor(source: any = {}) {
+        if ('string' === typeof source) source = JSON.parse(source);
+        this.content = source["content"];
+    }
+}
+export class MessageUpdateResponse {
+    index: number;
+    chatId: string;
+    message: Message;
+
+    constructor(source: any = {}) {
+        if ('string' === typeof source) source = JSON.parse(source);
+        this.index = source["index"];
+        this.chatId = source["chatId"];
+        this.message = this.convertValues(source["message"], Message);
+    }
+
+	convertValues(a: any, classs: any, asMap: boolean = false): any {
+	    if (!a) {
+	        return a;
+	    }
+	    if (Array.isArray(a)) {
+	        return (a as any[]).map(elem => this.convertValues(elem, classs));
+	    } else if ("object" === typeof a) {
+	        if (asMap) {
+	            for (const key of Object.keys(a)) {
+	                a[key] = new classs(a[key]);
+	            }
+	            return a;
+	        }
+	        return new classs(a);
+	    }
+	    return a;
+	}
+}
 export class Model {
     model: string;
     digest?: string;

app/ui/app/components.json

@@ -0,0 +1,20 @@
+{
+  "$schema": "https://ui.shadcn.com/schema.json",
+  "style": "default",
+  "rsc": false,
+  "tsx": true,
+  "tailwind": {
+    "config": "tailwind.config.js",
+    "css": "src/index.css",
+    "baseColor": "neutral",
+    "cssVariables": true,
+    "prefix": ""
+  },
+  "aliases": {
+    "components": "@/components",
+    "utils": "@/lib/utils",
+    "ui": "@/components/ui",
+    "lib": "@/lib",
+    "hooks": "@/hooks"
+  }
+}

app/ui/app/package-lock.json

@@ -26,6 +26,7 @@
         "rehype-sanitize": "^6.0.0",
         "remark-math": "^6.0.0",
         "streamdown": "^1.4.0",
+        "tailwind-merge": "^3.4.0",
         "unist-builder": "^4.0.0",
         "unist-util-parents": "^3.0.0"
       },
@@ -12110,9 +12111,9 @@
       "integrity": "sha512-Cat63mxsVJlzYvN51JmVXIgNoUokrIaT2zLclCXjRd8boZ0004U4KCs/sToJ75C6sdlByWxpYnb5Boif1VSFew=="
     },
     "node_modules/tailwind-merge": {
-      "version": "3.3.1",
-      "resolved": "https://registry.npmjs.org/tailwind-merge/-/tailwind-merge-3.3.1.tgz",
-      "integrity": "sha512-gBXpgUm/3rp1lMZZrM/w7D8GKqshif0zAymAhbCyIt8KMe+0v9DQ7cdYLR4FHH/cKpdTXb+A/tKKU3eolfsI+g==",
+      "version": "3.4.0",
+      "resolved": "https://registry.npmjs.org/tailwind-merge/-/tailwind-merge-3.4.0.tgz",
+      "integrity": "sha512-uSaO4gnW+b3Y2aWoWfFpX62vn2sR3skfhbjsEnaBI81WD1wBLlHZe5sWf0AqjksNdYTbGBEd0UasQMT3SNV15g==",
       "license": "MIT",
       "funding": {
         "type": "github",

app/ui/app/package.json

@@ -35,6 +35,7 @@
     "rehype-sanitize": "^6.0.0",
     "remark-math": "^6.0.0",
     "streamdown": "^1.4.0",
+    "tailwind-merge": "^3.4.0",
     "unist-builder": "^4.0.0",
     "unist-util-parents": "^3.0.0"
   },

app/ui/app/src/api.ts

@@ -11,6 +11,7 @@ import {
   ChatRequest,
   Settings,
   User,
+  Message,
 } from "@/gotypes";
 import { parseJsonlFromResponse } from "./util/jsonl-parsing";
 import { ollamaClient as ollama } from "./lib/ollama-client";
@@ -300,6 +301,40 @@ export async function deleteChat(chatId: string): Promise<void> {
   }
 }
 
+export async function updateChatMessage(
+  chatId: string,
+  index: number,
+  content: string,
+): Promise<{
+  index: number;
+  chatId: string;
+  message: Message;
+}> {
+  const response = await fetch(
+    `${API_BASE}/api/v1/chat/${chatId}/messages/${index}`,
+    {
+      method: "PATCH",
+      headers: {
+        "Content-Type": "application/json",
+      },
+      body: JSON.stringify({ content }),
+    },
+  );
+
+  if (!response.ok) {
+    const errorMessage = await response.text();
+    throw new Error(errorMessage || "Failed to update message");
+  }
+
+  const data = await response.json();
+
+  return {
+    index: data.index,
+    chatId: data.chatId,
+    message: new Message(data.message),
+  };
+}
+
 // Get upstream information for model staleness checking
 export async function getModelUpstreamInfo(
   model: Model,

app/ui/app/src/components/Chat.tsx

@@ -13,6 +13,7 @@ import {
   useChatError,
   useShouldShowStaleDisplay,
   useDismissStaleModel,
+  useUpdateChatMessage,
 } from "@/hooks/useChats";
 import { useHealth } from "@/hooks/useHealth";
 import { useMessageAutoscroll } from "@/hooks/useMessageAutoscroll";
@@ -47,6 +48,12 @@ export default function Chat({ chatId }: { chatId: string }) {
     index: number;
     originalMessage: Message;
   } | null>(null);
+  const [editingAssistantIndex, setEditingAssistantIndex] = useState<
+    number | null
+  >(null);
+  const [assistantEditError, setAssistantEditError] = useState<string | null>(
+    null,
+  );
   const prevChatIdRef = useRef<string>(chatId);
 
   const chatFormCallbackRef = useRef<
@@ -98,9 +105,14 @@ export default function Chat({ chatId }: { chatId: string }) {
   // Clear editing state when navigating to a different chat
   useEffect(() => {
     setEditingMessage(null);
+    setEditingAssistantIndex(null);
+    setAssistantEditError(null);
   }, [chatId]);
 
   const sendMessageMutation = useSendMessage(chatId);
+  const updateAssistantMessageMutation = useUpdateChatMessage(
+    chatId === "new" ? "" : chatId,
+  );
 
   const { containerRef, handleNewUserMessage, spacerHeight } =
     useMessageAutoscroll({
@@ -186,6 +198,44 @@ export default function Chat({ chatId }: { chatId: string }) {
     }
   };
 
+  const handleAssistantEditStart = (index: number) => {
+    setAssistantEditError(null);
+    setEditingAssistantIndex(index);
+  };
+
+  const handleAssistantEditCancel = () => {
+    if (updateAssistantMessageMutation.isPending) {
+      return;
+    }
+    setAssistantEditError(null);
+    setEditingAssistantIndex(null);
+  };
+
+  const handleAssistantEditSave = async (index: number, content: string) => {
+    if (updateAssistantMessageMutation.isPending) {
+      return;
+    }
+
+    const trimmedContent = content.trim();
+    if (!trimmedContent) {
+      setAssistantEditError("Response cannot be empty.");
+      return;
+    }
+
+    try {
+      setAssistantEditError(null);
+      await updateAssistantMessageMutation.mutateAsync({
+        index,
+        content: trimmedContent,
+      });
+      setEditingAssistantIndex(null);
+    } catch (error) {
+      setAssistantEditError(
+        error instanceof Error ? error.message : "Failed to update message.",
+      );
+    }
+  };
+
   const clearChatError = () => {
     queryClient.setQueryData(
       ["chatError", chatId === "new" ? "" : chatId],
@@ -236,6 +286,12 @@ export default function Chat({ chatId }: { chatId: string }) {
               editingMessageIndex={editingMessage?.index}
               error={chatError}
               browserToolResult={browserToolResult}
+              onAssistantEditStart={handleAssistantEditStart}
+              onAssistantEditSave={handleAssistantEditSave}
+              onAssistantEditCancel={handleAssistantEditCancel}
+              assistantEditingIndex={editingAssistantIndex}
+              assistantEditIsSaving={updateAssistantMessageMutation.isPending}
+              assistantEditError={assistantEditError}
             />
           </section>
 

app/ui/app/src/components/Message.tsx

@@ -3,8 +3,16 @@ import Thinking from "./Thinking";
 import StreamingMarkdownContent from "./StreamingMarkdownContent";
 import { ImageThumbnail } from "./ImageThumbnail";
 import { isImageFile } from "@/utils/imageUtils";
-import CopyButton from "./CopyButton";
-import React, { useState, useMemo, useRef } from "react";
+import React, { useState, useMemo, useRef, useEffect } from "react";
+import {
+  CheckIcon,
+  PencilSquareIcon,
+  Square2StackIcon,
+} from "@heroicons/react/24/outline";
+import {
+  MessageActions,
+  MessageAction,
+} from "@/components/ai-elements/message";
 
 const Message = React.memo(
   ({
@@ -15,6 +23,12 @@ const Message = React.memo(
     isFaded,
     browserToolResult,
     lastToolQuery,
+    onAssistantEditStart,
+    onAssistantEditSave,
+    onAssistantEditCancel,
+    assistantEditingIndex,
+    assistantEditIsSaving,
+    assistantEditError,
   }: {
     message: MessageType;
     onEditMessage?: (content: string, index: number) => void;
@@ -24,6 +38,15 @@ const Message = React.memo(
     // TODO(drifkin): this type isn't right
     browserToolResult?: BrowserToolResult;
     lastToolQuery?: string;
+    onAssistantEditStart?: (index: number) => void;
+    onAssistantEditSave?: (
+      index: number,
+      content: string,
+    ) => void | Promise<void>;
+    onAssistantEditCancel?: () => void;
+    assistantEditingIndex?: number | null;
+    assistantEditIsSaving?: boolean;
+    assistantEditError?: string | null;
   }) => {
     if (message.role === "user") {
       return (
@@ -42,6 +65,13 @@ const Message = React.memo(
           isFaded={isFaded}
           browserToolResult={browserToolResult}
           lastToolQuery={lastToolQuery}
+          messageIndex={messageIndex}
+          onAssistantEditStart={onAssistantEditStart}
+          onAssistantEditSave={onAssistantEditSave}
+          onAssistantEditCancel={onAssistantEditCancel}
+          assistantEditingIndex={assistantEditingIndex}
+          assistantEditIsSaving={assistantEditIsSaving}
+          assistantEditError={assistantEditError}
         />
       );
     }
@@ -53,7 +83,10 @@ const Message = React.memo(
       prevProps.messageIndex === nextProps.messageIndex &&
       prevProps.isStreaming === nextProps.isStreaming &&
       prevProps.isFaded === nextProps.isFaded &&
-      prevProps.browserToolResult === nextProps.browserToolResult
+      prevProps.browserToolResult === nextProps.browserToolResult &&
+      prevProps.assistantEditingIndex === nextProps.assistantEditingIndex &&
+      prevProps.assistantEditIsSaving === nextProps.assistantEditIsSaving &&
+      prevProps.assistantEditError === nextProps.assistantEditError
     );
   },
 );
@@ -880,6 +913,13 @@ function OtherRoleMessage({
   isFaded,
   browserToolResult,
   lastToolQuery,
+  messageIndex,
+  onAssistantEditStart,
+  onAssistantEditSave,
+  onAssistantEditCancel,
+  assistantEditingIndex,
+  assistantEditIsSaving,
+  assistantEditError,
 }: {
   message: MessageType;
   previousMessage?: MessageType;
@@ -888,8 +928,106 @@ function OtherRoleMessage({
   // TODO(drifkin): this type isn't right
   browserToolResult?: BrowserToolResult;
   lastToolQuery?: string;
+  messageIndex?: number;
+  onAssistantEditStart?: (index: number) => void;
+  onAssistantEditSave?: (
+    index: number,
+    content: string,
+  ) => void | Promise<void>;
+  onAssistantEditCancel?: () => void;
+  assistantEditingIndex?: number | null;
+  assistantEditIsSaving?: boolean;
+  assistantEditError?: string | null;
 }) {
   const messageRef = useRef<HTMLDivElement>(null);
+  const [draftContent, setDraftContent] = useState(message.content || "");
+  const [isCopied, setIsCopied] = useState(false);
+  const copyResetTimeoutRef = useRef<number | undefined>(undefined);
+
+  const isAssistantMessage = message.role === "assistant";
+  const isEditingAssistant =
+    isAssistantMessage &&
+    assistantEditingIndex !== null &&
+    assistantEditingIndex !== undefined &&
+    messageIndex !== undefined &&
+    assistantEditingIndex === messageIndex;
+
+  useEffect(() => {
+    if (isEditingAssistant) {
+      setDraftContent(message.content || "");
+    }
+  }, [isEditingAssistant, message.content]);
+
+  useEffect(() => {
+    setIsCopied(false);
+  }, [message.content]);
+
+  useEffect(() => {
+    return () => {
+      if (copyResetTimeoutRef.current !== undefined) {
+        window.clearTimeout(copyResetTimeoutRef.current);
+      }
+    };
+  }, []);
+
+  const handleAssistantEditStart = () => {
+    if (onAssistantEditStart && messageIndex !== undefined) {
+      onAssistantEditStart(messageIndex);
+    }
+  };
+
+  const handleAssistantEditSave = async () => {
+    if (!onAssistantEditSave || messageIndex === undefined) {
+      return;
+    }
+    await onAssistantEditSave(messageIndex, draftContent);
+  };
+
+  const handleAssistantEditCancel = () => {
+    onAssistantEditCancel?.();
+  };
+
+  const handleCopy = async () => {
+    const contentToCopy = message.content || "";
+    if (!contentToCopy) {
+      return;
+    }
+
+    const scheduleReset = () => {
+      if (copyResetTimeoutRef.current !== undefined) {
+        window.clearTimeout(copyResetTimeoutRef.current);
+      }
+      copyResetTimeoutRef.current = window.setTimeout(() => {
+        setIsCopied(false);
+      }, 2000);
+    };
+
+    try {
+      if (messageRef.current) {
+        const cloned = messageRef.current.cloneNode(true) as HTMLElement;
+        await navigator.clipboard.write([
+          new ClipboardItem({
+            "text/html": new Blob([cloned.innerHTML], { type: "text/html" }),
+            "text/plain": new Blob([contentToCopy], { type: "text/plain" }),
+          }),
+        ]);
+      } else {
+        await navigator.clipboard.writeText(contentToCopy);
+      }
+
+      setIsCopied(true);
+      scheduleReset();
+    } catch (error) {
+      console.error("Clipboard API failed, falling back to plain text", error);
+      try {
+        await navigator.clipboard.writeText(contentToCopy);
+        setIsCopied(true);
+        scheduleReset();
+      } catch (fallbackError) {
+        console.error("Fallback copy also failed:", fallbackError);
+      }
+    }
+  };
 
   return (
     <div
@@ -918,17 +1056,19 @@ function OtherRoleMessage({
 
           if (
             message.role !== "tool" &&
+            !isEditingAssistant &&
             (!message.content || !message.content.trim())
           ) {
             return null;
           }
 
-          // Render appropriate content
+          const refForContent = isEditingAssistant ? null : messageRef;
+
           return (
             <div
               className="max-w-full prose dark:prose-invert assistant-message-content break-words"
               id="message-container"
-              ref={messageRef}
+              ref={refForContent}
             >
               {message.role === "tool" ? (
                 <ToolRoleContent
@@ -936,6 +1076,42 @@ function OtherRoleMessage({
                   browserToolResult={browserToolResult}
                   lastToolQuery={lastToolQuery}
                 />
+              ) : isEditingAssistant ? (
+                <>
+                  <textarea
+                    value={draftContent}
+                    onChange={(event) => setDraftContent(event.target.value)}
+                    disabled={assistantEditIsSaving}
+                    autoFocus
+                    rows={Math.min(draftContent.split("\n").length, 20)}
+                    className="w-full max-h-[500px] overflow-y-auto resize-none rounded-2xl border border-neutral-200 bg-white p-4 text-sm leading-relaxed text-neutral-900 transition-colors dark:border-neutral-700 dark:bg-neutral-900 dark:text-neutral-100"
+                  />
+                  <div className="mt-3 flex flex-wrap items-center gap-2">
+                    <button
+                      type="button"
+                      onClick={handleAssistantEditSave}
+                      disabled={assistantEditIsSaving}
+                      tabIndex={0}
+                      className="rounded-2xl px-2 py-1 text-sm font-medium text-black transition-colors hover:text-neutral-700 focus-visible:ring-2 focus-visible:ring-neutral-300 focus-visible:ring-offset-2 disabled:cursor-not-allowed disabled:opacity-60 dark:text-neutral-200 dark:hover:text-neutral-400 dark:focus-visible:ring-neutral-500 cursor-pointer"
+                    >
+                      {assistantEditIsSaving ? "Saving…" : "Save"}
+                    </button>
+                    <button
+                      type="button"
+                      onClick={handleAssistantEditCancel}
+                      disabled={assistantEditIsSaving}
+                      tabIndex={0}
+                      className="rounded-2xl px-2 py-1 text-sm text-neutral-500 transition-colors hover:text-neutral-700 focus-visible:ring-2 focus-visible:ring-neutral-300 focus-visible:ring-offset-2 disabled:cursor-not-allowed disabled:opacity-60 dark:text-neutral-400 dark:hover:text-neutral-200 dark:focus-visible:ring-neutral-500 cursor-pointer"
+                    >
+                      Cancel
+                    </button>
+                  </div>
+                  {assistantEditError && (
+                    <p className="mt-2 text-sm text-red-500 dark:text-red-400">
+                      {assistantEditError}
+                    </p>
+                  )}
+                </>
               ) : (
                 <StreamingMarkdownContent
                   content={message.content}
@@ -972,18 +1148,28 @@ function OtherRoleMessage({
         message.content &&
         message.content.trim() &&
         (!message.tool_calls || message.tool_calls.length === 0) &&
-        !message.tool_call && (
-          <div className="-ml-1">
-            <CopyButton
-              content={message.content || ""}
-              copyRef={messageRef as React.RefObject<HTMLElement>}
-              removeClasses={["copy-button"]}
-              size="md"
-              showLabels={false}
-              className="copy-button z-10 text-neutral-500 dark:text-neutral-400"
-              title="Copy"
-            />
-          </div>
+        !message.tool_call &&
+        !isEditingAssistant && (
+          <MessageActions>
+            <MessageAction
+              onClick={handleCopy}
+              title={isCopied ? "Copied" : "Copy"}
+              aria-label={isCopied ? "Copied" : "Copy"}
+            >
+              {isCopied ? (
+                <CheckIcon className="h-4.5 w-4.5" />
+              ) : (
+                <Square2StackIcon className="h-4.5 w-4.5" />
+              )}
+            </MessageAction>
+            <MessageAction
+              onClick={handleAssistantEditStart}
+              title="Edit"
+              aria-label="Edit"
+            >
+              <PencilSquareIcon className="h-4.5 w-4.5" />
+            </MessageAction>
+          </MessageActions>
         )}
     </div>
   );

app/ui/app/src/components/MessageList.tsx

@@ -14,6 +14,12 @@ export default function MessageList({
   editingMessageIndex,
   error,
   browserToolResult,
+  onAssistantEditStart,
+  onAssistantEditSave,
+  onAssistantEditCancel,
+  assistantEditingIndex,
+  assistantEditIsSaving,
+  assistantEditError,
 }: {
   messages: MessageType[];
   spacerHeight: number;
@@ -24,6 +30,15 @@ export default function MessageList({
   editingMessageIndex?: number;
   error?: ErrorEvent | null;
   browserToolResult?: any;
+  onAssistantEditStart?: (index: number) => void;
+  onAssistantEditSave?: (
+    index: number,
+    content: string,
+  ) => void | Promise<void>;
+  onAssistantEditCancel?: () => void;
+  assistantEditingIndex?: number | null;
+  assistantEditIsSaving?: boolean;
+  assistantEditError?: string | null;
 }) {
   const [showDots, setShowDots] = React.useState(false);
   const isDownloadingModel = downloadProgress && !downloadProgress.done;
@@ -101,6 +116,12 @@ export default function MessageList({
               }
               browserToolResult={browserToolResult}
               lastToolQuery={lastToolQuery}
+              onAssistantEditStart={onAssistantEditStart}
+              onAssistantEditSave={onAssistantEditSave}
+              onAssistantEditCancel={onAssistantEditCancel}
+              assistantEditingIndex={assistantEditingIndex}
+              assistantEditIsSaving={assistantEditIsSaving}
+              assistantEditError={assistantEditError}
             />
           </div>
         );

app/ui/app/src/components/ai-elements/message.tsx

@@ -0,0 +1,227 @@
+import { cn } from "@/lib/utils";
+import type { HTMLAttributes, ReactElement } from "react";
+import { createContext, memo, useContext, useEffect, useState } from "react";
+import { Streamdown } from "streamdown";
+
+export type MessageProps = HTMLAttributes<HTMLDivElement> & {
+  from: "user" | "assistant" | "system";
+};
+
+export const Message = ({ className, from, ...props }: MessageProps) => (
+  <div
+    className={cn(
+      "group flex w-full max-w-[80%] flex-col gap-2",
+      from === "user" ? "is-user ml-auto justify-end" : "is-assistant",
+      className,
+    )}
+    {...props}
+  />
+);
+
+export type MessageContentProps = HTMLAttributes<HTMLDivElement>;
+
+export const MessageContent = ({
+  children,
+  className,
+  ...props
+}: MessageContentProps) => (
+  <div
+    className={cn(
+      "is-user:dark flex w-fit flex-col gap-2 overflow-hidden text-sm",
+      "group-[.is-user]:ml-auto group-[.is-user]:rounded-lg group-[.is-user]:bg-neutral-100 dark:group-[.is-user]:bg-neutral-800 group-[.is-user]:px-4 group-[.is-user]:py-3 group-[.is-user]:text-neutral-900 dark:group-[.is-user]:text-neutral-100",
+      "group-[.is-assistant]:text-neutral-900 dark:group-[.is-assistant]:text-neutral-100",
+      className,
+    )}
+    {...props}
+  >
+    {children}
+  </div>
+);
+
+export type MessageActionsProps = HTMLAttributes<HTMLDivElement>;
+
+export const MessageActions = ({
+  className,
+  children,
+  ...props
+}: MessageActionsProps) => (
+  <div className={cn("flex items-center gap-1", className)} {...props}>
+    {children}
+  </div>
+);
+
+export type MessageActionProps =
+  React.ButtonHTMLAttributes<HTMLButtonElement> & {
+    tooltip?: string;
+    label?: string;
+    active?: boolean;
+  };
+
+export const MessageAction = ({
+  tooltip,
+  children,
+  label,
+  active = false,
+  className,
+  ...props
+}: MessageActionProps) => {
+  return (
+    <button
+      type="button"
+      title={tooltip || label}
+      aria-label={label || tooltip}
+      className={cn(
+        "flex h-7 w-7 items-center justify-center rounded-md transition-colors focus-visible:ring-2 focus-visible:ring-neutral-300 focus-visible:ring-offset-2 focus-visible:ring-offset-white dark:focus-visible:ring-neutral-500 dark:focus-visible:ring-offset-neutral-900",
+        active
+          ? "bg-neutral-900 text-white dark:bg-white dark:text-neutral-900"
+          : "text-neutral-600 hover:bg-neutral-100 hover:text-neutral-900 dark:text-neutral-600 dark:hover:bg-neutral-800 dark:hover:text-neutral-100",
+        className,
+      )}
+      {...props}
+    >
+      {children}
+      {(label || tooltip) && (
+        <span className="sr-only">{label || tooltip}</span>
+      )}
+    </button>
+  );
+};
+
+type MessageBranchContextType = {
+  currentBranch: number;
+  totalBranches: number;
+  goToPrevious: () => void;
+  goToNext: () => void;
+  branches: ReactElement[];
+  setBranches: (branches: ReactElement[]) => void;
+};
+
+const MessageBranchContext = createContext<MessageBranchContextType | null>(
+  null,
+);
+
+const useMessageBranch = () => {
+  const context = useContext(MessageBranchContext);
+
+  if (!context) {
+    throw new Error(
+      "MessageBranch components must be used within MessageBranch",
+    );
+  }
+
+  return context;
+};
+
+export type MessageBranchProps = HTMLAttributes<HTMLDivElement> & {
+  defaultBranch?: number;
+  onBranchChange?: (branchIndex: number) => void;
+};
+
+export const MessageBranch = ({
+  defaultBranch = 0,
+  onBranchChange,
+  className,
+  ...props
+}: MessageBranchProps) => {
+  const [currentBranch, setCurrentBranch] = useState(defaultBranch);
+  const [branches, setBranches] = useState<ReactElement[]>([]);
+
+  const handleBranchChange = (newBranch: number) => {
+    setCurrentBranch(newBranch);
+    onBranchChange?.(newBranch);
+  };
+
+  const goToPrevious = () => {
+    const newBranch =
+      currentBranch > 0 ? currentBranch - 1 : branches.length - 1;
+    handleBranchChange(newBranch);
+  };
+
+  const goToNext = () => {
+    const newBranch =
+      currentBranch < branches.length - 1 ? currentBranch + 1 : 0;
+    handleBranchChange(newBranch);
+  };
+
+  const contextValue: MessageBranchContextType = {
+    currentBranch,
+    totalBranches: branches.length,
+    goToPrevious,
+    goToNext,
+    branches,
+    setBranches,
+  };
+
+  return (
+    <MessageBranchContext.Provider value={contextValue}>
+      <div
+        className={cn("grid w-full gap-2 [&>div]:pb-0", className)}
+        {...props}
+      />
+    </MessageBranchContext.Provider>
+  );
+};
+
+export type MessageBranchContentProps = HTMLAttributes<HTMLDivElement>;
+
+export const MessageBranchContent = ({
+  children,
+  ...props
+}: MessageBranchContentProps) => {
+  const { currentBranch, setBranches, branches } = useMessageBranch();
+  const childrenArray = Array.isArray(children) ? children : [children];
+
+  useEffect(() => {
+    if (branches.length !== childrenArray.length) {
+      setBranches(childrenArray);
+    }
+  }, [childrenArray, branches, setBranches]);
+
+  return childrenArray.map((branch, index) => (
+    <div
+      className={cn(
+        "grid gap-2 overflow-hidden [&>div]:pb-0",
+        index === currentBranch ? "block" : "hidden",
+      )}
+      key={branch.key}
+      {...props}
+    >
+      {branch}
+    </div>
+  ));
+};
+
+export type MessageResponseProps = React.ComponentProps<typeof Streamdown>;
+
+export const MessageResponse = memo(
+  ({ className, ...props }: MessageResponseProps) => (
+    <Streamdown
+      className={cn(
+        "size-full [&>*:first-child]:mt-0 [&>*:last-child]:mb-0",
+        className,
+      )}
+      {...props}
+    />
+  ),
+  (prevProps, nextProps) => prevProps.children === nextProps.children,
+);
+
+MessageResponse.displayName = "MessageResponse";
+
+export type MessageToolbarProps = HTMLAttributes<HTMLDivElement>;
+
+export const MessageToolbar = ({
+  className,
+  children,
+  ...props
+}: MessageToolbarProps) => (
+  <div
+    className={cn(
+      "mt-4 flex w-full items-center justify-between gap-4",
+      className,
+    )}
+    {...props}
+  >
+    {children}
+  </div>
+);

app/ui/app/src/hooks/useChats.ts

@@ -1,5 +1,11 @@
 import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query";
-import { getChats, getChat, sendMessage, type ChatEventUnion } from "../api";
+import {
+  getChats,
+  getChat,
+  sendMessage,
+  type ChatEventUnion,
+  updateChatMessage,
+} from "../api";
 import { Chat, ErrorEvent, Model } from "@/gotypes";
 import { Message } from "@/gotypes";
 import { useSelectedModel } from "./useSelectedModel";
@@ -713,6 +719,42 @@ export const useSendMessage = (chatId: string) => {
   });
 };
 
+export const useUpdateChatMessage = (chatId: string) => {
+  const queryClient = useQueryClient();
+
+  return useMutation({
+    mutationKey: ["updateChatMessage", chatId],
+    mutationFn: ({ index, content }: { index: number; content: string }) =>
+      updateChatMessage(chatId, index, content),
+    onSuccess: ({ index, message }) => {
+      queryClient.setQueryData(
+        ["chat", chatId],
+        (old: { chat: Chat } | undefined) => {
+          if (!old?.chat?.messages) {
+            return old;
+          }
+
+          const updatedMessages = [...old.chat.messages];
+          if (index < 0 || index >= updatedMessages.length) {
+            return old;
+          }
+
+          updatedMessages[index] = message;
+
+          return {
+            ...old,
+            chat: new Chat({
+              ...old.chat,
+              messages: updatedMessages,
+            }),
+          };
+        },
+      );
+      queryClient.invalidateQueries({ queryKey: ["chats"] });
+    },
+  });
+};
+
 export const useCancelMessage = () => {
   const {
     abortControllers,

app/ui/app/src/lib/utils.ts

@@ -0,0 +1,6 @@
+import { clsx, type ClassValue } from "clsx";
+import { twMerge } from "tailwind-merge";
+
+export function cn(...inputs: ClassValue[]) {
+  return twMerge(clsx(inputs));
+}

app/ui/responses/types.go

@@ -25,6 +25,16 @@ type ChatResponse struct {
 	Chat store.Chat `json:"chat"`
 }
 
+type MessageUpdateRequest struct {
+	Content string `json:"content"`
+}
+
+type MessageUpdateResponse struct {
+	Index   int           `json:"index"`
+	ChatID  string        `json:"chatId"`
+	Message store.Message `json:"message"`
+}
+
 type Model struct {
 	Model      string     `json:"model"`
 	Digest     string     `json:"digest,omitempty"`

app/ui/ui.go

@@ -162,7 +162,7 @@ func (s *Server) Handler() http.Handler {
 			// Add CORS headers for dev work
 			if CORS() {
 				w.Header().Set("Access-Control-Allow-Origin", "*")
-				w.Header().Set("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS")
+				w.Header().Set("Access-Control-Allow-Methods", "GET, POST, PUT, PATCH, DELETE, OPTIONS")
 				w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization, X-Requested-With")
 				w.Header().Set("Access-Control-Allow-Credentials", "true")
 
@@ -246,6 +246,9 @@ func (s *Server) Handler() http.Handler {
 	mux.Handle("OPTIONS /", handle(func(w http.ResponseWriter, r *http.Request) error {
 		return nil
 	}))
+	mux.Handle("OPTIONS /api/v1/chat/{id}/messages/{index}", handle(func(w http.ResponseWriter, r *http.Request) error {
+		return nil
+	}))
 
 	// API routes - handle first to take precedence
 	mux.Handle("GET /api/v1/chats", handle(s.listChats))
@@ -254,6 +257,7 @@ func (s *Server) Handler() http.Handler {
 	mux.Handle("DELETE /api/v1/chat/{id}", handle(s.deleteChat))
 	mux.Handle("POST /api/v1/create-chat", handle(s.createChat))
 	mux.Handle("PUT /api/v1/chat/{id}/rename", handle(s.renameChat))
+	mux.Handle("PATCH /api/v1/chat/{id}/messages/{index}", handle(s.updateChatMessage))
 
 	mux.Handle("GET /api/v1/inference-compute", handle(s.getInferenceCompute))
 	mux.Handle("POST /api/v1/model/upstream", handle(s.modelUpstream))
@@ -1280,6 +1284,57 @@ func (s *Server) renameChat(w http.ResponseWriter, r *http.Request) error {
 	return nil
 }
 
+func (s *Server) updateChatMessage(w http.ResponseWriter, r *http.Request) error {
+	chatID := r.PathValue("id")
+	if chatID == "" {
+		return fmt.Errorf("chat ID is required")
+	}
+
+	indexParam := r.PathValue("index")
+	msgIndex, err := strconv.Atoi(indexParam)
+	if err != nil {
+		return fmt.Errorf("invalid message index: %w", err)
+	}
+
+	var req responses.MessageUpdateRequest
+	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+		return fmt.Errorf("invalid request body: %w", err)
+	}
+
+	if strings.TrimSpace(req.Content) == "" {
+		return fmt.Errorf("message content cannot be empty")
+	}
+
+	chat, err := s.Store.ChatWithOptions(chatID, true)
+	if err != nil {
+		return fmt.Errorf("failed to load chat: %w", err)
+	}
+
+	if msgIndex < 0 || msgIndex >= len(chat.Messages) {
+		return fmt.Errorf("message index out of range")
+	}
+
+	message := chat.Messages[msgIndex]
+	if message.Role != "assistant" {
+		return fmt.Errorf("only assistant messages can be updated")
+	}
+
+	message.Content = req.Content
+	message.UpdatedAt = time.Now()
+	chat.Messages[msgIndex] = message
+
+	if err := s.Store.SetChat(*chat); err != nil {
+		return fmt.Errorf("failed to update chat message: %w", err)
+	}
+
+	w.Header().Set("Content-Type", "application/json")
+	return json.NewEncoder(w).Encode(responses.MessageUpdateResponse{
+		Index:   msgIndex,
+		ChatID:  chatID,
+		Message: message,
+	})
+}
+
 func (s *Server) deleteChat(w http.ResponseWriter, r *http.Request) error {
 	cid := r.PathValue("id")
 	if cid == "" {