From 168572dd76c36ba6350d726dc878502c571fc080 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Thu, 8 Oct 2026 11:08:38 -0400 Subject: [PATCH] Q63: model calls are priced and agent and department budgets limit real spend Co-Authored-By: Claude Opus 5.5 --- src/main/app/ipc-handlers.ts | 1 + src/main/database/migrations/072-spend.ts | 18 ++ src/main/database/schema.ts | 2 + src/main/services/heartbeat-scheduler.ts | 10 ++ src/main/services/provider-request-store.ts | 21 ++- src/main/services/provider-response-runner.ts | 26 +++ src/main/services/provider-service.ts | 23 ++- .../services/provider-tool-turns/router.ts | 4 +- .../services/provider-tool-turns/turns.ts | 8 +- src/main/services/run-control.ts | 11 +- src/main/services/spend.ts | 155 ++++++++++++++++++ src/preload/index.ts | 1 + src/renderer/src/i18n/zh.ts | 8 + .../src/pages/settings/ProvidersPanel.tsx | 47 ++++++ src/shared/commands/provider.ts | 9 + src/shared/desktop-api.ts | 2 + src/shared/domain.ts | 4 + src/shared/ipc.ts | 1 + tests/integration/spend-budget.test.ts | 84 ++++++++++ tests/unit/domain-services.test.ts | 2 +- tests/unit/spend.test.ts | 29 ++++ 21 files changed, 454 insertions(+), 12 deletions(-) create mode 100644 src/main/database/migrations/072-spend.ts create mode 100644 src/main/services/spend.ts create mode 100644 tests/integration/spend-budget.test.ts create mode 100644 tests/unit/spend.test.ts diff --git a/src/main/app/ipc-handlers.ts b/src/main/app/ipc-handlers.ts index 179b40a3..5f5647ff 100644 --- a/src/main/app/ipc-handlers.ts +++ b/src/main/app/ipc-handlers.ts @@ -211,6 +211,7 @@ export function registerIpcHandlers(deps: IpcDeps): void { ipcMain.handle(IPC_CHANNELS.saveProviderCredential, (_event, input) => mutate(() => providers.saveCredential(input))); ipcMain.handle(IPC_CHANNELS.revokeProviderCredential, (_event, input) => mutate(() => providers.revokeCredential(input))); ipcMain.handle(IPC_CHANNELS.updateProviderPolicy, (_event, input) => mutate(() => providers.updatePolicy(input))); + ipcMain.handle(IPC_CHANNELS.setProviderPrices, (_event, input) => mutate(() => providers.setPrices(input))); ipcMain.handle(IPC_CHANNELS.checkProviderHealth, async (_event, input) => { try { await providers.checkHealth(input.orgId, input.providerId); diff --git a/src/main/database/migrations/072-spend.ts b/src/main/database/migrations/072-spend.ts new file mode 100644 index 00000000..49fa4dd5 --- /dev/null +++ b/src/main/database/migrations/072-spend.ts @@ -0,0 +1,18 @@ +export const spendSchema = ` + alter table providers add column input_cents_per_mtok real; + alter table providers add column cached_input_cents_per_mtok real; + alter table providers add column output_cents_per_mtok real; + + alter table provider_requests add column agent_id text references agents(id) on delete set null; + alter table provider_requests add column task_run_id text references task_runs(id) on delete set null; + alter table provider_requests add column department text; + alter table provider_requests add column input_tokens integer; + alter table provider_requests add column output_tokens integer; + alter table provider_requests add column cached_input_tokens integer; + alter table provider_requests add column reasoning_tokens integer; + alter table provider_requests add column cost_cents real; + create index if not exists idx_provider_requests_agent_created on provider_requests(agent_id, created_at); + + insert into schema_migrations(version, name, applied_at) + values (72, 'spend', datetime('now')); +`; diff --git a/src/main/database/schema.ts b/src/main/database/schema.ts index 0b61c513..6a57680a 100644 --- a/src/main/database/schema.ts +++ b/src/main/database/schema.ts @@ -70,6 +70,7 @@ import { taskReviewLinkSchema } from "./migrations/068-task-review-link"; import { goalStallSchema } from "./migrations/069-goal-stall"; import { runActionsSchema } from "./migrations/070-run-actions"; import { backgroundPollFailuresSchema } from "./migrations/071-background-poll-failures"; +import { spendSchema } from "./migrations/072-spend"; interface Migration { version: number; @@ -149,6 +150,7 @@ const MIGRATIONS: Migration[] = [ { version: 69, name: "goal_stall", sql: goalStallSchema }, { version: 70, name: "run_actions", sql: runActionsSchema }, { version: 71, name: "background_poll_failures", sql: backgroundPollFailuresSchema }, + { version: 72, name: "spend", sql: spendSchema }, ]; const SCHEMA_VERSION = MIGRATIONS[MIGRATIONS.length - 1].version; diff --git a/src/main/services/heartbeat-scheduler.ts b/src/main/services/heartbeat-scheduler.ts index 35c218f4..9f060b81 100644 --- a/src/main/services/heartbeat-scheduler.ts +++ b/src/main/services/heartbeat-scheduler.ts @@ -4,6 +4,7 @@ import { createId } from "../../shared/ids"; import type { ProviderCooldowns } from "./provider-cooldowns"; import { assertExecutionApproval, createIsolationContext, loadWorkspace, requireAgentProvider, requireAssignedAgent, requireTask } from "./task-run-preparation"; import { AuditService } from "./audit-service"; +import { budgetBlock, isPriced, providerPrices } from "./spend"; /** Sessions that made no progress before a task stalls. */ const TASK_RETRY_BUDGET = 3; @@ -14,6 +15,8 @@ const TASK_SESSION_CEILING = 12; export class HeartbeatScheduler { /** Providers currently parked on their monthly quota, for transition logging. */ private quotaBlockedProviders = new Map(); + /** Agents parked on a spent monthly budget, for log-once. */ + private readonly budgetParkedAgents = new Set(); constructor( private readonly db: DatabaseSync, @@ -61,6 +64,13 @@ export class HeartbeatScheduler { for (const row of rows) { if (cooling.has(row.provider_id)) continue; if (row.interval_sec > 0 && pacedAgents.has(row.agent_id)) continue; + const overBudget = budgetBlock(this.db, row.agent_id, isPriced(providerPrices(this.db, row.provider_id)), new Date(this.now())); + if (overBudget) { + if (!this.budgetParkedAgents.has(row.agent_id)) console.info(`[task-queue] ${overBudget} Its work waits.`); + this.budgetParkedAgents.add(row.agent_id); + continue; + } + this.budgetParkedAgents.delete(row.agent_id); const sessions = (sessionsOf.all(row.task_id) as Array<{ error_message: string | null; progressed: number }>) .filter((run) => !isTransientProviderFailure(run.error_message)); const stuck = sessions.filter((run) => run.progressed === 0).length; diff --git a/src/main/services/provider-request-store.ts b/src/main/services/provider-request-store.ts index 310a0d63..91bc9935 100644 --- a/src/main/services/provider-request-store.ts +++ b/src/main/services/provider-request-store.ts @@ -1,6 +1,7 @@ import type { DatabaseSync } from "node:sqlite"; import { PROVIDER_EXECUTION_MODE, PROVIDER_REQUEST_STATE, type ProviderExecutionMode } from "../../shared/domain"; import { createId } from "../../shared/ids"; +import type { TokenUsage } from "./spend"; import { currentProviderUsageWindow } from "../../shared/provider-runtime-domain"; export interface ProviderRequestOwner { @@ -19,6 +20,11 @@ export interface CreateProviderRequestInput { inputCharCount: number; executionMode: ProviderExecutionMode; createdAt: string; + /** The agent and task run the call works for, when it is part of a run. */ + agentId?: string | null; + taskRunId?: string | null; + /** The agent's department at call time, which its spend is charged to. */ + department?: string | null; } export interface ProviderRequestProgressInput { @@ -104,9 +110,9 @@ export class ProviderRequestStore { id, org_id, provider_id, provider_key, workspace_id, model, prompt_sha256, response_text, raw_response_json, state, error_message, input_char_count, output_char_count, execution_mode, policy_decision, external_response_id, - stream_event_count, cancelled_at, created_at, completed_at + stream_event_count, cancelled_at, created_at, completed_at, agent_id, task_run_id, department ) - values (?, ?, ?, ?, ?, ?, ?, null, null, ?, null, ?, 0, ?, ?, null, 0, null, ?, null) + values (?, ?, ?, ?, ?, ?, ?, null, null, ?, null, ?, 0, ?, ?, null, 0, null, ?, null, ?, ?, ?) `).run( requestId, input.orgId, @@ -120,10 +126,21 @@ export class ProviderRequestStore { input.executionMode, policyDecision(input.provider, input.workspaceId), input.createdAt, + input.agentId ?? null, + input.taskRunId ?? null, + input.department ?? null, ); return requestId; } + /** Records the tokens a call used and what it cost (null when its provider has no prices). */ + recordSpend(requestId: string, usage: TokenUsage, costCents: number | null): void { + this.db.prepare(` + update provider_requests set input_tokens = ?, output_tokens = ?, cached_input_tokens = ?, reasoning_tokens = ?, cost_cents = ? + where id = ? + `).run(usage.inputTokens, usage.outputTokens, usage.cachedInputTokens, usage.reasoningTokens, costCents, requestId); + } + /** Terminal states (succeeded, failed, cancelled) are final: these transitions return false once one is reached. */ markStreaming(input: ProviderRequestProgressInput): boolean { return this.db.prepare(` diff --git a/src/main/services/provider-response-runner.ts b/src/main/services/provider-response-runner.ts index 00eb9756..5dec13cb 100644 --- a/src/main/services/provider-response-runner.ts +++ b/src/main/services/provider-response-runner.ts @@ -16,6 +16,8 @@ import { hashPrompt, readProviderResponseId, requireStreamText } from "./provide import { type AgentTurnMessage, type AgentTurnResult, type ToolDefinition, type ToolTurnFn } from "./provider-tool-turn"; import { selectToolTurn, supportsToolTurn } from "./provider-tool-turns/router"; import { resolveTurnCredential } from "./provider-tool-turns/credential-resolver"; +import { currentRunId } from "./run-control"; +import { budgetBlock, isPriced, providerPrices, readTokenUsage, requestCostCents } from "./spend"; import { HEALTH_STATE_HEALTHY, WORKSPACE_SCOPES, @@ -47,6 +49,9 @@ export interface ProviderResponseRunnerDeps { } export class ProviderResponseRunner { + /** Providers already logged as not reporting token usage. */ + private readonly unreportedUsage = new Set(); + constructor(private readonly deps: ProviderResponseRunnerDeps) {} async createResponse(input: ProviderResponseCreateInput): Promise { @@ -116,9 +121,19 @@ export class ProviderResponseRunner { if (!model) throw new Error("Provider model is required."); const owner: ProviderRow & { model: string } = { ...provider, model }; const turn = this.deps.toolTurn ?? selectToolTurn(provider.provider_key); + // A call made for a run is spent by its agent: refuse before calling once its budget is gone. + const taskRunId = currentRunId(); + const agent = taskRunId + ? this.deps.db.prepare("select a.id, a.department from task_runs r join agents a on a.id = r.agent_id where r.id = ?").get(taskRunId) as { id: string; department: string } + : null; + const blocked = agent ? budgetBlock(this.deps.db, agent.id, isPriced(providerPrices(this.deps.db, provider.id)), new Date()) : null; + if (blocked) throw new Error(blocked); const credential = await resolveTurnCredential({ authBroker: this.deps.authBroker, integrationKeys: this.deps.integrationKeys, orgId: input.orgId, provider }); const promptChars = JSON.stringify(input.messages).length; const requestId = this.deps.requestStore.create({ + agentId: agent?.id ?? null, + taskRunId, + department: agent?.department ?? null, orgId: input.orgId, provider: owner, model, @@ -149,6 +164,7 @@ export class ProviderResponseRunner { action: "response_created", }); this.deps.requestStore.recordUsage(input.orgId, provider.id, promptChars, result.text.length); + this.recordSpend(requestId, owner, result.rawResponseJson); return { ...result, requestId }; } catch (error) { this.deps.responseObserver.recordFailure({ @@ -164,6 +180,16 @@ export class ProviderResponseRunner { } } + private recordSpend(requestId: string, provider: { id: string; provider_key: string }, rawResponseJson: string): void { + const usage = readTokenUsage(rawResponseJson); + if (!usage) { + if (!this.unreportedUsage.has(provider.provider_key)) console.warn(`[spend] ${provider.provider_key} returned no token usage; its calls cost nothing on record.`); + this.unreportedUsage.add(provider.provider_key); + return; + } + this.deps.requestStore.recordSpend(requestId, usage, requestCostCents(usage, providerPrices(this.deps.db, provider.id))); + } + async streamResponse(input: StreamProviderResponseInput, onEvent: (event: ProviderStreamRendererEvent) => void): Promise { const execution = this.prepareExecution(input); const requestId = this.createRequestRecord(execution); diff --git a/src/main/services/provider-service.ts b/src/main/services/provider-service.ts index 5882d45f..8eb402d8 100644 --- a/src/main/services/provider-service.ts +++ b/src/main/services/provider-service.ts @@ -1,6 +1,6 @@ import type { DatabaseSync } from "node:sqlite"; import { createHash, randomBytes } from "node:crypto"; -import type { CancelProviderResponseInput, LaunchProviderSignInInput, ListProviderModelsInput, RevokeProviderCredentialInput, SaveProviderCredentialInput, StreamProviderResponseInput, UpdateProviderPolicyInput } from "../../shared/commands"; +import type { CancelProviderResponseInput, LaunchProviderSignInInput, ListProviderModelsInput, RevokeProviderCredentialInput, SaveProviderCredentialInput, StreamProviderResponseInput, UpdateProviderPolicyInput, SetProviderPricesInput } from "../../shared/commands"; import type { ProviderModelEntry } from "../../shared/provider-models"; import { PROVIDER_AUTH_TYPE, PROVIDER_DATA_POLICY, type ProviderAuthType } from "../../shared/domain"; import type { ProviderStreamRendererEvent } from "../../shared/provider-request-domain"; @@ -259,6 +259,27 @@ export class ProviderService { }); } + /** Sets the per-million-token prices that turn recorded usage into spend. */ + setPrices(input: SetProviderPricesInput): void { + const cents = (dollars: number | null, label: string): number | null => { + if (dollars === null) return null; + if (!Number.isFinite(dollars) || dollars < 0) throw new Error(`${label} price must be a number of dollars per million tokens, zero or more.`); + return dollars * 100; + }; + const inputCents = cents(input.inputDollarsPerMtok, "Input"); + const outputCents = cents(input.outputDollarsPerMtok, "Output"); + if ((inputCents === null) !== (outputCents === null)) throw new Error("Set both the input and the output price, or neither."); + const result = this.db.prepare(` + update providers set input_cents_per_mtok = ?, cached_input_cents_per_mtok = ?, output_cents_per_mtok = ? + where id = ? and org_id = ? + `).run(inputCents, cents(input.cachedInputDollarsPerMtok, "Cached input"), outputCents, input.providerId, input.orgId); + if (result.changes === 0) throw new Error("Provider was not found."); + this.audit.record({ + orgId: input.orgId, entityType: "provider", entityId: input.providerId, actorType: "employee", actorId: null, + action: "prices_updated", payload: { input: input.inputDollarsPerMtok, cachedInput: input.cachedInputDollarsPerMtok, output: input.outputDollarsPerMtok }, + }); + } + updatePolicy(input: UpdateProviderPolicyInput): void { const monthlyRequestLimit = Math.trunc(input.monthlyRequestLimit); if (monthlyRequestLimit < 1) throw new Error("Monthly request limit must be at least 1."); diff --git a/src/main/services/provider-tool-turns/router.ts b/src/main/services/provider-tool-turns/router.ts index ca63439a..ee83c90a 100644 --- a/src/main/services/provider-tool-turns/router.ts +++ b/src/main/services/provider-tool-turns/router.ts @@ -19,8 +19,8 @@ import { const TURNS: Record = { openai: runOpenAiResponsesTurn, anthropic: runAnthropicTurn, - xai: makeChatCompletionsTurn("https://api.x.ai/v1/chat/completions", "xAI"), - groq: makeChatCompletionsTurn("https://api.groq.com/openai/v1/chat/completions", "Groq"), + xai: makeChatCompletionsTurn("https://api.x.ai/v1/chat/completions", "xAI", true), + groq: makeChatCompletionsTurn("https://api.groq.com/openai/v1/chat/completions", "Groq", true), mistral: makeChatCompletionsTurn("https://api.mistral.ai/v1/chat/completions", "Mistral"), cerebras: makeChatCompletionsTurn("https://api.cerebras.ai/v1/chat/completions", "Cerebras"), openrouter: makeChatCompletionsTurn("https://openrouter.ai/api/v1/chat/completions", "OpenRouter"), diff --git a/src/main/services/provider-tool-turns/turns.ts b/src/main/services/provider-tool-turns/turns.ts index 296628b8..1c7c6f94 100644 --- a/src/main/services/provider-tool-turns/turns.ts +++ b/src/main/services/provider-tool-turns/turns.ts @@ -60,20 +60,22 @@ export async function runOpenAiResponsesTurn(input: ProviderToolTurnInput): Prom * "Provider request timed out." even though it is working. Streamed turns * emit chunks continuously and only a truly dead connection times out. */ -async function runChatCompletionsStream(endpoint: string, headers: Record, input: ProviderToolTurnInput): Promise { +async function runChatCompletionsStream(endpoint: string, headers: Record, input: ProviderToolTurnInput, requestUsage = false): Promise { const body: JsonRecord = { model: input.model, messages: toChatMessages(input.instructions, input.messages), tools: chatToolDefs(input.tools), tool_choice: "auto", stream: true, + // Some providers only report a stream's token usage when asked; the others send it unasked and may reject the field. + ...(requestUsage ? { stream_options: { include_usage: true } } : {}), }; const response = await postEventStream(endpoint, headers, body); return accumulateChatStream(response.messages); } -export function makeChatCompletionsTurn(endpoint: string, label: string): (input: ProviderToolTurnInput) => Promise { - return async (input) => runChatCompletionsStream(endpoint, { authorization: `Bearer ${requireApiKey(input.credential, label)}` }, input); +export function makeChatCompletionsTurn(endpoint: string, label: string, requestUsage = false): (input: ProviderToolTurnInput) => Promise { + return async (input) => runChatCompletionsStream(endpoint, { authorization: `Bearer ${requireApiKey(input.credential, label)}` }, input, requestUsage); } // ---- Anthropic Messages (api key) ---- diff --git a/src/main/services/run-control.ts b/src/main/services/run-control.ts index 11a83240..c2fd66ba 100644 --- a/src/main/services/run-control.ts +++ b/src/main/services/run-control.ts @@ -3,7 +3,7 @@ import { AsyncLocalStorage } from "node:async_hooks"; /** The error message a run ends with when the operator stops it. */ export const RUN_STOPPED_MESSAGE = "Stopped by the operator."; -const runSignals = new AsyncLocalStorage(); +const runScope = new AsyncLocalStorage<{ runId: string; signal: AbortSignal }>(); const controllers = new Map(); const settling = new Map>(); @@ -11,7 +11,7 @@ const settling = new Map>(); export async function withRunControl(runId: string, operation: () => Promise): Promise { const controller = new AbortController(); controllers.set(runId, controller); - const running = runSignals.run(controller.signal, operation); + const running = runScope.run({ runId, signal: controller.signal }, operation); settling.set(runId, running.catch(() => undefined)); try { return await running; @@ -49,7 +49,12 @@ export function stopAllRuns(): number { /** The stop signal of the run this code executes in, if any. */ export function currentRunSignal(): AbortSignal | undefined { - return runSignals.getStore(); + return runScope.getStore()?.signal; +} + +/** The task run this code executes in, if any. */ +export function currentRunId(): string | null { + return runScope.getStore()?.runId ?? null; } /** Why the current run was stopped, or null while it may continue. */ diff --git a/src/main/services/spend.ts b/src/main/services/spend.ts new file mode 100644 index 00000000..63fef677 --- /dev/null +++ b/src/main/services/spend.ts @@ -0,0 +1,155 @@ +import type { DatabaseSync } from "node:sqlite"; + +/** Tokens one model call consumed, as the provider reported them. */ +export interface TokenUsage { + /** All input tokens, cached ones included. */ + inputTokens: number; + outputTokens: number; + cachedInputTokens: number; + reasoningTokens: number; +} + +/** Prices in cents per million tokens; null when the operator has not set them. */ +export interface TokenPrices { + inputCentsPerMtok: number | null; + cachedInputCentsPerMtok: number | null; + outputCentsPerMtok: number | null; +} + +/** Reads token usage from a provider's raw response: one JSON body, or the list of a stream's event payloads. */ +export function readTokenUsage(rawResponseJson: string): TokenUsage | null { + const parsed = JSON.parse(rawResponseJson) as unknown; + const bodies = Array.isArray(parsed) ? parsed.map(parseEvent) : [parsed]; + for (let index = bodies.length - 1; index >= 0; index -= 1) { + const usage = usageIn(bodies[index]); + if (usage) return usage; + } + return null; +} + +/** True when calls to this provider cost money on record. */ +export function isPriced(prices: TokenPrices): boolean { + return prices.inputCentsPerMtok !== null && prices.outputCentsPerMtok !== null; +} + +/** What a call cost in cents, or null when its provider has no prices. */ +export function requestCostCents(usage: TokenUsage, prices: TokenPrices): number | null { + if (prices.inputCentsPerMtok === null || prices.outputCentsPerMtok === null) return null; + const cachedPrice = prices.cachedInputCentsPerMtok ?? prices.inputCentsPerMtok; + const fresh = usage.inputTokens - usage.cachedInputTokens; + return (fresh * prices.inputCentsPerMtok + usage.cachedInputTokens * cachedPrice + usage.outputTokens * prices.outputCentsPerMtok) / 1_000_000; +} + +export function providerPrices(db: DatabaseSync, providerId: string): TokenPrices { + const row = db.prepare("select input_cents_per_mtok, cached_input_cents_per_mtok, output_cents_per_mtok from providers where id = ?").get(providerId) as + { input_cents_per_mtok: number | null; cached_input_cents_per_mtok: number | null; output_cents_per_mtok: number | null } | undefined; + if (!row) throw new Error(`Provider ${providerId} was not found.`); + return { inputCentsPerMtok: row.input_cents_per_mtok, cachedInputCentsPerMtok: row.cached_input_cents_per_mtok, outputCentsPerMtok: row.output_cents_per_mtok }; +} + +export function monthStartIso(now: Date): string { + return new Date(Date.UTC(now.getUTCFullYear(), now.getUTCMonth(), 1)).toISOString(); +} + +/** Why this agent may not make a paid call this month, or null when it may; free (unpriced) calls are never blocked. */ +export function budgetBlock(db: DatabaseSync, agentId: string, paid: boolean, now: Date): string | null { + if (!paid) return null; + const agent = db.prepare("select org_id, name, department, budget_cents from agents where id = ?").get(agentId) as + { org_id: string; name: string; department: string; budget_cents: number } | undefined; + if (!agent) throw new Error(`Agent ${agentId} was not found.`); + const since = monthStartIso(now); + const agentSpend = (db.prepare("select coalesce(sum(cost_cents), 0) as cents from provider_requests where agent_id = ? and created_at >= ?") + .get(agentId, since) as { cents: number }).cents; + if (agentSpend >= agent.budget_cents) { + return `${agent.name} has spent its monthly budget (${dollars(agentSpend)} of ${dollars(agent.budget_cents)}). Raise the budget or wait for next month.`; + } + const department = db.prepare("select monthly_budget_limit_cents from org_departments where org_id = ? and name = ?") + .get(agent.org_id, agent.department) as { monthly_budget_limit_cents: number } | undefined; + if (!department) return null; + // Charges stay with the department that made them, even after the agent moves. + const departmentSpend = (db.prepare(` + select coalesce(sum(cost_cents), 0) as cents from provider_requests where org_id = ? and department = ? and created_at >= ? + `).get(agent.org_id, agent.department, since) as { cents: number }).cents; + if (departmentSpend >= department.monthly_budget_limit_cents) { + return `The ${agent.department} department has spent its monthly budget (${dollars(departmentSpend)} of ${dollars(department.monthly_budget_limit_cents)}).`; + } + return null; +} + +function dollars(cents: number): string { + return `$${(cents / 100).toFixed(2)}`; +} + +function parseEvent(raw: unknown): unknown { + if (typeof raw !== "string") return raw; + try { + return JSON.parse(raw) as unknown; + } catch { + // Stream sentinels such as [DONE] carry no usage. + return null; + } +} + +function usageIn(body: unknown): TokenUsage | null { + if (!isRecord(body)) return null; + const nested = isRecord(body.response) ? body.response : null; + const firstChoice = Array.isArray(body.choices) && isRecord(body.choices[0]) ? body.choices[0] : null; + const openAiStyle = record(body.usage) ?? record(nested?.usage) ?? record(isRecord(body.x_groq) ? body.x_groq.usage : null) ?? record(firstChoice?.usage); + if (openAiStyle) return fromOpenAiOrAnthropic(openAiStyle); + const gemini = record(body.usageMetadata) ?? record(nested?.usageMetadata); + if (gemini) return fromGemini(gemini); + return null; +} + +function fromOpenAiOrAnthropic(usage: Record): TokenUsage | null { + // Bedrock Converse + if (typeof usage.inputTokens === "number") { + return { + inputTokens: usage.inputTokens + count(usage.cacheReadInputTokens) + count(usage.cacheWriteInputTokens), + outputTokens: count(usage.outputTokens), + cachedInputTokens: count(usage.cacheReadInputTokens), + reasoningTokens: 0, + }; + } + // Chat Completions + if (typeof usage.prompt_tokens === "number") { + return { + inputTokens: usage.prompt_tokens, + outputTokens: count(usage.completion_tokens), + cachedInputTokens: count(record(usage.prompt_tokens_details)?.cached_tokens), + reasoningTokens: count(record(usage.completion_tokens_details)?.reasoning_tokens), + }; + } + if (typeof usage.input_tokens !== "number") return null; + // Anthropic reports cache reads and writes beside input_tokens; Responses reports them inside it. + const anthropicCache = count(usage.cache_read_input_tokens) + count(usage.cache_creation_input_tokens); + return { + inputTokens: usage.input_tokens + anthropicCache, + outputTokens: count(usage.output_tokens), + cachedInputTokens: count(usage.cache_read_input_tokens) + count(record(usage.input_tokens_details)?.cached_tokens), + reasoningTokens: count(record(usage.output_tokens_details)?.reasoning_tokens), + }; +} + +function fromGemini(usage: Record): TokenUsage | null { + if (typeof usage.promptTokenCount !== "number") return null; + const thoughts = count(usage.thoughtsTokenCount); + return { + inputTokens: usage.promptTokenCount, + outputTokens: count(usage.candidatesTokenCount) + thoughts, + cachedInputTokens: count(usage.cachedContentTokenCount), + reasoningTokens: thoughts, + }; +} + +function count(value: unknown): number { + return typeof value === "number" && Number.isFinite(value) ? value : 0; +} + +function record(value: unknown): Record | null { + return isRecord(value) ? value : null; +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} diff --git a/src/preload/index.ts b/src/preload/index.ts index 234d3300..3d69502d 100644 --- a/src/preload/index.ts +++ b/src/preload/index.ts @@ -61,6 +61,7 @@ const api: DesktopApi = { checkProviderHealth: (input) => invoke(IPC_CHANNELS.checkProviderHealth, input), listProviderModels: (input) => invoke(IPC_CHANNELS.listProviderModels, input), updateProviderPolicy: (input) => invoke(IPC_CHANNELS.updateProviderPolicy, input), + setProviderPrices: (input) => invoke(IPC_CHANNELS.setProviderPrices, input), openExternalUrl: (url) => invoke(IPC_CHANNELS.openExternalUrl, url), createProviderResponse: (input) => invoke(IPC_CHANNELS.createProviderResponse, input), streamProviderResponse: (input) => invoke(IPC_CHANNELS.streamProviderResponse, input), diff --git a/src/renderer/src/i18n/zh.ts b/src/renderer/src/i18n/zh.ts index 8b6e97b2..9ccb52f5 100644 --- a/src/renderer/src/i18n/zh.ts +++ b/src/renderer/src/i18n/zh.ts @@ -32,6 +32,14 @@ export const zh: Record = { "superseded": "已被取代", "stalled": "已停滞", "Resume": "恢复", + "Prices (US$ per million tokens)": "价格（美元 / 百万 token）", + "Cached input": "缓存输入", + "Unpriced: spend is not counted": "未定价：不计入花费", + "{label} {kind} price": "{label} {kind}价格", + "input": "输入", + "cached input": "缓存输入", + "output": "输出", + "Save {label} prices": "保存 {label} 价格", "Stop": "停止", "Stopped": "已停止", "stopped": "已停止", diff --git a/src/renderer/src/pages/settings/ProvidersPanel.tsx b/src/renderer/src/pages/settings/ProvidersPanel.tsx index 0694e2af..72a1f3a6 100644 --- a/src/renderer/src/pages/settings/ProvidersPanel.tsx +++ b/src/renderer/src/pages/settings/ProvidersPanel.tsx @@ -26,6 +26,7 @@ interface ProvidersPanelProps { actions: SnapshotActions; } +const PRICE_GRID = "minmax(150px, 1.4fr) 110px 110px 110px 92px"; const POLICY_GRID = "minmax(150px, 1.2fr) minmax(118px, 0.8fr) minmax(140px, 1fr) 112px 124px 124px 92px"; type WorkspaceScope = "organization" | "workspace"; @@ -81,6 +82,7 @@ export function ProvidersPanel({ orgId, employeeId, providers, authProfiles, cat providerRateLimits={providerRateLimits} actions={actions} /> + )} @@ -160,6 +162,51 @@ function ProviderPolicyTable({ ); } +/** Prices turn recorded token usage into spend, which agent and department budgets are checked against. */ +function ProviderPricesTable({ orgId, providers, actions }: { orgId: string; providers: ProviderRecord[]; actions: SnapshotActions }) { + return ( + <> +
{t("Prices (US$ per million tokens)")}
+
+
+
{t("Provider")}
{t("Input")}
{t("Cached input")}
{t("Output")}
+
+ {providers.map((provider) => )} +
+ + ); +} + +function ProviderPricesRow({ orgId, provider, actions }: { orgId: string; provider: ProviderRecord; actions: SnapshotActions }) { + const asDollars = (cents: number | null) => (cents === null ? "" : String(cents / 100)); + const [input, setInput] = useState(asDollars(provider.inputCentsPerMtok)); + const [cached, setCached] = useState(asDollars(provider.cachedInputCentsPerMtok)); + const [output, setOutput] = useState(asDollars(provider.outputCentsPerMtok)); + const [error, setError] = useState(null); + const price = (value: string) => (value.trim() === "" ? null : Number(value)); + const save = () => { + setError(null); + actions.run(() => window.agentCompany.setProviderPrices({ + orgId, providerId: provider.id, inputDollarsPerMtok: price(input), cachedInputDollarsPerMtok: price(cached), outputDollarsPerMtok: price(output), + })).catch((failure: unknown) => setError(failure instanceof Error ? failure.message : String(failure))); + }; + const field = (label: string, value: string, onChange: (next: string) => void) => ( + onChange(event.target.value)} /> + ); + return ( +
+
{provider.label}{provider.inputCentsPerMtok === null &&
{t("Unpriced: spend is not counted")}
}
+ {field(t("input"), input, setInput)} + {field(t("cached input"), cached, setCached)} + {field(t("output"), output, setOutput)} +
+ + {error &&
{error}
} +
+
+ ); +} + function ProviderPolicyRow({ orgId, provider, diff --git a/src/shared/commands/provider.ts b/src/shared/commands/provider.ts index ca821507..a128a967 100644 --- a/src/shared/commands/provider.ts +++ b/src/shared/commands/provider.ts @@ -64,6 +64,15 @@ export interface UpdateProviderPolicyInput { contextWindowTokens: number; } +/** List prices in dollars per million tokens; null leaves a price unset. */ +export interface SetProviderPricesInput { + orgId: string; + providerId: string; + inputDollarsPerMtok: number | null; + cachedInputDollarsPerMtok: number | null; + outputDollarsPerMtok: number | null; +} + export interface CreateProviderResponseInput { orgId: string; providerId: string; diff --git a/src/shared/desktop-api.ts b/src/shared/desktop-api.ts index 16538f8a..777bbb23 100644 --- a/src/shared/desktop-api.ts +++ b/src/shared/desktop-api.ts @@ -89,6 +89,7 @@ import type { StreamProviderResponseInput, SyncEmployeeIdentityInput, UpdateProviderPolicyInput, + SetProviderPricesInput, UpdateTaskStatusInput, } from "./commands"; import type { ProviderStreamRendererEvent } from "./provider-request-domain"; @@ -164,6 +165,7 @@ export interface DesktopApi { checkProviderHealth(input: CheckProviderHealthInput): Promise>; listProviderModels(input: ListProviderModelsInput): Promise>; updateProviderPolicy(input: UpdateProviderPolicyInput): Promise>; + setProviderPrices(input: SetProviderPricesInput): Promise>; openExternalUrl(url: string): Promise>; createProviderResponse(input: CreateProviderResponseInput): Promise>; streamProviderResponse(input: StreamProviderResponseInput): Promise>; diff --git a/src/shared/domain.ts b/src/shared/domain.ts index e5175bc5..c3dac344 100644 --- a/src/shared/domain.ts +++ b/src/shared/domain.ts @@ -296,6 +296,10 @@ export interface ProviderRecord { workspaceScope: string; /** Model context window the run budget is computed against. */ contextWindowTokens: number; + /** Prices in cents per million tokens; null until the operator sets them. */ + inputCentsPerMtok: number | null; + cachedInputCentsPerMtok: number | null; + outputCentsPerMtok: number | null; createdAt: string; } diff --git a/src/shared/ipc.ts b/src/shared/ipc.ts index 51727091..b18415d3 100644 --- a/src/shared/ipc.ts +++ b/src/shared/ipc.ts @@ -34,6 +34,7 @@ export const IPC_CHANNELS = { checkProviderHealth: "agentcompany:check-provider-health", listProviderModels: "agentcompany:list-provider-models", updateProviderPolicy: "agentcompany:update-provider-policy", + setProviderPrices: "agentcompany:set-provider-prices", openExternalUrl: "agentcompany:open-external-url", createProviderResponse: "agentcompany:create-provider-response", streamProviderResponse: "agentcompany:stream-provider-response", diff --git a/tests/integration/spend-budget.test.ts b/tests/integration/spend-budget.test.ts new file mode 100644 index 00000000..3fa46c30 --- /dev/null +++ b/tests/integration/spend-budget.test.ts @@ -0,0 +1,84 @@ +import { randomUUID } from "node:crypto"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { describe, expect, test } from "vitest"; +import { DatabaseConnection } from "../../src/main/database/connection"; +import { OnboardingService } from "../../src/main/services/onboarding-service"; +import { ProviderService } from "../../src/main/services/provider-service"; +import { providerAdapters } from "../../src/main/services/provider-adapters"; +import { TaskExecutionService } from "../../src/main/services/task-execution-service"; +import { claimNextQueuedRun } from "../../src/main/app/task-queue-worker"; +import type { AgentTurnResult } from "../../src/main/services/provider-tool-turn"; + +const testVault = { + encryptSecret: (secret: string) => Buffer.from(secret, "utf8"), + decryptSecret: (ciphertext: Buffer) => ciphertext.toString("utf8"), +}; + +describe("money budgets", () => { + test("a run's model calls cost real money, and an agent past its monthly budget is refused and parked", async () => { + let calls = 0; + const turn = async (): Promise => { + calls += 1; + return { text: "Worked.", toolCalls: [], rawResponseJson: JSON.stringify({ usage: { input_tokens: 2_000_000, output_tokens: 1_000_000 } }) }; + }; + const connection = new DatabaseConnection(join(tmpdir(), `ac-spend-${randomUUID()}.db`)); + const providers = new ProviderService(connection.db, providerAdapters, testVault, undefined, turn); + const onboarding = new OnboardingService({ db: connection.db, saveProviderCredential: (input) => providers.saveCredential(input) }); + const execution = new TaskExecutionService(connection.db, providers); + let orgId = ""; + let ceoId = ""; + connection.transaction(() => { + const summary = onboarding.bootstrap({ + companyName: "FlowPilot", companyDescription: "CSM analytics.", employeeName: "Yang", workspacePath: "", + executionMode: "ai_ceo", autoApproveHires: true, reviewDeliverables: false, + goalTitle: "Ship the MVP.", goalDescription: "A working product.", + providerKey: "openai", providerSecret: "sk-test-key", providerModel: "gpt-5.4", + }); + orgId = summary.orgId; + ceoId = summary.ceoAgentId; + }); + const providerId = (connection.db.prepare("select model_provider_id as id from agents where id = ?").get(ceoId) as { id: string }).id; + providers.setPrices({ orgId, providerId, inputDollarsPerMtok: 2.5, cachedInputDollarsPerMtok: null, outputDollarsPerMtok: 10 }); + connection.db.prepare("update agents set budget_cents = 2000 where id = ?").run(ceoId); + + const first = claimNextQueuedRun(connection)!; + await execution.executeRun(orgId, first.id); + expect(connection.db.prepare("select agent_id, task_run_id, input_tokens, output_tokens, cost_cents from provider_requests where task_run_id = ?").get(first.id)) + .toEqual({ agent_id: ceoId, task_run_id: first.id, input_tokens: 2_000_000, output_tokens: 1_000_000, cost_cents: 500 + 1000 }); + + const second = connection.transaction(() => execution.enqueueHeartbeatRuns()); + expect(second).toHaveLength(1); + connection.db.prepare("update task_runs set state = 'running' where id = ?").run(second[0]); + await execution.executeRun(orgId, second[0]); + expect(calls).toBe(2); + + // $30 spent of a $20 budget: the next call is refused before it is made, and the heartbeat parks the agent. + connection.db.prepare("update task_runs set created_at = '2000-01-01T00:00:00.000Z'").run(); + expect(connection.transaction(() => execution.enqueueHeartbeatRuns())).toEqual([]); + const third = execution.runTask({ orgId, taskId: (connection.db.prepare("select task_id from task_runs where id = ?").get(first.id) as { task_id: string }).task_id }); + await expect(third).rejects.toThrow("CEO has spent its monthly budget ($30.00 of $20.00). Raise the budget or wait for next month."); + expect(calls).toBe(2); + + // A $0 budget allows no paid call, and unpriced calls are never blocked by money. + connection.db.prepare("update agents set budget_cents = 0 where id = ?").run(ceoId); + connection.db.prepare("delete from provider_requests").run(); + const kickoffTask = (connection.db.prepare("select task_id from task_runs where id = ?").get(first.id) as { task_id: string }).task_id; + connection.db.prepare("update task_runs set state = 'failed' where state in ('queued', 'running')").run(); + await expect(execution.runTask({ orgId, taskId: kickoffTask })).rejects.toThrow("CEO has spent its monthly budget ($0.00 of $0.00)."); + providers.setPrices({ orgId, providerId, inputDollarsPerMtok: null, cachedInputDollarsPerMtok: null, outputDollarsPerMtok: null }); + await execution.runTask({ orgId, taskId: kickoffTask }); + expect(calls).toBe(3); + + // Department spend stays with the department that made it, even after the agent moves. + providers.setPrices({ orgId, providerId, inputDollarsPerMtok: 2.5, cachedInputDollarsPerMtok: null, outputDollarsPerMtok: 10 }); + connection.db.prepare("update agents set budget_cents = 1000000 where id = ?").run(ceoId); + connection.db.prepare("update org_departments set monthly_budget_limit_cents = 1000 where org_id = ? and name = 'Executive'").run(orgId); + connection.db.prepare("update provider_requests set cost_cents = 1500").run(); + connection.db.prepare("update agents set department = 'Design' where id = ?").run(ceoId); + expect(connection.db.prepare("select department from provider_requests").get()).toEqual({ department: "Executive" }); + connection.db.prepare("update agents set department = 'Executive' where id = ?").run(ceoId); + await expect(execution.runTask({ orgId, taskId: kickoffTask })).rejects.toThrow("The Executive department has spent its monthly budget ($15.00 of $10.00)."); + connection.close(); + }); +}); diff --git a/tests/unit/domain-services.test.ts b/tests/unit/domain-services.test.ts index e9d458b0..1e29c6fd 100644 --- a/tests/unit/domain-services.test.ts +++ b/tests/unit/domain-services.test.ts @@ -24,7 +24,7 @@ describe("domain services", () => { const connection = createConnection(); const rows = connection.db.prepare("select version from schema_migrations order by version asc").all() as Array<{ version: number }>; - expect(rows.map((row) => row.version)).toEqual([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71]); + expect(rows.map((row) => row.version)).toEqual([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72]); connection.close(); }); diff --git a/tests/unit/spend.test.ts b/tests/unit/spend.test.ts new file mode 100644 index 00000000..795a7d39 --- /dev/null +++ b/tests/unit/spend.test.ts @@ -0,0 +1,29 @@ +import { describe, expect, test } from "vitest"; +import { readTokenUsage, requestCostCents } from "../../src/main/services/spend"; + +describe("token usage and cost", () => { + test("reads usage from every provider shape, cached and reasoning tokens included", () => { + expect(readTokenUsage(JSON.stringify({ usage: { input_tokens: 1000, output_tokens: 200, input_tokens_details: { cached_tokens: 400 }, output_tokens_details: { reasoning_tokens: 50 } } }))) + .toEqual({ inputTokens: 1000, outputTokens: 200, cachedInputTokens: 400, reasoningTokens: 50 }); + expect(readTokenUsage(JSON.stringify({ usage: { input_tokens: 600, output_tokens: 80, cache_read_input_tokens: 300, cache_creation_input_tokens: 100 } }))) + .toEqual({ inputTokens: 1000, outputTokens: 80, cachedInputTokens: 300, reasoningTokens: 0 }); + expect(readTokenUsage(JSON.stringify({ usage: { inputTokens: 500, outputTokens: 70, cacheReadInputTokens: 100 } }))) + .toEqual({ inputTokens: 600, outputTokens: 70, cachedInputTokens: 100, reasoningTokens: 0 }); + expect(readTokenUsage(JSON.stringify({ usageMetadata: { promptTokenCount: 900, candidatesTokenCount: 60, thoughtsTokenCount: 40, cachedContentTokenCount: 200 } }))) + .toEqual({ inputTokens: 900, outputTokens: 100, cachedInputTokens: 200, reasoningTokens: 40 }); + const chatStream = ["{\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}", "{\"choices\":[],\"usage\":{\"prompt_tokens\":120,\"completion_tokens\":30,\"prompt_tokens_details\":{\"cached_tokens\":20}}}", "[DONE]"]; + expect(readTokenUsage(JSON.stringify(chatStream))).toEqual({ inputTokens: 120, outputTokens: 30, cachedInputTokens: 20, reasoningTokens: 0 }); + const responsesStream = ["{\"type\":\"response.created\"}", "{\"type\":\"response.completed\",\"response\":{\"usage\":{\"input_tokens\":10,\"output_tokens\":5}}}"]; + expect(readTokenUsage(JSON.stringify(responsesStream))).toEqual({ inputTokens: 10, outputTokens: 5, cachedInputTokens: 0, reasoningTokens: 0 }); + const moonshotStream = ["{\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\",\"usage\":{\"prompt_tokens\":40,\"completion_tokens\":9}}]}"]; + expect(readTokenUsage(JSON.stringify(moonshotStream))).toEqual({ inputTokens: 40, outputTokens: 9, cachedInputTokens: 0, reasoningTokens: 0 }); + expect(readTokenUsage(JSON.stringify({ output: [] }))).toBeNull(); + }); + + test("cost prices fresh input, cached input and output separately; no prices, no cost", () => { + const usage = { inputTokens: 1_000_000, outputTokens: 500_000, cachedInputTokens: 400_000, reasoningTokens: 0 }; + expect(requestCostCents(usage, { inputCentsPerMtok: 250, cachedInputCentsPerMtok: 125, outputCentsPerMtok: 1000 })).toBe(150 + 50 + 500); + expect(requestCostCents(usage, { inputCentsPerMtok: 250, cachedInputCentsPerMtok: null, outputCentsPerMtok: 1000 })).toBe(250 + 500); + expect(requestCostCents(usage, { inputCentsPerMtok: null, cachedInputCentsPerMtok: null, outputCentsPerMtok: 1000 })).toBeNull(); + }); +});