fix: require explicit pass-through billing access

packages/common/billing/ptbAccess.ts

@@ -0,0 +1,4 @@
+export const PTB_BLOCKED_FEATURE = "ptb_blocked";
+
+export const PTB_BLOCKED_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 { isPtbBlocked } from "../../lib/billing/ptbAccess";
+import { PTB_BLOCKED_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 blockedResult = await isPtbBlocked(request.authParams.organizationId);
+    if (blockedResult.error) {
+      console.error(
+        "Error checking pass-through billing access",
+        blockedResult.error
+      );
+      this.setStatus(503);
+      throw new Error("Unable to verify pass-through billing access");
+    }
+    if (blockedResult.data) {
+      this.setStatus(403);
+      throw new Error(PTB_BLOCKED_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,53 @@
+import { beforeEach, describe, expect, jest, test } from "@jest/globals";
+import { PTB_BLOCKED_FEATURE } from "../../../../../../packages/common/billing/ptbAccess";
+
+jest.mock("../../shared/db/dbExecute", () => ({
+  dbExecute: jest.fn(),
+}));
+
+import { dbExecute } from "../../shared/db/dbExecute";
+import { isPtbBlocked } from "../ptbAccess";
+
+describe("isPtbBlocked", () => {
+  const mockDbExecute = dbExecute as jest.MockedFunction<typeof dbExecute>;
+
+  beforeEach(() => {
+    jest.clearAllMocks();
+  });
+
+  test("returns true when the organization has the block flag", async () => {
+    mockDbExecute.mockResolvedValueOnce({
+      data: [{ id: "flag-1" }],
+      error: null,
+    });
+
+    const result = await isPtbBlocked("org-1");
+
+    expect(result).toEqual({ data: true, error: null });
+    expect(mockDbExecute).toHaveBeenCalledWith(expect.any(String), [
+      "org-1",
+      PTB_BLOCKED_FEATURE,
+    ]);
+  });
+
+  test("returns false when the organization is not blocked", async () => {
+    mockDbExecute.mockResolvedValueOnce({ data: [], error: null });
+
+    await expect(isPtbBlocked("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(isPtbBlocked("org-1")).resolves.toEqual({
+      data: null,
+      error: "database unavailable",
+    });
+  });
+});

valhalla/jawn/src/lib/billing/ptbAccess.ts

@@ -0,0 +1,21 @@
+import { PTB_BLOCKED_FEATURE } from "../../../../../packages/common/billing/ptbAccess";
+import { dbExecute } from "../shared/db/dbExecute";
+import { err, ok, Result } from "../../packages/common/result";
+
+export async function isPtbBlocked(
+  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_BLOCKED_FEATURE],
+  );
+
+  if (result.error) {
+    return err(result.error);
+  }
+
+  return ok(Boolean(result.data?.length));
+}

worker/src/lib/ai-gateway/AttemptExecutor.ts

@@ -24,6 +24,8 @@ 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_BLOCKED_MESSAGE } 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 +95,15 @@ export class AttemptExecutor {
       AttemptError
     >
   > {
+    const featureFlagManager = new FeatureFlagManager(this.env);
+    if (await featureFlagManager.isPtbBlocked(props.orgId)) {
+      return err({
+        type: "request_failed",
+        message: PTB_BLOCKED_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,8 @@ 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_BLOCKED_MESSAGE } from "../../../../packages/common/billing/ptbAccess";
 
 // Constants
 const CACHE_TTL_MS = 60 * 1000; // 1 minutes cache for auto-topoff settings
@@ -109,6 +111,11 @@ export class AutoTopoffManager {
     orgId: string,
     effectiveBalanceCents: number
   ): Promise<boolean> {
+    const featureFlagManager = new FeatureFlagManager(this.env);
+    if (await featureFlagManager.isPtbBlocked(orgId)) {
+      return false;
+    }
+
     // Get settings
     const settingsResult = await this.getAutoTopoffSettings(orgId);
     if (settingsResult.error || !settingsResult.data) {
@@ -186,6 +193,11 @@ export class AutoTopoffManager {
    */
   async initiateTopoff(orgId: string): Promise<Result<string, string>> {
     try {
+      const featureFlagManager = new FeatureFlagManager(this.env);
+      if (await featureFlagManager.isPtbBlocked(orgId)) {
+        return err(PTB_BLOCKED_MESSAGE);
+      }
+
       // Get settings
       const settingsResult = await this.getAutoTopoffSettings(orgId);
       if (settingsResult.error || !settingsResult.data) {

worker/src/lib/managers/FeatureFlagManager.ts

@@ -2,6 +2,7 @@ import { SupabaseClient, createClient } from "@supabase/supabase-js";
 import { Database } from "../../../supabase/database.types";
 import { SecureCacheEnv, getAndStoreInCache } from "../util/cache/secureCache";
 import { Result, ok, err } from "../util/results";
+import { PTB_BLOCKED_FEATURE } from "../../../../packages/common/billing/ptbAccess";
 
 export class FeatureFlagManager {
   private supabaseClient: SupabaseClient<Database>;
@@ -34,6 +35,23 @@ export class FeatureFlagManager {
     return features.data.includes(feature);
   }
 
+  /**
+   * Billing access checks fail closed so a database or cache error cannot
+   * create charges or spend Helicone provider credits.
+   */
+  async isPtbBlocked(orgId: string): Promise<boolean> {
+    const features = await this.getFeatureFlags(orgId);
+    if (features.error || !features.data) {
+      console.error(
+        `Unable to verify pass-through billing access for org ${orgId}:`,
+        features.error
+      );
+      return true;
+    }
+
+    return features.data.includes(PTB_BLOCKED_FEATURE);
+  }
+
   /**
    * Get all feature flags for an organization
    * Cached for 5 minutes to reduce database load

worker/src/lib/managers/StripeManager.ts

@@ -4,6 +4,7 @@ import { Wallet } from "../durable-objects/Wallet";
 import { createClient } from "@supabase/supabase-js";
 import { Database } from "../../../supabase/database.types";
 import { AutoTopoffManager } from "./AutoTopoffManager";
+import { FeatureFlagManager } from "./FeatureFlagManager";
 
 export class StripeManager {
   private webhookSecret: string;
@@ -176,6 +177,14 @@ export class StripeManager {
       return err("Unable to get org id from payment intent");
     }
 
+    const featureFlagManager = new FeatureFlagManager(this.env);
+    if (await featureFlagManager.isPtbBlocked(orgId)) {
+      console.error(
+        `Ignoring successful pass-through billing payment for blocked org ${orgId}: ${paymentIntent.id}`
+      );
+      return ok(undefined);
+    }
+
     const walletId = this.wallet.idFromName(orgId);
     const walletStub = this.wallet.get(walletId);
 

worker/test/ai-gateway/ptb-access.spec.ts

@@ -0,0 +1,79 @@
+import { env, runInDurableObject } from "cloudflare:test";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { PTB_BLOCKED_MESSAGE } from "../../../packages/common/billing/ptbAccess";
+import { AutoTopoffManager } from "../../src/lib/managers/AutoTopoffManager";
+import { StripeManager } from "../../src/lib/managers/StripeManager";
+import "../setup";
+
+vi.mock("../../src/lib/managers/FeatureFlagManager", () => ({
+  FeatureFlagManager: class {
+    async isPtbBlocked() {
+      return true;
+    }
+  },
+}));
+
+const ORG_ID = "test-org-id";
+
+describe("blocked PTB organization 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_BLOCKED_MESSAGE,
+    });
+  });
+
+  it("does not credit a successful payment webhook", async () => {
+    const walletId = env.WALLET.idFromName(ORG_ID);
+    const walletStub = env.WALLET.get(walletId);
+    const creditsBefore = await runInDurableObject(
+      walletStub,
+      async (wallet: any) => (await wallet.getWalletState(ORG_ID)).totalCredits
+    );
+    const consoleSpy = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => undefined);
+    const manager = new StripeManager(
+      "whsec_test",
+      "sk_test_example",
+      env.WALLET,
+      env
+    );
+
+    const result = await manager.handleEvent({
+      id: "evt_blocked_payment",
+      type: "payment_intent.succeeded",
+      data: {
+        object: {
+          id: "pi_blocked_payment",
+          customer: "cus_blocked",
+          currency: "usd",
+          metadata: {
+            productId: "prod_cloud_credits",
+            creditsAmountCents: "10000",
+          },
+        },
+      },
+    } as any);
+
+    const creditsAfter = await runInDurableObject(
+      walletStub,
+      async (wallet: any) => (await wallet.getWalletState(ORG_ID)).totalCredits
+    );
+
+    expect(result).toEqual({ data: undefined, error: null });
+    expect(creditsAfter).toBe(creditsBefore);
+    expect(consoleSpy).toHaveBeenCalledWith(
+      expect.stringContaining(
+        "Ignoring successful pass-through billing payment"
+      )
+    );
+  });
+});

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 suspended", async () => {
+    setSupabaseTestCase({
+      byokEnabled: false,
+      creditsEnabled: true,
+      featureFlags: ["credits", "ptb_blocked"],
+    });
+
+    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,13 @@ 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"] : []);
             return {
               ...chainObj,
               then: (resolve: any) =>
                 resolve({
-                  data: hasCredits ? [{ feature: "credits" }] : [],
+                  data: features.map((feature) => ({ feature })),
                   error: null,
                 }),
             };