fix: OIDC/SSO Token + Media issues

.changeset/fix-media-error-handling.md

@@ -0,0 +1,5 @@
+---
+default: patch
+---
+
+Fix unhandled promise rejections from encrypted media decrypt failures, and wire OIDC token refresh to propagate new tokens to the service worker so authenticated media requests no longer 401 after token expiry.

src/app/hooks/useAsyncCallback.test.tsx

@@ -30,7 +30,7 @@ describe('useAsyncCallback', () => {
     );
 
     await act(async () => {
-      await result.current[1]().catch(() => {});
+      await expect(result.current[1]()).rejects.toBe(boom);
     });
 
     expect(result.current[0]).toEqual({ status: AsyncStatus.Error, error: boom });

src/app/hooks/useAsyncCallback.ts

@@ -73,6 +73,9 @@ export const useAsync = <TData, TError, TArgs extends unknown[]>(
             });
           });
         }
+        // Re-throw so .then()/.catch() callers see the rejection and success
+        // handlers are skipped. Fire-and-forget unhandled-rejection warnings are
+        // suppressed at the useAsyncCallback level via a no-op .catch wrapper.
         throw e;
       }
 
@@ -102,7 +105,19 @@ export const useAsyncCallback = <TData, TError, TArgs extends unknown[]>(
     status: AsyncStatus.Idle,
   });
 
-  const callback = useAsync(asyncCallback, setState);
+  const innerCallback = useAsync(asyncCallback, setState);
+
+  // Re-throw preserves rejection for callers that await/chain; the no-op .catch
+  // suppresses "Uncaught (in promise)" for fire-and-forget call sites (e.g.
+  // loadSrc() in a useEffect) without swallowing the error from intentional callers.
+  const callback = useCallback(
+    (...args: TArgs): Promise<TData> => {
+      const p = innerCallback(...args);
+      p.catch(() => {});
+      return p;
+    },
+    [innerCallback]
+  ) as AsyncCallback<TArgs, TData>;
 
   return [state, callback, setState];
 };

src/app/pages/auth/login/loginUtil.ts

@@ -156,6 +156,8 @@ export const useLoginComplete = (data?: CustomLoginResponse) => {
         userId: loginRes.user_id,
         deviceId: loginRes.device_id,
         accessToken: loginRes.access_token,
+        ...(loginRes.refresh_token != null && { refreshToken: loginRes.refresh_token }),
+        ...(loginRes.expires_in_ms != null && { expiresInMs: loginRes.expires_in_ms }),
       };
       setSessions({ type: 'PUT', session: newSession });
       setActiveSessionId(loginRes.user_id);

src/app/pages/client/ClientRoot.tsx

@@ -201,11 +201,21 @@ export function ClientRoot({ children }: ClientRootProps) {
       }
       await clearMismatchedStores();
       log.log('initClient for', activeSession.userId);
-      const newMx = await initClient(activeSession);
+      const newMx = await initClient(activeSession, (newAccessToken, newRefreshToken) => {
+        setSessions({
+          type: 'PUT',
+          session: {
+            ...activeSession,
+            accessToken: newAccessToken,
+            ...(newRefreshToken !== undefined && { refreshToken: newRefreshToken }),
+          },
+        });
+        pushSessionToSW(activeSession.baseUrl, newAccessToken, activeSession.userId);
+      });
       loadedUserIdRef.current = activeSession.userId;
-      pushSessionToSW(activeSession.baseUrl, activeSession.accessToken);
+      pushSessionToSW(activeSession.baseUrl, activeSession.accessToken, activeSession.userId);
       return newMx;
-    }, [activeSession, activeSessionId, setActiveSessionId])
+    }, [activeSession, activeSessionId, setActiveSessionId, setSessions])
   );
 
   const mx = loadState.status === AsyncStatus.Success ? loadState.data : undefined;

src/client/initMatrix.ts

@@ -267,7 +267,10 @@ export const clearMismatchedStores = async (): Promise<void> => {
   );
 };
 
