From a9e4690d2dc45442b5e95f79595516076c7afb4d Mon Sep 17 00:00:00 2001
From: EYHN <cneyhn@gmail.com>
Date: Thu, 17 Apr 2025 13:31:53 +0000
Subject: [PATCH] fix(core): fix ai with multiple server (#11570)

---
 .../blocksuite/ai/provider/setup-provider.tsx | 25 ++++++++++++++++++-
 .../providers/workspace-side-effects.tsx      |  6 ++++-
 .../core/src/modules/cloud/services/auth.ts   | 20 +--------------
 3 files changed, 30 insertions(+), 21 deletions(-)

diff --git a/packages/frontend/core/src/blocksuite/ai/provider/setup-provider.tsx b/packages/frontend/core/src/blocksuite/ai/provider/setup-provider.tsx
index 62f4fedfd2..7d8b21aead 100644
--- a/packages/frontend/core/src/blocksuite/ai/provider/setup-provider.tsx
+++ b/packages/frontend/core/src/blocksuite/ai/provider/setup-provider.tsx
@@ -1,4 +1,5 @@
 import { toggleGeneralAIOnboarding } from '@affine/core/components/affine/ai-onboarding/apis';
+import type { AuthAccountInfo, AuthService } from '@affine/core/modules/cloud';
 import type { GlobalDialogService } from '@affine/core/modules/dialogs';
 import {
   type ChatHistoryOrder,
@@ -14,6 +15,16 @@ import type { PromptKey } from './prompt';
 import { textToText, toImage } from './request';
 import { setupTracker } from './tracker';
 
+function toAIUserInfo(account: AuthAccountInfo | null) {
+  if (!account) return null;
+  return {
+    avatarUrl: account.avatar ?? '',
+    email: account.email ?? '',
+    id: account.id,
+    name: account.label,
+  };
+}
+
 const filterStyleToPromptName = new Map<string, PromptKey>(
   Object.entries({
     'Clay style': 'workflow:image-clay',
@@ -33,7 +44,8 @@ const processTypeToPromptName = new Map<string, PromptKey>(
 
 export function setupAIProvider(
   client: CopilotClient,
-  globalDialogService: GlobalDialogService
+  globalDialogService: GlobalDialogService,
+  authService: AuthService
 ) {
   async function createSession({
     workspaceId,
@@ -58,6 +70,16 @@ export function setupAIProvider(
     });
   }
 
+  AIProvider.provide('userInfo', () => {
+    return toAIUserInfo(authService.session.account$.value);
+  });
+
+  const accountSubscription = authService.session.account$.subscribe(
+    account => {
+      AIProvider.slots.userInfo.next(toAIUserInfo(account));
+    }
+  );
+
   //#region actions
   AIProvider.provide('chat', async options => {
     const { input, contexts, attachments, networkSearch, retry } = options;
@@ -774,5 +796,6 @@ Could you make a new website based on these notes and send back just the html fi
 
   return () => {
     disposeRequestLoginHandler.unsubscribe();
+    accountSubscription.unsubscribe();
   };
 }
diff --git a/packages/frontend/core/src/components/providers/workspace-side-effects.tsx b/packages/frontend/core/src/components/providers/workspace-side-effects.tsx
index 59052a7e28..71bb110ef7 100644
--- a/packages/frontend/core/src/components/providers/workspace-side-effects.tsx
+++ b/packages/frontend/core/src/components/providers/workspace-side-effects.tsx
@@ -12,6 +12,7 @@ import { useRegisterFindInPageCommands } from '@affine/core/components/hooks/aff
 import { useRegisterWorkspaceCommands } from '@affine/core/components/hooks/use-register-workspace-commands';
 import { OverCapacityNotification } from '@affine/core/components/over-capacity';
 import {
+  AuthService,
   EventSourceService,
   FetchService,
   GraphQLService,
@@ -148,6 +149,7 @@ export const WorkspaceSideEffects = () => {
   const graphqlService = useService(GraphQLService);
   const eventSourceService = useService(EventSourceService);
   const fetchService = useService(FetchService);
+  const authService = useService(AuthService);
 
   useEffect(() => {
     const dispose = setupAIProvider(
@@ -156,7 +158,8 @@ export const WorkspaceSideEffects = () => {
         fetchService.fetch,
         eventSourceService.eventSource
       ),
-      globalDialogService
+      globalDialogService,
+      authService
     );
     return () => {
       dispose();
@@ -167,6 +170,7 @@ export const WorkspaceSideEffects = () => {
     workspaceDialogService,
     graphqlService,
     globalDialogService,
+    authService,
   ]);
 
   useRegisterWorkspaceCommands();
diff --git a/packages/frontend/core/src/modules/cloud/services/auth.ts b/packages/frontend/core/src/modules/cloud/services/auth.ts
index 61e9b0dd71..1c417b3644 100644
--- a/packages/frontend/core/src/modules/cloud/services/auth.ts
+++ b/packages/frontend/core/src/modules/cloud/services/auth.ts
@@ -1,4 +1,3 @@
-import { AIProvider } from '@affine/core/blocksuite/ai';
 import { UserFriendlyError } from '@affine/error';
 import type { OAuthProviderType } from '@affine/graphql';
 import { track } from '@affine/track';
@@ -8,7 +7,7 @@ import { distinctUntilChanged, map, skip } from 'rxjs';
 
 import { ApplicationFocused } from '../../lifecycle';
 import type { UrlService } from '../../url';
-import { type AuthAccountInfo, AuthSession } from '../entities/session';
+import { AuthSession } from '../entities/session';
 import { AccountChanged } from '../events/account-changed';
 import { AccountLoggedIn } from '../events/account-logged-in';
 import { AccountLoggedOut } from '../events/account-logged-out';
@@ -16,16 +15,6 @@ import { ServerStarted } from '../events/server-started';
 import type { AuthStore } from '../stores/auth';
 import type { FetchService } from './fetch';
 
-function toAIUserInfo(account: AuthAccountInfo | null) {
-  if (!account) return null;
-  return {
-    avatarUrl: account.avatar ?? '',
-    email: account.email ?? '',
-    id: account.id,
-    name: account.label,
-  };
-}
-
 @OnEvent(ApplicationFocused, e => e.onApplicationFocused)
 @OnEvent(ServerStarted, e => e.onServerStarted)
 export class AuthService extends Service {
@@ -38,11 +27,6 @@ export class AuthService extends Service {
   ) {
     super();
 
-    // TODO(@forehalo): make AIProvider a standalone service passed to AI elements by props
-    AIProvider.provide('userInfo', () => {
-      return toAIUserInfo(this.session.account$.value);
-    });
-
     this.session.account$
       .pipe(
         map(a => ({
@@ -53,8 +37,6 @@ export class AuthService extends Service {
         skip(1) // skip the initial value
       )
       .subscribe(({ account }) => {
-        AIProvider.slots.userInfo.next(toAIUserInfo(account));
-
         if (account === null) {
           this.eventBus.emit(AccountLoggedOut, account);
         } else {