fix(workflow): preserve latest collaboration session (#39646)

web/app/components/workflow/collaboration/core/__tests__/collaboration-manager.runtime-loading.spec.ts

@@ -16,6 +16,15 @@ const createMockSocket = (): Socket =>
     off: vi.fn(),
   }) as unknown as Socket
 
+const createDeferred = <T>() => {
+  let resolve!: (value: T) => void
+  const promise = new Promise<T>((resolvePromise) => {
+    resolve = resolvePromise
+  })
+
+  return { promise, resolve }
+}
+
 const loadCollaborationModules = async () => {
   const [{ CollaborationManager }, { webSocketClient }] = await Promise.all([
     import('../collaboration-manager'),
@@ -40,10 +49,11 @@ describe('CollaborationManager CRDT runtime loading', () => {
     expect(manager.isConnected()).toBe(false)
   })
 
-  it('does not create connection state when the runtime fails to load', async () => {
+  it('does not create connection state when the runtime fails to load and allows a retry', async () => {
     const { CollaborationManager, webSocketClient } = await loadCollaborationModules()
     const manager = new CollaborationManager()
     const runtimeError = new Error('runtime-load-failed')
+    const retryError = new Error('runtime-retry-failed')
     const loadRuntimeSpy = vi
       .spyOn(
         manager as unknown as {
@@ -51,7 +61,8 @@ describe('CollaborationManager CRDT runtime loading', () => {
         },
         'loadCrdtRuntime',
       )
-      .mockRejectedValue(runtimeError)
+      .mockRejectedValueOnce(runtimeError)
+      .mockRejectedValueOnce(retryError)
     const connectSpy = vi.spyOn(webSocketClient, 'connect')
 
     await expect(manager.connect('app-runtime-failure')).rejects.toBe(runtimeError)
@@ -59,11 +70,22 @@ describe('CollaborationManager CRDT runtime loading', () => {
     expect(loadRuntimeSpy).toHaveBeenCalledTimes(1)
     expect(connectSpy).not.toHaveBeenCalled()
     expect(manager.isConnected()).toBe(false)
+
+    await expect(manager.connect('app-runtime-failure')).rejects.toBe(retryError)
+
+    expect(loadRuntimeSpy).toHaveBeenCalledTimes(2)
+    expect(connectSpy).not.toHaveBeenCalled()
   })
 
   it('initializes one session for concurrent consumers of the same app', async () => {
     const { CollaborationManager, webSocketClient } = await loadCollaborationModules()
     const manager = new CollaborationManager()
+    const loadRuntimeSpy = vi.spyOn(
+      manager as unknown as {
+        loadCrdtRuntime: () => Promise<(typeof import('../crdt-runtime'))['crdtRuntime']>
+      },
+      'loadCrdtRuntime',
+    )
     const socket = createMockSocket()
     const connectSpy = vi.spyOn(webSocketClient, 'connect').mockReturnValue(socket)
     const disconnectSpy = vi
@@ -76,6 +98,7 @@ describe('CollaborationManager CRDT runtime loading', () => {
     ])
 
     expect(firstConnectionId).not.toBe(secondConnectionId)
+    expect(loadRuntimeSpy).toHaveBeenCalledTimes(1)
     expect(connectSpy).toHaveBeenCalledTimes(1)
     expect(loroModuleState.evaluations).toBe(1)
 
@@ -85,4 +108,61 @@ describe('CollaborationManager CRDT runtime loading', () => {
     manager.disconnect(secondConnectionId)
     expect(disconnectSpy).toHaveBeenCalledWith('app-concurrent')
   })
+
+  it('keeps the latest app session when the app changes during runtime loading', async () => {
+    const { CollaborationManager, webSocketClient } = await loadCollaborationModules()
+    const { crdtRuntime } = await import('../crdt-runtime')
+    const manager = new CollaborationManager()
+    const runtime = createDeferred<typeof crdtRuntime>()
+    vi.spyOn(
+      manager as unknown as {
+        loadCrdtRuntime: () => Promise<typeof crdtRuntime>
+      },
+      'loadCrdtRuntime',
+    ).mockReturnValue(runtime.promise)
+    const socket = createMockSocket()
+    const connectSpy = vi.spyOn(webSocketClient, 'connect').mockReturnValue(socket)
+    const disconnectSpy = vi
+      .spyOn(webSocketClient, 'disconnect')
+      .mockImplementation(() => undefined)
+
+    const firstConnection = manager.connect('app-first')
+    const secondConnection = manager.connect('app-second')
+
+    runtime.resolve(crdtRuntime)
+    const [firstConnectionId, secondConnectionId] = await Promise.all([
+      firstConnection,
+      secondConnection,
+    ])
+
+    expect(connectSpy).toHaveBeenCalledTimes(1)
+    expect(connectSpy).toHaveBeenCalledWith('app-second')
+
+    manager.disconnect(firstConnectionId)
+    expect(disconnectSpy).not.toHaveBeenCalled()
+
+    manager.disconnect(secondConnectionId)
+    expect(disconnectSpy).toHaveBeenCalledWith('app-second')
+  })
+
+  it('does not initialize a pending session after the manager is destroyed', async () => {
+    const { CollaborationManager, webSocketClient } = await loadCollaborationModules()
+    const { crdtRuntime } = await import('../crdt-runtime')
+    const manager = new CollaborationManager()
+    const runtime = createDeferred<typeof crdtRuntime>()
+    vi.spyOn(
+      manager as unknown as {
+        loadCrdtRuntime: () => Promise<typeof crdtRuntime>
+      },
+      'loadCrdtRuntime',
+    ).mockReturnValue(runtime.promise)
+    const connectSpy = vi.spyOn(webSocketClient, 'connect')
+
+    const connection = manager.connect('app-destroyed')
+    manager.destroy()
+    runtime.resolve(crdtRuntime)
+
+    await connection
+    expect(connectSpy).not.toHaveBeenCalled()
+  })
 })

web/app/components/workflow/collaboration/core/__tests__/collaboration-manager.socket-and-subscriptions.spec.ts

@@ -907,13 +907,17 @@ describe('CollaborationManager socket and subscription behavior', () => {
     expect(secondConnectionId).toBeTruthy()
     expect(disconnectSpy).not.toHaveBeenCalled()
 
-    await manager.connect('app-2', reactFlowStore)
+    const thirdConnectionId = await manager.connect('app-2', reactFlowStore)
     expect(disconnectSpy).toHaveBeenCalledWith('app-1')
     expect(internals.currentAppId).toBe('app-2')
 
     internals.isLeader = true
     manager.disconnect(secondConnectionId)
     manager.disconnect(firstConnectionId)
+    expect(disconnectSpy).not.toHaveBeenCalledWith('app-2')
+    expect(internals.currentAppId).toBe('app-2')
+
+    manager.disconnect(thirdConnectionId)
     expect(disconnectSpy).toHaveBeenCalledWith('app-2')
     expect(eventEmitSpy).toHaveBeenCalledWith('leaderChange', false)
     expect(internals.currentAppId).toBeNull()

web/app/components/workflow/collaboration/core/collaboration-manager.ts

@@ -156,6 +156,9 @@ const toUint8Array = (value: unknown): Uint8Array | null => {
 
 export class CollaborationManager {
   private crdtRuntime: CrdtRuntime | null = null
+  private crdtRuntimePromise: Promise<CrdtRuntime> | null = null
+  private targetAppId: string | null = null
+  private connectGeneration = 0
   private doc: LoroDoc | null = null
   private undoManager: UndoManager | null = null
   private provider: CRDTProvider | null = null
@@ -538,6 +541,20 @@ export class CollaborationManager {
     return crdtRuntime
   }
 
+  private async ensureCrdtRuntime(): Promise<void> {
+    if (this.crdtRuntime) return
+
+    const runtimePromise = this.crdtRuntimePromise ?? this.loadCrdtRuntime()
+    this.crdtRuntimePromise = runtimePromise
+
+    try {
+      this.crdtRuntime = await runtimePromise
+    } catch (error) {
+      if (this.crdtRuntimePromise === runtimePromise) this.crdtRuntimePromise = null
+      throw error
+    }
+  }
+
   private getCrdtRuntime(): CrdtRuntime {
     if (!this.crdtRuntime) throw new Error('CRDT runtime not initialized')
     return this.crdtRuntime
@@ -639,21 +656,31 @@ export class CollaborationManager {
   }
 
   async connect(appId: string, reactFlowStore?: ReactFlowStore): Promise<string> {
-    this.crdtRuntime ??= await this.loadCrdtRuntime()
-
     const connectionId = Math.random().toString(36).substring(2, 11)
+    if (this.targetAppId !== appId) {
+      this.targetAppId = appId
+      this.connectGeneration += 1
+    }
+    const connectGeneration = this.connectGeneration
 
-    this.activeConnections.add(connectionId)
+    if (!this.crdtRuntime) await this.ensureCrdtRuntime()
+
+    if (connectGeneration !== this.connectGeneration || this.targetAppId !== appId)
+      return connectionId
 
     if (this.currentAppId === appId && this.doc) {
       // Already connected to the same app, only update store if provided and we don't have one
       if (reactFlowStore && !this.reactFlowStore) this.reactFlowStore = reactFlowStore
+      this.activeConnections.add(connectionId)
 
       return connectionId
     }
 
     // Only disconnect if switching to a different app
-    if (this.currentAppId && this.currentAppId !== appId) this.forceDisconnect()
+    if (this.currentAppId && this.currentAppId !== appId)
+      this.forceDisconnect({ preserveConnectIntent: true })
+
+    this.activeConnections.add(connectionId)
 
     this.hasEstablishedConnection = false
     this.currentAppId = appId
@@ -685,7 +712,7 @@ export class CollaborationManager {
   }
 
   disconnect = (connectionId?: string): void => {
-    if (connectionId) this.activeConnections.delete(connectionId)
+    if (connectionId && !this.activeConnections.delete(connectionId)) return
 
     // Only disconnect when no more connections
     if (this.activeConnections.size === 0) this.forceDisconnect()
@@ -699,7 +726,9 @@ export class CollaborationManager {
     this.pendingWorkflowSyncRequests.clear()
   }
 
-  private forceDisconnect = (): void => {
+  private forceDisconnect = ({
+    preserveConnectIntent = false,
+  }: { preserveConnectIntent?: boolean } = {}): void => {
     if (this.currentAppId) webSocketClient.disconnect(this.currentAppId)
 
     this.clearInitialSyncRetry()
@@ -740,6 +769,10 @@ export class CollaborationManager {
     if (wasLeader) this.eventEmitter.emit('leaderChange', false)
 
     this.activeConnections.clear()
+    if (!preserveConnectIntent) {
+      this.targetAppId = null
+      this.connectGeneration += 1
+    }
     this.eventEmitter.removeAllListeners()
   }