feat(chat): implement branching functionality and context management

apps/mesh/src/web/components/chat/chat-input-context.tsx

@@ -0,0 +1,136 @@
+import { CornerUpLeft, X } from "@untitledui/icons";
+import { createContext, use, type PropsWithChildren } from "react";
+import {
+  useInputValue,
+  type InputController,
+  type BranchContext,
+} from "../../hooks/use-persisted-chat";
+
+interface ChatInputContextValue {
+  inputController: InputController;
+  branchContext: BranchContext | null;
+  clearBranchContext: () => void;
+  onGoToOriginalMessage: () => void;
+}
+
+const ChatInputContext = createContext<ChatInputContextValue | null>(null);
+
+function useChatInputContext() {
+  const ctx = use(ChatInputContext);
+  if (!ctx) {
+    throw new Error(
+      "useChatInputContext must be used within ChatInputProvider",
+    );
+  }
+  return ctx;
+}
+
+export function ChatInputProvider({
+  children,
+  inputController,
+  branchContext,
+  clearBranchContext,
+  onGoToOriginalMessage,
+}: PropsWithChildren<ChatInputContextValue>) {
+  return (
+    <ChatInputContext
+      value={{
+        inputController,
+        branchContext,
+        clearBranchContext,
+        onGoToOriginalMessage,
+      }}
+    >
+      {children}
+    </ChatInputContext>
+  );
+}
+
+/**
+ * Branch preview banner - shows when editing a message from a branch.
+ * Uses context internally, no props needed.
+ */
+export function BranchPreview() {
+  const {
+    branchContext,
+    clearBranchContext,
+    onGoToOriginalMessage,
+    inputController,
+  } = useChatInputContext();
+
+  if (!branchContext) return null;
+
+  return (
+    <button
+      type="button"
+      onClick={onGoToOriginalMessage}
+      className="flex items-start gap-2 px-2 py-2 rounded-lg border border-dashed border-muted-foreground/30 bg-muted/50 text-sm hover:bg-muted transition-colors cursor-pointer text-left w-full"
+      title="Click to view original message"
+    >
+      <CornerUpLeft
+        size={14}
+        className="text-muted-foreground mt-0.5 shrink-0"
+      />
+      <div className="flex-1 min-w-0">
+        <div className="text-xs text-muted-foreground mb-1">
+          Editing message (click to view original):
+        </div>
+        <div className="text-muted-foreground/70 line-clamp-2">
+          {branchContext.originalMessageText}
+        </div>
+      </div>
+      <span
+        role="button"
+        tabIndex={0}
+        onClick={(e) => {
+          e.stopPropagation();
+          clearBranchContext();
+          inputController.setValue("");
+        }}
+        onKeyDown={(e) => {
+          if (e.key === "Enter" || e.key === " ") {
+            e.preventDefault();
+            e.stopPropagation();
+            clearBranchContext();
+            inputController.setValue("");
+          }
+        }}
+        className="text-muted-foreground hover:text-foreground transition-colors shrink-0"
+        title="Cancel editing"
+      >
+        <X size={14} />
+      </span>
+    </button>
+  );
+}
+
+/**
+ * Hook to get controlled input value and handlers from context.
+ * Only the component using this hook re-renders on keystroke.
+ */
+export function useChatInputState() {
+  const { inputController, branchContext, clearBranchContext } =
+    useChatInputContext();
+  const [inputValue, setInputValue] = useInputValue(inputController);
+
+  const handleInputChange = (value: string) => {
+    setInputValue(value);
+  };
+
+  const handleSubmit = async (
+    text: string,
+    onSubmit: (text: string) => Promise<void>,
+  ) => {
+    setInputValue("");
+    await onSubmit(text);
+  };
+
+  return {
+    inputValue,
+    setInputValue,
+    handleInputChange,
+    handleSubmit,
+    branchContext,
+    clearBranchContext,
+  };
+}

apps/mesh/src/web/components/chat/chat.tsx

