fix(ui): stop rebuilding the socket on every token refresh

invokeai/frontend/web/src/features/auth/store/authSlice.ts

@@ -22,7 +22,8 @@ const zAuthState = z.object({
 type User = z.infer<typeof zUser>;
 type AuthState = z.infer<typeof zAuthState>;
 
-const getTokenUserId = (token: string | null): string | null => {
+/** The token's claims, or null when it is absent or not a readable JWT. */
+const decodeTokenPayload = (token: string | null): Record<string, unknown> | null => {
   if (!token) {
     return null;
   }
@@ -32,13 +33,51 @@ const getTokenUserId = (token: string | null): string | null => {
       return null;
     }
     const normalizedPayload = encodedPayload.replace(/-/g, '+').replace(/_/g, '/');
-    const payload = JSON.parse(atob(normalizedPayload.padEnd(Math.ceil(normalizedPayload.length / 4) * 4, '=')));
-    return typeof payload.user_id === 'string' ? payload.user_id : null;
+    return JSON.parse(atob(normalizedPayload.padEnd(Math.ceil(normalizedPayload.length / 4) * 4, '=')));
   } catch {
     return null;
   }
 };
 
+const getTokenUserId = (token: string | null): string | null => {
+  const payload = decodeTokenPayload(token);
+  return typeof payload?.user_id === 'string' ? payload.user_id : null;
+};
+
+/**
+ * The session a token belongs to, as opposed to the bytes it happens to be made of.
+ *
+ * The sliding-window middleware mints a replacement token on every mutating request, so a live
+ * session's token changes bytes on a fixed cadence (throttled to once a minute by
+ * `acceptRefreshedToken`) while the login behind it never changes. Anything that must be rebuilt
+ * when the *session* changes — the socket, most notably — keys on this rather than on the token,
+ * so a routine refresh does not tear down work that belongs to the same user.
+ *
+ * The revocation epoch is part of the identity, not incidental to it. A password change bumps
+ * `token_epoch` on the user record, and the server force-disconnects every socket that
+ * authenticated under the superseded epoch (`sockets.py`, `_handle_user_access_changed`). A
+ * server-initiated disconnect is terminal for socket.io — the client sets `skipReconnect` and
+ * never retries (`socket.io-client`, `Socket.ondisconnect` -> `Manager._close`) — so the only way
+ * back is to build a new socket, and the replacement token the server hands out carries the new
+ * epoch. Keying on `user_id` alone would leave that socket dead until a full page reload, with
+ * `$isConnected` stuck false and Invoke disabled with it.
+ *
+ * A token carrying no user id falls back to its own bytes: with no identity to compare, byte
+ * equality is the only safe answer, and consumers keep their pre-existing behaviour.
+ */
+export const getTokenSessionKey = (token: string | null): string | null => {
+  if (!token) {
+    return null;
+  }
+  const payload = decodeTokenPayload(token);
+  if (typeof payload?.user_id !== 'string') {
+    return token;
+  }
+  // Absent on tokens minted before the claim existed; the server reads those as epoch 0 too.
+  const epoch = typeof payload.token_epoch === 'number' ? payload.token_epoch : 0;
+  return `${payload.user_id}:${epoch}`;
+};
+
 export const tokensBelongToSameUser = (first: string | null, second: string | null): boolean => {
   const firstUserId = getTokenUserId(first);
   return firstUserId !== null && firstUserId === getTokenUserId(second);
@@ -164,5 +203,6 @@ export const authSliceConfig: SliceConfig<typeof authSlice> = {
 export const selectIsAuthenticated = (state: { auth: AuthState }) => state.auth.isAuthenticated;
 export const selectCurrentUser = (state: { auth: AuthState }) => state.auth.user;
 export const selectAuthToken = (state: { auth: AuthState }) => state.auth.token;
+export const selectAuthSessionKey = (state: { auth: AuthState }) => getTokenSessionKey(state.auth.token);
 export const selectIsAuthLoading = (state: { auth: AuthState }) => state.auth.isLoading;
 export const selectSessionExpired = (state: { auth: AuthState }) => state.auth.sessionExpired;

invokeai/frontend/web/src/features/auth/store/refreshedTokenConsumers.test.ts

@@ -4,7 +4,14 @@ import { fileURLToPath } from 'node:url';
 
 import { describe, expect, it } from 'vitest';
 
-import { authSliceConfig, currentUserUpdated, externalTokenAdopted, setCredentials, tokenRefreshed } from './authSlice';
+import {
+  authSliceConfig,
+  currentUserUpdated,
+  externalTokenAdopted,
+  getTokenSessionKey,
+  setCredentials,
+  tokenRefreshed,
+} from './authSlice';
 
 const user = {
   user_id: 'user',
@@ -14,17 +21,37 @@ const user = {
   is_active: true,
 };
 
-const tokenFor = (userId: string) =>
-  `header.${Buffer.from(JSON.stringify({ user_id: userId })).toString('base64url')}.signature`;
+const tokenFor = (userId: string, nonce = 0, epoch = 0) =>
+  `header.${Buffer.from(JSON.stringify({ user_id: userId, nonce, token_epoch: epoch })).toString('base64url')}.signature`;
 
 describe('refreshed token consumers', () => {
-  it('updates the Redux token used by socket reconnects', () => {
+  it('updates the Redux token the next request will carry', () => {
     let state = authSliceConfig.slice.reducer(undefined, setCredentials({ token: 'old', user }));
     state = authSliceConfig.slice.reducer(state, tokenRefreshed('new'));
 
     expect(state.token).toBe('new');
   });
 
+  it('keeps one session key across a refresh, so the socket is not rebuilt for it', () => {
+    // The middleware mints new bytes for the same login on every mutating request. Consumers that
+    // key work to a session — `useSocketIO` — must not see that as a new session.
+    expect(getTokenSessionKey(tokenFor('user', 1))).toBe(getTokenSessionKey(tokenFor('user', 2)));
+    expect(getTokenSessionKey(tokenFor('user', 1))).not.toBe(getTokenSessionKey(tokenFor('other', 1)));
+    expect(getTokenSessionKey(null)).toBeNull();
+  });
+
+  it('changes the session key when a revoked epoch supersedes the token', () => {
+    // A password change bumps the epoch and the server drops every socket authenticated under the
+    // old one; socket.io does not retry that, so the replacement token must count as a new session.
+    expect(getTokenSessionKey(tokenFor('user', 1, 0))).not.toBe(getTokenSessionKey(tokenFor('user', 2, 1)));
+  });
+
+  it('falls back to the token itself when it carries no user id', () => {
+    // Nothing to compare identities with, so byte equality is the only safe answer.
+    expect(getTokenSessionKey('opaque-token')).toBe('opaque-token');
+    expect(getTokenSessionKey('')).toBeNull();
+  });
+
   it('clears the previous user when adopting a token from another tab', () => {
     let state = authSliceConfig.slice.reducer(undefined, setCredentials({ token: tokenFor(user.user_id), user }));
     state = authSliceConfig.slice.reducer(state, externalTokenAdopted(tokenFor('other-user')));