Refresh stale IMS access token on collab WebSocket auth failures

blocks/edit/prose/index.js

@@ -23,6 +23,7 @@ import {
 } from 'da-y-wrapper';
 
 import { getSchema } from 'da-parser';
+import { getNx } from '../../../scripts/utils.js';
 import { COLLAB_ORIGIN, DA_ORIGIN } from '../../shared/constants.js';
 import { daFetch, getAuthToken } from '../../shared/utils.js';
 import { getDiffClass, checkForLocNodes, addActiveView } from './diff/diff-utils.js';
@@ -54,6 +55,42 @@ export async function createConnection(path) {
   // (exponential backoff starting with 100ms) and then every 30s.
   provider.maxBackoffTime = 30000;
 
+  // y-websocket re-reads provider.protocols on each reconnect, so swapping in a
+  // fresh IMS token here is enough; no provider rebuild needed.
+  // Server close codes: 4401 = token expired (refresh + retry), 4403 = forbidden (stop).
+  let lastSentToken = token || null;
+  provider.on('connection-close', async (event) => {
+    if (event?.code === 4403) {
+      provider.shouldConnect = false;
+      return;
+    }
+    if (event?.code === 4401) {
+      // Force imslib to attempt a refresh before deciding to give up.
+      try { await window.adobeIMS?.refreshToken?.(); } catch { /* ignore */ }
+      const fresh = await getAuthToken();
+      if (!fresh || fresh === lastSentToken) {
+        // No new token to try — retrying would loop on the same 4401 forever.
+        provider.shouldConnect = false;
+        // If the user expected to be signed in, route them through IMS sign-in
+        // so collab can recover. Matches daFetch's 401 handling. Anonymous users
+        // (no nx-ims flag) hitting a private doc are left disconnected.
+        if (localStorage.getItem('nx-ims')) {
+          try {
+            const { handleSignIn } = await import(`${getNx()}/utils/ims.js`);
+            handleSignIn();
+          } catch { /* nothing to do */ }
+        }
+        return;
+      }
+      provider.protocols = ['yjs', fresh];
+      lastSentToken = fresh;
+      return;
+    }
+    const fresh = await getAuthToken();
+    provider.protocols = fresh ? ['yjs', fresh] : ['yjs'];
+    lastSentToken = fresh;
+  });
+
   return { wsProvider: provider, ydoc };
 }
 

blocks/shared/utils.js

@@ -24,6 +24,11 @@ export async function getAuthToken() {
   if (!localStorage.getItem('nx-ims')) {
     return null;
   }
+  // imslib auto-refreshes its internal token; reading it live avoids returning the
+  // page-load snapshot that nx's loadIms() captured once in its onReady handler.
+  if (window.adobeIMS?.getAccessToken) {
+    return window.adobeIMS.getAccessToken()?.token || null;
+  }
   const ims = await initIms();
   return ims?.accessToken?.token || null;
 }

test/fixtures/nx/utils/ims.js

@@ -5,6 +5,8 @@ export async function loadIms() {
 }
 
 export function handleSignIn() {
-  // Mock implementation
-  return Promise.resolve();
+  // Mirrors the real nx handleSignIn so tests can observe the effect via
+  // window.adobeIMS.signIn().
+  localStorage.setItem('nx-ims', true);
+  if (window.adobeIMS?.signIn) window.adobeIMS.signIn();
 }

test/unit/blocks/edit/prose/index.test.js

