feat(webui): add assistant reply fork-from-here

2026-06-10 04:26:06 +08:00
parent 4a58b83acc
commit 03bca4c0a9
30 changed files with 1358 additions and 36 deletions
@@ -172,6 +172,7 @@ interface ThreadComposerProps {
  workspaceError?: string | null;
  onWorkspaceScopeChange?: (scope: WorkspaceScopePayload) => void;
  pendingQueueKey?: string | null;
+  externalError?: string | null;
 }

 const COMMAND_ICONS: Record<string, LucideIcon> = {
@@ -765,6 +766,7 @@ export function ThreadComposer({
  workspaceError = null,
  onWorkspaceScopeChange,
  pendingQueueKey = null,
+  externalError = null,
 }: ThreadComposerProps) {
  const { t } = useTranslation();
  const [value, setValue] = useState("");
@@ -782,6 +784,7 @@ export function ThreadComposer({
  const chipRefs = useRef(new Map<string, HTMLButtonElement>());
  const queuedPromptCounterRef = useRef(0);
  const draggedQueuedPromptIdRef = useRef<string | null>(null);
+  const previousPendingQueueKeyRef = useRef(pendingQueueKey);
  const wasStreamingRef = useRef(isStreaming);
  const skipNextQueuedFlushRef = useRef(false);
  const skipQueuedPromptPersistRef = useRef(false);
@@ -1128,6 +1131,28 @@ export function ThreadComposer({
    });
  }, []);

+  // Runs before paint so switching sessions never flashes stale draft text.
+  useLayoutEffect(() => {
+    if (previousPendingQueueKeyRef.current === pendingQueueKey) return;
+    previousPendingQueueKeyRef.current = pendingQueueKey;
+    setValue("");
+    setInlineError(null);
+    setSlashMenuDismissed(false);
+    setCliAppMenuDismissed(false);
+    setCursorPosition(0);
+    clear();
+    requestAnimationFrame(() => {
+      const el = textareaRef.current;
+      if (!el) return;
+      el.style.height = "auto";
+      el.style.height = `${Math.min(el.scrollHeight, 260)}px`;
+    });
+  }, [clear, pendingQueueKey]);
+
+  useEffect(() => {
+    if (externalError) setInlineError(externalError);
+  }, [externalError]);
+
  const appendTranscription = useCallback((text: string) => {
    const transcript = text.trim();
    if (!transcript) return;
@@ -8,6 +8,7 @@ import type { CliAppInfo, McpPresetInfo, UIMessage } from "@/lib/types";

 interface ThreadMessagesProps {
  messages: UIMessage[];
+  allMessages?: UIMessage[];
  /** When true, agent turn still in flight — keeps activity timeline expanded. */
  isStreaming?: boolean;
  hiddenMessageCount?: number;
@@ -15,6 +16,7 @@ interface ThreadMessagesProps {
  cliApps?: CliAppInfo[];
  mcpPresets?: McpPresetInfo[];
  onOpenFilePreview?: (path: string) => void;
+  onForkFromMessage?: (beforeUserIndex: number) => void;
 }

 export type DisplayUnit = TurnUnit;
@@ -62,15 +64,21 @@ export function assistantCopyFlags(units: DisplayUnit[]): boolean[] {

 export function ThreadMessages({
  messages,
+  allMessages,
  isStreaming = false,
  hiddenMessageCount = 0,
  onLoadEarlier,
  cliApps = [],
  mcpPresets = [],
  onOpenFilePreview,
+  onForkFromMessage,
 }: ThreadMessagesProps) {
  const { t } = useTranslation();
  const units = useMemo(() => buildDisplayUnits(messages, isStreaming), [isStreaming, messages]);
+  const assistantForkIndexById = useMemo(
+    () => assistantForkIndexByMessageId(allMessages ?? messages),
+    [allMessages, messages],
+  );
  const copyFlags = useMemo(() => assistantCopyFlags(units), [units]);
  const liveActivityClusterIndices = useMemo(
    () => isStreaming ? currentActivityClusterIndices(units) : new Set<number>(),
@@ -137,6 +145,16 @@ export function ThreadMessages({
                cliApps={cliApps}
                mcpPresets={mcpPresets}
                onOpenFilePreview={onOpenFilePreview}
+                onForkFromHere={
+                  onForkFromMessage
+                    ? forkHandlerForAssistantMessage(
+                        unit.message,
+                        copyFlags[index],
+                        assistantForkIndexById,
+                        onForkFromMessage,
+                      )
+                    : undefined
+                }
              />
            )}
          </div>
@@ -146,6 +164,34 @@ export function ThreadMessages({
  );
 }

+function assistantForkIndexByMessageId(messages: UIMessage[]): Map<string, number> {
+  const out = new Map<string, number>();
+  let nextUserIndex = 0;
+  for (const message of messages) {
+    if (message.role === "user") {
+      nextUserIndex += 1;
+    } else if (message.role === "assistant") {
+      out.set(message.id, nextUserIndex);
+    }
+  }
+  return out;
+}
+
+function forkHandlerForAssistantMessage(
+  message: UIMessage,
+  canForkAssistant: boolean,
+  assistantForkIndexById: Map<string, number>,
+  onForkFromMessage: NonNullable<ThreadMessagesProps["onForkFromMessage"]>,
+): (() => void) | undefined {
+  if (message.role === "assistant" && canForkAssistant) {
+    const beforeUserIndex = assistantForkIndexById.get(message.id);
+    return beforeUserIndex === undefined
+      ? undefined
+      : () => onForkFromMessage(beforeUserIndex);
+  }
+  return undefined;
+}
+
 function currentActivityClusterIndices(units: DisplayUnit[]): Set<number> {
  const indices = new Set<number>();
  let markedCurrentActivity = false;
@@ -77,6 +77,7 @@ interface ThreadShellProps {
  onGoHome?: () => void;
  onNewChat?: () => void;
  onCreateChat?: (workspaceScope?: WorkspaceScopePayload | null) => Promise<string | null>;
+  onForkChat?: (sourceChatId: string, beforeUserIndex: number) => Promise<string | null>;
  onTurnEnd?: () => void;
  theme?: "light" | "dark";
  onToggleTheme?: () => void;
@@ -226,6 +227,7 @@ export function ThreadShell({
  title,
  onToggleSidebar,
  onCreateChat,
+  onForkChat,
  onTurnEnd,
  theme = "light",
  onToggleTheme = () => {},
@@ -275,6 +277,8 @@ export function ThreadShell({
  const [filePreviewPath, setFilePreviewPath] = useState<string | null>(null);
  const [filePreviewClosing, setFilePreviewClosing] = useState(false);
  const [filePreviewWidth, setFilePreviewWidth] = useState(FILE_PREVIEW_DEFAULT_WIDTH);
+  const [forkError, setForkError] = useState<string | null>(null);
+  const [forkHydratingChatId, setForkHydratingChatId] = useState<string | null>(null);
  const shellRef = useRef<HTMLElement | null>(null);
  const filePreviewWidthRef = useRef(FILE_PREVIEW_DEFAULT_WIDTH);
  const filePreviewCloseTimerRef = useRef<number | null>(null);
@@ -283,6 +287,7 @@ export function ThreadShell({
  const messageCacheRef = useRef<Map<string, UIMessage[]>>(new Map());
  /** Last chatId we associated with the in-memory thread (for cache-on-switch). */
  const prevChatIdForCacheRef = useRef<string | null>(null);
+  const prevChatIdForComposerRef = useRef<string | null>(chatId);
  /** Skip one message-cache write right after chatId changes (messages may not match yet). */
  const skipLayoutCacheRef = useRef(false);
  const appliedHistoryVersionRef = useRef<Map<string, number>>(new Map());
@@ -334,6 +339,12 @@ export function ThreadShell({
    };
  }, []);

+  useEffect(() => {
+    if (prevChatIdForComposerRef.current === chatId) return;
+    prevChatIdForComposerRef.current = chatId;
+    setForkError(null);
+  }, [chatId]);
+
  const displayMessages = useMemo(() => projectWebuiThreadMessages(messages), [messages]);

  const showHeroComposer = messages.length === 0 && !loading;
@@ -443,6 +454,12 @@ export function ThreadShell({
    setMessages(projectWebuiThreadMessages(historical));
  }, [chatId, historical, setMessages]);

+  useEffect(() => {
+    if (!chatId || loading || forkHydratingChatId !== chatId) return;
+    setForkHydratingChatId(null);
+    setScrollToBottomSignal((value) => value + 1);
+  }, [chatId, forkHydratingChatId, loading]);
+
  useLayoutEffect(() => {
    if (chatId) {
      const prev = prevChatIdForCacheRef.current;
@@ -521,6 +538,7 @@ export function ThreadShell({

  const handleThreadSend = useCallback(
    (content: string, images?: SendImage[], options?: SendOptions) => {
+      setForkError(null);
      setScrollToBottomSignal((value) => value + 1);
      send(content, images, withWorkspaceScope(options));
    },
@@ -615,6 +633,26 @@ export function ThreadShell({
    };
  }, [filePreviewPath]);

+  const handleForkFromMessage = useCallback(
+    async (beforeUserIndex: number) => {
+      if (!chatId || !onForkChat) return;
+      setForkError(null);
+      const forkedChatId = await onForkChat(chatId, beforeUserIndex);
+      if (!forkedChatId) {
+        setForkError(t("thread.fork.failed", {
+          defaultValue: "Could not fork this chat. Try again.",
+        }));
+        return;
+      }
+      messageCacheRef.current.delete(forkedChatId);
+      appliedHistoryVersionRef.current.delete(forkedChatId);
+      pendingCanonicalHydrateRef.current.add(forkedChatId);
+      setForkHydratingChatId(forkedChatId);
+      setForkError(null);
+    },
+    [chatId, onForkChat, t],
+  );
+
  const composer = (
    <>
      {streamError ? (
@@ -626,7 +664,7 @@ export function ThreadShell({
      {session ? (
        <ThreadComposer
          onSend={handleThreadSend}
-          disabled={!chatId}
+          disabled={!chatId || forkHydratingChatId === chatId}
          isStreaming={isStreaming}
          placeholder={
            showHeroComposer
@@ -653,6 +691,7 @@ export function ThreadShell({
          workspaceError={workspaceError}
          onWorkspaceScopeChange={onWorkspaceScopeChange}
          pendingQueueKey={chatId}
+          externalError={forkError}
        />
      ) : (
        <ThreadComposer
@@ -736,7 +775,9 @@ export function ThreadShell({
          showScrollToBottomButton={!!session}
          cliApps={cliApps}
          mcpPresets={mcpPresets}
+          allMessages={displayMessages}
          onOpenFilePreview={historyKey ? handleOpenFilePreview : undefined}
+          onForkFromMessage={onForkChat ? handleForkFromMessage : undefined}
        />
      </div>
      {filePreviewPath && historyKey ? (
@@ -29,6 +29,7 @@ export interface ThreadViewportHandle {

 interface ThreadViewportProps {
  messages: UIMessage[];
+  allMessages?: UIMessage[];
  isStreaming: boolean;
  composer: ReactNode;
  emptyState?: ReactNode;
@@ -38,6 +39,7 @@ interface ThreadViewportProps {
  cliApps?: CliAppInfo[];
  mcpPresets?: McpPresetInfo[];
  onOpenFilePreview?: (path: string) => void;
+  onForkFromMessage?: (beforeUserIndex: number) => void;
 }

 const NEAR_BOTTOM_PX = 48;
@@ -61,6 +63,7 @@ export function windowMessages(messages: UIMessage[], visibleCount: number): UIM

 export const ThreadViewport = forwardRef<ThreadViewportHandle, ThreadViewportProps>(function ThreadViewport({
  messages,
+  allMessages,
  isStreaming,
  composer,
  emptyState,
@@ -70,6 +73,7 @@ export const ThreadViewport = forwardRef<ThreadViewportHandle, ThreadViewportPro
  cliApps = [],
  mcpPresets = [],
  onOpenFilePreview,
+  onForkFromMessage,
 }, ref) {
  const { t } = useTranslation();
  const scrollRef = useRef<HTMLDivElement>(null);
@@ -289,12 +293,14 @@ export const ThreadViewport = forwardRef<ThreadViewportHandle, ThreadViewportPro
              <div className="mx-auto w-full max-w-[49.5rem]">
                <ThreadMessages
                  messages={visibleMessages}
+                  allMessages={allMessages ?? messages}
                  isStreaming={isStreaming}
                  hiddenMessageCount={hiddenMessageCount}
                  onLoadEarlier={loadEarlierMessages}
                  cliApps={cliApps}
                  mcpPresets={mcpPresets}
                  onOpenFilePreview={onOpenFilePreview}
+                  onForkFromMessage={onForkFromMessage}
                />
              </div>
            </div>