📉 perf: cache OpenID JWT user documents (#14187)

* feat(auth): cache OpenID JWT user documents

* fix(auth): invalidate cached auth users on role changes

api/cache/getLogStores.js

@@ -63,6 +63,7 @@ const namespaces = {
     CacheKeys.OPENID_EXCHANGED_TOKENS,
     Time.TEN_MINUTES,
   ),
+  [CacheKeys.AUTH_USER_DOC]: standardCache(CacheKeys.AUTH_USER_DOC),
   [CacheKeys.ADMIN_OAUTH_EXCHANGE]: standardCache(
     CacheKeys.ADMIN_OAUTH_EXCHANGE,
     Time.THIRTY_SECONDS,

api/strategies/openIdJwtStrategy.js

@@ -1,7 +1,7 @@
 const cookies = require('cookie');
 const jwksRsa = require('jwks-rsa');
 const { logger } = require('@librechat/data-schemas');
-const { SystemRoles } = require('librechat-data-provider');
+const { CacheKeys, SystemRoles } = require('librechat-data-provider');
 const { Strategy: JwtStrategy, ExtractJwt } = require('passport-jwt');
 const {
   isEnabled,
@@ -9,10 +9,16 @@ const {
   getOpenIdEmail,
   getOpenIdIssuer,
   normalizeOpenIdIssuer,
+  buildAuthUserDocCacheKey,
+  getAuthUserDocCacheMode,
+  getCachedAuthUserDoc,
+  invalidateCachedAuthUserDoc,
+  setCachedAuthUserDoc,
   getHttpsProxyAgent,
   math,
 } = require('@librechat/api');
 const { updateUser, findUser } = require('~/models');
+const getLogStores = require('~/cache/getLogStores');
 
 const getOpenIdJwtAudience = () => {
   const parsedAudience = (process.env.OPENID_AUDIENCE ?? '')
@@ -47,6 +53,8 @@ const isOpenIdIssuerAllowed = (payload, openIdConfig) => {
   return actualIssuer === expectedIssuer || issuerMatchesTemplate(expectedIssuer, actualIssuer);
 };
 
+const getAuthUserDocCacheStore = () => getLogStores(CacheKeys.AUTH_USER_DOC);
+
 /**
  * @function openIdJwtLogin
  * @param {import('openid-client').Configuration} openIdConfig - Configuration object for the JWT strategy.
@@ -100,15 +108,31 @@ const openIdJwtLogin = (openIdConfig) => {
         const authHeader = req.headers.authorization;
         const rawToken = authHeader?.replace('Bearer ', '');
         const openidIssuer = getOpenIdIssuer(payload, openIdConfig);
-
-        const { user, error, migration } = await findOpenIDUser({
-          findUser,
-          email: payload ? getOpenIdEmail(payload) : undefined,
-          openidId: payload?.sub,
-          openidIssuer,
-          idOnTheSource: payload?.oid,
-          strategyName: 'openIdJwtLogin',
+        const authUserCacheKey = buildAuthUserDocCacheKey({
+          strategy: 'openid-jwt',
+          subject: payload?.sub,
+          issuer: openidIssuer,
         });
+        const authUserCacheMode = getAuthUserDocCacheMode();
+        const authUserCacheStore =
+          authUserCacheMode !== 'off' && authUserCacheKey ? getAuthUserDocCacheStore() : undefined;
+        const cachedUser =
+          authUserCacheMode !== 'off' && authUserCacheStore && authUserCacheKey
+            ? await getCachedAuthUserDoc(authUserCacheStore, authUserCacheKey)
+            : undefined;
+
+        const servedCachedUser = authUserCacheMode === 'on' && cachedUser;
+        const lookupResult = servedCachedUser
+          ? { user: cachedUser, error: null, migration: false }
+          : await findOpenIDUser({
+              findUser,
+              email: payload ? getOpenIdEmail(payload) : undefined,
+              openidId: payload?.sub,
+              openidIssuer,
+              idOnTheSource: payload?.oid,
+              strategyName: 'openIdJwtLogin',
+            });
+        const { user, error, migration } = lookupResult;
 
         if (error) {
           done(null, false, { message: error });
@@ -137,6 +161,17 @@ const openIdJwtLogin = (openIdConfig) => {
             await updateUser(user.id, updateData);
           }
 
+          if (authUserCacheStore && authUserCacheKey) {
+            if (Object.keys(updateData).length > 0) {
+              await invalidateCachedAuthUserDoc(authUserCacheStore, {
+                userId: user.id,
+                cacheKey: authUserCacheKey,
+              });
+            } else if (!servedCachedUser) {
+              await setCachedAuthUserDoc(authUserCacheStore, authUserCacheKey, user);
+            }
+          }
+
           /** Read tokens from session (server-side) to avoid large cookie issues */
           const sessionTokens = req.session?.openidTokens;
           let accessToken = sessionTokens?.accessToken;

api/strategies/openIdJwtStrategy.spec.js

@@ -3,6 +3,12 @@ const { SystemRoles } = require('librechat-data-provider');
 // --- Capture JwtStrategy inputs ---
 let capturedStrategyOptions;
 let capturedVerifyCallback;
+const mockAuthUserDocCacheStore = {
+  get: jest.fn(),
+  set: jest.fn(),
+  delete: jest.fn(),
+};
+const mockGetLogStores = jest.fn(() => mockAuthUserDocCacheStore);
 jest.mock('passport-jwt', () => ({
   Strategy: jest.fn((opts, verifyCallback) => {
     capturedStrategyOptions = opts;
@@ -28,6 +34,11 @@ jest.mock('@librechat/api', () => ({
   getOpenIdEmail: jest.requireActual('@librechat/api').getOpenIdEmail,
   getOpenIdIssuer: jest.fn(() => 'https://issuer.example.com'),
   normalizeOpenIdIssuer: jest.requireActual('@librechat/api').normalizeOpenIdIssuer,
+  buildAuthUserDocCacheKey: jest.fn(() => 'auth-user-doc-key'),
+  getAuthUserDocCacheMode: jest.fn(() => 'off'),
+  getCachedAuthUserDoc: jest.fn(),
+  invalidateCachedAuthUserDoc: jest.fn(),
+  setCachedAuthUserDoc: jest.fn(),
   getHttpsProxyAgent: jest.fn(() => undefined),
   math: jest.fn((val, fallback) => fallback),
 }));
@@ -43,14 +54,35 @@ jest.mock('~/server/services/Files/strategies', () => ({
 jest.mock('~/server/services/Config', () => ({
   getAppConfig: jest.fn().mockResolvedValue({}),
 }));
-jest.mock('~/cache/getLogStores', () =>
-  jest.fn().mockReturnValue({ get: jest.fn(), set: jest.fn() }),
-);
+jest.mock('~/cache/getLogStores', () => mockGetLogStores);
 
-const { findOpenIDUser } = require('@librechat/api');
+const {
+  buildAuthUserDocCacheKey,
+  findOpenIDUser,
+  getAuthUserDocCacheMode,
+  getCachedAuthUserDoc,
+  invalidateCachedAuthUserDoc,
+  setCachedAuthUserDoc,
+} = require('@librechat/api');
 const openIdJwtLogin = require('./openIdJwtStrategy');
 const { findUser, updateUser } = require('~/models');
 
+function resetAuthUserDocCacheMocks() {
+  mockAuthUserDocCacheStore.get.mockResolvedValue(undefined);
+  mockAuthUserDocCacheStore.set.mockResolvedValue(undefined);
+  mockAuthUserDocCacheStore.delete.mockResolvedValue(undefined);
+  mockGetLogStores.mockReturnValue(mockAuthUserDocCacheStore);
+  buildAuthUserDocCacheKey.mockReturnValue('auth-user-doc-key');
+  getAuthUserDocCacheMode.mockReturnValue('off');
+  getCachedAuthUserDoc.mockResolvedValue(undefined);
+  invalidateCachedAuthUserDoc.mockResolvedValue(undefined);
+  setCachedAuthUserDoc.mockResolvedValue(undefined);
+}
+
+beforeEach(() => {
+  resetAuthUserDocCacheMocks();
+});
+
 function withEnv(env, callback) {
   const previous = Object.fromEntries(Object.keys(env).map((key) => [key, process.env[key]]));
   Object.entries(env).forEach(([key, value]) => {
@@ -348,6 +380,103 @@ describe('openIdJwtStrategy – token source handling', () => {
   });
 });
 
+describe('openIdJwtStrategy – auth user document cache', () => {
+  const payload = {
+    sub: 'oidc-123',
+    email: 'test@example.com',
+    iss: 'https://issuer.example.com',
+    exp: 9999999999,
+  };
+
+  const req = { headers: { authorization: 'Bearer tok' }, session: {} };
+
+  const baseUser = {
+    _id: { toString: () => 'user-abc' },
+    role: SystemRoles.USER,
+    provider: 'openid',
+    email: 'test@example.com',
+  };
+
+  beforeEach(() => {
+    jest.clearAllMocks();
+    resetAuthUserDocCacheMocks();
+    updateUser.mockResolvedValue({});
+    openIdJwtLogin(mockOpenIdConfig);
+  });
+
+  it('does not initialize the cache store while cache mode is off', async () => {
+    findOpenIDUser.mockResolvedValue({ user: { ...baseUser }, error: null, migration: false });
+
+    await invokeVerify(req, payload);
+
+    expect(findOpenIDUser).toHaveBeenCalled();
+    expect(mockGetLogStores).not.toHaveBeenCalled();
+    expect(getCachedAuthUserDoc).not.toHaveBeenCalled();
+    expect(setCachedAuthUserDoc).not.toHaveBeenCalled();
+  });
+
+  it('uses the cached user document in on mode without a database lookup', async () => {
+    const cachedUser = {
+      _id: 'cached-user',
+      role: SystemRoles.USER,
+      provider: 'openid',
+      email: 'cached@example.com',
+    };
+    getAuthUserDocCacheMode.mockReturnValue('on');
+    getCachedAuthUserDoc.mockResolvedValue(cachedUser);
+
+    const { user } = await invokeVerify(req, payload);
+
+    expect(buildAuthUserDocCacheKey).toHaveBeenCalledWith({
+      strategy: 'openid-jwt',
+      subject: payload.sub,
+      issuer: 'https://issuer.example.com',
+    });
+    expect(findOpenIDUser).not.toHaveBeenCalled();
+    expect(user).toMatchObject({
+      id: 'cached-user',
+      email: 'cached@example.com',
+      idOnTheSource: null,
+    });
+    expect(setCachedAuthUserDoc).not.toHaveBeenCalled();
+    expect(invalidateCachedAuthUserDoc).not.toHaveBeenCalled();
+  });
+
+  it('populates the cache after a miss with the fresh user document', async () => {
+    getAuthUserDocCacheMode.mockReturnValue('on');
+    getCachedAuthUserDoc.mockResolvedValue(undefined);
+    findOpenIDUser.mockResolvedValue({ user: { ...baseUser }, error: null, migration: false });
+
+    await invokeVerify(req, payload);
+
+    expect(findOpenIDUser).toHaveBeenCalled();
+    expect(setCachedAuthUserDoc).toHaveBeenCalledWith(
+      mockAuthUserDocCacheStore,
+      'auth-user-doc-key',
+      expect.objectContaining({ id: 'user-abc' }),
+    );
+    expect(invalidateCachedAuthUserDoc).not.toHaveBeenCalled();
+  });
+
+  it('invalidates instead of populating when login mutates the user', async () => {
+    getAuthUserDocCacheMode.mockReturnValue('on');
+    findOpenIDUser.mockResolvedValue({
+      user: { ...baseUser, role: undefined },
+      error: null,
+      migration: false,
+    });
+
+    await invokeVerify(req, payload);
+
+    expect(updateUser).toHaveBeenCalledWith('user-abc', { role: SystemRoles.USER });
+    expect(setCachedAuthUserDoc).not.toHaveBeenCalled();
+    expect(invalidateCachedAuthUserDoc).toHaveBeenCalledWith(mockAuthUserDocCacheStore, {
+      userId: 'user-abc',
+      cacheKey: 'auth-user-doc-key',
+    });
+  });
+});
+
 describe('openIdJwtStrategy – idOnTheSource boundary coercion', () => {
   const payload = {
     sub: 'oidc-123',