From 2cdac7382acf6a8927c9056a2acd053c211db97d Mon Sep 17 00:00:00 2001 From: Markus Ylisiurunen <8409947+markusylisiurunen@users.noreply.github.com> Date: Fri, 17 Apr 2026 02:50:04 +0300 Subject: [PATCH] fix(ai): trust requested Codex service tier (#3307) --- .../src/providers/openai-codex-responses.ts | 12 +++ .../src/providers/openai-responses-shared.ts | 8 +- packages/ai/test/openai-codex-stream.test.ts | 89 +++++++++++++++++++ 3 files changed, 108 insertions(+), 1 deletion(-) diff --git a/packages/ai/src/providers/openai-codex-responses.ts b/packages/ai/src/providers/openai-codex-responses.ts index 248443ec..08a58779 100644 --- a/packages/ai/src/providers/openai-codex-responses.ts +++ b/packages/ai/src/providers/openai-codex-responses.ts @@ -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 { 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) { diff --git a/packages/ai/src/providers/openai-responses-shared.ts b/packages/ai/src/providers/openai-responses-shared.ts index bb245382..c43f926e 100644 --- a/packages/ai/src/providers/openai-responses-shared.ts +++ b/packages/ai/src/providers/openai-responses-shared.ts @@ -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( } 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 diff --git a/packages/ai/test/openai-codex-stream.test.ts b/packages/ai/test/openai-codex-stream.test.ts index 58496988..43f005fa 100644 --- a/packages/ai/test/openai-codex-stream.test.ts +++ b/packages/ai/test/openai-codex-stream.test.ts @@ -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({ + 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;