@@ -137,10 +137,12 @@ function ChatMessages({
   messages,
   status,
   minHeightOffset = 240,
+  onBranchFromMessage,
 }: {
   messages: ChatMessage[];
   status?: ChatStatus;
   minHeightOffset?: number;
+  onBranchFromMessage?: (messageId: string, messageText: string) => void;
 }) {
   const sentinelRef = useRef<HTMLDivElement>(null);
   useChatAutoScroll({ messageCount: messages.length, sentinelRef });
@@ -151,6 +153,7 @@ function ChatMessages({
         message.role === "user" ? (
           <MessageUser
             key={message.id}
+            onBranchFromMessage={onBranchFromMessage}
             message={message as UIMessage<Metadata>}
           />
         ) : message.role === "assistant" ? (
@@ -202,15 +205,23 @@ function ChatInput({
   placeholder,
   usageMessages,
   children,
+  value,
+  onValueChange,
 }: PropsWithChildren<{
   onSubmit: (text: string) => Promise<void>;
   onStop: () => void;
   disabled: boolean;
   isStreaming: boolean;
   placeholder: string;
   usageMessages?: ChatMessage[];
+  value?: string;
+  onValueChange?: (value: string) => void;
 }>) {
-  const [input, setInput] = useState("");
+  const [internalInput, setInternalInput] = useState("");
+
+  // Use controlled value if provided, otherwise use internal state
+  const input = value !== undefined ? value : internalInput;
+  const setInput = onValueChange ?? setInternalInput;
 
   const modelSelector = findChild(children, ChatInputModelSelector);
   const gatewaySelector = findChild(children, ChatInputGatewaySelector);

apps/mesh/src/web/components/chat/message-user.tsx

@@ -1,7 +1,22 @@
+import {
+  AlertDialog,
+  AlertDialogAction,
+  AlertDialogCancel,
+  AlertDialogContent,
+  AlertDialogDescription,
+  AlertDialogFooter,
+  AlertDialogHeader,
+  AlertDialogTitle,
+} from "@deco/ui/components/alert-dialog.tsx";
+import {
+  Tooltip,
+  TooltipContent,
+  TooltipTrigger,
+} from "@deco/ui/components/tooltip.tsx";
 import { Button } from "@deco/ui/components/button.tsx";
 import { cn } from "@deco/ui/lib/utils.ts";
 import { Metadata } from "@deco/ui/types/chat-metadata.ts";
-import { ChevronDown, ChevronUp } from "@untitledui/icons";
+import { ChevronDown, ChevronUp, ReverseLeft } from "@untitledui/icons";
 import { type UIMessage } from "ai";
 import { useContext, useRef, useState } from "react";
 import { MessageListContext } from "./message-list.tsx";
@@ -12,17 +27,20 @@ export interface MessageProps<T extends Metadata> {
   status?: "streaming" | "submitted" | "ready" | "error";
   className?: string;
   pairIndex?: number;
+  onBranchFromMessage?: (messageId: string, messageText: string) => void;
 }
 
 export function MessageUser<T extends Metadata>({
   message,
   className,
   pairIndex,
+  onBranchFromMessage,
 }: MessageProps<T>) {
   const { id, parts } = message;
   const messageRef = useRef<HTMLDivElement>(null);
   const messageListContext = useContext(MessageListContext);
   const [isExpanded, setIsExpanded] = useState(false);
+  const [showBranchDialog, setShowBranchDialog] = useState(false);
 
   // Early return if no parts
   if (!parts || parts.length === 0) {
@@ -38,64 +56,119 @@ export function MessageUser<T extends Metadata>({
 
   const isLongMessage = totalTextLength > 60;
 
+  // Extract the full text from all text parts
+  const messageText = parts
+    .filter((part) => part.type === "text")
+    .map((part) => (part as { type: "text"; text: string }).text)
+    .join("\n");
+
   const handleClick = () => {
     if (pairIndex !== undefined) {
       messageListContext?.scrollToPair(pairIndex);
     }
   };
 
+  const handleBranchClick = (e: React.MouseEvent) => {
+    e.stopPropagation();
+    setShowBranchDialog(true);
+  };
+
+  const handleConfirmBranch = () => {
+    setShowBranchDialog(false);
+    onBranchFromMessage?.(id, messageText);
+  };
+
+  const canBranch = Boolean(onBranchFromMessage);
+
   return (
-    <div
-      ref={messageRef}
-      className={cn(
-        "message-block w-full min-w-0 group relative flex items-start gap-4 px-2 text-foreground flex-row-reverse",
-        className,
-      )}
-    >
-      {" "}
+    <>
       <div
-        onClick={handleClick}
-        className="w-full border min-w-0 shadow-[0_3px_6px_-1px_rgba(0,0,0,0.1)] rounded-lg text-[0.9375rem] wrap-break-word overflow-wrap-anywhere bg-muted px-4 py-2 cursor-pointer transition-colors"
+        ref={messageRef}
+        className={cn(
+          "message-block w-full min-w-0 group relative flex items-start gap-4 px-2 text-foreground flex-row-reverse",
+          className,
+        )}
       >
+        {" "}
         <div
-          className={cn(
-            isLongMessage &&
-              !isExpanded &&
-              "overflow-hidden relative max-h-[60px]",
-          )}
+          onClick={handleClick}
+          className="w-full border min-w-0 shadow-[0_3px_6px_-1px_rgba(0,0,0,0.1)] rounded-lg text-[0.9375rem] wrap-break-word overflow-wrap-anywhere bg-muted px-4 py-2 cursor-pointer transition-colors"
         >
-          {parts.map((part, index) => {
-            if (part.type === "text") {
-              return (
-                <MessageTextPart key={`${id}-${index}`} id={id} part={part} />
-              );
-            }
-            return null;
-          })}
-          {isLongMessage && !isExpanded && (
-            <div className="absolute bottom-0 left-0 right-0 h-12 bg-linear-to-t from-muted to-transparent pointer-events-none" />
+          <div
+            className={cn(
+              isLongMessage &&
+                !isExpanded &&
+                "overflow-hidden relative max-h-[60px]",
+            )}
+          >
+            {parts.map((part, index) => {
+              if (part.type === "text") {
+                return (
+                  <MessageTextPart key={`${id}-${index}`} id={id} part={part} />
+                );
+              }
+              return null;
+            })}
+            {isLongMessage && !isExpanded && (
+              <div className="absolute bottom-0 left-0 right-0 h-12 bg-linear-to-t from-muted to-transparent pointer-events-none" />
+            )}
+          </div>
+          {isLongMessage && (
+            <div className="flex justify-center">
+              <Button
+                onClick={(e) => {
+                  e.stopPropagation();
+                  setIsExpanded(!isExpanded);
+                }}
+                variant="ghost"
+                size="xs"
+                className="text-xs w-full text-muted-foreground hover:text-foreground"
+              >
+                {isExpanded ? (
+                  <ChevronUp className="text-sm" />
+                ) : (
+                  <ChevronDown className="text-sm" />
+                )}
+              </Button>
+            </div>
+          )}
+          {canBranch && (
+            <div className="flex justify-end">
+              <Tooltip>
+                <TooltipTrigger asChild>
+                  <Button
+                    onClick={handleBranchClick}
+                    variant="ghost"
+                    size="xs"
+                    className="opacity-0 group-hover:opacity-100 hover:bg-gray-200/70 rounded-md transition-opacity text-muted-foreground hover:text-foreground"
+                  >
+                    <ReverseLeft size={16} className="p-0.5" />
+                  </Button>
+                </TooltipTrigger>
+                <TooltipContent side="bottom">Edit from here</TooltipContent>
+              </Tooltip>
+            </div>
           )}
         </div>
-        {isLongMessage && (
-          <div className="flex justify-center">
-            <Button
-              onClick={(e) => {
-                e.stopPropagation();
-                setIsExpanded(!isExpanded);
-              }}
-              variant="ghost"
-              size="xs"
-              className="text-xs w-full text-muted-foreground hover:text-foreground"
-            >
-              {isExpanded ? (
-                <ChevronUp className="text-sm" />
-              ) : (
-                <ChevronDown className="text-sm" />
-              )}
-            </Button>
-          </div>
-        )}
       </div>
-    </div>
+
+      <AlertDialog open={showBranchDialog} onOpenChange={setShowBranchDialog}>
+        <AlertDialogContent>
+          <AlertDialogHeader>
+            <AlertDialogTitle>Edit from here?</AlertDialogTitle>
+            <AlertDialogDescription>
+              This will create a new conversation branch from this point. The
+              original conversation will remain unchanged.
+            </AlertDialogDescription>
+          </AlertDialogHeader>
+          <AlertDialogFooter>
+            <AlertDialogCancel>Cancel</AlertDialogCancel>
+            <AlertDialogAction onClick={handleConfirmBranch}>
+              Confirm
+            </AlertDialogAction>
+          </AlertDialogFooter>
+        </AlertDialogContent>
+      </AlertDialog>
+    </>
   );
 }