fix(auth): amortize the entitlement check and bound cache eviction cost

Self-review follow-ups on the cache-coherence change:

- Splitting the entitlement check out of the combined policy cache meant it ran
  on every session refresh. A policy save or org-wide revoke makes every member
  session refresh at once, so a large org would have stampeded the billing
  tables with one plan check per session. Entitlement is now cached on its own
  TTL, independent of the org security record (a policy save invalidates the
  policy but says nothing about the plan). Staleness is harmless in both
  directions, and a failed check falls back to the last known decision.
- Pruning trimmed exactly to the cap, so the next insert was over again and
  every subsequent write rescanned the whole map. Prune to a low water mark
  instead, which keeps eviction amortized O(1).
- The session create hook no longer writes an `expiresAt: undefined` key when
  the clamp has nothing to narrow; it only ever overrides a real value.
- Drop the redundant `isBillingEnabled` short-circuit — the same condition is
  already the first branch of `resolveOrganizationEnterprisePlan`.

Author: waleedlatif1

apps/sim/lib/auth/auth.ts

@@ -695,7 +695,10 @@ export const auth = betterAuth({
             return {
               data: {
                 ...session,
-                expiresAt,
+                // Only ever narrows an existing expiry — never introduces an
+                // `expiresAt: undefined` key that would blank out the value
+                // Better Auth already put on the row.
+                ...(expiresAt ? { expiresAt } : {}),
                 ...(membershipOrgId ? { activeOrganizationId: membershipOrgId } : {}),
               },
             }

apps/sim/lib/auth/security-policy.ts

@@ -60,18 +60,28 @@ interface OrgSecurityCacheEntry extends OrgSecurityRecord {
 const orgSecurityCache = new Map<string, OrgSecurityCacheEntry>()
 
 /**
- * Evicts down to `maxEntries`, expired entries first. Map iteration follows
+ * Fraction of the cap an over-capacity prune evicts down to. Trimming to a low
+ * water mark rather than exactly to the cap is what keeps eviction amortized
+ * O(1): pruning back to the cap would leave the very next insert over again,
+ * so every subsequent write would rescan the whole map.
+ */
+const PRUNE_TARGET_RATIO = 0.9
+
+/**
+ * Evicts down to a low water mark, expired entries first. Map iteration follows
  * insertion order and {@link touch} re-inserts on every refresh, so whatever
  * remains at the head is the least recently refreshed.
  */
 function prune(cache: Map<string, { fetchedAt: number }>, maxEntries: number): void {
   if (cache.size <= maxEntries) return
+  const target = Math.floor(maxEntries * PRUNE_TARGET_RATIO)
   const cutoff = Date.now() - SECURITY_POLICY_CACHE_TTL_MS
   for (const [key, entry] of cache) {
+    if (cache.size <= target) break
     if (entry.fetchedAt < cutoff) cache.delete(key)
   }
   for (const key of cache.keys()) {
-    if (cache.size <= maxEntries) break
+    if (cache.size <= target) break
     cache.delete(key)
   }
 }

apps/sim/lib/auth/session-policy.test.ts

@@ -147,6 +147,23 @@ describe('getSessionPolicy', () => {
     })
   })
 
+  it('amortizes the entitlement check across repeated resolutions', async () => {
+    const orgId = nextOrgId()
+    const bounded = {
+      version: 1,
+      sessionPolicySettings: { maxSessionHours: 8, idleTimeoutHours: null },
+    }
+    queueTableRows(organization, [bounded])
+    queueTableRows(organization, [bounded])
+
+    // A policy save makes every member session refresh at once; the plan check
+    // must not run once per refreshing session.
+    await getSessionPolicy(orgId)
+    await getSessionPolicy(orgId, { bypassCache: true })
+
+    expect(mockResolveEnterprisePlan).toHaveBeenCalledTimes(1)
+  })
+
   it('keeps enforcing when the entitlement check itself fails', async () => {
     queueTableRows(organization, [
       { version: 1, sessionPolicySettings: { maxSessionHours: 8, idleTimeoutHours: null } },

apps/sim/lib/auth/session-policy.ts

@@ -8,7 +8,6 @@ import {
   invalidateMembershipCache,
 } from '@/lib/auth/security-policy'
 import { resolveOrganizationEnterprisePlan } from '@/lib/billing/core/subscription'
-import { isBillingEnabled } from '@/lib/core/config/env-flags'
 
 const logger = createLogger('SessionPolicy')
 
@@ -24,6 +23,20 @@ const NO_POLICY: ResolvedSessionPolicy = {
   idleTimeoutHours: null,
 }
 
+/**
+ * Entitlement is billing state, so it is cached independently of the org
+ * security record: a policy save must invalidate the policy, but it says
+ * nothing about the plan. The TTL matters because saving a policy or revoking
+ * org-wide makes every member session refresh at once — without it, each of
+ * those refreshes would re-run the plan check and stampede the billing tables.
+ * Staleness is harmless in both directions: a just-downgraded org keeps
+ * enforcing for one TTL, and a just-upgraded one starts one TTL late.
+ */
+const ENTITLEMENT_CACHE_TTL_MS = 60 * 1000
+const MAX_ENTITLEMENT_CACHE_ENTRIES = 5_000
+
+const entitlementCache = new Map<string, { entitled: boolean; fetchedAt: number }>()
+
 /**
  * Whether stored session bounds should be enforced for this org. Mirroring
  * data-retention's plan-gated effective settings, a hosted org that leaves the
@@ -38,15 +51,27 @@ const NO_POLICY: ResolvedSessionPolicy = {
  * `lib/logs/execution/logger.ts`.
  */
 async function isPolicyEnforced(organizationId: string, hasBounds: boolean): Promise<boolean> {
-  if (!hasBounds || !isBillingEnabled) return true
+  // Orgs with nothing stored are the common case and need no plan check at all.
+  if (!hasBounds) return true
+
+  const cached = entitlementCache.get(organizationId)
+  if (cached && Date.now() - cached.fetchedAt < ENTITLEMENT_CACHE_TTL_MS) {
+    return cached.entitled
+  }
+
   try {
-    return await resolveOrganizationEnterprisePlan(organizationId)
+    const entitled = await resolveOrganizationEnterprisePlan(organizationId)
+    // Entitled orgs are few, so a wholesale clear at the cap is enough — no
+    // eviction policy needed for a set this small and this slow-moving.
+    if (entitlementCache.size >= MAX_ENTITLEMENT_CACHE_ENTRIES) entitlementCache.clear()
+    entitlementCache.set(organizationId, { entitled, fetchedAt: Date.now() })
+    return entitled
   } catch (error) {
     logger.error('Enterprise entitlement check failed; enforcing stored session policy', {
       organizationId,
       error,
     })
-    return true
+    return cached?.entitled ?? true
   }
 }