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.
| AI Gateway: OpenAI's format, open models, and your own providers | 1 | //! Pricing, admission and the log's wording for the AI Gateway (`gateway.rs`). |
| 2 | ||
| 3 | use super::*; | |
| 4 | ||
| 5 | fn priced(model: &str, name: &str, prices: [i64; 5]) -> GatewayModel { | |
| 6 | GatewayModel { | |
| 7 | model: model.into(), | |
| 8 | name: name.into(), | |
| 9 | provider: "anthropic".into(), | |
| 10 | kind: "chat".into(), | |
| 11 | input_micros: prices[0], | |
| 12 | output_micros: prices[1], | |
| 13 | cache_read_micros: prices[2], | |
| 14 | cache_write_micros: prices[3], | |
| 15 | cache_write_1h_micros: prices[4], | |
| 16 | threshold: 0, | |
| 17 | over_input_micros: 0, | |
| 18 | over_output_micros: 0, | |
| 19 | over_cache_read_micros: 0, | |
| 20 | over_cache_write_micros: 0, | |
| 21 | over_cache_write_1h_micros: 0, | |
| 22 | } | |
| 23 | } | |
| 24 | ||
| 25 | fn sonnet() -> GatewayModel { | |
| 26 | priced("claude-sonnet-5-5", "Claude Sonnet 5.5", [2_000_000, 10_000_000, 100_000, 2_500_000, 4_000_000]) | |
| 27 | } | |
| 28 | ||
| 29 | /// Claude Haiku 5.5, priced by prompt length at 100,000 tokens. | |
| 30 | fn haiku() -> GatewayModel { | |
| 31 | GatewayModel { | |
| 32 | threshold: 100_000, | |
| 33 | over_input_micros: 500_000, | |
| 34 | over_output_micros: 2_500_000, | |
| 35 | over_cache_read_micros: 50_000, | |
| 36 | over_cache_write_micros: 625_000, | |
| 37 | over_cache_write_1h_micros: 1_000_000, | |
| 38 | ..priced("claude-haiku-5-5", "Claude Haiku 5.5", [100_000, 500_000, 10_000, 125_000, 200_000]) | |
| 39 | } | |
| 40 | } | |
| 41 | ||
| 42 | fn used(input: u64, output: u64, cache_read: u64, cache_write: u64, cache_write_1h: u64) -> Used { | |
| 43 | Used { input, output, cache_read, cache_write, cache_write_1h } | |
| 44 | } | |
| 45 | ||
| 46 | #[test] | |
| 47 | fn a_request_costs_its_tokens_at_the_models_prices() { | |
| 48 | // A million of each, the writes five-minute ones: $2 + $10 + $0.10 + $2.50. | |
| 49 | assert_eq!(cost_micros(&sonnet(), &used(1_000_000, 1_000_000, 1_000_000, 1_000_000, 0)), 14_600_000); | |
| 50 | // 1,000 input and 500 output: $0.002 + $0.005. | |
| 51 | assert_eq!(cost_micros(&sonnet(), &used(1_000, 500, 0, 0, 0)), 7_000); | |
| 52 | // Nothing used costs nothing. | |
| 53 | assert_eq!(cost_micros(&sonnet(), &Used::default()), 0); | |
| 54 | } | |
| 55 | ||
| 56 | #[test] | |
| 57 | fn hour_long_cache_writes_cost_twice_input() { | |
| 58 | // A million writes, all an hour long: $4, not the five-minute $2.50. | |
| 59 | assert_eq!(cost_micros(&sonnet(), &used(0, 0, 0, 1_000_000, 1_000_000)), 4_000_000); | |
| 60 | // Split: 600k five-minute ($1.50) and 400k hour-long ($1.60). | |
| 61 | assert_eq!(cost_micros(&sonnet(), &used(0, 0, 0, 1_000_000, 400_000)), 3_100_000); | |
| 62 | // More hour-long writes than writes is never more than the writes. | |
| 63 | assert_eq!(cost_micros(&sonnet(), &used(0, 0, 0, 1_000, 5_000)), 4_000); | |
| 64 | // A model with no hour-long price charges them as five-minute ones. | |
| 65 | let flat = GatewayModel { cache_write_1h_micros: 0, ..sonnet() }; | |
| 66 | assert_eq!(cost_micros(&flat, &used(0, 0, 0, 1_000_000, 1_000_000)), 2_500_000); | |
| 67 | } | |
| 68 | ||
| 69 | #[test] | |
| 70 | fn a_short_prompt_on_haiku_5_5_is_charged_the_lower_prices() { | |
| 71 | let short = used(60_000, 2_000, 30_000, 10_000, 0); | |
| 72 | assert!(!over_threshold(&haiku(), &short)); | |
| 73 | // 60k at $0.10, 2k at $0.50, 30k at $0.01 and 10k at $0.125 a million. | |
| 74 | assert_eq!(cost_micros(&haiku(), &short), 6_000 + 1_000 + 300 + 1_250); | |
| 75 | // Exactly the threshold is still the lower price. | |
| 76 | assert!(!over_threshold(&haiku(), &used(100_000, 50_000, 0, 0, 0))); | |
| 77 | } | |
| 78 | ||
| 79 | #[test] | |
| 80 | fn a_long_prompt_on_haiku_5_5_puts_the_whole_request_at_the_higher_prices() { | |
| 81 | // 100,001 prompt tokens, counting cache reads and writes. | |
| 82 | let long = used(40_001, 2_000, 50_000, 10_000, 4_000); | |
| 83 | assert!(over_threshold(&haiku(), &long)); | |
| 84 | // 40,001 at $0.50, 2k at $2.50, 50k at $0.05, 6k at $0.625 and 4k at $1 a million, rounded up. | |
| 85 | assert_eq!(cost_micros(&haiku(), &long), 20_001 + 5_000 + 2_500 + 3_750 + 4_000); | |
| 86 | // A model with one price is never over. | |
| 87 | assert!(!over_threshold(&sonnet(), &used(900_000, 0, 0, 0, 0))); | |
| 88 | } | |
| 89 | ||
| 90 | #[test] | |
| 91 | fn a_fraction_of_a_millionth_rounds_up() { | |
| 92 | // One cache-read token: 0.1 millionths. | |
| 93 | assert_eq!(cost_micros(&sonnet(), &used(0, 0, 1, 0, 0)), 1); | |
| 94 | // A negative price in the table is never a credit: it counts as nothing. | |
| 95 | let odd = GatewayModel { input_micros: -5, ..sonnet() }; | |
| 96 | assert_eq!(cost_micros(&odd, &used(10, 0, 0, 0, 0)), 0); | |
| 97 | } | |
| 98 | ||
| 99 | #[test] | |
| 100 | fn the_markup_is_the_price_books_and_zero_in_beta_charges_the_cost() { | |
| 101 | let cost = cost_micros(&sonnet(), &used(12_000, 800, 40_000, 0, 0)); | |
| 102 | assert_eq!(margin_on(cost, 0), cost); | |
| 103 | assert_eq!(margin_on(1_000, 20), 1_200); | |
| 104 | } | |
| 105 | ||
| 106 | #[test] | |
| 107 | fn the_workspaces_own_provider_is_counted_and_never_charged() { | |
| 108 | let tokens = used(50_000, 2_000, 0, 0, 0); | |
| 109 | assert_eq!(request_cost(false, Some(&sonnet()), &tokens), 120_000); | |
| 110 | assert_eq!(request_cost(true, Some(&sonnet()), &tokens), 0); | |
| 111 | assert_eq!(request_cost(true, Some(&haiku()), &used(500_000, 0, 0, 0, 0)), 0); | |
| 112 | // A model g1t has no price for is never charged a guess. | |
| 113 | assert_eq!(request_cost(false, None, &tokens), 0); | |
| 114 | } | |
| 115 | ||
| 116 | #[test] | |
| 117 | fn an_open_model_is_priced_like_any_other() { | |
| 118 | let oss = GatewayModel { | |
| 119 | provider: "workers-ai".into(), | |
| 120 | ..priced("@cf/openai/gpt-oss-120b", "gpt-oss-120b", [350_000, 750_000, 350_000, 350_000, 350_000]) | |
| 121 | }; | |
| 122 | // 200k in, 10k out: $0.07 + $0.0075. | |
| 123 | assert_eq!(request_cost(false, Some(&oss), &used(200_000, 10_000, 0, 0, 0)), 77_500); | |
| 124 | // Embeddings: input only. | |
| 125 | let bge = GatewayModel { kind: "embeddings".into(), ..priced("@cf/baai/bge-m3", "BGE M3", [12_000, 0, 12_000, 12_000, 12_000]) }; | |
| 126 | assert_eq!(request_cost(false, Some(&bge), &used(1_000_000, 0, 0, 0, 0)), 12_000); | |
| 127 | } | |
| 128 | ||
| 129 | #[test] | |
| 130 | fn out_of_credit_on_the_plan_is_refused_and_the_rest_by_plan() { | |
| 131 | assert_eq!(plan_standing(PlanKind::Paid, None), Standing::Admitted); | |
| 132 | assert_eq!(plan_standing(PlanKind::Paid, Some(false)), Standing::OutOfCredit { reload_failed: false }); | |
| 133 | assert_eq!(plan_standing(PlanKind::Paid, Some(true)), Standing::OutOfCredit { reload_failed: true }); | |
| 134 | assert_eq!(plan_standing(PlanKind::Free, None), Standing::NoPlan); | |
| 135 | // Comped and invoiced workspaces need no credit. | |
| 136 | assert_eq!(plan_standing(PlanKind::Internal, Some(false)), Standing::Admitted); | |
| 137 | assert_eq!(plan_standing(PlanKind::Enterprise, Some(false)), Standing::Admitted); | |
| 138 | assert!(refusal("acme", &plan_standing(PlanKind::Paid, Some(false))).is_some()); | |
| 139 | } | |
| 140 | ||
| 141 | #[test] | |
| 142 | fn only_an_admitted_workspace_is_let_through() { | |
| 143 | assert_eq!(refusal("acme", &Standing::Admitted), None); | |
| 144 | let out = refusal("acme", &Standing::OutOfCredit { reload_failed: false }).unwrap(); | |
| 145 | assert!(out.contains("out of AI credit") && out.contains("/acme/-/billing#ai-credit"), "{out}"); | |
| 146 | assert!(!out.contains("Auto-reload")); | |
| 147 | let failed = refusal("acme", &Standing::OutOfCredit { reload_failed: true }).unwrap(); | |
| 148 | assert!(failed.contains("Auto-reload was turned off")); | |
| 149 | let no_plan = refusal("acme", &Standing::NoPlan).unwrap(); | |
| 150 | assert!(no_plan.contains("g1t plan") && no_plan.contains("own model provider"), "{no_plan}"); | |
| 151 | assert_eq!(refusal("acme", &Standing::Stopped("Over the limit.".into())).as_deref(), Some("Over the limit.")); | |
| 152 | } | |
| 153 | ||
| 154 | #[test] | |
| 155 | fn a_ledger_line_names_the_model_tokens_and_token() { | |
| 156 | let tokens = used(12_000, 800, 0, 0, 0); | |
| 157 | assert_eq!(describe("Claude Sonnet 5.5", &tokens, false, Some("ci")), "AI Gateway: Claude Sonnet 5.5, 12,800 tokens, token ci"); | |
| 158 | assert_eq!(describe("Claude Haiku 4.5", &used(5, 0, 0, 0, 0), false, None), "AI Gateway: Claude Haiku 4.5, 5 tokens"); | |
| 159 | assert_eq!(describe("Claude Haiku 4.5", &used(5, 0, 0, 0, 0), false, Some(" ")), "AI Gateway: Claude Haiku 4.5, 5 tokens"); | |
| 160 | assert_eq!( | |
| 161 | describe("Claude Haiku 5.5", &used(120_000, 1_000, 0, 0, 0), true, None), | |
| 162 | "AI Gateway: Claude Haiku 5.5, 121,000 tokens, long-prompt price" | |
| 163 | ); | |
| 164 | } | |
| 165 | ||
| 166 | #[test] | |
| 167 | fn a_format_is_one_of_two() { | |
| 168 | assert_eq!(format_of("openai"), "openai"); | |
| 169 | assert_eq!(format_of("OpenAI"), "openai"); | |
| 170 | assert_eq!(format_of("anthropic"), "anthropic"); | |
| 171 | assert_eq!(format_of(""), "anthropic"); | |
| 172 | assert_eq!(format_of("gopher"), "anthropic"); | |
| 173 | } | |
| 174 | ||
| 175 | #[test] | |
| 176 | fn request_ids_are_the_proxys() { | |
| 177 | assert!(valid_id("gw_01kkr2m4c8f1t7qh3d6n9w5p0x")); | |
| 178 | assert!(valid_id("gw_a-b_c")); | |
| 179 | assert!(!valid_id("run_1")); | |
| 180 | assert!(!valid_id("gw_x'; DROP TABLE ledger")); | |
| 181 | assert!(!valid_id(&format!("gw_{}", "a".repeat(80)))); | |
| 182 | } | |
| 183 | ||
| 184 | #[test] | |
| 185 | fn the_migrations_price_every_model_they_offer() { | |
| 186 | let first = include_str!("../migrations/0045_gateway.sql"); | |
| 187 | for model in ["claude-opus-5-5", "claude-sonnet-5-5", "claude-haiku-4-5"] { | |
| 188 | assert!(first.contains(&format!("('{model}', ")), "{model}"); | |
| 189 | } | |
| Merge branch 'main' into worktree-agent-a69aeabc4b0deeb97 | 190 | let second = include_str!("../migrations/0047_gateway_formats.sql"); |
| AI Gateway: OpenAI's format, open models, and your own providers | 191 | // Sonnet 5.5's cache reads are $0.10, 0.05 times input; 0045 had $0.20. |
| 192 | assert!(second.contains("SET cache_read_micros = 100000, updated_at = '2026-10-07T00:00:00Z' WHERE model = 'claude-sonnet-5-5'")); | |
| 193 | // Hour-long cache writes at twice input. | |
| 194 | assert!(second.contains("cache_write_1h_micros = 8000000 WHERE model = 'claude-opus-5-5'")); | |
| 195 | assert!(second.contains("cache_write_1h_micros = 4000000 WHERE model = 'claude-sonnet-5-5'")); | |
| 196 | assert!(second.contains("cache_write_1h_micros = 2000000 WHERE model IN ('claude-haiku-4-5', 'claude-haiku-4-5-20251001')")); | |
| 197 | // Haiku 5.5: $0.10 / $0.50 to 100,000 prompt tokens, $0.50 / $2.50 over. | |
| 198 | assert!(second.contains( | |
| 199 | "('claude-haiku-5-5', 'Claude Haiku 5.5', 'anthropic', 'chat', 100000, 500000, 10000, 125000, 200000,\n 100000, 500000, 2500000, 50000, 625000, 1000000, 0," | |
| 200 | )); | |
| 201 | // Open models at Workers AI's prices. | |
| 202 | for (model, input, output) in [ | |
| 203 | ("@cf/openai/gpt-oss-120b", 350_000, 750_000), | |
| 204 | ("@cf/openai/gpt-oss-20b", 200_000, 300_000), | |
| 205 | ("@cf/zai-org/glm-5.3-flash", 150_000, 500_000), | |
| 206 | ("@cf/moonshotai/kimi-k2.6", 950_000, 4_000_000), | |
| 207 | ] { | |
| 208 | let row = second.lines().find(|line| line.contains(&format!("('{model}'"))).unwrap_or_else(|| panic!("{model}")); | |
| 209 | assert!(row.contains(&format!("'workers-ai', 'chat', {input}, {output}, ")), "{row}"); | |
| 210 | } | |
| 211 | assert!(second.contains("('@cf/baai/bge-m3', 'BGE M3', 'workers-ai', 'embeddings', 12000, 0,")); | |
| 212 | } |
This file's history is long; its oldest lines are credited to the oldest commit read.