fix(ios): ios local access limit for self-hosted workspace (#15338)

## Summary - Route iOS nbstore worker token reads through the main-thread MessagePort and skip Capacitor Auth on `/socket.io` polling so self-host WebSocket/XHR sync no longer hangs on Connect timeout. - Harden workspace `flavour:id` routing, DocSyncPeer abort/status handling, and `resetSync` so local selfhost edits push and Mac browsers can receive them. - Soften root-doc readiness waits and session-exchange throttling to keep mobile selfhost sign-in/sync stable under retries.  ## Summary by CodeRabbit ## Summary by CodeRabbit * **New Features** * Improved workspace switching across local and remote environments, preserving workspace type and server context. * Added support for page navigation with query parameters. * Added iOS local-network permission messaging for self-hosted workspaces. * **Bug Fixes** * Improved document synchronization, reset handling, prioritized document refreshes, and retry behavior. * Prevented authentication headers and refresh attempts for socket connection requests. * Improved workspace reopening and routing when multiple workspace types share an ID. * Fixed handling of unlimited data query limits.  --------- Co-authored-by: DarkSky <darksky2048@gmail.com>
2026-08-28 12:19:27 +08:00 · 2026-08-24 09:33:58 +08:00
parent b530198a3b
commit cd6593c659
15 changed files with 448 additions and 78 deletions
@@ -113,6 +113,12 @@ class TestDocStorage implements DocStorage {
  }
 }

