WIP call-record fix1: run ledger owned by the store (does not build until verified)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Meta Vens 2026-10-10 16:45:18 +02:00
commit bd8dc756d9
3 changed files with 524 additions and 144 deletions

View file

@ -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<InferenceEvent>,
receipt_digest: String,
usage: Option<TokenUsage>,
@ -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<TokenUsage> {
@ -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<ProviderCallRecord, CallRecordError> {
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<String, InferenceFault> {
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<T: Serialize + for<'de> Deserialize<'de>>(
row: &CultCacheEnvelope,
) -> Result<T, CallRecordError> {
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<DecodedRow, CallRecordError> {
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<String, RunTally>,
/// Runs with a recording port alive; one ledger has one owner.
owned: BTreeSet<String>,
}
/// 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<StoreState>,
}
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<Path>) -> Result<Self, CallRecordError> {
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<String, RunTally> = BTreeMap::new();
let mut counts: BTreeMap<String, u64> = BTreeMap::new();
let mut charged: BTreeMap<String, u64> = 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<u64, CallRecordError> {
/// Tokens this run has reserved or been charged: its persisted ledger.
pub fn spent(&self, run: &RunId) -> Result<u64, CallRecordError> {
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<dyn InferencePort>,
store: Arc<CallRecordStore>,
run: RunId,
cap: u64,
ledger: Mutex<Ledger>,
}
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<dyn InferencePort>,
store: Arc<CallRecordStore>,
run: RunId,
cap_tokens: u64,
) -> Result<Self, CallRecordError> {
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<InferenceOutput, InferenceFault>,
) -> Result<InferenceOutput, InferenceFault> {
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<InferenceOutput, InferenceFault> {
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",
));

View file

@ -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")));
}

View file

@ -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,
};