fix(core): ui state (#14933)

#### PR Dependency Tree * **PR #14933** 👈 This tree was auto-generated by [Charcoal](https://github.com/danerwilliams/charcoal)  ## Summary by CodeRabbit * **New Features** * Added draft tab option to AI chat interface * Introduced "Current document" session history view in chat history popover * Added control to show/hide "New Chat" button * **Improvements** * Enhanced chat history preservation when switching between sessions * Prevented duplicate session creation requests * Improved message handling during session transitions and generation [![Review Change Stack](https://storage.googleapis.com/coderabbit_public_assets/review-stack-in-coderabbit-ui.svg)](https://app.coderabbit.ai/change-stack/toeverything/AFFiNE/pull/14933)
2026-08-18 10:31:50 +08:00 · 2026-05-09 23:33:37 +08:00
parent fcc45a3f44
commit 417d31cabe
11 changed files with 487 additions and 53 deletions
@@ -3,7 +3,11 @@ import { type CopilotChatHistoryFragment } from '@affine/graphql';
 import { describe, expect, test, vi } from 'vitest';

 import {
+  canCreateNewDocPanelSession,
+  filterDocPanelTabs,
  getChatContentKey,
+  hasSessionMessages,
+  isSessionAvailableInDocPanel,
  resolveInitialSession,
  type SessionService,
  shouldResetChatPanelOnUserInfoChange,
@@ -89,6 +93,22 @@ describe('getChatContentKey', () => {
      },
      expected: 'session-2',
    },
+    {
+      name: 'keeps generating draft session on the doc key',
+      input: {
+        docId: 'doc-1',
+        hasPinned: false,
+        isGenerating: true,
+        previousSessionDocId: 'doc-1',
+        previousSessionId: 'session-1',
+        session: {
+          sessionId: 'session-2',
+          docId: 'doc-1',
+          messages: [],
+        },
+      },
+      expected: 'doc-1',
+    },
  ] satisfies {
    name: string;
    input: Parameters<typeof getChatContentKey>[0];
@@ -145,6 +165,55 @@ describe('shouldResetChatPanelOnUserInfoChange', () => {
  });
 });

