fix(ai): trust requested Codex service tier (#3307)

This commit is contained in:
Markus Ylisiurunen
2026-04-17 02:50:04 +03:00
committed by GitHub
parent e9808b5854
commit 2cdac7382a
3 changed files with 108 additions and 1 deletions

View File

@@ -385,6 +385,16 @@ function applyServiceTierPricing(usage: Usage, serviceTier: ResponseCreateParams
usage.cost.total = usage.cost.input + usage.cost.output + usage.cost.cacheRead + usage.cost.cacheWrite;
}
function resolveCodexServiceTier(
responseServiceTier: ResponseCreateParamsStreaming["service_tier"] | undefined,
requestServiceTier: ResponseCreateParamsStreaming["service_tier"] | undefined,
): ResponseCreateParamsStreaming["service_tier"] | undefined {
if (responseServiceTier === "default" && (requestServiceTier === "flex" || requestServiceTier === "priority")) {
return requestServiceTier;
}
return responseServiceTier ?? requestServiceTier;
}
function resolveCodexUrl(baseUrl?: string): string {
const raw = baseUrl && baseUrl.trim().length > 0 ? baseUrl : DEFAULT_CODEX_BASE_URL;
const normalized = raw.replace(/\/+$/, "");
@@ -413,6 +423,7 @@ async function processStream(
): Promise<void> {
await processResponsesStream(mapCodexEvents(parseSSE(response)), output, stream, model, {
serviceTier: options?.serviceTier,
resolveServiceTier: resolveCodexServiceTier,
applyServiceTierPricing,
});
}
@@ -847,6 +858,7 @@ async function processWebSocketStream(
stream.push({ type: "start", partial: output });
await processResponsesStream(mapCodexEvents(parseWebSocket(socket, options?.signal)), output, stream, model, {
serviceTier: options?.serviceTier,
resolveServiceTier: resolveCodexServiceTier,
applyServiceTierPricing,
});
if (options?.signal?.aborted) {

View File

@@ -65,6 +65,10 @@ function parseTextSignature(
export interface OpenAIResponsesStreamOptions {
serviceTier?: ResponseCreateParamsStreaming["service_tier"];
resolveServiceTier?: (
responseServiceTier: ResponseCreateParamsStreaming["service_tier"] | undefined,
requestServiceTier: ResponseCreateParamsStreaming["service_tier"] | undefined,
) => ResponseCreateParamsStreaming["service_tier"] | undefined;
applyServiceTierPricing?: (
usage: Usage,
serviceTier: ResponseCreateParamsStreaming["service_tier"] | undefined,
@@ -491,7 +495,9 @@ export async function processResponsesStream<TApi extends Api>(
}
calculateCost(model, output.usage);
if (options?.applyServiceTierPricing) {
const serviceTier = response?.service_tier ?? options.serviceTier;
const serviceTier = options.resolveServiceTier
? options.resolveServiceTier(response?.service_tier, options.serviceTier)
: (response?.service_tier ?? options.serviceTier);
options.applyServiceTierPricing(output.usage, serviceTier);
}
// Map status to stop reason

View File

@@ -499,6 +499,95 @@ describe("openai-codex streaming", () => {
await streamResult.result();
});
it.each([
["flex", 0.5],
["priority", 2],
] as const)("uses the client-sent %s service tier when Codex echoes default", async (serviceTier, multiplier) => {
const tempDir = mkdtempSync(join(tmpdir(), "pi-codex-stream-"));
process.env.PI_CODING_AGENT_DIR = tempDir;
const token = mockToken();
const sse = `${[
`data: ${JSON.stringify({
type: "response.output_item.added",
item: { type: "message", id: "msg_1", role: "assistant", status: "in_progress", content: [] },
})}`,
`data: ${JSON.stringify({ type: "response.content_part.added", part: { type: "output_text", text: "" } })}`,
`data: ${JSON.stringify({ type: "response.output_text.delta", delta: "Hello" })}`,
`data: ${JSON.stringify({
type: "response.output_item.done",
item: {
type: "message",
id: "msg_1",
role: "assistant",
status: "completed",
content: [{ type: "output_text", text: "Hello" }],
},
})}`,
`data: ${JSON.stringify({
type: "response.completed",
response: {
status: "completed",
service_tier: "default",
usage: {
input_tokens: 1000000,
output_tokens: 1000000,
total_tokens: 2000000,
input_tokens_details: { cached_tokens: 0 },
},
},
})}`,
].join("\n\n")}\n\n`;
const encoder = new TextEncoder();
const stream = new ReadableStream<Uint8Array>({
start(controller) {
controller.enqueue(encoder.encode(sse));
controller.close();
},
});
global.fetch = vi.fn(async (input: string | URL) => {
const url = typeof input === "string" ? input : input.toString();
if (url === "https://api.github.com/repos/openai/codex/releases/latest") {
return new Response(JSON.stringify({ tag_name: "rust-v0.0.0" }), { status: 200 });
}
if (url.startsWith("https://raw.githubusercontent.com/openai/codex/")) {
return new Response("PROMPT", { status: 200, headers: { etag: '"etag"' } });
}
if (url === "https://chatgpt.com/backend-api/codex/responses") {
return new Response(stream, {
status: 200,
headers: { "content-type": "text/event-stream" },
});
}
return new Response("not found", { status: 404 });
}) as typeof fetch;
const model: Model<"openai-codex-responses"> = {
id: "gpt-5.1-codex",
name: "GPT-5.1 Codex",
api: "openai-codex-responses",
provider: "openai-codex",
baseUrl: "https://chatgpt.com/backend-api",
reasoning: true,
input: ["text"],
cost: { input: 1, output: 2, cacheRead: 0, cacheWrite: 0 },
contextWindow: 400000,
maxTokens: 128000,
};
const context: Context = {
systemPrompt: "You are a helpful assistant.",
messages: [{ role: "user", content: "Say hello", timestamp: Date.now() }],
};
const result = await streamOpenAICodexResponses(model, context, { apiKey: token, serviceTier }).result();
expect(result.usage.cost.input).toBe(1 * multiplier);
expect(result.usage.cost.output).toBe(2 * multiplier);
expect(result.usage.cost.total).toBe(3 * multiplier);
});
it("does not set session_id/x-client-request-id headers when sessionId is not provided", async () => {
const tempDir = mkdtempSync(join(tmpdir(), "pi-codex-stream-"));
process.env.PI_CODING_AGENT_DIR = tempDir;