Sign in

pulls/pr_01m47d15m3e54sn21z27rpy5n9

public

pr_01m47d15m3e54sn21z27rpy5n9/services/billing/src/lib.rs

601 lines21,413 bytesCodeBlame

Pick any line to see why it is the way it is: the commit, the pull request and issue it came from, and what the agent was thinking.

Agents as a team: lifecycle, merge queue, billing and a new shell1//! The billing service: what agents cost, charged to the workspace they
2//! worked for.
3//!
4//! A workspace buys credit with a card. Before the runner starts an agent
5//! it asks here, and is refused if the workspace has none. When the
6//! agent's sandbox finishes it reports what the model cost, and that plus
7//! g1t's margin comes off the balance. Every change is a ledger entry, and
8//! a balance is always the sum of its ledger.
9//!
10//! Without a card processor configured the service says so and charges
11//! nothing, so that g1t still runs where billing has not been set up.
12//!
13//! Reached only through service bindings; see `g1t_contracts::billing` for
14//! the methods and their arguments.
15
16mod stripe;
17
18use g1t_contracts::billing::*;
19use g1t_contracts::time::rfc3339;
20use g1t_contracts::{FailureCode, Outcome, Role, new_id};
21use g1t_kit::{args, now_ms, reply, rpc_method};
22use serde::Deserialize;
23use sha2::{Digest, Sha256};
24use worker::wasm_bindgen::JsValue;
25use worker::{Context, D1Database, Env, Request, Response, Result, event};
26
27use stripe::Stripe;
28
29const MIN_TOP_UP_CENTS: u32 = 500;
30const MAX_TOP_UP_CENTS: u32 = 50_000;
31const LEDGER_PAGE: u32 = 100;
32/// A run's reported cost is believed up to this much. A sandbox cannot
33/// spend more in the time it has, so anything above is a fault.
34const MAX_RUN_COST_USD: f64 = 100.0;
35
36/// What a run is charged: its cost plus the margin, rounded up to a whole
37/// millionth of a dollar.
38pub fn charge_micros(cost_usd: f64, margin_percent: u32) -> i64 {
39 let cost_micros = (cost_usd.clamp(0.0, MAX_RUN_COST_USD) * MICROS_PER_DOLLAR as f64).ceil();
40 (cost_micros * f64::from(100 + margin_percent) / 100.0).ceil() as i64
41}
42
43fn hash(token: &str) -> String {
44 hex::encode(Sha256::digest(token.as_bytes()))
45}
46
47fn optional(value: Option<&str>) -> JsValue {
48 value.map_or(JsValue::NULL, JsValue::from)
49}
50
51#[derive(Deserialize)]
52struct AccountRow {
53 balance_micros: i64,
54 customer_id: Option<String>,
55}
56
57#[derive(Deserialize)]
58struct LedgerRow {
59 id: String,
60 kind: EntryKind,
61 amount_micros: i64,
62 description: String,
63 repo: Option<String>,
64 number: Option<u32>,
65 task: Option<String>,
66 model: Option<String>,
67 created_by: Option<String>,
68 created_at: String,
69}
70
71impl From<LedgerRow> for LedgerEntry {
72 fn from(row: LedgerRow) -> Self {
73 LedgerEntry {
74 id: row.id,
75 kind: row.kind,
76 amount_micros: row.amount_micros,
77 description: row.description,
78 repo: row.repo,
79 number: row.number,
80 task: row.task,
81 model: row.model,
82 created_by: row.created_by,
83 created_at: row.created_at,
84 }
85 }
86}
87
88#[derive(Deserialize)]
89struct RunRow {
90 workspace: String,
91 repo: String,
92 number: u32,
93 task: String,
94 model: String,
95 token_hash: String,
96}
97
98#[derive(Deserialize)]
99struct CheckoutRow {
100 workspace: String,
101 created_by: String,
102}
103
104/// A row an `UPDATE … RETURNING` touched.
105#[derive(Deserialize)]
106struct Touched {
107 #[allow(dead_code)]
108 id: String,
109}
110
111struct Billing {
112 db: D1Database,
113 /// Absent when no card processor is configured.
114 stripe: Option<Stripe>,
115 margin_percent: u32,
116}
117
118impl Billing {
119 fn status(&self) -> Status {
120 Status {
121 enabled: self.stripe.is_some(),
122 live: self.stripe.as_ref().is_some_and(Stripe::live),
123 }
124 }
125
126 async fn row(&self, workspace: &str) -> Result<Option<AccountRow>> {
127 self.db
128 .prepare("SELECT balance_micros, customer_id FROM accounts WHERE workspace = ?")
129 .bind(&[workspace.into()])?
130 .first::<AccountRow>(None)
131 .await
132 }
133
134 async fn standing(&self, workspace: &str) -> Result<Account> {
135 Ok(Account {
136 workspace: workspace.to_owned(),
137 balance_micros: self
138 .row(workspace)
139 .await?
140 .map_or(0, |row| row.balance_micros),
141 status: self.status(),
142 margin_percent: self.margin_percent,
143 })
144 }
145
146 /// Adds a ledger entry and moves the balance by the same amount, as
147 /// one write.
148 #[allow(clippy::too_many_arguments)]
149 async fn enter(
150 &self,
151 workspace: &str,
152 kind: EntryKind,
153 amount_micros: i64,
154 description: &str,
155 reference: &str,
156 run: Option<&RunRow>,
157 cost_micros: Option<i64>,
158 created_by: Option<&str>,
159 customer: Option<&str>,
160 ) -> Result<()> {
161 let now = now_ms();
162 let timestamp = rfc3339(now);
163 let kind = match kind {
164 EntryKind::TopUp => "top_up",
165 EntryKind::Usage => "usage",
166 };
167 self.db
168 .batch(vec![
169 self.db
170 .prepare(
171 "INSERT INTO ledger
172 (id, workspace, kind, amount_micros, description, repo, number, task,
173 model, cost_micros, reference, created_by, created_at)
174 VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
175 )
176 .bind(&[
177 new_id("led", now).into(),
178 workspace.into(),
179 kind.into(),
180 // D1 takes numbers as doubles, which hold every
181 // amount this service will see exactly.
182 (amount_micros as f64).into(),
183 description.into(),
184 optional(run.map(|run| run.repo.as_str())),
185 run.map_or(JsValue::NULL, |run| run.number.into()),
186 optional(run.map(|run| run.task.as_str())),
187 optional(run.map(|run| run.model.as_str())),
188 cost_micros.map_or(JsValue::NULL, |cost| (cost as f64).into()),
189 reference.into(),
190 optional(created_by),
191 timestamp.as_str().into(),
192 ])?,
193 self.db
194 .prepare(
195 "INSERT INTO accounts (workspace, balance_micros, customer_id, created_at)
196 VALUES (?1, ?2, ?3, ?4)
197 ON CONFLICT (workspace) DO UPDATE SET
198 balance_micros = balance_micros + ?2,
199 customer_id = COALESCE(?3, customer_id)",
200 )
201 .bind(&[
202 workspace.into(),
203 (amount_micros as f64).into(),
204 optional(customer),
205 timestamp.as_str().into(),
206 ])?,
207 ])
208 .await?;
209 Ok(())
210 }
211
212 async fn account(&self, a: AccountArgs) -> Result<Outcome<Account>> {
213 let workspace = a.workspace.to_lowercase();
214 if !a.viewer.is_some_and(|viewer| viewer.is_member(&workspace)) {
215 return Ok(members_only());
216 }
217 Ok(Outcome::Ok(self.standing(&workspace).await?))
218 }
219
220 async fn ledger(&self, a: AccountArgs) -> Result<Outcome<Vec<LedgerEntry>>> {
221 let workspace = a.workspace.to_lowercase();
222 if !a.viewer.is_some_and(|viewer| viewer.is_member(&workspace)) {
223 return Ok(members_only());
224 }
225 let rows = self
226 .db
227 .prepare("SELECT * FROM ledger WHERE workspace = ? ORDER BY id DESC LIMIT ?")
228 .bind(&[workspace.into(), LEDGER_PAGE.into()])?
229 .all()
230 .await?
231 .results::<LedgerRow>()?;
232 Ok(Outcome::Ok(
233 rows.into_iter().map(LedgerEntry::from).collect(),
234 ))
235 }
236
Usage, like a hosting provider's: what agents cost, per day, task, repository and pull request237 async fn usage(&self, a: UsageArgs) -> Result<Outcome<Usage>> {
238 let workspace = a.workspace.to_lowercase();
239 if !a.viewer.is_some_and(|viewer| viewer.is_member(&workspace)) {
240 return Ok(members_only());
241 }
242 #[derive(serde::Deserialize)]
243 struct SliceRow {
244 key: Option<String>,
245 micros: Option<i64>,
246 runs: Option<u32>,
247 }
248 let slices = |key: &str, limit: u32| {
249 format!(
250 "SELECT {key} AS key, -SUM(amount_micros) AS micros, COUNT(*) AS runs FROM ledger
251 WHERE workspace = ?1 AND kind = 'usage' AND created_at >= ?2
252 GROUP BY 1 ORDER BY micros DESC LIMIT {limit}"
253 )
254 };
255 let query = |sql: String| {
256 let db = &self.db;
257 let workspace = workspace.clone();
258 let since = a.since.clone();
259 async move {
260 let rows = db
261 .prepare(sql)
262 .bind(&[workspace.into(), since.into()])?
263 .all()
264 .await?
265 .results::<SliceRow>()?;
266 Ok::<Vec<UsageSlice>, worker::Error>(
267 rows.into_iter()
268 .map(|row| UsageSlice {
269 key: row.key.unwrap_or_else(|| "other".to_owned()),
270 micros: row.micros.unwrap_or_default(),
271 runs: row.runs.unwrap_or_default(),
272 })
273 .collect(),
274 )
275 }
276 };
277 #[derive(serde::Deserialize)]
278 struct Totals {
279 spent: Option<i64>,
280 cost: Option<i64>,
281 runs: Option<u32>,
282 added: Option<i64>,
283 }
284 let totals = self
285 .db
286 .prepare(
287 "SELECT
288 -SUM(CASE WHEN kind = 'usage' THEN amount_micros END) AS spent,
289 SUM(CASE WHEN kind = 'usage' THEN cost_micros END) AS cost,
290 SUM(CASE WHEN kind = 'usage' THEN 1 ELSE 0 END) AS runs,
291 SUM(CASE WHEN kind = 'top_up' THEN amount_micros END) AS added
292 FROM ledger WHERE workspace = ?1 AND created_at >= ?2",
293 )
294 .bind(&[workspace.as_str().into(), a.since.as_str().into()])?
295 .first::<Totals>(None)
296 .await?;
297 let totals = totals.unwrap_or(Totals {
298 spent: None,
299 cost: None,
300 runs: None,
301 added: None,
302 });
303 Ok(Outcome::Ok(Usage {
304 spent_micros: totals.spent.unwrap_or_default(),
305 cost_micros: totals.cost.unwrap_or_default(),
306 runs: totals.runs.unwrap_or_default(),
307 added_micros: totals.added.unwrap_or_default(),
308 by_day: query(slices("substr(created_at, 1, 10) || '/' || COALESCE(task, 'other')", 400)).await?,
309 by_task: query(slices("task", 20)).await?,
310 by_repo: query(slices("repo", 20)).await?,
311 by_pull: query(slices("repo || '#' || number", 10)).await?,
312 by_model: query(slices("model", 10)).await?,
313 since: a.since,
314 }))
315 }
316
Agents as a team: lifecycle, merge queue, billing and a new shell317 async fn checkout(&self, a: CheckoutArgs) -> Result<Outcome<Checkout>> {
318 let workspace = a.workspace.to_lowercase();
319 if a.actor.role_in(&workspace) != Some(Role::Owner) {
320 return Ok(Outcome::fail(
321 FailureCode::Forbidden,
322 "Only an owner can add credit to a workspace.",
323 ));
324 }
325 let Some(stripe) = &self.stripe else {
326 return Ok(Outcome::fail(
327 FailureCode::Conflict,
328 "Payments are not set up on this g1t yet.",
329 ));
330 };
331 if !(MIN_TOP_UP_CENTS..=MAX_TOP_UP_CENTS).contains(&a.amount_cents) {
332 return Ok(Outcome::fail(
333 FailureCode::Invalid,
334 format!(
335 "Add between ${} and ${} at a time.",
336 MIN_TOP_UP_CENTS / 100,
337 MAX_TOP_UP_CENTS / 100
338 ),
339 ));
340 }
341 let customer = self.row(&workspace).await?.and_then(|row| row.customer_id);
342 let session = stripe
343 .start_checkout(
344 &workspace,
345 a.amount_cents,
346 customer.as_deref(),
347 &a.return_url,
348 )
349 .await?;
350 let Some(url) = session.url else {
351 return Err(worker::Error::RustError(
352 "the card processor returned no payment page".into(),
353 ));
354 };
355 self.db
356 .prepare(
357 "INSERT INTO checkouts (id, workspace, amount_cents, created_by, created_at)
358 VALUES (?, ?, ?, ?, ?)",
359 )
360 .bind(&[
361 session.id.into(),
362 workspace.into(),
363 a.amount_cents.into(),
364 a.actor.username.into(),
365 rfc3339(now_ms()).into(),
366 ])?
367 .run()
368 .await?;
369 Ok(Outcome::Ok(Checkout { url }))
370 }
371
372 /// Credits a payment if the processor says it was made and it has not
373 /// been credited before. The amount credited is what the processor
374 /// says was paid, not what anyone here remembers asking for.
375 async fn confirm(&self, a: ConfirmArgs) -> Result<Outcome<Account>> {
376 let workspace = a.workspace.to_lowercase();
377 if !a.viewer.is_some_and(|viewer| viewer.is_member(&workspace)) {
378 return Ok(members_only());
379 }
380 let (Some(stripe), Some(checkout)) = (
381 &self.stripe,
382 self.db
383 .prepare(
384 "SELECT workspace, created_by FROM checkouts
385 WHERE id = ? AND workspace = ? AND status = 'open'",
386 )
387 .bind(&[a.session.as_str().into(), workspace.as_str().into()])?
388 .first::<CheckoutRow>(None)
389 .await?,
390 ) else {
391 // Unknown, someone else's, or already credited: nothing to do.
392 return Ok(Outcome::Ok(self.standing(&workspace).await?));
393 };
394 let session = stripe.session(&a.session).await?;
395 let paid = session
396 .amount_total
397 .filter(|_| session.payment_status == "paid");
398 if let Some(cents) = paid {
399 // Only whoever flips it from open to paid enters the credit.
400 let claimed = self
401 .db
402 .prepare(
403 "UPDATE checkouts SET status = 'paid' WHERE id = ? AND status = 'open'
404 RETURNING id",
405 )
406 .bind(&[a.session.as_str().into()])?
407 .first::<Touched>(None)
408 .await?;
409 if claimed.is_some() {
410 self.enter(
411 &checkout.workspace,
412 EntryKind::TopUp,
413 i64::from(cents) * MICROS_PER_DOLLAR / 100,
414 "Credit added by card",
415 &session.id,
416 None,
417 None,
418 Some(&checkout.created_by),
419 session.customer.as_deref(),
420 )
421 .await?;
422 }
423 }
424 Ok(Outcome::Ok(self.standing(&workspace).await?))
425 }
426
427 /// A refusal if the workspace has no credit to start an agent with.
428 async fn out_of_credit<T>(&self, workspace: &str) -> Result<Option<Outcome<T>>> {
429 let balance = self
430 .row(workspace)
431 .await?
432 .map_or(0, |row| row.balance_micros);
433 Ok((balance <= 0).then(|| {
434 Outcome::fail(
435 FailureCode::PaymentRequired,
436 format!(
437 "The {workspace} workspace has no agent credit. An owner can add some under Billing on the workspace's page."
438 ),
439 )
440 }))
441 }
442
443 async fn can_start(&self, a: CanStartArgs) -> Result<Outcome<bool>> {
444 if self.stripe.is_none() {
445 return Ok(Outcome::Ok(true));
446 }
447 Ok(self
448 .out_of_credit(&a.workspace.to_lowercase())
449 .await?
450 .unwrap_or(Outcome::Ok(true)))
451 }
452
453 async fn start_run(&self, a: StartRunArgs) -> Result<Outcome<Option<RunTicket>>> {
454 if self.stripe.is_none() {
455 return Ok(Outcome::Ok(None));
456 }
457 let workspace = a.workspace.to_lowercase();
458 if let Some(refused) = self.out_of_credit(&workspace).await? {
459 return Ok(refused);
460 }
461 let now = now_ms();
462 let run_id = new_id("run", now);
463 let mut bytes = [0u8; 32];
464 getrandom::getrandom(&mut bytes).expect("no source of randomness");
465 let token = hex::encode(bytes);
466 self.db
467 .prepare(
468 "INSERT INTO runs (id, workspace, repo, number, task, model, token_hash, created_at)
469 VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
470 )
471 .bind(&[
472 run_id.as_str().into(),
473 workspace.into(),
474 format!("{}/{}", a.repo.namespace, a.repo.name).into(),
475 a.number.into(),
476 a.task.into(),
477 a.model.into(),
478 hash(&token).into(),
479 rfc3339(now).into(),
480 ])?
481 .run()
482 .await?;
483 Ok(Outcome::Ok(Some(RunTicket { run_id, token })))
484 }
485
486 async fn finish_run(&self, a: FinishRunArgs) -> Result<Outcome<bool>> {
487 let run = self
488 .db
489 .prepare(
490 "SELECT workspace, repo, number, task, model, token_hash FROM runs
491 WHERE id = ? AND finished_at IS NULL",
492 )
493 .bind(&[a.run_id.as_str().into()])?
494 .first::<RunRow>(None)
495 .await?;
496 let Some(run) = run.filter(|run| run.token_hash == hash(&a.token)) else {
497 return Ok(Outcome::fail(FailureCode::NotFound, "Run not found."));
498 };
499 if !a.cost_usd.is_finite() || a.cost_usd < 0.0 {
500 return Ok(Outcome::fail(FailureCode::Invalid, "That is not a cost."));
501 }
502 // Only whoever closes the run charges for it.
503 let claimed = self
504 .db
505 .prepare(
506 "UPDATE runs SET finished_at = ? WHERE id = ? AND finished_at IS NULL RETURNING id",
507 )
508 .bind(&[rfc3339(now_ms()).into(), a.run_id.as_str().into()])?
509 .first::<Touched>(None)
510 .await?;
511 if claimed.is_none() {
512 return Ok(Outcome::Ok(false));
513 }
514 let charge = charge_micros(a.cost_usd, self.margin_percent);
515 let description = match run.task.as_str() {
516 "plan" => format!("Planning for {}", run.repo),
517 "review" => format!("Review of {}#{}", run.repo, run.number),
518 "update" => format!("Catching up {}#{}", run.repo, run.number),
519 _ => format!("Work on {}#{}", run.repo, run.number),
520 };
521 self.enter(
522 &run.workspace,
523 EntryKind::Usage,
524 -charge,
525 &description,
526 &a.run_id,
527 Some(&run),
528 Some(charge_micros(a.cost_usd, 0)),
529 None,
530 None,
531 )
532 .await?;
533 Ok(Outcome::Ok(true))
534 }
535}
536
537fn members_only<T>() -> Outcome<T> {
538 Outcome::fail(
539 FailureCode::Forbidden,
540 "Only members can see a workspace's billing.",
541 )
542}
543
544#[event(fetch)]
545async fn fetch(mut request: Request, env: Env, _ctx: Context) -> Result<Response> {
546 let Some(method) = rpc_method(&request) else {
547 return Response::error("Not found", 404);
548 };
549 let body: serde_json::Value = request.json().await?;
550 let billing = Billing {
551 db: env.d1("DB")?,
552 stripe: env
553 .secret("STRIPE_SECRET_KEY")
554 .ok()
555 .map(|key| key.to_string())
556 .filter(|key| !key.is_empty())
557 .map(Stripe::new),
558 margin_percent: env
559 .var("MARGIN_PERCENT")
560 .ok()
561 .and_then(|percent| percent.to_string().parse().ok())
562 .unwrap_or(20),
563 };
564 match method.as_str() {
565 "status" => reply(&billing.status()),
566 "account" => reply(&billing.account(args(body)?).await?),
567 "ledger" => reply(&billing.ledger(args(body)?).await?),
Usage, like a hosting provider's: what agents cost, per day, task, repository and pull request568 "usage" => reply(&billing.usage(args(body)?).await?),
Agents as a team: lifecycle, merge queue, billing and a new shell569 "checkout" => reply(&billing.checkout(args(body)?).await?),
570 "confirm" => reply(&billing.confirm(args(body)?).await?),
571 "can_start" => reply(&billing.can_start(args(body)?).await?),
572 "start_run" => reply(&billing.start_run(args(body)?).await?),
573 "finish_run" => reply(&billing.finish_run(args(body)?).await?),
574 _ => Response::error("Unknown method", 404),
575 }
576}
577
578#[cfg(test)]
579mod tests {
580 use super::*;
581
582 #[test]
583 fn a_run_is_charged_its_cost_plus_the_margin() {
584 // $0.05 at 20% is six cents.
585 assert_eq!(charge_micros(0.05, 20), 60_000);
586 assert_eq!(charge_micros(1.0, 20), 1_200_000);
587 assert_eq!(charge_micros(0.05, 0), 50_000);
588 }
589
590 #[test]
591 fn fractions_of_a_millionth_round_up_and_nothing_costs_less_than_nothing() {
592 assert_eq!(charge_micros(0.000_000_4, 20), 2);
593 assert_eq!(charge_micros(0.0, 20), 0);
594 assert_eq!(charge_micros(-3.0, 20), 0);
595 }
596
597 #[test]
598 fn an_absurd_cost_is_capped() {
599 assert_eq!(charge_micros(1e9, 20), 120 * MICROS_PER_DOLLAR);
600 }
601}