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
@@ -5,14 +5,29 @@ import {
  useRef,
  useState,
  type ReactNode,
+  type SVGProps,
 } from "react";
-import { Check, ChevronRight, Clock3, Copy, ImageIcon, Sparkles, Wrench } from "lucide-react";
+import {
+  Check,
+  ChevronRight,
+  Clock3,
+  Copy,
+  ImageIcon,
+  Sparkles,
+  Wrench,
+} from "lucide-react";
 import { useTranslation } from "react-i18next";

 import { AttachmentTile } from "@/components/AttachmentTile";
 import { CliAppMentionText } from "@/components/CliAppMentionText";
 import { ImageLightbox } from "@/components/ImageLightbox";
 import { MarkdownText, preloadMarkdownText } from "@/components/MarkdownText";
+import {
+  Tooltip,
+  TooltipContent,
+  TooltipProvider,
+  TooltipTrigger,
+} from "@/components/ui/tooltip";
 import { cn } from "@/lib/utils";
 import { copyTextToClipboard } from "@/lib/clipboard";
 import { formatTurnLatency } from "@/lib/format";
@@ -34,6 +49,7 @@ interface MessageBubbleProps {
  cliApps?: CliAppInfo[];
  mcpPresets?: McpPresetInfo[];
  onOpenFilePreview?: (path: string) => void;
+  onForkFromHere?: () => void;
 }

 /**
@@ -51,6 +67,7 @@ export function MessageBubble({
  cliApps = [],
  mcpPresets = [],
  onOpenFilePreview,
+  onForkFromHere,
 }: MessageBubbleProps) {
  const { t } = useTranslation();
  const [copied, setCopied] = useState(false);
@@ -73,7 +90,7 @@ export function MessageBubble({
    };
  }, []);

-  const onCopyAssistantReply = useCallback(() => {
+  const onCopyMessage = useCallback(() => {
    void copyTextToClipboard(message.content).then((ok) => {
      if (!ok) return;
      setCopied(true);
@@ -97,6 +114,11 @@ export function MessageBubble({
    const hasImages = images.length > 0;
    const hasMedia = media.length > 0;
    const hasText = message.content.trim().length > 0;
+    const showUserActions = hasText;
+    const timeLabel = formatMessageClock(message.createdAt);
+    const copyLabel = copied
+      ? t("message.copiedMessage", { defaultValue: "Copied message" })
+      : t("message.copyMessage", { defaultValue: "Copy message" });
    return (
      <div
        className={cn(
@@ -122,6 +144,43 @@ export function MessageBubble({
            />
          </p>
        ) : null}
+        {showUserActions ? (
+          <TooltipProvider delayDuration={180} skipDelayDuration={80}>
+            <div
+              className={cn(
+                "mt-0.5 flex h-8 items-center justify-end gap-1 self-end",
+                "text-[13px] text-muted-foreground/65 opacity-0 transition-opacity duration-150",
+                "group-focus-within:opacity-100 group-hover:opacity-100",
+              )}
+            >
+              {hasText ? (
+                <MessageActionTooltip label={copyLabel}>
+                  <button
+                    type="button"
+                    onClick={onCopyMessage}
+                    aria-label={copyLabel}
+                    className={cn(
+                      "inline-flex h-7 w-7 shrink-0 items-center justify-center rounded-full",
+                      "transition-colors hover:bg-muted/55 hover:text-foreground",
+                      "focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring",
+                    )}
+                  >
+                    {copied ? (
+                      <Check className="h-3.5 w-3.5" aria-hidden />
+                    ) : (
+                      <Copy className="h-3.5 w-3.5" aria-hidden />
+                    )}
+                  </button>
+                </MessageActionTooltip>
+              ) : null}
+              {timeLabel ? (
+                <span className="ml-1 shrink-0 select-none tabular-nums" title={timeLabel}>
+                  {timeLabel}
+                </span>
+              ) : null}
+            </div>
+          </TooltipProvider>
+        ) : null}
      </div>
    );
  }
@@ -138,13 +197,16 @@ export function MessageBubble({

  const showAssistantActions = message.role === "assistant" && !message.isStreaming && !empty;
  const showCopyButton = showAssistantCopyAction && showAssistantActions;
+  const showForkButton = showAssistantActions && !!onForkFromHere;
+  const copyReplyLabel = copied ? t("message.copiedReply") : t("message.copyReply");
+  const forkLabel = t("message.forkFromHere");
  const latencyMs = message.latencyMs;
  const showLatencyFooter =
    message.role === "assistant"
    && latencyMs != null
    && !message.isStreaming
    && (!empty || hasReasoning || media.length > 0);
-  const showAssistantFooterRow = showCopyButton || showLatencyFooter;
+  const showAssistantFooterRow = showCopyButton || showForkButton || showLatencyFooter;
  return (
    <div className={cn("w-full text-[15px]", baseAnim)} style={{ lineHeight: "var(--cjk-line-height)" }}>
      {hasReasoning ? (
@@ -173,35 +235,54 @@ export function MessageBubble({
          </MarkdownText>
          {media.length > 0 ? <MessageMedia media={media} align="left" /> : null}
          {showAssistantFooterRow ? (
-            <div className="mt-2 flex min-h-8 flex-wrap items-center gap-x-2 gap-y-1 text-muted-foreground">
-              {showCopyButton ? (
-                <button
-                  type="button"
-                  onClick={onCopyAssistantReply}
-                  aria-label={copied ? t("message.copiedReply") : t("message.copyReply")}
-                  title={copied ? t("message.copiedReply") : t("message.copyReply")}
-                  className={cn(
-                    "inline-flex h-8 w-8 shrink-0 items-center justify-center rounded-full",
-                    "transition-colors hover:bg-muted/55 hover:text-foreground",
-                    "focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring",
-                  )}
-                >
-                  {copied ? (
-                    <Check className="h-4 w-4" aria-hidden />
-                  ) : (
-                    <Copy className="h-4 w-4" aria-hidden />
-                  )}
-                </button>
-              ) : null}
-              {showLatencyFooter ? (
-                <span
-                  className="text-[11px] leading-none text-muted-foreground/70 tabular-nums"
-                  title={t("message.turnLatencyTitle")}
-                >
-                  {formatTurnLatency(latencyMs)}
-                </span>
-              ) : null}
-            </div>
+            <TooltipProvider delayDuration={180} skipDelayDuration={80}>
+              <div className="mt-2 flex min-h-8 flex-wrap items-center gap-x-2 gap-y-1 text-muted-foreground">
+                {showCopyButton ? (
+                  <MessageActionTooltip label={copyReplyLabel}>
+                    <button
+                      type="button"
+                      onClick={onCopyMessage}
+                      aria-label={copyReplyLabel}
+                      className={cn(
+                        "inline-flex h-8 w-8 shrink-0 items-center justify-center rounded-full",
+                        "transition-colors hover:bg-muted/55 hover:text-foreground",
+                        "focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring",
+                      )}
+                    >
+                      {copied ? (
+                        <Check className="h-4 w-4" aria-hidden />
+                      ) : (
+                        <Copy className="h-4 w-4" aria-hidden />
+                      )}
+                    </button>
+                  </MessageActionTooltip>
+                ) : null}
+                {showForkButton ? (
+                  <MessageActionTooltip label={forkLabel}>
+                    <button
+                      type="button"
+                      onClick={onForkFromHere}
+                      aria-label={forkLabel}
+                      className={cn(
+                        "inline-flex h-8 w-8 shrink-0 items-center justify-center rounded-full",
+                        "transition-colors hover:bg-muted/55 hover:text-foreground",
+                        "focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring",
+                      )}
+                    >
+                      <ForkFromHereIcon className="h-4 w-4" aria-hidden />
+                    </button>
+                  </MessageActionTooltip>
+                ) : null}
+                {showLatencyFooter ? (
+                  <span
+                    className="text-[11px] leading-none text-muted-foreground/70 tabular-nums"
+                    title={t("message.turnLatencyTitle")}
+                  >
+                    {formatTurnLatency(latencyMs)}
+                  </span>
+                ) : null}
+              </div>
+            </TooltipProvider>
          ) : null}
        </>
      )}
@@ -209,6 +290,27 @@ export function MessageBubble({
  );
 }

+function MessageActionTooltip({
+  label,
+  children,
+}: {
+  label: string;
+  children: ReactNode;
+}) {
+  return (
+    <Tooltip>
+      <TooltipTrigger asChild>{children}</TooltipTrigger>
+      <TooltipContent
+        side="top"
+        align="center"
+        className="rounded-full border-border/70 bg-background px-2.5 py-1 text-[12px] font-medium text-foreground shadow-[0_8px_24px_rgba(15,23,42,0.13)] dark:border-white/10 dark:bg-neutral-900 dark:text-white"
+      >
+        {label}
+      </TooltipContent>
+    </Tooltip>
+  );
+}
+
 function AutomationSourceBadge({ label, triggerLabel }: { label: string; triggerLabel: string }) {
  return (
    <div
@@ -228,6 +330,39 @@ function AutomationSourceBadge({ label, triggerLabel }: { label: string; trigger
  );
 }

+function formatMessageClock(createdAt: number): string {
+  if (!Number.isFinite(createdAt) || createdAt <= 0) return "";
+  try {
+    return new Intl.DateTimeFormat(undefined, {
+      hour: "2-digit",
+      minute: "2-digit",
+    }).format(new Date(createdAt));
+  } catch {
+    return "";
+  }
+}
+
+function ForkFromHereIcon({ className, ...props }: SVGProps<SVGSVGElement>) {
+  // Tabler Icons "arrow-fork" (MIT, Copyright Paweł Kuna).
+  return (
+    <svg
+      className={className}
+      viewBox="0 0 24 24"
+      fill="none"
+      stroke="currentColor"
+      strokeWidth="2"
+      strokeLinecap="round"
+      strokeLinejoin="round"
+      {...props}
+    >
+      <path d="M16 3h5v5" />
+      <path d="M8 3h-5v5" />
+      <path d="M21 3l-7.536 7.536a5 5 0 0 0 -1.464 3.534v6.93" />
+      <path d="M3 3l7.536 7.536a5 5 0 0 1 1.464 3.534v.93" />
+    </svg>
+  );
+}
+
 function mergeMcpMentionPresets(
  presets: McpPresetInfo[],
  attachments: UIMcpPresetAttachment[] | undefined,
@@ -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>