From 91ee5e05bbc7855d48fa584593c6358a08c27cbc Mon Sep 17 00:00:00 2001
From: pengx17 <pengxiao@outlook.com>
Date: Tue, 30 Apr 2024 15:08:26 +0000
Subject: [PATCH] fix: ai chat session handling (#6751)

fix AFF-999
upstream: https://github.com/toeverything/blocksuite/pull/6932
---
 .../block-suite-editor/ai/provider.tsx        | 23 +++++++++++++++----
 .../core/src/modules/cloud/services/auth.ts   | 20 +++++++++-------
 2 files changed, 30 insertions(+), 13 deletions(-)

diff --git a/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/provider.tsx b/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/provider.tsx
index a0ec359c09..ae16126822 100644
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/provider.tsx
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/provider.tsx
@@ -3,6 +3,7 @@ import { authAtom, openSettingModalAtom } from '@affine/core/atoms';
 import { mixpanel } from '@affine/core/utils';
 import { getBaseUrl } from '@affine/graphql';
 import { Trans } from '@affine/i18n';
+import { UnauthorizedError } from '@blocksuite/blocks';
 import { assertExists } from '@blocksuite/global/utils';
 import { AIProvider } from '@blocksuite/presets';
 import { getCurrentStore } from '@toeverything/infra';
@@ -71,11 +72,17 @@ const provideAction = <T extends AIAction>(
 
 export function setupAIProvider() {
   // a single workspace should have only a single chat session
-  // workspace-id:doc-id -> chat session id
+  // user-id:workspace-id:doc-id -> chat session id
   const chatSessions = new Map<string, Promise<string>>();
 
   async function getChatSessionId(workspaceId: string, docId: string) {
-    const storeKey = `${workspaceId}:${docId}`;
+    const userId = (await AIProvider.userInfo)?.id;
+
+    if (!userId) {
+      throw new UnauthorizedError();
+    }
+
+    const storeKey = `${userId}:${workspaceId}:${docId}`;
     if (!chatSessions.has(storeKey)) {
       chatSessions.set(
         storeKey,
@@ -85,9 +92,15 @@ export function setupAIProvider() {
         })
       );
     }
-    const sessionId = await chatSessions.get(storeKey);
-    assertExists(sessionId);
-    return sessionId;
+    try {
+      const sessionId = await chatSessions.get(storeKey);
+      assertExists(sessionId);
+      return sessionId;
+    } catch (err) {
+      // do not cache the error
+      chatSessions.delete(storeKey);
+      throw err;
+    }
   }
 
   //#region actions
diff --git a/packages/frontend/core/src/modules/cloud/services/auth.ts b/packages/frontend/core/src/modules/cloud/services/auth.ts
index cc7ae434b6..e469b97b11 100644
--- a/packages/frontend/core/src/modules/cloud/services/auth.ts
+++ b/packages/frontend/core/src/modules/cloud/services/auth.ts
@@ -24,6 +24,16 @@ export const AccountLoggedIn = createEvent<AuthAccountInfo>('AccountLoggedIn');
 export const AccountLoggedOut =
   createEvent<AuthAccountInfo>('AccountLoggedOut');
 
+function toAIUserInfo(account: AuthAccountInfo | null) {
+  if (!account) return null;
+  return {
+    avatarUrl: account.avatar ?? '',
+    email: account.email ?? '',
+    id: account.id,
+    name: account.label,
+  };
+}
+
 @OnEvent(ApplicationStarted, e => e.onApplicationStart)
 @OnEvent(ApplicationFocused, e => e.onApplicationFocused)
 export class AuthService extends Service {
@@ -36,14 +46,7 @@ export class AuthService extends Service {
     super();
 
     AIProvider.provide('userInfo', () => {
-      const account = this.session.account$.value;
-      if (!account) return null;
-      return {
-        avatarUrl: account.avatar ?? '',
-        email: account.email ?? '',
-        id: account.id,
-        name: account.label,
-      };
+      return toAIUserInfo(this.session.account$.value);
     });
 
     this.session.account$
@@ -62,6 +65,7 @@ export class AuthService extends Service {
           this.eventBus.emit(AccountLoggedIn, account);
         }
         this.eventBus.emit(AccountChanged, account);
+        AIProvider.slots.userInfo.emit(toAIUserInfo(account));
       });
   }