diff --git a/apps/sim/app/api/billing/update-cost/route.test.ts b/apps/sim/app/api/billing/update-cost/route.test.ts index d69abad26b3..15893913840 100644 --- a/apps/sim/app/api/billing/update-cost/route.test.ts +++ b/apps/sim/app/api/billing/update-cost/route.test.ts @@ -1040,32 +1040,38 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => { }) }) - it('never pauses a blocked payer with the usage card', async () => { - billingAttributionMockFns.mockCheckAccountBillingBlocks.mockResolvedValue({ - blocked: true, - scope: 'payer', + it("offers the card for its admitted payer's plan, not the actor's current one", async () => { + billingCoreMockFns.mockGetOrganizationSubscription.mockResolvedValue({ + id: 'sub-account-org', + referenceId: 'account-org', + plan: 'team', + status: 'active', + seats: 4, + }) + billingPlanMockFns.mockGetHighestPrioritySubscription.mockResolvedValue({ + id: 'sub-personal', + referenceId: 'user-1', + plan: 'pro', + status: 'active', }) const body = await (await POST(directCallback())).json() - expect(body.usageExceeded).toBe(false) + expect(body.usageUpgrade).toMatchObject({ + action: 'increase_limit', + message: expect.stringContaining("organization's usage limit"), + }) }) - it('keeps the exceeded verdict with the plan-upgrade card when the card read outlasts the callback budget', async () => { - billingPlanMockFns.mockGetHighestPrioritySubscription.mockImplementation(async () => { - await sleep(1500) - return { plan: 'pro' } + it('never pauses a blocked payer with the usage card', async () => { + billingAttributionMockFns.mockCheckAccountBillingBlocks.mockResolvedValue({ + blocked: true, + scope: 'payer', }) - const startedAt = Date.now() const body = await (await POST(directCallback())).json() - expect(body).toMatchObject({ - success: true, - usageExceeded: true, - usageUpgrade: { action: 'upgrade_plan' }, - }) - expect(Date.now() - startedAt).toBeLessThan(1400) + expect(body.usageExceeded).toBe(false) }) }) diff --git a/apps/sim/app/api/billing/update-cost/route.ts b/apps/sim/app/api/billing/update-cost/route.ts index 2a7bdfda6cd..bec648abdc6 100644 --- a/apps/sim/app/api/billing/update-cost/route.ts +++ b/apps/sim/app/api/billing/update-cost/route.ts @@ -86,7 +86,7 @@ function invalidBillingProtocolResponse(requestId: string, span: Span): NextResp * 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 - * verdict always pauses the run; a card read past that budget falls back to the plan-upgrade card. + * verdict always pauses the run. */ async function readUsageStanding( userId: string, @@ -99,10 +99,9 @@ async function readUsageStanding( ? () => readMidRunAccountUsageVerdict(accountDecision) : null if (!isHosted || !readVerdict) return { usageExceeded: false } - const deadlineAt = Date.now() + USAGE_STANDING_TIMEOUT_MS let verdict: MidRunUsageVerdict try { - verdict = await withinDeadline(readVerdict, deadlineAt) + verdict = await withinDeadline(readVerdict, Date.now() + USAGE_STANDING_TIMEOUT_MS) } catch { logger.warn('Usage standing read outlasted the callback budget; answering not exceeded') return { usageExceeded: false } @@ -114,9 +113,8 @@ async function readUsageStanding( usageExceeded: true, usageUpgrade: await resolveUsageUpgradePayload( userId, - billingAttribution, - verdict.scope, - deadlineAt + billingAttribution ?? verdict.payer, + verdict.scope ), } } 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 f9a88e61d3f..9c2188b9ac8 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 @@ -140,7 +140,10 @@ vi.mock('@/lib/workspaces/permissions/utils', () => permissionsMock) vi.mock('@/lib/workspaces/utils', () => workspacesUtilsMock) -import { validateCopilotApiKeyBodySchema } from '@/lib/api/contracts/copilot' +import { + validateCopilotApiKeyBodySchema, + 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 { POST } from '@/app/api/copilot/api-keys/validate/route' @@ -248,6 +251,17 @@ describe('POST /api/copilot/api-keys/validate billing protocols', () => { expect(mockCheckServerSideUsageLimits).not.toHaveBeenCalled() }) + it("sends a new turn's usage refusal as the empty 402 its contract declares", async () => { + mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'payer' }) + + const res = await POST(request(SELF_HOSTED_VALIDATE_BODY)) + + expect(res.status).toBe(402) + expect(await res.text()).toBe('') + const refusalSchema = validateCopilotApiKeyContract.response.statusSchemas?.[402] + expect(refusalSchema?.safeParse(undefined).success).toBe(true) + }) + it('preserves the actor member cap for markerless self-hosted admission', async () => { mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: true, @@ -593,6 +607,28 @@ describe('validation lifecycle purposes', () => { }) }) + it("refuses a direct-v1 continuation with the card for its admitted payer's plan", async () => { + mockCheckUsageStatus.mockResolvedValue({ isExceeded: true, currentUsage: 12, limit: 10 }) + mockGetOrganizationSubscription.mockResolvedValue({ + id: 'sub-account-org', + referenceId: 'account-org', + plan: 'team', + status: 'active', + seats: 4, + }) + mockGetHighestPrioritySubscription.mockResolvedValue(null) + + const response = await POST(request(body, directHeaders)) + + expect(response.status).toBe(402) + await expect(response.json()).resolves.toMatchObject({ + usageUpgrade: { + action: 'increase_limit', + message: expect.stringContaining("organization's usage limit"), + }, + }) + }) + 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/app/api/copilot/api-keys/validate/route.ts b/apps/sim/app/api/copilot/api-keys/validate/route.ts index 8723a9d5c90..0a1373dea56 100644 --- a/apps/sim/app/api/copilot/api-keys/validate/route.ts +++ b/apps/sim/app/api/copilot/api-keys/validate/route.ts @@ -453,7 +453,7 @@ export const POST = withRouteHandler((req: NextRequest) => span.setAttribute(TraceAttr.HttpStatusCode, 402) const usageUpgrade = await resolveUsageUpgradePayload( userId, - billing?.kind === 'attributed' ? billing.attribution : undefined, + billing?.kind === 'attributed' ? billing.attribution : verdict.payer, verdict.scope ) return NextResponse.json( diff --git a/apps/sim/lib/api/contracts/copilot.ts b/apps/sim/lib/api/contracts/copilot.ts index 66ba38ec119..5989ca03b4e 100644 --- a/apps/sim/lib/api/contracts/copilot.ts +++ b/apps/sim/lib/api/contracts/copilot.ts @@ -433,7 +433,7 @@ export const validateCopilotApiKeyContract = defineRouteContract({ status: [200, 402], statusSchemas: { 200: validateCopilotApiKeyResponseSchema, - 402: validateCopilotApiKeyRefusalSchema, + 402: validateCopilotApiKeyRefusalSchema.optional(), }, }, error: validateCopilotApiKeyErrorSchema, diff --git a/apps/sim/lib/billing/core/mid-run-usage.ts b/apps/sim/lib/billing/core/mid-run-usage.ts index 36ea034f81e..96393ba26bc 100644 --- a/apps/sim/lib/billing/core/mid-run-usage.ts +++ b/apps/sim/lib/billing/core/mid-run-usage.ts @@ -18,6 +18,7 @@ import { USAGE_GATE_SETTLE_TIMEOUT_MS, USAGE_GATE_TTL_MS, } from '@/lib/billing/core/usage-gate-cache' +import type { UsageUpgradePayer } from '@/lib/billing/usage-upgrade' import { coalesceLocally } from '@/lib/concurrency/singleflight' import { isBillingEnabled, isHosted } from '@/lib/core/config/env-flags' @@ -26,7 +27,8 @@ const logger = createLogger('MidRunUsage') /** * A run's standing while it is under way, read through the execution usage gate: * - `exceeded`: the payer (or the actor's member cap) spent its limit; the run pauses with the - * upgrade card. + * upgrade card. A direct-v1 verdict carries the payer it read, so the card names that payer's + * plan rather than the actor's. * - `blocked`: the account is blocked (payment failed, dispute); the run is refused as a blocked * account, never with the upgrade card. * - `unknown`: usage, or the payer's current period, could not be read. Admission fails closed @@ -35,7 +37,11 @@ const logger = createLogger('MidRunUsage') */ export type MidRunUsageVerdict = | { status: 'within' } - | { status: 'exceeded'; scope?: AttributedUsageLimitsResult['scope'] } + | { + status: 'exceeded' + scope?: AttributedUsageLimitsResult['scope'] + payer?: UsageUpgradePayer + } | { status: 'blocked'; message?: string } | { status: 'unknown' } @@ -213,7 +219,13 @@ export async function readMidRunAccountUsageVerdict( USAGE_GATE_SETTLE_TIMEOUT_MS ) if (usage.unavailable) return { status: 'unknown' } - if (usage.isExceeded) return { status: 'exceeded', scope: 'payer' } + if (usage.isExceeded) { + return { + status: 'exceeded', + scope: 'payer', + payer: { billingEntity: payer, payerSubscription: subscription }, + } + } const within: MidRunUsageVerdict = { status: 'within' } accountVerdictCache.set(key, within) return within diff --git a/apps/sim/lib/billing/usage-upgrade.ts b/apps/sim/lib/billing/usage-upgrade.ts index 95cf5813539..ceee5e1fe9e 100644 --- a/apps/sim/lib/billing/usage-upgrade.ts +++ b/apps/sim/lib/billing/usage-upgrade.ts @@ -1,14 +1,11 @@ import { createLogger } from '@sim/logger' import { getErrorMessage } from '@sim/utils/errors' import type { UsageUpgradePayload } from '@/lib/api/contracts/subscription' -import type { - AttributedUsageLimitsResult, - BillingAttributionSnapshot, -} from '@/lib/billing/core/billing-attribution' +import type { AttributedUsageLimitsResult } from '@/lib/billing/core/billing-attribution' import { getHighestPrioritySubscription } from '@/lib/billing/core/plan' +import type { BillingEntity } from '@/lib/billing/core/usage-log' import { isEnterprise, isPaid } from '@/lib/billing/plan-helpers' import { isOrgScopedSubscription } from '@/lib/billing/subscriptions/utils' -import { withinDeadline } from '@/lib/core/utils/deadline' const logger = createLogger('UsageUpgrade') @@ -18,19 +15,26 @@ const UPGRADE_PLAN_MESSAGE = const MEMBER_CAP_MESSAGE = "You've reached the usage limit your organization set for you this billing period. Only an organization owner or admin can raise it — please ask them to continue." +/** + * The payer a run is billed to, as the upgrade card needs it: its billing entity and its + * subscription's plan. An attribution snapshot is one; a direct-v1 run's mid-run verdict carries one. + */ +export interface UsageUpgradePayer { + readonly billingEntity: Readonly + readonly payerSubscription: { readonly plan: string } | null +} + /** * The upgrade card for a payer over its usage limit: a plan upgrade for a free payer, a limit * increase for a paid one, with copy naming who can raise an organization's limit. A member - * over the cap their organization set gets copy naming who can raise that cap. An attributed - * run reads the plan from its admission snapshot without a query; otherwise the actor's current - * subscription decides, and a lookup that fails or outlasts `deadlineAt` falls back to the - * plan-upgrade card. + * over the cap their organization set gets copy naming who can raise that cap. A known payer + * decides the card without a query; otherwise the actor's current subscription decides, and a + * lookup that fails falls back to the plan-upgrade card. */ export async function resolveUsageUpgradePayload( userId: string, - billingAttribution?: BillingAttributionSnapshot, - scope?: AttributedUsageLimitsResult['scope'], - deadlineAt?: number + payer?: UsageUpgradePayer, + scope?: AttributedUsageLimitsResult['scope'] ): Promise { if (scope === 'member') { return { reason: 'usage_limit', action: 'increase_limit', message: MEMBER_CAP_MESSAGE } @@ -38,13 +42,11 @@ export async function resolveUsageUpgradePayload( let plan: string | undefined let orgScoped = false try { - if (billingAttribution) { - plan = billingAttribution.payerSubscription?.plan - orgScoped = billingAttribution.billingEntity.type === 'organization' + if (payer) { + plan = payer.payerSubscription?.plan + orgScoped = payer.billingEntity.type === 'organization' } else { - const subscription = await (deadlineAt === undefined - ? getHighestPrioritySubscription(userId) - : withinDeadline(() => getHighestPrioritySubscription(userId), deadlineAt)) + const subscription = await getHighestPrioritySubscription(userId) plan = subscription?.plan orgScoped = isOrgScopedSubscription(subscription, userId) }