Files
codecraft/server/src/budget.rs
T
Richard ZinkandClaude Opus 5.5 bfa1f641d7 opencode prices and Go limits, settings out of .env
- Fetch prices for opencode (Zen and Go) incl. limits, image input, reasoning,
  long-context tiers and peak-hour/date factors; costs apply them per request
- Record the cost opencode reports per response; track Go allowance in 5h/7d/month
  windows against a provider plan limit (budget table, status overlay)
- Settings → Config file: move values of ${NAME} references into the config
- Examples, presets and forms keep keys in the config instead of .env

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-11-09 10:19:15 +01:00

432 lines
18 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
//! Budgets: monthly limits tracked from recorded costs, plus the real balance
//! for providers that expose one (OpenRouter, fal.ai), and price discovery.
use crate::app::App;
use crate::config::{ProviderConfig, expand_env};
use anyhow::{Context, Result, bail};
use chrono::{Datelike, TimeZone};
use serde_json::{Value, json};
use std::time::Duration;
/// Start of the current month in local time (ms).
pub fn month_start() -> i64 {
let now = chrono::Local::now();
chrono::Local
.with_ymd_and_hms(now.year(), now.month(), 1, 0, 0, 0)
.single()
.map(|d| d.timestamp_millis())
.unwrap_or(0)
}
fn is_openrouter(p: &ProviderConfig) -> bool {
p.base_url.contains("openrouter.ai")
}
/// Month-to-date cost from Anthropic's cost report (organisation-wide, USD).
async fn anthropic_cost(p: &ProviderConfig) -> Result<Value> {
let admin = expand_env(&p.admin_key);
let start = chrono::DateTime::from_timestamp_millis(month_start()).unwrap_or_default();
let end = chrono::Utc::now() + chrono::Duration::days(1);
let base = if p.base_url.is_empty() { "https://api.anthropic.com/v1".to_string() } else { expand_env(&p.base_url).trim_end_matches('/').to_string() };
let mut total_cents = 0.0;
let mut page: Option<String> = None;
for _ in 0..10 {
let mut url = reqwest::Url::parse(&format!("{base}/organizations/cost_report"))?;
url.query_pairs_mut()
.append_pair("starting_at", &start.format("%Y-%m-%dT00:00:00Z").to_string())
.append_pair("ending_at", &end.format("%Y-%m-%dT00:00:00Z").to_string())
.append_pair("limit", "31");
if let Some(pg) = &page {
url.query_pairs_mut().append_pair("page", pg);
}
let resp = crate::provider::http_client()
.get(url)
.header("x-api-key", &admin)
.header("anthropic-version", "2026-10-21")
.timeout(Duration::from_secs(20))
.send()
.await?;
if !resp.status().is_success() {
return Err(crate::provider::http_error(resp).await);
}
let v: Value = resp.json().await?;
for bucket in v["data"].as_array().cloned().unwrap_or_default() {
for r in bucket["results"].as_array().cloned().unwrap_or_default() {
// Amounts are decimal strings in cents.
total_cents += r["amount"].as_str().and_then(|a| a.parse::<f64>().ok()).or_else(|| r["amount"].as_f64()).unwrap_or(0.0);
}
}
if v["has_more"].as_bool() != Some(true) {
break;
}
page = v["next_page"].as_str().map(String::from);
}
Ok(json!({"currency": "USD", "spent_month": total_cents / 100.0, "scope": "organization"}))
}
/// Balance or cost as reported by the provider, if it has an API for it.
async fn remote_balance(p: &ProviderConfig) -> Option<Result<Value>> {
if p.kind == crate::config::ProviderKind::Anthropic && !expand_env(&p.admin_key).is_empty() {
return Some(anthropic_cost(p).await);
}
if !is_openrouter(p) {
return None;
}
Some(
async {
let base = expand_env(&p.base_url).trim_end_matches('/').to_string();
let resp = crate::provider::http_client()
.get(format!("{base}/key"))
.bearer_auth(expand_env(&p.api_key))
.timeout(Duration::from_secs(15))
.send()
.await?
.error_for_status()?;
let v: Value = resp.json().await?;
let d = &v["data"];
Ok(json!({
"currency": "USD",
"limit": d["limit"],
"used": d["usage"],
"remaining": d.get("limit_remaining").cloned().unwrap_or_else(|| match (d["limit"].as_f64(), d["usage"].as_f64()) {
(Some(l), Some(u)) => json!(l - u),
_ => Value::Null,
}),
}))
}
.await,
)
}
async fn fal_balance(admin_key: &str) -> Result<Value> {
let resp = crate::provider::http_client()
.get("https://api.fal.ai/v1/account/billing?expand=credits")
.header("Authorization", format!("Key {admin_key}"))
.timeout(Duration::from_secs(15))
.send()
.await?;
if !resp.status().is_success() {
return Err(crate::provider::http_error(resp).await);
}
let v: Value = resp.json().await?;
Ok(json!({
"currency": v["credits"]["currency"],
"remaining": v["credits"]["current_balance"],
}))
}
fn entry(name: &str, budget: Option<f64>, spent: f64, warn: f64) -> Value {
let share = budget.filter(|b| *b > 0.0).map(|b| spent / b);
json!({
"provider": name,
"budget": budget,
"spent": spent,
"remaining": budget.map(|b| b - spent),
"share": share,
"warning": share.is_some_and(|s| s >= warn),
"exceeded": share.is_some_and(|s| s >= 1.0),
})
}
/// Budget overview for the current month.
/// Subscription windows of opencode Go, as share of the plan's monthly limit: rolling
/// 5 hours (20 %) and 7 days (50 %), and the calendar month in UTC (100 %). Only usage
/// through CodeCraft is counted.
pub fn plan_usage(app: &App, provider: &str, limit: f64) -> Result<Value> {
let now = chrono::Utc::now();
let month = chrono::Utc.with_ymd_and_hms(now.year(), now.month(), 1, 0, 0, 0).single().map(|d| d.timestamp_millis()).unwrap_or(0);
let ms = now.timestamp_millis();
let window = |name: &str, from: i64, share: f64| -> Result<Value> {
let used = app.store.allowance_since(provider, from)?;
let cap = limit * share;
Ok(json!({"window": name, "used": used, "limit": cap, "ratio": if cap > 0.0 { used / cap } else { 0.0 }}))
};
Ok(json!({
"limit": limit,
"windows": [
window("5h", ms - 5 * 3_600_000, 0.2)?,
window("week", ms - 7 * 86_400_000, 0.5)?,
window("month", month, 1.0)?,
],
}))
}
pub async fn overview(app: &App) -> Result<Value> {
let cfg = app.config();
let spent = app.store.cost_by_provider(month_start())?;
let warn = cfg.server.budget_warning;
let mut rows = Vec::new();
let mut tasks = Vec::new();
for (name, p) in &cfg.providers {
if !p.enabled {
continue;
}
let mut row = entry(name, p.monthly_budget, spent.get(name).copied().unwrap_or(0.0), warn);
if let Some(limit) = p.plan_limit.filter(|l| *l > 0.0) {
row["plan"] = plan_usage(app, name, limit)?;
}
rows.push(row);
let p = p.clone();
tasks.push(tokio::spawn(async move { remote_balance(&p).await }));
}
for (row, task) in rows.iter_mut().zip(tasks) {
if let Ok(Some(res)) = task.await {
row["remote"] = match res {
Ok(v) => v,
Err(e) => json!({"error": format!("{e:#}")}),
};
}
}
if cfg.images.enabled {
let mut row = entry("fal", cfg.images.monthly_budget, spent.get("fal").copied().unwrap_or(0.0), warn);
let admin = expand_env(&cfg.images.admin_key);
if !admin.is_empty() {
row["remote"] = match fal_balance(&admin).await {
Ok(v) => v,
Err(e) => json!({"error": format!("{e:#}")}),
};
}
row["images"] = json!(true);
rows.push(row);
}
let total: f64 = spent.values().sum();
let global = (cfg.server.monthly_budget > 0.0).then_some(cfg.server.monthly_budget);
Ok(json!({
"month_start": month_start(),
"currency": cfg.server.currency,
"total": entry("total", global, total, warn),
"providers": rows,
}))
}
/// Per-token prices from the provider's model list, converted to per 1M tokens.
fn is_opencode(p: &ProviderConfig) -> bool {
p.base_url.contains("opencode.ai")
}
async fn get_json(url: &str) -> Result<Value> {
Ok(crate::provider::http_client().get(url).timeout(Duration::from_secs(30)).send().await?.error_for_status()?.json().await?)
}
/// One opencode condition object → (days, hours, from, until); None if it uses
/// something we do not understand.
type When = (Vec<String>, Vec<(String, String)>, Option<String>, Option<String>);
fn opencode_when(w: &Value) -> Option<When> {
let obj = w.as_object()?;
let (mut days, mut hours, mut from, mut until) = (Vec::new(), Vec::new(), None, None);
// A pair of utc_hours bounds, given as {"all": [{"utc_hours": {"gte"}}, {"utc_hours": {"lt"}}]}
// or directly as {"utc_hours": {"gte", "lt"}}.
let range = |v: &Value| -> Option<(String, String)> {
if let Some(h) = v.get("utc_hours") {
return Some((h["gte"].as_str()?.to_string(), h["lt"].as_str()?.to_string()));
}
let all = v.get("all")?.as_array()?;
let get = |k: &str| all.iter().find_map(|x| x["utc_hours"][k].as_str()).map(String::from);
Some((get("gte")?, get("lt")?))
};
for (k, v) in obj {
match k.as_str() {
"utc_days" => days = v["in"].as_array()?.iter().map(|d| d.as_str().map(String::from)).collect::<Option<_>>()?,
"any" => hours = v.as_array()?.iter().map(range).collect::<Option<_>>()?,
"all" | "utc_hours" => hours = vec![range(w)?],
"date" => {
from = v["gte"].as_str().map(String::from);
until = v["lt"].as_str().map(String::from);
if from.is_none() && until.is_none() {
return None;
}
}
_ => return None,
}
}
Some((days, hours, from, until))
}
/// Model row from opencode's model list (prices per 1M tokens, tiers, time factors).
fn opencode_model(m: &Value) -> Value {
let cost = &m["cost"]["text"];
let (input, output) = (&cost["input"], &cost["output"]);
// Context tiers: input and output overrides share their thresholds.
let mut tiers: std::collections::BTreeMap<u64, serde_json::Map<String, Value>> = Default::default();
for (side, field) in [(input, "input_price"), (output, "output_price")] {
for o in side["overrides"].as_array().into_iter().flatten() {
let Some(gt) = o["when"]["context"]["gt"].as_u64() else { continue };
let t = tiers.entry(gt).or_default();
t.insert(field.into(), o["price"].clone());
if field == "input_price" {
if o["cache"]["read"].is_number() {
t.insert("cache_read_price".into(), o["cache"]["read"].clone());
}
if o["cache"]["write"].is_number() {
t.insert("cache_write_price".into(), o["cache"]["write"].clone());
}
}
}
}
let price_tiers: Vec<Value> = tiers.into_iter().map(|(gt, mut t)| {
t.insert("above_context".into(), json!(gt));
Value::Object(t)
}).collect();
let mut notes = Vec::new();
let price_multipliers: Vec<Value> = m["cost"]["multipliers"]
.as_array()
.into_iter()
.flatten()
.filter_map(|x| match opencode_when(&x["when"]) {
Some((days, hours, from, until)) => Some(json!({"factor": x["factor"], "days": days, "hours_utc": hours, "from": from, "until": until})),
None => {
notes.push(format!("unsupported price condition: {}", x["when"]));
None
}
})
.map(|mut v| {
// Drop empty conditions so the config stays short.
v.as_object_mut().unwrap().retain(|_, x| !(x.is_null() || x.as_array().is_some_and(|a| a.is_empty())));
v
})
.collect();
let images = m["modalities"]["input"].as_array().is_some_and(|a| a.iter().any(|x| x.as_str().is_some_and(|s| s.starts_with("image"))));
json!({
"id": m["id"],
"context": m["limit"]["context"],
"max_output": m["limit"]["output"],
"vision": images,
"reasoning": m["capabilities"]["reasoning"]["supported"].as_bool().unwrap_or(false),
"input_price": input["price"],
"output_price": output["price"],
"cache_read_price": input["cache"]["read"],
"cache_write_price": input["cache"]["write"],
"price_tiers": price_tiers,
"price_multipliers": price_multipliers,
"notes": notes,
})
}
/// opencode (Zen and Go) publishes prices in its model list. The Anthropic-protocol
/// endpoint lists models without them, so fall back to the gateway's own /v1/models.
async fn opencode_prices(base: &str) -> Result<Value> {
let mut urls = vec![format!("{base}/models")];
for suffix in ["/openai/v1", "/anthropic/v1"] {
if let Some(root) = base.strip_suffix(suffix) {
urls.push(format!("{root}/v1/models"));
}
}
for url in &urls {
let Ok(v) = get_json(url).await else { continue };
let models: Vec<Value> = v["data"].as_array().into_iter().flatten().filter(|m| m.get("cost").is_some()).map(opencode_model).collect();
if !models.is_empty() {
return Ok(json!({"currency": "USD", "models": models}));
}
}
bail!("opencode did not return prices for this endpoint")
}
pub async fn discover_prices(app: &App, provider: &str) -> Result<Value> {
let cfg = app.config();
let p = cfg.providers.get(provider).context("unknown provider")?;
let base = expand_env(&p.base_url).trim_end_matches('/').to_string();
if is_opencode(p) {
return opencode_prices(&base).await;
}
if !is_openrouter(p) {
bail!("this provider does not publish prices via its API – enter them manually");
}
let v: Value = crate::provider::http_client()
.get(format!("{base}/models"))
.timeout(Duration::from_secs(30))
.send()
.await?
.error_for_status()?
.json()
.await?;
let per_m = |x: &Value| x.as_str().and_then(|s| s.parse::<f64>().ok()).or_else(|| x.as_f64()).map(|f| (f * 1_000_000.0 * 10_000.0).round() / 10_000.0);
let list: Vec<Value> = v["data"]
.as_array()
.cloned()
.unwrap_or_default()
.iter()
.map(|m| {
json!({
"id": m["id"],
"context": m["context_length"],
"input_price": per_m(&m["pricing"]["prompt"]),
"output_price": per_m(&m["pricing"]["completion"]),
"cache_read_price": per_m(&m["pricing"]["input_cache_read"]),
"cache_write_price": per_m(&m["pricing"]["input_cache_write"]),
})
})
.collect();
Ok(json!({"currency": "USD", "models": list}))
}
#[cfg(test)]
mod tests {
#[test]
fn opencode_cost_conversion() {
let m: serde_json::Value = serde_json::from_str(r#"{"id":"grok","limit":{"context":500000,"output":8000},
"modalities":{"input":["text","image/png"]},"capabilities":{"reasoning":{"supported":true}},
"cost":{"text":{"input":{"price":2,"cache":{"read":0.5},"overrides":[{"when":{"context":{"gt":200000}},"price":4,"cache":{"read":1}}]},
"output":{"price":6,"overrides":[{"when":{"context":{"gt":200000}},"price":12}]}},
"multipliers":[{"when":{"utc_days":{"in":["mon","tue"]},"any":[{"all":[{"utc_hours":{"gte":"01:00"}},{"utc_hours":{"lt":"04:00"}}]}]},"factor":2},
{"when":{"date":{"gte":"2026-12-01"}},"factor":2},
{"when":{"region":"eu"},"factor":3}]}}"#).unwrap();
let r = super::opencode_model(&m);
assert_eq!(r["input_price"], 2);
assert_eq!(r["cache_read_price"], 0.5);
assert_eq!(r["vision"], true);
assert_eq!(r["price_tiers"][0], serde_json::json!({"above_context": 200000, "input_price": 4, "cache_read_price": 1, "output_price": 12}));
assert_eq!(r["price_multipliers"][0], serde_json::json!({"factor": 2, "days": ["mon", "tue"], "hours_utc": [["01:00", "04:00"]]}));
assert_eq!(r["price_multipliers"][1], serde_json::json!({"factor": 2, "from": "2026-12-01"}));
assert_eq!(r["notes"].as_array().unwrap().len(), 1, "unknown condition is reported, not guessed");
// The result must deserialize into the model config.
let mc: crate::config::ModelConfig = serde_json::from_value(serde_json::json!({"id": "grok", "price_tiers": r["price_tiers"], "price_multipliers": r["price_multipliers"]})).unwrap();
assert_eq!(mc.price_tiers[0].above_context, 200000);
}
#[test]
fn budget_entry_flags() {
let e = super::entry("x", Some(10.0), 8.5, 0.8);
assert_eq!(e["warning"], true);
assert_eq!(e["exceeded"], false);
assert_eq!(e["remaining"], 1.5);
let none = super::entry("y", None, 3.0, 0.8);
assert_eq!(none["warning"], false);
}
}
/// Refuse a request when a budget or the owner's quota is used up.
/// `provider` is the provider name ("fal" for images).
pub fn check(app: &App, provider: &str, owner_id: Option<&str>) -> Result<()> {
let cfg = app.config();
let from = month_start();
if let Some(user) = owner_id.and_then(|id| app.store.user(id).ok().flatten()) {
if let Some(q) = user.monthly_quota.filter(|q| *q > 0.0) {
let spent = app.store.user_cost_since(&user.id, from)?;
if spent >= q {
bail!(
"monthly quota of {q:.2} {cur} for user {} is used up ({spent:.2} {cur} spent) – ask an admin to raise it",
user.username,
cur = cfg.server.currency
);
}
}
}
if cfg.server.budget_action != "block" {
return Ok(());
}
let spent = app.store.cost_by_provider(from)?;
let total: f64 = spent.values().sum();
if cfg.server.monthly_budget > 0.0 && total >= cfg.server.monthly_budget {
bail!("the monthly budget of {:.2} {} is used up", cfg.server.monthly_budget, cfg.server.currency);
}
let limit = if provider == "fal" { cfg.images.monthly_budget } else { cfg.providers.get(provider).and_then(|p| p.monthly_budget) };
if let Some(b) = limit.filter(|b| *b > 0.0) {
let s = spent.get(provider).copied().unwrap_or(0.0);
if s >= b {
bail!("the monthly budget of {b:.2} {} for {provider} is used up", cfg.server.currency);
}
}
Ok(())
}