diff --git a/apps/sim/app/api/billing/update-cost/route.ts b/apps/sim/app/api/billing/update-cost/route.ts index f7d2ef16fe3..3c3668cb732 100644 --- a/apps/sim/app/api/billing/update-cost/route.ts +++ b/apps/sim/app/api/billing/update-cost/route.ts @@ -82,8 +82,8 @@ function invalidBillingProtocolResponse(requestId: string, span: Span): NextResp * once it crosses the limit instead of at its next admission, with the card the worker writes * to its log. The payer is the attributed run's, or the one a direct-v1 run was admitted with. * A duplicate callback answers too: it is often a retry whose first answer was lost. An - * admission is cached per payer and actor for the gate TTL and a refusal is always re-read, so - * steady-state steps cost no ledger read. The charge is + * admission is cached per payer and actor (briefly once near the limit) and a refusal is always + * re-read, so most steps cost no ledger read. The charge is * already recorded when this runs; a gate that cannot answer reports not-exceeded and leaves the * refusal to the next step or re-check rather than ending a paying run on a database blip, * and so does a verdict read that outlasts {@link USAGE_STANDING_TIMEOUT_MS}. An exceeded diff --git a/apps/sim/app/api/copilot/api-keys/validate/route.test.ts b/apps/sim/app/api/copilot/api-keys/validate/route.test.ts index 2a7ebe9d956..04c6a3b2b4f 100644 --- a/apps/sim/app/api/copilot/api-keys/validate/route.test.ts +++ b/apps/sim/app/api/copilot/api-keys/validate/route.test.ts @@ -6,6 +6,7 @@ import { schemaMock, setEnvFlags, } from '@sim/testing' +import { createRouteContext } from '@sim/testing/helpers/http' import { billingCoreMock, billingCoreMockFns } from '@sim/testing/mocks/billing-core.mock' import { billingPlanMock, billingPlanMockFns } from '@sim/testing/mocks/billing-plan.mock' import { @@ -145,7 +146,10 @@ import { validateCopilotApiKeyContract, } from '@/lib/api/contracts/copilot' import { resetMidRunUsageCaches } from '@/lib/billing/core/mid-run-usage' -import { resetUsageGateCache } from '@/lib/billing/core/usage-gate-cache' +import { + NEAR_LIMIT_USAGE_GATE_TTL_MS, + resetUsageGateCache, +} from '@/lib/billing/core/usage-gate-cache' import { POST } from '@/app/api/copilot/api-keys/validate/route' const { mockGetWorkspaceBillingSettings } = workspacesUtilsMockFns @@ -629,6 +633,25 @@ describe('validation lifecycle purposes', () => { }) }) + it('re-reads a direct-v1 account near its limit, so a continuation past it is refused', async () => { + vi.useFakeTimers({ toFake: ['setTimeout', 'clearTimeout'] }) + try { + mockCheckUsageStatus.mockResolvedValue({ isExceeded: false, currentUsage: 4.2, limit: 5 }) + expect((await POST(request(body, directHeaders), createRouteContext({}))).status).toBe(200) + + mockCheckUsageStatus.mockResolvedValue({ isExceeded: true, currentUsage: 5.4, limit: 5 }) + /** `lru-cache` captures `performance` at load: spy its clock, then fire its 1 ms debounce. */ + const start = performance.now() + vi.spyOn(performance, 'now').mockReturnValue(start + NEAR_LIMIT_USAGE_GATE_TTL_MS + 1) + await vi.advanceTimersByTimeAsync(1) + + expect((await POST(request(body, directHeaders), createRouteContext({}))).status).toBe(402) + } finally { + vi.runOnlyPendingTimers() + vi.useRealTimers() + } + }) + it('admits a direct-v1 continuation whose usage cannot be read', async () => { mockCheckUsageStatus.mockResolvedValue({ isExceeded: true, unavailable: true }) expect((await POST(request(body, directHeaders))).status).toBe(200) diff --git a/apps/sim/lib/billing/core/mid-run-usage.ts b/apps/sim/lib/billing/core/mid-run-usage.ts index 3db0c1d342d..cae9c27e6e0 100644 --- a/apps/sim/lib/billing/core/mid-run-usage.ts +++ b/apps/sim/lib/billing/core/mid-run-usage.ts @@ -17,6 +17,7 @@ import { checkExecutionUsageLimits, USAGE_GATE_SETTLE_TIMEOUT_MS, USAGE_GATE_TTL_MS, + usageAdmissionTtlMs, } from '@/lib/billing/core/usage-gate-cache' import type { UsageUpgradePayer } from '@/lib/billing/usage-upgrade' import { coalesceLocally } from '@/lib/concurrency/singleflight' @@ -141,7 +142,7 @@ export async function readMidRunUsageVerdict( } /** - * Admitted direct-v1 verdicts, served for the execution gate's TTL like + * Admitted direct-v1 verdicts, served for {@link usageAdmissionTtlMs} like * {@link checkExecutionUsageLimits} serves attributed ones: the worker re-validates a run on * every resume leg, and each uncached read sums the payer's ledger for the period. Only a * `within` verdict is stored, so a refusal or an unreadable ledger is always read again. @@ -227,7 +228,7 @@ export async function readMidRunAccountUsageVerdict( } } const within: MidRunUsageVerdict = { status: 'within' } - accountVerdictCache.set(key, within) + accountVerdictCache.set(key, within, { ttl: usageAdmissionTtlMs(usage) }) return within } catch (error) { logger.warn('Mid-run account usage read failed; continuing the run', { diff --git a/apps/sim/lib/billing/core/usage-gate-cache.test.ts b/apps/sim/lib/billing/core/usage-gate-cache.test.ts index 79f35c65b3b..aeb4290683d 100644 --- a/apps/sim/lib/billing/core/usage-gate-cache.test.ts +++ b/apps/sim/lib/billing/core/usage-gate-cache.test.ts @@ -5,7 +5,7 @@ import { billingAttributionMockFns, } from '@sim/testing/mocks/billing-attribution.mock' import { sleep } from '@sim/utils/helpers' -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' vi.mock('@/lib/billing/core/attribution', () => billingAttributionMock) @@ -14,6 +14,7 @@ import { checkExecutionUsageLimits, checkIngestionUsageLimits, checkSearchUsageLimits, + NEAR_LIMIT_USAGE_GATE_TTL_MS, resetUsageGateCache, USAGE_GATE_SETTLE_TIMEOUT_MS, USAGE_GATE_TTL_MS, @@ -21,6 +22,30 @@ import { const mockCheck = billingAttributionMockFns.mockCheckAttributedUsageLimits +/** Fakes `setTimeout` for {@link advanceClockPast}; install it before the cache is first read. */ +function useFakeTimeouts(): void { + vi.useFakeTimers({ toFake: ['setTimeout', 'clearTimeout'] }) +} + +/** + * Fires the cache's pending clock debounce before restoring real timers; dropped unfired, it would + * pin the module-level cache to this test's spied clock for every later test. + */ +function restoreRealTimers(): void { + vi.runOnlyPendingTimers() + vi.useRealTimers() +} + +/** + * `lru-cache` captures `performance` at load, so fake timers cannot move its clock: spy the clock + * forward, then fire the faked 1 ms timer behind which it debounces `performance.now()`. + */ +async function advanceClockPast(ms: number): Promise { + const start = performance.now() + vi.spyOn(performance, 'now').mockReturnValue(start + ms + 1) + await vi.advanceTimersByTimeAsync(1) +} + const ATTRIBUTION: BillingAttributionSnapshot = { actorUserId: 'member-1', workspaceId: null, @@ -74,17 +99,19 @@ describe('checkIngestionUsageLimits', () => { }) it('keeps a refusal for the same bounded window, then reads fresh', async () => { - mockCheck.mockResolvedValueOnce({ isExceeded: true, scope: 'payer', message: 'over' }) - expect((await checkIngestionUsageLimits(ATTRIBUTION)).isExceeded).toBe(true) - expect((await checkIngestionUsageLimits(ATTRIBUTION)).isExceeded).toBe(true) - expect(mockCheck).toHaveBeenCalledTimes(1) + useFakeTimeouts() + try { + mockCheck.mockResolvedValueOnce({ isExceeded: true, scope: 'payer', message: 'over' }) + expect((await checkIngestionUsageLimits(ATTRIBUTION)).isExceeded).toBe(true) + expect((await checkIngestionUsageLimits(ATTRIBUTION)).isExceeded).toBe(true) + expect(mockCheck).toHaveBeenCalledTimes(1) - /** `lru-cache` reads `performance.now()` and debounces it behind a real 1 ms timer. */ - const start = performance.now() - vi.spyOn(performance, 'now').mockReturnValue(start + USAGE_GATE_TTL_MS + 1) - await sleep(5) - expect((await checkIngestionUsageLimits(ATTRIBUTION)).isExceeded).toBe(false) - expect(mockCheck).toHaveBeenCalledTimes(2) + await advanceClockPast(USAGE_GATE_TTL_MS) + expect((await checkIngestionUsageLimits(ATTRIBUTION)).isExceeded).toBe(false) + expect(mockCheck).toHaveBeenCalledTimes(2) + } finally { + restoreRealTimers() + } }) it('separates answers by actor, period, period source, payer and plan', async () => { @@ -212,3 +239,72 @@ describe('checkExecutionUsageLimits', () => { } }) }) + +describe('admission lifetime by headroom', () => { + beforeEach(() => { + useFakeTimeouts() + resetUsageGateCache() + mockCheck.mockReset() + }) + + afterEach(() => { + restoreRealTimers() + }) + + it('re-reads a free-plan admission after the near-limit window, so spend past the limit is refused', async () => { + mockCheck + .mockResolvedValueOnce({ isExceeded: false, payerUsage: { currentUsage: 1, limit: 5 } }) + .mockResolvedValueOnce({ + isExceeded: true, + scope: 'payer', + message: 'over', + payerUsage: { currentUsage: 5.4, limit: 5 }, + }) + expect((await checkExecutionUsageLimits(ATTRIBUTION)).isExceeded).toBe(false) + expect((await checkExecutionUsageLimits(ATTRIBUTION)).isExceeded).toBe(false) + expect(mockCheck).toHaveBeenCalledTimes(1) + + await advanceClockPast(NEAR_LIMIT_USAGE_GATE_TTL_MS) + expect((await checkExecutionUsageLimits(ATTRIBUTION)).isExceeded).toBe(true) + expect(mockCheck).toHaveBeenCalledTimes(2) + }) + + it('re-reads a large pool in the last fifth of its limit', async () => { + mockCheck.mockResolvedValue({ + isExceeded: false, + payerUsage: { currentUsage: 900, limit: 1000 }, + }) + await checkExecutionUsageLimits(ATTRIBUTION) + await advanceClockPast(NEAR_LIMIT_USAGE_GATE_TTL_MS) + await checkExecutionUsageLimits(ATTRIBUTION) + expect(mockCheck).toHaveBeenCalledTimes(2) + }) + + it('re-reads when the acting member is near their cap though the pool has room', async () => { + mockCheck.mockResolvedValue({ + isExceeded: false, + payerUsage: { currentUsage: 100, limit: 1000 }, + memberUsage: { currentUsage: 48, limit: 50 }, + }) + await checkSearchUsageLimits(ATTRIBUTION) + await advanceClockPast(NEAR_LIMIT_USAGE_GATE_TTL_MS) + await checkSearchUsageLimits(ATTRIBUTION) + expect(mockCheck).toHaveBeenCalledTimes(2) + }) + + it('holds an admission with ample headroom for the full TTL', async () => { + mockCheck.mockResolvedValue({ + isExceeded: false, + payerUsage: { currentUsage: 100, limit: 1000 }, + memberUsage: { currentUsage: 10, limit: null }, + }) + await checkExecutionUsageLimits(ATTRIBUTION) + await advanceClockPast(NEAR_LIMIT_USAGE_GATE_TTL_MS) + await checkExecutionUsageLimits(ATTRIBUTION) + expect(mockCheck).toHaveBeenCalledTimes(1) + + await advanceClockPast(USAGE_GATE_TTL_MS) + await checkExecutionUsageLimits(ATTRIBUTION) + expect(mockCheck).toHaveBeenCalledTimes(2) + }) +}) diff --git a/apps/sim/lib/billing/core/usage-gate-cache.ts b/apps/sim/lib/billing/core/usage-gate-cache.ts index 230cf09e7d7..30677b92921 100644 --- a/apps/sim/lib/billing/core/usage-gate-cache.ts +++ b/apps/sim/lib/billing/core/usage-gate-cache.ts @@ -13,14 +13,50 @@ import { coalesceLocally } from '@/lib/concurrency/singleflight' * The gate sums the payer's usage ledger for the billing period, which grows * with the payer's activity, so a large organization scans its whole period on * every uncached call. Bulk ingestion re-checks per document and knowledge - * search checks per query. Staleness is bounded by this TTL and fails in the - * harmless direction: a payer who crosses their limit keeps going for at most - * this long, which charges nobody wrongly. Five minutes: the sum is a few - * hundred milliseconds for a busy payer, and a minute made every search after - * a pause pay it. + * search checks per query. A payer who crosses their limit keeps going for at + * most this long, so only an admission with ample headroom is held this long + * (see {@link usageAdmissionTtlMs}). Five minutes: the sum is a few hundred + * milliseconds for a busy payer, and a minute made every search after a pause + * pay it. */ export const USAGE_GATE_TTL_MS = 5 * 60 * 1000 +/** + * How long an admission near its limit stays usable. Every spend landing + * inside this window is spend past the limit nobody bills (a free plan) or + * past a cap the payer set, so it is short; it stays above zero so a burst of + * executions or model steps from one payer shares a read instead of summing + * the ledger on each. Staleness is this plus any lag in the usage read itself. + */ +export const NEAR_LIMIT_USAGE_GATE_TTL_MS = 10 * 1000 + +/** + * An admission is near its limit when the remaining headroom is under the + * larger of this floor and {@link NEAR_LIMIT_HEADROOM_SHARE} of the limit. The + * floor puts every free-plan payer permanently near its limit; the share keeps + * a large pool from caching its last stretch. + */ +const NEAR_LIMIT_HEADROOM_FLOOR_USD = 25 +const NEAR_LIMIT_HEADROOM_SHARE = 0.2 + +type UsageStanding = NonNullable + +function isNearLimit(usage: UsageStanding | undefined): boolean { + if (!usage || usage.limit === null) return false + const headroom = usage.limit - usage.currentUsage + return headroom < Math.max(NEAR_LIMIT_HEADROOM_FLOOR_USD, usage.limit * NEAR_LIMIT_HEADROOM_SHARE) +} + +/** + * How long an admission may be served from cache: the full gate TTL while + * every standing (the payer's, the acting member's) has ample headroom, + * otherwise {@link NEAR_LIMIT_USAGE_GATE_TTL_MS}. An admission that carries no + * figures (billing off) has no limit to approach. + */ +export function usageAdmissionTtlMs(...standings: Array): number { + return standings.some(isNearLimit) ? NEAR_LIMIT_USAGE_GATE_TTL_MS : USAGE_GATE_TTL_MS +} + /** * How long a coalesced usage read may take before its callers give up on it. The read's cost is * its ledger aggregates, each of which the database ends at @@ -78,7 +114,8 @@ function gateKey(attribution: BillingAttributionSnapshot): string { * `coalesceLocally` collapses concurrent misses onto one ledger read and bounds * a hung read at {@link USAGE_GATE_SETTLE_TIMEOUT_MS}. The write stays on the * value this caller received, so a producer that timed out and later resolved - * cannot overwrite a fresher answer. + * cannot overwrite a fresher answer. An admission is stored for + * {@link usageAdmissionTtlMs}, a refusal for the full TTL. * * There is deliberately no invalidator: usage and limit changes land in other * processes (execution workers, Stripe webhooks), so the TTL is the real bound. @@ -96,7 +133,13 @@ async function checkUsageLimitsThroughCache( () => checkAttributedUsageLimits(attribution), USAGE_GATE_SETTLE_TIMEOUT_MS ) - if (cacheRefusals || !result.isExceeded) gateCache.set(key, result) + if (cacheRefusals || !result.isExceeded) { + gateCache.set(key, result, { + ttl: result.isExceeded + ? USAGE_GATE_TTL_MS + : usageAdmissionTtlMs(result.payerUsage, result.memberUsage), + }) + } return result } diff --git a/packages/testing/src/mocks/billing-usage-gate-cache.mock.ts b/packages/testing/src/mocks/billing-usage-gate-cache.mock.ts index 4142506127e..75ebe916263 100644 --- a/packages/testing/src/mocks/billing-usage-gate-cache.mock.ts +++ b/packages/testing/src/mocks/billing-usage-gate-cache.mock.ts @@ -5,7 +5,7 @@ import { vi } from 'vitest' * * The three `check*UsageLimits` gates are bare `vi.fn()` (resolve `undefined`); set * `mockResolvedValue({ isExceeded: false })` for an admitted path. `resetUsageGateCache` is a - * no-op. + * no-op, and `usageAdmissionTtlMs` returns the full gate TTL. * * @example * ```ts @@ -21,6 +21,8 @@ export const billingUsageGateCacheMockFns = { mockResetUsageGateCache: vi.fn((): void => {}), } +const USAGE_GATE_TTL_MS = 5 * 60 * 1000 + /** * Static mock module for `@/lib/billing/core/usage-gate-cache`. Constants carry the real * values (`USAGE_GATE_SETTLE_TIMEOUT_MS` = 2 x the 60 s ledger statement timeout + 15 s). @@ -31,10 +33,11 @@ export const billingUsageGateCacheMockFns = { * ``` */ export const billingUsageGateCacheMock = { - USAGE_GATE_TTL_MS: 5 * 60 * 1000, + USAGE_GATE_TTL_MS, USAGE_GATE_SETTLE_TIMEOUT_MS: 2 * 60_000 + 15_000, checkIngestionUsageLimits: billingUsageGateCacheMockFns.mockCheckIngestionUsageLimits, checkSearchUsageLimits: billingUsageGateCacheMockFns.mockCheckSearchUsageLimits, checkExecutionUsageLimits: billingUsageGateCacheMockFns.mockCheckExecutionUsageLimits, resetUsageGateCache: billingUsageGateCacheMockFns.mockResetUsageGateCache, + usageAdmissionTtlMs: vi.fn((): number => USAGE_GATE_TTL_MS), } diff --git a/scripts/test-patterns-baseline.json b/scripts/test-patterns-baseline.json index 30969c0dc8f..0d7ed8685bb 100644 --- a/scripts/test-patterns-baseline.json +++ b/scripts/test-patterns-baseline.json @@ -151,7 +151,6 @@ "real-wait\tapps/sim/hooks/queries/workspace.test.tsx\tsleep() #1", "real-wait\tapps/sim/lib/billing/core/reporting-usage-cache.test.ts\tsleep() #1", "real-wait\tapps/sim/lib/billing/core/reporting-usage-cache.test.ts\tsleep() #2", - "real-wait\tapps/sim/lib/billing/core/usage-gate-cache.test.ts\tsleep() #1", "real-wait\tapps/sim/lib/concurrency/leader-lock.test.ts\tsleep() #1", "real-wait\tapps/sim/lib/concurrency/singleflight.test.ts\tsleep() #1", "real-wait\tapps/sim/lib/core/async-jobs/backends/database.test.ts\tsleep() #1",