-const buildClient = async (session: Session): Promise<MatrixClient> => {
+const buildClient = async (
+  session: Session,
+  onTokenRefresh?: (newAccessToken: string, newRefreshToken?: string) => void
+): Promise<MatrixClient> => {
   const storeName = getSessionStoreName(session);
 
   const indexedDBStore = new IndexedDBStore({
@@ -278,6 +281,11 @@ const buildClient = async (session: Session): Promise<MatrixClient> => {
 
   const legacyCryptoStore = new IndexedDBCryptoStore(global.indexedDB, storeName.crypto);
 
+  // Pre-allocate a slot for the MatrixClient reference used by tokenRefreshFunction.
+  // The refresh function is only ever invoked after startClient, so mxRef is always
+  // assigned before it is called.
+  let mxRef!: MatrixClient;
+
   const mx = createClient({
     baseUrl: session.baseUrl,
     accessToken: session.accessToken,
@@ -288,13 +296,32 @@ const buildClient = async (session: Session): Promise<MatrixClient> => {
     timelineSupport: true,
     cryptoCallbacks: cryptoCallbacks as any,
     verificationMethods: ['m.sas.v1'],
+    ...(session.refreshToken && {
+      refreshToken: session.refreshToken,
+      tokenRefreshFunction: async (oldRefreshToken: string) => {
+        const res = await mxRef.refreshToken(oldRefreshToken);
+        onTokenRefresh?.(res.access_token, res.refresh_token);
+        return {
+          accessToken: res.access_token,
+          refreshToken: res.refresh_token ?? oldRefreshToken,
+          expiry:
+            typeof res.expires_in_ms === 'number'
+              ? new Date(Date.now() + res.expires_in_ms)
+              : undefined,
+        };
+      },
+    }),
   });
 
+  mxRef = mx;
   await indexedDBStore.startup();
   return mx;
 };
 
-export const initClient = async (session: Session): Promise<MatrixClient> => {
+export const initClient = async (
+  session: Session,
+  onTokenRefresh?: (newAccessToken: string, newRefreshToken?: string) => void
+): Promise<MatrixClient> => {
   const storeName = getSessionStoreName(session);
   debugLog.info('sync', 'Initializing Matrix client', {
     userId: session.userId,
@@ -338,7 +365,7 @@ export const initClient = async (session: Session): Promise<MatrixClient> => {
 
   let mx: MatrixClient;
   try {
-    mx = await buildClient(session);
+    mx = await buildClient(session, onTokenRefresh);
   } catch (err) {
     if (!isMismatch(err)) {
       debugLog.error('sync', 'Failed to build client', { error: err });
@@ -347,7 +374,7 @@ export const initClient = async (session: Session): Promise<MatrixClient> => {
     log.warn('initClient: mismatch on buildClient — wiping and retrying:', err);
     debugLog.warn('sync', 'Client build mismatch - wiping stores and retrying', { error: err });
     await wipeAllStores();
-    mx = await buildClient(session);
+    mx = await buildClient(session, onTokenRefresh);
   }
 
   try {
@@ -361,7 +388,7 @@ export const initClient = async (session: Session): Promise<MatrixClient> => {
     debugLog.warn('sync', 'Crypto init mismatch - wiping stores and retrying', { error: err });
     mx.stopClient();
     await wipeAllStores();
-    mx = await buildClient(session);
+    mx = await buildClient(session, onTokenRefresh);
     await mx.initRustCrypto({ cryptoDatabasePrefix: storeName.rustCryptoPrefix });
   }
 

src/sw.ts

@@ -152,7 +152,7 @@
     sessions.set(clientId, info);
     // A real session has arrived — discard the preloaded fallback.
     preloadedSession = undefined;
     console.debug('[SW] setSession: stored', clientId, baseUrl);
     // Persist so push-event fetches work after iOS restarts the SW.
     persistSession(info).catch(() => undefined);
   } else {
@@ -579,6 +579,53 @@
   };
 }
 
+/**
+ * Fetch a media URL, retrying once with the most-current in-memory session on 401.
+ *
+ * There is a timing window between when the SDK refreshes its access token
+ * (tokenRefreshFunction resolves) and when the resulting pushSessionToSW()
+ * postMessage is processed by the SW. Media requests that land in this window
+ * are sent with the stale token and receive 401. By the time the retry runs,
+ * the setSession message will normally have been processed and sessions will
+ * hold the new token.
+ *
+ * A second timing window exists at startup: preloadedSession may hold a stale
+ * token but the live setSession from the page hasn't arrived yet. In that case
+ * the in-memory check yields no fresher token, so we ask the live client tab
+ * directly (requestSessionWithTimeout) before giving up.
+ */
+async function fetchMediaWithRetry(
+  url: string,
+  token: string,
+  redirect: RequestRedirect,
+  clientId: string
+): Promise<Response> {
+  const response = await fetch(url, { ...fetchConfig(token), redirect });
+  if (response.status !== 401) return response;
+
+  const updated =
+    (clientId ? sessions.get(clientId) : undefined) ??
+    [...sessions.values()].find((s) => validMediaRequest(url, s.baseUrl)) ??
+    preloadedSession;
+
+  if (updated && updated.accessToken !== token && validMediaRequest(url, updated.baseUrl)) {
+    return fetch(url, { ...fetchConfig(updated.accessToken), redirect });
+  }
+
+  // No fresher cached token available — ask the live client tab for its current
+  // session. This covers the startup window where preloadedSession is stale but
+  // the page hasn't yet sent setSession. Use a short timeout so we don't block
+  // image rendering for long on unresponsive clients.
+  if (clientId) {
+    const live = await requestSessionWithTimeout(clientId, 1500);
+    if (live && live.accessToken !== token && validMediaRequest(url, live.baseUrl)) {
+      return fetch(url, { ...fetchConfig(live.accessToken), redirect });
+    }
+  }
+
+  return response;
+}
+
 self.addEventListener('message', (event: ExtendableMessageEvent) => {
   if (event.data.type === 'togglePush') {
     const token = event.data?.token;
@@ -609,7 +656,7 @@
 
   const session = clientId ? sessions.get(clientId) : undefined;
   if (session && validMediaRequest(url, session.baseUrl)) {
-    event.respondWith(fetch(url, { ...fetchConfig(session.accessToken), redirect }));
+    event.respondWith(fetchMediaWithRetry(url, session.accessToken, redirect, clientId));
     return;
   }
 
@@ -629,7 +676,7 @@
       ? preloadedSession
       : undefined);
   if (byBaseUrl) {
-    event.respondWith(fetch(url, { ...fetchConfig(byBaseUrl.accessToken), redirect }));
+    event.respondWith(fetchMediaWithRetry(url, byBaseUrl.accessToken, redirect, clientId));
     return;
   }
 
@@ -639,7 +686,7 @@
     event.respondWith(
       loadPersistedSession().then((persisted) => {
         if (persisted && validMediaRequest(url, persisted.baseUrl)) {
-          return fetch(url, { ...fetchConfig(persisted.accessToken), redirect });
+          return fetchMediaWithRetry(url, persisted.accessToken, redirect, '');
         }
         return fetch(event.request);
       })
@@ -651,13 +698,13 @@
     requestSessionWithTimeout(clientId).then(async (s) => {
       // Primary: session received from the live client window.
       if (s && validMediaRequest(url, s.baseUrl)) {
-        return fetch(url, { ...fetchConfig(s.accessToken), redirect });
+        return fetchMediaWithRetry(url, s.accessToken, redirect, clientId);
       }
       // Fallback: try the persisted session (helps when SW restarts on iOS and
       // the client window hasn't responded to requestSession yet).
       const persisted = await loadPersistedSession();
       if (persisted && validMediaRequest(url, persisted.baseUrl)) {
-        return fetch(url, { ...fetchConfig(persisted.accessToken), redirect });
+        return fetchMediaWithRetry(url, persisted.accessToken, redirect, clientId);
       }
       console.warn(
         '[SW fetch] No valid session for media request',