feat(billing): classify gateway billing denials and rework the usage-limit modal for usage-based billing (#3485)

Author: adboio

packages/core/src/billing/usageLimitContent.test.ts

@@ -0,0 +1,57 @@
+import { describe, expect, it } from "vitest";
+import { usageLimitContent } from "./usageLimitContent";
+
+describe("usageLimitContent", () => {
+  it("offers a payment method for the model gate", () => {
+    const content = usageLimitContent({
+      cause: "model_gate",
+      resetLabel: null,
+      subscribed: false,
+    });
+    expect(content.title).toBe("Unlock premium models");
+    expect(content.description).toContain("This model isn't");
+    expect(content.actionLabel).toBe("Add payment method");
+  });
+
+  it.each([
+    // Confirmed-free org: allocation used up, the fix is adding a card.
+    [false, "Free usage used up", "Add payment method"],
+    // Subscribed org: the fix is raising the spend limit.
+    [true, "Organization usage limit reached", "Manage billing"],
+    // Unknown subscription state must not read as free.
+    [undefined, "Organization usage limit reached", "Manage billing"],
+  ] as const)(
+    "org_limit with subscribed=%s -> %s / %s",
+    (subscribed, title, actionLabel) => {
+      const content = usageLimitContent({
+        cause: "org_limit",
+        resetLabel: null,
+        subscribed,
+      });
+      expect(content.title).toBe(title);
+      expect(content.actionLabel).toBe(actionLabel);
+    },
+  );
+
+  it("includes the reset hint in the free-tier copy when available", () => {
+    const content = usageLimitContent({
+      cause: "org_limit",
+      resetLabel: "Resets in 3h",
+      subscribed: false,
+    });
+    expect(content.title).toBe("Free usage used up");
+    expect(content.description).toContain("Resets in 3h");
+  });
+
+  it("renders generic copy without a billing CTA when the cause is unknown", () => {
+    const content = usageLimitContent({
+      cause: null,
+      resetLabel: "Resets in 2h",
+      subscribed: true,
+    });
+    expect(content.title).toBe("Usage limit reached");
+    expect(content.description).toContain("Resets in 2h");
+    expect(content.actionLabel).toBeNull();
+    expect(content.dismissLabel).toBe("Got it");
+  });
+});

packages/core/src/billing/usageLimitContent.ts

@@ -0,0 +1,57 @@
+import type { GatewayLimitCause } from "@posthog/shared";
+
+export interface UsageLimitContent {
+  title: string;
+  description: string;
+  actionLabel: string | null;
+  dismissLabel: string;
+}
+
+export function usageLimitContent(args: {
+  cause: GatewayLimitCause | null;
+  resetLabel: string | null;
+  subscribed: boolean | undefined;
+}): UsageLimitContent {
+  const { cause, resetLabel, subscribed } = args;
+
+  if (cause === "model_gate") {
+    return {
+      title: "Unlock premium models",
+      description:
+        "This model isn't included in the free tier. Add a payment method to your organization to unlock all models — you only pay for what you use. You can keep working now by switching to an included model.",
+      actionLabel: "Add payment method",
+      dismissLabel: "Not now",
+    };
+  }
+
+  if (cause === "org_limit") {
+    if (subscribed === false) {
+      return {
+        title: "Free usage used up",
+        description: `Your organization has used up its included PostHog Code usage.${
+          resetLabel ? ` ${resetLabel}.` : ""
+        } Add a payment method to keep going — you only pay for what you use.`,
+        actionLabel: "Add payment method",
+        dismissLabel: "Not now",
+      };
+    }
+    return {
+      title: "Organization usage limit reached",
+      description:
+        "Your organization has reached its PostHog Code spend limit for this billing period. Raise or remove the limit in your PostHog billing settings to keep going.",
+      actionLabel: "Manage billing",
+      dismissLabel: "Got it",
+    };
+  }
+
+  // Not a billing denial (e.g. an upstream provider's own rate limit) —
+  // don't send the user to billing for something billing can't fix.
+  return {
+    title: "Usage limit reached",
+    description: `PostHog Code hit a usage limit.${
+      resetLabel ? ` ${resetLabel}.` : ""
+    } Please try again shortly.`,
+    actionLabel: null,
+    dismissLabel: "Got it",
+  };
+}

packages/core/src/llm-gateway/llm-gateway.test.ts

@@ -1,4 +1,5 @@
 import { beforeEach, describe, expect, it, vi } from "vitest";
+import type { AuthService } from "../auth/auth";
 import type {
   LlmGatewayAuth,
   LlmGatewayEndpoints,
@@ -43,8 +44,25 @@ function createService(
   };
   const logger = { ...log, scope: () => log };
 
-  const service = new LlmGatewayService(host, logger);
-  return { service, auth, endpoints, log };
+  const orgListeners: Array<(state: { currentOrgId: string | null }) => void> =
+    [];
+  const authService = {
+    getState: () => ({ currentOrgId: "org-1" }),
+    on: (
+      _event: string,
+      listener: (state: { currentOrgId: string | null }) => void,
+    ) => {
+      orgListeners.push(listener);
+    },
+  } as unknown as AuthService;
+  const emitAuthState = (currentOrgId: string | null) => {
+    for (const listener of orgListeners) {
+      listener({ currentOrgId });
+    }
+  };
+
+  const service = new LlmGatewayService(host, logger, authService);
+  return { service, auth, endpoints, log, emitAuthState };
 }
 
 const SUCCESS_BODY = {
@@ -157,6 +175,137 @@ describe("LlmGatewayService.prompt", () => {
     });
   });
 
