Skip to content

g1t/services/billing/src/gateway.rs

512 lines22,207 bytesCodeBlame
1//! The AI Gateway: a workspace's own model requests, sent with one of its
2//! access tokens to the model proxy at `models.g1t.sh/anthropic` (Anthropic's
3//! Messages format) or `models.g1t.sh/openai/v1` (OpenAI's Chat Completions
4//! and Embeddings formats).
5//!
6//! - **Admission.** Before a request goes to g1t's models the proxy asks
7//! `gateway_admit`. A workspace over its spend limit is refused, as one is
8//! for anything else. On the plan it needs AI credit or included usage
9//! left, as an agent run does (auto-reload is tried first). A workspace
10//! with no plan is refused: the gateway on g1t's key is paid for from AI
11//! credit, which comes with the plan. A 100% discount and an enterprise
12//! need nothing more. On the workspace's own provider key nothing is
13//! asked: those requests cost g1t nothing.
14//! - **Charging.** Each request that used tokens on g1t's models is
15//! charged its tokens at the model's list price (`gateway_models`: Claude
16//! on Anthropic, open models on Workers AI) by kind, five-minute and
17//! hour-long cache writes apart, and at the long-prompt prices when the
18//! model is priced by prompt length and the prompt is longer, plus
19//! the price book's `gateway_models` markup (0 while the gateway is in
20//! beta), on a ledger line of its own (task `gateway`, one request each).
21//! The plan's included usage pays first, then AI credit (`grants::replay`
22//! counts gateway lines as model usage). Not an agent run, so never the
23//! agent rate, and never trial credit or g1t's pools.
24//! - **On the workspace's own provider.** Any of its model connections (an
25//! Anthropic or OpenAI key, any compatible endpoint), chosen by the model
26//! a request names. Logged with its tokens and charged nothing.
27//! - **The log.** Every request is kept for `RETENTION_DAYS`, with its
28//! model, format, provider, tokens by kind, cost, status and the token
29//! that sent it; never
30//! its prompt or answer.
31
32use g1t_contracts::billing::{
33 GatewayAdmitArgs, GatewayModel, GatewayRequest, GatewayRequests, GatewayRequestsArgs, LimitState, PlanKind, RecordGatewayArgs,
34};
35use g1t_contracts::time::rfc3339;
36use g1t_contracts::{FailureCode, Outcome, new_id};
37use g1t_kit::now_ms;
38use serde::Deserialize;
39use worker::Result;
40use worker::wasm_bindgen::JsValue;
41
42use crate::credits::{Eligible, month_of};
43use crate::features::thousands;
44use crate::{Billing, margin_on, members_only, optional};
45
46/// How long the log keeps a request.
47pub(crate) const RETENTION_DAYS: u32 = 30;
48/// A page of the log: this many when not asked, and at most.
49const PAGE: u32 = 50;
50const MAX_PAGE: u32 = 200;
51const DAY_MS: u64 = 86_400_000;
52
53/// The tokens of one request, by kind. `cache_write` counts every cache
54/// write; `cache_write_1h` those of them that live an hour.
55#[derive(Clone, Copy, Debug, Default, PartialEq)]
56pub(crate) struct Used {
57 pub input: u64,
58 pub output: u64,
59 pub cache_read: u64,
60 pub cache_write: u64,
61 pub cache_write_1h: u64,
62}
63
64impl Used {
65 /// The prompt's length, which a model priced by it is priced by.
66 pub(crate) fn prompt(&self) -> u64 {
67 self.input + self.cache_read + self.cache_write
68 }
69
70 pub(crate) fn total(&self) -> u64 {
71 self.prompt() + self.output
72 }
73}
74
75/// Whether a request is charged at a model's over-threshold prices: its
76/// prompt is longer than the model's threshold.
77pub(crate) fn over_threshold(model: &GatewayModel, used: &Used) -> bool {
78 model.threshold > 0 && used.prompt() > model.threshold
79}
80
81/// What a request's tokens cost at a model's prices per million, rounded
82/// up to a whole millionth of a dollar. A prompt longer than the model's
83/// threshold puts the whole request at the over-threshold prices.
84pub(crate) fn cost_micros(model: &GatewayModel, used: &Used) -> i64 {
85 let over = over_threshold(model, used);
86 let pick = |base: i64, above: i64| if over { above } else { base };
87 let hour = used.cache_write_1h.min(used.cache_write);
88 let tokens = [used.input, used.output, used.cache_read, used.cache_write - hour, hour];
89 let five_minutes = pick(model.cache_write_micros, model.over_cache_write_micros);
90 let prices = [
91 pick(model.input_micros, model.over_input_micros),
92 pick(model.output_micros, model.over_output_micros),
93 pick(model.cache_read_micros, model.over_cache_read_micros),
94 five_minutes,
95 // A model with no hour-long price charges those writes as five-minute ones.
96 match pick(model.cache_write_1h_micros, model.over_cache_write_1h_micros) {
97 0 => five_minutes,
98 price => price,
99 },
100 ];
101 let millionths: u128 = tokens
102 .iter()
103 .zip(prices)
104 .map(|(n, price)| u128::from(*n) * u128::from(price.max(0).unsigned_abs()))
105 .sum();
106 i64::try_from(millionths.div_ceil(1_000_000)).unwrap_or(i64::MAX)
107}
108
109/// Where a workspace stands for a request on g1t's models.
110#[derive(Clone, Debug, PartialEq)]
111pub(crate) enum Standing {
112 /// Over its spend limit, with the limit's own message.
113 Stopped(String),
114 /// No plan.
115 NoPlan,
116 /// On the plan with no AI credit or included usage left.
117 OutOfCredit { reload_failed: bool },
118 Admitted,
119}
120
121/// Where a workspace within its limit stands, by its plan: `exhausted` is
122/// `credit_exhausted`'s answer (asked only on the plan). A workspace with
123/// no plan has no AI credit to spend; a 100% discount and an enterprise
124/// need none.
125pub(crate) fn plan_standing(plan: PlanKind, exhausted: Option<bool>) -> Standing {
126 match plan {
127 PlanKind::Free => Standing::NoPlan,
128 PlanKind::Internal | PlanKind::Enterprise => Standing::Admitted,
129 PlanKind::Paid => match exhausted {
130 Some(reload_failed) => Standing::OutOfCredit { reload_failed },
131 None => Standing::Admitted,
132 },
133 }
134}
135
136/// What one request costs g1t, to charge: its tokens at its model's prices
137/// on g1t's key; nothing on the workspace's own key, or for a model g1t
138/// has no price for.
139pub(crate) fn request_cost(own_key: bool, model: Option<&GatewayModel>, used: &Used) -> i64 {
140 match model {
141 Some(model) if !own_key => cost_micros(model, used),
142 _ => 0,
143 }
144}
145
146/// Why a request is refused, in words for whoever sent it, or None.
147pub(crate) fn refusal(workspace: &str, standing: &Standing) -> Option<String> {
148 match standing {
149 Standing::Admitted => None,
150 Standing::Stopped(message) => Some(message.clone()),
151 Standing::NoPlan => Some(format!(
152 "The AI Gateway on g1t's models is paid for from AI credit, which comes with the g1t plan. An owner can start the plan for {workspace} at /{workspace}/-/billing, or connect the workspace's own model provider (an Anthropic or OpenAI key, or any compatible endpoint) under Integrations to use the gateway at no charge."
153 )),
154 Standing::OutOfCredit { reload_failed } => {
155 let reload = if *reload_failed { " Auto-reload was turned off after its last charge failed." } else { "" };
156 Some(format!(
157 "The {workspace} workspace is out of AI credit and has used this month's included usage, so the AI Gateway refuses requests to g1t's models.{reload} An owner can buy AI credit or turn on auto-reload at /{workspace}/-/billing#ai-credit."
158 ))
159 }
160 }
161}
162
163/// What a ledger line for one request says.
164pub(crate) fn describe(model_name: &str, used: &Used, over: bool, token_name: Option<&str>) -> String {
165 let by = token_name.map(str::trim).filter(|name| !name.is_empty()).map_or(String::new(), |name| format!(", token {name}"));
166 let long = if over { ", long-prompt price" } else { "" };
167 format!("AI Gateway: {model_name}, {} tokens{long}{by}", thousands(used.total()))
168}
169
170/// A request's format as the log keeps it: `anthropic` or `openai`.
171pub(crate) fn format_of(format: &str) -> &'static str {
172 if format.eq_ignore_ascii_case("openai") { "openai" } else { "anthropic" }
173}
174
175/// Whether a request id is one the proxy makes: `gw_` and up to 64 letters,
176/// digits, `_` and `-`.
177pub(crate) fn valid_id(id: &str) -> bool {
178 id.starts_with("gw_") && id.len() <= 64 && id.chars().all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-')
179}
180
181#[derive(Deserialize)]
182struct ModelRow {
183 model: String,
184 name: String,
185 provider: String,
186 #[serde(default)]
187 kind: Option<String>,
188 input_micros: i64,
189 output_micros: i64,
190 cache_read_micros: i64,
191 cache_write_micros: i64,
192 #[serde(default)]
193 cache_write_1h_micros: Option<i64>,
194 #[serde(default)]
195 threshold: Option<f64>,
196 #[serde(default)]
197 over_input_micros: Option<i64>,
198 #[serde(default)]
199 over_output_micros: Option<i64>,
200 #[serde(default)]
201 over_cache_read_micros: Option<i64>,
202 #[serde(default)]
203 over_cache_write_micros: Option<i64>,
204 #[serde(default)]
205 over_cache_write_1h_micros: Option<i64>,
206}
207
208impl From<ModelRow> for GatewayModel {
209 fn from(row: ModelRow) -> Self {
210 GatewayModel {
211 model: row.model,
212 name: row.name,
213 provider: row.provider,
214 kind: row.kind.unwrap_or_else(|| "chat".to_owned()),
215 input_micros: row.input_micros,
216 output_micros: row.output_micros,
217 cache_read_micros: row.cache_read_micros,
218 cache_write_micros: row.cache_write_micros,
219 cache_write_1h_micros: row.cache_write_1h_micros.unwrap_or(0),
220 threshold: row.threshold.unwrap_or(0.0).max(0.0) as u64,
221 over_input_micros: row.over_input_micros.unwrap_or(0),
222 over_output_micros: row.over_output_micros.unwrap_or(0),
223 over_cache_read_micros: row.over_cache_read_micros.unwrap_or(0),
224 over_cache_write_micros: row.over_cache_write_micros.unwrap_or(0),
225 over_cache_write_1h_micros: row.over_cache_write_1h_micros.unwrap_or(0),
226 }
227 }
228}
229
230#[derive(Deserialize)]
231struct RequestRow {
232 id: String,
233 created_at: String,
234 model: String,
235 token_id: String,
236 token_name: Option<String>,
237 input: f64,
238 output: f64,
239 cache_read: f64,
240 cache_write: f64,
241 #[serde(default)]
242 cache_write_1h: Option<f64>,
243 cost_micros: f64,
244 charged_micros: f64,
245 status: f64,
246 own_key: f64,
247 #[serde(default)]
248 format: Option<String>,
249 #[serde(default)]
250 provider: Option<String>,
251 #[serde(default)]
252 connection: Option<String>,
253 streamed: f64,
254 duration_ms: f64,
255 error: Option<String>,
256}
257
258impl From<RequestRow> for GatewayRequest {
259 fn from(row: RequestRow) -> Self {
260 let n = |v: f64| v.max(0.0) as u64;
261 GatewayRequest {
262 id: row.id,
263 created_at: row.created_at,
264 model: row.model,
265 token_id: row.token_id,
266 token_name: row.token_name,
267 input: n(row.input),
268 output: n(row.output),
269 cache_read: n(row.cache_read),
270 cache_write: n(row.cache_write),
271 cache_write_hour: n(row.cache_write_1h.unwrap_or(0.0)),
272 cost_micros: row.cost_micros as i64,
273 charged_micros: row.charged_micros as i64,
274 status: row.status as u16,
275 own_key: row.own_key != 0.0,
276 format: format_of(row.format.as_deref().unwrap_or_default()).to_owned(),
277 provider: row.provider.unwrap_or_default(),
278 connection: row.connection,
279 streamed: row.streamed != 0.0,
280 duration_ms: n(row.duration_ms),
281 error: row.error,
282 }
283 }
284}
285
286/// D1 takes numbers as doubles; counts and amounts here fit exactly.
287fn number(n: u64) -> JsValue {
288 JsValue::from_f64(n as f64)
289}
290
291impl Billing {
292 /// `gateway_models`: what the gateway offers on g1t's key, with prices.
293 pub(crate) async fn gateway_models(&self) -> Result<Vec<GatewayModel>> {
294 Ok(self
295 .db
296 .prepare("SELECT * FROM gateway_models ORDER BY position, model")
297 .all()
298 .await?
299 .results::<ModelRow>()?
300 .into_iter()
301 .map(GatewayModel::from)
302 .collect())
303 }
304
305 async fn gateway_model(&self, model: &str) -> Result<Option<GatewayModel>> {
306 Ok(self
307 .db
308 .prepare("SELECT * FROM gateway_models WHERE model = ?")
309 .bind(&[model.into()])?
310 .first::<ModelRow>(None)
311 .await?
312 .map(GatewayModel::from))
313 }
314
315 /// Where a workspace stands for a request on g1t's models.
316 pub(crate) async fn gateway_standing(&self, workspace: &str) -> Result<Standing> {
317 if self.stripe.is_none() || self.free {
318 return Ok(Standing::Admitted);
319 }
320 let limit = self.limit_of(workspace).await?;
321 if limit.state == LimitState::Stopped {
322 return Ok(Standing::Stopped(limit.message.unwrap_or_else(|| "This workspace is over its limit.".to_owned())));
323 }
324 if let Some(Outcome::Fail(failure)) = self.out_of_credit::<bool>(workspace).await? {
325 return Ok(Standing::Stopped(failure.message));
326 }
327 let account = self.account_of(workspace).await?;
328 let plan = self.plan_kind_for(workspace, &account).await?;
329 // Only a workspace paying on the plan needs credit to spend.
330 let exhausted = if plan == PlanKind::Paid { self.credit_exhausted(workspace).await? } else { None };
331 Ok(plan_standing(plan, exhausted))
332 }
333
334 /// `gateway_admit`.
335 pub(crate) async fn gateway_admit(&self, a: GatewayAdmitArgs) -> Result<Outcome<bool>> {
336 let workspace = a.workspace.trim().to_lowercase();
337 if workspace.is_empty() {
338 return Ok(Outcome::fail(FailureCode::Invalid, "Name the workspace."));
339 }
340 let standing = self.gateway_standing(&workspace).await?;
341 Ok(match refusal(&workspace, &standing) {
342 Some(why) => Outcome::fail(FailureCode::PaymentRequired, why),
343 None => Outcome::Ok(true),
344 })
345 }
346
347 /// `record_gateway`: logs a request once, and charges it once when it
348 /// used tokens on g1t's models.
349 pub(crate) async fn record_gateway(&self, a: RecordGatewayArgs) -> Result<Outcome<bool>> {
350 if !valid_id(&a.id) {
351 return Ok(Outcome::fail(FailureCode::Invalid, "A gateway request's id is gw_ and up to 64 letters and digits."));
352 }
353 let workspace = a.workspace.trim().to_lowercase();
354 if workspace.is_empty() {
355 return Ok(Outcome::fail(FailureCode::Invalid, "Name the workspace."));
356 }
357 let used = Used {
358 input: a.input,
359 output: a.output,
360 cache_read: a.cache_read,
361 cache_write: a.cache_write,
362 cache_write_1h: a.cache_write_hour.min(a.cache_write),
363 };
364 let model_name: String = a.model.trim().chars().take(200).collect();
365 let token_name = a.token_name.as_deref().map(|name| name.trim().chars().take(100).collect::<String>());
366 let provider: String = a.provider.trim().chars().take(40).collect();
367 let connection = a.connection.as_deref().map(|name| name.trim().chars().take(80).collect::<String>());
368 let priced = if a.own_key { None } else { self.gateway_model(&model_name).await? };
369 let cost = request_cost(a.own_key, priced.as_ref(), &used);
370 let now = now_ms();
371 let timestamp = rfc3339(now);
372 // Claimed first: the same request recorded twice is one row and one charge.
373 let claimed = self
374 .db
375 .prepare(
376 "INSERT OR IGNORE INTO gateway_requests
377 (id, workspace, created_at, token_id, token_name, model, input, output, cache_read, cache_write,
378 cache_write_1h, cost_micros, charged_micros, status, own_key, format, provider, connection,
379 streamed, duration_ms, error)
380 VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 0, ?, ?, ?, ?, ?, ?, ?, ?)
381 RETURNING id",
382 )
383 .bind(&[
384 a.id.as_str().into(),
385 workspace.as_str().into(),
386 timestamp.as_str().into(),
387 a.token_id.chars().take(64).collect::<String>().into(),
388 optional(token_name.as_deref()),
389 model_name.as_str().into(),
390 number(a.input),
391 number(a.output),
392 number(a.cache_read),
393 number(a.cache_write),
394 number(used.cache_write_1h),
395 (cost as f64).into(),
396 f64::from(a.status).into(),
397 f64::from(u8::from(a.own_key)).into(),
398 format_of(&a.format).into(),
399 provider.as_str().into(),
400 optional(connection.as_deref()),
401 f64::from(u8::from(a.streamed)).into(),
402 number(a.duration_ms),
403 optional(a.error.as_deref().map(|e| e.chars().take(500).collect::<String>()).as_deref()),
404 ])?
405 .first::<serde_json::Value>(None)
406 .await?;
407 if claimed.is_none() {
408 return Ok(Outcome::Ok(false));
409 }
410 // Without a card processor there is no bill to put it on.
411 let Some(model) = priced.filter(|_| cost > 0 && self.stripe.is_some()) else {
412 return Ok(Outcome::Ok(true));
413 };
414 let base = margin_on(cost, self.gateway_markup().await?);
415 let (charge, terms_note, discount) = self.charged(&workspace, base).await?;
416 // Included usage pays first, then AI credit. Never the trial or
417 // g1t's pools, and g1t never covers the rest.
418 let eligible = Eligible { trial: false, repo: None, cover_rest: false };
419 let drawn = self.draw(&workspace, charge, &month_of(&timestamp), &eligible).await?;
420 let owed = charge - drawn.total();
421 let over = over_threshold(&model, &used);
422 let description = format!("{}{terms_note}{}", describe(&model.name, &used, over, token_name.as_deref()), drawn.note());
423 self.db
424 .batch(vec![
425 self.db
426 .prepare(
427 "INSERT INTO ledger
428 (id, workspace, kind, amount_micros, description, task, model, cost_micros, reference, created_at,
429 billed_to, credit_micros, trial_micros, oss_micros, given_micros, price_version, quantity)
430 VALUES (?, ?, 'usage', ?, ?, 'gateway', ?, ?, ?, ?, 'g1t', ?, ?, ?, ?, ?, 1)",
431 )
432 .bind(&[
433 new_id("led", now).into(),
434 workspace.as_str().into(),
435 (-(owed as f64)).into(),
436 description.as_str().into(),
437 model.model.as_str().into(),
438 (cost as f64).into(),
439 a.id.as_str().into(),
440 timestamp.as_str().into(),
441 (drawn.credit as f64).into(),
442 (drawn.trial as f64).into(),
443 (drawn.oss as f64).into(),
444 (drawn.given as f64).into(),
445 optional(self.version_now("gateway_models").await?.as_deref()),
446 ])?,
447 self.db
448 .prepare(
449 "INSERT INTO accounts (workspace, balance_micros, created_at)
450 VALUES (?1, ?2, ?3)
451 ON CONFLICT (workspace) DO UPDATE SET balance_micros = balance_micros + ?2",
452 )
453 .bind(&[workspace.as_str().into(), (-(owed as f64)).into(), timestamp.as_str().into()])?,
454 self.db
455 .prepare("UPDATE gateway_requests SET charged_micros = ? WHERE id = ?")
456 .bind(&[(charge as f64).into(), a.id.as_str().into()])?,
457 ])
458 .await?;
459 self.record_discount(&a.id, discount).await?;
460 self.count_spend(&workspace, cost, owed, &drawn).await;
461 Ok(Outcome::Ok(true))
462 }
463
464 /// `gateway_requests`: the log, newest first, for members.
465 pub(crate) async fn gateway_requests(&self, a: GatewayRequestsArgs) -> Result<Outcome<GatewayRequests>> {
466 let workspace = a.workspace.to_lowercase();
467 if !a.viewer.is_some_and(|viewer| viewer.is_member(&workspace)) {
468 return Ok(members_only());
469 }
470 let limit = a.limit.unwrap_or(PAGE).clamp(1, MAX_PAGE);
471 let mut binds: Vec<JsValue> = vec![workspace.as_str().into()];
472 let older = match a.before.as_deref().map(str::trim).filter(|id| !id.is_empty()) {
473 Some(before) => {
474 binds.push(before.into());
475 " AND (created_at, id) < (SELECT created_at, id FROM gateway_requests WHERE id = ?2)"
476 }
477 None => "",
478 };
479 binds.push(f64::from(limit + 1).into());
480 let at = binds.len();
481 let mut rows = self
482 .db
483 .prepare(format!(
484 "SELECT * FROM gateway_requests WHERE workspace = ?1{older} ORDER BY created_at DESC, id DESC LIMIT ?{at}"
485 ))
486 .bind(&binds)?
487 .all()
488 .await?
489 .results::<RequestRow>()?;
490 let more = rows.len() > limit as usize;
491 rows.truncate(limit as usize);
492 let requests: Vec<GatewayRequest> = rows.into_iter().map(GatewayRequest::from).collect();
493 let next = if more { requests.last().map(|r| r.id.clone()) } else { None };
494 Ok(Outcome::Ok(GatewayRequests { requests, next, retention_days: RETENTION_DAYS }))
495 }
496
497 /// Daily: requests older than the log keeps are deleted.
498 pub(crate) async fn forget_gateway_requests(&self) -> Result<()> {
499 let cutoff = rfc3339(now_ms().saturating_sub(u64::from(RETENTION_DAYS) * DAY_MS));
500 self.db
501 .prepare("DELETE FROM gateway_requests WHERE created_at < ?")
502 .bind(&[cutoff.into()])?
503 .run()
504 .await?;
505 Ok(())
506 }
507}
508
509#[cfg(test)]
510#[path = "gateway_tests.rs"]
511mod tests;
512