+class TimestampBlindDocStorage extends IndexedDBDocStorage {
+  override async getDocTimestamps(): Promise<DocClocks> {
+    return {};
+  }
+}
+
 class PermissionDeniedRemoteDocStorage implements DocStorage {
  readonly storageType = 'doc' as const;
  readonly connection = new DummyConnection();
@@ -306,7 +312,7 @@ test('doc', async () => {
    type: 'workspace',
  });

-  const peerBDoc = new IndexedDBDocStorage({
+  const peerBDoc = new TimestampBlindDocStorage({
    id: 'ws1',
    flavour: 'b',
    type: 'workspace',
@@ -340,6 +346,26 @@ test('doc', async () => {
    docId: 'doc1',
    bin: update,
  });
+  const prioritizedDocId = 'prioritized-doc';
+  const localPrioritizedDoc = new YDoc();
+  localPrioritizedDoc.getMap('test').set('local', true);
+  const localPrioritizedClock = await peerA.get('doc').pushDocUpdate({
+    docId: prioritizedDocId,
+    bin: encodeStateAsUpdate(localPrioritizedDoc),
+  });
+  await peerASync.setPeerPushedClock('b', localPrioritizedClock);
+  const remotePrioritizedDoc = new YDoc();
+  remotePrioritizedDoc.getMap('test').set('remote', true);
+  await peerB.get('doc').pushDocUpdate({
+    docId: prioritizedDocId,
+    bin: encodeStateAsUpdate(remotePrioritizedDoc),
+  });
+  const rootDoc = new YDoc();
+  rootDoc.getMap('meta').set('name', 'Self-host workspace');
+  await peerB.get('doc').pushDocUpdate({
+    docId: 'ws1',
+    bin: encodeStateAsUpdate(rootDoc),
+  });

  const sync = new Sync({
    local: peerA,
@@ -348,6 +374,7 @@ test('doc', async () => {
      c: peerC,
    },
  });
+  const removeRootPriority = sync.doc.addPriority('ws1', 100);
  sync.start();

  await new Promise(resolve => setTimeout(resolve, 1000));
@@ -366,8 +393,33 @@ test('doc', async () => {
        hello: 'world',
      },
    });
+
+    const root = await peerA.get('doc').getDoc('ws1');
+    expectYjsEqual(root!.bin, {
+      meta: {
+        name: 'Self-host workspace',
+      },
+    });
+
+    const prioritized = await peerA.get('doc').getDoc(prioritizedDocId);
+    expectYjsEqual(prioritized!.bin, {
+      test: {
+        local: true,
+      },
+    });
  }

+  const removeDocPriority = sync.doc.addPriority(prioritizedDocId, 100);
+  await vi.waitFor(async () => {
+    const prioritized = await peerA.get('doc').getDoc(prioritizedDocId);
+    expectYjsEqual(prioritized!.bin, {
+      test: {
+        local: true,
+        remote: true,
+      },
+    });
+  });
+
  doc.getMap('test').set('foo', 'bar');
  const update2 = encodeStateAsUpdate(doc);
  await peerC.get('doc').pushDocUpdate({
@@ -394,6 +446,13 @@ test('doc', async () => {
      },
    });
  }
+
+  removeDocPriority();
+  removeRootPriority();
+  sync.stop();
+  peerA.disconnect();
+  peerB.disconnect();
+  peerC.disconnect();
 });

 test('blob', async () => {
@@ -24,6 +24,7 @@ import {
 import { createNode } from './node-builder';

 const SQLITE_INDEXER_VERSION_OFFSET = 1;
+const NATIVE_INDEXER_MAX_LIMIT = 0xffffffff;

 export class SqliteIndexerStorage extends IndexerStorageBase {
  static readonly identifier = 'SqliteIndexerStorage';
@@ -83,7 +84,7 @@ export class SqliteIndexerStorage extends IndexerStorageBase {
      String(table),
      toNativeQuery(query),
      String(field),
-      limit,
+      toNativeLimit(limit),
      skip,
      options?.hits
        ? toNativeOptions(options.hits, hitLimit, hitSkip)
@@ -240,7 +241,7 @@ function toNativeOptions(
 ): NativeIndexSearchOptions {
  const highlights = options?.highlights?.map(item => String(item.field)) ?? [];
  return {
-    limit,
+    limit: toNativeLimit(limit),
    offset,
    fields: [
      ...new Set([...(options?.fields?.map(String) ?? []), ...highlights]),
@@ -248,3 +249,7 @@ function toNativeOptions(
    highlights,
  };
 }
+
+function toNativeLimit(limit: number) {
+  return limit === Infinity ? NATIVE_INDEXER_MAX_LIMIT : limit;
+}
@@ -34,6 +34,8 @@ export interface DocSyncDocState {
  errorMessage: string | null;
 }

+const RESET_SYNC_CONNECT_TIMEOUT_MS = 30_000;
+
 export interface DocSync {
  readonly state$: Observable<DocSyncState>;
  docState$(docId: string): Observable<DocSyncDocState>;
@@ -50,6 +52,8 @@ export class DocSyncImpl implements DocSync {
      new DocSyncPeer(peerId, this.storages.local, this.sync, remote)
  );
  private abort: AbortController | null = null;
+  private running: Promise<void> = Promise.resolve();
+  private resetting: Promise<void> | null = null;

  private readonly _state$ = combineLatest(
    this.peers.map(peer => peer.peerState$)
@@ -155,12 +159,16 @@ export class DocSyncImpl implements DocSync {
    if (this.abort) {
      this.abort.abort(MANUALLY_STOP);
    }
+    const previous = this.running;
    const abort = new AbortController();
    this.abort = abort;
-    Promise.allSettled(
-      this.peers.map(peer => peer.mainLoop(abort.signal))
-    ).catch(error => {
-      console.error(error);
+    this.running = previous.then(async () => {
+      if (abort.signal.aborted) {
+        return;
+      }
+      await Promise.allSettled(
+        this.peers.map(peer => peer.mainLoop(abort.signal))
+      );
    });
  }

@@ -174,12 +182,50 @@ export class DocSyncImpl implements DocSync {
    return () => undo.forEach(fn => fn());
  }

-  async resetSync() {
+  resetSync() {
+    if (this.resetting) {
+      return this.resetting;
+    }
+    const resetting = this.performReset().finally(() => {
+      if (this.resetting === resetting) {
+        this.resetting = null;
+      }
+    });
+    this.resetting = resetting;
+    return resetting;
+  }
+
+  private async performReset() {
    const running = this.abort !== null;
+    const activeRun = this.running;
+    const shouldConnectSyncStorage =
+      this.sync.connection.status === 'idle' ||
+      this.sync.connection.status === 'closed';
    this.stop();
-    await this.sync.clearClocks();
-    if (running) {
-      this.start();
+    await activeRun;
+    if (shouldConnectSyncStorage) {
+      this.sync.connection.connect();
+    }
+    const abort = new AbortController();
+    const timeoutId = setTimeout(() => {
+      abort.abort(new Error('Connect to remote timeout'));
+    }, RESET_SYNC_CONNECT_TIMEOUT_MS) as ReturnType<typeof setTimeout> & {
+      unref?: () => void;
+    };
+    timeoutId.unref?.();
+    try {
+      await this.sync.connection.waitForConnected(abort.signal);
+      await this.sync.clearClocks();
+    } catch (error) {
+      console.error('Failed to reset sync', error);
+      throw error;
+    } finally {
+      clearTimeout(timeoutId);
+      if (running) {
+        this.start();
+      } else if (shouldConnectSyncStorage) {
+        this.sync.connection.disconnect();
+      }
    }
  }
 }
@@ -280,9 +280,14 @@ export class DocSyncPeer {
          (await this.syncMetadata.getPeerPulledRemoteClock(this.peerId, docId))
            ?.timestamp ?? null;
        const remoteClock = this.status.remoteClocks.get(docId);
+        const hasRemoteClock = remoteClock.getTime() > 0;
+        const hasPulled = pulled !== null && pulled.getTime() > 0;
        if (
-          remoteClock &&
-          (pulled === null || pulled.getTime() < remoteClock.getTime())
+          hasRemoteClock
+            ? !hasPulled ||
+              (pulled !== null && pulled.getTime() < remoteClock.getTime())
+            : (this.prioritySettings.get(docId) ?? 0) > 0 &&
+              (!clock || !hasPulled)
        ) {
          await this.jobs.pull(docId, signal);
        }
@@ -503,10 +508,7 @@ export class DocSyncPeer {
      if (!this.status.docs.has(docId)) {
        this.status.docs.add(docId);
        this.statusUpdatedSubject$.next(docId);
-        this.schedule({
-          type: 'connect',
-          docId,
-        });
+        this.schedule({ type: 'connect', docId });
      }
    },
  };
@@ -756,6 +758,11 @@ export class DocSyncPeer {
      for (const docId of this.status.remoteClocks.keys()) {
        this.actions.addDoc(docId);
      }
+      for (const [docId, priority] of this.prioritySettings) {
+        if (priority > 0) {
+          this.actions.addDoc(docId);
+        }
+      }

      // begin to process jobs

@@ -890,13 +897,22 @@ export class DocSyncPeer {

  addPriority(id: string, priority: number) {
    const oldPriority = this.prioritySettings.get(id) ?? 0;
-    this.prioritySettings.set(id, priority);
-    this.status.jobDocQueue.setPriority(id, oldPriority + priority);
+    const newPriority = oldPriority + priority;
+    this.prioritySettings.set(id, newPriority);
+    this.status.jobDocQueue.setPriority(id, newPriority);
+    if (oldPriority <= 0 && newPriority > 0 && this.status.syncing) {
+      if (!this.status.docs.has(id)) {
+        this.actions.addDoc(id);
+      } else {
+        this.schedule({ type: 'connect', docId: id });
+      }
+    }

    return () => {
      const currentPriority = this.prioritySettings.get(id) ?? 0;
-      this.prioritySettings.set(id, currentPriority - priority);
-      this.status.jobDocQueue.setPriority(id, currentPriority - priority);
+      const restoredPriority = currentPriority - priority;
+      this.prioritySettings.set(id, restoredPriority);
+      this.status.jobDocQueue.setPriority(id, restoredPriority);
    };
  }

@@ -39,6 +39,8 @@
 	<true/>
 	<key>NSCameraUsageDescription</key>
 	<string>AFFiNE requires access to the camera to capture images and insert them into your documents</string>
+	<key>NSLocalNetworkUsageDescription</key>
+	<string>AFFiNE needs local network access to connect to your self-hosted workspace on this device's Wi-Fi network.</string>
 	<key>NSPhotoLibraryUsageDescription</key>
 	<string>AFFiNE requires access to select photos from your photo library and insert them into your documents</string>
 	<key>NSUserTrackingUsageDescription</key>
@@ -37,6 +37,18 @@
        }
      }
    },
+    "NSLocalNetworkUsageDescription" : {
+      "comment" : "Privacy - Local Network Usage Description",
+      "extractionState" : "extracted_with_value",
+      "localizations" : {
+        "en" : {
+          "stringUnit" : {
+            "state" : "new",
+            "value" : "AFFiNE needs local network access to connect to your self-hosted workspace on this device's Wi-Fi network."
+          }
+        }
+      }
+    },
    "NSPhotoLibraryUsageDescription" : {
      "comment" : "Privacy - Photo Library Usage Description",
      "extractionState" : "extracted_with_value",
@@ -335,15 +335,31 @@ registerNativeImageFilesPicker(async () => {
 });

 // ------ some apis for native ------
-(window as any).getCurrentServerBaseUrl = () => {
+const getCurrentServerForNative = () => {
  const globalContextService = frameworkProvider.get(GlobalContextService);
-  const currentServerId = globalContextService.globalContext.serverId.get();
+  const globalContext = globalContextService.globalContext;
+  const currentServerId = globalContext.serverId.get();
+  const currentWorkspaceFlavour = globalContext.workspaceFlavour.get();
  const serversService = frameworkProvider.get(ServersService);
  const defaultServerService = frameworkProvider.get(DefaultServerService);
-  const currentServer =
+
+  if (currentWorkspaceFlavour && currentWorkspaceFlavour !== 'local') {
+    const workspaceServer = serversService.server$(
+      currentWorkspaceFlavour
+    ).value;
+    if (workspaceServer) {
+      return workspaceServer;
+    }
+  }
+
+  return (
    (currentServerId ? serversService.server$(currentServerId).value : null) ??
-    defaultServerService.server;
-  return currentServer.baseUrl;
+    defaultServerService.server
+  );
+};
+
+(window as any).getCurrentServerBaseUrl = () => {
+  return getCurrentServerForNative().baseUrl;
 };
 (window as any).getCurrentI18nLocale = () => {
  return I18n.language;
@@ -365,11 +381,15 @@ registerNativeImageFilesPicker(async () => {
 };
 (window as any).waitForSelectedSources = async (documentIds: string[]) => {
  const globalContextService = frameworkProvider.get(GlobalContextService);
-  const currentWorkspaceId =
-    globalContextService.globalContext.workspaceId.get();
+  const globalContext = globalContextService.globalContext;
+  const currentWorkspaceId = globalContext.workspaceId.get();
+  const currentWorkspaceFlavour = globalContext.workspaceFlavour.get();
  const workspacesService = frameworkProvider.get(WorkspacesService);
  const workspaceRef = currentWorkspaceId
-    ? workspacesService.openByWorkspaceId(currentWorkspaceId)
+    ? workspacesService.openByWorkspaceId(
+        currentWorkspaceId,
+        currentWorkspaceFlavour
+      )
    : null;
  if (!workspaceRef) {
    throw new Error('Current workspace is unavailable');
@@ -408,13 +428,7 @@ registerNativeImageFilesPicker(async () => {
  return true;
 };
 const getCurrentNativeSignInContext = () => {
-  const globalContextService = frameworkProvider.get(GlobalContextService);
-  const currentServerId = globalContextService.globalContext.serverId.get();
-  const serversService = frameworkProvider.get(ServersService);
-  const defaultServerService = frameworkProvider.get(DefaultServerService);
-  const currentServer =
-    (currentServerId ? serversService.server$(currentServerId).value : null) ??
-    defaultServerService.server;
+  const currentServer = getCurrentServerForNative();
  const authService = currentServer.scope.get(AuthService);
  return { authService, currentServer };
 };
@@ -531,12 +545,16 @@ const showNativeSignIn = async () => {
 };
 (window as any).getCurrentDocContentInMarkdown = async () => {
  const globalContextService = frameworkProvider.get(GlobalContextService);
-  const currentWorkspaceId =
-    globalContextService.globalContext.workspaceId.get();
-  const currentDocId = globalContextService.globalContext.docId.get();
+  const globalContext = globalContextService.globalContext;
+  const currentWorkspaceId = globalContext.workspaceId.get();
+  const currentWorkspaceFlavour = globalContext.workspaceFlavour.get();
+  const currentDocId = globalContext.docId.get();
  const workspacesService = frameworkProvider.get(WorkspacesService);
  const workspaceRef = currentWorkspaceId
-    ? workspacesService.openByWorkspaceId(currentWorkspaceId)
+    ? workspacesService.openByWorkspaceId(
+        currentWorkspaceId,
+        currentWorkspaceFlavour
+      )
    : null;
  if (!workspaceRef) {
    return;
@@ -589,11 +607,15 @@ const showNativeSignIn = async () => {
  title: string
 ) => {
  const globalContextService = frameworkProvider.get(GlobalContextService);
-  const currentWorkspaceId =
-    globalContextService.globalContext.workspaceId.get();
+  const globalContext = globalContextService.globalContext;
+  const currentWorkspaceId = globalContext.workspaceId.get();
+  const currentWorkspaceFlavour = globalContext.workspaceFlavour.get();
  const workspacesService = frameworkProvider.get(WorkspacesService);
  const workspaceRef = currentWorkspaceId
-    ? workspacesService.openByWorkspaceId(currentWorkspaceId)
+    ? workspacesService.openByWorkspaceId(
+        currentWorkspaceId,
+        currentWorkspaceFlavour
+      )
    : null;

  try {
@@ -96,6 +96,30 @@ describe('shouldRefreshAccessToken', () => {
 });

 describe('auth request fetch', () => {
+  test.each([
+    ['/socket.io', true],
+    ['/socket.io/', true],
+    ['/socket.io/?EIO=4&transport=polling', true],
+    ['/socket.ioevil', false],
+  ])('handles the socket auth boundary for %s', async (path, skipped) => {
+    const provider = {
+      getValidAccessToken: vi.fn(async () => 'access-token'),
+      refreshAccessToken: vi.fn(async () => 'refreshed-token'),
+    };
+    const rawFetch = vi.fn<typeof fetch>(async request => {
+      expect((request as Request).headers.has('Authorization')).toBe(!skipped);
+      return new Response(JSON.stringify({ code: 'ACCESS_TOKEN_EXPIRED' }), {
+        status: 401,
+        headers: { 'content-type': 'application/json' },
+      });
+    });
+
+    await createAuthFetch(provider, rawFetch)(`https://example.com${path}`);
+
+    expect(provider.getValidAccessToken).toHaveBeenCalledTimes(skipped ? 0 : 1);
+    expect(provider.refreshAccessToken).toHaveBeenCalledTimes(skipped ? 0 : 1);
+  });
+
  test('injects the endpoint token', async () => {
    const provider = {
      getValidAccessToken: vi.fn(async () => 'access-token'),
@@ -164,6 +188,30 @@ describe('auth request fetch', () => {
 });

 describe('auth request XMLHttpRequest', () => {
+  test.each([
+    ['/socket.io', true],
+    ['/socket.io/', true],
+    ['/socket.io/?EIO=4&transport=polling', true],
+    ['/socket.ioevil', false],
+  ])('handles the socket auth boundary for %s', async (path, skipped) => {
+    const xhrCalls = stubXMLHttpRequest();
+    const provider = {
+      getValidAccessToken: vi.fn(async () => 'access-token'),
+      refreshAccessToken: vi.fn(async () => 'refreshed-token'),
+    };
+    installAuthRequestProxy(provider);
+    const xhr = new XMLHttpRequest();
+
+    xhr.open('GET', `https://example.com${path}`);
+    xhr.send();
+    await vi.waitFor(() => expect(xhrCalls.send).toHaveBeenCalledOnce());
+    xhrCalls.respond(401, JSON.stringify({ code: 'ACCESS_TOKEN_EXPIRED' }));
+    await Promise.resolve();
+
+    expect(provider.getValidAccessToken).toHaveBeenCalledTimes(skipped ? 0 : 1);
+    expect(provider.refreshAccessToken).toHaveBeenCalledTimes(skipped ? 0 : 1);
+  });
+
  test('does not send after abort while waiting for a token', async () => {
    let resolveToken: (token: string | null) => void = () => {};
    const token = new Promise<string | null>(resolve => {
@@ -19,6 +19,30 @@ function authEndpointForUrl(url: string | URL) {
  }
 }

+function shouldSkipStoredAuthToken(url: string | URL) {
+  try {
+    const { pathname } = new URL(
+      url,
+      globalThis.location?.origin ?? 'http://localhost'
+    );
+    if (pathname === '/socket.io' || pathname.startsWith('/socket.io/')) {
+      return true;
+    }
+    return [
+      '/api/auth/captcha',
+      '/api/auth/magic-link',
+      '/api/auth/open-app/sign-in',
+      '/api/auth/preflight',
+      '/api/auth/session/exchange',
+      '/api/auth/sign-in',
+      '/api/oauth/callback',
+      '/api/oauth/preflight',
+    ].includes(pathname);
+  } catch {
+    return false;
+  }
+}
+
 export function createAuthFetch(
  provider: AuthRequestProvider,
  rawFetch: typeof globalThis.fetch
@@ -26,7 +50,8 @@ export function createAuthFetch(
  return async (input: RequestInfo | URL, init?: RequestInit) => {
    const request = new Request(input, init);
    const retry = request.clone();
-    const endpoint = authEndpointForUrl(request.url);
+    const skipStoredAuth = shouldSkipStoredAuthToken(request.url);
+    const endpoint = skipStoredAuth ? null : authEndpointForUrl(request.url);
    const token = endpoint
      ? await provider.getValidAccessToken(endpoint)
      : null;
@@ -70,6 +95,7 @@ export function installAuthRequestProxy(provider: AuthRequestProvider) {
    private requestBody?: Document | XMLHttpRequestBodyInit | null;
    private replaying = false;
    private hasReplayed = false;
+    private skipStoredAuth = false;
    private sendVersion = 0;

    constructor() {
@@ -87,6 +113,7 @@ export function installAuthRequestProxy(provider: AuthRequestProvider) {
            this.status !== 401 ||
            this.replaying ||
            this.hasReplayed ||
+            this.skipStoredAuth ||
            !this.request?.async
          ) {
            return;
@@ -123,6 +150,7 @@ export function installAuthRequestProxy(provider: AuthRequestProvider) {
      this.requestBody = undefined;
      this.replaying = false;
      this.hasReplayed = false;
+      this.skipStoredAuth = shouldSkipStoredAuthToken(url.toString());
      xhrRequestUrls.set(this, url.toString());
      return super.open(
        method,
@@ -141,9 +169,11 @@ export function installAuthRequestProxy(provider: AuthRequestProvider) {
    override send(body?: Document | XMLHttpRequestBodyInit | null): void {
      this.requestBody = body;
      const requestUrl = xhrRequestUrls.get(this);
-      const endpoint = authEndpointForUrl(
-        requestUrl ?? globalThis.location.href
-      );
+      const targetUrl = requestUrl ?? globalThis.location.href;
+      this.skipStoredAuth = shouldSkipStoredAuthToken(targetUrl);
+      const endpoint = this.skipStoredAuth
+        ? null
+        : authEndpointForUrl(targetUrl);
      const sendVersion = this.sendVersion;

      const sendWithToken = (token: string | null) => {
@@ -168,7 +198,7 @@ export function installAuthRequestProxy(provider: AuthRequestProvider) {

    private async replayWithFreshToken() {
      const request = this.request;
-      if (!request) return this.failReplay();
+      if (!request || this.skipStoredAuth) return this.failReplay();
      const endpoint = authEndpointForUrl(request.url);
      if (!endpoint) return this.failReplay();
      const sendVersion = this.sendVersion;
@@ -22,6 +22,16 @@ export type WorkspaceSettingsRouteOptions = {
  scrollAnchor?: string;
 };

+export type NavigateToPageOptions = Omit<NavigateOptions, 'replace'> & {
+  search?: string | URLSearchParams;
+};
+
+const normalizeSearch = (search?: string | URLSearchParams) => {
+  const value = search?.toString();
+  if (!value) return '';
+  return value.startsWith('?') ? value : `?${value}`;
+};
+
 export function buildWorkspaceSettingsPath(
  workspaceId: string,
  options?: WorkspaceSettingsRouteOptions
@@ -84,11 +94,17 @@ export function useNavigateHelper() {
    (
      workspaceId: string,
      pageId: string,
-      logic: RouteLogic = RouteLogic.PUSH
+      logic: RouteLogic = RouteLogic.PUSH,
+      options?: NavigateToPageOptions
    ) => {
-      return navigate(`/workspace/${workspaceId}/${pageId}`, {
-        replace: logic === RouteLogic.REPLACE,
-      });
+      const { search, ...navigateOptions } = options ?? {};
+      return navigate(
+        `/workspace/${workspaceId}/${pageId}${normalizeSearch(search)}`,
+        {
+          ...navigateOptions,
+          replace: logic === RouteLogic.REPLACE,
+        }
+      );
    },
    [navigate]
  );
@@ -176,8 +192,13 @@ export function useNavigateHelper() {
  );

  const openPage = useCallback(
-    (workspaceId: string, pageId: string, logic?: RouteLogic) => {
-      return jumpToPage(workspaceId, pageId, logic);
+    (
+      workspaceId: string,
+      pageId: string,
+      logic?: RouteLogic,
+      options?: NavigateToPageOptions
+    ) => {
+      return jumpToPage(workspaceId, pageId, logic, options);
    },
    [jumpToPage]
  );
@@ -1,6 +1,9 @@
 import { Divider, IconButton, Menu, MenuItem } from '@affine/component';
 import { useAsyncCallback } from '@affine/core/components/hooks/affine-async-hooks';
-import { useNavigateHelper } from '@affine/core/components/hooks/use-navigate-helper';
+import {
+  RouteLogic,
+  useNavigateHelper,
+} from '@affine/core/components/hooks/use-navigate-helper';
 import { useWorkspaceInfo } from '@affine/core/components/hooks/use-workspace-info';
 import { WorkspaceAvatar } from '@affine/core/components/workspace-avatar';
 import {
@@ -13,7 +16,6 @@ import { GlobalDialogService } from '@affine/core/modules/dialogs';
 import { GlobalContextService } from '@affine/core/modules/global-context';
 import {
  type WorkspaceMetadata,
-  WorkspaceService,
  WorkspacesService,
 } from '@affine/core/modules/workspace';
 import { useI18n } from '@affine/i18n';
@@ -45,7 +47,7 @@ const WorkspaceItem = ({
    <li className={styles.wsItem}>
      <button className={clsx(styles.wsCard, className)} {...attrs}>
        <WorkspaceAvatar
-          key={workspace.id}
+          key={`${workspace.flavour}:${workspace.id}`}
          meta={workspace}
          rounded={6}
          data-testid="workspace-avatar"
@@ -71,7 +73,7 @@ export const WorkspaceList = (props: WorkspaceListProps) => {

  return workspaceList.map(item => (
    <WorkspaceItem
-      key={item.id}
+      key={`${item.flavour}:${item.id}`}
      workspace={item}
      onClick={() => props.onClick(item)}
    />
@@ -276,12 +278,18 @@ const AddServer = () => {
 };

 export const SelectorMenu = ({ onClose }: { onClose?: () => void }) => {
-  const currentWorkspace = useService(WorkspaceService).workspace;
  const workspacesService = useService(WorkspacesService);
  const workspaces = useLiveData(workspacesService.list.workspaces$);
  const serversService = useService(ServersService);
+  const globalContextService = useService(GlobalContextService);
  const { jumpToPage } = useNavigateHelper();

+  const currentWorkspaceId = useLiveData(
+    globalContextService.globalContext.workspaceId.$
+  );
+  const currentWorkspaceFlavour = useLiveData(
+    globalContextService.globalContext.workspaceFlavour.$
+  );
  const servers = useLiveData(serversService.servers$);
  const affineCloudServer = useMemo(
    () => servers.find(s => s.id === 'affine-cloud') as Server,
@@ -311,12 +319,29 @@ export const SelectorMenu = ({ onClose }: { onClose?: () => void }) => {
  const handleClickWorkspace = useCallback(
    (workspaceMetadata: WorkspaceMetadata) => {
      const id = workspaceMetadata.id;
-      if (id !== currentWorkspace?.id) {
-        jumpToPage(id, 'home');
+      const isCurrentWorkspace =
+        id === currentWorkspaceId &&
+        workspaceMetadata.flavour === currentWorkspaceFlavour;
+      if (!isCurrentWorkspace) {
+        const server = servers.find(
+          server => server.id === workspaceMetadata.flavour
+        );
+        if (workspaceMetadata.flavour !== 'local' && !server) {
+          return;
+        }
+        const searchParams = new URLSearchParams({
+          flavour: workspaceMetadata.flavour,
+        });
+        if (workspaceMetadata.flavour !== 'local' && server) {
+          searchParams.set('server', server.baseUrl);
+        }
+        jumpToPage(id, 'home', RouteLogic.PUSH, {
+          search: searchParams,
+        });
      }
      onClose?.();
    },
-    [currentWorkspace.id, jumpToPage, onClose]
+    [currentWorkspaceFlavour, currentWorkspaceId, jumpToPage, onClose, servers]
  );

  return (
@@ -102,9 +102,54 @@ export const Component = () => {
  const [workspaceNotFound, setWorkspaceNotFound] = useState(false);
  const listLoading = useLiveData(workspacesService.list.isRevalidating$);
  const workspaces = useLiveData(workspacesService.list.workspaces$);
+
+  const serverSearchParam = searchParams.get('server');
+  const flavourSearchParam = searchParams.get('flavour');
+  const serverFromSearchParams = useLiveData(
+    serverSearchParam
+      ? serversService.serverByBaseUrl$(serverSearchParam)
+      : undefined
+  );
  const meta = useMemo(() => {
-    return workspaces.find(({ id }) => id === params.workspaceId);
-  }, [workspaces, params.workspaceId]);
+    const workspaceId = params.workspaceId;
+    if (!workspaceId) {
+      return undefined;
+    }
+
+    const findByFlavour = (flavour: string) =>
+      workspaces.find(
+        workspace =>
+          workspace.id === workspaceId && workspace.flavour === flavour
+      );
+
+    if (flavourSearchParam) {
+      return findByFlavour(flavourSearchParam);
+    }
+
+    if (serverSearchParam) {
+      if (!serverFromSearchParams) {
+        return undefined;
+      }
+      return findByFlavour(serverFromSearchParams.id);
+    }
+
+    const lastWorkspaceFlavour = localStorage.getItem('last_workspace_flavour');
+    if (lastWorkspaceFlavour) {
+      const lastWorkspace = findByFlavour(lastWorkspaceFlavour);
+      if (lastWorkspace) {
+        return lastWorkspace;
+      }
+    }
+
+    const matches = workspaces.filter(({ id }) => id === workspaceId);
+    return matches.length === 1 ? matches[0] : undefined;
+  }, [
+    flavourSearchParam,
+    params.workspaceId,
+    serverSearchParam,
+    serverFromSearchParams,
+    workspaces,
+  ]);

  // if listLoading is false, we can show 404 page, otherwise we should show loading page.
  useEffect(() => {
@@ -135,19 +180,16 @@ export const Component = () => {
    return;
  }, [listLoading, meta, workspaceNotFound, workspacesService]);

-  // server search params
-  const serverFromSearchParams = useLiveData(
-    searchParams.has('server')
-      ? serversService.serverByBaseUrl$(searchParams.get('server') as string)
-      : undefined
-  );
  // server from workspace
  const serverFromWorkspace = useLiveData(
    meta?.flavour && meta.flavour !== 'local'
      ? serversService.server$(meta?.flavour)
      : undefined
  );
-  const server = serverFromWorkspace ?? serverFromSearchParams;
+  const server =
+    meta?.flavour === 'local'
+      ? undefined
+      : (serverFromWorkspace ?? serverFromSearchParams);

  if (workspaceNotFound) {
    if (
@@ -86,6 +86,7 @@ export const WorkspaceLayout = ({
        })
      );
      localStorage.setItem('last_workspace_id', workspace.id);
+      localStorage.setItem('last_workspace_flavour', workspace.flavour);
      globalContextService.globalContext.workspaceId.set(workspace.id);
      if (workspaceServer) {
        globalContextService.globalContext.serverId.set(workspaceServer.id);
@@ -41,7 +41,8 @@ export function useBindWorkbenchToBrowserRouter(

      const newBrowserLocation = viewLocationToBrowserLocation(
        update.location,
-        basename
+        basename,
+        browserLocation.search
      );

      navigate(newBrowserLocation, {
@@ -97,12 +98,44 @@ function browserLocationToViewLocation(
  };
 }

+function preserveWorkspaceContextSearch(
+  nextSearch: string,
+  currentSearch: string
+) {
+  const nextParams = new URLSearchParams(nextSearch);
+  const currentParams = new URLSearchParams(currentSearch);
+  const currentFlavour = currentParams.get('flavour');
+  const nextFlavour = nextParams.get('flavour');
+
+  if (
+    !nextParams.has('flavour') &&
+    !nextParams.has('server') &&
+    currentFlavour
+  ) {
+    nextParams.set('flavour', currentFlavour);
+  }
+
+  const resolvedNextFlavour = nextParams.get('flavour');
+  const shouldPreserveServer =
+    resolvedNextFlavour !== 'local' &&
+    (!nextFlavour || !currentFlavour || nextFlavour === currentFlavour);
+  const currentServer = currentParams.get('server');
+  if (!nextParams.has('server') && currentServer && shouldPreserveServer) {
+    nextParams.set('server', currentServer);
+  }
+
+  const search = nextParams.toString();
+  return search ? `?${search}` : '';
+}
+
 function viewLocationToBrowserLocation(
  location: Location,
-  basename: string
+  basename: string,
+  currentSearch: string
 ): Location {
  return {
    ...location,
    pathname: `${basename}${location.pathname}`,
+    search: preserveWorkspaceContextSearch(location.search, currentSearch),
  };
 }
@@ -4,6 +4,7 @@ import { ObjectPool, Service } from '@toeverything/infra';

 import type { Workspace } from '../entities/workspace';
 import { WorkspaceInitialized } from '../events';
+import type { WorkspaceMetadata } from '../metadata';
 import type { WorkspaceOpenOptions } from '../open-options';
 import { WorkspaceScope } from '../scopes/workspace';
 import type { WorkspaceFlavoursService } from './flavours';
@@ -13,6 +14,9 @@ import { WorkspaceService } from './workspace';

 const logger = new DebugLogger('affine:workspace-repository');

+const getWorkspacePoolKey = (metadata: WorkspaceMetadata) =>
+  `${metadata.flavour}:${metadata.id}`;
+
 export class WorkspaceRepositoryService extends Service {
  constructor(
    private readonly flavoursService: WorkspaceFlavoursService,
@@ -58,7 +62,7 @@ export class WorkspaceRepositoryService extends Service {
      };
    }

-    const exist = this.pool.get(options.metadata.id);
+    const exist = this.pool.get(getWorkspacePoolKey(options.metadata));
    if (exist) {
      return {
        workspace: exist.obj,
@@ -68,7 +72,7 @@ export class WorkspaceRepositoryService extends Service {

    const workspace = this.instantiate(options, customEngineWorkerInitOptions);

-    const ref = this.pool.put(workspace.meta.id, workspace);
+    const ref = this.pool.put(getWorkspacePoolKey(workspace.meta), workspace);

    return {
      workspace: ref.obj,
@@ -76,9 +80,13 @@ export class WorkspaceRepositoryService extends Service {
    };
  };

-  openByWorkspaceId = (workspaceId: string) => {
-    const workspaceMetadata =
-      this.workspacesListService.list.workspace$(workspaceId).value;
+  openByWorkspaceId = (workspaceId: string, flavour?: string | null) => {
+    const workspaceMetadata = flavour
+      ? this.workspacesListService.list.workspaces$.value.find(
+          workspace =>
+            workspace.id === workspaceId && workspace.flavour === flavour
+        )
+      : this.workspacesListService.list.workspace$(workspaceId).value;
    return workspaceMetadata && this.open({ metadata: workspaceMetadata });
  };