+  it("surfaces a FastAPI bare-string detail as the error message", async () => {
+    const fetchMock = vi.fn().mockResolvedValue(
+      createJsonResponse(
+        {
+          detail: "OAuth application not authorized for product 'posthog_code'",
+        },
+        403,
+      ),
+    );
+    const { service } = createService(fetchMock);
+
+    await expect(
+      service.prompt([{ role: "user", content: "hi" }]),
+    ).rejects.toMatchObject({
+      name: "LlmGatewayError",
+      message: "OAuth application not authorized for product 'posthog_code'",
+      type: "unknown_error",
+      statusCode: 403,
+    });
+  });
+
+  // The free-tier model gate's 403 body, as the gateway serves it.
+  const MODEL_GATE_BODY = {
+    error: {
+      message:
+        "Model 'claude-haiku-4-5' needs a paid PostHog plan. Models available on the free tier: @cf/zai-org/glm-5.2. Add a payment method to your organization to unlock all models. (rate_limit)",
+      type: "permission_error",
+      code: "model_gate",
+    },
+  };
+
+  it("retries once on the free-tier model when the model gate 403s", async () => {
+    const fetchMock = vi
+      .fn()
+      .mockResolvedValueOnce(createJsonResponse(MODEL_GATE_BODY, 403))
+      .mockResolvedValueOnce(createJsonResponse(SUCCESS_BODY));
+    const { service } = createService(fetchMock);
+
+    const result = await service.prompt([{ role: "user", content: "hi" }], {
+      model: "claude-haiku-4-5",
+    });
+
+    expect(result.content).toBe("hello world");
+    expect(fetchMock).toHaveBeenCalledTimes(2);
+    const firstBody = JSON.parse(fetchMock.mock.calls[0][1].body as string);
+    const retryBody = JSON.parse(fetchMock.mock.calls[1][1].body as string);
+    expect(firstBody.model).toBe("claude-haiku-4-5");
+    expect(retryBody.model).toBe("@cf/zai-org/glm-5.2");
+  });
+
+  it("routes straight to the free-tier model once the org is known unsubscribed", async () => {
+    const fetchMock = vi
+      .fn()
+      .mockResolvedValueOnce(createJsonResponse(MODEL_GATE_BODY, 403))
+      .mockImplementation(async () => createJsonResponse(SUCCESS_BODY));
+    const { service } = createService(fetchMock);
+
+    // First call learns "unsubscribed" from the gate's 403.
+    await service.prompt([{ role: "user", content: "hi" }], {
+      model: "claude-haiku-4-5",
+    });
+    // Second call must not burn a round trip on the gate.
+    await service.prompt([{ role: "user", content: "hi" }], {
+      model: "claude-haiku-4-5",
+    });
+
+    expect(fetchMock).toHaveBeenCalledTimes(3);
+    const thirdBody = JSON.parse(fetchMock.mock.calls[2][1].body as string);
+    expect(thirdBody.model).toBe("@cf/zai-org/glm-5.2");
+  });
+
+  it("forgets the learned subscription state when the organization changes", async () => {
+    const fetchMock = vi
+      .fn()
+      .mockResolvedValueOnce(createJsonResponse(MODEL_GATE_BODY, 403))
+      .mockImplementation(async () => createJsonResponse(SUCCESS_BODY));
+    const { service, emitAuthState } = createService(fetchMock);
+
+    await service.prompt([{ role: "user", content: "hi" }], {
+      model: "claude-haiku-4-5",
+    });
+    emitAuthState("org-2");
+    await service.prompt([{ role: "user", content: "hi" }], {
+      model: "claude-haiku-4-5",
+    });
+
+    const bodyAfterSwitch = JSON.parse(
+      fetchMock.mock.calls[2][1].body as string,
+    );
+    expect(bodyAfterSwitch.model).toBe("claude-haiku-4-5");
+  });
+
+  it("keeps the learned subscription state across same-org auth changes", async () => {
+    const fetchMock = vi
+      .fn()
+      .mockResolvedValueOnce(createJsonResponse(MODEL_GATE_BODY, 403))
+      .mockImplementation(async () => createJsonResponse(SUCCESS_BODY));
+    const { service, emitAuthState } = createService(fetchMock);
+
+    await service.prompt([{ role: "user", content: "hi" }], {
+      model: "claude-haiku-4-5",
+    });
+    emitAuthState("org-1");
+    await service.prompt([{ role: "user", content: "hi" }], {
+      model: "claude-haiku-4-5",
+    });
+
+    expect(fetchMock).toHaveBeenCalledTimes(3);
+    const bodyAfterRefresh = JSON.parse(
+      fetchMock.mock.calls[2][1].body as string,
+    );
+    expect(bodyAfterRefresh.model).toBe("@cf/zai-org/glm-5.2");
+  });
+
+  it("does not retry non-gate 403s on the free-tier model", async () => {
+    const fetchMock = vi
+      .fn()
+      .mockResolvedValue(
+        createJsonResponse(
+          { detail: "Product 'posthog_code' requires OAuth authentication" },
+          403,
+        ),
+      );
+    const { service } = createService(fetchMock);
+
+    await expect(
+      service.prompt([{ role: "user", content: "hi" }]),
+    ).rejects.toMatchObject({ statusCode: 403 });
+    expect(fetchMock).toHaveBeenCalledTimes(1);
+  });
+
   it("throws a timeout LlmGatewayError when the request aborts via the internal timeout", async () => {
     const fetchMock = vi.fn((_url: string, init?: RequestInit) => {
       return new Promise<Response>((_resolve, reject) => {
@@ -214,6 +363,40 @@ describe("LlmGatewayService.fetchUsage", () => {
       statusCode: 503,
     });
   });
+
+  it("parses the usage-based billing fields when present", async () => {
+    const fetchMock = vi.fn().mockResolvedValue(
+      createJsonResponse({
+        ...USAGE_BODY,
+        ai_credits: { exhausted: true },
+        code_usage_subscribed: true,
+      }),
+    );
+    const { service } = createService(fetchMock);
+
+    const usage = await service.fetchUsage();
+
+    expect(usage.ai_credits?.exhausted).toBe(true);
+    expect(usage.code_usage_subscribed).toBe(true);
+  });
+
+  it("feeds code_usage_subscribed into helper model routing", async () => {
+    const fetchMock = vi
+      .fn()
+      .mockResolvedValueOnce(
+        createJsonResponse({ ...USAGE_BODY, code_usage_subscribed: false }),
+      )
+      .mockResolvedValue(createJsonResponse(SUCCESS_BODY));
+    const { service } = createService(fetchMock);
+
+    await service.fetchUsage();
+    await service.prompt([{ role: "user", content: "hi" }], {
+      model: "claude-haiku-4-5",
+    });
+
+    const promptBody = JSON.parse(fetchMock.mock.calls[1][1].body as string);
+    expect(promptBody.model).toBe("@cf/zai-org/glm-5.2");
+  });
 });
 
 describe("LlmGatewayService.invalidatePlanCache", () => {