diff --git a/crates/ghostlight/src/verse/calls.rs b/crates/ghostlight/src/verse/calls.rs index 21d5c72..cfd3460 100644 --- a/crates/ghostlight/src/verse/calls.rs +++ b/crates/ghostlight/src/verse/calls.rs @@ -10,33 +10,37 @@ //! Lanes and runners never count tokens. A request is held against the cap by //! its bound: the serialized provider request's bytes (a token is at least one //! byte, so bytes bound the prompt) plus the request's own output allowance. -//! The bound is exact for a port that makes one provider round per call (the -//! hosted and local lanes). A port that loops tools inside one call can spend -//! more than one round's bound; its settled usage is still charged, so the next -//! call is refused, but this decorator cannot stop the overshoot of that call. //! -//! Charging rule, one path: a reply that reports usage is charged prompt plus -//! completion (cached prompt tokens are part of the prompt); a reply that -//! reports none is charged its bound; a fault is charged its bound, except a -//! connect fault, which never reached the provider and costs nothing. +//! The run's spend is one persisted ledger row per run, owned by the store and +//! written where it is spent. A call reserves its bound in that row before the +//! provider is called; a reply settles the reservation to what it reported +//! (prompt plus completion, cached prompt tokens being part of the prompt; none +//! reported is charged the bound); a fault leaves the reservation charged. A +//! restart reads the row, so no restart returns budget. A reply that reports +//! more than the call's bound is charged in full and faults the call, so the +//! run's spend never passes the cap unnoticed. One port at a time owns a run. use crate::controllers::{ - InferenceEvent, InferenceFault, InferenceFaultClass, InferenceOutput, InferencePort, - InferencePurpose, InferenceRequest, PreparedInference, REQUEST_EXPIRY, TokenUsage, - ToolResultOracle, unix_ms, + InferenceEvent, InferenceFault, InferenceOutput, InferencePort, InferencePurpose, + InferenceRequest, PreparedInference, REQUEST_EXPIRY, TokenUsage, ToolResultOracle, unix_ms, }; use async_trait::async_trait; use chrono::Utc; +use codex_connector::CodexProviderRequest; use cultcache_rs::{CacheBackingStore, CultCacheEnvelope, OwnedRedbMessagePackBackingStore}; use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; use std::{ - collections::{BTreeMap, VecDeque}, + collections::{BTreeMap, BTreeSet, VecDeque}, path::Path, sync::{Arc, Mutex}, }; use thiserror::Error; const CALL_RECORD_ROW: &str = "verse_provider_call.v1"; +const RUN_SPEND_ROW: &str = "verse_run_spend.v1"; +/// The run ledger document's schema. +const RUN_SPEND_SCHEMA: &str = "ghostlight.verse_run_spend.v1"; /// The call record document's schema. A store holding any other row type or /// schema is refused at open: this store is not migrated from, or shared with, /// another document. @@ -54,6 +58,10 @@ pub enum CallRecordError { OutOfSequence, #[error("replay needs one run's records in dense sequence from 1")] NotReplayable, + #[error("another recording port already owns this run")] + RunInUse, + #[error("the call could pass the run's token cap")] + OverCap, } /// A run's identity as a key segment. @@ -79,7 +87,7 @@ impl RunId { } } -/// One provider call of one run: the request's digest (never its text, so no +/// One provider call of one run: the request's content digest (never its text, so no /// prompt or credential is stored), the reply exactly as the lane received it, /// and what the call was charged. Keyed (run id, sequence). #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] @@ -88,7 +96,7 @@ pub struct ProviderCallRecord { sequence: u64, purpose: InferencePurpose, model: String, - request_sha256: String, + content_sha256: String, events: Vec, receipt_digest: String, usage: Option, @@ -112,9 +120,9 @@ impl ProviderCallRecord { &self.model } - /// Lowercase hex of the provider request's digest. - pub fn request_sha256(&self) -> &str { - &self.request_sha256 + /// Lowercase hex of the request's content digest (see `content_digest`). + pub fn content_sha256(&self) -> &str { + &self.content_sha256 } pub fn usage(&self) -> Option { @@ -140,49 +148,110 @@ fn row_key(run: &str, sequence: u64) -> String { format!("{run}/{sequence:020}") } -fn hex(digest: &[u8; 32]) -> String { +fn hex(digest: &[u8]) -> String { digest.iter().map(|byte| format!("{byte:02x}")).collect() } -fn decode_row(row: &CultCacheEnvelope) -> Result { - if row.r#type != CALL_RECORD_ROW || row.schema_id.as_deref() != Some(CALL_RECORD_SCHEMA) { - return Err(CallRecordError::Corrupt); +/// What a request asks of the provider, without the ids that name it. The +/// request and conversation ids are derived from the command id, and a command +/// id from the world it ran in, so keying on them would tie a recorded run to +/// one world's random identity. Everything else the provider is sent +/// (model, instructions, input, tools, shape) is in the digest. +fn content_digest(request: &CodexProviderRequest) -> Result { + let bytes = rmp_serde::to_vec(&( + &request.schema_id, + &request.model, + &request.instructions, + &request.input, + &request.reasoning_effort, + &request.reasoning_summary, + &request.service_tier, + &request.output_format_name, + &request.previous_response_id, + &request.tools, + &request.tool_choice, + request.parallel_tool_calls, + &request.output_schema_json, + request.max_output_tokens, + &request.prompt_cache_key, + )) + .map_err(|_| InferenceFault::new("a provider request could not be measured"))?; + Ok(hex(&Sha256::digest(bytes))) +} + +/// A run's persisted ledger: tokens reserved or charged against its cap. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +struct RunSpend { + run_id: String, + spent: u64, +} + +enum DecodedRow { + Call(ProviderCallRecord), + Spend(RunSpend), +} + +fn canonical Deserialize<'de>>( + row: &CultCacheEnvelope, +) -> Result { + let value: T = rmp_serde::from_slice(&row.payload).map_err(|_| CallRecordError::Corrupt)?; + let again = rmp_serde::to_vec_named(&value).map_err(|_| CallRecordError::Corrupt)?; + if again == row.payload { + Ok(value) + } else { + Err(CallRecordError::Corrupt) } - let record: ProviderCallRecord = - rmp_serde::from_slice(&row.payload).map_err(|_| CallRecordError::Corrupt)?; - let canonical = rmp_serde::to_vec_named(&record).map_err(|_| CallRecordError::Corrupt)?; - if canonical != row.payload - || RunId::new(&record.run_id).is_err() - || record.sequence == 0 - || row.key != row_key(&record.run_id, record.sequence) - { - return Err(CallRecordError::Corrupt); +} + +fn decode_row(row: &CultCacheEnvelope) -> Result { + match (row.r#type.as_str(), row.schema_id.as_deref()) { + (CALL_RECORD_ROW, Some(CALL_RECORD_SCHEMA)) => { + let record: ProviderCallRecord = canonical(row)?; + if RunId::new(&record.run_id).is_err() + || record.sequence == 0 + || row.key != row_key(&record.run_id, record.sequence) + { + return Err(CallRecordError::Corrupt); + } + Ok(DecodedRow::Call(record)) + } + (RUN_SPEND_ROW, Some(RUN_SPEND_SCHEMA)) => { + let spend: RunSpend = canonical(row)?; + if RunId::new(&spend.run_id).is_err() || row.key != spend.run_id { + return Err(CallRecordError::Corrupt); + } + Ok(DecodedRow::Spend(spend)) + } + _ => Err(CallRecordError::Corrupt), } - Ok(record) } #[derive(Default)] struct RunTally { last_sequence: u64, - charged: u64, + /// Tokens reserved or charged: the persisted ledger, read at open. + spent: u64, } struct StoreState { store: OwnedRedbMessagePackBackingStore, runs: BTreeMap, + /// Runs with a recording port alive; one ledger has one owner. + owned: BTreeSet, } -/// The call records of every run in one CultCache file. Records are appended -/// and never changed; a run's sequence is dense from 1. +/// The call records and spend ledgers of every run in one CultCache file. +/// Records are appended and never changed; a run's sequence is dense from 1. pub struct CallRecordStore { state: Mutex, } impl CallRecordStore { /// Opens the file, creating it empty when absent. Every row must be a - /// canonical call record of this schema with its key and a dense sequence; - /// anything else, an earlier or later schema version included, is refused - /// and nothing is migrated. + /// canonical call record or run ledger of this schema with its key, a + /// dense sequence, and a ledger that covers its records' charges; anything + /// else, an earlier or later schema version included, is refused and + /// nothing is migrated. pub fn open(path: impl AsRef) -> Result { let store = OwnedRedbMessagePackBackingStore::new(path.as_ref()) .map_err(|_| CallRecordError::Store)?; @@ -192,20 +261,37 @@ impl CallRecordStore { let rows = store.pull_all().map_err(|_| CallRecordError::Store)?; let mut runs: BTreeMap = BTreeMap::new(); let mut counts: BTreeMap = BTreeMap::new(); + let mut charged: BTreeMap = BTreeMap::new(); for row in &rows { - let record = decode_row(row)?; - *counts.entry(record.run_id.clone()).or_default() += 1; - let tally = runs.entry(record.run_id).or_default(); - tally.last_sequence = tally.last_sequence.max(record.sequence); - tally.charged = tally.charged.saturating_add(record.charged); + match decode_row(row)? { + DecodedRow::Call(record) => { + *counts.entry(record.run_id.clone()).or_default() += 1; + let total = charged.entry(record.run_id.clone()).or_default(); + *total = total.saturating_add(record.charged); + let tally = runs.entry(record.run_id).or_default(); + tally.last_sequence = tally.last_sequence.max(record.sequence); + } + DecodedRow::Spend(spend) => { + runs.entry(spend.run_id).or_default().spent = spend.spent; + } + } } // Keys are unique, so a run is dense from 1 exactly when its rows - // number its highest sequence. - if runs.iter().any(|(run, tally)| counts[run] != tally.last_sequence) { - return Err(CallRecordError::Corrupt); + // number its highest sequence. Every record was reserved before it was + // made, so the ledger covers what the records charged. + for (run, tally) in &runs { + let rows = counts.get(run).copied().unwrap_or(0); + let recorded = charged.get(run).copied().unwrap_or(0); + if rows != tally.last_sequence || tally.spent < recorded { + return Err(CallRecordError::Corrupt); + } } Ok(Self { - state: Mutex::new(StoreState { store, runs }), + state: Mutex::new(StoreState { + store, + runs, + owned: BTreeSet::new(), + }), }) } @@ -223,17 +309,90 @@ impl CallRecordStore { + 1) } - /// Tokens this run's stored calls were charged. - pub fn charged(&self, run: &RunId) -> Result { + /// Tokens this run has reserved or been charged: its persisted ledger. + pub fn spent(&self, run: &RunId) -> Result { Ok(self .lock()? .runs .get(run.as_str()) - .map_or(0, |tally| tally.charged)) + .map_or(0, |tally| tally.spent)) + } + + /// Takes the run's ledger for one port, or refuses when another holds it. + fn own(&self, run: &RunId) -> Result<(), CallRecordError> { + if self.lock()?.owned.insert(run.as_str().to_owned()) { + Ok(()) + } else { + Err(CallRecordError::RunInUse) + } + } + + fn disown(&self, run: &RunId) { + if let Ok(mut state) = self.state.lock() { + state.owned.remove(run.as_str()); + } + } + + /// Persists `spent` as the run's ledger, then adopts it in memory. A write + /// that fails leaves the previous figure. + fn write_spent(state: &mut StoreState, run: &str, spent: u64) -> Result<(), CallRecordError> { + let spend = RunSpend { + run_id: run.to_owned(), + spent, + }; + let row = CultCacheEnvelope { + key: run.to_owned(), + r#type: RUN_SPEND_ROW.into(), + payload: rmp_serde::to_vec_named(&spend).map_err(|_| CallRecordError::Store)?, + stored_at: Utc::now().to_rfc3339(), + schema_id: Some(RUN_SPEND_SCHEMA.into()), + }; + state.store.push(&row).map_err(|_| CallRecordError::Store)?; + state.runs.entry(run.to_owned()).or_default().spent = spent; + Ok(()) + } + + /// Charges `bound` to the run's persisted ledger, or refuses with + /// `OverCap` when the ledger plus the bound would pass `cap`. The check and + /// the write hold one lock, so calls on the wire hold their bounds. + fn reserve(&self, run: &RunId, bound: u64, cap: u64) -> Result<(), CallRecordError> { + let mut state = self.lock()?; + let held = state.runs.get(run.as_str()).map_or(0, |tally| tally.spent); + let after = held.saturating_add(bound); + if after > cap { + return Err(CallRecordError::OverCap); + } + Self::write_spent(&mut state, run.as_str(), after) + } + + /// Replaces a reservation of `bound` with what the call was charged. With + /// a record, the record is stored first, so a crash between the two writes + /// leaves the larger reservation charged. + fn settle( + &self, + run: &RunId, + bound: u64, + charged: u64, + record: Option<&ProviderCallRecord>, + ) -> Result<(), CallRecordError> { + let mut state = self.lock()?; + if let Some(record) = record { + Self::push_record(&mut state, record)?; + } + let held = state.runs.get(run.as_str()).map_or(0, |tally| tally.spent); + let settled = held.saturating_sub(bound).saturating_add(charged); + Self::write_spent(&mut state, run.as_str(), settled) } fn append(&self, record: &ProviderCallRecord) -> Result<(), CallRecordError> { let mut state = self.lock()?; + Self::push_record(&mut state, record) + } + + fn push_record( + state: &mut StoreState, + record: &ProviderCallRecord, + ) -> Result<(), CallRecordError> { let last = state .runs .get(&record.run_id) @@ -249,9 +408,11 @@ impl CallRecordStore { schema_id: Some(CALL_RECORD_SCHEMA.into()), }; state.store.push(&row).map_err(|_| CallRecordError::Store)?; - let tally = state.runs.entry(record.run_id.clone()).or_default(); - tally.last_sequence = record.sequence; - tally.charged = tally.charged.saturating_add(record.charged); + state + .runs + .entry(record.run_id.clone()) + .or_default() + .last_sequence = record.sequence; Ok(()) } @@ -265,54 +426,46 @@ impl CallRecordStore { .map_err(|_| CallRecordError::Store)?; let mut records = Vec::new(); for row in rows.iter().filter(|row| row.key.starts_with(&prefix)) { - records.push(decode_row(row)?); + if let DecodedRow::Call(record) = decode_row(row)? { + records.push(record); + } } records.sort_by_key(|record| record.sequence); Ok(records) } } -struct Ledger { - /// Tokens charged by finished calls, restored from the store at open. - settled: u64, - /// Bounds held by calls still on the wire. - in_flight: u64, -} - -/// Decorates a port for one run: refuses a call that could pass the cap, -/// charges the real cost, and stores every reply before returning it. +/// Decorates a port for one run: reserves each call's bound in the run's +/// ledger before the provider is called, settles it to the real cost, and +/// stores every reply before returning it. pub struct RecordingInferencePort { inner: Arc, store: Arc, run: RunId, cap: u64, - ledger: Mutex, } impl RecordingInferencePort { - /// The run's spend resumes from its stored records. + /// Takes ownership of the run's ledger, which resumes from the store. A + /// second port on a run whose first is alive is refused with `RunInUse`. pub fn open( inner: Arc, store: Arc, run: RunId, cap_tokens: u64, ) -> Result { - let settled = store.charged(&run)?; + store.own(&run)?; Ok(Self { inner, store, run, cap: cap_tokens, - ledger: Mutex::new(Ledger { - settled, - in_flight: 0, - }), }) } - /// Tokens charged to the run so far. + /// Tokens reserved or charged to the run so far. pub fn spent(&self) -> u64 { - self.ledger.lock().map_or(u64::MAX, |ledger| ledger.settled) + self.store.spent(&self.run).unwrap_or(u64::MAX) } /// The most this request can cost: its serialized bytes plus its output @@ -329,62 +482,49 @@ impl RecordingInferencePort { Ok((payload.len() as u64).saturating_add(u64::from(allowance))) } - fn reserve(&self, bound: u64) -> Result<(), InferenceFault> { - let mut ledger = self - .ledger - .lock() - .map_err(|_| InferenceFault::integrity_violation("the run's ledger was poisoned"))?; - let held = ledger.settled.saturating_add(ledger.in_flight); - if held.saturating_add(bound) > self.cap { - return Err(InferenceFault::budget_exhausted()); - } - ledger.in_flight += bound; - Ok(()) - } - fn settle( &self, bound: u64, identity: (InferencePurpose, String, String), result: Result, ) -> Result { - let mut ledger = self - .ledger - .lock() - .map_err(|_| InferenceFault::integrity_violation("the run's ledger was poisoned"))?; - ledger.in_flight = ledger.in_flight.saturating_sub(bound); - let output = match result { - Ok(output) => output, - Err(fault) => { - if fault.class() != InferenceFaultClass::Connect { - ledger.settled = ledger.settled.saturating_add(bound); - } - return Err(fault); - } - }; + // A fault leaves its reservation charged: the request may have run. + let output = result?; let charged = match output.usage { Some(usage) => usage.prompt.saturating_add(usage.completion), None => bound, }; - ledger.settled = ledger.settled.saturating_add(charged); - let (purpose, model, request_sha256) = identity; - let sequence = self.store.next_sequence(&self.run).map_err(unrecorded)?; + if charged > bound { + self.store + .settle(&self.run, bound, charged, None) + .map_err(unrecorded)?; + return Err(InferenceFault::budget_exhausted()); + } + let (purpose, model, content_sha256) = identity; let record = ProviderCallRecord { run_id: self.run.as_str().to_owned(), - sequence, + sequence: self.store.next_sequence(&self.run).map_err(unrecorded)?, purpose, model, - request_sha256, + content_sha256, events: output.events.clone(), receipt_digest: output.receipt_digest.clone(), usage: output.usage, charged, }; - self.store.append(&record).map_err(unrecorded)?; + self.store + .settle(&self.run, bound, charged, Some(&record)) + .map_err(unrecorded)?; Ok(output) } } +impl Drop for RecordingInferencePort { + fn drop(&mut self) { + self.store.disown(&self.run); + } +} + fn unrecorded(_: CallRecordError) -> InferenceFault { InferenceFault::integrity_violation("a provider call could not be recorded") } @@ -397,12 +537,17 @@ impl InferencePort for RecordingInferencePort { async fn infer(&self, request: PreparedInference) -> Result { let bound = Self::bound(&request)?; - self.reserve(bound)?; let identity = ( request.purpose, request.invocation.request.model.clone(), - hex(&request.invocation.provider_request_sha256), + content_digest(&request.invocation.request)?, ); + self.store + .reserve(&self.run, bound, self.cap) + .map_err(|error| match error { + CallRecordError::OverCap => InferenceFault::budget_exhausted(), + _ => unrecorded(error), + })?; let result = self.inner.infer(request).await; self.settle(bound, identity, result) } @@ -412,7 +557,8 @@ impl InferencePort for RecordingInferencePort { } } -/// Serves a run's stored replies, in sequence, to the same requests. It holds +/// Serves a run's stored replies, in sequence, to the same requests, matched +/// by what they ask of the provider and not by the ids that name them. It holds /// no provider and makes no call: a request that is not the next recorded one /// is a fault. pub struct ReplayInferencePort { @@ -454,13 +600,13 @@ impl InferencePort for ReplayInferencePort { .pending .lock() .map_err(|_| InferenceFault::integrity_violation("the replay queue was poisoned"))?; - let digest = hex(&request.invocation.provider_request_sha256); + let digest = content_digest(&request.invocation.request)?; let Some(next) = pending.front() else { return Err(InferenceFault::integrity_violation( "replay has no recorded call left", )); }; - if next.request_sha256 != digest { + if next.content_sha256 != digest { return Err(InferenceFault::integrity_violation( "the request is not the next recorded call", )); diff --git a/crates/ghostlight/src/verse/calls/tests.rs b/crates/ghostlight/src/verse/calls/tests.rs index eaecd1b..8de27c4 100644 --- a/crates/ghostlight/src/verse/calls/tests.rs +++ b/crates/ghostlight/src/verse/calls/tests.rs @@ -1,6 +1,8 @@ use super::*; use crate::CommandId; -use crate::controllers::{InferenceFaultDisposition, RequestShape, tool_request}; +use crate::controllers::{ + InferenceFaultClass, InferenceFaultDisposition, RequestShape, tool_request, +}; use codex_connector::{CodexInputItem, CodexToolDefinition}; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -171,7 +173,7 @@ async fn recording_port_stores_each_call_before_returning() { let prepared = port .prepare(request_for(CommandId::new(), "tell", 1_200)) .unwrap(); - let digest = hex(&prepared.invocation.provider_request_sha256); + let digest = content_digest(&prepared.invocation.request).unwrap(); let returned = port.infer(prepared).await.unwrap(); assert_eq!(json(&returned), json(expected)); let stored = store.records(&run).unwrap(); @@ -183,7 +185,7 @@ async fn recording_port_stores_each_call_before_returning() { let record = &stored[index]; assert_eq!(record.sequence(), index as u64 + 1); assert_eq!(record.run_id(), "run-a"); - assert_eq!(record.request_sha256(), digest); + assert_eq!(record.content_sha256(), digest); assert_eq!(record.model(), MODEL); assert_eq!(record.purpose(), InferencePurpose::Persona); assert_eq!(record.usage(), expected.usage); @@ -251,6 +253,7 @@ async fn cap_refuses_before_the_provider_is_called() { .is_empty() ); assert_eq!(port.spent(), 0); + drop(port); // Exactly the cap is allowed. let exact = recording(&inner, &store, "run-cap", bound); @@ -389,7 +392,7 @@ async fn a_reply_with_no_usage_is_charged_its_bound() { } #[tokio::test] -async fn a_fault_is_charged_its_bound_unless_it_never_connected() { +async fn a_fault_keeps_its_reservation_charged() { let directory = tempfile::tempdir().unwrap(); let (store, _) = store_in(&directory); let prepared = prepared_for("fault", 1_200); @@ -397,23 +400,21 @@ async fn a_fault_is_charged_its_bound_unless_it_never_connected() { let inner = ScriptedPort::new(vec![ Err(InferenceFault::retryable("down").classed(InferenceFaultClass::Connect)), Err(InferenceFault::retryable("slow").classed(InferenceFaultClass::Timeout)), + Err(InferenceFault::new("junk").classed(InferenceFaultClass::BadReply)), ]); let port = recording(&inner, &store, "run-fault", u64::MAX); - port.infer(prepared.clone()) - .await - .expect_err("connect fault"); - assert_eq!(port.spent(), 0, "nothing was sent"); - let fault = port.infer(prepared).await.expect_err("timeout fault"); - assert_eq!( - fault.class(), + for (index, class) in [ + InferenceFaultClass::Connect, InferenceFaultClass::Timeout, - "the fault passes through" - ); - assert_eq!( - port.spent(), - bound, - "a request that may have run is charged its bound" - ); + InferenceFaultClass::BadReply, + ] + .into_iter() + .enumerate() + { + let fault = port.infer(prepared.clone()).await.expect_err("a fault"); + assert_eq!(fault.class(), class, "the fault passes through"); + assert_eq!(port.spent(), bound * (index as u64 + 1)); + } assert!( store .records(&RunId::new("run-fault").unwrap()) @@ -422,6 +423,151 @@ async fn a_fault_is_charged_its_bound_unless_it_never_connected() { ); } +#[tokio::test] +async fn faulting_calls_across_restarts_never_pass_the_cap() { + let directory = tempfile::tempdir().unwrap(); + let (store, path) = store_in(&directory); + let prepared = prepared_for("loop", 1_200); + let bound = bound_of(&prepared); + let cap = bound * 3 + bound / 2; + drop(store); + let mut reached_the_provider = 0; + for attempt in 0..12 { + let store = Arc::new(CallRecordStore::open(&path).unwrap()); + let class = if attempt % 2 == 0 { + InferenceFaultClass::BadReply + } else { + InferenceFaultClass::Timeout + }; + let inner = ScriptedPort::new(vec![Err(InferenceFault::new("down").classed(class))]); + let port = recording(&inner, &store, "run-loop", cap); + let fault = port.infer(prepared.clone()).await.expect_err("a fault"); + reached_the_provider += inner.calls(); + if inner.calls() == 0 { + assert_eq!(fault.class(), InferenceFaultClass::BudgetExhausted); + } + assert!(port.spent() <= cap, "{} > {cap}", port.spent()); + } + assert_eq!(reached_the_provider, 3, "a restart returns no budget"); +} + +#[tokio::test] +async fn a_call_abandoned_on_the_wire_stays_charged_after_a_restart() { + let directory = tempfile::tempdir().unwrap(); + let (store, path) = store_in(&directory); + let prepared = prepared_for("crash", 1_200); + let bound = bound_of(&prepared); + let gate = Arc::new(tokio::sync::Semaphore::new(0)); + let inner = ScriptedPort::gated(vec![Ok(reply("never", "r"))], gate); + let port = Arc::new(recording(&inner, &store, "run-crash", bound)); + let call = { + let port = port.clone(); + let prepared = prepared.clone(); + tokio::spawn(async move { port.infer(prepared).await }) + }; + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while inner.calls() == 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("the call reached the provider"); + call.abort(); + let _ = call.await; + drop((port, store)); + + let store = Arc::new(CallRecordStore::open(&path).unwrap()); + let inner = ScriptedPort::new(vec![Ok(reply("again", "r2"))]); + let reopened = recording(&inner, &store, "run-crash", bound); + assert_eq!(reopened.spent(), bound); + let fault = reopened + .infer(prepared) + .await + .expect_err("the cap is spent"); + assert_eq!(fault.class(), InferenceFaultClass::BudgetExhausted); + assert_eq!(inner.calls(), 0); +} + +#[tokio::test] +async fn one_run_has_one_ledger_and_one_owner() { + let directory = tempfile::tempdir().unwrap(); + let (store, _) = store_in(&directory); + let prepared = prepared_for("owner", 1_200); + let bound = bound_of(&prepared); + let cap = bound + bound / 2; + let inner = ScriptedPort::new(vec![Ok(reply("a", "r1")), Ok(reply("b", "r2"))]); + let first = recording(&inner, &store, "run-own", cap); + let second = RecordingInferencePort::open( + inner.clone(), + store.clone(), + RunId::new("run-own").unwrap(), + cap, + ); + assert_eq!(second.err(), Some(CallRecordError::RunInUse)); + // Another run is another ledger. + drop(recording(&inner, &store, "run-other", cap)); + first.infer(prepared.clone()).await.unwrap(); + assert_eq!(first.spent(), bound); + // The owner going away frees the run, and its spend carries over. + drop(first); + let next = recording(&inner, &store, "run-own", cap); + assert_eq!(next.spent(), bound); + next.infer(prepared) + .await + .expect_err("the cap holds across owners"); + assert_eq!(inner.calls(), 1); +} + +#[tokio::test] +async fn usage_over_the_bound_is_charged_in_full_and_faults_the_call() { + let directory = tempfile::tempdir().unwrap(); + let (store, _) = store_in(&directory); + let prepared = prepared_for("liar", 1_200); + let bound = bound_of(&prepared); + let cap = bound + 200; + let run = RunId::new("run-over").unwrap(); + let inner = ScriptedPort::new(vec![ + Ok(reply("big", "r1").with_usage(usage(bound, bound, None))), + Ok(reply("never", "r2")), + ]); + let port = recording(&inner, &store, "run-over", cap); + let fault = port.infer(prepared.clone()).await.expect_err("over bound"); + assert_eq!(fault.class(), InferenceFaultClass::BudgetExhausted); + assert_eq!( + fault.disposition(), + InferenceFaultDisposition::IntegrityViolation + ); + assert_eq!(port.spent(), bound * 2, "charged in full"); + assert!(store.records(&run).unwrap().is_empty()); + // The run is stopped: the next call is refused before the provider. + let fault = port.infer(prepared).await.expect_err("the cap is passed"); + assert_eq!(fault.class(), InferenceFaultClass::BudgetExhausted); + assert_eq!(inner.calls(), 1); +} + +#[tokio::test] +async fn usage_exactly_at_the_bound_is_a_normal_reply() { + let directory = tempfile::tempdir().unwrap(); + let (store, _) = store_in(&directory); + let prepared = prepared_for("exact", 1_200); + let bound = bound_of(&prepared); + let inner = ScriptedPort::new(vec![Ok(reply("fits", "r1").with_usage(usage( + bound - 10, + 10, + None, + )))]); + let port = recording(&inner, &store, "run-exact", bound); + port.infer(prepared).await.expect("a reply at its bound"); + assert_eq!(port.spent(), bound); + assert_eq!( + store + .records(&RunId::new("run-exact").unwrap()) + .unwrap() + .len(), + 1 + ); +} + #[tokio::test] async fn a_reopened_run_resumes_its_spend() { let directory = tempfile::tempdir().unwrap(); @@ -477,7 +623,7 @@ async fn replay_serves_stored_responses_and_refuses_strangers() { let replay = ReplayInferencePort::new("verse-replay", stored).unwrap(); // A request the run never made is a fault and consumes nothing. let stranger = replay - .prepare(request_for(CommandId::new(), "a", 1_200)) + .prepare(request_for(CommandId::new(), "z", 1_200)) .unwrap(); let fault = replay.infer(stranger).await.expect_err("a stranger"); assert_eq!( @@ -488,10 +634,7 @@ async fn replay_serves_stored_responses_and_refuses_strangers() { let second_first = replay .prepare(request_for(commands[1], "b", 1_200)) .unwrap(); - replay - .infer(second_first) - .await - .expect_err("out of order"); + replay.infer(second_first).await.expect_err("out of order"); // In order, each request gets exactly what the run received. for ((command, prompt), expected) in commands.iter().zip(["a", "b", "c"]).zip(&outputs) { let prepared = replay @@ -513,7 +656,7 @@ fn sample_record(run: &str, sequence: u64) -> ProviderCallRecord { sequence, purpose: InferencePurpose::Persona, model: MODEL.into(), - request_sha256: "ab".into(), + content_sha256: "ab".into(), events: vec![InferenceEvent::Text("t".into())], receipt_digest: "r".into(), usage: None, @@ -551,7 +694,7 @@ async fn a_record_holds_no_request_text() { .unwrap() .pull_all() .unwrap(); - assert_eq!(raw.len(), 1); + assert_eq!(raw.len(), 2, "one record and the run's ledger"); for row in raw { let bytes = [row.key.as_bytes(), row.payload.as_slice()].concat(); assert!( @@ -565,6 +708,19 @@ async fn a_record_holds_no_request_text() { assert!(!format!("{refusal} {refusal:?}").contains(canary)); } +fn spend_row(run: &str, spent: u64) -> CultCacheEnvelope { + envelope( + RUN_SPEND_ROW, + Some(RUN_SPEND_SCHEMA), + run, + rmp_serde::to_vec_named(&RunSpend { + run_id: run.into(), + spent, + }) + .unwrap(), + ) +} + fn envelope( row_type: &str, schema: Option<&str>, @@ -640,6 +796,37 @@ fn the_store_refuses_any_row_that_is_not_a_canonical_v1_record() { ), ("sequence zero", record_at(0)), ("a hole in the run", record_at(2)), + ("a record with no ledger", record_at(1)), + ( + "a ledger at the wrong key", + envelope( + RUN_SPEND_ROW, + Some(RUN_SPEND_SCHEMA), + "other", + rmp_serde::to_vec_named(&RunSpend { + run_id: "run-x".into(), + spent: 1, + }) + .unwrap(), + ), + ), + ( + "a ledger that is not canonical", + envelope( + RUN_SPEND_ROW, + Some(RUN_SPEND_SCHEMA), + "run-x", + [ + rmp_serde::to_vec_named(&RunSpend { + run_id: "run-x".into(), + spent: 1, + }) + .unwrap(), + vec![0xc0], + ] + .concat(), + ), + ), ]; for (label, row) in cases { let directory = tempfile::tempdir().unwrap(); @@ -651,17 +838,28 @@ fn the_store_refuses_any_row_that_is_not_a_canonical_v1_record() { let error = CallRecordStore::open(&path).err(); assert_eq!(error, Some(CallRecordError::Corrupt), "{label}"); } - // The same row, well formed, opens. + // A ledger under the charges of its records is refused too. let directory = tempfile::tempdir().unwrap(); let path = directory.path().join("calls.redb"); - OwnedRedbMessagePackBackingStore::new(&path) - .unwrap() - .push(&v1(&key, payload)) - .unwrap(); + let mut backing = OwnedRedbMessagePackBackingStore::new(&path).unwrap(); + backing.push(&v1(&key, payload.clone())).unwrap(); + backing.push(&spend_row("run-x", 2)).unwrap(); + drop(backing); + assert_eq!( + CallRecordStore::open(&path).err(), + Some(CallRecordError::Corrupt) + ); + // The same rows, well formed and covered by the ledger, open. + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("calls.redb"); + let mut backing = OwnedRedbMessagePackBackingStore::new(&path).unwrap(); + backing.push(&v1(&key, payload)).unwrap(); + backing.push(&spend_row("run-x", 3)).unwrap(); + drop(backing); let store = CallRecordStore::open(&path).unwrap(); let run = RunId::new("run-x").unwrap(); assert_eq!(store.records(&run).unwrap(), vec![good]); - assert_eq!(store.charged(&run).unwrap(), 3); + assert_eq!(store.spent(&run).unwrap(), 3); } #[test] @@ -694,3 +892,39 @@ fn a_run_id_is_a_key_safe_segment() { assert_eq!(RunId::new(bad), Err(CallRecordError::InvalidRunId)); } } + +#[tokio::test] +async fn replay_keys_the_request_content_not_the_command_identity() { + let directory = tempfile::tempdir().unwrap(); + let (store, _) = store_in(&directory); + let run = RunId::new("run-world").unwrap(); + let inner = ScriptedPort::new(ok_all(&[reply("one", "r1"), reply("two", "r2")])); + let live = recording(&inner, &store, "run-world", u64::MAX); + for prompt in ["first", "second"] { + live.infer( + live.prepare(request_for(CommandId::new(), prompt, 1_200)) + .unwrap(), + ) + .await + .unwrap(); + } + // A fresh genesis names every command anew; the content is what repeats. + let replay = ReplayInferencePort::new("verse-replay", store.records(&run).unwrap()).unwrap(); + let under_new_names = |prompt: &str, max: u32| { + replay + .prepare(request_for(CommandId::new(), prompt, max)) + .unwrap() + }; + // The same words under a different output allowance are another request. + replay + .infer(under_new_names("first", 1_201)) + .await + .expect_err("a different shape"); + let one = replay.infer(under_new_names("first", 1_200)).await.unwrap(); + let two = replay + .infer(under_new_names("second", 1_200)) + .await + .unwrap(); + assert_eq!(json(&one), json(&reply("one", "r1"))); + assert_eq!(json(&two), json(&reply("two", "r2"))); +} diff --git a/crates/ghostlight/src/verse/mod.rs b/crates/ghostlight/src/verse/mod.rs index 5c1e308..7c7b6fb 100644 --- a/crates/ghostlight/src/verse/mod.rs +++ b/crates/ghostlight/src/verse/mod.rs @@ -5,6 +5,6 @@ mod calls; pub use calls::{ - CALL_RECORD_SCHEMA, CallRecordError, CallRecordStore, ProviderCallRecord, RecordingInferencePort, - ReplayInferencePort, RunId, + CALL_RECORD_SCHEMA, CallRecordError, CallRecordStore, ProviderCallRecord, + RecordingInferencePort, ReplayInferencePort, RunId, };