From d855b12c98af64ce98cd9b26bf9609370f9935e9 Mon Sep 17 00:00:00 2001 From: Mux Date: Fri, 21 Aug 2026 12:51:50 -0500 Subject: [PATCH] =?UTF-8?q?=F0=9F=A4=96=20fix:=20track=20estimated=20costs?= =?UTF-8?q?=20for=20new=20ChatGPT=20OAuth=20usage?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/common/utils/tokens/displayUsage.test.ts | 23 +++++- src/node/services/agentStatusService.ts | 1 - src/node/services/aiService.test.ts | 19 ++--- src/node/services/aiService.ts | 42 +---------- .../services/memoryConsolidationService.ts | 3 - .../services/providerModelFactory.test.ts | 70 ------------------- src/node/services/providerModelFactory.ts | 22 +----- src/node/services/workspaceStatusGenerator.ts | 6 +- 8 files changed, 36 insertions(+), 150 deletions(-) diff --git a/src/common/utils/tokens/displayUsage.test.ts b/src/common/utils/tokens/displayUsage.test.ts index d59c78d8c7c..c15afa174d7 100644 --- a/src/common/utils/tokens/displayUsage.test.ts +++ b/src/common/utils/tokens/displayUsage.test.ts @@ -388,7 +388,7 @@ describe("createDisplayUsage", () => { }); }); - describe("Subscription-covered usage costs", () => { + describe("Costs-included usage costs", () => { test("returns $0 costs when providerMetadata.mux.costsIncluded is true", () => { const usage: LanguageModelV2Usage = { inputTokens: 1000, // OpenAI includes cached tokens @@ -413,7 +413,7 @@ describe("createDisplayUsage", () => { expect(result!.reasoning.cost_usd).toBe(0); }); - test("gpt-5.3-codex routed through ChatGPT subscription is always zero-cost", () => { + test("returns $0 when Codex usage is marked as costs-included", () => { const usage: LanguageModelV2Usage = { inputTokens: 1500, // includes cached input tokens outputTokens: 450, @@ -462,6 +462,25 @@ describe("createDisplayUsage", () => { expect(result!.reasoning.cost_usd).toBeGreaterThan(0); }); + test("prices Codex usage when no costs-included marker is present", () => { + const usage: LanguageModelV2Usage = { + inputTokens: 1500, + outputTokens: 450, + reasoningTokens: 150, + totalTokens: 1950, + cachedInputTokens: 500, + }; + + const result = createDisplayUsage(usage, "openai:gpt-5.3-codex"); + + expect(result).toBeDefined(); + expect(result!.costsIncluded).toBeUndefined(); + expect(result!.input.cost_usd).toBeGreaterThan(0); + expect(result!.cached.cost_usd).toBeGreaterThan(0); + expect(result!.output.cost_usd).toBeGreaterThan(0); + expect(result!.reasoning.cost_usd).toBeGreaterThan(0); + }); + test("returns $0 costs even when model pricing is unknown", () => { const usage: LanguageModelV2Usage = { inputTokens: 100, diff --git a/src/node/services/agentStatusService.ts b/src/node/services/agentStatusService.ts index c7b3e5b1625..2bcf15dcb69 100644 --- a/src/node/services/agentStatusService.ts +++ b/src/node/services/agentStatusService.ts @@ -363,7 +363,6 @@ export class AgentStatusService { usage, usageOptions.providerMetadata, { - costsIncluded: usageOptions.costsIncluded, analyticsSource: "workspace_status", // Creation-time identity from the generator's pinned snapshot. metadataModel: usageOptions.metadataModel, diff --git a/src/node/services/aiService.test.ts b/src/node/services/aiService.test.ts index fd89220b9d1..6814728202f 100644 --- a/src/node/services/aiService.test.ts +++ b/src/node/services/aiService.test.ts @@ -3242,7 +3242,7 @@ describe("AIService.streamMessage compaction boundary slicing", () => { expect(typeof sessionUsageDeltaRecord.timestamp).toBe("number"); }); - it("zeros advisor tool usage costs for costs-included models before persisting", async () => { + it("tracks advisor tool usage costs for ChatGPT OAuth models", async () => { using muxHome = new DisposableTempDir("ai-service-tool-model-usage-costs-included"); const projectPath = path.join(muxHome.path, "project"); await fs.mkdir(projectPath, { recursive: true }); @@ -3331,18 +3331,19 @@ describe("AIService.streamMessage compaction boundary slicing", () => { }, timestamp: Date.now(), }; - const expectedDisplayUsage = createDisplayUsage(event.usage, event.model, { - ...(event.providerMetadata ?? {}), - mux: { costsIncluded: true }, - }); + const expectedDisplayUsage = createDisplayUsage( + event.usage, + event.model, + event.providerMetadata + ); expect(expectedDisplayUsage).toBeDefined(); if (!expectedDisplayUsage) { throw new Error("Expected tool usage event to produce display usage"); } - expect(expectedDisplayUsage.costsIncluded).toBe(true); - expect(expectedDisplayUsage.input.cost_usd).toBe(0); - expect(expectedDisplayUsage.output.cost_usd).toBe(0); - expect(expectedDisplayUsage.reasoning.cost_usd).toBe(0); + expect(expectedDisplayUsage.costsIncluded).toBeUndefined(); + expect(expectedDisplayUsage.input.cost_usd).toBeGreaterThan(0); + expect(expectedDisplayUsage.output.cost_usd).toBeGreaterThan(0); + expect(expectedDisplayUsage.reasoning.cost_usd).toBeGreaterThan(0); reportModelUsage(event); await Promise.resolve(); diff --git a/src/node/services/aiService.ts b/src/node/services/aiService.ts index 2cc19d35335..9499c5a4f99 100644 --- a/src/node/services/aiService.ts +++ b/src/node/services/aiService.ts @@ -167,7 +167,7 @@ import { import type { PTCEventWithParent } from "@/node/services/tools/code_execution"; import { MockAiStreamPlayer } from "./mock/mockAiStreamPlayer"; import { DEVTOOLS_RUN_METADATA_ID_HEADER } from "./devToolsHeaderCapture"; -import { ProviderModelFactory, modelCostsIncluded } from "./providerModelFactory"; +import { ProviderModelFactory } from "./providerModelFactory"; import { prepareMessagesForProvider } from "./messagePipeline"; import { getLegacyModeForAgentMetadata, resolveAgentForStream } from "./agentResolution"; import { buildPlanInstructions, buildStreamSystemContext } from "./streamContextBuilder"; @@ -341,29 +341,6 @@ function mergeProviderExtrasUnderMux( return merged; } -function markProviderMetadataCostsIncluded( - providerMetadata: Record | undefined, - costsIncluded: boolean | undefined -): Record | undefined { - if (!costsIncluded) { - return providerMetadata; - } - - const muxMetadata = providerMetadata?.mux; - const existingMux = - muxMetadata && typeof muxMetadata === "object" - ? (muxMetadata as Record) - : undefined; - - return { - ...(providerMetadata ?? {}), - mux: { - ...(existingMux ?? {}), - costsIncluded: true, - }, - }; -} - const WORKFLOW_CONTINUATION_RETRY_DELAY_MS = 1_000; const WORKSPACE_BUSY_IDLE_ONLY_SEND_MESSAGE = "Workspace is busy; idle-only send was skipped."; @@ -2300,10 +2277,6 @@ export class AIService extends EventEmitter { return; } }; - // Tool-side generateText() results do not consistently echo mux.costsIncluded in - // providerMetadata, so remember the resolved billing mode from model creation and - // re-stamp it before converting usage into display/session costs. - const toolModelCostsIncludedByModelString = new Map(); // Creation-time pricing identity for tool-created models (advisor): a // Coder catalog refresh can remove/retag the instance while the tool // request runs, and resolving the identity from live config at @@ -2559,10 +2532,6 @@ export class AIService extends EventEmitter { `Failed to create advisor model: ${getErrorMessage(advisorModel.error)}` ); } - toolModelCostsIncludedByModelString.set( - advisorModelString, - modelCostsIncluded(advisorModel.data) - ); // Same effective-route rule as createModelWithPinnedMetadata: // a coder: selection whose gateway is unavailable falls away // to a direct provider inside createModel, and identity or @@ -2668,10 +2637,7 @@ export class AIService extends EventEmitter { assert(eventModel.length > 0, "tool model usage event model must be non-empty"); // Persist tool-side model usage under its own model bucket so session costs keep // advisor/system-side pricing separate from the parent chat model. - const providerMetadata = markProviderMetadataCostsIncluded( - event.providerMetadata, - toolModelCostsIncludedByModelString.get(eventModel) - ); + const providerMetadata = event.providerMetadata; // Prefer the creation-time identity captured when the tool model // was created; models not created through the tool runtime fall // back to live resolution (their identity is not coder-scoped). @@ -3944,9 +3910,6 @@ export class AIService extends EventEmitter { initialMetadataPatch: { routedThroughGateway: next.routedThroughGateway, ...(next.routeProvider != null ? { routeProvider: next.routeProvider } : {}), - // Explicit undefined clears a stale costsIncluded when falling - // back from a subscription-routed model to an API model. - costsIncluded: modelCostsIncluded(next.model) ? true : undefined, systemMessageTokens: nextSystemTokens, }, }); @@ -4126,7 +4089,6 @@ export class AIService extends EventEmitter { ...(routeProvider != null ? { routeProvider } : {}), ...(muxMetadata !== undefined ? { muxMetadata } : {}), ...(acpPromptId != null ? { acpPromptId } : {}), - ...(modelCostsIncluded(modelResult.data.model) ? { costsIncluded: true } : {}), }, streamProviderOptions, maxOutputTokens, diff --git a/src/node/services/memoryConsolidationService.ts b/src/node/services/memoryConsolidationService.ts index 702f7062a54..d98bcd16c95 100644 --- a/src/node/services/memoryConsolidationService.ts +++ b/src/node/services/memoryConsolidationService.ts @@ -17,7 +17,6 @@ import * as path from "node:path"; import writeFileAtomic from "write-file-atomic"; import { z } from "zod"; import type { LanguageModel } from "ai"; -import { modelCostsIncluded } from "@/node/services/providerModelFactory"; import type { SessionUsageService } from "@/node/services/sessionUsageService"; import type { CompactionCompletionMetadata } from "@/common/types/compaction"; import type { Result } from "@/common/types/result"; @@ -515,7 +514,6 @@ export class MemoryConsolidationService extends EventEmitter { usage, providerMetadata, { - costsIncluded: modelCostsIncluded(modelResult.data.model), analyticsSource: "memory_consolidation", // Creation-time identity: a catalog refresh mid-run must not // re-attribute this spend (see ModelFactoryLike). @@ -662,7 +660,6 @@ export class MemoryConsolidationService extends EventEmitter { usage, providerMetadata, { - costsIncluded: modelCostsIncluded(modelResult.data.model), analyticsSource: "memory_harvest", // Creation-time identity (see ModelFactoryLike). metadataModel: modelResult.data.metadataModel, diff --git a/src/node/services/providerModelFactory.test.ts b/src/node/services/providerModelFactory.test.ts index b87d64f2d4a..60e99f33336 100644 --- a/src/node/services/providerModelFactory.test.ts +++ b/src/node/services/providerModelFactory.test.ts @@ -17,7 +17,6 @@ import { buildAIProviderRequestHeaders, classifyCopilotInitiator, countAnthropicCacheBreakpoints, - modelCostsIncluded, XUM_AI_PROVIDER_USER_AGENT, normalizeCodexResponsesBody, markCodexOauthRoutedResponse, @@ -1199,7 +1198,6 @@ describe("ProviderModelFactory OpenAI WebSocket transport", () => { return; } expect(hasLanguageModelCleanup(result.data)).toBe(false); - expect(modelCostsIncluded(result.data)).toBe(true); }); }); @@ -1295,74 +1293,6 @@ describe("ProviderModelFactory OpenAI WebSocket transport", () => { }); }); -describe("ProviderModelFactory modelCostsIncluded", () => { - it("marks gpt-5.3-codex as subscription-covered when routed through Codex OAuth", async () => { - await withTempConfig(async (config, factory) => { - config.saveProvidersConfig({ - openai: { - codexOauth: { - type: "oauth", - access: "test-access-token", - refresh: "test-refresh-token", - expires: Date.now() + 60_000, - accountId: "test-account-id", - }, - }, - }); - - const result = await factory.createModel(KNOWN_MODELS.GPT_53_CODEX.id); - expect(result.success).toBe(true); - if (!result.success) { - return; - } - - expect(modelCostsIncluded(result.data)).toBe(true); - }); - }); - - it("routes a custom OpenAI model through Codex OAuth when it inherits from a compatible model", async () => { - await withTempConfig(async (config, factory) => { - config.saveProvidersConfig({ - openai: { - codexOauth: { - type: "oauth", - access: "test-access-token", - refresh: "test-refresh-token", - expires: Date.now() + 60_000, - accountId: "test-account-id", - }, - models: [{ id: "team-codex", mappedToModel: KNOWN_MODELS.GPT_53_CODEX.id }], - }, - }); - - const result = await factory.createModel("openai:team-codex"); - expect(result.success).toBe(true); - if (!result.success) { - return; - } - - expect(modelCostsIncluded(result.data)).toBe(true); - }); - }); - - it("does not mark gpt-5.3-codex as subscription-covered when routed through API key", async () => { - await withTempConfig(async (config, factory) => { - config.saveProvidersConfig({ - openai: { - apiKey: "sk-test", - }, - }); - - const result = await factory.createModel(KNOWN_MODELS.GPT_53_CODEX.id); - expect(result.success).toBe(true); - if (!result.success) { - return; - } - - expect(modelCostsIncluded(result.data)).toBe(false); - }); - }); -}); describe("ProviderModelFactory routing", () => { it("honors non-mux gateway routes end-to-end", async () => { await withTempConfig(async (config, factory) => { diff --git a/src/node/services/providerModelFactory.ts b/src/node/services/providerModelFactory.ts index 50dad21c6a1..5d8e689bb31 100644 --- a/src/node/services/providerModelFactory.ts +++ b/src/node/services/providerModelFactory.ts @@ -737,24 +737,6 @@ function parseAnthropicCacheTtl(value: unknown): AnthropicCacheTtl | undefined { return undefined; } -// --------------------------------------------------------------------------- -// Model cost tracking -// --------------------------------------------------------------------------- - -const MUX_MODEL_COSTS_INCLUDED = Symbol("mux:modelCostsIncluded"); - -type LanguageModelWithMuxCostsIncluded = LanguageModel & { - [MUX_MODEL_COSTS_INCLUDED]?: true; -}; - -function markModelCostsIncluded(model: LanguageModel): void { - (model as LanguageModelWithMuxCostsIncluded)[MUX_MODEL_COSTS_INCLUDED] = true; -} - -export function modelCostsIncluded(model: LanguageModel): boolean { - return (model as LanguageModelWithMuxCostsIncluded)[MUX_MODEL_COSTS_INCLUDED] === true; -} - const CODEX_ALLOWED_PARAMS = new Set([ "model", "input", @@ -1601,8 +1583,8 @@ export class ProviderModelFactory { // Skip Codex OAuth routing for chatCompletions — the Codex endpoint // only accepts Responses API format, so chat-completions requests would fail. if (shouldRouteThroughCodexOauth && effectiveWireFormat !== "chatCompletions") { - markModelCostsIncluded(model); - + // Keep OAuth usage on the normal model-pricing path. ChatGPT OAuth does + // not expose an API charge, but cost analytics still need an estimate. // Codex OAuth requires store=false and must override any request-level // setting to avoid unresolved item_reference lookups. injectModelOpenAIStore(false, "force"); diff --git a/src/node/services/workspaceStatusGenerator.ts b/src/node/services/workspaceStatusGenerator.ts index bfd059d8932..ec04a5a329a 100644 --- a/src/node/services/workspaceStatusGenerator.ts +++ b/src/node/services/workspaceStatusGenerator.ts @@ -1,6 +1,5 @@ import { streamText, tool } from "ai"; import type { LanguageModelV2Usage } from "@ai-sdk/provider"; -import { modelCostsIncluded } from "./providerModelFactory"; import type { AIService } from "./aiService"; import { log } from "./log"; import { runLanguageModelCleanup } from "./languageModelCleanup"; @@ -136,14 +135,12 @@ export async function generateWorkspaceStatus( /** * Best-effort cost telemetry: status generation bypasses StreamManager, * so the caller records the successful candidate's usage into - * session-usage.json. costsIncluded reflects subscription-covered routing - * (Codex OAuth) so those tokens are priced at $0. + * session-usage.json. */ recordUsage?: ( modelString: string, usage: LanguageModelV2Usage, options: { - costsIncluded: boolean; /** * Step-accumulated provider metadata. Anthropic reports billed * cache-write tokens only here (cacheCreationInputTokens), not in @@ -235,7 +232,6 @@ export async function generateWorkspaceStatus( if (settled !== undefined) { const [usage, steps] = settled; await options.recordUsage(modelString, usage, { - costsIncluded: modelCostsIncluded(modelResult.data.model), providerMetadata: accumulateStepsProviderMetadata(steps), metadataModel: modelResult.data.metadataModel, });