From 06a47bb13f99dc0ab9e7d35dabb9da839c4be810 Mon Sep 17 00:00:00 2001 From: Meta Vens Date: Sat, 31 Oct 2026 04:43:18 +0100 Subject: [PATCH] connector-usage-only: delete the dollar authorities (WIP, does not yet build) Co-Authored-By: Claude Opus 5.5 --- .gitattributes | 2 +- README.md | 8 +- crates/claude-connector/Cargo.toml | 3 + crates/claude-connector/src/backend/tests.rs | 81 +- crates/claude-connector/src/daemon.rs | 66 +- crates/claude-connector/src/lease.rs | 247 +- crates/claude-connector/src/lib.rs | 2 +- crates/claude-connector/src/main.rs | 58 +- crates/claude-connector/src/quota.rs | 2 +- crates/claude-connector/src/stream.rs | 6 +- .../claude-connector/tests/config_document.rs | 1684 +----------- .../tests/fixtures/lease_v1/grant.bin | Bin 109 -> 0 bytes .../tests/fixtures/lease_v1/refusal.bin | Bin 104 -> 0 bytes .../tests/fixtures/lease_v1/request.bin | Bin 66 -> 0 bytes .../tests/fixtures/lease_v1/settle.bin | Bin 336 -> 0 bytes .../tests/fixtures/lease_v1/settled.bin | Bin 81 -> 0 bytes crates/claude-connector/tests/lease.rs | 1306 ++------- crates/model-connector/src/client.rs | 34 +- crates/model-connector/src/contract.rs | 83 +- crates/model-connector/src/envelope.rs | 94 +- crates/model-connector/src/fixtures.rs | 43 +- crates/model-connector/src/readout.rs | 55 +- crates/model-connector/src/serve.rs | 127 +- crates/model-connector/src/service.rs | 2342 ++--------------- crates/model-connector/src/service/lease.rs | 874 ++---- .../fixtures/{wire_v3 => wire_v4}/README.md | 8 +- .../carriage_invocation.frame | 0 .../input_reasoning_state.mp | 0 .../{wire_v3 => wire_v4}/provider_request.mp | Bin .../transport_envelope.mp | Bin .../transport_invocation.mp | Bin .../transport_receipt_failed.mp | Bin .../transport_result_refused.mp | 0 .../transport_result_transported.mp | Bin .../tests/published_schemas.rs | 12 +- .../model-connector/tests/reasoning_state.rs | 22 +- .../{wire_golden_v3.rs => wire_golden_v4.rs} | 18 +- schemas/spend_readout.json | 112 - schemas/transport_envelope.json | 4 +- 39 files changed, 1263 insertions(+), 6030 deletions(-) delete mode 100644 crates/claude-connector/tests/fixtures/lease_v1/grant.bin delete mode 100644 crates/claude-connector/tests/fixtures/lease_v1/refusal.bin delete mode 100644 crates/claude-connector/tests/fixtures/lease_v1/request.bin delete mode 100644 crates/claude-connector/tests/fixtures/lease_v1/settle.bin delete mode 100644 crates/claude-connector/tests/fixtures/lease_v1/settled.bin rename crates/model-connector/tests/fixtures/{wire_v3 => wire_v4}/README.md (74%) rename crates/model-connector/tests/fixtures/{wire_v3 => wire_v4}/carriage_invocation.frame (100%) rename crates/model-connector/tests/fixtures/{wire_v3 => wire_v4}/input_reasoning_state.mp (100%) rename crates/model-connector/tests/fixtures/{wire_v3 => wire_v4}/provider_request.mp (100%) rename crates/model-connector/tests/fixtures/{wire_v3 => wire_v4}/transport_envelope.mp (100%) rename crates/model-connector/tests/fixtures/{wire_v3 => wire_v4}/transport_invocation.mp (100%) rename crates/model-connector/tests/fixtures/{wire_v3 => wire_v4}/transport_receipt_failed.mp (100%) rename crates/model-connector/tests/fixtures/{wire_v3 => wire_v4}/transport_result_refused.mp (100%) rename crates/model-connector/tests/fixtures/{wire_v3 => wire_v4}/transport_result_transported.mp (100%) rename crates/model-connector/tests/{wire_golden_v3.rs => wire_golden_v4.rs} (96%) delete mode 100644 schemas/spend_readout.json diff --git a/.gitattributes b/.gitattributes index 0a1d199..56862bb 100644 --- a/.gitattributes +++ b/.gitattributes @@ -1,2 +1,2 @@ schemas/*.json text eol=lf -crates/model-connector/tests/fixtures/wire_v3/* -text +crates/model-connector/tests/fixtures/wire_v4/* -text diff --git a/README.md b/README.md index c2e23c7..34c8730 100644 --- a/README.md +++ b/README.md @@ -21,11 +21,11 @@ Crate `crates/model-connector`: - `src/envelope.rs` - the AES-256-GCM `TransportEnvelope` (binary ciphertext), its identity-binding AAD, the length-prefixed framing, and the CultNet carriage: each frame is a `cultnet.operation_request.v0` / `operation_response.v0` for service - `gamecult.model.connector`, operation `invoke` (or `spend_readout`), its `message_id` the - request id and its payload the base64 of the sealed envelope. `spend_readout` answers a - `SpendReadout` (`gamecult.model.spend_readout.v1`): each caller's spent and reserved + `gamecult.model.connector`, operation `invoke` (or `usage_readout`), its `message_id` the + request id and its payload the base64 of the sealed envelope. `usage_readout` answers a + `UsageReadout` (`gamecult.model.usage_readout.v1`): each caller's spent and reserved micro-USD against its ceiling in its current window, and the backend's plan quota; - `ConnectorClient::spend_readout` asks. + `ConnectorClient::usage_readout` asks. - `schemas/` - one JSON Schema per v3 document, generated from the types. - `src/client.rs` - `ConnectorClient`, the one-request caller side. - `src/service.rs` (feature `daemon`) - `TransportService`: caller admission, durable diff --git a/crates/claude-connector/Cargo.toml b/crates/claude-connector/Cargo.toml index 6592252..7b310e5 100644 --- a/crates/claude-connector/Cargo.toml +++ b/crates/claude-connector/Cargo.toml @@ -25,3 +25,6 @@ zeroize = "1" [target.'cfg(unix)'.dependencies] libc = "0.2" + +[dev-dependencies] +rmpv = "1" diff --git a/crates/claude-connector/src/backend/tests.rs b/crates/claude-connector/src/backend/tests.rs index fc2be99..227b9f9 100644 --- a/crates/claude-connector/src/backend/tests.rs +++ b/crates/claude-connector/src/backend/tests.rs @@ -5,8 +5,7 @@ use std::sync::{Arc, Mutex}; use std::time::Duration; use model_connector::{ - CallerAdmission, ModelPrice, PriceTable, Rates, ReasoningEffort, SpendCeiling, SpendWindowMs, - StopReason, ToolDefinition, TransportAdmission, TransportDisposition, TransportInvocation, + CallerAdmission, PlanQuotaSnapshot, ReasoningEffort, StopReason, ToolDefinition, TransportAdmission, TransportDisposition, TransportInvocation, TransportKey, TransportReceipt, TransportService, decrypt_result, encrypt_invocation, }; @@ -460,25 +459,6 @@ struct Door { service: TransportService, } -fn prices() -> PriceTable { - PriceTable::from([MODEL, OTHER_MODEL].map(|model| { - ( - model.to_string(), - 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, - }, - ) - })) -} - impl Door { fn new() -> Self { let dir = std::env::temp_dir().join(format!( @@ -496,13 +476,11 @@ impl Door { 4, FRAME / 2, 32_768, - SpendCeiling(1 << 40), - SpendWindowMs(3_600_000), ) .unwrap() }); let service = - TransportService::open(&dir.join("replay.cc"), PROVIDER, prices(), callers, 5_000) + TransportService::open(&dir.join("replay.cc"), PROVIDER, callers, 5_000) .unwrap(); Self { dir, service } } @@ -521,9 +499,12 @@ impl Door { ) -> (model_connector::ExecutionClaim, TransportInvocation) { let key = TransportKey::from_connection_secret(&format!("{caller}-key")).unwrap(); let invocation = - TransportInvocation::new(caller, 2_000, [7; 32], 10_000_000, request).unwrap(); + TransportInvocation::new(caller, 2_000, [7; 32], request).unwrap(); let envelope = encrypt_invocation(&invocation, &key).unwrap(); - match self.service.begin(&envelope, 1_000, FRAME).unwrap() { + match self + .service + .begin(&envelope, 1_000, &PlanQuotaSnapshot::default(), FRAME) + .unwrap() { TransportAdmission::Execute(claim) => (*claim, invocation), TransportAdmission::Reply(_) => panic!("the door answered instead of admitting"), } @@ -1458,11 +1439,6 @@ fn usage_sums_model_usage_across_models() { }; assert_eq!(*summed, usage(12, 37, 1083, 24, Some(8))); // The receipt is priced from usage alone: 12*3 + 37*0.3 + 1083*6 + 24*15 micro-USD, rounded up. - assert_eq!(parts(&cheap).1.cost_micro_usd, 36 + 12 + 6498 + 360); - assert_eq!( - parts(&cheap).1.cost_micro_usd, - parts(&dear).1.cost_micro_usd - ); // No model reported thinking tokens: none is claimed. let plain = r#"{"claude-haiku-5-5":{"inputTokens":1,"outputTokens":1,"cacheReadInputTokens":0,"cacheCreationInputTokens":0}}"#; @@ -1599,7 +1575,6 @@ fn no_result_event_is_a_failure_with_unknown_usage() { }; assert_eq!(zero_billed, Some(usage_zero())); - let mut costs = Vec::new(); for (script, expected) in ends { let rig = rig(vec![script]); let result = rig.run(request(vec![user("hi")])); @@ -1607,20 +1582,7 @@ fn no_result_event_is_a_failure_with_unknown_usage() { assert_eq!(kind, expected); assert_eq!(usage, None, "{message}"); assert!(result.transcript_items().is_empty()); - costs.push(parts(&result).1.cost_micro_usd); } - // Unknown usage is charged the reservation, which is far above a typical turn. - assert!(costs.iter().all(|cost| *cost >= 15_000), "{costs:?}"); - assert_eq!( - parts(&rig(vec![]).run({ - let mut bad = request(vec![user("a")]); - bad.tools.push(a_tool()); - bad - })) - .1 - .cost_micro_usd, - 0 - ); } #[test] @@ -1719,7 +1681,6 @@ fn a_port_that_cannot_spawn_is_a_provider_failure() { Some(usage_zero()), "a child that never ran bills nothing" ); - assert_eq!(parts(&result).1.cost_micro_usd, 0); assert!( !message.contains(dir.path().to_str().unwrap()), "no path in a message" @@ -1743,8 +1704,6 @@ fn backend_end_to_end_through_the_service() { assert_eq!(receipt.provider, "claude-code"); assert_eq!(receipt.model, MODEL); // USAGE priced at the table: 2*3 + 7*0.3 + 1043*6 + 4*15 micro-USD, rounded up. - assert_eq!(receipt.cost_micro_usd, 6 + 3 + 6258 + 60); - assert_eq!(receipt.overrun_micro_usd, 0); assert_eq!(events.len(), 2); let items = opened.transcript_items(); assert_eq!(items[0], assistant("hello back")); @@ -1757,13 +1716,6 @@ fn backend_end_to_end_through_the_service() { const USAGE_TURN_2: &str = r#"{"claude-haiku-5-5":{"inputTokens":3,"outputTokens":5,"cacheReadInputTokens":1100,"cacheCreationInputTokens":9}}"#; -/// What a call whose usage is unknown is billed: the kit's reservation. -fn reservation() -> u64 { - let result = rig(vec![Script::lines(vec![])]).run(request(vec![user("hi")])); - assert_eq!(failure_of(&result).2, None); - parts(&result).1.cost_micro_usd -} - #[test] fn a_resumed_turn_is_charged_only_its_own_usage() { let rig = rig(vec![ @@ -1788,13 +1740,10 @@ fn a_resumed_turn_is_charged_only_its_own_usage() { }; assert_eq!(*billed, usage(3, 1100, 9, 5, None)); // 3*3 + 1100*0.3 + 9*6 + 5*15 micro-USD. - assert_eq!(parts(&second).1.cost_micro_usd, 9 + 330 + 54 + 75); } #[test] -fn usage_the_result_does_not_report_is_billed_the_reservation() { - let reservation = reservation(); - assert!(reservation >= 15_000, "{reservation}"); +fn usage_the_result_does_not_report_stays_unknown() { let success = |fields: &str| { format!( r#"{{"type":"result","subtype":"success","is_error":false,"stop_reason":"end_turn","modelUsage":{CUMULATIVE}{fields}}}"# @@ -1838,7 +1787,6 @@ fn usage_the_result_does_not_report_is_billed_the_reservation() { let (kind, _, usage) = failure_of(&result); assert_eq!(kind, expected, "{line}"); assert_eq!(usage, None, "{line}"); - assert_eq!(parts(&result).1.cost_micro_usd, reservation, "{line}"); } for line in cases { let result = rig(vec![Script::lines(vec![ @@ -1850,7 +1798,6 @@ fn usage_the_result_does_not_report_is_billed_the_reservation() { let (kind, _, usage) = failure_of(&result); assert_eq!(kind, FailureKind::Provider, "{line}"); assert_eq!(usage, None, "{line}"); - assert_eq!(parts(&result).1.cost_micro_usd, reservation, "{line}"); assert!(result.transcript_items().is_empty(), "{line}"); } } @@ -1865,7 +1812,6 @@ fn door_set_is_exactly_401_403_429() { // with a bogus token. let retry = r#"{"type":"system","subtype":"api_retry","error":"authentication_failed"}"#; let synthetic = r#"{"type":"assistant","api_error_status":401,"error":"authentication_failed","message":{"id":"e48001b7-cf4f-49b1-9810-db8fd7f73c73","model":"","stop_reason":"stop_sequence","content":[{"type":"text","text":"Failed to authenticate. API Error: 401 OAuth access token is invalid."}]}}"#; - let reservation = reservation(); let statuses: [Option; 9] = [ Some(400), Some(401), @@ -1904,10 +1850,8 @@ fn door_set_is_exactly_401_403_429() { let case = format!("{status:?} {shape}"); if matches!(status, Some(401 | 403 | 429)) && shape == "empty" { assert_eq!(usage, Some(usage_zero()), "{case}"); - assert_eq!(parts(&result).1.cost_micro_usd, 0, "{case}"); } else { assert_eq!(usage, None, "{case}"); - assert_eq!(parts(&result).1.cost_micro_usd, reservation, "{case}"); } } } @@ -1979,7 +1923,6 @@ fn prompt_file_lives_in_prompt_dir_exclusively() { let result = rig.run(request(vec![user("hi")])); let (kind, _, usage) = failure_of(&result); assert_eq!((kind, usage), (FailureKind::Provider, Some(usage_zero()))); - assert_eq!(parts(&result).1.cost_micro_usd, 0); assert_eq!(rig.spawns(), 0); } @@ -1996,11 +1939,6 @@ fn pin_is_checked_before_every_spawn() { let result = rig.run(request(vec![user("hi")])); let (kind, message, usage) = failure_of(&result); assert_eq!((kind, usage), (FailureKind::Provider, Some(usage_zero()))); - assert_eq!( - parts(&result).1.cost_micro_usd, - 0, - "billed nothing, not the reservation" - ); assert!(!message.contains(rig.dir.path().to_str().unwrap())); assert_eq!(rig.spawns(), 1, "no child for a binary off its pin"); assert!(!rig.backend.binary_matches_pin()); @@ -2695,7 +2633,6 @@ fn not_logged_in_line(status: &str, api_ms: u64, extra: &str) -> String { #[test] fn a_result_that_made_no_api_call_bills_nothing() { - let reservation = reservation(); let empty = r#","model_usage_since_last_result":{}"#; let cases = [ // (status field, api ms, usage field, billed nothing) @@ -2723,10 +2660,8 @@ fn a_result_that_made_no_api_call_bills_nothing() { let (_, _, usage) = failure_of(&result); if free { assert_eq!(usage, Some(usage_zero()), "{line}"); - assert_eq!(parts(&result).1.cost_micro_usd, 0, "{line}"); } else { assert_eq!(usage, None, "{line}"); - assert_eq!(parts(&result).1.cost_micro_usd, reservation, "{line}"); } } } diff --git a/crates/claude-connector/src/daemon.rs b/crates/claude-connector/src/daemon.rs index 4b2dd6a..eaf26d0 100644 --- a/crates/claude-connector/src/daemon.rs +++ b/crates/claude-connector/src/daemon.rs @@ -3,7 +3,7 @@ //! run with; nothing in the connector's environment, flags or JSON can change //! them. Validation errors name a field and never echo a value. -use std::collections::{BTreeMap, BTreeSet}; +use std::collections::BTreeSet; use std::net::{SocketAddr, TcpListener}; use std::path::{Path, PathBuf}; use std::sync::{Arc, Mutex}; @@ -11,7 +11,7 @@ use std::time::Duration; use cultcache_rs::{CultCache, DatabaseEntry, SingleFileMessagePackBackingStore}; use model_connector::{ - CallerAdmission, CallerConfig, LeaseTurn, MIN_PAYLOAD_BYTES, ModelPrice, ProviderBackend, + CallerAdmission, CallerConfig, LeaseTurn, MIN_PAYLOAD_BYTES, ProviderBackend, ServeError, ServeLimits, ServiceError, TransportService, load_admissions, serve_service, }; use serde::{Deserialize, Serialize}; @@ -31,8 +31,8 @@ const MODEL_PREFIX: &str = "claude-"; /// A unix account that may lease the subscription token for a Self run or a /// Persona turn, the unit its process must run in (`mind-self@` -/// or `mind-persona@`, by `turn`), the configured caller whose -/// spend window its leases reserve in, and what each lease reserves there. +/// or `mind-persona@`, by `turn`), and the configured caller +/// whose name its leases carry. #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] pub struct LeaseCallerConfig { pub unix_user: String, @@ -41,11 +41,10 @@ pub struct LeaseCallerConfig { /// or `mind-persona@`, its process must run in. pub turn: LeaseTurn, pub unit_instance: String, - pub reservation_micro_usd: u64, } #[derive(Clone, Debug, PartialEq, Eq, DatabaseEntry)] -#[cultcache(type = "gamecult.model_connector.claude_connector_config.v1")] +#[cultcache(type = "gamecult.model_connector.claude_connector_config.v2")] pub struct ClaudeConnectorConfig { /// A loopback socket address. #[cultcache(key = 0)] @@ -86,14 +85,10 @@ pub struct ClaudeConnectorConfig { #[cultcache(key = 16)] pub max_expiry_skew_ms: u64, /// Every caller of this connector: its own connection key and epoch, models, - /// concurrency, payload and output bounds, and spend ceiling and window. + /// concurrency, payload and output bounds. /// Ids are unique. Handles are bound to the caller that was given them. #[cultcache(key = 17)] pub callers: Vec, - /// What each allowed model costs. The keys are the known list of models: - /// ids exactly as the CLI reports them in its init event. - #[cultcache(key = 18)] - pub prices: BTreeMap, /// Where the Self-run lease socket is bound (a systemd RuntimeDirectory), /// or none when this connector lends no token. #[cultcache(key = 19)] @@ -143,7 +138,6 @@ impl From for ConfigError { ServeError::ConnectionKey(_) => Self::Unusable("connection_key_file"), ServeError::Service(error) => match error { ServiceError::ReplayStore(_) => Self::Unusable("replay_store"), - ServiceError::UnpricedModel(_) => Self::Invalid("prices"), ServiceError::InvalidAdmission | ServiceError::DuplicateCaller | ServiceError::SharedCallerKey => Self::Invalid("callers"), @@ -228,8 +222,6 @@ impl ClaudeConnectorConfig { if caller.max_concurrent_requests == 0 || caller.max_payload_bytes < MIN_PAYLOAD_BYTES || caller.max_output_tokens == 0 - || caller.spend_ceiling_micro_usd == 0 - || caller.spend_window_ms == 0 { return Err(ConfigError::Invalid("caller bounds")); } @@ -240,19 +232,11 @@ impl ClaudeConnectorConfig { for model in &caller.allowed_models { if !has_model_id_form(model) || !seen.insert(model.as_str()) - || !self.prices.contains_key(model) { return Err(ConfigError::Invalid("allowed_models")); } } } - if !self - .prices - .iter() - .all(|(model, price)| has_model_id_form(model) && is_fully_priced(price)) - { - return Err(ConfigError::Invalid("prices")); - } self.validate_leases()?; Ok(listen) } @@ -278,11 +262,7 @@ impl ClaudeConnectorConfig { && self .callers .iter() - .find(|caller| caller.caller_runtime_id == lease.caller_id) - .is_some_and(|caller| { - (1..=caller.spend_ceiling_micro_usd) - .contains(&lease.reservation_micro_usd) - }) + .any(|caller| caller.caller_runtime_id == lease.caller_id) && is_unit_instance(&lease.unit_instance) && users.insert(lease.unix_user.as_str()) }) @@ -336,8 +316,7 @@ fn is_trimmed_nonempty(text: &str) -> bool { /// The loose shape of a model id, only to catch typos: /// `claude--`, then an optional `-` and an optional /// `-YYYYMMDD`, the family lowercase letters. No shape can tell an alias from -/// a canonical id, so the known list is the price table's keys (ruling -/// claude-model-ids-from-price-table); an alias that slips in fails its first +/// a canonical id, so the known list is `allowed_models` itself; an alias that slips in fails its first /// call, when the stream's `init.model` differs from the request. fn has_model_id_form(model: &str) -> bool { let Some(rest) = model.strip_prefix(MODEL_PREFIX) else { @@ -365,25 +344,6 @@ fn has_model_id_form(model: &str) -> bool { && after_ok } -/// A price is a fuse only if every component it bills is nonzero: a zero rate -/// lets a call cost nothing against the caller's ceiling. -fn is_fully_priced(price: &ModelPrice) -> bool { - let long = price.long_context.map(|long| long.rates); - [Some(price.rates), long] - .into_iter() - .flatten() - .all(|rates| { - [ - rates.input, - rates.cache_write_short, - rates.cache_write_long, - rates.cache_read, - rates.output, - ] - .iter() - .all(|rate| *rate > 0) - }) -} /// One file name under `$CREDENTIALS_DIRECTORY`, never a path. fn is_credential_name(name: &str) -> bool { @@ -509,7 +469,6 @@ fn lease_callers_of(config: &ClaudeConnectorConfig) -> Result, caller_id: lease.caller_id.clone(), turn: lease.turn, unit_instance: lease.unit_instance.clone(), - reservation_micro_usd: lease.reservation_micro_usd, }) .ok_or(ConfigError::Unusable("lease_callers")) }) @@ -531,6 +490,7 @@ fn start_leases( config: &ClaudeConnectorConfig, socket: &Path, service: Arc>, + quota: Arc, ) -> Result<(), ConfigError> { use std::os::unix::fs::{FileTypeExt, PermissionsExt}; use std::os::unix::net::UnixListener; @@ -559,6 +519,7 @@ fn start_leases( callers, limits, service, + quota, tokens, Arc::new(SystemdUnits::system()), ); @@ -589,7 +550,6 @@ pub fn serve_config(path: &Path) -> Result<(), ConfigError> { TransportService::open( &config.replay_store, backend.provider(), - config.prices.clone(), admissions, config.max_expiry_skew_ms, ) @@ -597,7 +557,7 @@ pub fn serve_config(path: &Path) -> Result<(), ConfigError> { )); #[cfg(target_os = "linux")] if let Some(socket) = &config.lease_socket { - start_leases(&config, socket, service.clone())?; + start_leases(&config, socket, service.clone(), backend.quota())?; } serve_service(listener, config.serve_limits(), service, backend)?; Ok(()) @@ -643,10 +603,6 @@ mod tests { ServeError::Service(ServiceError::ReplayStore(format!("{CANARY}/replay.cc"))), "replay_store", ), - ( - ServeError::Service(ServiceError::UnpricedModel(CANARY.to_string())), - "prices", - ), ( ServeError::Service(ServiceError::InvalidAdmission), "callers", diff --git a/crates/claude-connector/src/lease.rs b/crates/claude-connector/src/lease.rs index 9235a17..78b19ac 100644 --- a/crates/claude-connector/src/lease.rs +++ b/crates/claude-connector/src/lease.rs @@ -1,15 +1,14 @@ //! The Self-run credential lease. The connector is the only holder of the -//! subscription token and the one owner of Self-run admission, reservation and -//! the spend ledger (rulings self-runs-through-connector, -//! self-run-credential-lease and self-run-lease-billing-and-binding). A Self -//! run asks over a unix socket for a lease. The connector decides who may ask -//! (the peer's uid, and the systemd unit invocation its process runs in), -//! reserves its own configured amount in the caller's spend window, and only -//! then hands back one pipe fd that holds the token. The run settles with the -//! per-model usage it reports, and the connector bills that at true cost, -//! uncapped; a run that vanishes or lets its lease expire is billed the whole -//! reservation. Nothing the run says about itself decides who it is, what it -//! reserves, or what the plan allows. +//! subscription token and the one owner of Self-run admission (rulings +//! self-runs-through-connector and self-run-credential-lease). A Self run asks +//! over a unix socket for a lease. The connector decides who may ask (the +//! peer's uid, and the systemd unit invocation its process runs in), and hands +//! back one pipe fd that holds the token, with the plan's usage windows. The +//! lease lasts while the holder keeps the socket open: it ends when the holder +//! closes it, when its unit invocation ends, or when its expiry passes. Nothing +//! is reserved and nothing is billed; the only usage-derived refusal is the +//! provider's own, repeated until its own reset. Nothing the run says about +//! itself decides who it is or what the plan allows. //! //! The token exists in this process only as the value the token source returns //! and in the pipe. It is not in an environment, an argument, a file, a log or @@ -28,12 +27,11 @@ use std::thread; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use model_connector::{ - LeaseAdmission, LeaseGrant, LeaseRefusal, LeaseSettlement, LeaseTurn, LeaseUse, ModelUsage, - TransportService, + LeaseAdmission, LeaseGrant, LeaseRefusal, LeaseTurn, PlanQuota, TransportService, }; use serde::{Deserialize, Serialize}; -use crate::TokenSource; +use crate::{QuotaCell, TokenSource}; /// The largest frame either side sends or accepts, in bytes. It is wire /// contract: a Huginn client sizes against it, so it is pinned as a literal. @@ -50,12 +48,9 @@ pub enum LeaseMessage { /// holds the token, its write end closed. ClaudeLeaseGrant(LeaseGranted), ClaudeLeaseRefusal(LeaseRefused), - ClaudeLeaseSettle(LeaseSettle), - ClaudeLeaseSettled(LeaseSettled), } -/// A request names only when the lease ends. Who asks, and what is reserved, -/// are the connector's to know. +/// A request names only when the lease ends. Who asks is the connector's to know. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct LeaseRequest { pub expires_at_unix_ms: u64, @@ -64,16 +59,17 @@ pub struct LeaseRequest { #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct LeaseGranted { pub lease_id: String, - pub reserved_micro_usd: u64, + /// The plan's usage windows as the connector last saw them, so the run + /// can read its own burn. + pub plan_quota: Vec, } #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] pub enum LeaseRefusalKind { - SpendCeiling, - CostBound, UnknownCaller, Expired, Busy, + ProviderRejected, } #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] @@ -82,20 +78,6 @@ pub struct LeaseRefused { pub detail: String, } -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct LeaseSettle { - pub lease_id: String, - /// What the run used of each model, as the run reports it. - pub usage: Vec, -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct LeaseSettled { - pub billed_micro_usd: u64, - /// How far the bill passed the reservation; zero when it did not. - pub overrun_micro_usd: u64, -} - /// A frame: the length, then the named-map MessagePack of `message`. pub fn encode_frame(message: &LeaseMessage) -> Vec { // Every field of the message is plain data, so encoding cannot fail. @@ -180,8 +162,8 @@ fn read_message(stream: &mut impl Bounded, deadline: Instant) -> io::Result` of `mind-self@.service` or /// `mind-persona@.service`. pub unit_instance: String, - /// What each of this caller's leases reserves, whatever the run asks. - pub reservation_micro_usd: u64, } impl LeaseCaller { @@ -400,6 +380,7 @@ struct Lender { callers: Vec, limits: LeaseLimits, service: Arc>, + quota: Arc, tokens: Arc, units: Arc, } @@ -415,6 +396,7 @@ pub fn serve_leases( callers: Vec, limits: LeaseLimits, service: Arc>, + quota: Arc, tokens: Arc, units: Arc, ) -> io::Result<()> { @@ -422,6 +404,7 @@ pub fn serve_leases( callers, limits, service, + quota, tokens, units, }); @@ -484,9 +467,9 @@ impl Lender { "the expiry is not in the future or is past the lease bound", ); } - // The reservation is claimed in the ledger before the token is read, - // let alone lent. A ledger that fails closes the connection with no - // answer and no lease. + // The lease is claimed before the token is read, let alone lent. The + // plan the refusal reads is the one this process last observed. + let plan = self.quota.snapshot(); let admission = self .service .lock() @@ -495,120 +478,134 @@ impl Lender { &caller.caller_id, caller.turn, &invocation.invocation_id, - caller.reservation_micro_usd, request.expires_at_unix_ms, now, + &plan, ) - .map_err(|_| io::Error::other("ledger"))?; + .map_err(|_| io::Error::other("lease"))?; let grant = match admission { LeaseAdmission::Granted(grant) => grant, LeaseAdmission::Refused(why) => { let (kind, detail) = match why { - LeaseRefusal::SpendCeiling => ( - LeaseRefusalKind::SpendCeiling, - "the window cannot hold this reservation", - ), - LeaseRefusal::CostBound => ( - LeaseRefusalKind::CostBound, - "the reservation must be above zero", - ), LeaseRefusal::UnknownCaller => ( LeaseRefusalKind::UnknownCaller, "this account is bound to no admitted caller", ), LeaseRefusal::Busy => ( LeaseRefusalKind::Busy, - "a lease is open for this caller or this unit invocation leased", + "a lease is open for this caller or this unit invocation", + ), + LeaseRefusal::ProviderRejected => ( + LeaseRefusalKind::ProviderRejected, + "the provider is refusing until its own reset", ), }; return refuse(&mut stream, kind, detail); } }; - self.hold(stream, caller, &grant, request.expires_at_unix_ms) + self.hold( + stream, + caller, + &grant, + &invocation, + peer.process, + request.expires_at_unix_ms, + plan.windows, + ) } - /// Everything after a reservation exists: lend the token, then settle the - /// lease exactly once whatever happens. + /// Everything after a lease exists: lend the token, then keep the lease + /// while its holder keeps the socket and its unit invocation, and end it + /// exactly once whatever happens. + // One argument per fact the held lease binds. + #[allow(clippy::too_many_arguments)] fn hold( &self, mut stream: UnixStream, caller: &LeaseCaller, grant: &LeaseGrant, + invocation: &PeerUnit, + process: Option, expires_at_unix_ms: u64, + plan_quota: Vec, ) -> io::Result<()> { - // A token that cannot be read or piped was never lent: nothing was - // used, so nothing is billed. + let held = self.keep( + &mut stream, + caller, + grant, + invocation, + process.as_ref(), + expires_at_unix_ms, + plan_quota, + ); + self.service + .lock() + .map_err(|_| io::Error::other("service lock"))? + .end_lease(&caller.caller_id, &grant.lease_id); + held + } + + // One argument per fact the held lease binds. + #[allow(clippy::too_many_arguments)] + fn keep( + &self, + stream: &mut UnixStream, + caller: &LeaseCaller, + grant: &LeaseGrant, + invocation: &PeerUnit, + process: Option<&OwnedFd>, + expires_at_unix_ms: u64, + plan_quota: Vec, + ) -> io::Result<()> { + // A token that cannot be read or piped was never lent: the lease ends. let Some(read_end) = self .tokens .read() .and_then(|token| token_pipe(token.as_bytes()).ok()) else { - self.settle(caller, grant, LeaseUse::NeverLent)?; return Ok(()); }; let frame = encode_frame(&LeaseMessage::ClaudeLeaseGrant(LeaseGranted { lease_id: grant.lease_id.clone(), - reserved_micro_usd: grant.reserved_micro_usd, + plan_quota, })); - let sent = send_with_fd(&stream, &frame, read_end.as_raw_fd()); + let sent = send_with_fd(stream, &frame, read_end.as_raw_fd()); drop(read_end); - if sent.is_err() { - // The descriptor may have been delivered: the reservation is billed. - self.settle(caller, grant, LeaseUse::Reservation)?; - return Ok(()); - } - let outcome = self.await_settle(&mut stream, grant, expires_at_unix_ms); - match outcome { - Some(settle) => { - let settlement = self.settle(caller, grant, LeaseUse::Reported(&settle.usage))?; - let _ = stream.write_all(&encode_frame(&LeaseMessage::ClaudeLeaseSettled( - LeaseSettled { - billed_micro_usd: settlement.billed_micro_usd, - overrun_micro_usd: settlement.overrun_micro_usd, - }, - ))); + sent?; + // The holder says nothing more. The lease ends when it closes the + // socket, when its unit invocation is no longer the one that leased, + // or at the expiry. + let mut byte = [0_u8; 1]; + loop { + let left = expires_at_unix_ms.saturating_sub(unix_ms()); + if left == 0 { + return Ok(()); } - None => { - self.settle(caller, grant, LeaseUse::Reservation)?; + stream.set_read_timeout(Some(Duration::from_millis(left.min(HOLD_POLL_MS))))?; + match stream.read(&mut byte) { + Ok(0) => return Ok(()), + Ok(_) => {} + Err(error) + if matches!( + error.kind(), + io::ErrorKind::WouldBlock | io::ErrorKind::TimedOut + ) => {} + Err(error) if error.kind() == io::ErrorKind::Interrupted => {} + Err(_) => return Ok(()), + } + let still = process + .and_then(|process| self.units.invocation_in(process.as_fd(), &caller.unit())) + .is_some_and(|found| found == *invocation); + if !still { + return Ok(()); } } - Ok(()) - } - - /// The run's settle, if its whole frame arrives for this lease before the - /// lease expires. Anything else (end of file, silence, a frame still - /// trickling at the expiry, a malformed frame, the wrong lease) is no - /// report. - fn await_settle( - &self, - stream: &mut UnixStream, - grant: &LeaseGrant, - expires_at_unix_ms: u64, - ) -> Option { - let remaining = expires_at_unix_ms.saturating_sub(unix_ms()); - let deadline = Instant::now() + Duration::from_millis(remaining); - match read_message(stream, deadline).ok()? { - LeaseMessage::ClaudeLeaseSettle(settle) if settle.lease_id == grant.lease_id => { - Some(settle) - } - _ => None, - } - } - - fn settle( - &self, - caller: &LeaseCaller, - grant: &LeaseGrant, - usage: LeaseUse<'_>, - ) -> io::Result { - self.service - .lock() - .map_err(|_| io::Error::other("service lock"))? - .settle_lease(&caller.caller_id, &grant.lease_id, usage) - .map_err(|_| io::Error::other("ledger")) } } +/// How often a held lease looks again at its holder's unit invocation. +const HOLD_POLL_MS: u64 = 500; + fn refuse(stream: &mut UnixStream, kind: LeaseRefusalKind, detail: &str) -> io::Result<()> { stream.write_all(&encode_frame(&LeaseMessage::ClaudeLeaseRefusal( LeaseRefused { @@ -622,12 +619,12 @@ fn refuse(stream: &mut UnixStream, kind: LeaseRefusalKind, detail: &str) -> io:: mod tests { use super::*; - /// A settle frame whose body is exactly `body_bytes` long. - fn settle_frame_of(body_bytes: usize) -> Vec { + /// A refusal frame whose body is exactly `body_bytes` long. + fn frame_of(body_bytes: usize) -> Vec { let frame_with = |pad: usize| { - encode_frame(&LeaseMessage::ClaudeLeaseSettle(LeaseSettle { - lease_id: "x".repeat(pad), - usage: Vec::new(), + encode_frame(&LeaseMessage::ClaudeLeaseRefusal(LeaseRefused { + kind: LeaseRefusalKind::Busy, + detail: "x".repeat(pad), })) }; // The string header is the same width between 256 and 65535 bytes. @@ -646,10 +643,10 @@ mod tests { #[test] fn a_frame_is_one_whole_message_within_the_bound() { - let at_bound = settle_frame_of(MAX_LEASE_FRAME_BYTES); + let at_bound = frame_of(MAX_LEASE_FRAME_BYTES); assert_eq!(at_bound.len(), 4 + MAX_LEASE_FRAME_BYTES); assert!(decode_frame(&at_bound).is_some()); - let past_bound = settle_frame_of(MAX_LEASE_FRAME_BYTES + 1); + let past_bound = frame_of(MAX_LEASE_FRAME_BYTES + 1); assert!(decode_frame(&past_bound).is_none()); let mut trailing = at_bound.clone(); trailing.push(0); @@ -661,9 +658,9 @@ mod tests { #[test] fn a_stream_frame_over_the_bound_is_refused_before_it_is_read() { - let past_bound = settle_frame_of(MAX_LEASE_FRAME_BYTES + 1); + let past_bound = frame_of(MAX_LEASE_FRAME_BYTES + 1); assert!(within(&past_bound).is_err()); - let at_bound = settle_frame_of(MAX_LEASE_FRAME_BYTES); + let at_bound = frame_of(MAX_LEASE_FRAME_BYTES); assert!(within(&at_bound).is_ok()); assert!(within(&[0_u8, 0, 0, 0]).is_err()); // Short of the length it announced is end of file, not a message. @@ -673,10 +670,10 @@ mod tests { #[test] fn a_frame_that_trickles_past_its_deadline_is_refused_however_small_each_read() { let (mut reader, mut writer) = UnixStream::pair().unwrap(); - let frame = encode_frame(&LeaseMessage::ClaudeLeaseSettle(LeaseSettle { - lease_id: "lease:x".to_string(), - usage: Vec::new(), - })); + let frame = encode_frame(&LeaseMessage::ClaudeLeaseRefusal(LeaseRefused { + kind: LeaseRefusalKind::Busy, + detail: "lease:x".to_string(), + })); let trickle = thread::spawn(move || { for byte in frame { if writer.write_all(&[byte]).is_err() { diff --git a/crates/claude-connector/src/lib.rs b/crates/claude-connector/src/lib.rs index 6ed8757..4a9a42b 100644 --- a/crates/claude-connector/src/lib.rs +++ b/crates/claude-connector/src/lib.rs @@ -28,7 +28,7 @@ pub use launch::ChildEnv; #[cfg(target_os = "linux")] pub use lease::{ LeaseCaller, LeaseGranted, LeaseLimits, LeaseMessage, LeaseRefusalKind, LeaseRefused, - LeaseRequest, LeaseSettle, LeaseSettled, MAX_LEASE_FRAME_BYTES, PeerCredentials, PeerUnit, + LeaseRequest, MAX_LEASE_FRAME_BYTES, PeerCredentials, PeerUnit, PeerUnits, decode_frame, encode_frame, peer_credentials, serve_leases, user_uid, }; #[cfg(unix)] diff --git a/crates/claude-connector/src/main.rs b/crates/claude-connector/src/main.rs index ac435a1..273a4a1 100644 --- a/crates/claude-connector/src/main.rs +++ b/crates/claude-connector/src/main.rs @@ -18,7 +18,7 @@ use claude_connector::{ ClaudeConnectorConfig, ConfigError, LeaseCallerConfig, MIN_SECRET_BYTES, check_ready, load_config, write_config, }; -use model_connector::{CallerConfig, LeaseTurn, ModelPrice, Rates, frame_bytes_for_payload}; +use model_connector::{CallerConfig, LeaseTurn, frame_bytes_for_payload}; fn main() -> ExitCode { let args: Vec = std::env::args().skip(1).collect(); @@ -72,7 +72,6 @@ fn check(path: &str) -> Result<(), Failure> { println!("claude connector configuration is valid and ready to serve"); println!("callers: {}", config.callers.len()); println!("lease_callers: {}", config.lease_callers.len()); - println!("prices: {}", config.prices.len()); Ok(()) } @@ -85,17 +84,15 @@ struct Flags { /// The flags that describe one caller, shared by `config init` (which writes /// the first caller) and `config add-caller`. -const CALLER: [&str; 7] = [ +const CALLER: [&str; 5] = [ "--caller-id", "--connection-key-file", "--connection-key-epoch", - "--spend-ceiling-micro-usd", - "--spend-window-ms", "--max-concurrent-requests", "--max-output-tokens", ]; -const INIT: [&str; 22] = [ +const INIT: [&str; 21] = [ "--listen", "--replay-store", "--claude-binary", @@ -117,15 +114,9 @@ const INIT: [&str; 22] = [ "--lease-max-ms", "--turn", "--lease-unit-instance", - "--lease-reservation-micro-usd", ]; -const LEASE_CALLER: [&str; 4] = [ - "--caller-id", - "--turn", - "--lease-unit-instance", - "--lease-reservation-micro-usd", -]; +const LEASE_CALLER: [&str; 3] = ["--caller-id", "--turn", "--lease-unit-instance"]; impl Flags { fn parse( @@ -213,8 +204,6 @@ impl Flags { max_concurrent_requests: self.number_or("--max-concurrent-requests", 4)?, max_payload_bytes: self.number_or("--max-payload-bytes", 1 << 20)?, max_output_tokens: self.number_or("--max-output-tokens", 32_768)?, - spend_ceiling_micro_usd: self.number("--spend-ceiling-micro-usd")?, - spend_window_ms: self.number("--spend-window-ms")?, }) } } @@ -230,33 +219,6 @@ fn sha256(hex: &str) -> Option<[u8; 32]> { Some(digest) } -/// MODEL:INPUT:OUTPUT:CACHE_WRITE:CACHE_READ:OVERHEAD, rates in micro-USD per -/// million tokens. The one cache-write rate prices both retentions. -fn price(text: &str) -> Option<(String, ModelPrice)> { - let mut parts = text.split(':'); - let model = parts.next()?.to_string(); - let mut next = || parts.next()?.parse::().ok(); - let (input, output, cache_write, cache_read, overhead_tokens) = - (next()?, next()?, next()?, next()?, next()?); - if parts.next().is_some() { - return None; - } - Some(( - model, - ModelPrice { - rates: Rates { - input, - cache_write_short: cache_write, - cache_write_long: cache_write, - cache_read, - output, - }, - long_context: None, - overhead_tokens, - }, - )) -} - /// Writes 32 random bytes as 64 hex characters, mode 0600, unless the file is /// already there: the secret is never read back or shown, and a file that /// exists is the secret handles already depend on, so it is kept as it is when @@ -322,16 +284,9 @@ fn init(path: &str, args: &[&str]) -> Result<(), Failure> { let flags = Flags::parse( args, &[&INIT, &CALLER], - &["--allowed-model", "--price", "--lease-user"], + &["--allowed-model", "--lease-user"], )?; let first = flags.caller()?; - let mut prices = BTreeMap::new(); - for text in flags.list("--price") { - let (model, price) = price(text).ok_or(Failure::Malformed("--price"))?; - if prices.insert(model, price).is_some() { - return Err(Failure::Malformed("--price")); - } - } let lease_callers = flags .list("--lease-user") .iter() @@ -341,7 +296,6 @@ fn init(path: &str, args: &[&str]) -> Result<(), Failure> { caller_id: first.caller_runtime_id.clone(), turn: flags.turn()?, unit_instance: flags.text("--lease-unit-instance")?, - reservation_micro_usd: flags.number("--lease-reservation-micro-usd")?, }) }) .collect::, Failure>>()?; @@ -366,7 +320,6 @@ fn init(path: &str, args: &[&str]) -> Result<(), Failure> { socket_timeout_ms: flags.number_or("--socket-timeout-ms", 10_000)?, max_expiry_skew_ms: 300_000, callers: vec![first], - prices, lease_socket: flags.single.get("--lease-socket").map(PathBuf::from), lease_callers, lease_max_ms: flags.number_or("--lease-max-ms", 3_600_000)?, @@ -412,7 +365,6 @@ fn add_lease_caller(path: &str, args: &[&str]) -> Result<(), Failure> { caller_id: flags.text("--caller-id")?, turn: flags.turn()?, unit_instance: flags.text("--lease-unit-instance")?, - reservation_micro_usd: flags.number("--lease-reservation-micro-usd")?, }); write_config(path.as_ref(), &config)?; Ok(()) diff --git a/crates/claude-connector/src/quota.rs b/crates/claude-connector/src/quota.rs index e4ea4b3..450d3e2 100644 --- a/crates/claude-connector/src/quota.rs +++ b/crates/claude-connector/src/quota.rs @@ -1,7 +1,7 @@ //! The plan's own usage windows, as the CLI reports them. The subscription //! decides how much a plan may burn; the CLI's `rate_limit_event` is the only //! place that number appears, so the latest observation of each window is kept -//! here and read out beside the connector's own spend ledger. +//! here and published as the usage readout and lent with each lease. use std::collections::BTreeMap; use std::sync::Mutex; diff --git a/crates/claude-connector/src/stream.rs b/crates/claude-connector/src/stream.rs index f712892..c5e2e41 100644 --- a/crates/claude-connector/src/stream.rs +++ b/crates/claude-connector/src/stream.rs @@ -14,7 +14,7 @@ use crate::quota::QuotaCell; const STRUCTURED_OUTPUT_TOOL: &str = "StructuredOutput"; /// A call that cannot complete, with the stage that failed. `usage` is `None` -/// (the kit bills its reservation) unless a result event reported it or the +/// (unknown) unless a result event reported it or the /// call never reached the provider (zero). #[derive(Debug, PartialEq, Eq)] pub(crate) struct Failure { @@ -25,7 +25,7 @@ pub(crate) struct Failure { impl Failure { /// A failure after the child ran: what the provider billed is unknown, so - /// the kit bills its reservation. + /// the usage is unknown. pub(crate) fn cost_unknown(kind: FailureKind, message: &'static str) -> Self { Self { kind, @@ -356,7 +356,7 @@ fn fail(kind: FailureKind, message: &'static str) -> Step { /// provider cannot have billed (`empty_is_zero`: a refusal at the door, 401, /// 403 or 429, or a call never made); elsewhere the CLI also zeroes usage on /// crashes after real work. A call never made (`absent_is_zero`) may also -/// omit the field. `None` is billed the reservation by the kit. +/// omit the field. `None` is unknown usage. fn read_usage(reported: Option, empty_is_zero: bool, absent_is_zero: bool) -> Option { let Some(reported) = reported else { return absent_is_zero.then(|| sum_usage(std::iter::empty())); diff --git a/crates/claude-connector/tests/config_document.rs b/crates/claude-connector/tests/config_document.rs index 8ca6828..865ea00 100644 --- a/crates/claude-connector/tests/config_document.rs +++ b/crates/claude-connector/tests/config_document.rs @@ -1,29 +1,21 @@ //! The claude-connector's configuration document and binary: what the -//! document accepts, what it refuses, and that no refusal or output echoes a -//! byte of a value that was typed, stored or read from a credential. +//! document accepts, what it refuses, and that a document of the dollar era +//! fails closed. use std::collections::BTreeMap; use std::path::{Path, PathBuf}; use std::process::{Command, Output}; use std::sync::atomic::{AtomicUsize, Ordering}; -use std::time::Duration; use claude_connector::{ ClaudeConnectorConfig, ConfigError, LeaseCallerConfig, load_config, write_config, }; -use cultcache_rs::{ - CacheBackingStore, CultCache, CultCacheEnvelope, DatabaseEntry, - SingleFileMessagePackBackingStore, -}; -use model_connector::{CallerConfig, LeaseTurn, MIN_PAYLOAD_BYTES, ModelPrice, Rates}; +use cultcache_rs::{CultCache, DatabaseEntry, SingleFileMessagePackBackingStore}; +use model_connector::{CallerConfig, LeaseTurn}; -const CANARY: &str = "canary-9f3b7c1d"; const OPUS: &str = "claude-opus-5-1"; const HAIKU: &str = "claude-haiku-5-5"; -/// An edit that should make the valid document refuse. -type Change = Box; - static SEQUENCE: AtomicUsize = AtomicUsize::new(0); struct Scratch(PathBuf); @@ -49,20 +41,6 @@ impl Drop for Scratch { } } -fn price() -> ModelPrice { - ModelPrice { - rates: Rates { - input: 3_000_000, - cache_write_short: 3_750_000, - cache_write_long: 3_750_000, - cache_read: 300_000, - output: 15_000_000, - }, - long_context: None, - overhead_tokens: 100, - } -} - fn valid() -> ClaudeConnectorConfig { ClaudeConnectorConfig { listen: "127.0.0.1:7443".to_string(), @@ -91,17 +69,13 @@ fn valid() -> ClaudeConnectorConfig { max_concurrent_requests: 4, max_payload_bytes: 1 << 20, max_output_tokens: 32_768, - spend_ceiling_micro_usd: 50_000_000, - spend_window_ms: 3_600_000, }], - prices: BTreeMap::from([(HAIKU.to_string(), price()), (OPUS.to_string(), price())]), lease_socket: Some("/run/claude-connector/lease.sock".into()), lease_callers: vec![LeaseCallerConfig { unix_user: "mind-eureka-self".to_string(), caller_id: "persona-a-caller".to_string(), turn: LeaseTurn::SelfRun, unit_instance: "eureka".to_string(), - reservation_micro_usd: 5_000_000, }], lease_max_ms: 3_600_000, } @@ -124,509 +98,119 @@ fn config_round_trips_through_cultcache() { } #[test] -fn a_second_record_or_the_wrong_type_is_refused() { +fn a_valid_document_validates() { + assert!(valid().validate().is_ok()); +} + +/// A document of the dollar era: the first type id, carrying a price table. +#[derive(Clone, Debug, PartialEq, Eq, DatabaseEntry)] +#[cultcache(type = "gamecult.model_connector.claude_connector_config.v1")] +struct DollarEraConfig { + #[cultcache(key = 0)] + listen: String, + #[cultcache(key = 18)] + prices: BTreeMap, +} + +#[test] +fn a_dollar_era_config_is_refused_at_load() { let scratch = Scratch::new(); - let path = scratch.path("two.cc"); - let config = valid(); - write_config(&path, &config).unwrap(); - let mut cache = CultCache::new(); - cache - .register_entry_type::() - .unwrap(); - cache - .add_generic_backing_store(SingleFileMessagePackBackingStore::new(&path)) - .unwrap(); - cache.pull_all_backing_stores().unwrap(); - cache.put("extra", &config).unwrap(); + let path = scratch.path("old.cc"); + { + let mut cache = CultCache::new(); + cache.register_entry_type::().unwrap(); + cache + .add_generic_backing_store(SingleFileMessagePackBackingStore::new(&path)) + .unwrap(); + cache + .put( + "runtime".to_string(), + &DollarEraConfig { + listen: "127.0.0.1:7443".to_string(), + prices: BTreeMap::from([(OPUS.to_string(), 5_000_000)]), + }, + ) + .unwrap(); + } assert!(matches!(load_config(&path), Err(ConfigError::Store))); - - #[derive(Clone, Debug, PartialEq, Eq, DatabaseEntry)] - #[cultcache(type = "gamecult.other.document.v1")] - struct Other { - #[cultcache(key = 0)] - value: u32, - } - let other = scratch.path("other.cc"); - let mut cache = CultCache::new(); - cache.register_entry_type::().unwrap(); - cache - .add_generic_backing_store(SingleFileMessagePackBackingStore::new(&other)) - .unwrap(); - cache.pull_all_backing_stores().unwrap(); - cache.put("runtime", &Other { value: 1 }).unwrap(); - assert!(matches!(load_config(&other), Err(ConfigError::Store))); - - assert!(matches!( - load_config(&scratch.path("absent.cc")), - Err(ConfigError::Store) - )); -} - -fn store_raw(path: &Path, key: &str, ty: &str, schema: Option<&str>, payload: Vec) { - let mut store = SingleFileMessagePackBackingStore::new(path); - store - .push(&CultCacheEnvelope { - key: key.to_string(), - r#type: ty.to_string(), - payload, - stored_at: "2026-10-29T00:00:00Z".to_string(), - schema_id: schema.map(str::to_string), - }) - .unwrap(); } #[test] -fn the_envelope_must_be_exactly_the_documents_key_type_schema_and_encoding() { - let scratch = Scratch::new(); - let canonical = rmp_serde::to_vec(&valid()).unwrap(); - let ty = ClaudeConnectorConfig::TYPE; - let mut trailing = canonical.clone(); - trailing.push(0xc0); - - let cases = vec![ - ("wrong key", "other", ty, Some(ty), canonical.clone()), - ( - "wrong type", - "runtime", - "gamecult.other.v1", - Some(ty), - canonical.clone(), - ), - ( - "wrong schema", - "runtime", - ty, - Some("gamecult.other.v1"), - canonical.clone(), - ), - ("not canonical", "runtime", ty, Some(ty), trailing), - ]; - for (name, key, ty, schema, payload) in cases { - let path = scratch.path(&format!("{}.cc", name.replace(' ', "-"))); - store_raw(&path, key, ty, schema, payload); - assert!( - matches!(load_config(&path), Err(ConfigError::Store)), - "{name}" - ); - } - - // The same bytes with every part right load. - let path = scratch.path("right.cc"); - store_raw(&path, "runtime", ty, Some(ty), canonical); - assert_eq!(load_config(&path).unwrap(), valid()); -} - -#[test] -fn bounds_sit_exactly_where_the_connector_needs_them() { - let mut config = valid(); - config.max_total_bytes = config.max_line_bytes; - config.validate().unwrap(); - config.callers[0].max_payload_bytes = MIN_PAYLOAD_BYTES; - config.validate().unwrap(); - config.callers[0].max_payload_bytes = MIN_PAYLOAD_BYTES - 1; - assert_eq!(refused(&config), "caller bounds"); - - for set in [ - (|c: &mut ClaudeConnectorConfig| c.max_connections = 0) as fn(&mut ClaudeConnectorConfig), - |c| c.socket_timeout_ms = 0, - |c| c.max_expiry_skew_ms = 0, - |c| c.max_frame_bytes = 0, - ] { - let mut config = valid(); - set(&mut config); - assert_eq!(refused(&config), "serve limits"); - } - let mut config = valid(); - config.prices.insert("claude-".to_string(), price()); - assert_eq!(refused(&config), "prices"); -} - -#[test] -fn a_write_refuses_an_invalid_document_and_leaves_no_file() { - let scratch = Scratch::new(); - let path = scratch.path("bad.cc"); - let mut config = valid(); - config.callers[0].allowed_models = vec!["opus".to_string()]; - assert!(write_config(&path, &config).is_err()); - assert!(!path.exists()); -} - -#[test] -fn config_refuses_aliases_unpriced_models_and_second_callers() { - let valid_config = valid(); - assert_eq!(valid_config.validate().unwrap().port(), 7443); - - let cases: Vec<(&str, Change)> = vec![ - ( - "alias allowed", - Box::new(|c| { - c.callers[0].allowed_models = vec!["opus".to_string()]; - }), - ), - ( - "alias price key", - Box::new(|c| { - c.prices.insert("sonnet".to_string(), price()); - }), - ), - ( - "bare prefix", - Box::new(|c| { - c.callers[0].allowed_models = vec!["claude-".to_string()]; - }), - ), - ( - "family only", - Box::new(|c| { - c.callers[0].allowed_models = vec!["claude-haiku".to_string()]; - c.prices.insert("claude-haiku".to_string(), price()); - }), - ), - ( - "latest alias", - Box::new(|c| { - c.callers[0].allowed_models = vec!["claude-sonnet-latest".to_string()]; - c.prices.insert("claude-sonnet-latest".to_string(), price()); - }), - ), - ( - "uppercase family", - Box::new(|c| { - c.callers[0].allowed_models = vec!["claude-Opus-5-1".to_string()]; - c.prices.insert("claude-Opus-5-1".to_string(), price()); - }), - ), - ( - "short date", - Box::new(|c| { - c.callers[0].allowed_models = vec!["claude-opus-5-1-2026".to_string()]; - c.prices.insert("claude-opus-5-1-2026".to_string(), price()); - }), - ), - ( - "trailing part", - Box::new(|c| { - c.callers[0].allowed_models = vec!["claude-opus-5-1-20260101-x".to_string()]; - c.prices - .insert("claude-opus-5-1-20260101-x".to_string(), price()); - }), - ), - ( - "unpriced allowed", - Box::new(|c| { - c.prices.remove(OPUS); - }), - ), - ( - "duplicate allowed", - Box::new(|c| { - c.callers[0].allowed_models = vec![HAIKU.to_string(), HAIKU.to_string()]; - }), - ), - ( - "no allowed", - Box::new(|c| c.callers[0].allowed_models.clear()), - ), - ( - "duplicate caller id", - Box::new(|c| { - let again = c.callers[0].clone(); - c.callers.push(again); - }), - ), - ("no caller", Box::new(|c| c.callers.clear())), - ]; - let fields = [ - "allowed_models", - "prices", - "allowed_models", - "allowed_models", - "allowed_models", - "allowed_models", - "allowed_models", - "allowed_models", - "allowed_models", - "allowed_models", - "allowed_models", - "callers", - "callers", - ]; - for ((name, change), field) in cases.iter().zip(fields) { - let mut config = valid(); - change(&mut config); - assert_eq!(refused(&config), field, "{name}"); - } -} - -#[test] -fn a_priced_id_of_any_loose_form_is_allowed_and_the_price_table_decides() { - // Real ids come with a minor, without one, and with a date either way. - for id in [ - "claude-opus-5-1-20260101", +fn the_known_models_are_the_allowed_models_by_form_alone() { + for good in [ + "claude-opus-5-1", + "claude-haiku-5-5", "claude-opus-5", - "claude-opus-5-20260101", - "claude-sonnet-4-5", + "claude-opus-5-1-20260101", ] { let mut config = valid(); - config.callers[0].allowed_models = vec![id.to_string()]; - config.prices.insert(id.to_string(), price()); - config.validate().unwrap_or_else(|e| panic!("{id}: {e}")); - - // The same id with no price is not on the known list. - config.prices.remove(id); - assert_eq!(refused(&config), "allowed_models", "{id}"); + config.callers[0].allowed_models = vec![good.to_string()]; + assert!(config.validate().is_ok(), "{good}"); } + for bad in [ + "sonnet", + "claude-", + "claude-sonnet-latest", + "claude-Opus-5-1", + "claude-opus-5-1-2026", + "claude-opus-5-1-20260101-x", + ] { + let mut config = valid(); + config.callers[0].allowed_models = vec![bad.to_string()]; + assert_eq!(refused(&config), "allowed_models", "{bad}"); + } + let mut twice = valid(); + twice.callers[0].allowed_models = vec![OPUS.to_string(), OPUS.to_string()]; + assert_eq!(refused(&twice), "allowed_models"); + let mut none = valid(); + none.callers[0].allowed_models.clear(); + assert_eq!(refused(&none), "allowed_models"); } #[test] -fn a_price_with_a_zero_billed_component_is_refused() { - type Zero = fn(&mut Rates); - let zeroed: Vec<(&str, Zero)> = vec![ - ("input", |r| r.input = 0), - ("cache_write_short", |r| r.cache_write_short = 0), - ("cache_write_long", |r| r.cache_write_long = 0), - ("cache_read", |r| r.cache_read = 0), - ("output", |r| r.output = 0), +fn caller_bounds_are_required_nonzero() { + let edits: [fn(&mut ClaudeConnectorConfig); 3] = [ + |c| c.callers[0].max_concurrent_requests = 0, + |c| c.callers[0].max_output_tokens = 0, + |c| c.callers[0].max_payload_bytes = 0, ]; - for (name, zero) in &zeroed { + for edit in edits { let mut config = valid(); - zero(&mut config.prices.get_mut(HAIKU).unwrap().rates); - assert_eq!(refused(&config), "prices", "{name}"); - - let mut config = valid(); - let mut long = price().rates; - zero(&mut long); - config.prices.get_mut(HAIKU).unwrap().long_context = Some(model_connector::LongContext { - above_prompt_tokens: 200_000, - rates: long, - }); - assert_eq!(refused(&config), "prices", "long context {name}"); - } - // A zero overhead is not a rate. - let mut config = valid(); - config.prices.get_mut(HAIKU).unwrap().overhead_tokens = 0; - config.validate().unwrap(); -} - -#[test] -fn config_refuses_exposed_listeners_relative_paths_and_empty_bounds() { - let cases: Vec<(&str, &str, Change)> = vec![ - ( - "wildcard listen", - "listen", - Box::new(|c| c.listen = "0.0.0.0:7443".into()), - ), - ( - "lan listen", - "listen", - Box::new(|c| c.listen = "192.168.1.9:7443".into()), - ), - ( - "hostname listen", - "listen", - Box::new(|c| c.listen = "localhost:7443".into()), - ), - ( - "replay", - "replay_store", - Box::new(|c| c.replay_store = "replay.cc".into()), - ), - ( - "binary", - "claude_binary", - Box::new(|c| c.claude_binary = "claude".into()), - ), - ("work", "work_dir", Box::new(|c| c.work_dir = "work".into())), - ( - "config", - "config_dir", - Box::new(|c| c.config_dir = "cfg".into()), - ), - ( - "prompts", - "prompt_dir", - Box::new(|c| c.prompt_dir = "prompts".into()), - ), - ( - "key file", - "connection_key_file", - Box::new(|c| { - c.callers[0].connection_key_file = "caller.key".into(); - }), - ), - ( - "zero pin", - "claude_binary_sha256", - Box::new(|c| c.claude_binary_sha256 = [0; 32]), - ), - ( - "blank instance", - "instance", - Box::new(|c| c.instance = " ".into()), - ), - ( - "credential path", - "token_credential", - Box::new(|c| c.token_credential = "a/b".into()), - ), - ( - "credential parent", - "token_credential", - Box::new(|c| c.token_credential = "..".into()), - ), - ( - "credential empty", - "token_credential", - Box::new(|c| c.token_credential = String::new()), - ), - ( - "zero deadline", - "child limits", - Box::new(|c| c.deadline_ms = 0), - ), - ( - "total below line", - "child limits", - Box::new(|c| c.max_total_bytes = c.max_line_bytes - 1), - ), - ( - "zero frame", - "serve limits", - Box::new(|c| c.max_frame_bytes = 0), - ), - ( - "zero epoch", - "connection_key_epoch", - Box::new(|c| c.callers[0].connection_key_epoch = 0), - ), - ( - "tiny payload", - "caller bounds", - Box::new(|c| c.callers[0].max_payload_bytes = 16), - ), - ( - "no ceiling", - "caller bounds", - Box::new(|c| c.callers[0].spend_ceiling_micro_usd = 0), - ), - ]; - for (name, field, change) in cases { - let mut config = valid(); - change(&mut config); - assert_eq!(refused(&config), field, "{name}"); + edit(&mut config); + assert_eq!(refused(&config), "caller bounds"); } } #[test] -fn no_refusal_echoes_a_value() { - let canary_path: PathBuf = format!("{CANARY}/relative").into(); - let canaried: Vec = vec![ - Box::new(|c| c.listen = format!("{CANARY}:1")), - Box::new(|c| c.instance = format!(" {CANARY} ")), - Box::new(|c| c.token_credential = format!("{CANARY}/x")), - Box::new(|c| c.callers[0].allowed_models = vec![format!("{CANARY}-model")]), - Box::new(|c| { - c.prices.insert(format!("{CANARY}-alias"), price()); - }), - Box::new(|c| c.callers[0].caller_runtime_id = format!(" {CANARY}")), - Box::new({ - let path = canary_path.clone(); - move |c| c.replay_store = path.clone() - }), - Box::new({ - let path = canary_path.clone(); - move |c| c.callers[0].connection_key_file = path.clone() - }), - Box::new(|c| { - let mut second = c.callers[0].clone(); - second.caller_runtime_id = format!(" {CANARY}"); - c.callers.push(second); - }), - ]; - for change in canaried { - let mut config = valid(); - change(&mut config); - let error = config.validate().unwrap_err(); - for shown in [error.to_string(), format!("{error:?}")] { - assert!(!shown.contains(CANARY), "{shown}"); - } - let scratch = Scratch::new(); - let write = write_config(&scratch.path("x.cc"), &config).unwrap_err(); - assert!(!format!("{write} {write:?}").contains(CANARY)); - } +fn a_lease_caller_names_a_configured_caller_and_a_socket_exists_exactly_when_it_leases() { + let mut stranger = valid(); + stranger.lease_callers[0].caller_id = "nobody".to_string(); + assert_eq!(refused(&stranger), "lease_callers"); + + let mut no_callers = valid(); + no_callers.lease_callers.clear(); + assert_eq!(refused(&no_callers), "lease_callers"); + + let mut no_socket = valid(); + no_socket.lease_socket = None; + assert_eq!(refused(&no_socket), "lease_callers"); + no_socket.lease_callers.clear(); + assert!(no_socket.validate().is_ok()); + + let mut twice = valid(); + twice.lease_callers.push(twice.lease_callers[0].clone()); + assert_eq!(refused(&twice), "lease_callers"); } -// ---- the binary ---- - -/// How long any connector verb may take in these tests. `check` and `init` -/// return at once and `serve` is meant to refuse; a `serve` that starts -/// instead would run forever, so it fails a named test here, not the suite. -const VERB_DEADLINE: Duration = Duration::from_secs(20); - -fn run(args: &[&str], envs: &[(&str, &Path)]) -> Output { - let mut command = Command::new(env!("CARGO_BIN_EXE_claude-connector")); - command.args(args).env_clear(); - for (name, value) in envs { - command.env(name, value); - } - run_within(command, VERB_DEADLINE).unwrap_or_else(|| { - panic!( - "`{}` was still running after {VERB_DEADLINE:?}: it started", - args.join(" ") - ) - }) +fn run(args: &[&str]) -> Output { + Command::new(env!("CARGO_BIN_EXE_claude-connector")) + .args(args) + .env_clear() + .output() + .unwrap() } -/// The command's output, or `None` when it outlived `deadline` and was killed. -fn run_within(mut command: Command, deadline: Duration) -> Option { - use std::process::Stdio; - use std::time::Instant; - - let mut child = command - .stdin(Stdio::null()) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .spawn() - .unwrap(); - let started = Instant::now(); - while child.try_wait().unwrap().is_none() { - if started.elapsed() > deadline { - child.kill().unwrap(); - child.wait().unwrap(); - return None; - } - std::thread::sleep(Duration::from_millis(20)); - } - Some(child.wait_with_output().unwrap()) -} - -#[cfg(unix)] -#[test] -fn a_process_past_its_deadline_is_killed_and_reported_as_none() { - let mut forever = Command::new("sleep"); - forever.arg("60"); - let started = std::time::Instant::now(); - assert!(run_within(forever, Duration::from_millis(200)).is_none()); - assert!(started.elapsed() < Duration::from_secs(10)); - let mut quick = Command::new("true"); - quick.env_clear(); - assert!( - run_within(quick, Duration::from_secs(10)) - .unwrap() - .status - .success() - ); -} - -fn shown(output: &Output) -> String { - format!( - "{}{}", - String::from_utf8_lossy(&output.stdout), - String::from_utf8_lossy(&output.stderr) - ) -} - -fn init_flags<'a>(scratch: &'a Scratch, extra: &[&'a str]) -> Vec { +fn init_flags(scratch: &Scratch, extra: &[&str]) -> Vec { let path = |name: &str| scratch.path(name).to_string_lossy().into_owned(); let mut flags: Vec = [ "--listen", @@ -657,14 +241,8 @@ fn init_flags<'a>(scratch: &'a Scratch, extra: &[&'a str]) -> Vec { &path("caller.key"), "--connection-key-epoch", "1", - "--spend-ceiling-micro-usd", - "50000000", - "--spend-window-ms", - "3600000", "--allowed-model", HAIKU, - "--price", - &format!("{HAIKU}:1000000:5000000:1250000:100000:100"), ] .map(String::from) .to_vec(); @@ -676,1078 +254,32 @@ fn init(scratch: &Scratch, flags: &[String]) -> Output { let config = scratch.path("connector.cc"); let mut args = vec!["config", "init", config.to_str().unwrap()]; args.extend(flags.iter().map(String::as_str)); - run(&args, &[]) + run(&args) } #[test] fn config_init_writes_the_flags_into_the_document() { let scratch = Scratch::new(); let made = init(&scratch, &init_flags(&scratch, &[])); - assert!(made.status.success(), "{}", shown(&made)); + assert!(made.status.success(), "{:?}", made); let config = load_config(&scratch.path("connector.cc")).unwrap(); assert_eq!(config.callers[0].allowed_models, [HAIKU]); - assert_eq!(config.prices[HAIKU].rates.output, 5_000_000); - assert_eq!(config.prices[HAIKU].rates.cache_write_long, 1_250_000); - assert_eq!(config.prices[HAIKU].overhead_tokens, 100); assert_eq!(config.claude_binary_sha256, [0xab; 32]); + assert!(config.lease_callers.is_empty()); } -/// A scratch directory in which `serve` would start: every directory, the -/// pinned binary, the connection key and the credential exist. Returns the -/// `config init` flags and the credential directory. -#[cfg(unix)] -fn prepare(scratch: &Scratch) -> (Vec, PathBuf) { - use sha2::{Digest, Sha256}; - use std::os::unix::fs::PermissionsExt; - - for name in ["work", "home", "prompts", "credentials"] { - std::fs::create_dir_all(scratch.path(name)).unwrap(); - } - std::fs::set_permissions( - scratch.path("prompts"), - std::fs::Permissions::from_mode(0o700), - ) - .unwrap(); - std::fs::write(scratch.path("credentials/claude-token"), CANARY).unwrap(); - std::fs::write( - scratch.path("credentials/claude-handle-secret"), - "s".repeat(64), - ) - .unwrap(); - std::fs::write(scratch.path("caller.key"), "caller-secret").unwrap(); - std::fs::write(scratch.path("claude"), b"pinned claude").unwrap(); - let sha: String = Sha256::digest(b"pinned claude") - .iter() - .map(|byte| format!("{byte:02x}")) - .collect(); - let mut flags = init_flags(scratch, &[]); - let at = flags - .iter() - .position(|flag| flag == "--claude-sha256") - .unwrap(); - flags[at + 1] = sha; - (flags, scratch.path("credentials")) -} - -#[cfg(unix)] -fn check_and_serve(config: &Path, credentials: Option<&Path>) -> (Output, Output) { - let envs: Vec<(&str, &Path)> = credentials - .map(|dir| ("CREDENTIALS_DIRECTORY", dir)) - .into_iter() - .collect(); - let path = config.to_str().unwrap(); - ( - run(&["config", "check", path], &envs), - run(&["serve", "--config", path], &envs), - ) -} - -#[cfg(unix)] #[test] -fn check_passes_a_document_serve_can_start_and_reports_counts_only() { - let scratch = Scratch::new(); - let (flags, credentials) = prepare(&scratch); - let made = init(&scratch, &flags); - assert!(made.status.success(), "{}", shown(&made)); - let checked = run( - &[ - "config", - "check", - scratch.path("connector.cc").to_str().unwrap(), - ], - &[("CREDENTIALS_DIRECTORY", &credentials)], - ); - assert!(checked.status.success(), "{}", shown(&checked)); - let text = shown(&checked); - assert!(text.contains("callers: 1") && text.contains("lease_callers: 0")); - assert!(text.contains("prices: 1")); - for secret in ["persona-a", HAIKU, CANARY, "caller-secret"] { - assert!(!text.contains(secret), "{text}"); - } -} - -#[cfg(unix)] -#[test] -fn check_refuses_exactly_what_serve_refuses_and_names_the_field() { - use std::os::unix::fs::PermissionsExt; - - type Break = fn(&Scratch, &Path); - let cases: Vec<(&str, &str, Break)> = vec![ - ("open prompt dir", "prompt directory", |s, _| { - std::fs::set_permissions(s.path("prompts"), std::fs::Permissions::from_mode(0o755)) - .unwrap(); - }), - ("no key file", "connection_key_file", |s, _| { - std::fs::remove_file(s.path("caller.key")).unwrap(); - }), - ("binary off its pin", "claude_binary_sha256", |s, _| { - std::fs::write(s.path("claude"), b"other bytes").unwrap(); - }), - ("no credential", "token_credential", |_, credentials| { - std::fs::remove_file(credentials.join("claude-token")).unwrap(); - }), - ("no work dir", "work_dir", |s, _| { - std::fs::remove_dir(s.path("work")).unwrap(); - }), - ("no config dir", "config_dir", |s, _| { - std::fs::remove_dir(s.path("home")).unwrap(); - }), - ("frame too small", "max_frame_bytes", |s, _| { - let path = s.path("connector.cc"); - let mut config = load_config(&path).unwrap(); - config.max_frame_bytes = 4099; - write_config(&path, &config).unwrap(); - }), - ("two callers on one key", "callers", |s, _| { - std::fs::write(s.path("twin.key"), "caller-secret").unwrap(); - let config = s.path("connector.cc"); - let twin = s.path("twin.key"); - let added = run( - &[ - "config", - "add-caller", - config.to_str().unwrap(), - "--caller-id", - "twin", - "--connection-key-file", - twin.to_str().unwrap(), - "--connection-key-epoch", - "1", - "--allowed-model", - HAIKU, - "--spend-ceiling-micro-usd", - "1000", - "--spend-window-ms", - "60000", - ], - &[], - ); - assert!(added.status.success(), "{}", shown(&added)); - }), - ("too few connections", "connection capacity", |s, _| { - let path = s.path("connector.cc"); - let mut config = load_config(&path).unwrap(); - config.max_connections = 1; - write_config(&path, &config).unwrap(); - }), +fn the_dollar_flags_are_not_flags() { + let dollar_flags: [&[&str]; 4] = [ + &["--spend-ceiling-micro-usd", "50000000"], + &["--spend-window-ms", "3600000"], + &["--price", "claude-haiku-5-5:1:1:1:1:1"], + &["--lease-reservation-micro-usd", "5000000"], ]; - for (name, field, break_it) in cases { + for extra in dollar_flags { let scratch = Scratch::new(); - let (flags, credentials) = prepare(&scratch); - let made = init(&scratch, &flags); - assert!(made.status.success(), "{name}: {}", shown(&made)); - break_it(&scratch, &credentials); - let (checked, served) = check_and_serve(&scratch.path("connector.cc"), Some(&credentials)); - assert!(!checked.status.success(), "{name} check passed"); - assert!(!served.status.success(), "{name} serve started"); - assert!( - shown(&checked).contains(field), - "{name}: {}", - shown(&checked) - ); - assert_eq!(shown(&checked), shown(&served), "{name}"); - for secret in [CANARY, "caller-secret", "other bytes"] { - assert!(!shown(&checked).contains(secret), "{name}"); - } - } - - // No credential directory at all is the same refusal. - let scratch = Scratch::new(); - let (flags, _) = prepare(&scratch); - assert!(init(&scratch, &flags).status.success()); - let (checked, served) = check_and_serve(&scratch.path("connector.cc"), None); - assert!(!checked.status.success() && !served.status.success()); - assert!(shown(&checked).contains("credentials_directory")); - assert!(!shown(&checked).contains("token_credential")); - assert_eq!(shown(&checked), shown(&served)); - - // A credential directory without the credential is the other fault, and - // the token is checked last: a broken pin is named first. - let scratch = Scratch::new(); - let (flags, credentials) = prepare(&scratch); - assert!(init(&scratch, &flags).status.success()); - std::fs::remove_file(credentials.join("claude-token")).unwrap(); - std::fs::write(scratch.path("claude"), b"other bytes").unwrap(); - let (checked, _) = check_and_serve(&scratch.path("connector.cc"), Some(&credentials)); - assert!( - shown(&checked).contains("claude_binary_sha256"), - "{}", - shown(&checked) - ); - std::fs::write(scratch.path("claude"), b"pinned claude").unwrap(); - let (checked, _) = check_and_serve(&scratch.path("connector.cc"), Some(&credentials)); - assert!( - shown(&checked).contains("token_credential"), - "{}", - shown(&checked) - ); - assert!(!shown(&checked).contains("credentials_directory")); -} - -#[test] -fn usage_says_check_runs_in_the_units_environment() { - let output = run(&["nonsense"], &[]); - assert!(!output.status.success()); - assert!( - shown(&output).contains("unit's environment"), - "{}", - shown(&output) - ); -} - -#[cfg(unix)] -#[test] -fn serve_names_a_corrupt_replay_store_without_its_path() { - let scratch = Scratch::new(); - let (mut flags, credentials) = prepare(&scratch); - let store = scratch.path(&format!("{CANARY}-store/replay.cc")); - std::fs::create_dir_all(store.parent().unwrap()).unwrap(); - std::fs::write(&store, b"not a cultcache store").unwrap(); - let at = flags - .iter() - .position(|flag| flag == "--replay-store") - .unwrap(); - flags[at + 1] = store.to_string_lossy().into_owned(); - assert!(init(&scratch, &flags).status.success()); - let (_, served) = check_and_serve(&scratch.path("connector.cc"), Some(&credentials)); - assert!(!served.status.success()); - let text = shown(&served); - assert!(text.contains("replay_store"), "{text}"); - assert!( - !text.contains(CANARY) && !text.contains("replay.cc"), - "{text}" - ); -} - -#[test] -fn the_binary_refuses_without_echoing_what_it_was_given() { - let spaced = format!(" {CANARY}"); - // Each pair replaces the value of an existing flag; the others are appended. - let replaced: Vec<(&str, &str)> = vec![ - ("--listen", CANARY), - ("--claude-sha256", CANARY), - ("--claude-sha256", "abcd"), - ("--price", CANARY), - ("--instance", &spaced), - ("--connection-key-epoch", CANARY), - ("--allowed-model", CANARY), - ]; - let appended: Vec> = vec![ - vec![CANARY, CANARY], - vec!["--listen"], - vec!["--listen", "127.0.0.1:1"], - ]; - let mut attempts: Vec> = Vec::new(); - for (name, value) in replaced { - let scratch = Scratch::new(); - let mut flags = init_flags(&scratch, &[]); - let at = flags.iter().position(|flag| flag == name).unwrap(); - flags[at + 1] = value.to_string(); - attempts.push(flags); - } - for extra in appended { - let scratch = Scratch::new(); - attempts.push(init_flags(&scratch, &extra)); - } - for flags in attempts { - let scratch = Scratch::new(); - let made = init(&scratch, &flags); - assert!(!made.status.success(), "{flags:?}"); - assert!(!shown(&made).contains(CANARY), "{}", shown(&made)); - assert!(!scratch.path("connector.cc").exists()); - } - - // A missing flag is refused by name. - let scratch = Scratch::new(); - let mut flags = init_flags(&scratch, &[]); - let at = flags.iter().position(|f| f == "--instance").unwrap(); - flags.drain(at..at + 2); - let made = init(&scratch, &flags); - assert!(!made.status.success()); - assert!(shown(&made).contains("--instance")); -} - -#[cfg(unix)] -#[test] -fn serve_refuses_an_unsafe_prompt_dir_and_prints_no_credential() { - use std::os::unix::fs::{PermissionsExt, symlink}; - - let scratch = Scratch::new(); - let credentials = scratch.path("credentials"); - std::fs::create_dir_all(&credentials).unwrap(); - std::fs::write(credentials.join("claude-token"), CANARY).unwrap(); - std::fs::write(scratch.path("caller.key"), "caller-secret").unwrap(); - let made = init(&scratch, &init_flags(&scratch, &[])); - assert!(made.status.success(), "{}", shown(&made)); - let config = scratch.path("connector.cc"); - let serve = |config: &Path| { - run( - &["serve", "--config", config.to_str().unwrap()], - &[("CREDENTIALS_DIRECTORY", &credentials)], - ) - }; - - // No prompt directory, then one others can enter, then a link to a good one. - let prompts = scratch.path("prompts"); - let refused = serve(&config); - assert!(!refused.status.success()); - assert!( - shown(&refused).contains("prompt directory"), - "{}", - shown(&refused) - ); - std::fs::create_dir_all(&prompts).unwrap(); - std::fs::set_permissions(&prompts, std::fs::Permissions::from_mode(0o755)).unwrap(); - let open = serve(&config); - assert!( - shown(&open).contains("prompt directory"), - "{}", - shown(&open) - ); - assert!(claude_connector::check_prompt_dir(&prompts).is_err()); - std::fs::set_permissions(&prompts, std::fs::Permissions::from_mode(0o700)).unwrap(); - assert!(claude_connector::check_prompt_dir(&prompts).is_ok()); - let linked = scratch.path("linked"); - symlink(&prompts, &linked).unwrap(); - assert!(claude_connector::check_prompt_dir(&linked).is_err()); - - // A safe directory but no claude binary: the backend is not ready, and - // nothing the credential holds is printed. - let not_ready = serve(&config); - assert!(!not_ready.status.success()); - for output in [&refused, &open, ¬_ready] { - let text = shown(output); - assert!( - !text.contains(CANARY) && !text.contains("caller-secret"), - "{text}" - ); + let made = init(&scratch, &init_flags(&scratch, extra)); + assert!(!made.status.success(), "{extra:?} was accepted"); + assert!(!Path::new(&scratch.path("connector.cc")).exists(), "{extra:?}"); } } - -#[cfg(unix)] -#[test] -fn a_private_regular_file_is_not_a_prompt_dir() { - use std::os::unix::fs::PermissionsExt; - let scratch = Scratch::new(); - let file = scratch.path("file"); - std::fs::write(&file, "x").unwrap(); - std::fs::set_permissions(&file, std::fs::Permissions::from_mode(0o700)).unwrap(); - assert!(claude_connector::check_prompt_dir(&file).is_err()); - assert!(claude_connector::check_prompt_dir(&scratch.path("absent")).is_err()); -} - -#[test] -fn leases_are_configured_together_or_not_at_all() { - let cases: Vec<(&str, Change)> = vec![ - ( - "lease_socket", - Box::new(|c| c.lease_socket = Some("lease.sock".into())), - ), - ("lease_callers", Box::new(|c| c.lease_callers.clear())), - ("lease_callers", Box::new(|c| c.lease_socket = None)), - ( - "lease_callers", - Box::new(|c| c.lease_callers[0].caller_id = "another-caller".to_string()), - ), - ( - "lease_callers", - Box::new(|c| c.lease_callers[0].unix_user = " ".to_string()), - ), - ( - "lease_callers", - Box::new(|c| { - let again = c.lease_callers[0].clone(); - c.lease_callers.push(again); - }), - ), - ("lease_max_ms", Box::new(|c| c.lease_max_ms = 0)), - ( - "lease_callers", - Box::new(|c| c.lease_callers[0].reservation_micro_usd = 0), - ), - ( - "lease_callers", - Box::new(|c| c.lease_callers[0].reservation_micro_usd = 50_000_001), - ), - ( - "lease_callers", - Box::new(|c| c.lease_callers[0].unit_instance = String::new()), - ), - ( - "lease_callers", - Box::new(|c| c.lease_callers[0].unit_instance = "a/b".to_string()), - ), - ]; - for (field, change) in cases { - let mut config = valid(); - change(&mut config); - assert_eq!(refused(&config), field); - } - // The whole ceiling is a reservation the caller can hold. - let mut whole = valid(); - whole.lease_callers[0].reservation_micro_usd = 50_000_000; - whole.validate().unwrap(); - let mut none = valid(); - none.lease_socket = None; - none.lease_callers.clear(); - none.validate().unwrap(); -} - -#[test] -fn config_init_binds_each_lease_user_to_the_configured_caller() { - let scratch = Scratch::new(); - let made = init( - &scratch, - &init_flags( - &scratch, - &[ - "--lease-socket", - "/run/claude-connector/lease.sock", - "--lease-user", - "mind-eureka-self", - "--lease-max-ms", - "120000", - "--turn", - "self", - "--lease-unit-instance", - "eureka", - "--lease-reservation-micro-usd", - "5000000", - ], - ), - ); - assert!(made.status.success(), "{}", shown(&made)); - let config = load_config(&scratch.path("connector.cc")).unwrap(); - assert_eq!( - config.lease_callers, - [LeaseCallerConfig { - unix_user: "mind-eureka-self".to_string(), - caller_id: "persona-a-caller".to_string(), - turn: LeaseTurn::SelfRun, - unit_instance: "eureka".to_string(), - reservation_micro_usd: 5_000_000, - }] - ); - assert_eq!(config.lease_max_ms, 120_000); -} - -#[cfg(target_os = "linux")] -#[test] -fn config_check_refuses_a_lease_account_the_system_does_not_have() { - let scratch = Scratch::new(); - let (mut flags, credentials) = prepare(&scratch); - let socket = scratch.path("lease.sock"); - flags.extend( - [ - "--lease-socket", - socket.to_str().unwrap(), - "--lease-user", - "no-such-account-4f1c", - "--turn", - "self", - "--lease-unit-instance", - "eureka", - "--lease-reservation-micro-usd", - "5000000", - ] - .map(String::from), - ); - let made = init(&scratch, &flags); - assert!(made.status.success(), "{}", shown(&made)); - let config = scratch.path("connector.cc"); - let checked = run( - &["config", "check", config.to_str().unwrap()], - &[("CREDENTIALS_DIRECTORY", &credentials)], - ); - assert!(!checked.status.success()); - assert!( - shown(&checked).contains("lease_callers"), - "{}", - shown(&checked) - ); - assert!(!shown(&checked).contains("no-such-account")); -} - -#[cfg(target_os = "linux")] -#[test] -fn serve_binds_the_lease_socket_and_answers_a_stranger_with_a_refusal() { - use std::io::{Read, Write}; - use std::os::unix::fs::PermissionsExt; - use std::os::unix::net::UnixStream; - use std::process::Stdio; - use std::time::Instant; - - use claude_connector::{ - LeaseMessage, LeaseRefusalKind, LeaseRefused, LeaseRequest, decode_frame, encode_frame, - }; - - let scratch = Scratch::new(); - let (mut flags, credentials) = prepare(&scratch); - let socket = scratch.path("lease.sock"); - // The test runs as someone other than `nobody`, so it is a stranger. - flags.extend( - [ - "--lease-socket", - socket.to_str().unwrap(), - "--lease-user", - "nobody", - "--turn", - "self", - "--lease-unit-instance", - "eureka", - "--lease-reservation-micro-usd", - "5000000", - ] - .map(String::from), - ); - let made = init(&scratch, &flags); - assert!(made.status.success(), "{}", shown(&made)); - let config = scratch.path("connector.cc"); - let mut server = Command::new(env!("CARGO_BIN_EXE_claude-connector")) - .args(["serve", "--config", config.to_str().unwrap()]) - .env_clear() - .env("CREDENTIALS_DIRECTORY", &credentials) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .spawn() - .unwrap(); - let started = Instant::now(); - while !socket.exists() && started.elapsed() < Duration::from_secs(10) { - std::thread::sleep(Duration::from_millis(50)); - } - let outcome = std::panic::catch_unwind(|| { - assert_eq!( - std::fs::metadata(&socket).unwrap().permissions().mode() & 0o777, - 0o660 - ); - let mut stream = UnixStream::connect(&socket).unwrap(); - stream - .set_read_timeout(Some(Duration::from_secs(5))) - .unwrap(); - stream - .write_all(&encode_frame(&LeaseMessage::ClaudeLeaseRequest( - LeaseRequest { - expires_at_unix_ms: u64::MAX / 2, - }, - ))) - .unwrap(); - let mut reply = vec![0_u8; 4096]; - let read = stream.read(&mut reply).unwrap(); - assert!(matches!( - decode_frame(&reply[..read]), - Some(LeaseMessage::ClaudeLeaseRefusal(LeaseRefused { - kind: LeaseRefusalKind::UnknownCaller, - .. - })) - )); - }); - // The server is still serving: it did not return when it started. - let running = server.try_wait().unwrap().is_none(); - server.kill().unwrap(); - server.wait().unwrap(); - outcome.unwrap(); - assert!(running); -} - -// ---- many callers ---- - -fn second_caller() -> CallerConfig { - CallerConfig { - caller_runtime_id: "heartbeat".to_string(), - connection_key_file: "/etc/claude-connector/heartbeat.key".into(), - connection_key_epoch: 3, - allowed_models: vec![HAIKU.to_string()], - max_concurrent_requests: 1, - max_payload_bytes: 64 << 10, - max_output_tokens: 1024, - spend_ceiling_micro_usd: 1_000_000, - spend_window_ms: 60_000, - } -} - -#[test] -fn several_callers_round_trip_and_each_is_validated_alone() { - let scratch = Scratch::new(); - let path = scratch.path("connector.cc"); - let mut config = valid(); - config.callers.push(second_caller()); - write_config(&path, &config).unwrap(); - assert_eq!(load_config(&path).unwrap(), config); - - let broken: Vec<(&str, Change)> = vec![ - ( - "caller bounds", - Box::new(|c| c.callers[1].spend_ceiling_micro_usd = 0), - ), - ( - "connection_key_epoch", - Box::new(|c| c.callers[1].connection_key_epoch = 0), - ), - ( - "connection_key_file", - Box::new(|c| c.callers[1].connection_key_file = "heartbeat.key".into()), - ), - ( - "allowed_models", - Box::new(|c| c.callers[1].allowed_models = vec!["claude-opus-9".to_string()]), - ), - ( - "callers", - Box::new(|c| c.callers[1].caller_runtime_id = c.callers[0].caller_runtime_id.clone()), - ), - ]; - for (field, change) in &broken { - let mut config = valid(); - config.callers.push(second_caller()); - change(&mut config); - assert_eq!(refused(&config), *field); - } -} - -#[test] -fn a_lease_caller_bills_one_configured_caller_within_that_callers_ceiling() { - let mut config = valid(); - config.callers.push(second_caller()); - let lease = |user: &str, caller: &str, instance: &str, reservation: u64| LeaseCallerConfig { - unix_user: user.to_string(), - caller_id: caller.to_string(), - turn: LeaseTurn::SelfRun, - unit_instance: instance.to_string(), - reservation_micro_usd: reservation, - }; - // A lease caller of each caller, each within its own caller's ceiling. - config - .lease_callers - .push(lease("mind-heartbeat", "heartbeat", "heartbeat", 1_000_000)); - config.validate().unwrap(); - // The first caller's ceiling does not stretch to the second's lease, nor the reverse. - config.lease_callers[1].reservation_micro_usd = 1_000_001; - assert_eq!(refused(&config), "lease_callers"); - config.lease_callers[1].reservation_micro_usd = 1; - config.lease_callers[0].reservation_micro_usd = 50_000_001; - assert_eq!(refused(&config), "lease_callers"); - config.lease_callers[0].reservation_micro_usd = 50_000_000; - // A lease caller naming no configured caller, or sharing another's account. - config.lease_callers[1].caller_id = "nobody".to_string(); - assert_eq!(refused(&config), "lease_callers"); - config.lease_callers[1].caller_id = "heartbeat".to_string(); - config.lease_callers[1].unix_user = config.lease_callers[0].unix_user.clone(); - assert_eq!(refused(&config), "lease_callers"); -} - -#[test] -fn add_caller_and_add_lease_caller_extend_the_document_and_grow_the_serve_limits() { - let scratch = Scratch::new(); - let made = init( - &scratch, - &init_flags( - &scratch, - &[ - "--lease-socket", - "/run/claude-connector/lease.sock", - "--lease-user", - "mind-eureka-self", - "--turn", - "self", - "--lease-unit-instance", - "eureka", - "--lease-reservation-micro-usd", - "5000000", - ], - ), - ); - assert!(made.status.success(), "{}", shown(&made)); - let path = scratch.path("connector.cc"); - let path_text = path.to_str().unwrap(); - let before = load_config(&path).unwrap(); - - let key = scratch.path("heartbeat.key"); - let added = run( - &[ - "config", - "add-caller", - path_text, - "--caller-id", - "heartbeat", - "--connection-key-file", - key.to_str().unwrap(), - "--connection-key-epoch", - "2", - "--allowed-model", - HAIKU, - "--spend-ceiling-micro-usd", - "2000000", - "--spend-window-ms", - "60000", - "--max-concurrent-requests", - "30", - "--max-payload-bytes", - "4194304", - ], - &[], - ); - assert!(added.status.success(), "{}", shown(&added)); - let leased = run( - &[ - "config", - "add-lease-caller", - path_text, - "--lease-user", - "mind-heartbeat", - "--caller-id", - "heartbeat", - "--turn", - "self", - "--lease-unit-instance", - "heartbeat", - "--lease-reservation-micro-usd", - "1000000", - ], - &[], - ); - assert!(leased.status.success(), "{}", shown(&leased)); - - let after = load_config(&path).unwrap(); - assert_eq!(after.callers[0], before.callers[0]); - assert_eq!(after.callers.len(), 2); - let heartbeat = &after.callers[1]; - assert_eq!(heartbeat.caller_runtime_id, "heartbeat"); - assert_eq!(heartbeat.connection_key_epoch, 2); - assert_eq!(heartbeat.spend_ceiling_micro_usd, 2_000_000); - assert_eq!(heartbeat.max_payload_bytes, 4 << 20); - assert_eq!(after.lease_callers[0], before.lease_callers[0]); - assert_eq!(after.lease_callers[1].caller_id, "heartbeat"); - // The frame carries the largest payload and the connections every slot. - assert!(after.max_frame_bytes > before.max_frame_bytes); - assert!(after.max_connections >= 34); - assert_eq!(after.prices, before.prices); - assert_eq!( - after.handle_secret_credential, - before.handle_secret_credential - ); - - // Repeating either verb is refused and leaves the document as it was. - let again = run( - &[ - "config", - "add-caller", - path_text, - "--caller-id", - "heartbeat", - "--connection-key-file", - key.to_str().unwrap(), - "--connection-key-epoch", - "2", - "--allowed-model", - HAIKU, - "--spend-ceiling-micro-usd", - "2000000", - "--spend-window-ms", - "60000", - ], - &[], - ); - assert!(!again.status.success()); - assert_eq!(load_config(&path).unwrap(), after); - let unleased = run( - &[ - "config", - "add-lease-caller", - path_text, - "--lease-user", - "mind-heartbeat", - "--caller-id", - "heartbeat", - "--turn", - "self", - "--lease-unit-instance", - "heartbeat", - "--lease-reservation-micro-usd", - "1000000", - ], - &[], - ); - assert!(!unleased.status.success()); - assert_eq!(load_config(&path).unwrap(), after); -} - -#[test] -fn add_lease_caller_needs_a_document_that_lends_a_token() { - let scratch = Scratch::new(); - let made = init(&scratch, &init_flags(&scratch, &[])); - assert!(made.status.success(), "{}", shown(&made)); - let path = scratch.path("connector.cc"); - let leased = run( - &[ - "config", - "add-lease-caller", - path.to_str().unwrap(), - "--lease-user", - "mind-heartbeat", - "--caller-id", - "persona-a-caller", - "--turn", - "self", - "--lease-unit-instance", - "heartbeat", - "--lease-reservation-micro-usd", - "1000000", - ], - &[], - ); - assert!(!leased.status.success()); - assert!(load_config(&path).unwrap().lease_callers.is_empty()); -} - -#[cfg(unix)] -#[test] -fn init_creates_the_handle_secret_once_private_and_never_shows_it() { - use std::os::unix::fs::PermissionsExt; - - let scratch = Scratch::new(); - let secret = scratch.path("handle.secret"); - let made = init(&scratch, &init_flags(&scratch, &[])); - assert!(made.status.success(), "{}", shown(&made)); - let metadata = std::fs::metadata(&secret).unwrap(); - assert_eq!(metadata.permissions().mode() & 0o777, 0o600); - let text = std::fs::read_to_string(&secret).unwrap(); - assert_eq!(text.len(), 64); - assert!(text.bytes().all(|b| b.is_ascii_hexdigit())); - assert!(!shown(&made).contains(&text)); - - // A second init keeps the secret the first one made, since handles depend on it. - std::fs::remove_file(scratch.path("connector.cc")).unwrap(); - let again = init(&scratch, &init_flags(&scratch, &[])); - assert!(again.status.success(), "{}", shown(&again)); - assert_eq!(std::fs::read_to_string(&secret).unwrap(), text); - - // A refused document leaves no secret file behind. - let fresh = Scratch::new(); - let flags = init_flags(&fresh, &["--max-payload-bytes", "1"]); - let refused = init(&fresh, &flags); - assert!(!refused.status.success()); - assert!(!fresh.path("handle.secret").exists()); -} - -#[cfg(unix)] -#[test] -fn check_wants_the_handle_secret_credential_and_names_it() { - let scratch = Scratch::new(); - let (flags, credentials) = prepare(&scratch); - let made = init(&scratch, &flags); - assert!(made.status.success(), "{}", shown(&made)); - let config = scratch.path("connector.cc"); - let path = config.to_str().unwrap(); - for body in [None, Some("too-short")] { - let secret = credentials.join("claude-handle-secret"); - match body { - None => std::fs::remove_file(&secret).unwrap(), - Some(text) => std::fs::write(&secret, text).unwrap(), - } - let checked = run( - &["config", "check", path], - &[("CREDENTIALS_DIRECTORY", &credentials)], - ); - assert!(!checked.status.success()); - let text = shown(&checked); - assert!(text.contains("handle_secret_credential"), "{text}"); - assert!(!text.contains("too-short"), "{text}"); - } -} - -/// A document written before the handle secret had a field cannot name where -/// the secret is, so it is not a document this connector serves from: it is -/// refused as a store fault, to be authored again with `config init`. -#[test] -fn a_document_without_the_handle_secret_field_is_refused_as_a_store() { - let scratch = Scratch::new(); - let path = scratch.path("old.cc"); - let mut payload = rmp_serde::to_vec(&valid()).unwrap(); - // array16(23) -> array16(22), and drop the trailing 20-byte fixstr. - assert_eq!(payload[..3], [0xdc, 0, 23]); - payload[2] = 22; - payload.truncate(payload.len() - 21); - store_raw( - &path, - "runtime", - ClaudeConnectorConfig::TYPE, - Some(ClaudeConnectorConfig::TYPE), - payload, - ); - assert!(matches!(load_config(&path), Err(ConfigError::Store))); -} - -#[test] -fn the_handle_secret_credential_is_a_name_of_its_own() { - for name in ["claude-token", "a/b", "", ".."] { - let mut config = valid(); - config.handle_secret_credential = name.to_string(); - assert_eq!(refused(&config), "handle_secret_credential", "{name:?}"); - } -} - -#[cfg(unix)] -#[test] -fn init_refuses_a_secret_file_it_cannot_create_and_writes_no_document() { - let scratch = Scratch::new(); - let mut flags = init_flags(&scratch, &[]); - let at = flags - .iter() - .position(|f| f == "--handle-secret-file") - .unwrap(); - flags[at + 1] = scratch - .path("absent/handle.secret") - .to_string_lossy() - .into_owned(); - let made = init(&scratch, &flags); - assert!(!made.status.success()); - assert!(!scratch.path("connector.cc").exists()); -} - -#[cfg(unix)] -#[test] -fn config_init_keeps_an_existing_secret_file_only_if_it_is_usable() { - use std::os::unix::fs::{PermissionsExt, symlink}; - - let write = |path: &Path, bytes: usize, mode: u32| { - std::fs::write(path, "k".repeat(bytes)).unwrap(); - std::fs::set_permissions(path, std::fs::Permissions::from_mode(mode)).unwrap(); - }; - type Plant = fn(&Scratch, &Path); - let refused: Vec<(&str, Plant)> = vec![ - ("world readable", |_, p| { - std::fs::write(p, "k".repeat(64)).unwrap(); - std::fs::set_permissions(p, std::fs::Permissions::from_mode(0o644)).unwrap(); - }), - ("group writable", |_, p| { - std::fs::write(p, "k".repeat(64)).unwrap(); - std::fs::set_permissions(p, std::fs::Permissions::from_mode(0o660)).unwrap(); - }), - ("read only", |_, p| { - std::fs::write(p, "k".repeat(64)).unwrap(); - std::fs::set_permissions(p, std::fs::Permissions::from_mode(0o400)).unwrap(); - }), - ("one byte short", |_, p| { - std::fs::write(p, "k".repeat(31)).unwrap(); - std::fs::set_permissions(p, std::fs::Permissions::from_mode(0o600)).unwrap(); - }), - ("empty", |_, p| { - std::fs::write(p, "").unwrap(); - std::fs::set_permissions(p, std::fs::Permissions::from_mode(0o600)).unwrap(); - }), - ("a directory", |_, p| { - std::fs::create_dir(p).unwrap(); - std::fs::set_permissions(p, std::fs::Permissions::from_mode(0o600)).unwrap(); - }), - ("a link to a good file", |s, p| { - let target = s.path("elsewhere.secret"); - std::fs::write(&target, "k".repeat(64)).unwrap(); - std::fs::set_permissions(&target, std::fs::Permissions::from_mode(0o600)).unwrap(); - symlink(&target, p).unwrap(); - }), - ]; - for (name, plant) in refused { - let scratch = Scratch::new(); - let flags = init_flags(&scratch, &[]); - let secret = scratch.path("handle.secret"); - plant(&scratch, &secret); - let before = std::fs::symlink_metadata(&secret).unwrap(); - let made = init(&scratch, &flags); - assert!(!made.status.success(), "{name} was accepted"); - assert!(shown(&made).contains("--handle-secret-file"), "{name}"); - assert!(!shown(&made).contains("kkkk"), "{name}"); - assert!( - !scratch.path("connector.cc").exists(), - "{name} left a document" - ); - let after = std::fs::symlink_metadata(&secret).unwrap(); - assert_eq!( - before.permissions().mode(), - after.permissions().mode(), - "{name}" - ); - assert_eq!(before.len(), after.len(), "{name}"); - assert_eq!( - before.file_type().is_symlink(), - after.file_type().is_symlink(), - "{name}" - ); - } - - // The smallest usable file is accepted and left exactly as it was. - let scratch = Scratch::new(); - let flags = init_flags(&scratch, &[]); - let secret = scratch.path("handle.secret"); - write(&secret, 32, 0o600); - let made = init(&scratch, &flags); - assert!(made.status.success(), "{}", shown(&made)); - assert_eq!(std::fs::read_to_string(&secret).unwrap(), "k".repeat(32)); - assert!(scratch.path("connector.cc").exists()); -} - -#[test] -fn a_persona_lease_row_validates_and_an_account_leases_for_one_turn_only() { - let mut config = valid(); - let mut persona = config.lease_callers[0].clone(); - persona.unix_user = "mind-eureka-persona".to_string(); - persona.turn = LeaseTurn::PersonaTurn; - config.lease_callers.push(persona.clone()); - config.validate().unwrap(); - // The same account under the other turn is a second row for one uid. - config.lease_callers[1].unix_user = config.lease_callers[0].unix_user.clone(); - assert_eq!(refused(&config), "lease_callers"); -} - -#[test] -fn add_lease_caller_takes_the_turn_and_refuses_an_unknown_one() { - let scratch = Scratch::new(); - let made = init( - &scratch, - &init_flags( - &scratch, - &[ - "--lease-socket", - "/run/claude-connector/lease.sock", - "--lease-user", - "mind-eureka-self", - "--turn", - "self", - "--lease-unit-instance", - "eureka", - "--lease-reservation-micro-usd", - "5000000", - ], - ), - ); - assert!(made.status.success(), "{}", shown(&made)); - let path = scratch.path("connector.cc"); - let add = |turn: &str| { - run( - &[ - "config", - "add-lease-caller", - path.to_str().unwrap(), - "--lease-user", - "mind-eureka-persona", - "--caller-id", - "persona-a-caller", - "--turn", - turn, - "--lease-unit-instance", - "eureka", - "--lease-reservation-micro-usd", - "1000000", - ], - &[], - ) - }; - let unknown = add("heartbeat"); - assert!(!unknown.status.success()); - assert_eq!(load_config(&path).unwrap().lease_callers.len(), 1); - let added = add("persona"); - assert!(added.status.success(), "{}", shown(&added)); - let config = load_config(&path).unwrap(); - assert_eq!(config.lease_callers[1].turn, LeaseTurn::PersonaTurn); - assert_eq!(config.lease_callers[0].turn, LeaseTurn::SelfRun); -} diff --git a/crates/claude-connector/tests/fixtures/lease_v1/grant.bin b/crates/claude-connector/tests/fixtures/lease_v1/grant.bin deleted file mode 100644 index 554e3a64cb35edbf81585cd63ecbe95744d580df..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 109 zcmZQzV90D*lAW2CvMD(yu{0$$J|{J?I5j@KC^4^OX<1@UX=>97h)`zAO*Jsf%D~Xb r*u>P#+`=+3DLExI4V}NKD784Xs4O)lJ~uPDC_lcmIOQAzr}_Z^$PFl0 diff --git a/crates/claude-connector/tests/fixtures/lease_v1/refusal.bin b/crates/claude-connector/tests/fixtures/lease_v1/refusal.bin deleted file mode 100644 index 22a4a25e3ceba32fc73bbfaf38dd854feff2550d..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 104 zcmXBL!3~2j3KM{AsLOPL()t?<|E=k6otQfW-a1l#;wZPj&fC#$7NP yxM8fAdt=)cwvG#YhZkJ#P*Z>Gf?&1(Ei!05ds%5>_Y#zg2}RaM6%BRa(BcQ7qADT) diff --git a/crates/claude-connector/tests/fixtures/lease_v1/request.bin b/crates/claude-connector/tests/fixtures/lease_v1/request.bin deleted file mode 100644 index fddaa3785fe2e36c3d51f2ec7bb63649d2decda3..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 66 zcmZQzV6bailAW2CvL!huu{0$$J|{J?I5oZ~wXig`xMXQrVoqsl3W!){%1t#e%gVse z$k@cx%-q51Y}lBbn4FOsUzD1d0uwAX*aQ(QFUl-Qg^4nLS({&4f~1#OVFOTCaeiKA lUOH4!3e@fmP`fh{GqX!^TFy8D({ctZRx_YlY=dF#0RUJliN*i` diff --git a/crates/claude-connector/tests/fixtures/lease_v1/settled.bin b/crates/claude-connector/tests/fixtures/lease_v1/settled.bin deleted file mode 100644 index 710d4b817e73b57863b3e0adaf6b682b017db04c..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 81 zcmZQzVDN2PlAW2CvL!huu{0$$J|{J?I5oaFwWK5`HDzg8Voqsl(}twX93U?~H#4~? bKfbg$>, + quota: Arc, tokens: Arc, /// The fake behind the listener, when it has one. units: Option>, @@ -126,129 +120,6 @@ impl Drop for Harness { let _ = std::fs::remove_dir_all(&self.dir); } } - -fn price() -> ModelPrice { - // One micro-USD per token in every class. - ModelPrice { - rates: Rates { - input: 1_000_000, - cache_write_short: 1_000_000, - cache_write_long: 1_000_000, - cache_read: 1_000_000, - output: 1_000_000, - }, - long_context: None, - overhead_tokens: 0, - } -} - -/// What a harness is built from. -struct Setup { - ceiling: u64, - reservation: u64, - uid: u32, - turn: LeaseTurn, - token: Option<&'static str>, - max_connections: usize, -} - -fn setup(ceiling: u64, reservation: u64) -> Setup { - Setup { - ceiling, - reservation, - uid: own_uid(), - turn: LeaseTurn::SelfRun, - token: Some(TOKEN), - max_connections: 8, - } -} - -/// A listener serving `uid` as the caller `eureka`, each of whose leases -/// reserves `reservation`, with a spend window so long that it never rolls -/// over during a test. Peers start in invocation 1 of the Self unit. -fn start(setup: Setup) -> Harness { - let units = Arc::new(FakeUnits { - peer: Mutex::new(None), - asked_about: Mutex::new(Vec::new()), - }); - let harness = start_with(setup, units.clone(), Some(units), INSTANCE); - harness.enter(1); - harness -} - -fn start_with( - setup: Setup, - units: Arc, - fake: Option>, - instance: &str, -) -> Harness { - let dir = std::env::temp_dir().join(format!( - "claude-lease-{}-{}", - std::process::id(), - SEQUENCE.fetch_add(1, Ordering::SeqCst) - )); - std::fs::create_dir_all(&dir).unwrap(); - let socket = dir.join("lease.sock"); - let service = Arc::new(Mutex::new( - TransportService::open( - &dir.join("replay.cc"), - "claude-code", - BTreeMap::from([(MODEL.to_string(), price())]), - [CallerAdmission::new( - CALLER, - "lease-test-connection-key", - 1, - [MODEL.to_string()], - 1, - 65_536, - 32_768, - SpendCeiling(setup.ceiling), - SpendWindowMs(1 << 50), - ) - .unwrap()], - 300_000, - ) - .unwrap(), - )); - let tokens = Arc::new(FakeTokens { - token: setup.token, - reads: AtomicUsize::new(0), - probe: Mutex::new(None), - }); - let listener = UnixListener::bind(&socket).unwrap(); - let serving = (service.clone(), tokens.clone(), units); - let caller = LeaseCaller { - uid: setup.uid, - caller_id: CALLER.to_string(), - turn: setup.turn, - unit_instance: instance.to_string(), - reservation_micro_usd: setup.reservation, - }; - std::thread::spawn(move || { - let _ = serve_leases( - &listener, - vec![caller], - LeaseLimits { - socket_timeout: Duration::from_secs(2), - max_lease_ms: 3_600_000, - max_connections: setup.max_connections, - }, - serving.0, - serving.1, - serving.2, - ); - }); - Harness { - dir, - socket, - service, - tokens, - units: fake, - } -} - -/// One frame, and the descriptors that arrived with it. An empty frame is end -/// of file. fn receive(stream: &UnixStream) -> (Vec, Vec) { use std::os::fd::AsRawFd; let mut buffer = vec![0_u8; 70_000]; @@ -357,49 +228,6 @@ fn granted_ok(answer: &Answer) -> bool { } /// Asks, as a fresh unit invocation each time, until the caller's previous -/// lease has been settled by the listener. -fn ask_after_busy(harness: &Harness, first_invocation: u64) -> Answer { - let deadline = Instant::now() + Duration::from_secs(5); - for attempt in 0.. { - harness.enter(first_invocation + attempt); - let answer = ask(harness, 60_000); - if refusal_of(&answer) != Some(LeaseRefusalKind::Busy) || Instant::now() > deadline { - return answer; - } - std::thread::sleep(Duration::from_millis(20)); - } - unreachable!() -} - -fn used(model: &str, tokens: u64) -> ModelUsage { - ModelUsage { - model: model.to_string(), - usage: Usage { - input_tokens: 0, - cache_read_tokens: 0, - cache_write_tokens: 0, - output_tokens: tokens, - reasoning_tokens: None, - }, - } -} - -fn settle_frame(answer: &Answer, usage: Vec) -> Vec { - encode_frame(&LeaseMessage::ClaudeLeaseSettle(LeaseSettle { - lease_id: granted(answer).lease_id, - usage, - })) -} - -fn settle(answer: &mut Answer, usage: Vec) -> LeaseSettled { - let frame = settle_frame(answer, usage); - answer.stream.write_all(&frame).unwrap(); - let (raw, _) = receive(&answer.stream); - match decode_frame(&raw) { - Some(LeaseMessage::ClaudeLeaseSettled(settled)) => settled, - other => panic!("expected a settled reply, got {other:?}"), - } -} fn contains(haystack: &[u8], needle: &str) -> bool { haystack @@ -407,917 +235,311 @@ fn contains(haystack: &[u8], needle: &str) -> bool { .any(|window| window == needle.as_bytes()) } -/// Asks the ledger directly for a reservation of `cost`, as an invocation -/// nobody else uses. A grant is left open. -fn probe_reserve(harness: &Harness, name: &str, cost: u64) -> LeaseAdmission { - harness - .service - .lock() - .unwrap() - .begin_lease(CALLER, LeaseTurn::SelfRun, name, cost, FAR, now_ms()) - .unwrap() +/// What a harness is built from. +struct Setup { + uid: u32, + turn: LeaseTurn, + token: Option<&'static str>, + max_connections: usize, } -fn ceiling_refused() -> LeaseAdmission { - LeaseAdmission::Refused(LeaseRefusal::SpendCeiling) +fn setup() -> Setup { + Setup { + uid: own_uid(), + turn: LeaseTurn::SelfRun, + token: Some(TOKEN), + max_connections: 8, + } +} + +/// A listener serving `uid` as the caller `eureka`. Peers start in invocation +/// 1 of the Self unit. +fn start(setup: Setup) -> Harness { + let units = Arc::new(FakeUnits { + peer: Mutex::new(None), + asked_about: Mutex::new(Vec::new()), + }); + let dir = std::env::temp_dir().join(format!( + "claude-lease-{}-{}", + std::process::id(), + SEQUENCE.fetch_add(1, Ordering::SeqCst) + )); + std::fs::create_dir_all(&dir).unwrap(); + let socket = dir.join("lease.sock"); + let service = Arc::new(Mutex::new( + TransportService::open( + &dir.join("replay.cc"), + "claude-code", + [CallerAdmission::new( + CALLER, + "lease-test-connection-key", + 1, + [MODEL.to_string()], + 1, + 65_536, + 32_768, + ) + .unwrap()], + 300_000, + ) + .unwrap(), + )); + let tokens = Arc::new(FakeTokens { + token: setup.token, + reads: AtomicUsize::new(0), + probe: Mutex::new(None), + }); + let quota = Arc::new(QuotaCell::new()); + let listener = UnixListener::bind(&socket).unwrap(); + let serving = (service.clone(), quota.clone(), tokens.clone(), units.clone()); + let caller = LeaseCaller { + uid: setup.uid, + caller_id: CALLER.to_string(), + turn: setup.turn, + unit_instance: INSTANCE.to_string(), + }; + std::thread::spawn(move || { + let _ = serve_leases( + &listener, + vec![caller], + LeaseLimits { + socket_timeout: Duration::from_secs(2), + max_lease_ms: 3_600_000, + max_connections: setup.max_connections, + }, + serving.0, + serving.1, + serving.2, + serving.3, + ); + }); + let harness = Harness { + dir, + socket, + service, + quota, + tokens, + units: Some(units), + }; + harness.enter(1); + harness +} + +fn window(name: &str, permille: u32, status: &str, resets_at_unix_ms: u64) -> PlanQuota { + PlanQuota { + window: name.to_string(), + utilization_permille: Some(permille), + resets_at_unix_ms, + status: status.to_string(), + observed_at_unix_ms: 0, + } +} + +/// Asks as fresh unit invocations until a lease is granted, or panics. +fn ask_until_granted(harness: &Harness, first_invocation: u64) -> Answer { + let deadline = Instant::now() + Duration::from_secs(8); + let mut n = first_invocation; + loop { + harness.enter(n); + let answer = ask(harness, 60_000); + if granted_ok(&answer) { + return answer; + } + assert!(Instant::now() < deadline, "no lease was ever granted"); + n += 1; + std::thread::sleep(Duration::from_millis(100)); + } } #[test] fn a_grant_passes_one_fd_holding_the_token_then_end_of_file() { - let harness = start(setup(1_000, 600)); + let harness = start(setup()); let mut answer = ask(&harness, 60_000); - assert_eq!(granted(&answer).reserved_micro_usd, 600); + assert_eq!(granted(&answer).lease_id, format!("lease:{}", invocation(1))); assert_eq!(answer.fds.len(), 1); - // Reading to end of file proves the write end is closed: nobody holds - // the pipe open for another write. let mut read = Vec::new(); File::from(answer.fds.remove(0)) .read_to_end(&mut read) .unwrap(); assert_eq!(read, TOKEN.as_bytes()); assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 1); + assert!(!contains(&answer.raw, TOKEN), "the token is in the pipe only"); } #[test] -fn the_lease_is_named_for_the_unit_invocation_not_for_anything_the_peer_says() { - let harness = start(setup(1_000, 600)); - harness.enter(0xabc); +fn a_grant_carries_the_plan_usage_as_the_connector_last_saw_it() { + let harness = start(setup()); + let reset = now_ms() + 3_600_000; + harness + .quota + .observe([window("five_hour", 420, "allowed", reset)]); let answer = ask(&harness, 60_000); - assert_eq!( - granted(&answer).lease_id, - format!("lease:{}", invocation(0xabc)) - ); + let quota = granted(&answer).plan_quota; + assert_eq!(quota.len(), 1); + assert_eq!(quota[0].window, "five_hour"); + assert_eq!(quota[0].utilization_permille, Some(420)); + assert_eq!(quota[0].status, "allowed"); + assert_eq!(quota[0].resets_at_unix_ms, reset); } #[test] -fn the_token_is_in_the_pipe_and_nowhere_else() { - let harness = start(setup(1_000, 600)); - let mut answer = ask(&harness, 60_000); - let mut frames = vec![answer.raw.clone()]; - let settled = settle(&mut answer, vec![used(MODEL, 10)]); - frames.push(encode_frame(&LeaseMessage::ClaudeLeaseSettled(settled))); - harness.enter(2); - let refused = ask(&harness, 1); - frames.push(refused.raw.clone()); - for frame in &frames { - assert!(!contains(frame, TOKEN)); - } - assert!(std::env::vars().all(|(name, value)| !name.contains(TOKEN) && !value.contains(TOKEN))); - assert!(std::env::args().all(|argument| !argument.contains(TOKEN))); - // The one place it is: the pipe. - let mut read = String::new(); - File::from(answer.fds.remove(0)) - .read_to_string(&mut read) - .unwrap(); - assert_eq!(read, TOKEN); -} - -#[test] -fn a_peer_uid_outside_the_lease_callers_is_refused_and_nothing_is_reserved_or_read() { - let harness = start(Setup { - uid: own_uid() + 1, - ..setup(1_000, 600) - }); - let answer = ask(&harness, 60_000); - assert_eq!(refusal_of(&answer), Some(LeaseRefusalKind::UnknownCaller)); - assert!(answer.fds.is_empty()); - assert!(!contains(&answer.raw, TOKEN)); - assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 0); - // The ledger was not touched: the caller's whole ceiling is still free. - assert!(matches!( - probe_reserve(&harness, "probe", 1_000), - LeaseAdmission::Granted(_) - )); -} - -#[test] -fn the_configured_uid_outside_the_self_unit_is_refused_and_nothing_is_reserved_or_read() { - let harness = start(setup(1_000, 600)); - for (label, peer) in [ - ("another instance", Some(("mind-self@other.service", 1))), - ("a user service", Some(("huginn.service", 1))), - ("no unit", None), - ] { - *harness.units.as_ref().unwrap().peer.lock().unwrap() = peer.map(|(unit, n)| PeerUnit { - unit: unit.to_string(), - invocation_id: invocation(n), - }); - let answer = ask(&harness, 60_000); - assert_eq!( - refusal_of(&answer), - Some(LeaseRefusalKind::UnknownCaller), - "{label}" - ); - assert!(answer.fds.is_empty(), "{label}"); - } - assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 0); - assert!(matches!( - probe_reserve(&harness, "probe", 1_000), - LeaseAdmission::Granted(_) - )); -} - -#[test] -fn one_invocation_of_the_self_unit_leases_once_and_the_next_invocation_leases_again() { - let harness = start(setup(1_000_000, 600)); - let mut first = ask(&harness, 60_000); - assert!(granted_ok(&first)); - settle(&mut first, vec![used(MODEL, 10)]); - // Settled, the caller holds nothing; the invocation has still leased. - let again = ask(&harness, 60_000); - assert_eq!(refusal_of(&again), Some(LeaseRefusalKind::Busy)); - assert!(again.fds.is_empty()); - harness.enter(2); - assert!(granted_ok(&ask(&harness, 60_000))); - assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 2); -} - -#[test] -fn the_reservation_is_claimed_before_the_token_is_read_or_lent() { - let harness = start(setup(1_000, 600)); - let seen = Arc::new(Mutex::new(None)); - let (service, sink) = (harness.service.clone(), seen.clone()); - *harness.tokens.probe.lock().unwrap() = Some(Box::new(move || { - // From inside the token read: the caller already holds a lease. - let admission = service - .lock() - .unwrap() - .begin_lease(CALLER, LeaseTurn::SelfRun, "probe", 1, FAR, now_ms()) - .unwrap(); - *sink.lock().unwrap() = Some(admission); - })); - let answer = ask(&harness, 60_000); - assert_eq!(answer.fds.len(), 1); - assert_eq!( - *seen.lock().unwrap(), - Some(LeaseAdmission::Refused(LeaseRefusal::Busy)) - ); -} - -#[test] -fn a_reservation_the_window_cannot_hold_is_refused_and_lends_nothing() { - let harness = start(setup(1_000, 1_001)); - let answer = ask(&harness, 60_000); - assert_eq!(refusal_of(&answer), Some(LeaseRefusalKind::SpendCeiling)); - assert!(answer.fds.is_empty()); - assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 0); - // Exactly the ceiling fits. - let harness = start(setup(1_000, 1_000)); - assert!(granted_ok(&ask(&harness, 60_000))); -} - -#[test] -fn a_free_lease_and_a_bad_expiry_are_refused() { - let harness = start(setup(1_000, 0)); - assert_eq!( - refusal_of(&ask(&harness, 60_000)), - Some(LeaseRefusalKind::CostBound) - ); - let harness = start(setup(1_000, 100)); - for ttl in [0, 3_600_001 + 1_000] { - let mut stream = connect(&harness); - let expires_at_unix_ms = if ttl == 0 { 1 } else { now_ms() + ttl }; - send( - &mut stream, - &LeaseMessage::ClaudeLeaseRequest(LeaseRequest { expires_at_unix_ms }), - ); - let (raw, fds) = receive(&stream); - assert!(fds.is_empty()); - assert!(matches!( - decode_frame(&raw), - Some(LeaseMessage::ClaudeLeaseRefusal(LeaseRefused { - kind: LeaseRefusalKind::Expired, - .. - })) - )); - } - assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 0); -} - -#[test] -fn a_second_open_lease_for_the_caller_is_busy_even_from_another_invocation() { - let harness = start(setup(1_000_000, 100)); +fn a_lease_ends_when_its_holder_closes() { + let harness = start(setup()); let first = ask(&harness, 60_000); assert!(granted_ok(&first)); harness.enter(2); let second = ask(&harness, 60_000); assert_eq!(refusal_of(&second), Some(LeaseRefusalKind::Busy)); assert!(second.fds.is_empty()); - assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 1); -} - -#[test] -fn settlement_bills_the_true_cost_of_each_models_use_and_the_window_holds_it() { - let harness = start(setup(1_000, 600)); - let mut first = ask(&harness, 60_000); - // One micro-USD a token: 250 of the priced model, and 100 of a model the - // table has never heard of, which bills at the dearest known price. - let settled = settle( - &mut first, - vec![used(MODEL, 250), used("claude-new-9", 100)], - ); - assert_eq!( - (settled.billed_micro_usd, settled.overrun_micro_usd), - (350, 0) - ); - // 350 spent: 650 remain, not 651. - assert_eq!(probe_reserve(&harness, "over", 651), ceiling_refused()); - assert!(matches!( - probe_reserve(&harness, "fits", 650), - LeaseAdmission::Granted(_) - )); -} - -#[test] -fn usage_above_the_reservation_is_billed_in_full_and_the_overrun_is_reported_and_counted() { - let harness = start(setup(1_000, 100)); - let mut first = ask(&harness, 60_000); - let settled = settle(&mut first, vec![used(MODEL, 700), used(MODEL, 5_000)]); - assert_eq!(settled.billed_micro_usd, 5_700); - assert_eq!(settled.overrun_micro_usd, 5_600); - // The window is past its ceiling: nothing more reserves in it. - assert_eq!(probe_reserve(&harness, "next", 1), ceiling_refused()); -} - -#[test] -fn a_plan_quota_the_run_reports_is_not_part_of_the_protocol() { - // A settle in the shape the first protocol used, with a plan_quota the - // run vouches for: the extra key is ignored and the usage is billed. - let harness = start(setup(1_000, 600)); - let mut first = ask(&harness, 60_000); - let legacy = serde_json::json!({ - "kind": "claude_lease_settle", - "value": { - "lease_id": granted(&first).lease_id, - "usage": [{"model": MODEL, "usage": { - "input_tokens": 0, "cache_read_tokens": 0, "cache_write_tokens": 0, - "output_tokens": 40, "reasoning_tokens": null}}], - "plan_quota": [{"window": "five_hour", "utilization_permille": 0, - "resets_at_unix_ms": 1, "status": "allowed", "observed_at_unix_ms": 1}] - } - }); - let body = rmp_serde::to_vec_named(&legacy).unwrap(); - let mut frame = (body.len() as u32).to_be_bytes().to_vec(); - frame.extend_from_slice(&body); - first.stream.write_all(&frame).unwrap(); - let (raw, _) = receive(&first.stream); - assert!(matches!( - decode_frame(&raw), - Some(LeaseMessage::ClaudeLeaseSettled(LeaseSettled { - billed_micro_usd: 40, - overrun_micro_usd: 0 - })) - )); -} - -#[test] -fn a_connection_dropped_without_settling_is_billed_the_reservation() { - let harness = start(setup(1_000, 600)); - let first = ask(&harness, 60_000); - assert!(granted_ok(&first)); drop(first); - let over = ask_after_busy(&harness, 2); - assert_eq!(refusal_of(&over), Some(LeaseRefusalKind::SpendCeiling)); - assert_eq!(probe_reserve(&harness, "over", 401), ceiling_refused()); - assert!(matches!( - probe_reserve(&harness, "fits", 400), - LeaseAdmission::Granted(_) - )); + let third = ask_until_granted(&harness, 3); + assert!(granted_ok(&third)); } #[test] -fn a_lease_that_expires_unsettled_is_closed_and_billed_the_reservation() { - let harness = start(setup(1_000, 600)); - let first = ask(&harness, 400); - assert!(granted_ok(&first)); - let started = Instant::now(); - // The listener, not the client, ends the lease at its expiry. - let (raw, fds) = receive(&first.stream); - assert!(raw.is_empty() && fds.is_empty()); - assert!(started.elapsed() >= Duration::from_millis(250)); - harness.enter(2); - let over = ask(&harness, 60_000); - assert_eq!(refusal_of(&over), Some(LeaseRefusalKind::SpendCeiling)); -} - -#[test] -fn a_settle_trickled_past_the_expiry_is_no_report_and_bills_the_reservation() { - let harness = start(setup(1_000, 600)); - let mut answer = ask(&harness, 1_500); - let frame = settle_frame(&answer, vec![used(MODEL, 10)]); - let chunk = frame.len() / 4 + 1; - // Each write lands inside any per-read wait; the frame does not land - // inside the lease. - for (i, part) in frame.chunks(chunk).enumerate() { - if i > 0 { - std::thread::sleep(Duration::from_millis(1_200)); - } - let _ = answer.stream.write_all(part); - } - let (raw, _) = receive(&answer.stream); - assert!(raw.is_empty(), "{:?}", decode_frame(&raw)); - assert_eq!(probe_reserve(&harness, "over", 401), ceiling_refused()); -} - -#[test] -fn a_request_trickled_past_the_socket_timeout_is_dropped_without_a_reservation() { - let harness = start(setup(1_000, 600)); - let mut stream = connect(&harness); - let started = Instant::now(); - // Three bytes of a length, 0.9 s apart, then silence: every wait is inside - // the 2 s socket timeout, the frame is not. - for byte in [0_u8, 0, 0] { - stream.write_all(&[byte]).unwrap(); - std::thread::sleep(Duration::from_millis(900)); - } - let (raw, fds) = receive(&stream); - assert!(raw.is_empty() && fds.is_empty()); - assert!( - started.elapsed() < Duration::from_millis(3_300), - "{:?}", - started.elapsed() - ); - assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 0); - assert!(matches!( - probe_reserve(&harness, "probe", 1_000), - LeaseAdmission::Granted(_) - )); -} - -#[test] -fn a_settle_for_another_lease_is_no_report_and_bills_the_reservation() { - let harness = start(setup(1_000, 600)); - let mut first = ask(&harness, 60_000); - send( - &mut first.stream, - &LeaseMessage::ClaudeLeaseSettle(LeaseSettle { - lease_id: "lease:someone-else".to_string(), - usage: vec![used(MODEL, 1)], - }), - ); - let (raw, _) = receive(&first.stream); - assert!(raw.is_empty()); - harness.enter(2); - let over = ask(&harness, 60_000); - assert_eq!(refusal_of(&over), Some(LeaseRefusalKind::SpendCeiling)); -} - -#[test] -fn a_settle_that_reports_nothing_is_billed_the_reservation() { - let harness = start(setup(1_000, 600)); - let mut first = ask(&harness, 60_000); - let settled = settle(&mut first, vec![]); - assert_eq!( - (settled.billed_micro_usd, settled.overrun_micro_usd), - (600, 0) - ); - // 600 spent: the whole ceiling no longer reserves. - assert_eq!(probe_reserve(&harness, "all", 1_000), ceiling_refused()); -} - -#[test] -fn a_token_that_cannot_be_read_is_never_lent_and_bills_nothing() { - let harness = start(Setup { - token: None, - ..setup(1_000, 1_000) - }); +fn a_lease_ends_when_its_unit_invocation_ends() { + let harness = start(setup()); let first = ask(&harness, 60_000); - assert!(first.raw.is_empty() && first.fds.is_empty()); - assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 1); - // Nothing was billed: the whole ceiling reserves again and gets as far - // as the token read, not a refusal. - let again = ask_after_busy(&harness, 2); - assert!(again.message.is_none() && again.fds.is_empty()); - assert!(harness.tokens.reads.load(Ordering::SeqCst) >= 2); + assert!(granted_ok(&first)); + // The invocation is over: the system no longer places the peer in it. The + // holder still has its socket open, and the lease does not outlive it. + *harness.units.as_ref().unwrap().peer.lock().unwrap() = None; + std::thread::sleep(Duration::from_millis(1_200)); + harness.enter(2); + assert!(granted_ok(&ask(&harness, 60_000))); + drop(first); } #[test] -fn a_malformed_or_oversized_frame_is_dropped_without_a_reservation() { - let harness = start(setup(1_000, 1_000)); - for bytes in [ - vec![0xff, 0xff, 0xff, 0xff], - vec![0, 0, 0, 0], - vec![0, 0, 0, 3, 0xc1, 0xc1, 0xc1], - ] { - let mut stream = connect(&harness); - stream.write_all(&bytes).unwrap(); - let started = Instant::now(); - let (raw, fds) = receive(&stream); - assert!(raw.is_empty() && fds.is_empty()); - // Refused at once, not after the socket timeout of a read that was - // never going to complete. - assert!(started.elapsed() < Duration::from_millis(1_000)); - } +fn a_lease_ends_at_its_expiry_though_its_holder_is_still_there() { + let harness = start(setup()); + let first = ask(&harness, 700); + assert!(granted_ok(&first)); + let next = ask_until_granted(&harness, 2); + assert!(granted_ok(&next)); + drop(first); +} + +#[test] +fn a_rejected_window_refuses_the_lease_until_its_reset() { + let harness = start(setup()); + let reset = now_ms() + 1_500; + harness + .quota + .observe([window("five_hour", 1_000, "rejected", reset)]); + let refused = ask(&harness, 60_000); + assert_eq!(refusal_of(&refused), Some(LeaseRefusalKind::ProviderRejected)); + assert!(refused.fds.is_empty()); assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 0); + // At the provider's own reset the refusal lifts, with no one lifting it. + let lifted = ask_until_granted(&harness, 2); + assert!(granted_ok(&lifted)); + assert!(now_ms() >= reset); +} + +#[test] +fn high_utilization_refuses_nothing() { + let harness = start(setup()); + let reset = now_ms() + 3_600_000; + harness.quota.observe([ + window("five_hour", 999, "allowed", reset), + window("seven_day", 1_000, "allowed_warning", reset), + ]); assert!(granted_ok(&ask(&harness, 60_000))); } #[test] -fn every_message_kind_is_the_pinned_bytes_of_its_fixture() { - let expires = 1_793_307_000_000; - let lease = format!("lease:{}", "0123456789abcdef0123456789abcdef"); - let messages = [ +fn a_peer_uid_outside_the_lease_callers_is_refused_and_nothing_is_read() { + let harness = start(Setup { + uid: own_uid().wrapping_add(1), + ..setup() + }); + let answer = ask(&harness, 60_000); + assert_eq!(refusal_of(&answer), Some(LeaseRefusalKind::UnknownCaller)); + assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 0); +} + +#[test] +fn the_configured_uid_outside_the_self_unit_is_refused_and_nothing_is_read() { + let harness = start(setup()); + harness.run_in("mind-persona@eureka.service", 1); + let answer = ask(&harness, 60_000); + assert_eq!(refusal_of(&answer), Some(LeaseRefusalKind::UnknownCaller)); + assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 0); +} + +#[test] +fn an_expiry_in_the_past_or_past_the_bound_is_refused() { + let harness = start(setup()); + let past = ask(&harness, 0); + assert_eq!(refusal_of(&past), Some(LeaseRefusalKind::Expired)); + let far = ask(&harness, 3_600_001 + 60_000); + assert_eq!(refusal_of(&far), Some(LeaseRefusalKind::Expired)); + assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 0); +} + +#[test] +fn a_token_that_cannot_be_read_is_never_lent_and_frees_the_lease() { + let harness = start(Setup { + token: None, + ..setup() + }); + let first = ask(&harness, 60_000); + assert!(first.raw.is_empty() && first.fds.is_empty()); + harness.enter(2); + let again = ask(&harness, 60_000); + assert!(again.message.is_none() && again.fds.is_empty(), "not Busy"); + assert!(harness.tokens.reads.load(Ordering::SeqCst) >= 2); +} + +#[test] +fn the_protocol_has_no_settle_frame() { + let grant = encode_frame(&LeaseMessage::ClaudeLeaseGrant(LeaseGranted { + lease_id: "lease:x".to_string(), + plan_quota: Vec::new(), + })); + let value: rmpv::Value = rmp_serde::from_slice(&grant[4..]).unwrap(); + assert_eq!( + value, + rmpv::Value::Map(vec![ + ("kind".into(), "claude_lease_grant".into()), + ( + "value".into(), + rmpv::Value::Map(vec![ + ("lease_id".into(), "lease:x".into()), + ("plan_quota".into(), rmpv::Value::Array(Vec::new())), + ]) + ), + ]) + ); + // A frame of the old kind is no message at all. + let old = rmpv::Value::Map(vec![ + ("kind".into(), "claude_lease_settle".into()), ( - &include_bytes!("fixtures/lease_v1/request.bin")[..], - LeaseMessage::ClaudeLeaseRequest(LeaseRequest { - expires_at_unix_ms: expires, - }), + "value".into(), + rmpv::Value::Map(vec![ + ("lease_id".into(), "lease:x".into()), + ("usage".into(), rmpv::Value::Array(Vec::new())), + ]), ), - ( - &include_bytes!("fixtures/lease_v1/grant.bin")[..], - LeaseMessage::ClaudeLeaseGrant(LeaseGranted { - lease_id: lease.clone(), - reserved_micro_usd: 600_000, - }), - ), - ( - &include_bytes!("fixtures/lease_v1/refusal.bin")[..], - LeaseMessage::ClaudeLeaseRefusal(LeaseRefused { - kind: LeaseRefusalKind::SpendCeiling, - detail: "the window cannot hold this reservation".to_string(), - }), - ), - ( - &include_bytes!("fixtures/lease_v1/settle.bin")[..], - LeaseMessage::ClaudeLeaseSettle(LeaseSettle { - lease_id: lease, - usage: vec![ - ModelUsage { - model: "claude-opus-5-1".to_string(), - usage: Usage { - input_tokens: 1_200, - cache_read_tokens: 30_000, - cache_write_tokens: 500, - output_tokens: 800, - reasoning_tokens: Some(100), - }, - }, - ModelUsage { - model: "claude-haiku-5-1".to_string(), - usage: Usage { - input_tokens: 400, - cache_read_tokens: 0, - cache_write_tokens: 0, - output_tokens: 60, - reasoning_tokens: None, - }, - }, - ], - }), - ), - ( - &include_bytes!("fixtures/lease_v1/settled.bin")[..], - LeaseMessage::ClaudeLeaseSettled(LeaseSettled { - billed_micro_usd: 123_456, - overrun_micro_usd: 23_456, - }), - ), - ]; - for (bytes, message) in messages { - assert_eq!(encode_frame(&message), bytes, "{message:?}"); - assert_eq!(decode_frame(bytes), Some(message)); - } + ]); + let mut body = Vec::new(); + rmpv::encode::write_value(&mut body, &old).unwrap(); + let mut frame = (body.len() as u32).to_be_bytes().to_vec(); + frame.extend(body); + assert_eq!(decode_frame(&frame), None); } #[test] fn connections_past_the_limit_are_dropped_and_finished_ones_free_their_place() { let harness = start(Setup { - max_connections: 2, - ..setup(1_000, 0) + max_connections: 1, + ..setup() }); - let held: Vec = (0..2).map(|_| connect(&harness)).collect(); - // Let the listener take both before the third arrives. - std::thread::sleep(Duration::from_millis(300)); - let third = ask(&harness, 60_000); - assert!(third.raw.is_empty() && third.fds.is_empty()); - drop(held); - // Finished connections give their places back: more refusals in a row - // than the limit holds, each answered. - let first = ask_until_answered(&harness); - assert_eq!(refusal_of(&first), Some(LeaseRefusalKind::CostBound)); - for _ in 0..4 { - let again = ask(&harness, 60_000); - assert_eq!(refusal_of(&again), Some(LeaseRefusalKind::CostBound)); - } -} - -/// A free lease, asked until the listener answers instead of dropping. -fn ask_until_answered(harness: &Harness) -> Answer { - let deadline = Instant::now() + Duration::from_secs(5); - loop { - let answer = ask(harness, 60_000); - if !answer.raw.is_empty() || Instant::now() > deadline { - return answer; - } - std::thread::sleep(Duration::from_millis(20)); - } -} - -#[test] -fn an_account_resolves_to_its_uid_and_a_missing_one_to_none() { - use claude_connector::user_uid; - assert_eq!(user_uid("root"), Some(0)); - assert_eq!(user_uid("nobody"), Some(65534)); - assert_eq!(user_uid("no-such-account-4f1c"), None); - assert_eq!(user_uid("bad name"), None); -} - -#[test] -fn a_token_up_to_one_atomic_pipe_write_is_lent_and_a_longer_one_is_not() { - let fits: &'static str = Box::leak("t".repeat(4096).into_boxed_str()); - let harness = start(Setup { - token: Some(fits), - ..setup(1_000, 100) - }); - let mut answer = ask(&harness, 60_000); - let mut read = Vec::new(); - File::from(answer.fds.remove(0)) - .read_to_end(&mut read) - .unwrap(); - assert_eq!(read.len(), 4096); - - let too_long: &'static str = Box::leak("t".repeat(4097).into_boxed_str()); - let harness = start(Setup { - token: Some(too_long), - ..setup(1_000, 100) - }); - let answer = ask(&harness, 60_000); - assert!(answer.raw.is_empty() && answer.fds.is_empty()); -} - -// ---- a second real uid (Soul s1, kept as a regression) ---- - -fn describe(answer: &Answer) -> String { - let kind = match &answer.message { - Some(LeaseMessage::ClaudeLeaseGrant(_)) => "grant".to_string(), - Some(LeaseMessage::ClaudeLeaseRefusal(r)) => format!("refusal:{:?}", r.kind), - Some(other) => format!("other:{other:?}"), - None => format!("none(raw={})", answer.raw.len()), - }; - format!("kind={kind} fds={}", answer.fds.len()) -} - -/// The child half of the second-uid test: connects to `LEASE_SOCK` as whoever -/// it was started as and prints one line saying what came back. The fake -/// canary token is never printed, only whether it matched. -#[test] -#[ignore = "run as a child of a_second_real_uid_cannot_lease"] -fn lease_probe_client() { - let socket = PathBuf::from(std::env::var("LEASE_SOCK").unwrap()); - let mut answer = ask_at(&socket, 60_000); - let shown = describe(&answer); - let mut token_matches = false; - if !answer.fds.is_empty() { - let mut read = Vec::new(); - File::from(answer.fds.remove(0)) - .read_to_end(&mut read) - .unwrap(); - token_matches = read == TOKEN.as_bytes(); - } - println!( - "LEASEPROBE uid={} {} token_matches_canary={token_matches}", - own_uid(), - shown - ); -} - -/// Runs the probe client as `reuid`, with primary group `regid` and -/// supplementary `groups`, and returns its LEASEPROBE line. The marker is -/// found anywhere on a line: the test harness prints it after its own prefix. -fn probe_as(socket: &std::path::Path, reuid: u32, regid: u32, groups: &str) -> String { - let exe = std::env::current_exe().unwrap(); - let out = std::process::Command::new("setpriv") - .arg(format!("--reuid={reuid}")) - .arg(format!("--regid={regid}")) - .arg(format!("--groups={groups}")) - .arg(&exe) - .args([ - "--exact", - "lease_probe_client", - "--ignored", - "--nocapture", - "--test-threads=1", - ]) - .env("LEASE_SOCK", socket) - .output() - .unwrap(); - let stdout = String::from_utf8_lossy(&out.stdout).to_string(); - stdout - .lines() - .find_map(|line| line.find("LEASEPROBE").map(|at| line[at..].to_string())) - .unwrap_or_else(|| { - format!( - "NO-LINE status={:?} stdout={stdout:?} stderr={}", - out.status, - String::from_utf8_lossy(&out.stderr) - ) - }) -} - -fn have(tool: &str) -> bool { - std::process::Command::new(tool) - .arg("--version") - .output() - .is_ok() -} - -#[test] -fn a_second_real_uid_cannot_lease() { - use std::os::unix::fs::PermissionsExt; - if own_uid() != 0 || !have("setpriv") { - eprintln!("SKIPPED a_second_real_uid_cannot_lease: needs root and setpriv"); - return; - } - // Numeric uids need no account: 4242 is the configured Self user, 4243 a - // second account that shares its groups. - let harness = start(Setup { - uid: 4242, - ..setup(1_000_000, 600) - }); - std::fs::set_permissions(&harness.dir, std::fs::Permissions::from_mode(0o755)).unwrap(); - std::fs::set_permissions(&harness.socket, std::fs::Permissions::from_mode(0o666)).unwrap(); - // Another uid, in the caller's group and the lease group. - let a = probe_as(&harness.socket, 4243, 4243, "4242,4244"); - // Another uid whose primary gid is the caller's. - let b = probe_as(&harness.socket, 4243, 4242, "4242"); - // Root, which is not configured. - let root = describe(&ask(&harness, 60_000)); - assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 0); - // The configured uid does lease. - let c = probe_as(&harness.socket, 4242, 4242, "4242"); - assert!(a.contains("kind=refusal:UnknownCaller fds=0"), "{a}"); - assert!(b.contains("kind=refusal:UnknownCaller fds=0"), "{b}"); - assert!(root.contains("kind=refusal:UnknownCaller fds=0"), "{root}"); - assert!( - c.contains("kind=grant fds=1") && c.contains("token_matches_canary=true"), - "{c}" - ); - assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 1); -} - -// ---- a real systemd, where there is one ---- - -/// The child half of the systemd test, run inside a transient service: leases, -/// settles, and leases again from the same invocation. -#[test] -#[ignore = "run inside a transient unit by a_real_systemd_unit_invocation_leases_once"] -fn lease_in_unit_client() { - let socket = PathBuf::from(std::env::var("LEASE_SOCK").unwrap()); - let mut first = ask_at(&socket, 60_000); - let first_kind = describe(&first); - if granted_ok(&first) { - settle(&mut first, Vec::new()); - } - let second = ask_at(&socket, 60_000); - println!( - "LEASEPROBE first[{first_kind}] second[{}]", - describe(&second) - ); -} - -fn run_in_unit(socket: &std::path::Path) -> String { - let exe = std::env::current_exe().unwrap(); - let out = std::process::Command::new("systemd-run") - .args([ - "--unit=mind-self@leasetest.service", - "--pipe", - "--wait", - "--collect", - "--quiet", - ]) - .arg(format!("--setenv=LEASE_SOCK={}", socket.display())) - .arg(&exe) - .args([ - "--exact", - "lease_in_unit_client", - "--ignored", - "--nocapture", - "--test-threads=1", - ]) - .output() - .unwrap(); - let stdout = String::from_utf8_lossy(&out.stdout).to_string(); - stdout - .lines() - .find_map(|line| line.find("LEASEPROBE").map(|at| line[at..].to_string())) - .unwrap_or_else(|| { - format!( - "NO-LINE status={:?} stdout={stdout:?} stderr={}", - out.status, - String::from_utf8_lossy(&out.stderr) - ) - }) -} - -#[test] -fn a_real_systemd_unit_invocation_leases_once() { - if own_uid() != 0 - || !std::path::Path::new("/run/systemd/system").exists() - || !have("systemd-run") - { - eprintln!( - "SKIPPED a_real_systemd_unit_invocation_leases_once: needs root, systemd as init and systemd-run" - ); - return; - } - // The listener asks the production resolver, not the fake. - let harness = start_with( - setup(1_000_000, 600), - Arc::new(SystemdUnits::system()), - None, - "leasetest", - ); - let leased_then_busy = "first[kind=grant fds=1] second[kind=refusal:Busy fds=0]"; - // Two invocations of the same unit: each leases once, as its own. - let one = run_in_unit(&harness.socket); - let two = run_in_unit(&harness.socket); - assert!(one.contains(leased_then_busy), "{one}"); - assert!(two.contains(leased_then_busy), "{two}"); - assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 2); - // A process outside any such unit is not a Self run, though its uid is - // the configured one. - let outside = ask(&harness, 60_000); - assert_eq!(refusal_of(&outside), Some(LeaseRefusalKind::UnknownCaller)); - assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 2); -} - -#[test] -fn the_peer_is_named_by_a_process_handle_not_a_pid_number() { - let (one, _other) = UnixStream::pair().unwrap(); - let peer = claude_connector::peer_credentials(&one).unwrap(); - assert_eq!(peer.uid, own_uid()); - let process = peer.process.expect("SO_PEERPIDFD on this kernel"); - let info = - std::fs::read_to_string(format!("/proc/self/fdinfo/{}", process.as_raw_fd())).unwrap(); - assert!( - info.lines() - .any(|line| line.strip_prefix("Pid:").map(str::trim) - == Some(&std::process::id().to_string())), - "{info}" - ); -} - -#[test] -fn the_unit_is_asked_of_the_callers_unit_through_the_peers_handle() { - let harness = start(setup(1_000, 600)); - let answer = ask(&harness, 60_000); - assert!(refusal_of(&answer).is_none() && answer.fds.len() == 1); - assert_eq!( - *harness.units.as_ref().unwrap().asked_about.lock().unwrap(), - vec![(std::process::id().to_string(), UNIT.to_string())] - ); -} - -const PERSONA_UNIT: &str = "mind-persona@eureka.service"; - -/// A listener whose one lease caller is a Persona account; peers start in -/// invocation 1 of the Persona unit. -fn start_persona(setup: Setup) -> Harness { - let units = Arc::new(FakeUnits { - peer: Mutex::new(None), - asked_about: Mutex::new(Vec::new()), - }); - let harness = start_with( - Setup { - turn: LeaseTurn::PersonaTurn, - ..setup - }, - units.clone(), - Some(units), - INSTANCE, - ); - harness.run_in(PERSONA_UNIT, 1); - harness -} - -fn lease_on(harness: &Harness, turn: LeaseTurn, invocation: &str) -> LeaseAdmission { - harness - .service - .lock() - .unwrap() - .begin_lease(CALLER, turn, invocation, 100, FAR, now_ms()) - .unwrap() -} - -fn lease_id(admission: &LeaseAdmission) -> String { - match admission { - LeaseAdmission::Granted(grant) => grant.lease_id.clone(), - other => panic!("expected a grant, got {other:?}"), - } -} - -fn busy() -> LeaseAdmission { - LeaseAdmission::Refused(LeaseRefusal::Busy) -} - -#[test] -fn persona_and_self_of_one_caller_lease_together() { - let harness = start(setup(1_000_000, 100)); - let own = lease_on(&harness, LeaseTurn::SelfRun, "self-run"); - assert_eq!(lease_id(&own), "lease:self-run"); - let persona = lease_on(&harness, LeaseTurn::PersonaTurn, "persona-run"); - assert_eq!(lease_id(&persona), "lease:persona-run:1"); - // Each turn kind still holds one: a second Self run is busy while the - // first is open, a second Persona lease of the same invocation too. - assert_eq!(lease_on(&harness, LeaseTurn::SelfRun, "self-run-2"), busy()); - assert_eq!( - lease_on(&harness, LeaseTurn::PersonaTurn, "persona-run"), - busy() - ); - // Settling the Persona lease frees nothing of the Self lease's. - harness - .service - .lock() - .unwrap() - .settle_lease(CALLER, "lease:persona-run:1", LeaseUse::Reservation) - .unwrap(); - assert_eq!(lease_on(&harness, LeaseTurn::SelfRun, "self-run-2"), busy()); -} - -#[test] -fn persona_invocation_leases_sequentially() { - let harness = start(setup(1_000_000, 100)); - for expected in 1..=3 { - let granted = lease_on(&harness, LeaseTurn::PersonaTurn, "persona-run"); - let id = lease_id(&granted); - assert_eq!(id, format!("lease:persona-run:{expected}")); - harness - .service - .lock() - .unwrap() - .settle_lease(CALLER, &id, LeaseUse::Reservation) - .unwrap(); - } - // Another invocation counts from one: the sequence is per invocation. - let other = lease_on(&harness, LeaseTurn::PersonaTurn, "persona-other"); - assert_eq!(lease_id(&other), "lease:persona-other:1"); -} - -#[test] -fn persona_invocation_second_open_lease_busy() { - let harness = start_persona(setup(1_000_000, 100)); let first = ask(&harness, 60_000); assert!(granted_ok(&first)); let second = ask(&harness, 60_000); - assert_eq!(refusal_of(&second), Some(LeaseRefusalKind::Busy)); - assert!(second.fds.is_empty()); - assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 1); -} - -#[test] -fn persona_turns_of_a_live_invocation_each_lend_a_token_over_the_socket() { - let harness = start_persona(setup(1_000_000, 100)); - for _ in 0..2 { - let mut answer = ask(&harness, 60_000); - assert!(granted_ok(&answer)); - settle(&mut answer, vec![used(MODEL, 10)]); - } - assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 2); -} - -#[test] -fn self_invocation_still_leases_once() { - let harness = start(setup(1_000_000, 100)); - let id = lease_id(&lease_on(&harness, LeaseTurn::SelfRun, "self-run")); - harness - .service - .lock() - .unwrap() - .settle_lease(CALLER, &id, LeaseUse::Reservation) - .unwrap(); - assert_eq!(lease_on(&harness, LeaseTurn::SelfRun, "self-run"), busy()); -} - -#[test] -fn persona_row_refuses_mind_self_unit() { - let harness = start_persona(setup(1_000_000, 100)); - harness.run_in(UNIT, 1); - let answer = ask(&harness, 60_000); - assert_eq!(refusal_of(&answer), Some(LeaseRefusalKind::UnknownCaller)); - assert!(answer.fds.is_empty()); - assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 0); - // And a Self row refuses the Persona unit. - let harness = start(setup(1_000_000, 100)); - harness.run_in(PERSONA_UNIT, 1); - let answer = ask(&harness, 60_000); - assert_eq!(refusal_of(&answer), Some(LeaseRefusalKind::UnknownCaller)); - assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 0); -} - -#[test] -fn a_persona_lease_reserves_and_settles_in_the_callers_window() { - let harness = start(setup(1_000, 600)); - let id = lease_id(&lease_on(&harness, LeaseTurn::PersonaTurn, "persona-run")); - // 100 reserved by the helper; a Self probe of 901 cannot fit, 900 can. - assert_eq!(probe_reserve(&harness, "p1", 901), ceiling_refused()); - harness - .service - .lock() - .unwrap() - .settle_lease(CALLER, &id, LeaseUse::NeverLent) - .unwrap(); - assert!(matches!( - probe_reserve(&harness, "p2", 1_000), - LeaseAdmission::Granted(_) - )); + assert!(second.message.is_none() && second.fds.is_empty()); + drop(first); + assert!(granted_ok(&ask_until_granted(&harness, 2))); } diff --git a/crates/model-connector/src/client.rs b/crates/model-connector/src/client.rs index 47b6c94..881ef01 100644 --- a/crates/model-connector/src/client.rs +++ b/crates/model-connector/src/client.rs @@ -9,11 +9,11 @@ use rand_core::{OsRng, RngCore}; use crate::contract::{TransportInvocation, TransportKey, TransportResult}; use crate::envelope::{ MAX_IDENTITY_BYTES, ServiceError, TransportFrameError, decode_result_frame, - decode_spend_readout_response_frame, decrypt_result, decrypt_spend_readout, - encode_invocation_frame, encode_spend_readout_request_frame, encrypt_invocation, - encrypt_spend_readout_request, read_transport_frame, write_transport_frame, + decode_usage_readout_response_frame, decrypt_result, decrypt_usage_readout, + encode_invocation_frame, encode_usage_readout_request_frame, encrypt_invocation, + encrypt_usage_readout_request, read_transport_frame, write_transport_frame, }; -use crate::readout::SpendReadout; +use crate::readout::UsageReadout; #[derive(Clone)] pub struct ConnectorClient { @@ -62,25 +62,25 @@ impl ConnectorClient { decrypt_result(&envelope, &self.security, invocation).map_err(Into::into) } - /// Reads the connector's spend ledger and plan quota as `caller_runtime_id`, + /// Reads the plan's usage windows as `caller_runtime_id`, /// which must be a caller whose connection key this client holds. - pub fn spend_readout( + pub fn usage_readout( &self, caller_runtime_id: &str, - ) -> Result { + ) -> Result { let mut nonce = [0_u8; 16]; OsRng.fill_bytes(&mut nonce); let request_id = format!( - "spend-readout-{}", + "usage-readout-{}", nonce.iter().map(|b| format!("{b:02x}")).collect::() ); - let sealed = encrypt_spend_readout_request(caller_runtime_id, &request_id, &self.security)?; - let response = self.exchange(&encode_spend_readout_request_frame(&sealed)?)?; - let envelope = decode_spend_readout_response_frame(&response)?; + let sealed = encrypt_usage_readout_request(caller_runtime_id, &request_id, &self.security)?; + let response = self.exchange(&encode_usage_readout_request_frame(&sealed)?)?; + let envelope = decode_usage_readout_response_frame(&response)?; if envelope.caller_runtime_id != caller_runtime_id || envelope.request_id != request_id { return Err(ServiceError::OuterIdentity.into()); } - decrypt_spend_readout(&envelope, &self.security).map_err(Into::into) + decrypt_usage_readout(&envelope, &self.security).map_err(Into::into) } /// Sends one frame on a fresh connection and reads the one answer. @@ -274,17 +274,17 @@ mod fixture_tests { let server = std::thread::spawn(move || { let (mut stream, _) = listener.accept().unwrap(); let request = read_transport_frame(&mut stream, 64 * 1024).unwrap(); - let RequestFrame::SpendReadout(asked) = decode_request_frame(&request).unwrap() + let RequestFrame::UsageReadout(asked) = decode_request_frame(&request).unwrap() else { panic!("the client sent something other than a readout request"); }; let request_id = answer_request.unwrap_or(&asked.request_id); - let readout = SpendReadout::new(1, Vec::new(), PlanQuotaSnapshot::default()); + let readout = UsageReadout::new(1, PlanQuotaSnapshot::default()); let sealed = - encrypt_spend_readout(answer_caller, request_id, &readout, &key).unwrap(); + encrypt_usage_readout(answer_caller, request_id, &readout, &key).unwrap(); write_transport_frame( &mut stream, - &encode_spend_readout_response_frame(&sealed).unwrap(), + &encode_usage_readout_response_frame(&sealed).unwrap(), 64 * 1024, ) .unwrap(); @@ -296,7 +296,7 @@ mod fixture_tests { Some(Duration::from_secs(1)), ) .unwrap(); - let answer = client.spend_readout("caller-1"); + let answer = client.usage_readout("caller-1"); if accepted { assert_eq!(answer.unwrap().at_unix_ms, 1); } else { diff --git a/crates/model-connector/src/contract.rs b/crates/model-connector/src/contract.rs index f22b8f9..d039774 100644 --- a/crates/model-connector/src/contract.rs +++ b/crates/model-connector/src/contract.rs @@ -9,11 +9,11 @@ use zeroize::Zeroize; use crate::envelope::{ServiceError, binary_schema}; -pub const PROVIDER_REQUEST_SCHEMA_ID: &str = "gamecult.model.provider_request.v3"; -pub const INVOCATION_SCHEMA_ID: &str = "gamecult.model.transport_invocation.v3"; -pub const RESULT_SCHEMA_ID: &str = "gamecult.model.transport_result.v3"; -pub const RECEIPT_SCHEMA_ID: &str = "gamecult.model.transport_receipt.v3"; -pub const ENVELOPE_SCHEMA_ID: &str = "gamecult.model.transport_envelope.v3"; +pub const PROVIDER_REQUEST_SCHEMA_ID: &str = "gamecult.model.provider_request.v4"; +pub const INVOCATION_SCHEMA_ID: &str = "gamecult.model.transport_invocation.v4"; +pub const RESULT_SCHEMA_ID: &str = "gamecult.model.transport_result.v4"; +pub const RECEIPT_SCHEMA_ID: &str = "gamecult.model.transport_receipt.v4"; +pub const ENVELOPE_SCHEMA_ID: &str = "gamecult.model.transport_envelope.v4"; #[derive(Clone, PartialEq, Eq)] pub struct TransportKey(pub(crate) [u8; 32]); @@ -190,10 +190,10 @@ pub enum FailureKind { Provider, } -/// Token accounting for one request, as the backend reports it; the service -/// prices it. Cache reads and writes are counted apart from plain input. +/// Token accounting for one request, as the backend reports 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. +/// (absent when the provider does not report it) and is never counted apart. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] #[serde(deny_unknown_fields)] pub struct Usage { @@ -252,10 +252,6 @@ 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, } @@ -264,7 +260,6 @@ 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)?; @@ -274,7 +269,6 @@ impl TransportInvocation { expires_at_unix_ms, native_request_sha256, provider_request_sha256, - max_cost_micro_usd, request, }) } @@ -287,9 +281,6 @@ 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); @@ -485,15 +476,6 @@ 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, } @@ -558,12 +540,6 @@ 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, @@ -573,15 +549,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 provider itself is refusing right now: a plan window it reported + // `rejected` has not reached its own reset. Plain comments, as above. + ProviderRejected, } -/// The one wire codec of every v3 document: named MessagePack (struct fields +/// The one wire codec of every v4 document: named MessagePack (struct fields /// are keyed by name), with unknown fields refused by the types themselves. pub(crate) fn wire_encode(document: &T) -> Result, ContractError> { rmp_serde::to_vec_named(document).map_err(|_| ContractError::Encoding) @@ -787,28 +760,16 @@ 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() { + fn the_invocations_identity_covers_its_request() { let one = invocation(); let mut other = one.clone(); - other.max_cost_micro_usd += 1; + other.request.max_output_tokens += 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); + assert_eq!(decoded, one); } #[test] @@ -869,7 +830,7 @@ mod tests { } #[test] - fn v3_documents_carry_their_v3_ids() { + fn v4_documents_carry_their_v4_ids() { assert_eq!( [ PROVIDER_REQUEST_SCHEMA_ID, @@ -879,11 +840,11 @@ mod tests { ENVELOPE_SCHEMA_ID ], [ - "gamecult.model.provider_request.v3", - "gamecult.model.transport_invocation.v3", - "gamecult.model.transport_result.v3", - "gamecult.model.transport_receipt.v3", - "gamecult.model.transport_envelope.v3" + "gamecult.model.provider_request.v4", + "gamecult.model.transport_invocation.v4", + "gamecult.model.transport_result.v4", + "gamecult.model.transport_receipt.v4", + "gamecult.model.transport_envelope.v4" ] ); let invocation = invocation(); @@ -922,7 +883,7 @@ mod tests { let bytes = wire_encode(document).unwrap(); let mut value = rmpv::decode::read_value(&mut bytes.as_slice()).unwrap(); let rmpv::Value::Map(entries) = &mut value else { - panic!("a v3 document is a named map"); + panic!("a v4 document is a named map"); }; entries.push((rmpv::Value::from("unexpected_field"), rmpv::Value::from(1))); let mut out = Vec::new(); diff --git a/crates/model-connector/src/envelope.rs b/crates/model-connector/src/envelope.rs index 70d8ed2..2fb4543 100644 --- a/crates/model-connector/src/envelope.rs +++ b/crates/model-connector/src/envelope.rs @@ -17,14 +17,14 @@ use crate::contract::{ ContractError, ENVELOPE_SCHEMA_ID, TransportInvocation, TransportKey, TransportResult, canonical_invocation_bytes, require_id, wire_decode, wire_encode, }; -use crate::readout::SpendReadout; +use crate::readout::UsageReadout; #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] pub enum EnvelopeKind { Invocation, Result, - SpendReadoutRequest, - SpendReadoutResult, + UsageReadoutRequest, + UsageReadoutResult, } #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] @@ -54,8 +54,8 @@ pub(crate) fn binary_schema(_: &mut schemars::SchemaGenerator) -> schemars::Sche pub const CONNECTOR_SERVICE_ID: &str = "gamecult.model.connector"; /// The operation that runs one sealed invocation. pub const INVOKE_OPERATION: &str = "invoke"; -/// The operation that reads the spend ledger and the plan's quota. -pub const SPEND_READOUT_OPERATION: &str = "spend_readout"; +/// The operation that reads the plan's usage. +pub const USAGE_READOUT_OPERATION: &str = "usage_readout"; const PAYLOAD_ENCODING: &str = "messagepack-base64"; const OK_STATUS: &str = "ok"; @@ -78,19 +78,19 @@ fn seal_payload(envelope: &TransportEnvelope) -> Result { #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) enum Operation { Invoke, - SpendReadout, + UsageReadout, } impl Operation { fn name(self) -> &'static str { match self { Operation::Invoke => INVOKE_OPERATION, - Operation::SpendReadout => SPEND_READOUT_OPERATION, + Operation::UsageReadout => USAGE_READOUT_OPERATION, } } - fn from_name(name: &str) -> Option { - [Operation::Invoke, Operation::SpendReadout] + pub(crate) fn from_name(name: &str) -> Option { + [Operation::Invoke, Operation::UsageReadout] .into_iter() .find(|operation| operation.name() == name) } @@ -141,10 +141,10 @@ pub(crate) fn encode_invocation_frame( request_frame(Operation::Invoke, envelope) } -pub(crate) fn encode_spend_readout_request_frame( +pub(crate) fn encode_usage_readout_request_frame( envelope: &TransportEnvelope, ) -> Result, ServiceError> { - request_frame(Operation::SpendReadout, envelope) + request_frame(Operation::UsageReadout, envelope) } #[cfg(any(feature = "daemon", test))] @@ -153,10 +153,10 @@ pub(crate) fn encode_result_frame(envelope: &TransportEnvelope) -> Result Result, ServiceError> { - response_frame(Operation::SpendReadout, envelope) + response_frame(Operation::UsageReadout, envelope) } /// Opens the carriage of one frame. Everything the carriage says is checked @@ -193,7 +193,7 @@ fn read_carriage(bytes: &[u8]) -> Result { #[cfg(any(feature = "daemon", test))] pub(crate) enum RequestFrame { Invoke(TransportEnvelope), - SpendReadout(TransportEnvelope), + UsageReadout(TransportEnvelope), } #[cfg(any(feature = "daemon", test))] @@ -215,7 +215,7 @@ pub(crate) fn decode_request_frame(bytes: &[u8]) -> Result Ok(RequestFrame::Invoke(envelope)), - (Operation::SpendReadout, envelope) => Ok(RequestFrame::SpendReadout(envelope)), + (Operation::UsageReadout, envelope) => Ok(RequestFrame::UsageReadout(envelope)), } } _ => Err(ServiceError::Carriage), @@ -226,7 +226,7 @@ pub(crate) fn decode_request_frame(bytes: &[u8]) -> Result Result { match decode_request_frame(bytes)? { RequestFrame::Invoke(envelope) => Ok(envelope), - RequestFrame::SpendReadout(_) => Err(ServiceError::Carriage), + RequestFrame::UsageReadout(_) => Err(ServiceError::Carriage), } } @@ -263,10 +263,10 @@ pub(crate) fn decode_result_frame(bytes: &[u8]) -> Result Result { - decode_response_frame(bytes, Operation::SpendReadout) + decode_response_frame(bytes, Operation::UsageReadout) } /// The most bytes one frame holding a sealed envelope of `ciphertext_len` @@ -416,9 +416,9 @@ pub fn decrypt_result( Ok(result) } -/// A spend readout request carries no content: the sealed envelope's +/// A usage readout request carries no content: the sealed envelope's /// authenticated identity and kind are the whole request. -pub(crate) fn encrypt_spend_readout_request( +pub(crate) fn encrypt_usage_readout_request( caller_runtime_id: &str, request_id: &str, security: &TransportKey, @@ -426,45 +426,45 @@ pub(crate) fn encrypt_spend_readout_request( encrypt_message( caller_runtime_id, request_id, - EnvelopeKind::SpendReadoutRequest, + EnvelopeKind::UsageReadoutRequest, &[], security, ) } #[cfg(any(feature = "daemon", test))] -pub(crate) fn decrypt_spend_readout_request( +pub(crate) fn decrypt_usage_readout_request( envelope: &TransportEnvelope, security: &TransportKey, ) -> Result<(), ServiceError> { - validate_envelope(envelope, EnvelopeKind::SpendReadoutRequest)?; + validate_envelope(envelope, EnvelopeKind::UsageReadoutRequest)?; decrypt_message(envelope, security).map(drop) } #[cfg(feature = "daemon")] -pub(crate) fn encrypt_spend_readout( +pub(crate) fn encrypt_usage_readout( caller_runtime_id: &str, request_id: &str, - readout: &SpendReadout, + readout: &UsageReadout, security: &TransportKey, ) -> Result { let plaintext = wire_encode(readout).map_err(|_| ServiceError::Encoding)?; encrypt_message( caller_runtime_id, request_id, - EnvelopeKind::SpendReadoutResult, + EnvelopeKind::UsageReadoutResult, &plaintext, security, ) } -pub(crate) fn decrypt_spend_readout( +pub(crate) fn decrypt_usage_readout( envelope: &TransportEnvelope, security: &TransportKey, -) -> Result { - validate_envelope(envelope, EnvelopeKind::SpendReadoutResult)?; +) -> Result { + validate_envelope(envelope, EnvelopeKind::UsageReadoutResult)?; let plaintext = decrypt_message(envelope, security)?; - let readout: SpendReadout = wire_decode(&plaintext).map_err(|_| ServiceError::Encoding)?; + let readout: UsageReadout = wire_decode(&plaintext).map_err(|_| ServiceError::Encoding)?; readout.validate()?; Ok(readout) } @@ -607,8 +607,6 @@ 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")] @@ -1072,7 +1070,7 @@ mod tests { fn golden_frame() -> Vec { std::fs::read( std::path::Path::new(env!("CARGO_MANIFEST_DIR")) - .join("tests/fixtures/wire_v3/carriage_invocation.frame"), + .join("tests/fixtures/wire_v4/carriage_invocation.frame"), ) .expect("carriage_invocation.frame is committed") } @@ -1109,7 +1107,7 @@ mod tests { assert_eq!(text(&request, "operation"), "invoke"); assert_eq!( text(&request, "payloadSchema"), - "gamecult.model.transport_envelope.v3" + "gamecult.model.transport_envelope.v4" ); assert_eq!(text(&request, "payloadEncoding"), "messagepack-base64"); assert_eq!(text(&request, "messageId"), envelope.request_id); @@ -1123,7 +1121,7 @@ mod tests { assert_eq!(text(&response, "status"), "ok"); assert_eq!( text(&response, "payloadSchema"), - "gamecult.model.transport_envelope.v3" + "gamecult.model.transport_envelope.v4" ); assert_eq!(text(&response, "payloadEncoding"), "messagepack-base64"); assert_eq!(text(&response, "messageId"), envelope.request_id); @@ -1287,23 +1285,23 @@ mod tests { #[test] fn a_readout_request_is_not_an_invocation_and_names_its_own_operation() { let key = TransportKey::from_connection_secret("readout-test-key").unwrap(); - let sealed = encrypt_spend_readout_request("caller-1", "readout-1", &key).unwrap(); - let frame = encode_spend_readout_request_frame(&sealed).unwrap(); + let sealed = encrypt_usage_readout_request("caller-1", "readout-1", &key).unwrap(); + let frame = encode_usage_readout_request_frame(&sealed).unwrap(); assert!(matches!( decode_request_frame(&frame), - Ok(RequestFrame::SpendReadout(found)) if found == sealed + Ok(RequestFrame::UsageReadout(found)) if found == sealed )); assert_eq!(decode_invocation_frame(&frame), Err(ServiceError::Carriage)); // Sealed as one kind, the authenticated data refuses it as another. - assert!(decrypt_spend_readout_request(&sealed, &key).is_ok()); + assert!(decrypt_usage_readout_request(&sealed, &key).is_ok()); assert_eq!( - decrypt_spend_readout(&sealed, &key).unwrap_err(), + decrypt_usage_readout(&sealed, &key).unwrap_err(), ServiceError::Envelope ); let mut relabelled = sealed.clone(); - relabelled.message_kind = EnvelopeKind::SpendReadoutResult; + relabelled.message_kind = EnvelopeKind::UsageReadoutResult; assert_eq!( - decrypt_spend_readout(&relabelled, &key).unwrap_err(), + decrypt_usage_readout(&relabelled, &key).unwrap_err(), ServiceError::Encryption ); } @@ -1313,12 +1311,12 @@ mod tests { fn a_response_is_read_as_its_own_operation_only() { use crate::readout::PlanQuotaSnapshot; let key = TransportKey::from_connection_secret("readout-test-key").unwrap(); - let readout = SpendReadout::new(1, Vec::new(), PlanQuotaSnapshot::default()); - let sealed = encrypt_spend_readout("caller-1", "readout-1", &readout, &key).unwrap(); - let readout_frame = encode_spend_readout_response_frame(&sealed).unwrap(); + let readout = UsageReadout::new(1, PlanQuotaSnapshot::default()); + let sealed = encrypt_usage_readout("caller-1", "readout-1", &readout, &key).unwrap(); + let readout_frame = encode_usage_readout_response_frame(&sealed).unwrap(); let invoke_frame = encode_result_frame(&sealed).unwrap(); assert_eq!( - decode_spend_readout_response_frame(&readout_frame).unwrap(), + decode_usage_readout_response_frame(&readout_frame).unwrap(), sealed ); assert_eq!(decode_result_frame(&invoke_frame).unwrap(), sealed); @@ -1327,9 +1325,9 @@ mod tests { Err(ServiceError::Carriage) ); assert_eq!( - decode_spend_readout_response_frame(&invoke_frame), + decode_usage_readout_response_frame(&invoke_frame), Err(ServiceError::Carriage) ); - assert_eq!(decrypt_spend_readout(&sealed, &key).unwrap(), readout); + assert_eq!(decrypt_usage_readout(&sealed, &key).unwrap(), readout); } } diff --git a/crates/model-connector/src/fixtures.rs b/crates/model-connector/src/fixtures.rs index 079a99c..1bfc37a 100644 --- a/crates/model-connector/src/fixtures.rs +++ b/crates/model-connector/src/fixtures.rs @@ -133,16 +133,9 @@ 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], MAX_COST, request()).unwrap() + TransportInvocation::new("epiphany-yggdrasil", 2_000, [7; 32], request()).unwrap() } pub(crate) fn receipt(invocation: &TransportInvocation) -> TransportReceipt { @@ -155,8 +148,6 @@ 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, @@ -170,33 +161,6 @@ 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( @@ -212,8 +176,6 @@ pub(crate) fn caller( max_concurrent_requests, 64 * 1024, 32_768, - SpendCeiling(CEILING), - SpendWindowMs(WINDOW_MS), ) .unwrap() } @@ -223,7 +185,6 @@ pub(crate) fn open_service(store: &TestStore, max_concurrent_requests: usize) -> TransportService::open( &store.path, "fixture-sse", - fixture_prices(), [ caller( "epiphany-yggdrasil", @@ -249,5 +210,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], MAX_COST, request).unwrap() + TransportInvocation::new(caller_runtime_id, 2_000, [7; 32], request).unwrap() } diff --git a/crates/model-connector/src/readout.rs b/crates/model-connector/src/readout.rs index 36d4cb9..197e65a 100644 --- a/crates/model-connector/src/readout.rs +++ b/crates/model-connector/src/readout.rs @@ -1,13 +1,14 @@ -//! The spend readout: what each caller has spent and holds in its current -//! window, and the plan's own usage as the backend last saw it. The ledger is -//! the replay store's; a readout only reads it. +//! The usage readout: the subscription plan's own usage windows as the backend +//! last saw them. The connector meters nothing in dollars; it publishes what +//! the provider reports, and the minds read it and decide their own burn. use schemars::JsonSchema; use serde::{Deserialize, Serialize}; use crate::contract::ContractError; -pub const SPEND_READOUT_SCHEMA_ID: &str = "gamecult.model.spend_readout.v1"; +pub const USAGE_READOUT_SCHEMA_ID: &str = "gamecult.model.usage_readout.v1"; + /// One window of the subscription plan, as the provider reported it. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] @@ -32,40 +33,25 @@ pub struct PlanQuotaSnapshot { pub windows: Vec, } -/// One caller's current spend window in micro-USD. `spent` is settled charges -/// plus what a crash left counted at its full reservation; `reserved` is what -/// calls in flight hold, and both count against the ceiling. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] #[serde(deny_unknown_fields)] -pub struct CallerSpend { - pub caller_id: String, - pub window_ms: u64, - pub ceiling_micro_usd: u64, - pub spent_micro_usd: u64, - pub reserved_micro_usd: u64, -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] -#[serde(deny_unknown_fields)] -pub struct SpendReadout { +pub struct UsageReadout { pub schema_id: String, pub at_unix_ms: u64, - pub callers: Vec, pub plan_quota: Vec, } -impl SpendReadout { - pub fn new(at_unix_ms: u64, callers: Vec, plan: PlanQuotaSnapshot) -> Self { +impl UsageReadout { + pub fn new(at_unix_ms: u64, plan: PlanQuotaSnapshot) -> Self { Self { - schema_id: SPEND_READOUT_SCHEMA_ID.to_string(), + schema_id: USAGE_READOUT_SCHEMA_ID.to_string(), at_unix_ms, - callers, plan_quota: plan.windows, } } pub(crate) fn validate(&self) -> Result<(), ContractError> { - if self.schema_id != SPEND_READOUT_SCHEMA_ID { + if self.schema_id != USAGE_READOUT_SCHEMA_ID { return Err(ContractError::Schema); } Ok(()) @@ -77,16 +63,9 @@ mod tests { use super::*; use crate::contract::{wire_decode, wire_encode}; - fn readout() -> SpendReadout { - SpendReadout::new( + fn readout() -> UsageReadout { + UsageReadout::new( 1_000, - vec![CallerSpend { - caller_id: "eureka".into(), - window_ms: 3_600_000, - ceiling_micro_usd: 50_000_000, - spent_micro_usd: 7, - reserved_micro_usd: 3, - }], PlanQuotaSnapshot { windows: vec![PlanQuota { window: "five_hour".into(), @@ -107,21 +86,21 @@ mod tests { panic!("a readout is a map, not a positional array"); }; let names: Vec<&str> = fields.iter().filter_map(|(k, _)| k.as_str()).collect(); - assert_eq!(names, ["schema_id", "at_unix_ms", "callers", "plan_quota"]); - assert_eq!(wire_decode::(&bytes).unwrap(), readout()); + assert_eq!(names, ["schema_id", "at_unix_ms", "plan_quota"]); + assert_eq!(wire_decode::(&bytes).unwrap(), readout()); } #[test] fn an_unobserved_plan_is_none_not_zero() { - let bare = SpendReadout::new(1, Vec::new(), PlanQuotaSnapshot::default()); - let decoded: SpendReadout = wire_decode(&wire_encode(&bare).unwrap()).unwrap(); + let bare = UsageReadout::new(1, PlanQuotaSnapshot::default()); + let decoded: UsageReadout = wire_decode(&wire_encode(&bare).unwrap()).unwrap(); assert!(decoded.plan_quota.is_empty()); } #[test] fn a_readout_of_another_schema_is_refused() { let mut other = readout(); - other.schema_id = "gamecult.model.spend_readout.v2".into(); + other.schema_id = "gamecult.model.usage_readout.v2".into(); assert!(other.validate().is_err()); assert!(readout().validate().is_ok()); } diff --git a/crates/model-connector/src/serve.rs b/crates/model-connector/src/serve.rs index cd3663f..270fa1b 100644 --- a/crates/model-connector/src/serve.rs +++ b/crates/model-connector/src/serve.rs @@ -13,13 +13,12 @@ use zeroize::Zeroizing; use crate::contract::TransportResult; use crate::envelope::{ RequestFrame, ServiceError, TransportFrameError, decode_request_frame, encode_result_frame, - encode_spend_readout_response_frame, frame_bytes_for_payload, read_transport_frame, + encode_usage_readout_response_frame, frame_bytes_for_payload, read_transport_frame, write_transport_frame, }; use crate::readout::PlanQuotaSnapshot; use crate::service::{ - CallerAdmission, PriceTable, ProviderCall, SpendCeiling, SpendWindowMs, TransportAdmission, - TransportService, distinct_admissions, + CallerAdmission, ProviderCall, TransportAdmission, TransportService, distinct_admissions, }; /// The only provider seam of the connector kit. A connector binary implements @@ -34,8 +33,9 @@ pub trait ProviderBackend: Send + Sync { /// `serve` on: it names the provider on every receipt and decides which /// reasoning state is foreign. fn provider(&self) -> &str; - /// What the provider's own plan allows and has used, for the spend - /// readout. A provider with no plan to report keeps the default. + /// What the provider's own plan allows and has used, for the usage + /// readout and the provider's own refusals. A provider with no plan to report + /// keeps the default. fn plan_quota(&self) -> PlanQuotaSnapshot { PlanQuotaSnapshot::default() } @@ -123,12 +123,6 @@ 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> { @@ -149,8 +143,6 @@ 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)?, ); @@ -159,15 +151,12 @@ 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> { @@ -176,7 +165,6 @@ pub fn serve( let service = Arc::new(Mutex::new(TransportService::open( replay_store, backend.provider(), - prices, admissions, limits.max_expiry_skew_ms, )?)); @@ -184,8 +172,8 @@ pub fn serve( } /// `serve` over a service the caller opened and keeps a handle to, so another -/// front of the same ledger (the Self-run credential lease) reserves in the -/// same windows. The limits must already have been validated against the +/// front of the same service (the Self-run credential lease) shares its +/// replay store and plan. The limits must already have been validated against the /// service's admissions. pub fn serve_service( listener: TcpListener, @@ -197,7 +185,7 @@ pub fn serve_service( accept_loop(|| listener.accept(), limits, service, backend) } -/// What must be true before a serve loop touches the ledger: the listener is +/// What must be true before a serve loop touches the store: the listener is /// on loopback and the backend can take work. fn admit_listener(listener: &TcpListener, backend: &dyn ProviderBackend) -> Result<(), ServeError> { if !listener @@ -272,20 +260,25 @@ fn answer_frame( ) -> Result, ServeError> { let envelope = match decode_request_frame(request).map_err(|_| ServeError::FrameEncoding)? { RequestFrame::Invoke(envelope) => envelope, - RequestFrame::SpendReadout(envelope) => { + RequestFrame::UsageReadout(envelope) => { let plan = backend.plan_quota(); let response = service .lock() .map_err(|_| ServeError::ServicePoisoned)? - .spend_readout(&envelope, unix_ms()?, plan, limits.max_frame_bytes)?; - return encode_spend_readout_response_frame(&response) + .usage_readout(&envelope, unix_ms()?, plan, limits.max_frame_bytes)?; + return encode_usage_readout_response_frame(&response) .map_err(|_| ServeError::FrameEncoding); } }; let admission = service .lock() .map_err(|_| ServeError::ServicePoisoned)? - .begin(&envelope, unix_ms()?, limits.max_frame_bytes)?; + .begin( + &envelope, + unix_ms()?, + &backend.plan_quota(), + limits.max_frame_bytes, + )?; let response = match admission { TransportAdmission::Reply(response) => response, TransportAdmission::Execute(claim) => { @@ -390,7 +383,7 @@ mod tests { TransportInvocation, TransportOutcome, TransportResult, }; use crate::fixtures::{ - CEILING, MAX_COST, TestStore, WINDOW_MS, caller, fixture_prices, invocation_for, receipt, + TestStore, caller, invocation_for, receipt, request, }; @@ -472,7 +465,6 @@ mod tests { listener, limits, admissions, - fixture_prices(), &path, backend, )); @@ -506,7 +498,6 @@ mod tests { listener, limits(), admissions, - fixture_prices(), &path, backend, ) @@ -528,8 +519,6 @@ 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() } @@ -559,7 +548,7 @@ mod tests { } #[test] - fn serve_reads_out_spend_and_plan_quota_without_executing_anything() { + fn serve_reads_out_plan_usage_without_executing_anything() { let store = TestStore::new(); let inner = CountingBackend::new(true); let plan = PlanQuotaSnapshot { @@ -580,16 +569,9 @@ mod tests { assert_eq!(inner.executions.load(Ordering::SeqCst), 1); let before = now_ms(); - let readout = client.spend_readout(CALLER).unwrap(); + let readout = client.usage_readout(CALLER).unwrap(); assert!(readout.at_unix_ms >= before && readout.at_unix_ms <= now_ms()); assert_eq!(readout.plan_quota, plan.windows); - let [row] = readout.callers.as_slice() else { - panic!("one admitted caller"); - }; - assert_eq!(row.caller_id, CALLER); - assert_eq!((row.ceiling_micro_usd, row.window_ms), (CEILING, WINDOW_MS)); - assert!(row.spent_micro_usd > 0, "the settled call was charged"); - assert_eq!(row.reserved_micro_usd, 0); assert_eq!( inner.executions.load(Ordering::SeqCst), 1, @@ -604,39 +586,62 @@ mod tests { Some(Duration::from_secs(5)), ) .unwrap(); - assert!(stranger.spend_readout(CALLER).is_err()); + assert!(stranger.usage_readout(CALLER).is_err()); } #[test] fn a_backend_with_no_plan_reads_out_none() { let store = TestStore::new(); let (client, _) = start(CountingBackend::new(true), &store); - let readout = client.spend_readout(CALLER).unwrap(); + let readout = client.usage_readout(CALLER).unwrap(); assert!(readout.plan_quota.is_empty()); } + fn window_plan(permille: u32, status: &str, resets_at_unix_ms: u64) -> PlanQuotaSnapshot { + PlanQuotaSnapshot { + windows: vec![crate::readout::PlanQuota { + window: "five_hour".into(), + utilization_permille: Some(permille), + resets_at_unix_ms, + status: status.into(), + observed_at_unix_ms: 1, + }], + } + } + #[test] - fn serve_refuses_a_request_over_its_max_cost_without_executing_it() { + fn serve_repeats_the_providers_rejection_and_executes_nothing_while_it_stands() { 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(); + let inner = CountingBackend::new(true); + let backend = Arc::new(PlanBackend { + inner: inner.clone(), + plan: window_plan(1_000, "rejected", now_ms() + 3_600_000), + }); + let (client, _) = start_with(backend, &store, vec![caller(CALLER, KEY, 4)]); + let refused = client.execute(&fresh("request-rejected")).unwrap(); assert_eq!( refused.disposition, - TransportDisposition::Refused(Refusal::CostBound) + TransportDisposition::Refused(Refusal::ProviderRejected) ); - assert_eq!(backend.executions.load(Ordering::SeqCst), 0); + assert_eq!(inner.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(); + #[test] + fn serve_admits_a_full_window_the_provider_still_allows() { + let store = TestStore::new(); + let inner = CountingBackend::new(true); + let backend = Arc::new(PlanBackend { + inner: inner.clone(), + plan: window_plan(999, "allowed", now_ms() + 3_600_000), + }); + let (client, _) = start_with(backend, &store, vec![caller(CALLER, KEY, 4)]); + let admitted = client.execute(&fresh("request-full")).unwrap(); assert!(matches!( admitted.disposition, TransportDisposition::Transported { .. } )); + assert_eq!(inner.executions.load(Ordering::SeqCst), 1); + } assert_eq!(backend.executions.load(Ordering::SeqCst), 1); } @@ -676,7 +681,6 @@ mod tests { CALLER, invocation.expires_at_unix_ms, [7; 32], - MAX_COST, conflicting_request, ) .unwrap(); @@ -761,7 +765,6 @@ mod tests { let mut service = TransportService::open( &store.path, "fixture-sse", - fixture_prices(), vec![caller(CALLER, KEY, 4)], 300_000, ) @@ -770,7 +773,9 @@ mod tests { crate::envelope::encrypt_invocation(&invocation, &crate::fixtures::security(KEY)) .unwrap(); let TransportAdmission::Execute(claim) = - service.begin(&envelope, now_ms(), 1 << 20).unwrap() + service + .begin(&envelope, now_ms(), &PlanQuotaSnapshot::default(), 1 << 20) + .unwrap() else { panic!("a fresh request must execute"); }; @@ -793,7 +798,6 @@ mod tests { listener, narrowed, vec![caller(CALLER, KEY, 4)], - fixture_prices(), &path, served, ) @@ -1001,8 +1005,6 @@ mod tests { slots, payload, 32_768, - SpendCeiling(CEILING), - SpendWindowMs(WINDOW_MS), ) .unwrap() } @@ -1110,7 +1112,7 @@ mod tests { ..limits() }; let admissions = vec![admission(CALLER, KEY, 4, bound)]; - thread::spawn(move || serve(listener, tight, admissions, fixture_prices(), &path, served)); + thread::spawn(move || serve(listener, tight, admissions, &path, served)); let client = ConnectorClient::new(endpoint, KEY, frame, Some(Duration::from_secs(5))).unwrap(); let reply = client.execute(&invocation).unwrap(); @@ -1225,7 +1227,6 @@ mod tests { let mut service = TransportService::open( &store.path, "fixture-sse", - fixture_prices(), vec![caller(CALLER, KEY, 4)], 300_000, ) @@ -1236,7 +1237,9 @@ mod tests { ) .unwrap(); let TransportAdmission::Execute(claim) = - service.begin(&envelope, now_ms(), 1 << 20).unwrap() + service + .begin(&envelope, now_ms(), &PlanQuotaSnapshot::default(), 1 << 20) + .unwrap() else { panic!("a fresh request must execute"); }; @@ -1328,7 +1331,6 @@ mod tests { TransportService::open( &store.path, "fixture-sse", - fixture_prices(), vec![caller(CALLER, KEY, 4)], 300_000, ) @@ -1391,7 +1393,6 @@ mod tests { listener, limits(), admissions, - fixture_prices(), &path, served, ) diff --git a/crates/model-connector/src/service.rs b/crates/model-connector/src/service.rs index ea3dbeb..3e52dfb 100644 --- a/crates/model-connector/src/service.rs +++ b/crates/model-connector/src/service.rs @@ -1,4 +1,4 @@ -use std::collections::{BTreeMap, HashMap, HashSet}; +use std::collections::{HashMap, HashSet}; use std::path::Path; use cultcache_rs::{CacheBackingStore, CultCache, DatabaseEntry, OwnedRedbMessagePackBackingStore}; @@ -7,45 +7,16 @@ use sha2::{Digest, Sha256}; use zeroize::Zeroizing; use crate::contract::{ - 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, + ContractError, FailureKind, InputItem, ProviderRequest, RECEIPT_SCHEMA_ID, RESULT_SCHEMA_ID, + Refusal, TransportDisposition, TransportEvent, TransportInvocation, TransportKey, + TransportOutcome, TransportReceipt, TransportResult, canonical_invocation_bytes, require_id, }; use crate::envelope::{ EnvelopeKind, MAX_IDENTITY_BYTES, ServiceError, TransportEnvelope, decrypt_invocation, - decrypt_spend_readout_request, encode_result_frame, encode_spend_readout_response_frame, - encrypt_result, encrypt_spend_readout, validate_envelope, + decrypt_usage_readout_request, encode_result_frame, encode_usage_readout_response_frame, + encrypt_result, encrypt_usage_readout, validate_envelope, }; -use crate::readout::{CallerSpend, PlanQuotaSnapshot, SpendReadout}; - -/// 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); +use crate::readout::{PlanQuotaSnapshot, UsageReadout}; /// The rule that admitted callers are distinct: no two share a connection key /// (a key names its caller), and no two share a caller id. `TransportService` @@ -73,14 +44,9 @@ 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, @@ -89,10 +55,7 @@ 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::>(); @@ -105,8 +68,6 @@ 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); } @@ -120,8 +81,6 @@ impl CallerAdmission { max_concurrent_requests, max_payload_bytes, max_output_tokens, - spend_ceiling_micro_usd, - spend_window_ms, }) } @@ -138,206 +97,24 @@ impl CallerAdmission { } } -/// What one million tokens of each class cost, in micro-USD. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -pub struct Rates { - pub input: u64, - pub cache_write_short: u64, - pub cache_write_long: u64, - pub cache_read: u64, - pub output: u64, +/// The status a window carries while the provider itself refuses. +const REJECTED: &str = "rejected"; + +/// Whether the provider is refusing right now: some window of `plan` is +/// `rejected` and its own reset lies in the future. The one usage-derived +/// refusal there is. It repeats the provider's word with the provider's +/// expiry, and no utilization figure is ever compared to anything. +pub(crate) fn provider_rejects(plan: &PlanQuotaSnapshot, now_unix_ms: u64) -> bool { + plan.windows + .iter() + .any(|window| window.status == REJECTED && window.resets_at_unix_ms > now_unix_ms) } -/// The rates a prompt strictly above `above_prompt_tokens` is priced at. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -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, Serialize, Deserialize)] -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, request)| { - request.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); - } -} - -/// A request between claim and settlement: the window it was admitted in and -/// the reservation it holds there. +/// A Self run's or Persona turn's credential lease while its holder keeps it. #[derive(Debug, Clone, Copy)] -struct LiveRequest { - window: u64, - reserved_micro_usd: u64, - /// A Self run's credential lease, not a provider call: it holds its - /// reservation in the window like a call does but takes no request slot. - lease: Option, +struct LiveLease { + turn: LeaseTurn, + expires_at_unix_ms: u64, } #[derive(Debug, Clone, PartialEq, Eq, Hash)] @@ -348,27 +125,17 @@ struct RequestIdentity { #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] enum ReplayState { - /// 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. + /// Admitted and possibly executing. + Active, Completed { response: TransportEnvelope, - charged_micro_usd: u64, - overrun_micro_usd: u64, - }, - /// A credential lease that settled: it has no answer to replay, only the - /// charge it left in its window and how far that passed the reservation. - LeaseSettled { - charged_micro_usd: u64, - overrun_micro_usd: u64, }, } +/// Idempotency only: what was asked and what was answered. Nothing here is a +/// charge. #[derive(Clone, Debug, PartialEq, Eq, DatabaseEntry)] -#[cultcache(type = "gamecult.model.replay_record.v2")] +#[cultcache(type = "gamecult.model.replay_record.v3")] struct ReplayRecord { #[cultcache(key = 0)] caller_runtime_id: String, @@ -382,8 +149,6 @@ 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 @@ -395,8 +160,6 @@ struct ReplayRecord { pub struct ProviderCall { request: ProviderRequest, receipt: TransportReceipt, - price: ModelPrice, - reserved_micro_usd: u64, } impl ProviderCall { @@ -418,69 +181,20 @@ impl ProviderCall { } /// The result of this call. The whole receipt (schema, ids, both digests, - /// model, provider tag, drop count and cost) is derived here, so a backend + /// model, provider tag and drop count) 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, { - let cost_micro_usd = self.cost_of(&outcome); + TransportResult::transported( + events, 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, - } + ) } } @@ -505,39 +219,30 @@ pub enum TransportAdmission { pub struct TransportService { callers: HashMap, - /// The requests between claim and settlement, each with the window it was - /// admitted in. - live: HashMap, + /// The requests between claim and completion. + live: HashSet, + /// The credential leases their holders keep, in memory only: a lease + /// decides `Busy` and nothing else, and ends when its holder lets go or + /// its expiry passes. + leases: 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. `prices` must price every model any - /// caller is allowed. + /// what provider every receipt names. 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, ) @@ -547,7 +252,6 @@ impl TransportService { pub fn with_store( store: impl CacheBackingStore + 'static, provider: &str, - prices: PriceTable, admissions: impl IntoIterator, max_expiry_skew_ms: u64, ) -> Result { @@ -566,14 +270,6 @@ 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 @@ -582,11 +278,12 @@ impl TransportService { replay_store .add_generic_backing_store(store) .map_err(replay_store_error)?; + // A record of any other type id (the dollar era's v2 included) has no + // registered schema, so the pull fails and the store is refused. 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)? @@ -598,7 +295,6 @@ 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); } @@ -608,103 +304,36 @@ 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, .. - } - | ReplayState::LeaseSettled { - 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, + live: HashSet::new(), + leases: HashMap::new(), 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) - } - - /// Answers one sealed spend readout request: for each admitted caller, what - /// its current window holds against its ceiling, with the backend's plan - /// quota beside it. Read-only: nothing is reserved, recorded or pruned. - /// Any admitted caller reads every caller's row; the spend is the - /// operator's and the callers are all hers. `max_frame_bytes` bounds the - /// reply like it bounds `complete`'s. - pub fn spend_readout( + /// Answers one sealed usage readout request with the backend's plan quota + /// as it stands. Read-only: nothing is recorded. Any admitted caller reads + /// it; the usage is the operator's and the callers are all hers. + /// `max_frame_bytes` bounds the reply like it bounds `complete`'s. + pub fn usage_readout( &self, envelope: &TransportEnvelope, now_unix_ms: u64, plan: PlanQuotaSnapshot, max_frame_bytes: usize, ) -> Result { - validate_envelope(envelope, EnvelopeKind::SpendReadoutRequest)?; + validate_envelope(envelope, EnvelopeKind::UsageReadoutRequest)?; let asker = self .callers .get(&envelope.caller_runtime_id) .ok_or(ServiceError::CallerNotAdmitted)?; - decrypt_spend_readout_request(envelope, &asker.security)?; - let mut callers: Vec = self - .callers - .values() - .map(|caller| { - // The window admission would use right now. - let window = admission_clock( - &self.window_spend, - &caller.caller_runtime_id, - now_unix_ms, - caller.spend_window_ms, - ) / caller.spend_window_ms; - let held = self.window_spend(&caller.caller_runtime_id, window); - let reserved: u64 = self - .live - .iter() - .filter(|(identity, request)| { - request.window == window - && identity.caller_runtime_id == caller.caller_runtime_id - }) - .fold(0, |sum, (_, request)| { - sum.saturating_add(request.reserved_micro_usd) - }); - CallerSpend { - caller_id: caller.caller_runtime_id.clone(), - window_ms: caller.spend_window_ms, - ceiling_micro_usd: caller.spend_ceiling_micro_usd, - spent_micro_usd: held.saturating_sub(reserved), - reserved_micro_usd: reserved, - } - }) - .collect(); - callers.sort_by(|a, b| a.caller_id.cmp(&b.caller_id)); - let readout = SpendReadout::new(now_unix_ms, callers, plan); - let sealed = encrypt_spend_readout( + decrypt_usage_readout_request(envelope, &asker.security)?; + let readout = UsageReadout::new(now_unix_ms, plan); + let sealed = encrypt_usage_readout( &envelope.caller_runtime_id, &envelope.request_id, &readout, @@ -716,7 +345,9 @@ impl TransportService { Ok(sealed) } - /// Admits one request. `max_frame_bytes` is the frame every reply must + /// Admits one request. `plan` is the backend's plan quota as it stands: a + /// provider that is refusing right now is repeated as `ProviderRejected` + /// for new work. `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 /// with the typed `result_size` failure, which is recorded in its place. @@ -727,6 +358,7 @@ impl TransportService { &mut self, envelope: &TransportEnvelope, now_unix_ms: u64, + plan: &PlanQuotaSnapshot, max_frame_bytes: usize, ) -> Result { validate_envelope(envelope, EnvelopeKind::Invocation)?; @@ -737,8 +369,6 @@ impl TransportService { 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( @@ -783,25 +413,15 @@ impl TransportService { return self.reply_refusal(&invocation, Refusal::ReplayConflict, &security); } return match &record.state { - ReplayState::Completed { - response, - charged_micro_usd, - overrun_micro_usd, - } => { + ReplayState::Completed { response } => { 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( @@ -809,8 +429,6 @@ impl TransportService { &ReplayRecord { state: ReplayState::Completed { response: stand_in.clone(), - charged_micro_usd: *charged_micro_usd, - overrun_micro_usd: *overrun_micro_usd, }, ..record.clone() }, @@ -818,11 +436,8 @@ impl TransportService { .map_err(replay_store_error)?; Ok(TransportAdmission::Reply(stand_in)) } - ReplayState::LeaseSettled { .. } => { - self.reply_refusal(&invocation, Refusal::ReplayConflict, &security) - } - ReplayState::Active { .. } => { - let refusal = if self.live.contains_key(&identity) { + ReplayState::Active => { + let refusal = if self.live.contains(&identity) { Refusal::InFlight } else { Refusal::Indeterminate @@ -831,87 +446,35 @@ 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 + // The provider's own refusal binds new work only: a request already + // recorded is answered whatever the provider says now. + if provider_rejects(plan, now_unix_ms) { + return self.reply_refusal(&invocation, Refusal::ProviderRejected, &security); + } + if self .live .iter() - .filter(|(key, request)| { - request.lease.is_none() && key.caller_runtime_id == invocation.caller_runtime_id - }) - .count(); - if active_for_caller >= max_concurrent_requests { + .filter(|key| key.caller_runtime_id == invocation.caller_runtime_id) + .count() + >= 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. + // Claiming the request is one write: a store 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, invocation_sha256, expires_at_unix_ms: invocation.expires_at_unix_ms, - state: ReplayState::Active { reserved_micro_usd }, - admitted_at_unix_ms, + state: ReplayState::Active, }; self.replay_store .put(replay_key(&identity), &replay_record) .map_err(replay_store_error)?; - add_spend( - &mut self.window_spend, - &identity.caller_runtime_id, - window, - 0, - reserved_micro_usd, - ); - self.live.insert( - identity, - LiveRequest { - window, - reserved_micro_usd, - lease: None, - }, - ); + self.live.insert(identity); Ok(TransportAdmission::Execute(Box::new(ExecutionClaim { call: ProviderCall { receipt: TransportReceipt { @@ -924,8 +487,6 @@ 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(), @@ -933,8 +494,6 @@ impl TransportService { }, }, request, - price, - reserved_micro_usd, }, invocation, invocation_sha256, @@ -957,7 +516,7 @@ impl TransportService { caller_runtime_id: invocation.caller_runtime_id.clone(), request_id: invocation.request_id().to_string(), }; - if !self.live.contains_key(&identity) { + if !self.live.contains(&identity) { return Err(ServiceError::NoActiveRequest); } let active = self @@ -965,10 +524,8 @@ impl TransportService { .get::(&replay_key(&identity)) .map_err(replay_store_error)? .ok_or(ServiceError::NoActiveRequest)?; - let ReplayState::Active { reserved_micro_usd } = active.state else { - return Err(ServiceError::ActiveRequestMismatch); - }; - if active.invocation_sha256 != claim.invocation_sha256 { + if active.state != ReplayState::Active || active.invocation_sha256 != claim.invocation_sha256 + { return Err(ServiceError::ActiveRequestMismatch); } let caller = self @@ -977,47 +534,42 @@ impl TransportService { .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).inspect_err(|_| { + let response = if self.admits_result(&result, &claim) { + let sealed = encrypt_result(&result, &security).inspect_err(|_| { self.live.remove(&identity); - })? + })?; + if fits_frame(&sealed, max_frame_bytes)? { + sealed + } else { + // Nothing is recorded when even the stand-in cannot be + // carried, but the caller's slot is not held by a request + // that can no longer finish. + self.seal_stand_in( + invocation, + &security, + max_frame_bytes, + FailureKind::ResultSize, + ) + .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); })? }; - if admissible && !fits_frame(&response, max_frame_bytes)? { - // Nothing is recorded when even the stand-in cannot be carried, - // but the caller's slot is not held by a request that can no - // longer finish. - response = self - .seal_stand_in( - invocation, - &security, - max_frame_bytes, - FailureKind::ResultSize, - charge, - ) - .inspect_err(|_| { - self.live.remove(&identity); - })?; - } let replay_record = ReplayRecord { caller_runtime_id: identity.caller_runtime_id.clone(), request_id: identity.request_id.clone(), @@ -1026,36 +578,18 @@ impl TransportService { 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. + // A failed write frees the caller's slot and answers the typed store + // failure; the Active record stays, so a retry is told the request is + // indeterminate. self.replay_store .put(replay_key(&identity), &replay_record) .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) } @@ -1066,7 +600,6 @@ impl TransportService { security: &TransportKey, max_frame_bytes: usize, kind: FailureKind, - charge: Charge, ) -> Result { let (_, foreign_reasoning_dropped) = partition_reasoning(&invocation.request, &self.provider) @@ -1080,7 +613,6 @@ impl TransportService { invocation.provider_request_sha256, foreign_reasoning_dropped, kind, - charge, ), security, )?; @@ -1092,8 +624,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, drop - /// count and cost the service derived. + /// bound to the caller's invocation, and carrying the provider tag and + /// drop count the service derived. fn admits_result(&self, result: &TransportResult, claim: &ExecutionClaim) -> bool { let TransportDisposition::Transported { receipt, .. } = &result.disposition else { return false; @@ -1101,8 +633,6 @@ 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( @@ -1122,12 +652,12 @@ fn fits_frame(envelope: &TransportEnvelope, max_frame_bytes: usize) -> Result Result { - encode_spend_readout_response_frame(envelope).map(|bytes| bytes.len() <= max_frame_bytes) + encode_usage_readout_response_frame(envelope).map(|bytes| bytes.len() <= max_frame_bytes) } /// The provider named on the stand-in failure, fixed so its size is bounded. @@ -1144,12 +674,7 @@ 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, @@ -1158,21 +683,11 @@ 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, @@ -1181,7 +696,6 @@ 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", @@ -1202,8 +716,6 @@ 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(), @@ -1214,12 +726,6 @@ fn stand_in_failure( } } -/// The one place foreignness is decided. A reasoning item of another provider -/// in the in-flight exchange (after the last user text; the whole input when -/// there is none) is refused, because dropping it would break the provider's -/// own tool-exchange rules. Earlier foreign items are removed and counted. -/// Foreignness is the provider tag alone: another model of this provider -/// passes through for the provider to judge. fn partition_reasoning( request: &ProviderRequest, provider: &str, @@ -1272,9 +778,7 @@ fn replay_store_error(error: impl std::fmt::Display) -> ServiceError { } mod lease; -pub use lease::{ - LeaseAdmission, LeaseGrant, LeaseRefusal, LeaseSettlement, LeaseTurn, LeaseUse, ModelUsage, -}; +pub use lease::{LeaseAdmission, LeaseGrant, LeaseRefusal, LeaseTurn}; fn contract_refusal(error: &ContractError) -> Refusal { match error { @@ -1300,7 +804,6 @@ mod tests { TransportService::open( &store.path, "fixture-sse", - fixture_prices(), [ caller("epiphany-yggdrasil", "shared-key", 1), caller("ghostlight-yggdrasil", "shared-key", 1), @@ -1315,7 +818,7 @@ mod tests { let mut service = open_service(&store, 1); assert!(matches!( service - .begin( + .begin_unplanned( &encrypt_invocation(&epiphany, &security("epiphany-distinct-test-key")) .unwrap(), 1_000, @@ -1326,7 +829,7 @@ mod tests { )); assert!(matches!( service - .begin( + .begin_unplanned( &encrypt_invocation(&ghostlight, &security("ghostlight-distinct-test-key")) .unwrap(), 1_000, @@ -1344,7 +847,7 @@ mod tests { let security = security("epiphany-distinct-test-key"); let mut service = open_service(&store, 1); let claim = match service - .begin( + .begin_unplanned( &encrypt_invocation(&invocation, &security).unwrap(), 1_000, FRAME, @@ -1375,7 +878,7 @@ mod tests { let mut service = open_service(&store, 1); let replay = match service - .begin( + .begin_unplanned( &encrypt_invocation(&invocation, &security).unwrap(), 1_000, FRAME, @@ -1398,12 +901,11 @@ mod tests { "epiphany-yggdrasil", 2_000, [7; 32], - MAX_COST, changed_request, ) .unwrap(); let refusal = match service - .begin( + .begin_unplanned( &encrypt_invocation(&conflicting, &security).unwrap(), 1_000, FRAME, @@ -1430,13 +932,13 @@ mod tests { let mut service = open_service(&store, 1); let first_envelope = encrypt_invocation(&first, &security).unwrap(); assert!(matches!( - service.begin(&first_envelope, 1_000, FRAME).unwrap(), + service.begin_unplanned(&first_envelope, 1_000, FRAME).unwrap(), TransportAdmission::Execute(_) )); for (invocation, expected) in [(&first, Refusal::InFlight), (&second, Refusal::Capacity)] { let response = match service - .begin( + .begin_unplanned( &encrypt_invocation(invocation, &security).unwrap(), 1_000, FRAME, @@ -1464,7 +966,7 @@ mod tests { let mut service = open_service(&store, 1); assert!(matches!( service - .begin( + .begin_unplanned( &encrypt_invocation(&first, &security).unwrap(), 1_000, FRAME @@ -1476,7 +978,7 @@ mod tests { let mut restarted = open_service(&store, 1); let ambiguous = match restarted - .begin( + .begin_unplanned( &encrypt_invocation(&first, &security).unwrap(), 1_000, FRAME, @@ -1494,7 +996,7 @@ mod tests { ); assert!(matches!( restarted - .begin( + .begin_unplanned( &encrypt_invocation(&second, &security).unwrap(), 1_000, FRAME @@ -1507,7 +1009,6 @@ mod tests { TransportService::open( &store.path, "fixture-sse", - fixture_prices(), [CallerAdmission::new( "epiphany-yggdrasil", "rotated-connection-key", @@ -1516,8 +1017,6 @@ mod tests { 1, 64 * 1024, 32_768, - SpendCeiling(CEILING), - SpendWindowMs(WINDOW_MS) ) .unwrap()], 5_000, @@ -1537,9 +1036,9 @@ mod tests { let mut request = invocation_for(CALLER, request_id).request; edit(&mut request); let invocation = - TransportInvocation::new(CALLER, 2_000, [7; 32], MAX_COST, request).unwrap(); + TransportInvocation::new(CALLER, 2_000, [7; 32], request).unwrap(); let envelope = encrypt_invocation(&invocation, &security(KEY)).unwrap(); - (service.begin(&envelope, 1_000, FRAME).unwrap(), invocation) + (service.begin_unplanned(&envelope, 1_000, FRAME).unwrap(), invocation) } fn refusal_of(admission: TransportAdmission, invocation: &TransportInvocation) -> Refusal { @@ -1617,19 +1116,16 @@ mod tests { 4, max_payload_bytes, 32_768, - SpendCeiling(CEILING), - SpendWindowMs(WINDOW_MS), ) .unwrap(); TransportService::open( &store.path, "fixture-sse", - fixture_prices(), [admission], 5_000, ) .unwrap() - .begin(&envelope, 1_000, FRAME) + .begin_unplanned(&envelope, 1_000, FRAME) }; assert!(matches!(admit(size), Ok(TransportAdmission::Execute(_)))); let refused = admit(size - 1).unwrap(); @@ -1644,14 +1140,14 @@ mod tests { let mut expired = invocation_for(CALLER, "expired"); expired.expires_at_unix_ms = 10; let reply = service - .begin(&raw_envelope(&expired), 1_000, FRAME) + .begin_unplanned(&raw_envelope(&expired), 1_000, FRAME) .unwrap(); assert_eq!(refusal_of(reply, &expired), Refusal::Expired); let mut substituted = invocation_for(CALLER, "substituted"); substituted.provider_request_sha256 = [9; 32]; let reply = service - .begin(&raw_envelope(&substituted), 1_000, FRAME) + .begin_unplanned(&raw_envelope(&substituted), 1_000, FRAME) .unwrap(); assert_eq!( refusal_of(reply, &substituted), @@ -1661,7 +1157,7 @@ mod tests { let mut malformed = invocation_for(CALLER, "malformed"); malformed.native_request_sha256 = [0; 32]; let reply = service - .begin(&raw_envelope(&malformed), 1_000, FRAME) + .begin_unplanned(&raw_envelope(&malformed), 1_000, FRAME) .unwrap(); assert_eq!(refusal_of(reply, &malformed), Refusal::Malformed); } @@ -1673,10 +1169,7 @@ mod tests { connection_key_epoch: 1, invocation_sha256: [1; 32], expires_at_unix_ms: 5_000, - state: ReplayState::Active { - reserved_micro_usd: 1, - }, - admitted_at_unix_ms: 1_000, + state: ReplayState::Active, }; edit(&mut record); let mut cache = CultCache::new(); @@ -1703,7 +1196,6 @@ mod tests { TransportService::open( &control.path, "fixture-sse", - fixture_prices(), [caller(CALLER, KEY, 1)], 5_000 ) @@ -1711,11 +1203,10 @@ mod tests { ); type Edit = fn(&mut ReplayRecord); - let cases: [(&str, Edit); 4] = [ + let cases: [(&str, Edit); 3] = [ ("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(); @@ -1730,7 +1221,6 @@ mod tests { TransportService::open( &store.path, "fixture-sse", - fixture_prices(), [caller(CALLER, KEY, 1)], 5_000 ), @@ -1771,8 +1261,6 @@ 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!( @@ -1792,8 +1280,6 @@ mod tests { 4, 64 * 1024, 32_768, - SpendCeiling(CEILING), - SpendWindowMs(WINDOW_MS), ) }; let at = "i".repeat(MAX_IDENTITY_BYTES); @@ -1812,7 +1298,7 @@ mod tests { let mut service = open_service(&store, 8); let ok = invocation_for(CALLER, &at); assert!(matches!( - service.begin(&raw_envelope(&ok), 1_000, FRAME).unwrap(), + service.begin_unplanned(&raw_envelope(&ok), 1_000, FRAME).unwrap(), TransportAdmission::Execute(_) )); } @@ -1830,7 +1316,7 @@ mod tests { let id = "i".repeat(length); let long = invocation_for(CALLER, &id); let TransportAdmission::Reply(reply) = - service.begin(&raw_envelope(&long), 1_000, FRAME).unwrap() + service.begin_unplanned(&raw_envelope(&long), 1_000, FRAME).unwrap() else { panic!("an over-long id must be refused, not executed ({length})"); }; @@ -1861,18 +1347,18 @@ mod tests { // The frame carries the id twice (the message id and the envelope's // request id) and the base64 of the envelope, so bracket it: the // ciphertext plus the id is a floor, twice the ciphertext and both - // ids plus the envelope's other fields is a ceiling. + // ids plus the envelope's other fields is a frame_cap. let floor = sealed.ciphertext.len() + id.len(); - let ceiling = 2 * (sealed.ciphertext.len() + 2 * id.len()) + 512; + let frame_cap = 2 * (sealed.ciphertext.len() + 2 * id.len()) + 512; let store = TestStore::new(); let mut service = open_service(&store, 1); assert!(matches!( - service.begin(&raw_envelope(&long), 1_000, ceiling).unwrap(), + service.begin_unplanned(&raw_envelope(&long), 1_000, frame_cap).unwrap(), TransportAdmission::Reply(_) )); assert_eq!( service - .begin(&raw_envelope(&long), 1_000, floor - 1) + .begin_unplanned(&raw_envelope(&long), 1_000, floor - 1) .unwrap_err(), ServiceError::FrameBound ); @@ -1890,9 +1376,9 @@ mod tests { #[test] fn a_readout_frame_check_is_inclusive_of_the_serialized_envelope() { - let readout = SpendReadout::new(1, Vec::new(), PlanQuotaSnapshot::default()); - let envelope = encrypt_spend_readout(CALLER, "sized", &readout, &security(KEY)).unwrap(); - let size = encode_spend_readout_response_frame(&envelope) + let readout = UsageReadout::new(1, PlanQuotaSnapshot::default()); + let envelope = encrypt_usage_readout(CALLER, "sized", &readout, &security(KEY)).unwrap(); + let size = encode_usage_readout_response_frame(&envelope) .unwrap() .len(); assert_eq!(readout_fits_frame(&envelope, size), Ok(true)); @@ -1905,7 +1391,7 @@ mod tests { ) -> (Box, TransportInvocation) { let invocation = invocation_for(CALLER, request_id); match service - .begin(&raw_envelope(&invocation), 1_000, FRAME) + .begin_unplanned(&raw_envelope(&invocation), 1_000, FRAME) .unwrap() { TransportAdmission::Execute(claim) => (claim, invocation), @@ -1943,7 +1429,6 @@ mod tests { invocation.native_request_sha256, invocation.provider_request_sha256, 0, - FIXTURE_USAGE_COST, ) ); let TransportDisposition::Transported { receipt, .. } = &stand_in.disposition else { @@ -1953,7 +1438,7 @@ mod tests { // The recorded replay is the stand-in's exact bytes. let TransportAdmission::Reply(replay) = service - .begin(&raw_envelope(&invocation), 1_000, FRAME) + .begin_unplanned(&raw_envelope(&invocation), 1_000, FRAME) .unwrap() else { panic!("expected a replay"); @@ -1988,7 +1473,7 @@ mod tests { // longer in flight. let (_next, _) = claim_for(&mut service, "after-wedge"); let TransportAdmission::Reply(reply) = service - .begin(&raw_envelope(&invocation), 1_000, FRAME) + .begin_unplanned(&raw_envelope(&invocation), 1_000, FRAME) .unwrap() else { panic!("expected a refusal"); @@ -2013,7 +1498,7 @@ mod tests { let mut service = open_service(&store, 8); let TransportAdmission::Reply(same) = service - .begin(&raw_envelope(&invocation), 1_000, FRAME) + .begin_unplanned(&raw_envelope(&invocation), 1_000, FRAME) .unwrap() else { panic!("expected a replay"); @@ -2023,7 +1508,7 @@ mod tests { let mut service = open_service(&store, 8); let TransportAdmission::Reply(smaller) = service - .begin(&raw_envelope(&invocation), 1_000, 2000) + .begin_unplanned(&raw_envelope(&invocation), 1_000, 2000) .unwrap() else { panic!("expected a replay"); @@ -2038,12 +1523,11 @@ 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. let TransportAdmission::Reply(again) = service - .begin(&raw_envelope(&invocation), 1_000, FRAME) + .begin_unplanned(&raw_envelope(&invocation), 1_000, FRAME) .unwrap() else { panic!("expected a replay"); @@ -2059,7 +1543,7 @@ mod tests { let big = text_result(claim.call(), "y"); service.complete(*claim, big, FRAME).unwrap(); assert_eq!( - service.begin(&raw_envelope(&invocation), 1_000, 10).err(), + service.begin_unplanned(&raw_envelope(&invocation), 1_000, 10).err(), Some(ServiceError::FrameBound) ); } @@ -2077,7 +1561,7 @@ mod tests { // Still active, never completed, and no longer live: the retry is told the // outcome is indeterminate. let TransportAdmission::Reply(reply) = service - .begin(&raw_envelope(&invocation), 1_000, FRAME) + .begin_unplanned(&raw_envelope(&invocation), 1_000, FRAME) .unwrap() else { panic!("expected a refusal"); @@ -2277,7 +1761,6 @@ mod tests { TransportService::open( &store.path, tag, - fixture_prices(), [caller(CALLER, KEY, 1)], 5_000 ), @@ -2290,7 +1773,6 @@ mod tests { TransportService::open( &store.path, "fixture-sse", - fixture_prices(), [caller(CALLER, KEY, 1)], 5_000 ) @@ -2338,7 +1820,7 @@ mod tests { )); // A retry of the same id replays the failure, not ReplayConflict. let TransportAdmission::Reply(replay) = service - .begin(&raw_envelope(&invocation), 1_000, FRAME) + .begin_unplanned(&raw_envelope(&invocation), 1_000, FRAME) .unwrap() else { panic!("a retry must replay the recorded failure"); @@ -2413,7 +1895,6 @@ mod tests { invocation.native_request_sha256, invocation.provider_request_sha256, 1, - FIXTURE_USAGE_COST, ); assert_eq!( decrypt_result(&response, &security(KEY), &invocation).unwrap(), @@ -2428,7 +1909,7 @@ mod tests { let big = text_result(claim.call(), &"y".repeat(4000)); service.complete(*claim, big, FRAME).unwrap(); let TransportAdmission::Reply(smaller) = service - .begin(&raw_envelope(&invocation), 1_000, 2000) + .begin_unplanned(&raw_envelope(&invocation), 1_000, 2000) .unwrap() else { panic!("expected a replay"); @@ -2444,1435 +1925,174 @@ mod tests { assert_eq!(receipt.foreign_reasoning_dropped, 1); } - // ---- spend ---- + // ---- usage readout, the provider's refusal, the dollar era ---- - 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, + impl TransportService { + /// `begin` while the provider reports no plan at all. + fn begin_unplanned( + &mut self, + envelope: &TransportEnvelope, + now_unix_ms: u64, + max_frame_bytes: usize, + ) -> Result { + self.begin( + envelope, + now_unix_ms, + &PlanQuotaSnapshot::default(), + max_frame_bytes, + ) } } - 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 plan_of(permille: u32, status: &str, resets_at_unix_ms: u64) -> PlanQuotaSnapshot { + PlanQuotaSnapshot { + windows: vec![crate::readout::PlanQuota { + window: "five_hour".to_string(), + utilization_permille: Some(permille), + resets_at_unix_ms, + status: status.to_string(), + observed_at_unix_ms: 1_000, + }], } } - 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( + fn begin_under( service: &mut TransportService, request_id: &str, now: u64, - max_cost: u64, + plan: &PlanQuotaSnapshot, ) -> (TransportAdmission, TransportInvocation) { - let invocation = invocation_at(request_id, now, max_cost); + let mut request = invocation_for(CALLER, request_id).request; + request.request_id = request_id.to_string(); + let invocation = TransportInvocation::new(CALLER, now + 1_000, [7; 32], request).unwrap(); let admission = service - .begin(&raw_envelope(&invocation), now, FRAME) + .begin(&raw_envelope(&invocation), now, plan, 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 { .. } | ReplayState::LeaseSettled { .. } => { - 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) - ); - } - - // ---- spend readout ---- - fn readout_as( service: &TransportService, caller: &str, key: &str, - now: u64, + plan: PlanQuotaSnapshot, frame: usize, - ) -> Result { - let sealed = encrypt_spend_readout_request(caller, "readout-1", &security(key)).unwrap(); - let answer = service.spend_readout(&sealed, now, PlanQuotaSnapshot::default(), frame)?; + ) -> Result { + let sealed = encrypt_usage_readout_request(caller, "readout-1", &security(key)).unwrap(); + let answer = service.usage_readout(&sealed, 1_500, plan, frame)?; assert_eq!(answer.request_id, "readout-1"); - Ok(decrypt_spend_readout(&answer, &security(key)).unwrap()) - } - - fn row<'a>(readout: &'a SpendReadout, caller: &str) -> &'a CallerSpend { - readout - .callers - .iter() - .find(|row| row.caller_id == caller) - .expect("the caller has a row") + Ok(decrypt_usage_readout(&answer, &security(key)).unwrap()) } #[test] - fn spend_readout_reports_window() { + fn a_rejected_window_refuses_invoke_until_its_reset() { let store = TestStore::new(); let mut service = open_service(&store, 4); - let claim = claim_at(&mut service, "request-1", 1_000); + let rejected = plan_of(1_000, "rejected", 10_000); + let (refused, invocation) = begin_under(&mut service, "request-1", 9_999, &rejected); + assert_eq!(refusal_of(refused, &invocation), Refusal::ProviderRejected); + // At the provider's own reset the refusal lifts, and it was never + // recorded: the same id is admitted. + assert!(matches!( + begin_under(&mut service, "request-1", 10_000, &rejected).0, + TransportAdmission::Execute(_) + )); + } - // In flight: the reservation is held, nothing is spent. - let readout = readout_as(&service, CALLER, KEY, 1_500, FRAME).unwrap(); + #[test] + fn high_utilization_refuses_nothing() { + let store = TestStore::new(); + let mut service = open_service(&store, 4); + for (n, status) in ["allowed", "allowed_warning"].into_iter().enumerate() { + let plan = plan_of(999, status, 1_000_000); + assert!(matches!( + begin_under(&mut service, &format!("request-{n}"), 1_000, &plan).0, + TransportAdmission::Execute(_) + )); + } + } + + #[test] + fn a_recorded_answer_replays_while_the_provider_rejects() { + let store = TestStore::new(); + let mut service = open_service(&store, 4); + let (claim, invocation) = claim_for(&mut service, "request-1"); + let result = text_result(claim.call(), "done"); + let answer = service.complete(*claim, result, FRAME).unwrap(); + let rejected = plan_of(1_000, "rejected", 1_000_000); + let again = service + .begin(&raw_envelope(&invocation), 1_000, &rejected, FRAME) + .unwrap(); + assert_eq!(again, TransportAdmission::Reply(answer)); + } + + #[test] + fn a_dollar_era_replay_store_is_refused_at_open() { + use cultcache_rs::{CultCache, DatabaseEntry, OwnedRedbMessagePackBackingStore}; + + // A record under the dollar era's type id, in the dollar era's shape. + #[derive(Clone, Debug, PartialEq, Eq, DatabaseEntry)] + #[cultcache(type = "gamecult.model.replay_record.v2")] + struct DollarEraRecord { + #[cultcache(key = 0)] + caller_runtime_id: String, + #[cultcache(key = 1)] + charged: u64, + } + + let store = TestStore::new(); + { + let mut cache = CultCache::new(); + cache.register_entry_type::().unwrap(); + cache + .add_generic_backing_store( + OwnedRedbMessagePackBackingStore::new(&store.path).unwrap(), + ) + .unwrap(); + cache + .put( + "old".to_string(), + &DollarEraRecord { + caller_runtime_id: CALLER.to_string(), + charged: 75, + }, + ) + .unwrap(); + } + assert!(matches!( + TransportService::open( + &store.path, + "fixture-sse", + [caller(CALLER, KEY, 1)], + 5_000 + ), + Err(ServiceError::ReplayStore(_)) + )); + } + + #[test] + fn the_usage_readout_is_the_plan_as_it_stands_and_changes_nothing() { + let store = TestStore::new(); + let mut service = open_service(&store, 4); + let (claim, _) = claim_for(&mut service, "request-1"); + let plan = plan_of(80, "allowed", 1_793_307_000_000); + // In flight, completed, asked by either caller: the same plan. + let readout = readout_as(&service, CALLER, KEY, plan.clone(), FRAME).unwrap(); assert_eq!(readout.at_unix_ms, 1_500); - let mine = row(&readout, CALLER); - assert_eq!(mine.window_ms, WINDOW_MS); - assert_eq!(mine.ceiling_micro_usd, CEILING); - assert_eq!( - (mine.spent_micro_usd, mine.reserved_micro_usd), - (0, reservation()) - ); - - // Settled: the charge replaces the reservation. + assert_eq!(readout.plan_quota, plan.windows); let result = text_result(claim.call(), "done"); service.complete(*claim, result, FRAME).unwrap(); - let readout = readout_as(&service, CALLER, KEY, 1_500, FRAME).unwrap(); - let mine = row(&readout, CALLER); - assert_eq!( - (mine.spent_micro_usd, mine.reserved_micro_usd), - (FIXTURE_USAGE_COST, 0) - ); - - // Every admitted caller has a row, in id order; the idle one is empty. - let ids: Vec<&str> = readout - .callers - .iter() - .map(|r| r.caller_id.as_str()) - .collect(); - assert_eq!(ids, ["epiphany-yggdrasil", "ghostlight-yggdrasil"]); - let idle = row(&readout, "ghostlight-yggdrasil"); - assert_eq!((idle.spent_micro_usd, idle.reserved_micro_usd), (0, 0)); - } - - #[test] - fn any_admitted_caller_reads_every_callers_row() { - 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(); - // Ghostlight asks, with its own key, and sees epiphany's spend. - let readout = readout_as( + let ghostlight = readout_as( &service, "ghostlight-yggdrasil", "ghostlight-distinct-test-key", - 1_500, + plan.clone(), FRAME, ) .unwrap(); - assert_eq!(row(&readout, CALLER).spent_micro_usd, FIXTURE_USAGE_COST); - assert_eq!(row(&readout, "ghostlight-yggdrasil").spent_micro_usd, 0); + assert_eq!(ghostlight, readout); // A key that belongs to no admitted caller reads nothing. - assert!(readout_as(&service, "ghostlight-yggdrasil", KEY, 1_500, FRAME).is_err()); - } - - #[test] - fn a_claim_a_crash_left_reads_as_spent_not_reserved() { - let store = TestStore::new(); - let mut service = open_service(&store, 4); - let _claim = claim_at(&mut service, "request-1", 1_000); - drop(service); - let service = open_service(&store, 4); - let readout = readout_as(&service, CALLER, KEY, 1_500, FRAME).unwrap(); - let mine = row(&readout, CALLER); - assert_eq!( - (mine.spent_micro_usd, mine.reserved_micro_usd), - (reservation(), 0) - ); - } - - #[test] - fn a_readout_reads_the_ledger_and_changes_nothing_in_it() { - let store = TestStore::new(); - let mut service = open_service(&store, 4); - let claim = claim_at(&mut service, "request-1", 1_000); - // Two windows later: the caller's current window is empty, and - // looking must not prune the window the live claim settles into. - let later = 2 * WINDOW_MS + 5; - let readout = readout_as(&service, CALLER, KEY, later, FRAME).unwrap(); - let mine = row(&readout, CALLER); - assert_eq!((mine.spent_micro_usd, mine.reserved_micro_usd), (0, 0)); - 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!(readout_as(&service, "ghostlight-yggdrasil", KEY, plan, FRAME).is_err()); } #[test] @@ -3880,11 +2100,11 @@ mod tests { let store = TestStore::new(); let service = open_service(&store, 4); let sealed = |caller: &str, key: &str| { - encrypt_spend_readout_request(caller, "readout-1", &security(key)).unwrap() + encrypt_usage_readout_request(caller, "readout-1", &security(key)).unwrap() }; assert_eq!( service - .spend_readout( + .usage_readout( &sealed("nobody", KEY), 1, PlanQuotaSnapshot::default(), @@ -3895,7 +2115,7 @@ mod tests { ); assert_eq!( service - .spend_readout( + .usage_readout( &sealed(CALLER, "another-key"), 1, PlanQuotaSnapshot::default(), @@ -3904,10 +2124,10 @@ mod tests { .unwrap_err(), ServiceError::Encryption ); - let invocation = invocation_at("request-1", 1_000, MAX_COST); + let invocation = invocation_for(CALLER, "request-1"); assert_eq!( service - .spend_readout( + .usage_readout( &raw_envelope(&invocation), 1, PlanQuotaSnapshot::default(), @@ -3923,9 +2143,19 @@ mod tests { let store = TestStore::new(); let service = open_service(&store, 4); assert_eq!( - readout_as(&service, CALLER, KEY, 1_000, 64).unwrap_err(), + readout_as(&service, CALLER, KEY, PlanQuotaSnapshot::default(), 64).unwrap_err(), ServiceError::FrameBound ); - assert!(readout_as(&service, CALLER, KEY, 1_000, FRAME).is_ok()); + assert!(readout_as(&service, CALLER, KEY, PlanQuotaSnapshot::default(), FRAME).is_ok()); + } + + #[test] + fn the_connector_service_has_exactly_two_operations() { + assert_eq!(Operation::from_name("readout"), None); + assert_eq!(Operation::from_name("invoke"), Some(Operation::Invoke)); + assert_eq!( + Operation::from_name("usage_readout"), + Some(Operation::UsageReadout) + ); } } diff --git a/crates/model-connector/src/service/lease.rs b/crates/model-connector/src/service/lease.rs index 03f34dc..f677ae0 100644 --- a/crates/model-connector/src/service/lease.rs +++ b/crates/model-connector/src/service/lease.rs @@ -1,27 +1,23 @@ -//! A credential lease: a Self run's or a Persona turn's reservation in the caller's spend window. -//! The lease is one record of the same ledger a provider call reserves in, not -//! a second ledger: it is claimed before anything is lent, counts against the -//! window's ceiling while open, and settles to the true cost of what the run -//! reports it used, per model and uncapped: a charge above the reservation is -//! kept as the lease's overrun, and still counts against the window. A lease -//! nothing settled is billed the whole reservation, as a call a crash left -//! behind is. +//! A credential lease: a Self run's or a Persona turn's hold on the token the +//! connector lends. The lease is memory bound to its holder's open connection, +//! not a record and not a charge: it decides `Busy` while it is held and +//! nothing else. It ends when its holder lets go (`end_lease`), when the unit +//! invocation it served ends, or when its expiry passes. Nothing is billed and +//! nothing is reserved. The only usage-derived refusal is the provider's own +//! (`ProviderRejected`), repeated with the provider's own expiry. use super::*; -/// The request-id namespace of leases in a caller's replay records. +/// The request-id namespace of leases. const LEASE_PREFIX: &str = "lease:"; -/// The most leases one Persona invocation may take in its life. -const MAX_PERSONA_LEASES: u32 = 1024; - -/// Room kept in an identity for `:`, `n` at most `MAX_PERSONA_LEASES`. -const SEQUENCE_SUFFIX_BYTES: usize = 5; +/// What a Persona turn's lease id adds to its invocation id, so a Persona +/// lease and a Self lease of one caller and one invocation never share an id. +const PERSONA_SUFFIX: &str = ":persona"; /// Which kind of turn a lease serves. It is a property of the lease caller's -/// configured row, never of a request. A Self run leases once per unit -/// invocation; a Persona turn leases once per turn, so one live invocation -/// leases in sequence, one open lease at a time. +/// configured row, never of a request. A caller holds one Self lease at a +/// time; a Persona turn holds one lease per live invocation, one at a time. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] pub enum LeaseTurn { SelfRun, @@ -32,70 +28,17 @@ pub enum LeaseTurn { #[derive(Debug, Clone, PartialEq, Eq)] pub struct LeaseGrant { pub lease_id: String, - pub reserved_micro_usd: u64, } -/// What a run used of one model. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct ModelUsage { - pub model: String, - pub usage: Usage, -} - -/// What a settled lease was billed. `overrun_micro_usd` is how far the charge -/// passed the reservation. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct LeaseSettlement { - pub billed_micro_usd: u64, - pub overrun_micro_usd: u64, -} - -/// What a lease's settlement is billed on. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum LeaseUse<'a> { - /// No report: the whole reservation. A run that cannot say what it used - /// pays what was set aside for it. - Reservation, - /// The run's report. A report that names no use (an empty list) is no - /// report, and bills the reservation. - Reported(&'a [ModelUsage]), - /// The token was never lent, so nothing was used and nothing is billed. - NeverLent, -} - -/// What `usage` costs at the table's prices: each entry at its model's price, -/// and a model with no price at the dearest price the table holds, so an -/// unknown name never bills less than any known one. `None` when the table -/// prices nothing. -fn lease_cost(prices: &PriceTable, usage: &[ModelUsage]) -> Option { - let dearest = |usage: &Usage| { - prices - .values() - .map(|price| cost_micro_usd(price, usage, None)) - .max() - }; - usage - .iter() - .map(|entry| match prices.get(&entry.model) { - Some(price) => Some(cost_micro_usd(price, &entry.usage, None)), - None => dearest(&entry.usage), - }) - .try_fold(0_u64, |total, cost| Some(total.saturating_add(cost?))) -} - -/// Why no lease was reserved. Nothing was written to the ledger. +/// Why no lease was granted. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum LeaseRefusal { - /// The window's spend plus open reservations plus this one passes the - /// caller's ceiling. - SpendCeiling, - /// The reservation asked for is zero: a lease that bills nothing is a - /// free pass. - CostBound, UnknownCaller, - /// A Self run: the caller holds a Self lease already, or this invocation - /// leased before. A Persona turn: this invocation holds a lease now. + /// A Self run: the caller holds a Self lease now. A Persona turn: this + /// invocation holds a lease now. Busy, + /// The provider is refusing right now and has said until when. + ProviderRejected, } #[derive(Debug, Clone, PartialEq, Eq)] @@ -105,264 +48,86 @@ pub enum LeaseAdmission { } impl TransportService { - /// Reserves `max_cost_micro_usd` in `caller_runtime_id`'s current window - /// for the unit invocation `invocation_id`. A Self run leases once per - /// invocation, as `lease:`, and one Self lease is open per - /// caller. A Persona turn takes the next free sequence of its invocation - /// as `lease::`, one open at a time; a Persona lease never - /// blocks or is blocked by the caller's Self lease. A ledger that cannot - /// be written answers `Err` and reserves nothing. A lease takes no - /// request slot. + /// Grants `caller_runtime_id` a lease for the unit invocation + /// `invocation_id` until `expires_at_unix_ms`. A Self run's lease is + /// `lease:` and one is open per caller; a Persona turn's is + /// `lease::persona`, one open per invocation; neither blocks + /// the other. A lease takes no request slot. `plan` is the backend's plan + /// quota as it stands: a provider that is refusing right now refuses the + /// lease. A lease past its expiry no longer holds anything. pub fn begin_lease( &mut self, caller_runtime_id: &str, turn: LeaseTurn, invocation_id: &str, - max_cost_micro_usd: u64, expires_at_unix_ms: u64, now_unix_ms: u64, + plan: &PlanQuotaSnapshot, ) -> Result { - let Some(caller) = self.callers.get(caller_runtime_id) else { + if !self.callers.contains_key(caller_runtime_id) { return Ok(LeaseAdmission::Refused(LeaseRefusal::UnknownCaller)); - }; - let spend_ceiling = caller.spend_ceiling_micro_usd; - let spend_window_ms = caller.spend_window_ms; - let connection_key_epoch = caller.connection_key_epoch; + } require_id(invocation_id, "invocation_id")?; let suffix = match turn { - LeaseTurn::SelfRun => 0, - LeaseTurn::PersonaTurn => SEQUENCE_SUFFIX_BYTES, + LeaseTurn::SelfRun => "", + LeaseTurn::PersonaTurn => PERSONA_SUFFIX, }; - if invocation_id.len() + LEASE_PREFIX.len() + suffix > MAX_IDENTITY_BYTES { + let lease_id = format!("{LEASE_PREFIX}{invocation_id}{suffix}"); + if lease_id.len() > MAX_IDENTITY_BYTES { return Err(ContractError::Invalid("invocation_id").into()); } - if max_cost_micro_usd == 0 { - return Ok(LeaseAdmission::Refused(LeaseRefusal::CostBound)); - } - let identity = match turn { - LeaseTurn::SelfRun => { - let identity = RequestIdentity { - caller_runtime_id: caller_runtime_id.to_string(), - request_id: format!("{LEASE_PREFIX}{invocation_id}"), - }; - let holding = self.live.iter().any(|(held, request)| { - request.lease == Some(LeaseTurn::SelfRun) - && held.caller_runtime_id == caller_runtime_id - }); - if holding - || self - .replay_store - .get::(&replay_key(&identity)) - .map_err(replay_store_error)? - .is_some() - { - return Ok(LeaseAdmission::Refused(LeaseRefusal::Busy)); - } - identity - } - LeaseTurn::PersonaTurn => { - let mut free = None; - for sequence in 1..=MAX_PERSONA_LEASES { - let identity = RequestIdentity { - caller_runtime_id: caller_runtime_id.to_string(), - request_id: format!("{LEASE_PREFIX}{invocation_id}:{sequence}"), - }; - if self.live.contains_key(&identity) { - return Ok(LeaseAdmission::Refused(LeaseRefusal::Busy)); - } - if self - .replay_store - .get::(&replay_key(&identity)) - .map_err(replay_store_error)? - .is_none() - { - free = Some(identity); - break; - } - } - let Some(identity) = free else { - return Ok(LeaseAdmission::Refused(LeaseRefusal::Busy)); - }; - identity - } - }; - let key = replay_key(&identity); - let admitted_at_unix_ms = admission_clock( - &self.window_spend, - 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, - caller_runtime_id, - window, - ); - if self - .window_spend(caller_runtime_id, window) - .saturating_add(max_cost_micro_usd) - > spend_ceiling - { - return Ok(LeaseAdmission::Refused(LeaseRefusal::SpendCeiling)); - } - let mut digest = Sha256::new(); - digest.update(b"claude_lease_request"); - digest.update((invocation_id.len() as u64).to_be_bytes()); - digest.update(invocation_id.as_bytes()); - digest.update(max_cost_micro_usd.to_be_bytes()); - digest.update(expires_at_unix_ms.to_be_bytes()); - let record = ReplayRecord { - caller_runtime_id: caller_runtime_id.to_string(), - request_id: identity.request_id.clone(), - connection_key_epoch, - invocation_sha256: digest.finalize().into(), - expires_at_unix_ms, - state: ReplayState::Active { - reserved_micro_usd: max_cost_micro_usd, - }, - admitted_at_unix_ms, - }; - // Reserving and claiming are one write, as for a call: a ledger that - // cannot be written lends nothing. - self.replay_store - .put(key, &record) - .map_err(replay_store_error)?; - add_spend( - &mut self.window_spend, - caller_runtime_id, - window, - 0, - max_cost_micro_usd, - ); - let lease_id = identity.request_id.clone(); - self.live.insert( - identity, - LiveRequest { - window, - reserved_micro_usd: max_cost_micro_usd, - lease: Some(turn), - }, - ); - Ok(LeaseAdmission::Granted(LeaseGrant { - lease_id, - reserved_micro_usd: max_cost_micro_usd, - })) - } - - /// Settles an open lease and returns what it was billed. A report bills - /// the true cost of every model's reported use, uncapped (the overrun is - /// recorded); no report, or an empty one, bills the whole reservation; - /// only a token never lent bills nothing. A failed ledger - /// write frees the caller's lease and keeps the reservation counted, as a - /// call's does. - pub fn settle_lease( - &mut self, - caller_runtime_id: &str, - lease_id: &str, - usage: LeaseUse<'_>, - ) -> Result { + self.leases + .retain(|_, lease| lease.expires_at_unix_ms > now_unix_ms); let identity = RequestIdentity { caller_runtime_id: caller_runtime_id.to_string(), - request_id: lease_id.to_string(), + request_id: lease_id.clone(), }; - if self - .live - .get(&identity) - .is_none_or(|request| request.lease.is_none()) - { - return Err(ServiceError::NoActiveRequest); + let busy = match turn { + LeaseTurn::SelfRun => self.leases.iter().any(|(held, lease)| { + lease.turn == LeaseTurn::SelfRun && held.caller_runtime_id == caller_runtime_id + }), + LeaseTurn::PersonaTurn => self.leases.contains_key(&identity), + }; + if busy { + return Ok(LeaseAdmission::Refused(LeaseRefusal::Busy)); } - let key = replay_key(&identity); - let active = self - .replay_store - .get::(&key) - .map_err(replay_store_error)? - .ok_or(ServiceError::NoActiveRequest)?; - let ReplayState::Active { reserved_micro_usd } = active.state else { - return Err(ServiceError::ActiveRequestMismatch); - }; - let caller = self - .callers - .get(caller_runtime_id) - .ok_or(ServiceError::CallerNotAdmitted)?; - let spend_window_ms = caller.spend_window_ms; - let charged_micro_usd = match usage { - LeaseUse::Reservation | LeaseUse::Reported([]) => reserved_micro_usd, - LeaseUse::NeverLent => 0, - LeaseUse::Reported(usage) => { - lease_cost(&self.prices, usage).unwrap_or(reserved_micro_usd) - } - }; - let overrun_micro_usd = charged_micro_usd.saturating_sub(reserved_micro_usd); - let settled = ReplayRecord { - state: ReplayState::LeaseSettled { - charged_micro_usd, - overrun_micro_usd, + if provider_rejects(plan, now_unix_ms) { + return Ok(LeaseAdmission::Refused(LeaseRefusal::ProviderRejected)); + } + self.leases.insert( + identity, + LiveLease { + turn, + expires_at_unix_ms, }, - ..active.clone() - }; - self.replay_store - .put(key, &settled) - .map_err(replay_store_error) - .inspect_err(|_| { - self.live.remove(&identity); - })?; - add_spend( - &mut self.window_spend, - 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, caller_runtime_id); - Ok(LeaseSettlement { - billed_micro_usd: charged_micro_usd, - overrun_micro_usd, - }) + Ok(LeaseAdmission::Granted(LeaseGrant { lease_id })) + } + + /// Ends a lease its holder let go of. Ending one that is not held is a + /// no-op: a holder may close after its expiry already freed the lease. + pub fn end_lease(&mut self, caller_runtime_id: &str, lease_id: &str) { + self.leases.remove(&RequestIdentity { + caller_runtime_id: caller_runtime_id.to_string(), + request_id: lease_id.to_string(), + }); } } #[cfg(test)] mod tests { - use std::sync::atomic::Ordering; - use super::*; use crate::fixtures::*; + use crate::readout::PlanQuota; const CALLER: &str = "eureka"; const KEY: &str = "lease-distinct-test-key"; const NOW: u64 = 7_200_500; - fn open_with(store: impl CacheBackingStore + 'static, ceiling: u64) -> TransportService { - TransportService::with_store( - store, - "fixture-sse", - fixture_prices(), - [CallerAdmission::new( - CALLER, - KEY, - 1, - ["gpt-5.4".to_string()], - 1, - 65_536, - 32_768, - SpendCeiling(ceiling), - SpendWindowMs(WINDOW_MS), - ) - .unwrap()], - 5_000, - ) - .unwrap() - } - - fn open(store: &TestStore, ceiling: u64) -> TransportService { + fn open(store: &TestStore) -> TransportService { TransportService::open( &store.path, "fixture-sse", - fixture_prices(), [CallerAdmission::new( CALLER, KEY, @@ -371,8 +136,6 @@ mod tests { 1, 65_536, 32_768, - SpendCeiling(ceiling), - SpendWindowMs(WINDOW_MS), ) .unwrap()], 5_000, @@ -380,12 +143,37 @@ mod tests { .unwrap() } - fn lease(service: &mut TransportService, run: &str, cost: u64) -> LeaseAdmission { + fn none() -> PlanQuotaSnapshot { + PlanQuotaSnapshot::default() + } + + fn window(permille: u32, status: &str, resets_at: u64) -> PlanQuotaSnapshot { + PlanQuotaSnapshot { + windows: vec![PlanQuota { + window: "five_hour".to_string(), + utilization_permille: Some(permille), + resets_at_unix_ms: resets_at, + status: status.to_string(), + observed_at_unix_ms: NOW, + }], + } + } + + fn lease_in( + service: &mut TransportService, + turn: LeaseTurn, + run: &str, + plan: &PlanQuotaSnapshot, + ) -> LeaseAdmission { service - .begin_lease(CALLER, LeaseTurn::SelfRun, run, cost, NOW + 60_000, NOW) + .begin_lease(CALLER, turn, run, NOW + 60_000, NOW, plan) .unwrap() } + fn lease(service: &mut TransportService, run: &str) -> LeaseAdmission { + lease_in(service, LeaseTurn::SelfRun, run, &none()) + } + fn granted(admission: LeaseAdmission) -> LeaseGrant { match admission { LeaseAdmission::Granted(grant) => grant, @@ -393,376 +181,226 @@ mod tests { } } - fn held(service: &TransportService) -> u64 { - service.window_spend(CALLER, NOW / WINDOW_MS) - } - - fn output(tokens: u64) -> Vec { - used("gpt-5.4", tokens) - } - - fn used(model: &str, tokens: u64) -> Vec { - vec![ModelUsage { - model: model.to_string(), - usage: Usage { - input_tokens: 0, - cache_read_tokens: 0, - cache_write_tokens: 0, - output_tokens: tokens, - reasoning_tokens: None, - }, - }] - } - - fn billed(settlement: LeaseSettlement) -> u64 { - settlement.billed_micro_usd + #[test] + fn a_lease_is_named_by_its_invocation() { + let store = TestStore::new(); + let mut service = open(&store); + assert_eq!(granted(lease(&mut service, "run-1")).lease_id, "lease:run-1"); } #[test] - fn a_lease_reserves_in_the_callers_window_and_names_what_it_holds() { + fn a_caller_holds_one_self_lease_at_a_time() { let store = TestStore::new(); - let mut service = open(&store, 1_000); - let grant = granted(lease(&mut service, "run-1", 600)); - assert_eq!(grant.reserved_micro_usd, 600); - assert_eq!(held(&service), 600); + let mut service = open(&store); + granted(lease(&mut service, "run-1")); assert_eq!( - billed( - service - .settle_lease(CALLER, &grant.lease_id, LeaseUse::Reservation) - .unwrap() - ), - 600 - ); - } - - #[test] - fn the_window_ceiling_counts_open_and_settled_reservations_inclusively() { - let store = TestStore::new(); - let mut service = open(&store, 1_000); - let first = granted(lease(&mut service, "run-1", 600)); - service - .settle_lease(CALLER, &first.lease_id, LeaseUse::Reservation) - .unwrap(); - assert_eq!( - lease(&mut service, "run-2", 401), - LeaseAdmission::Refused(LeaseRefusal::SpendCeiling) - ); - assert_eq!(held(&service), 600); - granted(lease(&mut service, "run-3", 400)); - assert_eq!(held(&service), 1_000); - } - - #[test] - fn a_caller_holds_one_lease_at_a_time_and_an_invocation_id_is_leased_once() { - let store = TestStore::new(); - let mut service = open(&store, 1_000_000); - let first = granted(lease(&mut service, "run-1", 100)); - assert_eq!( - lease(&mut service, "run-2", 100), + lease(&mut service, "run-2"), LeaseAdmission::Refused(LeaseRefusal::Busy) ); - service - .settle_lease(CALLER, &first.lease_id, LeaseUse::Reservation) - .unwrap(); assert_eq!( - lease(&mut service, "run-1", 100), + lease(&mut service, "run-1"), LeaseAdmission::Refused(LeaseRefusal::Busy) ); - granted(lease(&mut service, "run-2", 100)); } #[test] - fn a_stranger_and_a_free_lease_are_refused_without_a_ledger_write() { + fn a_lease_ends_when_its_holder_closes() { let store = TestStore::new(); - let mut service = open(&store, 1_000); + let mut service = open(&store); + let grant = granted(lease(&mut service, "run-1")); + service.end_lease(CALLER, &grant.lease_id); + // The same invocation may lease again: nothing remembers a closed + // lease, and nothing was billed for it. + granted(lease(&mut service, "run-1")); + service.end_lease(CALLER, "lease:not-held"); + assert_eq!( + lease(&mut service, "run-2"), + LeaseAdmission::Refused(LeaseRefusal::Busy), + "ending a lease that is not held frees nothing else" + ); + } + + #[test] + fn a_lease_ends_at_its_expiry_without_its_holder() { + let store = TestStore::new(); + let mut service = open(&store); + granted( + service + .begin_lease(CALLER, LeaseTurn::SelfRun, "run-1", NOW + 10, NOW, &none()) + .unwrap(), + ); assert_eq!( service - .begin_lease("stranger", LeaseTurn::SelfRun, "run-1", 100, NOW + 1, NOW) + .begin_lease(CALLER, LeaseTurn::SelfRun, "run-2", NOW + 99, NOW + 9, &none()) + .unwrap(), + LeaseAdmission::Refused(LeaseRefusal::Busy), + "one millisecond before its expiry the lease still holds" + ); + granted( + service + .begin_lease(CALLER, LeaseTurn::SelfRun, "run-2", NOW + 99, NOW + 10, &none()) + .unwrap(), + ); + } + + #[test] + fn a_stranger_is_refused() { + let store = TestStore::new(); + let mut service = open(&store); + assert_eq!( + service + .begin_lease("stranger", LeaseTurn::SelfRun, "run-1", NOW + 1, NOW, &none()) .unwrap(), LeaseAdmission::Refused(LeaseRefusal::UnknownCaller) ); - assert_eq!( - lease(&mut service, "run-1", 0), - LeaseAdmission::Refused(LeaseRefusal::CostBound) - ); - assert_eq!(held(&service), 0); - // Nothing was claimed: the same run id leases once it asks sensibly. - granted(lease(&mut service, "run-1", 100)); } #[test] - fn settlement_bills_the_true_cost_uncapped_and_records_the_overrun() { + fn a_persona_holds_one_lease_per_invocation_and_never_blocks_the_self_lease() { let store = TestStore::new(); - let mut service = open(&store, 1_000_000); - // Output at 15 USD per million tokens: 10 tokens bill 150 micro-USD. - let small = granted(lease(&mut service, "run-1", 10_000)); - let settled = service - .settle_lease(CALLER, &small.lease_id, LeaseUse::Reported(&output(10)[..])) - .unwrap(); - assert_eq!( - (settled.billed_micro_usd, settled.overrun_micro_usd), - (150, 0) - ); - assert_eq!(held(&service), 150); - // 1000 tokens bill 15_000: the reservation does not cap it, and the - // 5_000 past it is recorded as the overrun. - let big = granted(lease(&mut service, "run-2", 10_000)); - let settled = service - .settle_lease( - CALLER, - &big.lease_id, - LeaseUse::Reported(&output(1_000)[..]), - ) - .unwrap(); - assert_eq!( - (settled.billed_micro_usd, settled.overrun_micro_usd), - (15_000, 5_000) - ); - assert_eq!(held(&service), 15_150); - } - - #[test] - fn settlement_prices_each_model_at_its_own_rate_and_an_unknown_model_at_the_dearest() { - let store = TestStore::new(); - let cheap = ModelPrice { - rates: Rates { - input: 1_000_000, - cache_write_short: 1_000_000, - cache_write_long: 1_000_000, - cache_read: 1_000_000, - output: 2_000_000, - }, - ..fixture_price() + let mut service = open(&store); + let persona = |service: &mut TransportService, run: &str| { + lease_in(service, LeaseTurn::PersonaTurn, run, &none()) }; - let prices = PriceTable::from([ - ("gpt-5.4".to_string(), fixture_price()), - ("cheap".to_string(), cheap), - ]); - let mut service = TransportService::open( - &store.path, - "fixture-sse", - prices, - [CallerAdmission::new( + let first = granted(persona(&mut service, "turn-1")); + assert_eq!(first.lease_id, "lease:turn-1:persona"); + assert_eq!( + persona(&mut service, "turn-1"), + LeaseAdmission::Refused(LeaseRefusal::Busy) + ); + granted(persona(&mut service, "turn-2")); + granted(lease(&mut service, "turn-1")); + service.end_lease(CALLER, &first.lease_id); + granted(persona(&mut service, "turn-1")); + } + + #[test] + fn a_rejected_window_refuses_invoke_and_lease_until_its_reset() { + let store = TestStore::new(); + let mut service = open(&store); + let rejected = window(1_000, "rejected", NOW + 5_000); + assert_eq!( + lease_in(&mut service, LeaseTurn::SelfRun, "run-1", &rejected), + LeaseAdmission::Refused(LeaseRefusal::ProviderRejected) + ); + let call = |service: &mut TransportService, now: u64, plan: &PlanQuotaSnapshot| { + let mut asked = request(); + asked.request_id = format!("call-{now}"); + let invocation = + TransportInvocation::new(CALLER, now + 1_000, [7; 32], asked).unwrap(); + let sealed = crate::envelope::encrypt_message( CALLER, - KEY, - 1, - ["gpt-5.4".to_string()], - 1, - 65_536, - 32_768, - SpendCeiling(1_000_000), - SpendWindowMs(WINDOW_MS), + invocation.request_id(), + crate::envelope::EnvelopeKind::Invocation, + &crate::contract::wire_encode(&invocation).unwrap(), + &security(KEY), ) - .unwrap()], - 5_000, - ) - .unwrap(); - let mixed: Vec = [("gpt-5.4", 10), ("cheap", 10), ("unheard-of", 10)] - .into_iter() - .flat_map(|(model, tokens)| used(model, tokens)) - .collect(); - let grant = granted(lease(&mut service, "run-1", 10_000)); - // 10 tokens at 15, 10 at 2, and 10 at the dearest (15) per million. + .unwrap(); + let TransportAdmission::Reply(reply) = + service.begin(&sealed, now, plan, 1 << 20).unwrap() + else { + return None; + }; + let result = crate::envelope::decrypt_result(&reply, &security(KEY)).unwrap(); + match result.disposition { + TransportDisposition::Refused { refusal } => Some(refusal), + other => panic!("unexpected {other:?}"), + } + }; assert_eq!( - billed( - service - .settle_lease(CALLER, &grant.lease_id, LeaseUse::Reported(&mixed[..])) - .unwrap() - ), - 150 + 20 + 150 + call(&mut service, NOW, &rejected), + Some(Refusal::ProviderRejected) ); - // A run that reports no use pays the reservation, not nothing. - let grant = granted(lease(&mut service, "run-2", 10_000)); assert_eq!( - billed( - service - .settle_lease(CALLER, &grant.lease_id, LeaseUse::Reported(&[])) - .unwrap() - ), - 10_000 + call(&mut service, NOW + 4_999, &rejected), + Some(Refusal::ProviderRejected) ); - } - - #[test] - fn settling_an_empty_report_cannot_free_the_reservation() { - let store = TestStore::new(); - let mut service = open(&store, 1_000); - let grant = granted(lease(&mut service, "run-1", 600)); - assert_eq!( - billed( - service - .settle_lease(CALLER, &grant.lease_id, LeaseUse::Reported(&[])) - .unwrap() - ), - 600 - ); - assert_eq!(held(&service), 600); + assert_eq!(call(&mut service, NOW + 5_000, &rejected), None); assert!(matches!( - lease(&mut service, "run-2", 1_000), - LeaseAdmission::Refused(LeaseRefusal::SpendCeiling) + service + .begin_lease( + CALLER, + LeaseTurn::SelfRun, + "run-1", + NOW + 9_000, + NOW + 5_000, + &rejected + ) + .unwrap(), + LeaseAdmission::Granted(_) )); } #[test] - fn a_token_never_lent_bills_nothing_and_frees_the_window() { + fn high_utilization_refuses_nothing() { let store = TestStore::new(); - let mut service = open(&store, 1_000); - let grant = granted(lease(&mut service, "run-1", 600)); - let settled = service - .settle_lease(CALLER, &grant.lease_id, LeaseUse::NeverLent) - .unwrap(); - assert_eq!( - (settled.billed_micro_usd, settled.overrun_micro_usd), - (0, 0) - ); - assert_eq!(held(&service), 0); - granted(lease(&mut service, "run-2", 1_000)); - } - - #[test] - fn a_lease_settles_once() { - let store = TestStore::new(); - let mut service = open(&store, 1_000_000); - let grant = granted(lease(&mut service, "run-1", 100)); - service - .settle_lease(CALLER, &grant.lease_id, LeaseUse::Reservation) - .unwrap(); - assert!(matches!( - service.settle_lease(CALLER, &grant.lease_id, LeaseUse::Reservation), - Err(ServiceError::NoActiveRequest) + let mut service = open(&store); + let full = window(999, "allowed", NOW + 5_000); + granted(lease_in(&mut service, LeaseTurn::SelfRun, "run-1", &full)); + let warned = window(1_000, "allowed_warning", NOW + 5_000); + granted(lease_in( + &mut service, + LeaseTurn::PersonaTurn, + "turn-1", + &warned, )); - assert_eq!(held(&service), 100); } #[test] - fn a_lease_left_open_by_a_crash_counts_at_its_reservation_after_restart() { + fn an_invocation_id_is_bounded_so_its_lease_id_fits_an_identity() { let store = TestStore::new(); - let mut service = open(&store, 1_000_000); - granted(lease(&mut service, "run-1", 700)); - drop(service); - let service = open(&store, 1_000_000); - assert_eq!(held(&service), 700); + let mut service = open(&store); + let longest = "r".repeat(MAX_IDENTITY_BYTES - LEASE_PREFIX.len()); + let grant = granted(lease(&mut service, &longest)); + assert_eq!(grant.lease_id.len(), MAX_IDENTITY_BYTES); + assert!(lease_try(&mut service, LeaseTurn::SelfRun, &format!("{longest}r")).is_err()); } #[test] - fn a_ledger_that_cannot_be_written_lends_nothing() { - let flaky = FlakyStore::new(1_000); - let allowance = flaky.allowance(); - let mut service = open_with(flaky, 1_000_000); - allowance.store(0, Ordering::SeqCst); - assert!( - service - .begin_lease(CALLER, LeaseTurn::SelfRun, "run-1", 100, NOW + 1, NOW) - .is_err() - ); - assert_eq!(held(&service), 0); - allowance.store(1_000, Ordering::SeqCst); - granted(lease(&mut service, "run-1", 100)); + fn a_persona_invocation_id_leaves_room_for_its_suffix() { + let store = TestStore::new(); + let mut service = open(&store); + let longest = "r".repeat(MAX_IDENTITY_BYTES - LEASE_PREFIX.len() - PERSONA_SUFFIX.len()); + let grant = granted(lease_in(&mut service, LeaseTurn::PersonaTurn, &longest, &none())); + assert_eq!(grant.lease_id.len(), MAX_IDENTITY_BYTES); + assert!(lease_try(&mut service, LeaseTurn::PersonaTurn, &format!("{longest}r")).is_err()); } - #[test] - fn a_failed_settle_write_frees_the_lease_and_keeps_the_reservation_counted() { - let flaky = FlakyStore::new(1_000); - let allowance = flaky.allowance(); - let mut service = open_with(flaky, 1_000_000); - let grant = granted(lease(&mut service, "run-1", 100)); - allowance.store(0, Ordering::SeqCst); - assert!( - service - .settle_lease(CALLER, &grant.lease_id, LeaseUse::Reported(&output(1)[..])) - .is_err() - ); - assert_eq!(held(&service), 100); - allowance.store(1_000, Ordering::SeqCst); - granted(lease(&mut service, "run-2", 100)); - } - - fn begin_call(service: &mut TransportService, id: &str) -> TransportAdmission { - let mut call = request(); - call.request_id = id.to_string(); - let invocation = - TransportInvocation::new(CALLER, NOW + 1_000, [7; 32], MAX_COST, call).unwrap(); - let sealed = crate::envelope::encrypt_message( - CALLER, - id, - crate::envelope::EnvelopeKind::Invocation, - &crate::contract::wire_encode(&invocation).unwrap(), - &security(KEY), - ) - .unwrap(); - service.begin(&sealed, NOW, 1 << 20).unwrap() + fn lease_try( + service: &mut TransportService, + turn: LeaseTurn, + run: &str, + ) -> Result { + service.begin_lease(CALLER, turn, run, NOW + 60_000, NOW, &none()) } #[test] fn a_lease_takes_no_request_slot_from_the_callers_calls() { // The caller admits one concurrent call; an open lease must not use it. let store = TestStore::new(); - let mut service = open(&store, 1 << 40); - granted(lease(&mut service, "run-1", 100)); - match begin_call(&mut service, "call-1") { - TransportAdmission::Execute(_) => {} - TransportAdmission::Reply(_) => panic!("the lease held the call's slot"), - } - } - - #[test] - fn a_call_in_flight_does_not_make_the_caller_busy_for_a_lease() { - let store = TestStore::new(); - let mut service = open(&store, 1 << 40); + let mut service = open(&store); + granted(lease(&mut service, "run-1")); + let mut call = request(); + call.request_id = "call-1".to_string(); + let invocation = TransportInvocation::new(CALLER, NOW + 1_000, [7; 32], call).unwrap(); + let sealed = crate::envelope::encrypt_message( + CALLER, + "call-1", + crate::envelope::EnvelopeKind::Invocation, + &crate::contract::wire_encode(&invocation).unwrap(), + &security(KEY), + ) + .unwrap(); assert!(matches!( - begin_call(&mut service, "call-1"), + service.begin(&sealed, NOW, &none(), 1 << 20).unwrap(), TransportAdmission::Execute(_) )); - granted(lease(&mut service, "run-1", 100)); - } - - #[test] - fn an_invocation_id_is_bounded_so_its_lease_id_fits_an_identity() { - let store = TestStore::new(); - let mut service = open(&store, 1 << 40); - let longest = "r".repeat(MAX_IDENTITY_BYTES - LEASE_PREFIX.len()); - let grant = granted(lease(&mut service, &longest, 100)); - assert_eq!(grant.lease_id.len(), MAX_IDENTITY_BYTES); - service - .settle_lease(CALLER, &grant.lease_id, LeaseUse::Reservation) - .unwrap(); - assert!( - service - .begin_lease( - CALLER, - LeaseTurn::SelfRun, - &format!("{longest}r"), - 100, - NOW + 1, - NOW - ) - .is_err() - ); - } - - #[test] - fn a_persona_invocation_id_leaves_room_for_its_sequence() { - let store = TestStore::new(); - let mut service = open(&store, 1 << 40); - let longest = "r".repeat(MAX_IDENTITY_BYTES - LEASE_PREFIX.len() - SEQUENCE_SUFFIX_BYTES); - let persona = |service: &mut TransportService, id: &str| { - service.begin_lease(CALLER, LeaseTurn::PersonaTurn, id, 100, NOW + 60_000, NOW) - }; - let grant = granted(persona(&mut service, &longest).unwrap()); - assert!(grant.lease_id.len() <= MAX_IDENTITY_BYTES); - assert!(persona(&mut service, &format!("{longest}r")).is_err()); - } - - #[test] - fn a_settled_lease_leaves_its_charge_not_its_reservation_across_a_restart() { - let store = TestStore::new(); - let mut service = open(&store, 1_000_000); - let grant = granted(lease(&mut service, "run-1", 10_000)); - service - .settle_lease(CALLER, &grant.lease_id, LeaseUse::Reported(&output(10)[..])) - .unwrap(); - drop(service); - let service = open(&store, 1_000_000); - assert_eq!(held(&service), 150); + // And a call in flight does not make the caller busy for a lease. + granted(lease_in( + &mut service, + LeaseTurn::PersonaTurn, + "turn-1", + &none(), + )); } } diff --git a/crates/model-connector/tests/fixtures/wire_v3/README.md b/crates/model-connector/tests/fixtures/wire_v4/README.md similarity index 74% rename from crates/model-connector/tests/fixtures/wire_v3/README.md rename to crates/model-connector/tests/fixtures/wire_v4/README.md index 2751c8f..523ca3c 100644 --- a/crates/model-connector/tests/fixtures/wire_v3/README.md +++ b/crates/model-connector/tests/fixtures/wire_v4/README.md @@ -1,6 +1,6 @@ -# wire_v3 fixtures +# wire_v4 fixtures -The named-MessagePack bytes of the v3 documents and one sealed carriage frame, +The named-MessagePack bytes of the v4 documents and one sealed carriage frame, produced by the kit at the revision that added them. They pin drift; the types in `src/contract.rs` own the shape, and `schemas/` publishes it. @@ -10,7 +10,7 @@ produced by the kit at the revision that added them. They pin drift; the types i the plaintext documents. - `carriage_invocation.frame`: a `cultnet.operation_request.v0` for service `gamecult.model.connector`, operation `invoke`, whose payload is the base64 of an envelope - sealed under the fixed test key in `tests/wire_golden_v3.rs`. + sealed under the fixed test key in `tests/wire_golden_v4.rs`. -Generator: `tests/wire_golden_v3.rs` itself. Run it with `WIRE_V3_REGENERATE=1` to rewrite +Generator: `tests/wire_golden_v4.rs` itself. Run it with `WIRE_V4_REGENERATE=1` to rewrite the files from the kit, then review the diff: a changed byte is a wire change. diff --git a/crates/model-connector/tests/fixtures/wire_v3/carriage_invocation.frame b/crates/model-connector/tests/fixtures/wire_v4/carriage_invocation.frame similarity index 100% rename from crates/model-connector/tests/fixtures/wire_v3/carriage_invocation.frame rename to crates/model-connector/tests/fixtures/wire_v4/carriage_invocation.frame diff --git a/crates/model-connector/tests/fixtures/wire_v3/input_reasoning_state.mp b/crates/model-connector/tests/fixtures/wire_v4/input_reasoning_state.mp similarity index 100% rename from crates/model-connector/tests/fixtures/wire_v3/input_reasoning_state.mp rename to crates/model-connector/tests/fixtures/wire_v4/input_reasoning_state.mp diff --git a/crates/model-connector/tests/fixtures/wire_v3/provider_request.mp b/crates/model-connector/tests/fixtures/wire_v4/provider_request.mp similarity index 100% rename from crates/model-connector/tests/fixtures/wire_v3/provider_request.mp rename to crates/model-connector/tests/fixtures/wire_v4/provider_request.mp diff --git a/crates/model-connector/tests/fixtures/wire_v3/transport_envelope.mp b/crates/model-connector/tests/fixtures/wire_v4/transport_envelope.mp similarity index 100% rename from crates/model-connector/tests/fixtures/wire_v3/transport_envelope.mp rename to crates/model-connector/tests/fixtures/wire_v4/transport_envelope.mp diff --git a/crates/model-connector/tests/fixtures/wire_v3/transport_invocation.mp b/crates/model-connector/tests/fixtures/wire_v4/transport_invocation.mp similarity index 100% rename from crates/model-connector/tests/fixtures/wire_v3/transport_invocation.mp rename to crates/model-connector/tests/fixtures/wire_v4/transport_invocation.mp diff --git a/crates/model-connector/tests/fixtures/wire_v3/transport_receipt_failed.mp b/crates/model-connector/tests/fixtures/wire_v4/transport_receipt_failed.mp similarity index 100% rename from crates/model-connector/tests/fixtures/wire_v3/transport_receipt_failed.mp rename to crates/model-connector/tests/fixtures/wire_v4/transport_receipt_failed.mp diff --git a/crates/model-connector/tests/fixtures/wire_v3/transport_result_refused.mp b/crates/model-connector/tests/fixtures/wire_v4/transport_result_refused.mp similarity index 100% rename from crates/model-connector/tests/fixtures/wire_v3/transport_result_refused.mp rename to crates/model-connector/tests/fixtures/wire_v4/transport_result_refused.mp diff --git a/crates/model-connector/tests/fixtures/wire_v3/transport_result_transported.mp b/crates/model-connector/tests/fixtures/wire_v4/transport_result_transported.mp similarity index 100% rename from crates/model-connector/tests/fixtures/wire_v3/transport_result_transported.mp rename to crates/model-connector/tests/fixtures/wire_v4/transport_result_transported.mp diff --git a/crates/model-connector/tests/published_schemas.rs b/crates/model-connector/tests/published_schemas.rs index cb21b01..9e19b8c 100644 --- a/crates/model-connector/tests/published_schemas.rs +++ b/crates/model-connector/tests/published_schemas.rs @@ -1,4 +1,4 @@ -//! `schemas/` at the repository root publishes one JSON Schema per v3 document, +//! `schemas/` at the repository root publishes one JSON Schema per v4 document, //! generated from the types. This test regenerates them and fails on any //! difference; run it with `MODEL_CONNECTOR_REGENERATE_SCHEMAS=1` to rewrite the //! files from the types. @@ -28,7 +28,7 @@ fn published_schemas_are_the_types() { ("transport_result", render::()), ("transport_receipt", render::()), ("transport_envelope", render::()), - ("spend_readout", render::()), + ("usage_readout", render::()), ]; let dir = schemas_dir(); if std::env::var_os(REGENERATE).is_some() { @@ -80,7 +80,7 @@ fn the_published_ciphertext_is_what_the_wire_sends() { assert_eq!(&published["properties"]["ciphertext"], ciphertext); let envelope = TransportEnvelope { - schema_id: "gamecult.model.transport_envelope.v3".to_string(), + schema_id: "gamecult.model.transport_envelope.v4".to_string(), caller_runtime_id: "caller".to_string(), request_id: "request".to_string(), message_kind: EnvelopeKind::Invocation, @@ -111,9 +111,7 @@ fn the_refusal_schema_is_one_flat_string_enum() { .iter() .map(|value| value.as_str().unwrap()) .collect(); - assert_eq!(variants.len(), 15); - assert!(variants.contains(&"ClockBehindAdmissionWindow")); + assert_eq!(variants.len(), 13); assert!(variants.contains(&"ForeignReasoningState")); - assert!(variants.contains(&"CostBound")); - assert!(variants.contains(&"SpendCeiling")); + assert!(variants.contains(&"ProviderRejected")); } diff --git a/crates/model-connector/tests/reasoning_state.rs b/crates/model-connector/tests/reasoning_state.rs index 8c9c4af..170dd8e 100644 --- a/crates/model-connector/tests/reasoning_state.rs +++ b/crates/model-connector/tests/reasoning_state.rs @@ -106,8 +106,6 @@ 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 { @@ -116,21 +114,7 @@ fn start_with(backend: Arc, max_concurrent: usize) -> Conne socket_timeout_ms: 5_000, max_expiry_skew_ms: 300_000, }; - 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)); + thread::spawn(move || serve(listener, limits, vec![admission], &store, backend)); ConnectorClient::new(endpoint, KEY, FRAME, Some(Duration::from_secs(5))).unwrap() } @@ -143,7 +127,7 @@ fn invocation(request_id: &str, input: Vec) -> TransportInvocation { 4_096, ); request.input = input; - TransportInvocation::new(CALLER, now_ms() + 60_000, [7; 32], 10_000_000, request).unwrap() + TransportInvocation::new(CALLER, now_ms() + 60_000, [7; 32], request).unwrap() } fn user(text: &str) -> InputItem { @@ -332,7 +316,7 @@ fn a_refused_completion_reaches_the_caller_as_a_typed_failure_and_frees_the_slot 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, + |receipt| receipt.request_id = "another-request".to_string(), ]; 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_v4.rs similarity index 96% rename from crates/model-connector/tests/wire_golden_v3.rs rename to crates/model-connector/tests/wire_golden_v4.rs index 9364035..e440fe6 100644 --- a/crates/model-connector/tests/wire_golden_v3.rs +++ b/crates/model-connector/tests/wire_golden_v4.rs @@ -1,5 +1,5 @@ -//! Pins the bytes of the v3 wire so a drift is a failing test. The fixtures in -//! `fixtures/wire_v3` are produced by this file (see its README); they pin the +//! Pins the bytes of the v4 wire so a drift is a failing test. The fixtures in +//! `fixtures/wire_v4` are produced by this file (see its README); they pin the //! kit against itself and are not a second contract: the types in //! `model_connector::contract` own the shape. @@ -17,10 +17,10 @@ use serde::de::DeserializeOwned; const KEY: &str = "golden-distinct-test-key"; const CALLER: &str = "golden-caller"; -const REGENERATE: &str = "WIRE_V3_REGENERATE"; +const REGENERATE: &str = "WIRE_V4_REGENERATE"; fn fixtures() -> PathBuf { - Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/wire_v3") + Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/wire_v4") } fn named(value: &T) -> Vec { @@ -28,7 +28,7 @@ fn named(value: &T) -> Vec { } /// The value serializes to the committed bytes and the committed bytes decode -/// back to it. With `WIRE_V3_REGENERATE` set, the committed bytes are rewritten +/// back to it. With `WIRE_V4_REGENERATE` set, the committed bytes are rewritten /// from the value first. fn assert_wire( name: &str, @@ -81,7 +81,7 @@ fn request() -> ProviderRequest { } fn invocation() -> TransportInvocation { - TransportInvocation::new(CALLER, 1_800_000_000_000, [7; 32], 5_000_000, request()).unwrap() + TransportInvocation::new(CALLER, 1_800_000_000_000, [7; 32], request()).unwrap() } fn receipt(outcome: TransportOutcome) -> TransportReceipt { @@ -95,8 +95,6 @@ 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, } } @@ -116,7 +114,7 @@ fn completed() -> TransportOutcome { } #[test] -fn wire_golden_v3() { +fn wire_golden_v4() { let invocation = invocation(); assert_wire("provider_request.mp", &invocation.request); assert_wire("transport_invocation.mp", &invocation); @@ -233,7 +231,7 @@ fn the_golden_carriage_frame_decrypts_and_validates_under_the_fixed_key() { }; assert_eq!(service_id, "gamecult.model.connector"); assert_eq!(operation, "invoke"); - assert_eq!(payload_schema, "gamecult.model.transport_envelope.v3"); + assert_eq!(payload_schema, "gamecult.model.transport_envelope.v4"); assert_eq!(message_id, invocation.request_id()); let envelope: TransportEnvelope = rmp_serde::from_slice(&STANDARD.decode(payload).unwrap()).unwrap(); diff --git a/schemas/spend_readout.json b/schemas/spend_readout.json deleted file mode 100644 index 5a5d23d..0000000 --- a/schemas/spend_readout.json +++ /dev/null @@ -1,112 +0,0 @@ -{ - "$schema": "https://json-schema.org/draft/2020-12/schema", - "title": "SpendReadout", - "type": "object", - "properties": { - "at_unix_ms": { - "type": "integer", - "format": "uint64", - "minimum": 0 - }, - "callers": { - "type": "array", - "items": { - "$ref": "#/$defs/CallerSpend" - } - }, - "plan_quota": { - "type": "array", - "items": { - "$ref": "#/$defs/PlanQuota" - } - }, - "schema_id": { - "type": "string" - } - }, - "additionalProperties": false, - "required": [ - "schema_id", - "at_unix_ms", - "callers", - "plan_quota" - ], - "$defs": { - "CallerSpend": { - "description": "One caller's current spend window in micro-USD. `spent` is settled charges\nplus what a crash left counted at its full reservation; `reserved` is what\ncalls in flight hold, and both count against the ceiling.", - "type": "object", - "properties": { - "caller_id": { - "type": "string" - }, - "ceiling_micro_usd": { - "type": "integer", - "format": "uint64", - "minimum": 0 - }, - "reserved_micro_usd": { - "type": "integer", - "format": "uint64", - "minimum": 0 - }, - "spent_micro_usd": { - "type": "integer", - "format": "uint64", - "minimum": 0 - }, - "window_ms": { - "type": "integer", - "format": "uint64", - "minimum": 0 - } - }, - "additionalProperties": false, - "required": [ - "caller_id", - "window_ms", - "ceiling_micro_usd", - "spent_micro_usd", - "reserved_micro_usd" - ] - }, - "PlanQuota": { - "description": "One window of the subscription plan, as the provider reported it.", - "type": "object", - "properties": { - "observed_at_unix_ms": { - "description": "When this window was last reported. Its age in a readout is the\nreadout's `at_unix_ms` minus this.", - "type": "integer", - "format": "uint64", - "minimum": 0 - }, - "resets_at_unix_ms": { - "type": "integer", - "format": "uint64", - "minimum": 0 - }, - "status": { - "type": "string" - }, - "utilization_permille": { - "description": "How much of the window was used when it was observed, in thousandths;\n`None` once the window has reset (`status` is then `expired`), because\nthe old reading says nothing about the new window.", - "type": [ - "integer", - "null" - ], - "format": "uint32", - "minimum": 0 - }, - "window": { - "type": "string" - } - }, - "additionalProperties": false, - "required": [ - "window", - "resets_at_unix_ms", - "status", - "observed_at_unix_ms" - ] - } - } -} diff --git a/schemas/transport_envelope.json b/schemas/transport_envelope.json index 8b062e3..3b1e7db 100644 --- a/schemas/transport_envelope.json +++ b/schemas/transport_envelope.json @@ -47,8 +47,8 @@ "enum": [ "Invocation", "Result", - "SpendReadoutRequest", - "SpendReadoutResult" + "UsageReadoutRequest", + "UsageReadoutResult" ] } }