diff --git a/packages/agent/src/gateway-models.test.ts b/packages/agent/src/gateway-models.test.ts index 0c3b7a764a..ac671d151e 100644 --- a/packages/agent/src/gateway-models.test.ts +++ b/packages/agent/src/gateway-models.test.ts @@ -10,6 +10,7 @@ import { isBlockedModelId, isCloudflareModel, pickAllowedModel, + resetGatewayModelCaches, } from "./gateway-models"; const model = (id: string, owned_by = ""): GatewayModel => ({ @@ -240,6 +241,55 @@ describe("gateway models cache", () => { expect(cached[0]?.allowed).toBe(false); }); + // The stale-catalog bug: a model can be marked restricted (or dropped) for a + // token, then become available for that same token after an entitlement + // change. resetGatewayModelCaches() must let the next fetch re-read the + // gateway so the model reappears without a token change (i.e. without a full + // sign-out/sign-in — the workaround affected users had to use). + const codexResponse = (allowed: boolean) => + new Response( + JSON.stringify({ + object: "list", + data: [ + { + id: "gpt-5.5", + owned_by: "openai", + context_window: 400000, + supports_streaming: true, + supports_vision: false, + allowed, + }, + ], + }), + { status: 200, headers: { "Content-Type": "application/json" } }, + ); + + it.each([ + { name: "fetchGatewayModels", fn: fetchGatewayModels }, + { name: "fetchModelsList", fn: fetchModelsList }, + ])( + "$name refreshes the catalog for the same token after a reset", + async ({ fn }) => { + const fetchMock = vi + .spyOn(globalThis, "fetch") + .mockResolvedValueOnce(codexResponse(false)) + .mockResolvedValueOnce(codexResponse(true)); + const gatewayUrl = "https://gateway.reset-refresh-test"; + + const before = await fn({ gatewayUrl, authToken: "tok-a" }); + // Same token would otherwise be served the stale cached list. + const stillCached = await fn({ gatewayUrl, authToken: "tok-a" }); + + resetGatewayModelCaches(); + const afterReset = await fn({ gatewayUrl, authToken: "tok-a" }); + + expect(fetchMock).toHaveBeenCalledTimes(2); + expect(before[0]?.allowed).toBe(false); + expect(stillCached[0]?.allowed).toBe(false); + expect(afterReset[0]?.allowed).toBe(true); + }, + ); + it("corrects stale GLM 5.2 context-window metadata", async () => { vi.spyOn(globalThis, "fetch").mockResolvedValue( new Response( diff --git a/packages/agent/src/gateway-models.ts b/packages/agent/src/gateway-models.ts index 3acb52d534..ee48c03ead 100644 --- a/packages/agent/src/gateway-models.ts +++ b/packages/agent/src/gateway-models.ts @@ -236,6 +236,21 @@ export async function fetchModelsList( } } +/** + * Drop the in-process model caches so the next fetch re-reads the gateway. + * + * The caches are otherwise only busted by TTL expiry or a change of the + * `(url, token)` key, so an entitlement change that keeps the same token — a + * plan upgrade, an org switch that reuses the access token, a model rollout — + * would keep serving a stale catalog until the TTL lapsed. Wiring this to + * auth-state changes (see AgentService) lets a model that becomes available + * surface without waiting out the TTL or forcing a full sign-out/sign-in. + */ +export function resetGatewayModelCaches(): void { + gatewayModelsCache = null; + modelsListCache = null; +} + /** * The model a session should start on: the preferred id when present and * allowed, else the newest allowed model — a free-tier org must not default diff --git a/packages/core/src/auth/auth.ts b/packages/core/src/auth/auth.ts index ec0512ed4e..5bc396ff05 100644 --- a/packages/core/src/auth/auth.ts +++ b/packages/core/src/auth/auth.ts @@ -120,6 +120,16 @@ export class AuthService extends TypedEventEmitter { getState(): AuthState { return { ...this.state }; } + /** + * Subscribe to auth-state changes; returns an unsubscribe. A thin wrapper + * over the `StateChanged` event for consumers that only need "something + * changed" (e.g. dropping identity-scoped caches) without depending on the + * event key. + */ + onAuthStateChanged(listener: () => void): () => void { + this.on(AuthServiceEvent.StateChanged, listener); + return () => this.off(AuthServiceEvent.StateChanged, listener); + } async login(region: CloudRegion): Promise { await this.authenticateWithFlow( () => this.oauthFlow.startFlow(region), diff --git a/packages/workspace-server/src/services/agent/agent.test.ts b/packages/workspace-server/src/services/agent/agent.test.ts index 7d479b71df..757836e9cb 100644 --- a/packages/workspace-server/src/services/agent/agent.test.ts +++ b/packages/workspace-server/src/services/agent/agent.test.ts @@ -164,6 +164,7 @@ function createMockDependencies() { getPluginPath: vi.fn(() => "/mock/plugin"), }, agentAuthAdapter: { + onAuthStateChanged: vi.fn(() => () => {}), ensureGatewayProxy: vi.fn().mockResolvedValue("http://127.0.0.1:9999"), configureProcessEnv: vi.fn().mockResolvedValue(undefined), createPosthogConfig: vi.fn((credentials) => ({ diff --git a/packages/workspace-server/src/services/agent/agent.ts b/packages/workspace-server/src/services/agent/agent.ts index 9af21e716d..1ad1b63021 100644 --- a/packages/workspace-server/src/services/agent/agent.ts +++ b/packages/workspace-server/src/services/agent/agent.ts @@ -38,6 +38,7 @@ import { isCloudflareModel, isOpenAIModel, pickAllowedModel, + resetGatewayModelCaches, } from "@posthog/agent/gateway-models"; import { getLlmGatewayUrl } from "@posthog/agent/posthog-api"; import { @@ -441,6 +442,14 @@ export class AgentService extends TypedEventEmitter { this.onAgentLog = makeOnAgentLog(loggerFactory); powerManager.onResume(() => this.checkIdleDeadlines()); + + // A stale model catalog must not outlive an entitlement change. Dropping + // the cache on every auth-state change means the next picker/config read + // re-reads the gateway, so a model that becomes available (or an org + // switch) surfaces without a full sign-out/sign-in. The refetch is cheap + // and still TTL-cached, and this service is a process-lifetime singleton so + // the subscription never needs tearing down. + this.agentAuthAdapter.onAuthStateChanged(() => resetGatewayModelCaches()); } private getClaudeCliPath(): string { diff --git a/packages/workspace-server/src/services/agent/auth-adapter.ts b/packages/workspace-server/src/services/agent/auth-adapter.ts index 0cf747e951..0ae1f09486 100644 --- a/packages/workspace-server/src/services/agent/auth-adapter.ts +++ b/packages/workspace-server/src/services/agent/auth-adapter.ts @@ -151,6 +151,11 @@ export class AgentAuthAdapter { } } + /** Subscribe to auth-state changes; returns an unsubscribe. */ + onAuthStateChanged(listener: () => void): () => void { + return this.authService.onAuthStateChanged(listener); + } + async configureProcessEnv({ credentials, proxyUrl, diff --git a/packages/workspace-server/src/services/agent/ports.ts b/packages/workspace-server/src/services/agent/ports.ts index 13a0d4d36f..a7dfdba917 100644 --- a/packages/workspace-server/src/services/agent/ports.ts +++ b/packages/workspace-server/src/services/agent/ports.ts @@ -69,4 +69,6 @@ export interface AgentAuth { input: string | Request, init?: RequestInit, ): Promise; + /** Subscribe to auth-state changes; returns an unsubscribe. */ + onAuthStateChanged(listener: () => void): () => void; }