feat(server): realtime notification & task status (#14934)

#### PR Dependency Tree * **PR #14934** 👈 This tree was auto-generated by [Charcoal](https://github.com/danerwilliams/charcoal)  ## Summary by CodeRabbit * **New Features** * Full realtime platform added: live notifications, comments, embedding progress, and transcription task updates via realtime subscriptions. * **Chores** * Frontend switched from polling/GraphQL queries to realtime channels; legacy query fields marked deprecated and client libs updated to use realtime APIs. [![Review Change Stack](https://storage.googleapis.com/coderabbit_public_assets/review-stack-in-coderabbit-ui.svg)](https://app.coderabbit.ai/change-stack/toeverything/AFFiNE/pull/14934)  #### PR Dependency Tree * **PR #14934** 👈 * **PR #14936** This tree was auto-generated by [Charcoal](https://github.com/danerwilliams/charcoal)
2026-07-25 06:18:45 +08:00 · 2026-05-10 23:21:50 +08:00
parent 417d31cabe
commit 8cf00738c2
70 changed files with 2378 additions and 283 deletions
@@ -19,6 +19,7 @@ export { InvitationService } from './services/invitation';
 export { InvoicesService } from './services/invoices';
 export type { PublicUserInfo } from './services/public-user';
 export { PublicUserService } from './services/public-user';
+export { RealtimeService } from './services/realtime';
 export { SelfhostGenerateLicenseService } from './services/selfhost-generate-license';
 export { SelfhostLicenseService } from './services/selfhost-license';
 export { ServerService } from './services/server';
@@ -41,6 +42,7 @@ import { type Framework } from '@toeverything/infra';

 import { GlobalCache, GlobalState } from '../storage/providers/global';
 import { GlobalStateService } from '../storage/services/global';
+import { GlobalContextService } from '../global-context';
 import { UrlService } from '../url';
 import { WorkspaceScope, WorkspaceService } from '../workspace';
 import { CloudDocMeta } from './entities/cloud-doc-meta';
@@ -69,6 +71,7 @@ import { FetchService } from './services/fetch';
 import { GraphQLService } from './services/graphql';
 import { InvoicesService } from './services/invoices';
 import { PublicUserService } from './services/public-user';
+import { RealtimeService } from './services/realtime';
 import { SelfhostGenerateLicenseService } from './services/selfhost-generate-license';
 import { SelfhostLicenseService } from './services/selfhost-license';
 import { ServerService } from './services/server';
@@ -100,6 +103,7 @@ import { DocCreatedByService } from './services/doc-created-by';
 import { DocUpdatedByService } from './services/doc-updated-by';
 import { DocCreatedByUpdatedBySyncService } from './services/doc-created-by-updated-by-sync';
 import { WorkspacePermissionService } from '../permissions';
+import { NbstoreService } from '../storage';
 import { DocScope, DocService, DocsService } from '../doc';
 import { DocCreatedByUpdatedBySyncStore } from './stores/doc-created-by-updated-by-sync';
 import { GlobalDialogService } from '../dialogs';
@@ -111,6 +115,11 @@ export function configureCloudModule(framework: Framework) {

  framework
    .service(ServersService, [ServerListStore, ServerConfigStore])
+    .service(RealtimeService, [
+      GlobalContextService,
+      ServersService,
+      NbstoreService,
+    ])
    .service(DefaultServerService, [ServersService])
    .store(ServerListStore, [GlobalStateService])
    .store(ServerConfigStore)
@@ -0,0 +1,54 @@
+import { shallowEqual } from '@affine/component';
+import { ServerDeploymentType } from '@affine/graphql';
+import { LiveData, OnEvent, Service } from '@toeverything/infra';
+
+import type { GlobalContextService } from '../../global-context';
+import { ApplicationStarted } from '../../lifecycle';
+import type { NbstoreService } from '../../storage';
+import type { Server } from '../entities/server';
+import type { ServersService } from './servers';
+
+@OnEvent(ApplicationStarted, service => service.onApplicationStarted)
+export class RealtimeService extends Service {
+  private readonly currentServer$ =
+    this.globalContextService.globalContext.serverId.$.selector(id =>
+      id
+        ? this.serversService.server$(id)
+        : new LiveData<Server | undefined>(undefined)
+    )
+      .flat()
+      .selector(
+        server =>
+          [
+            server,
+            server?.account$,
+            server?.config$.selector(
+              c => c.type === ServerDeploymentType.Selfhosted
+            ),
+          ] as const
+      )
+      .flat()
+      .map(([server, account, selfHosted]) => ({
+        endpoint: server?.baseUrl ?? '',
+        authenticated: !!account,
+        isSelfHosted: !!selfHosted,
+      }))
+      .distinctUntilChanged(shallowEqual);
+
+  constructor(
+    private readonly globalContextService: GlobalContextService,
+    private readonly serversService: ServersService,
+    private readonly nbstoreService: NbstoreService
+  ) {
+    super();
+
+    const subscription = this.currentServer$.subscribe(context => {
+      this.nbstoreService.realtime.configure(context).catch(error => {
+        console.error('Failed to configure realtime context', error);
+      });
+    });
+    this.disposables.push(() => subscription.unsubscribe());
+  }
+
+  onApplicationStarted() {}
+}
@@ -4,7 +4,6 @@ import {
  deleteCommentMutation,
  deleteReplyMutation,
  type DocMode,
-  listCommentChangesQuery,
  type ListCommentsQuery,
  listCommentsQuery,
  resolveCommentMutation,
@@ -13,9 +12,11 @@ import {
  uploadCommentAttachmentMutation,
 } from '@affine/graphql';
 import { Entity } from '@toeverything/infra';
+import type { Observable } from 'rxjs';

 import type { DefaultServerService, WorkspaceServerService } from '../../cloud';
 import { GraphQLService } from '../../cloud/services/graphql';
+import type { NbstoreService } from '../../storage';
 import type { WorkspaceService } from '../../workspace';
 import type {
  DocComment,
@@ -75,7 +76,8 @@ export class DocCommentStore extends Entity<{
  constructor(
    private readonly workspaceService: WorkspaceService,
    private readonly workspaceServerService: WorkspaceServerService,
-    private readonly defaultServerService: DefaultServerService
+    private readonly defaultServerService: DefaultServerService,
+    private readonly nbstoreService: NbstoreService
  ) {
    super();
  }
@@ -132,51 +134,35 @@ export class DocCommentStore extends Entity<{
    };
  }

-  // pool every 30s
  async listCommentChanges({
    after,
  }: {
    after?: string;
  }): Promise<DocCommentChangeListResult> {
-    const graphql = this.graphqlService;
-    if (!graphql) {
-      throw new Error('GraphQL service not found');
-    }
-
-    const response = await graphql.gql({
-      query: listCommentChangesQuery,
-      variables: {
-        pagination: {
-          after,
-        },
-        workspaceId: this.currentWorkspaceId,
-        docId: this.props.docId,
-      },
-    });
-
-    const commentChanges = response.workspace?.commentChanges;
-    if (!commentChanges) {
-      return {
-        changes: [],
-        startCursor: '',
-        endCursor: after ?? '',
-        hasNextPage: false,
-      };
-    }
-
+    const commentChanges = await this.nbstoreService.realtime.request(
+      'comment.changes.get',
+      { after, workspaceId: this.currentWorkspaceId, docId: this.props.docId }
+    );
    return {
-      changes: commentChanges.edges.map(edge => ({
-        id: edge.node.id,
-        action: edge.node.action,
-        comment: normalizeComment(edge.node.item),
-        commentId: edge.node.commentId || undefined,
+      changes: commentChanges.changes.map((change: any) => ({
+        id: change.id,
+        action: change.action,
+        comment: normalizeComment(change.item as GQLCommentType),
+        commentId: change.commentId,
      })),
-      startCursor: commentChanges.pageInfo.startCursor || '',
-      endCursor: commentChanges.pageInfo.endCursor || '',
-      hasNextPage: commentChanges.pageInfo.hasNextPage,
+      startCursor: commentChanges.startCursor,
+      endCursor: commentChanges.endCursor,
+      hasNextPage: commentChanges.hasNextPage,
    };
  }

+  subscribeCommentChanged(): Observable<unknown> {
+    return this.nbstoreService.realtime.subscribe('comment.changed', {
+      workspaceId: this.currentWorkspaceId,
+      docId: this.props.docId,
+    });
+  }
+
  async createComment(commentInput: {
    content: DocCommentContent;
    mentions?: string[];
@@ -22,9 +22,9 @@ import {
  first,
  of,
  Subject,
+  type Subscription,
  switchMap,
  tap,
-  timer,
 } from 'rxjs';

 import { type DocDisplayMetaService } from '../../doc-display-meta';
@@ -93,8 +93,11 @@ export class DocCommentEntity extends Entity<{
  private readonly commentDeleted$ = new Subject<CommentId>();
  readonly commentHighlighted$ = new LiveData<CommentId | null>(null);

-  private pollingDisposable?: DisposeCallback;
+  private realtimeDisposable?: DisposeCallback;
+  private realtimeSubscription?: Subscription;
  private startCursor?: string;
+  private fetchingCommentChanges = false;
+  private pendingCommentChangeFetch = false;

  async addComment(
    selections?: BaseSelection[],
@@ -486,84 +489,53 @@ export class DocCommentEntity extends Entity<{
    return () => subscription.unsubscribe();
  }

-  // Start polling comments every 30s
-  // 1. when comments$ is empty, fetch all comments
-  // 2. when comments$ is not empty, fetch changes (using end cursor)
-  // 3. loop. when doc is not loaded, skip
  start(): void {
-    if (this.pollingDisposable) {
-      this.pollingDisposable();
+    if (this.realtimeDisposable) {
+      this.realtimeDisposable();
    }
+    this.realtimeSubscription?.unsubscribe();

-    // Initial load
    this.revalidate();
    this.revalidateCommentsInEditor();

-    // Set up polling every 10 seconds
-    const polling$ = timer(10000, 10000).pipe(
-      switchMap(() => {
-        // If we have comments, fetch changes; otherwise fetch all
-        if (this.comments$.value.length > 0) {
-          return fromPromise(async () => {
-            const res = await this.store.listCommentChanges({
-              after: this.startCursor,
-            });
-            return res;
-          }).pipe(
-            tap(result => {
-              if (result) {
-                this.handleCommentChanges(result);
-                this.startCursor = result.endCursor;
-              }
-            }),
-            catchError(error => {
-              console.error('Failed to fetch comment changes:', error);
-              return of(null);
-            })
-          );
-        } else {
-          return fromPromise(async () => {
-            const allComments: DocComment[] = [];
-            let cursor = '';
-            let firstResult: DocCommentListResult | null = null;
-
-            // Fetch all pages of comments
-            while (true) {
-              const result = await this.store.listComments({ after: cursor });
-              if (!firstResult) {
-                firstResult = result;
-                // Store the startCursor from the first page for future polling
-                this.startCursor = result.startCursor;
-              }
-              allComments.push(...result.comments);
-              cursor = result.endCursor;
-              if (!result.hasNextPage) {
-                break;
-              }
-            }
-
-            // Update state with all comments
-            this.comments$.setValue(allComments);
-
-            return allComments;
-          }).pipe(
-            tap(() => this.revalidateCommentsInEditor()),
-            catchError(error => {
-              console.error('Failed to fetch comments:', error);
-              return of(null);
-            })
-          );
-        }
-      })
-    );
-
-    const subscription = polling$.subscribe();
-    this.pollingDisposable = () => subscription.unsubscribe();
+    this.realtimeSubscription = this.store.subscribeCommentChanged().subscribe({
+      next: () => {
+        this.fetchCommentChanges().catch(() => {});
+      },
+      error: error => {
+        console.error('Failed to subscribe comment changes:', error);
+      },
+    });
+    this.realtimeDisposable = () => this.realtimeSubscription?.unsubscribe();
  }

  stop(): void {
-    if (this.pollingDisposable) {
-      this.pollingDisposable();
+    if (this.realtimeDisposable) {
+      this.realtimeDisposable();
+    }
+    this.realtimeSubscription?.unsubscribe();
+  }
+
+  private async fetchCommentChanges() {
+    if (this.fetchingCommentChanges) {
+      this.pendingCommentChangeFetch = true;
+      return;
+    }
+
+    this.fetchingCommentChanges = true;
+    try {
+      do {
+        this.pendingCommentChangeFetch = false;
+        const result = await this.store.listCommentChanges({
+          after: this.startCursor,
+        });
+        this.handleCommentChanges(result);
+        this.startCursor = result.endCursor;
+      } while (this.pendingCommentChangeFetch);
+    } catch (error) {
+      console.error('Failed to fetch comment changes:', error);
+    } finally {
+      this.fetchingCommentChanges = false;
    }
  }

@@ -688,7 +660,7 @@ export class DocCommentEntity extends Entity<{
          const result = await this.store.listComments({ after: cursor });
          if (!firstResult) {
            firstResult = result;
-            // Store the startCursor from the first page for polling
+            // Store the startCursor from the first page for incremental changes
            this.startCursor = result.startCursor;
          }
          allComments.push(...result.comments);
@@ -2,6 +2,7 @@ import type { Framework } from '@toeverything/infra';

 import { DefaultServerService, WorkspaceServerService } from '../cloud';
 import { DocDisplayMetaService } from '../doc-display-meta';
+import { NbstoreService } from '../storage';
 import { WorkbenchService } from '../workbench';
 import { WorkspaceScope, WorkspaceService } from '../workspace';
 import { DocCommentEntity } from './entities/doc-comment';
@@ -25,5 +26,6 @@ export function configureCommentModule(framework: Framework) {
      WorkspaceService,
      WorkspaceServerService,
      DefaultServerService,
+      NbstoreService,
    ]);
 }
@@ -1,5 +1,4 @@
 import {
-  getTranscriptTaskQuery,
  retryTranscriptTaskMutation,
  settleTranscriptTaskMutation,
  submitTranscriptTaskMutation,
@@ -10,6 +9,7 @@ import { describe, expect, test, vi } from 'vitest';
 import { DefaultServerService } from '../../cloud/services/default-server';
 import { GraphQLService } from '../../cloud/services/graphql';
 import { WorkspaceServerService } from '../../cloud/services/workspace-server';
+import { NbstoreService } from '../../storage';
 import { WorkspaceService } from '../../workspace';
 import { AudioTranscriptionJobStore } from './audio-transcription-job-store';

@@ -30,6 +30,10 @@ function createStore(
      get: (key: unknown) => (key === GraphQLService ? { gql } : null),
    },
  };
+  const realtime = {
+    request: vi.fn().mockResolvedValue({ task: { id: 'task-2' } }),
+    subscribe: vi.fn(),
+  };
  framework
    .service(WorkspaceService, {
      workspace: { id: 'workspace-1' },
@@ -42,10 +46,14 @@ function createStore(
    .service(DefaultServerService, {
      server: null,
    } as unknown as DefaultServerService)
+    .service(NbstoreService, {
+      realtime,
+    } as unknown as NbstoreService)
    .entity(AudioTranscriptionJobStore, [
      WorkspaceService,
      WorkspaceServerService,
      DefaultServerService,
+      NbstoreService,
    ]);
  return framework.provider().createEntity(AudioTranscriptionJobStore, {
    blobId: 'blob-1',
@@ -60,13 +68,6 @@ describe('AudioTranscriptionJobStore transcript task API', () => {
      .fn()
      .mockResolvedValueOnce({ submitTranscriptTask: { id: 'task-1' } })
      .mockResolvedValueOnce({ retryTranscriptTask: { id: 'task-2' } })
-      .mockResolvedValueOnce({
-        currentUser: {
-          copilot: {
-            transcriptTask: { id: 'task-2' },
-          },
-        },
-      })
      .mockResolvedValueOnce({ settleTranscriptTask: { id: 'task-2' } });
    const store = createStore(gql, async () => ({
      files: [file],
@@ -102,17 +103,6 @@ describe('AudioTranscriptionJobStore transcript task API', () => {
    );
    expect(gql).toHaveBeenNthCalledWith(
      3,
-      expect.objectContaining({
-        query: getTranscriptTaskQuery,
-        variables: {
-          workspaceId: 'workspace-1',
-          taskId: 'task-2',
-          blobId: 'blob-1',
-        },
-      })
-    );
-    expect(gql).toHaveBeenNthCalledWith(
-      4,
      expect.objectContaining({
        query: settleTranscriptTaskMutation,
        variables: {
@@ -1,13 +1,14 @@
 import {
-  getTranscriptTaskQuery,
  retryTranscriptTaskMutation,
  settleTranscriptTaskMutation,
  submitTranscriptTaskMutation,
+  type TranscriptionResultType,
 } from '@affine/graphql';
 import { Entity } from '@toeverything/infra';

 import type { DefaultServerService, WorkspaceServerService } from '../../cloud';
 import { GraphQLService } from '../../cloud/services/graphql';
+import type { NbstoreService } from '../../storage';
 import type { WorkspaceService } from '../../workspace';

 export class AudioTranscriptionJobStore extends Entity<{
@@ -20,7 +21,8 @@ export class AudioTranscriptionJobStore extends Entity<{
  constructor(
    private readonly workspaceService: WorkspaceService,
    private readonly workspaceServerService: WorkspaceServerService,
-    private readonly defaultServerService: DefaultServerService
+    private readonly defaultServerService: DefaultServerService,
+    private readonly nbstoreService: NbstoreService
  ) {
    super();
  }
@@ -79,27 +81,27 @@ export class AudioTranscriptionJobStore extends Entity<{
    return response.retryTranscriptTask;
  };

-  getTranscriptTask = async (blobId: string, taskId?: string) => {
-    const graphqlService = this.graphqlService;
-    if (!graphqlService) {
-      throw new Error('No graphql service available');
-    }
+  getTranscriptTask = async (
+    blobId: string,
+    taskId?: string
+  ): Promise<TranscriptionResultType | null> => {
    const currentWorkspaceId = this.currentWorkspaceId;
    if (!currentWorkspaceId) {
      throw new Error('No current workspace id');
    }
-    const response = await graphqlService.gql({
-      query: getTranscriptTaskQuery,
-      variables: {
-        workspaceId: currentWorkspaceId,
-        taskId,
-        blobId,
-      },
-    });
-    if (!response.currentUser?.copilot?.transcriptTask) {
-      return null;
-    }
-    return response.currentUser.copilot.transcriptTask;
+    const response = await this.nbstoreService.realtime.request(
+      'copilot.transcript.task.get',
+      { workspaceId: currentWorkspaceId, taskId, blobId },
+      { timeoutMs: 10000 }
+    );
+    return response.task as TranscriptionResultType | null;
+  };
+
+  subscribeTranscriptTask = (taskId: string) => {
+    return this.nbstoreService.realtime.subscribe(
+      'copilot.transcript.task.changed',
+      { workspaceId: this.currentWorkspaceId, taskId }
+    );
  };
  settleTranscriptTask = async (taskId: string) => {
    const graphqlService = this.graphqlService;
@@ -4,6 +4,7 @@ import { DebugLogger } from '@affine/debug';
 import { UserFriendlyError } from '@affine/error';
 import { AiJobStatus } from '@affine/graphql';
 import { Entity, LiveData } from '@toeverything/infra';
+import type { Subscription } from 'rxjs';

 import type { DefaultServerService, WorkspaceServerService } from '../../cloud';
 import { AuthService } from '../../cloud/services/auth';
@@ -59,10 +60,13 @@ export class AudioTranscriptionJob extends Entity<{
    super();
    this.disposables.push(() => {
      this.disposed = true;
+      this.rejectTaskWait(new Error('Job disposed'));
    });
  }

  disposed = false;
+  private taskSubscription?: Subscription;
+  private taskWaitReject?: (error: unknown) => void;

  private readonly _status$ = new LiveData<TranscriptionStatus>({
    status: 'waiting-for-job',
@@ -190,38 +194,77 @@ export class AudioTranscriptionJob extends Entity<{
  }

  private async untilTaskReadyOrSettled() {
-    while (
-      !this.disposed &&
-      this.props.blockProps.jobId &&
-      this.props.blockProps.createdBy === this.currentUserId
+    const taskId = this.props.blockProps.jobId;
+    if (!taskId || this.props.blockProps.createdBy !== this.currentUserId) {
+      return;
+    }
+
+    await this.checkTranscriptTask(taskId);
+
+    this.rejectTaskWait(new Error('Transcript task wait replaced'));
+
+    await new Promise<void>((resolve, reject) => {
+      this.taskWaitReject = reject;
+      this.taskSubscription = this.store
+        .subscribeTranscriptTask(taskId)
+        .subscribe({
+          next: event => {
+            if (
+              'type' in event ||
+              (event.status as AiJobStatus) === AiJobStatus.finished
+            ) {
+              this.checkTranscriptTask(taskId).then(() => {
+                if (this.status$.value.status === AiJobStatus.finished) {
+                  resolve();
+                }
+              }, reject);
+            } else if ((event.status as AiJobStatus) === AiJobStatus.failed) {
+              reject(
+                UserFriendlyError.fromAny(
+                  event.error ?? 'Transcription job failed'
+                )
+              );
+            }
+          },
+          error: reject,
+        });
+    }).finally(() => {
+      this.taskWaitReject = undefined;
+      this.taskSubscription?.unsubscribe();
+      this.taskSubscription = undefined;
+    });
+  }
+
+  private rejectTaskWait(error: unknown) {
+    const reject = this.taskWaitReject;
+    this.taskWaitReject = undefined;
+    this.taskSubscription?.unsubscribe();
+    this.taskSubscription = undefined;
+    reject?.(error);
+  }
+
+  private async checkTranscriptTask(taskId: string) {
+    if (
+      this.disposed ||
+      this.props.blockProps.jobId !== taskId ||
+      this.props.blockProps.createdBy !== this.currentUserId
    ) {
-      logger.debug('Polling job status', {
-        jobId: this.props.blockProps.jobId,
+      return;
+    }
+    const job = await this.store.getTranscriptTask(this.props.blobId, taskId);
+
+    if (!job || job.status === AiJobStatus.failed) {
+      logger.debug('Job failed during realtime status check', {
+        jobId: taskId,
      });
-      const job = await this.store.getTranscriptTask(
-        this.props.blobId,
-        this.props.blockProps.jobId
-      );
+      throw UserFriendlyError.fromAny('Transcription job failed');
+    }

-      if (!job || job?.status === 'failed') {
-        logger.debug('Job failed during polling', {
-          jobId: this.props.blockProps.jobId,
-        });
-        throw UserFriendlyError.fromAny('Transcription job failed');
-      }
-
-      if (job?.status === AiJobStatus.finished) {
-        logger.debug('Transcript task is ready to settle', {
-          jobId: this.props.blockProps.jobId,
-        });
-        this._status$.value = {
-          status: AiJobStatus.finished,
-        };
-        return;
-      }
-
-      // Add delay between polling attempts
-      await new Promise(resolve => setTimeout(resolve, 3000));
+    if (job.status === AiJobStatus.finished) {
+      logger.debug('Transcript task is ready to settle', { jobId: taskId });
+      this._status$.value = {
+        status: AiJobStatus.finished,
+      };
    }
  }

@@ -1,7 +1,7 @@
 import type { Framework } from '@toeverything/infra';

 import { DefaultServerService, WorkspaceServerService } from '../cloud';
-import { GlobalState, GlobalStateService } from '../storage';
+import { GlobalState, GlobalStateService, NbstoreService } from '../storage';
 import { WorkbenchService } from '../workbench';
 import { WorkspaceScope, WorkspaceService } from '../workspace';
 import { AudioAttachmentBlock } from './entities/audio-attachment-block';
@@ -35,6 +35,7 @@ export function configureMediaModule(framework: Framework) {
      WorkspaceService,
      WorkspaceServerService,
      DefaultServerService,
+      NbstoreService,
    ])
    .service(AudioAttachmentService)
    .service(AudioMediaManagerService, [
@@ -12,7 +12,7 @@ import {
  ServerScope,
  ServerService,
 } from '../cloud';
-import { GlobalSessionState } from '../storage';
+import { GlobalSessionState, NbstoreService } from '../storage';
 import { NotificationCountService } from './services/count';
 import { NotificationListService } from './services/list';
 import { NotificationService } from './services/notification';
@@ -22,7 +22,11 @@ export function configureNotificationModule(framework: Framework) {
  framework
    .scope(ServerScope)
    .service(NotificationService, [NotificationStore])
-    .service(NotificationCountService, [NotificationStore, AuthService])
+    .service(NotificationCountService, [
+      NotificationStore,
+      AuthService,
+      NbstoreService,
+    ])
    .service(NotificationListService, [
      NotificationStore,
      NotificationCountService,
@@ -10,11 +10,13 @@ import {
  Service,
  smartRetry,
 } from '@toeverything/infra';
-import { switchMap, tap, timer } from 'rxjs';
+import type { Subscription } from 'rxjs';
+import { tap } from 'rxjs';

 import { AccountChanged, type AuthService } from '../../cloud';
 import { ServerStarted } from '../../cloud/events/server-started';
 import { ApplicationFocused } from '../../lifecycle';
+import type { NbstoreService } from '../../storage';
 import type { NotificationStore } from '../stores/notification';

@OnEvent(ApplicationFocused, s => s.handleApplicationFocused)
@@ -23,7 +25,8 @@ import type { NotificationStore } from '../stores/notification';
 export class NotificationCountService extends Service {
  constructor(
    private readonly store: NotificationStore,
-    private readonly authService: AuthService
+    private readonly authService: AuthService,
+    private readonly nbstoreService: NbstoreService
  ) {
    super();
  }
@@ -33,17 +36,17 @@ export class NotificationCountService extends Service {
  readonly count$ = LiveData.from(this.store.watchNotificationCountCache(), 0);
  readonly isLoading$ = new LiveData(false);
  readonly error$ = new LiveData<any>(null);
+  private subscription?: Subscription;

  revalidate = effect(
-    switchMap(() => {
-      return timer(0, 30000); // revalidate every 30 seconds
-    }),
    exhaustMapWithTrailing(() => {
-      return fromPromise(signal => {
+      return fromPromise(() => {
        if (!this.loggedIn$.value) {
          return Promise.resolve(0);
        }
-        return this.store.getNotificationCount(signal);
+        return this.nbstoreService.realtime
+          .request('notification.count.get', {}, { timeoutMs: 10000 })
+          .then(result => result.count);
      }).pipe(
        tap(result => {
          this.setCount(result ?? 0);
@@ -63,10 +66,12 @@ export class NotificationCountService extends Service {
  }

  handleServerStarted() {
+    this.subscribe();
    this.revalidate();
  }

  handleAccountChanged() {
+    this.subscribe();
    this.revalidate();
  }

@@ -77,5 +82,27 @@ export class NotificationCountService extends Service {
  override dispose(): void {
    super.dispose();
    this.revalidate.unsubscribe();
+    this.subscription?.unsubscribe();
+  }
+
+  private subscribe() {
+    this.subscription?.unsubscribe();
+    if (!this.loggedIn$.value) {
+      return;
+    }
+    this.subscription = this.nbstoreService.realtime
+      .subscribe('notification.count.changed', {})
+      .subscribe({
+        next: event => {
+          if ('type' in event) {
+            this.revalidate();
+          } else {
+            this.setCount(event.count);
+          }
+        },
+        error: error => {
+          this.error$.setValue(error);
+        },
+      });
  }
 }
@@ -3,7 +3,6 @@ import {
  type ListNotificationsQuery,
  listNotificationsQuery,
  mentionUserMutation,
-  notificationCountQuery,
  type PaginationInput,
  readAllNotificationsMutation,
  readNotificationMutation,
@@ -52,17 +51,6 @@ export class NotificationStore extends Store {
    );
  }

-  async getNotificationCount(signal?: AbortSignal) {
-    const result = await this.gqlService.gql({
-      query: notificationCountQuery,
-      context: {
-        signal,
-      },
-    });
-
-    return result.currentUser?.notifications.totalCount;
-  }
-
  async listNotification(pagination: PaginationInput, signal?: AbortSignal) {
    const result = await this.gqlService.gql({
      query: listNotificationsQuery,
@@ -1,10 +1,13 @@
 import type {
  StoreClient,
+  StoreManagerClient,
  WorkerInitOptions,
 } from '@affine/nbstore/worker/client';
 import { createIdentifier } from '@toeverything/infra';

 export interface NbstoreProvider {
+  readonly realtime: StoreManagerClient['realtime'];
+
  /**
   * Open a nbstore with the given options, if the store with the given key already exists, it will be returned.
   *
@@ -11,4 +11,8 @@ export class NbstoreService extends Service {
  openStore(key: string, options: WorkerInitOptions) {
    return this.nbstoreProvider.openStore(key, options);
  }
+
+  get realtime() {
+    return this.nbstoreProvider.realtime;
+  }
 }
@@ -10,8 +10,8 @@ import {
  onStart,
  smartRetry,
 } from '@toeverything/infra';
-import { EMPTY, interval, Subject } from 'rxjs';
-import { exhaustMap, mergeMap, switchMap, takeUntil } from 'rxjs/operators';
+import { EMPTY, type Subscription } from 'rxjs';
+import { exhaustMap, mergeMap } from 'rxjs/operators';

 import type { EmbeddingStore } from '../stores/embedding';
 import type { LocalAttachmentFile } from '../types';
@@ -26,8 +26,7 @@ export class EmbeddingProgress extends Entity {
  error$ = new LiveData<any>(null);
  loading$ = new LiveData(true);

-  private readonly EMBEDDING_PROGRESS_POLL_INTERVAL = 3000;
-  private readonly stopEmbeddingProgress$ = new Subject<void>();
+  private progressSubscription?: Subscription;
  uploadingAttachments$ = new LiveData<LocalAttachmentFile[]>([]);

  constructor(
@@ -37,50 +36,49 @@ export class EmbeddingProgress extends Entity {
    super();
  }

-  startEmbeddingProgressPolling() {
-    this.stopEmbeddingProgressPolling();
+  startEmbeddingProgress() {
+    this.stopEmbeddingProgress();
+    this.progressSubscription = this.store
+      .subscribeEmbeddingProgress(this.workspaceService.workspace.id)
+      .subscribe({
+        next: () => this.getEmbeddingProgress(),
+        error: error => this.error$.setValue(error),
+      });
    this.getEmbeddingProgress();
  }

-  stopEmbeddingProgressPolling() {
-    this.stopEmbeddingProgress$.next();
+  stopEmbeddingProgress() {
+    this.progressSubscription?.unsubscribe();
+    this.progressSubscription = undefined;
  }

  getEmbeddingProgress = effect(
    exhaustMap(() => {
-      return interval(this.EMBEDDING_PROGRESS_POLL_INTERVAL).pipe(
-        takeUntil(this.stopEmbeddingProgress$),
-        switchMap(() =>
-          fromPromise(signal =>
-            this.store.getEmbeddingProgress(
-              this.workspaceService.workspace.id,
-              signal
-            )
-          ).pipe(
-            smartRetry(),
-            mergeMap(value => {
-              this.progress$.next(value);
-              if (value && value.embedded === value.total && value.total > 0) {
-                this.stopEmbeddingProgressPolling();
-              }
-              return EMPTY;
-            }),
-            catchErrorInto(this.error$, error => {
-              logger.error(
-                'Failed to fetch workspace embedding progress',
-                error
-              );
-            }),
-            onStart(() => this.loading$.setValue(true)),
-            onComplete(() => this.loading$.setValue(false))
-          )
+      return fromPromise(signal =>
+        this.store.getEmbeddingProgress(
+          this.workspaceService.workspace.id,
+          signal
        )
+      ).pipe(
+        smartRetry(),
+        mergeMap(value => {
+          this.progress$.next(value);
+          if (value && value.embedded === value.total && value.total > 0) {
+            this.stopEmbeddingProgress();
+          }
+          return EMPTY;
+        }),
+        catchErrorInto(this.error$, error => {
+          logger.error('Failed to fetch workspace embedding progress', error);
+        }),
+        onStart(() => this.loading$.setValue(true)),
+        onComplete(() => this.loading$.setValue(false))
      );
    })
  );

  override dispose(): void {
-    this.stopEmbeddingProgress$.next();
+    this.progressSubscription?.unsubscribe();
    this.getEmbeddingProgress.unsubscribe();
  }
 }
@@ -1,4 +1,5 @@
 import { WorkspaceServerService } from '@affine/core/modules/cloud';
+import { NbstoreService } from '@affine/core/modules/storage';
 import {
  WorkspaceScope,
  WorkspaceService,
@@ -16,7 +17,7 @@ export function configureIndexerEmbeddingModule(framework: Framework) {
  framework
    .scope(WorkspaceScope)
    .service(EmbeddingService)
-    .store(EmbeddingStore, [WorkspaceServerService])
+    .store(EmbeddingStore, [WorkspaceServerService, NbstoreService])
    .entity(EmbeddingEnabled, [WorkspaceService, EmbeddingStore])
    .entity(AdditionalAttachments, [WorkspaceService, EmbeddingStore])
    .entity(IgnoredDocs, [WorkspaceService, EmbeddingStore])
@@ -1,11 +1,11 @@
 import type { WorkspaceServerService } from '@affine/core/modules/cloud';
+import type { NbstoreService } from '@affine/core/modules/storage';
 import {
  addWorkspaceEmbeddingFilesMutation,
  addWorkspaceEmbeddingIgnoredDocsMutation,
  getAllWorkspaceEmbeddingIgnoredDocsQuery,
  getWorkspaceConfigQuery,
  getWorkspaceEmbeddingFilesQuery,
-  getWorkspaceEmbeddingStatusQuery,
  type PaginationInput,
  removeWorkspaceEmbeddingFilesMutation,
  removeWorkspaceEmbeddingIgnoredDocsMutation,
@@ -14,7 +14,10 @@ import {
 import { Store } from '@toeverything/infra';

 export class EmbeddingStore extends Store {
-  constructor(private readonly workspaceServerService: WorkspaceServerService) {
+  constructor(
+    private readonly workspaceServerService: WorkspaceServerService,
+    private readonly nbstoreService: NbstoreService
+  ) {
    super();
  }

@@ -178,17 +181,17 @@ export class EmbeddingStore extends Store {
  }

  async getEmbeddingProgress(workspaceId: string, signal?: AbortSignal) {
-    if (!this.workspaceServerService.server) {
-      throw new Error('No Server');
-    }
+    return await this.nbstoreService.realtime.request(
+      'workspace.embedding.progress.get',
+      { workspaceId },
+      { signal, timeoutMs: 10000 }
+    );
+  }

-    const data = await this.workspaceServerService.server.gql({
-      query: getWorkspaceEmbeddingStatusQuery,
-      variables: {
-        workspaceId,
-      },
-      context: { signal },
-    });
-    return data.queryWorkspaceEmbeddingStatus;
+  subscribeEmbeddingProgress(workspaceId: string) {
+    return this.nbstoreService.realtime.subscribe(
+      'workspace.embedding.progress.changed',
+      { workspaceId }
+    );
  }
 }
@@ -66,7 +66,7 @@ const EmbeddingCloud: React.FC<{ disabled: boolean }> = ({ disabled }) => {
        .setEnabled(checked)
        .then(() => {
          if (checked) {
-            embeddingService.embeddingProgress.startEmbeddingProgressPolling();
+            embeddingService.embeddingProgress.startEmbeddingProgress();
          }
        })
        .catch(error => {
@@ -91,8 +91,7 @@ const EmbeddingCloud: React.FC<{ disabled: boolean }> = ({ disabled }) => {
        docType: file.type,
      });
      embeddingService.additionalAttachments.addAttachments([file]);
-      // Restart polling to track progress of newly uploaded files
-      embeddingService.embeddingProgress.startEmbeddingProgressPolling();
+      embeddingService.embeddingProgress.startEmbeddingProgress();
    },
    [embeddingService.additionalAttachments, embeddingService.embeddingProgress]
  );
@@ -168,7 +167,7 @@ const EmbeddingCloud: React.FC<{ disabled: boolean }> = ({ disabled }) => {
  ]);

  useEffect(() => {
-    embeddingService.embeddingProgress.startEmbeddingProgressPolling();
+    embeddingService.embeddingProgress.startEmbeddingProgress();
    embeddingService.embeddingEnabled.getEnabled();
    embeddingService.additionalAttachments.getAttachments({
      first: COUNT_PER_PAGE,
@@ -178,7 +177,7 @@ const EmbeddingCloud: React.FC<{ disabled: boolean }> = ({ disabled }) => {
    embeddingService.embeddingProgress.getEmbeddingProgress();

    return () => {
-      embeddingService.embeddingProgress.stopEmbeddingProgressPolling();
+      embeddingService.embeddingProgress.stopEmbeddingProgress();
    };
  }, [
    embeddingService.embeddingProgress,