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-11 05:58:56 +08:00 · 2026-05-09 23:33:37 +08:00
parent fcc45a3f44
commit 417d31cabe
11 changed files with 487 additions and 53 deletions
@@ -1,10 +1,17 @@
 /**
 * @vitest-environment happy-dom
 */
-import { describe, expect, test, vi } from 'vitest';
+import { afterEach, describe, expect, test, vi } from 'vitest';

+import { AIProvider } from '../../provider';
 import { AIChatContent } from './ai-chat-content';

+const originalHistories = AIProvider.histories;
+
+afterEach(() => {
+  AIProvider.provide('histories', originalHistories as any);
+});
+
 describe('AIChatContent pinned scroll tracking', () => {
  test('records scroll position from the chat messages host', async () => {
    let scrollEndHandler: (() => void) | undefined;
@@ -39,3 +46,104 @@ describe('AIChatContent pinned scroll tracking', () => {
    expect((content as any).lastScrollTop).toBe(256);
  });
 });
+
+describe('AIChatContent history loading', () => {
+  test('replaces messages when the active session changes', async () => {
+    const histories = {
+      chats: vi.fn(async (_workspaceId: string, sessionId: string) => [
+        {
+          messages: [
+            {
+              id: `${sessionId}-message`,
+              role: 'user',
+              content: sessionId,
+              createdAt: '2026-01-01T00:00:00.000Z',
+            },
+          ],
+        },
+      ]),
+      actions: vi.fn(async () => []),
+      cleanup: vi.fn(),
+      ids: vi.fn(),
+    };
+    AIProvider.provide('histories', histories as any);
+
+    const content: {
+      updateHistoryCounter: number;
+      historyKey: string | undefined;
+      workspaceId: string;
+      docId: string;
+      session: { sessionId: string };
+      chatContextValue: { messages: unknown[]; status?: string };
+      updateContext: (context: { messages: unknown[] }) => void;
+    } = {
+      updateHistoryCounter: 0,
+      historyKey: undefined,
+      workspaceId: 'ws-1',
+      docId: 'doc-1',
+      session: { sessionId: 'session-1' },
+      chatContextValue: { messages: [] },
+      updateContext(context: { messages: unknown[] }) {
+        this.chatContextValue = {
+          ...this.chatContextValue,
+          ...context,
+        };
+      },
+    };
+
+    await (AIChatContent.prototype as any).updateHistory.call(content);
+    expect(
+      content.chatContextValue.messages.map((message: any) => message.id)
+    ).toEqual(['session-1-message']);
+
+    content.session = { sessionId: 'session-2' };
+    await (AIChatContent.prototype as any).updateHistory.call(content);
+
+    expect(
+      content.chatContextValue.messages.map((message: any) => message.id)
+    ).toEqual(['session-2-message']);
+  });
+
+  test('does not overwrite in-flight optimistic messages when a session is created', async () => {
+    const histories = {
+      chats: vi.fn(async () => [{ messages: [] }]),
+      actions: vi.fn(async () => []),
+      cleanup: vi.fn(),
+      ids: vi.fn(),
+    };
+    AIProvider.provide('histories', histories as any);
+
+    const optimisticMessages = [
+      {
+        id: '',
+        role: 'user',
+        content: 'hello',
+        createdAt: '2026-01-01T00:00:00.000Z',
+      },
+      {
+        id: '',
+        role: 'assistant',
+        content: '',
+        createdAt: '2026-01-01T00:00:01.000Z',
+      },
+    ];
+    const updateContext = vi.fn();
+    const content = {
+      updateHistoryCounter: 0,
+      historyKey: 'ws-1:doc-1:',
+      workspaceId: 'ws-1',
+      docId: 'doc-1',
+      session: { sessionId: 'session-1' },
+      chatContextValue: {
+        messages: optimisticMessages,
+        status: 'loading',
+      },
+      updateContext,
+    };
+
+    await (AIChatContent.prototype as any).updateHistory.call(content);
+
+    expect(updateContext).not.toHaveBeenCalled();
+    expect(content.chatContextValue.messages).toBe(optimisticMessages);
+  });
+});
@@ -213,6 +213,8 @@ export class AIChatContent extends SignalWatcher(
  // request counter to track the latest request
  private updateHistoryCounter = 0;

+  private historyKey: string | undefined;
+
  private lastScrollTop: number | undefined;

  get messages() {
@@ -230,13 +232,19 @@ export class AIChatContent extends SignalWatcher(
    return false;
  }

-  private readonly updateHistory = async () => {
+  private async updateHistory() {
    const currentRequest = ++this.updateHistoryCounter;
    if (!AIProvider.histories) {
      return;
    }

    const sessionId = this.session?.sessionId;
+    const nextHistoryKey = `${this.workspaceId}:${this.docId ?? ''}:${
+      sessionId ?? ''
+    }`;
+    const previousHistoryKey = this.historyKey;
+    const preserveCurrentMessages = previousHistoryKey === nextHistoryKey;
+    this.historyKey = nextHistoryKey;
    const [histories, actions] = await Promise.all([
      sessionId
        ? AIProvider.histories.chats(this.workspaceId, sessionId)
@@ -251,9 +259,18 @@ export class AIChatContent extends SignalWatcher(
      return;
    }

-    const messages: HistoryMessage[] = this.chatContextValue.messages
-      .slice()
-      .filter(isChatMessage);
+    if (
+      !preserveCurrentMessages &&
+      (this.chatContextValue.status === 'loading' ||
+        this.chatContextValue.status === 'transmitting') &&
+      this.chatContextValue.messages.length
+    ) {
+      return;
+    }
+
+    const messages: HistoryMessage[] = preserveCurrentMessages
+      ? this.chatContextValue.messages.slice().filter(isChatMessage)
+      : [];

    const chatActions = (actions || []) as ChatAction[];
    messages.push(...chatActions);
@@ -267,7 +284,7 @@ export class AIChatContent extends SignalWatcher(
          new Date(a.createdAt).getTime() - new Date(b.createdAt).getTime()
      ),
    });
-  };
+  }

  private readonly updateActions = async () => {
    if (!this.docId || !AIProvider.histories || !this.showActions) {
@@ -337,6 +354,17 @@ export class AIChatContent extends SignalWatcher(
  }

  protected override updated(changedProperties: PropertyValues) {
+    const historySourceChanged =
+      (changedProperties.has('session') &&
+        changedProperties.get('session') !== undefined) ||
+      (changedProperties.has('docId') &&
+        changedProperties.get('docId') !== undefined) ||
+      (changedProperties.has('workspaceId') &&
+        changedProperties.get('workspaceId') !== undefined);
+    if (historySourceChanged) {
+      this.initChatContent().catch(console.error);
+    }
+
    // restore pinned chat scroll position
    if (
      changedProperties.has('host') &&
@@ -32,6 +32,9 @@ export class AIChatTabs extends WithDisposable(ShadowlessElement) {
  @property({ attribute: false })
  accessor activeSessionId: string | undefined;

+  @property({ attribute: false })
+  accessor showDraftTab = false;
+
  @property({ attribute: false })
  accessor onSelectTab!: (sessionId: string) => void;

@@ -129,10 +132,11 @@ export class AIChatTabs extends WithDisposable(ShadowlessElement) {
  `;

  override render() {
-    if (!this.sessions.length) return html``;
+    if (!this.sessions.length && !this.showDraftTab) return html``;
    return html`
      <div class="ai-chat-tabs" data-testid="ai-chat-tabs">
        <div class="tabs-scroll" @wheel=${this._handleWheel}>
+          ${this.showDraftTab ? this._renderDraftTab() : null}
          ${repeat(
            this.sessions,
            session => session.sessionId,
@@ -177,6 +181,19 @@ export class AIChatTabs extends WithDisposable(ShadowlessElement) {
    `;
  }

+  private _renderDraftTab() {
+    return html`
+      <div
+        class="tab"
+        data-active="true"
+        data-testid="ai-chat-draft-tab"
+        title=${DEFAULT_TAB_TITLE}
+      >
+        <span class="tab-title">${DEFAULT_TAB_TITLE}</span>
+      </div>
+    `;
+  }
+
  private readonly _handleSelect = (sessionId: string) => {
    if (sessionId === this.activeSessionId) return;
    this.onSelectTab(sessionId);
@@ -33,6 +33,9 @@ export class AIChatToolbar extends WithDisposable(ShadowlessElement) {
  @property({ attribute: false })
  accessor onNewSession!: () => void;

+  @property({ attribute: false })
+  accessor canCreateNewSession = true;
+
  @property({ attribute: false })
  accessor onTogglePin!: () => Promise<void>;

@@ -97,14 +100,16 @@ export class AIChatToolbar extends WithDisposable(ShadowlessElement) {
    const pinned = this.session?.pinned;
    return html`
      <div class="ai-chat-toolbar">
-        <div
-          class="chat-toolbar-icon"
-          @click=${this.onPlusClick}
-          data-testid="ai-panel-new-chat"
-        >
-          ${PlusIcon()}
-          <affine-tooltip>New Chat</affine-tooltip>
-        </div>
+        ${this.canCreateNewSession
+          ? html` <div
+              class="chat-toolbar-icon"
+              @click=${this.onPlusClick}
+              data-testid="ai-panel-new-chat"
+            >
+              ${PlusIcon()}
+              <affine-tooltip>New Chat</affine-tooltip>
+            </div>`
+          : null}
        <div
          class="chat-toolbar-icon"
          @click=${this.onPinClick}
@@ -202,6 +207,7 @@ export class AIChatToolbar extends WithDisposable(ShadowlessElement) {
        <ai-session-history
          .session=${this.session}
          .workspaceId=${this.workspaceId}
+          .docId=${this.docId}
          .docDisplayConfig=${this.docDisplayConfig}
          .onSessionClick=${this.onSessionClick}
          .onSessionDelete=${this.onSessionDelete}
@@ -159,6 +159,9 @@ export class AISessionHistory extends WithDisposable(ShadowlessElement) {
  @property({ attribute: false })
  accessor workspaceId!: string;

+  @property({ attribute: false })
+  accessor docId: string | undefined;
+
  @property({ attribute: false })
  accessor docDisplayConfig!: DocDisplayConfig;

@@ -179,6 +182,11 @@ export class AISessionHistory extends WithDisposable(ShadowlessElement) {
  @state()
  private accessor sessions: BlockSuitePresets.AIRecentSession[] | undefined;

+  @state()
+  private accessor currentDocSessions:
+    | BlockSuitePresets.AIRecentSession[]
+    | undefined;
+
  @state()
  private accessor loadingMore = false;

@@ -256,6 +264,18 @@ export class AISessionHistory extends WithDisposable(ShadowlessElement) {
    this.loadingMore = false;
  }

+  private async getCurrentDocSessions() {
+    if (!this.docId) {
+      this.currentDocSessions = [];
+      return;
+    }
+    this.currentDocSessions =
+      (await AIProvider.session?.getSessions(this.workspaceId, this.docId, {
+        action: false,
+        fork: false,
+      })) || [];
+  }
+
  private readonly onScroll = () => {
    if (!this.hasMore || this.loadingMore) {
      return;
@@ -271,6 +291,7 @@ export class AISessionHistory extends WithDisposable(ShadowlessElement) {
  override connectedCallback() {
    super.connectedCallback();
    this.selectedSessionId = this.session?.sessionId ?? undefined;
+    this.getCurrentDocSessions().catch(console.error);
    this.getRecentSessions().catch(console.error);
  }

@@ -305,7 +326,11 @@ export class AISessionHistory extends WithDisposable(ShadowlessElement) {
              @click=${(e: MouseEvent) => {
                e.stopPropagation();
                this.selectedSessionId = session.sessionId;
-                this.onSessionClick(session.sessionId);
+                if (session.docId) {
+                  this.onDocClick(session.docId, session.sessionId);
+                } else {
+                  this.onSessionClick(session.sessionId);
+                }
              }}
              aria-selected=${this.selectedSessionId === session.sessionId}
              data-session-id=${session.sessionId}
@@ -374,12 +399,23 @@ export class AISessionHistory extends WithDisposable(ShadowlessElement) {
      return this.renderLoading();
    }

-    if (this.sessions.length === 0) {
+    const currentDocSessions = this.currentDocSessions ?? [];
+    const currentDocSessionIds = new Set(
+      currentDocSessions.map(session => session.sessionId)
+    );
+    const otherSessions = this.sessions.filter(
+      session =>
+        !currentDocSessionIds.has(session.sessionId) &&
+        (!this.docId || session.docId !== this.docId)
+    );
+
+    if (currentDocSessions.length === 0 && otherSessions.length === 0) {
      return this.renderEmpty();
    }

-    const groupedSessions = this.groupSessionsByTime(this.sessions);
+    const groupedSessions = this.groupSessionsByTime(otherSessions);
    return html`
+      ${this.renderSessionGroup('Current document', currentDocSessions)}
      ${this.renderSessionGroup('Today', groupedSessions.today)}
      ${this.renderSessionGroup('Last 7 days', groupedSessions.last7Days)}
      ${this.renderSessionGroup('Last 30 days', groupedSessions.last30Days)}
@@ -13,6 +13,7 @@ export type ConfigureAIChatToolbarOptions = {
  docDisplayConfig: DocDisplayConfig;
  notificationService: NotificationService;
  onNewSession: () => void;
+  canCreateNewSession?: boolean;
  onTogglePin: () => Promise<void>;
  onOpenSession: (sessionId: string) => void;
  onOpenDoc: (docId: string, sessionId: string) => void;
@@ -36,6 +37,7 @@ export function configureAIChatToolbar(
  tool.docDisplayConfig = options.docDisplayConfig;
  tool.notificationService = options.notificationService;
  tool.onNewSession = options.onNewSession;
+  tool.canCreateNewSession = options.canCreateNewSession ?? true;
  tool.onTogglePin = options.onTogglePin;
  tool.onOpenSession = options.onOpenSession;
  tool.onOpenDoc = options.onOpenDoc;