fix(billing): show the limit modal for any classified cause; scope the subscription memo to the org

Greptile P1s: an org_limit classification only opened the modal when the
message also carried generic rate-limit wording, and the singleton's learned
code_usage_subscribed bit survived org switches, pinning helpers of a
subscribed org to the free-tier model. The modal now keys off the classified
cause directly, and the memo clears when the auth org changes.

Generated-By: PostHog Code
Task-Id: 1039ea23-9930-44e2-9888-b05cf8b129ec

Author: adboio

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 = {
@@ -228,6 +246,49 @@ describe("LlmGatewayService.prompt", () => {
     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()

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

@@ -5,6 +5,9 @@ import {
   type PosthogProperties,
 } from "@posthog/shared/posthog-property-headers";
 import { inject, injectable } from "inversify";
+import type { AuthService } from "../auth/auth";
+import { AUTH_SERVICE } from "../auth/auth.module";
+import { AuthServiceEvent } from "../auth/schemas";
 import {
   LLM_GATEWAY_HOST,
   type LlmGatewayAuth,
@@ -47,10 +50,18 @@ export class LlmGatewayService {
     host: LlmGatewayHost,
     @inject(ROOT_LOGGER)
     logger: RootLogger,
+    @inject(AUTH_SERVICE)
+    authService: AuthService,
   ) {
     this.auth = host;
     this.endpoints = host;
     this.log = logger.scope("llm-gateway");
+    let orgId = authService.getState().currentOrgId;
+    authService.on(AuthServiceEvent.StateChanged, (state) => {
+      if (state.currentOrgId === orgId) return;
+      orgId = state.currentOrgId;
+      this.lastKnownCodeUsageSubscribed = null;
+    });
   }
 
   private readonly auth: LlmGatewayAuth;

packages/core/src/sessions/sessionService.ts

@@ -2715,10 +2715,7 @@ export class SessionService {
 
       const limitCause = classifyGatewayLimitError(errorMessage, errorDetails);
 
-      if (
-        limitCause === "model_gate" ||
-        isRateLimitError(errorMessage, errorDetails)
-      ) {
+      if (limitCause !== null || isRateLimitError(errorMessage, errorDetails)) {
         this.d.log.warn("Gateway limit reached, showing usage limit modal", {
           taskRunId: session.taskRunId,
           cause: limitCause,

packages/core/src/sessions/sessionServicePromptRecovery.test.ts

@@ -43,6 +43,7 @@ function createHarness() {
     clearOptimisticItems: vi.fn(),
   };
   const promptMutate = vi.fn();
+  const usageLimitShow = vi.fn();
   const deps = {
     store,
     h: { extractSkillButtonId: () => undefined },
@@ -51,7 +52,7 @@ function createHarness() {
     track: vi.fn(),
     getIsOnline: () => true,
     addDirectoryDialog: { open: false },
-    usageLimit: { show: vi.fn() },
+    usageLimit: { show: usageLimitShow },
     trpc: {
       agent: {
         prompt: { mutate: promptMutate },
@@ -74,7 +75,7 @@ function createHarness() {
     },
     "tryAutoRecoverLocalSession",
   );
-  return { service, sessions, store, promptMutate, recoverSpy };
+  return { service, sessions, store, promptMutate, recoverSpy, usageLimitShow };
 }
 
 type RecoverSpy = ReturnType<typeof createHarness>["recoverSpy"];
@@ -131,3 +132,46 @@ describe("SessionService prompt recovery on fatal session errors", () => {
     },
   );
 });
+
+describe("SessionService gateway billing denials", () => {
+  it.each([
+    {
+      case: "a model-gate 403",
+      message:
+        'Internal error: API Error: 403 {"error":{"message":"Model \'claude-fable-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"}}',
+      expectedShow: { cause: "model_gate" },
+    },
+    {
+      case: "an org-limit 429",
+      message:
+        "Rate limit exceeded: Your team has reached its PostHog Code usage limit for this billing period.",
+      expectedShow: { cause: "org_limit" },
+    },
+    // The classified cause alone must open the modal — the org-limit prose
+    // is not guaranteed to carry generic rate-limit wording.
+    {
+      case: "an org-limit message without rate-limit wording",
+      message:
+        "Your team has reached its PostHog Code usage limit for this billing period.",
+      expectedShow: { cause: "org_limit" },
+    },
+    {
+      case: "an unclassified rate limit",
+      message: "[429] Too many requests",
+      expectedShow: undefined,
+    },
+  ])(
+    "shows the usage-limit modal and stops the prompt for $case",
+    async ({ message, expectedShow }) => {
+      const { service, promptMutate, usageLimitShow, recoverSpy } =
+        createHarness();
+      promptMutate.mockRejectedValue(new Error(message));
+
+      const result = await service.sendPrompt(TASK_ID, "hello");
+
+      expect(result).toEqual({ stopReason: "rate_limited" });
+      expect(usageLimitShow).toHaveBeenCalledWith(expectedShow);
+      expect(recoverSpy).not.toHaveBeenCalled();
+    },
+  );
+});