diff --git a/packages/common/billing/ptbAccess.ts b/packages/common/billing/ptbAccess.ts new file mode 100644 index 0000000000..1e7505d022 --- /dev/null +++ b/packages/common/billing/ptbAccess.ts @@ -0,0 +1,4 @@ +export const PTB_ENABLED_FEATURE = "ptb_enabled"; + +export const PTB_DISABLED_MESSAGE = + "Pass-through billing is disabled for this organization. Contact support@helicone.ai for help."; diff --git a/packages/common/index.ts b/packages/common/index.ts index 34a48de04e..0864b1f3ed 100644 --- a/packages/common/index.ts +++ b/packages/common/index.ts @@ -2,4 +2,7 @@ export * from './stripe/feeCalculator'; // Export attribution tracking utilities -export * from './attribution'; \ No newline at end of file +export * from './attribution'; + +// Export shared pass-through billing access controls +export * from './billing/ptbAccess'; diff --git a/valhalla/jawn/src/controllers/public/stripeController.ts b/valhalla/jawn/src/controllers/public/stripeController.ts index 98310a0f81..10d5ab4d2a 100644 --- a/valhalla/jawn/src/controllers/public/stripeController.ts +++ b/valhalla/jawn/src/controllers/public/stripeController.ts @@ -16,6 +16,8 @@ import type { JawnAuthenticatedRequest } from "../../types/request"; import { isError } from "../../packages/common/result"; import express from "express"; import Stripe from "stripe"; +import { hasPtbAccess } from "../../lib/billing/ptbAccess"; +import { PTB_DISABLED_MESSAGE } from "../../../../../packages/common/billing/ptbAccess"; export interface UpgradeToProRequest { addons?: { @@ -198,6 +200,20 @@ export class StripeController extends Controller { @Request() request: JawnAuthenticatedRequest, @Body() body: CreateCloudGatewayCheckoutSessionRequest ): Promise<{ checkoutUrl: string }> { + const accessResult = await hasPtbAccess(request.authParams.organizationId); + if (accessResult.error) { + console.error( + "Error checking pass-through billing access", + accessResult.error + ); + this.setStatus(503); + throw new Error("Unable to verify pass-through billing access"); + } + if (!accessResult.data) { + this.setStatus(403); + throw new Error(PTB_DISABLED_MESSAGE); + } + const stripeManager = new StripeManager(request.authParams); if (body.amount < 5) { this.setStatus(400); diff --git a/valhalla/jawn/src/lib/billing/__tests__/ptbAccess.test.ts b/valhalla/jawn/src/lib/billing/__tests__/ptbAccess.test.ts new file mode 100644 index 0000000000..c4e53e1bb1 --- /dev/null +++ b/valhalla/jawn/src/lib/billing/__tests__/ptbAccess.test.ts @@ -0,0 +1,54 @@ +import { beforeEach, describe, expect, jest, test } from "@jest/globals"; +import { PTB_ENABLED_FEATURE } from "../../../../../../packages/common/billing/ptbAccess"; + +jest.mock("../../shared/db/dbExecute", () => ({ + dbExecute: jest.fn(), +})); + +import { dbExecute } from "../../shared/db/dbExecute"; +import { hasPtbAccess } from "../ptbAccess"; + +describe("hasPtbAccess", () => { + const mockDbExecute = dbExecute as jest.MockedFunction; + + beforeEach(() => { + jest.clearAllMocks(); + }); + + test("returns true when the organization has the allow flag", async () => { + mockDbExecute.mockResolvedValueOnce({ + data: [{ id: "flag-1" }], + error: null, + }); + + await expect(hasPtbAccess("org-1")).resolves.toEqual({ + data: true, + error: null, + }); + expect(mockDbExecute).toHaveBeenCalledWith(expect.any(String), [ + "org-1", + PTB_ENABLED_FEATURE, + ]); + }); + + test("returns false when the organization has no PTB access flag", async () => { + mockDbExecute.mockResolvedValueOnce({ data: [], error: null }); + + await expect(hasPtbAccess("org-1")).resolves.toEqual({ + data: false, + error: null, + }); + }); + + test("returns the database error so checkout can fail closed", async () => { + mockDbExecute.mockResolvedValueOnce({ + data: null, + error: "database unavailable", + }); + + await expect(hasPtbAccess("org-1")).resolves.toEqual({ + data: null, + error: "database unavailable", + }); + }); +}); diff --git a/valhalla/jawn/src/lib/billing/ptbAccess.ts b/valhalla/jawn/src/lib/billing/ptbAccess.ts new file mode 100644 index 0000000000..fb2b9165c9 --- /dev/null +++ b/valhalla/jawn/src/lib/billing/ptbAccess.ts @@ -0,0 +1,21 @@ +import { PTB_ENABLED_FEATURE } from "../../../../../packages/common/billing/ptbAccess"; +import { dbExecute } from "../shared/db/dbExecute"; +import { err, ok, Result } from "../../packages/common/result"; + +export async function hasPtbAccess( + orgId: string, +): Promise> { + const result = await dbExecute<{ id: string }>( + `SELECT id + FROM feature_flags + WHERE org_id = $1 AND feature = $2 + LIMIT 1`, + [orgId, PTB_ENABLED_FEATURE], + ); + + if (result.error) { + return err(result.error); + } + + return ok(Boolean(result.data?.length)); +} diff --git a/worker/src/lib/ai-gateway/AttemptExecutor.ts b/worker/src/lib/ai-gateway/AttemptExecutor.ts index 0d640e7749..573eab59ec 100644 --- a/worker/src/lib/ai-gateway/AttemptExecutor.ts +++ b/worker/src/lib/ai-gateway/AttemptExecutor.ts @@ -24,6 +24,11 @@ import { GatewayMetrics } from "./GatewayMetrics"; import { Attempt, AttemptError, EscrowInfo, PendingEscrow } from "./types"; import { toChatCompletions } from "@helicone-package/llm-mapper/transform/providers/responses/request/toChatCompletions"; import { WalletKVSync } from "./WalletKVSync"; +import { FeatureFlagManager } from "../managers/FeatureFlagManager"; +import { + PTB_DISABLED_MESSAGE, + PTB_ENABLED_FEATURE, +} from "../../../../packages/common/billing/ptbAccess"; // Minimum balance (in cents) to allow optimistic execution without waiting for escrow const ALLOWABLE_BALANCE_TO_SKIP_CHECK = 450; // $4.50 in cents @@ -93,6 +98,17 @@ export class AttemptExecutor { AttemptError > > { + const featureFlagManager = new FeatureFlagManager(this.env); + if ( + !(await featureFlagManager.hasFeature(props.orgId, PTB_ENABLED_FEATURE)) + ) { + return err({ + type: "request_failed", + message: PTB_DISABLED_MESSAGE, + statusCode: 403, + }); + } + const walletSpanId = props.traceContext?.sampled ? this.tracer.startSpan( "ai_gateway.ptb.credit_validation.reserve_escrow", diff --git a/worker/src/lib/managers/AutoTopoffManager.ts b/worker/src/lib/managers/AutoTopoffManager.ts index a9341b36dc..63224d9ee8 100644 --- a/worker/src/lib/managers/AutoTopoffManager.ts +++ b/worker/src/lib/managers/AutoTopoffManager.ts @@ -4,6 +4,11 @@ import { Result, err, ok } from "../util/results"; import Stripe from "stripe"; import { getAndStoreInCache, removeFromCache } from "../util/cache/secureCache"; import { randomUUID } from "crypto"; +import { FeatureFlagManager } from "./FeatureFlagManager"; +import { + PTB_DISABLED_MESSAGE, + PTB_ENABLED_FEATURE, +} from "../../../../packages/common/billing/ptbAccess"; // Constants const CACHE_TTL_MS = 60 * 1000; // 1 minutes cache for auto-topoff settings @@ -109,6 +114,11 @@ export class AutoTopoffManager { orgId: string, effectiveBalanceCents: number ): Promise { + const featureFlagManager = new FeatureFlagManager(this.env); + if (!(await featureFlagManager.hasFeature(orgId, PTB_ENABLED_FEATURE))) { + return false; + } + // Get settings const settingsResult = await this.getAutoTopoffSettings(orgId); if (settingsResult.error || !settingsResult.data) { @@ -186,6 +196,11 @@ export class AutoTopoffManager { */ async initiateTopoff(orgId: string): Promise> { try { + const featureFlagManager = new FeatureFlagManager(this.env); + if (!(await featureFlagManager.hasFeature(orgId, PTB_ENABLED_FEATURE))) { + return err(PTB_DISABLED_MESSAGE); + } + // Get settings const settingsResult = await this.getAutoTopoffSettings(orgId); if (settingsResult.error || !settingsResult.data) { diff --git a/worker/test/ai-gateway/ptb-access.spec.ts b/worker/test/ai-gateway/ptb-access.spec.ts new file mode 100644 index 0000000000..e2f02ca212 --- /dev/null +++ b/worker/test/ai-gateway/ptb-access.spec.ts @@ -0,0 +1,31 @@ +import { env } from "cloudflare:test"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { PTB_DISABLED_MESSAGE } from "../../../packages/common/billing/ptbAccess"; +import { AutoTopoffManager } from "../../src/lib/managers/AutoTopoffManager"; +import "../setup"; + +vi.mock("../../src/lib/managers/FeatureFlagManager", () => ({ + FeatureFlagManager: class { + async hasFeature() { + return false; + } + }, +})); + +const ORG_ID = "test-org-id"; + +describe("organization without PTB access", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("does not schedule or initiate automatic top-offs", async () => { + const manager = new AutoTopoffManager(env); + + await expect(manager.shouldTriggerTopoff(ORG_ID, 0)).resolves.toBe(false); + await expect(manager.initiateTopoff(ORG_ID)).resolves.toEqual({ + data: null, + error: PTB_DISABLED_MESSAGE, + }); + }); +}); diff --git a/worker/test/ai-gateway/ptb-validation.spec.ts b/worker/test/ai-gateway/ptb-validation.spec.ts index ca9e48adcf..91b2535781 100644 --- a/worker/test/ai-gateway/ptb-validation.spec.ts +++ b/worker/test/ai-gateway/ptb-validation.spec.ts @@ -98,6 +98,28 @@ describe("PTB request validation", () => { }); }); + it("blocks PTB before calling a provider when the org is not allowlisted", async () => { + setSupabaseTestCase({ + byokEnabled: false, + creditsEnabled: true, + featureFlags: ["credits"], + }); + + const { response } = await runGatewayTest({ + model: "gpt-4o-mini/openai", + request: { + messages: [{ role: "user", content: "Hello" }], + }, + expected: { + providers: [], + finalStatus: 403, + responseContains: "Pass-through billing is disabled", + }, + }); + + expect(response.status).toBe(403); + }); + it("returns 400 when PTB payload contains web_search_options", async () => { setSupabaseTestCase({ byokEnabled: false, creditsEnabled: true }); diff --git a/worker/test/providers/base.test-config.ts b/worker/test/providers/base.test-config.ts index cd119bc522..cc5eb0f231 100644 --- a/worker/test/providers/base.test-config.ts +++ b/worker/test/providers/base.test-config.ts @@ -18,6 +18,7 @@ export interface TestCase { byokEnabled?: boolean; currentCredits?: number; creditsEnabled?: boolean; + featureFlags?: string[]; orgId?: string; } diff --git a/worker/test/setup.ts b/worker/test/setup.ts index ff949da294..5c4865510d 100644 --- a/worker/test/setup.ts +++ b/worker/test/setup.ts @@ -83,11 +83,14 @@ vi.mock("@supabase/supabase-js", () => ({ chainObj.eq = vi.fn((field: string, value: any) => { if (field === "org_id") { const hasCredits = currentTestCase?.creditsEnabled === true; + const features = + currentTestCase?.featureFlags ?? + (hasCredits ? ["credits", "ptb_enabled"] : ["ptb_enabled"]); return { ...chainObj, then: (resolve: any) => resolve({ - data: hasCredits ? [{ feature: "credits" }] : [], + data: features.map((feature) => ({ feature })), error: null, }), };