@@ -74,7 +74,11 @@ describe('prose/index createConnection', () => {
     window.localStorage.removeItem('nx-ims');
   });
   afterEach(() => {
-    if (savedNxIms) window.localStorage.setItem('nx-ims', savedNxIms);
+    if (savedNxIms) {
+      window.localStorage.setItem('nx-ims', savedNxIms);
+    } else {
+      window.localStorage.removeItem('nx-ims');
+    }
   });
 
   it('Returns a wsProvider and a Y.Doc with maxBackoffTime configured', async () => {
@@ -87,6 +91,171 @@ describe('prose/index createConnection', () => {
     result.wsProvider.destroy?.();
     result.ydoc.destroy();
   });
+
+  it('Refreshes protocols with the live IMS token on connection-close', async () => {
+    window.localStorage.setItem('nx-ims', 'true');
+    const savedIMS = window.adobeIMS;
+    let tokenIndex = 0;
+    const tokens = ['T-initial', 'T-rotated'];
+    window.adobeIMS = {
+      getAccessToken: () => {
+        const t = tokens[tokenIndex];
+        tokenIndex += 1;
+        return { token: t };
+      },
+    };
+
+    try {
+      const { wsProvider, ydoc } = await createConnection('https://admin.da.live/source/org/repo/page.html');
+      expect(wsProvider.protocols).to.deep.equal(['yjs', 'T-initial']);
+
+      // Simulate a server-signalled auth failure
+      wsProvider.emit('connection-close', [{ code: 4401, reason: 'auth' }, wsProvider]);
+      await new Promise((r) => { setTimeout(r, 0); });
+
+      expect(wsProvider.protocols).to.deep.equal(['yjs', 'T-rotated']);
+      expect(wsProvider.shouldConnect).to.equal(true);
+
+      wsProvider.disconnect();
+      wsProvider.destroy?.();
+      ydoc.destroy();
+    } finally {
+      if (savedIMS === undefined) delete window.adobeIMS; else window.adobeIMS = savedIMS;
+    }
+  });
+
+  it('Stops reconnecting on a 4403 forbidden close', async () => {
+    window.localStorage.setItem('nx-ims', 'true');
+    const savedIMS = window.adobeIMS;
+    window.adobeIMS = { getAccessToken: () => ({ token: 'T-initial' }) };
+
+    try {
+      const { wsProvider, ydoc } = await createConnection('https://admin.da.live/source/org/repo/page.html');
+      expect(wsProvider.shouldConnect).to.equal(true);
+
+      wsProvider.emit('connection-close', [{ code: 4403, reason: 'forbidden' }, wsProvider]);
+      await new Promise((r) => { setTimeout(r, 0); });
+
+      expect(wsProvider.shouldConnect).to.equal(false);
+
+      wsProvider.disconnect();
+      wsProvider.destroy?.();
+      ydoc.destroy();
+    } finally {
+      if (savedIMS === undefined) delete window.adobeIMS; else window.adobeIMS = savedIMS;
+    }
+  });
+
+  it('Stops reconnecting on 4401 when imslib cannot produce a token, triggers sign-in', async () => {
+    window.localStorage.setItem('nx-ims', 'true');
+    const savedIMS = window.adobeIMS;
+    let refreshCalls = 0;
+    let signInCalls = 0;
+    window.adobeIMS = {
+      getAccessToken: () => ({ token: 'T-initial' }),
+      refreshToken: async () => { refreshCalls += 1; },
+      signIn: () => { signInCalls += 1; },
+    };
+
+    try {
+      const { wsProvider, ydoc } = await createConnection('https://admin.da.live/source/org/repo/page.html');
+
+      // After construction, simulate imslib losing the token (SSO expired)
+      window.adobeIMS.getAccessToken = () => null;
+
+      wsProvider.emit('connection-close', [{ code: 4401, reason: 'auth' }, wsProvider]);
+      // Allow the dynamic import + signIn call to settle
+      await new Promise((r) => { setTimeout(r, 50); });
+
+      expect(refreshCalls).to.equal(1);
+      expect(wsProvider.shouldConnect).to.equal(false);
+      expect(signInCalls).to.equal(1);
+
+      wsProvider.disconnect();
+      wsProvider.destroy?.();
+      ydoc.destroy();
+    } finally {
+      if (savedIMS === undefined) delete window.adobeIMS; else window.adobeIMS = savedIMS;
+    }
+  });
+
+  it('Anonymous user hitting a private doc bails on 4401 without sign-in redirect', async () => {
+    window.localStorage.removeItem('nx-ims');
+    const savedIMS = window.adobeIMS;
+    let signInCalls = 0;
+    window.adobeIMS = {
+      getAccessToken: () => null,
+      refreshToken: async () => {},
+      signIn: () => { signInCalls += 1; },
+    };
+
+    try {
+      const { wsProvider, ydoc } = await createConnection('https://admin.da.live/source/org/repo/page.html');
+      expect(wsProvider.protocols).to.deep.equal(['yjs']);
+
+      wsProvider.emit('connection-close', [{ code: 4401, reason: 'auth' }, wsProvider]);
+      await new Promise((r) => { setTimeout(r, 50); });
+
+      expect(wsProvider.shouldConnect).to.equal(false);
+      expect(signInCalls).to.equal(0);
+
+      wsProvider.disconnect();
+      wsProvider.destroy?.();
+      ydoc.destroy();
+    } finally {
+      if (savedIMS === undefined) delete window.adobeIMS; else window.adobeIMS = savedIMS;
+    }
+  });
+
+  it('Stops reconnecting on 4401 when imslib hands back the same stale token', async () => {
+    window.localStorage.setItem('nx-ims', 'true');
+    const savedIMS = window.adobeIMS;
+    window.adobeIMS = {
+      getAccessToken: () => ({ token: 'T-same' }),
+      refreshToken: async () => {},
+    };
+
+    try {
+      const { wsProvider, ydoc } = await createConnection('https://admin.da.live/source/org/repo/page.html');
+      expect(wsProvider.protocols).to.deep.equal(['yjs', 'T-same']);
+
+      wsProvider.emit('connection-close', [{ code: 4401, reason: 'auth' }, wsProvider]);
+      await new Promise((r) => { setTimeout(r, 0); });
+
+      // No new token to try — don't loop.
+      expect(wsProvider.shouldConnect).to.equal(false);
+
+      wsProvider.disconnect();
+      wsProvider.destroy?.();
+      ydoc.destroy();
+    } finally {
+      if (savedIMS === undefined) delete window.adobeIMS; else window.adobeIMS = savedIMS;
+    }
+  });
+
+  it('Non-auth close with no token reconnects as anonymous', async () => {
+    window.localStorage.removeItem('nx-ims');
+    const savedIMS = window.adobeIMS;
+    delete window.adobeIMS;
+
+    try {
+      const { wsProvider, ydoc } = await createConnection('https://admin.da.live/source/org/repo/page.html');
+      expect(wsProvider.protocols).to.deep.equal(['yjs']);
+
+      // Simulate a generic network drop — no custom code
+      wsProvider.emit('connection-close', [{ code: 1006 }, wsProvider]);
+      await new Promise((r) => { setTimeout(r, 0); });
+
+      expect(wsProvider.protocols).to.deep.equal(['yjs']);
+      expect(wsProvider.shouldConnect).to.equal(true);
+
+      wsProvider.disconnect();
+      wsProvider.destroy?.();
+      ydoc.destroy();
+    } finally {
+      if (savedIMS === undefined) delete window.adobeIMS; else window.adobeIMS = savedIMS;
+    }
+  });
 });
 
 describe('prose/index createAwarenessStatusWidget', () => {

test/unit/blocks/shared/utils.test.js

@@ -11,6 +11,7 @@ import {
   getSidekickConfig,
   sanitizeName,
   fetchDaConfigs,
+  getAuthToken,
 } from '../../../../blocks/shared/utils.js';
 
 describe('getSheetByIndex', () => {
@@ -178,6 +179,57 @@ describe('sanitizeName', () => {
   });
 });
 
+describe('getAuthToken', () => {
+  let savedLocalStorage;
+  let savedAdobeIMS;
+
+  beforeEach(() => {
+    savedLocalStorage = window.localStorage.getItem('nx-ims');
+    savedAdobeIMS = window.adobeIMS;
+  });
+
+  afterEach(() => {
+    if (savedLocalStorage) {
+      window.localStorage.setItem('nx-ims', savedLocalStorage);
+    } else {
+      window.localStorage.removeItem('nx-ims');
+    }
+    if (savedAdobeIMS === undefined) {
+      delete window.adobeIMS;
+    } else {
+      window.adobeIMS = savedAdobeIMS;
+    }
+  });
+
+  it('Returns null when nx-ims is not set', async () => {
+    window.localStorage.removeItem('nx-ims');
+    window.adobeIMS = { getAccessToken: () => ({ token: 'should-not-see' }) };
+    expect(await getAuthToken()).to.be.null;
+  });
+
+  it('Reads the live token from window.adobeIMS.getAccessToken when present', async () => {
+    window.localStorage.setItem('nx-ims', 'true');
+    let calls = 0;
+    const tokens = ['T1', 'T2'];
+    window.adobeIMS = {
+      getAccessToken: () => {
+        const t = tokens[calls];
+        calls += 1;
+        return { token: t };
+      },
+    };
+
+    expect(await getAuthToken()).to.equal('T1');
+    expect(await getAuthToken()).to.equal('T2');
+  });
+
+  it('Returns null when getAccessToken returns null (signed out)', async () => {
+    window.localStorage.setItem('nx-ims', 'true');
+    window.adobeIMS = { getAccessToken: () => null };
+    expect(await getAuthToken()).to.be.null;
+  });
+});
+
 describe('daFetch', () => {
   let savedFetch;
   let savedLocalStorage;