GameCult/ModelConnector
Watch
9
0
Fork
You've already forked ModelConnector
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

spend_readout: a CultNet operation beside invoke reports each caller's spent and reserved micro-USD against its ceiling and the backend's plan quota (gamecult.model.spend_readout.v1); ProviderBackend::plan_quota; ConnectorClient::spend_readout

Browse source 
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
This commit is contained in:
Meta Vens 2026-10-29 18:58:37 +01:00
commit 244531426f
8 changed files with 725 additions and 54 deletions
Download patch file Download diff file Expand all files Collapse all files

44
crates/model-connector/src/client.rs
View file

@ -4,11 +4,16 @@ use std::time::Duration;
use thiserror::Error;
use zeroize::Zeroizing;
use rand_core::{OsRng, RngCore};
use crate::contract::{TransportInvocation, TransportKey, TransportResult};
use crate::envelope::{
MAX_IDENTITY_BYTES, ServiceError, TransportFrameError, decode_result_frame, decrypt_result,
encode_invocation_frame, encrypt_invocation, read_transport_frame, write_transport_frame,
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,
};
use crate::readout::SpendReadout;
#[derive(Clone)]
pub struct ConnectorClient {
@ -52,6 +57,34 @@ impl ConnectorClient {
return Err(ConnectorClientError::RequestIdTooLong);
}
let request = encode_invocation_frame(&encrypt_invocation(invocation, &self.security)?)?;
let response = self.exchange(&request)?;
let envelope = decode_result_frame(&response)?;
decrypt_result(&envelope, &self.security, invocation).map_err(Into::into)
}
/// Reads the connector's spend ledger and plan quota as `caller_runtime_id`,
/// which must be a caller whose connection key this client holds.
pub fn spend_readout(
&self,
caller_runtime_id: &str,
) -> Result<SpendReadout, ConnectorClientError> {
let mut nonce = [0_u8; 16];
OsRng.fill_bytes(&mut nonce);
let request_id = format!(
"spend-readout-{}",
nonce.iter().map(|b| format!("{b:02x}")).collect::<String>()
);
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)?;
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)
}
/// Sends one frame on a fresh connection and reads the one answer.
fn exchange(&self, request: &[u8]) -> Result<Vec<u8>, ConnectorClientError> {
let connect_timeout = self
.response_timeout
.unwrap_or(Duration::from_secs(10))
@ -62,12 +95,9 @@ impl ConnectorClient {
.set_read_timeout(self.response_timeout)
.and_then(|()| stream.set_write_timeout(self.response_timeout))
.map_err(ConnectorClientError::Connection)?;
write_transport_frame(&mut stream, &request, self.max_frame_bytes)
write_transport_frame(&mut stream, request, self.max_frame_bytes)
.map_err(client_frame_error)?;
let response =
read_transport_frame(&mut stream, self.max_frame_bytes).map_err(client_frame_error)?;
let envelope = decode_result_frame(&response)?;
decrypt_result(&envelope, &self.security, invocation).map_err(Into::into)
read_transport_frame(&mut stream, self.max_frame_bytes).map_err(client_frame_error)
}
}

235
crates/model-connector/src/envelope.rs
View file

