feat(core): migrate more pull to realtime (#14936)

#### PR Dependency Tree


* **PR #14936** 👈

This tree was auto-generated by
[Charcoal](https://github.com/danerwilliams/charcoal)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Refactor**
* Consolidated realtime subscription patterns for consistent, more
reliable live updates across comments, notifications, transcription
tasks, and embedding progress.
* Standardized realtime room naming and subscription keys for
deterministic delivery.

* **New Features**
* Introduced a reusable live-query mechanism powering realtime snapshot
+ event workflows used by comments, notifications, transcript tasks, and
embedding progress.

* **Tests**
* Added tests covering live-query behavior and deterministic
subscription key generation.

[![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/14936)
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

packages/frontend/core/src/modules/cloud/index.ts

@@ -6,6 +6,11 @@ export { AccountLoggedIn } from './events/account-logged-in';
 export { AccountLoggedOut } from './events/account-logged-out';
 export { AuthProvider } from './provider/auth';
 export { ValidatorProvider } from './provider/validator';
+export {
+  RealtimeLiveQuery,
+  type RealtimeLiveQueryEventResult,
+  type RealtimeLiveQueryOptions,
+} from './realtime/live-query';
 export { ServerScope } from './scopes/server';
 export { AccessTokenService } from './services/access-token';
 export { AuthService } from './services/auth';

packages/frontend/core/src/modules/cloud/realtime/live-query.spec.ts

@@ -0,0 +1,156 @@
+import { Subject } from 'rxjs';
+import { describe, expect, test, vi } from 'vitest';
+
+import { RealtimeLiveQuery } from './live-query';
+
+describe('RealtimeLiveQuery', () => {
+  test('requests snapshot when subscription is ready', async () => {
+    const events$ = new Subject<{ type: 'ready' } | { count: number }>();
+    const applySnapshot = vi.fn();
+    const query = new RealtimeLiveQuery({
+      request: vi.fn().mockResolvedValue({ count: 1 }),
+      subscribe: () => events$,
+      applySnapshot,
+    });
+
+    query.start();
+    events$.next({ type: 'ready' });
+
+    await vi.waitFor(() =>
+      expect(applySnapshot).toHaveBeenCalledWith({ count: 1 })
+    );
+    query.dispose();
+  });
+
+  test('applies event without revalidating when applyEvent returns applied', async () => {
+    const events$ = new Subject<{ type: 'ready' } | { count: number }>();
+    const request = vi.fn().mockResolvedValue({ count: 1 });
+    const applyEvent = vi.fn().mockReturnValue('applied');
+    const query = new RealtimeLiveQuery({
+      request,
+      subscribe: () => events$,
+      applySnapshot: vi.fn(),
+      applyEvent,
+    });
+
+    query.start();
+    events$.next({ count: 2 });
+
+    expect(applyEvent).toHaveBeenCalledWith({ count: 2 });
+    expect(request).not.toHaveBeenCalled();
+    query.dispose();
+  });
+
+  test('applies business events that include a type field', () => {
+    const events$ = new Subject<
+      { type: 'ready' } | { type: 'updated'; count: number }
+    >();
+    const request = vi.fn().mockResolvedValue({ count: 1 });
+    const applyEvent = vi.fn().mockReturnValue('applied');
+    const query = new RealtimeLiveQuery({
+      request,
+      subscribe: () => events$,
+      applySnapshot: vi.fn(),
+      applyEvent,
+    });
+
+    query.start();
+    events$.next({ type: 'updated', count: 2 });
+
+    expect(applyEvent).toHaveBeenCalledWith({ type: 'updated', count: 2 });
+    expect(request).not.toHaveBeenCalled();
+    query.dispose();
+  });
+
+  test('revalidates event when applyEvent asks for it', async () => {
+    const events$ = new Subject<{ type: 'ready' } | { changed: true }>();
+    const applySnapshot = vi.fn();
+    const query = new RealtimeLiveQuery({
+      request: vi.fn().mockResolvedValue({ changes: [] }),
+      subscribe: () => events$,
+      applySnapshot,
+      applyEvent: () => 'revalidate',
+    });
+
+    query.start();
+    events$.next({ changed: true });
+
+    await vi.waitFor(() =>
+      expect(applySnapshot).toHaveBeenCalledWith({ changes: [] })
+    );
+    query.dispose();
+  });
+
+  test('passes subscription and request errors to onError', async () => {
+    const events$ = new Subject<{ type: 'ready' }>();
+    const onError = vi.fn();
+    const requestError = new Error('request failed');
+    const query = new RealtimeLiveQuery({
+      request: vi.fn().mockRejectedValue(requestError),
+      subscribe: () => events$,
+      applySnapshot: vi.fn(),
+      onError,
+    });
+
+    query.start();
+    events$.next({ type: 'ready' });
+    await vi.waitFor(() => expect(onError).toHaveBeenCalledWith(requestError));
+
+    const subscriptionError = new Error('subscribe failed');
+    events$.error(subscriptionError);
+    expect(onError).toHaveBeenCalledWith(subscriptionError);
+    query.dispose();
+  });
+
+  test('dispose aborts in-flight request and ignores stale result', async () => {
+    const events$ = new Subject<{ type: 'ready' }>();
+    const applySnapshot = vi.fn();
+    let resolveRequest: (value: { count: number }) => void = () => {};
+    const query = new RealtimeLiveQuery({
+      request: vi.fn(
+        () =>
+          new Promise<{ count: number }>(resolve => {
+            resolveRequest = resolve;
+          })
+      ),
+      subscribe: () => events$,
+      applySnapshot,
+    });
+
+    query.start();
+    events$.next({ type: 'ready' });
+    query.dispose();
+    resolveRequest({ count: 1 });
+
+    await Promise.resolve();
+    expect(applySnapshot).not.toHaveBeenCalled();
+  });
+
+  test('new request supersedes older in-flight result', async () => {
+    const events$ = new Subject<{ type: 'ready' }>();
+    const applySnapshot = vi.fn();
+    const resolvers: Array<(value: { count: number }) => void> = [];
+    const query = new RealtimeLiveQuery({
+      request: vi.fn(
+        () =>
+          new Promise<{ count: number }>(resolve => {
+            resolvers.push(resolve);
+          })
+      ),
+      subscribe: () => events$,
+      applySnapshot,
+    });
+
+    query.start();
+    events$.next({ type: 'ready' });
+    events$.next({ type: 'ready' });
+    resolvers[0]({ count: 1 });
+    resolvers[1]({ count: 2 });
+
+    await vi.waitFor(() =>
+      expect(applySnapshot).toHaveBeenCalledWith({ count: 2 })
+    );
+    expect(applySnapshot).toHaveBeenCalledTimes(1);
+    query.dispose();
+  });
+});

packages/frontend/core/src/modules/cloud/realtime/live-query.ts

@@ -0,0 +1,100 @@
+import type { RealtimeSubscriptionReady } from '@affine/realtime';
+import type { Observable, Subscription } from 'rxjs';
+
+export type RealtimeLiveQueryEventResult = 'applied' | 'revalidate';
+
+export type RealtimeLiveQueryOptions<TSnapshot, TEvent extends object> = {
+  request: (signal: AbortSignal) => Promise<TSnapshot>;
+  subscribe: () => Observable<TEvent | RealtimeSubscriptionReady>;
+  applySnapshot: (snapshot: TSnapshot) => void;
+  applyEvent?: (event: TEvent) => RealtimeLiveQueryEventResult;
+  onError?: (error: unknown) => void;
+};
+
+function isReadyEvent<TEvent extends object>(
+  event: TEvent | RealtimeSubscriptionReady
+): event is RealtimeSubscriptionReady {
+  return 'type' in event && event.type === 'ready';
+}
+
+export class RealtimeLiveQuery<TSnapshot, TEvent extends object> {
+  private subscription?: Subscription;
+  private requestController?: AbortController;
+  private generation = 0;
+  private started = false;
+
+  constructor(
+    private readonly options: RealtimeLiveQueryOptions<TSnapshot, TEvent>
+  ) {}
+
+  start() {
+    this.stop();
+    this.started = true;
+    const generation = this.generation;
+    this.subscription = this.options.subscribe().subscribe({
+      next: event => {
+        if (isReadyEvent(event)) {
+          this.revalidate();
+          return;
+        }
+        if (!this.options.applyEvent) {
+          this.revalidate();
+          return;
+        }
+        if (this.options.applyEvent(event) === 'revalidate') {
+          this.revalidate();
+        }
+      },
+      error: error => {
+        if (this.generation === generation) {
+          this.options.onError?.(error);
+        }
+      },
+    });
+  }
+
+  revalidate() {
+    if (!this.started) {
+      return;
+    }
+    this.requestController?.abort();
+    const controller = new AbortController();
+    this.requestController = controller;
+    const generation = this.generation;
+    this.options.request(controller.signal).then(
+      snapshot => {
+        if (
+          this.started &&
+          this.generation === generation &&
+          this.requestController === controller &&
+          !controller.signal.aborted
+        ) {
+          this.options.applySnapshot(snapshot);
+        }
+      },
+      error => {
+        if (
+          this.started &&
+          this.generation === generation &&
+          this.requestController === controller &&
+          !controller.signal.aborted
+        ) {
+          this.options.onError?.(error);
+        }
+      }
+    );
+  }
+
+  stop() {
+    this.started = false;
+    this.generation += 1;
+    this.subscription?.unsubscribe();
+    this.subscription = undefined;
+    this.requestController?.abort();
+    this.requestController = undefined;
+  }
+
+  dispose() {
+    this.stop();
+  }
+}

packages/frontend/core/src/modules/comment/entities/doc-comment-store.ts

@@ -11,6 +11,10 @@ import {
   updateReplyMutation,
   uploadCommentAttachmentMutation,
 } from '@affine/graphql';
+import type {
+  RealtimeSubscriptionReady,
+  RealtimeTopicEventOf,
+} from '@affine/realtime';
 import { Entity } from '@toeverything/infra';
 import type { Observable } from 'rxjs';
 
@@ -144,11 +148,11 @@ export class DocCommentStore extends Entity<{
       { after, workspaceId: this.currentWorkspaceId, docId: this.props.docId }
     );
     return {
-      changes: commentChanges.changes.map((change: any) => ({
+      changes: commentChanges.changes.map(change => ({
         id: change.id,
         action: change.action,
         comment: normalizeComment(change.item as GQLCommentType),
-        commentId: change.commentId,
+        commentId: change.commentId ?? undefined,
       })),
       startCursor: commentChanges.startCursor,
       endCursor: commentChanges.endCursor,
@@ -156,7 +160,9 @@ export class DocCommentStore extends Entity<{
     };
   }
 
-  subscribeCommentChanged(): Observable<unknown> {
+  subscribeCommentChanged(): Observable<
+    RealtimeTopicEventOf<'comment.changed'> | RealtimeSubscriptionReady
+  > {
     return this.nbstoreService.realtime.subscribe('comment.changed', {
       workspaceId: this.currentWorkspaceId,
       docId: this.props.docId,

packages/frontend/core/src/modules/comment/entities/doc-comment.ts

@@ -16,17 +16,9 @@ import {
   onStart,
 } from '@toeverything/infra';
 import { nanoid } from 'nanoid';
-import {
-  catchError,
-  filter,
-  first,
-  of,
-  Subject,
-  type Subscription,
-  switchMap,
-  tap,
-} from 'rxjs';
+import { catchError, filter, first, of, Subject, switchMap, tap } from 'rxjs';
 
+import { RealtimeLiveQuery } from '../../cloud/realtime/live-query';
 import { type DocDisplayMetaService } from '../../doc-display-meta';
 import { GlobalContextService } from '../../global-context';
 import type { SnapshotHelper } from '../services/snapshot-helper';
@@ -93,11 +85,19 @@ export class DocCommentEntity extends Entity<{
   private readonly commentDeleted$ = new Subject<CommentId>();
   readonly commentHighlighted$ = new LiveData<CommentId | null>(null);
 
-  private realtimeDisposable?: DisposeCallback;
-  private realtimeSubscription?: Subscription;
+  private readonly liveQuery = new RealtimeLiveQuery({
+    request: () => this.store.listCommentChanges({ after: this.startCursor }),
+    subscribe: () => this.store.subscribeCommentChanged(),
+    applySnapshot: result => {
+      this.handleCommentChanges(result);
+      this.startCursor = result.endCursor;
+    },
+    onError: error => {
+      console.error('Failed to sync comment changes:', error);
+    },
+  });
   private startCursor?: string;
-  private fetchingCommentChanges = false;
-  private pendingCommentChangeFetch = false;
+  private startVersion = 0;
 
   async addComment(
     selections?: BaseSelection[],
@@ -490,52 +490,34 @@ export class DocCommentEntity extends Entity<{
   }
 
   start(): void {
-    if (this.realtimeDisposable) {
-      this.realtimeDisposable();
-    }
-    this.realtimeSubscription?.unsubscribe();
-
-    this.revalidate();
     this.revalidateCommentsInEditor();
-
-    this.realtimeSubscription = this.store.subscribeCommentChanged().subscribe({
-      next: () => {
-        this.fetchCommentChanges().catch(() => {});
-      },
-      error: error => {
-        console.error('Failed to subscribe comment changes:', error);
-      },
+    this.startLiveQueryAfterInitialLoad().catch(error => {
+      console.error('Failed to start comment realtime:', error);
     });
-    this.realtimeDisposable = () => this.realtimeSubscription?.unsubscribe();
   }
 
   stop(): void {
-    if (this.realtimeDisposable) {
-      this.realtimeDisposable();
-    }
-    this.realtimeSubscription?.unsubscribe();
+    this.startVersion++;
+    this.liveQuery.stop();
+    this.loading$.setValue(false);
   }
 
-  private async fetchCommentChanges() {
-    if (this.fetchingCommentChanges) {
-      this.pendingCommentChangeFetch = true;
-      return;
-    }
-
-    this.fetchingCommentChanges = true;
+  private async startLiveQueryAfterInitialLoad() {
+    const version = ++this.startVersion;
+    this.liveQuery.stop();
+    this.loading$.setValue(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);
+      const allComments = await this.loadAllComments();
+      if (version !== this.startVersion) {
+        return;
+      }
+      this.revalidateCommentsInEditor();
+      this.comments$.setValue(allComments);
+      this.liveQuery.start();
     } finally {
-      this.fetchingCommentChanges = false;
+      if (version === this.startVersion) {
+        this.loading$.setValue(false);
+      }
     }
   }
 
@@ -650,31 +632,9 @@ export class DocCommentEntity extends Entity<{
 
   revalidate = effect(
     switchMap(() => {
-      return fromPromise(async () => {
-        const allComments: DocComment[] = [];
-        let cursor = '';
-        let firstResult: DocCommentListResult | null = null;
-        this.revalidateCommentsInEditor();
-        // 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 incremental changes
-            this.startCursor = result.startCursor;
-          }
-          allComments.push(...result.comments);
-          cursor = result.endCursor;
-          if (!result.hasNextPage) {
-            break;
-          }
-        }
-
-        return allComments;
-      }).pipe(
+      return fromPromise(() => this.loadAllComments()).pipe(
         tap(allComments => {
           this.revalidateCommentsInEditor();
-          // Update state with all comments
           this.comments$.setValue(allComments);
         }),
         onStart(() => this.loading$.setValue(true)),
@@ -688,6 +648,27 @@ export class DocCommentEntity extends Entity<{
     })
   );
 
+  private async loadAllComments() {
+    const allComments: DocComment[] = [];
+    let cursor = '';
+    let firstResult: DocCommentListResult | null = null;
+    this.revalidateCommentsInEditor();
+    while (true) {
+      const result = await this.store.listComments({ after: cursor });
+      if (!firstResult) {
+        firstResult = result;
+        this.startCursor = result.startCursor;
+      }
+      allComments.push(...result.comments);
+      cursor = result.endCursor;
+      if (!result.hasNextPage) {
+        break;
+      }
+    }
+
+    return allComments;
+  }
+
   private readonly revalidateCommentsInEditor = () => {
     this.commentsInEditor$.setValue(this.getCommentsInEditor());
   };

packages/frontend/core/src/modules/media/entities/audio-transcription-job.ts

@@ -1,10 +1,11 @@
 import { shallowEqual } from '@affine/component';
 import type { TranscriptionBlockProps } from '@affine/core/blocksuite/ai/blocks/transcription-block/model';
+import { RealtimeLiveQuery } from '@affine/core/modules/cloud/realtime/live-query';
 import { DebugLogger } from '@affine/debug';
 import { UserFriendlyError } from '@affine/error';
-import { AiJobStatus } from '@affine/graphql';
+import { AiJobStatus, type TranscriptionResultType } from '@affine/graphql';
+import type { RealtimeTopicEventOf } from '@affine/realtime';
 import { Entity, LiveData } from '@toeverything/infra';
-import type { Subscription } from 'rxjs';
 
 import type { DefaultServerService, WorkspaceServerService } from '../../cloud';
 import { AuthService } from '../../cloud/services/auth';
@@ -65,7 +66,10 @@ export class AudioTranscriptionJob extends Entity<{
   }
 
   disposed = false;
-  private taskSubscription?: Subscription;
+  private taskLiveQuery?: RealtimeLiveQuery<
+    TranscriptionResultType | null,
+    RealtimeTopicEventOf<'copilot.transcript.task.changed'>
+  >;
   private taskWaitReject?: (error: unknown) => void;
 
   private readonly _status$ = new LiveData<TranscriptionStatus>({
@@ -205,45 +209,56 @@ export class AudioTranscriptionJob extends Entity<{
 
     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'
-                )
-              );
+      this.taskLiveQuery = new RealtimeLiveQuery({
+        request: () => this.store.getTranscriptTask(this.props.blobId, taskId),
+        subscribe: () => this.store.subscribeTranscriptTask(taskId),
+        applySnapshot: job => {
+          this.applyTranscriptTaskSnapshot(taskId, job).then(() => {
+            if (this.status$.value.status === AiJobStatus.finished) {
+              resolve();
             }
-          },
-          error: reject,
-        });
+          }, reject);
+        },
+        applyEvent: event => {
+          if (event.status === AiJobStatus.failed) {
+            reject(
+              UserFriendlyError.fromAny(
+                event.error ?? 'Transcription job failed'
+              )
+            );
+            return 'applied';
+          }
+          return event.status === AiJobStatus.finished
+            ? 'revalidate'
+            : 'applied';
+        },
+        onError: reject,
+      });
+      this.taskLiveQuery.start();
     }).finally(() => {
       this.taskWaitReject = undefined;
-      this.taskSubscription?.unsubscribe();
-      this.taskSubscription = undefined;
+      this.taskLiveQuery?.dispose();
+      this.taskLiveQuery = undefined;
     });
   }
 
   private rejectTaskWait(error: unknown) {
     const reject = this.taskWaitReject;
     this.taskWaitReject = undefined;
-    this.taskSubscription?.unsubscribe();
-    this.taskSubscription = undefined;
+    this.taskLiveQuery?.dispose();
+    this.taskLiveQuery = undefined;
     reject?.(error);
   }
 
   private async checkTranscriptTask(taskId: string) {
+    const job = await this.store.getTranscriptTask(this.props.blobId, taskId);
+    await this.applyTranscriptTaskSnapshot(taskId, job);
+  }
+
+  private async applyTranscriptTaskSnapshot(
+    taskId: string,
+    job: TranscriptionResultType | null
+  ) {
     if (
       this.disposed ||
       this.props.blockProps.jobId !== taskId ||
@@ -251,7 +266,6 @@ export class AudioTranscriptionJob extends Entity<{
     ) {
       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', {

packages/frontend/core/src/modules/notification/services/count.ts

@@ -1,20 +1,8 @@
-import {
-  catchErrorInto,
-  effect,
-  exhaustMapWithTrailing,
-  fromPromise,
-  LiveData,
-  onComplete,
-  OnEvent,
-  onStart,
-  Service,
-  smartRetry,
-} from '@toeverything/infra';
-import type { Subscription } from 'rxjs';
-import { tap } from 'rxjs';
+import { LiveData, OnEvent, Service } from '@toeverything/infra';
 
 import { AccountChanged, type AuthService } from '../../cloud';
 import { ServerStarted } from '../../cloud/events/server-started';
+import { RealtimeLiveQuery } from '../../cloud/realtime/live-query';
 import { ApplicationFocused } from '../../lifecycle';
 import type { NbstoreService } from '../../storage';
 import type { NotificationStore } from '../stores/notification';
@@ -36,30 +24,25 @@ 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;
+  private readonly liveQuery = new RealtimeLiveQuery({
+    request: signal => this.requestCount(signal),
+    subscribe: () =>
+      this.nbstoreService.realtime.subscribe('notification.count.changed', {}),
+    applySnapshot: result => this.setCount(result.count),
+    applyEvent: event => {
+      this.setCount(event.count);
+      return 'applied';
+    },
+    onError: error => this.error$.setValue(error),
+  });
 
-  revalidate = effect(
-    exhaustMapWithTrailing(() => {
-      return fromPromise(() => {
-        if (!this.loggedIn$.value) {
-          return Promise.resolve(0);
-        }
-        return this.nbstoreService.realtime
-          .request('notification.count.get', {}, { timeoutMs: 10000 })
-          .then(result => result.count);
-      }).pipe(
-        tap(result => {
-          this.setCount(result ?? 0);
-        }),
-        smartRetry(),
-        catchErrorInto(this.error$),
-        onStart(() => {
-          this.isLoading$.setValue(true);
-        }),
-        onComplete(() => this.isLoading$.setValue(false))
-      );
-    })
-  );
+  revalidate = () => {
+    if (!this.loggedIn$.value) {
+      this.setCount(0);
+      return;
+    }
+    this.liveQuery.revalidate();
+  };
 
   handleApplicationFocused() {
     this.revalidate();
@@ -76,33 +59,34 @@ export class NotificationCountService extends Service {
   }
 
   setCount(count: number) {
+    this.error$.setValue(null);
     this.store.setNotificationCountCache(count);
   }
 
   override dispose(): void {
     super.dispose();
-    this.revalidate.unsubscribe();
-    this.subscription?.unsubscribe();
+    this.liveQuery.dispose();
   }
 
   private subscribe() {
-    this.subscription?.unsubscribe();
     if (!this.loggedIn$.value) {
+      this.liveQuery.stop();
+      this.setCount(0);
       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);
-        },
-      });
+    this.liveQuery.start();
+  }
+
+  private async requestCount(signal: AbortSignal) {
+    this.isLoading$.setValue(true);
+    try {
+      return await this.nbstoreService.realtime.request(
+        'notification.count.get',
+        {},
+        { signal, timeoutMs: 10000 }
+      );
+    } finally {
+      this.isLoading$.setValue(false);
+    }
   }
 }

packages/frontend/core/src/modules/workspace-indexer-embedding/entities/embedding-progress.ts

@@ -1,17 +1,8 @@
+import { RealtimeLiveQuery } from '@affine/core/modules/cloud/realtime/live-query';
 import type { WorkspaceService } from '@affine/core/modules/workspace';
+import type { RealtimeTopicEventOf } from '@affine/realtime';
 import { logger } from '@sentry/react';
-import {
-  catchErrorInto,
-  effect,
-  Entity,
-  fromPromise,
-  LiveData,
-  onComplete,
-  onStart,
-  smartRetry,
-} from '@toeverything/infra';
-import { EMPTY, type Subscription } from 'rxjs';
-import { exhaustMap, mergeMap } from 'rxjs/operators';
+import { Entity, LiveData } from '@toeverything/infra';
 
 import type { EmbeddingStore } from '../stores/embedding';
 import type { LocalAttachmentFile } from '../types';
@@ -26,8 +17,31 @@ export class EmbeddingProgress extends Entity {
   error$ = new LiveData<any>(null);
   loading$ = new LiveData(true);
 
-  private progressSubscription?: Subscription;
   uploadingAttachments$ = new LiveData<LocalAttachmentFile[]>([]);
+  private started = false;
+  private readonly liveQuery = new RealtimeLiveQuery<
+    Progress,
+    RealtimeTopicEventOf<'workspace.embedding.progress.changed'>
+  >({
+    request: signal => this.requestProgress(signal),
+    subscribe: () =>
+      this.store.subscribeEmbeddingProgress(this.workspaceService.workspace.id),
+    applySnapshot: progress => this.applyProgress(progress),
+    applyEvent: event => {
+      if (
+        typeof event.embedded === 'number' &&
+        typeof event.total === 'number'
+      ) {
+        this.applyProgress({ embedded: event.embedded, total: event.total });
+        return 'applied';
+      }
+      return 'revalidate';
+    },
+    onError: error => {
+      this.error$.setValue(error);
+      logger.error('Failed to fetch workspace embedding progress', { error });
+    },
+  });
 
   constructor(
     private readonly workspaceService: WorkspaceService,
@@ -37,48 +51,46 @@ export class EmbeddingProgress extends Entity {
   }
 
   startEmbeddingProgress() {
-    this.stopEmbeddingProgress();
-    this.progressSubscription = this.store
-      .subscribeEmbeddingProgress(this.workspaceService.workspace.id)
-      .subscribe({
-        next: () => this.getEmbeddingProgress(),
-        error: error => this.error$.setValue(error),
-      });
-    this.getEmbeddingProgress();
+    this.started = true;
+    this.liveQuery.start();
   }
 
   stopEmbeddingProgress() {
-    this.progressSubscription?.unsubscribe();
-    this.progressSubscription = undefined;
+    this.started = false;
+    this.liveQuery.stop();
   }
 
-  getEmbeddingProgress = effect(
-    exhaustMap(() => {
-      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))
-      );
-    })
-  );
+  getEmbeddingProgress = () => {
+    if (this.started) {
+      this.liveQuery.revalidate();
+      return;
+    }
+    this.requestProgress(new AbortController().signal).then(
+      progress => this.applyProgress(progress),
+      error => this.error$.setValue(error)
+    );
+  };
 
   override dispose(): void {
-    this.progressSubscription?.unsubscribe();
-    this.getEmbeddingProgress.unsubscribe();
+    this.liveQuery.dispose();
+  }
+
+  private applyProgress(value: Progress | null) {
+    this.progress$.next(value);
+    if (value && value.embedded === value.total && value.total > 0) {
+      this.stopEmbeddingProgress();
+    }
+  }
+
+  private async requestProgress(signal: AbortSignal) {
+    this.loading$.setValue(true);
+    try {
+      return await this.store.getEmbeddingProgress(
+        this.workspaceService.workspace.id,
+        signal
+      );
+    } finally {
+      this.loading$.setValue(false);
+    }
   }
 }