Skip to content

Commit 9233471

Browse files
authored
feat(billing): honor the gateway's model marks — authed fetches, safe defaults, locked pickers
Third slice of the cutover stack: - Gateway /v1/models fetches send the user's token wherever one exists (Claude adapter, workspace-server preview/models, cloud agent, pi harness) so the gateway can mark plan-restricted models; the models caches key on auth presence so authed and anonymous responses never cross-serve. - pickAllowedModel keeps session defaults off restricted models — a free-tier org lands on the free model instead of a premium default that 403s its first message. Applied on the Claude-family default paths; codex free tier has no allowed models, so its sends surface the PR-1 modal. - Restricted models stay listed and render dimmed with a lock (ModelRadioItem); picking one opens the upgrade gate instead of selecting. Marks ride option _meta (posthog.code/restrictedModel); no codex model/list threading — its live options stay unmarked by design. - The pi harness drops restricted models (no locked rendering there) and falls back to the first served model when the preferred one is gone. Generated-By: PostHog Code Task-Id: 1039ea23-9930-44e2-9888-b05cf8b129ec
1 parent 66ffc18 commit 9233471

17 files changed

Lines changed: 325 additions & 74 deletions

File tree

packages/agent/src/adapters/base-acp-agent.ts

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import type {
1616
WriteTextFileRequest,
1717
WriteTextFileResponse,
1818
} from "@agentclientprotocol/sdk";
19+
import { restrictedModelMeta } from "@posthog/shared";
1920
import {
2021
DEFAULT_GATEWAY_MODEL,
2122
fetchGatewayModels,
@@ -25,6 +26,7 @@ import {
2526
isAnthropicModel,
2627
isCloudflareModel,
2728
isCloudflareModelId,
29+
pickAllowedModel,
2830
} from "../gateway-models";
2931
import { Logger } from "../utils/logger";
3032
/**
@@ -136,22 +138,30 @@ export abstract class BaseAcpAgent implements Agent {
136138
async getModelConfigOptions(
137139
currentModelOverride?: string,
138140
gatewayUrl?: string,
141+
gatewayAuthToken?: string,
139142
): Promise<{
140143
currentModelId: string;
141144
options: SessionConfigSelectOption[];
142145
}> {
146+
// Authenticated so the gateway can mark plan-restricted models —
147+
// anonymous fetches see everything allowed.
143148
this.gatewayModels = await fetchGatewayModels(
144-
gatewayUrl ? { gatewayUrl } : undefined,
149+
gatewayUrl ? { gatewayUrl, authToken: gatewayAuthToken } : undefined,
145150
);
146151

147-
const options = this.gatewayModels
152+
const adapterModels = this.gatewayModels
148153
// Cloudflare models are servable on the Claude adapter too — the gateway translates the
149154
// `@cf/` path onto its Anthropic-Messages surface — so include them alongside Anthropic models.
150-
.filter((model) => isAnthropicModel(model) || isCloudflareModel(model))
155+
.filter((model) => isAnthropicModel(model) || isCloudflareModel(model));
156+
157+
const options = adapterModels
151158
.map((model) => ({
152159
value: model.id,
153160
name: formatGatewayModelName(model),
154161
description: `Context: ${model.context_window.toLocaleString()} tokens`,
162+
// Locked models stay listed so the picker can gate them instead of
163+
// silently dropping them.
164+
...(model.allowed ? {} : { _meta: restrictedModelMeta() }),
155165
}))
156166
// Sort oldest-to-newest so the picker is deterministic and the newest
157167
// model lands at the end of the list, closest to the trigger.
@@ -186,6 +196,11 @@ export abstract class BaseAcpAgent implements Agent {
186196
}
187197
}
188198

199+
// Never auto-select a model the org's plan can't use — it would 403 on
200+
// the first message. An explicit user pick still goes through the
201+
// picker's upgrade gate.
202+
currentModelId = pickAllowedModel(adapterModels, currentModelId);
203+
189204
if (!options.some((opt) => opt.value === currentModelId)) {
190205
options.unshift({
191206
value: currentModelId,

packages/agent/src/adapters/claude/claude-agent.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2097,6 +2097,7 @@ export class ClaudeAcpAgent extends BaseAcpAgent {
20972097
this.getModelConfigOptions(
20982098
settingsManager.getSettings().model || meta?.model || undefined,
20992099
this.options?.gatewayEnv?.anthropicBaseUrl,
2100+
this.options?.gatewayEnv?.anthropicAuthToken,
21002101
),
21012102
...(meta?.taskRunId
21022103
? [

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

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import {
88
isAnthropicModel,
99
isBlockedModelId,
1010
isCloudflareModel,
11+
pickAllowedModel,
1112
} from "./gateway-models";
1213

1314
const model = (id: string, owned_by = ""): GatewayModel => ({
@@ -16,6 +17,7 @@ const model = (id: string, owned_by = ""): GatewayModel => ({
1617
context_window: 128000,
1718
supports_streaming: true,
1819
supports_vision: false,
20+
allowed: true,
1921
});
2022

2123
describe("formatGatewayModelName", () => {
@@ -27,6 +29,7 @@ describe("formatGatewayModelName", () => {
2729
context_window: 200000,
2830
supports_streaming: true,
2931
supports_vision: true,
32+
allowed: true,
3033
}),
3134
).toBe("Claude Opus 4.8");
3235
});
@@ -39,6 +42,7 @@ describe("formatGatewayModelName", () => {
3942
context_window: 200000,
4043
supports_streaming: true,
4144
supports_vision: true,
45+
allowed: true,
4246
}),
4347
).toBe("gpt-5.5");
4448
});
@@ -51,6 +55,7 @@ describe("formatGatewayModelName", () => {
5155
context_window: 200000,
5256
supports_streaming: true,
5357
supports_vision: true,
58+
allowed: true,
5459
}),
5560
).toBe("gpt-5.5");
5661
});
@@ -63,6 +68,7 @@ describe("formatGatewayModelName", () => {
6368
context_window: 128000,
6469
supports_streaming: true,
6570
supports_vision: false,
71+
allowed: true,
6672
}),
6773
).toBe("glm-5.2");
6874
});
@@ -189,3 +195,46 @@ describe("isCloudflareModel", () => {
189195
expect(isAnthropicModel(glm)).toBe(false);
190196
});
191197
});
198+
199+
describe("pickAllowedModel", () => {
200+
const entry = (id: string, allowed: boolean) => ({ id, allowed });
201+
202+
it.each([
203+
[
204+
"keeps an allowed preferred model",
205+
[entry("claude-opus-4-8", true)],
206+
"claude-opus-4-8",
207+
"claude-opus-4-8",
208+
],
209+
[
210+
"keeps a preferred model absent from the list",
211+
[entry("claude-opus-4-8", true)],
212+
"claude-sonnet-5",
213+
"claude-sonnet-5",
214+
],
215+
[
216+
"moves a restricted preferred model to the newest allowed one",
217+
[
218+
entry("claude-opus-4-8", false),
219+
entry("claude-sonnet-4-6", true),
220+
entry("@cf/zai-org/glm-5.2", true),
221+
],
222+
"claude-opus-4-8",
223+
"@cf/zai-org/glm-5.2",
224+
],
225+
[
226+
"keeps the preferred model when everything is restricted",
227+
[entry("claude-opus-4-8", false)],
228+
"claude-opus-4-8",
229+
"claude-opus-4-8",
230+
],
231+
[
232+
"keeps the preferred model when the list is empty",
233+
[],
234+
"claude-opus-4-8",
235+
"claude-opus-4-8",
236+
],
237+
] as const)("%s", (_name, models, preferred, expected) => {
238+
expect(pickAllowedModel(models, preferred)).toBe(expected);
239+
});
240+
});

packages/agent/src/gateway-models.ts

Lines changed: 84 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -4,15 +4,22 @@ export interface GatewayModel {
44
context_window: number;
55
supports_streaming: boolean;
66
supports_vision: boolean;
7+
// Free-tier model gate: authenticated fetches mark models outside the
8+
// caller's plan `allowed: false`. Anonymous fetches and older gateways
9+
// don't mark, so absence means allowed.
10+
allowed: boolean;
11+
restriction_reason?: string | null;
712
}
813

914
interface GatewayModelsResponse {
1015
object: "list";
11-
data: GatewayModel[];
16+
data: Array<Omit<GatewayModel, "allowed"> & { allowed?: boolean }>;
1217
}
1318

1419
export interface FetchGatewayModelsOptions {
1520
gatewayUrl: string;
21+
/** Bearer token; required for accurate free-tier marks. */
22+
authToken?: string;
1623
}
1724

