Add basic cost calculations for testing

package-lock.json

@@ -65,6 +65,7 @@
                 "express": "5.2.1",
                 "express-rate-limit": "8.5.2",
                 "glob": "13.0.6",
+                "gpt-tokenizer": "^3.4.0",
                 "helmet": "8.2.0",
                 "http-errors": "2.0.1",
                 "input-otp": "1.4.2",
@@ -12760,6 +12761,12 @@
                 "url": "https://github.com/sponsors/ljharb"
             }
         },
+        "node_modules/gpt-tokenizer": {
+            "version": "3.4.0",
+            "resolved": "https://registry.npmjs.org/gpt-tokenizer/-/gpt-tokenizer-3.4.0.tgz",
+            "integrity": "sha512-wxFLnhIXTDjYebd9A9pGl3e31ZpSypbpIJSOswbgop5jLte/AsZVDvjlbEuVFlsqZixVKqbcoNmRlFDf6pz/UQ==",
+            "license": "MIT"
+        },
         "node_modules/graceful-fs": {
             "version": "4.2.11",
             "resolved": "https://registry.npmjs.org/graceful-fs/-/graceful-fs-4.2.11.tgz",

package.json

@@ -33,9 +33,9 @@
     },
     "dependencies": {
         "@asteasolutions/zod-to-openapi": "8.5.0",
+        "@aws-sdk/client-s3": "3.1056.0",
         "@devolutions/iron-remote-desktop": "https://static.pangolin.net/packages/devolutions-iron-remote-desktop-0.0.0.tgz",
         "@devolutions/iron-remote-desktop-rdp": "https://static.pangolin.net/packages/devolutions-iron-remote-desktop-rdp-0.0.1.tgz",
-        "@aws-sdk/client-s3": "3.1056.0",
         "@headlessui/react": "2.2.10",
         "@hookform/resolvers": "5.4.0",
         "@monaco-editor/react": "4.7.0",
@@ -88,6 +88,7 @@
         "express": "5.2.1",
         "express-rate-limit": "8.5.2",
         "glob": "13.0.6",
+        "gpt-tokenizer": "^3.4.0",
         "helmet": "8.2.0",
         "http-errors": "2.0.1",
         "input-otp": "1.4.2",

server/lib/aiModelPricing.ts

@@ -0,0 +1,230 @@
+import fs from "node:fs";
+import path from "node:path";
+import { APP_PATH } from "@server/lib/consts";
+import type { AiProviderType } from "@server/lib/aiProviderDefaults";
+import type { AiUsage } from "@server/lib/aiUsageExtraction";
+import logger from "@server/logger";
+
+// config/models.json is a runtime asset (same category as config.yml or the
+// MaxMind DBs) - not part of the source tree. Its shape mirrors litellm's
+// public model_prices_and_context_window.json: a flat list of
+// { id, name, provider, input_cost_per_token, output_cost_per_token,
+// cache_read_input_token_cost, output_cost_per_reasoning_token }, where
+// `provider` is litellm's provider bucket, not our AiProviderType.
+const MODELS_JSON_PATH = path.join(APP_PATH, "models.json");
+
+export type AiModelPricingEntry = {
+    id: string;
+    name: string;
+    provider: string;
+    input_cost_per_token: number | null;
+    output_cost_per_token: number | null;
+    cache_read_input_token_cost: number | null;
+    output_cost_per_reasoning_token: number | null;
+};
+
+export type AiModelPricing = {
+    inputCostPerToken: number | null;
+    outputCostPerToken: number | null;
+    cacheReadInputTokenCost: number | null;
+    outputCostPerReasoningToken: number | null;
+    // True when the match came from a different provider bucket than the one
+    // mapped to this provider's type (e.g. an openRouter/custom model id that
+    // only matched by stripping a "vendor/" prefix against the whole table).
+    // Costs found this way are a best-effort approximation, not a guarantee
+    // the upstream provider bills at the same rate.
+    approximate: boolean;
+};
+
+// Which litellm provider buckets to search for each of our provider types.
+// Several of our provider types (openRouter, vercelAiGateway, custom) proxy
+// arbitrary underlying models and have no dedicated bucket in the pricing
+// data, so they fall back to a global search across all buckets.
+const PROVIDER_PRICING_BUCKETS: Record<
+    Exclude<AiProviderType, "custom">,
+    string[]
+> = {
+    openai: ["openai"],
+    anthropic: ["anthropic"],
+    googleGemini: ["gemini"],
+    vertexAi: [
+        "vertex_ai-language-models",
+        "vertex_ai",
+        "vertex_ai-anthropic_models",
+        "vertex_ai-mistral_models",
+        "vertex_ai-deepseek_models",
+        "vertex_ai-ai21_models",
+        "vertex_ai-llama_models",
+        "vertex_ai-minimax_models",
+        "vertex_ai-moonshot_models",
+        "vertex_ai-zai_models",
+        "vertex_ai-openai_models",
+        "vertex_ai-qwen_models",
+        "vertex_ai-text-models"
+    ],
+    bedrock: ["bedrock_converse", "bedrock", "bedrock_mantle"],
+    microsoftFoundry: ["azure", "azure_ai", "azure_text"],
+    openRouter: [],
+    vercelAiGateway: []
+};
+
+let modelsById: Map<string, AiModelPricingEntry[]> | null = null;
+
+function loadModels(): Map<string, AiModelPricingEntry[]> {
+    if (modelsById) {
+        return modelsById;
+    }
+
+    const byId = new Map<string, AiModelPricingEntry[]>();
+    try {
+        if (fs.existsSync(MODELS_JSON_PATH)) {
+            const raw = fs.readFileSync(MODELS_JSON_PATH, "utf-8");
+            const parsed = JSON.parse(raw) as { data: AiModelPricingEntry[] };
+            for (const entry of parsed.data ?? []) {
+                for (const key of [entry.id, entry.name]) {
+                    if (!key) continue;
+                    const list = byId.get(key) ?? [];
+                    list.push(entry);
+                    byId.set(key, list);
+                }
+            }
+        } else {
+            logger.debug(
+                `AI model pricing file not found at ${MODELS_JSON_PATH}; cost calculation will fall back to unknown pricing`
+            );
+        }
+    } catch (error) {
+        logger.warn("Failed to load AI model pricing file", { error });
+    }
+
+    modelsById = byId;
+    return byId;
+}
+
+function stripVendorPrefix(modelId: string): string | null {
+    const idx = modelId.indexOf("/");
+    if (idx === -1 || idx === modelId.length - 1) {
+        return null;
+    }
+    return modelId.slice(idx + 1);
+}
+
+function toPricing(
+    entry: AiModelPricingEntry,
+    approximate: boolean
+): AiModelPricing {
+    return {
+        inputCostPerToken: entry.input_cost_per_token,
+        outputCostPerToken: entry.output_cost_per_token,
+        cacheReadInputTokenCost: entry.cache_read_input_token_cost,
+        outputCostPerReasoningToken: entry.output_cost_per_reasoning_token,
+        approximate
+    };
+}
+
+function findInBuckets(
+    byId: Map<string, AiModelPricingEntry[]>,
+    modelId: string,
+    buckets: string[] | null
+): AiModelPricingEntry | null {
+    const candidates = [modelId, stripVendorPrefix(modelId)].filter(
+        (v): v is string => v != null
+    );
+
+    for (const key of candidates) {
+        const entries = byId.get(key);
+        if (!entries) continue;
+        const match = buckets
+            ? entries.find((e) => buckets.includes(e.provider))
+            : entries[0];
+        if (match) {
+            return match;
+        }
+    }
+    return null;
+}
+
+/**
+ * Looks up per-token pricing for a model, scoped first to the litellm
+ * provider bucket(s) that correspond to our provider type, then falling
+ * back to a global search across all buckets (marked `approximate`) for
+ * provider types that proxy arbitrary underlying models.
+ */
+export function getModelPricing(
+    providerType: AiProviderType,
+    modelId: string | undefined
+): AiModelPricing | null {
+    if (!modelId) {
+        return null;
+    }
+
+    const byId = loadModels();
+    const buckets =
+        providerType === "custom"
+            ? []
+            : PROVIDER_PRICING_BUCKETS[providerType];
+
+    if (buckets && buckets.length > 0) {
+        const scoped = findInBuckets(byId, modelId, buckets);
+        if (scoped) {
+            return toPricing(scoped, false);
+        }
+    }
+
+    const fallback = findInBuckets(byId, modelId, null);
+    if (fallback) {
+        return toPricing(fallback, true);
+    }
+
+    return null;
+}
+
+export type AiCostBreakdown = {
+    promptCost: number;
+    cacheReadCost: number;
+    cacheWriteCost: number;
+    completionCost: number;
+    reasoningCost: number;
+    totalCost: number;
+};
+
+/**
+ * Computes a $ cost breakdown for a usage record given a model's pricing.
+ * Cache writes and reasoning tokens fall back to the normal input/output
+ * rate respectively when the pricing data has no dedicated rate for them
+ * (the models.json schema here has no cache-write field at all, and only
+ * some models report a distinct reasoning rate).
+ */
+export function calculateAiCost(
+    pricing: AiModelPricing | null,
+    usage: AiUsage
+): AiCostBreakdown | null {
+    if (!pricing) {
+        return null;
+    }
+
+    const inputRate = pricing.inputCostPerToken ?? 0;
+    const outputRate = pricing.outputCostPerToken ?? 0;
+    const cacheReadRate = pricing.cacheReadInputTokenCost ?? inputRate;
+    const reasoningRate = pricing.outputCostPerReasoningToken ?? outputRate;
+
+    const promptCost = usage.promptTokens * inputRate;
+    const cacheReadCost = usage.cacheReadTokens * cacheReadRate;
+    const cacheWriteCost = usage.cacheWriteTokens * inputRate;
+    const completionCost = usage.completionTokens * outputRate;
+    const reasoningCost = usage.reasoningTokens * reasoningRate;
+
+    return {
+        promptCost,
+        cacheReadCost,
+        cacheWriteCost,
+        completionCost,
+        reasoningCost,
+        totalCost:
+            promptCost +
+            cacheReadCost +
+            cacheWriteCost +
+            completionCost +
+            reasoningCost
+    };
+}

server/lib/aiUsageExtraction.ts

@@ -0,0 +1,468 @@
+import { encode } from "gpt-tokenizer";
+import type { AiCapability } from "@server/lib/aiCapabilities";
+import logger from "@server/logger";
+
+export type AiUsage = {
+    // Input tokens billed at the normal input rate (i.e. NOT already
+    // covered by cacheReadTokens/cacheWriteTokens below).
+    promptTokens: number;
+    cacheReadTokens: number;
+    cacheWriteTokens: number;
+    // Output tokens billed at the normal output rate (i.e. NOT already
+    // covered by reasoningTokens below).
+    completionTokens: number;
+    reasoningTokens: number;
+    // True when these numbers are our own best-guess estimate (the upstream
+    // response didn't report usage), rather than provider-reported figures.
+    estimated: boolean;
+};
+
+function emptyUsage(): AiUsage {
+    return {
+        promptTokens: 0,
+        cacheReadTokens: 0,
+        cacheWriteTokens: 0,
+        completionTokens: 0,
+        reasoningTokens: 0,
+        estimated: false
+    };
+}
+
+/**
+ * Scans raw (possibly binary-framed, e.g. Bedrock's vnd.amazon.eventstream)
+ * text for `"fieldName":123` occurrences and returns the last value seen for
+ * each field. Used as a best-effort fallback for response shapes we can't
+ * fully parse as JSON/SSE (streaming Bedrock, raw predict passthroughs).
+ */
+function scanNumericFields(
+    text: string,
+    fields: string[]
+): Record<string, number> {
+    const out: Record<string, number> = {};
+    for (const field of fields) {
+        const re = new RegExp(`"${field}"\\s*:\\s*(\\d+)`, "g");
+        let match: RegExpExecArray | null;
+        while ((match = re.exec(text)) !== null) {
+            out[field] = Number(match[1]);
+        }
+    }
+    return out;
+}
+
+function sseDataFrames(text: string): string[] {
+    const frames: string[] = [];
+    for (const rawFrame of text.split(/\r?\n\r?\n/)) {
+        for (const line of rawFrame.split(/\r?\n/)) {
+            if (!line.startsWith("data:")) continue;
+            const data = line.slice("data:".length).trim();
+            if (data && data !== "[DONE]") {
+                frames.push(data);
+            }
+        }
+    }
+    return frames;
+}
+
+function tryParseJson(text: string): any | null {
+    try {
+        return JSON.parse(text);
+    } catch {
+        return null;
+    }
+}
+
+function extractOpenAiChat(text: string, isStream: boolean): AiUsage | null {
+    let usage: any = null;
+
+    if (isStream) {
+        for (const frame of sseDataFrames(text)) {
+            const parsed = tryParseJson(frame);
+            if (parsed?.usage) {
+                usage = parsed.usage;
+            }
+        }
+    } else {
+        usage = tryParseJson(text)?.usage ?? null;
+    }
+
+    if (!usage) {
+        return null;
+    }
+
+    const cacheReadTokens = usage.prompt_tokens_details?.cached_tokens ?? 0;
+    const reasoningTokens =
+        usage.completion_tokens_details?.reasoning_tokens ?? 0;
+
+    return {
+        promptTokens: Math.max(0, (usage.prompt_tokens ?? 0) - cacheReadTokens),
+        cacheReadTokens,
+        cacheWriteTokens: 0,
+        completionTokens: Math.max(
+            0,
+            (usage.completion_tokens ?? 0) - reasoningTokens
+        ),
+        reasoningTokens,
+        estimated: false
+    };
+}
+
+function extractOpenAiResponses(
+    text: string,
+    isStream: boolean
+): AiUsage | null {
+    let usage: any = null;
+
+    if (isStream) {
+        for (const frame of sseDataFrames(text)) {
+            const parsed = tryParseJson(frame);
+            if (parsed?.type === "response.completed" && parsed?.response?.usage) {
+                usage = parsed.response.usage;
+            } else if (parsed?.usage) {
+                usage = parsed.usage;
+            }
+        }
+    } else {
+        const parsed = tryParseJson(text);
+        usage = parsed?.usage ?? parsed?.response?.usage ?? null;
+    }
+
+    if (!usage) {
+        return null;
+    }
+
+    const cacheReadTokens = usage.input_tokens_details?.cached_tokens ?? 0;
+    const reasoningTokens = usage.output_tokens_details?.reasoning_tokens ?? 0;
+
+    return {
+        promptTokens: Math.max(0, (usage.input_tokens ?? 0) - cacheReadTokens),
+        cacheReadTokens,
+        cacheWriteTokens: 0,
+        completionTokens: Math.max(
+            0,
+            (usage.output_tokens ?? 0) - reasoningTokens
+        ),
+        reasoningTokens,
+        estimated: false
+    };
+}
+
+function extractAnthropicMessages(
+    text: string,
+    isStream: boolean
+): AiUsage | null {
+    let inputTokens = 0;
+    let cacheReadTokens = 0;
+    let cacheWriteTokens = 0;
+    let outputTokens = 0;
+    let found = false;
+
+    const applyUsage = (usage: any) => {
+        if (!usage) return;
+        found = true;
+        if (typeof usage.input_tokens === "number") {
+            inputTokens = usage.input_tokens;
+        }
+        if (typeof usage.cache_read_input_tokens === "number") {
+            cacheReadTokens = usage.cache_read_input_tokens;
+        }
+        if (typeof usage.cache_creation_input_tokens === "number") {
+            cacheWriteTokens = usage.cache_creation_input_tokens;
+        }
+        if (typeof usage.output_tokens === "number") {
+            outputTokens = usage.output_tokens;
+        }
+    };
+
+    if (isStream) {
+        for (const frame of sseDataFrames(text)) {
+            const parsed = tryParseJson(frame);
+            if (!parsed) continue;
+            applyUsage(parsed.message?.usage);
+            applyUsage(parsed.usage);
+        }
+    } else {
+        applyUsage(tryParseJson(text)?.usage);
+    }
+
+    if (!found) {
+        return null;
+    }
+
+    return {
+        promptTokens: inputTokens,
+        cacheReadTokens,
+        cacheWriteTokens,
+        completionTokens: outputTokens,
+        // Anthropic bills extended-thinking output at the normal output
+        // rate, so there's no separate reasoning bucket to report.
+        reasoningTokens: 0,
+        estimated: false
+    };
+}
+
+function extractGoogleGenerateContent(
+    text: string,
+    _isStream: boolean
+): AiUsage | null {
+    // Both the plain-JSON-array stream format and the SSE (?alt=sse) format
+    // repeat a cumulative `usageMetadata` object per chunk; the regex scan
+    // below naturally picks up the last (most complete) one either way.
+    const fields = scanNumericFields(text, [
+        "promptTokenCount",
+        "candidatesTokenCount",
+        "cachedContentTokenCount",
+        "thoughtsTokenCount"
+    ]);
+
+    if (fields.promptTokenCount === undefined) {
+        return null;
+    }
+
+    const cacheReadTokens = fields.cachedContentTokenCount ?? 0;
+    const reasoningTokens = fields.thoughtsTokenCount ?? 0;
+
+    return {
+        promptTokens: Math.max(0, fields.promptTokenCount - cacheReadTokens),
+        cacheReadTokens,
+        cacheWriteTokens: 0,
+        completionTokens: fields.candidatesTokenCount ?? 0,
+        reasoningTokens,
+        estimated: false
+    };
+}
+
+function extractBedrockConverse(
+    text: string,
+    _isStream: boolean
+): AiUsage | null {
+    // Non-streaming responses are plain JSON; converse-stream frames the
+    // final `metadata` event's usage object inside binary event-stream
+    // framing, but the JSON text survives intact inside that binary
+    // envelope, so the same field scan works for both.
+    const parsed = tryParseJson(text);
+    const usage = parsed?.usage;
+    if (usage) {
+        const cacheReadTokens = usage.cacheReadInputTokens ?? 0;
+        return {
+            promptTokens: Math.max(0, (usage.inputTokens ?? 0) - cacheReadTokens),
+            cacheReadTokens,
+            cacheWriteTokens: usage.cacheWriteInputTokens ?? 0,
+            completionTokens: usage.outputTokens ?? 0,
+            reasoningTokens: 0,
+            estimated: false
+        };
+    }
+
+    const fields = scanNumericFields(text, [
+        "inputTokens",
+        "outputTokens",
+        "cacheReadInputTokens",
+        "cacheWriteInputTokens"
+    ]);
+    if (fields.inputTokens === undefined) {
+        return null;
+    }
+    const cacheReadTokens = fields.cacheReadInputTokens ?? 0;
+    return {
+        promptTokens: Math.max(0, fields.inputTokens - cacheReadTokens),
+        cacheReadTokens,
+        cacheWriteTokens: fields.cacheWriteInputTokens ?? 0,
+        completionTokens: fields.outputTokens ?? 0,
+        reasoningTokens: 0,
+        estimated: false
+    };
+}
+
+function extractBedrockModelInvoke(
+    text: string,
+    _isStream: boolean,
+    headers: Headers
+): AiUsage | null {
+    // Non-streaming invoke reports counts via response headers regardless
+    // of the underlying model's payload format.
+    const headerInput = headers.get("x-amzn-bedrock-input-token-count");
+    const headerOutput = headers.get("x-amzn-bedrock-output-token-count");
+    if (headerInput !== null || headerOutput !== null) {
+        return {
+            promptTokens: Number(headerInput ?? 0),
+            cacheReadTokens: 0,
+            cacheWriteTokens: 0,
+            completionTokens: Number(headerOutput ?? 0),
+            reasoningTokens: 0,
+            estimated: false
+        };
+    }
+
+    // invoke-with-response-stream has no equivalent headers; the model's
+    // own usage shape (frequently Anthropic-style on Bedrock) is embedded
+    // inside binary event-stream framing, so fall back to a couple of
+    // known field-name shapes via regex.
+    const anthropicStyle = extractAnthropicMessages(text, true);
+    if (anthropicStyle) {
+        return anthropicStyle;
+    }
+
+    const fields = scanNumericFields(text, [
+        "inputTokenCount",
+        "outputTokenCount"
+    ]);
+    if (fields.inputTokenCount === undefined) {
+        return null;
+    }
+    return {
+        promptTokens: fields.inputTokenCount,
+        cacheReadTokens: 0,
+        cacheWriteTokens: 0,
+        completionTokens: fields.outputTokenCount ?? 0,
+        reasoningTokens: 0,
+        estimated: false
+    };
+}
+
+const EXTRACTORS: Record<
+    AiCapability,
+    (text: string, isStream: boolean, headers: Headers) => AiUsage | null
+> = {
+    openai_chat: extractOpenAiChat,
+    openai_responses: extractOpenAiResponses,
+    anthropic_messages: extractAnthropicMessages,
+    gemini_generate_content: extractGoogleGenerateContent,
+    google_generate_content: extractGoogleGenerateContent,
+    // rawPredict is a passthrough to whatever the underlying publisher
+    // model speaks (often Anthropic-shaped on Vertex); try that, then give
+    // up to the token-count estimate.
+    google_raw_predict: (text, isStream) =>
+        extractAnthropicMessages(text, isStream),
+    bedrock_model_invoke: extractBedrockModelInvoke,
+    bedrock_converse: extractBedrockConverse
+};
+
+/**
+ * Attempts to pull provider-reported token usage out of an upstream AI
+ * gateway response. Returns null if the response didn't contain (or we
+ * couldn't find) usage data, in which case callers should fall back to
+ * `estimateUsage`.
+ */
+export function extractUsage(
+    capability: AiCapability,
+    responseText: string,
+    isStream: boolean,
+    headers: Headers
+): AiUsage | null {
+    try {
+        return EXTRACTORS[capability](responseText, isStream, headers);
+    } catch (error) {
+        logger.debug("Failed to extract AI usage from response", {
+            capability,
+            error
+        });
+        return null;
+    }
+}
+
+/**
+ * Best-guess token estimate for when the provider doesn't report usage.
+ * Uses OpenAI's BPE tokenizer as a stand-in for whatever tokenizer the
+ * actual model uses - close enough for an approximate cost figure, not
+ * exact for non-OpenAI models.
+ */
+export function estimateUsage(
+    promptText: string,
+    completionText: string
+): AiUsage {
+    const usage = emptyUsage();
+    usage.estimated = true;
+    try {
+        usage.promptTokens = promptText ? encode(promptText).length : 0;
+    } catch (error) {
+        logger.debug("Failed to estimate prompt tokens", { error });
+    }
+    try {
+        usage.completionTokens = completionText
+            ? encode(completionText).length
+            : 0;
+    } catch (error) {
+        logger.debug("Failed to estimate completion tokens", { error });
+    }
+    return usage;
+}
+
+/**
+ * OpenAI's Chat Completions API only includes a `usage` field in a
+ * streaming response when the request opts in via `stream_options:
+ * {include_usage: true}` - unlike the Responses API, Anthropic, Gemini and
+ * Bedrock, which report usage in a streaming response by default. Returns
+ * whether we need to inject that option ourselves to be able to track cost.
+ */
+export function needsStreamUsageInjection(
+    capability: AiCapability,
+    body: any
+): boolean {
+    return (
+        capability === "openai_chat" &&
+        body?.stream === true &&
+        body?.stream_options?.include_usage !== true
+    );
+}
+
+/**
+ * Returns a shallow-cloned body with `stream_options.include_usage`
+ * injected, for capabilities/requests where `needsStreamUsageInjection`
+ * is true. Leaves the original body untouched.
+ */
+export function withStreamUsageOption(body: any): any {
+    return {
+        ...body,
+        stream_options: { ...body.stream_options, include_usage: true }
+    };
+}
+
+/**
+ * When we injected stream_options.include_usage ourselves (the caller
+ * didn't ask for it), OpenAI appends an extra terminal SSE frame with an
+ * empty `choices: []` array carrying only the usage data. Callers that
+ * don't expect that shape (most minimal SSE parsers assume a non-empty
+ * choices array) shouldn't see it, so it's stripped back out of the bytes
+ * forwarded to the client.
+ */
+export function stripInjectedUsageFrame(sseText: string): string {
+    const parts = sseText.split(/(\r?\n\r?\n)/);
+    let out = "";
+    for (let i = 0; i < parts.length; i += 2) {
+        const frame = parts[i];
+        const separator = parts[i + 1] ?? "";
+        const dataLine = frame
+            .split(/\r?\n/)
+            .find((line) => line.startsWith("data:"));
+        if (dataLine) {
+            const data = dataLine.slice("data:".length).trim();
+            const parsed = data !== "[DONE]" ? tryParseJson(data) : null;
+            if (parsed && Array.isArray(parsed.choices) && parsed.choices.length === 0 && parsed.usage) {
+                continue;
+            }
+        }
+        out += frame + separator;
+    }
+    return out;
+}
+
+/**
+ * Best-effort extraction of the model the upstream provider actually
+ * served, which some gateways/routers echo back and which may differ from
+ * the model the caller requested (e.g. an alias resolving to a dated
+ * snapshot). Falls back to the caller's requested model when absent.
+ */
+export function extractResponseModel(responseText: string): string | null {
+    const match = responseText.match(/"model"\s*:\s*"([^"]+)"/);
+    return match ? match[1] : null;
+}
+
+export function isUsageEmpty(usage: AiUsage): boolean {
+    return (
+        usage.promptTokens === 0 &&
+        usage.cacheReadTokens === 0 &&
+        usage.cacheWriteTokens === 0 &&
+        usage.completionTokens === 0 &&
+        usage.reasoningTokens === 0
+    );
+}

