From dda3103d1b9c00332b19c09cceabbbe77993b825 Mon Sep 17 00:00:00 2001
From: EYHN <cneyhn@gmail.com>
Date: Tue, 25 Mar 2025 08:15:22 +0000
Subject: [PATCH] feat(nbstore): allow polling protocol (#11160)

---
 .../common/nbstore/src/impls/cloud/awareness.ts   |  6 +++++-
 packages/common/nbstore/src/impls/cloud/doc.ts    |  3 ++-
 packages/common/nbstore/src/impls/cloud/socket.ts | 15 +++++++++------
 .../modules/userspace/entities/user-db-engine.ts  |  4 ++++
 .../src/modules/workspace-engine/impls/cloud.ts   |  7 +++++++
 5 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/packages/common/nbstore/src/impls/cloud/awareness.ts b/packages/common/nbstore/src/impls/cloud/awareness.ts
index 472bf095f4..17f9280d48 100644
--- a/packages/common/nbstore/src/impls/cloud/awareness.ts
+++ b/packages/common/nbstore/src/impls/cloud/awareness.ts
@@ -10,6 +10,7 @@ import {
 } from './socket';
 
 interface CloudAwarenessStorageOptions {
+  isSelfHosted: boolean;
   serverBaseUrl: string;
   type: SpaceType;
   id: string;
@@ -22,7 +23,10 @@ export class CloudAwarenessStorage extends AwarenessStorageBase {
     super();
   }
 
-  connection = new SocketConnection(this.options.serverBaseUrl);
+  connection = new SocketConnection(
+    this.options.serverBaseUrl,
+    this.options.isSelfHosted
+  );
 
   private get socket() {
     return this.connection.inner.socket;
diff --git a/packages/common/nbstore/src/impls/cloud/doc.ts b/packages/common/nbstore/src/impls/cloud/doc.ts
index b5fee8b37c..c2fe4a06d7 100644
--- a/packages/common/nbstore/src/impls/cloud/doc.ts
+++ b/packages/common/nbstore/src/impls/cloud/doc.ts
@@ -18,6 +18,7 @@ import {
 
 interface CloudDocStorageOptions extends DocStorageOptions {
   serverBaseUrl: string;
+  isSelfHosted: boolean;
   type: SpaceType;
 }
 
@@ -185,7 +186,7 @@ class CloudDocStorageConnection extends SocketConnection {
     private readonly options: CloudDocStorageOptions,
     private readonly onServerUpdate: ServerEventsMap['space:broadcast-doc-update']
   ) {
-    super(options.serverBaseUrl);
+    super(options.serverBaseUrl, options.isSelfHosted);
   }
 
   idConverter: IdConverter | null = null;
diff --git a/packages/common/nbstore/src/impls/cloud/socket.ts b/packages/common/nbstore/src/impls/cloud/socket.ts
index 238ef83b9d..95b5d52965 100644
--- a/packages/common/nbstore/src/impls/cloud/socket.ts
+++ b/packages/common/nbstore/src/impls/cloud/socket.ts
@@ -165,10 +165,10 @@ class SocketManager {
   socket: Socket;
   refCount = 0;
 
-  constructor(endpoint: string) {
+  constructor(endpoint: string, isSelfHosted: boolean) {
     this.socketIOManager = new SocketIOManager(endpoint, {
       autoConnect: false,
-      transports: ['websocket'],
+      transports: isSelfHosted ? ['websocket', 'polling'] : ['websocket'], // self-hosted server may not support websocket
       secure: new URL(endpoint).protocol === 'https:',
       // we will handle reconnection by ourselves
       reconnection: false,
@@ -205,10 +205,10 @@ class SocketManager {
 }
 
 const SOCKET_MANAGER_CACHE = new Map<string, SocketManager>();
-function getSocketManager(endpoint: string) {
+function getSocketManager(endpoint: string, isSelfHosted: boolean) {
   let manager = SOCKET_MANAGER_CACHE.get(endpoint);
   if (!manager) {
-    manager = new SocketManager(endpoint);
+    manager = new SocketManager(endpoint, isSelfHosted);
     SOCKET_MANAGER_CACHE.set(endpoint, manager);
   }
   return manager;
@@ -218,9 +218,12 @@ export class SocketConnection extends AutoReconnectConnection<{
   socket: Socket;
   disconnect: () => void;
 }> {
-  manager = getSocketManager(this.endpoint);
+  manager = getSocketManager(this.endpoint, this.isSelfHosted);
 
-  constructor(private readonly endpoint: string) {
+  constructor(
+    private readonly endpoint: string,
+    private readonly isSelfHosted: boolean
+  ) {
     super();
   }
 
diff --git a/packages/frontend/core/src/modules/userspace/entities/user-db-engine.ts b/packages/frontend/core/src/modules/userspace/entities/user-db-engine.ts
index d4c67d7000..479ed5b5de 100644
--- a/packages/frontend/core/src/modules/userspace/entities/user-db-engine.ts
+++ b/packages/frontend/core/src/modules/userspace/entities/user-db-engine.ts
@@ -1,3 +1,4 @@
+import { ServerDeploymentType } from '@affine/graphql';
 import {
   IndexedDBDocStorage,
   IndexedDBDocSyncStorage,
@@ -64,6 +65,9 @@ export class UserDBEngine extends Entity<{
                 id: this.userId,
                 serverBaseUrl: serverService.server.baseUrl,
                 type: 'userspace',
+                isSelfHosted:
+                  serverService.server.config$.value.type ===
+                  ServerDeploymentType.Selfhosted,
               },
             },
           },
diff --git a/packages/frontend/core/src/modules/workspace-engine/impls/cloud.ts b/packages/frontend/core/src/modules/workspace-engine/impls/cloud.ts
index 6ac997029d..8e58005426 100644
--- a/packages/frontend/core/src/modules/workspace-engine/impls/cloud.ts
+++ b/packages/frontend/core/src/modules/workspace-engine/impls/cloud.ts
@@ -4,6 +4,7 @@ import {
   deleteWorkspaceMutation,
   getWorkspaceInfoQuery,
   getWorkspacesQuery,
+  ServerDeploymentType,
 } from '@affine/graphql';
 import type {
   BlobStorage,
@@ -468,6 +469,9 @@ class CloudWorkspaceFlavourProvider implements WorkspaceFlavourProvider {
               type: 'workspace',
               id: workspaceId,
               serverBaseUrl: this.server.serverMetadata.baseUrl,
+              isSelfHosted:
+                this.server.config$.value.type ===
+                ServerDeploymentType.Selfhosted,
             },
           },
           blob: {
@@ -483,6 +487,9 @@ class CloudWorkspaceFlavourProvider implements WorkspaceFlavourProvider {
               type: 'workspace',
               id: workspaceId,
               serverBaseUrl: this.server.serverMetadata.baseUrl,
+              isSelfHosted:
+                this.server.config$.value.type ===
+                ServerDeploymentType.Selfhosted,
             },
           },
         },