diff --git a/src/backend/controllers/events/limits.ts b/src/backend/controllers/events/limits.ts index ad3a1a9e2..3431bd6ee 100644 --- a/src/backend/controllers/events/limits.ts +++ b/src/backend/controllers/events/limits.ts @@ -21,7 +21,7 @@ import type { RouteRateLimit } from '../../core/http/types'; import { DEFAULT_FREE_SUBSCRIPTION, DEFAULT_TEMP_SUBSCRIPTION, - FREE_SUBSCRIPTION_IDS, + isFreeSubscription, } from '../../services/metering/consts.js'; // -- Shared event limits --------------------------------------------- @@ -69,7 +69,7 @@ export const limitFor = ( if (subscriptionId === null) return tier.limit; const own = tier.bySubscription[subscriptionId]; if (typeof own === 'number') return own; - return FREE_SUBSCRIPTION_IDS.has(subscriptionId) + return isFreeSubscription(subscriptionId) ? (tier.bySubscription[DEFAULT_FREE_SUBSCRIPTION] ?? tier.limit) : tier.limit; }; diff --git a/src/backend/core/http/middleware/rateLimit.js b/src/backend/core/http/middleware/rateLimit.js index af02e32ce..7e5ae37b8 100644 --- a/src/backend/core/http/middleware/rateLimit.js +++ b/src/backend/core/http/middleware/rateLimit.js @@ -23,7 +23,7 @@ import { withSpan } from '../../../util/span.js'; import { HttpError } from '../HttpError.js'; import { DEFAULT_FREE_SUBSCRIPTION, - FREE_SUBSCRIPTION_IDS, + isFreeSubscription, } from '../../../services/metering/consts.js'; /** @@ -873,7 +873,7 @@ export const CONCURRENT_SLOT_TTL_MS = ORPHAN_SAFETY_TTL_MS; function overrideFor(bySubscription, subscriptionId) { const own = bySubscription[subscriptionId]; if (typeof own === 'number') return own; - return FREE_SUBSCRIPTION_IDS.has(subscriptionId) + return isFreeSubscription(subscriptionId) ? bySubscription[DEFAULT_FREE_SUBSCRIPTION] : undefined; } diff --git a/src/backend/drivers/ai-chat/ChatCompletionDriver.test.ts b/src/backend/drivers/ai-chat/ChatCompletionDriver.test.ts index 0cbcceb90..18b537942 100644 --- a/src/backend/drivers/ai-chat/ChatCompletionDriver.test.ts +++ b/src/backend/drivers/ai-chat/ChatCompletionDriver.test.ts @@ -915,6 +915,46 @@ describe('ChatCompletionDriver.complete credit gate and max_tokens cap', () => { }); }); + it('admits a plan outside the free set to subscriber-only models', async () => { + vi.spyOn(FakeChatProvider.prototype, 'models').mockResolvedValueOnce([ + { + id: 'subonly-paid', + aliases: [], + costs_currency: 'usd-cents', + costs: { 'input-tokens': 100, 'output-tokens': 100 }, + max_tokens: 8192, + subscriberOnly: true, + }, + ]); + const d = await makeDriver(); + vi.spyOn(server.services.metering, 'getRemainingUsage').mockResolvedValue( + 1_000_000, + ); + vi.spyOn( + server.services.metering, + 'getActorSubscription', + ).mockResolvedValue({ id: 'some-paid-tier' } as never); + const completeSpy = vi + .spyOn(FakeChatProvider.prototype, 'complete') + .mockResolvedValueOnce({ + message: { + role: 'assistant', + content: [{ type: 'text', text: 'ok' }], + }, + usage: { input_tokens: 1, output_tokens: 1 }, + finish_reason: 'stop', + } as never); + + await withTestActor(() => + d.complete({ + model: 'subonly-paid', + messages: [{ role: 'user', content: 'hi' }], + }), + ); + + expect(completeSpy).toHaveBeenCalledOnce(); + }); + it('rejects subscriber-only models for the default free subscription', async () => { vi.spyOn(FakeChatProvider.prototype, 'models').mockResolvedValueOnce([ { diff --git a/src/backend/drivers/ai-chat/ChatCompletionDriver.ts b/src/backend/drivers/ai-chat/ChatCompletionDriver.ts index 005f8c4f6..6ac50c2af 100644 --- a/src/backend/drivers/ai-chat/ChatCompletionDriver.ts +++ b/src/backend/drivers/ai-chat/ChatCompletionDriver.ts @@ -23,7 +23,7 @@ import { EventMap } from '../../clients/event/types.js'; import type { Actor } from '../../core/actor.js'; import { Context } from '../../core/context.js'; import { HttpError, isHttpError } from '../../core/http/HttpError.js'; -import { FREE_SUBSCRIPTION_IDS } from '../../services/metering/consts.js'; +import { isFreeSubscription } from '../../services/metering/consts.js'; import type { CreditHold } from '../../services/metering/types.js'; import { NO_CREDIT_HOLD } from '../../services/metering/types.js'; import type { DriverStreamResult } from '../meta.js'; @@ -977,8 +977,7 @@ export class ChatCompletionDriver extends PuterDriver { if (model.subscriberOnly) { const subscription = await metering.getActorSubscription(actor); - // Every free plan, not two named ones. - if (FREE_SUBSCRIPTION_IDS.has(subscription.id)) { + if (isFreeSubscription(subscription.id)) { throw new HttpError( 403, `The model ${model.id} is only available to subscribers. Please subscribe to access this model.`, diff --git a/src/backend/services/metering/MeteringService.ts b/src/backend/services/metering/MeteringService.ts index 91362c81c..33ea5b86a 100644 --- a/src/backend/services/metering/MeteringService.ts +++ b/src/backend/services/metering/MeteringService.ts @@ -26,7 +26,6 @@ import { DEFAULT_FREE_SUBSCRIPTION, DEFAULT_TEMP_SUBSCRIPTION, DETAIL_PATH_COUNTER, - FREE_SUBSCRIPTION_IDS, GLOBAL_APP_KEY, METRICS_PREFIX, METRICS_V2_PREFIX, @@ -37,6 +36,7 @@ import { UNLIMITED_SUBSCRIPTION, USAGE_DETAIL_SHARD_COUNT, V1_CLAIM_THROUGH_MONTH, + isFreeSubscription, } from './consts'; import { EGRESS_COSTS } from './costs'; import type { @@ -2078,7 +2078,7 @@ export class MeteringService extends PuterService { month: string = this.monthYearString(), ): Promise { if ( - FREE_SUBSCRIPTION_IDS.has(policy.id) || + isFreeSubscription(policy.id) || policy.id === UNLIMITED_SUBSCRIPTION || !(policy.monthUsageAllowance > 0) ) { diff --git a/src/backend/services/metering/consts.ts b/src/backend/services/metering/consts.ts index a36a0e4bf..3df76bca5 100644 --- a/src/backend/services/metering/consts.ts +++ b/src/backend/services/metering/consts.ts @@ -78,5 +78,9 @@ export const FREE_SUBSCRIPTION_IDS: ReadonlySet = new Set([ ORG_SEAT_FREE_SUBSCRIPTION, ]); +/** Whether a policy id is a free plan. Check free-ness here, not on the set. */ +export const isFreeSubscription = (id: string): boolean => + FREE_SUBSCRIPTION_IDS.has(id); + // WARNING: DO NOT USE THESE IN PROD export const UNLIMITED_SUBSCRIPTION = 'unlimited'; diff --git a/src/backend/services/metering/enforcement.test.ts b/src/backend/services/metering/enforcement.test.ts index 2d0fa90cb..d26d1ff16 100644 --- a/src/backend/services/metering/enforcement.test.ts +++ b/src/backend/services/metering/enforcement.test.ts @@ -21,6 +21,7 @@ import { describe, expect, it, vi } from 'vitest'; import { SYSTEM_ACTOR, type Actor } from '../../core/actor.js'; import { HttpError } from '../../core/http/HttpError.js'; import type { IConfig } from '../../types'; +import { FREE_SUBSCRIPTION_IDS, isFreeSubscription } from './consts.js'; import { actorHasSubscription, actorOnPaidPlan, @@ -260,6 +261,28 @@ describe('subscriptionSatisfies', () => { }); }); +describe('isFreeSubscription', () => { + it('is true for exactly the free policies', () => { + for (const id of FREE_SUBSCRIPTION_IDS) { + expect(isFreeSubscription(id)).toBe(true); + } + for (const id of ['professional', 'unlimited', 'some-paid-tier', '']) { + expect(isFreeSubscription(id)).toBe(false); + } + // Untyped callers can pass an id that never resolved. + expect(isFreeSubscription(undefined as never)).toBe(false); + expect(isFreeSubscription(null as never)).toBe(false); + }); + + it('is the complement of a `true` requirement', () => { + for (const id of [...FREE_SUBSCRIPTION_IDS, 'professional', '']) { + expect(subscriptionSatisfies(id, true)).toBe( + !isFreeSubscription(id), + ); + } + }); +}); + describe('validateSubscriptionRequirement', () => { it('passes booleans and non-empty id lists through', () => { expect(validateSubscriptionRequirement(true, 'x')).toBe(true); diff --git a/src/backend/services/metering/enforcement.ts b/src/backend/services/metering/enforcement.ts index 81e36e9e6..7b9ada955 100644 --- a/src/backend/services/metering/enforcement.ts +++ b/src/backend/services/metering/enforcement.ts @@ -21,7 +21,7 @@ import type { Actor } from '../../core/actor'; import { isSystemActor } from '../../core/actor'; import { HttpError } from '../../core/http/HttpError.js'; import type { IConfig } from '../../types'; -import { FREE_SUBSCRIPTION_IDS } from './consts.js'; +import { isFreeSubscription } from './consts.js'; import type { MeteringService } from './MeteringService'; // -- Credit enforcement ---------------------------------------------- @@ -148,7 +148,7 @@ export const subscriptionSatisfies = ( if (Array.isArray(requirement)) { return requirement.length === 0 || requirement.includes(id); } - return !FREE_SUBSCRIPTION_IDS.has(id); + return !isFreeSubscription(id); }; /**