1825
export const DEFAULT_GATEWAY_MODEL = "claude-opus-4-8";
@@ -42,12 +49,19 @@ export function isBlockedModelId(modelId: string): boolean {
4249
return BLOCKED_MODELS.has(modelId.toLowerCase());
4350
}
4451

52+
interface ModelsListEntry {
53+
id?: string;
54+
owned_by?: string;
55+
allowed?: boolean;
56+
restriction_reason?: string | null;
57+
}
58+
4559
type ModelsListResponse =
4660
| {
47-
data?: Array<{ id?: string; owned_by?: string }>;
48-
models?: Array<{ id?: string; owned_by?: string }>;
61+
data?: ModelsListEntry[];
62+
models?: ModelsListEntry[];
4963
}
50-
| Array<{ id?: string; owned_by?: string }>;
64+
| ModelsListEntry[];
5165

5266
const CACHE_TTL = 10 * 60 * 1000; // 10 minutes
5367

@@ -57,11 +71,29 @@ const CACHE_TTL = 10 * 60 * 1000; // 10 minutes
5771
// the callers fall through to `return []`.
5872
const GATEWAY_FETCH_TIMEOUT_MS = 10_000;
5973

60-
let gatewayModelsCache: {
61-
models: GatewayModel[];
74+
// Authed and anonymous responses differ (free-tier marks are authed-only),
75+
// so cache entries are keyed on auth presence.
76+
interface ModelsCache<T> {
77+
models: T[];
6278
expiry: number;
6379
url: string;
64-
} | null = null;
80+
authed: boolean;
81+
}
82+
83+
function readModelsCache<T>(
84+
cache: ModelsCache<T> | null,
85+
url: string,
86+
authed: boolean,
87+
): T[] | null {
88+
if (!cache || cache.url !== url || cache.authed !== authed) return null;
89+
return Date.now() < cache.expiry ? cache.models : null;
90+
}
91+
92+
function authHeaders(authToken?: string): Record<string, string> | undefined {
93+
return authToken ? { Authorization: `Bearer ${authToken}` } : undefined;
94+
}
95+
96+
let gatewayModelsCache: ModelsCache<GatewayModel> | null = null;
6597

