improvement(org): resolve an account's token once per Home widget read

Author: waleedlatif1

apps/sim/lib/organizations/application/home-widgets.ts

@@ -1,4 +1,3 @@
-import { AuditAction, AuditResourceType, recordAudit } from '@sim/audit'
 import { createLogger } from '@sim/logger'
 import { compareStrings } from '@sim/utils/string'
 import type {
@@ -20,14 +19,14 @@ import {
   findCredentialGroupProviderFromProviderId,
 } from '@/lib/credential-groups/providers'
 import { listViewerOrganizationAccounts } from '@/lib/credential-groups/viewer-accounts'
-import {
-  ManagedOAuthCredentialError,
-  rejectManagedOAuthToken,
-  resolveManagedOAuthToken,
-} from '@/lib/credentials/managed-oauth'
+import { ManagedOAuthCredentialError } from '@/lib/credentials/managed-oauth'
 import { getOwnOrganizationManagedOAuthCredentials } from '@/lib/credentials/organization-managed'
 import { isKnowledgeMemberAccessAvailable } from '@/lib/knowledge/access/availability'
-import { recordCredentialAccess } from '@/lib/oauth/token-resolution'
+import {
+  type AccountTokenRead,
+  createAccountTokenRead,
+  type ProviderTokenInput,
+} from '@/lib/organizations/home/account-token'
 import {
   HOME_CALENDAR_EVENT_LIMIT,
   HOME_EMAIL_THREAD_LIMIT,
@@ -71,77 +70,6 @@ export const organizationHomeOperations = {
   }),
 } as const
 
