Sign in

flagon-io/g1t

public

g1t/services/billing/src/lib.rs

520 lines18,105 bytesCodeBlame
1//! 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
237 async fn checkout(&self, a: CheckoutArgs) -> Result<Outcome<Checkout>> {
238 let workspace = a.workspace.to_lowercase();
239 if a.actor.role_in(&workspace) != Some(Role::Owner) {
240 return Ok(Outcome::fail(
241 FailureCode::Forbidden,
242 "Only an owner can add credit to a workspace.",
243 ));
244 }
245 let Some(stripe) = &self.stripe else {
246 return Ok(Outcome::fail(
247 FailureCode::Conflict,
248 "Payments are not set up on this g1t yet.",
249 ));
250 };
251 if !(MIN_TOP_UP_CENTS..=MAX_TOP_UP_CENTS).contains(&a.amount_cents) {
252 return Ok(Outcome::fail(
253 FailureCode::Invalid,
254 format!(
255 "Add between ${} and ${} at a time.",
256 MIN_TOP_UP_CENTS / 100,
257 MAX_TOP_UP_CENTS / 100
258 ),
259 ));
260 }
261 let customer = self.row(&workspace).await?.and_then(|row| row.customer_id);
262 let session = stripe
263 .start_checkout(
264 &workspace,
265 a.amount_cents,
266 customer.as_deref(),
267 &a.return_url,
268 )
269 .await?;
270 let Some(url) = session.url else {
271 return Err(worker::Error::RustError(
272 "the card processor returned no payment page".into(),
273 ));
274 };
275 self.db
276 .prepare(
277 "INSERT INTO checkouts (id, workspace, amount_cents, created_by, created_at)
278 VALUES (?, ?, ?, ?, ?)",
279 )
280 .bind(&[
281 session.id.into(),
282 workspace.into(),
283 a.amount_cents.into(),
284 a.actor.username.into(),
285 rfc3339(now_ms()).into(),
286 ])?
287 .run()
288 .await?;
289 Ok(Outcome::Ok(Checkout { url }))
290 }
291
292 /// Credits a payment if the processor says it was made and it has not
293 /// been credited before. The amount credited is what the processor
294 /// says was paid, not what anyone here remembers asking for.
295 async fn confirm(&self, a: ConfirmArgs) -> Result<Outcome<Account>> {
296 let workspace = a.workspace.to_lowercase();
297 if !a.viewer.is_some_and(|viewer| viewer.is_member(&workspace)) {
298 return Ok(members_only());
299 }
300 let (Some(stripe), Some(checkout)) = (
301 &self.stripe,
302 self.db
303 .prepare(
304 "SELECT workspace, created_by FROM checkouts
305 WHERE id = ? AND workspace = ? AND status = 'open'",
306 )
307 .bind(&[a.session.as_str().into(), workspace.as_str().into()])?
308 .first::<CheckoutRow>(None)
309 .await?,
310 ) else {
311 // Unknown, someone else's, or already credited: nothing to do.
312 return Ok(Outcome::Ok(self.standing(&workspace).await?));
313 };
314 let session = stripe.session(&a.session).await?;
315 let paid = session
316 .amount_total
317 .filter(|_| session.payment_status == "paid");
318 if let Some(cents) = paid {
319 // Only whoever flips it from open to paid enters the credit.
320 let claimed = self
321 .db
322 .prepare(
323 "UPDATE checkouts SET status = 'paid' WHERE id = ? AND status = 'open'
324 RETURNING id",
325 )
326 .bind(&[a.session.as_str().into()])?
327 .first::<Touched>(None)
328 .await?;
329 if claimed.is_some() {
330 self.enter(
331 &checkout.workspace,
332 EntryKind::TopUp,
333 i64::from(cents) * MICROS_PER_DOLLAR / 100,
334 "Credit added by card",
335 &session.id,
336 None,
337 None,
338 Some(&checkout.created_by),
339 session.customer.as_deref(),
340 )
341 .await?;
342 }
343 }
344 Ok(Outcome::Ok(self.standing(&workspace).await?))
345 }
346
347 /// A refusal if the workspace has no credit to start an agent with.
348 async fn out_of_credit<T>(&self, workspace: &str) -> Result<Option<Outcome<T>>> {
349 let balance = self
350 .row(workspace)
351 .await?
352 .map_or(0, |row| row.balance_micros);
353 Ok((balance <= 0).then(|| {
354 Outcome::fail(
355 FailureCode::PaymentRequired,
356 format!(
357 "The {workspace} workspace has no agent credit. An owner can add some under Billing on the workspace's page."
358 ),
359 )
360 }))
361 }
362
363 async fn can_start(&self, a: CanStartArgs) -> Result<Outcome<bool>> {
364 if self.stripe.is_none() {
365 return Ok(Outcome::Ok(true));
366 }
367 Ok(self
368 .out_of_credit(&a.workspace.to_lowercase())
369 .await?
370 .unwrap_or(Outcome::Ok(true)))
371 }
372
373 async fn start_run(&self, a: StartRunArgs) -> Result<Outcome<Option<RunTicket>>> {
374 if self.stripe.is_none() {
375 return Ok(Outcome::Ok(None));
376 }
377 let workspace = a.workspace.to_lowercase();
378 if let Some(refused) = self.out_of_credit(&workspace).await? {
379 return Ok(refused);
380 }
381 let now = now_ms();
382 let run_id = new_id("run", now);
383 let mut bytes = [0u8; 32];
384 getrandom::getrandom(&mut bytes).expect("no source of randomness");
385 let token = hex::encode(bytes);
386 self.db
387 .prepare(
388 "INSERT INTO runs (id, workspace, repo, number, task, model, token_hash, created_at)
389 VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
390 )
391 .bind(&[
392 run_id.as_str().into(),
393 workspace.into(),
394 format!("{}/{}", a.repo.namespace, a.repo.name).into(),
395 a.number.into(),
396 a.task.into(),
397 a.model.into(),
398 hash(&token).into(),
399 rfc3339(now).into(),
400 ])?
401 .run()
402 .await?;
403 Ok(Outcome::Ok(Some(RunTicket { run_id, token })))
404 }
405
406 async fn finish_run(&self, a: FinishRunArgs) -> Result<Outcome<bool>> {
407 let run = self
408 .db
409 .prepare(
410 "SELECT workspace, repo, number, task, model, token_hash FROM runs
411 WHERE id = ? AND finished_at IS NULL",
412 )
413 .bind(&[a.run_id.as_str().into()])?
414 .first::<RunRow>(None)
415 .await?;
416 let Some(run) = run.filter(|run| run.token_hash == hash(&a.token)) else {
417 return Ok(Outcome::fail(FailureCode::NotFound, "Run not found."));
418 };
419 if !a.cost_usd.is_finite() || a.cost_usd < 0.0 {
420 return Ok(Outcome::fail(FailureCode::Invalid, "That is not a cost."));
421 }
422 // Only whoever closes the run charges for it.
423 let claimed = self
424 .db
425 .prepare(
426 "UPDATE runs SET finished_at = ? WHERE id = ? AND finished_at IS NULL RETURNING id",
427 )
428 .bind(&[rfc3339(now_ms()).into(), a.run_id.as_str().into()])?
429 .first::<Touched>(None)
430 .await?;
431 if claimed.is_none() {
432 return Ok(Outcome::Ok(false));
433 }
434 let charge = charge_micros(a.cost_usd, self.margin_percent);
435 let description = match run.task.as_str() {
436 "plan" => format!("Planning for {}", run.repo),
437 "review" => format!("Review of {}#{}", run.repo, run.number),
438 "update" => format!("Catching up {}#{}", run.repo, run.number),
439 _ => format!("Work on {}#{}", run.repo, run.number),
440 };
441 self.enter(
442 &run.workspace,
443 EntryKind::Usage,
444 -charge,
445 &description,
446 &a.run_id,
447 Some(&run),
448 Some(charge_micros(a.cost_usd, 0)),
449 None,
450 None,
451 )
452 .await?;
453 Ok(Outcome::Ok(true))
454 }
455}
456
457fn members_only<T>() -> Outcome<T> {
458 Outcome::fail(
459 FailureCode::Forbidden,
460 "Only members can see a workspace's billing.",
461 )
462}
463
464#[event(fetch)]
465async fn fetch(mut request: Request, env: Env, _ctx: Context) -> Result<Response> {
466 let Some(method) = rpc_method(&request) else {
467 return Response::error("Not found", 404);
468 };
469 let body: serde_json::Value = request.json().await?;
470 let billing = Billing {
471 db: env.d1("DB")?,
472 stripe: env
473 .secret("STRIPE_SECRET_KEY")
474 .ok()
475 .map(|key| key.to_string())
476 .filter(|key| !key.is_empty())
477 .map(Stripe::new),
478 margin_percent: env
479 .var("MARGIN_PERCENT")
480 .ok()
481 .and_then(|percent| percent.to_string().parse().ok())
482 .unwrap_or(20),
483 };
484 match method.as_str() {
485 "status" => reply(&billing.status()),
486 "account" => reply(&billing.account(args(body)?).await?),
487 "ledger" => reply(&billing.ledger(args(body)?).await?),
488 "checkout" => reply(&billing.checkout(args(body)?).await?),
489 "confirm" => reply(&billing.confirm(args(body)?).await?),
490 "can_start" => reply(&billing.can_start(args(body)?).await?),
491 "start_run" => reply(&billing.start_run(args(body)?).await?),
492 "finish_run" => reply(&billing.finish_run(args(body)?).await?),
493 _ => Response::error("Unknown method", 404),
494 }
495}
496
497#[cfg(test)]
498mod tests {
499 use super::*;
500
501 #[test]
502 fn a_run_is_charged_its_cost_plus_the_margin() {
503 // $0.05 at 20% is six cents.
504 assert_eq!(charge_micros(0.05, 20), 60_000);
505 assert_eq!(charge_micros(1.0, 20), 1_200_000);
506 assert_eq!(charge_micros(0.05, 0), 50_000);
507 }
508
509 #[test]
510 fn fractions_of_a_millionth_round_up_and_nothing_costs_less_than_nothing() {
511 assert_eq!(charge_micros(0.000_000_4, 20), 2);
512 assert_eq!(charge_micros(0.0, 20), 0);
513 assert_eq!(charge_micros(-3.0, 20), 0);
514 }
515
516 #[test]
517 fn an_absurd_cost_is_capped() {
518 assert_eq!(charge_micros(1e9, 20), 120 * MICROS_PER_DOLLAR);
519 }
520}