refactor(webui): shrink fork implementation

2026-06-10 04:26:06 +08:00
parent 1f926e3769
commit 916525f94a
24 changed files with 134 additions and 879 deletions
@@ -172,7 +172,6 @@ interface ThreadComposerProps {
  workspaceError?: string | null;
  onWorkspaceScopeChange?: (scope: WorkspaceScopePayload) => void;
  pendingQueueKey?: string | null;
-  externalError?: string | null;
 }

 const COMMAND_ICONS: Record<string, LucideIcon> = {
@@ -766,7 +765,6 @@ export function ThreadComposer({
  workspaceError = null,
  onWorkspaceScopeChange,
  pendingQueueKey = null,
-  externalError = null,
 }: ThreadComposerProps) {
  const { t } = useTranslation();
  const [value, setValue] = useState("");
@@ -1149,10 +1147,6 @@ export function ThreadComposer({
    });
  }, [clear, pendingQueueKey]);

-  useEffect(() => {
-    if (externalError) setInlineError(externalError);
-  }, [externalError]);
-
  const appendTranscription = useCallback((text: string) => {
    const transcript = text.trim();
    if (!transcript) return;
@@ -8,10 +8,10 @@ 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;
+  hiddenUserMessageCount?: number;
  onLoadEarlier?: () => void;
  cliApps?: CliAppInfo[];
  mcpPresets?: McpPresetInfo[];
@@ -65,9 +65,9 @@ export function assistantCopyFlags(units: DisplayUnit[]): boolean[] {

 export function ThreadMessages({
  messages,
-  allMessages,
  isStreaming = false,
  hiddenMessageCount = 0,
+  hiddenUserMessageCount = 0,
  onLoadEarlier,
  cliApps = [],
  mcpPresets = [],
@@ -81,15 +81,12 @@ export function ThreadMessages({
    () => unitIndexAfterMessageCount(units, forkBoundaryMessageCount),
    [forkBoundaryMessageCount, units],
  );
-  const assistantForkIndexById = useMemo(
-    () => assistantForkIndexByMessageId(allMessages ?? messages),
-    [allMessages, messages],
-  );
  const copyFlags = useMemo(() => assistantCopyFlags(units), [units]);
  const liveActivityClusterIndices = useMemo(
    () => isStreaming ? currentActivityClusterIndices(units) : new Set<number>(),
    [isStreaming, units],
  );
+  let nextUserIndex = hiddenUserMessageCount;

  return (
    <div className="flex w-full flex-col">
@@ -123,6 +120,11 @@ export function ThreadMessages({
          unit.type === "message" && unit.message.role === "user"
            ? unit.message.id
            : undefined;
+        const forkIndex =
+          unit.type === "message" && unit.message.role === "assistant" && copyFlags[index]
+            ? nextUserIndex
+            : undefined;
+        if (unit.type === "message" && unit.message.role === "user") nextUserIndex += 1;

        return (
          <Fragment key={unitKey(unit, index)}>
@@ -149,20 +151,15 @@ export function ThreadMessages({
                  mcpPresets={mcpPresets}
                  onOpenFilePreview={onOpenFilePreview}
                  onForkFromHere={
-                    onForkFromMessage
-                      ? forkHandlerForAssistantMessage(
-                          unit.message,
-                          copyFlags[index],
-                          assistantForkIndexById,
-                          onForkFromMessage,
-                        )
+                    onForkFromMessage && forkIndex !== undefined
+                      ? () => onForkFromMessage(forkIndex)
                      : undefined
                  }
                />
              )}
            </div>
            {index === forkBoundaryAfterUnitIndex ? (
-              <ForkBoundaryDivider label={t("thread.fork.fromHistory")} />
+              <ForkBoundaryDivider label={t("thread.forkedFromHistory")} />
            ) : null}
          </Fragment>
        );
@@ -195,34 +192,6 @@ function ForkBoundaryDivider({ label }: { label: string }) {
  );
 }

-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;
@@ -278,8 +278,6 @@ 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);
@@ -288,7 +286,6 @@ 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());
@@ -340,12 +337,6 @@ 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;
@@ -455,12 +446,6 @@ 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;
@@ -539,7 +524,6 @@ export function ThreadShell({

  const handleThreadSend = useCallback(
    (content: string, images?: SendImage[], options?: SendOptions) => {
-      setForkError(null);
      setScrollToBottomSignal((value) => value + 1);
      send(content, images, withWorkspaceScope(options));
    },
@@ -637,21 +621,13 @@ export function ThreadShell({
  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;
-      }
+      if (!forkedChatId) return;
      messageCacheRef.current.delete(forkedChatId);
      appliedHistoryVersionRef.current.delete(forkedChatId);
      pendingCanonicalHydrateRef.current.add(forkedChatId);
-      setForkHydratingChatId(forkedChatId);
-      setForkError(null);
    },
-    [chatId, onForkChat, t],
+    [chatId, onForkChat],
  );

  const composer = (
@@ -665,7 +641,7 @@ export function ThreadShell({
      {session ? (
        <ThreadComposer
          onSend={handleThreadSend}
-          disabled={!chatId || forkHydratingChatId === chatId}
+          disabled={!chatId}
          isStreaming={isStreaming}
          placeholder={
            showHeroComposer
@@ -692,7 +668,6 @@ export function ThreadShell({
          workspaceError={workspaceError}
          onWorkspaceScopeChange={onWorkspaceScopeChange}
          pendingQueueKey={chatId}
-          externalError={forkError}
        />
      ) : (
        <ThreadComposer
@@ -776,7 +751,6 @@ export function ThreadShell({
          showScrollToBottomButton={!!session}
          cliApps={cliApps}
          mcpPresets={mcpPresets}
-          allMessages={displayMessages}
          forkBoundaryMessageCount={forkBoundaryMessageCount}
          onOpenFilePreview={historyKey ? handleOpenFilePreview : undefined}
          onForkFromMessage={onForkChat ? handleForkFromMessage : undefined}
@@ -29,7 +29,6 @@ export interface ThreadViewportHandle {

 interface ThreadViewportProps {
  messages: UIMessage[];
-  allMessages?: UIMessage[];
  isStreaming: boolean;
  composer: ReactNode;
  emptyState?: ReactNode;
@@ -64,7 +63,6 @@ export function windowMessages(messages: UIMessage[], visibleCount: number): UIM

 export const ThreadViewport = forwardRef<ThreadViewportHandle, ThreadViewportProps>(function ThreadViewport({
  messages,
-  allMessages,
  isStreaming,
  composer,
  emptyState,
@@ -100,6 +98,10 @@ export const ThreadViewport = forwardRef<ThreadViewportHandle, ThreadViewportPro
    [messages, visibleMessageCount],
  );
  const hiddenMessageCount = messages.length - visibleMessages.length;
+  const hiddenUserMessageCount =
+    hiddenMessageCount > 0
+      ? messages.slice(0, hiddenMessageCount).filter((message) => message.role === "user").length
+      : 0;
  const visibleForkBoundaryMessageCount =
    forkBoundaryMessageCount !== null && forkBoundaryMessageCount > hiddenMessageCount
      ? forkBoundaryMessageCount - hiddenMessageCount
@@ -299,9 +301,9 @@ 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}
+                  hiddenUserMessageCount={hiddenUserMessageCount}
                  onLoadEarlier={loadEarlierMessages}
                  cliApps={cliApps}
                  mcpPresets={mcpPresets}