fix: require explicit pass-through billing access

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.";

packages/common/index.ts

@@ -2,4 +2,7 @@
 export * from './stripe/feeCalculator';
 
 // Export attribution tracking utilities
-export * from './attribution';
+export * from './attribution';
+
+// Export shared pass-through billing access controls
+export * from './billing/ptbAccess';

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);

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<typeof dbExecute>;
+
+  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",
+    });
+  });
+});

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<Result<boolean, string>> {
+  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));
+}

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",

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<boolean> {
+    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<Result<string, string>> {
     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) {

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,
+    });
+  });
+});

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 });
 

worker/test/providers/base.test-config.ts

@@ -18,6 +18,7 @@ export interface TestCase {
   byokEnabled?: boolean;
   currentCredits?: number;
   creditsEnabled?: boolean;
+  featureFlags?: string[];
   orgId?: string;
 }
 

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,
                 }),
             };