diff --git a/src/backend/clients/event/types.ts b/src/backend/clients/event/types.ts index d42253d0d..c4531b253 100644 --- a/src/backend/clients/event/types.ts +++ b/src/backend/clients/event/types.ts @@ -767,6 +767,12 @@ export type EventMap = { // `ai.cost.factor..:`. Emitted once per model // per batch; the last listener to set `factor` wins. [K in `ai.cost.factor.${string}`]: AiCostFactorEvent; +} & { + // A plan gate is about to refuse a caller. Keyed per surface the way the + // lifecycles above are: `subscription.gate.route..`, + // `subscription.gate.driver..`. Emitted via `emitAndWait`; + // a listener that sets `allow` waives the plan requirement for that call. + [K in `subscription.gate.${string}`]: SubscriptionGateEvent; } & { [K in `pubsub.login.${string}`]: { authtoken: string }; } & { @@ -787,6 +793,18 @@ export type EventMap = { 'outer.pubsub.metering.credits-changed': { userUuid: string }; } & IExtensionEventMap; +/** Payload for `subscription.gate.` events. */ +export type SubscriptionGateEvent = { + actor: Actor; + /** The key after `subscription.gate.`, e.g. `driver.puter-kvstore.get`. */ + surface: string; + /** What the gate asked for: any paid plan (`true`) or these policy ids. */ + requirement: boolean | readonly string[]; + /** The request being gated, when there is one. */ + req?: ExpressRequest; + allow: boolean; +}; + /** Payload for `ai.cost.factor..` events. */ export type AiCostFactorEvent = { /** Driver doing the pricing, e.g. `ai-chat`. */ diff --git a/src/backend/controllers/drivers/DriverController.subscription.test.ts b/src/backend/controllers/drivers/DriverController.subscription.test.ts index e483cb1cd..317d452ba 100644 --- a/src/backend/controllers/drivers/DriverController.subscription.test.ts +++ b/src/backend/controllers/drivers/DriverController.subscription.test.ts @@ -21,6 +21,7 @@ import { EventEmitter } from 'node:events'; import type { Request, RequestHandler, Response } from 'express'; import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { EventClient } from '../../clients/event/EventClient.js'; import { runWithContext } from '../../core/context.js'; import { isHttpError } from '../../core/http/HttpError.js'; import { configureRateLimit } from '../../core/http/middleware/rateLimit.js'; @@ -97,8 +98,8 @@ const syntheticDriver = { premium: async () => 'premium-ok', }; -const buildController = (plan: string | null) => { - const clients = { alarm: { create: () => {} } }; +const buildController = (plan: string | null, events?: EventClient) => { + const clients = { alarm: { create: () => {} }, ...(events ? { event: events } : {}) }; const services = plan ? { metering: { @@ -115,8 +116,13 @@ const buildController = (plan: string | null) => { ); }; -const call = (plan: string | null, method: string, tag: string) => { - const handler = captureCallHandler(buildController(plan)); +const call = ( + plan: string | null, + method: string, + tag: string, + events?: EventClient, +) => { + const handler = captureCallHandler(buildController(plan, events)); const res = new StubRes(); return runWithContext({}, () => handler(makeReq(method, tag), res as unknown as Response, () => {}), @@ -157,6 +163,17 @@ describe('DriverController — per-method subscription requirement', () => { expect(res.body).toMatchObject({ success: true, result: 'premium-ok' }); }); + it('lets a subscription.gate listener waive one method for a free account', async () => { + const events = new EventClient({} as any); + events.on('subscription.gate.driver.test-iface.paid', (_key, event) => { + (event as { allow: boolean }).allow = true; + }); + + const res = await call('user_free', 'paid', 'waived', events); + expect(res.body).toMatchObject({ success: true, result: 'paid-ok' }); + await expect402(call('user_free', 'premium', 'not-waived', events)); + }); + it('enforces nothing on a deployment with no metering service', async () => { const res = await call(null, 'paid', 'no-metering'); expect(res.body).toMatchObject({ success: true, result: 'paid-ok' }); diff --git a/src/backend/controllers/drivers/DriverController.ts b/src/backend/controllers/drivers/DriverController.ts index 32a7c9eab..f3f302595 100644 --- a/src/backend/controllers/drivers/DriverController.ts +++ b/src/backend/controllers/drivers/DriverController.ts @@ -390,6 +390,11 @@ export class DriverController extends PuterController { req.actor, subscriptionRequirement, this.config, + { + events: this.clients.event, + surface: `driver.${ifaceName}.${method}`, + req, + }, ); } diff --git a/src/backend/core/http/middleware/subscription.test.ts b/src/backend/core/http/middleware/subscription.test.ts index 90712ccbe..829b76573 100644 --- a/src/backend/core/http/middleware/subscription.test.ts +++ b/src/backend/core/http/middleware/subscription.test.ts @@ -19,6 +19,7 @@ import type { Request, Response } from 'express'; import { describe, expect, it, vi } from 'vitest'; +import { EventClient } from '../../../clients/event/EventClient'; import type { IConfig } from '../../../types'; import { isHttpError } from '../HttpError'; import { requireSubscriptionGate } from './subscription'; @@ -79,3 +80,65 @@ describe('requireSubscriptionGate', () => { expect(got).toBeInstanceOf(Error); }); }); + +describe('subscription.gate hook', () => { + const surface = 'route.post.puterai.anthropic.v1.messages'; + + it('lets a listener on the surface waive the plan for that call', async () => { + const events = new EventClient({} as IConfig); + const seen: unknown[] = []; + events.on(`subscription.gate.${surface}`, (_key, event) => { + seen.push(event); + (event as { allow: boolean }).allow = true; + }); + + const got = await runGate( + requireSubscriptionGate(metering('user_free'), config, true, { + events, + surface, + }), + 'user_free', + ); + + expect(got).toBeUndefined(); + expect(seen).toHaveLength(1); + expect(seen[0]).toMatchObject({ + surface, + requirement: true, + actor: { user: { uuid: 'u-1' } }, + }); + }); + + it('still refuses when nothing answers, or a listener throws', async () => { + const events = new EventClient({} as IConfig); + events.on('subscription.gate.*', () => { + throw new Error('listener bug'); + }); + + const got = await runGate( + requireSubscriptionGate(metering('user_free'), config, true, { + events, + surface, + }), + 'user_free', + ); + + expect(got).toMatchObject({ statusCode: 402 }); + }); + + it('is not asked about a caller the plan already covers', async () => { + const events = new EventClient({} as IConfig); + const listener = vi.fn(); + events.on('subscription.gate.*', listener); + + await runGate( + requireSubscriptionGate(metering('business'), config, true, { + events, + surface, + }), + 'business', + ); + + expect(listener).not.toHaveBeenCalled(); + }); +}); diff --git a/src/backend/core/http/middleware/subscription.ts b/src/backend/core/http/middleware/subscription.ts index e2d4f5a40..df426b70f 100644 --- a/src/backend/core/http/middleware/subscription.ts +++ b/src/backend/core/http/middleware/subscription.ts @@ -20,6 +20,7 @@ import type { RequestHandler } from 'express'; import { assertActorHasSubscription, + type SubscriptionGateHook, type SubscriptionMetering, type SubscriptionRequirement, } from '../../../services/metering/enforcement.js'; @@ -40,6 +41,7 @@ export const requireSubscriptionGate = ( metering: SubscriptionMetering | undefined, config: IConfig, requirement: SubscriptionRequirement, + hook?: Omit, ): RequestHandler => { return (req, _res, next) => { assertActorHasSubscription( @@ -47,6 +49,7 @@ export const requireSubscriptionGate = ( req.actor, requirement, config, + hook && { ...hook, req }, ).then( () => next(), (err) => next(err), diff --git a/src/backend/server.ts b/src/backend/server.ts index c2ab43f7a..4a8470884 100644 --- a/src/backend/server.ts +++ b/src/backend/server.ts @@ -85,7 +85,10 @@ import { createEgressMeteringMiddleware } from './core/http/middleware/egressMet import { createLocalWorkerProxyMiddleware } from './core/http/middleware/localWorkerProxy'; import { createPuterSiteMiddleware } from './core/http/middleware/puterSite'; import { PuterRouter } from './core/http/PuterRouter'; -import { createRouteLifecycleMiddleware } from './core/http/routeLifecycle'; +import { + createRouteLifecycleMiddleware, + routeEventKeyBase, +} from './core/http/routeLifecycle'; import { PREFIX_METADATA_KEY, type RouteDescriptor } from './core/http/types'; import type { AuthService } from './services/auth/AuthService'; import { puterDrivers } from './drivers'; @@ -1147,6 +1150,18 @@ export class PuterServer { this.services.metering, this.#config, subscriptionRequirement!, + { + events: this.clients.event, + surface: routeEventKeyBase( + route.method, + route.path !== undefined + ? PuterServer.#joinPath( + routerPrefix, + route.path, + ) + : undefined, + ), + }, ), ); } diff --git a/src/backend/services/metering/enforcement.ts b/src/backend/services/metering/enforcement.ts index 336b745c0..bf126cdca 100644 --- a/src/backend/services/metering/enforcement.ts +++ b/src/backend/services/metering/enforcement.ts @@ -17,6 +17,8 @@ * along with this program. If not, see . */ +import type { Request } from 'express'; +import type { EventClient } from '../../clients/event/EventClient'; import type { Actor } from '../../core/actor'; import { isSystemActor } from '../../core/actor'; import { HttpError } from '../../core/http/HttpError.js'; @@ -215,6 +217,38 @@ export const actorOnPaidPlan = async ( return subscriptionSatisfies(subscription.id, true); }; +/** + * Where a plan gate sits, so a `subscription.gate.` listener can waive + * it for a call it recognizes. `surface` is the key segment after + * `subscription.gate.`, e.g. `route.post.puterai.openai.v1.chat.completions`. + */ +export interface SubscriptionGateHook { + events: Pick | undefined; + surface: string; + req?: Request; +} + +const waivedByHook = async ( + hook: SubscriptionGateHook | undefined, + actor: Actor, + requirement: SubscriptionRequirement, +): Promise => { + if (!hook?.events) return false; + const event = { + actor, + surface: hook.surface, + requirement, + req: hook.req, + allow: false, + }; + await hook.events.emitAndWait( + `subscription.gate.${hook.surface}`, + event, + {}, + ); + return event.allow === true; +}; + /** * Reject a caller whose plan doesn't cover the surface they're calling. * @@ -224,13 +258,15 @@ export const actorOnPaidPlan = async ( * metering wired — there are no plans to be on. * * The answer comes from the metering service's per-actor subscription cache, so - * this normally costs a map lookup. + * this normally costs a map lookup. A caller the plan doesn't cover is offered + * to `hook` listeners before being refused. */ export const assertActorHasSubscription = async ( metering: SubscriptionMetering, actor: Actor | undefined, requirement: SubscriptionRequirement, config: EnforcementConfig, + hook?: SubscriptionGateHook, ): Promise => { if (subscriptionCheckWaived(metering, actor, requirement, config)) return; @@ -242,6 +278,7 @@ export const assertActorHasSubscription = async ( const subscription = await metering!.getActorSubscription(actor); if (subscriptionSatisfies(subscription.id, requirement)) return; + if (await waivedByHook(hook, actor, requirement)) return; throw new HttpError(402, 'A subscription is required for this action', { legacyCode: 'subscription_required', diff --git a/src/backend/services/share/ShareService.ts b/src/backend/services/share/ShareService.ts index 3d915f9f0..feed605a3 100644 --- a/src/backend/services/share/ShareService.ts +++ b/src/backend/services/share/ShareService.ts @@ -2811,6 +2811,7 @@ export class ShareService extends PuterService { actor, true, this.config, + { events: this.clients.event, surface: 'share.anyone' }, ); // Moving an existing link to another mode is not new reach.