fix(billing): key the gateway models cache on the auth token

Restriction marks are identity-scoped, but the cache keyed on token
presence — an org switch in the same process served the previous org's
marks for up to the 10-minute TTL, mislocking pickers and letting
pickAllowedModel swap session defaults off stale data. Entries are now
keyed on the exact token; a rotation costs one refetch.

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

Author: adboio

packages/agent/src/gateway-models.test.ts

@@ -173,6 +173,60 @@ describe("gateway model fetch timeout", () => {
   );
 });
 
+describe("gateway models cache", () => {
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  const modelsResponse = (allowed: boolean) =>
+    new Response(
+      JSON.stringify({
+        object: "list",
+        data: [
+          {
+            id: "claude-opus-4-8",
+            owned_by: "anthropic",
+            context_window: 200000,
+            supports_streaming: true,
+            supports_vision: true,
+            allowed,
+          },
+        ],
+      }),
+      { status: 200, headers: { "Content-Type": "application/json" } },
+    );
+
+  // Restriction marks are org-scoped: an org switch swaps the token in the
+  // same process, and the old org's marks must not be served to the new one.
+  it("does not serve one token's marks to another token", async () => {
+    const fetchMock = vi
+      .spyOn(globalThis, "fetch")
+      .mockResolvedValueOnce(modelsResponse(false))
+      .mockResolvedValueOnce(modelsResponse(true));
+    const gatewayUrl = "https://gateway.token-key-test";
+
+    const first = await fetchGatewayModels({ gatewayUrl, authToken: "tok-a" });
+    const second = await fetchGatewayModels({ gatewayUrl, authToken: "tok-b" });
+
+    expect(fetchMock).toHaveBeenCalledTimes(2);
+    expect(first[0]?.allowed).toBe(false);
+    expect(second[0]?.allowed).toBe(true);
+  });
+
+  it("serves the cached list to the same token without refetching", async () => {
+    const fetchMock = vi
+      .spyOn(globalThis, "fetch")
+      .mockResolvedValue(modelsResponse(false));
+    const gatewayUrl = "https://gateway.token-cache-hit-test";
+
+    await fetchGatewayModels({ gatewayUrl, authToken: "tok-a" });
+    const cached = await fetchGatewayModels({ gatewayUrl, authToken: "tok-a" });
+
+    expect(fetchMock).toHaveBeenCalledTimes(1);
+    expect(cached[0]?.allowed).toBe(false);
+  });
+});
+
 describe("isCloudflareModel", () => {
   it.each([
     { id: "@cf/zai-org/glm-5.2", owned_by: "cloudflare", expected: true },

packages/agent/src/gateway-models.ts

@@ -71,21 +71,23 @@ const CACHE_TTL = 10 * 60 * 1000; // 10 minutes
 // the callers fall through to `return []`.
 const GATEWAY_FETCH_TIMEOUT_MS = 10_000;
 
-// Authed and anonymous responses differ (free-tier marks are authed-only),
-// so cache entries are keyed on auth presence.
+// Restriction marks are identity-scoped (free-tier marks are authed-only and
+// differ per org), so cache entries are keyed on the exact token — an org
+// switch in the same process must never be served the old org's marks. A
+// token rotation just costs one refetch.
 interface ModelsCache<T> {
   models: T[];
   expiry: number;
   url: string;
-  authed: boolean;
+  token: string | null;
 }
 
 function readModelsCache<T>(
   cache: ModelsCache<T> | null,
   url: string,
-  authed: boolean,
+  token: string | null,
 ): T[] | null {
-  if (!cache || cache.url !== url || cache.authed !== authed) return null;
+  if (!cache || cache.url !== url || cache.token !== token) return null;
   return Date.now() < cache.expiry ? cache.models : null;
 }
 
@@ -103,8 +105,8 @@ export async function fetchGatewayModels(
     return [];
   }
 
-  const authed = Boolean(options?.authToken);
-  const cached = readModelsCache(gatewayModelsCache, gatewayUrl, authed);
+  const token = options?.authToken ?? null;
+  const cached = readModelsCache(gatewayModelsCache, gatewayUrl, token);
   if (cached) return cached;
 
   const modelsUrl = `${gatewayUrl}/v1/models`;
@@ -127,7 +129,7 @@ export async function fetchGatewayModels(
       models,
       expiry: Date.now() + CACHE_TTL,
       url: gatewayUrl,
-      authed,
+      token,
     };
     return models;
   } catch {
@@ -182,8 +184,8 @@ export async function fetchModelsList(
     return [];
   }
 
-  const authed = Boolean(options?.authToken);
-  const cached = readModelsCache(modelsListCache, gatewayUrl, authed);
+  const token = options?.authToken ?? null;
+  const cached = readModelsCache(modelsListCache, gatewayUrl, token);
   if (cached) return cached;
 
   try {
@@ -215,7 +217,7 @@ export async function fetchModelsList(
       models: results,
       expiry: Date.now() + CACHE_TTL,
       url: gatewayUrl,
-      authed,
+      token,
     };
     return results;
   } catch {