server/routers/aiGateway/pipeline.ts

@@ -19,6 +19,7 @@ import config from "@server/lib/config";
 import { decrypt } from "@server/lib/crypto";
 import {
     AiProviderAuthType,
+    AiProviderType,
     applyAiProviderAuthHeaders,
     applyAiProviderCustomHeaders,
     authTypeRequiresApiKey
@@ -49,6 +50,17 @@ import {
     mostSpecificMatchingAllow
 } from "@server/lib/aiModelKeyMatch";
 import { aiGatewayUpstreamFetch } from "@server/lib/aiGatewayUpstreamFetch";
+import { getModelPricing, calculateAiCost } from "@server/lib/aiModelPricing";
+import {
+    extractUsage,
+    estimateUsage,
+    isUsageEmpty,
+    needsStreamUsageInjection,
+    withStreamUsageOption,
+    stripInjectedUsageFrame,
+    extractResponseModel,
+    type AiUsage
+} from "@server/lib/aiUsageExtraction";
 
 // Short-lived local caches so a burst of requests from the same IP/user
 // doesn't hit the database on every single request. None of this is
@@ -500,6 +512,48 @@ async function selectProvider(
     };
 }
 
+function logAiUsageAndCost(args: {
+    capability: AiCapability;
+    provider: AiProvider;
+    requestedModel: string | undefined;
+    requestBody: unknown;
+    responseText: string;
+    isStream: boolean;
+    headers: Headers;
+}): void {
+    const { capability, provider, requestedModel, requestBody, responseText, isStream, headers } =
+        args;
+
+    let usage: AiUsage | null = extractUsage(
+        capability,
+        responseText,
+        isStream,
+        headers
+    );
+    if (!usage || isUsageEmpty(usage)) {
+        usage = estimateUsage(JSON.stringify(requestBody ?? ""), responseText);
+    }
+
+    const model = extractResponseModel(responseText) ?? requestedModel;
+    const pricing = getModelPricing(provider.type as AiProviderType, model);
+    const cost = calculateAiCost(pricing, usage);
+
+    logger.info("AI gateway request usage", {
+        capability,
+        providerId: provider.providerId,
+        providerType: provider.type,
+        model,
+        estimated: usage.estimated,
+        promptTokens: usage.promptTokens,
+        cacheReadTokens: usage.cacheReadTokens,
+        cacheWriteTokens: usage.cacheWriteTokens,
+        completionTokens: usage.completionTokens,
+        reasoningTokens: usage.reasoningTokens,
+        pricingApproximate: pricing?.approximate ?? null,
+        totalCostUsd: cost?.totalCost ?? null
+    });
+}
+
 export async function handleAiGatewayProxy(
     req: Request,
     res: Response,
@@ -636,14 +690,25 @@ export async function handleAiGatewayProxy(
         applyAiProviderAuthHeaders(headers, authType, apiKey);
         applyRequestUserHeaders(headers, requestUser);
 
-        const body = JSON.stringify(req.body);
+        // OpenAI's Chat Completions API only reports usage in a streaming
+        // response when asked to via stream_options.include_usage - inject
+        // it ourselves when the caller didn't, so we can still track cost,
+        // and strip the extra frame it adds back out of what we forward.
+        const injectedUsageOurselves = needsStreamUsageInjection(
+            capability,
+            req.body
+        );
+        const outboundBody = injectedUsageOurselves
+            ? withStreamUsageOption(req.body)
+            : req.body;
+        const body = JSON.stringify(outboundBody);
 
         logger.debug("AI gateway upstream request", {
             capability,
             url: targetUrl,
             method: "POST",
             headers,
-            body: req.body,
+            body: outboundBody,
             skipTlsVerification: provider.skipTlsVerification
         });
 
@@ -701,11 +766,31 @@ export async function handleAiGatewayProxy(
         if (isStream && upstreamRes.body) {
             res.flushHeaders();
             const reader = upstreamRes.body.getReader();
+            const decoder = new TextDecoder();
+            let fullText = "";
+            // Frame-boundary buffer, only used when we need to filter the
+            // usage-only frame we injected out of what reaches the client.
+            let sseCarry = "";
             try {
                 while (!abortController.signal.aborted) {
                     const { done, value } = await reader.read();
                     if (done) break;
-                    res.write(value);
+                    const chunkText = decoder.decode(value, { stream: true });
+                    fullText += chunkText;
+                    if (injectedUsageOurselves) {
+                        sseCarry += chunkText;
+                        const lastBoundary = sseCarry.lastIndexOf("\n\n");
+                        if (lastBoundary !== -1) {
+                            const toEmit = sseCarry.slice(0, lastBoundary + 2);
+                            sseCarry = sseCarry.slice(lastBoundary + 2);
+                            res.write(stripInjectedUsageFrame(toEmit));
+                        }
+                    } else {
+                        res.write(value);
+                    }
+                }
+                if (injectedUsageOurselves && sseCarry) {
+                    res.write(stripInjectedUsageFrame(sseCarry));
                 }
             } finally {
                 await reader.cancel().catch(() => {});
@@ -714,11 +799,31 @@ export async function handleAiGatewayProxy(
             if (!res.writableEnded) {
                 res.end();
             }
+            if (!abortController.signal.aborted) {
+                logAiUsageAndCost({
+                    capability,
+                    provider,
+                    requestedModel,
+                    requestBody: req.body,
+                    responseText: fullText,
+                    isStream: true,
+                    headers: upstreamRes.headers
+                });
+            }
             return;
         }
 
         res.off("close", onClientClose);
         const text = await upstreamRes.text();
+        logAiUsageAndCost({
+            capability,
+            provider,
+            requestedModel,
+            requestBody: req.body,
+            responseText: text,
+            isStream: false,
+            headers: upstreamRes.headers
+        });
         return res.send(text);
     } catch (error) {
         logger.error(error);