Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 18 additions & 3 deletions packages/agent/src/adapters/base-acp-agent.ts
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ import type {
WriteTextFileRequest,
WriteTextFileResponse,
} from "@agentclientprotocol/sdk";
import { restrictedModelMeta } from "@posthog/shared";
import {
DEFAULT_GATEWAY_MODEL,
fetchGatewayModels,
Expand All @@ -25,6 +26,7 @@ import {
isAnthropicModel,
isCloudflareModel,
isCloudflareModelId,
pickAllowedModel,
} from "../gateway-models";
import { Logger } from "../utils/logger";
/**
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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,
Expand Down
1 change: 1 addition & 0 deletions packages/agent/src/adapters/claude/claude-agent.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
? [
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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");
}
}
Expand Down
103 changes: 103 additions & 0 deletions packages/agent/src/gateway-models.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import {
isAnthropicModel,
isBlockedModelId,
isCloudflareModel,
pickAllowedModel,
} from "./gateway-models";

const model = (id: string, owned_by = ""): GatewayModel => ({
Expand All @@ -16,6 +17,7 @@ const model = (id: string, owned_by = ""): GatewayModel => ({
context_window: 128000,
supports_streaming: true,
supports_vision: false,
allowed: true,
});

describe("formatGatewayModelName", () => {
Expand All @@ -27,6 +29,7 @@ describe("formatGatewayModelName", () => {
context_window: 200000,
supports_streaming: true,
supports_vision: true,
allowed: true,
}),
).toBe("Claude Opus 4.8");
});
Expand All @@ -39,6 +42,7 @@ describe("formatGatewayModelName", () => {
context_window: 200000,
supports_streaming: true,
supports_vision: true,
allowed: true,
}),
).toBe("gpt-5.5");
});
Expand All @@ -51,6 +55,7 @@ describe("formatGatewayModelName", () => {
context_window: 200000,
supports_streaming: true,
supports_vision: true,
allowed: true,
}),
).toBe("gpt-5.5");
});
Expand All @@ -63,6 +68,7 @@ describe("formatGatewayModelName", () => {
context_window: 128000,
supports_streaming: true,
supports_vision: false,
allowed: true,
}),
).toBe("glm-5.2");
});
Expand Down Expand Up @@ -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 },
Expand All @@ -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);
});
});
Loading
Loading