+describe('doc panel tabs', () => {
+  const sessions = [
+    { sessionId: 'current-doc-session', docId: 'doc-1' },
+    { sessionId: 'workspace-session', docId: null },
+    { sessionId: 'other-doc-session', docId: 'doc-2' },
+  ];
+
+  test('allows only current doc or workspace sessions', () => {
+    expect(filterDocPanelTabs(sessions, 'doc-1')).toEqual([
+      sessions[0],
+      sessions[1],
+    ]);
+  });
+
+  test('rejects other document sessions', () => {
+    expect(isSessionAvailableInDocPanel(sessions[2], 'doc-1')).toBe(false);
+  });
+});
+
+describe('new session guard', () => {
+  test('allows a new session only after the current chat has messages', () => {
+    expect(
+      canCreateNewDocPanelSession({
+        hasContextMessages: false,
+        session: { messages: [] },
+        status: 'idle',
+      })
+    ).toBe(false);
+    expect(
+      canCreateNewDocPanelSession({
+        hasContextMessages: true,
+        session: { messages: [] },
+        status: 'idle',
+      })
+    ).toBe(true);
+    expect(hasSessionMessages({ messages: [{ id: 'message-1' }] })).toBe(true);
+  });
+
+  test('does not allow a new session while generating', () => {
+    expect(
+      canCreateNewDocPanelSession({
+        hasContextMessages: true,
+        session: null,
+        status: 'loading',
+      })
+    ).toBe(false);
+  });
+});
+
 test('returns undefined without session service or doc', async () => {
  await expect(
    resolveInitialSession({ sessionService: null, doc, workbench: null })
@@ -49,6 +49,12 @@ interface ChatContentKeySession {
  messages?: readonly unknown[] | null;
 }

+type TabSession = {
+  sessionId: string;
+  docId?: string | null;
+  messages?: readonly unknown[] | null;
+};
+
 export const shouldResetChatPanelOnUserInfoChange = ({
  previousUserId,
  nextUserId,
@@ -62,12 +68,14 @@ export const shouldResetChatPanelOnUserInfoChange = ({
 export const getChatContentKey = ({
  docId,
  hasPinned,
+  isGenerating,
  previousSessionDocId,
  previousSessionId,
  session,
 }: {
  docId?: string | null;
  hasPinned: boolean;
+  isGenerating?: boolean;
  previousSessionDocId?: string | null;
  previousSessionId?: string | null;
  session?: ChatContentKeySession | null;
@@ -80,6 +88,7 @@ export const getChatContentKey = ({

  const sessionDocId = session.docId ?? docId ?? null;
  const hasSessionHistory = !!session.messages?.length;
+  const shouldPreserveTransientMessages = isGenerating && !hasSessionHistory;
  const sessionSwitchedWithinDoc = !!(
    previousSessionId &&
    previousSessionId !== sessionId &&
@@ -89,11 +98,51 @@ export const getChatContentKey = ({
    sessionDocId === docId
  );

-  return hasPinned || hasSessionHistory || sessionSwitchedWithinDoc
+  return hasPinned ||
+    hasSessionHistory ||
+    (sessionSwitchedWithinDoc && !shouldPreserveTransientMessages)
    ? sessionId
    : fallbackKey;
 };

+export const isSessionAvailableInDocPanel = (
+  session: TabSession,
+  docId?: string | null
+) => {
+  return !session.docId || session.docId === docId;
+};
+
+export const filterDocPanelTabs = <T extends TabSession>(
+  sessions: T[],
+  docId?: string | null
+) => {
+  return sessions.filter(session =>
+    isSessionAvailableInDocPanel(session, docId)
+  );
+};
+
+export const hasSessionMessages = (
+  session?: Pick<TabSession, 'messages'> | null
+) => {
+  return !!session?.messages?.length;
+};
+
+export const canCreateNewDocPanelSession = ({
+  hasContextMessages,
+  session,
+  status,
+}: {
+  hasContextMessages: boolean;
+  session?: Pick<TabSession, 'messages'> | null;
+  status?: string | null;
+}) => {
+  return (
+    (hasContextMessages || hasSessionMessages(session)) &&
+    status !== 'loading' &&
+    status !== 'transmitting'
+  );
+};
+
 export const getSessionIdFromUrl = (workbench?: WorkbenchLike | null) => {
  if (!workbench) {
    return undefined;
@@ -52,7 +52,10 @@ import {
 } from '../../chat-panel-utils';
 import * as styles from './chat.css';
 import {
+  canCreateNewDocPanelSession,
+  filterDocPanelTabs,
  getChatContentKey,
+  isSessionAvailableInDocPanel,
  resolveInitialSession,
  shouldResetChatPanelOnUserInfoChange,
  type WorkbenchLike,
@@ -108,6 +111,14 @@ export const EditorChatPanel = ({ editor, onLoad }: SidebarTabProps) => {
  const prevSessionDocIdRef = useRef<string | null>(null);
  const lastDocIdRef = useRef<string | null>(null);
  const sessionLoadSeqRef = useRef(0);
+  const creatingSessionRef = useRef<{
+    docId: string;
+    promise: Promise<CopilotChatHistoryFragment | undefined>;
+  } | null>(null);
+  const creatingFreshSessionRef = useRef<{
+    docId: string;
+    promise: Promise<void>;
+  } | null>(null);
  const userIdRef = useRef<string | null | undefined>(undefined);

  const doc = editor?.doc;
@@ -117,6 +128,7 @@ export const EditorChatPanel = ({ editor, onLoad }: SidebarTabProps) => {
  const [sessionServiceReady, setSessionServiceReady] = useState(
    () => !!AIProvider.session
  );
+  const [hasContextMessages, setHasContextMessages] = useState(false);

  useEffect(() => {
    if (sessionServiceReady) return;
@@ -142,6 +154,15 @@ export const EditorChatPanel = ({ editor, onLoad }: SidebarTabProps) => {

  const { openTabs, setOpenTabs } =
    useAIChatOpenTabs<CopilotChatHistoryFragment>(loadSession);
+  const visibleOpenTabs = useMemo(
+    () => filterDocPanelTabs(openTabs, doc?.id),
+    [doc?.id, openTabs]
+  );
+  const canCreateNewSession = canCreateNewDocPanelSession({
+    hasContextMessages,
+    session,
+    status,
+  });

  const appSidebarConfig = useMemo<AppSidebarConfig>(() => {
    return {
@@ -200,18 +221,32 @@ export const EditorChatPanel = ({ editor, onLoad }: SidebarTabProps) => {
      if (session || !AIProvider.session || !doc) {
        return session ?? undefined;
      }
+      if (creatingSessionRef.current?.docId === doc.id) {
+        return creatingSessionRef.current.promise;
+      }
      const requestSeq = ++sessionLoadSeqRef.current;
-      const nextSession = await AIProvider.session.createSessionWithHistory({
-        docId: doc.id,
-        workspaceId: doc.workspace.id,
-        promptName: 'Chat With AFFiNE AI',
-        reuseLatestChat: false,
-        ...options,
-      });
-      if (requestSeq !== sessionLoadSeqRef.current) return undefined;
-      setSession(nextSession ?? null);
-      setHasPinned(!!nextSession?.pinned);
-      return nextSession ?? undefined;
+      let promise: Promise<CopilotChatHistoryFragment | undefined>;
+      promise = AIProvider.session
+        .createSessionWithHistory({
+          docId: doc.id,
+          workspaceId: doc.workspace.id,
+          promptName: 'Chat With AFFiNE AI',
+          reuseLatestChat: false,
+          ...options,
+        })
+        .then(nextSession => {
+          if (requestSeq !== sessionLoadSeqRef.current) return undefined;
+          setSession(nextSession ?? null);
+          setHasPinned(!!nextSession?.pinned);
+          return nextSession ?? undefined;
+        })
+        .finally(() => {
+          if (creatingSessionRef.current?.promise === promise) {
+            creatingSessionRef.current = null;
+          }
+        });
+      creatingSessionRef.current = { docId: doc.id, promise };
+      return promise;
    },
    [doc, session]
  );
@@ -236,29 +271,44 @@ export const EditorChatPanel = ({ editor, onLoad }: SidebarTabProps) => {
  );

  const newSession = useCallback(async () => {
+    if (!canCreateNewSession) {
+      return;
+    }
+    if (doc && creatingFreshSessionRef.current?.docId === doc.id) {
+      return creatingFreshSessionRef.current.promise;
+    }
    resetPanel();
    const requestSeq = sessionLoadSeqRef.current;
    setSession(null);
+    setHasContextMessages(false);

    if (!AIProvider.session || !doc) {
      return;
    }

-    try {
-      const nextSession = await AIProvider.session.createSessionWithHistory({
+    let promise: Promise<void>;
+    promise = AIProvider.session
+      .createSessionWithHistory({
        docId: doc.id,
        workspaceId: doc.workspace.id,
        promptName: 'Chat With AFFiNE AI',
        reuseLatestChat: false,
+      })
+      .then(nextSession => {
+        if (requestSeq === sessionLoadSeqRef.current) {
+          setSession(nextSession ?? null);
+          setHasPinned(!!nextSession?.pinned);
+        }
+      })
+      .catch(console.error)
+      .finally(() => {
+        if (creatingFreshSessionRef.current?.promise === promise) {
+          creatingFreshSessionRef.current = null;
+        }
      });
-      if (requestSeq === sessionLoadSeqRef.current) {
-        setSession(nextSession ?? null);
-        setHasPinned(!!nextSession?.pinned);
-      }
-    } catch (error) {
-      console.error(error);
-    }
-  }, [doc, resetPanel]);
+    creatingFreshSessionRef.current = { docId: doc.id, promise };
+    return promise;
+  }, [canCreateNewSession, doc, resetPanel]);

  const openSession = useCallback(
    async (sessionId: string) => {
@@ -277,13 +327,20 @@ export const EditorChatPanel = ({ editor, onLoad }: SidebarTabProps) => {
          setOpenTabs(prev => prev.filter(tab => tab.sessionId !== sessionId));
          return;
        }
+        if (!isSessionAvailableInDocPanel(nextSession, doc.id)) {
+          setOpenTabs([]);
+          workbench.open(`/${nextSession.docId}?sessionId=${sessionId}`, {
+            at: 'active',
+          });
+          return;
+        }
        setSession(nextSession);
        setHasPinned(!!nextSession.pinned);
      } catch (error) {
        console.error(error);
      }
    },
-    [doc, session?.sessionId, setOpenTabs]
+    [doc, session?.sessionId, setOpenTabs, workbench]
  );

  const openDoc = useCallback(
@@ -304,9 +361,17 @@ export const EditorChatPanel = ({ editor, onLoad }: SidebarTabProps) => {
        workbench.open(`/${docId}`, { at: 'active' });
        return;
      }
+      setOpenTabs([]);
      workbench.open(`/${docId}?sessionId=${sessionId}`, { at: 'active' });
    },
-    [doc, openSession, session?.pinned, session?.sessionId, workbench]
+    [
+      doc,
+      openSession,
+      session?.pinned,
+      session?.sessionId,
+      setOpenTabs,
+      workbench,
+    ]
  );

  const deleteSession = useMemo(
@@ -325,10 +390,12 @@ export const EditorChatPanel = ({ editor, onLoad }: SidebarTabProps) => {
        isActiveSession: sessionToDelete =>
          sessionToDelete.sessionId === session?.sessionId,
        onActiveSessionDeleted: () => {
-          newSession().catch(console.error);
+          resetPanel();
+          setSession(null);
+          setHasContextMessages(false);
        },
      }),
-    [newSession, notificationService, session?.sessionId, t]
+    [notificationService, resetPanel, session?.sessionId, t]
  );

  const closeTab = useCallback(
@@ -338,17 +405,20 @@ export const EditorChatPanel = ({ editor, onLoad }: SidebarTabProps) => {
        const idx = prev.findIndex(tab => tab.sessionId === sessionId);
        if (idx === -1) return prev;
        const next = prev.filter(tab => tab.sessionId !== sessionId);
-        fallback = next[idx] ?? next[idx - 1] ?? next[0];
+        const visibleNext = filterDocPanelTabs(next, doc?.id);
+        fallback = visibleNext[idx] ?? visibleNext[idx - 1] ?? visibleNext[0];
        return next;
      });
      if (session?.sessionId !== sessionId) return;
      if (fallback) {
        openSession(fallback.sessionId).catch(console.error);
      } else {
-        newSession().catch(console.error);
+        resetPanel();
+        setSession(null);
+        setHasContextMessages(false);
      }
    },
-    [newSession, openSession, session?.sessionId, setOpenTabs]
+    [doc?.id, openSession, resetPanel, session?.sessionId, setOpenTabs]
  );

  const togglePin = useCallback(async () => {
@@ -387,7 +457,12 @@ export const EditorChatPanel = ({ editor, onLoad }: SidebarTabProps) => {

  const onContextChange = useCallback(
    (context: Partial<ChatContextValue>) => {
-      setStatus(context.status ?? 'idle');
+      if (context.status) {
+        setStatus(context.status);
+      }
+      if (context.messages) {
+        setHasContextMessages(context.messages.length > 0);
+      }
      if (context.status === 'success') {
        rebindSession().catch(console.error);
      }
@@ -468,6 +543,7 @@ export const EditorChatPanel = ({ editor, onLoad }: SidebarTabProps) => {
      !session?.pinned
    ) {
      resetPanel();
+      setHasContextMessages(false);
    }
    lastDocIdRef.current = docId;
  }, [doc?.id, resetPanel, session?.pinned]);
@@ -490,6 +566,7 @@ export const EditorChatPanel = ({ editor, onLoad }: SidebarTabProps) => {
  const contentKey = getChatContentKey({
    docId: doc?.id,
    hasPinned,
+    isGenerating: status === 'loading' || status === 'transmitting',
    previousSessionDocId: prevSessionDocIdRef.current,
    previousSessionId: prevSessionIdRef.current,
    session,
@@ -597,6 +674,7 @@ export const EditorChatPanel = ({ editor, onLoad }: SidebarTabProps) => {
      workspaceId: doc.workspace.id,
      docId: doc.id,
      status,
+      canCreateNewSession,
      docDisplayConfig,
      notificationService,
      onNewSession: () => {
@@ -620,6 +698,7 @@ export const EditorChatPanel = ({ editor, onLoad }: SidebarTabProps) => {
    }
  }, [
    chatToolbar,
+    canCreateNewSession,
    deleteSession,
    doc,
    docDisplayConfig,
@@ -647,15 +726,17 @@ export const EditorChatPanel = ({ editor, onLoad }: SidebarTabProps) => {
      chatTabsContainerRef.current.append(tabs);
      setChatTabs(tabs);
    }
-    tabs.sessions = openTabs;
+    tabs.sessions = visibleOpenTabs;
    tabs.activeSessionId = session?.sessionId;
+    tabs.showDraftTab =
+      visibleOpenTabs.length === 0 && !session?.sessionId && !!doc;
    tabs.onSelectTab = (sessionId: string) => {
      openSession(sessionId).catch(console.error);
    };
    tabs.onCloseTab = (sessionId: string) => {
      closeTab(sessionId);
    };
-  }, [chatTabs, closeTab, doc, openSession, openTabs, session]);
+  }, [chatTabs, closeTab, doc, openSession, session, visibleOpenTabs]);

  useEffect(() => {
    if (!editor?.host || !chatContent) {