diff --git a/src/CodexAcpServer.ts b/src/CodexAcpServer.ts index 6e0c394b..68b29947 100644 --- a/src/CodexAcpServer.ts +++ b/src/CodexAcpServer.ts @@ -94,6 +94,8 @@ import { createUserMessageChunk, } from "./ContentChunks"; import {sameThreadGoalSnapshot, type ThreadGoalSnapshot, toThreadGoalSnapshot,} from "./ThreadGoalSnapshot"; +import {OpenAiPricingProvider, type PricingProvider} from "./PricingProvider"; +import {SessionCostTracker} from "./SessionCostTracker"; import {randomUUID} from "node:crypto"; import {once} from "node:events"; import { @@ -144,6 +146,7 @@ export interface SessionState { sessionTitle: string | null; sessionTitleSource: "unset" | "fallback" | "explicit" | "unknown"; sessionFailure?: SessionFailure; + costTracker: SessionCostTracker; } export type SessionFailureCategory = @@ -240,6 +243,7 @@ export class CodexAcpServer { private readonly getExitCode: () => number | null; private readonly getRecentStderr: () => string; private readonly sessionFailureEpoch: string; + private readonly pricingProvider: PricingProvider; private availableCommands: CodexCommands; private clientInfo: acp.Implementation | null; private clientCapabilities: acp.ClientCapabilities | null; @@ -266,6 +270,7 @@ export class CodexAcpServer { getExitCode?: () => number | null, getRecentStderr?: () => string, codexProcessState?: CodexProcessState, + pricingProvider: PricingProvider = new OpenAiPricingProvider(), ) { this.sessions = new Map(); this.pendingMcpStartupSessions = new Map(); @@ -284,6 +289,7 @@ export class CodexAcpServer { this.getExitCode = getExitCode ?? (() => this.codexProcessState?.connection.process.exitCode ?? null); this.getRecentStderr = getRecentStderr ?? (() => this.codexProcessState?.stderr ?? ""); this.sessionFailureEpoch = randomUUID(); + this.pricingProvider = pricingProvider; this.clientInfo = null; this.clientCapabilities = null; this.terminalOutputMode = "terminal_output_delta"; @@ -600,6 +606,11 @@ export class CodexAcpServer { const sessionMcpServers = this.resolveSessionMcpServers(requestedMcpServers, "sessionId" in request); const currentModel = this.findCurrentModel(models, currentModelId); const currentModelSupportsFast = modelSupportsFast(currentModel); + const costTracker = await this.createSessionCostTracker( + models, + authProvider, + "sessionId" in request, + ); const sessionState: SessionState = { sessionId: sessionId, currentModelId: currentModelId, @@ -626,6 +637,7 @@ export class CodexAcpServer { goalRevision: 0, sessionTitle: null, sessionTitleSource: "sessionId" in request ? "unknown" : "unset", + costTracker, }; this.sessions.set(sessionId, sessionState); resumeSubscribed = false; @@ -666,6 +678,18 @@ export class CodexAcpServer { return authProvider === null || authProvider === "openai"; } + private async createSessionCostTracker( + models: readonly Model[], + authProvider: string | null, + baselineInitialUsage: boolean, + ): Promise { + if (!this.authProviderUsesOpenAiAccount(authProvider)) { + return SessionCostTracker.disabled(); + } + const pricing = await this.pricingProvider.getPricing(models); + return new SessionCostTracker(pricing, baselineInitialUsage); + } + private authProvidersMatch(a: string | null, b: string | null): boolean { if (this.authProviderUsesOpenAiAccount(a) && this.authProviderUsesOpenAiAccount(b)) { return true; @@ -1598,6 +1622,7 @@ export class CodexAcpServer { const sessionMcpServers = this.resolveSessionMcpServers(requestedMcpServers, true); const currentModel = this.findCurrentModel(models, currentModelId); const currentModelSupportsFast = modelSupportsFast(currentModel); + const costTracker = await this.createSessionCostTracker(models, authProvider, true); const sessionState: SessionState = { sessionId: sessionId, currentModelId: currentModelId, @@ -1624,6 +1649,7 @@ export class CodexAcpServer { goalRevision: 0, sessionTitle: null, sessionTitleSource: "unset", + costTracker, }; this.sessions.set(sessionId, sessionState); subscribed = false; diff --git a/src/CodexEventHandler.ts b/src/CodexEventHandler.ts index 6418b624..e5858e1d 100644 --- a/src/CodexEventHandler.ts +++ b/src/CodexEventHandler.ts @@ -1233,10 +1233,20 @@ export class CodexEventHandler { return null; } + const cost = this.sessionState.totalTokenUsage === null || this.sessionState.lastTokenUsage === null + ? null + : this.sessionState.costTracker.update( + this.sessionState.totalTokenUsage, + this.sessionState.lastTokenUsage, + this.sessionState.currentModelId, + this.sessionState.fastModeEnabled, + ); + return { sessionUpdate: "usage_update", used, size, + ...(cost === null ? {} : {cost}), }; } diff --git a/src/PricingProvider.ts b/src/PricingProvider.ts new file mode 100644 index 00000000..94346068 --- /dev/null +++ b/src/PricingProvider.ts @@ -0,0 +1,144 @@ +import type {Model} from "./app-server/v2"; +import {logger} from "./Logger"; + +export interface TokenRates { + input: number; + cachedInput: number; + output: number; +} + +export interface TierPricing { + shortContext: TokenRates; + longContext?: TokenRates; +} + +export interface ModelPricing { + standard: TierPricing; + fast?: TierPricing; +} + +export type ModelPricingSnapshot = ReadonlyMap; + +export interface PricingProvider { + getPricing(models: readonly Model[]): Promise; +} + +const OPENAI_PRICING_URL = "https://developers.openai.com/api/docs/pricing.md"; +const PRICING_FETCH_TIMEOUT_MS = 5_000; + +export class OpenAiPricingProvider implements PricingProvider { + private pricingDocument: Promise | null = null; + + async getPricing(models: readonly Model[]): Promise { + if (models.length === 0) return new Map(); + + try { + const markdown = await this.getPricingDocument(); + return parseModelPricing(markdown, models.map(model => model.id)); + } catch (error) { + logger.error("Failed to load OpenAI model pricing", error); + return new Map(); + } + } + + private async getPricingDocument(): Promise { + if (this.pricingDocument === null) { + this.pricingDocument = this.fetchPricingDocument().catch(error => { + this.pricingDocument = null; + throw error; + }); + } + return await this.pricingDocument; + } + + private async fetchPricingDocument(): Promise { + const response = await fetch(OPENAI_PRICING_URL, { + headers: {accept: "text/markdown"}, + signal: AbortSignal.timeout(PRICING_FETCH_TIMEOUT_MS), + }); + if (!response.ok) { + throw new Error(`OpenAI pricing request failed with HTTP ${response.status}`); + } + return await response.text(); + } +} + +function parseModelPricing(markdown: string, modelIds: readonly string[]): ModelPricingSnapshot { + const requestedModels = new Map(); + for (const modelId of modelIds) { + const pricingId = pricingModelId(modelId); + const matchingIds = requestedModels.get(pricingId) ?? []; + matchingIds.push(modelId); + requestedModels.set(pricingId, matchingIds); + } + + const standard = parsePricingSection(markdown, "Standard pricing data", requestedModels); + const fast = parsePricingSection(markdown, "Fast pricing data", requestedModels); + const result = new Map(); + for (const [modelId, standardPricing] of standard) { + const fastPricing = fast.get(modelId); + result.set(modelId, { + standard: standardPricing, + ...(fastPricing === undefined ? {} : {fast: fastPricing}), + }); + } + return result; +} + +function parsePricingSection( + markdown: string, + heading: string, + requestedModels: ReadonlyMap, +): Map { + const headingText = `### ${heading}`; + const start = markdown.indexOf(headingText); + if (start < 0) return new Map(); + const nextHeading = markdown.indexOf("\n### ", start + headingText.length); + const section = markdown.slice(start, nextHeading < 0 ? undefined : nextHeading); + const result = new Map(); + + for (const line of section.split(/\r?\n/)) { + if (!line.startsWith("|")) continue; + const cells = line.split("|").slice(1, -1).map(cell => cell.trim()); + if (cells.length < 9 || cells[0] === "Model" || cells[0]?.startsWith("---")) continue; + + const documentedModel = cells[0]?.replace(/\s+\([^)]*\)\s*$/, ""); + if (documentedModel === undefined) continue; + const matchingModelIds = requestedModels.get(documentedModel); + if (matchingModelIds === undefined) continue; + + const shortContext = parseTokenRates(cells[1], cells[2], cells[4]); + if (shortContext === null) continue; + const longContext = parseTokenRates(cells[5], cells[6], cells[8]); + const pricing: TierPricing = { + shortContext, + ...(longContext === null ? {} : {longContext}), + }; + for (const modelId of matchingModelIds) { + result.set(modelId, pricing); + } + } + return result; +} + +function parseTokenRates( + inputValue: string | undefined, + cachedInputValue: string | undefined, + outputValue: string | undefined, +): TokenRates | null { + const input = parseUsdRate(inputValue); + const cachedInput = parseUsdRate(cachedInputValue); + const output = parseUsdRate(outputValue); + if (input === null || cachedInput === null || output === null) return null; + return {input, cachedInput, output}; +} + +function parseUsdRate(value: string | undefined): number | null { + if (value === undefined || value === "-") return null; + const amount = Number(value.replace(/[$,]/g, "")); + return Number.isFinite(amount) && amount >= 0 ? amount : null; +} + +function pricingModelId(modelId: string): string { + return modelId.toLowerCase().replace(/-\d{4}-\d{2}-\d{2}$/, ""); +} diff --git a/src/SessionCostTracker.ts b/src/SessionCostTracker.ts new file mode 100644 index 00000000..409884e2 --- /dev/null +++ b/src/SessionCostTracker.ts @@ -0,0 +1,93 @@ +import type {Cost} from "@agentclientprotocol/sdk"; +import type {TokenCount} from "./TokenCount"; +import type {ModelPricingSnapshot, TierPricing, TokenRates} from "./PricingProvider"; + +const LONG_CONTEXT_THRESHOLD = 272_000; +const TOKENS_PER_MILLION = 1_000_000; + +export class SessionCostTracker { + private previousTotalUsage: TokenCount | null = null; + private amountUsd = 0; + private available: boolean; + + constructor( + private readonly pricing: ModelPricingSnapshot, + private baselineInitialUsage = false, + ) { + this.available = pricing.size > 0; + } + + static disabled(): SessionCostTracker { + return new SessionCostTracker(new Map()); + } + + update( + totalUsage: TokenCount, + lastUsage: TokenCount, + currentModelId: string, + fastModeEnabled: boolean, + ): Cost | null { + const delta = usageDelta(totalUsage, this.previousTotalUsage); + this.previousTotalUsage = {...totalUsage}; + if (!this.available || delta === null) { + this.available = false; + return null; + } + + const modelId = currentModelId.replace(/\[[^\]]*]$/, ""); + const modelPricing = this.pricing.get(modelId); + const tierPricing = fastModeEnabled ? modelPricing?.fast : modelPricing?.standard; + const rates = selectRates(tierPricing, lastUsage); + if (rates === null) { + this.available = false; + return null; + } + + if (this.baselineInitialUsage) { + this.baselineInitialUsage = false; + return {amount: this.amountUsd, currency: "USD"}; + } + + const incrementalCost = ( + delta.inputTokens * rates.input + + delta.cachedInputTokens * rates.cachedInput + + delta.outputTokens * rates.output + ) / TOKENS_PER_MILLION; + if (!Number.isFinite(incrementalCost) || incrementalCost < 0) { + this.available = false; + return null; + } + + this.amountUsd += incrementalCost; + return {amount: this.amountUsd, currency: "USD"}; + } +} + +function selectRates(pricing: TierPricing | undefined, lastUsage: TokenCount): TokenRates | null { + if (pricing === undefined) return null; + const inputTokens = lastUsage.inputTokens + lastUsage.cachedInputTokens; + if (inputTokens > LONG_CONTEXT_THRESHOLD) { + return pricing.longContext ?? null; + } + return pricing.shortContext; +} + +function usageDelta(current: TokenCount, previous: TokenCount | null): TokenCount | null { + if (previous === null) return {...current}; + if ( + current.totalTokens < previous.totalTokens + || current.inputTokens < previous.inputTokens + || current.cachedInputTokens < previous.cachedInputTokens + || current.outputTokens < previous.outputTokens + || current.reasoningOutputTokens < previous.reasoningOutputTokens + ) { + return null; + } + return { + totalTokens: current.totalTokens - previous.totalTokens, + inputTokens: current.inputTokens - previous.inputTokens, + cachedInputTokens: current.cachedInputTokens - previous.cachedInputTokens, + outputTokens: current.outputTokens - previous.outputTokens, + reasoningOutputTokens: current.reasoningOutputTokens - previous.reasoningOutputTokens, + }; +} diff --git a/src/__tests__/SessionCostTracker.test.ts b/src/__tests__/SessionCostTracker.test.ts new file mode 100644 index 00000000..ed2ad38f --- /dev/null +++ b/src/__tests__/SessionCostTracker.test.ts @@ -0,0 +1,60 @@ +import {describe, expect, it} from "vitest"; +import type {ModelPricingSnapshot} from "../PricingProvider"; +import {SessionCostTracker} from "../SessionCostTracker"; + +const pricing: ModelPricingSnapshot = new Map([ + ["gpt-test", { + standard: { + shortContext: {input: 10, cachedInput: 2, output: 20}, + longContext: {input: 10, cachedInput: 2, output: 20}, + }, + }], +]); + +describe("SessionCostTracker", () => { + it("accumulates USD cost from cumulative token usage", () => { + const tracker = new SessionCostTracker(pricing); + + expect(tracker.update( + usage(1_000_000, 500_000, 250_000), + usage(1_000_000, 500_000, 250_000), + "gpt-test[medium]", + false, + )).toEqual({amount: 16, currency: "USD"}); + + expect(tracker.update( + usage(1_100_000, 550_000, 300_000), + usage(100_000, 50_000, 50_000), + "gpt-test[medium]", + false, + )).toEqual({amount: 18.1, currency: "USD"}); + }); + + it("baselines existing usage when a session is resumed", () => { + const tracker = new SessionCostTracker(pricing, true); + + expect(tracker.update( + usage(1_000_000, 500_000, 250_000), + usage(1_000_000, 500_000, 250_000), + "gpt-test[medium]", + false, + )).toEqual({amount: 0, currency: "USD"}); + + expect(tracker.update( + usage(1_100_000, 550_000, 300_000), + usage(100_000, 50_000, 50_000), + "gpt-test[medium]", + false, + )).toEqual({amount: 2.1, currency: "USD"}); + }); +}); + +function usage(inputTokens: number, cachedInputTokens: number, outputTokens: number) { + return { + totalTokens: inputTokens + cachedInputTokens + outputTokens, + inputTokens, + cachedInputTokens, + outputTokens, + reasoningOutputTokens: 0, + }; +} diff --git a/src/__tests__/acp-test-utils.ts b/src/__tests__/acp-test-utils.ts index 8d6c9071..236e89b4 100644 --- a/src/__tests__/acp-test-utils.ts +++ b/src/__tests__/acp-test-utils.ts @@ -4,6 +4,7 @@ import {CodexAcpClient} from '../CodexAcpClient'; import {CodexAppServerClient, type CodexConnectionEvent} from '../CodexAppServerClient'; import {startCodexConnection} from "../CodexJsonRpcConnection"; import {CodexAcpServer, type SessionState} from "../CodexAcpServer"; +import {SessionCostTracker} from "../SessionCostTracker"; import type {AcpClientConnection} from "../ACPSessionConnection"; import type {ServerNotification} from "../app-server"; import type {MessageConnection} from "vscode-jsonrpc/node"; @@ -403,6 +404,7 @@ export function createTestSessionState(overrides?: Partial): Sessi goalRevision: 0, sessionTitle: null, sessionTitleSource: "unknown", + costTracker: SessionCostTracker.disabled(), ...overrides, }; }