diff --git a/apps/mesh/src/web/components/chat/chat-input-context.tsx b/apps/mesh/src/web/components/chat/chat-input-context.tsx new file mode 100644 index 0000000000..88315cee2a --- /dev/null +++ b/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(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) { + return ( + + {children} + + ); +} + +/** + * 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 ( + + ); +} + +/** + * 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, + ) => { + setInputValue(""); + await onSubmit(text); + }; + + return { + inputValue, + setInputValue, + handleInputChange, + handleSubmit, + branchContext, + clearBranchContext, + }; +} diff --git a/apps/mesh/src/web/components/chat/chat.tsx b/apps/mesh/src/web/components/chat/chat.tsx index a1a95d7015..c8104c07d2 100644 --- a/apps/mesh/src/web/components/chat/chat.tsx +++ b/apps/mesh/src/web/components/chat/chat.tsx @@ -10,7 +10,7 @@ import type { ReactNode, RefObject, } from "react"; -import { Children, isValidElement, useRef, useState } from "react"; +import { Children, isValidElement, useRef } from "react"; import { toast } from "sonner"; import { GatewaySelector } from "./gateway-selector"; import { MessageAssistant } from "./message-assistant.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(null); useChatAutoScroll({ messageCount: messages.length, sentinelRef }); @@ -151,6 +153,7 @@ function ChatMessages({ message.role === "user" ? ( } /> ) : message.role === "assistant" ? ( @@ -202,6 +205,8 @@ function ChatInput({ placeholder, usageMessages, children, + value, + onValueChange, }: PropsWithChildren<{ onSubmit: (text: string) => Promise; onStop: () => void; @@ -209,9 +214,9 @@ function ChatInput({ isStreaming: boolean; placeholder: string; usageMessages?: ChatMessage[]; + value?: string; + onValueChange?: (value: string) => void; }>) { - const [input, setInput] = useState(""); - const modelSelector = findChild(children, ChatInputModelSelector); const gatewaySelector = findChild(children, ChatInputGatewaySelector); const rest = filterChildren(children, [ @@ -221,13 +226,13 @@ function ChatInput({ const handleSubmit = async (e?: React.FormEvent) => { e?.preventDefault(); - if (!input?.trim() || isStreaming) { + if (!value?.trim() || isStreaming) { return; } - const text = input.trim(); + const text = value.trim(); try { await onSubmit(text); - setInput(""); + onValueChange?.(""); } catch (error) { console.error("Failed to send message:", error); const message = @@ -279,8 +284,8 @@ function ChatInput({ return ( {})} onSubmit={handleSubmit} onStop={onStop} disabled={disabled} diff --git a/apps/mesh/src/web/components/chat/message-user.tsx b/apps/mesh/src/web/components/chat/message-user.tsx index 64f18cdfa8..37ebb54180 100644 --- a/apps/mesh/src/web/components/chat/message-user.tsx +++ b/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 { status?: "streaming" | "submitted" | "ready" | "error"; className?: string; pairIndex?: number; + onBranchFromMessage?: (messageId: string, messageText: string) => void; } export function MessageUser({ message, className, pairIndex, + onBranchFromMessage, }: MessageProps) { const { id, parts } = message; const messageRef = useRef(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({ 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 ( -
- {" "} + <>
+ {" "}
- {parts.map((part, index) => { - if (part.type === "text") { - return ( - - ); - } - return null; - })} - {isLongMessage && !isExpanded && ( -
+
+ {parts.map((part, index) => { + if (part.type === "text") { + return ( + + ); + } + return null; + })} + {isLongMessage && !isExpanded && ( +
+ )} +
+ {isLongMessage && ( +
+ +
+ )} + {canBranch && ( +
+ + + + + Edit from here + +
)}
- {isLongMessage && ( -
- -
- )}
-
+ + + + + Edit from here? + + This will create a new conversation branch from this point. The + original conversation will remain unchanged. + + + + Cancel + + Confirm + + + + + ); } diff --git a/apps/mesh/src/web/components/chat/side-panel-chat.tsx b/apps/mesh/src/web/components/chat/side-panel-chat.tsx index 02b0b91bab..06b9fe5da3 100644 --- a/apps/mesh/src/web/components/chat/side-panel-chat.tsx +++ b/apps/mesh/src/web/components/chat/side-panel-chat.tsx @@ -16,6 +16,11 @@ import { useConnections, } from "../../hooks/collections/use-connection"; import { useBindingConnections } from "../../hooks/use-binding"; +import { + ChatInputProvider, + BranchPreview, + useChatInputState, +} from "./chat-input-context"; import { useThreads } from "../../hooks/use-chat-store"; import { useGatewayPrompts, @@ -52,6 +57,57 @@ function OpenRouterIllustration() { ); } +function ChatInputField({ + onSubmit, + onStop, + isStreaming, + disabled, + placeholder, + usageMessages, + selectedGatewayId, + onGatewayChange, + selectedModel, + onModelChange, +}: { + onSubmit: (text: string) => Promise; + onStop: () => void; + isStreaming: boolean; + disabled: boolean; + placeholder: string; + usageMessages: ReturnType["messages"]; + selectedGatewayId?: string; + onGatewayChange: (gatewayId: string) => void; + selectedModel?: { id: string; connectionId: string }; + onModelChange: (model: ModelChangePayload) => void; +}) { + const { inputValue, handleInputChange, handleSubmit, branchContext } = + useChatInputState(); + + return ( + handleSubmit(text, onSubmit)} + onStop={onStop} + disabled={disabled} + isStreaming={isStreaming} + placeholder={branchContext ? "Edit your message..." : placeholder} + usageMessages={usageMessages} + value={inputValue} + onValueChange={handleInputChange} + > + id && onGatewayChange(id)} + /> + + + ); +} + /** * Ice breakers component that uses suspense to fetch gateway prompts */ @@ -143,8 +199,26 @@ export function ChatPanel() { onToolCall, onCreateThread: (thread) => createThread({ id: thread.id, title: thread.title }), + onThreadChange: setActiveThreadId, }); + // Destructure branching-related values from the hook + const { + inputController, + branchContext, + clearBranchContext, + branchFromMessage, + } = chat; + + // Handle clicking on the branch preview to go back to original thread + const handleGoToOriginalMessage = () => { + if (!branchContext) return; + setActiveThreadId(branchContext.originalThreadId); + // Clear the branch context since we're going back + clearBranchContext(); + inputController.setValue(""); + }; + const handleSendMessage = async (text: string) => { if (!selectedModel) { toast.error("No model configured"); @@ -167,6 +241,9 @@ export function ChatPanel() { }; await chat.sendMessage(text, metadata); + + // Clear editing state after successful send + clearBranchContext(); }; const handleModelChange = (m: ModelChangePayload) => { @@ -434,41 +511,42 @@ export function ChatPanel() { messages={chat.messages} status={chat.status} minHeightOffset={240} + onBranchFromMessage={branchFromMessage} /> )} -
- - +
+ + { - if (!gatewayId) return; - setSelectedGatewayState({ gatewayId }); - }} - /> - + setSelectedGatewayState({ gatewayId }) + } selectedModel={effectiveSelectedModelState ?? undefined} onModelChange={handleModelChange} /> - -
+
+ ); diff --git a/apps/mesh/src/web/components/details/assistant/index.tsx b/apps/mesh/src/web/components/details/assistant/index.tsx index 605622ebff..993159ace9 100644 --- a/apps/mesh/src/web/components/details/assistant/index.tsx +++ b/apps/mesh/src/web/components/details/assistant/index.tsx @@ -15,6 +15,11 @@ import { type GatewayPrompt, } from "@/web/hooks/use-gateway-prompts"; import { useLocalStorage } from "@/web/hooks/use-local-storage"; +import { + ChatInputProvider, + BranchPreview, + useChatInputState, +} from "@/web/components/chat/chat-input-context"; import { usePersistedChat } from "@/web/hooks/use-persisted-chat"; import { LOCALSTORAGE_KEYS } from "@/web/lib/localstorage-keys"; import { useProjectContext } from "@/web/providers/project-context-provider"; @@ -206,8 +211,41 @@ function AssistantEditForm({ form }: { form: AssistantForm }) { ); } +function AssistantInputField({ + onSubmit, + onStop, + isStreaming, + disabled, + usageMessages, +}: { + onSubmit: (text: string) => Promise; + onStop: () => void; + isStreaming: boolean; + disabled: boolean; + usageMessages: ReturnType["messages"]; +}) { + const { inputValue, handleInputChange, handleSubmit, branchContext } = + useChatInputState(); + + return ( + handleSubmit(text, onSubmit)} + onStop={onStop} + disabled={disabled} + isStreaming={isStreaming} + placeholder={ + branchContext ? "Edit your message..." : "Ask anything or @ for context" + } + usageMessages={usageMessages} + value={inputValue} + onValueChange={handleInputChange} + /> + ); +} + interface AssistantChatPanelProps { activeThreadId: string; + setActiveThreadId: (id: string) => void; mode: "chat" | "edit"; assistant: Assistant; form: AssistantForm; @@ -215,6 +253,7 @@ interface AssistantChatPanelProps { function AssistantChatPanel({ activeThreadId, + setActiveThreadId, mode, assistant, form, @@ -223,8 +262,17 @@ function AssistantChatPanel({ const chat = usePersistedChat({ threadId: activeThreadId, systemPrompt: assistant.system_prompt, + onThreadChange: setActiveThreadId, }); + // Destructure branching-related values from the hook + const { + inputController, + branchContext, + clearBranchContext, + branchFromMessage, + } = chat; + // Chat config is valid when gateway and model are both configured const hasChatConfig = Boolean(assistant.gateway_id) && @@ -246,6 +294,18 @@ function AssistantChatPanel({ }; await chat.sendMessage(text, metadata); + + // Clear editing state after successful send + clearBranchContext(); + }; + + // Handle clicking on the branch preview to go back to original thread + const handleGoToOriginalMessage = () => { + if (!branchContext) return; + setActiveThreadId(branchContext.originalThreadId); + // Clear the branch context since we're going back + clearBranchContext(); + inputController.setValue(""); }; const emptyState = ( @@ -333,6 +393,7 @@ function AssistantChatPanel({ messages={chat.messages} status={chat.status} minHeightOffset={240} + onBranchFromMessage={branchFromMessage} /> )}
@@ -346,18 +407,25 @@ function AssistantChatPanel({ )} > -
- -
+ +
+ + +
+
@@ -635,6 +703,7 @@ function AssistantDetailContent({ > => @@ -52,6 +65,20 @@ export interface UsePersistedChatOptions { onError?: (error: Error) => void; /** Called when a tool is invoked during chat */ onToolCall?: (event: { toolCall: { toolName: string } }) => void; + /** Called when the active thread changes (for branching) */ + onThreadChange?: (newThreadId: string) => void; +} + +/** + * Input controller for managing input state without re-rendering parent + */ +export interface InputController { + /** Get the current input value */ + getValue: () => string; + /** Set the input value (triggers subscriber) */ + setValue: (value: string) => void; + /** Subscribe to value changes (returns unsubscribe function) */ + subscribe: (callback: (value: string) => void) => () => void; } /** @@ -69,6 +96,34 @@ export interface PersistedChatResult { sendMessage: (text: string, metadata: Metadata) => Promise; /** Stop the current streaming response */ stop: () => void; + /** Set messages directly (for reverting, clearing, etc.) */ + setMessages: (messages: ChatMessage[]) => void; + /** Input controller for managing input state without re-rendering parent */ + inputController: InputController; + /** Current branch context if branching is in progress */ + branchContext: BranchContext | null; + /** Clear the branch context */ + clearBranchContext: () => void; + /** + * Branch from a specific message - creates a new thread with messages + * before the specified message, and sets up input for editing. + */ + branchFromMessage: (messageId: string, messageText: string) => Promise; +} + +/** + * Hook to subscribe to an InputController's value. + * Only re-renders the component using this hook, not the parent. + */ +export function useInputValue( + controller: InputController, +): [string, (value: string) => void] { + const value = useSyncExternalStore( + controller.subscribe, + controller.getValue, + controller.getValue, + ); + return [value, controller.setValue]; } /** @@ -86,8 +141,14 @@ export interface PersistedChatResult { export function usePersistedChat( options: UsePersistedChatOptions, ): PersistedChatResult { - const { threadId, systemPrompt, onCreateThread, onError, onToolCall } = - options; + const { + threadId, + systemPrompt, + onCreateThread, + onError, + onThreadChange, + onToolCall, + } = options; const { org: { slug: orgSlug }, @@ -98,6 +159,32 @@ export function usePersistedChat( const threadActions = useThreadActions(); const messageActions = useMessageActions(); + // Input controller using ref + pub/sub pattern (no re-renders on input change) + const inputControllerRef = useRef(null); + if (!inputControllerRef.current) { + let value = ""; + const subscribers = new Set<(value: string) => void>(); + inputControllerRef.current = { + getValue: () => value, + setValue: (newValue: string) => { + value = newValue; + for (const callback of subscribers) { + callback(newValue); + } + }, + subscribe: (callback: (value: string) => void) => { + subscribers.add(callback); + return () => subscribers.delete(callback); + }, + }; + } + const inputController = inputControllerRef.current; + + // State to track if we're editing from a branch (shows the original message preview) + const [branchContext, setBranchContext] = useState( + null, + ); + // Load persisted messages for this thread const persistedMessages = useThreadMessages(threadId) as unknown as Message[]; @@ -215,10 +302,67 @@ export function usePersistedChat( ); }; + // Branch from a specific message - creates a new thread with messages + // before the specified message, and sets up input for editing. + const branchFromMessage = async (messageId: string, messageText: string) => { + // Find the index of the message to branch from + const messageIndex = chat.messages.findIndex((m) => m.id === messageId); + if (messageIndex === -1) return; + + // Save the original thread context before switching + const originalThreadId = threadId; + + // Get messages to copy (before the clicked message, excluding system) + const messagesToCopy = chat.messages + .slice(0, messageIndex) + .filter((m) => m.role !== "system"); + + // Create a new thread + const newThreadId = crypto.randomUUID(); + + // Copy messages to the new thread with new IDs and updated thread_id + if (messagesToCopy.length > 0) { + const copiedMessages = messagesToCopy.map((msg) => ({ + ...msg, + id: crypto.randomUUID(), + metadata: { + ...msg.metadata, + thread_id: newThreadId, + created_at: msg.metadata?.created_at || new Date().toISOString(), + }, + })); + + // Insert copied messages into IndexedDB + await messageActions.insertMany.mutateAsync( + copiedMessages as unknown as Message[], + ); + } + + // Switch to the new thread + onThreadChange?.(newThreadId); + + // Set the message text in the input for editing + inputController.setValue(messageText); + + // Track the original context for the preview (allows navigating back) + setBranchContext({ + originalThreadId, + originalMessageId: messageId, + originalMessageText: messageText, + }); + }; + + const clearBranchContext = () => setBranchContext(null); + return { messages: chat.messages, status: chat.status, sendMessage, stop: chat.stop.bind(chat), + setMessages: chat.setMessages, + inputController, + branchContext, + clearBranchContext, + branchFromMessage, }; }