@ -17,11 +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;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
pub enum EnvelopeKind {
Invocation,
Result,
SpendReadoutRequest,
SpendReadoutResult,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
@ -49,8 +52,10 @@ pub(crate) fn binary_schema(_: &mut schemars::SchemaGenerator) -> schemars::Sche
/// The CultNet service every frame of this connector addresses.
pub const CONNECTOR_SERVICE_ID: &str = "gamecult.model.connector";
/// The one operation of that service: run one sealed invocation.
/// 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";
const PAYLOAD_ENCODING: &str = "messagepack-base64";
const OK_STATUS: &str = "ok";
@ -69,16 +74,39 @@ fn seal_payload(envelope: &TransportEnvelope) -> Result<String, ServiceError> {
Ok(STANDARD.encode(wire_encode(envelope)?))
}
/// The CultNet frame an invocation travels in: a `cultnet.operation_request.v0`
/// for [`CONNECTOR_SERVICE_ID`] and [`INVOKE_OPERATION`], its `message_id` the
/// The two operations of the connector service.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum Operation {
Invoke,
SpendReadout,
}
impl Operation {
fn name(self) -> &'static str {
match self {
Operation::Invoke => INVOKE_OPERATION,
Operation::SpendReadout => SPEND_READOUT_OPERATION,
}
}
fn from_name(name: &str) -> Option<Self> {
[Operation::Invoke, Operation::SpendReadout]
.into_iter()
.find(|operation| operation.name() == name)
}
}
/// The CultNet frame a request travels in: a `cultnet.operation_request.v0`
/// for [`CONNECTOR_SERVICE_ID`] and the named operation, its `message_id` the
/// request id and its payload the base64 of the named sealed envelope.
pub(crate) fn encode_invocation_frame(
fn request_frame(
operation: Operation,
envelope: &TransportEnvelope,
) -> Result<Vec<u8>, ServiceError> {
carriage_bytes(&CultNetMessage::OperationRequest {
message_id: envelope.request_id.clone(),
service_id: CONNECTOR_SERVICE_ID.to_string(),
operation: INVOKE_OPERATION.to_string(),
operation: operation.name().to_string(),
payload_schema: ENVELOPE_SCHEMA_ID.to_string(),
payload_encoding: PAYLOAD_ENCODING.to_string(),
payload: seal_payload(envelope)?,
@ -87,14 +115,17 @@ pub(crate) fn encode_invocation_frame(
})
}
/// The CultNet frame a result travels in: a `cultnet.operation_response.v0`
/// with status `ok`, otherwise shaped like the invocation's.
/// The CultNet frame a response travels in: a `cultnet.operation_response.v0`
/// with status `ok`, otherwise shaped like the request's.
#[cfg(any(feature = "daemon", test))]
pub(crate) fn encode_result_frame(envelope: &TransportEnvelope) -> Result<Vec<u8>, ServiceError> {
fn response_frame(
operation: Operation,
envelope: &TransportEnvelope,
) -> Result<Vec<u8>, ServiceError> {
carriage_bytes(&CultNetMessage::OperationResponse {
message_id: envelope.request_id.clone(),
service_id: CONNECTOR_SERVICE_ID.to_string(),
operation: INVOKE_OPERATION.to_string(),
operation: operation.name().to_string(),
status: OK_STATUS.to_string(),
payload_schema: ENVELOPE_SCHEMA_ID.to_string(),
payload_encoding: PAYLOAD_ENCODING.to_string(),
@ -104,20 +135,43 @@ pub(crate) fn encode_result_frame(envelope: &TransportEnvelope) -> Result<Vec<u8
})
}
pub(crate) fn encode_invocation_frame(
envelope: &TransportEnvelope,
) -> Result<Vec<u8>, ServiceError> {
request_frame(Operation::Invoke, envelope)
}
pub(crate) fn encode_spend_readout_request_frame(
envelope: &TransportEnvelope,
) -> Result<Vec<u8>, ServiceError> {
request_frame(Operation::SpendReadout, envelope)
}
#[cfg(any(feature = "daemon", test))]
pub(crate) fn encode_result_frame(envelope: &TransportEnvelope) -> Result<Vec<u8>, ServiceError> {
response_frame(Operation::Invoke, envelope)
}
#[cfg(any(feature = "daemon", test))]
pub(crate) fn encode_spend_readout_response_frame(
envelope: &TransportEnvelope,
) -> Result<Vec<u8>, ServiceError> {
response_frame(Operation::SpendReadout, envelope)
}
/// Opens the carriage of one frame. Everything the carriage says is checked
/// here, before any decryption: the service, the operation, the payload schema,
/// and that the message id is the sealed envelope's request id.
/// here, before any decryption: the service, that the operation is one this
/// service has, the payload schema, and that the message id is the sealed
/// envelope's request id.
fn open_carriage(
message_id: &str,
service_id: &str,
operation: &str,
payload_schema: &str,
payload: &str,
) -> Result<TransportEnvelope, ServiceError> {
if service_id != CONNECTOR_SERVICE_ID
|| operation != INVOKE_OPERATION
|| payload_schema != ENVELOPE_SCHEMA_ID
{
) -> Result<(Operation, TransportEnvelope), ServiceError> {
let operation = Operation::from_name(operation).ok_or(ServiceError::Carriage)?;
if service_id != CONNECTOR_SERVICE_ID || payload_schema != ENVELOPE_SCHEMA_ID {
return Err(ServiceError::Carriage);
}
let bytes = STANDARD
@ -127,7 +181,7 @@ fn open_carriage(
if envelope.request_id != message_id {
return Err(ServiceError::OuterIdentity);
}
Ok(envelope)
Ok((operation, envelope))
}
fn read_carriage(bytes: &[u8]) -> Result<CultNetMessage, ServiceError> {
@ -135,8 +189,15 @@ fn read_carriage(bytes: &[u8]) -> Result<CultNetMessage, ServiceError> {
.map_err(|_| ServiceError::Carriage)
}
/// A request frame, by the operation it names.
#[cfg(any(feature = "daemon", test))]
pub(crate) fn decode_invocation_frame(bytes: &[u8]) -> Result<TransportEnvelope, ServiceError> {
pub(crate) enum RequestFrame {
Invoke(TransportEnvelope),
SpendReadout(TransportEnvelope),
}
#[cfg(any(feature = "daemon", test))]
pub(crate) fn decode_request_frame(bytes: &[u8]) -> Result<RequestFrame, ServiceError> {
match read_carriage(bytes)? {
CultNetMessage::OperationRequest {
message_id,
@ -145,18 +206,34 @@ pub(crate) fn decode_invocation_frame(bytes: &[u8]) -> Result<TransportEnvelope,
payload_schema,
payload,
..
} => open_carriage(
&message_id,
&service_id,
&operation,
&payload_schema,
&payload,
),
} => {
match open_carriage(
&message_id,
&service_id,
&operation,
&payload_schema,
&payload,
)? {
(Operation::Invoke, envelope) => Ok(RequestFrame::Invoke(envelope)),
(Operation::SpendReadout, envelope) => Ok(RequestFrame::SpendReadout(envelope)),
}
}
_ => Err(ServiceError::Carriage),
}
}
pub(crate) fn decode_result_frame(bytes: &[u8]) -> Result<TransportEnvelope, ServiceError> {
#[cfg(test)]
pub(crate) fn decode_invocation_frame(bytes: &[u8]) -> Result<TransportEnvelope, ServiceError> {
match decode_request_frame(bytes)? {
RequestFrame::Invoke(envelope) => Ok(envelope),
RequestFrame::SpendReadout(_) => Err(ServiceError::Carriage),
}
}
fn decode_response_frame(
bytes: &[u8],
expected: Operation,
) -> Result<TransportEnvelope, ServiceError> {
match read_carriage(bytes)? {
CultNetMessage::OperationResponse {
message_id,
@ -166,17 +243,32 @@ pub(crate) fn decode_result_frame(bytes: &[u8]) -> Result<TransportEnvelope, Ser
payload_schema,
payload,
..
} if status == OK_STATUS => open_carriage(
&message_id,
&service_id,
&operation,
&payload_schema,
&payload,
),
} if status == OK_STATUS => {
match open_carriage(
&message_id,
&service_id,
&operation,
&payload_schema,
&payload,
)? {
(found, envelope) if found == expected => Ok(envelope),
_ => Err(ServiceError::Carriage),
}
}
_ => Err(ServiceError::Carriage),
}
}
pub(crate) fn decode_result_frame(bytes: &[u8]) -> Result<TransportEnvelope, ServiceError> {
decode_response_frame(bytes, Operation::Invoke)
}
pub(crate) fn decode_spend_readout_response_frame(
bytes: &[u8],
) -> Result<TransportEnvelope, ServiceError> {
decode_response_frame(bytes, Operation::SpendReadout)
}
/// The most bytes one frame holding a sealed envelope of `ciphertext_len`
/// ciphertext bytes can occupy, with ids at their longest: the frame a caller
/// payload bound (`CallerAdmission`'s `max_payload_bytes`, a bound on
@ -324,6 +416,59 @@ pub fn decrypt_result(
Ok(result)
}
/// A spend readout request carries no content: the sealed envelope's
/// authenticated identity and kind are the whole request.
pub(crate) fn encrypt_spend_readout_request(
caller_runtime_id: &str,
request_id: &str,
security: &TransportKey,
) -> Result<TransportEnvelope, ServiceError> {
encrypt_message(
caller_runtime_id,
request_id,
EnvelopeKind::SpendReadoutRequest,
&[],
security,
)
}
#[cfg(any(feature = "daemon", test))]
pub(crate) fn decrypt_spend_readout_request(
envelope: &TransportEnvelope,
security: &TransportKey,
) -> Result<(), ServiceError> {
validate_envelope(envelope, EnvelopeKind::SpendReadoutRequest)?;
decrypt_message(envelope, security).map(drop)
}
#[cfg(any(feature = "daemon", test))]
pub(crate) fn encrypt_spend_readout(
caller_runtime_id: &str,
request_id: &str,
readout: &SpendReadout,
security: &TransportKey,
) -> Result<TransportEnvelope, ServiceError> {
let plaintext = wire_encode(readout).map_err(|_| ServiceError::Encoding)?;
encrypt_message(
caller_runtime_id,
request_id,
EnvelopeKind::SpendReadoutResult,
&plaintext,
security,
)
}
pub(crate) fn decrypt_spend_readout(
envelope: &TransportEnvelope,
security: &TransportKey,
) -> Result<SpendReadout, ServiceError> {
validate_envelope(envelope, EnvelopeKind::SpendReadoutResult)?;
let plaintext = decrypt_message(envelope, security)?;
let readout: SpendReadout = wire_decode(&plaintext).map_err(|_| ServiceError::Encoding)?;
readout.validate()?;
Ok(readout)
}
#[derive(Debug)]
pub(crate) enum TransportFrameError {
Connection(std::io::Error),
@ -1138,4 +1283,28 @@ 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();
assert!(matches!(
decode_request_frame(&frame),
Ok(RequestFrame::SpendReadout(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_eq!(
decrypt_spend_readout(&sealed, &key).unwrap_err(),
ServiceError::Envelope
);
let mut relabelled = sealed.clone();
relabelled.message_kind = EnvelopeKind::SpendReadoutResult;
assert_eq!(
decrypt_spend_readout(&relabelled, &key).unwrap_err(),
ServiceError::Encryption
);
}
}

2
crates/model-connector/src/lib.rs
View file

@ -1,6 +1,7 @@
mod client;
mod contract;
mod envelope;
mod readout;
#[cfg(test)]
mod fixtures;
@ -13,6 +14,7 @@ mod service;
pub use client::*;
pub use contract::*;
pub use envelope::*;
pub use readout::*;
#[cfg(feature = "daemon")]
pub use serve::*;

136
crates/model-connector/src/readout.rs Normal file
View file

@ -0,0 +1,136 @@
//! 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.
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use crate::contract::ContractError;
pub const SPEND_READOUT_SCHEMA_ID: &str = "gamecult.model.spend_readout.v1";
/// One window of the subscription plan, as the provider reported it.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct PlanQuota {
pub window: String,
/// How much of the window is used, in thousandths.
pub utilization_permille: u32,
pub resets_at_unix_ms: u64,
pub status: String,
}
/// What a backend last saw of the plan's windows and when it saw them. A
/// backend with no plan to report returns the default.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct PlanQuotaSnapshot {
pub windows: Vec<PlanQuota>,
pub observed_at_unix_ms: Option<u64>,
}
/// 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 schema_id: String,
pub at_unix_ms: u64,
pub callers: Vec<CallerSpend>,
pub plan_quota: Vec<PlanQuota>,
pub plan_quota_observed_at_unix_ms: Option<u64>,
}
impl SpendReadout {
pub fn new(at_unix_ms: u64, callers: Vec<CallerSpend>, plan: PlanQuotaSnapshot) -> Self {
Self {
schema_id: SPEND_READOUT_SCHEMA_ID.to_string(),
at_unix_ms,
callers,
plan_quota: plan.windows,
plan_quota_observed_at_unix_ms: plan.observed_at_unix_ms,
}
}
pub(crate) fn validate(&self) -> Result<(), ContractError> {
if self.schema_id != SPEND_READOUT_SCHEMA_ID {
return Err(ContractError::Schema);
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::contract::{wire_decode, wire_encode};
fn readout() -> SpendReadout {
SpendReadout::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(),
utilization_permille: 80,
resets_at_unix_ms: 1_793_307_000_000,
status: "allowed".into(),
}],
observed_at_unix_ms: Some(999),
},
)
}
#[test]
fn a_readout_is_a_named_map_that_decodes_to_itself() {
let bytes = wire_encode(&readout()).unwrap();
let value: rmpv::Value = rmp_serde::from_slice(&bytes).unwrap();
let rmpv::Value::Map(fields) = value else {
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",
"plan_quota_observed_at_unix_ms"
]
);
assert_eq!(wire_decode::<SpendReadout>(&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();
assert!(decoded.plan_quota.is_empty());
assert_eq!(decoded.plan_quota_observed_at_unix_ms, None);
}
#[test]
fn a_readout_of_another_schema_is_refused() {
let mut other = readout();
other.schema_id = "gamecult.model.spend_readout.v2".into();
assert!(other.validate().is_err());
assert!(readout().validate().is_ok());
}
}

114
crates/model-connector/src/serve.rs
View file

@ -12,9 +12,11 @@ use zeroize::Zeroizing;
use crate::contract::TransportResult;
use crate::envelope::{
ServiceError, TransportFrameError, decode_invocation_frame, encode_result_frame,
frame_bytes_for_payload, read_transport_frame, write_transport_frame,
RequestFrame, ServiceError, TransportFrameError, decode_request_frame, encode_result_frame,
encode_spend_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,
@ -32,6 +34,11 @@ 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.
fn plan_quota(&self) -> PlanQuotaSnapshot {
PlanQuotaSnapshot::default()
}
/// Executes one admitted call. The service has already bound it to its
/// caller and recorded its replay claim. The call cannot show a foreign
/// reasoning item and builds the result's receipt itself.
@ -232,17 +239,28 @@ fn serve_connection(
.and_then(|()| stream.set_write_timeout(Some(timeout)))
.map_err(ServeError::Connection)?;
let request = read_transport_frame(&mut stream, limits.max_frame_bytes).map_err(frame_error)?;
let response = execute_inference_frame(&request, service, backend, limits)?;
let response = answer_frame(&request, service, backend, limits)?;
write_transport_frame(&mut stream, &response, limits.max_frame_bytes).map_err(frame_error)
}
fn execute_inference_frame(
fn answer_frame(
request: &[u8],
service: &Mutex<TransportService>,
backend: &dyn ProviderBackend,
limits: &ServeLimits,
) -> Result<Vec<u8>, ServeError> {
let envelope = decode_invocation_frame(request).map_err(|_| ServeError::FrameEncoding)?;
let envelope = match decode_request_frame(request).map_err(|_| ServeError::FrameEncoding)? {
RequestFrame::Invoke(envelope) => envelope,
RequestFrame::SpendReadout(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)
.map_err(|_| ServeError::FrameEncoding);
}
};
let admission = service
.lock()
.map_err(|_| ServeError::ServicePoisoned)?
@ -495,6 +513,92 @@ mod tests {
.collect()
}
/// A counting backend that reports a plan, as a provider with one would.
struct PlanBackend {
inner: Arc<CountingBackend>,
plan: PlanQuotaSnapshot,
}
impl ProviderBackend for PlanBackend {
fn ready(&self) -> Result<(), &'static str> {
self.inner.ready()
}
fn provider(&self) -> &str {
self.inner.provider()
}
fn plan_quota(&self) -> PlanQuotaSnapshot {
self.plan.clone()
}
fn execute(&self, call: &ProviderCall) -> TransportResult {
self.inner.execute(call)
}
}
#[test]
fn serve_reads_out_spend_and_plan_quota_without_executing_anything() {
let store = TestStore::new();
let inner = CountingBackend::new(true);
let plan = PlanQuotaSnapshot {
windows: vec![crate::readout::PlanQuota {
window: "five_hour".into(),
utilization_permille: 80,
resets_at_unix_ms: 1_793_307_000_000,
status: "allowed".into(),
}],
observed_at_unix_ms: Some(1_793_300_000_000),
};
let backend = Arc::new(PlanBackend {
inner: inner.clone(),
plan: plan.clone(),
});
let (client, endpoint) = start_with(backend, &store, vec![caller(CALLER, KEY, 4)]);
client.execute(&fresh("request-1")).unwrap();
assert_eq!(inner.executions.load(Ordering::SeqCst), 1);
let before = now_ms();
let readout = client.spend_readout(CALLER).unwrap();
assert!(readout.at_unix_ms >= before && readout.at_unix_ms <= now_ms());
assert_eq!(readout.plan_quota, plan.windows);
assert_eq!(
readout.plan_quota_observed_at_unix_ms,
plan.observed_at_unix_ms
);
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,
"a readout runs nothing"
);
// A client without the caller's key is dropped, not answered.
let stranger = ConnectorClient::new(
endpoint,
"some-other-key",
FRAME,
Some(Duration::from_secs(5)),
)
.unwrap();
assert!(stranger.spend_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();
assert!(readout.plan_quota.is_empty());
assert_eq!(readout.plan_quota_observed_at_unix_ms, None);
}
#[test]
fn serve_refuses_a_request_over_its_max_cost_without_executing_it() {
let store = TestStore::new();

241
crates/model-connector/src/service.rs
View file

@ -14,8 +14,10 @@ use crate::contract::{
};
use crate::envelope::{
EnvelopeKind, MAX_IDENTITY_BYTES, ServiceError, TransportEnvelope, decrypt_invocation,
encode_result_frame, encrypt_result, validate_envelope,
decrypt_spend_readout_request, encode_result_frame, encode_spend_readout_response_frame,
encrypt_result, encrypt_spend_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:
@ -280,15 +282,15 @@ fn admission_clock(
/// it was admitted in, so nothing older is kept.
fn prune_windows(
window_spend: &mut HashMap<String, BTreeMap<u64, u64>>,
live: &HashMap<RequestIdentity, u64>,
live: &HashMap<RequestIdentity, LiveRequest>,
caller_runtime_id: &str,
floor: u64,
) {
if let Some(windows) = window_spend.get_mut(caller_runtime_id) {
windows.retain(|window, _| {
*window >= floor
|| live.iter().any(|(identity, live_window)| {
live_window == window && identity.caller_runtime_id == caller_runtime_id
|| live.iter().any(|(identity, request)| {
request.window == *window && identity.caller_runtime_id == caller_runtime_id
})
});
}
@ -298,7 +300,7 @@ fn prune_windows(
/// places that do not know the clock (open, settle).
fn prune_to_latest(
window_spend: &mut HashMap<String, BTreeMap<u64, u64>>,
live: &HashMap<RequestIdentity, u64>,
live: &HashMap<RequestIdentity, LiveRequest>,
caller_runtime_id: &str,
) {
let latest = window_spend
@ -309,6 +311,14 @@ fn prune_to_latest(
}
}
/// A request between claim and settlement: the window it was admitted in and
/// the reservation it holds there.
#[derive(Debug, Clone, Copy)]
struct LiveRequest {
window: u64,
reserved_micro_usd: u64,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct RequestIdentity {
caller_runtime_id: String,
@ -464,7 +474,7 @@ pub struct TransportService {
callers: HashMap<String, CallerAdmission>,
/// The requests between claim and settlement, each with the window it was
/// admitted in.
live: HashMap<RequestIdentity, u64>,
live: HashMap<RequestIdentity, LiveRequest>,
replay_store: CultCache,
max_expiry_skew_ms: u64,
provider: String,
@ -611,6 +621,70 @@ impl TransportService {
.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(
&self,
envelope: &TransportEnvelope,
now_unix_ms: u64,
plan: PlanQuotaSnapshot,
max_frame_bytes: usize,
) -> Result<TransportEnvelope, ServiceError> {
validate_envelope(envelope, EnvelopeKind::SpendReadoutRequest)?;
let asker = self
.callers
.get(&envelope.caller_runtime_id)
.ok_or(ServiceError::CallerNotAdmitted)?;
decrypt_spend_readout_request(envelope, &asker.security)?;
let mut callers: Vec<CallerSpend> = 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(
&envelope.caller_runtime_id,
&envelope.request_id,
&readout,
&asker.security,
)?;
if encode_spend_readout_response_frame(&sealed)?.len() > max_frame_bytes {
return Err(ServiceError::FrameBound);
}
Ok(sealed)
}
/// Admits one request. `max_frame_bytes` is the frame every reply must
/// fit, the same bound `complete` is given: a recorded result that no
/// longer fits it (the service reopened at a smaller frame) is answered
@ -794,7 +868,13 @@ impl TransportService {
0,
reserved_micro_usd,
);
self.live.insert(identity, window);
self.live.insert(
identity,
LiveRequest {
window,
reserved_micro_usd,
},
);
Ok(TransportAdmission::Execute(Box::new(ExecutionClaim {
call: ProviderCall {
receipt: TransportReceipt {
@ -3616,4 +3696,151 @@ mod tests {
<= sized(Refusal::ProviderDigestSubstitution)
);
}
// ---- spend readout ----
fn readout_as(
service: &TransportService,
caller: &str,
key: &str,
now: u64,
frame: usize,
) -> Result<SpendReadout, ServiceError> {
let sealed = encrypt_spend_readout_request(caller, "readout-1", &security(key)).unwrap();
let answer = service.spend_readout(&sealed, now, PlanQuotaSnapshot::default(), 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")
}
#[test]
fn spend_readout_reports_window() {
let store = TestStore::new();
let mut service = open_service(&store, 4);
let claim = claim_at(&mut service, "request-1", 1_000);
// In flight: the reservation is held, nothing is spent.
let readout = readout_as(&service, CALLER, KEY, 1_500, 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.
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 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);
}
#[test]
fn a_readout_is_refused_to_a_stranger_a_wrong_key_and_an_invocation() {
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()
};
assert_eq!(
service
.spend_readout(
&sealed("nobody", KEY),
1,
PlanQuotaSnapshot::default(),
FRAME
)
.unwrap_err(),
ServiceError::CallerNotAdmitted
);
assert_eq!(
service
.spend_readout(
&sealed(CALLER, "another-key"),
1,
PlanQuotaSnapshot::default(),
FRAME
)
.unwrap_err(),
ServiceError::Encryption
);
let invocation = invocation_at("request-1", 1_000, MAX_COST);
assert_eq!(
service
.spend_readout(
&raw_envelope(&invocation),
1,
PlanQuotaSnapshot::default(),
FRAME
)
.unwrap_err(),
ServiceError::Envelope
);
}
#[test]
fn a_readout_that_cannot_fit_the_frame_is_a_frame_bound_failure() {
let store = TestStore::new();
let service = open_service(&store, 4);
assert_eq!(
readout_as(&service, CALLER, KEY, 1_000, 64).unwrap_err(),
ServiceError::FrameBound
);
assert!(readout_as(&service, CALLER, KEY, 1_000, FRAME).is_ok());
}
}

3
crates/model-connector/tests/published_schemas.rs
View file

@ -28,6 +28,7 @@ fn published_schemas_are_the_types() {
("transport_result", render::<TransportResult>()),
("transport_receipt", render::<TransportReceipt>()),
("transport_envelope", render::<TransportEnvelope>()),
("spend_readout", render::<SpendReadout>()),
];
let dir = schemas_dir();
if std::env::var_os(REGENERATE).is_some() {
@ -53,7 +54,7 @@ fn published_schemas_are_the_types() {
expected.sort();
assert_eq!(
on_disk, expected,
"schemas/ holds exactly the five documents"
"schemas/ holds exactly the six documents"
);
}

4
schemas/transport_envelope.json
View file

@ -46,7 +46,9 @@
"type": "string",
"enum": [
"Invocation",
"Result"
"Result",
"SpendReadoutRequest",
"SpendReadoutResult"
]
}
}