From 665120eb9fb2bc2dcb15af9822f268fdef042fa4 Mon Sep 17 00:00:00 2001 From: hugocasa Date: Thu, 13 Aug 2026 19:37:34 +0200 Subject: [PATCH] fix: address review findings on AI cost tracking --- ...6b109ee9f016af5d6b9128b7c3a16d560615.json} | 7 +- backend/windmill-api/openapi.yaml | 23 ++++- backend/windmill-api/src/ai.rs | 79 ++++++++++++++++-- backend/windmill-api/src/workspaces.rs | 2 + .../copilot/chat/AIChatManager.svelte.ts | 62 +++++++------- .../lib/components/copilot/chat/chatLoop.ts | 29 +++---- .../src/lib/components/copilot/modelConfig.ts | 6 ++ .../lib/components/copilot/modelPricing.ts | 64 +++++++++----- .../workspaceSettings/AISettings.svelte | 6 +- .../workspaceSettings/AiUsagePanel.svelte | 25 ++++-- .../workspaceSettings/ModelPricing.svelte | 41 +++++---- frontend/src/lib/utils/aiUsageReporter.ts | Bin 5254 -> 5261 bytes 12 files changed, 233 insertions(+), 111 deletions(-) rename backend/.sqlx/{query-43ec1c4adf6e453848b532bf3fab5de8825d247ea0b337007ee43c9cdffac83d.json => query-9c4c5ced0473e5c6e4a08995e8ec6b109ee9f016af5d6b9128b7c3a16d560615.json} (86%) diff --git a/backend/.sqlx/query-43ec1c4adf6e453848b532bf3fab5de8825d247ea0b337007ee43c9cdffac83d.json b/backend/.sqlx/query-9c4c5ced0473e5c6e4a08995e8ec6b109ee9f016af5d6b9128b7c3a16d560615.json similarity index 86% rename from backend/.sqlx/query-43ec1c4adf6e453848b532bf3fab5de8825d247ea0b337007ee43c9cdffac83d.json rename to backend/.sqlx/query-9c4c5ced0473e5c6e4a08995e8ec6b109ee9f016af5d6b9128b7c3a16d560615.json index efda9b2a2b..2f747a12ec 100644 --- a/backend/.sqlx/query-43ec1c4adf6e453848b532bf3fab5de8825d247ea0b337007ee43c9cdffac83d.json +++ b/backend/.sqlx/query-9c4c5ced0473e5c6e4a08995e8ec6b109ee9f016af5d6b9128b7c3a16d560615.json @@ -1,6 +1,6 @@ { "db_name": "PostgreSQL", - "query": "SELECT\n (CASE $3::text\n WHEN 'day' THEN day::text\n WHEN 'user' THEN email\n WHEN 'session' THEN session_id\n ELSE ''\n END) AS \"key!\",\n provider AS \"provider!\",\n model AS \"model!\",\n SUM(input_tokens)::bigint AS \"input_tokens!\",\n SUM(cache_read_tokens)::bigint AS \"cache_read_tokens!\",\n SUM(cache_write_tokens)::bigint AS \"cache_write_tokens!\",\n SUM(output_tokens)::bigint AS \"output_tokens!\",\n SUM(reported_cost_nano_usd)::bigint AS \"reported_cost_nano_usd\",\n SUM(requests)::bigint AS \"requests!\"\n FROM ai_token_usage\n WHERE workspace_id = $1 AND day >= CURRENT_DATE - $2::int\n GROUP BY 1, provider, model\n ORDER BY 1 DESC, SUM(input_tokens + output_tokens) DESC\n LIMIT 1000", + "query": "SELECT\n (CASE $3::text\n WHEN 'day' THEN day::text\n WHEN 'user' THEN email\n WHEN 'session' THEN session_id\n ELSE ''\n END) AS \"key!\",\n provider AS \"provider!\",\n model AS \"model!\",\n SUM(input_tokens)::bigint AS \"input_tokens!\",\n SUM(cache_read_tokens)::bigint AS \"cache_read_tokens!\",\n SUM(cache_write_tokens)::bigint AS \"cache_write_tokens!\",\n SUM(output_tokens)::bigint AS \"output_tokens!\",\n SUM(reported_cost_nano_usd)::bigint AS \"reported_cost_nano_usd\",\n SUM(requests)::bigint AS \"requests!\"\n FROM ai_token_usage\n WHERE workspace_id = $1 AND day > CURRENT_DATE - $2::int\n GROUP BY 1, provider, model\n ORDER BY SUM(input_tokens + cache_read_tokens + cache_write_tokens + output_tokens) DESC\n LIMIT $4", "describe": { "columns": [ { @@ -53,7 +53,8 @@ "Left": [ "Text", "Int4", - "Text" + "Text", + "Int8" ] }, "nullable": [ @@ -68,5 +69,5 @@ null ] }, - "hash": "43ec1c4adf6e453848b532bf3fab5de8825d247ea0b337007ee43c9cdffac83d" + "hash": "9c4c5ced0473e5c6e4a08995e8ec6b109ee9f016af5d6b9128b7c3a16d560615" } diff --git a/backend/windmill-api/openapi.yaml b/backend/windmill-api/openapi.yaml index be6e1fd0d2..85184d4e3c 100644 --- a/backend/windmill-api/openapi.yaml +++ b/backend/windmill-api/openapi.yaml @@ -11886,9 +11886,18 @@ paths: content: application/json: schema: - type: array - items: - $ref: "#/components/schemas/AITokenUsageBucket" + type: object + required: + - buckets + - truncated + properties: + buckets: + type: array + items: + $ref: "#/components/schemas/AITokenUsageBucket" + truncated: + type: boolean + description: more buckets matched than were returned, so summing them under-reports /w/{workspace}/ai_skills/list: get: @@ -26062,12 +26071,20 @@ components: properties: input: type: number + minimum: 0 + maximum: 1000 output: type: number + minimum: 0 + maximum: 1000 cache_read: type: number + minimum: 0 + maximum: 1000 cache_write: type: number + minimum: 0 + maximum: 1000 required: - input - output diff --git a/backend/windmill-api/src/ai.rs b/backend/windmill-api/src/ai.rs index 24b7987acc..4676180816 100644 --- a/backend/windmill-api/src/ai.rs +++ b/backend/windmill-api/src/ai.rs @@ -6,7 +6,7 @@ use axum::routing::get; use axum::Json; use axum::{ body::Bytes, - extract::{Path, Query}, + extract::{DefaultBodyLimit, Path, Query}, response::IntoResponse, routing::post, Extension, Router, @@ -441,7 +441,39 @@ pub struct ModelPriceOverride { pub cache_write: Option, } +/// Far above any real per-million-token rate, so a value beyond it is a unit +/// mistake rather than a price. The floor matters more: a negative rate would make +/// spend subtract, and NaN/infinity would poison every total derived from it. +pub const MAX_MODEL_RATE: f64 = 1000.0; + +impl ModelPriceOverride { + pub fn validate(&self, key: &str) -> Result<()> { + for (field, rate) in [ + ("input", Some(self.input)), + ("output", Some(self.output)), + ("cache_read", self.cache_read), + ("cache_write", self.cache_write), + ] { + let Some(rate) = rate else { continue }; + if !rate.is_finite() || rate < 0.0 || rate > MAX_MODEL_RATE { + return Err(Error::BadRequest(format!( + "Price override for {}: {} must be between 0 and {}", + key, field, MAX_MODEL_RATE + ))); + } + } + Ok(()) + } +} + impl AIConfig { + pub fn validate_model_pricing(&self) -> Result<()> { + for (key, price) in self.model_pricing.iter().flatten() { + price.validate(key)?; + } + Ok(()) + } + pub fn has_providers(&self) -> bool { self.providers .as_ref() @@ -456,7 +488,16 @@ pub fn global_service() -> Router { pub fn workspaced_service() -> Router { let router = Router::new() .route("/proxy/{*ai}", post(proxy).get(proxy)) - .route("/usage", post(record_ai_usage).get(list_ai_usage)); + .route( + "/usage", + post(record_ai_usage) + .get(list_ai_usage) + // The handler caps how many events it *stores*, but Json deserializes + // the whole array first — without a body limit an authenticated member + // could make the server allocate and parse an arbitrarily large one. + // Sized well above a full batch of the shape below. + .layer(DefaultBodyLimit::max(AI_USAGE_BODY_LIMIT)), + ); #[cfg(feature = "bedrock")] let router = router.route("/check_bedrock_credentials", get(check_bedrock_credentials)); @@ -492,6 +533,8 @@ struct RecordAIUsagePayload { } const MAX_AI_USAGE_EVENTS: usize = 50; +/// 64 KiB — a 50-event batch is a few kB even with the longest model ids. +const AI_USAGE_BODY_LIMIT: usize = 64 * 1024; /// Well above any single conversation and far below an i64 overflow, so a client /// bug caps out at one absurd row instead of poisoning the running total. const MAX_TOKENS_PER_EVENT: i64 = 100_000_000; @@ -636,12 +679,24 @@ struct AITokenUsageBucket { requests: i64, } +/// Grouping by session (or by day over a long range) can produce more buckets than +/// a table is worth rendering, so the listing is capped. `truncated` says so +/// explicitly — a caller that sums the rows into a total must be able to tell that +/// the total is partial rather than silently under-reporting spend. +#[derive(Serialize)] +struct AITokenUsageListing { + buckets: Vec, + truncated: bool, +} + +const AI_USAGE_MAX_BUCKETS: i64 = 1000; + async fn list_ai_usage( authed: ApiAuthed, Extension(db): Extension, Path(w_id): Path, Query(query): Query, -) -> Result>> { +) -> Result> { require_admin(authed.is_admin, &authed.username)?; let days = query.days.unwrap_or(30).clamp(1, 365); @@ -653,7 +708,9 @@ async fn list_ai_usage( ))); } - let rows = sqlx::query_as!( + // Fetch one past the cap to detect truncation, and order by spend so a capped + // listing keeps the buckets worth looking at rather than an arbitrary slice. + let mut rows = sqlx::query_as!( AITokenUsageBucket, r#"SELECT (CASE $3::text @@ -671,18 +728,22 @@ async fn list_ai_usage( SUM(reported_cost_nano_usd)::bigint AS "reported_cost_nano_usd", SUM(requests)::bigint AS "requests!" FROM ai_token_usage - WHERE workspace_id = $1 AND day >= CURRENT_DATE - $2::int + WHERE workspace_id = $1 AND day > CURRENT_DATE - $2::int GROUP BY 1, provider, model - ORDER BY 1 DESC, SUM(input_tokens + output_tokens) DESC - LIMIT 1000"#, + ORDER BY SUM(input_tokens + cache_read_tokens + cache_write_tokens + output_tokens) DESC + LIMIT $4"#, &w_id, days, - group_by + group_by, + AI_USAGE_MAX_BUCKETS + 1 ) .fetch_all(&db) .await?; - Ok(Json(rows)) + let truncated = rows.len() as i64 > AI_USAGE_MAX_BUCKETS; + rows.truncate(AI_USAGE_MAX_BUCKETS as usize); + + Ok(Json(AITokenUsageListing { buckets: rows, truncated })) } /// Check if AWS Bedrock credentials are available from environment variables. diff --git a/backend/windmill-api/src/workspaces.rs b/backend/windmill-api/src/workspaces.rs index 1e26a20e05..a46b647d76 100644 --- a/backend/windmill-api/src/workspaces.rs +++ b/backend/windmill-api/src/workspaces.rs @@ -108,6 +108,8 @@ async fn edit_copilot_config( } } + ai_config.validate_model_pricing()?; + let mut tx = db.begin().await?; sqlx::query!( diff --git a/frontend/src/lib/components/copilot/chat/AIChatManager.svelte.ts b/frontend/src/lib/components/copilot/chat/AIChatManager.svelte.ts index e27508aad3..962b4218e3 100644 --- a/frontend/src/lib/components/copilot/chat/AIChatManager.svelte.ts +++ b/frontend/src/lib/components/copilot/chat/AIChatManager.svelte.ts @@ -56,7 +56,7 @@ import { import { dfs } from '$lib/components/flows/previousResults' import { SvelteMap, SvelteSet } from 'svelte/reactivity' import { createLongHash } from '$lib/editorLangUtils' -import type { UserDraftItemKind } from '$lib/gen' +import type { AIProvider, UserDraftItemKind } from '$lib/gen' import { maskKey } from '$lib/components/sessions/modifiedItemsMask' import { getStringError } from './utils' import { type PasteAttachment } from './pasteTokens' @@ -103,6 +103,7 @@ import { addModelTokenUsage, billedTokens, normalizeContextUsage, + type ChatTokenUsage, type ModelTokenUsageTotals } from './tokenUsage' import { logAiUsage } from '$lib/utils/aiUsageReporter' @@ -663,33 +664,27 @@ export class AIChatManager { return Object.keys(snapshot).length > 0 ? snapshot : undefined } - /** Fold a completed turn's usage into the conversation's running spend and - * report it for the workspace usage view. Only token counts leave the browser: - * rates are applied when the usage is read, so a corrected price also corrects - * everything already recorded. */ - private recordUsage(byModel: ModelTokenUsageTotals | undefined) { - // Accounting must never take a turn down with it: a path that reports no - // per-model breakdown simply records nothing. - for (const entry of Object.values(byModel ?? {})) { - this.usageByModel = addModelTokenUsage( - this.usageByModel, - entry.provider, - entry.model, - entry.usage - ) - const tokens = billedTokens(entry.usage) - logAiUsage({ - provider: entry.provider, - model: entry.model, - sessionId: this.sessionId, - inputTokens: tokens.input, - cacheReadTokens: tokens.cacheRead, - cacheWriteTokens: tokens.cacheWrite, - outputTokens: tokens.output, - costUsd: entry.usage.cost, - workspace: this.operatingWorkspace - }) - } + /** Fold one completed provider response into the conversation's running spend + * and report it for the workspace usage view. Called per response rather than + * per turn: a tool loop makes several, each separately billed, and a turn that + * fails partway through has still spent everything up to that point. + * + * Only token counts leave the browser — rates are applied when the usage is + * read, so a corrected price also corrects everything already recorded. */ + private recordUsage(usage: ChatTokenUsage, provider: AIProvider, model: string) { + this.usageByModel = addModelTokenUsage(this.usageByModel, provider, model, usage) + const tokens = billedTokens(usage) + logAiUsage({ + provider, + model, + sessionId: this.sessionId, + inputTokens: tokens.input, + cacheReadTokens: tokens.cacheRead, + cacheWriteTokens: tokens.cacheWrite, + outputTokens: tokens.output, + costUsd: usage.cost, + workspace: this.operatingWorkspace + }) } // Serialized, snapshot-at-write-time persistence: two rapid dock actions @@ -2384,6 +2379,14 @@ export class AIChatManager { } return undefined }, + onUsage: (usage, modelProvider) => { + // Accounting must never take a turn down with it. + try { + this.recordUsage(usage, modelProvider.provider, modelProvider.model) + } catch (e) { + console.error('Failed to record AI usage', e) + } + }, onBeforeIteration: async (tools, _helpers, modelProvider) => { this.lastIterationModel = modelProvider for (const tool of tools) { @@ -2393,9 +2396,6 @@ export class AIChatManager { } } }) - if (result.tokenUsage.total > 0) { - this.recordUsage(result.tokenUsageByModel) - } if (this.isSessionChat && this.sessionId && result.tokenUsage.total > 0) { logFeatureUsage('ai_session', 'tokens', { entityId: this.sessionId, diff --git a/frontend/src/lib/components/copilot/chat/chatLoop.ts b/frontend/src/lib/components/copilot/chat/chatLoop.ts index ceb9d421e7..e8530de3f1 100644 --- a/frontend/src/lib/components/copilot/chat/chatLoop.ts +++ b/frontend/src/lib/components/copilot/chat/chatLoop.ts @@ -21,13 +21,7 @@ import { } from './openai-responses' import type { Tool, ToolCallbacks } from './shared' import { sanitizeToolCallArguments } from './toolCallArguments' -import { - addChatTokenUsage, - addModelTokenUsage, - emptyChatTokenUsage, - type ChatTokenUsage, - type ModelTokenUsageTotals -} from './tokenUsage' +import { addChatTokenUsage, emptyChatTokenUsage, type ChatTokenUsage } from './tokenUsage' export interface ChatClients { openai: OpenAI @@ -83,14 +77,17 @@ export interface ChatLoopConfig { helpers: any, modelProvider: ReasoningProviderModel ) => Promise + /** Fired for each completed provider response, before the loop continues. The + * loop can fail or be aborted at any iteration, so spend has to be handed over + * as it happens — a callback only at the end would discard everything the + * earlier iterations were already billed for. */ + onUsage?: (usage: ChatTokenUsage, modelProvider: ReasoningProviderModel) => void } export interface ChatLoopResult { addedMessages: ChatCompletionMessageParam[] - /** Sum of usage across all loop iterations (suitable for cost accounting). */ + /** Sum of usage across all loop iterations. */ tokenUsage: ChatTokenUsage - /** The same usage split per model, so a turn that switched model prices correctly. */ - tokenUsageByModel: ModelTokenUsageTotals lastIterationUsage: ChatTokenUsage | null hitMaxIterations: boolean } @@ -333,7 +330,6 @@ export async function runChatLoop(config: ChatLoopConfig): Promise { tokenUsage = addChatTokenUsage(tokenUsage, usage) - if (iterationModel) { - tokenUsageByModel = addModelTokenUsage( - tokenUsageByModel, - iterationModel.provider, - iterationModel.model, - usage - ) + if (usage && iterationModel) { + config.onUsage?.(usage, iterationModel) } // Some providers/paths report no usage (prompt 0); keep the last real one. if (usage && usage.prompt > 0) { @@ -594,5 +585,5 @@ export async function runChatLoop(config: ChatLoopConfig): Promise(entries: [name: string, value: T][]): [Reg * The `provider:model` key the workspace AI settings use for their per-model maps * (`max_tokens_per_model`, `model_pricing`). A bare model id is not enough: the * same id can be served by more than one provider at different rates. + * + * Matched exactly, unlike the fuzzy tables above. Those tables generalize across + * every route to one model on purpose; a per-model *setting* must not, or an + * admin could not give two variants of a family different values — and the key is + * built from the exact id the provider config lists, which is the same string the + * chat sends. */ export function modelKey(provider: AIProvider | string, model: string): string { return `${provider}:${model}` diff --git a/frontend/src/lib/components/copilot/modelPricing.ts b/frontend/src/lib/components/copilot/modelPricing.ts index 8424047887..99d3ce8346 100644 --- a/frontend/src/lib/components/copilot/modelPricing.ts +++ b/frontend/src/lib/components/copilot/modelPricing.ts @@ -24,10 +24,12 @@ export type PricedTokens = { output: number } -// Rates that fall out of the input rate unless a provider prices them separately. -// Anthropic reads a cached prefix at a tenth of the input rate and writes one at -// 1.25x (5-minute TTL, the default the chat uses). Providers whose caching is -// automatic never report a cache write, so their write rate is unused. +// Fallbacks for entries that do not price their cache separately: Anthropic reads a +// cached prefix at a tenth of the input rate and writes one at 1.25x (5-minute TTL, +// the default the chat uses). The read ratio is NOT universal — OpenAI and Google +// discount a cached read far less — so every non-Anthropic entry below states its own +// `cacheRead` rather than inheriting this. Providers whose caching is automatic never +// report a cache write, so their write rate is unused. const CACHE_READ_RATIO = 0.1 const CACHE_WRITE_RATIO = 1.25 @@ -72,25 +74,26 @@ const MODEL_PRICES: [name: string, price: PriceEntry][] = [ ['claude-opus', { input: 5, output: 25 }], ['claude-sonnet', { input: 3, output: 15 }], ['claude-haiku', { input: 1, output: 5 }], - // OpenAI — cached input is a tenth of input, and there is no separate charge - // for writing the cache, so the write rate never applies (the OpenAI usage - // parsers report no cache-write tokens). The -mini/-nano entries must precede + // OpenAI — the cached-input discount varies by family (a tenth on gpt-5, a + // quarter on 4.1 and the o-series, half on 4o), so each entry carries its own + // rate. There is no charge for writing the cache and no usage field reporting + // one, so the write rate never applies. The -mini/-nano entries must precede // the family entry, which would otherwise claim them. - ['gpt-5-mini', { input: 0.25, output: 2 }], - ['gpt-5-nano', { input: 0.05, output: 0.4 }], - ['gpt-5', { input: 1.25, output: 10 }], - ['gpt-4.1-mini', { input: 0.4, output: 1.6 }], - ['gpt-4.1-nano', { input: 0.1, output: 0.4 }], - ['gpt-4.1', { input: 2, output: 8 }], - ['gpt-4o-mini', { input: 0.15, output: 0.6 }], - ['gpt-4o', { input: 2.5, output: 10 }], - ['o4-mini', { input: 1.1, output: 4.4 }], - ['o3-mini', { input: 1.1, output: 4.4 }], - ['o3', { input: 2, output: 8 }], - // Google - ['gemini-2.5-flash-lite', { input: 0.1, output: 0.4 }], - ['gemini-2.5-flash', { input: 0.3, output: 2.5 }], - ['gemini-2.5-pro', { input: 1.25, output: 10 }] + ['gpt-5-mini', { input: 0.25, output: 2, cacheRead: 0.025 }], + ['gpt-5-nano', { input: 0.05, output: 0.4, cacheRead: 0.005 }], + ['gpt-5', { input: 1.25, output: 10, cacheRead: 0.125 }], + ['gpt-4.1-mini', { input: 0.4, output: 1.6, cacheRead: 0.1 }], + ['gpt-4.1-nano', { input: 0.1, output: 0.4, cacheRead: 0.025 }], + ['gpt-4.1', { input: 2, output: 8, cacheRead: 0.5 }], + ['gpt-4o-mini', { input: 0.15, output: 0.6, cacheRead: 0.075 }], + ['gpt-4o', { input: 2.5, output: 10, cacheRead: 1.25 }], + ['o4-mini', { input: 1.1, output: 4.4, cacheRead: 0.275 }], + ['o3-mini', { input: 1.1, output: 4.4, cacheRead: 0.55 }], + ['o3', { input: 2, output: 8, cacheRead: 0.5 }], + // Google — a cached read is a quarter of input across the 2.5 family + ['gemini-2.5-flash-lite', { input: 0.1, output: 0.4, cacheRead: 0.025 }], + ['gemini-2.5-flash', { input: 0.3, output: 2.5, cacheRead: 0.075 }], + ['gemini-2.5-pro', { input: 1.25, output: 10, cacheRead: 0.31 }] ] const MODEL_PRICE_MATCHERS = buildModelMatchers( @@ -116,12 +119,27 @@ export function getKnownModelPrice(model: string): ModelPrice | undefined { * rate, so an admin who only knows their input/output pricing does not have to * invent the other two. */ +/** A rate that would make spend negative, infinite or NaN is not a price. The API + * validates what it stores, but an instance-level config is written as an untyped + * settings blob, so the reader refuses bad values rather than rendering nonsense. */ +function isUsableRate(rate: number | undefined): boolean { + return rate === undefined || (Number.isFinite(rate) && rate >= 0) +} + export function resolveModelPrice( provider: AIProvider | string, model: string, overrides: Record | undefined ): ResolvedModelPrice | undefined { - const override = overrides?.[modelKey(provider, model)] + const candidate = overrides?.[modelKey(provider, model)] + const override = + candidate && + isUsableRate(candidate.input) && + isUsableRate(candidate.output) && + isUsableRate(candidate.cache_read) && + isUsableRate(candidate.cache_write) + ? candidate + : undefined if (override) { return { source: 'override', diff --git a/frontend/src/lib/components/workspaceSettings/AISettings.svelte b/frontend/src/lib/components/workspaceSettings/AISettings.svelte index e5ed2c9f09..88a918ff9e 100644 --- a/frontend/src/lib/components/workspaceSettings/AISettings.svelte +++ b/frontend/src/lib/components/workspaceSettings/AISettings.svelte @@ -588,8 +588,6 @@ - -