Skip to content

g1t/services/billing/src/ai.rs

1,208 lines55,861 bytesCodeBlame
1//! Prepaid AI credit: what Agent and AI Gateway usage draws on, bought in
2//! advance so g1t never fronts a model's cost.
3//!
4//! - **Buying.** An owner buys AI credit on Stripe's page: one payment by
5//! card, $10 to $1,000, with Stripe's card fee as its own line when the
6//! `card_fee` cost setting is on (`card_fee_cents`, a gross-up of the
7//! price book's `card_fee_percent` and `card_fee_fixed`). The card is
8//! kept for auto-reload. The credit is entered once, whichever comes
9//! first: the person coming back (`confirm_ai_credit`) or Stripe's
10//! `checkout.session.completed` (`webhooks.rs`). Both claim the same
11//! `checkouts` row, and the grant's id is the page's id, so a payment is
12//! credited exactly once.
13//! - **What it is.** A `credit_grants` row of kind `purchased`, scope
14//! `models`, source `purchase`, expiring a year after purchase, and its
15//! ledger line (a payment: its reference is Stripe's id, never `crd…`).
16//! Model usage draws on it before anything else (`grants::replay`), and
17//! it counts as money paid, never as given.
18//! - **Auto-reload.** Off by default. When AI credit falls below the
19//! threshold, the saved card is charged off-session to bring it back to
20//! the target, at most the monthly maximum. Each attempt has its own
21//! idempotency key (`ai_reloads.id`), so a retry is the same payment. A
22//! failed charge turns auto-reload off and tells the owners.
23//! - **At $0.** A workspace on the plan with no AI credit and none of its
24//! included usage left cannot start a run on g1t's models: `start_run`
25//! refuses with what to do. Auto-reload, when on, is tried first. A 100%
26//! discount (Flagon) pays for everything, so nothing is needed; an
27//! enterprise is invoiced for models after use.
28//! - **Once on upgrading.** A workspace that starts the paid plan is given
29//! $5 of AI credit once (promotional: given, not revenue), expiring in a
30//! year.
31//! - **The agent rate.** Every agent run's tokens (input, output and
32//! cached, as the model proxy counts them) are charged at the price
33//! book's `agent_tokens` price per million, on a line of their own
34//! (`<run>/agent`), on top of the model at the provider's price
35//! (`agent_models`, no markup).
36
37use g1t_contracts::billing::{
38 AccountArgs, AiCredit, AiReload, BuyAiCreditArgs, CardFee, Checkout, ConfirmAiCreditArgs, CreditKind, EntryKind, PlanKind,
39 RunTokens, SetAiReloadArgs, MICROS_PER_DOLLAR,
40};
41use g1t_contracts::time::rfc3339;
42use g1t_contracts::{FailureCode, Outcome, Role};
43use g1t_kit::now_ms;
44use serde::Deserialize;
45use worker::Result;
46
47use crate::features::cents;
48use crate::{Billing, RunRow, members_only, optional};
49
50/// What a checkout row for AI credit is marked with.
51pub(crate) const AI_CREDIT: &str = "ai_credit";
52/// The amounts offered, in cents, and the bounds of a custom one.
53pub(crate) const PRESETS_CENTS: [u32; 4] = [1_000, 2_500, 5_000, 10_000];
54pub(crate) const MIN_CENTS: u32 = 1_000;
55pub(crate) const MAX_CENTS: u32 = 100_000;
56/// Bought credit lasts a year.
57pub(crate) const EXPIRES_DAYS: u64 = 365;
58/// Given once, on starting the paid plan.
59pub(crate) const UPGRADE_CREDIT_MICROS: i64 = 5_000_000;
60/// Auto-reload's bounds: a reload of at least $10, a target of at most
61/// $1,000, and at most $10,000 a month.
62const MIN_RELOAD_MICROS: i64 = 10 * MICROS_PER_DOLLAR;
63const MAX_TARGET_MICROS: i64 = 1_000 * MICROS_PER_DOLLAR;
64const MAX_MONTHLY_MICROS: i64 = 10_000 * MICROS_PER_DOLLAR;
65const DAY_MS: u64 = 24 * 60 * 60 * 1000;
66
67// ---------------------------------------------------------------------
68// The arithmetic, apart from the database so it can be tested.
69// ---------------------------------------------------------------------
70
71/// Whether an amount of AI credit can be bought, in cents.
72pub(crate) fn amount_ok(cents: u32) -> std::result::Result<(), String> {
73 if (MIN_CENTS..=MAX_CENTS).contains(&cents) && cents.is_multiple_of(100) {
74 Ok(())
75 } else {
76 Err(format!("Buy between ${} and ${} of AI credit, in whole dollars.", MIN_CENTS / 100, crate::group(MAX_CENTS / 100)))
77 }
78}
79
80/// The card fee on `credit_cents`, so that what is left after Stripe's fee
81/// is the credit: the total is `(credit + fixed) / (1 − percent)`, rounded
82/// up to the cent. None when the fee is off.
83pub(crate) fn card_fee_cents(credit_cents: u32, fee: &CardFee) -> u32 {
84 if !fee.on || credit_cents == 0 {
85 return 0;
86 }
87 let rate = fee.percent_micros / MICROS_PER_DOLLAR as f64;
88 if !(0.0..0.5).contains(&rate) {
89 return 0;
90 }
91 let total = ((f64::from(credit_cents) + f64::from(fee.fixed_cents)) / (1.0 - rate)).ceil();
92 (total as u32).saturating_sub(credit_cents)
93}
94
95/// What auto-reload should buy now, if anything: enough to bring the
96/// credit from `balance` back to the target, in whole dollars, within
97/// what is left of the monthly maximum, and never less than $10.
98pub(crate) fn reload_amount(reload: &AiReload, balance: i64, reloaded_this_month: i64) -> Option<i64> {
99 if !reload.enabled || reload.failed_at.is_some() || balance >= reload.threshold_micros {
100 return None;
101 }
102 let wanted = (reload.target_micros - balance).max(MIN_RELOAD_MICROS);
103 let wanted = (wanted + MICROS_PER_DOLLAR - 1) / MICROS_PER_DOLLAR * MICROS_PER_DOLLAR;
104 let room = (reload.monthly_max_micros - reloaded_this_month).max(0) / MICROS_PER_DOLLAR * MICROS_PER_DOLLAR;
105 let amount = wanted.min(room);
106 (amount >= MIN_RELOAD_MICROS).then_some(amount)
107}
108
109/// What is wrong with auto-reload's settings, if anything.
110pub(crate) fn reload_invalid(threshold: i64, target: i64, monthly_max: i64) -> Option<&'static str> {
111 if threshold < 0 || target <= 0 || monthly_max <= 0 {
112 return Some("Amounts are in dollars, more than $0.");
113 }
114 if target < threshold + MIN_RELOAD_MICROS {
115 return Some("Reload to at least $10 more than the amount it reloads below.");
116 }
117 if target > MAX_TARGET_MICROS {
118 return Some("Reload to at most $1,000.");
119 }
120 if monthly_max < target - threshold {
121 return Some("The monthly maximum has to cover at least one reload.");
122 }
123 if monthly_max > MAX_MONTHLY_MICROS {
124 return Some("The monthly maximum is at most $10,000.");
125 }
126 if [threshold, target, monthly_max].iter().any(|m| m % MICROS_PER_DOLLAR != 0) {
127 return Some("Use whole dollars.");
128 }
129 None
130}
131
132/// Whether a purchase's page was paid for what was asked: Stripe says it
133/// is paid, and what was paid covers the credit (the fee is Stripe's).
134pub(crate) fn purchase_paid(payment_status: &str, amount_total: Option<u32>, credit_cents: u32) -> std::result::Result<(), String> {
135 if payment_status != "paid" {
136 return Err("The payment is not finished yet. It is credited as soon as Stripe says it was paid.".to_owned());
137 }
138 if amount_total.unwrap_or(0) < credit_cents {
139 return Err("Stripe says less was paid than the credit asked for; nothing was credited. Write to support@g1t.sh.".to_owned());
140 }
141 Ok(())
142}
143
144/// The id of the AI credit given for starting the plan: one per workspace,
145/// so however often the plan is recorded, it is given once.
146pub(crate) fn upgrade_reference(workspace: &str) -> String {
147 format!("crd_upgrade_{}", workspace.to_lowercase())
148}
149
150/// The agent rate on `tokens`, at `per_million` micros a million, rounded
151/// up to a whole millionth of a dollar.
152pub(crate) fn agent_rate_micros(tokens: u64, per_million: f64) -> i64 {
153 if tokens == 0 || !per_million.is_finite() || per_million <= 0.0 {
154 return 0;
155 }
156 (tokens as f64 * per_million / 1_000_000.0).ceil() as i64
157}
158
159/// Whether a workspace's runs on g1t's models need AI credit (or included
160/// usage) to start: on the plan, paying full or part price. A 100%
161/// discount pays for all of it; an enterprise is invoiced after use; a
162/// free workspace runs on its trial, which has its own limits.
163pub(crate) fn needs_credit(plan: PlanKind) -> bool {
164 plan == PlanKind::Paid
165}
166
167/// The refusal at $0.
168pub(crate) fn out_of_credit_message(workspace: &str, reload_failed: bool) -> String {
169 let reload = if reload_failed { " Auto-reload was turned off after its last charge failed." } else { "" };
170 format!(
171 "The {workspace} workspace is out of AI credit and has used this month's included usage, so g1t does not start new runs on its models.{reload} An owner can buy AI credit or turn on auto-reload at /{workspace}/-/billing#ai-credit."
172 )
173}
174
175#[derive(Deserialize)]
176struct ReloadRow {
177 enabled: i64,
178 threshold_micros: i64,
179 target_micros: i64,
180 monthly_max_micros: i64,
181 failed_at: Option<String>,
182 error: Option<String>,
183}
184
185#[derive(Deserialize)]
186struct Sum {
187 micros: Option<f64>,
188}
189
190#[derive(Deserialize)]
191struct Open {
192 workspace: String,
193 created_by: String,
194 amount_cents: u32,
195 fee_cents: Option<u32>,
196}
197
198impl Billing {
199 // --- The price book ----------------------------------------------------
200
201 /// The card fee, as the price book and the `card_fee` setting have it.
202 pub(crate) async fn card_fee(&self) -> Result<CardFee> {
203 #[derive(Deserialize)]
204 struct Row {
205 value: String,
206 }
207 let on = self
208 .db
209 .prepare("SELECT value FROM cost_settings WHERE key = 'card_fee'")
210 .first::<Row>(None)
211 .await?
212 .is_none_or(|row| row.value.trim() != "off");
213 let percent = self.price("card_fee_percent").await?.map_or(29_000.0, |(_, price)| price);
214 let fixed = self.price("card_fee_fixed").await?.map_or(300_000.0, |(_, price)| price);
215 Ok(CardFee { on, percent_micros: percent, fixed_cents: (fixed / 10_000.0).round().max(0.0) as u32 })
216 }
217
218 /// The agent rate per million tokens, at price.
219 pub(crate) async fn agent_rate(&self) -> Result<f64> {
220 Ok(self.price("agent_tokens").await?.map_or(0.0, |(_, price)| price))
221 }
222
223 /// The markup on a price-book meter, in percent.
224 async fn markup_of(&self, meter: &str) -> Result<Option<u32>> {
225 #[derive(Deserialize)]
226 struct Row {
227 markup_percent: u32,
228 }
229 Ok(self
230 .db
231 .prepare("SELECT markup_percent FROM prices WHERE meter = ?")
232 .bind(&[meter.into()])?
233 .first::<Row>(None)
234 .await?
235 .map(|row| row.markup_percent))
236 }
237
238 /// The markup on a model's provider price for agent runs: the price
239 /// book's `agent_models` (0 from 2026-10-08), or `MARGIN_PERCENT`
240 /// where the price book has no row.
241 pub(crate) async fn model_markup(&self) -> Result<u32> {
242 Ok(self.markup_of("agent_models").await?.unwrap_or(self.margin_percent))
243 }
244
245 /// The markup on AI Gateway's provider price: 0 while it is in beta.
246 pub(crate) async fn gateway_markup(&self) -> Result<u32> {
247 Ok(self.markup_of("gateway_models").await?.unwrap_or(0))
248 }
249
250 // --- Balances ------------------------------------------------------------
251
252 /// AI credit left: the open grants scoped to models, by kind.
253 pub(crate) async fn ai_balance(&self, workspace: &str) -> Result<(i64, i64, i64)> {
254 let credits = self.credits_of(workspace).await?;
255 let models: Vec<_> = credits.grants.iter().filter(|g| g.scope == "models").collect();
256 let purchased = models.iter().filter(|g| g.kind == CreditKind::Purchased).map(|g| g.left_micros).sum();
257 let given = models.iter().filter(|g| g.kind != CreditKind::Purchased).map(|g| g.left_micros).sum();
258 Ok((purchased + given, purchased, given))
259 }
260
261 /// What is owed now, with the balance `balance`: credit scoped to
262 /// models is not money for anything else, so what is left of it is
263 /// owed on top of a balance it props up.
264 pub(crate) async fn owed_with(&self, workspace: &str, balance: i64) -> Result<i64> {
265 let (left, _, _) = self.ai_balance(workspace).await?;
266 Ok((left - balance).max(0))
267 }
268
269 async fn reload_settings(&self, workspace: &str) -> Result<AiReload> {
270 let row = self
271 .db
272 .prepare("SELECT enabled, threshold_micros, target_micros, monthly_max_micros, failed_at, error FROM ai_reload WHERE workspace = ?")
273 .bind(&[workspace.into()])?
274 .first::<ReloadRow>(None)
275 .await?;
276 let reloaded = self.reloaded_this_month(workspace).await?;
277 Ok(match row {
278 Some(row) => AiReload {
279 enabled: row.enabled != 0,
280 threshold_micros: row.threshold_micros,
281 target_micros: row.target_micros,
282 monthly_max_micros: row.monthly_max_micros,
283 reloaded_micros: reloaded,
284 failed_at: row.failed_at,
285 error: row.error,
286 },
287 // The suggestion the form starts from: below $10, back to $25,
288 // at most $100 a month.
289 None => AiReload {
290 enabled: false,
291 threshold_micros: 10 * MICROS_PER_DOLLAR,
292 target_micros: 25 * MICROS_PER_DOLLAR,
293 monthly_max_micros: 100 * MICROS_PER_DOLLAR,
294 reloaded_micros: reloaded,
295 failed_at: None,
296 error: None,
297 },
298 })
299 }
300
301 async fn reloaded_this_month(&self, workspace: &str) -> Result<i64> {
302 let month = &rfc3339(now_ms())[..7];
303 Ok(self
304 .db
305 .prepare("SELECT SUM(amount_micros) AS micros FROM ai_reloads WHERE workspace = ? AND month = ? AND status IN ('paid', 'pending')")
306 .bind(&[workspace.into(), month.into()])?
307 .first::<Sum>(None)
308 .await?
309 .and_then(|s| s.micros)
310 .unwrap_or(0.0) as i64)
311 }
312
313 /// What the plan's included usage has left this month, on the plan.
314 async fn included_left(&self, workspace: &str) -> Result<i64> {
315 let month = crate::credits::month_of(&rfc3339(now_ms()));
316 let used = self.allowance_used("plan_credit", workspace, &month).await?;
317 Ok(crate::credits::left(self.plans.plan_included_micros, used))
318 }
319
320 // --- The page --------------------------------------------------------------
321
322 /// `ai_credit`: the workspace's AI credit, for its members.
323 pub(crate) async fn ai_credit(&self, a: AccountArgs) -> Result<Outcome<AiCredit>> {
324 let workspace = a.workspace.to_lowercase();
325 if !a.viewer.is_some_and(|viewer| viewer.is_member(&workspace)) {
326 return Ok(members_only());
327 }
328 Ok(Outcome::Ok(self.ai_credit_of(&workspace).await?))
329 }
330
331 pub(crate) async fn ai_credit_of(&self, workspace: &str) -> Result<AiCredit> {
332 let account = self.account_of(workspace).await?;
333 let plan = self.plan_kind_for(workspace, &account).await?;
334 let credits = self.credits_of(workspace).await?;
335 let grants: Vec<_> = credits.grants.into_iter().filter(|g| g.scope == "models").collect();
336 let purchased: i64 = grants.iter().filter(|g| g.kind == CreditKind::Purchased).map(|g| g.left_micros).sum();
337 let given: i64 = grants.iter().filter(|g| g.kind != CreditKind::Purchased).map(|g| g.left_micros).sum();
338 let balance = purchased + given;
339 let reload = self.reload_settings(workspace).await?;
340 let blocked = self.stripe.is_some()
341 && !self.free
342 && needs_credit(plan)
343 && balance <= 0
344 && self.included_left(workspace).await? <= 0;
345 Ok(AiCredit {
346 balance_micros: balance,
347 purchased_micros: purchased,
348 given_micros: given,
349 grants,
350 free_via_discount: account.terms.full_discount(),
351 postpaid: plan == PlanKind::Enterprise,
352 blocked,
353 can_buy: self.stripe.is_some() && plan == PlanKind::Paid,
354 presets_cents: PRESETS_CENTS.to_vec(),
355 min_cents: MIN_CENTS,
356 max_cents: MAX_CENTS,
357 card_fee: self.card_fee().await?,
358 reload,
359 agent_rate_micros: self.agent_rate().await?,
360 model_markup_percent: self.model_markup().await?,
361 gateway_markup_percent: self.gateway_markup().await?,
362 upgrade_credit_micros: UPGRADE_CREDIT_MICROS,
363 expires_days: EXPIRES_DAYS as u32,
364 })
365 }
366
367 // --- Buying --------------------------------------------------------------
368
369 /// `buy_ai_credit`: Stripe's page for a purchase.
370 pub(crate) async fn buy_ai_credit(&self, a: BuyAiCreditArgs) -> Result<Outcome<Checkout>> {
371 let workspace = a.workspace.to_lowercase();
372 if a.actor.role_in(&workspace) != Some(Role::Owner) {
373 return Ok(Outcome::fail(FailureCode::Forbidden, "Only an owner can buy AI credit for the workspace."));
374 }
375 let Some(stripe) = &self.stripe else {
376 return Ok(Outcome::fail(FailureCode::Conflict, "Payments are not set up on this g1t."));
377 };
378 if let Err(why) = amount_ok(a.amount_cents) {
379 return Ok(Outcome::fail(FailureCode::Invalid, why));
380 }
381 let account = self.account_of(&workspace).await?;
382 if account.terms.full_discount() {
383 return Ok(Outcome::fail(FailureCode::Conflict, format!("{workspace}'s AI usage is free under its discount: there is nothing to buy.")));
384 }
385 match self.plan_kind_for(&workspace, &account).await? {
386 PlanKind::Paid => {}
387 PlanKind::Enterprise => {
388 return Ok(Outcome::fail(FailureCode::Conflict, format!("{workspace} is invoiced for AI usage after use, through its enterprise.")));
389 }
390 _ => {
391 return Ok(Outcome::fail(FailureCode::PaymentRequired, format!("AI credit is for workspaces on the g1t plan. Start the plan for {workspace} first; it comes with $5 of AI credit.")));
392 }
393 }
394 let fee = card_fee_cents(a.amount_cents, &self.card_fee().await?);
395 let customer = self.row(&workspace).await?.and_then(|row| row.customer_id);
396 let purchase = |customer| crate::stripe::CreditPurchase {
397 workspace: &workspace,
398 credit_cents: a.amount_cents,
399 fee_cents: fee,
400 customer,
401 return_url: &a.return_url,
402 };
403 let started = match stripe.start_credit_checkout(&purchase(customer.as_deref())).await {
404 Err(error) if customer.is_some() && crate::stripe::is_missing(&error) => {
405 self.forget_customer(&workspace).await?;
406 stripe.start_credit_checkout(&purchase(None)).await
407 }
408 other => other,
409 };
410 self.page_opened(started, &workspace, a.amount_cents, fee, &a.actor.username, Some(AI_CREDIT)).await
411 }
412
413 /// `confirm_ai_credit`: back from Stripe's page.
414 pub(crate) async fn confirm_ai_credit(&self, a: ConfirmAiCreditArgs) -> Result<Outcome<AiCredit>> {
415 let workspace = a.workspace.to_lowercase();
416 if !a.viewer.is_some_and(|viewer| viewer.is_member(&workspace)) {
417 return Ok(members_only());
418 }
419 let mine = self
420 .db
421 .prepare("SELECT workspace FROM checkouts WHERE id = ? AND workspace = ? AND feature = ?")
422 .bind(&[a.session.as_str().into(), workspace.as_str().into(), AI_CREDIT.into()])?
423 .first::<serde_json::Value>(None)
424 .await?;
425 if mine.is_some() {
426 match self.settle_ai_credit(&a.session).await {
427 Ok(Ok(_)) => {}
428 Ok(Err(why)) => return Ok(Outcome::fail(FailureCode::Conflict, why)),
429 Err(error) => return Ok(Outcome::fail(FailureCode::Conflict, crate::stripe::friendly(&error))),
430 }
431 }
432 Ok(Outcome::Ok(self.ai_credit_of(&workspace).await?))
433 }
434
435 /// Credits a purchase whose page is paid, once. What happened, or why
436 /// it was not credited (yet).
437 pub(crate) async fn settle_ai_credit(&self, session_id: &str) -> Result<std::result::Result<String, String>> {
438 let Some(open) = self
439 .db
440 .prepare("SELECT workspace, created_by, amount_cents, fee_cents FROM checkouts WHERE id = ? AND feature = ? AND status = 'open'")
441 .bind(&[session_id.into(), AI_CREDIT.into()])?
442 .first::<Open>(None)
443 .await?
444 else {
445 return Ok(Ok("ignored: already credited or not AI credit".to_owned()));
446 };
447 let Some(stripe) = &self.stripe else { return Ok(Ok("ignored: payments off".to_owned())) };
448 let session = stripe.session(session_id).await?;
449 if let Err(why) = purchase_paid(&session.payment_status, session.amount_total, open.amount_cents) {
450 return Ok(Err(why));
451 }
452 let claimed = self
453 .db
454 .prepare("UPDATE checkouts SET status = 'paid' WHERE id = ? AND status = 'open' RETURNING id")
455 .bind(&[session_id.into()])?
456 .first::<serde_json::Value>(None)
457 .await?;
458 if claimed.is_none() {
459 return Ok(Ok("ignored: credited meanwhile".to_owned()));
460 }
461 let micros = i64::from(open.amount_cents) * 10_000;
462 let fee = i64::from(open.fee_cents.unwrap_or(0)) * 10_000;
463 self.grant_purchased(&open.workspace, session_id, micros, fee, &open.created_by, session.customer.as_deref())
464 .await?;
465 let extras = crate::tax::Extras { tax_cents: session.tax_cents(), fee_cents: fee / 10_000 };
466 self.record_extras(&open.workspace, session_id, session.payment_intent.as_deref(), extras, None).await?;
467 Ok(Ok(format!("{}: {} of AI credit bought", open.workspace, cents(micros))))
468 }
469
470 /// Enters bought AI credit: its grant and its ledger line, once for
471 /// `reference` (Stripe's id for the payment).
472 pub(crate) async fn grant_purchased(
473 &self,
474 workspace: &str,
475 reference: &str,
476 micros: i64,
477 fee_micros: i64,
478 by: &str,
479 customer: Option<&str>,
480 ) -> Result<bool> {
481 let now = now_ms();
482 let expires = rfc3339(now + EXPIRES_DAYS * DAY_MS);
483 let fee = if fee_micros > 0 { format!(" (card fee {} paid to Stripe)", cents(fee_micros)) } else { String::new() };
484 let note = format!("AI credit bought{fee}");
485 let inserted = self
486 .db
487 .prepare(
488 "INSERT OR IGNORE INTO credit_grants (id, workspace, kind, scope, source, amount_micros, note, expires_at, created_by, created_at)
489 VALUES (?, ?, 'purchased', 'models', 'purchase', ?, ?, ?, ?, ?) RETURNING id",
490 )
491 .bind(&[
492 reference.into(),
493 workspace.into(),
494 (micros as f64).into(),
495 note.as_str().into(),
496 expires.as_str().into(),
497 by.into(),
498 rfc3339(now).into(),
499 ])?
500 .first::<serde_json::Value>(None)
501 .await?;
502 if inserted.is_none() {
503 return Ok(false);
504 }
505 let description = format!("AI credit bought: {}, until {}", cents(micros), &expires[..10]);
506 self.enter(workspace, EntryKind::TopUp, micros, &description, reference, None, None, Some(by), customer).await?;
507 self.db
508 .prepare("UPDATE ledger SET credit_kind = 'purchased' WHERE reference = ?")
509 .bind(&[reference.into()])?
510 .run()
511 .await?;
512 let account = self.account_of(workspace).await?;
513 self.audit(&account.id, "ai_credit", &format!("{workspace}: {} of AI credit bought{fee}", cents(micros)), by).await?;
514 Ok(true)
515 }
516
517 /// The $5 of AI credit a workspace gets once, on starting the paid
518 /// plan. Promotional: given, not revenue. Never twice, and never for a
519 /// workspace whose discount pays for everything anyway.
520 pub(crate) async fn grant_upgrade_credit(&self, workspace: &str) -> Result<()> {
521 let workspace = workspace.to_lowercase();
522 if self.terms_of(&workspace).await?.full_discount() {
523 return Ok(());
524 }
525 let reference = upgrade_reference(&workspace);
526 let now = now_ms();
527 let expires = rfc3339(now + EXPIRES_DAYS * DAY_MS);
528 let inserted = self
529 .db
530 .prepare(
531 "INSERT OR IGNORE INTO credit_grants (id, workspace, kind, scope, source, amount_micros, note, expires_at, created_by, created_at)
532 VALUES (?, ?, 'promotional', 'models', 'upgrade', ?, 'AI credit for starting the g1t plan', ?, 'g1t', ?) RETURNING id",
533 )
534 .bind(&[
535 reference.as_str().into(),
536 workspace.as_str().into(),
537 (UPGRADE_CREDIT_MICROS as f64).into(),
538 expires.as_str().into(),
539 rfc3339(now).into(),
540 ])?
541 .first::<serde_json::Value>(None)
542 .await?;
543 if inserted.is_none() {
544 return Ok(());
545 }
546 let description = format!("Credit from g1t (promotional, until {}): AI credit for starting the g1t plan", &expires[..10]);
547 self.enter(&workspace, EntryKind::TopUp, UPGRADE_CREDIT_MICROS, &description, &reference, None, None, Some("g1t"), None).await?;
548 self.db
549 .prepare("UPDATE ledger SET credit_kind = 'promotional' WHERE reference = ?")
550 .bind(&[reference.as_str().into()])?
551 .run()
552 .await?;
553 let account = self.account_of(&workspace).await?;
554 self.audit(&account.id, "credit", &format!("{} promotional AI credit to {workspace} for starting the plan", cents(UPGRADE_CREDIT_MICROS)), "g1t")
555 .await?;
556 Ok(())
557 }
558
559 // --- At $0 -----------------------------------------------------------------
560
561 /// Why a run on g1t's models cannot start for want of AI credit, if it
562 /// cannot. Auto-reload, when on, is tried first.
563 pub(crate) async fn ai_refusal(&self, workspace: &str) -> Result<Option<String>> {
564 if self.stripe.is_none() || self.free {
565 return Ok(None);
566 }
567 let account = self.account_of(workspace).await?;
568 if !needs_credit(self.plan_kind_for(workspace, &account).await?) {
569 return Ok(None);
570 }
571 if self.included_left(workspace).await? > 0 {
572 return Ok(None);
573 }
574 let (balance, _, _) = self.ai_balance(workspace).await?;
575 if balance > 0 {
576 return Ok(None);
577 }
578 let reload = self.reload_settings(workspace).await?;
579 if reload.enabled && reload.failed_at.is_none() {
580 if let Err(error) = self.reload_now(workspace).await {
581 worker::console_error!("{workspace}: auto-reload at a run's start failed: {error}");
582 }
583 if self.ai_balance(workspace).await?.0 > 0 {
584 return Ok(None);
585 }
586 }
587 let failed = self.reload_settings(workspace).await?.failed_at.is_some();
588 Ok(Some(out_of_credit_message(workspace, failed)))
589 }
590
591 // --- Auto-reload ---------------------------------------------------------
592
593 /// `set_ai_reload`: owners only.
594 pub(crate) async fn set_ai_reload(&self, a: SetAiReloadArgs) -> Result<Outcome<AiCredit>> {
595 let workspace = a.workspace.to_lowercase();
596 if a.actor.role_in(&workspace) != Some(Role::Owner) {
597 return Ok(Outcome::fail(FailureCode::Forbidden, "Only an owner can change auto-reload."));
598 }
599 if let Some(why) = reload_invalid(a.threshold_micros, a.target_micros, a.monthly_max_micros) {
600 return Ok(Outcome::fail(FailureCode::Invalid, why));
601 }
602 if a.enabled {
603 let credit = self.ai_credit_of(&workspace).await?;
604 if !credit.can_buy {
605 return Ok(Outcome::fail(FailureCode::Conflict, "Auto-reload is for workspaces on the g1t plan that buy AI credit."));
606 }
607 let has_card = match (&self.stripe, self.row(&workspace).await?.and_then(|row| row.customer_id)) {
608 (Some(stripe), Some(customer)) => stripe.default_payment_method(&customer).await.ok().flatten().is_some(),
609 _ => false,
610 };
611 if !has_card {
612 return Ok(Outcome::fail(FailureCode::Conflict, "Auto-reload charges the workspace's saved card, and it has none. Buy AI credit once, or add a card on Stripe's billing page, first."));
613 }
614 }
615 let now = rfc3339(now_ms());
616 self.db
617 .prepare(
618 "INSERT INTO ai_reload (workspace, enabled, threshold_micros, target_micros, monthly_max_micros, updated_by, updated_at, failed_at, error)
619 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, NULL, NULL)
620 ON CONFLICT (workspace) DO UPDATE SET enabled = ?2, threshold_micros = ?3, target_micros = ?4, monthly_max_micros = ?5,
621 updated_by = ?6, updated_at = ?7, failed_at = NULL, error = NULL",
622 )
623 .bind(&[
624 workspace.as_str().into(),
625 i32::from(a.enabled).into(),
626 (a.threshold_micros as f64).into(),
627 (a.target_micros as f64).into(),
628 (a.monthly_max_micros as f64).into(),
629 a.actor.username.as_str().into(),
630 now.as_str().into(),
631 ])?
632 .run()
633 .await?;
634 let account = self.account_of(&workspace).await?;
635 self.audit(
636 &account.id,
637 "ai_reload",
638 &format!(
639 "{workspace}: auto-reload {}: below {}, back to {}, at most {} a month",
640 if a.enabled { "on" } else { "off" },
641 cents(a.threshold_micros),
642 cents(a.target_micros),
643 cents(a.monthly_max_micros)
644 ),
645 &a.actor.username,
646 )
647 .await?;
648 // Below the threshold already: reload now rather than at the next run.
649 if a.enabled
650 && let Err(error) = self.reload_now(&workspace).await
651 {
652 worker::console_error!("{workspace}: auto-reload right after turning it on failed: {error}");
653 }
654 Ok(Outcome::Ok(self.ai_credit_of(&workspace).await?))
655 }
656
657 /// Every workspace with auto-reload on that is below its threshold,
658 /// reloaded: each cron run.
659 pub(crate) async fn reload_ai_credit(&self) -> Result<u32> {
660 if self.stripe.is_none() {
661 return Ok(0);
662 }
663 #[derive(Deserialize)]
664 struct Row {
665 workspace: String,
666 }
667 let due = self
668 .db
669 .prepare("SELECT workspace FROM ai_reload WHERE enabled = 1 AND failed_at IS NULL LIMIT 100")
670 .all()
671 .await?
672 .results::<Row>()?;
673 let mut done = 0;
674 for Row { workspace } in due {
675 match self.reload_now(&workspace).await {
676 Ok(Some(_)) => done += 1,
677 Ok(None) => {}
678 Err(error) => worker::console_error!("{workspace}: auto-reload failed: {error}"),
679 }
680 }
681 Ok(done)
682 }
683
684 /// Reloads the workspace's AI credit if it is below its threshold.
685 /// What was reloaded, or None.
686 pub(crate) async fn reload_now(&self, workspace: &str) -> Result<Option<i64>> {
687 let Some(stripe) = &self.stripe else { return Ok(None) };
688 let reload = self.reload_settings(workspace).await?;
689 let (balance, _, _) = self.ai_balance(workspace).await?;
690 let Some(amount) = reload_amount(&reload, balance, reload.reloaded_micros) else {
691 return Ok(None);
692 };
693 let Some(customer) = self.row(workspace).await?.and_then(|row| row.customer_id) else {
694 self.reload_failed(workspace, None, amount, "the workspace has no saved card").await?;
695 return Ok(None);
696 };
697 let method = match stripe.default_payment_method(&customer).await {
698 Ok(Some(method)) => method,
699 Ok(None) => {
700 self.reload_failed(workspace, None, amount, "the workspace has no saved card").await?;
701 return Ok(None);
702 }
703 Err(error) => return Err(error),
704 };
705 let month = rfc3339(now_ms())[..7].to_owned();
706 // One key per attempt: the month and how many reloads came before.
707 // A retry after a crash is the same attempt, so the same payment.
708 #[derive(Deserialize)]
709 struct Count {
710 n: Option<f64>,
711 }
712 let before = self
713 .db
714 .prepare("SELECT COUNT(*) AS n FROM ai_reloads WHERE workspace = ? AND month = ? AND status = 'paid'")
715 .bind(&[workspace.into(), month.as_str().into()])?
716 .first::<Count>(None)
717 .await?
718 .and_then(|c| c.n)
719 .unwrap_or(0.0) as u32;
720 let key = format!("reload/{workspace}/{month}/{}", before + 1);
721 let credit_cents = (amount / 10_000) as u32;
722 let fee_cents = card_fee_cents(credit_cents, &self.card_fee().await?);
723 self.db
724 .prepare(
725 "INSERT OR IGNORE INTO ai_reloads (id, workspace, month, amount_micros, fee_micros, status, created_at)
726 VALUES (?, ?, ?, ?, ?, 'pending', ?)",
727 )
728 .bind(&[
729 key.as_str().into(),
730 workspace.into(),
731 month.as_str().into(),
732 (amount as f64).into(),
733 (f64::from(fee_cents) * 10_000.0).into(),
734 rfc3339(now_ms()).into(),
735 ])?
736 .run()
737 .await?;
738 let untaxed = crate::stripe::SavedCharge {
739 workspace,
740 customer: &customer,
741 payment_method: &method.id,
742 credit_cents,
743 fee_cents,
744 tax_cents: 0,
745 tax_calculation: None,
746 key: &key,
747 };
748 // No Checkout page or invoice works the tax out here: Stripe Tax's
749 // calculation does, for the customer's saved address. Without one,
750 // nothing is charged and the owners are asked for it.
751 let calculation = match stripe.tax_calculation(&untaxed).await {
752 Ok(calculation) => calculation,
753 Err(error) if crate::stripe::is_tax_location_error(&error) => {
754 self.tax_address_needed(workspace).await?;
755 self.reload_failed(workspace, Some(&key), amount, "Stripe needs the workspace's billing address to work out tax; add it under Invoice details")
756 .await?;
757 return Ok(None);
758 }
759 Err(error) => return Err(error),
760 };
761 let tax_cents = u32::try_from(calculation.tax_amount_exclusive.max(0)).unwrap_or(0);
762 let charge = crate::stripe::SavedCharge { tax_cents, tax_calculation: Some(&calculation.id), ..untaxed };
763 let paid = match stripe.charge_saved(&charge).await {
764 Ok(intent) if intent["status"].as_str() == Some("succeeded") => intent["id"].as_str().map(str::to_owned),
765 Ok(intent) => {
766 let status = intent["status"].as_str().unwrap_or("unknown").to_owned();
767 self.reload_failed(workspace, Some(&key), amount, &format!("the card needs the bank's approval ({status})")).await?;
768 return Ok(None);
769 }
770 Err(error) if crate::stripe::is_card_error(&error) => {
771 self.reload_failed(workspace, Some(&key), amount, &crate::stripe::friendly(&error)).await?;
772 return Ok(None);
773 }
774 Err(error) => return Err(error),
775 };
776 let Some(intent) = paid else { return Ok(None) };
777 self.db
778 .prepare("UPDATE ai_reloads SET status = 'paid', payment_intent = ? WHERE id = ?")
779 .bind(&[intent.as_str().into(), key.as_str().into()])?
780 .run()
781 .await?;
782 self.grant_purchased(workspace, &intent, amount, i64::from(fee_cents) * 10_000, "g1t", Some(&customer)).await?;
783 // Recorded with Stripe Tax once paid, so it is reported and filed;
784 // a failure is logged and the payment stands.
785 let transaction = match stripe.record_tax(&calculation.id, &intent).await {
786 Ok(id) => Some(id),
787 Err(error) => {
788 worker::console_error!("{workspace}: the tax on auto-reload {intent} was not recorded with Stripe Tax: {error}");
789 None
790 }
791 };
792 let extras = crate::tax::Extras { tax_cents: i64::from(tax_cents), fee_cents: i64::from(fee_cents) };
793 self.record_extras(workspace, &intent, Some(&intent), extras, transaction.as_deref()).await?;
794 self.tax_address_given(workspace).await?;
795 Ok(Some(amount))
796 }
797
798 /// A reload that could not be charged: auto-reload is turned off, and
799 /// the owners are told.
800 async fn reload_failed(&self, workspace: &str, key: Option<&str>, amount: i64, why: &str) -> Result<()> {
801 let now = rfc3339(now_ms());
802 if let Some(key) = key {
803 self.db
804 .prepare("UPDATE ai_reloads SET status = 'failed', error = ? WHERE id = ?")
805 .bind(&[why.into(), key.into()])?
806 .run()
807 .await?;
808 }
809 self.db
810 .prepare("UPDATE ai_reload SET enabled = 0, failed_at = ?, error = ? WHERE workspace = ?")
811 .bind(&[now.as_str().into(), why.into(), workspace.into()])?
812 .run()
813 .await?;
814 let account = self.account_of(workspace).await?;
815 self.audit(&account.id, "ai_reload_failed", &format!("{workspace}: auto-reload of {} failed ({why}); turned off", cents(amount)), "g1t")
816 .await?;
817 if let Some(identity) = &self.identity {
818 crate::limits::notify_with(
819 identity,
820 workspace,
821 &format!("g1t: auto-reload for {workspace} failed and is off"),
822 &format!(
823 "g1t tried to reload {} of AI credit for {workspace} and could not: {why}. Auto-reload is off until an owner turns it on again. Until then, runs on g1t's models stop once the AI credit is spent."
824 , cents(amount)),
825 "Open billing",
826 &format!("https://g1t.sh/{workspace}/-/billing#ai-credit"),
827 "You get this because you own this workspace on g1t. AI credit is explained at https://docs.g1t.sh/guides/usage-and-billing/#ai-credit",
828 )
829 .await;
830 }
831 Ok(())
832 }
833
834 // --- The agent rate --------------------------------------------------------
835
836 /// How much each kind of token counts toward the agent rate: the price
837 /// book's `agent_token_weight_*` meters, a token's weight in millionths
838 /// (1,000,000 is 1). A meter missing counts 1.
839 pub(crate) async fn token_weights(&self) -> Result<TokenWeights> {
840 let mut weights = TokenWeights::default();
841 for (meter, slot) in [
842 ("agent_token_weight_input", &mut weights.input),
843 ("agent_token_weight_output", &mut weights.output),
844 ("agent_token_weight_cache_read", &mut weights.cache_read),
845 ("agent_token_weight_cache_write", &mut weights.cache_write),
846 ] {
847 if let Some((cost, _)) = self.price(meter).await? {
848 *slot = weight_of(cost);
849 }
850 }
851 Ok(weights)
852 }
853
854 /// Charges a run's agent rate for the tokens counted since it was last
855 /// charged, once each: when the run reports and again when it is
856 /// settled, so tokens counted late are charged too. `reported` is what
857 /// the run's harness counted; the rate is charged on no fewer. Tokens
858 /// are weighted by kind (`token_weights`), and `runs.agent_tokens`
859 /// keeps the weighted tokens charged so far.
860 ///
861 /// On the workspace's own provider the model is not g1t's to charge,
862 /// but the agent rate is, on its own meter (`agent_tokens_own`) and its
863 /// own line (`<run>/agent-own`), so Usage and the statement show it as
864 /// the agent rate on the workspace's own model key.
865 pub(crate) async fn charge_agent_rate(&self, run_id: &str, run: &RunRow, reported: Option<RunTokens>) -> Result<()> {
866 if self.stripe.is_none() {
867 return Ok(());
868 }
869 #[derive(Deserialize)]
870 struct Row {
871 session_id: Option<String>,
872 agent_tokens: Option<f64>,
873 created_at: String,
874 }
875 let Some(row) = self
876 .db
877 .prepare("SELECT session_id, agent_tokens, created_at FROM runs WHERE id = ?")
878 .bind(&[run_id.into()])?
879 .first::<Row>(None)
880 .await?
881 else {
882 return Ok(());
883 };
884 #[derive(Deserialize)]
885 struct Kinds {
886 input: Option<f64>,
887 output: Option<f64>,
888 cache_read: Option<f64>,
889 cache_write: Option<f64>,
890 }
891 let counted = match &row.session_id {
892 Some(session) => self
893 .db
894 .prepare(
895 "SELECT SUM(input) AS input, SUM(output) AS output, SUM(cache_read) AS cache_read, SUM(cache_write) AS cache_write
896 FROM token_usage WHERE workspace = ? AND session = ? AND day >= ?",
897 )
898 .bind(&[run.workspace.as_str().into(), session.as_str().into(), row.created_at[..10].into()])?
899 .first::<Kinds>(None)
900 .await?
901 .map(|k| {
902 let n = |v: Option<f64>| v.unwrap_or(0.0).max(0.0) as u64;
903 RunTokens { input: n(k.input), output: n(k.output), cache_read: n(k.cache_read), cache_write: n(k.cache_write) }
904 })
905 .unwrap_or_default(),
906 None => RunTokens::default(),
907 };
908 let weights = self.token_weights().await?;
909 let charged = row.agent_tokens.unwrap_or(0.0) as u64;
910 let Some(total) = tokens_to_charge(
911 weighted(&counted, &weights),
912 reported.map_or(0, |tokens| weighted(&tokens, &weights)),
913 charged,
914 ) else {
915 return Ok(());
916 };
917 // Claimed first: two callers never charge the same tokens.
918 let claimed = self
919 .db
920 .prepare("UPDATE runs SET agent_tokens = ?1 WHERE id = ?2 AND agent_tokens = ?3 RETURNING id")
921 .bind(&[(total as f64).into(), run_id.into(), (charged as f64).into()])?
922 .first::<serde_json::Value>(None)
923 .await?;
924 if claimed.is_none() {
925 return Ok(());
926 }
927 let own = run.own_provider();
928 let meter = agent_rate_meter(own);
929 let tokens = total - charged;
930 let per_million = self.price(meter).await?.map_or(0.0, |(_, price)| price);
931 let base = agent_rate_micros(tokens, per_million);
932 // Before the rate takes effect: counted, and nothing charged.
933 if base == 0 {
934 return Ok(());
935 }
936 let (charge, terms_note, discount) = self.charged(&run.workspace, base).await?;
937 let now = rfc3339(now_ms());
938 let eligible = crate::credits::eligible_for(Some(g1t_contracts::billing::ComputeKind::Agent), None);
939 let drawn = self.draw(&run.workspace, charge, &crate::credits::month_of(&now), &eligible).await?;
940 let reference = agent_rate_reference(run_id, own, charged, total);
941 let what = match run.task.as_str() {
942 "plan" => format!("planning for {}", run.repo),
943 "review" => format!("the review of {}#{}", run.repo, run.number),
944 "update" => format!("catching up {}#{}", run.repo, run.number),
945 _ => format!("work on {}#{}", run.repo, run.number),
946 };
947 let label = if own { "g1t agent rate, your own model key" } else { "g1t agent rate" };
948 let counted_as = if weights.is_flat() { "tokens".to_owned() } else { format!("weighted tokens ({})", weights.describe()) };
949 let description = format!("{label}: {} {counted_as} for {what}{terms_note}{}", crate::features::thousands(tokens), drawn.note());
950 // g1t's own charge, even on the workspace's provider: it counts
951 // toward limits and spend like any other.
952 let line = RunRow { billed_to: None, ..run.clone() };
953 self.enter(&run.workspace, EntryKind::Usage, -(charge - drawn.total()), &description, &reference, Some(&line), Some(0), None, None)
954 .await?;
955 self.db
956 .prepare("UPDATE ledger SET quantity = ?, price_version = ? WHERE reference = ?")
957 .bind(&[(tokens as f64).into(), optional(self.version_now(meter).await?.as_deref()), reference.as_str().into()])?
958 .run()
959 .await?;
960 self.record_drawn(&reference, &drawn).await?;
961 self.record_discount(&reference, discount).await?;
962 self.count_spend(&run.workspace, 0, charge - drawn.total(), &drawn).await;
963 Ok(())
964 }
965}
966
967/// How much each kind of token counts toward the agent rate. All 1 by
968/// default: every token counts once.
969#[derive(Clone, Copy, Debug, PartialEq)]
970pub(crate) struct TokenWeights {
971 pub input: f64,
972 pub output: f64,
973 pub cache_read: f64,
974 pub cache_write: f64,
975}
976
977impl Default for TokenWeights {
978 fn default() -> Self {
979 TokenWeights { input: 1.0, output: 1.0, cache_read: 1.0, cache_write: 1.0 }
980 }
981}
982
983impl TokenWeights {
984 /// Every token counts once.
985 pub(crate) fn is_flat(&self) -> bool {
986 *self == TokenWeights::default()
987 }
988
989 /// `input ×1, output ×1, cache reads ×0.1, cache writes ×1`.
990 pub(crate) fn describe(&self) -> String {
991 let w = |v: f64| {
992 let text = format!("{v:.3}");
993 text.trim_end_matches('0').trim_end_matches('.').to_owned()
994 };
995 format!(
996 "input ×{}, output ×{}, cache reads ×{}, cache writes ×{}",
997 w(self.input),
998 w(self.output),
999 w(self.cache_read),
1000 w(self.cache_write)
1001 )
1002 }
1003}
1004
1005/// A weight from its price-book figure, in millionths; never below 0.
1006pub(crate) fn weight_of(micros: f64) -> f64 {
1007 if micros.is_finite() { (micros / 1_000_000.0).max(0.0) } else { 1.0 }
1008}
1009
1010/// A run's tokens as the agent rate counts them, each kind at its weight,
1011/// rounded down to a whole token.
1012pub(crate) fn weighted(tokens: &RunTokens, weights: &TokenWeights) -> u64 {
1013 let sum = tokens.input as f64 * weights.input
1014 + tokens.output as f64 * weights.output
1015 + tokens.cache_read as f64 * weights.cache_read
1016 + tokens.cache_write as f64 * weights.cache_write;
1017 sum.max(0.0).floor() as u64
1018}
1019
1020/// The price-book meter a run's agent rate is on: its own for runs on the
1021/// workspace's own model key, so it can be priced and shown apart.
1022pub(crate) fn agent_rate_meter(own_provider: bool) -> &'static str {
1023 if own_provider { "agent_tokens_own" } else { "agent_tokens" }
1024}
1025
1026/// The tokens a run's agent rate covers now: the more of what the proxy
1027/// counted and what the harness reported, when that is more than was
1028/// charged already. None when there is nothing new.
1029pub(crate) fn tokens_to_charge(counted: u64, reported: u64, charged: u64) -> Option<u64> {
1030 let total = counted.max(reported);
1031 (total > charged).then_some(total)
1032}
1033
1034/// The ledger reference of an agent-rate line: `<run>/agent` the first
1035/// time, `<run>/agent/<tokens>` for tokens counted later; `agent-own` on
1036/// the workspace's own model key.
1037pub(crate) fn agent_rate_reference(run_id: &str, own_provider: bool, charged: u64, total: u64) -> String {
1038 let part = if own_provider { "agent-own" } else { "agent" };
1039 if charged == 0 { format!("{run_id}/{part}") } else { format!("{run_id}/{part}/{total}") }
1040}
1041
1042#[cfg(test)]
1043mod tests {
1044 use super::*;
1045
1046 fn fee(on: bool) -> CardFee {
1047 CardFee { on, percent_micros: 29_000.0, fixed_cents: 30 }
1048 }
1049
1050 #[test]
1051 fn the_card_fee_is_stripes_fee_grossed_up_and_off_when_switched_off() {
1052 // $25 of credit: ($25 + $0.30) / 0.971 = $26.06, so a $1.06 fee.
1053 assert_eq!(card_fee_cents(2_500, &fee(true)), 106);
1054 // What is left after Stripe's 2.9% + 30¢ is at least the credit.
1055 for credit in [1_000u32, 2_500, 5_000, 10_000, 100_000] {
1056 let total = credit + card_fee_cents(credit, &fee(true));
1057 let net = f64::from(total) - (f64::from(total) * 0.029).round() - 30.0;
1058 assert!(net >= f64::from(credit) - 1.0, "{credit}: {total} leaves {net}");
1059 }
1060 assert_eq!(card_fee_cents(2_500, &fee(false)), 0);
1061 assert_eq!(card_fee_cents(0, &fee(true)), 0);
1062 }
1063
1064 #[test]
1065 fn amounts_are_whole_dollars_from_ten_to_a_thousand() {
1066 assert!(amount_ok(1_000).is_ok() && amount_ok(100_000).is_ok() && amount_ok(2_500).is_ok());
1067 assert!(amount_ok(999).is_err() && amount_ok(100_100).is_err() && amount_ok(1_050).is_err());
1068 }
1069
1070 fn reload(threshold: i64, target: i64, max: i64) -> AiReload {
1071 AiReload { enabled: true, threshold_micros: threshold, target_micros: target, monthly_max_micros: max, ..AiReload::default() }
1072 }
1073
1074 const D: i64 = MICROS_PER_DOLLAR;
1075
1076 #[test]
1077 fn auto_reload_tops_up_to_the_target_below_the_threshold_within_the_monthly_maximum() {
1078 let r = reload(10 * D, 25 * D, 100 * D);
1079 // Above the threshold: nothing.
1080 assert_eq!(reload_amount(&r, 10 * D, 0), None);
1081 // Below it: back to the target, in whole dollars.
1082 assert_eq!(reload_amount(&r, 9 * D, 0), Some(16 * D));
1083 assert_eq!(reload_amount(&r, 9_500_000, 0), Some(16 * D));
1084 // Owing more than the target is still a reload to the target.
1085 assert_eq!(reload_amount(&r, -3 * D, 0), Some(28 * D));
1086 // Never less than $10.
1087 assert_eq!(reload_amount(&reload(10 * D, 12 * D, 100 * D), 9 * D, 0), Some(10 * D));
1088 // The monthly maximum caps it, and below $10 of room nothing is done.
1089 assert_eq!(reload_amount(&r, 0, 90 * D), Some(10 * D));
1090 assert_eq!(reload_amount(&r, 0, 95 * D), None);
1091 assert_eq!(reload_amount(&r, 0, 100 * D), None);
1092 }
1093
1094 #[test]
1095 fn a_failed_or_disabled_reload_does_nothing() {
1096 let off = AiReload { enabled: false, ..reload(10 * D, 25 * D, 100 * D) };
1097 assert_eq!(reload_amount(&off, 0, 0), None);
1098 let failed = AiReload { failed_at: Some("2026-10-08T00:00:00Z".into()), ..reload(10 * D, 25 * D, 100 * D) };
1099 assert_eq!(reload_amount(&failed, 0, 0), None);
1100 }
1101
1102 #[test]
1103 fn auto_reload_settings_are_checked() {
1104 assert_eq!(reload_invalid(10 * D, 25 * D, 100 * D), None);
1105 assert!(reload_invalid(10 * D, 15 * D, 100 * D).is_some());
1106 assert!(reload_invalid(10 * D, 2_000 * D, 10_000 * D).is_some());
1107 assert!(reload_invalid(10 * D, 25 * D, 10 * D).is_some());
1108 assert!(reload_invalid(10 * D, 25 * D, 20_000 * D).is_some());
1109 assert!(reload_invalid(10 * D, 25_500_000, 100 * D).is_some());
1110 assert!(reload_invalid(-1, 25 * D, 100 * D).is_some());
1111 }
1112
1113 #[test]
1114 fn the_agent_rate_is_per_million_tokens_rounded_up() {
1115 // $0.25 a million: 2 million tokens are 50 cents.
1116 assert_eq!(agent_rate_micros(2_000_000, 250_000.0), 500_000);
1117 assert_eq!(agent_rate_micros(1, 250_000.0), 1);
1118 assert_eq!(agent_rate_micros(0, 250_000.0), 0);
1119 // Before it takes effect the price book says 0.
1120 assert_eq!(agent_rate_micros(5_000_000, 0.0), 0);
1121 }
1122
1123 #[test]
1124 fn own_key_runs_are_charged_the_agent_rate_on_their_own_meter_and_line() {
1125 assert_eq!(agent_rate_meter(true), "agent_tokens_own");
1126 assert_eq!(agent_rate_meter(false), "agent_tokens");
1127 assert_eq!(agent_rate_reference("run_1", true, 0, 900), "run_1/agent-own");
1128 assert_eq!(agent_rate_reference("run_1", true, 900, 1_200), "run_1/agent-own/1200");
1129 assert_eq!(agent_rate_reference("run_1", false, 0, 900), "run_1/agent");
1130 assert_eq!(agent_rate_reference("run_1", false, 900, 1_200), "run_1/agent/1200");
1131 }
1132
1133 fn tokens(input: u64, output: u64, cache_read: u64, cache_write: u64) -> RunTokens {
1134 RunTokens { input, output, cache_read, cache_write }
1135 }
1136
1137 #[test]
1138 fn every_token_counts_once_by_default() {
1139 let flat = TokenWeights::default();
1140 assert!(flat.is_flat());
1141 assert_eq!(weighted(&tokens(1_000, 500, 90_000, 5_000), &flat), 96_500);
1142 assert_eq!(weighted(&RunTokens::default(), &flat), 0);
1143 }
1144
1145 #[test]
1146 fn cache_reads_can_count_for_a_tenth() {
1147 let tenth = TokenWeights { cache_read: weight_of(100_000.0), ..TokenWeights::default() };
1148 assert!(!tenth.is_flat());
1149 // 1,000 + 500 + 9,000 + 5,000.
1150 assert_eq!(weighted(&tokens(1_000, 500, 90_000, 5_000), &tenth), 15_500);
1151 // Rounded down to a whole token.
1152 assert_eq!(weighted(&tokens(0, 0, 15, 0), &tenth), 1);
1153 assert_eq!(tenth.describe(), "input ×1, output ×1, cache reads ×0.1, cache writes ×1");
1154 assert_eq!(TokenWeights::default().describe(), "input ×1, output ×1, cache reads ×1, cache writes ×1");
1155 }
1156
1157 #[test]
1158 fn a_weight_is_never_below_nothing() {
1159 assert_eq!(weight_of(1_000_000.0), 1.0);
1160 assert_eq!(weight_of(1_250_000.0), 1.25);
1161 assert_eq!(weight_of(-5.0), 0.0);
1162 assert_eq!(weight_of(f64::NAN), 1.0);
1163 }
1164
1165 #[test]
1166 fn the_rate_covers_the_more_of_what_was_counted_and_reported_once() {
1167 // The proxy counted nothing (no session, or its reports were lost):
1168 // the harness's count is charged.
1169 assert_eq!(tokens_to_charge(0, 5_000, 0), Some(5_000));
1170 // The proxy counted more: its count.
1171 assert_eq!(tokens_to_charge(6_000, 5_000, 0), Some(6_000));
1172 // Charged already: only what is new, and nothing twice.
1173 assert_eq!(tokens_to_charge(6_000, 5_000, 6_000), None);
1174 assert_eq!(tokens_to_charge(7_000, 0, 6_000), Some(7_000));
1175 assert_eq!(tokens_to_charge(0, 0, 0), None);
1176 }
1177
1178 #[test]
1179 fn only_workspaces_paying_on_the_plan_need_ai_credit() {
1180 assert!(needs_credit(PlanKind::Paid));
1181 assert!(!needs_credit(PlanKind::Internal));
1182 assert!(!needs_credit(PlanKind::Enterprise));
1183 assert!(!needs_credit(PlanKind::Free));
1184 assert!(out_of_credit_message("acme", false).contains("/acme/-/billing#ai-credit"));
1185 assert!(out_of_credit_message("acme", true).contains("Auto-reload was turned off"));
1186 }
1187
1188 #[test]
1189 fn a_purchase_is_credited_only_once_paid_and_only_for_what_was_paid() {
1190 assert!(purchase_paid("paid", Some(2_606), 2_500).is_ok());
1191 assert!(purchase_paid("unpaid", Some(2_606), 2_500).unwrap_err().contains("not finished"));
1192 assert!(purchase_paid("paid", Some(2_000), 2_500).is_err());
1193 assert!(purchase_paid("paid", None, 2_500).is_err());
1194 // The grant's id is the page's id and the ledger's reference is
1195 // unique, and the checkout row is claimed open → paid before either
1196 // is written: the webhook and the person coming back credit once.
1197 let source = include_str!("ai.rs");
1198 assert!(source.contains("UPDATE checkouts SET status = 'paid' WHERE id = ? AND status = 'open' RETURNING id"));
1199 assert!(source.contains("INSERT OR IGNORE INTO credit_grants"));
1200 }
1201
1202 #[test]
1203 fn the_upgrade_credit_is_one_grant_per_workspace() {
1204 assert_eq!(upgrade_reference("Acme"), upgrade_reference("acme"));
1205 assert!(upgrade_reference("acme").starts_with("crd"), "given, never a payment");
1206 assert_eq!(UPGRADE_CREDIT_MICROS, 5_000_000);
1207 }
1208}