-interface ProviderTokenInput {
-  context: OrganizationMembershipContext
-  credentialId: string
-  providerId: string
-}
-
-function resolveToken({ context, credentialId, providerId }: ProviderTokenInput) {
-  return resolveManagedOAuthToken({
-    credentialId,
-    organizationId: context.organizationId,
-    expectedProviderId: providerId,
-    requiredScopes: [],
-  })
-}
-
-function isTokenRejected(error: unknown) {
-  return error instanceof HomeProviderReadError && error.failure === 'token_rejected'
-}
-
-/**
- * One read with the member's token, and one more with a fresh token if the first is refused. A
- * concurrent read of the same account may already have rotated it, so the fresh token is resolved
- * whether or not this call was the one that rejected the old token.
- */
-async function readWithToken<T>(
-  input: ProviderTokenInput,
-  read: (accessToken: string) => Promise<T[]>
-): Promise<T[]> {
-  const { context, credentialId, providerId } = input
-  const { accessToken } = await resolveToken(input)
-  recordCredentialAccess({
-    actorId: context.userId,
-    workspaceId: null,
-    resourceId: credentialId,
-    providerId,
-    credentialType: 'oauth',
-  })
-  try {
-    return await read(accessToken)
-  } catch (error) {
-    if (!isTokenRejected(error)) throw error
-  }
-  const rejected = await rejectManagedOAuthToken({
-    credentialId,
-    organizationId: context.organizationId,
-    expectedProviderId: providerId,
-    requiredScopes: [],
-    rejectedAccessToken: accessToken,
-  })
-  if (rejected)
-    recordAudit({
-      actorId: context.userId,
-      action: AuditAction.CREDENTIAL_UPDATED,
-      resourceType: AuditResourceType.CREDENTIAL,
-      resourceId: credentialId,
-      description: 'Marked provider-rejected managed OAuth credential as needing reauthorization',
-      metadata: {
-        organizationId: context.organizationId,
-        provider: providerId,
-        credentialType: 'managed_oauth',
-      },
-    })
-  const retry = await resolveToken(input)
-  try {
-    return await read(retry.accessToken)
-  } catch (error) {
-    if (isTokenRejected(error)) throw new HomeProviderReadError('needs_reauth', 401)
-    throw error
-  }
-}
-
 const REAUTH_ERROR_CODES: ReadonlySet<ManagedOAuthCredentialError['code']> = new Set([
   'MANAGED_CREDENTIAL_NEEDS_REAUTH',
   'MANAGED_CREDENTIAL_REVOKED',
@@ -181,6 +109,7 @@ interface WidgetReadResult<T> {
 interface ViewerAccountRead extends ProviderTokenInput {
   accountId: string | null
   readableAccounts: number
+  withToken: AccountTokenRead
 }
 
 type ViewerAccountReader<T> = (accessToken: string, account: ViewerAccountRead) => Promise<T[]>
@@ -254,16 +183,16 @@ async function readViewerAccounts<T>(
       if (!ownedIds.has(credentialId)) return settled('not_offered')
       if (!toRead.has(credentialId)) return settled('not_read')
       try {
+        const tokenInput = { context, credentialId, providerId }
         const account = {
-          context,
-          credentialId,
+          ...tokenInput,
           accountId,
-          providerId,
           readableAccounts: toRead.size,
+          withToken: createAccountTokenRead(tokenInput),
         }
         return settled(
           'ready',
-          await readWithToken(account, (accessToken) => read(accessToken, account))
+          await account.withToken((accessToken) => read(accessToken, account))
         )
       } catch (error) {
         return settled(failedState(error, { organizationId, credentialId, providerId }))
@@ -277,7 +206,7 @@ async function readViewerAccounts<T>(
 }
 
 interface GmailAccountCursor extends GmailThreadCursor {
-  account: ProviderTokenInput
+  account: ViewerAccountRead
 }
 
 /**
@@ -298,7 +227,7 @@ async function readRemainingGmailThreads(
         const cursor = cursors.get(credentialId)
         if (!cursor) return
         try {
-          const threads = await readWithToken(cursor.account, (accessToken) =>
+          const threads = await cursor.account.withToken((accessToken) =>
             readMoreGmailThreads({ accessToken, signal, cursor, count })
           )
           rows.push(...threads.map((thread) => ({ ...thread, credentialId })))

apps/sim/lib/organizations/home/account-token.test.ts

@@ -0,0 +1,88 @@
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+const mocks = vi.hoisted(() => ({
+  resolve: vi.fn(),
+  reject: vi.fn(),
+  recordAccess: vi.fn(),
+}))
+
+vi.mock('@/lib/credentials/managed-oauth', () => ({
+  resolveManagedOAuthToken: mocks.resolve,
+  rejectManagedOAuthToken: mocks.reject,
+}))
+vi.mock('@/lib/oauth/token-resolution', () => ({ recordCredentialAccess: mocks.recordAccess }))
+
+import type { OrganizationMembershipContext } from '@/lib/core/application/organization-authorization'
+import { createAccountTokenRead } from '@/lib/organizations/home/account-token'
+import { HomeProviderReadError } from '@/lib/organizations/home/provider-reads'
+
+const input = {
+  context: { organizationId: 'org', userId: 'viewer' } as OrganizationMembershipContext,
+  credentialId: 'credential',
+  providerId: 'google-email',
+}
+const refused = () => new HomeProviderReadError('token_rejected', 401)
+
+describe('one request reading an account several times', () => {
+  beforeEach(() => {
+    let issued = 0
+    mocks.resolve.mockReset().mockImplementation(async () => {
+      issued += 1
+      return { accessToken: `token-${issued}` }
+    })
+    mocks.reject.mockReset().mockResolvedValue(false)
+    mocks.recordAccess.mockReset()
+  })
+
+  it('resolves and records the token once across reads', async () => {
+    const withToken = createAccountTokenRead(input)
+    const used = [await withToken(async (token) => token), await withToken(async (token) => token)]
+    expect(used).toEqual(['token-1', 'token-1'])
+    expect(mocks.resolve).toHaveBeenCalledOnce()
+    expect(mocks.recordAccess).toHaveBeenCalledOnce()
+  })
+
+  it('swaps a refused token mid-request and keeps the fresh one for later reads', async () => {
+    const withToken = createAccountTokenRead(input)
+    await withToken(async (token) => token)
+    const seen: string[] = []
+    const retried = await withToken(async (token) => {
+      seen.push(token)
+      if (token === 'token-1') throw refused()
+      return token
+    })
+    expect(seen).toEqual(['token-1', 'token-2'])
+    expect(retried).toBe('token-2')
+    expect(mocks.reject).toHaveBeenCalledWith(
+      expect.objectContaining({ rejectedAccessToken: 'token-1' })
+    )
+    await expect(withToken(async (token) => token)).resolves.toBe('token-2')
+    expect(mocks.resolve).toHaveBeenCalledTimes(2)
+  })
+
+  it('asks to reconnect when the fresh token is refused too', async () => {
+    const withToken = createAccountTokenRead(input)
+    await expect(
+      withToken(async () => {
+        throw refused()
+      })
+    ).rejects.toMatchObject({ failure: 'needs_reauth' })
+  })
+
+  it('does not swap the token for any other failure', async () => {
+    const withToken = createAccountTokenRead(input)
+    const failure = new HomeProviderReadError('unavailable', 500)
+    await expect(
+      withToken(async () => {
+        throw failure
+      })
+    ).rejects.toBe(failure)
+    expect(mocks.reject).not.toHaveBeenCalled()
+    await expect(withToken(async (token) => token)).resolves.toBe('token-1')
+  })
+
+  it('resolves a new token for a new request', async () => {
+    await createAccountTokenRead(input)(async (token) => token)
+    await expect(createAccountTokenRead(input)(async (token) => token)).resolves.toBe('token-2')
+  })
+})

apps/sim/lib/organizations/home/account-token.ts

@@ -0,0 +1,85 @@
+import { AuditAction, AuditResourceType, recordAudit } from '@sim/audit'
+import type { OrganizationMembershipContext } from '@/lib/core/application/organization-authorization'
+import { rejectManagedOAuthToken, resolveManagedOAuthToken } from '@/lib/credentials/managed-oauth'
+import { recordCredentialAccess } from '@/lib/oauth/token-resolution'
+import { HomeProviderReadError } from '@/lib/organizations/home/provider-reads'
+
+export interface ProviderTokenInput {
+  context: OrganizationMembershipContext
+  credentialId: string
+  providerId: string
+}
+
+/** Runs a provider read with an account's token, however many reads one widget request needs. */
+export type AccountTokenRead = <T>(read: (accessToken: string) => Promise<T>) => Promise<T>
+
+function resolveToken({ context, credentialId, providerId }: ProviderTokenInput) {
+  return resolveManagedOAuthToken({
+    credentialId,
+    organizationId: context.organizationId,
+    expectedProviderId: providerId,
+    requiredScopes: [],
+  })
+}
+
+function isTokenRejected(error: unknown) {
+  return error instanceof HomeProviderReadError && error.failure === 'token_rejected'
+}
+
+/**
+ * One widget request's use of an account's token: resolved and recorded once, and kept until the
+ * provider refuses it. A refused read runs once more with a fresh token, resolved whether or not
+ * this call was the one that rejected the old token, since a concurrent read of the same account may
+ * already have rotated it. Create one per request, so a revoked grant applies on the next request.
+ */
+export function createAccountTokenRead(input: ProviderTokenInput): AccountTokenRead {
+  const { context, credentialId, providerId } = input
+  let currentToken: string | null = null
+
+  return async (read) => {
+    if (currentToken === null) {
+      currentToken = (await resolveToken(input)).accessToken
+      recordCredentialAccess({
+        actorId: context.userId,
+        workspaceId: null,
+        resourceId: credentialId,
+        providerId,
+        credentialType: 'oauth',
+      })
+    }
+    try {
+      return await read(currentToken)
+    } catch (error) {
+      if (!isTokenRejected(error)) throw error
+    }
+    const rejected = await rejectManagedOAuthToken({
+      credentialId,
+      organizationId: context.organizationId,
+      expectedProviderId: providerId,
+      requiredScopes: [],
+      rejectedAccessToken: currentToken,
+    })
+    if (rejected)
+      recordAudit({
+        actorId: context.userId,
+        action: AuditAction.CREDENTIAL_UPDATED,
+        resourceType: AuditResourceType.CREDENTIAL,
+        resourceId: credentialId,
+        description: 'Marked provider-rejected managed OAuth credential as needing reauthorization',
+        metadata: {
+          organizationId: context.organizationId,
+          provider: providerId,
+          credentialType: 'managed_oauth',
+        },
+      })
+    currentToken = null
+    const fresh = await resolveToken(input)
+    currentToken = fresh.accessToken
+    try {
+      return await read(fresh.accessToken)
+    } catch (error) {
+      if (isTokenRejected(error)) throw new HomeProviderReadError('needs_reauth', 401)
+      throw error
+    }
+  }
+}