From e0c0832ae0248e9eeaa8fabef86638b4fcb7253a Mon Sep 17 00:00:00 2001 From: Meta Vens Date: Wed, 28 Oct 2026 18:47:06 +0100 Subject: [PATCH] connector-spend (B3c): reserve-then-settle spend in the replay ledger, one charge owner, a windowed per-caller ceiling (squash of eureka-body/connector-spend 1d51e71..45d0927) The reservation and the replay claim are one write in the connector's replay store; a failed ledger write runs nothing. Amounts are integer micro-USD with checked arithmetic, and the service derives every cost from usage and the configured price. One function, ProviderCall::charge, decides the charge: a result bound to its request and caller pays its true cost with a truthful overrun, anything else pays the reservation (ruling spend-bills-on-identity-not-validity). Each caller has a spend ceiling over fixed epoch-aligned windows in its root-owned config (connector-spend-caller-window-ceiling); the admission window is monotonic, so a clock stepping back cannot reopen budget, and a clock-locked caller is refused ClockBehindAdmissionWindow (spend-window-clock-ahead-fails-closed). Refusals CostBound, SpendCeiling and ClockBehindAdmissionWindow are typed. Spec eureka-body:cut_spec:cut-connector-spend.r1; reports h1-h5; Soul s5 holds (verdict eureka-body:verdict:cut-connector-spend.s5). Co-Authored-By: Claude Opus 5.5 --- Cargo.lock | 1 + crates/model-connector/Cargo.toml | 3 + crates/model-connector/src/contract.rs | 117 +- crates/model-connector/src/envelope.rs | 2 + crates/model-connector/src/fixtures.rs | 122 +- crates/model-connector/src/serve.rs | 109 +- crates/model-connector/src/service.rs | 1957 ++++++++++++++++- .../wire_v3/carriage_invocation.frame | 2 +- .../fixtures/wire_v3/transport_invocation.mp | Bin 799 -> 823 bytes .../wire_v3/transport_receipt_failed.mp | Bin 385 -> 517 bytes .../wire_v3/transport_result_transported.mp | Bin 914 -> 953 bytes .../tests/published_schemas.rs | 5 +- .../model-connector/tests/reasoning_state.rs | 23 +- .../model-connector/tests/wire_golden_v3.rs | 11 +- schemas/transport_invocation.json | 7 + schemas/transport_receipt.json | 27 +- schemas/transport_result.json | 32 +- 17 files changed, 2342 insertions(+), 76 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 3c3ff8d..a3cbdf7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -490,6 +490,7 @@ name = "model-connector" version = "0.1.0" dependencies = [ "aes-gcm", + "anyhow", "base64", "cultcache-rs", "cultnet-rs", diff --git a/crates/model-connector/Cargo.toml b/crates/model-connector/Cargo.toml index c54547b..33d139f 100644 --- a/crates/model-connector/Cargo.toml +++ b/crates/model-connector/Cargo.toml @@ -26,3 +26,6 @@ serde_json = "1" sha2 = "0.10" thiserror = "2" zeroize = "1" + +[dev-dependencies] +anyhow = "1" diff --git a/crates/model-connector/src/contract.rs b/crates/model-connector/src/contract.rs index 3062fe4..f22b8f9 100644 --- a/crates/model-connector/src/contract.rs +++ b/crates/model-connector/src/contract.rs @@ -190,9 +190,10 @@ pub enum FailureKind { Provider, } -/// Token accounting for one completed request. Cache reads and writes are -/// counted apart from plain input; `reasoning_tokens` is absent when the -/// provider does not report it. +/// Token accounting for one request, as the backend reports it; the service +/// prices it. Cache reads and writes are counted apart from plain input. +/// `output_tokens` includes reasoning tokens: `reasoning_tokens` is a detail +/// (absent when the provider does not report it) and is never priced apart. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] #[serde(deny_unknown_fields)] pub struct Usage { @@ -251,6 +252,10 @@ pub struct TransportInvocation { pub expires_at_unix_ms: u64, pub native_request_sha256: [u8; 32], pub provider_request_sha256: [u8; 32], + /// The most this request may cost, in micro-USD: the caller's own bound + /// on it. Required nonzero. It is part of the invocation's identity, so a + /// changed budget under a recorded request id is a `ReplayConflict`. + pub max_cost_micro_usd: u64, pub request: ProviderRequest, } @@ -259,6 +264,7 @@ impl TransportInvocation { caller_runtime_id: impl Into, expires_at_unix_ms: u64, native_request_sha256: [u8; 32], + max_cost_micro_usd: u64, request: ProviderRequest, ) -> Result { let provider_request_sha256 = provider_request_sha256(&request)?; @@ -268,6 +274,7 @@ impl TransportInvocation { expires_at_unix_ms, native_request_sha256, provider_request_sha256, + max_cost_micro_usd, request, }) } @@ -280,6 +287,9 @@ impl TransportInvocation { if self.native_request_sha256 == [0; 32] { return Err(ContractError::Invalid("native_request_sha256")); } + if self.max_cost_micro_usd == 0 { + return Err(ContractError::Invalid("max_cost_micro_usd")); + } self.request.validate()?; if self.provider_request_sha256 != provider_request_sha256(&self.request)? { return Err(ContractError::ProviderDigest); @@ -475,6 +485,15 @@ pub struct TransportReceipt { /// How many foreign `ReasoningState` items the service removed from the /// request before the backend saw it. pub foreign_reasoning_dropped: u32, + /// What the request cost, in micro-USD. The service derives it from the + /// outcome's usage and its price table; a backend never writes it. + pub cost_micro_usd: u64, + /// How far `cost_micro_usd` exceeds what the service reserved for the + /// request: zero unless the backend reported usage above the worst case + /// the request allowed, which is a breach of the backend's contract. The + /// cost is still the true one and the caller's window is charged all of + /// it. The service derives this; a backend never writes it. + pub overrun_micro_usd: u64, pub outcome: TransportOutcome, } @@ -524,6 +543,10 @@ pub enum TransportOutcome { Failed { kind: FailureKind, message: String, + /// What the provider billed for the failed request: `Some` when it + /// reported usage (`Some` of zeros when it billed nothing), `None` + /// when unknown. + usage: Option, }, } @@ -535,6 +558,12 @@ pub enum Refusal { ModelNotAllowed, OutputBound, PayloadBound, + // The worst case of the request exceeds the invocation's own + // `max_cost_micro_usd`. Plain comments: a doc comment would split the + // published schema into a union. + CostBound, + // The caller's spend window cannot hold the request's reservation. + SpendCeiling, Capacity, InFlight, Indeterminate, @@ -544,6 +573,12 @@ pub enum Refusal { // exchange, where dropping it would break the provider's own rules. A plain // comment: a doc comment would split the published schema into a union. ForeignReasoningState, + // The caller's newest admission window starts after the clock's reading, + // and the request does not fit what that window holds: the host's clock + // is behind the window and the caller stays refused until it catches up. + // A plain comment: a doc comment would split the published schema into a + // union. + ClockBehindAdmissionWindow, } /// The one wire codec of every v3 document: named MessagePack (struct fields @@ -751,6 +786,31 @@ mod tests { ); } + #[test] + fn max_cost_zero_is_malformed() { + let mut invocation = invocation(); + invocation.max_cost_micro_usd = 0; + assert_eq!( + invocation.validate(1_000, 2_000), + Err(ContractError::Invalid("max_cost_micro_usd")) + ); + invocation.max_cost_micro_usd = 1; + assert_eq!(invocation.validate(1_000, 2_000), Ok(())); + } + + #[test] + fn max_cost_is_part_of_the_invocations_identity() { + let one = invocation(); + let mut other = one.clone(); + other.max_cost_micro_usd += 1; + assert_ne!( + canonical_invocation_bytes(&one).unwrap(), + canonical_invocation_bytes(&other).unwrap() + ); + let decoded: TransportInvocation = wire_decode(&wire_encode(&one).unwrap()).unwrap(); + assert_eq!(decoded.max_cost_micro_usd, one.max_cost_micro_usd); + } + #[test] fn typed_tool_round_survives_messagepack() { let invocation = invocation(); @@ -946,6 +1006,7 @@ mod tests { sent.outcome = TransportOutcome::Failed { kind, message: "failed".to_string(), + usage: None, }; assert_eq!(sent.validate_against(&invocation), Ok(())); let decoded: TransportReceipt = wire_decode(&wire_encode(&sent).unwrap()).unwrap(); @@ -955,6 +1016,7 @@ mod tests { blank.outcome = TransportOutcome::Failed { kind: FailureKind::Provider, message: String::new(), + usage: None, }; assert_eq!( blank.validate_against(&invocation), @@ -993,6 +1055,54 @@ mod tests { } } + #[test] + fn validate_against_refuses_each_binding_field_of_the_receipt() { + let invocation = invocation(); + let good = receipt(&invocation); + assert_eq!(good.validate_against(&invocation), Ok(())); + + let mut schema = good.clone(); + schema.schema_id = "other".to_string(); + assert_eq!( + schema.validate_against(&invocation), + Err(ContractError::Schema) + ); + let mut provider_digest = good.clone(); + provider_digest.provider_request_sha256 = [9; 32]; + assert_eq!( + provider_digest.validate_against(&invocation), + Err(ContractError::ProviderDigest) + ); + let mut native_digest = good.clone(); + native_digest.native_request_sha256 = [9; 32]; + assert_eq!( + native_digest.validate_against(&invocation), + Err(ContractError::NativeDigest) + ); + let mut blank_provider = good.clone(); + blank_provider.provider = " ".to_string(); + assert_eq!( + blank_provider.validate_against(&invocation), + Err(ContractError::Invalid("receipt.provider")) + ); + let mut blank_response = good.clone(); + blank_response.outcome = TransportOutcome::Completed { + provider_response_id: Some(String::new()), + stop_reason: StopReason::EndTurn, + usage: Usage { + input_tokens: 0, + cache_read_tokens: 0, + cache_write_tokens: 0, + output_tokens: 0, + reasoning_tokens: None, + }, + }; + assert_eq!( + blank_response.validate_against(&invocation), + Err(ContractError::Invalid("receipt.provider_response_id")) + ); + } + fn reasoning(state: &[u8]) -> TransportEventPayload { TransportEventPayload::ReasoningState { state: state.to_vec(), @@ -1067,6 +1177,7 @@ mod tests { failed.outcome = TransportOutcome::Failed { kind: FailureKind::RateLimited, message: "slow down".to_string(), + usage: None, }; let partial = TransportResult::transported( vec![event(0, text("half")), event(1, reasoning(&[1, 2, 3]))], diff --git a/crates/model-connector/src/envelope.rs b/crates/model-connector/src/envelope.rs index 56bb1c2..40cfd09 100644 --- a/crates/model-connector/src/envelope.rs +++ b/crates/model-connector/src/envelope.rs @@ -462,6 +462,8 @@ pub enum ServiceError { SharedCallerKey, #[error("connector replay store failed: {0}")] ReplayStore(String), + #[error("an allowed model has no price: {0}")] + UnpricedModel(String), #[error("connector replay record is malformed")] InvalidReplayRecord, #[error("caller key changed while durable replay records still exist")] diff --git a/crates/model-connector/src/fixtures.rs b/crates/model-connector/src/fixtures.rs index b96fea8..079a99c 100644 --- a/crates/model-connector/src/fixtures.rs +++ b/crates/model-connector/src/fixtures.rs @@ -4,8 +4,15 @@ use crate::contract::*; #[cfg(feature = "daemon")] use crate::service::*; +#[cfg(feature = "daemon")] +use std::collections::BTreeMap; #[cfg(feature = "daemon")] use std::sync::atomic::{AtomicUsize, Ordering}; +#[cfg(feature = "daemon")] +use std::sync::{Arc, Mutex}; + +#[cfg(feature = "daemon")] +use cultcache_rs::{CacheBackingStore, CultCacheEnvelope, PushAllOptions}; #[cfg(feature = "daemon")] static TEST_STORE_SEQUENCE: AtomicUsize = AtomicUsize::new(0); @@ -37,6 +44,76 @@ impl Drop for TestStore { } } +/// An in-memory ledger that accepts a fixed number of writes and then refuses +/// every one, so a test can make the replay store fail where it chooses. A +/// store reopened from it sees what was written, like a restart. +#[cfg(feature = "daemon")] +pub(crate) struct FlakyStore { + held: Arc>>, + pushes_allowed: Arc, +} + +#[cfg(feature = "daemon")] +impl FlakyStore { + pub(crate) fn new(pushes_allowed: usize) -> Self { + Self { + held: Arc::default(), + pushes_allowed: Arc::new(AtomicUsize::new(pushes_allowed)), + } + } + + /// The remaining writes this store accepts, shared with it after it has + /// been handed to a service. + pub(crate) fn allowance(&self) -> Arc { + self.pushes_allowed.clone() + } + + /// A store over the same records that accepts `pushes_allowed` more writes. + pub(crate) fn reopened(&self, pushes_allowed: usize) -> Self { + Self { + held: self.held.clone(), + pushes_allowed: Arc::new(AtomicUsize::new(pushes_allowed)), + } + } +} + +#[cfg(feature = "daemon")] +impl CacheBackingStore for FlakyStore { + fn pull_all(&self) -> anyhow::Result> { + Ok(self.held.lock().unwrap().values().cloned().collect()) + } + + fn push(&mut self, entry: &CultCacheEnvelope) -> anyhow::Result<()> { + let allowed = self.pushes_allowed.load(Ordering::SeqCst); + if allowed == 0 { + return Err(anyhow::anyhow!("ledger write refused")); + } + self.pushes_allowed.store(allowed - 1, Ordering::SeqCst); + self.held + .lock() + .unwrap() + .insert(entry.key.clone(), entry.clone()); + Ok(()) + } + + fn delete(&mut self, entry: &CultCacheEnvelope) -> anyhow::Result<()> { + self.held.lock().unwrap().remove(&entry.key); + Ok(()) + } + + fn push_all( + &mut self, + entries: &[CultCacheEnvelope], + _options: PushAllOptions, + ) -> anyhow::Result<()> { + *self.held.lock().unwrap() = entries + .iter() + .map(|entry| (entry.key.clone(), entry.clone())) + .collect(); + Ok(()) + } +} + pub(crate) fn request() -> ProviderRequest { let mut request = ProviderRequest::new( "request-1", @@ -56,8 +133,16 @@ pub(crate) fn request() -> ProviderRequest { request } +/// Generous next to the fixture request's worst case (about 64,000), so it +/// binds only in the tests that are about spend. +pub(crate) const MAX_COST: u64 = 10_000_000; + +/// What the fixture usage (10 input and 3 output tokens) costs under +/// [`fixture_price`]: 10 x 3,000,000 + 3 x 15,000,000 = 75,000,000 / 1e6. +pub(crate) const FIXTURE_USAGE_COST: u64 = 75; + pub(crate) fn invocation() -> TransportInvocation { - TransportInvocation::new("epiphany-yggdrasil", 2_000, [7; 32], request()).unwrap() + TransportInvocation::new("epiphany-yggdrasil", 2_000, [7; 32], MAX_COST, request()).unwrap() } pub(crate) fn receipt(invocation: &TransportInvocation) -> TransportReceipt { @@ -70,6 +155,8 @@ pub(crate) fn receipt(invocation: &TransportInvocation) -> TransportReceipt { model: invocation.request.model.clone(), provider: "fixture-sse".to_string(), foreign_reasoning_dropped: 0, + cost_micro_usd: FIXTURE_USAGE_COST, + overrun_micro_usd: 0, outcome: TransportOutcome::Completed { provider_response_id: Some("response-1".to_string()), stop_reason: StopReason::EndTurn, @@ -83,6 +170,34 @@ pub(crate) fn receipt(invocation: &TransportInvocation) -> TransportReceipt { }, } } +/// A price with a distinct nonzero rate per class, so a swapped rate changes a cost. +#[cfg(feature = "daemon")] +pub(crate) fn fixture_price() -> ModelPrice { + ModelPrice { + rates: Rates { + input: 3_000_000, + cache_write_short: 3_750_000, + cache_write_long: 6_000_000, + cache_read: 300_000, + output: 15_000_000, + }, + long_context: None, + overhead_tokens: 100, + } +} + +#[cfg(feature = "daemon")] +pub(crate) fn fixture_prices() -> PriceTable { + PriceTable::from([("gpt-5.4".to_string(), fixture_price())]) +} + +/// Far above anything a fixture request reserves: the ceiling binds only in +/// the tests that are about it. +#[cfg(feature = "daemon")] +pub(crate) const CEILING: u64 = 1 << 40; +#[cfg(feature = "daemon")] +pub(crate) const WINDOW_MS: u64 = 3_600_000; + #[cfg(feature = "daemon")] pub(crate) fn caller( caller_runtime_id: &str, @@ -97,6 +212,8 @@ pub(crate) fn caller( max_concurrent_requests, 64 * 1024, 32_768, + SpendCeiling(CEILING), + SpendWindowMs(WINDOW_MS), ) .unwrap() } @@ -106,6 +223,7 @@ pub(crate) fn open_service(store: &TestStore, max_concurrent_requests: usize) -> TransportService::open( &store.path, "fixture-sse", + fixture_prices(), [ caller( "epiphany-yggdrasil", @@ -131,5 +249,5 @@ pub(crate) fn security(key: &str) -> TransportKey { pub(crate) fn invocation_for(caller_runtime_id: &str, request_id: &str) -> TransportInvocation { let mut request = request(); request.request_id = request_id.to_string(); - TransportInvocation::new(caller_runtime_id, 2_000, [7; 32], request).unwrap() + TransportInvocation::new(caller_runtime_id, 2_000, [7; 32], MAX_COST, request).unwrap() } diff --git a/crates/model-connector/src/serve.rs b/crates/model-connector/src/serve.rs index 40ece2d..6d87f6e 100644 --- a/crates/model-connector/src/serve.rs +++ b/crates/model-connector/src/serve.rs @@ -15,7 +15,10 @@ use crate::envelope::{ ServiceError, TransportFrameError, decode_invocation_frame, encode_result_frame, frame_bytes_for_payload, read_transport_frame, write_transport_frame, }; -use crate::service::{CallerAdmission, ProviderCall, TransportAdmission, TransportService}; +use crate::service::{ + CallerAdmission, PriceTable, ProviderCall, SpendCeiling, SpendWindowMs, TransportAdmission, + TransportService, +}; /// The only provider seam of the connector kit. A connector binary implements /// it over one model provider; the kit names no provider. @@ -111,6 +114,12 @@ pub struct CallerConfig { pub max_concurrent_requests: usize, pub max_payload_bytes: usize, pub max_output_tokens: u32, + /// The most the caller may spend in one window, in micro-USD. Required + /// nonzero. A fuse set once, generously above what the dial allows. + pub spend_ceiling_micro_usd: u64, + /// The length of the caller's spend window, in milliseconds. Windows are + /// fixed and aligned to the unix epoch. Required nonzero. + pub spend_window_ms: u64, } pub fn load_admissions(callers: &[CallerConfig]) -> Result, ServeError> { @@ -131,6 +140,8 @@ pub fn load_admissions(callers: &[CallerConfig]) -> Result, caller.max_concurrent_requests, caller.max_payload_bytes, caller.max_output_tokens, + SpendCeiling(caller.spend_ceiling_micro_usd), + SpendWindowMs(caller.spend_window_ms), ) .map_err(ServeError::Service)?, ); @@ -139,12 +150,15 @@ pub fn load_admissions(callers: &[CallerConfig]) -> Result, } /// Serves admitted callers on a loopback listener until the listener fails. +/// `prices` is the connector binary's config document's table of what every +/// allowed model costs (B4 owns that document). /// The caller binds; a listener that is not bound to loopback is refused, and /// peers that are not loopback are dropped. pub fn serve( listener: TcpListener, limits: ServeLimits, admissions: Vec, + prices: PriceTable, replay_store: &Path, backend: Arc, ) -> Result<(), ServeError> { @@ -163,6 +177,7 @@ pub fn serve( let service = Arc::new(Mutex::new(TransportService::open( replay_store, backend.provider(), + prices, admissions, limits.max_expiry_skew_ms, )?)); @@ -336,7 +351,10 @@ mod tests { FailureKind, Refusal, TransportDisposition, TransportEvent, TransportEventPayload, TransportInvocation, TransportOutcome, TransportResult, }; - use crate::fixtures::{TestStore, caller, invocation_for, receipt, request}; + use crate::fixtures::{ + CEILING, MAX_COST, TestStore, WINDOW_MS, caller, fixture_prices, invocation_for, receipt, + request, + }; const CALLER: &str = "epiphany-yggdrasil"; const KEY: &str = "epiphany-distinct-test-key"; @@ -412,7 +430,14 @@ mod tests { let (sender, receiver) = std::sync::mpsc::channel(); let path = path.to_path_buf(); thread::spawn(move || { - let _ = sender.send(serve(listener, limits, admissions, &path, backend)); + let _ = sender.send(serve( + listener, + limits, + admissions, + fixture_prices(), + &path, + backend, + )); }); receiver.recv_timeout(Duration::from_secs(2)).ok() } @@ -438,7 +463,16 @@ mod tests { let listener = TcpListener::bind("127.0.0.1:0").unwrap(); let endpoint = listener.local_addr().unwrap(); let path = store.path.clone(); - thread::spawn(move || serve(listener, limits(), admissions, &path, backend)); + thread::spawn(move || { + serve( + listener, + limits(), + admissions, + fixture_prices(), + &path, + backend, + ) + }); let client = ConnectorClient::new(endpoint, KEY, FRAME, Some(Duration::from_secs(5))).unwrap(); (client, endpoint) @@ -456,10 +490,37 @@ mod tests { max_concurrent_requests: *capacity, max_payload_bytes: 1_048_576, max_output_tokens: 16_384, + spend_ceiling_micro_usd: CEILING, + spend_window_ms: WINDOW_MS, }) .collect() } + #[test] + fn serve_refuses_a_request_over_its_max_cost_without_executing_it() { + let store = TestStore::new(); + let backend = CountingBackend::new(true); + let (client, _) = start(backend.clone(), &store); + let mut invocation = fresh("request-costly"); + invocation.max_cost_micro_usd = 1; + let refused = client.execute(&invocation).unwrap(); + assert_eq!( + refused.disposition, + TransportDisposition::Refused(Refusal::CostBound) + ); + assert_eq!(backend.executions.load(Ordering::SeqCst), 0); + + // The refusal was not recorded: the same id with a budget that holds + // the worst case is admitted and executed. + invocation.max_cost_micro_usd = MAX_COST; + let admitted = client.execute(&invocation).unwrap(); + assert!(matches!( + admitted.disposition, + TransportDisposition::Transported { .. } + )); + assert_eq!(backend.executions.load(Ordering::SeqCst), 1); + } + #[test] fn serve_end_to_end_replays_without_a_second_execution() { let store = TestStore::new(); @@ -496,6 +557,7 @@ mod tests { CALLER, invocation.expires_at_unix_ms, [7; 32], + MAX_COST, conflicting_request, ) .unwrap(); @@ -580,6 +642,7 @@ mod tests { let mut service = TransportService::open( &store.path, "fixture-sse", + fixture_prices(), vec![caller(CALLER, KEY, 4)], 300_000, ) @@ -611,6 +674,7 @@ mod tests { listener, narrowed, vec![caller(CALLER, KEY, 4)], + fixture_prices(), &path, served, ) @@ -693,7 +757,13 @@ mod tests { fn the_smallest_frame_also_carries_the_failure_that_stands_in_for_a_result() { let long = "x".repeat(crate::envelope::MAX_IDENTITY_BYTES); let plaintext = crate::contract::wire_encode(&crate::service::result_size_failure( - &long, &long, &long, [0xff; 32], [0xff; 32], 0, + &long, + &long, + &long, + [0xff; 32], + [0xff; 32], + 0, + u64::MAX, )) .unwrap() .len(); @@ -781,7 +851,18 @@ mod tests { } fn admission(id: &str, key: &str, slots: usize, payload: usize) -> CallerAdmission { - CallerAdmission::new(id, key, 1, ["gpt-5.4".to_string()], slots, payload, 32_768).unwrap() + CallerAdmission::new( + id, + key, + 1, + ["gpt-5.4".to_string()], + slots, + payload, + 32_768, + SpendCeiling(CEILING), + SpendWindowMs(WINDOW_MS), + ) + .unwrap() } /// `Some(field)` when serve refused the configuration, `None` while it serves. @@ -887,7 +968,7 @@ mod tests { ..limits() }; let admissions = vec![admission(CALLER, KEY, 4, bound)]; - thread::spawn(move || serve(listener, tight, admissions, &path, served)); + thread::spawn(move || serve(listener, tight, admissions, fixture_prices(), &path, served)); let client = ConnectorClient::new(endpoint, KEY, frame, Some(Duration::from_secs(5))).unwrap(); let reply = client.execute(&invocation).unwrap(); @@ -987,6 +1068,7 @@ mod tests { && receipt.outcome == TransportOutcome::Failed { kind: FailureKind::ResultSize, message: "provider result exceeded its byte bound".to_string(), + usage: None, } ) } @@ -1001,6 +1083,7 @@ mod tests { let mut service = TransportService::open( &store.path, "fixture-sse", + fixture_prices(), vec![caller(CALLER, KEY, 4)], 300_000, ) @@ -1103,6 +1186,7 @@ mod tests { TransportService::open( &store.path, "fixture-sse", + fixture_prices(), vec![caller(CALLER, KEY, 4)], 300_000, ) @@ -1160,7 +1244,16 @@ mod tests { let backend = CountingBackend::new(true); let path = store.path.clone(); let served = backend.clone(); - thread::spawn(move || serve(listener, limits(), admissions, &path, served)); + thread::spawn(move || { + serve( + listener, + limits(), + admissions, + fixture_prices(), + &path, + served, + ) + }); let client = ConnectorClient::new(endpoint, KEY, FRAME, Some(Duration::from_secs(5))).unwrap(); client.execute(&fresh("request-key")).unwrap(); diff --git a/crates/model-connector/src/service.rs b/crates/model-connector/src/service.rs index 9a2ea2a..8d4fd80 100644 --- a/crates/model-connector/src/service.rs +++ b/crates/model-connector/src/service.rs @@ -1,21 +1,50 @@ -use std::collections::{HashMap, HashSet}; +use std::collections::{BTreeMap, HashMap, HashSet}; use std::path::Path; -use cultcache_rs::{CultCache, DatabaseEntry, OwnedRedbMessagePackBackingStore}; +use cultcache_rs::{CacheBackingStore, CultCache, DatabaseEntry, OwnedRedbMessagePackBackingStore}; use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; use zeroize::Zeroizing; use crate::contract::{ - ContractError, FailureKind, InputItem, ProviderRequest, RECEIPT_SCHEMA_ID, RESULT_SCHEMA_ID, - Refusal, TransportDisposition, TransportEvent, TransportInvocation, TransportKey, - TransportOutcome, TransportReceipt, TransportResult, canonical_invocation_bytes, require_id, + CacheRetention, ContractError, FailureKind, InputItem, ProviderRequest, RECEIPT_SCHEMA_ID, + RESULT_SCHEMA_ID, Refusal, TransportDisposition, TransportEvent, TransportInvocation, + TransportKey, TransportOutcome, TransportReceipt, TransportResult, Usage, + canonical_invocation_bytes, require_id, wire_encode, }; use crate::envelope::{ EnvelopeKind, MAX_IDENTITY_BYTES, ServiceError, TransportEnvelope, decrypt_invocation, encode_result_frame, encrypt_result, validate_envelope, }; +/// What a caller may spend in one window, in micro-USD. A distinct type from +/// [`SpendWindowMs`] so the pair cannot be passed swapped: +/// +/// ``` +/// use model_connector::{CallerAdmission, SpendCeiling, SpendWindowMs}; +/// CallerAdmission::new( +/// "caller-1", "a-connection-key", 1, ["model-1".to_string()], 1, 1, 1, +/// SpendCeiling(50_000_000), SpendWindowMs(3_600_000), +/// ) +/// .unwrap(); +/// ``` +/// +/// ```compile_fail +/// use model_connector::{CallerAdmission, SpendCeiling, SpendWindowMs}; +/// CallerAdmission::new( +/// "caller-1", "a-connection-key", 1, ["model-1".to_string()], 1, 1, 1, +/// SpendWindowMs(3_600_000), SpendCeiling(50_000_000), +/// ) +/// .unwrap(); +/// ``` +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct SpendCeiling(pub u64); + +/// How long a spend window lasts, in milliseconds. Windows are fixed and +/// aligned to the unix epoch. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct SpendWindowMs(pub u64); + pub struct CallerAdmission { caller_runtime_id: String, connection_key_epoch: u32, @@ -24,9 +53,14 @@ pub struct CallerAdmission { max_concurrent_requests: usize, max_payload_bytes: usize, max_output_tokens: u32, + spend_ceiling_micro_usd: u64, + spend_window_ms: u64, } impl CallerAdmission { + // The ceiling and window ride beside the other limits the root-owned caller + // config sets, one parameter each. + #[allow(clippy::too_many_arguments)] pub fn new( caller_runtime_id: impl Into, connection_key: impl Into, @@ -35,7 +69,10 @@ impl CallerAdmission { max_concurrent_requests: usize, max_payload_bytes: usize, max_output_tokens: u32, + spend_ceiling: SpendCeiling, + spend_window: SpendWindowMs, ) -> Result { + let (spend_ceiling_micro_usd, spend_window_ms) = (spend_ceiling.0, spend_window.0); let caller_runtime_id = caller_runtime_id.into(); require_id(&caller_runtime_id, "caller_runtime_id")?; let allowed_models = allowed_models.into_iter().collect::>(); @@ -48,6 +85,8 @@ impl CallerAdmission { || max_concurrent_requests == 0 || max_payload_bytes == 0 || max_output_tokens == 0 + || spend_ceiling_micro_usd == 0 + || spend_window_ms == 0 { return Err(ServiceError::InvalidAdmission); } @@ -61,6 +100,8 @@ impl CallerAdmission { max_concurrent_requests, max_payload_bytes, max_output_tokens, + spend_ceiling_micro_usd, + spend_window_ms, }) } @@ -77,6 +118,197 @@ impl CallerAdmission { } } +/// What one million tokens of each class cost, in micro-USD. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Rates { + pub input: u64, + pub cache_write_short: u64, + pub cache_write_long: u64, + pub cache_read: u64, + pub output: u64, +} + +/// The rates a prompt strictly above `above_prompt_tokens` is priced at. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct LongContext { + pub above_prompt_tokens: u64, + pub rates: Rates, +} + +/// One model's price. Explicit zero rates are admitted (a local model is +/// free); a missing price is not. `overhead_tokens` is added to the byte +/// bound of a request when its worst case is reserved. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ModelPrice { + pub rates: Rates, + pub long_context: Option, + pub overhead_tokens: u64, +} + +/// The price of every model a caller may be allowed, supplied at open by the +/// connector binary's config document. +pub type PriceTable = BTreeMap; + +const TOKENS_PER_RATE: u128 = 1_000_000; + +/// Rounds up to whole micro-USD; `None` when the amount does not fit a u64. +fn checked_ceil_micro_usd(rate_weighted_tokens: u128) -> Option { + let whole = rate_weighted_tokens / TOKENS_PER_RATE; + let rounded = whole + u128::from(!rate_weighted_tokens.is_multiple_of(TOKENS_PER_RATE)); + u64::try_from(rounded).ok() +} + +fn ceil_micro_usd(rate_weighted_tokens: u128) -> u64 { + checked_ceil_micro_usd(rate_weighted_tokens).unwrap_or(u64::MAX) +} + +impl Rates { + /// The dearer of each class. + fn dearest(self, other: Rates) -> Rates { + Rates { + input: self.input.max(other.input), + cache_write_short: self.cache_write_short.max(other.cache_write_short), + cache_write_long: self.cache_write_long.max(other.cache_write_long), + cache_read: self.cache_read.max(other.cache_read), + output: self.output.max(other.output), + } + } +} + +fn rates_for(price: &ModelPrice, prompt_tokens: u128) -> &Rates { + match &price.long_context { + Some(long) if prompt_tokens > u128::from(long.above_prompt_tokens) => &long.rates, + _ => &price.rates, + } +} + +/// A cache write nobody asked to keep short is priced dear. +fn write_rate(rates: &Rates, retention: Option) -> u64 { + match retention { + Some(CacheRetention::Short) => rates.cache_write_short, + Some(CacheRetention::Long) | None => rates.cache_write_long, + } +} + +/// What `usage` costs under `price`, in micro-USD: the sum over token classes +/// of tokens times rate, rounded up once, saturating at `u64::MAX`. +pub(crate) fn cost_micro_usd( + price: &ModelPrice, + usage: &Usage, + retention: Option, +) -> u64 { + let prompt = u128::from(usage.input_tokens) + + u128::from(usage.cache_read_tokens) + + u128::from(usage.cache_write_tokens); + let rates = rates_for(price, prompt); + let total = [ + u128::from(usage.input_tokens) * u128::from(rates.input), + u128::from(usage.cache_read_tokens) * u128::from(rates.cache_read), + u128::from(usage.cache_write_tokens) * u128::from(write_rate(rates, retention)), + u128::from(usage.output_tokens) * u128::from(rates.output), + ] + .into_iter() + .fold(0u128, u128::saturating_add); + ceil_micro_usd(total) +} + +/// The worst case `request` can cost, in micro-USD: its encoded length plus +/// the model's overhead as input tokens (a token is at least one byte) at the +/// dearest input rate the cache retention allows, plus `max_output_tokens` at +/// the output rate, rounded up once. A prompt that may cross the long-context +/// threshold is bounded at the dearer of the two rate sets, so a backend whose +/// usage prices above this broke its contract. `None` when the worst case does +/// not fit a u64: such a request can never be afforded. +pub(crate) fn reservation_micro_usd( + price: &ModelPrice, + request: &ProviderRequest, +) -> Result, ContractError> { + let bound = (wire_encode(request)?.len() as u128) + u128::from(price.overhead_tokens); + let rates = match &price.long_context { + Some(long) if bound > u128::from(long.above_prompt_tokens) => { + price.rates.dearest(long.rates) + } + _ => price.rates, + }; + let input_rate = rates + .input + .max(rates.cache_read) + .max(write_rate(&rates, request.cache_retention)); + let total = bound.checked_mul(u128::from(input_rate)).and_then(|input| { + input.checked_add(u128::from(request.max_output_tokens) * u128::from(rates.output)) + }); + Ok(total.and_then(checked_ceil_micro_usd)) +} + +/// Moves `held` of a caller's window: releases `release`, then adds `charge`. +fn add_spend( + window_spend: &mut HashMap>, + caller_runtime_id: &str, + window: u64, + release: u64, + charge: u64, +) { + let held = window_spend + .entry(caller_runtime_id.to_string()) + .or_default() + .entry(window) + .or_insert(0); + *held = held.saturating_sub(release).saturating_add(charge); +} + +/// The clock a request is admitted at: `now`, held monotonic per caller. A +/// caller's window never moves back, so a clock that steps back (or a boot +/// with the clock behind a recorded window) cannot reopen spent budget: the +/// request is admitted into the caller's newest recorded window. The record +/// stores this value, so the window re-derives from the records at open. +fn admission_clock( + window_spend: &HashMap>, + caller_runtime_id: &str, + now_unix_ms: u64, + spend_window_ms: u64, +) -> u64 { + let newest = window_spend + .get(caller_runtime_id) + .and_then(|windows| windows.keys().next_back().copied()) + .unwrap_or(0); + now_unix_ms.max(newest.saturating_mul(spend_window_ms)) +} + +/// Drops the windows of `caller_runtime_id` that nothing can still read: +/// those before `floor` that hold no live request's reservation. Admission +/// reads the current window alone, and a live request settles into the window +/// it was admitted in, so nothing older is kept. +fn prune_windows( + window_spend: &mut HashMap>, + live: &HashMap, + caller_runtime_id: &str, + floor: u64, +) { + if let Some(windows) = window_spend.get_mut(caller_runtime_id) { + windows.retain(|window, _| { + *window >= floor + || live.iter().any(|(identity, live_window)| { + live_window == window && identity.caller_runtime_id == caller_runtime_id + }) + }); + } +} + +/// `prune_windows` with the caller's newest window as the floor, for the +/// places that do not know the clock (open, settle). +fn prune_to_latest( + window_spend: &mut HashMap>, + live: &HashMap, + caller_runtime_id: &str, +) { + let latest = window_spend + .get(caller_runtime_id) + .and_then(|windows| windows.keys().next_back().copied()); + if let Some(latest) = latest { + prune_windows(window_spend, live, caller_runtime_id, latest); + } +} + #[derive(Debug, Clone, PartialEq, Eq, Hash)] struct RequestIdentity { caller_runtime_id: String, @@ -85,12 +317,21 @@ struct RequestIdentity { #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] enum ReplayState { - Active, - Completed { response: TransportEnvelope }, + /// Admitted and possibly executing: the worst case is held against the + /// caller's window until the request settles (an indeterminate one is + /// billed at it). + Active { reserved_micro_usd: u64 }, + /// `overrun_micro_usd` is what the charge exceeded the reservation by: a + /// backend contract breach, kept so a replayed stand-in still says so. + Completed { + response: TransportEnvelope, + charged_micro_usd: u64, + overrun_micro_usd: u64, + }, } #[derive(Clone, Debug, PartialEq, Eq, DatabaseEntry)] -#[cultcache(type = "gamecult.model.replay_record.v1")] +#[cultcache(type = "gamecult.model.replay_record.v2")] struct ReplayRecord { #[cultcache(key = 0)] caller_runtime_id: String, @@ -104,6 +345,8 @@ struct ReplayRecord { expires_at_unix_ms: u64, #[cultcache(key = 5)] state: ReplayState, + #[cultcache(key = 6)] + admitted_at_unix_ms: u64, } /// What a backend is handed: the request with every foreign reasoning item @@ -115,6 +358,8 @@ struct ReplayRecord { pub struct ProviderCall { request: ProviderRequest, receipt: TransportReceipt, + price: ModelPrice, + reserved_micro_usd: u64, } impl ProviderCall { @@ -130,20 +375,69 @@ impl ProviderCall { } /// The result of this call. The whole receipt (schema, ids, both digests, - /// model, provider tag and drop count) is derived here, so a backend + /// model, provider tag, drop count and cost) is derived here, so a backend /// writes only the events and the outcome. pub fn transported( &self, events: Vec, outcome: TransportOutcome, ) -> TransportResult { - TransportResult::transported( - events, + TransportResult::transported(events, { + let cost_micro_usd = self.cost_of(&outcome); TransportReceipt { + cost_micro_usd, + overrun_micro_usd: self.overrun_of(cost_micro_usd), outcome, ..self.receipt.clone() + } + }) + } + + /// What a completed `result` is billed: the one owner of the charge + /// decision. A transported result bound to `invocation` by request id and + /// caller (on the result and on its receipt) costs what its own outcome's + /// usage derives, with the overrun past the reservation recorded + /// truthfully, whether or not it is otherwise malformed. Any other result + /// (refused, or bound to another request or caller) costs the whole + /// reservation. Validity decides what the caller is told, never this. + fn charge(&self, invocation: &TransportInvocation, result: &TransportResult) -> Charge { + match &result.disposition { + TransportDisposition::Transported { receipt, .. } + if result.request_id == invocation.request_id() + && result.caller_runtime_id == invocation.caller_runtime_id + && receipt.request_id == invocation.request_id() + && receipt.caller_runtime_id == invocation.caller_runtime_id => + { + let cost_micro_usd = self.cost_of(&receipt.outcome); + Charge { + cost_micro_usd, + overrun_micro_usd: self.overrun_of(cost_micro_usd), + } + } + _ => Charge { + cost_micro_usd: self.reserved_micro_usd, + overrun_micro_usd: 0, }, - ) + } + } + + /// How far `cost_micro_usd` exceeds what was reserved. The reservation is + /// the worst case the kit can compute, so a nonzero overrun means the + /// backend broke its contract (usage above its own request's bounds). + fn overrun_of(&self, cost_micro_usd: u64) -> u64 { + cost_micro_usd.saturating_sub(self.reserved_micro_usd) + } + + /// What `outcome` costs: its usage priced, or the whole reservation when a + /// failure does not say what it billed (unknown is billed as the worst case). + fn cost_of(&self, outcome: &TransportOutcome) -> u64 { + match outcome { + TransportOutcome::Completed { usage, .. } + | TransportOutcome::Failed { + usage: Some(usage), .. + } => cost_micro_usd(&self.price, usage, self.request.cache_retention), + TransportOutcome::Failed { usage: None, .. } => self.reserved_micro_usd, + } } } @@ -168,19 +462,49 @@ pub enum TransportAdmission { pub struct TransportService { callers: HashMap, - live: HashSet, + /// The requests between claim and settlement, each with the window it was + /// admitted in. + live: HashMap, replay_store: CultCache, max_expiry_skew_ms: u64, provider: String, + prices: PriceTable, + /// Per caller, what each spend window holds: the reservations of its + /// Active records and the charges of its Completed ones. A cache of the + /// replay records, re-derived at open and moved only after a record's put + /// returns Ok; the records are the ledger. A request belongs to the window + /// it was admitted in, so settlement never moves it to another. Only what + /// admission and settlement can still read is kept: the caller's newest + /// window and any older one a live request is reserved in. + window_spend: HashMap>, } impl TransportService { /// `provider` is the tag of the backend this service fronts. The service /// owns it from here on: it decides which reasoning state is foreign and - /// what provider every receipt names. + /// what provider every receipt names. `prices` must price every model any + /// caller is allowed. pub fn open( replay_store_path: &Path, provider: &str, + prices: PriceTable, + admissions: impl IntoIterator, + max_expiry_skew_ms: u64, + ) -> Result { + Self::with_store( + OwnedRedbMessagePackBackingStore::new(replay_store_path).map_err(replay_store_error)?, + provider, + prices, + admissions, + max_expiry_skew_ms, + ) + } + + /// `open` over any backing store, so a test can inject one that fails. + pub fn with_store( + store: impl CacheBackingStore + 'static, + provider: &str, + prices: PriceTable, admissions: impl IntoIterator, max_expiry_skew_ms: u64, ) -> Result { @@ -204,21 +528,27 @@ impl TransportService { if callers.is_empty() { return Err(ServiceError::InvalidAdmission); } + if let Some(model) = callers + .values() + .flat_map(|caller| caller.allowed_models.iter()) + .filter(|model| !prices.contains_key(*model)) + .min() + { + return Err(ServiceError::UnpricedModel(model.clone())); + } let mut replay_store = CultCache::new(); replay_store .register_entry_type::() .map_err(replay_store_error)?; replay_store - .add_generic_backing_store( - OwnedRedbMessagePackBackingStore::new(replay_store_path) - .map_err(replay_store_error)?, - ) + .add_generic_backing_store(store) .map_err(replay_store_error)?; replay_store .pull_all_backing_stores() .map_err(replay_store_error)?; + let mut window_spend: HashMap> = HashMap::new(); for (key, record) in replay_store .get_all_with_keys::() .map_err(replay_store_error)? @@ -230,6 +560,7 @@ impl TransportService { if key != replay_key(&identity) || record.invocation_sha256 == [0; 32] || record.expires_at_unix_ms == 0 + || record.admitted_at_unix_ms == 0 { return Err(ServiceError::InvalidReplayRecord); } @@ -239,16 +570,47 @@ impl TransportService { if record.connection_key_epoch != caller.connection_key_epoch { return Err(ServiceError::ReplayCallerKeyMismatch); } + // An Active record left by a crash counts at its full reservation: + // indeterminate is billed as the worst case. + let held = match &record.state { + ReplayState::Active { reserved_micro_usd } => *reserved_micro_usd, + ReplayState::Completed { + charged_micro_usd, .. + } => *charged_micro_usd, + }; + add_spend( + &mut window_spend, + &record.caller_runtime_id, + record.admitted_at_unix_ms / caller.spend_window_ms, + 0, + held, + ); + } + let live = HashMap::new(); + let spenders: Vec = window_spend.keys().cloned().collect(); + for caller_runtime_id in &spenders { + prune_to_latest(&mut window_spend, &live, caller_runtime_id); } Ok(Self { callers, - live: HashSet::new(), + live, replay_store, max_expiry_skew_ms, provider: provider.to_string(), + prices, + window_spend, }) } + /// What `caller_runtime_id`'s spend window `window` holds. + pub(crate) fn window_spend(&self, caller_runtime_id: &str, window: u64) -> u64 { + self.window_spend + .get(caller_runtime_id) + .and_then(|windows| windows.get(&window)) + .copied() + .unwrap_or(0) + } + /// Admits one request. `max_frame_bytes` is the frame every reply must /// fit, the same bound `complete` is given: a recorded result that no /// longer fits it (the service reopened at a smaller frame) is answered @@ -269,6 +631,9 @@ impl TransportService { .ok_or(ServiceError::CallerNotAdmitted)?; let security = caller.security.clone(); let max_concurrent_requests = caller.max_concurrent_requests; + let connection_key_epoch = caller.connection_key_epoch; + let spend_ceiling = caller.spend_ceiling_micro_usd; + let spend_window_ms = caller.spend_window_ms; let invocation = decrypt_invocation(envelope, &security)?; if envelope.request_id.len() > MAX_IDENTITY_BYTES { let refusal = encrypt_result( @@ -298,7 +663,6 @@ impl TransportService { Ok(partition) => partition, Err(refusal) => return self.reply_refusal(&invocation, refusal, &security), }; - let identity = RequestIdentity { caller_runtime_id: invocation.caller_runtime_id.clone(), request_id: invocation.request_id().to_string(), @@ -314,15 +678,25 @@ impl TransportService { return self.reply_refusal(&invocation, Refusal::ReplayConflict, &security); } return match &record.state { - ReplayState::Completed { response } => { + ReplayState::Completed { + response, + charged_micro_usd, + overrun_micro_usd, + } => { if fits_frame(response, max_frame_bytes)? { return Ok(TransportAdmission::Reply(response.clone())); } + // The stand-in keeps the recorded charge: the request cost + // what it cost whatever frame its answer travels in. let stand_in = self.seal_stand_in( &invocation, &security, max_frame_bytes, FailureKind::ResultSize, + Charge { + cost_micro_usd: *charged_micro_usd, + overrun_micro_usd: *overrun_micro_usd, + }, )?; self.replay_store .put( @@ -330,6 +704,8 @@ impl TransportService { &ReplayRecord { state: ReplayState::Completed { response: stand_in.clone(), + charged_micro_usd: *charged_micro_usd, + overrun_micro_usd: *overrun_micro_usd, }, ..record.clone() }, @@ -337,8 +713,8 @@ impl TransportService { .map_err(replay_store_error)?; Ok(TransportAdmission::Reply(stand_in)) } - ReplayState::Active => { - let refusal = if self.live.contains(&identity) { + ReplayState::Active { .. } => { + let refusal = if self.live.contains_key(&identity) { Refusal::InFlight } else { Refusal::Indeterminate @@ -347,27 +723,78 @@ impl TransportService { } }; } + // The worst case is priced at today's table, so it is checked only for + // new work: a request already recorded is answered whatever it would + // cost now. + let price = *self + .prices + .get(&invocation.request.model) + .ok_or_else(|| ServiceError::UnpricedModel(invocation.request.model.clone()))?; + let reserved_micro_usd = match reservation_micro_usd(&price, &request)? { + Some(reserved) if reserved <= invocation.max_cost_micro_usd => reserved, + _ => return self.reply_refusal(&invocation, Refusal::CostBound, &security), + }; let active_for_caller = self .live - .iter() + .keys() .filter(|key| key.caller_runtime_id == invocation.caller_runtime_id) .count(); if active_for_caller >= max_concurrent_requests { return self.reply_refusal(&invocation, Refusal::Capacity, &security); } + // A request belongs to the window it is admitted in. The ceiling is + // checked after the replay lookup, so a recorded answer is replayed + // whatever the caller has left. + let admitted_at_unix_ms = admission_clock( + &self.window_spend, + &identity.caller_runtime_id, + now_unix_ms, + spend_window_ms, + ); + let window = admitted_at_unix_ms / spend_window_ms; + prune_windows( + &mut self.window_spend, + &self.live, + &identity.caller_runtime_id, + window, + ); + if self + .window_spend(&identity.caller_runtime_id, window) + .saturating_add(reserved_micro_usd) + > spend_ceiling + { + // A clock reading behind the caller's newest window is named: the + // refusal holds until the clock reaches that window's start. + let refusal = if now_unix_ms < window.saturating_mul(spend_window_ms) { + Refusal::ClockBehindAdmissionWindow + } else { + Refusal::SpendCeiling + }; + return self.reply_refusal(&invocation, refusal, &security); + } + // Reserving and claiming the request are one write: a ledger that + // cannot be written means nothing runs. let replay_record = ReplayRecord { caller_runtime_id: identity.caller_runtime_id.clone(), request_id: identity.request_id.clone(), - connection_key_epoch: caller.connection_key_epoch, + connection_key_epoch, invocation_sha256, expires_at_unix_ms: invocation.expires_at_unix_ms, - state: ReplayState::Active, + state: ReplayState::Active { reserved_micro_usd }, + admitted_at_unix_ms, }; self.replay_store .put(replay_key(&identity), &replay_record) .map_err(replay_store_error)?; - self.live.insert(identity); + add_spend( + &mut self.window_spend, + &identity.caller_runtime_id, + window, + 0, + reserved_micro_usd, + ); + self.live.insert(identity, window); Ok(TransportAdmission::Execute(Box::new(ExecutionClaim { call: ProviderCall { receipt: TransportReceipt { @@ -380,12 +807,17 @@ impl TransportService { provider: self.provider.clone(), foreign_reasoning_dropped, // Replaced by the backend's outcome in `transported`. + cost_micro_usd: 0, + overrun_micro_usd: 0, outcome: TransportOutcome::Failed { kind: FailureKind::Provider, message: String::new(), + usage: None, }, }, request, + price, + reserved_micro_usd, }, invocation, invocation_sha256, @@ -408,7 +840,7 @@ impl TransportService { caller_runtime_id: invocation.caller_runtime_id.clone(), request_id: invocation.request_id().to_string(), }; - if !self.live.contains(&identity) { + if !self.live.contains_key(&identity) { return Err(ServiceError::NoActiveRequest); } let active = self @@ -416,9 +848,10 @@ impl TransportService { .get::(&replay_key(&identity)) .map_err(replay_store_error)? .ok_or(ServiceError::NoActiveRequest)?; - if active.invocation_sha256 != claim.invocation_sha256 - || !matches!(active.state, ReplayState::Active) - { + let ReplayState::Active { reserved_micro_usd } = active.state else { + return Err(ServiceError::ActiveRequestMismatch); + }; + if active.invocation_sha256 != claim.invocation_sha256 { return Err(ServiceError::ActiveRequestMismatch); } let caller = self @@ -426,20 +859,27 @@ impl TransportService { .get(&identity.caller_runtime_id) .ok_or(ServiceError::CallerNotAdmitted)?; let security = caller.security.clone(); + let connection_key_epoch = caller.connection_key_epoch; + let spend_window_ms = caller.spend_window_ms; // A result the service refuses (not a transported result, not bound to // this invocation, or carrying a provider tag or drop count the call // did not derive) is answered like an oversized one: a typed failure // is recorded and the slot released, so the caller gets an answer and // a retry replays it. let admissible = self.admits_result(&result, &claim); + let charge = claim.call.charge(invocation, &result); + let charged_micro_usd = charge.cost_micro_usd; let mut response = if admissible { - encrypt_result(&result, &security)? + encrypt_result(&result, &security).inspect_err(|_| { + self.live.remove(&identity); + })? } else { self.seal_stand_in( invocation, &security, max_frame_bytes, FailureKind::Provider, + charge, ) .inspect_err(|_| { self.live.remove(&identity); @@ -455,29 +895,50 @@ impl TransportService { &security, max_frame_bytes, FailureKind::ResultSize, + charge, ) .inspect_err(|_| { self.live.remove(&identity); })?; } - let caller = self - .callers - .get(&identity.caller_runtime_id) - .ok_or(ServiceError::CallerNotAdmitted)?; let replay_record = ReplayRecord { caller_runtime_id: identity.caller_runtime_id.clone(), request_id: identity.request_id.clone(), - connection_key_epoch: caller.connection_key_epoch, + connection_key_epoch, invocation_sha256: claim.invocation_sha256, expires_at_unix_ms: invocation.expires_at_unix_ms, state: ReplayState::Completed { response: response.clone(), + charged_micro_usd, + overrun_micro_usd: charge.overrun_micro_usd, }, + admitted_at_unix_ms: active.admitted_at_unix_ms, }; + // A failed settle write frees the caller's slot and answers the typed + // store failure; the Active record stays, so the reservation stays + // counted and a retry is told the request is indeterminate. self.replay_store .put(replay_key(&identity), &replay_record) - .map_err(replay_store_error)?; + .map_err(replay_store_error) + .inspect_err(|_| { + self.live.remove(&identity); + })?; + // Settled: the window the request was admitted in holds the charge in + // place of the reservation. A failed put above leaves the reservation + // counted, as the store does. + add_spend( + &mut self.window_spend, + &identity.caller_runtime_id, + active.admitted_at_unix_ms / spend_window_ms, + reserved_micro_usd, + charged_micro_usd, + ); self.live.remove(&identity); + prune_to_latest( + &mut self.window_spend, + &self.live, + &identity.caller_runtime_id, + ); Ok(response) } @@ -488,6 +949,7 @@ impl TransportService { security: &TransportKey, max_frame_bytes: usize, kind: FailureKind, + charge: Charge, ) -> Result { let (_, foreign_reasoning_dropped) = partition_reasoning(&invocation.request, &self.provider) @@ -501,6 +963,7 @@ impl TransportService { invocation.provider_request_sha256, foreign_reasoning_dropped, kind, + charge, ), security, )?; @@ -512,8 +975,8 @@ impl TransportService { } /// Whether `result` is the answer to this claim's call: transported, - /// bound to the caller's invocation, and carrying the provider tag and drop - /// count the service derived. + /// bound to the caller's invocation, and carrying the provider tag, drop + /// count and cost the service derived. fn admits_result(&self, result: &TransportResult, claim: &ExecutionClaim) -> bool { let TransportDisposition::Transported { receipt, .. } = &result.disposition else { return false; @@ -521,6 +984,8 @@ impl TransportService { result.validate_against(&claim.invocation).is_ok() && receipt.provider == self.provider && receipt.foreign_reasoning_dropped == claim.call.receipt.foreign_reasoning_dropped + && receipt.cost_micro_usd == claim.call.cost_of(&receipt.outcome) + && receipt.overrun_micro_usd == claim.call.overrun_of(receipt.cost_micro_usd) } fn reply_refusal( @@ -554,7 +1019,12 @@ pub(crate) fn result_size_failure( native_request_sha256: [u8; 32], provider_request_sha256: [u8; 32], foreign_reasoning_dropped: u32, + cost_micro_usd: u64, ) -> TransportResult { + let charge = Charge { + cost_micro_usd, + overrun_micro_usd: 0, + }; stand_in_failure( caller_runtime_id, request_id, @@ -563,11 +1033,21 @@ pub(crate) fn result_size_failure( provider_request_sha256, foreign_reasoning_dropped, FailureKind::ResultSize, + charge, ) } +/// What a request cost, and by how much that exceeded its reservation. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct Charge { + cost_micro_usd: u64, + overrun_micro_usd: u64, +} + /// The stand-in failure for `kind`: `ResultSize` for a result that cannot be /// carried, `Provider` for one the service refused. +// One argument per fact the stand-in repeats from the invocation it answers. +#[allow(clippy::too_many_arguments)] fn stand_in_failure( caller_runtime_id: &str, request_id: &str, @@ -576,6 +1056,7 @@ fn stand_in_failure( provider_request_sha256: [u8; 32], foreign_reasoning_dropped: u32, kind: FailureKind, + charge: Charge, ) -> TransportResult { let message = match kind { FailureKind::ResultSize => "provider result exceeded its byte bound", @@ -596,9 +1077,12 @@ fn stand_in_failure( model: model.to_string(), provider: STAND_IN_PROVIDER.to_string(), foreign_reasoning_dropped, + cost_micro_usd: charge.cost_micro_usd, + overrun_micro_usd: charge.overrun_micro_usd, outcome: TransportOutcome::Failed { kind, message: message.to_string(), + usage: None, }, }), }, @@ -686,6 +1170,7 @@ mod tests { TransportService::open( &store.path, "fixture-sse", + fixture_prices(), [ caller("epiphany-yggdrasil", "shared-key", 1), caller("ghostlight-yggdrasil", "shared-key", 1), @@ -779,9 +1264,14 @@ mod tests { let mut changed_request = request(); changed_request.instructions.push_str(" Different cargo."); - let conflicting = - TransportInvocation::new("epiphany-yggdrasil", 2_000, [7; 32], changed_request) - .unwrap(); + let conflicting = TransportInvocation::new( + "epiphany-yggdrasil", + 2_000, + [7; 32], + MAX_COST, + changed_request, + ) + .unwrap(); let refusal = match service .begin( &encrypt_invocation(&conflicting, &security).unwrap(), @@ -887,6 +1377,7 @@ mod tests { TransportService::open( &store.path, "fixture-sse", + fixture_prices(), [CallerAdmission::new( "epiphany-yggdrasil", "rotated-connection-key", @@ -895,6 +1386,8 @@ mod tests { 1, 64 * 1024, 32_768, + SpendCeiling(CEILING), + SpendWindowMs(WINDOW_MS) ) .unwrap()], 5_000, @@ -913,7 +1406,8 @@ mod tests { ) -> (TransportAdmission, TransportInvocation) { let mut request = invocation_for(CALLER, request_id).request; edit(&mut request); - let invocation = TransportInvocation::new(CALLER, 2_000, [7; 32], request).unwrap(); + let invocation = + TransportInvocation::new(CALLER, 2_000, [7; 32], MAX_COST, request).unwrap(); let envelope = encrypt_invocation(&invocation, &security(KEY)).unwrap(); (service.begin(&envelope, 1_000, FRAME).unwrap(), invocation) } @@ -993,11 +1487,19 @@ mod tests { 4, max_payload_bytes, 32_768, + SpendCeiling(CEILING), + SpendWindowMs(WINDOW_MS), ) .unwrap(); - TransportService::open(&store.path, "fixture-sse", [admission], 5_000) - .unwrap() - .begin(&envelope, 1_000, FRAME) + TransportService::open( + &store.path, + "fixture-sse", + fixture_prices(), + [admission], + 5_000, + ) + .unwrap() + .begin(&envelope, 1_000, FRAME) }; assert!(matches!(admit(size), Ok(TransportAdmission::Execute(_)))); let refused = admit(size - 1).unwrap(); @@ -1041,7 +1543,10 @@ mod tests { connection_key_epoch: 1, invocation_sha256: [1; 32], expires_at_unix_ms: 5_000, - state: ReplayState::Active, + state: ReplayState::Active { + reserved_micro_usd: 1, + }, + admitted_at_unix_ms: 1_000, }; edit(&mut record); let mut cache = CultCache::new(); @@ -1068,6 +1573,7 @@ mod tests { TransportService::open( &control.path, "fixture-sse", + fixture_prices(), [caller(CALLER, KEY, 1)], 5_000 ) @@ -1075,10 +1581,11 @@ mod tests { ); type Edit = fn(&mut ReplayRecord); - let cases: [(&str, Edit); 3] = [ + let cases: [(&str, Edit); 4] = [ ("wrong-key", |_| {}), ("zero-digest", |record| record.invocation_sha256 = [0; 32]), ("zero-expiry", |record| record.expires_at_unix_ms = 0), + ("zero-admission", |record| record.admitted_at_unix_ms = 0), ]; for (name, edit) in cases { let store = TestStore::new(); @@ -1093,6 +1600,7 @@ mod tests { TransportService::open( &store.path, "fixture-sse", + fixture_prices(), [caller(CALLER, KEY, 1)], 5_000 ), @@ -1133,6 +1641,8 @@ mod tests { .unwrap(); record.state = ReplayState::Completed { response: encrypt_result(&result, &security(KEY)).unwrap(), + charged_micro_usd: 0, + overrun_micro_usd: 0, }; service.replay_store.put(key, &record).unwrap(); assert_eq!( @@ -1144,7 +1654,17 @@ mod tests { #[test] fn identities_are_bounded_at_admission_and_on_the_request() { let admit = |caller_id: String, model: String| { - CallerAdmission::new(caller_id, KEY, 1, [model], 4, 64 * 1024, 32_768) + CallerAdmission::new( + caller_id, + KEY, + 1, + [model], + 4, + 64 * 1024, + 32_768, + SpendCeiling(CEILING), + SpendWindowMs(WINDOW_MS), + ) }; let at = "i".repeat(MAX_IDENTITY_BYTES); let over = "i".repeat(MAX_IDENTITY_BYTES + 1); @@ -1282,6 +1802,7 @@ mod tests { invocation.native_request_sha256, invocation.provider_request_sha256, 0, + FIXTURE_USAGE_COST, ) ); let TransportDisposition::Transported { receipt, .. } = &stand_in.disposition else { @@ -1303,7 +1824,7 @@ mod tests { fn a_result_whose_sealed_envelope_fits_is_recorded_as_the_backend_made_it() { let store = TestStore::new(); let mut service = open_service(&store, 8); - let (claim, invocation) = claim_for(&mut service, "fits"); + let (claim, invocation) = claim_for(&mut service, "request-4"); let small = text_result(claim.call(), "done"); let response = service.complete(*claim, small.clone(), 1 << 20).unwrap(); assert_eq!( @@ -1376,6 +1897,7 @@ mod tests { invocation.native_request_sha256, invocation.provider_request_sha256, 0, + FIXTURE_USAGE_COST, ) ); // The stand-in is what is recorded now, byte for byte, at any frame. @@ -1611,15 +2133,27 @@ mod tests { for tag in ["", " ", STAND_IN_PROVIDER] { assert!( matches!( - TransportService::open(&store.path, tag, [caller(CALLER, KEY, 1)], 5_000), + TransportService::open( + &store.path, + tag, + fixture_prices(), + [caller(CALLER, KEY, 1)], + 5_000 + ), Err(ServiceError::InvalidAdmission) ), "{tag:?}" ); } assert!( - TransportService::open(&store.path, "fixture-sse", [caller(CALLER, KEY, 1)], 5_000) - .is_ok() + TransportService::open( + &store.path, + "fixture-sse", + fixture_prices(), + [caller(CALLER, KEY, 1)], + 5_000 + ) + .is_ok() ); } @@ -1738,6 +2272,7 @@ mod tests { invocation.native_request_sha256, invocation.provider_request_sha256, 1, + FIXTURE_USAGE_COST, ); assert_eq!( decrypt_result(&response, &security(KEY), &invocation).unwrap(), @@ -1767,4 +2302,1318 @@ mod tests { assert_eq!(receipt.provider, STAND_IN_PROVIDER); assert_eq!(receipt.foreign_reasoning_dropped, 1); } + + // ---- spend ---- + + fn usage(input: u64, cache_read: u64, cache_write: u64, output: u64) -> Usage { + Usage { + input_tokens: input, + cache_read_tokens: cache_read, + cache_write_tokens: cache_write, + output_tokens: output, + reasoning_tokens: None, + } + } + + const MTOK: u64 = 1_000_000; + + fn flat_price(rate: u64) -> ModelPrice { + ModelPrice { + rates: Rates { + input: rate, + cache_write_short: rate, + cache_write_long: rate, + cache_read: rate, + output: rate, + }, + long_context: None, + overhead_tokens: 0, + } + } + + fn dearest_input(rates: &Rates, retention: Option) -> u64 { + let write = match retention { + Some(CacheRetention::Short) => rates.cache_write_short, + Some(CacheRetention::Long) | None => rates.cache_write_long, + }; + rates.input.max(rates.cache_read).max(write) + } + + /// ceil(((encoded length + overhead) x dearest input rate + max_output_tokens + /// x output rate) / 1e6), spelled out apart from the service. + fn expected_reservation(request: &ProviderRequest, rates: &Rates, overhead: u64) -> u64 { + let bound = wire_encode(request).unwrap().len() as u128 + u128::from(overhead); + let total = bound * u128::from(dearest_input(rates, request.cache_retention)) + + u128::from(request.max_output_tokens) * u128::from(rates.output); + total.div_ceil(1_000_000) as u64 + } + + /// What every fixture request (request ids of one length) reserves. + fn reservation() -> u64 { + expected_reservation(&request(), &fixture_price().rates, 100) + } + + fn invocation_at(request_id: &str, now: u64, max_cost: u64) -> TransportInvocation { + let mut request = request(); + request.request_id = request_id.to_string(); + TransportInvocation::new(CALLER, now + 1_000, [7; 32], max_cost, request).unwrap() + } + + fn begin_at( + service: &mut TransportService, + request_id: &str, + now: u64, + max_cost: u64, + ) -> (TransportAdmission, TransportInvocation) { + let invocation = invocation_at(request_id, now, max_cost); + let admission = service + .begin(&raw_envelope(&invocation), now, FRAME) + .unwrap(); + (admission, invocation) + } + + fn claim_at(service: &mut TransportService, request_id: &str, now: u64) -> Box { + match begin_at(service, request_id, now, MAX_COST).0 { + TransportAdmission::Execute(claim) => claim, + TransportAdmission::Reply(_) => panic!("{request_id} should have executed"), + } + } + + fn refusal_at( + service: &mut TransportService, + request_id: &str, + now: u64, + max_cost: u64, + ) -> Refusal { + let (admission, invocation) = begin_at(service, request_id, now, max_cost); + refusal_of(admission, &invocation) + } + + fn open_limited(store: &TestStore, ceiling: u64, window_ms: u64) -> TransportService { + TransportService::open( + &store.path, + "fixture-sse", + fixture_prices(), + [CallerAdmission::new( + CALLER, + KEY, + 1, + ["gpt-5.4".to_string()], + 8, + 65_536, + 32_768, + SpendCeiling(ceiling), + SpendWindowMs(window_ms), + ) + .unwrap()], + 5_000, + ) + .unwrap() + } + + fn record_of(service: &TransportService, request_id: &str) -> Option { + service + .replay_store + .get::(&replay_key(&RequestIdentity { + caller_runtime_id: CALLER.to_string(), + request_id: request_id.to_string(), + })) + .unwrap() + } + + fn charged_of(service: &TransportService, request_id: &str) -> u64 { + match record_of(service, request_id).unwrap().state { + ReplayState::Completed { + charged_micro_usd, .. + } => charged_micro_usd, + ReplayState::Active { .. } => panic!("{request_id} is not completed"), + } + } + + fn receipt_of( + response: &TransportEnvelope, + invocation: &TransportInvocation, + ) -> TransportReceipt { + let TransportDisposition::Transported { receipt, .. } = + decrypt_result(response, &security(KEY), invocation) + .unwrap() + .disposition + else { + panic!("expected a transported result"); + }; + *receipt + } + + #[test] + fn cost_prices_each_usage_class_at_its_own_rate() { + let price = fixture_price(); + let cost = |usage: Usage| cost_micro_usd(&price, &usage, Some(CacheRetention::Long)); + assert_eq!(cost(usage(MTOK, 0, 0, 0)), 3_000_000); + assert_eq!(cost(usage(0, MTOK, 0, 0)), 300_000); + assert_eq!(cost(usage(0, 0, MTOK, 0)), 6_000_000); + assert_eq!(cost(usage(0, 0, 0, MTOK)), 15_000_000); + assert_eq!( + cost(usage(2 * MTOK, 3 * MTOK, 5 * MTOK, 7 * MTOK)), + 6_000_000 + 900_000 + 30_000_000 + 105_000_000 + ); + // Reasoning tokens are a detail of the output tokens, never priced apart. + let mut reasoning_only = usage(0, 0, 0, 0); + reasoning_only.reasoning_tokens = Some(MTOK); + assert_eq!(cost(reasoning_only), 0); + let mut detailed = usage(0, 0, 0, MTOK); + detailed.reasoning_tokens = Some(MTOK); + assert_eq!(cost(detailed), 15_000_000); + } + + #[test] + fn cost_rounds_up_once_per_request() { + let price = flat_price(10_000); + let cost = |usage: Usage| cost_micro_usd(&price, &usage, None); + assert_eq!(cost(usage(0, 1, 0, 0)), 1); + assert_eq!(cost(usage(0, 0, 0, 0)), 0); + // Four classes at one token each: 40,000 / 1e6 rounds up once, to 1. + assert_eq!(cost(usage(1, 1, 1, 1)), 1); + assert_eq!(cost(usage(100, 0, 0, 0)), 1); + assert_eq!(cost(usage(101, 0, 0, 0)), 2); + } + + #[test] + fn cache_writes_price_by_retention() { + let price = fixture_price(); + let written = usage(0, 0, MTOK, 0); + assert_eq!( + cost_micro_usd(&price, &written, Some(CacheRetention::Short)), + 3_750_000 + ); + assert_eq!( + cost_micro_usd(&price, &written, Some(CacheRetention::Long)), + 6_000_000 + ); + // A write nobody asked to keep short is priced dear. + assert_eq!(cost_micro_usd(&price, &written, None), 6_000_000); + } + + #[test] + fn long_context_rates_apply_strictly_above_the_threshold() { + let long = Rates { + input: 9_000_000, + cache_write_short: 9_500_000, + cache_write_long: 12_000_000, + cache_read: 900_000, + output: 45_000_000, + }; + let price = |above_prompt_tokens: u64, overhead_tokens: u64| ModelPrice { + long_context: Some(LongContext { + above_prompt_tokens, + rates: long, + }), + overhead_tokens, + ..fixture_price() + }; + // Prompt tokens are input + cache reads + cache writes. + let at_threshold = cost_micro_usd(&price(100, 0), &usage(1, 1, 98, 0), None); + assert_eq!(at_threshold, 592); // 3e6 + 3e5 + 98 x 6e6 = 591.3e6, base rates + let above = cost_micro_usd(&price(100, 0), &usage(1, 1, 99, 0), None); + assert_eq!(above, 1198); // 9e6 + 9e5 + 99 x 12e6 = 1197.9e6, long rates + assert_eq!( + cost_micro_usd(&price(100, 0), &usage(100, 0, 0, 0), None), + 300 + ); + assert_eq!( + cost_micro_usd(&price(100, 0), &usage(101, 0, 0, 0), None), + 909 + ); + + // The reservation switches at the same strict threshold, on its byte bound. + let request = request(); + let length = wire_encode(&request).unwrap().len() as u64; + let bound = length + 7; + assert_eq!( + reservation_micro_usd(&price(bound, 7), &request) + .unwrap() + .unwrap(), + expected_reservation(&request, &fixture_price().rates, 7) + ); + assert_eq!( + reservation_micro_usd(&price(bound - 1, 7), &request) + .unwrap() + .unwrap(), + expected_reservation(&request, &long, 7) + ); + } + + #[test] + fn cost_saturates_instead_of_wrapping() { + let price = fixture_price(); + for huge in [ + usage(u64::MAX, 0, 0, 0), + usage(0, 0, u64::MAX, 0), + usage(0, 0, 0, u64::MAX), + usage(u64::MAX, u64::MAX, u64::MAX, u64::MAX), + ] { + assert_eq!(cost_micro_usd(&price, &huge, None), u64::MAX, "{huge:?}"); + } + // A class that does not overflow is exact: 300,000 x u64::MAX / 1e6, rounded up. + assert_eq!( + cost_micro_usd(&price, &usage(0, u64::MAX, 0, 0), None), + 5_534_023_222_112_865_485 + ); + } + + #[test] + fn reservation_is_the_byte_bound_plus_max_output() { + let rate_sets = [ + (fixture_price().rates, 100), + // The dearest input rate is a cache read here... + ( + Rates { + input: 1_000_000, + cache_write_short: 2_000_000, + cache_write_long: 3_000_000, + cache_read: 8_000_000, + output: 5_000_000, + }, + 7, + ), + // ...and plain input here. + ( + Rates { + input: 9_000_000, + cache_write_short: 1_000_000, + cache_write_long: 2_000_000, + cache_read: 3_000_000, + output: 4_000_000, + }, + 0, + ), + ]; + for (rates, overhead_tokens) in rate_sets { + let price = ModelPrice { + rates, + long_context: None, + overhead_tokens, + }; + for retention in [ + None, + Some(CacheRetention::Short), + Some(CacheRetention::Long), + ] { + let mut request = request(); + request.cache_retention = retention; + assert_eq!( + reservation_micro_usd(&price, &request).unwrap().unwrap(), + expected_reservation(&request, &rates, overhead_tokens), + "{rates:?} {retention:?}" + ); + } + } + // More output costs at least its price, and the bound follows the encoding + // (a wider integer is more bytes). + let price = fixture_price(); + let mut reserved = Vec::new(); + for max_output_tokens in [1, 1_000, 1_001_000] { + let mut request = request(); + request.max_output_tokens = max_output_tokens; + let amount = reservation_micro_usd(&price, &request).unwrap().unwrap(); + assert_eq!( + amount, + expected_reservation(&request, &price.rates, 100), + "{max_output_tokens}" + ); + reserved.push(amount); + } + assert!(reserved[2] - reserved[1] >= 15_000_000); + } + + #[test] + fn worst_case_over_max_cost_is_refused_cost_bound_inclusively() { + let store = TestStore::new(); + let mut service = open_service(&store, 1); + let reserved = reservation(); + let (admission, _) = begin_at(&mut service, "request-1", 1_000, reserved); + assert!(matches!(admission, TransportAdmission::Execute(_))); + // The only slot is taken: a request refused for its cost is refused + // for that, before the replay lookup and the capacity check. + assert_eq!( + refusal_at(&mut service, "request-2", 1_000, reserved - 1), + Refusal::CostBound + ); + assert!(record_of(&service, "request-2").is_none()); + assert_eq!(service.window_spend(CALLER, 0), reserved); + assert_eq!( + refusal_at(&mut service, "request-3", 1_000, reserved), + Refusal::Capacity + ); + } + + #[test] + fn max_cost_zero_is_malformed() { + let store = TestStore::new(); + let mut service = open_service(&store, 4); + assert_eq!( + refusal_at(&mut service, "request-1", 1_000, 0), + Refusal::Malformed + ); + assert!(record_of(&service, "request-1").is_none()); + } + + #[test] + fn a_changed_max_cost_is_a_replay_conflict() { + let store = TestStore::new(); + let mut service = open_service(&store, 4); + let claim = claim_at(&mut service, "request-1", 1_000); + let result = text_result(claim.call(), "done"); + service.complete(*claim, result, FRAME).unwrap(); + assert_eq!( + refusal_at(&mut service, "request-1", 1_000, MAX_COST + 1), + Refusal::ReplayConflict + ); + } + + #[test] + fn spend_ceiling_refuses_inclusively_and_resets_next_window() { + let reserved = reservation(); + let window = 10_000; + let store = TestStore::new(); + let mut service = open_limited(&store, 2 * reserved, window); + claim_at(&mut service, "request-1", 1_000); + // Windows are aligned to the epoch: the last millisecond of window 0. + claim_at(&mut service, "request-2", 9_999); + assert_eq!(service.window_spend(CALLER, 0), 2 * reserved); + assert_eq!( + refusal_at(&mut service, "request-3", 9_999, MAX_COST), + Refusal::SpendCeiling + ); + assert!(record_of(&service, "request-3").is_none()); + // The first millisecond of window 1: the same request id is admitted. + claim_at(&mut service, "request-3", 10_000); + assert_eq!(service.window_spend(CALLER, 0), 2 * reserved); + assert_eq!(service.window_spend(CALLER, 1), reserved); + + // One micro-USD below the two reservations refuses the second. + let store = TestStore::new(); + let mut service = open_limited(&store, 2 * reserved - 1, window); + claim_at(&mut service, "request-1", 1_000); + assert_eq!( + refusal_at(&mut service, "request-2", 1_000, MAX_COST), + Refusal::SpendCeiling + ); + } + + #[test] + fn a_recorded_answer_replays_past_an_exhausted_ceiling() { + let reserved = reservation(); + let store = TestStore::new(); + let mut service = open_limited(&store, reserved, 3_600_000); + let (admission, _) = begin_at(&mut service, "request-1", 1_000, MAX_COST); + let TransportAdmission::Execute(claim) = admission else { + panic!("the first request fits the ceiling exactly"); + }; + let result = text_result(claim.call(), "done"); + let response = service.complete(*claim, result, FRAME).unwrap(); + // Spent 75 of a ceiling that holds one reservation: nothing else fits. + assert_eq!( + refusal_at(&mut service, "request-2", 1_000, MAX_COST), + Refusal::SpendCeiling + ); + // The recorded answer is replayed whatever the caller has left. + let (replay, _) = begin_at(&mut service, "request-1", 1_000, MAX_COST); + assert_eq!(replay, TransportAdmission::Reply(response)); + } + + #[test] + fn settlement_replaces_the_reservation_with_the_charge() { + let store = TestStore::new(); + let mut service = open_service(&store, 4); + let claim = claim_at(&mut service, "request-1", 1_000); + assert_eq!(service.window_spend(CALLER, 0), reservation()); + let result = text_result(claim.call(), "done"); + service.complete(*claim, result, FRAME).unwrap(); + assert_eq!(service.window_spend(CALLER, 0), FIXTURE_USAGE_COST); + assert_eq!(charged_of(&service, "request-1"), FIXTURE_USAGE_COST); + } + + #[test] + fn failed_with_usage_charges_usage_and_failed_without_keeps_the_reservation() { + let failed = |usage: Option| TransportOutcome::Failed { + kind: FailureKind::Provider, + message: "the provider refused".to_string(), + usage, + }; + let store = TestStore::new(); + let mut service = open_service(&store, 4); + + let claim = claim_at(&mut service, "request-1", 1_000); + let result = claim + .call() + .transported(Vec::new(), failed(Some(usage(0, 0, 0, 0)))); + let invocation = invocation_at("request-1", 1_000, MAX_COST); + let response = service.complete(*claim, result, FRAME).unwrap(); + assert_eq!(receipt_of(&response, &invocation).cost_micro_usd, 0); + assert_eq!(charged_of(&service, "request-1"), 0); + assert_eq!(service.window_spend(CALLER, 0), 0); + + let claim = claim_at(&mut service, "request-2", 1_000); + let result = claim + .call() + .transported(Vec::new(), failed(Some(usage(MTOK, 0, 0, 0)))); + service.complete(*claim, result, FRAME).unwrap(); + assert_eq!(charged_of(&service, "request-2"), 3_000_000); + + let claim = claim_at(&mut service, "request-3", 1_000); + let result = claim.call().transported(Vec::new(), failed(None)); + let invocation = invocation_at("request-3", 1_000, MAX_COST); + let response = service.complete(*claim, result, FRAME).unwrap(); + assert_eq!( + receipt_of(&response, &invocation).cost_micro_usd, + reservation() + ); + assert_eq!(charged_of(&service, "request-3"), reservation()); + assert_eq!(service.window_spend(CALLER, 0), 3_000_000 + reservation()); + } + + #[test] + fn a_backend_cost_misreport_is_answered_with_the_stand_in() { + let edits: [fn(u64) -> u64; 3] = [|cost| cost + 1, |cost| cost - 1, |_| 0]; + for (index, edit) in edits.into_iter().enumerate() { + let store = TestStore::new(); + let mut service = open_service(&store, 4); + let id = format!("request-{index}"); + let claim = claim_at(&mut service, &id, 1_000); + let mut result = text_result(claim.call(), "done"); + let TransportDisposition::Transported { receipt, .. } = &mut result.disposition else { + unreachable!(); + }; + receipt.cost_micro_usd = edit(receipt.cost_micro_usd); + let invocation = invocation_at(&id, 1_000, MAX_COST); + let response = service.complete(*claim, result, FRAME).unwrap(); + let receipt = receipt_of(&response, &invocation); + assert_eq!(receipt.provider, STAND_IN_PROVIDER, "{index}"); + // A result the service refused is still billed what its usage costs. + assert_eq!(receipt.cost_micro_usd, FIXTURE_USAGE_COST, "{index}"); + assert_eq!(charged_of(&service, &id), FIXTURE_USAGE_COST, "{index}"); + assert_eq!( + service.window_spend(CALLER, 0), + FIXTURE_USAGE_COST, + "{index}" + ); + } + } + + #[test] + fn result_size_stand_in_keeps_the_real_charge() { + let store = TestStore::new(); + let mut service = open_service(&store, 4); + let claim = claim_at(&mut service, "request-1", 1_000); + let big = text_result(claim.call(), &"y".repeat(4000)); + let invocation = invocation_at("request-1", 1_000, MAX_COST); + let response = service.complete(*claim, big, 2000).unwrap(); + let receipt = receipt_of(&response, &invocation); + assert_eq!(receipt.provider, STAND_IN_PROVIDER); + assert_eq!(receipt.cost_micro_usd, FIXTURE_USAGE_COST); + assert_eq!(charged_of(&service, "request-1"), FIXTURE_USAGE_COST); + assert_eq!(service.window_spend(CALLER, 0), FIXTURE_USAGE_COST); + + // The replay path: a wide record re-measured at a smaller frame. + let claim = claim_at(&mut service, "request-2", 1_000); + let big = text_result(claim.call(), &"y".repeat(4000)); + service.complete(*claim, big, FRAME).unwrap(); + drop(service); + let mut service = open_service(&store, 4); + let invocation = invocation_at("request-2", 1_000, MAX_COST); + let TransportAdmission::Reply(smaller) = service + .begin(&raw_envelope(&invocation), 1_000, 2000) + .unwrap() + else { + panic!("expected a replay"); + }; + let receipt = receipt_of(&smaller, &invocation); + assert_eq!(receipt.provider, STAND_IN_PROVIDER); + assert_eq!(receipt.cost_micro_usd, FIXTURE_USAGE_COST); + assert_eq!(charged_of(&service, "request-2"), FIXTURE_USAGE_COST); + assert_eq!(service.window_spend(CALLER, 0), 2 * FIXTURE_USAGE_COST); + } + + #[test] + fn a_ledger_write_failure_executes_nothing() { + let refusing = FlakyStore::new(0); + let healthy = refusing.reopened(5); + let mut service = open_flaky(refusing); + let invocation = invocation_for(CALLER, "request-1"); + assert!(matches!( + service.begin(&raw_envelope(&invocation), 1_000, FRAME), + Err(ServiceError::ReplayStore(_)) + )); + assert_eq!(service.window_spend(CALLER, 0), 0); + assert!(service.live.is_empty()); + assert!(record_of(&service, "request-1").is_none()); + drop(service); + + // Nothing was written, so the same ledger admits the request afresh. + let mut service = open_flaky(healthy); + claim_at(&mut service, "request-1", 1_000); + } + + #[test] + fn a_settle_write_failure_keeps_the_reservation_counted() { + let one_write = FlakyStore::new(1); + let restart = one_write.reopened(5); + let mut service = open_flaky(one_write); + let claim = claim_at(&mut service, "request-1", 1_000); + let result = text_result(claim.call(), "done"); + assert!(matches!( + service.complete(*claim, result, FRAME), + Err(ServiceError::ReplayStore(_)) + )); + assert_eq!(service.window_spend(CALLER, 0), reservation()); + assert!(matches!( + record_of(&service, "request-1").unwrap().state, + ReplayState::Active { .. } + )); + drop(service); + + // A restart bills the indeterminate request at its whole reservation. + let service = open_flaky(restart); + assert_eq!(service.window_spend(CALLER, 0), reservation()); + } + + #[test] + fn restart_rederives_window_spend_from_the_records() { + let reserved = reservation(); + let window = 10_000; + let store = TestStore::new(); + let mut service = open_limited(&store, 1 << 40, window); + // An earlier window: completed. + let claim = claim_at(&mut service, "request-1", 1_000); + let result = text_result(claim.call(), "done"); + service.complete(*claim, result, FRAME).unwrap(); + // The current window: one completed, one left Active by a crash. + let claim = claim_at(&mut service, "request-2", 10_500); + let result = text_result(claim.call(), "done"); + service.complete(*claim, result, FRAME).unwrap(); + claim_at(&mut service, "request-3", 10_600); + drop(service); + + let mut service = open_limited(&store, FIXTURE_USAGE_COST + 2 * reserved, window); + // Only the current window is kept: the earlier one is pruned. + assert_eq!(service.window_spend(CALLER, 0), 0); + assert_eq!( + service.window_spend(CALLER, 1), + FIXTURE_USAGE_COST + reserved + ); + // The ceiling sees the rebuilt window: one more reservation fits it + // exactly, and the next does not. The earlier window counts for nothing. + claim_at(&mut service, "request-4", 11_000); + assert_eq!( + refusal_at(&mut service, "request-5", 11_100, MAX_COST), + Refusal::SpendCeiling + ); + } + + #[test] + fn open_refuses_an_unpriced_model_and_admits_a_zero_price() { + let store = TestStore::new(); + assert_eq!( + TransportService::open( + &store.path, + "fixture-sse", + PriceTable::from([("another-model".to_string(), fixture_price())]), + [caller(CALLER, KEY, 1)], + 5_000, + ) + .err(), + Some(ServiceError::UnpricedModel("gpt-5.4".to_string())) + ); + + let free = PriceTable::from([("gpt-5.4".to_string(), flat_price(0))]); + let mut service = TransportService::open( + &store.path, + "fixture-sse", + free, + [caller(CALLER, KEY, 1)], + 5_000, + ) + .unwrap(); + let claim = claim_at(&mut service, "request-1", 1_000); + assert_eq!(service.window_spend(CALLER, 0), 0); + let result = text_result(claim.call(), "done"); + service.complete(*claim, result, FRAME).unwrap(); + assert_eq!(charged_of(&service, "request-1"), 0); + assert_eq!(service.window_spend(CALLER, 0), 0); + } + + #[test] + fn a_ceiling_and_a_window_are_required_nonzero() { + let admit = |ceiling: u64, window_ms: u64| { + CallerAdmission::new( + CALLER, + KEY, + 1, + ["gpt-5.4".to_string()], + 4, + 64 * 1024, + 32_768, + SpendCeiling(ceiling), + SpendWindowMs(window_ms), + ) + }; + assert!(admit(1, 1).is_ok()); + assert!(matches!(admit(0, 1), Err(ServiceError::InvalidAdmission))); + assert!(matches!(admit(1, 0), Err(ServiceError::InvalidAdmission))); + } + + // ---- settle past the reservation, pruning, extremes ---- + + fn overrun_result(call: &ProviderCall, text: &str) -> TransportResult { + let mut outcome = receipt(&invocation_for(CALLER, "outcome")).outcome; + let TransportOutcome::Completed { usage, .. } = &mut outcome else { + unreachable!(); + }; + // 2 million output tokens at 15,000,000 per million: 30,000,000. + *usage = self::usage(0, 0, 0, 2 * MTOK); + call.transported( + vec![TransportEvent { + sequence: 0, + payload: TransportEventPayload::TextDelta { + text: text.to_string(), + }, + }], + outcome, + ) + } + + const OVERRUN_COST: u64 = 30_000_000; + + #[test] + fn an_overrun_settles_at_the_true_cost_and_the_receipt_says_so() { + let reserved = reservation(); + let store = TestStore::new(); + let mut service = open_limited(&store, 2 * reserved, 3_600_000); + let claim = claim_at(&mut service, "request-1", 1_000); + let result = overrun_result(claim.call(), "done"); + let invocation = invocation_at("request-1", 1_000, MAX_COST); + let response = service.complete(*claim, result, FRAME).unwrap(); + let receipt = receipt_of(&response, &invocation); + // The provider's own answer, not a stand-in, carrying the true cost + // and the marker. + assert_eq!(receipt.provider, "fixture-sse"); + assert_eq!(receipt.cost_micro_usd, OVERRUN_COST); + assert_eq!(receipt.overrun_micro_usd, OVERRUN_COST - reserved); + assert_eq!(charged_of(&service, "request-1"), OVERRUN_COST); + assert_eq!(service.window_spend(CALLER, 0), OVERRUN_COST); + // The ceiling holds for every later request in the window. + assert_eq!( + refusal_at(&mut service, "request-2", 1_000, MAX_COST), + Refusal::SpendCeiling + ); + // A result inside its reservation carries no marker. + drop(service); + let mut service = open_service(&store, 4); + let claim = claim_at(&mut service, "request-3", 1_000); + let result = text_result(claim.call(), "done"); + let invocation = invocation_at("request-3", 1_000, MAX_COST); + let response = service.complete(*claim, result, FRAME).unwrap(); + assert_eq!(receipt_of(&response, &invocation).overrun_micro_usd, 0); + } + + #[test] + fn an_overrun_survives_the_result_size_stand_in_and_its_replay() { + let reserved = reservation(); + let store = TestStore::new(); + let mut service = open_service(&store, 4); + let claim = claim_at(&mut service, "request-1", 1_000); + let big = overrun_result(claim.call(), &"y".repeat(4000)); + let invocation = invocation_at("request-1", 1_000, MAX_COST); + let response = service.complete(*claim, big, 2000).unwrap(); + let receipt = receipt_of(&response, &invocation); + assert_eq!(receipt.provider, STAND_IN_PROVIDER); + assert_eq!(receipt.cost_micro_usd, OVERRUN_COST); + assert_eq!(receipt.overrun_micro_usd, OVERRUN_COST - reserved); + + let claim = claim_at(&mut service, "request-2", 1_000); + let big = overrun_result(claim.call(), &"y".repeat(4000)); + service.complete(*claim, big, FRAME).unwrap(); + drop(service); + let mut service = open_service(&store, 4); + let invocation = invocation_at("request-2", 1_000, MAX_COST); + let TransportAdmission::Reply(smaller) = service + .begin(&raw_envelope(&invocation), 1_000, 2000) + .unwrap() + else { + panic!("expected a replay"); + }; + let receipt = receipt_of(&smaller, &invocation); + assert_eq!(receipt.provider, STAND_IN_PROVIDER); + assert_eq!(receipt.cost_micro_usd, OVERRUN_COST); + assert_eq!(receipt.overrun_micro_usd, OVERRUN_COST - reserved); + } + + #[test] + fn a_backend_overrun_misreport_is_answered_with_the_stand_in() { + let reserved = reservation(); + for (index, overrun) in [1, 0].into_iter().enumerate() { + let store = TestStore::new(); + let mut service = open_service(&store, 4); + let id = format!("request-{index}"); + let claim = claim_at(&mut service, &id, 1_000); + // The first result is inside its reservation and claims an overrun; + // the second overran and claims none. + let mut result = if index == 0 { + text_result(claim.call(), "done") + } else { + overrun_result(claim.call(), "done") + }; + let TransportDisposition::Transported { receipt, .. } = &mut result.disposition else { + unreachable!(); + }; + receipt.overrun_micro_usd = overrun; + let invocation = invocation_at(&id, 1_000, MAX_COST); + let response = service.complete(*claim, result, FRAME).unwrap(); + let receipt = receipt_of(&response, &invocation); + // The service prices the usage itself: the stand-in and the ledger + // carry the true cost and the true overrun, whatever was claimed. + let (cost, true_overrun) = if index == 0 { + (FIXTURE_USAGE_COST, 0) + } else { + (OVERRUN_COST, OVERRUN_COST - reserved) + }; + assert_eq!(receipt.provider, STAND_IN_PROVIDER, "{index}"); + assert_eq!(receipt.cost_micro_usd, cost, "{index}"); + assert_eq!(receipt.overrun_micro_usd, true_overrun, "{index}"); + assert_eq!(charged_of(&service, &id), cost, "{index}"); + assert_eq!(service.window_spend(CALLER, 0), cost, "{index}"); + } + } + + #[test] + fn the_window_map_stays_bounded_live_and_after_reopen() { + let window = 10; + let store = TestStore::new(); + let mut service = open_limited(&store, 1 << 40, window); + let held = + |service: &TransportService| service.window_spend.get(CALLER).map_or(0, |w| w.len()); + for index in 0..300u64 { + let id = format!("request-{index}"); + let claim = claim_at(&mut service, &id, 1_000 + index * window); + let result = text_result(claim.call(), "done"); + service.complete(*claim, result, FRAME).unwrap(); + assert_eq!(held(&service), 1, "{index}"); + } + // A request still in flight keeps its own window until it settles. + let early = claim_at(&mut service, "request-e", 1_000 + 300 * window); + claim_at(&mut service, "request-l", 1_000 + 305 * window); + assert_eq!(held(&service), 2); + let result = text_result(early.call(), "done"); + service.complete(*early, result, FRAME).unwrap(); + assert_eq!(held(&service), 1); + drop(service); + let service = open_limited(&store, 1 << 40, window); + assert_eq!(held(&service), 1); + assert_eq!( + service.window_spend(CALLER, 1_000 / window + 305), + reservation() + ); + } + + #[test] + fn a_reservation_that_overflows_is_refused_cost_bound_not_wrapped() { + for price in [ + ModelPrice { + overhead_tokens: u64::MAX, + ..flat_price(u64::MAX) + }, + ModelPrice { + overhead_tokens: u64::MAX, + ..flat_price(1 << 40) + }, + ] { + assert_eq!(reservation_micro_usd(&price, &request()).unwrap(), None); + let store = TestStore::new(); + let mut service = TransportService::open( + &store.path, + "fixture-sse", + PriceTable::from([("gpt-5.4".to_string(), price)]), + [caller(CALLER, KEY, 4)], + 5_000, + ) + .unwrap(); + // Not even a caller that bounds its request at u64::MAX can afford it. + assert_eq!( + refusal_at(&mut service, "request-1", 1_000, u64::MAX), + Refusal::CostBound + ); + assert!(record_of(&service, "request-1").is_none()); + } + } + + #[test] + fn a_recorded_answer_replays_after_a_price_raise() { + let store = TestStore::new(); + let mut service = open_service(&store, 4); + let (admission, _) = begin_at(&mut service, "request-1", 1_000, reservation()); + let TransportAdmission::Execute(claim) = admission else { + panic!("expected an execution claim"); + }; + let result = text_result(claim.call(), "done"); + service.complete(*claim, result, FRAME).unwrap(); + drop(service); + + let mut dear = fixture_price(); + dear.rates.output *= 2; + let mut service = TransportService::open( + &store.path, + "fixture-sse", + PriceTable::from([("gpt-5.4".to_string(), dear)]), + [caller(CALLER, KEY, 4)], + 5_000, + ) + .unwrap(); + let (admission, invocation) = begin_at(&mut service, "request-1", 1_000, reservation()); + let TransportAdmission::Reply(response) = admission else { + panic!("expected the recorded answer"); + }; + assert_eq!( + receipt_of(&response, &invocation).cost_micro_usd, + FIXTURE_USAGE_COST + ); + // New work is still checked against today's prices. + assert_eq!( + refusal_at(&mut service, "request-2", 1_000, reservation()), + Refusal::CostBound + ); + } + + #[test] + fn a_settle_write_failure_releases_the_slot_and_keeps_the_reservation() { + let store = FlakyStore::new(1); + let allowance = store.allowance(); + let mut service = open_flaky(store); + let claim = claim_at(&mut service, "request-1", 1_000); + let result = text_result(claim.call(), "done"); + assert!(matches!( + service.complete(*claim, result, FRAME), + Err(ServiceError::ReplayStore(_)) + )); + assert!(service.live.is_empty()); + assert_eq!(service.window_spend(CALLER, 0), reservation()); + // The store recovers: the caller's only slot is free again, and the + // unsettled request is told it is indeterminate, not in flight. + allowance.store(5, std::sync::atomic::Ordering::SeqCst); + claim_at(&mut service, "request-2", 1_000); + assert_eq!(service.window_spend(CALLER, 0), 2 * reservation()); + let (admission, invocation) = begin_at(&mut service, "request-1", 1_000, MAX_COST); + assert_eq!(refusal_of(admission, &invocation), Refusal::Indeterminate); + } + + fn open_flaky(store: FlakyStore) -> TransportService { + TransportService::with_store( + store, + "fixture-sse", + fixture_prices(), + [caller(CALLER, KEY, 1)], + 5_000, + ) + .unwrap() + } + + fn run_window(service: &mut TransportService, prefix: &str, now: u64) -> u32 { + let mut admitted = 0; + for i in 0..20u64 { + let (admission, _) = begin_at(service, &format!("{prefix}-{i:02}"), now, MAX_COST); + let TransportAdmission::Execute(claim) = admission else { + break; + }; + admitted += 1; + let result = text_result(claim.call(), "done"); + service.complete(*claim, result, FRAME).unwrap(); + } + admitted + } + + #[test] + fn a_clock_stepping_back_cannot_reopen_the_callers_budget() { + let window = 10_000; + // Room for the reservation plus two settled requests: three admitted. + let ceiling = reservation() + 2 * FIXTURE_USAGE_COST; + let control_store = TestStore::new(); + let mut control = open_limited(&control_store, ceiling, window); + assert_eq!(run_window(&mut control, "back", 20_000), 3); + + let store = TestStore::new(); + let mut service = open_limited(&store, ceiling, window); + let claim = claim_at(&mut service, "ahead", 50_000); + let result = text_result(claim.call(), "done"); + service.complete(*claim, result, FRAME).unwrap(); + // The clock now reads window 2, behind the recorded window 5: the + // request belongs to window 5, which already holds one. + assert_eq!(run_window(&mut service, "back", 20_000), 2); + assert_eq!(service.window_spend(CALLER, 5), 3 * FIXTURE_USAGE_COST); + assert_eq!(service.window_spend(CALLER, 2), 0); + drop(service); + + let mut service = open_limited(&store, ceiling, window); + assert_eq!(service.window_spend(CALLER, 5), 3 * FIXTURE_USAGE_COST); + assert_eq!(run_window(&mut service, "again", 20_100), 0); + assert_eq!( + refusal_at(&mut service, "last", 20_200, MAX_COST), + Refusal::ClockBehindAdmissionWindow + ); + } + + #[test] + fn a_clock_jumping_forward_opens_a_fresh_window() { + let window = 10_000; + let ceiling = reservation() + 2 * FIXTURE_USAGE_COST; + let store = TestStore::new(); + let mut service = open_limited(&store, ceiling, window); + assert_eq!(run_window(&mut service, "early", 20_000), 3); + assert_eq!(run_window(&mut service, "later", 90_000), 3); + assert_eq!(service.window_spend(CALLER, 9), 3 * FIXTURE_USAGE_COST); + assert_eq!(service.window_spend(CALLER, 2), 0); + } + + /// Hand-written rates: the long tier is cheaper on input and dearer on + /// output, so neither tier alone bounds both. + fn mixed_long_tier(above_prompt_tokens: u64) -> ModelPrice { + ModelPrice { + rates: Rates { + input: 3_000_000, + cache_write_short: 3_750_000, + cache_write_long: 6_000_000, + cache_read: 300_000, + output: 15_000_000, + }, + long_context: Some(LongContext { + above_prompt_tokens, + rates: Rates { + input: 1_000_000, + cache_write_short: 1_250_000, + cache_write_long: 2_000_000, + cache_read: 100_000, + output: 30_000_000, + }, + }), + overhead_tokens: 100, + } + } + + #[test] + fn the_reservation_is_the_dearer_of_both_tiers_class_by_class() { + let request = request(); + let bound = wire_encode(&request).unwrap().len() as u64 + 100; + // Above the threshold: input at the base tier's dearer rate, output at + // the long tier's. + let dearest = Rates { + input: 3_000_000, + cache_write_short: 3_750_000, + cache_write_long: 6_000_000, + cache_read: 300_000, + output: 30_000_000, + }; + assert_eq!( + reservation_micro_usd(&mixed_long_tier(bound - 1), &request).unwrap(), + Some(expected_reservation(&request, &dearest, 100)) + ); + // At the threshold the base tier alone applies. + assert_eq!( + reservation_micro_usd(&mixed_long_tier(bound), &request).unwrap(), + Some(expected_reservation(&request, &fixture_price().rates, 100)) + ); + } + + fn xorshift(state: &mut u64) -> u64 { + *state ^= *state << 13; + *state ^= *state >> 7; + *state ^= *state << 17; + *state + } + + fn random_rates(state: &mut u64) -> Rates { + Rates { + input: xorshift(state) % 50_000_000, + cache_write_short: xorshift(state) % 50_000_000, + cache_write_long: xorshift(state) % 50_000_000, + cache_read: xorshift(state) % 50_000_000, + output: xorshift(state) % 200_000_000, + } + } + + #[test] + fn the_property_test_generator_is_xorshift64_within_the_rate_bounds() { + let mut s = 0x9E37_79B9_7F4A_7C15u64; + let first: Vec = (0..3).map(|_| xorshift(&mut s)).collect(); + assert_eq!( + first, + [0xdc1b77ae0bf34dad, 0x64f0eeb9026e6076, 0x7b07ce91e5906136] + ); + for _ in 0..200 { + let rates = random_rates(&mut s); + for (rate, bound) in [ + (rates.input, 50_000_000), + (rates.cache_write_short, 50_000_000), + (rates.cache_write_long, 50_000_000), + (rates.cache_read, 50_000_000), + (rates.output, 200_000_000), + ] { + assert!(rate < bound); + } + } + } + + /// Any usage within the bound the reservation assumes (prompt tokens no + /// more than encoded bytes plus overhead, output no more than the cap) + /// costs no more than the reservation, whatever the price shape. + #[test] + fn no_usage_within_the_bound_prices_above_the_reservation() { + let mut s = 0x9E37_79B9_7F4A_7C15u64; + for _ in 0..5_000 { + let mut request = request(); + request.max_output_tokens = (xorshift(&mut s) % 100_000) as u32; + request.cache_retention = match xorshift(&mut s) % 3 { + 0 => None, + 1 => Some(CacheRetention::Short), + _ => Some(CacheRetention::Long), + }; + request.instructions = "x".repeat((xorshift(&mut s) % 3000) as usize); + let overhead = xorshift(&mut s) % 5_000; + let bound = wire_encode(&request).unwrap().len() as u64 + overhead; + let long_context = (xorshift(&mut s).is_multiple_of(2)).then(|| LongContext { + above_prompt_tokens: xorshift(&mut s) % (2 * bound + 1), + rates: random_rates(&mut s), + }); + let price = ModelPrice { + rates: random_rates(&mut s), + long_context, + overhead_tokens: overhead, + }; + let reserved = reservation_micro_usd(&price, &request).unwrap().unwrap(); + let prompt = if xorshift(&mut s).is_multiple_of(4) { + bound + } else { + xorshift(&mut s) % (bound + 1) + }; + let a = xorshift(&mut s) % (prompt + 1); + let b = xorshift(&mut s) % (prompt - a + 1); + let max_output = u64::from(request.max_output_tokens); + let output = if xorshift(&mut s).is_multiple_of(4) { + max_output + } else { + xorshift(&mut s) % (max_output + 1) + }; + let cost = cost_micro_usd( + &price, + &usage(a, b, prompt - a - b, output), + request.cache_retention, + ); + assert!( + cost <= reserved, + "{price:?} {a} {b} {output}: {cost} > {reserved}" + ); + } + } + + #[test] + fn another_callers_live_request_does_not_keep_this_callers_old_window() { + let store = TestStore::new(); + let mut service = open_service(&store, 4); + // Ghostlight holds a request live in window 0 throughout. + let other = invocation_for("ghostlight-yggdrasil", "request-1"); + let envelope = + encrypt_invocation(&other, &security("ghostlight-distinct-test-key")).unwrap(); + let TransportAdmission::Execute(_held) = service.begin(&envelope, 1_000, FRAME).unwrap() + else { + panic!("ghostlight should have executed"); + }; + for (id, now) in [("first", 1_000), ("second", WINDOW_MS + 1_000)] { + let claim = claim_at(&mut service, id, now); + let result = text_result(claim.call(), "done"); + service.complete(*claim, result, FRAME).unwrap(); + } + assert_eq!(service.window_spend(CALLER, 0), 0); + assert_eq!(service.window_spend(CALLER, 1), FIXTURE_USAGE_COST); + assert_eq!( + service.window_spend("ghostlight-yggdrasil", 0), + reservation() + ); + } + + #[test] + fn a_result_bound_to_another_request_is_charged_the_reservation() { + let store = TestStore::new(); + let mut service = open_service(&store, 4); + let one = claim_at(&mut service, "request-1", 1_000); + let two = claim_at(&mut service, "request-2", 1_000); + // request-2's result reports no usage at all; request-1 must not be + // billed what another request's backend claims to have spent. + let foreign = two.call().transported( + Vec::new(), + TransportOutcome::Failed { + kind: FailureKind::Provider, + message: String::new(), + usage: Some(usage(0, 0, 0, 0)), + }, + ); + let invocation = invocation_at("request-1", 1_000, MAX_COST); + let response = service.complete(*one, foreign, FRAME).unwrap(); + assert_eq!( + receipt_of(&response, &invocation).provider, + STAND_IN_PROVIDER + ); + assert_eq!(charged_of(&service, "request-1"), reservation()); + // request-1 is charged the reservation; request-2 still holds its own. + assert_eq!(service.window_spend(CALLER, 0), 2 * reservation()); + drop(two); + } + + #[test] + fn a_bound_malformed_overrun_is_charged_its_true_cost_and_reaches_the_ceiling() { + let reserved = reservation(); + for case in 0..3 { + let store = TestStore::new(); + let mut service = open_limited(&store, 2 * reserved, 3_600_000); + let claim = claim_at(&mut service, "request-1", 1_000); + let mut result = overrun_result(claim.call(), "done"); + let TransportDisposition::Transported { events, receipt } = &mut result.disposition + else { + unreachable!(); + }; + match case { + 0 => events[0].sequence = 1, + 1 => { + let TransportOutcome::Completed { usage, .. } = &receipt.outcome else { + unreachable!(); + }; + receipt.outcome = TransportOutcome::Failed { + kind: FailureKind::Provider, + message: String::new(), + usage: Some(usage.clone()), + }; + } + _ => receipt.provider = String::new(), + } + let invocation = invocation_at("request-1", 1_000, MAX_COST); + assert!(result.validate_against(&invocation).is_err(), "{case}"); + let response = service.complete(*claim, result, FRAME).unwrap(); + // The caller is told the result was refused; the bill is its true cost. + assert_eq!( + receipt_of(&response, &invocation).provider, + STAND_IN_PROVIDER, + "{case}" + ); + assert_eq!(charged_of(&service, "request-1"), OVERRUN_COST, "{case}"); + assert_eq!(service.window_spend(CALLER, 0), OVERRUN_COST, "{case}"); + assert_eq!( + refusal_at(&mut service, "request-2", 1_000, MAX_COST), + Refusal::SpendCeiling, + "{case}" + ); + } + } + + #[test] + fn each_binding_check_alone_sends_a_result_to_the_reservation() { + let reserved = reservation(); + // Every other field is valid and the usage overruns: a guard that + // ignored the one changed identity field would charge 30,000,000. + for case in 0..4 { + let store = TestStore::new(); + let mut service = open_limited(&store, 2 * reserved, 3_600_000); + let claim = claim_at(&mut service, "request-1", 1_000); + let mut result = overrun_result(claim.call(), "done"); + let TransportDisposition::Transported { receipt, .. } = &mut result.disposition else { + unreachable!(); + }; + match case { + 0 => result.request_id = "request-2".to_string(), + 1 => result.caller_runtime_id = "another-caller".to_string(), + 2 => receipt.request_id = "request-2".to_string(), + _ => receipt.caller_runtime_id = "another-caller".to_string(), + } + service.complete(*claim, result, FRAME).unwrap(); + assert_eq!(charged_of(&service, "request-1"), reserved, "{case}"); + assert_eq!(service.window_spend(CALLER, 0), reserved, "{case}"); + } + } + + #[test] + fn a_result_that_is_not_transported_is_charged_the_reservation() { + let store = TestStore::new(); + let mut service = open_service(&store, 4); + let claim = claim_at(&mut service, "request-1", 1_000); + let invocation = invocation_at("request-1", 1_000, MAX_COST); + let refused = TransportResult::refused(&invocation, Refusal::Capacity); + service.complete(*claim, refused, FRAME).unwrap(); + assert_eq!(charged_of(&service, "request-1"), reservation()); + assert_eq!(service.window_spend(CALLER, 0), reservation()); + } + + #[test] + fn a_stepped_back_clock_admits_into_the_newest_window_not_an_older_live_one() { + let window = 10_000; + // Two live reservations fill window 9 (the newest); one live + // reservation leaves window 2 (older, still held) room for another. + // Request ids share a length, so every reservation is the same size. + let ceiling = reservation() + reservation() + reservation() / 2; + let store = TestStore::new(); + let mut service = open_limited(&store, ceiling, window); + let old = claim_at(&mut service, "request-a", 20_000); + let newest = [ + claim_at(&mut service, "request-b", 90_000), + claim_at(&mut service, "request-c", 90_000), + ]; + assert_eq!(service.window_spend(CALLER, 2), reservation()); + assert_eq!(service.window_spend(CALLER, 9), 2 * reservation()); + // The clock reads window 2 again. Window 2 would hold the request; + // the caller's window is the newest, which cannot. + assert_eq!( + refusal_at(&mut service, "request-d", 20_500, MAX_COST), + Refusal::ClockBehindAdmissionWindow + ); + assert_eq!(service.window_spend(CALLER, 2), reservation()); + assert_eq!(service.window_spend(CALLER, 9), 2 * reservation()); + drop((old, newest)); + } + + #[test] + fn a_clock_behind_the_callers_window_is_refused_by_name_until_it_catches_up() { + let window = 10_000; + let ceiling = reservation() + 2 * FIXTURE_USAGE_COST; + let store = TestStore::new(); + let mut service = open_limited(&store, ceiling, window); + assert_eq!(run_window(&mut service, "fill", 50_000), 3); + // Behind the full window: the cause is named, no slot is taken and + // nothing is reserved. + assert_eq!( + refusal_at(&mut service, "behind", 20_000, MAX_COST), + Refusal::ClockBehindAdmissionWindow + ); + assert!(service.live.is_empty()); + assert_eq!(service.window_spend(CALLER, 5), 3 * FIXTURE_USAGE_COST); + // At the window's first instant, and inside it, the clock is not + // behind: an honest ceiling. + assert_eq!( + refusal_at(&mut service, "at-start", 50_000, MAX_COST), + Refusal::SpendCeiling + ); + assert_eq!( + refusal_at(&mut service, "inside", 50_500, MAX_COST), + Refusal::SpendCeiling + ); + // Once the clock reaches the next window the caller is admitted. + assert_eq!(run_window(&mut service, "later", 60_000), 3); + } + + #[test] + fn the_clock_behind_refusal_is_no_wider_than_the_widest_other_refusal() { + let invocation = invocation_at("request-1", 1_000, MAX_COST); + let sized = |refusal| { + wire_encode(&TransportResult::refused(&invocation, refusal)) + .unwrap() + .len() + }; + assert!( + sized(Refusal::ClockBehindAdmissionWindow) + <= sized(Refusal::ProviderDigestSubstitution) + ); + } } diff --git a/crates/model-connector/tests/fixtures/wire_v3/carriage_invocation.frame b/crates/model-connector/tests/fixtures/wire_v3/carriage_invocation.frame index 6051b64..b531c71 100644 --- a/crates/model-connector/tests/fixtures/wire_v3/carriage_invocation.frame +++ b/crates/model-connector/tests/fixtures/wire_v3/carriage_invocation.frame @@ -1 +1 @@ -‡©messageId®golden-request©operation¦invoke§payloadÚ$hqlzY2hlbWFfaWTZJGdhbWVjdWx0Lm1vZGVsLnRyYW5zcG9ydF9lbnZlbG9wZS52M7FjYWxsZXJfcnVudGltZV9pZK1nb2xkZW4tY2FsbGVyqnJlcXVlc3RfaWSuZ29sZGVuLXJlcXVlc3SsbWVzc2FnZV9raW5kqkludm9jYXRpb26lbm9uY2WcS2XMwjgRzMrM9QTM037MsczmqmNpcGhlcnRleHTFAy8eYmQGV5JCOLETwqQE+wIz/6yyPltHTPywbakMD/mgkJUxAvKWLNljXa4RAVVITAb3GxklnGF8zoOX/9qGSM9r0JZ126s+0LqvsRNhllG1S0pfG49fQJy99pevHReOSvX/cNWs2cHppYNdJUbVVAkDuioEXwtEs/vRCsmD/ab0bKUX0XRzrNbmjMU4+dHUCus8yt+iNTK9GdjXbEQAweeKNehirXEgnngLnwylo/MDoI4G237du3PmwhMzyIHztoIbwD0Bv9x7WqY+jj90vqEBsOHl+FmYc+F8+OiP6NKTDNRHypUtdo32G/IgfMfbsP/XHiNpPHtLxDqZVIcYdYqLAvhDfLBReLzQVLuNJnE0a2P8xLVAkfbmQZNsxRqS14RjYVtRoDcacxH1h0afo9OWMSK8Sp8eEBAhAJ7ovhCMklc7xWB3o4ny4cmDSExx5zWRLdI14rC1CZ7nq8k3V8fIIHctSp6am35A5yyiTwDH0m3h7/hvXrzhOqtAwHd0jkfNv+qOXZCEPdtbzOXOz0fMVZwvTAJWhoEYmhfKcGvXu8SmEtZcZe0WjF4xpFIFCFsSxh8HcrDFvxKDa58bnimo6DE2VeEMolpaBvnIGEst6qNZ9VQQnQIpCGbbtFd7jMrncnfcq+36NI10qW30r8sqAqgORp8cAGP5fy2X9SQDmH13SXgSfIdr7uxAHhK0X1r0iopJn4XtVfBiTd+IfoLFN+KO/JTZvLc48e6Jl9QRnxkL/UCgT/5wENd9VRxP+oO1CXG3DthI882rVbZDkhhwxupM9J/EURONajVyVSkeju+r9vzHCgR5odRW8cwqY1RXoq6L6mhj4ph01MBkUugz8MFwBsem+8QpnGatyVbZB99knWjK/01o5J6SZihcXNowRdon6BTydFByJaV5bRvgIK7IGRct3BmC8xidyxN9p7QAZYMPscxhQ0hAM615w76Cy82kx6eAzREBQX3COKApXSlPA/xlOKa/iUUMWHvFKyf/n0KrOlVvHjoDzPflHi3hpxxWcQGt8CxxNdSQMufeqce/5/BCeEwRIJ99+k7CEqIqpFIs7cAsI1GpF+NO5w==¯payloadEncoding²messagepack-base64­payloadSchemaÙ$gamecult.model.transport_envelope.v3­schemaVersion¼cultnet.operation_request.v0©serviceId¸gamecult.model.connector \ No newline at end of file +‡©messageId®golden-request©operation¦invoke§payloadÚDhqlzY2hlbWFfaWTZJGdhbWVjdWx0Lm1vZGVsLnRyYW5zcG9ydF9lbnZlbG9wZS52M7FjYWxsZXJfcnVudGltZV9pZK1nb2xkZW4tY2FsbGVyqnJlcXVlc3RfaWSuZ29sZGVuLXJlcXVlc3SsbWVzc2FnZV9raW5kqkludm9jYXRpb26lbm9uY2WczO7MxRrMxw7MhTh7SMzwzNB2qmNpcGhlcnRleHTFA0d65EskLG6kIwJ/G+XDK03gZOGtLFi6vZHJbI+/KpRLwnLT0TPF7RQa0KD98CrW+VwkKY3sRa4+Z9+/5OdPq1YmDFwbxvMVjZsFjXXw84tzvuLKHsb/3kQSEvAlFE4K6FXHvznLcMo4CnGUzoXwNGALf+xDTqFG0aFbdG4xGynV/kttsL6f34F4+l6xiXMUy87Q8/1UyOK/YaVxMMmlT3tdTllhq8YU8Dvfq0KuMaj/f9Anya8A5bJwlxafZPtc53RoH1/oNjLdaDqvwrTEnTB9TvKlnEOxKyyKhmmwGaC03qqPJXRaj8cDQO6EHDbML4J0zMN8TvssQISd5UCeVMkvDGjmvEOmRUoPk1/SnSrIHdha4+6rU4lUSumloLmlx5vv++Mq4AqzOdd/vPB6Ov/AoF32X27o4RI6H3xCKKHjKC4m3+cVjXCrOcGojuRf9yeDFkgbiPaC31hmSc8rAm6HO/NcaIHvx+G+6AaurAEhDLdF6DdrXb4xwjuv6yWQKihNvLXpnwmhRFkSeD5lS6FHW+CoOIdq/RBnFFyDrboTPWFoxQYIrBuqFoHDOajIz7dXYif0sf9sB9JBrDhlcbTpiGcyURfzDhW55NOPqKFGi6ropaJ1/SyLP3f7E57kA22aQ2Op7o36Tu06Pnzo2wI7UsL/Qz27F+AGAukQCvTJyjxARiGr3TRpChjVXx6GXhLrG1T9P+Mrhyrt9vwwhboN8qz1uvDZwKhGWBKXPXkxluq/NY1DwB3JTCCZfEwukWjMZbqiINLEKvW9BStcBR1kfXcK6sQvn/F0g9r81e9brbe+w3U24GERnHwgruOqytTeax/xPbQ7uISV48g5L8DAYf788KhUXwI43QDHNBdBUyo0rCsgoWJKeO05fcAgKgYYGsyLtMW10Vstw3q4POyn3o7rrBz4S2cnl2iZgr0ixg/KgF0V+loEVuj/6rYrcW0SLCHwV1+w1ClFMn0ZPsrz28EoZMA617G0Rrch66cS5uI9/Zf/wLX3z6n7xoPLNXk51ZsLSZ2bqEQkcBJbW/dT5MdH6anWVogRDvIW1coUbkIq9l+pkQSuelWsfakX4UQwYuUD0UwNbqxHiGwlEFj0f1bdhEj26Q==¯payloadEncoding²messagepack-base64­payloadSchemaÙ$gamecult.model.transport_envelope.v3­schemaVersion¼cultnet.operation_request.v0©serviceId¸gamecult.model.connector \ No newline at end of file diff --git a/crates/model-connector/tests/fixtures/wire_v3/transport_invocation.mp b/crates/model-connector/tests/fixtures/wire_v3/transport_invocation.mp index 95447869386911a8ebd04061ff2f8126e337e8a1..dff0ca28b054caeaea7a24b39f301fb7aba54c9d 100644 GIT binary patch delta 38 ucmbQwww;Z!eInyWiA}kQ74gaW#U=5%naM@@@ukHn=NNpv9X5W_WdZ;uY7a30 delta 14 VcmdnaHlK~LZ6f2x&0LJSOaLVh1X=(9 diff --git a/crates/model-connector/tests/fixtures/wire_v3/transport_receipt_failed.mp b/crates/model-connector/tests/fixtures/wire_v3/transport_receipt_failed.mp index 0c5feb28476fc1102939f7a750b06510305c5243..385a3188e4656e1de2360443cb81bdd64346d8da 100644 GIT binary patch delta 165 zcmZo~ diff --git a/crates/model-connector/tests/fixtures/wire_v3/transport_result_transported.mp b/crates/model-connector/tests/fixtures/wire_v3/transport_result_transported.mp index b85b43448f0a2afb42cb4d04961e2b3f3e75634b..79e6ebd3cd6a3a79292b265afe20da18941a2f1a 100644 GIT binary patch delta 58 zcmbQlzLR~!c1FhT%{v$^8P(S%=NFg6=Vm4s<;Ry6r<`LDy4Joizbv(=s5B2D969+T G(?S5Pu^34J delta 18 acmdnVK8bz9c1Fg|%{v$^87E6HF9ZNU7zRNA diff --git a/crates/model-connector/tests/published_schemas.rs b/crates/model-connector/tests/published_schemas.rs index 8a1bf47..15ce197 100644 --- a/crates/model-connector/tests/published_schemas.rs +++ b/crates/model-connector/tests/published_schemas.rs @@ -110,6 +110,9 @@ fn the_refusal_schema_is_one_flat_string_enum() { .iter() .map(|value| value.as_str().unwrap()) .collect(); - assert_eq!(variants.len(), 12); + assert_eq!(variants.len(), 15); + assert!(variants.contains(&"ClockBehindAdmissionWindow")); assert!(variants.contains(&"ForeignReasoningState")); + assert!(variants.contains(&"CostBound")); + assert!(variants.contains(&"SpendCeiling")); } diff --git a/crates/model-connector/tests/reasoning_state.rs b/crates/model-connector/tests/reasoning_state.rs index 559bd2f..f42e0f1 100644 --- a/crates/model-connector/tests/reasoning_state.rs +++ b/crates/model-connector/tests/reasoning_state.rs @@ -106,6 +106,8 @@ fn start_with(backend: Arc, max_concurrent: usize) -> Conne max_concurrent, 64 * 1024, 32_768, + SpendCeiling(1 << 40), + SpendWindowMs(3_600_000), ) .unwrap(); let limits = ServeLimits { @@ -114,7 +116,21 @@ fn start_with(backend: Arc, max_concurrent: usize) -> Conne socket_timeout_ms: 5_000, max_expiry_skew_ms: 300_000, }; - thread::spawn(move || serve(listener, limits, vec![admission], &store, backend)); + let prices = PriceTable::from([( + "claude-opus-4-5".to_string(), + ModelPrice { + rates: Rates { + input: 5_000_000, + cache_write_short: 6_250_000, + cache_write_long: 10_000_000, + cache_read: 500_000, + output: 25_000_000, + }, + long_context: None, + overhead_tokens: 100, + }, + )]); + thread::spawn(move || serve(listener, limits, vec![admission], prices, &store, backend)); ConnectorClient::new(endpoint, KEY, FRAME, Some(Duration::from_secs(5))).unwrap() } @@ -127,7 +143,7 @@ fn invocation(request_id: &str, input: Vec) -> TransportInvocation { 4_096, ); request.input = input; - TransportInvocation::new(CALLER, now_ms() + 60_000, [7; 32], request).unwrap() + TransportInvocation::new(CALLER, now_ms() + 60_000, [7; 32], 10_000_000, request).unwrap() } fn user(text: &str) -> InputItem { @@ -313,9 +329,10 @@ impl ProviderBackend for TamperingBackend { #[test] fn a_refused_completion_reaches_the_caller_as_a_typed_failure_and_frees_the_slot() { - let edits: [fn(&mut TransportReceipt); 2] = [ + let edits: [fn(&mut TransportReceipt); 3] = [ |receipt| receipt.provider = "another-provider".to_string(), |receipt| receipt.foreign_reasoning_dropped += 1, + |receipt| receipt.cost_micro_usd += 1, ]; for edit in edits { // One concurrent request: a wedged slot would refuse the next id. diff --git a/crates/model-connector/tests/wire_golden_v3.rs b/crates/model-connector/tests/wire_golden_v3.rs index dddf9ef..9364035 100644 --- a/crates/model-connector/tests/wire_golden_v3.rs +++ b/crates/model-connector/tests/wire_golden_v3.rs @@ -81,7 +81,7 @@ fn request() -> ProviderRequest { } fn invocation() -> TransportInvocation { - TransportInvocation::new(CALLER, 1_800_000_000_000, [7; 32], request()).unwrap() + TransportInvocation::new(CALLER, 1_800_000_000_000, [7; 32], 5_000_000, request()).unwrap() } fn receipt(outcome: TransportOutcome) -> TransportReceipt { @@ -95,6 +95,8 @@ fn receipt(outcome: TransportOutcome) -> TransportReceipt { model: invocation.request.model.clone(), provider: "anthropic-messages".to_string(), foreign_reasoning_dropped: 2, + cost_micro_usd: 1_234_567, + overrun_micro_usd: 89, outcome, } } @@ -168,6 +170,13 @@ fn wire_golden_v3() { let failed = receipt(TransportOutcome::Failed { kind: FailureKind::RateLimited, message: "provider asked to slow down".to_string(), + usage: Some(Usage { + input_tokens: 120, + cache_read_tokens: 0, + cache_write_tokens: 0, + output_tokens: 0, + reasoning_tokens: None, + }), }); assert_eq!(failed.validate_against(&invocation), Ok(())); assert_wire("transport_receipt_failed.mp", &failed); diff --git a/schemas/transport_invocation.json b/schemas/transport_invocation.json index 5435b2e..7c32147 100644 --- a/schemas/transport_invocation.json +++ b/schemas/transport_invocation.json @@ -11,6 +11,12 @@ "format": "uint64", "minimum": 0 }, + "max_cost_micro_usd": { + "description": "The most this request may cost, in micro-USD: the caller's own bound\non it. Required nonzero. It is part of the invocation's identity, so a\nchanged budget under a recorded request id is a `ReplayConflict`.", + "type": "integer", + "format": "uint64", + "minimum": 0 + }, "native_request_sha256": { "type": "array", "items": { @@ -47,6 +53,7 @@ "expires_at_unix_ms", "native_request_sha256", "provider_request_sha256", + "max_cost_micro_usd", "request" ], "$defs": { diff --git a/schemas/transport_receipt.json b/schemas/transport_receipt.json index 722a42a..01ff920 100644 --- a/schemas/transport_receipt.json +++ b/schemas/transport_receipt.json @@ -6,6 +6,12 @@ "caller_runtime_id": { "type": "string" }, + "cost_micro_usd": { + "description": "What the request cost, in micro-USD. The service derives it from the\noutcome's usage and its price table; a backend never writes it.", + "type": "integer", + "format": "uint64", + "minimum": 0 + }, "foreign_reasoning_dropped": { "description": "How many foreign `ReasoningState` items the service removed from the\nrequest before the backend saw it.", "type": "integer", @@ -29,6 +35,12 @@ "outcome": { "$ref": "#/$defs/TransportOutcome" }, + "overrun_micro_usd": { + "description": "How far `cost_micro_usd` exceeds what the service reserved for the\nrequest: zero unless the backend reported usage above the worst case\nthe request allowed, which is a breach of the backend's contract. The\ncost is still the true one and the caller's window is charged all of\nit. The service derives this; a backend never writes it.", + "type": "integer", + "format": "uint64", + "minimum": 0 + }, "provider": { "type": "string" }, @@ -60,6 +72,8 @@ "model", "provider", "foreign_reasoning_dropped", + "cost_micro_usd", + "overrun_micro_usd", "outcome" ], "$defs": { @@ -132,6 +146,17 @@ }, "message": { "type": "string" + }, + "usage": { + "description": "What the provider billed for the failed request: `Some` when it\nreported usage (`Some` of zeros when it billed nothing), `None`\nwhen unknown.", + "anyOf": [ + { + "$ref": "#/$defs/Usage" + }, + { + "type": "null" + } + ] } }, "additionalProperties": false, @@ -149,7 +174,7 @@ ] }, "Usage": { - "description": "Token accounting for one completed request. Cache reads and writes are\ncounted apart from plain input; `reasoning_tokens` is absent when the\nprovider does not report it.", + "description": "Token accounting for one request, as the backend reports it; the service\nprices it. Cache reads and writes are counted apart from plain input.\n`output_tokens` includes reasoning tokens: `reasoning_tokens` is a detail\n(absent when the provider does not report it) and is never priced apart.", "type": "object", "properties": { "cache_read_tokens": { diff --git a/schemas/transport_result.json b/schemas/transport_result.json index acdf6b6..39e010f 100644 --- a/schemas/transport_result.json +++ b/schemas/transport_result.json @@ -46,12 +46,15 @@ "ModelNotAllowed", "OutputBound", "PayloadBound", + "CostBound", + "SpendCeiling", "Capacity", "InFlight", "Indeterminate", "ReplayConflict", "Malformed", - "ForeignReasoningState" + "ForeignReasoningState", + "ClockBehindAdmissionWindow" ] }, "StopReason": { @@ -250,6 +253,17 @@ }, "message": { "type": "string" + }, + "usage": { + "description": "What the provider billed for the failed request: `Some` when it\nreported usage (`Some` of zeros when it billed nothing), `None`\nwhen unknown.", + "anyOf": [ + { + "$ref": "#/$defs/Usage" + }, + { + "type": "null" + } + ] } }, "additionalProperties": false, @@ -272,6 +286,12 @@ "caller_runtime_id": { "type": "string" }, + "cost_micro_usd": { + "description": "What the request cost, in micro-USD. The service derives it from the\noutcome's usage and its price table; a backend never writes it.", + "type": "integer", + "format": "uint64", + "minimum": 0 + }, "foreign_reasoning_dropped": { "description": "How many foreign `ReasoningState` items the service removed from the\nrequest before the backend saw it.", "type": "integer", @@ -295,6 +315,12 @@ "outcome": { "$ref": "#/$defs/TransportOutcome" }, + "overrun_micro_usd": { + "description": "How far `cost_micro_usd` exceeds what the service reserved for the\nrequest: zero unless the backend reported usage above the worst case\nthe request allowed, which is a breach of the backend's contract. The\ncost is still the true one and the caller's window is charged all of\nit. The service derives this; a backend never writes it.", + "type": "integer", + "format": "uint64", + "minimum": 0 + }, "provider": { "type": "string" }, @@ -326,11 +352,13 @@ "model", "provider", "foreign_reasoning_dropped", + "cost_micro_usd", + "overrun_micro_usd", "outcome" ] }, "Usage": { - "description": "Token accounting for one completed request. Cache reads and writes are\ncounted apart from plain input; `reasoning_tokens` is absent when the\nprovider does not report it.", + "description": "Token accounting for one request, as the backend reports it; the service\nprices it. Cache reads and writes are counted apart from plain input.\n`output_tokens` includes reasoning tokens: `reasoning_tokens` is a detail\n(absent when the provider does not report it) and is never priced apart.", "type": "object", "properties": { "cache_read_tokens": {