feat(core): integrate realtime features (#15003)

packages/frontend/core/src/blocksuite/ai/runtime/request/copilot-client.spec.ts

@@ -31,4 +31,22 @@ describe('CopilotClient action streams', () => {
       '/api/copilot/actions/session-1/stream?messageId=message-1&actionId=mindmap.generate&actionVersion=v1&runId=run-1&retry=true'
     );
   });
+
+  test('loads embedding status through realtime request', async () => {
+    const gql = vi.fn();
+    const request = vi.fn().mockResolvedValue({ total: 4, embedded: 2 });
+    const client = new CopilotClient(gql, vi.fn(), { request });
+
+    await expect(client.getEmbeddingStatus('workspace-1')).resolves.toEqual({
+      total: 4,
+      embedded: 2,
+    });
+
+    expect(gql).not.toHaveBeenCalled();
+    expect(request).toHaveBeenCalledWith(
+      'workspace.embedding.progress.get',
+      { workspaceId: 'workspace-1' },
+      { timeoutMs: 10000 }
+    );
+  });
 });

packages/frontend/core/src/blocksuite/ai/runtime/request/copilot-client.ts

@@ -1,5 +1,6 @@
 import { showAILoginRequiredAtom } from '@affine/core/components/affine/auth/ai-login-required';
 import type { AIToolsConfig } from '@affine/core/modules/ai-button';
+import type { NbstoreService } from '@affine/core/modules/storage';
 import { UserFriendlyError } from '@affine/error';
 import {
   addContextBlobMutation,
@@ -17,7 +18,6 @@ import {
   getCopilotRecentSessionsQuery,
   getCopilotSessionQuery,
   getCopilotSessionsQuery,
-  getWorkspaceEmbeddingStatusQuery,
   type GraphQLQuery,
   listContextObjectQuery,
   listContextQuery,
@@ -98,7 +98,8 @@ export class CopilotClient {
     readonly eventSource: (
       url: string,
       eventSourceInitDict?: EventSourceInit
-    ) => EventSource
+    ) => EventSource,
+    readonly realtime?: Pick<NbstoreService['realtime'], 'request'>
   ) {}
 
   async createSession(
@@ -546,11 +547,15 @@ export class CopilotClient {
     return queryString.toString();
   }
 
-  getEmbeddingStatus(workspaceId: string) {
-    return this.gql({
-      query: getWorkspaceEmbeddingStatusQuery,
-      variables: { workspaceId },
-    }).then(res => res.queryWorkspaceEmbeddingStatus);
+  async getEmbeddingStatus(workspaceId: string) {
+    if (!this.realtime) {
+      throw new Error('Realtime client is required');
+    }
+    return await this.realtime.request(
+      'workspace.embedding.progress.get',
+      { workspaceId },
+      { timeoutMs: 10000 }
+    );
   }
 
   addContextBlob(options: OptionsField<typeof addContextBlobMutation>) {

packages/frontend/core/src/blocksuite/ai/runtime/request/service.ts

@@ -1,3 +1,4 @@
+import type { NbstoreService } from '@affine/core/modules/storage';
 import {
   ContextCategories,
   type CopilotChatHistoryFragment,
@@ -390,7 +391,8 @@ export function createAIRequestService(
   eventSource: (
     url: string,
     eventSourceInitDict?: EventSourceInit
-  ) => EventSource
+  ) => EventSource,
+  realtime: Pick<NbstoreService['realtime'], 'request'>
 ) {
-  return new AIRequestService(new CopilotClient(gql, eventSource));
+  return new AIRequestService(new CopilotClient(gql, eventSource, realtime));
 }