diff --git a/packages/agent/src/adapters/base-acp-agent.ts b/packages/agent/src/adapters/base-acp-agent.ts index 874c68de77..249b3655e9 100644 --- a/packages/agent/src/adapters/base-acp-agent.ts +++ b/packages/agent/src/adapters/base-acp-agent.ts @@ -16,6 +16,7 @@ import type { WriteTextFileRequest, WriteTextFileResponse, } from "@agentclientprotocol/sdk"; +import { restrictedModelMeta } from "@posthog/shared"; import { DEFAULT_GATEWAY_MODEL, fetchGatewayModels, @@ -25,6 +26,7 @@ import { isAnthropicModel, isCloudflareModel, isCloudflareModelId, + pickAllowedModel, } from "../gateway-models"; import { Logger } from "../utils/logger"; /** @@ -136,22 +138,30 @@ export abstract class BaseAcpAgent implements Agent { async getModelConfigOptions( currentModelOverride?: string, gatewayUrl?: string, + gatewayAuthToken?: string, ): Promise<{ currentModelId: string; options: SessionConfigSelectOption[]; }> { + // Authenticated so the gateway can mark plan-restricted models — + // anonymous fetches see everything allowed. this.gatewayModels = await fetchGatewayModels( - gatewayUrl ? { gatewayUrl } : undefined, + gatewayUrl ? { gatewayUrl, authToken: gatewayAuthToken } : undefined, ); - const options = this.gatewayModels + const adapterModels = this.gatewayModels // Cloudflare models are servable on the Claude adapter too — the gateway translates the // `@cf/` path onto its Anthropic-Messages surface — so include them alongside Anthropic models. - .filter((model) => isAnthropicModel(model) || isCloudflareModel(model)) + .filter((model) => isAnthropicModel(model) || isCloudflareModel(model)); + + const options = adapterModels .map((model) => ({ value: model.id, name: formatGatewayModelName(model), description: `Context: ${model.context_window.toLocaleString()} tokens`, + // Locked models stay listed so the picker can gate them instead of + // silently dropping them. + ...(model.allowed ? {} : { _meta: restrictedModelMeta() }), })) // Sort oldest-to-newest so the picker is deterministic and the newest // model lands at the end of the list, closest to the trigger. @@ -186,6 +196,11 @@ export abstract class BaseAcpAgent implements Agent { } } + // Never auto-select a model the org's plan can't use — it would 403 on + // the first message. An explicit user pick still goes through the + // picker's upgrade gate. + currentModelId = pickAllowedModel(adapterModels, currentModelId); + if (!options.some((opt) => opt.value === currentModelId)) { options.unshift({ value: currentModelId, diff --git a/packages/agent/src/adapters/claude/claude-agent.ts b/packages/agent/src/adapters/claude/claude-agent.ts index 6e3b1bae86..402d8dbdd4 100644 --- a/packages/agent/src/adapters/claude/claude-agent.ts +++ b/packages/agent/src/adapters/claude/claude-agent.ts @@ -2097,6 +2097,7 @@ export class ClaudeAcpAgent extends BaseAcpAgent { this.getModelConfigOptions( settingsManager.getSettings().model || meta?.model || undefined, this.options?.gatewayEnv?.anthropicBaseUrl, + this.options?.gatewayEnv?.anthropicAuthToken, ), ...(meta?.taskRunId ? [ diff --git a/packages/agent/src/adapters/codex-app-server/codex-app-server-agent.test.ts b/packages/agent/src/adapters/codex-app-server/codex-app-server-agent.test.ts index e91cf2b80c..3e567b09c7 100644 --- a/packages/agent/src/adapters/codex-app-server/codex-app-server-agent.test.ts +++ b/packages/agent/src/adapters/codex-app-server/codex-app-server-agent.test.ts @@ -4,6 +4,7 @@ import type { NewSessionRequest, PromptRequest, } from "@agentclientprotocol/sdk"; +import { RequestError } from "@agentclientprotocol/sdk"; import { describe, expect, it } from "vitest"; import type { AppServerClientHandlers, @@ -1596,6 +1597,45 @@ describe("CodexAppServerAgent", () => { expect((await done).stopReason).toBe("refusal"); }); + it("rejects the prompt when the fatal error is a gateway billing denial", async () => { + // A silent refusal would hide the free-tier gate: the host only shows the + // upgrade modal when the prompt rejects with the gateway's message. + const stub = makeStubRpc({ "thread/start": { thread: { id: "t" } } }); + const { client } = makeFakeClient(); + const agent = new CodexAppServerAgent(client, { + processOptions: { binaryPath: "/x/codex" }, + rpcFactory: stub.factory, + }); + + await agent.newSession({ cwd: "/r" } as unknown as NewSessionRequest); + const done = agent.prompt({ + sessionId: "t", + prompt: [{ type: "text", text: "go" }], + } as unknown as PromptRequest); + stub.emit("error", { + willRetry: false, + error: { + message: + "unexpected status 403 Forbidden: Model 'gpt-5.5' needs a paid PostHog plan. Models available on the free tier: @cf/zai-org/glm-5.2. (rate_limit)", + }, + }); + + const err = await done.then( + () => { + throw new Error("prompt resolved instead of rejecting"); + }, + (e: unknown) => e, + ); + // Must be a RequestError with the gateway text in its message — a plain + // Error crosses the ACP boundary as a bare "Internal error", which the + // host classifies as fatal and answers with a kill/respawn loop. + expect(err).toBeInstanceOf(RequestError); + expect((err as RequestError).message).toContain("Internal error: "); + expect((err as RequestError).message).toContain( + "needs a paid PostHog plan", + ); + }); + it("ends the turn without turn/start when no prompt block is usable", async () => { const stub = makeStubRpc({ "thread/start": { thread: { id: "t" } } }); const { client } = makeFakeClient(); diff --git a/packages/agent/src/adapters/codex-app-server/codex-app-server-agent.ts b/packages/agent/src/adapters/codex-app-server/codex-app-server-agent.ts index 7ed57f3339..4c4035a53b 100644 --- a/packages/agent/src/adapters/codex-app-server/codex-app-server-agent.ts +++ b/packages/agent/src/adapters/codex-app-server/codex-app-server-agent.ts @@ -19,7 +19,12 @@ import type { SetSessionConfigOptionResponse, StopReason, } from "@agentclientprotocol/sdk"; -import { mcpToolKey, posthogToolMeta } from "@posthog/shared"; +import { RequestError } from "@agentclientprotocol/sdk"; +import { + classifyGatewayLimitError, + mcpToolKey, + posthogToolMeta, +} from "@posthog/shared"; import { type NativeGoalState, POSTHOG_NOTIFICATIONS, @@ -1245,11 +1250,27 @@ export class CodexAppServerAgent extends BaseAcpAgent { if (method === APP_SERVER_NOTIFICATIONS.ERROR) { // A non-retried fatal error: resolve the turn so prompt() returns rather than hangs. - const willRetry = (params as { willRetry?: boolean })?.willRetry; + const { willRetry, error } = (params ?? {}) as { + willRetry?: boolean; + error?: { message?: string }; + }; if (willRetry === false) { this.logger.warn("codex app-server fatal error notification", { params, }); + const message = error?.message ?? ""; + // A gateway billing denial rejects the prompt so the host classifies + // it and shows the upgrade gate. It must be a RequestError: a plain + // Error serializes to a bare "Internal error" at the ACP boundary, + // which the host reads as fatal and answers with a respawn loop. + if (classifyGatewayLimitError(message) !== null) { + if (this.compactionActive) { + this.compactionActive = false; + this.emitCompactionBoundary(); + } + this.turns.fail(RequestError.internalError(undefined, message)); + return; + } void this.finalizeTurn("refusal"); } } diff --git a/packages/agent/src/gateway-models.test.ts b/packages/agent/src/gateway-models.test.ts index c804652b67..bd5d34dab0 100644 --- a/packages/agent/src/gateway-models.test.ts +++ b/packages/agent/src/gateway-models.test.ts @@ -8,6 +8,7 @@ import { isAnthropicModel, isBlockedModelId, isCloudflareModel, + pickAllowedModel, } from "./gateway-models"; const model = (id: string, owned_by = ""): GatewayModel => ({ @@ -16,6 +17,7 @@ const model = (id: string, owned_by = ""): GatewayModel => ({ context_window: 128000, supports_streaming: true, supports_vision: false, + allowed: true, }); describe("formatGatewayModelName", () => { @@ -27,6 +29,7 @@ describe("formatGatewayModelName", () => { context_window: 200000, supports_streaming: true, supports_vision: true, + allowed: true, }), ).toBe("Claude Opus 4.8"); }); @@ -39,6 +42,7 @@ describe("formatGatewayModelName", () => { context_window: 200000, supports_streaming: true, supports_vision: true, + allowed: true, }), ).toBe("gpt-5.5"); }); @@ -51,6 +55,7 @@ describe("formatGatewayModelName", () => { context_window: 200000, supports_streaming: true, supports_vision: true, + allowed: true, }), ).toBe("gpt-5.5"); }); @@ -63,6 +68,7 @@ describe("formatGatewayModelName", () => { context_window: 128000, supports_streaming: true, supports_vision: false, + allowed: true, }), ).toBe("glm-5.2"); }); @@ -167,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 }, @@ -189,3 +249,46 @@ describe("isCloudflareModel", () => { expect(isAnthropicModel(glm)).toBe(false); }); }); + +describe("pickAllowedModel", () => { + const entry = (id: string, allowed: boolean) => ({ id, allowed }); + + it.each([ + [ + "keeps an allowed preferred model", + [entry("claude-opus-4-8", true)], + "claude-opus-4-8", + "claude-opus-4-8", + ], + [ + "keeps a preferred model absent from the list", + [entry("claude-opus-4-8", true)], + "claude-sonnet-5", + "claude-sonnet-5", + ], + [ + "moves a restricted preferred model to the newest allowed one", + [ + entry("claude-opus-4-8", false), + entry("claude-sonnet-4-6", true), + entry("@cf/zai-org/glm-5.2", true), + ], + "claude-opus-4-8", + "@cf/zai-org/glm-5.2", + ], + [ + "keeps the preferred model when everything is restricted", + [entry("claude-opus-4-8", false)], + "claude-opus-4-8", + "claude-opus-4-8", + ], + [ + "keeps the preferred model when the list is empty", + [], + "claude-opus-4-8", + "claude-opus-4-8", + ], + ] as const)("%s", (_name, models, preferred, expected) => { + expect(pickAllowedModel(models, preferred)).toBe(expected); + }); +}); diff --git a/packages/agent/src/gateway-models.ts b/packages/agent/src/gateway-models.ts index 81db1d8b0a..62a379665d 100644 --- a/packages/agent/src/gateway-models.ts +++ b/packages/agent/src/gateway-models.ts @@ -4,15 +4,22 @@ export interface GatewayModel { context_window: number; supports_streaming: boolean; supports_vision: boolean; + // Free-tier model gate: authenticated fetches mark models outside the + // caller's plan `allowed: false`. Anonymous fetches and older gateways + // don't mark, so absence means allowed. + allowed: boolean; + restriction_reason?: string | null; } interface GatewayModelsResponse { object: "list"; - data: GatewayModel[]; + data: Array & { allowed?: boolean }>; } export interface FetchGatewayModelsOptions { gatewayUrl: string; + /** Bearer token; required for accurate free-tier marks. */ + authToken?: string; } export const DEFAULT_GATEWAY_MODEL = "claude-opus-4-8"; @@ -42,12 +49,19 @@ export function isBlockedModelId(modelId: string): boolean { return BLOCKED_MODELS.has(modelId.toLowerCase()); } +interface ModelsListEntry { + id?: string; + owned_by?: string; + allowed?: boolean; + restriction_reason?: string | null; +} + type ModelsListResponse = | { - data?: Array<{ id?: string; owned_by?: string }>; - models?: Array<{ id?: string; owned_by?: string }>; + data?: ModelsListEntry[]; + models?: ModelsListEntry[]; } - | Array<{ id?: string; owned_by?: string }>; + | ModelsListEntry[]; const CACHE_TTL = 10 * 60 * 1000; // 10 minutes @@ -57,11 +71,31 @@ const CACHE_TTL = 10 * 60 * 1000; // 10 minutes // the callers fall through to `return []`. const GATEWAY_FETCH_TIMEOUT_MS = 10_000; -let gatewayModelsCache: { - models: GatewayModel[]; +// 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 { + models: T[]; expiry: number; url: string; -} | null = null; + token: string | null; +} + +function readModelsCache( + cache: ModelsCache | null, + url: string, + token: string | null, +): T[] | null { + if (!cache || cache.url !== url || cache.token !== token) return null; + return Date.now() < cache.expiry ? cache.models : null; +} + +function authHeaders(authToken?: string): Record | undefined { + return authToken ? { Authorization: `Bearer ${authToken}` } : undefined; +} + +let gatewayModelsCache: ModelsCache | null = null; export async function fetchGatewayModels( options?: FetchGatewayModelsOptions, @@ -71,18 +105,15 @@ export async function fetchGatewayModels( return []; } - if ( - gatewayModelsCache && - gatewayModelsCache.url === gatewayUrl && - Date.now() < gatewayModelsCache.expiry - ) { - return gatewayModelsCache.models; - } + const token = options?.authToken ?? null; + const cached = readModelsCache(gatewayModelsCache, gatewayUrl, token); + if (cached) return cached; const modelsUrl = `${gatewayUrl}/v1/models`; try { const response = await fetch(modelsUrl, { + headers: authHeaders(options?.authToken), signal: AbortSignal.timeout(GATEWAY_FETCH_TIMEOUT_MS), }); @@ -91,11 +122,14 @@ export async function fetchGatewayModels( } const data = (await response.json()) as GatewayModelsResponse; - const models = (data.data ?? []).filter((m) => !isBlockedModelId(m.id)); + const models = (data.data ?? []) + .filter((m) => !isBlockedModelId(m.id)) + .map((m) => ({ ...m, allowed: m.allowed !== false })); gatewayModelsCache = { models, expiry: Date.now() + CACHE_TTL, url: gatewayUrl, + token, }; return models; } catch { @@ -136,13 +170,11 @@ export function isCloudflareModel(model: GatewayModel): boolean { export interface ModelInfo { id: string; owned_by?: string; + allowed: boolean; + restriction_reason?: string | null; } -let modelsListCache: { - models: ModelInfo[]; - expiry: number; - url: string; -} | null = null; +let modelsListCache: ModelsCache | null = null; export async function fetchModelsList( options?: FetchGatewayModelsOptions, @@ -152,17 +184,14 @@ export async function fetchModelsList( return []; } - if ( - modelsListCache && - modelsListCache.url === gatewayUrl && - Date.now() < modelsListCache.expiry - ) { - return modelsListCache.models; - } + const token = options?.authToken ?? null; + const cached = readModelsCache(modelsListCache, gatewayUrl, token); + if (cached) return cached; try { const modelsUrl = `${gatewayUrl}/v1/models`; const response = await fetch(modelsUrl, { + headers: authHeaders(options?.authToken), signal: AbortSignal.timeout(GATEWAY_FETCH_TIMEOUT_MS), }); if (!response.ok) { @@ -177,12 +206,18 @@ export async function fetchModelsList( const id = model?.id ? String(model.id) : ""; if (!id) continue; if (isBlockedModelId(id)) continue; - results.push({ id, owned_by: model?.owned_by }); + results.push({ + id, + owned_by: model?.owned_by, + allowed: model?.allowed !== false, + restriction_reason: model?.restriction_reason ?? null, + }); } modelsListCache = { models: results, expiry: Date.now() + CACHE_TTL, url: gatewayUrl, + token, }; return results; } catch { @@ -190,6 +225,29 @@ export async function fetchModelsList( } } +/** + * 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 + * onto a model that 403s its first message. Falls back to the preferred id + * when the list is empty (fetch failed) or nothing is allowed (all locked — + * the picker gate communicates that state better than a silent swap). + */ +export function pickAllowedModel( + models: ReadonlyArray>, + preferred: string, +): string { + if (models.length === 0) return preferred; + const preferredEntry = models.find((m) => m.id === preferred); + if (!preferredEntry || preferredEntry.allowed) return preferred; + const allowed = models.filter((m) => m.allowed); + if (allowed.length === 0) return preferred; + return allowed.reduce((best, candidate) => + getClaudeModelRecency(candidate.id) >= getClaudeModelRecency(best.id) + ? candidate + : best, + ).id; +} + const PROVIDER_NAMES: Record = { anthropic: "Anthropic", openai: "OpenAI", diff --git a/packages/agent/src/server/agent-server.ts b/packages/agent/src/server/agent-server.ts index 01e69aa342..7de8e1eb78 100644 --- a/packages/agent/src/server/agent-server.ts +++ b/packages/agent/src/server/agent-server.ts @@ -1279,8 +1279,11 @@ export class AgentServer { }); if (this.config.repoReadyFile && gatewayEnv.anthropicBaseUrl) { + // Authed so this cache-warm matches the session's own authed fetch + // (the models cache is keyed on auth presence). void fetchGatewayModels({ gatewayUrl: gatewayEnv.anthropicBaseUrl, + authToken: gatewayEnv.anthropicAuthToken, }).catch(() => {}); } diff --git a/packages/harness/src/extensions/posthog-provider/models.ts b/packages/harness/src/extensions/posthog-provider/models.ts index 3ff27d7096..3c96c3bfd5 100644 --- a/packages/harness/src/extensions/posthog-provider/models.ts +++ b/packages/harness/src/extensions/posthog-provider/models.ts @@ -12,6 +12,9 @@ export interface GatewayModel { display_name?: string; context_window?: number; supports_vision?: boolean; + // Free-tier model gate: authenticated fetches mark models outside the + // caller's plan. Absence (anonymous fetch or older gateway) means allowed. + allowed?: boolean; } type ModelFamily = "anthropic" | "openai" | "cloudflare"; @@ -181,12 +184,14 @@ export function fallbackModelConfigs( async function fetchGatewayModels( region: CloudRegion, + apiKey?: string, ): Promise { if (process.env.PI_OFFLINE || process.env.HARNESS_STATIC_MODELS) { return []; } try { const response = await fetch(`${getLlmGatewayUrl(region)}/v1/models`, { + headers: apiKey ? { Authorization: `Bearer ${apiKey}` } : undefined, signal: AbortSignal.timeout(MODELS_FETCH_TIMEOUT_MS), }); if (!response.ok) { @@ -201,12 +206,17 @@ async function fetchGatewayModels( export async function resolveModelConfigs( region: CloudRegion, + apiKey?: string, ): Promise { - const live = await fetchGatewayModels(region); + const live = await fetchGatewayModels(region, apiKey); if (live.length === 0) { return fallbackModelConfigs(region); } - return live - .filter((model) => Boolean(model.id)) - .map((model) => toModelConfig(model, region)); + const withIds = live.filter((model) => Boolean(model.id)); + // pi has no locked-model rendering, so restricted models are dropped. The + // free tier always includes a servable model; guard against empty anyway. + const usable = withIds.filter((model) => model.allowed !== false); + return (usable.length > 0 ? usable : withIds).map((model) => + toModelConfig(model, region), + ); } diff --git a/packages/harness/src/extensions/posthog-provider/provider.ts b/packages/harness/src/extensions/posthog-provider/provider.ts index 856b8834bc..1620511118 100644 --- a/packages/harness/src/extensions/posthog-provider/provider.ts +++ b/packages/harness/src/extensions/posthog-provider/provider.ts @@ -70,6 +70,6 @@ export async function resolvePosthogProvider( options: PosthogProviderOptions = {}, ): Promise { const region = resolveRegion(options.region); - const models = await resolveModelConfigs(region); + const models = await resolveModelConfigs(region, options.apiKey); return buildPosthogProvider(models, options); } diff --git a/packages/harness/src/session.ts b/packages/harness/src/session.ts index c2d6cc0700..447859bd18 100644 --- a/packages/harness/src/session.ts +++ b/packages/harness/src/session.ts @@ -33,15 +33,18 @@ export async function createHarnessSession( const authStorage = AuthStorage.create(); const modelRegistry = ModelRegistry.create(authStorage); - modelRegistry.registerProvider( - POSTHOG_PROVIDER_NAME, - await resolvePosthogProvider(options), - ); + const provider = await resolvePosthogProvider(options); + modelRegistry.registerProvider(POSTHOG_PROVIDER_NAME, provider); - const model = modelRegistry.find( - POSTHOG_PROVIDER_NAME, - options.model ?? DEFAULT_MODEL, - ); + // Restricted models are dropped for free-tier orgs, so the preferred model + // can be absent — fall back to the first model the provider serves rather + // than letting pi error. + const fallbackModelId = provider.models?.[0]?.id; + const model = + modelRegistry.find(POSTHOG_PROVIDER_NAME, options.model ?? DEFAULT_MODEL) ?? + (fallbackModelId + ? modelRegistry.find(POSTHOG_PROVIDER_NAME, fallbackModelId) + : undefined); const resourceLoader = new DefaultResourceLoader({ cwd, diff --git a/packages/shared/src/analytics-events.ts b/packages/shared/src/analytics-events.ts index b1f313cebc..32756ce73f 100644 --- a/packages/shared/src/analytics-events.ts +++ b/packages/shared/src/analytics-events.ts @@ -974,7 +974,8 @@ export type UpgradePromptShownSurface = | "usage_limit_modal" | "upgrade_dialog" | "titlebar_card" - | "billing_announcement"; + | "billing_announcement" + | "model_picker"; export type UpgradePromptClickedSurface = | "usage_limit_modal" @@ -983,7 +984,8 @@ export type UpgradePromptClickedSurface = | "titlebar_card" | "plan_page_card" | "upgrade_dialog" - | "billing_announcement"; + | "billing_announcement" + | "model_picker"; export type UpgradePromptCause = "model_gate" | "org_limit"; diff --git a/packages/shared/src/index.ts b/packages/shared/src/index.ts index 203c46790c..c42ecffda1 100644 --- a/packages/shared/src/index.ts +++ b/packages/shared/src/index.ts @@ -133,7 +133,12 @@ export { mentionsToPlainText, splitMentionSegments, } from "./mentions"; -export { defaultEligibleModel } from "./models"; +export { + defaultEligibleModel, + isRestrictedModelOption, + RESTRICTED_MODEL_META_KEY, + restrictedModelMeta, +} from "./models"; export { getOauthClientIdFromRegion, OAUTH_SCOPE_VERSION, diff --git a/packages/shared/src/models.ts b/packages/shared/src/models.ts index f2dd69e951..65cb3f6ad9 100644 --- a/packages/shared/src/models.ts +++ b/packages/shared/src/models.ts @@ -9,3 +9,20 @@ export function defaultEligibleModel( const family = modelId.toLowerCase().split("/").pop() ?? ""; return family.startsWith("claude-fable") ? undefined : modelId; } + +/** + * ACP SessionConfigSelectOption `_meta` key for the free-tier model gate: + * adapters mark models the caller's org can't use so pickers render them + * locked behind an upgrade gate instead of omitting them. + */ +export const RESTRICTED_MODEL_META_KEY = "posthog.code/restrictedModel"; + +export function restrictedModelMeta(): Record { + return { [RESTRICTED_MODEL_META_KEY]: true }; +} + +export function isRestrictedModelOption( + meta: Record | null | undefined, +): boolean { + return meta?.[RESTRICTED_MODEL_META_KEY] === true; +} diff --git a/packages/ui/src/features/billing/modelGate.ts b/packages/ui/src/features/billing/modelGate.ts new file mode 100644 index 0000000000..38b22844e0 --- /dev/null +++ b/packages/ui/src/features/billing/modelGate.ts @@ -0,0 +1,30 @@ +import type { SessionConfigSelectOption } from "@agentclientprotocol/sdk"; +import { isRestrictedModelOption } from "@posthog/shared"; +import { ANALYTICS_EVENTS } from "@posthog/shared/analytics-events"; +import { track } from "../../shell/analytics"; +import { useUsageLimitStore } from "./usageLimitStore"; + +/** Whether a picker option is outside the org's plan (free-tier model gate). */ +export function isRestrictedModel( + option: Pick, +): boolean { + return isRestrictedModelOption(option._meta ?? undefined); +} + +/** + * Intercepts a pick of a plan-restricted model: opens the upgrade gate and + * returns true so the caller skips the selection. + */ +export function gateRestrictedModelPick( + options: SessionConfigSelectOption[], + value: string, +): boolean { + const picked = options.find((opt) => opt.value === value); + if (!picked || !isRestrictedModel(picked)) return false; + track(ANALYTICS_EVENTS.UPGRADE_PROMPT_SHOWN, { + surface: "model_picker", + cause: "model_gate", + }); + useUsageLimitStore.getState().show({ cause: "model_gate" }); + return true; +} diff --git a/packages/ui/src/features/sessions/components/ModelRadioItem.tsx b/packages/ui/src/features/sessions/components/ModelRadioItem.tsx new file mode 100644 index 0000000000..81a9fbd7db --- /dev/null +++ b/packages/ui/src/features/sessions/components/ModelRadioItem.tsx @@ -0,0 +1,31 @@ +import { Lock } from "@phosphor-icons/react"; +import { DropdownMenuRadioItem } from "@posthog/quill"; +import { isRestrictedModel } from "@posthog/ui/features/billing/modelGate"; + +/** + * Model picker entry. Plan-restricted models render dimmed with a lock; + * picking one is intercepted by the selector's change handler, which opens + * the upgrade gate instead of selecting. + */ +export function ModelRadioItem({ + model, +}: { + model: { + value: string; + name: string; + _meta?: Record | null; + }; +}) { + const restricted = isRestrictedModel(model); + return ( + + {model.name} + {restricted && ( + + )} + + ); +} diff --git a/packages/ui/src/features/sessions/components/ModelSelector.tsx b/packages/ui/src/features/sessions/components/ModelSelector.tsx index 6cce5064cd..2cc5173f8c 100644 --- a/packages/ui/src/features/sessions/components/ModelSelector.tsx +++ b/packages/ui/src/features/sessions/components/ModelSelector.tsx @@ -8,13 +8,14 @@ import { DropdownMenu, DropdownMenuContent, DropdownMenuRadioGroup, - DropdownMenuRadioItem, DropdownMenuSeparator, DropdownMenuTrigger, MenuLabel, } from "@posthog/quill"; import { type Adapter, GLM_MODEL_FLAG } from "@posthog/shared"; +import { gateRestrictedModelPick } from "@posthog/ui/features/billing/modelGate"; import { useFeatureFlag } from "@posthog/ui/features/feature-flags/useFeatureFlag"; +import { ModelRadioItem } from "@posthog/ui/features/sessions/components/ModelRadioItem"; import { stripGlmModelOption } from "@posthog/ui/features/sessions/modelOptionFilters"; import { flattenSelectOptions, @@ -63,6 +64,9 @@ export function ModelSelector({ if (!selectOption || options.length === 0) return null; const handleChange = (value: string) => { + // A plan-restricted model opens the upgrade gate instead of becoming + // the selection. + if (gateRestrictedModelPick(options, value)) return; onModelChange?.(value); if (!taskId) return; @@ -110,9 +114,7 @@ export function ModelSelector({ {index > 0 && } {group.name} {group.options.map((model) => ( - - {model.name} - + ))} ))} @@ -123,9 +125,7 @@ export function ModelSelector({ onValueChange={handleChange} > {options.map((model) => ( - - {model.name} - + ))} )} diff --git a/packages/ui/src/features/sessions/components/UnifiedModelSelector.tsx b/packages/ui/src/features/sessions/components/UnifiedModelSelector.tsx index 3b71f14d4b..06b3956f41 100644 --- a/packages/ui/src/features/sessions/components/UnifiedModelSelector.tsx +++ b/packages/ui/src/features/sessions/components/UnifiedModelSelector.tsx @@ -15,11 +15,12 @@ import { DropdownMenuContent, DropdownMenuItem, DropdownMenuRadioGroup, - DropdownMenuRadioItem, DropdownMenuSeparator, DropdownMenuTrigger, MenuLabel, } from "@posthog/quill"; +import { gateRestrictedModelPick } from "@posthog/ui/features/billing/modelGate"; +import { ModelRadioItem } from "@posthog/ui/features/sessions/components/ModelRadioItem"; import { flattenSelectOptions } from "@posthog/ui/features/sessions/sessionStore"; import { useRetainedConfigOption } from "@posthog/ui/features/sessions/useRetainedConfigOption"; import type { AgentAdapter } from "@posthog/ui/features/settings/settingsStore"; @@ -144,6 +145,13 @@ export function UnifiedModelSelector({ { + // A plan-restricted model opens the upgrade gate instead of + // becoming the selection. + if (gateRestrictedModelPick(options, value)) { + pendingValueRef.current = null; + setOpen(false); + return; + } pendingValueRef.current = value; setOpen(false); }} @@ -154,19 +162,12 @@ export function UnifiedModelSelector({ {index > 0 && } {group.name} {group.options.map((model) => ( - - {model.name} - + ))} )) : options.map((model) => ( - - {model.name} - + ))} )} diff --git a/packages/workspace-server/src/services/agent/agent.ts b/packages/workspace-server/src/services/agent/agent.ts index dfb440287c..9af21e716d 100644 --- a/packages/workspace-server/src/services/agent/agent.ts +++ b/packages/workspace-server/src/services/agent/agent.ts @@ -37,6 +37,7 @@ import { isAnthropicModel, isCloudflareModel, isOpenAIModel, + pickAllowedModel, } from "@posthog/agent/gateway-models"; import { getLlmGatewayUrl } from "@posthog/agent/posthog-api"; import { @@ -70,6 +71,7 @@ import { type ExecutionMode, isAuthError, resolveCloudInitialPermissionMode, + restrictedModelMeta, serializeError, TypedEventEmitter, } from "@posthog/shared"; @@ -2241,7 +2243,10 @@ For git operations while detached: async getGatewayModels(apiHost: string) { const gatewayUrl = getLlmGatewayUrl(apiHost); - const models = await fetchGatewayModels({ gatewayUrl }); + const models = await fetchGatewayModels({ + gatewayUrl, + authToken: (await this.agentAuthAdapter.gatewayAuthToken()) ?? undefined, + }); const mapped = models.map((model) => ({ modelId: model.id, @@ -2270,7 +2275,12 @@ For git operations while detached: adapter: Adapter = "claude", ): Promise { const gatewayUrl = getLlmGatewayUrl(apiHost); - const gatewayModels = await fetchGatewayModels({ gatewayUrl }); + // Authenticated so the gateway can mark plan-restricted models; falls + // back to an anonymous fetch (everything allowed) without auth. + const gatewayModels = await fetchGatewayModels({ + gatewayUrl, + authToken: (await this.agentAuthAdapter.gatewayAuthToken()) ?? undefined, + }); // The Claude adapter can also drive Cloudflare `@cf/` models the gateway serves over its // Anthropic-Messages surface, so the preview/default-model path must offer them too — otherwise an @@ -2281,13 +2291,15 @@ For git operations while detached: : (model: GatewayModel) => isAnthropicModel(model) || isCloudflareModel(model); - const modelOptions = gatewayModels - .filter((model) => modelFilter(model)) - .map((model) => ({ - value: model.id, - name: formatGatewayModelName(model), - description: `Context: ${model.context_window.toLocaleString()} tokens`, - })); + const adapterModels = gatewayModels.filter((model) => modelFilter(model)); + const modelOptions = adapterModels.map((model) => ({ + value: model.id, + name: formatGatewayModelName(model), + description: `Context: ${model.context_window.toLocaleString()} tokens`, + // Locked models stay listed so the picker can gate them instead of + // silently dropping them. + ...(model.allowed ? {} : { _meta: restrictedModelMeta() }), + })); // The gateway returns models in an arbitrary order. Sort Claude models // oldest-to-newest so the picker is deterministic and the newest model @@ -2306,9 +2318,12 @@ For git operations while detached: "") : DEFAULT_GATEWAY_MODEL; - const resolvedModelId = modelOptions.some((o) => o.value === defaultModel) + const preferredModelId = modelOptions.some((o) => o.value === defaultModel) ? defaultModel : (modelOptions[0]?.value ?? defaultModel); + // Never preselect a model the org's plan can't use — it would 403 on the + // first message. + const resolvedModelId = pickAllowedModel(adapterModels, preferredModelId); if (!modelOptions.some((o) => o.value === resolvedModelId)) { modelOptions.unshift({ diff --git a/packages/workspace-server/src/services/agent/auth-adapter.ts b/packages/workspace-server/src/services/agent/auth-adapter.ts index bd139edeca..0cf747e951 100644 --- a/packages/workspace-server/src/services/agent/auth-adapter.ts +++ b/packages/workspace-server/src/services/agent/auth-adapter.ts @@ -138,6 +138,19 @@ export class AgentAuthAdapter { return this.authProxy.start(getLlmGatewayUrl(apiHost)); } + /** + * Bearer token for direct gateway REST calls (the models fetch), so the + * gateway can mark plan-restricted models. Null when auth isn't available — + * callers fall back to an anonymous fetch. + */ + async gatewayAuthToken(): Promise { + try { + return await this.getValidToken(); + } catch { + return null; + } + } + async configureProcessEnv({ credentials, proxyUrl,