6698
export async function fetchGatewayModels(
6799
options?: FetchGatewayModelsOptions,
@@ -71,18 +103,15 @@ export async function fetchGatewayModels(
71103
return [];
72104
}
73105

74-
if (
75-
gatewayModelsCache &&
76-
gatewayModelsCache.url === gatewayUrl &&
77-
Date.now() < gatewayModelsCache.expiry
78-
) {
79-
return gatewayModelsCache.models;
80-
}
106+
const authed = Boolean(options?.authToken);
107+
const cached = readModelsCache(gatewayModelsCache, gatewayUrl, authed);
108+
if (cached) return cached;
81109

82110
const modelsUrl = `${gatewayUrl}/v1/models`;
83111

84112
try {
85113
const response = await fetch(modelsUrl, {
114+
headers: authHeaders(options?.authToken),
86115
signal: AbortSignal.timeout(GATEWAY_FETCH_TIMEOUT_MS),
87116
});
88117

@@ -91,11 +120,14 @@ export async function fetchGatewayModels(
91120
}
92121

93122
const data = (await response.json()) as GatewayModelsResponse;
94-
const models = (data.data ?? []).filter((m) => !isBlockedModelId(m.id));
123+
const models = (data.data ?? [])
124+
.filter((m) => !isBlockedModelId(m.id))
125+
.map((m) => ({ ...m, allowed: m.allowed !== false }));
95126
gatewayModelsCache = {
96127
models,
97128
expiry: Date.now() + CACHE_TTL,
98129
url: gatewayUrl,
130+
authed,
99131
};
100132
return models;
101133
} catch {
@@ -136,13 +168,11 @@ export function isCloudflareModel(model: GatewayModel): boolean {
136168
export interface ModelInfo {
137169
id: string;
138170
owned_by?: string;
171+
allowed: boolean;
172+
restriction_reason?: string | null;
139173
}
140174

141-
let modelsListCache: {
142-
models: ModelInfo[];
143-
expiry: number;
144-
url: string;
145-
} | null = null;
175+
let modelsListCache: ModelsCache<ModelInfo> | null = null;
146176

147177
export async function fetchModelsList(
148178
options?: FetchGatewayModelsOptions,
@@ -152,17 +182,14 @@ export async function fetchModelsList(
152182
return [];
153183
}
154184

155-
if (
156-
modelsListCache &&
157-
modelsListCache.url === gatewayUrl &&
158-
Date.now() < modelsListCache.expiry
159-
) {
160-
return modelsListCache.models;
161-
}
185+
const authed = Boolean(options?.authToken);
186+
const cached = readModelsCache(modelsListCache, gatewayUrl, authed);
187+
if (cached) return cached;
162188

163189
try {
164190
const modelsUrl = `${gatewayUrl}/v1/models`;
165191
const response = await fetch(modelsUrl, {
192+
headers: authHeaders(options?.authToken),
166193
signal: AbortSignal.timeout(GATEWAY_FETCH_TIMEOUT_MS),
167194
});
168195
if (!response.ok) {
@@ -177,19 +204,48 @@ export async function fetchModelsList(
177204
const id = model?.id ? String(model.id) : "";
178205
if (!id) continue;
179206
if (isBlockedModelId(id)) continue;
180-
results.push({ id, owned_by: model?.owned_by });
207+
results.push({
208+
id,
209+
owned_by: model?.owned_by,
210+
allowed: model?.allowed !== false,
211+
restriction_reason: model?.restriction_reason ?? null,
212+
});
181213
}
182214
modelsListCache = {
183215
models: results,
184216
expiry: Date.now() + CACHE_TTL,
185217
url: gatewayUrl,
218+
authed,
186219
};
187220
return results;
188221
} catch {
189222
return [];
190223
}
191224
}
192225

226+
/**
227+
* The model a session should start on: the preferred id when present and
228+
* allowed, else the newest allowed model — a free-tier org must not default
229+
* onto a model that 403s its first message. Falls back to the preferred id
230+
* when the list is empty (fetch failed) or nothing is allowed (all locked —
231+
* the picker gate communicates that state better than a silent swap).
232+
*/
233+
export function pickAllowedModel(
234+
models: ReadonlyArray<Pick<GatewayModel, "id" | "allowed">>,
235+
preferred: string,
236+
): string {
237+
if (models.length === 0) return preferred;
238+
const preferredEntry = models.find((m) => m.id === preferred);
239+
if (!preferredEntry || preferredEntry.allowed) return preferred;
240+
const allowed = models.filter((m) => m.allowed);
241+
if (allowed.length === 0) return preferred;
242+
return allowed.reduce((best, candidate) =>
243+
getClaudeModelRecency(candidate.id) >= getClaudeModelRecency(best.id)
244+
? candidate
245+
: best,
246+
).id;
247+
}
248+
193249
const PROVIDER_NAMES: Record<string, string> = {
194250
anthropic: "Anthropic",
195251
openai: "OpenAI",

packages/agent/src/server/agent-server.ts

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1279,8 +1279,11 @@ export class AgentServer {
12791279
});
12801280

12811281
if (this.config.repoReadyFile && gatewayEnv.anthropicBaseUrl) {
1282+
// Authed so this cache-warm matches the session's own authed fetch
1283+
// (the models cache is keyed on auth presence).
12821284
void fetchGatewayModels({
12831285
gatewayUrl: gatewayEnv.anthropicBaseUrl,
1286+
authToken: gatewayEnv.anthropicAuthToken,
12841287
}).catch(() => {});
12851288
}
12861289

0 commit comments

Comments
 (0)