diff --git a/packages/backend/server/src/__tests__/nestjs/throttler.spec.ts b/packages/backend/server/src/__tests__/nestjs/throttler.spec.ts index e8f6aac481..da4330b9d9 100644 --- a/packages/backend/server/src/__tests__/nestjs/throttler.spec.ts +++ b/packages/backend/server/src/__tests__/nestjs/throttler.spec.ts @@ -85,6 +85,18 @@ class NonThrottledController { } } +@UseGuards(CloudThrottlerGuard) +@Throttle('strict') +@Controller('/strict-throttled') +class StrictThrottledController { + @Public() + @SkipThrottle() + @Get('/skip') + skip() { + return 'skip'; + } +} + test.before(async t => { const app = await createTestingApp({ imports: [ @@ -100,7 +112,11 @@ test.before(async t => { }), AppModule, ], - controllers: [ThrottledController, NonThrottledController], + controllers: [ + ThrottledController, + NonThrottledController, + StrictThrottledController, + ], }); t.context.storage = app.get(ThrottlerStorage); @@ -244,6 +260,18 @@ test('should skip throttler for unauthenticated user when specified', async t => t.is(headers.reset, undefined!); }); +test('should skip class-level strict throttler when specified', async t => { + const { app } = t.context; + + const res = await app.GET('/strict-throttled/skip').expect(200); + + const headers = rateLimitHeaders(res); + + t.is(headers.limit, undefined!); + t.is(headers.remaining, undefined!); + t.is(headers.reset, undefined!); +}); + test('should use specified throttler for unauthenticated user', async t => { const { app } = t.context; diff --git a/packages/backend/server/src/base/throttler/decorators.ts b/packages/backend/server/src/base/throttler/decorators.ts index fe3da6c490..fa29a41df3 100644 --- a/packages/backend/server/src/base/throttler/decorators.ts +++ b/packages/backend/server/src/base/throttler/decorators.ts @@ -1,5 +1,8 @@ import { applyDecorators, SetMetadata } from '@nestjs/common'; -import { SkipThrottle, Throttle as RawThrottle } from '@nestjs/throttler'; +import { + SkipThrottle as RawSkipThrottle, + Throttle as RawThrottle, +} from '@nestjs/throttler'; import { ThrottlerType } from './config'; @@ -38,4 +41,11 @@ export function Throttle( ); } -export { SkipThrottle }; +export function SkipThrottle( + skip: Partial> = { + default: true, + strict: true, + } +): MethodDecorator & ClassDecorator { + return RawSkipThrottle(skip); +} diff --git a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-content/ai-chat-content.spec.ts b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-content/ai-chat-content.spec.ts index 47c3b2db30..fd0fb52ef5 100644 --- a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-content/ai-chat-content.spec.ts +++ b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-content/ai-chat-content.spec.ts @@ -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); + }); +}); diff --git a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-content/ai-chat-content.ts b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-content/ai-chat-content.ts index 1849d261c6..b6216d92af 100644 --- a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-content/ai-chat-content.ts +++ b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-content/ai-chat-content.ts @@ -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') && diff --git a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-toolbar/ai-chat-tabs.ts b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-toolbar/ai-chat-tabs.ts index dc17d3e15d..277de4481a 100644 --- a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-toolbar/ai-chat-tabs.ts +++ b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-toolbar/ai-chat-tabs.ts @@ -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`
+ ${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` +
+ ${DEFAULT_TAB_TITLE} +
+ `; + } + private readonly _handleSelect = (sessionId: string) => { if (sessionId === this.activeSessionId) return; this.onSelectTab(sessionId); diff --git a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-toolbar/ai-chat-toolbar.ts b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-toolbar/ai-chat-toolbar.ts index 06c3e5a41b..e4757abbcf 100644 --- a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-toolbar/ai-chat-toolbar.ts +++ b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-toolbar/ai-chat-toolbar.ts @@ -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; @@ -97,14 +100,16 @@ export class AIChatToolbar extends WithDisposable(ShadowlessElement) { const pinned = this.session?.pinned; return html`
-
- ${PlusIcon()} - New Chat -
+ ${this.canCreateNewSession + ? html`
+ ${PlusIcon()} + New Chat +
` + : null}
{ 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)} diff --git a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-toolbar/configure-ai-chat-toolbar.ts b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-toolbar/configure-ai-chat-toolbar.ts index 5194de06aa..ba89b6fd4c 100644 --- a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-toolbar/configure-ai-chat-toolbar.ts +++ b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-toolbar/configure-ai-chat-toolbar.ts @@ -13,6 +13,7 @@ export type ConfigureAIChatToolbarOptions = { docDisplayConfig: DocDisplayConfig; notificationService: NotificationService; onNewSession: () => void; + canCreateNewSession?: boolean; onTogglePin: () => Promise; 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; diff --git a/packages/frontend/core/src/desktop/pages/workspace/detail-page/tabs/chat-panel-session.spec.ts b/packages/frontend/core/src/desktop/pages/workspace/detail-page/tabs/chat-panel-session.spec.ts index 9c9ebe1457..d37ade8aa4 100644 --- a/packages/frontend/core/src/desktop/pages/workspace/detail-page/tabs/chat-panel-session.spec.ts +++ b/packages/frontend/core/src/desktop/pages/workspace/detail-page/tabs/chat-panel-session.spec.ts @@ -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[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 }) diff --git a/packages/frontend/core/src/desktop/pages/workspace/detail-page/tabs/chat-panel-session.ts b/packages/frontend/core/src/desktop/pages/workspace/detail-page/tabs/chat-panel-session.ts index 569dc44240..4875a5666b 100644 --- a/packages/frontend/core/src/desktop/pages/workspace/detail-page/tabs/chat-panel-session.ts +++ b/packages/frontend/core/src/desktop/pages/workspace/detail-page/tabs/chat-panel-session.ts @@ -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 = ( + sessions: T[], + docId?: string | null +) => { + return sessions.filter(session => + isSessionAvailableInDocPanel(session, docId) + ); +}; + +export const hasSessionMessages = ( + session?: Pick | null +) => { + return !!session?.messages?.length; +}; + +export const canCreateNewDocPanelSession = ({ + hasContextMessages, + session, + status, +}: { + hasContextMessages: boolean; + session?: Pick | null; + status?: string | null; +}) => { + return ( + (hasContextMessages || hasSessionMessages(session)) && + status !== 'loading' && + status !== 'transmitting' + ); +}; + export const getSessionIdFromUrl = (workbench?: WorkbenchLike | null) => { if (!workbench) { return undefined; diff --git a/packages/frontend/core/src/desktop/pages/workspace/detail-page/tabs/chat.tsx b/packages/frontend/core/src/desktop/pages/workspace/detail-page/tabs/chat.tsx index 825484296e..7f0d6b4218 100644 --- a/packages/frontend/core/src/desktop/pages/workspace/detail-page/tabs/chat.tsx +++ b/packages/frontend/core/src/desktop/pages/workspace/detail-page/tabs/chat.tsx @@ -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(null); const lastDocIdRef = useRef(null); const sessionLoadSeqRef = useRef(0); + const creatingSessionRef = useRef<{ + docId: string; + promise: Promise; + } | null>(null); + const creatingFreshSessionRef = useRef<{ + docId: string; + promise: Promise; + } | null>(null); const userIdRef = useRef(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(loadSession); + const visibleOpenTabs = useMemo( + () => filterDocPanelTabs(openTabs, doc?.id), + [doc?.id, openTabs] + ); + const canCreateNewSession = canCreateNewDocPanelSession({ + hasContextMessages, + session, + status, + }); const appSidebarConfig = useMemo(() => { 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; + 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; + 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) => { - 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) {