feat: enable refresh token rotation for salesforce

previously, we were not using the new refresh token if one was provided

Author: paustint

apps/api/src/app/routes/route.middleware.ts

@@ -361,9 +361,6 @@ export async function getOrgForRequest(
     // Refresh event will be fired when renewed access token
     // to store it in your storage for next request
     try {
-      if (!refreshToken) {
-        return;
-      }
       await salesforceOrgsDb.updateAccessToken_UNSAFE({ accessToken, refreshToken, org, userId: user.id });
     } catch (ex) {
       logger.error({ requestId, ...getExceptionLog(ex) }, '[ORG][REFRESH] Error saving refresh token');
@@ -378,6 +375,25 @@ export async function getOrgForRequest(
     }
   };
 
+  // Re-reads current tokens from DB so concurrent workers that lose the refresh token rotation race
+  // can retry with the tokens written by the worker that won.
+  const getFreshTokens = async () => {
+    try {
+      const freshOrg = await salesforceOrgsDb.findByUniqueId_UNSAFE(user.id, uniqueId);
+      if (!freshOrg) {
+        return null;
+      }
+      const [freshAccessToken, freshRefreshToken] = await sfdcEncService.decryptAccessToken({
+        encryptedAccessToken: freshOrg.accessToken,
+        userId: user.id,
+      });
+      return { accessToken: freshAccessToken, refreshToken: freshRefreshToken };
+    } catch (ex) {
+      logger.error({ requestId, ...getExceptionLog(ex) }, '[ORG][REFRESH] Error fetching fresh tokens for race condition check');
+      return null;
+    }
+  };
+
   const jetstreamConn = new ApiConnection(
     {
       apiRequestAdapter: getApiRequestFactoryFn(fetch),
@@ -392,6 +408,7 @@ export async function getOrgForRequest(
       logger,
       sfdcClientId: ENV.SFDC_CONSUMER_KEY,
       sfdcClientSecret: ENV.SFDC_CONSUMER_SECRET,
+      getFreshTokens,
     },
     handleRefresh,
     handleConnectionError,
@@ -574,7 +591,6 @@ export function setPermissionPolicy(_req: express.Request, res: express.Response
   next();
 }
 
-
 export function setCacheControlForApiRoutes(_req: express.Request, res: express.Response, next: express.NextFunction) {
   res.setHeader('Cache-Control', 'no-store, no-cache, must-revalidate, proxy-revalidate, max-age=0');
   next();

apps/jetstream-desktop/src/utils/route.utils.ts

@@ -204,9 +204,6 @@ export function initApiConnection(
     // Refresh event will be fired when renewed access token
     // to store it in your storage for next request
     try {
-      if (!refreshToken) {
-        return;
-      }
       await updateAccessTokens(org.uniqueId, { accessToken, refreshToken });
     } catch (ex) {
       logger.error('[ORG][REFRESH] Error saving refresh token', getErrorMessage(ex));
@@ -221,6 +218,26 @@ export function initApiConnection(
     }
   };
 
+  // Re-reads current tokens from the in-memory store so concurrent requests that lose the
+  // refresh token rotation race can retry with the tokens written by the request that won.
+  const getFreshTokens = async () => {
+    try {
+      const freshOrg = getSalesforceOrgById(org.uniqueId);
+      if (!freshOrg) {
+        return null;
+      }
+      const plaintext = decryptTokenPortable(freshOrg.accessToken);
+      const spaceIndex = plaintext.indexOf(' ');
+      if (spaceIndex === -1) {
+        return null;
+      }
+      return { accessToken: plaintext.slice(0, spaceIndex), refreshToken: plaintext.slice(spaceIndex + 1) };
+    } catch (ex) {
+      logger.error('[ORG][REFRESH] Error fetching fresh tokens for race condition check', getErrorMessage(ex));
+      return null;
+    }
+  };
+
   const jetstreamConn = new ApiConnection(
     {
       apiRequestAdapter: getApiRequestFactoryFn(fetch),
@@ -234,6 +251,7 @@ export function initApiConnection(
       logger: logger as any,
       enableLogging: false,
       sfdcClientId: ENV.DESKTOP_SFDC_CLIENT_ID,
+      getFreshTokens,
     },
     handleRefresh,
     handleConnectionError,

libs/salesforce-api/src/lib/callout-adapter.ts

@@ -1,4 +1,5 @@
 import { ERROR_MESSAGES, HTTP } from '@jetstream/shared/constants';
+import { getErrorMessageAndStackObj } from '@jetstream/shared/utils';
 import { parse } from '@jetstreamapp/simple-xml';
 import isObject from 'lodash/isObject';
 import { ApiRequestOptions, ApiRequestOutputType, BulkXmlErrorResponse, FetchFn, FetchResponse, Logger, SoapErrorResponse } from './types';
@@ -46,9 +47,14 @@ function parseXml(value: string) {
 export function getApiRequestFactoryFn(fetch: FetchFn) {
   return (
     logger: Logger,
-    onRefresh?: (accessToken: string) => void,
-    onConnectionError?: (accessToken: string) => void,
+    onRefresh?: (accessToken: string, refreshToken?: string) => void,
+    onConnectionError?: (error: string) => void,
+    /**
+     * Enable logging only applies to request/response data
+     * other logging for refresh flow and logic errors will still be logged
+     */
     enableLogging?: boolean,
+    getFreshTokens?: () => Promise<{ accessToken: string; refreshToken: string } | null>,
   ) => {
     const apiRequest = async <Response = unknown>(options: ApiRequestOptions, attemptRefresh = true): Promise<Response> => {
       // eslint-disable-next-line prefer-const
@@ -125,18 +131,37 @@ export function getApiRequestFactoryFn(fetch: FetchFn) {
             sessionInfo.refreshToken
           ) {
             try {
-              // if 401 and we have a refresh token, then attempt to refresh the token
-              const { access_token: newAccessToken } = await exchangeRefreshToken(fetch, sessionInfo);
-              onRefresh?.(newAccessToken);
+              logger.debug({ url, method, status: response.status }, '[TOKEN REFRESH] Attempting token refresh');
+              const { access_token: newAccessToken, refresh_token: newRefreshToken } = await exchangeRefreshToken(fetch, sessionInfo);
+              logger.debug({ url, method, tokenRotated: !!newRefreshToken }, '[TOKEN REFRESH] Token refresh successful');
+              onRefresh?.(newAccessToken, newRefreshToken);
               // replace token in body
               if (typeof options.body === 'string' && options.body.includes(accessToken)) {
                 // if the response is soap, we need to return the response as is
                 options.body = options.body.replace(accessToken, newAccessToken);
               }
 
               return apiRequest({ ...options, sessionInfo: { ...sessionInfo, accessToken: newAccessToken } }, false);
-            } catch {
-              logger.warn('Unable to refresh accessToken');
+            } catch (ex) {
+              logger.warn({ url, method, ...getErrorMessageAndStackObj(ex) }, '[TOKEN REFRESH] Unable to refresh accessToken');
+
+              // Check if another worker already refreshed (race condition on token rotation).
+              // If the DB has a different access token, a concurrent request won the race — retry with fresh tokens.
+              if (getFreshTokens) {
+                try {
+                  const freshTokens = await getFreshTokens();
+                  if (freshTokens && freshTokens.accessToken !== accessToken) {
+                    logger.info({ url, method }, '[TOKEN REFRESH] Concurrent refresh detected — retrying with tokens from another worker');
+                    return apiRequest({ ...options, sessionInfo: { ...sessionInfo, ...freshTokens } }, false);
+                  }
+                } catch (freshEx) {
+                  logger.warn(
+                    { url, method, ...getErrorMessageAndStackObj(freshEx) },
+                    '[TOKEN REFRESH] Failed to retrieve fresh tokens for race condition check',
+                  );
+                }
+              }
+
               responseText = ERROR_MESSAGES.SFDC_EXPIRED_TOKEN;
               onConnectionError?.(ERROR_MESSAGES.SFDC_EXPIRED_TOKEN);
             }
@@ -192,7 +217,10 @@ function handleSalesforceApiError(outputType: ApiRequestOutputType, responseText
   return output;
 }
 
-function exchangeRefreshToken(fetch: FetchFn, sessionInfo: ApiRequestOptions['sessionInfo']): Promise<{ access_token: string }> {
+function exchangeRefreshToken(
+  fetch: FetchFn,
+  sessionInfo: ApiRequestOptions['sessionInfo'],
+): Promise<{ access_token: string; refresh_token?: string }> {
   const searchParams = new URLSearchParams({
     grant_type: 'refresh_token',
   });
@@ -221,6 +249,6 @@ function exchangeRefreshToken(fetch: FetchFn, sessionInfo: ApiRequestOptions['se
     })
     .then((response) => response.json())
     .then((response) => {
-      return response as { access_token: string };
+      return response as { access_token: string; refresh_token?: string };
     });
 }

libs/salesforce-api/src/lib/connection.ts

@@ -23,6 +23,8 @@ export interface ApiConnectionOptions {
   sfdcClientId?: string;
   sfdcClientSecret?: string;
   logger: Logger;
+  /** Re-reads current tokens from the source of truth (e.g. DB) to handle concurrent refresh token rotation across workers */
+  getFreshTokens?: () => Promise<{ accessToken: string; refreshToken: string } | null>;
 }
 
 export class ApiConnection {
@@ -56,12 +58,19 @@ export class ApiConnection {
       sfdcClientId,
       sfdcClientSecret,
       logger,
+      getFreshTokens,
     }: ApiConnectionOptions,
     refreshCallback?: (accessToken: string, refreshToken: string) => void,
     onConnectionError?: (error: string) => void,
   ) {
     this.logger = logger;
-    this.apiRequest = apiRequestAdapter(logger, this.handleRefresh.bind(this), this.handleConnectionError.bind(this), enableLogging);
+    this.apiRequest = apiRequestAdapter(
+      logger,
+      this.handleRefresh.bind(this),
+      this.handleConnectionError.bind(this),
+      enableLogging,
+      getFreshTokens,
+    );
     this.refreshCallback = refreshCallback;
     this.onConnectionError = onConnectionError;
     this.sessionInfo = {
@@ -122,8 +131,11 @@ export class ApiConnection {
     this.sessionInfo.userId = userId ?? this.sessionInfo.userId;
   }
 
-  public handleRefresh(accessToken: string) {
+  public handleRefresh(accessToken: string, newRefreshToken?: string) {
     this.sessionInfo.accessToken = accessToken;
+    if (newRefreshToken) {
+      this.sessionInfo.refreshToken = newRefreshToken;
+    }
     // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
     this.refreshCallback?.(accessToken, this.sessionInfo.refreshToken!);
   }