feat: cleanup chat panel (#14259)

#### PR Dependency Tree * **PR #14258** * **PR #14259** 👈 This tree was auto-generated by [Charcoal](https://github.com/danerwilliams/charcoal)  ## Summary by CodeRabbit * **Refactor** * Split AI initialization into separate editor, app, and shared registries; removed legacy chat-panel and replaced it with a component-based editor chat, updating wiring and public exports. * Propagated server/subscription/model services into chat/playground components and improved session lifecycle and UI composition. * **Tests** * Added tests for AI effect registration and chat session resolution; extended DOM/test utilities and assertions. * **Chores** * Added happy-dom for runtime and test environments. <sub>✏️ Tip: You can customize this high-level summary in your review settings.</sub>
2026-07-21 20:16:26 +08:00 · 2026-01-26 04:02:07 +08:00
parent 02449026b9
commit 09cc2dceda
35 changed files with 1408 additions and 1059 deletions
@@ -21,6 +21,7 @@ import {
  EventSourceService,
  FetchService,
  GraphQLService,
+  ServerService,
  SubscriptionService,
 } from '@affine/core/modules/cloud';
 import { WorkspaceDialogService } from '@affine/core/modules/dialogs';
@@ -228,6 +229,7 @@ export const Component = () => {
    );
    content.aiDraftService = framework.get(AIDraftService);
    content.aiToolsConfigService = framework.get(AIToolsConfigService);
+    content.serverService = framework.get(ServerService);
    content.subscriptionService = framework.get(SubscriptionService);
    content.aiModelService = framework.get(AIModelService);
    content.onAISubscribe = handleAISubscribe;
@@ -1,6 +1,6 @@
 import { Scrollable } from '@affine/component';
 import { PageDetailLoading } from '@affine/component/page-detail-skeleton';
-import type { AIChatParams, ChatPanel } from '@affine/core/blocksuite/ai';
+import type { AIChatParams } from '@affine/core/blocksuite/ai';
 import { AIProvider } from '@affine/core/blocksuite/ai';
 import type { AffineEditorContainer } from '@affine/core/blocksuite/block-suite-editor';
 import { EditorOutlineViewer } from '@affine/core/blocksuite/outline-viewer';
@@ -103,7 +103,6 @@ const DetailPageImpl = memo(function DetailPageImpl() {

  const isSideBarOpen = useLiveData(workbench.sidebarOpen$);
  const { appSettings } = useAppSettingHelper();
-  const chatPanelRef = useRef<ChatPanel | null>(null);

  const peekView = useService(PeekViewService).peekView;

@@ -373,7 +372,7 @@ const DetailPageImpl = memo(function DetailPageImpl() {
          icon={<AiIcon />}
          unmountOnInactive={false}
        >
-          <EditorChatPanel editor={editorContainer} ref={chatPanelRef} />
+          <EditorChatPanel editor={editorContainer} />
        </ViewSidebarTab>
      )}

@@ -0,0 +1,147 @@
+/* eslint-disable rxjs/finnish */
+import type { CopilotChatHistoryFragment } from '@affine/graphql';
+import { describe, expect, test, vi } from 'vitest';
+
+import {
+  resolveInitialSession,
+  type SessionService,
+  type WorkbenchLike,
+} from './chat-panel-session';
+
+const createWorkbench = (search: string) => {
+  const updateQueryString = vi.fn();
+  const workbench = {
+    location$: { value: { search } },
+    activeView$: { value: { updateQueryString } },
+  } satisfies WorkbenchLike;
+
+  return { workbench, updateQueryString };
+};
+
+const doc = { id: 'doc-1', workspace: { id: 'ws-1' } };
+
+test('returns undefined without session service or doc', async () => {
+  await expect(
+    resolveInitialSession({ sessionService: null, doc, workbench: null })
+  ).resolves.toBeUndefined();
+  await expect(
+    resolveInitialSession({
+      sessionService: {
+        getSessions: vi.fn(),
+        getSession: vi.fn(),
+      },
+      doc: null,
+      workbench: null,
+    })
+  ).resolves.toBeUndefined();
+});
+
+describe('resolveInitialSession', () => {
+  test('prefers pinned session and clears sessionId from url', async () => {
+    const pinnedSession = {
+      sessionId: 'pinned-session',
+      pinned: true,
+    } as CopilotChatHistoryFragment;
+
+    const sessionService: SessionService = {
+      getSessions: vi.fn().mockResolvedValueOnce([pinnedSession]),
+      getSession: vi.fn(),
+    };
+
+    const { workbench, updateQueryString } = createWorkbench(
+      '?sessionId=from-url'
+    );
+
+    const result = await resolveInitialSession({
+      sessionService,
+      doc,
+      workbench,
+    });
+
+    expect(result).toBe(pinnedSession);
+    expect(updateQueryString).toHaveBeenCalledWith(
+      { sessionId: undefined },
+      { replace: true }
+    );
+    expect(sessionService.getSession).not.toHaveBeenCalled();
+  });
+
+  test('loads session from url when no pinned session', async () => {
+    const sessionFromUrl = {
+      sessionId: 'url-session',
+      pinned: false,
+    } as CopilotChatHistoryFragment;
+
+    const sessionService: SessionService = {
+      getSessions: vi.fn().mockResolvedValueOnce([]),
+      getSession: vi.fn().mockResolvedValueOnce(sessionFromUrl),
+    };
+
+    const { workbench, updateQueryString } = createWorkbench(
+      '?sessionId=url-session'
+    );
+
+    const result = await resolveInitialSession({
+      sessionService,
+      doc,
+      workbench,
+    });
+
+    expect(result).toBe(sessionFromUrl);
+    expect(sessionService.getSession).toHaveBeenCalledWith(
+      doc.workspace.id,
+      'url-session'
+    );
+    expect(updateQueryString).toHaveBeenCalledWith(
+      { sessionId: undefined },
+      { replace: true }
+    );
+  });
+
+  test('falls back to latest doc session', async () => {
+    const docSession = {
+      sessionId: 'doc-session',
+      pinned: false,
+    } as CopilotChatHistoryFragment;
+
+    const sessionService: SessionService = {
+      getSessions: vi
+        .fn()
+        .mockResolvedValueOnce([])
+        .mockResolvedValueOnce([docSession]),
+      getSession: vi.fn(),
+    };
+
+    const { workbench } = createWorkbench('');
+
+    const result = await resolveInitialSession({
+      sessionService,
+      doc,
+      workbench,
+    });
+
+    expect(result).toBe(docSession);
+    expect(sessionService.getSessions).toHaveBeenCalledWith(
+      doc.workspace.id,
+      doc.id,
+      { action: false, fork: false, limit: 1 }
+    );
+  });
+
+  test('returns null when url session is missing', async () => {
+    const sessionService: SessionService = {
+      getSessions: vi.fn().mockResolvedValueOnce([]),
+      getSession: vi.fn().mockResolvedValueOnce(null),
+    };
+
+    const { workbench } = createWorkbench('?sessionId=missing');
+
+    const result = await resolveInitialSession({
+      sessionService,
+      doc,
+      workbench,
+    });
+
+    expect(result).toBeNull();
+  });
+});
@@ -0,0 +1,108 @@
+/* eslint-disable rxjs/finnish */
+import type { CopilotChatHistoryFragment } from '@affine/graphql';
+
+type SessionListOptions = {
+  pinned?: boolean;
+  action?: boolean;
+  fork?: boolean;
+  limit?: number;
+};
+
+export interface SessionService {
+  getSessions: (
+    workspaceId: string,
+    docId?: string,
+    options?: SessionListOptions
+  ) => Promise<CopilotChatHistoryFragment[] | null | undefined>;
+  getSession: (
+    workspaceId: string,
+    sessionId: string
+  ) => Promise<CopilotChatHistoryFragment | null | undefined>;
+}
+
+export interface WorkbenchLike {
+  location$: {
+    value: {
+      search: string;
+    };
+  };
+  activeView$: {
+    value: {
+      updateQueryString: (
+        patch: Record<string, unknown>,
+        options?: { replace?: boolean }
+      ) => void;
+    };
+  };
+}
+
+export interface DocLike {
+  id: string;
+  workspace: {
+    id: string;
+  };
+}
+
+export const getSessionIdFromUrl = (workbench?: WorkbenchLike | null) => {
+  if (!workbench) {
+    return undefined;
+  }
+  const searchParams = new URLSearchParams(workbench.location$.value.search);
+  const sessionId = searchParams.get('sessionId');
+  if (sessionId) {
+    workbench.activeView$.value.updateQueryString(
+      { sessionId: undefined },
+      { replace: true }
+    );
+  }
+  return sessionId ?? undefined;
+};
+
+export const resolveInitialSession = async ({
+  sessionService,
+  doc,
+  workbench,
+}: {
+  sessionService?: SessionService | null;
+  doc?: DocLike | null;
+  workbench?: WorkbenchLike | null;
+}): Promise<CopilotChatHistoryFragment | null | undefined> => {
+  if (!sessionService || !doc) {
+    return undefined;
+  }
+
+  const sessionId = getSessionIdFromUrl(workbench);
+
+  const pinSessions = await sessionService.getSessions(
+    doc.workspace.id,
+    undefined,
+    {
+      pinned: true,
+      limit: 1,
+    }
+  );
+
+  if (Array.isArray(pinSessions) && pinSessions[0]) {
+    return pinSessions[0];
+  }
+
+  if (sessionId) {
+    const session = await sessionService.getSession(
+      doc.workspace.id,
+      sessionId
+    );
+    return session ?? null;
+  }
+
+  const docSessions = await sessionService.getSessions(
+    doc.workspace.id,
+    doc.id,
+    {
+      action: false,
+      fork: false,
+      limit: 1,
+    }
+  );
+
+  return docSessions?.[0] ?? null;
+};
@@ -1,6 +1,100 @@
-import { style } from '@vanilla-extract/css';
+import { cssVarV2 } from '@toeverything/theme/v2';
+import { globalStyle, style } from '@vanilla-extract/css';
+
 export const root = style({
  display: 'flex',
+  flexDirection: 'column',
  height: '100%',
  width: '100%',
+  userSelect: 'text',
+});
+
+export const container = style({
+  display: 'flex',
+  flexDirection: 'column',
+  height: '100%',
+});
+
+export const header = style({
+  background: 'var(--affine-background-primary-color)',
+  position: 'relative',
+  padding: '8px var(--h-padding, 16px)',
+  width: '100%',
+  height: '36px',
+  display: 'flex',
+  alignItems: 'center',
+  justifyContent: 'space-between',
+  zIndex: 1,
+});
+
+export const title = style({
+  fontSize: '14px',
+  fontWeight: 500,
+  color: 'var(--affine-text-secondary-color)',
+  display: 'flex',
+  alignItems: 'center',
+});
+
+export const playground = style({
+  cursor: 'pointer',
+  padding: '2px',
+  marginLeft: '8px',
+  marginRight: 'auto',
+  display: 'flex',
+  alignItems: 'center',
+  justifyContent: 'center',
+});
+
+export const content = style({
+  flexGrow: 1,
+  height: 0,
+  minHeight: 0,
+  display: 'flex',
+  flexDirection: 'column',
+  width: '100%',
+});
+
+export const loadingContainer = style({
+  position: 'relative',
+  padding: '44px 0 166px 0',
+  height: '100%',
+  display: 'flex',
+  alignItems: 'center',
+});
+
+export const loading = style({
+  width: '100%',
+  display: 'flex',
+  flexDirection: 'column',
+  alignItems: 'center',
+  gap: '12px',
+});
+
+export const loadingTitle = style({
+  fontWeight: 600,
+  fontSize: 'var(--affine-font-sm)',
+  color: 'var(--affine-text-secondary-color)',
+});
+
+export const loadingIcon = style({
+  width: '44px',
+  height: '44px',
+  color: 'var(--affine-icon-secondary)',
+});
+
+globalStyle(`${playground} svg`, {
+  width: '18px',
+  height: '18px',
+  color: 'var(--affine-text-secondary-color)',
+});
+
+globalStyle(`${playground}:hover svg`, {
+  color: cssVarV2('icon/activated'),
+});
+
+globalStyle(`${content} > ai-chat-content`, {
+  flexGrow: 1,
+  height: 0,
+  minHeight: 0,
+  width: '100%',
 });
@@ -1,5 +1,14 @@
 import { useConfirmModal } from '@affine/component';
-import { AIProvider, ChatPanel } from '@affine/core/blocksuite/ai';
+import { AIProvider } from '@affine/core/blocksuite/ai';
+import type { AppSidebarConfig } from '@affine/core/blocksuite/ai/chat-panel/chat-config';
+import {
+  AIChatContent,
+  type ChatContextValue,
+} from '@affine/core/blocksuite/ai/components/ai-chat-content';
+import type { ChatStatus } from '@affine/core/blocksuite/ai/components/ai-chat-messages';
+import { AIChatToolbar } from '@affine/core/blocksuite/ai/components/ai-chat-toolbar';
+import { createPlaygroundModal } from '@affine/core/blocksuite/ai/components/playground/modal';
+import { registerAIAppEffects } from '@affine/core/blocksuite/ai/effects/app';
 import type { AffineEditorContainer } from '@affine/core/blocksuite/block-suite-editor';
 import { NotificationServiceImpl } from '@affine/core/blocksuite/view-extensions/editor-view/notification-service';
 import { useAIChatConfig } from '@affine/core/components/hooks/affine/use-ai-chat-config';
@@ -12,48 +21,49 @@ import {
 import { AIModelService } from '@affine/core/modules/ai-button/services/models';
 import { ServerService, SubscriptionService } from '@affine/core/modules/cloud';
 import { WorkspaceDialogService } from '@affine/core/modules/dialogs';
+import { useSignalValue } from '@affine/core/modules/doc-info/utils';
 import { FeatureFlagService } from '@affine/core/modules/feature-flag';
 import { PeekViewService } from '@affine/core/modules/peek-view';
 import { AppThemeService } from '@affine/core/modules/theme';
 import { WorkbenchService } from '@affine/core/modules/workbench';
+import type {
+  ContextEmbedStatus,
+  CopilotChatHistoryFragment,
+  UpdateChatSessionInput,
+} from '@affine/graphql';
 import { RefNodeSlotsProvider } from '@blocksuite/affine/inlines/reference';
 import { DocModeProvider } from '@blocksuite/affine/shared/services';
 import { createSignalFromObservable } from '@blocksuite/affine/shared/utils';
+import { CenterPeekIcon, Logo1Icon } from '@blocksuite/icons/rc';
+import type { Signal } from '@preact/signals-core';
 import { useFramework, useService } from '@toeverything/infra';
-import { forwardRef, useEffect, useRef, useState } from 'react';
+import { html } from 'lit';
+import { useCallback, useEffect, useMemo, useRef, useState } from 'react';

 import * as styles from './chat.css';
+import {
+  resolveInitialSession,
+  type WorkbenchLike,
+} from './chat-panel-session';
+
+registerAIAppEffects();

 export interface SidebarTabProps {
  editor: AffineEditorContainer | null;
  onLoad?: ((component: HTMLElement) => void) | null;
 }

-// A wrapper for CopilotPanel
-export const EditorChatPanel = forwardRef(function EditorChatPanel(
-  { editor, onLoad }: SidebarTabProps,
-  ref: React.ForwardedRef<ChatPanel>
-) {
-  const chatPanelRef = useRef<ChatPanel | null>(null);
-  const containerRef = useRef<HTMLDivElement | null>(null);
-  const workbench = useService(WorkbenchService).workbench;
+export const EditorChatPanel = ({ editor, onLoad }: SidebarTabProps) => {
  const framework = useFramework();
+  const workbench = useService(WorkbenchService).workbench;

-  useEffect(() => {
-    if (onLoad && chatPanelRef.current) {
-      (chatPanelRef.current as ChatPanel).updateComplete
-        .then(() => {
-          if (ref) {
-            if (typeof ref === 'function') {
-              ref(chatPanelRef.current);
-            } else {
-              ref.current = chatPanelRef.current;
-            }
-          }
-        })
-        .catch(console.error);
-    }
-  }, [onLoad, ref]);
+  const { closeConfirmModal, openConfirmModal } = useConfirmModal();
+  const notificationService = useMemo(
+    () => new NotificationServiceImpl(closeConfirmModal, openConfirmModal),
+    [closeConfirmModal, openConfirmModal]
+  );
+  const specs = useAISpecs();
+  const handleAISubscribe = useAISubscribe();

  const {
    docDisplayConfig,
@@ -61,101 +71,477 @@ export const EditorChatPanel = forwardRef(function EditorChatPanel(
    reasoningConfig,
    playgroundConfig,
  } = useAIChatConfig();
-  const confirmModal = useConfirmModal();
-  const specs = useAISpecs();
-  const handleAISubscribe = useAISubscribe();
+  const playgroundVisible = useSignalValue(playgroundConfig.visible) ?? false;
+
+  const [session, setSession] = useState<
+    CopilotChatHistoryFragment | null | undefined
+  >(undefined);
+  const [embeddingProgress, setEmbeddingProgress] = useState<[number, number]>([
+    0, 0,
+  ]);
+  const [status, setStatus] = useState<ChatStatus>('idle');
+  const [hasPinned, setHasPinned] = useState(false);
+
+  const [chatContent, setChatContent] = useState<AIChatContent | null>(null);
+  const [chatToolbar, setChatToolbar] = useState<AIChatToolbar | null>(null);
+  const [isBodyProvided, setIsBodyProvided] = useState(false);
+  const [isHeaderProvided, setIsHeaderProvided] = useState(false);
+
+  const chatContainerRef = useRef<HTMLDivElement | null>(null);
+  const chatToolbarContainerRef = useRef<HTMLDivElement | null>(null);
+  const contentKeyRef = useRef<string | null>(null);
+  const lastDocIdRef = useRef<string | null>(null);
+
+  const doc = editor?.doc;
+  const host = editor?.host;
+
+  const appSidebarConfig = useMemo<AppSidebarConfig>(() => {
+    return {
+      getWidth: () =>
+        createSignalFromObservable<number | undefined>(
+          workbench.sidebarWidth$.asObservable(),
+          0
+        ),
+      isOpen: () =>
+        createSignalFromObservable<boolean | undefined>(
+          workbench.sidebarOpen$.asObservable(),
+          true
+        ),
+    };
+  }, [workbench]);
+
+  const [sidebarWidthSignal, setSidebarWidthSignal] =
+    useState<Signal<number | undefined>>();

  useEffect(() => {
-    if (!editor || !editor.host) return;
+    const { signal, cleanup } = appSidebarConfig.getWidth();
+    setSidebarWidthSignal(signal);
+    return cleanup;
+  }, [appSidebarConfig]);

-    if (!chatPanelRef.current) {
-      chatPanelRef.current = new ChatPanel();
-      chatPanelRef.current.host = editor.host;
-      chatPanelRef.current.doc = editor.doc;
+  const resetPanel = useCallback(() => {
+    setSession(undefined);
+    setEmbeddingProgress([0, 0]);
+    setHasPinned(false);
+  }, []);

-      const workbench = framework.get(WorkbenchService).workbench;
-      chatPanelRef.current.appSidebarConfig = {
-        getWidth: () => {
-          const width$ = workbench.sidebarWidth$;
-          return createSignalFromObservable(width$, 0);
-        },
-        isOpen: () => {
-          const open$ = workbench.sidebarOpen$;
-          return createSignalFromObservable(open$, true);
-        },
-      };
+  const initPanel = useCallback(async () => {
+    try {
+      const nextSession = await resolveInitialSession({
+        sessionService: AIProvider.session ?? undefined,
+        doc,
+        workbench: workbench as WorkbenchLike,
+      });

-      chatPanelRef.current.docDisplayConfig = docDisplayConfig;
-      chatPanelRef.current.searchMenuConfig = searchMenuConfig;
-      chatPanelRef.current.reasoningConfig = reasoningConfig;
-      chatPanelRef.current.playgroundConfig = playgroundConfig;
-      chatPanelRef.current.extensions = specs;
-      chatPanelRef.current.serverService = framework.get(ServerService);
-      chatPanelRef.current.affineFeatureFlagService =
-        framework.get(FeatureFlagService);
-      chatPanelRef.current.affineWorkspaceDialogService = framework.get(
-        WorkspaceDialogService
+      if (nextSession === undefined) {
+        return;
+      }
+
+      setSession(nextSession);
+      setHasPinned(!!nextSession?.pinned);
+    } catch (error) {
+      console.error(error);
+    }
+  }, [doc, workbench]);
+
+  const createSession = useCallback(
+    async (options: Partial<BlockSuitePresets.AICreateSessionOptions> = {}) => {
+      if (session || !AIProvider.session || !doc) {
+        return session ?? undefined;
+      }
+      const sessionId = await AIProvider.session.createSession({
+        docId: doc.id,
+        workspaceId: doc.workspace.id,
+        promptName: 'Chat With AFFiNE AI',
+        reuseLatestChat: false,
+        ...options,
+      });
+      if (sessionId) {
+        const nextSession = await AIProvider.session.getSession(
+          doc.workspace.id,
+          sessionId
+        );
+        setSession(nextSession ?? null);
+        return nextSession ?? undefined;
+      }
+      return session ?? undefined;
+    },
+    [doc, session]
+  );
+
+  const updateSession = useCallback(
+    async (options: UpdateChatSessionInput) => {
+      if (!AIProvider.session || !doc) {
+        return undefined;
+      }
+      await AIProvider.session.updateSession(options);
+      const nextSession = await AIProvider.session.getSession(
+        doc.workspace.id,
+        options.sessionId
      );
-      chatPanelRef.current.affineWorkbenchService =
-        framework.get(WorkbenchService);
-      chatPanelRef.current.affineThemeService = framework.get(AppThemeService);
-      chatPanelRef.current.peekViewService = framework.get(PeekViewService);
-      chatPanelRef.current.notificationService = new NotificationServiceImpl(
-        confirmModal.closeConfirmModal,
-        confirmModal.openConfirmModal
-      );
-      chatPanelRef.current.aiDraftService = framework.get(AIDraftService);
-      chatPanelRef.current.aiToolsConfigService =
-        framework.get(AIToolsConfigService);
-      chatPanelRef.current.subscriptionService =
-        framework.get(SubscriptionService);
-      chatPanelRef.current.aiModelService = framework.get(AIModelService);
-      chatPanelRef.current.onAISubscribe = handleAISubscribe;
+      setSession(nextSession ?? null);
+      return nextSession ?? undefined;
+    },
+    [doc]
+  );

-      containerRef.current?.append(chatPanelRef.current);
-    } else {
-      chatPanelRef.current.host = editor.host;
-      chatPanelRef.current.doc = editor.doc;
+  const newSession = useCallback(() => {
+    resetPanel();
+    requestAnimationFrame(() => {
+      setSession(null);
+    });
+  }, [resetPanel]);
+
+  const openSession = useCallback(
+    async (sessionId: string) => {
+      if (session?.sessionId === sessionId || !AIProvider.session || !doc) {
+        return;
+      }
+      resetPanel();
+      const nextSession = await AIProvider.session.getSession(
+        doc.workspace.id,
+        sessionId
+      );
+      setSession(nextSession ?? null);
+    },
+    [doc, resetPanel, session?.sessionId]
+  );
+
+  const openDoc = useCallback(
+    async (docId: string, sessionId?: string) => {
+      if (!doc) {
+        return;
+      }
+      if (doc.id === docId) {
+        if (session?.sessionId === sessionId || session?.pinned) {
+          return;
+        }
+        if (sessionId) {
+          await openSession(sessionId);
+        }
+        return;
+      }
+      if (session?.pinned || !sessionId) {
+        workbench.open(`/${docId}`, { at: 'active' });
+        return;
+      }
+      workbench.open(`/${docId}?sessionId=${sessionId}`, { at: 'active' });
+    },
+    [doc, openSession, session?.pinned, session?.sessionId, workbench]
+  );
+
+  const deleteSession = useCallback(
+    async (sessionToDelete: BlockSuitePresets.AIRecentSession) => {
+      if (!AIProvider.histories) {
+        return;
+      }
+      const confirm = await notificationService.confirm({
+        title: 'Delete this history?',
+        message:
+          'Do you want to delete this AI conversation history? Once deleted, it cannot be recovered.',
+        confirmText: 'Delete',
+        cancelText: 'Cancel',
+      });
+      if (confirm) {
+        await AIProvider.histories.cleanup(
+          sessionToDelete.workspaceId,
+          sessionToDelete.docId || undefined,
+          [sessionToDelete.sessionId]
+        );
+        if (sessionToDelete.sessionId === session?.sessionId) {
+          newSession();
+        }
+      }
+    },
+    [newSession, notificationService, session?.sessionId]
+  );
+
+  const togglePin = useCallback(async () => {
+    const pinned = !session?.pinned;
+    setHasPinned(true);
+    if (!session) {
+      await createSession({ pinned });
+      return;
+    }
+    setSession(prev => (prev ? { ...prev, pinned } : prev));
+    await updateSession({
+      sessionId: session.sessionId,
+      pinned,
+    });
+  }, [createSession, session, updateSession]);
+
+  const rebindSession = useCallback(async () => {
+    if (!session || !doc) {
+      return;
+    }
+    if (session.docId !== doc.id) {
+      await updateSession({
+        sessionId: session.sessionId,
+        docId: doc.id,
+      });
+    }
+  }, [doc, session, updateSession]);
+
+  const onEmbeddingProgressChange = useCallback(
+    (count: Record<ContextEmbedStatus, number>) => {
+      const total = count.finished + count.processing + count.failed;
+      setEmbeddingProgress([count.finished, total]);
+    },
+    []
+  );
+
+  const onContextChange = useCallback(
+    (context: Partial<ChatContextValue>) => {
+      setStatus(context.status ?? 'idle');
+      if (context.status === 'success') {
+        rebindSession().catch(console.error);
+      }
+    },
+    [rebindSession]
+  );
+
+  useEffect(() => {
+    if (session !== undefined) {
+      return;
+    }
+    if (chatContent) {
+      chatContent.remove();
+      setChatContent(null);
+    }
+    if (chatToolbar) {
+      chatToolbar.remove();
+      setChatToolbar(null);
+    }
+  }, [chatContent, chatToolbar, session]);
+
+  useEffect(() => {
+    const subscription = AIProvider.slots.userInfo.subscribe(() => {
+      resetPanel();
+      initPanel().catch(console.error);
+    });
+    return () => subscription.unsubscribe();
+  }, [initPanel, resetPanel]);
+
+  useEffect(() => {
+    const docId = doc?.id;
+    if (!docId) {
+      return;
+    }
+    if (
+      lastDocIdRef.current &&
+      lastDocIdRef.current !== docId &&
+      !session?.pinned
+    ) {
+      resetPanel();
+    }
+    lastDocIdRef.current = docId;
+  }, [doc?.id, resetPanel, session?.pinned]);
+
+  useEffect(() => {
+    if (!doc || session !== undefined) {
+      return;
    }

+    let cancelled = false;
+    let timerId: ReturnType<typeof setTimeout> | null = null;
+
+    const tryInit = () => {
+      if (cancelled || session !== undefined) {
+        return;
+      }
+      // Session service may be registered after the panel mounts.
+      if (AIProvider.session) {
+        initPanel().catch(console.error);
+        return;
+      }
+      timerId = setTimeout(tryInit, 200);
+    };
+
+    tryInit();
+
+    return () => {
+      cancelled = true;
+      if (timerId) {
+        clearTimeout(timerId);
+      }
+    };
+  }, [doc, initPanel, session]);
+
+  const contentKey = hasPinned
+    ? (session?.sessionId ?? doc?.id ?? 'chat-panel')
+    : (doc?.id ?? 'chat-panel');
+
+  useEffect(() => {
+    if (!chatContent) {
+      contentKeyRef.current = contentKey;
+      return;
+    }
+    if (contentKeyRef.current && contentKeyRef.current !== contentKey) {
+      chatContent.remove();
+      setChatContent(null);
+    }
+    contentKeyRef.current = contentKey;
+  }, [chatContent, contentKey]);
+
+  useEffect(() => {
+    if (!isBodyProvided || !chatContainerRef.current || !doc || !host) {
+      return;
+    }
+    if (session === undefined) {
+      return;
+    }
+
+    let content = chatContent;
+
+    if (!content) {
+      content = new AIChatContent();
+    }
+
+    content.host = host;
+    content.session = session;
+    content.createSession = createSession;
+    content.workspaceId = doc.workspace.id;
+    content.docId = doc.id;
+    content.reasoningConfig = reasoningConfig;
+    content.searchMenuConfig = searchMenuConfig;
+    content.docDisplayConfig = docDisplayConfig;
+    content.extensions = specs;
+    content.serverService = framework.get(ServerService);
+    content.affineFeatureFlagService = framework.get(FeatureFlagService);
+    content.affineWorkspaceDialogService = framework.get(
+      WorkspaceDialogService
+    );
+    content.affineThemeService = framework.get(AppThemeService);
+    content.notificationService = notificationService;
+    content.aiDraftService = framework.get(AIDraftService);
+    content.aiToolsConfigService = framework.get(AIToolsConfigService);
+    content.peekViewService = framework.get(PeekViewService);
+    content.subscriptionService = framework.get(SubscriptionService);
+    content.aiModelService = framework.get(AIModelService);
+    content.onAISubscribe = handleAISubscribe;
+    content.onEmbeddingProgressChange = onEmbeddingProgressChange;
+    content.onContextChange = onContextChange;
+    content.width = sidebarWidthSignal;
+    content.onOpenDoc = (docId: string, sessionId?: string) => {
+      openDoc(docId, sessionId).catch(console.error);
+    };
+
+    if (!chatContent) {
+      chatContainerRef.current.append(content);
+      setChatContent(content);
+      onLoad?.(content);
+    }
+  }, [
+    chatContent,
+    createSession,
+    doc,
+    docDisplayConfig,
+    framework,
+    handleAISubscribe,
+    host,
+    isBodyProvided,
+    notificationService,
+    onContextChange,
+    onEmbeddingProgressChange,
+    onLoad,
+    openDoc,
+    reasoningConfig,
+    searchMenuConfig,
+    session,
+    sidebarWidthSignal,
+    specs,
+  ]);
+
+  useEffect(() => {
+    if (!isHeaderProvided || !chatToolbarContainerRef.current || !doc) {
+      return;
+    }
+    if (session === undefined) {
+      return;
+    }
+
+    let tool = chatToolbar;
+
+    if (!tool) {
+      tool = new AIChatToolbar();
+    }
+
+    tool.session = session;
+    tool.workspaceId = doc.workspace.id;
+    tool.docId = doc.id;
+    tool.status = status;
+    tool.docDisplayConfig = docDisplayConfig;
+    tool.notificationService = notificationService;
+    tool.onNewSession = newSession;
+    tool.onTogglePin = togglePin;
+    tool.onOpenSession = (sessionId: string) => {
+      openSession(sessionId).catch(console.error);
+    };
+    tool.onOpenDoc = (docId: string, sessionId: string) => {
+      openDoc(docId, sessionId).catch(console.error);
+    };
+    tool.onSessionDelete = (
+      sessionToDelete: BlockSuitePresets.AIRecentSession
+    ) => {
+      deleteSession(sessionToDelete).catch(console.error);
+    };
+
+    if (!chatToolbar) {
+      chatToolbarContainerRef.current.append(tool);
+      setChatToolbar(tool);
+    }
+  }, [
+    chatToolbar,
+    deleteSession,
+    doc,
+    docDisplayConfig,
+    isHeaderProvided,
+    newSession,
+    notificationService,
+    openDoc,
+    openSession,
+    session,
+    status,
+    togglePin,
+  ]);
+
+  useEffect(() => {
+    if (!editor?.host || !chatContent) {
+      return;
+    }
    const docModeService = editor.host.std.get(DocModeProvider);
    const refNodeService = editor.host.std.getOptional(RefNodeSlotsProvider);
    const disposable = [
-      refNodeService?.docLinkClicked.subscribe(({ host }) => {
-        if (host === editor.host) {
-          (chatPanelRef.current as ChatPanel).doc = editor.doc;
+      refNodeService?.docLinkClicked.subscribe(({ host: clickedHost }) => {
+        if (clickedHost === editor.host) {
+          chatContent.docId = editor.doc.id;
        }
      }),
      docModeService?.onPrimaryModeChange(() => {
-        if (!editor.host) return;
-        (chatPanelRef.current as ChatPanel).host = editor.host;
+        if (!editor.host) {
+          return;
+        }
+        chatContent.host = editor.host;
      }, editor.doc.id),
    ];

-    return () => disposable.forEach(d => d?.unsubscribe());
-  }, [
-    docDisplayConfig,
-    editor,
-    framework,
-    searchMenuConfig,
-    reasoningConfig,
-    playgroundConfig,
-    confirmModal,
-    specs,
-    handleAISubscribe,
-  ]);
+    return () => disposable.forEach(item => item?.unsubscribe());
+  }, [chatContent, editor]);

  const [autoResized, setAutoResized] = useState(false);
  useEffect(() => {
-    // after auto expanded first time, do not auto expand again(even if user manually resized)
-    if (autoResized) return;
+    if (autoResized) {
+      return;
+    }
    const subscription = AIProvider.slots.previewPanelOpenChange.subscribe(
      open => {
-        if (!open) return;
+        if (!open) {
+          return;
+        }
        const sidebarWidth = workbench.sidebarWidth$.value;
-        const MIN_SIDEBAR_WIDTH = 1080;
-        if (!sidebarWidth || sidebarWidth < MIN_SIDEBAR_WIDTH) {
-          workbench.setSidebarWidth(MIN_SIDEBAR_WIDTH);
+        const minSidebarWidth = 1080;
+        if (!sidebarWidth || sidebarWidth < minSidebarWidth) {
+          workbench.setSidebarWidth(minSidebarWidth);
          setAutoResized(true);
        }
      }
@@ -165,5 +551,99 @@ export const EditorChatPanel = forwardRef(function EditorChatPanel(
    };
  }, [autoResized, workbench]);

-  return <div className={styles.root} ref={containerRef} />;
-});
+  const openPlayground = useCallback(() => {
+    if (!doc || !host) {
+      return;
+    }
+    const playgroundContent = html`
+      <playground-content
+        .host=${host}
+        .doc=${doc}
+        .reasoningConfig=${reasoningConfig}
+        .playgroundConfig=${playgroundConfig}
+        .appSidebarConfig=${appSidebarConfig}
+        .searchMenuConfig=${searchMenuConfig}
+        .docDisplayConfig=${docDisplayConfig}
+        .extensions=${specs}
+        .serverService=${framework.get(ServerService)}
+        .affineFeatureFlagService=${framework.get(FeatureFlagService)}
+        .affineThemeService=${framework.get(AppThemeService)}
+        .notificationService=${notificationService}
+        .affineWorkspaceDialogService=${framework.get(WorkspaceDialogService)}
+        .aiToolsConfigService=${framework.get(AIToolsConfigService)}
+        .subscriptionService=${framework.get(SubscriptionService)}
+        .aiModelService=${framework.get(AIModelService)}
+      ></playground-content>
+    `;
+
+    createPlaygroundModal(playgroundContent, 'AI Playground');
+  }, [
+    appSidebarConfig,
+    doc,
+    docDisplayConfig,
+    framework,
+    host,
+    notificationService,
+    playgroundConfig,
+    reasoningConfig,
+    searchMenuConfig,
+    specs,
+  ]);
+
+  const onChatContainerRef = useCallback((node: HTMLDivElement) => {
+    if (!node) {
+      return;
+    }
+    setIsBodyProvided(true);
+    chatContainerRef.current = node;
+  }, []);
+
+  const onChatToolContainerRef = useCallback((node: HTMLDivElement) => {
+    if (!node) {
+      return;
+    }
+    setIsHeaderProvided(true);
+    chatToolbarContainerRef.current = node;
+  }, []);
+
+  const isEmbedding =
+    embeddingProgress[1] > 0 && embeddingProgress[0] < embeddingProgress[1];
+  const [done, total] = embeddingProgress;
+  const isInitialized = session !== undefined;
+
+  return (
+    <div className={styles.root}>
+      {!isInitialized ? (
+        <div className={styles.loadingContainer}>
+          <div className={styles.loading}>
+            <Logo1Icon className={styles.loadingIcon} />
+            <div className={styles.loadingTitle}>
+              AFFiNE AI is loading history...
+            </div>
+          </div>
+        </div>
+      ) : (
+        <div className={styles.container}>
+          <div className={styles.header}>
+            <div className={styles.title}>
+              {isEmbedding ? (
+                <span data-testid="chat-panel-embedding-progress">
+                  Embedding {done}/{total}
+                </span>
+              ) : (
+                'AFFiNE AI'
+              )}
+            </div>
+            {playgroundVisible ? (
+              <div className={styles.playground} onClick={openPlayground}>
+                <CenterPeekIcon />
+              </div>
+            ) : null}
+            <div ref={onChatToolContainerRef} />
+          </div>
+          <div className={styles.content} ref={onChatContainerRef} />
+        </div>
+      )}
+    </div>
+  );
+};