claude-self-run-lease: Self runs borrow the one Claude token over a unix socket bound to the peer's uid, its pidfd and its unit's InvocationID; the connector sets the reservation, settle bills reported per-model usage uncapped, no report bills the reservation (squash of eureka-body/claude-self-run-lease 480cf94..b5f26ef; Soul s1-s3)
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
parent
480cf941c8
commit
d1f25b5ae6
15 changed files with 3421 additions and 28 deletions
|
|
@ -6,17 +6,22 @@
|
||||||
use std::collections::{BTreeMap, BTreeSet};
|
use std::collections::{BTreeMap, BTreeSet};
|
||||||
use std::net::{SocketAddr, TcpListener};
|
use std::net::{SocketAddr, TcpListener};
|
||||||
use std::path::{Path, PathBuf};
|
use std::path::{Path, PathBuf};
|
||||||
use std::sync::Arc;
|
use std::sync::{Arc, Mutex};
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
||||||
use cultcache_rs::{CultCache, DatabaseEntry, SingleFileMessagePackBackingStore};
|
use cultcache_rs::{CultCache, DatabaseEntry, SingleFileMessagePackBackingStore};
|
||||||
use model_connector::{
|
use model_connector::{
|
||||||
CallerAdmission, CallerConfig, MIN_PAYLOAD_BYTES, ModelPrice, ServeError, ServeLimits,
|
CallerAdmission, CallerConfig, MIN_PAYLOAD_BYTES, ModelPrice, ProviderBackend, ServeError,
|
||||||
ServiceError, load_admissions, serve,
|
ServeLimits, ServiceError, TransportService, load_admissions, serve_service,
|
||||||
};
|
};
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
use thiserror::Error;
|
use thiserror::Error;
|
||||||
|
|
||||||
use crate::backend::readiness;
|
use crate::backend::readiness;
|
||||||
|
#[cfg(target_os = "linux")]
|
||||||
|
use crate::lease::{LeaseCaller, LeaseLimits, serve_leases, user_uid};
|
||||||
|
#[cfg(target_os = "linux")]
|
||||||
|
use crate::unit::SystemdUnits;
|
||||||
use crate::{
|
use crate::{
|
||||||
ClaudeCodeBackend, ClaudeCodeConfig, ClaudeCodeLimits, CredentialTokenSource, StderrSink,
|
ClaudeCodeBackend, ClaudeCodeConfig, ClaudeCodeLimits, CredentialTokenSource, StderrSink,
|
||||||
};
|
};
|
||||||
|
|
@ -24,6 +29,18 @@ use crate::{
|
||||||
const CONFIG_KEY: &str = "runtime";
|
const CONFIG_KEY: &str = "runtime";
|
||||||
const MODEL_PREFIX: &str = "claude-";
|
const MODEL_PREFIX: &str = "claude-";
|
||||||
|
|
||||||
|
/// A unix account that may lease the subscription token for a Self run, the
|
||||||
|
/// Self unit (`mind-self@<unit_instance>.service`) its process must run in,
|
||||||
|
/// the caller whose spend window its leases reserve in, and what each lease
|
||||||
|
/// reserves there.
|
||||||
|
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||||
|
pub struct LeaseCallerConfig {
|
||||||
|
pub unix_user: String,
|
||||||
|
pub caller_id: String,
|
||||||
|
pub unit_instance: String,
|
||||||
|
pub reservation_micro_usd: u64,
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Clone, Debug, PartialEq, Eq, DatabaseEntry)]
|
#[derive(Clone, Debug, PartialEq, Eq, DatabaseEntry)]
|
||||||
#[cultcache(type = "gamecult.model_connector.claude_connector_config.v1")]
|
#[cultcache(type = "gamecult.model_connector.claude_connector_config.v1")]
|
||||||
pub struct ClaudeConnectorConfig {
|
pub struct ClaudeConnectorConfig {
|
||||||
|
|
@ -72,6 +89,17 @@ pub struct ClaudeConnectorConfig {
|
||||||
/// ids exactly as the CLI reports them in its init event.
|
/// ids exactly as the CLI reports them in its init event.
|
||||||
#[cultcache(key = 18)]
|
#[cultcache(key = 18)]
|
||||||
pub prices: BTreeMap<String, ModelPrice>,
|
pub prices: BTreeMap<String, ModelPrice>,
|
||||||
|
/// Where the Self-run lease socket is bound (a systemd RuntimeDirectory),
|
||||||
|
/// or none when this connector lends no token.
|
||||||
|
#[cultcache(key = 19)]
|
||||||
|
pub lease_socket: Option<PathBuf>,
|
||||||
|
/// The accounts allowed to lease, each bound to a configured caller.
|
||||||
|
/// Empty exactly when `lease_socket` is none.
|
||||||
|
#[cultcache(key = 20)]
|
||||||
|
pub lease_callers: Vec<LeaseCallerConfig>,
|
||||||
|
/// The furthest a lease's expiry may lie from now.
|
||||||
|
#[cultcache(key = 21)]
|
||||||
|
pub lease_max_ms: u64,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// A refusal. It names the field or stage that failed and nothing the
|
/// A refusal. It names the field or stage that failed and nothing the
|
||||||
|
|
@ -205,9 +233,39 @@ impl ClaudeConnectorConfig {
|
||||||
{
|
{
|
||||||
return Err(ConfigError::Invalid("prices"));
|
return Err(ConfigError::Invalid("prices"));
|
||||||
}
|
}
|
||||||
|
self.validate_leases(&caller.caller_runtime_id, caller.spend_ceiling_micro_usd)?;
|
||||||
Ok(listen)
|
Ok(listen)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn validate_leases(&self, caller_id: &str, spend_ceiling: u64) -> Result<(), ConfigError> {
|
||||||
|
let Some(socket) = &self.lease_socket else {
|
||||||
|
return if self.lease_callers.is_empty() {
|
||||||
|
Ok(())
|
||||||
|
} else {
|
||||||
|
Err(ConfigError::Invalid("lease_callers"))
|
||||||
|
};
|
||||||
|
};
|
||||||
|
if !socket.is_absolute() {
|
||||||
|
return Err(ConfigError::Invalid("lease_socket"));
|
||||||
|
}
|
||||||
|
if self.lease_max_ms == 0 {
|
||||||
|
return Err(ConfigError::Invalid("lease_max_ms"));
|
||||||
|
}
|
||||||
|
let mut users = BTreeSet::new();
|
||||||
|
if self.lease_callers.is_empty()
|
||||||
|
|| !self.lease_callers.iter().all(|lease| {
|
||||||
|
is_trimmed_nonempty(&lease.unix_user)
|
||||||
|
&& lease.caller_id == caller_id
|
||||||
|
&& is_unit_instance(&lease.unit_instance)
|
||||||
|
&& (1..=spend_ceiling).contains(&lease.reservation_micro_usd)
|
||||||
|
&& users.insert(lease.unix_user.as_str())
|
||||||
|
})
|
||||||
|
{
|
||||||
|
return Err(ConfigError::Invalid("lease_callers"));
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
fn backend_config(&self) -> ClaudeCodeConfig {
|
fn backend_config(&self) -> ClaudeCodeConfig {
|
||||||
ClaudeCodeConfig {
|
ClaudeCodeConfig {
|
||||||
claude_binary: self.claude_binary.clone(),
|
claude_binary: self.claude_binary.clone(),
|
||||||
|
|
@ -236,6 +294,14 @@ impl ClaudeConnectorConfig {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// The `<instance>` of a unit name, as plain as a systemd instance can be.
|
||||||
|
fn is_unit_instance(text: &str) -> bool {
|
||||||
|
!text.is_empty()
|
||||||
|
&& text
|
||||||
|
.bytes()
|
||||||
|
.all(|b| b.is_ascii_alphanumeric() || matches!(b, b'-' | b'_'))
|
||||||
|
}
|
||||||
|
|
||||||
fn is_trimmed_nonempty(text: &str) -> bool {
|
fn is_trimmed_nonempty(text: &str) -> bool {
|
||||||
!text.is_empty() && text.trim() == text
|
!text.is_empty() && text.trim() == text
|
||||||
}
|
}
|
||||||
|
|
@ -394,9 +460,85 @@ fn readiness_of(config: &ClaudeConnectorConfig) -> Result<Vec<CallerAdmission>,
|
||||||
&CredentialTokenSource::from_environment(backend.token_credential.clone()),
|
&CredentialTokenSource::from_environment(backend.token_credential.clone()),
|
||||||
)
|
)
|
||||||
.map_err(ConfigError::Unusable)?;
|
.map_err(ConfigError::Unusable)?;
|
||||||
|
if config.lease_socket.is_some() {
|
||||||
|
lease_callers_of(config)?;
|
||||||
|
}
|
||||||
Ok(admissions)
|
Ok(admissions)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// The accounts allowed to lease, resolved to the uids the kernel will report
|
||||||
|
/// for their connections. An account the system does not have is not ready.
|
||||||
|
#[cfg(target_os = "linux")]
|
||||||
|
fn lease_callers_of(config: &ClaudeConnectorConfig) -> Result<Vec<LeaseCaller>, ConfigError> {
|
||||||
|
config
|
||||||
|
.lease_callers
|
||||||
|
.iter()
|
||||||
|
.map(|lease| {
|
||||||
|
user_uid(&lease.unix_user)
|
||||||
|
.map(|uid| LeaseCaller {
|
||||||
|
uid,
|
||||||
|
caller_id: lease.caller_id.clone(),
|
||||||
|
unit_instance: lease.unit_instance.clone(),
|
||||||
|
reservation_micro_usd: lease.reservation_micro_usd,
|
||||||
|
})
|
||||||
|
.ok_or(ConfigError::Unusable("lease_callers"))
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Peer credentials identify a lease caller only on Linux.
|
||||||
|
#[cfg(not(target_os = "linux"))]
|
||||||
|
fn lease_callers_of(_: &ClaudeConnectorConfig) -> Result<(), ConfigError> {
|
||||||
|
Err(ConfigError::Unusable("lease_socket"))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Binds the lease socket (mode 0660; the unit's RuntimeDirectory and group
|
||||||
|
/// decide who can reach it, and the kernel's peer uid decides who may lease)
|
||||||
|
/// and serves it beside the RUDP loop. The lease listener is half of the
|
||||||
|
/// connector: if it stops, so does the process.
|
||||||
|
#[cfg(target_os = "linux")]
|
||||||
|
fn start_leases(
|
||||||
|
config: &ClaudeConnectorConfig,
|
||||||
|
socket: &Path,
|
||||||
|
service: Arc<Mutex<TransportService>>,
|
||||||
|
) -> Result<(), ConfigError> {
|
||||||
|
use std::os::unix::fs::{FileTypeExt, PermissionsExt};
|
||||||
|
use std::os::unix::net::UnixListener;
|
||||||
|
|
||||||
|
let callers = lease_callers_of(config)?;
|
||||||
|
// A socket a previous run left is replaced; anything else at the path is not ours.
|
||||||
|
if std::fs::symlink_metadata(socket).is_ok_and(|found| found.file_type().is_socket()) {
|
||||||
|
let _ = std::fs::remove_file(socket);
|
||||||
|
}
|
||||||
|
let listener = UnixListener::bind(socket).map_err(|_| ConfigError::Listen)?;
|
||||||
|
std::fs::set_permissions(socket, std::fs::Permissions::from_mode(0o660))
|
||||||
|
.map_err(|_| ConfigError::Listen)?;
|
||||||
|
let limits = LeaseLimits {
|
||||||
|
socket_timeout: Duration::from_millis(config.socket_timeout_ms),
|
||||||
|
max_lease_ms: config.lease_max_ms,
|
||||||
|
max_connections: config.max_connections,
|
||||||
|
};
|
||||||
|
let tokens = Arc::new(CredentialTokenSource::from_environment(
|
||||||
|
config.token_credential.clone(),
|
||||||
|
));
|
||||||
|
std::thread::Builder::new()
|
||||||
|
.name("claude-lease-listener".to_string())
|
||||||
|
.spawn(move || {
|
||||||
|
let _ = serve_leases(
|
||||||
|
&listener,
|
||||||
|
callers,
|
||||||
|
limits,
|
||||||
|
service,
|
||||||
|
tokens,
|
||||||
|
Arc::new(SystemdUnits::system()),
|
||||||
|
);
|
||||||
|
eprintln!("claude-connector: the lease listener stopped");
|
||||||
|
std::process::exit(1);
|
||||||
|
})
|
||||||
|
.map_err(|_| ConfigError::Serving("lease listener"))?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
/// Whether `serve` would start with this document and this environment.
|
/// Whether `serve` would start with this document and this environment.
|
||||||
pub fn check_ready(config: &ClaudeConnectorConfig) -> Result<(), ConfigError> {
|
pub fn check_ready(config: &ClaudeConnectorConfig) -> Result<(), ConfigError> {
|
||||||
readiness_of(config).map(drop)
|
readiness_of(config).map(drop)
|
||||||
|
|
@ -413,14 +555,21 @@ pub fn serve_config(path: &Path) -> Result<(), ConfigError> {
|
||||||
Box::new(LogSink(std::sync::Mutex::new(std::io::stderr()))),
|
Box::new(LogSink(std::sync::Mutex::new(std::io::stderr()))),
|
||||||
));
|
));
|
||||||
let listener = TcpListener::bind(listen).map_err(|_| ConfigError::Listen)?;
|
let listener = TcpListener::bind(listen).map_err(|_| ConfigError::Listen)?;
|
||||||
serve(
|
let service = Arc::new(Mutex::new(
|
||||||
listener,
|
TransportService::open(
|
||||||
config.serve_limits(),
|
&config.replay_store,
|
||||||
admissions,
|
backend.provider(),
|
||||||
config.prices.clone(),
|
config.prices.clone(),
|
||||||
&config.replay_store,
|
admissions,
|
||||||
backend,
|
config.max_expiry_skew_ms,
|
||||||
)?;
|
)
|
||||||
|
.map_err(|error| ConfigError::from(ServeError::Service(error)))?,
|
||||||
|
));
|
||||||
|
#[cfg(target_os = "linux")]
|
||||||
|
if let Some(socket) = &config.lease_socket {
|
||||||
|
start_leases(&config, socket, service.clone())?;
|
||||||
|
}
|
||||||
|
serve_service(listener, config.serve_limits(), service, backend)?;
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
729
crates/claude-connector/src/lease.rs
Normal file
729
crates/claude-connector/src/lease.rs
Normal file
|
|
@ -0,0 +1,729 @@
|
||||||
|
//! The Self-run credential lease. The connector is the only holder of the
|
||||||
|
//! subscription token and the one owner of Self-run admission, reservation and
|
||||||
|
//! the spend ledger (rulings self-runs-through-connector,
|
||||||
|
//! self-run-credential-lease and self-run-lease-billing-and-binding). A Self
|
||||||
|
//! run asks over a unix socket for a lease. The connector decides who may ask
|
||||||
|
//! (the peer's uid, and the systemd unit invocation its process runs in),
|
||||||
|
//! reserves its own configured amount in the caller's spend window, and only
|
||||||
|
//! then hands back one pipe fd that holds the token. The run settles with the
|
||||||
|
//! per-model usage it reports, and the connector bills that at true cost,
|
||||||
|
//! uncapped; a run that vanishes or lets its lease expire is billed the whole
|
||||||
|
//! reservation. Nothing the run says about itself decides who it is, what it
|
||||||
|
//! reserves, or what the plan allows.
|
||||||
|
//!
|
||||||
|
//! The token exists in this process only as the value the token source returns
|
||||||
|
//! and in the pipe. It is not in an environment, an argument, a file, a log or
|
||||||
|
//! a reply frame. The connector spawns nothing here: the Self run's `claude`
|
||||||
|
//! stays in the Self run's own unit.
|
||||||
|
|
||||||
|
use std::ffi::CString;
|
||||||
|
use std::io::{self, Read, Write};
|
||||||
|
use std::mem::{size_of, zeroed};
|
||||||
|
use std::os::fd::{AsFd, AsRawFd, BorrowedFd, FromRawFd, OwnedFd, RawFd};
|
||||||
|
use std::os::unix::net::{UnixListener, UnixStream};
|
||||||
|
use std::ptr;
|
||||||
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||||
|
use std::sync::{Arc, Mutex};
|
||||||
|
use std::thread;
|
||||||
|
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
|
||||||
|
|
||||||
|
use model_connector::{
|
||||||
|
LeaseAdmission, LeaseGrant, LeaseRefusal, LeaseSettlement, LeaseUse, ModelUsage,
|
||||||
|
TransportService,
|
||||||
|
};
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
|
||||||
|
use crate::TokenSource;
|
||||||
|
|
||||||
|
/// The largest frame either side sends or accepts, in bytes. It is wire
|
||||||
|
/// contract: a Huginn client sizes against it, so it is pinned as a literal.
|
||||||
|
pub const MAX_LEASE_FRAME_BYTES: usize = 65_536;
|
||||||
|
const _: () = assert!(MAX_LEASE_FRAME_BYTES == 65_536);
|
||||||
|
|
||||||
|
/// One message of the lease protocol: a u32 big-endian length, then this as a
|
||||||
|
/// MessagePack named map `{kind, value}`.
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||||
|
#[serde(tag = "kind", content = "value", rename_all = "snake_case")]
|
||||||
|
pub enum LeaseMessage {
|
||||||
|
ClaudeLeaseRequest(LeaseRequest),
|
||||||
|
/// Travels with exactly one SCM_RIGHTS fd: the read end of a pipe that
|
||||||
|
/// holds the token, its write end closed.
|
||||||
|
ClaudeLeaseGrant(LeaseGranted),
|
||||||
|
ClaudeLeaseRefusal(LeaseRefused),
|
||||||
|
ClaudeLeaseSettle(LeaseSettle),
|
||||||
|
ClaudeLeaseSettled(LeaseSettled),
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A request names only when the lease ends. Who asks, and what is reserved,
|
||||||
|
/// are the connector's to know.
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||||
|
pub struct LeaseRequest {
|
||||||
|
pub expires_at_unix_ms: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||||
|
pub struct LeaseGranted {
|
||||||
|
pub lease_id: String,
|
||||||
|
pub reserved_micro_usd: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||||
|
pub enum LeaseRefusalKind {
|
||||||
|
SpendCeiling,
|
||||||
|
CostBound,
|
||||||
|
UnknownCaller,
|
||||||
|
Expired,
|
||||||
|
Busy,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||||
|
pub struct LeaseRefused {
|
||||||
|
pub kind: LeaseRefusalKind,
|
||||||
|
pub detail: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||||
|
pub struct LeaseSettle {
|
||||||
|
pub lease_id: String,
|
||||||
|
/// What the run used of each model, as the run reports it.
|
||||||
|
pub usage: Vec<ModelUsage>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||||
|
pub struct LeaseSettled {
|
||||||
|
pub billed_micro_usd: u64,
|
||||||
|
/// How far the bill passed the reservation; zero when it did not.
|
||||||
|
pub overrun_micro_usd: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A frame: the length, then the named-map MessagePack of `message`.
|
||||||
|
pub fn encode_frame(message: &LeaseMessage) -> Vec<u8> {
|
||||||
|
// Every field of the message is plain data, so encoding cannot fail.
|
||||||
|
let body = rmp_serde::to_vec_named(message).unwrap_or_default();
|
||||||
|
let mut frame = Vec::with_capacity(4 + body.len());
|
||||||
|
frame.extend_from_slice(&(body.len() as u32).to_be_bytes());
|
||||||
|
frame.extend_from_slice(&body);
|
||||||
|
frame
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The message in one whole frame, or `None` when the frame is not exactly one
|
||||||
|
/// well-formed message within the bound.
|
||||||
|
pub fn decode_frame(frame: &[u8]) -> Option<LeaseMessage> {
|
||||||
|
let (length, body) = frame.split_first_chunk::<4>()?;
|
||||||
|
let length = u32::from_be_bytes(*length) as usize;
|
||||||
|
if length == 0 || length > MAX_LEASE_FRAME_BYTES || body.len() != length {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
rmp_serde::from_slice(body).ok()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A reader whose waiting can be bounded, so a frame has one deadline across
|
||||||
|
/// all of its reads.
|
||||||
|
trait Bounded: Read {
|
||||||
|
fn within(&mut self, left: Duration) -> io::Result<()>;
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Bounded for UnixStream {
|
||||||
|
fn within(&mut self, left: Duration) -> io::Result<()> {
|
||||||
|
self.set_read_timeout(Some(left))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Bounded for &[u8] {
|
||||||
|
fn within(&mut self, _: Duration) -> io::Result<()> {
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Fills `buffer` or fails: when `deadline` passes, however the bytes trickle.
|
||||||
|
fn read_exact_by(
|
||||||
|
stream: &mut impl Bounded,
|
||||||
|
buffer: &mut [u8],
|
||||||
|
deadline: Instant,
|
||||||
|
) -> io::Result<()> {
|
||||||
|
let mut filled = 0;
|
||||||
|
while filled < buffer.len() {
|
||||||
|
let left = deadline
|
||||||
|
.checked_duration_since(Instant::now())
|
||||||
|
.filter(|left| !left.is_zero())
|
||||||
|
.ok_or_else(|| io::Error::from(io::ErrorKind::TimedOut))?;
|
||||||
|
stream.within(left)?;
|
||||||
|
match stream.read(&mut buffer[filled..]) {
|
||||||
|
Ok(0) => return Err(io::Error::from(io::ErrorKind::UnexpectedEof)),
|
||||||
|
Ok(read) => filled += read,
|
||||||
|
Err(error)
|
||||||
|
if matches!(
|
||||||
|
error.kind(),
|
||||||
|
io::ErrorKind::Interrupted | io::ErrorKind::WouldBlock
|
||||||
|
) => {}
|
||||||
|
Err(error) => return Err(error),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// One whole frame, complete by `deadline`.
|
||||||
|
fn read_message(stream: &mut impl Bounded, deadline: Instant) -> io::Result<LeaseMessage> {
|
||||||
|
let invalid = || io::Error::from(io::ErrorKind::InvalidData);
|
||||||
|
let mut length = [0_u8; 4];
|
||||||
|
read_exact_by(stream, &mut length, deadline)?;
|
||||||
|
let length = u32::from_be_bytes(length) as usize;
|
||||||
|
if length == 0 || length > MAX_LEASE_FRAME_BYTES {
|
||||||
|
return Err(invalid());
|
||||||
|
}
|
||||||
|
let mut body = vec![0_u8; length];
|
||||||
|
read_exact_by(stream, &mut body, deadline)?;
|
||||||
|
rmp_serde::from_slice(&body).map_err(|_| invalid())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The most a pipe takes in one atomic write.
|
||||||
|
const PIPE_BUF: usize = 4096;
|
||||||
|
|
||||||
|
/// A unix account allowed to lease, the Self unit its process must run in,
|
||||||
|
/// and the caller whose spend window its leases reserve in.
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||||
|
pub struct LeaseCaller {
|
||||||
|
pub uid: u32,
|
||||||
|
pub caller_id: String,
|
||||||
|
/// The `<instance>` of `mind-self@<instance>.service`.
|
||||||
|
pub unit_instance: String,
|
||||||
|
/// What each of this caller's leases reserves, whatever the run asks.
|
||||||
|
pub reservation_micro_usd: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl LeaseCaller {
|
||||||
|
/// The one unit a lease to this caller may be requested from.
|
||||||
|
pub fn unit(&self) -> String {
|
||||||
|
format!("mind-self@{}.service", self.unit_instance)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The systemd unit invocation a process runs in.
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||||
|
pub struct PeerUnit {
|
||||||
|
pub unit: String,
|
||||||
|
pub invocation_id: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Which invocation of `unit` a process is running, as the system records it.
|
||||||
|
/// The lease listener asks this port; `SystemdUnits` answers it in production.
|
||||||
|
/// `None` when the process is not in `unit` or the system cannot say. The
|
||||||
|
/// process is named by the handle the kernel gave for the peer, never by a pid.
|
||||||
|
pub trait PeerUnits: Send + Sync {
|
||||||
|
fn invocation_in(&self, process: BorrowedFd<'_>, unit: &str) -> Option<PeerUnit>;
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The bounds of the lease listener.
|
||||||
|
#[derive(Debug, Clone, Copy)]
|
||||||
|
pub struct LeaseLimits {
|
||||||
|
/// How long a peer has to send its request frame, whole.
|
||||||
|
pub socket_timeout: Duration,
|
||||||
|
/// The furthest an expiry may lie from now.
|
||||||
|
pub max_lease_ms: u64,
|
||||||
|
pub max_connections: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Who is on the other end of a unix stream, as the kernel recorded them at
|
||||||
|
/// connect.
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub struct PeerCredentials {
|
||||||
|
pub uid: u32,
|
||||||
|
/// A handle on the peer's process (SO_PEERPIDFD, Linux 6.5 and later). It
|
||||||
|
/// pins the process's identity: a recycled pid cannot be read through it.
|
||||||
|
/// `None` when the kernel cannot give one, and then no unit is vouched
|
||||||
|
/// for.
|
||||||
|
pub process: Option<OwnedFd>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The uid and process handle of the process on the other end of `stream`.
|
||||||
|
pub fn peer_credentials(stream: &UnixStream) -> io::Result<PeerCredentials> {
|
||||||
|
// SAFETY: `credentials` and `length` are valid for the call, and `length`
|
||||||
|
// is the size of the buffer getsockopt writes into.
|
||||||
|
let (result, credentials) = unsafe {
|
||||||
|
let mut credentials: libc::ucred = zeroed();
|
||||||
|
let mut length = size_of::<libc::ucred>() as libc::socklen_t;
|
||||||
|
let result = libc::getsockopt(
|
||||||
|
stream.as_raw_fd(),
|
||||||
|
libc::SOL_SOCKET,
|
||||||
|
libc::SO_PEERCRED,
|
||||||
|
ptr::from_mut(&mut credentials).cast(),
|
||||||
|
&mut length,
|
||||||
|
);
|
||||||
|
(result, credentials)
|
||||||
|
};
|
||||||
|
if result != 0 {
|
||||||
|
return Err(io::Error::last_os_error());
|
||||||
|
}
|
||||||
|
// SAFETY: as above, with a `c_int` the kernel fills with a new descriptor
|
||||||
|
// that this function then owns.
|
||||||
|
let (result, pidfd) = unsafe {
|
||||||
|
let mut pidfd: libc::c_int = -1;
|
||||||
|
let mut length = size_of::<libc::c_int>() as libc::socklen_t;
|
||||||
|
let result = libc::getsockopt(
|
||||||
|
stream.as_raw_fd(),
|
||||||
|
libc::SOL_SOCKET,
|
||||||
|
libc::SO_PEERPIDFD,
|
||||||
|
ptr::from_mut(&mut pidfd).cast(),
|
||||||
|
&mut length,
|
||||||
|
);
|
||||||
|
(result, pidfd)
|
||||||
|
};
|
||||||
|
Ok(PeerCredentials {
|
||||||
|
uid: credentials.uid,
|
||||||
|
// SAFETY: on success `pidfd` is a descriptor nothing else owns.
|
||||||
|
process: (result == 0 && pidfd >= 0).then(|| unsafe { OwnedFd::from_raw_fd(pidfd) }),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The uid of the account `name`, if the system has one.
|
||||||
|
pub fn user_uid(name: &str) -> Option<u32> {
|
||||||
|
let name = CString::new(name).ok()?;
|
||||||
|
let mut buffer = vec![0 as libc::c_char; 16 * 1024];
|
||||||
|
// SAFETY: every pointer is valid for the call and `buffer` is as long as
|
||||||
|
// the length passed; `found` is read only when the call reports success.
|
||||||
|
unsafe {
|
||||||
|
let mut entry: libc::passwd = zeroed();
|
||||||
|
let mut found: *mut libc::passwd = ptr::null_mut();
|
||||||
|
let result = libc::getpwnam_r(
|
||||||
|
name.as_ptr(),
|
||||||
|
&mut entry,
|
||||||
|
buffer.as_mut_ptr(),
|
||||||
|
buffer.len(),
|
||||||
|
&mut found,
|
||||||
|
);
|
||||||
|
(result == 0 && !found.is_null()).then_some(entry.pw_uid)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A pipe whose read end holds `token` and whose write end is closed, so a
|
||||||
|
/// reader gets the token and then end of file. The token is smaller than
|
||||||
|
/// PIPE_BUF, so the write is whole and cannot block.
|
||||||
|
fn token_pipe(token: &[u8]) -> io::Result<OwnedFd> {
|
||||||
|
if token.len() > PIPE_BUF {
|
||||||
|
return Err(io::Error::from(io::ErrorKind::InvalidInput));
|
||||||
|
}
|
||||||
|
let mut ends = [0 as RawFd; 2];
|
||||||
|
// SAFETY: `ends` has room for the two descriptors pipe2 writes.
|
||||||
|
if unsafe { libc::pipe2(ends.as_mut_ptr(), libc::O_CLOEXEC) } != 0 {
|
||||||
|
return Err(io::Error::last_os_error());
|
||||||
|
}
|
||||||
|
// SAFETY: pipe2 just returned these two descriptors and nothing else owns them.
|
||||||
|
let (read, write) = unsafe { (OwnedFd::from_raw_fd(ends[0]), OwnedFd::from_raw_fd(ends[1])) };
|
||||||
|
// SAFETY: the pointer and length describe `token`.
|
||||||
|
let written = unsafe { libc::write(write.as_raw_fd(), token.as_ptr().cast(), token.len()) };
|
||||||
|
if usize::try_from(written).ok() != Some(token.len()) {
|
||||||
|
return Err(io::Error::last_os_error());
|
||||||
|
}
|
||||||
|
drop(write);
|
||||||
|
Ok(read)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Writes `frame` to `stream` with `fd` attached as SCM_RIGHTS.
|
||||||
|
fn send_with_fd(stream: &UnixStream, frame: &[u8], fd: RawFd) -> io::Result<()> {
|
||||||
|
let fd_size = size_of::<RawFd>() as u32;
|
||||||
|
// SAFETY: CMSG_SPACE is a pure size computation.
|
||||||
|
let space = unsafe { libc::CMSG_SPACE(fd_size) } as usize;
|
||||||
|
// u64 elements align the control buffer for a cmsghdr.
|
||||||
|
let mut control = vec![0_u64; space.div_ceil(size_of::<u64>())];
|
||||||
|
let mut vector = libc::iovec {
|
||||||
|
iov_base: frame.as_ptr().cast_mut().cast(),
|
||||||
|
iov_len: frame.len(),
|
||||||
|
};
|
||||||
|
// SAFETY: msghdr is plain data for which all-zero is valid; the pointers
|
||||||
|
// set below outlive the sendmsg call, and the cmsg header written lies
|
||||||
|
// inside `control`, which CMSG_SPACE sized for one descriptor.
|
||||||
|
let sent = unsafe {
|
||||||
|
let mut message: libc::msghdr = zeroed();
|
||||||
|
message.msg_iov = &mut vector;
|
||||||
|
message.msg_iovlen = 1;
|
||||||
|
message.msg_control = control.as_mut_ptr().cast();
|
||||||
|
message.msg_controllen = space as _;
|
||||||
|
let header = libc::CMSG_FIRSTHDR(&message);
|
||||||
|
(*header).cmsg_level = libc::SOL_SOCKET;
|
||||||
|
(*header).cmsg_type = libc::SCM_RIGHTS;
|
||||||
|
(*header).cmsg_len = libc::CMSG_LEN(fd_size) as _;
|
||||||
|
ptr::copy_nonoverlapping(
|
||||||
|
ptr::from_ref(&fd).cast::<u8>(),
|
||||||
|
libc::CMSG_DATA(header),
|
||||||
|
size_of::<RawFd>(),
|
||||||
|
);
|
||||||
|
loop {
|
||||||
|
let sent = libc::sendmsg(stream.as_raw_fd(), &message, libc::MSG_NOSIGNAL);
|
||||||
|
if sent >= 0 || io::Error::last_os_error().kind() != io::ErrorKind::Interrupted {
|
||||||
|
break sent;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let Ok(sent) = usize::try_from(sent) else {
|
||||||
|
return Err(io::Error::last_os_error());
|
||||||
|
};
|
||||||
|
// The descriptor travels with the first byte; any rest is plain data.
|
||||||
|
let mut stream = stream;
|
||||||
|
stream.write_all(&frame[sent..])
|
||||||
|
}
|
||||||
|
|
||||||
|
fn unix_ms() -> u64 {
|
||||||
|
SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.map_or(0, |elapsed| {
|
||||||
|
u64::try_from(elapsed.as_millis()).unwrap_or(u64::MAX)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
struct Permit(Arc<AtomicUsize>);
|
||||||
|
|
||||||
|
impl Permit {
|
||||||
|
fn acquire(active: &Arc<AtomicUsize>, limit: usize) -> Option<Self> {
|
||||||
|
active
|
||||||
|
.fetch_update(Ordering::AcqRel, Ordering::Acquire, |held| {
|
||||||
|
(held < limit).then_some(held + 1)
|
||||||
|
})
|
||||||
|
.ok()
|
||||||
|
.map(|_| Self(active.clone()))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Drop for Permit {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
self.0.fetch_sub(1, Ordering::AcqRel);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
struct Lender {
|
||||||
|
callers: Vec<LeaseCaller>,
|
||||||
|
limits: LeaseLimits,
|
||||||
|
service: Arc<Mutex<TransportService>>,
|
||||||
|
tokens: Arc<dyn TokenSource>,
|
||||||
|
units: Arc<dyn PeerUnits>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Serves leases on `listener` until it fails. Each connection is one lease:
|
||||||
|
/// request, then grant or refusal, then (after a grant) settle. A caller is
|
||||||
|
/// the account whose uid the kernel reports for the connection, and only a
|
||||||
|
/// process in that caller's Self unit may lease, once per unit invocation; no
|
||||||
|
/// field of any message names either.
|
||||||
|
pub fn serve_leases(
|
||||||
|
listener: &UnixListener,
|
||||||
|
callers: Vec<LeaseCaller>,
|
||||||
|
limits: LeaseLimits,
|
||||||
|
service: Arc<Mutex<TransportService>>,
|
||||||
|
tokens: Arc<dyn TokenSource>,
|
||||||
|
units: Arc<dyn PeerUnits>,
|
||||||
|
) -> io::Result<()> {
|
||||||
|
let lender = Arc::new(Lender {
|
||||||
|
callers,
|
||||||
|
limits,
|
||||||
|
service,
|
||||||
|
tokens,
|
||||||
|
units,
|
||||||
|
});
|
||||||
|
let active = Arc::new(AtomicUsize::new(0));
|
||||||
|
loop {
|
||||||
|
let (stream, _) = listener.accept()?;
|
||||||
|
let Some(permit) = Permit::acquire(&active, limits.max_connections) else {
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
let lender = lender.clone();
|
||||||
|
thread::Builder::new()
|
||||||
|
.name("claude-lease".to_string())
|
||||||
|
.spawn(move || {
|
||||||
|
let _permit = permit;
|
||||||
|
// A connection that fails is dropped; the lease it may have
|
||||||
|
// held is settled inside `lend`.
|
||||||
|
let _ = lender.lend(stream);
|
||||||
|
})?;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Lender {
|
||||||
|
fn lend(&self, mut stream: UnixStream) -> io::Result<()> {
|
||||||
|
stream.set_write_timeout(Some(self.limits.socket_timeout))?;
|
||||||
|
let peer = peer_credentials(&stream)?;
|
||||||
|
let request_deadline = Instant::now() + self.limits.socket_timeout;
|
||||||
|
let LeaseMessage::ClaudeLeaseRequest(request) =
|
||||||
|
read_message(&mut stream, request_deadline)?
|
||||||
|
else {
|
||||||
|
return Err(io::Error::from(io::ErrorKind::InvalidData));
|
||||||
|
};
|
||||||
|
let Some(caller) = self.callers.iter().find(|caller| caller.uid == peer.uid) else {
|
||||||
|
return refuse(
|
||||||
|
&mut stream,
|
||||||
|
LeaseRefusalKind::UnknownCaller,
|
||||||
|
"this account may not lease",
|
||||||
|
);
|
||||||
|
};
|
||||||
|
// The unit and its invocation come from the system's record of the
|
||||||
|
// peer's process, never from the peer.
|
||||||
|
let Some(invocation) = peer
|
||||||
|
.process
|
||||||
|
.as_ref()
|
||||||
|
.and_then(|process| self.units.invocation_in(process.as_fd(), &caller.unit()))
|
||||||
|
.filter(|found| found.unit == caller.unit())
|
||||||
|
else {
|
||||||
|
return refuse(
|
||||||
|
&mut stream,
|
||||||
|
LeaseRefusalKind::UnknownCaller,
|
||||||
|
"this process is not in the Self unit",
|
||||||
|
);
|
||||||
|
};
|
||||||
|
let now = unix_ms();
|
||||||
|
if request.expires_at_unix_ms <= now
|
||||||
|
|| request.expires_at_unix_ms > now.saturating_add(self.limits.max_lease_ms)
|
||||||
|
{
|
||||||
|
return refuse(
|
||||||
|
&mut stream,
|
||||||
|
LeaseRefusalKind::Expired,
|
||||||
|
"the expiry is not in the future or is past the lease bound",
|
||||||
|
);
|
||||||
|
}
|
||||||
|
// The reservation is claimed in the ledger before the token is read,
|
||||||
|
// let alone lent. A ledger that fails closes the connection with no
|
||||||
|
// answer and no lease.
|
||||||
|
let admission = self
|
||||||
|
.service
|
||||||
|
.lock()
|
||||||
|
.map_err(|_| io::Error::other("service lock"))?
|
||||||
|
.begin_lease(
|
||||||
|
&caller.caller_id,
|
||||||
|
&invocation.invocation_id,
|
||||||
|
caller.reservation_micro_usd,
|
||||||
|
request.expires_at_unix_ms,
|
||||||
|
now,
|
||||||
|
)
|
||||||
|
.map_err(|_| io::Error::other("ledger"))?;
|
||||||
|
let grant = match admission {
|
||||||
|
LeaseAdmission::Granted(grant) => grant,
|
||||||
|
LeaseAdmission::Refused(why) => {
|
||||||
|
let (kind, detail) = match why {
|
||||||
|
LeaseRefusal::SpendCeiling => (
|
||||||
|
LeaseRefusalKind::SpendCeiling,
|
||||||
|
"the window cannot hold this reservation",
|
||||||
|
),
|
||||||
|
LeaseRefusal::CostBound => (
|
||||||
|
LeaseRefusalKind::CostBound,
|
||||||
|
"the reservation must be above zero",
|
||||||
|
),
|
||||||
|
LeaseRefusal::UnknownCaller => (
|
||||||
|
LeaseRefusalKind::UnknownCaller,
|
||||||
|
"this account is bound to no admitted caller",
|
||||||
|
),
|
||||||
|
LeaseRefusal::Busy => (
|
||||||
|
LeaseRefusalKind::Busy,
|
||||||
|
"a lease is open for this caller or this unit invocation leased",
|
||||||
|
),
|
||||||
|
};
|
||||||
|
return refuse(&mut stream, kind, detail);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
self.hold(stream, caller, &grant, request.expires_at_unix_ms)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Everything after a reservation exists: lend the token, then settle the
|
||||||
|
/// lease exactly once whatever happens.
|
||||||
|
fn hold(
|
||||||
|
&self,
|
||||||
|
mut stream: UnixStream,
|
||||||
|
caller: &LeaseCaller,
|
||||||
|
grant: &LeaseGrant,
|
||||||
|
expires_at_unix_ms: u64,
|
||||||
|
) -> io::Result<()> {
|
||||||
|
// A token that cannot be read or piped was never lent: nothing was
|
||||||
|
// used, so nothing is billed.
|
||||||
|
let Some(read_end) = self
|
||||||
|
.tokens
|
||||||
|
.read()
|
||||||
|
.and_then(|token| token_pipe(token.as_bytes()).ok())
|
||||||
|
else {
|
||||||
|
self.settle(caller, grant, LeaseUse::NeverLent)?;
|
||||||
|
return Ok(());
|
||||||
|
};
|
||||||
|
let frame = encode_frame(&LeaseMessage::ClaudeLeaseGrant(LeaseGranted {
|
||||||
|
lease_id: grant.lease_id.clone(),
|
||||||
|
reserved_micro_usd: grant.reserved_micro_usd,
|
||||||
|
}));
|
||||||
|
let sent = send_with_fd(&stream, &frame, read_end.as_raw_fd());
|
||||||
|
drop(read_end);
|
||||||
|
if sent.is_err() {
|
||||||
|
// The descriptor may have been delivered: the reservation is billed.
|
||||||
|
self.settle(caller, grant, LeaseUse::Reservation)?;
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
let outcome = self.await_settle(&mut stream, grant, expires_at_unix_ms);
|
||||||
|
match outcome {
|
||||||
|
Some(settle) => {
|
||||||
|
let settlement = self.settle(caller, grant, LeaseUse::Reported(&settle.usage))?;
|
||||||
|
let _ = stream.write_all(&encode_frame(&LeaseMessage::ClaudeLeaseSettled(
|
||||||
|
LeaseSettled {
|
||||||
|
billed_micro_usd: settlement.billed_micro_usd,
|
||||||
|
overrun_micro_usd: settlement.overrun_micro_usd,
|
||||||
|
},
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
None => {
|
||||||
|
self.settle(caller, grant, LeaseUse::Reservation)?;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The run's settle, if its whole frame arrives for this lease before the
|
||||||
|
/// lease expires. Anything else (end of file, silence, a frame still
|
||||||
|
/// trickling at the expiry, a malformed frame, the wrong lease) is no
|
||||||
|
/// report.
|
||||||
|
fn await_settle(
|
||||||
|
&self,
|
||||||
|
stream: &mut UnixStream,
|
||||||
|
grant: &LeaseGrant,
|
||||||
|
expires_at_unix_ms: u64,
|
||||||
|
) -> Option<LeaseSettle> {
|
||||||
|
let remaining = expires_at_unix_ms.saturating_sub(unix_ms());
|
||||||
|
let deadline = Instant::now() + Duration::from_millis(remaining);
|
||||||
|
match read_message(stream, deadline).ok()? {
|
||||||
|
LeaseMessage::ClaudeLeaseSettle(settle) if settle.lease_id == grant.lease_id => {
|
||||||
|
Some(settle)
|
||||||
|
}
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn settle(
|
||||||
|
&self,
|
||||||
|
caller: &LeaseCaller,
|
||||||
|
grant: &LeaseGrant,
|
||||||
|
usage: LeaseUse<'_>,
|
||||||
|
) -> io::Result<LeaseSettlement> {
|
||||||
|
self.service
|
||||||
|
.lock()
|
||||||
|
.map_err(|_| io::Error::other("service lock"))?
|
||||||
|
.settle_lease(&caller.caller_id, &grant.lease_id, usage)
|
||||||
|
.map_err(|_| io::Error::other("ledger"))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn refuse(stream: &mut UnixStream, kind: LeaseRefusalKind, detail: &str) -> io::Result<()> {
|
||||||
|
stream.write_all(&encode_frame(&LeaseMessage::ClaudeLeaseRefusal(
|
||||||
|
LeaseRefused {
|
||||||
|
kind,
|
||||||
|
detail: detail.to_string(),
|
||||||
|
},
|
||||||
|
)))
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
/// A settle frame whose body is exactly `body_bytes` long.
|
||||||
|
fn settle_frame_of(body_bytes: usize) -> Vec<u8> {
|
||||||
|
let frame_with = |pad: usize| {
|
||||||
|
encode_frame(&LeaseMessage::ClaudeLeaseSettle(LeaseSettle {
|
||||||
|
lease_id: "x".repeat(pad),
|
||||||
|
usage: Vec::new(),
|
||||||
|
}))
|
||||||
|
};
|
||||||
|
// The string header is the same width between 256 and 65535 bytes.
|
||||||
|
let base = frame_with(300).len() - 4 - 300;
|
||||||
|
frame_with(body_bytes - base)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn within(frame: &[u8]) -> io::Result<LeaseMessage> {
|
||||||
|
read_message(&mut &frame[..], Instant::now() + Duration::from_secs(5))
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn the_frame_bound_is_sixty_five_thousand_five_hundred_thirty_six() {
|
||||||
|
assert_eq!(MAX_LEASE_FRAME_BYTES, 65_536);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_frame_is_one_whole_message_within_the_bound() {
|
||||||
|
let at_bound = settle_frame_of(MAX_LEASE_FRAME_BYTES);
|
||||||
|
assert_eq!(at_bound.len(), 4 + MAX_LEASE_FRAME_BYTES);
|
||||||
|
assert!(decode_frame(&at_bound).is_some());
|
||||||
|
let past_bound = settle_frame_of(MAX_LEASE_FRAME_BYTES + 1);
|
||||||
|
assert!(decode_frame(&past_bound).is_none());
|
||||||
|
let mut trailing = at_bound.clone();
|
||||||
|
trailing.push(0);
|
||||||
|
assert!(decode_frame(&trailing).is_none());
|
||||||
|
assert!(decode_frame(&at_bound[..at_bound.len() - 1]).is_none());
|
||||||
|
assert!(decode_frame(&[0, 0, 0, 0]).is_none());
|
||||||
|
assert!(decode_frame(&[0, 0]).is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_stream_frame_over_the_bound_is_refused_before_it_is_read() {
|
||||||
|
let past_bound = settle_frame_of(MAX_LEASE_FRAME_BYTES + 1);
|
||||||
|
assert!(within(&past_bound).is_err());
|
||||||
|
let at_bound = settle_frame_of(MAX_LEASE_FRAME_BYTES);
|
||||||
|
assert!(within(&at_bound).is_ok());
|
||||||
|
assert!(within(&[0_u8, 0, 0, 0]).is_err());
|
||||||
|
// Short of the length it announced is end of file, not a message.
|
||||||
|
assert!(within(&at_bound[..at_bound.len() - 1]).is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_frame_that_trickles_past_its_deadline_is_refused_however_small_each_read() {
|
||||||
|
let (mut reader, mut writer) = UnixStream::pair().unwrap();
|
||||||
|
let frame = encode_frame(&LeaseMessage::ClaudeLeaseSettle(LeaseSettle {
|
||||||
|
lease_id: "lease:x".to_string(),
|
||||||
|
usage: Vec::new(),
|
||||||
|
}));
|
||||||
|
let trickle = thread::spawn(move || {
|
||||||
|
for byte in frame {
|
||||||
|
if writer.write_all(&[byte]).is_err() {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
thread::sleep(Duration::from_millis(60));
|
||||||
|
}
|
||||||
|
});
|
||||||
|
// Every read lands well inside any per-read timeout; the frame as a
|
||||||
|
// whole does not.
|
||||||
|
let started = Instant::now();
|
||||||
|
let outcome = read_message(&mut reader, started + Duration::from_millis(400));
|
||||||
|
assert_eq!(outcome.unwrap_err().kind(), io::ErrorKind::TimedOut);
|
||||||
|
assert!(started.elapsed() < Duration::from_millis(1_000));
|
||||||
|
drop(reader);
|
||||||
|
trickle.join().unwrap();
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A reader that fails the way a reset connection does.
|
||||||
|
struct Reset;
|
||||||
|
|
||||||
|
impl Read for Reset {
|
||||||
|
fn read(&mut self, _: &mut [u8]) -> io::Result<usize> {
|
||||||
|
Err(io::Error::from(io::ErrorKind::ConnectionReset))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Bounded for Reset {
|
||||||
|
fn within(&mut self, _: Duration) -> io::Result<()> {
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_read_that_fails_ends_the_frame_at_once_not_at_the_deadline() {
|
||||||
|
let started = Instant::now();
|
||||||
|
let outcome = read_message(&mut Reset, started + Duration::from_secs(30));
|
||||||
|
assert_eq!(outcome.unwrap_err().kind(), io::ErrorKind::ConnectionReset);
|
||||||
|
assert!(started.elapsed() < Duration::from_secs(1));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_frame_that_arrives_in_pieces_within_its_deadline_is_read() {
|
||||||
|
let (mut reader, mut writer) = UnixStream::pair().unwrap();
|
||||||
|
let frame = encode_frame(&LeaseMessage::ClaudeLeaseRequest(LeaseRequest {
|
||||||
|
expires_at_unix_ms: 7,
|
||||||
|
}));
|
||||||
|
let expected = decode_frame(&frame);
|
||||||
|
let pieces = thread::spawn(move || {
|
||||||
|
for part in frame.chunks(3) {
|
||||||
|
writer.write_all(part).unwrap();
|
||||||
|
thread::sleep(Duration::from_millis(10));
|
||||||
|
}
|
||||||
|
});
|
||||||
|
let read = read_message(&mut reader, Instant::now() + Duration::from_secs(5)).unwrap();
|
||||||
|
assert_eq!(Some(read), expected);
|
||||||
|
pieces.join().unwrap();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -8,18 +8,32 @@ mod config;
|
||||||
mod daemon;
|
mod daemon;
|
||||||
mod handle;
|
mod handle;
|
||||||
mod launch;
|
mod launch;
|
||||||
|
#[cfg(target_os = "linux")]
|
||||||
|
mod lease;
|
||||||
mod process;
|
mod process;
|
||||||
mod quota;
|
mod quota;
|
||||||
mod stream;
|
mod stream;
|
||||||
|
#[cfg(target_os = "linux")]
|
||||||
|
mod unit;
|
||||||
|
|
||||||
pub use backend::{ClaudeCodeBackend, CredentialTokenSource, StderrSink, TokenSource};
|
pub use backend::{ClaudeCodeBackend, CredentialTokenSource, StderrSink, TokenSource};
|
||||||
pub use config::{ClaudeCodeConfig, ClaudeCodeLimits};
|
pub use config::{ClaudeCodeConfig, ClaudeCodeLimits};
|
||||||
pub use daemon::{ClaudeConnectorConfig, ConfigError, check_ready, load_config, write_config};
|
pub use daemon::{
|
||||||
|
ClaudeConnectorConfig, ConfigError, LeaseCallerConfig, check_ready, load_config, write_config,
|
||||||
|
};
|
||||||
#[cfg(unix)]
|
#[cfg(unix)]
|
||||||
pub use daemon::{check_prompt_dir, serve_config};
|
pub use daemon::{check_prompt_dir, serve_config};
|
||||||
pub use handle::{FsSessionStore, SessionStore};
|
pub use handle::{FsSessionStore, SessionStore};
|
||||||
pub use launch::ChildEnv;
|
pub use launch::ChildEnv;
|
||||||
|
#[cfg(target_os = "linux")]
|
||||||
|
pub use lease::{
|
||||||
|
LeaseCaller, LeaseGranted, LeaseLimits, LeaseMessage, LeaseRefusalKind, LeaseRefused,
|
||||||
|
LeaseRequest, LeaseSettle, LeaseSettled, MAX_LEASE_FRAME_BYTES, PeerCredentials, PeerUnit,
|
||||||
|
PeerUnits, decode_frame, encode_frame, peer_credentials, serve_leases, user_uid,
|
||||||
|
};
|
||||||
#[cfg(unix)]
|
#[cfg(unix)]
|
||||||
pub use process::StdProcessPort;
|
pub use process::StdProcessPort;
|
||||||
pub use process::{ProcessPort, RunningChild, StdoutRead};
|
pub use process::{ProcessPort, RunningChild, StdoutRead};
|
||||||
pub use quota::QuotaCell;
|
pub use quota::QuotaCell;
|
||||||
|
#[cfg(target_os = "linux")]
|
||||||
|
pub use unit::SystemdUnits;
|
||||||
|
|
|
||||||
|
|
@ -13,7 +13,7 @@ use std::path::PathBuf;
|
||||||
use std::process::ExitCode;
|
use std::process::ExitCode;
|
||||||
|
|
||||||
use claude_connector::{
|
use claude_connector::{
|
||||||
ClaudeConnectorConfig, ConfigError, check_ready, load_config, write_config,
|
ClaudeConnectorConfig, ConfigError, LeaseCallerConfig, check_ready, load_config, write_config,
|
||||||
};
|
};
|
||||||
use model_connector::{CallerConfig, ModelPrice, Rates, frame_bytes_for_payload};
|
use model_connector::{CallerConfig, ModelPrice, Rates, frame_bytes_for_payload};
|
||||||
|
|
||||||
|
|
@ -76,9 +76,10 @@ struct Flags {
|
||||||
single: BTreeMap<&'static str, String>,
|
single: BTreeMap<&'static str, String>,
|
||||||
models: Vec<String>,
|
models: Vec<String>,
|
||||||
prices: Vec<String>,
|
prices: Vec<String>,
|
||||||
|
lease_users: Vec<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
const SINGLE: [&str; 22] = [
|
const SINGLE: [&str; 26] = [
|
||||||
"--listen",
|
"--listen",
|
||||||
"--replay-store",
|
"--replay-store",
|
||||||
"--claude-binary",
|
"--claude-binary",
|
||||||
|
|
@ -101,6 +102,10 @@ const SINGLE: [&str; 22] = [
|
||||||
"--max-payload-bytes",
|
"--max-payload-bytes",
|
||||||
"--max-output-tokens",
|
"--max-output-tokens",
|
||||||
"--socket-timeout-ms",
|
"--socket-timeout-ms",
|
||||||
|
"--lease-socket",
|
||||||
|
"--lease-max-ms",
|
||||||
|
"--lease-unit-instance",
|
||||||
|
"--lease-reservation-micro-usd",
|
||||||
];
|
];
|
||||||
|
|
||||||
impl Flags {
|
impl Flags {
|
||||||
|
|
@ -109,6 +114,7 @@ impl Flags {
|
||||||
single: BTreeMap::new(),
|
single: BTreeMap::new(),
|
||||||
models: Vec::new(),
|
models: Vec::new(),
|
||||||
prices: Vec::new(),
|
prices: Vec::new(),
|
||||||
|
lease_users: Vec::new(),
|
||||||
};
|
};
|
||||||
let mut args = args.iter();
|
let mut args = args.iter();
|
||||||
while let Some(&name) = args.next() {
|
while let Some(&name) = args.next() {
|
||||||
|
|
@ -118,6 +124,7 @@ impl Flags {
|
||||||
match name {
|
match name {
|
||||||
"--allowed-model" => flags.models.push(value.to_string()),
|
"--allowed-model" => flags.models.push(value.to_string()),
|
||||||
"--price" => flags.prices.push(value.to_string()),
|
"--price" => flags.prices.push(value.to_string()),
|
||||||
|
"--lease-user" => flags.lease_users.push(value.to_string()),
|
||||||
_ => {
|
_ => {
|
||||||
let Some(known) = SINGLE.iter().find(|known| **known == name) else {
|
let Some(known) = SINGLE.iter().find(|known| **known == name) else {
|
||||||
return Err(Failure::Flags);
|
return Err(Failure::Flags);
|
||||||
|
|
@ -206,6 +213,19 @@ fn init(path: &str, args: &[&str]) -> Result<(), Failure> {
|
||||||
return Err(Failure::Malformed("--price"));
|
return Err(Failure::Malformed("--price"));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
let caller_id = flags.text("--caller-id")?;
|
||||||
|
let lease_callers = flags
|
||||||
|
.lease_users
|
||||||
|
.iter()
|
||||||
|
.map(|unix_user| {
|
||||||
|
Ok(LeaseCallerConfig {
|
||||||
|
unix_user: unix_user.clone(),
|
||||||
|
caller_id: caller_id.clone(),
|
||||||
|
unit_instance: flags.text("--lease-unit-instance")?,
|
||||||
|
reservation_micro_usd: flags.number("--lease-reservation-micro-usd")?,
|
||||||
|
})
|
||||||
|
})
|
||||||
|
.collect::<Result<Vec<_>, Failure>>()?;
|
||||||
let config = ClaudeConnectorConfig {
|
let config = ClaudeConnectorConfig {
|
||||||
listen: flags.text("--listen")?,
|
listen: flags.text("--listen")?,
|
||||||
replay_store: flags.path("--replay-store")?,
|
replay_store: flags.path("--replay-store")?,
|
||||||
|
|
@ -226,7 +246,7 @@ fn init(path: &str, args: &[&str]) -> Result<(), Failure> {
|
||||||
socket_timeout_ms: flags.number_or("--socket-timeout-ms", 10_000)?,
|
socket_timeout_ms: flags.number_or("--socket-timeout-ms", 10_000)?,
|
||||||
max_expiry_skew_ms: 300_000,
|
max_expiry_skew_ms: 300_000,
|
||||||
callers: vec![CallerConfig {
|
callers: vec![CallerConfig {
|
||||||
caller_runtime_id: flags.text("--caller-id")?,
|
caller_runtime_id: caller_id.clone(),
|
||||||
connection_key_file: flags.path("--connection-key-file")?,
|
connection_key_file: flags.path("--connection-key-file")?,
|
||||||
connection_key_epoch: flags.number("--connection-key-epoch")?,
|
connection_key_epoch: flags.number("--connection-key-epoch")?,
|
||||||
allowed_models: flags.models.clone(),
|
allowed_models: flags.models.clone(),
|
||||||
|
|
@ -237,6 +257,9 @@ fn init(path: &str, args: &[&str]) -> Result<(), Failure> {
|
||||||
spend_window_ms: flags.number("--spend-window-ms")?,
|
spend_window_ms: flags.number("--spend-window-ms")?,
|
||||||
}],
|
}],
|
||||||
prices,
|
prices,
|
||||||
|
lease_socket: flags.single.get("--lease-socket").map(PathBuf::from),
|
||||||
|
lease_callers,
|
||||||
|
lease_max_ms: flags.number_or("--lease-max-ms", 3_600_000)?,
|
||||||
};
|
};
|
||||||
write_config(path.as_ref(), &config)?;
|
write_config(path.as_ref(), &config)?;
|
||||||
Ok(())
|
Ok(())
|
||||||
|
|
|
||||||
361
crates/claude-connector/src/unit.rs
Normal file
361
crates/claude-connector/src/unit.rs
Normal file
|
|
@ -0,0 +1,361 @@
|
||||||
|
//! Which systemd unit invocation a process runs in, from the system's own
|
||||||
|
//! record: the process's cgroup says the unit, and the service manager says
|
||||||
|
//! that unit's current InvocationID. The lease listener asks this through the
|
||||||
|
//! `PeerUnits` port; the answer never comes from the peer.
|
||||||
|
//!
|
||||||
|
//! This is the one place the connector asks systemd anything. It runs
|
||||||
|
//! `systemctl show` for one read-only property of one unit and takes nothing
|
||||||
|
//! from the output but a 32-digit hex id.
|
||||||
|
|
||||||
|
use std::fs;
|
||||||
|
use std::io::Read;
|
||||||
|
use std::os::fd::{AsRawFd, BorrowedFd};
|
||||||
|
use std::path::PathBuf;
|
||||||
|
use std::process::{Command, Stdio};
|
||||||
|
use std::time::{Duration, Instant};
|
||||||
|
|
||||||
|
use crate::lease::{PeerUnit, PeerUnits};
|
||||||
|
|
||||||
|
/// Answers `PeerUnits` from `<proc>/<pid>/cgroup` and `systemctl show`.
|
||||||
|
pub struct SystemdUnits {
|
||||||
|
proc_root: PathBuf,
|
||||||
|
systemctl: PathBuf,
|
||||||
|
/// How long the manager has to answer before the lease is refused.
|
||||||
|
deadline: Duration,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl SystemdUnits {
|
||||||
|
/// The running system's `/proc` and `systemctl`.
|
||||||
|
pub fn system() -> Self {
|
||||||
|
Self {
|
||||||
|
proc_root: PathBuf::from("/proc"),
|
||||||
|
systemctl: PathBuf::from("/usr/bin/systemctl"),
|
||||||
|
deadline: Duration::from_secs(5),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The pid a process handle names now, from the handle's own record. The
|
||||||
|
/// handle holds the process's identity, so the number cannot have been
|
||||||
|
/// recycled; a process that has exited reads as `-1`.
|
||||||
|
fn pid_of(&self, process: BorrowedFd<'_>) -> Option<u32> {
|
||||||
|
let info = fs::read_to_string(
|
||||||
|
self.proc_root
|
||||||
|
.join("self/fdinfo")
|
||||||
|
.join(process.as_raw_fd().to_string()),
|
||||||
|
)
|
||||||
|
.ok()?;
|
||||||
|
info.lines()
|
||||||
|
.find_map(|line| line.strip_prefix("Pid:"))?
|
||||||
|
.trim()
|
||||||
|
.parse()
|
||||||
|
.ok()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn unit_in_cgroup(&self, pid: u32) -> Option<String> {
|
||||||
|
let cgroup = fs::read_to_string(self.proc_root.join(pid.to_string()).join("cgroup"));
|
||||||
|
unit_of_cgroup(&cgroup.ok()?)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `systemctl show` of `unit`'s InvocationID, killed and refused if it has
|
||||||
|
/// not exited by the deadline.
|
||||||
|
fn invocation_of(&self, unit: &str) -> Option<String> {
|
||||||
|
let mut child = Command::new(&self.systemctl)
|
||||||
|
.env_clear()
|
||||||
|
.args(["show", "--property=InvocationID", "--value", "--", unit])
|
||||||
|
.stdin(Stdio::null())
|
||||||
|
.stdout(Stdio::piped())
|
||||||
|
.stderr(Stdio::null())
|
||||||
|
.spawn()
|
||||||
|
.ok()?;
|
||||||
|
let give_up = Instant::now() + self.deadline;
|
||||||
|
let status = loop {
|
||||||
|
match child.try_wait() {
|
||||||
|
Ok(Some(status)) => break status,
|
||||||
|
Ok(None) if Instant::now() < give_up => {
|
||||||
|
std::thread::sleep(Duration::from_millis(10));
|
||||||
|
}
|
||||||
|
_ => {
|
||||||
|
let _ = child.kill();
|
||||||
|
let _ = child.wait();
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
};
|
||||||
|
if !status.success() {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
// The id is 33 bytes; anything past a small bound is not an answer.
|
||||||
|
let mut output = Vec::new();
|
||||||
|
child
|
||||||
|
.stdout
|
||||||
|
.take()?
|
||||||
|
.take(256)
|
||||||
|
.read_to_end(&mut output)
|
||||||
|
.ok()?;
|
||||||
|
invocation_id_of(std::str::from_utf8(&output).ok()?)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The service unit a cgroup file puts its process in: the last component of
|
||||||
|
/// the unified-hierarchy path, if it names a `.service`. A process in a
|
||||||
|
/// sub-cgroup, a scope, a slice or the root is in no service unit this
|
||||||
|
/// recognises.
|
||||||
|
pub fn unit_of_cgroup(cgroup: &str) -> Option<String> {
|
||||||
|
let path = cgroup.lines().find_map(|line| line.strip_prefix("0::"))?;
|
||||||
|
let unit = path.rsplit('/').next()?;
|
||||||
|
(unit.len() > ".service".len() && unit.ends_with(".service")).then(|| unit.to_string())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// An InvocationID as `systemctl show --value` prints it: 32 lowercase hex
|
||||||
|
/// digits and a newline. Anything else (including the empty answer of a unit
|
||||||
|
/// that is not running) is not an invocation.
|
||||||
|
pub fn invocation_id_of(output: &str) -> Option<String> {
|
||||||
|
let id = output.strip_suffix('\n').unwrap_or(output);
|
||||||
|
(id.len() == 32 && id.bytes().all(|b| matches!(b, b'0'..=b'9' | b'a'..=b'f')))
|
||||||
|
.then(|| id.to_string())
|
||||||
|
}
|
||||||
|
|
||||||
|
impl PeerUnits for SystemdUnits {
|
||||||
|
fn invocation_in(&self, process: BorrowedFd<'_>, unit: &str) -> Option<PeerUnit> {
|
||||||
|
let pid = self.pid_of(process)?;
|
||||||
|
// A process in any other unit is not asked about.
|
||||||
|
if self.unit_in_cgroup(pid).as_deref() != Some(unit) {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
let invocation_id = self.invocation_of(unit)?;
|
||||||
|
// The same process must still be in that unit now: one that left (or
|
||||||
|
// died) while the manager was asked is not vouched for.
|
||||||
|
(self.pid_of(process) == Some(pid) && self.unit_in_cgroup(pid).as_deref() == Some(unit))
|
||||||
|
.then(|| PeerUnit {
|
||||||
|
unit: unit.to_string(),
|
||||||
|
invocation_id,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use std::os::fd::{AsFd, AsRawFd, OwnedFd};
|
||||||
|
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_service_unit_is_the_last_cgroup_component() {
|
||||||
|
let service = "0::/system.slice/system-mind\\x2dself.slice/mind-self@huginn.service\n";
|
||||||
|
assert_eq!(
|
||||||
|
unit_of_cgroup(service).as_deref(),
|
||||||
|
Some("mind-self@huginn.service")
|
||||||
|
);
|
||||||
|
// Several hierarchies listed: the unified one decides.
|
||||||
|
let hybrid = "12:cpu:/foo\n0::/system.slice/mind-self@a.service\n";
|
||||||
|
assert_eq!(
|
||||||
|
unit_of_cgroup(hybrid).as_deref(),
|
||||||
|
Some("mind-self@a.service")
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_process_outside_a_service_unit_has_none() {
|
||||||
|
for cgroup in [
|
||||||
|
"0::/\n",
|
||||||
|
"0::/user.slice/user-1000.slice/session-3.scope\n",
|
||||||
|
"0::/user.slice/user-1000.slice/user@1000.service/app.slice/x.service/child\n",
|
||||||
|
"0::/system.slice/mind-self@huginn.service/sub\n",
|
||||||
|
"0::/system.slice/.service\n",
|
||||||
|
"1:name=systemd:/system.slice/a.service\n",
|
||||||
|
"",
|
||||||
|
] {
|
||||||
|
assert_eq!(unit_of_cgroup(cgroup), None, "{cgroup:?}");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn an_invocation_id_is_exactly_32_lowercase_hex_digits() {
|
||||||
|
let id = "0123456789abcdef0123456789abcdef";
|
||||||
|
assert_eq!(invocation_id_of(&format!("{id}\n")).as_deref(), Some(id));
|
||||||
|
assert_eq!(invocation_id_of(id).as_deref(), Some(id));
|
||||||
|
for bad in [
|
||||||
|
"",
|
||||||
|
"\n",
|
||||||
|
"0123456789ABCDEF0123456789ABCDEF\n",
|
||||||
|
"0123456789abcdef0123456789abcde\n",
|
||||||
|
"0123456789abcdef0123456789abcdef0\n",
|
||||||
|
"0123456789abcdef0123456789abcdeg\n",
|
||||||
|
"0123456789abcdef0123456789abcdef\n\n",
|
||||||
|
] {
|
||||||
|
assert_eq!(invocation_id_of(bad), None, "{bad:?}");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
const ID: &str = "0123456789abcdef0123456789abcdef";
|
||||||
|
const UNIT: &str = "mind-self@eureka.service";
|
||||||
|
|
||||||
|
/// A fake `/proc` with one process in `cgroup`, and a fake `systemctl`
|
||||||
|
/// that records its arguments, runs `side_effect`, and exits `status`
|
||||||
|
/// after printing `output`.
|
||||||
|
struct Fake {
|
||||||
|
dir: PathBuf,
|
||||||
|
units: SystemdUnits,
|
||||||
|
// A script being written while another test forks is "text file
|
||||||
|
// busy" when it is run: fakes take turns.
|
||||||
|
_turn: std::sync::MutexGuard<'static, ()>,
|
||||||
|
}
|
||||||
|
|
||||||
|
static TURN: std::sync::Mutex<()> = std::sync::Mutex::new(());
|
||||||
|
|
||||||
|
impl Fake {
|
||||||
|
fn new(name: &str, cgroup: &str, output: &str, status: i32, side_effect: &str) -> Self {
|
||||||
|
use std::os::unix::fs::PermissionsExt;
|
||||||
|
let turn = TURN
|
||||||
|
.lock()
|
||||||
|
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||||
|
let dir = std::env::temp_dir().join(format!("unit-{}-{name}", std::process::id()));
|
||||||
|
let _ = fs::remove_dir_all(&dir);
|
||||||
|
let side_effect = side_effect.replace("{d}", &dir.display().to_string());
|
||||||
|
fs::create_dir_all(dir.join("proc/42")).unwrap();
|
||||||
|
fs::write(dir.join("proc/42/cgroup"), cgroup).unwrap();
|
||||||
|
let script = dir.join("systemctl");
|
||||||
|
fs::write(
|
||||||
|
&script,
|
||||||
|
format!(
|
||||||
|
"#!/bin/sh\nprintf '%s' \"$*\" > {d}/args\n{side_effect}\nprintf '%s' '{output}'\nexit {status}\n",
|
||||||
|
d = dir.display()
|
||||||
|
),
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
fs::set_permissions(&script, fs::Permissions::from_mode(0o755)).unwrap();
|
||||||
|
let units = SystemdUnits {
|
||||||
|
proc_root: dir.join("proc"),
|
||||||
|
systemctl: script,
|
||||||
|
deadline: Duration::from_secs(5),
|
||||||
|
};
|
||||||
|
Self {
|
||||||
|
dir,
|
||||||
|
units,
|
||||||
|
_turn: turn,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A process handle that names process `pid`, as the kernel's own
|
||||||
|
/// record of the handle says (`-1` once the process has exited).
|
||||||
|
fn process(&self, pid: i64) -> OwnedFd {
|
||||||
|
let handle: OwnedFd = fs::File::open("/dev/null").unwrap().into();
|
||||||
|
let info = self.dir.join("proc/self/fdinfo");
|
||||||
|
fs::create_dir_all(&info).unwrap();
|
||||||
|
fs::write(
|
||||||
|
info.join(handle.as_raw_fd().to_string()),
|
||||||
|
format!("pos:\t0\nflags:\t02000002\nPid:\t{pid}\nNSpid:\t{pid}\n"),
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
handle
|
||||||
|
}
|
||||||
|
|
||||||
|
fn invocation_in(&self, pid: i64, unit: &str) -> Option<PeerUnit> {
|
||||||
|
self.units.invocation_in(self.process(pid).as_fd(), unit)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn args(&self) -> String {
|
||||||
|
fs::read_to_string(self.dir.join("args")).unwrap()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Drop for Fake {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
let _ = fs::remove_dir_all(&self.dir);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn eureka_cgroup() -> String {
|
||||||
|
format!("0::/system.slice/{UNIT}\n")
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_process_in_the_unit_is_that_unit_at_the_invocation_the_manager_reports() {
|
||||||
|
let fake = Fake::new("ok", &eureka_cgroup(), &format!("{ID}\n"), 0, "");
|
||||||
|
assert_eq!(
|
||||||
|
fake.invocation_in(42, UNIT),
|
||||||
|
Some(PeerUnit {
|
||||||
|
unit: UNIT.to_string(),
|
||||||
|
invocation_id: ID.to_string()
|
||||||
|
})
|
||||||
|
);
|
||||||
|
// The manager is asked for exactly the invocation of that unit.
|
||||||
|
assert_eq!(
|
||||||
|
fake.args(),
|
||||||
|
format!("show --property=InvocationID --value -- {UNIT}")
|
||||||
|
);
|
||||||
|
// A handle on a process that has exited, or on one with no cgroup
|
||||||
|
// record, is in no unit.
|
||||||
|
assert_eq!(fake.invocation_in(-1, UNIT), None);
|
||||||
|
assert_eq!(fake.invocation_in(43, UNIT), None);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_process_outside_the_unit_is_not_asked_about() {
|
||||||
|
let scope = Fake::new("scope", "0::/user.slice/session-3.scope\n", ID, 0, "");
|
||||||
|
assert_eq!(scope.invocation_in(42, UNIT), None);
|
||||||
|
assert!(!scope.dir.join("args").exists());
|
||||||
|
drop(scope);
|
||||||
|
// Another service, even a real one, is not the caller's unit: the
|
||||||
|
// manager is never spawned for it.
|
||||||
|
let other = Fake::new("other", "0::/system.slice/ssh.service\n", ID, 0, "");
|
||||||
|
assert_eq!(other.invocation_in(42, UNIT), None);
|
||||||
|
assert!(!other.dir.join("args").exists());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_manager_that_fails_or_names_no_invocation_vouches_for_nothing() {
|
||||||
|
let failed = Fake::new("fail", &eureka_cgroup(), ID, 1, "");
|
||||||
|
assert_eq!(failed.invocation_in(42, UNIT), None);
|
||||||
|
drop(failed);
|
||||||
|
let not_running = Fake::new("idle", &eureka_cgroup(), "\n", 0, "");
|
||||||
|
assert_eq!(not_running.invocation_in(42, UNIT), None);
|
||||||
|
let missing = SystemdUnits {
|
||||||
|
proc_root: not_running.dir.join("proc"),
|
||||||
|
systemctl: not_running.dir.join("no-such-systemctl"),
|
||||||
|
deadline: Duration::from_secs(5),
|
||||||
|
};
|
||||||
|
assert_eq!(
|
||||||
|
missing.invocation_in(not_running.process(42).as_fd(), UNIT),
|
||||||
|
None
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_manager_that_does_not_answer_in_time_is_killed_and_the_lease_refused() {
|
||||||
|
let mut fake = Fake::new("hung", &eureka_cgroup(), ID, 0, "exec sleep 30");
|
||||||
|
fake.units.deadline = Duration::from_millis(200);
|
||||||
|
let started = Instant::now();
|
||||||
|
assert_eq!(fake.invocation_in(42, UNIT), None);
|
||||||
|
assert!(started.elapsed() < Duration::from_secs(5));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_process_that_left_its_unit_while_the_manager_was_asked_is_not_vouched_for() {
|
||||||
|
let fake = Fake::new(
|
||||||
|
"moves",
|
||||||
|
&eureka_cgroup(),
|
||||||
|
ID,
|
||||||
|
0,
|
||||||
|
"printf '0::/system.slice/other.service\n' > {d}/proc/42/cgroup",
|
||||||
|
);
|
||||||
|
assert_eq!(fake.invocation_in(42, UNIT), None);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_process_that_exited_while_the_manager_was_asked_is_not_vouched_for() {
|
||||||
|
let fake = Fake::new("dies", &eureka_cgroup(), ID, 0, "");
|
||||||
|
let handle = fake.process(42);
|
||||||
|
let side = fake
|
||||||
|
.dir
|
||||||
|
.join("proc/self/fdinfo")
|
||||||
|
.join(handle.as_raw_fd().to_string());
|
||||||
|
let mut script = fs::read_to_string(fake.dir.join("systemctl")).unwrap();
|
||||||
|
script = script.replace(
|
||||||
|
"exit 0",
|
||||||
|
&format!("printf 'Pid:\t-1\n' > {}\nexit 0", side.display()),
|
||||||
|
);
|
||||||
|
fs::write(fake.dir.join("systemctl"), script).unwrap();
|
||||||
|
assert_eq!(fake.units.invocation_in(handle.as_fd(), UNIT), None);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -8,7 +8,9 @@ use std::process::{Command, Output};
|
||||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
||||||
use claude_connector::{ClaudeConnectorConfig, ConfigError, load_config, write_config};
|
use claude_connector::{
|
||||||
|
ClaudeConnectorConfig, ConfigError, LeaseCallerConfig, load_config, write_config,
|
||||||
|
};
|
||||||
use cultcache_rs::{
|
use cultcache_rs::{
|
||||||
CacheBackingStore, CultCache, CultCacheEnvelope, DatabaseEntry,
|
CacheBackingStore, CultCache, CultCacheEnvelope, DatabaseEntry,
|
||||||
SingleFileMessagePackBackingStore,
|
SingleFileMessagePackBackingStore,
|
||||||
|
|
@ -92,6 +94,14 @@ fn valid() -> ClaudeConnectorConfig {
|
||||||
spend_window_ms: 3_600_000,
|
spend_window_ms: 3_600_000,
|
||||||
}],
|
}],
|
||||||
prices: BTreeMap::from([(HAIKU.to_string(), price()), (OPUS.to_string(), price())]),
|
prices: BTreeMap::from([(HAIKU.to_string(), price()), (OPUS.to_string(), price())]),
|
||||||
|
lease_socket: Some("/run/claude-connector/lease.sock".into()),
|
||||||
|
lease_callers: vec![LeaseCallerConfig {
|
||||||
|
unix_user: "mind-eureka-self".to_string(),
|
||||||
|
caller_id: "persona-a-caller".to_string(),
|
||||||
|
unit_instance: "eureka".to_string(),
|
||||||
|
reservation_micro_usd: 5_000_000,
|
||||||
|
}],
|
||||||
|
lease_max_ms: 3_600_000,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -989,3 +999,210 @@ fn a_private_regular_file_is_not_a_prompt_dir() {
|
||||||
assert!(claude_connector::check_prompt_dir(&file).is_err());
|
assert!(claude_connector::check_prompt_dir(&file).is_err());
|
||||||
assert!(claude_connector::check_prompt_dir(&scratch.path("absent")).is_err());
|
assert!(claude_connector::check_prompt_dir(&scratch.path("absent")).is_err());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn leases_are_configured_together_or_not_at_all() {
|
||||||
|
let cases: Vec<(&str, Change)> = vec![
|
||||||
|
(
|
||||||
|
"lease_socket",
|
||||||
|
Box::new(|c| c.lease_socket = Some("lease.sock".into())),
|
||||||
|
),
|
||||||
|
("lease_callers", Box::new(|c| c.lease_callers.clear())),
|
||||||
|
("lease_callers", Box::new(|c| c.lease_socket = None)),
|
||||||
|
(
|
||||||
|
"lease_callers",
|
||||||
|
Box::new(|c| c.lease_callers[0].caller_id = "another-caller".to_string()),
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"lease_callers",
|
||||||
|
Box::new(|c| c.lease_callers[0].unix_user = " ".to_string()),
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"lease_callers",
|
||||||
|
Box::new(|c| {
|
||||||
|
let again = c.lease_callers[0].clone();
|
||||||
|
c.lease_callers.push(again);
|
||||||
|
}),
|
||||||
|
),
|
||||||
|
("lease_max_ms", Box::new(|c| c.lease_max_ms = 0)),
|
||||||
|
(
|
||||||
|
"lease_callers",
|
||||||
|
Box::new(|c| c.lease_callers[0].reservation_micro_usd = 0),
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"lease_callers",
|
||||||
|
Box::new(|c| c.lease_callers[0].reservation_micro_usd = 50_000_001),
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"lease_callers",
|
||||||
|
Box::new(|c| c.lease_callers[0].unit_instance = String::new()),
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"lease_callers",
|
||||||
|
Box::new(|c| c.lease_callers[0].unit_instance = "a/b".to_string()),
|
||||||
|
),
|
||||||
|
];
|
||||||
|
for (field, change) in cases {
|
||||||
|
let mut config = valid();
|
||||||
|
change(&mut config);
|
||||||
|
assert_eq!(refused(&config), field);
|
||||||
|
}
|
||||||
|
// The whole ceiling is a reservation the caller can hold.
|
||||||
|
let mut whole = valid();
|
||||||
|
whole.lease_callers[0].reservation_micro_usd = 50_000_000;
|
||||||
|
whole.validate().unwrap();
|
||||||
|
let mut none = valid();
|
||||||
|
none.lease_socket = None;
|
||||||
|
none.lease_callers.clear();
|
||||||
|
none.validate().unwrap();
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn config_init_binds_each_lease_user_to_the_configured_caller() {
|
||||||
|
let scratch = Scratch::new();
|
||||||
|
let made = init(
|
||||||
|
&scratch,
|
||||||
|
&init_flags(
|
||||||
|
&scratch,
|
||||||
|
&[
|
||||||
|
"--lease-socket",
|
||||||
|
"/run/claude-connector/lease.sock",
|
||||||
|
"--lease-user",
|
||||||
|
"mind-eureka-self",
|
||||||
|
"--lease-max-ms",
|
||||||
|
"120000",
|
||||||
|
"--lease-unit-instance",
|
||||||
|
"eureka",
|
||||||
|
"--lease-reservation-micro-usd",
|
||||||
|
"5000000",
|
||||||
|
],
|
||||||
|
),
|
||||||
|
);
|
||||||
|
assert!(made.status.success(), "{}", shown(&made));
|
||||||
|
let config = load_config(&scratch.path("connector.cc")).unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
config.lease_callers,
|
||||||
|
[LeaseCallerConfig {
|
||||||
|
unix_user: "mind-eureka-self".to_string(),
|
||||||
|
caller_id: "persona-a-caller".to_string(),
|
||||||
|
unit_instance: "eureka".to_string(),
|
||||||
|
reservation_micro_usd: 5_000_000,
|
||||||
|
}]
|
||||||
|
);
|
||||||
|
assert_eq!(config.lease_max_ms, 120_000);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(target_os = "linux")]
|
||||||
|
#[test]
|
||||||
|
fn config_check_refuses_a_lease_account_the_system_does_not_have() {
|
||||||
|
let scratch = Scratch::new();
|
||||||
|
let (mut flags, credentials) = prepare(&scratch);
|
||||||
|
let socket = scratch.path("lease.sock");
|
||||||
|
flags.extend(
|
||||||
|
[
|
||||||
|
"--lease-socket",
|
||||||
|
socket.to_str().unwrap(),
|
||||||
|
"--lease-user",
|
||||||
|
"no-such-account-4f1c",
|
||||||
|
"--lease-unit-instance",
|
||||||
|
"eureka",
|
||||||
|
"--lease-reservation-micro-usd",
|
||||||
|
"5000000",
|
||||||
|
]
|
||||||
|
.map(String::from),
|
||||||
|
);
|
||||||
|
let made = init(&scratch, &flags);
|
||||||
|
assert!(made.status.success(), "{}", shown(&made));
|
||||||
|
let config = scratch.path("connector.cc");
|
||||||
|
let checked = run(
|
||||||
|
&["config", "check", config.to_str().unwrap()],
|
||||||
|
&[("CREDENTIALS_DIRECTORY", &credentials)],
|
||||||
|
);
|
||||||
|
assert!(!checked.status.success());
|
||||||
|
assert!(
|
||||||
|
shown(&checked).contains("lease_callers"),
|
||||||
|
"{}",
|
||||||
|
shown(&checked)
|
||||||
|
);
|
||||||
|
assert!(!shown(&checked).contains("no-such-account"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(target_os = "linux")]
|
||||||
|
#[test]
|
||||||
|
fn serve_binds_the_lease_socket_and_answers_a_stranger_with_a_refusal() {
|
||||||
|
use std::io::{Read, Write};
|
||||||
|
use std::os::unix::fs::PermissionsExt;
|
||||||
|
use std::os::unix::net::UnixStream;
|
||||||
|
use std::process::Stdio;
|
||||||
|
use std::time::Instant;
|
||||||
|
|
||||||
|
use claude_connector::{
|
||||||
|
LeaseMessage, LeaseRefusalKind, LeaseRefused, LeaseRequest, decode_frame, encode_frame,
|
||||||
|
};
|
||||||
|
|
||||||
|
let scratch = Scratch::new();
|
||||||
|
let (mut flags, credentials) = prepare(&scratch);
|
||||||
|
let socket = scratch.path("lease.sock");
|
||||||
|
// The test runs as someone other than `nobody`, so it is a stranger.
|
||||||
|
flags.extend(
|
||||||
|
[
|
||||||
|
"--lease-socket",
|
||||||
|
socket.to_str().unwrap(),
|
||||||
|
"--lease-user",
|
||||||
|
"nobody",
|
||||||
|
"--lease-unit-instance",
|
||||||
|
"eureka",
|
||||||
|
"--lease-reservation-micro-usd",
|
||||||
|
"5000000",
|
||||||
|
]
|
||||||
|
.map(String::from),
|
||||||
|
);
|
||||||
|
let made = init(&scratch, &flags);
|
||||||
|
assert!(made.status.success(), "{}", shown(&made));
|
||||||
|
let config = scratch.path("connector.cc");
|
||||||
|
let mut server = Command::new(env!("CARGO_BIN_EXE_claude-connector"))
|
||||||
|
.args(["serve", "--config", config.to_str().unwrap()])
|
||||||
|
.env_clear()
|
||||||
|
.env("CREDENTIALS_DIRECTORY", &credentials)
|
||||||
|
.stdin(Stdio::null())
|
||||||
|
.stdout(Stdio::null())
|
||||||
|
.stderr(Stdio::null())
|
||||||
|
.spawn()
|
||||||
|
.unwrap();
|
||||||
|
let started = Instant::now();
|
||||||
|
while !socket.exists() && started.elapsed() < Duration::from_secs(10) {
|
||||||
|
std::thread::sleep(Duration::from_millis(50));
|
||||||
|
}
|
||||||
|
let outcome = std::panic::catch_unwind(|| {
|
||||||
|
assert_eq!(
|
||||||
|
std::fs::metadata(&socket).unwrap().permissions().mode() & 0o777,
|
||||||
|
0o660
|
||||||
|
);
|
||||||
|
let mut stream = UnixStream::connect(&socket).unwrap();
|
||||||
|
stream
|
||||||
|
.set_read_timeout(Some(Duration::from_secs(5)))
|
||||||
|
.unwrap();
|
||||||
|
stream
|
||||||
|
.write_all(&encode_frame(&LeaseMessage::ClaudeLeaseRequest(
|
||||||
|
LeaseRequest {
|
||||||
|
expires_at_unix_ms: u64::MAX / 2,
|
||||||
|
},
|
||||||
|
)))
|
||||||
|
.unwrap();
|
||||||
|
let mut reply = vec![0_u8; 4096];
|
||||||
|
let read = stream.read(&mut reply).unwrap();
|
||||||
|
assert!(matches!(
|
||||||
|
decode_frame(&reply[..read]),
|
||||||
|
Some(LeaseMessage::ClaudeLeaseRefusal(LeaseRefused {
|
||||||
|
kind: LeaseRefusalKind::UnknownCaller,
|
||||||
|
..
|
||||||
|
}))
|
||||||
|
));
|
||||||
|
});
|
||||||
|
// The server is still serving: it did not return when it started.
|
||||||
|
let running = server.try_wait().unwrap().is_none();
|
||||||
|
server.kill().unwrap();
|
||||||
|
server.wait().unwrap();
|
||||||
|
outcome.unwrap();
|
||||||
|
assert!(running);
|
||||||
|
}
|
||||||
|
|
|
||||||
BIN
crates/claude-connector/tests/fixtures/lease_v1/grant.bin
vendored
Normal file
BIN
crates/claude-connector/tests/fixtures/lease_v1/grant.bin
vendored
Normal file
Binary file not shown.
BIN
crates/claude-connector/tests/fixtures/lease_v1/refusal.bin
vendored
Normal file
BIN
crates/claude-connector/tests/fixtures/lease_v1/refusal.bin
vendored
Normal file
Binary file not shown.
BIN
crates/claude-connector/tests/fixtures/lease_v1/request.bin
vendored
Normal file
BIN
crates/claude-connector/tests/fixtures/lease_v1/request.bin
vendored
Normal file
Binary file not shown.
BIN
crates/claude-connector/tests/fixtures/lease_v1/settle.bin
vendored
Normal file
BIN
crates/claude-connector/tests/fixtures/lease_v1/settle.bin
vendored
Normal file
Binary file not shown.
BIN
crates/claude-connector/tests/fixtures/lease_v1/settled.bin
vendored
Normal file
BIN
crates/claude-connector/tests/fixtures/lease_v1/settled.bin
vendored
Normal file
Binary file not shown.
1166
crates/claude-connector/tests/lease.rs
Normal file
1166
crates/claude-connector/tests/lease.rs
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -171,15 +171,7 @@ pub fn serve(
|
||||||
backend: Arc<dyn ProviderBackend>,
|
backend: Arc<dyn ProviderBackend>,
|
||||||
) -> Result<(), ServeError> {
|
) -> Result<(), ServeError> {
|
||||||
limits.validate(&admissions)?;
|
limits.validate(&admissions)?;
|
||||||
if !listener
|
admit_listener(&listener, backend.as_ref())?;
|
||||||
.local_addr()
|
|
||||||
.map_err(ServeError::Listen)?
|
|
||||||
.ip()
|
|
||||||
.is_loopback()
|
|
||||||
{
|
|
||||||
return Err(ServeError::InvalidConfig("listener address"));
|
|
||||||
}
|
|
||||||
backend.ready().map_err(ServeError::BackendUnready)?;
|
|
||||||
let service = Arc::new(Mutex::new(TransportService::open(
|
let service = Arc::new(Mutex::new(TransportService::open(
|
||||||
replay_store,
|
replay_store,
|
||||||
backend.provider(),
|
backend.provider(),
|
||||||
|
|
@ -190,6 +182,34 @@ pub fn serve(
|
||||||
accept_loop(|| listener.accept(), limits, service, backend)
|
accept_loop(|| listener.accept(), limits, service, backend)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// `serve` over a service the caller opened and keeps a handle to, so another
|
||||||
|
/// front of the same ledger (the Self-run credential lease) reserves in the
|
||||||
|
/// same windows. The limits must already have been validated against the
|
||||||
|
/// service's admissions.
|
||||||
|
pub fn serve_service(
|
||||||
|
listener: TcpListener,
|
||||||
|
limits: ServeLimits,
|
||||||
|
service: Arc<Mutex<TransportService>>,
|
||||||
|
backend: Arc<dyn ProviderBackend>,
|
||||||
|
) -> Result<(), ServeError> {
|
||||||
|
admit_listener(&listener, backend.as_ref())?;
|
||||||
|
accept_loop(|| listener.accept(), limits, service, backend)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// What must be true before a serve loop touches the ledger: the listener is
|
||||||
|
/// on loopback and the backend can take work.
|
||||||
|
fn admit_listener(listener: &TcpListener, backend: &dyn ProviderBackend) -> Result<(), ServeError> {
|
||||||
|
if !listener
|
||||||
|
.local_addr()
|
||||||
|
.map_err(ServeError::Listen)?
|
||||||
|
.ip()
|
||||||
|
.is_loopback()
|
||||||
|
{
|
||||||
|
return Err(ServeError::InvalidConfig("listener address"));
|
||||||
|
}
|
||||||
|
backend.ready().map_err(ServeError::BackendUnready)
|
||||||
|
}
|
||||||
|
|
||||||
/// The accept loop over any source of connections, so a test can present a
|
/// The accept loop over any source of connections, so a test can present a
|
||||||
/// peer address that a loopback listener never reports.
|
/// peer address that a loopback listener never reports.
|
||||||
fn accept_loop(
|
fn accept_loop(
|
||||||
|
|
|
||||||
|
|
@ -317,6 +317,9 @@ fn prune_to_latest(
|
||||||
struct LiveRequest {
|
struct LiveRequest {
|
||||||
window: u64,
|
window: u64,
|
||||||
reserved_micro_usd: u64,
|
reserved_micro_usd: u64,
|
||||||
|
/// A Self run's credential lease, not a provider call: it holds its
|
||||||
|
/// reservation in the window like a call does but takes no request slot.
|
||||||
|
lease: bool,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
|
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
|
||||||
|
|
@ -338,6 +341,12 @@ enum ReplayState {
|
||||||
charged_micro_usd: u64,
|
charged_micro_usd: u64,
|
||||||
overrun_micro_usd: u64,
|
overrun_micro_usd: u64,
|
||||||
},
|
},
|
||||||
|
/// A credential lease that settled: it has no answer to replay, only the
|
||||||
|
/// charge it left in its window and how far that passed the reservation.
|
||||||
|
LeaseSettled {
|
||||||
|
charged_micro_usd: u64,
|
||||||
|
overrun_micro_usd: u64,
|
||||||
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Clone, Debug, PartialEq, Eq, DatabaseEntry)]
|
#[derive(Clone, Debug, PartialEq, Eq, DatabaseEntry)]
|
||||||
|
|
@ -586,6 +595,9 @@ impl TransportService {
|
||||||
ReplayState::Active { reserved_micro_usd } => *reserved_micro_usd,
|
ReplayState::Active { reserved_micro_usd } => *reserved_micro_usd,
|
||||||
ReplayState::Completed {
|
ReplayState::Completed {
|
||||||
charged_micro_usd, ..
|
charged_micro_usd, ..
|
||||||
|
}
|
||||||
|
| ReplayState::LeaseSettled {
|
||||||
|
charged_micro_usd, ..
|
||||||
} => *charged_micro_usd,
|
} => *charged_micro_usd,
|
||||||
};
|
};
|
||||||
add_spend(
|
add_spend(
|
||||||
|
|
@ -787,6 +799,9 @@ impl TransportService {
|
||||||
.map_err(replay_store_error)?;
|
.map_err(replay_store_error)?;
|
||||||
Ok(TransportAdmission::Reply(stand_in))
|
Ok(TransportAdmission::Reply(stand_in))
|
||||||
}
|
}
|
||||||
|
ReplayState::LeaseSettled { .. } => {
|
||||||
|
self.reply_refusal(&invocation, Refusal::ReplayConflict, &security)
|
||||||
|
}
|
||||||
ReplayState::Active { .. } => {
|
ReplayState::Active { .. } => {
|
||||||
let refusal = if self.live.contains_key(&identity) {
|
let refusal = if self.live.contains_key(&identity) {
|
||||||
Refusal::InFlight
|
Refusal::InFlight
|
||||||
|
|
@ -810,8 +825,10 @@ impl TransportService {
|
||||||
};
|
};
|
||||||
let active_for_caller = self
|
let active_for_caller = self
|
||||||
.live
|
.live
|
||||||
.keys()
|
.iter()
|
||||||
.filter(|key| key.caller_runtime_id == invocation.caller_runtime_id)
|
.filter(|(key, request)| {
|
||||||
|
!request.lease && key.caller_runtime_id == invocation.caller_runtime_id
|
||||||
|
})
|
||||||
.count();
|
.count();
|
||||||
if active_for_caller >= max_concurrent_requests {
|
if active_for_caller >= max_concurrent_requests {
|
||||||
return self.reply_refusal(&invocation, Refusal::Capacity, &security);
|
return self.reply_refusal(&invocation, Refusal::Capacity, &security);
|
||||||
|
|
@ -873,6 +890,7 @@ impl TransportService {
|
||||||
LiveRequest {
|
LiveRequest {
|
||||||
window,
|
window,
|
||||||
reserved_micro_usd,
|
reserved_micro_usd,
|
||||||
|
lease: false,
|
||||||
},
|
},
|
||||||
);
|
);
|
||||||
Ok(TransportAdmission::Execute(Box::new(ExecutionClaim {
|
Ok(TransportAdmission::Execute(Box::new(ExecutionClaim {
|
||||||
|
|
@ -1234,6 +1252,9 @@ fn replay_store_error(error: impl std::fmt::Display) -> ServiceError {
|
||||||
ServiceError::ReplayStore(error.to_string())
|
ServiceError::ReplayStore(error.to_string())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
mod lease;
|
||||||
|
pub use lease::{LeaseAdmission, LeaseGrant, LeaseRefusal, LeaseSettlement, LeaseUse, ModelUsage};
|
||||||
|
|
||||||
fn contract_refusal(error: &ContractError) -> Refusal {
|
fn contract_refusal(error: &ContractError) -> Refusal {
|
||||||
match error {
|
match error {
|
||||||
ContractError::Expiry => Refusal::Expired,
|
ContractError::Expiry => Refusal::Expired,
|
||||||
|
|
@ -2525,7 +2546,9 @@ mod tests {
|
||||||
ReplayState::Completed {
|
ReplayState::Completed {
|
||||||
charged_micro_usd, ..
|
charged_micro_usd, ..
|
||||||
} => charged_micro_usd,
|
} => charged_micro_usd,
|
||||||
ReplayState::Active { .. } => panic!("{request_id} is not completed"),
|
ReplayState::Active { .. } | ReplayState::LeaseSettled { .. } => {
|
||||||
|
panic!("{request_id} is not completed")
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
691
crates/model-connector/src/service/lease.rs
Normal file
691
crates/model-connector/src/service/lease.rs
Normal file
|
|
@ -0,0 +1,691 @@
|
||||||
|
//! A credential lease: a Self run's reservation in the caller's spend window.
|
||||||
|
//! The lease is one record of the same ledger a provider call reserves in, not
|
||||||
|
//! a second ledger: it is claimed before anything is lent, counts against the
|
||||||
|
//! window's ceiling while open, and settles to the true cost of what the run
|
||||||
|
//! reports it used, per model and uncapped: a charge above the reservation is
|
||||||
|
//! kept as the lease's overrun, and still counts against the window. A lease
|
||||||
|
//! nothing settled is billed the whole reservation, as a call a crash left
|
||||||
|
//! behind is.
|
||||||
|
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
/// The request-id namespace of leases in a caller's replay records.
|
||||||
|
const LEASE_PREFIX: &str = "lease:";
|
||||||
|
|
||||||
|
/// What an admitted lease holds.
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||||
|
pub struct LeaseGrant {
|
||||||
|
pub lease_id: String,
|
||||||
|
pub reserved_micro_usd: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// What a run used of one model.
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||||
|
pub struct ModelUsage {
|
||||||
|
pub model: String,
|
||||||
|
pub usage: Usage,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// What a settled lease was billed. `overrun_micro_usd` is how far the charge
|
||||||
|
/// passed the reservation.
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
|
pub struct LeaseSettlement {
|
||||||
|
pub billed_micro_usd: u64,
|
||||||
|
pub overrun_micro_usd: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// What a lease's settlement is billed on.
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
|
pub enum LeaseUse<'a> {
|
||||||
|
/// No report: the whole reservation. A run that cannot say what it used
|
||||||
|
/// pays what was set aside for it.
|
||||||
|
Reservation,
|
||||||
|
/// The run's report. A report that names no use (an empty list) is no
|
||||||
|
/// report, and bills the reservation.
|
||||||
|
Reported(&'a [ModelUsage]),
|
||||||
|
/// The token was never lent, so nothing was used and nothing is billed.
|
||||||
|
NeverLent,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// What `usage` costs at the table's prices: each entry at its model's price,
|
||||||
|
/// and a model with no price at the dearest price the table holds, so an
|
||||||
|
/// unknown name never bills less than any known one. `None` when the table
|
||||||
|
/// prices nothing.
|
||||||
|
fn lease_cost(prices: &PriceTable, usage: &[ModelUsage]) -> Option<u64> {
|
||||||
|
let dearest = |usage: &Usage| {
|
||||||
|
prices
|
||||||
|
.values()
|
||||||
|
.map(|price| cost_micro_usd(price, usage, None))
|
||||||
|
.max()
|
||||||
|
};
|
||||||
|
usage
|
||||||
|
.iter()
|
||||||
|
.map(|entry| match prices.get(&entry.model) {
|
||||||
|
Some(price) => Some(cost_micro_usd(price, &entry.usage, None)),
|
||||||
|
None => dearest(&entry.usage),
|
||||||
|
})
|
||||||
|
.try_fold(0_u64, |total, cost| Some(total.saturating_add(cost?)))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Why no lease was reserved. Nothing was written to the ledger.
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
|
pub enum LeaseRefusal {
|
||||||
|
/// The window's spend plus open reservations plus this one passes the
|
||||||
|
/// caller's ceiling.
|
||||||
|
SpendCeiling,
|
||||||
|
/// The reservation asked for is zero: a lease that bills nothing is a
|
||||||
|
/// free pass.
|
||||||
|
CostBound,
|
||||||
|
UnknownCaller,
|
||||||
|
/// The caller holds a lease already, or this run id was leased before.
|
||||||
|
Busy,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||||
|
pub enum LeaseAdmission {
|
||||||
|
Granted(LeaseGrant),
|
||||||
|
Refused(LeaseRefusal),
|
||||||
|
}
|
||||||
|
|
||||||
|
impl TransportService {
|
||||||
|
/// Reserves `max_cost_micro_usd` in `caller_runtime_id`'s current window
|
||||||
|
/// for the unit invocation `invocation_id`, which leases once. A ledger
|
||||||
|
/// that cannot be written answers `Err` and reserves nothing. A lease takes no request slot.
|
||||||
|
pub fn begin_lease(
|
||||||
|
&mut self,
|
||||||
|
caller_runtime_id: &str,
|
||||||
|
invocation_id: &str,
|
||||||
|
max_cost_micro_usd: u64,
|
||||||
|
expires_at_unix_ms: u64,
|
||||||
|
now_unix_ms: u64,
|
||||||
|
) -> Result<LeaseAdmission, ServiceError> {
|
||||||
|
let Some(caller) = self.callers.get(caller_runtime_id) else {
|
||||||
|
return Ok(LeaseAdmission::Refused(LeaseRefusal::UnknownCaller));
|
||||||
|
};
|
||||||
|
let spend_ceiling = caller.spend_ceiling_micro_usd;
|
||||||
|
let spend_window_ms = caller.spend_window_ms;
|
||||||
|
let connection_key_epoch = caller.connection_key_epoch;
|
||||||
|
require_id(invocation_id, "invocation_id")?;
|
||||||
|
if invocation_id.len() + LEASE_PREFIX.len() > MAX_IDENTITY_BYTES {
|
||||||
|
return Err(ContractError::Invalid("invocation_id").into());
|
||||||
|
}
|
||||||
|
if max_cost_micro_usd == 0 {
|
||||||
|
return Ok(LeaseAdmission::Refused(LeaseRefusal::CostBound));
|
||||||
|
}
|
||||||
|
let identity = RequestIdentity {
|
||||||
|
caller_runtime_id: caller_runtime_id.to_string(),
|
||||||
|
request_id: format!("{LEASE_PREFIX}{invocation_id}"),
|
||||||
|
};
|
||||||
|
let key = replay_key(&identity);
|
||||||
|
let holding = self
|
||||||
|
.live
|
||||||
|
.iter()
|
||||||
|
.any(|(held, request)| request.lease && held.caller_runtime_id == caller_runtime_id);
|
||||||
|
if holding
|
||||||
|
|| self
|
||||||
|
.replay_store
|
||||||
|
.get::<ReplayRecord>(&key)
|
||||||
|
.map_err(replay_store_error)?
|
||||||
|
.is_some()
|
||||||
|
{
|
||||||
|
return Ok(LeaseAdmission::Refused(LeaseRefusal::Busy));
|
||||||
|
}
|
||||||
|
let admitted_at_unix_ms = admission_clock(
|
||||||
|
&self.window_spend,
|
||||||
|
caller_runtime_id,
|
||||||
|
now_unix_ms,
|
||||||
|
spend_window_ms,
|
||||||
|
);
|
||||||
|
let window = admitted_at_unix_ms / spend_window_ms;
|
||||||
|
prune_windows(
|
||||||
|
&mut self.window_spend,
|
||||||
|
&self.live,
|
||||||
|
caller_runtime_id,
|
||||||
|
window,
|
||||||
|
);
|
||||||
|
if self
|
||||||
|
.window_spend(caller_runtime_id, window)
|
||||||
|
.saturating_add(max_cost_micro_usd)
|
||||||
|
> spend_ceiling
|
||||||
|
{
|
||||||
|
return Ok(LeaseAdmission::Refused(LeaseRefusal::SpendCeiling));
|
||||||
|
}
|
||||||
|
let mut digest = Sha256::new();
|
||||||
|
digest.update(b"claude_lease_request");
|
||||||
|
digest.update((invocation_id.len() as u64).to_be_bytes());
|
||||||
|
digest.update(invocation_id.as_bytes());
|
||||||
|
digest.update(max_cost_micro_usd.to_be_bytes());
|
||||||
|
digest.update(expires_at_unix_ms.to_be_bytes());
|
||||||
|
let record = ReplayRecord {
|
||||||
|
caller_runtime_id: caller_runtime_id.to_string(),
|
||||||
|
request_id: identity.request_id.clone(),
|
||||||
|
connection_key_epoch,
|
||||||
|
invocation_sha256: digest.finalize().into(),
|
||||||
|
expires_at_unix_ms,
|
||||||
|
state: ReplayState::Active {
|
||||||
|
reserved_micro_usd: max_cost_micro_usd,
|
||||||
|
},
|
||||||
|
admitted_at_unix_ms,
|
||||||
|
};
|
||||||
|
// Reserving and claiming are one write, as for a call: a ledger that
|
||||||
|
// cannot be written lends nothing.
|
||||||
|
self.replay_store
|
||||||
|
.put(key, &record)
|
||||||
|
.map_err(replay_store_error)?;
|
||||||
|
add_spend(
|
||||||
|
&mut self.window_spend,
|
||||||
|
caller_runtime_id,
|
||||||
|
window,
|
||||||
|
0,
|
||||||
|
max_cost_micro_usd,
|
||||||
|
);
|
||||||
|
let lease_id = identity.request_id.clone();
|
||||||
|
self.live.insert(
|
||||||
|
identity,
|
||||||
|
LiveRequest {
|
||||||
|
window,
|
||||||
|
reserved_micro_usd: max_cost_micro_usd,
|
||||||
|
lease: true,
|
||||||
|
},
|
||||||
|
);
|
||||||
|
Ok(LeaseAdmission::Granted(LeaseGrant {
|
||||||
|
lease_id,
|
||||||
|
reserved_micro_usd: max_cost_micro_usd,
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Settles an open lease and returns what it was billed. A report bills
|
||||||
|
/// the true cost of every model's reported use, uncapped (the overrun is
|
||||||
|
/// recorded); no report, or an empty one, bills the whole reservation;
|
||||||
|
/// only a token never lent bills nothing. A failed ledger
|
||||||
|
/// write frees the caller's lease and keeps the reservation counted, as a
|
||||||
|
/// call's does.
|
||||||
|
pub fn settle_lease(
|
||||||
|
&mut self,
|
||||||
|
caller_runtime_id: &str,
|
||||||
|
lease_id: &str,
|
||||||
|
usage: LeaseUse<'_>,
|
||||||
|
) -> Result<LeaseSettlement, ServiceError> {
|
||||||
|
let identity = RequestIdentity {
|
||||||
|
caller_runtime_id: caller_runtime_id.to_string(),
|
||||||
|
request_id: lease_id.to_string(),
|
||||||
|
};
|
||||||
|
if !self
|
||||||
|
.live
|
||||||
|
.get(&identity)
|
||||||
|
.is_some_and(|request| request.lease)
|
||||||
|
{
|
||||||
|
return Err(ServiceError::NoActiveRequest);
|
||||||
|
}
|
||||||
|
let key = replay_key(&identity);
|
||||||
|
let active = self
|
||||||
|
.replay_store
|
||||||
|
.get::<ReplayRecord>(&key)
|
||||||
|
.map_err(replay_store_error)?
|
||||||
|
.ok_or(ServiceError::NoActiveRequest)?;
|
||||||
|
let ReplayState::Active { reserved_micro_usd } = active.state else {
|
||||||
|
return Err(ServiceError::ActiveRequestMismatch);
|
||||||
|
};
|
||||||
|
let caller = self
|
||||||
|
.callers
|
||||||
|
.get(caller_runtime_id)
|
||||||
|
.ok_or(ServiceError::CallerNotAdmitted)?;
|
||||||
|
let spend_window_ms = caller.spend_window_ms;
|
||||||
|
let charged_micro_usd = match usage {
|
||||||
|
LeaseUse::Reservation | LeaseUse::Reported([]) => reserved_micro_usd,
|
||||||
|
LeaseUse::NeverLent => 0,
|
||||||
|
LeaseUse::Reported(usage) => {
|
||||||
|
lease_cost(&self.prices, usage).unwrap_or(reserved_micro_usd)
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let overrun_micro_usd = charged_micro_usd.saturating_sub(reserved_micro_usd);
|
||||||
|
let settled = ReplayRecord {
|
||||||
|
state: ReplayState::LeaseSettled {
|
||||||
|
charged_micro_usd,
|
||||||
|
overrun_micro_usd,
|
||||||
|
},
|
||||||
|
..active.clone()
|
||||||
|
};
|
||||||
|
self.replay_store
|
||||||
|
.put(key, &settled)
|
||||||
|
.map_err(replay_store_error)
|
||||||
|
.inspect_err(|_| {
|
||||||
|
self.live.remove(&identity);
|
||||||
|
})?;
|
||||||
|
add_spend(
|
||||||
|
&mut self.window_spend,
|
||||||
|
caller_runtime_id,
|
||||||
|
active.admitted_at_unix_ms / spend_window_ms,
|
||||||
|
reserved_micro_usd,
|
||||||
|
charged_micro_usd,
|
||||||
|
);
|
||||||
|
self.live.remove(&identity);
|
||||||
|
prune_to_latest(&mut self.window_spend, &self.live, caller_runtime_id);
|
||||||
|
Ok(LeaseSettlement {
|
||||||
|
billed_micro_usd: charged_micro_usd,
|
||||||
|
overrun_micro_usd,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use std::sync::atomic::Ordering;
|
||||||
|
|
||||||
|
use super::*;
|
||||||
|
use crate::fixtures::*;
|
||||||
|
|
||||||
|
const CALLER: &str = "eureka";
|
||||||
|
const KEY: &str = "lease-distinct-test-key";
|
||||||
|
const NOW: u64 = 7_200_500;
|
||||||
|
|
||||||
|
fn open_with(store: impl CacheBackingStore + 'static, ceiling: u64) -> TransportService {
|
||||||
|
TransportService::with_store(
|
||||||
|
store,
|
||||||
|
"fixture-sse",
|
||||||
|
fixture_prices(),
|
||||||
|
[CallerAdmission::new(
|
||||||
|
CALLER,
|
||||||
|
KEY,
|
||||||
|
1,
|
||||||
|
["gpt-5.4".to_string()],
|
||||||
|
1,
|
||||||
|
65_536,
|
||||||
|
32_768,
|
||||||
|
SpendCeiling(ceiling),
|
||||||
|
SpendWindowMs(WINDOW_MS),
|
||||||
|
)
|
||||||
|
.unwrap()],
|
||||||
|
5_000,
|
||||||
|
)
|
||||||
|
.unwrap()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn open(store: &TestStore, ceiling: u64) -> TransportService {
|
||||||
|
TransportService::open(
|
||||||
|
&store.path,
|
||||||
|
"fixture-sse",
|
||||||
|
fixture_prices(),
|
||||||
|
[CallerAdmission::new(
|
||||||
|
CALLER,
|
||||||
|
KEY,
|
||||||
|
1,
|
||||||
|
["gpt-5.4".to_string()],
|
||||||
|
1,
|
||||||
|
65_536,
|
||||||
|
32_768,
|
||||||
|
SpendCeiling(ceiling),
|
||||||
|
SpendWindowMs(WINDOW_MS),
|
||||||
|
)
|
||||||
|
.unwrap()],
|
||||||
|
5_000,
|
||||||
|
)
|
||||||
|
.unwrap()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn lease(service: &mut TransportService, run: &str, cost: u64) -> LeaseAdmission {
|
||||||
|
service
|
||||||
|
.begin_lease(CALLER, run, cost, NOW + 60_000, NOW)
|
||||||
|
.unwrap()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn granted(admission: LeaseAdmission) -> LeaseGrant {
|
||||||
|
match admission {
|
||||||
|
LeaseAdmission::Granted(grant) => grant,
|
||||||
|
LeaseAdmission::Refused(why) => panic!("refused {why:?}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn held(service: &TransportService) -> u64 {
|
||||||
|
service.window_spend(CALLER, NOW / WINDOW_MS)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn output(tokens: u64) -> Vec<ModelUsage> {
|
||||||
|
used("gpt-5.4", tokens)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn used(model: &str, tokens: u64) -> Vec<ModelUsage> {
|
||||||
|
vec![ModelUsage {
|
||||||
|
model: model.to_string(),
|
||||||
|
usage: Usage {
|
||||||
|
input_tokens: 0,
|
||||||
|
cache_read_tokens: 0,
|
||||||
|
cache_write_tokens: 0,
|
||||||
|
output_tokens: tokens,
|
||||||
|
reasoning_tokens: None,
|
||||||
|
},
|
||||||
|
}]
|
||||||
|
}
|
||||||
|
|
||||||
|
fn billed(settlement: LeaseSettlement) -> u64 {
|
||||||
|
settlement.billed_micro_usd
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_lease_reserves_in_the_callers_window_and_names_what_it_holds() {
|
||||||
|
let store = TestStore::new();
|
||||||
|
let mut service = open(&store, 1_000);
|
||||||
|
let grant = granted(lease(&mut service, "run-1", 600));
|
||||||
|
assert_eq!(grant.reserved_micro_usd, 600);
|
||||||
|
assert_eq!(held(&service), 600);
|
||||||
|
assert_eq!(
|
||||||
|
billed(
|
||||||
|
service
|
||||||
|
.settle_lease(CALLER, &grant.lease_id, LeaseUse::Reservation)
|
||||||
|
.unwrap()
|
||||||
|
),
|
||||||
|
600
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn the_window_ceiling_counts_open_and_settled_reservations_inclusively() {
|
||||||
|
let store = TestStore::new();
|
||||||
|
let mut service = open(&store, 1_000);
|
||||||
|
let first = granted(lease(&mut service, "run-1", 600));
|
||||||
|
service
|
||||||
|
.settle_lease(CALLER, &first.lease_id, LeaseUse::Reservation)
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
lease(&mut service, "run-2", 401),
|
||||||
|
LeaseAdmission::Refused(LeaseRefusal::SpendCeiling)
|
||||||
|
);
|
||||||
|
assert_eq!(held(&service), 600);
|
||||||
|
granted(lease(&mut service, "run-3", 400));
|
||||||
|
assert_eq!(held(&service), 1_000);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_caller_holds_one_lease_at_a_time_and_an_invocation_id_is_leased_once() {
|
||||||
|
let store = TestStore::new();
|
||||||
|
let mut service = open(&store, 1_000_000);
|
||||||
|
let first = granted(lease(&mut service, "run-1", 100));
|
||||||
|
assert_eq!(
|
||||||
|
lease(&mut service, "run-2", 100),
|
||||||
|
LeaseAdmission::Refused(LeaseRefusal::Busy)
|
||||||
|
);
|
||||||
|
service
|
||||||
|
.settle_lease(CALLER, &first.lease_id, LeaseUse::Reservation)
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
lease(&mut service, "run-1", 100),
|
||||||
|
LeaseAdmission::Refused(LeaseRefusal::Busy)
|
||||||
|
);
|
||||||
|
granted(lease(&mut service, "run-2", 100));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_stranger_and_a_free_lease_are_refused_without_a_ledger_write() {
|
||||||
|
let store = TestStore::new();
|
||||||
|
let mut service = open(&store, 1_000);
|
||||||
|
assert_eq!(
|
||||||
|
service
|
||||||
|
.begin_lease("stranger", "run-1", 100, NOW + 1, NOW)
|
||||||
|
.unwrap(),
|
||||||
|
LeaseAdmission::Refused(LeaseRefusal::UnknownCaller)
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
lease(&mut service, "run-1", 0),
|
||||||
|
LeaseAdmission::Refused(LeaseRefusal::CostBound)
|
||||||
|
);
|
||||||
|
assert_eq!(held(&service), 0);
|
||||||
|
// Nothing was claimed: the same run id leases once it asks sensibly.
|
||||||
|
granted(lease(&mut service, "run-1", 100));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn settlement_bills_the_true_cost_uncapped_and_records_the_overrun() {
|
||||||
|
let store = TestStore::new();
|
||||||
|
let mut service = open(&store, 1_000_000);
|
||||||
|
// Output at 15 USD per million tokens: 10 tokens bill 150 micro-USD.
|
||||||
|
let small = granted(lease(&mut service, "run-1", 10_000));
|
||||||
|
let settled = service
|
||||||
|
.settle_lease(CALLER, &small.lease_id, LeaseUse::Reported(&output(10)[..]))
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
(settled.billed_micro_usd, settled.overrun_micro_usd),
|
||||||
|
(150, 0)
|
||||||
|
);
|
||||||
|
assert_eq!(held(&service), 150);
|
||||||
|
// 1000 tokens bill 15_000: the reservation does not cap it, and the
|
||||||
|
// 5_000 past it is recorded as the overrun.
|
||||||
|
let big = granted(lease(&mut service, "run-2", 10_000));
|
||||||
|
let settled = service
|
||||||
|
.settle_lease(
|
||||||
|
CALLER,
|
||||||
|
&big.lease_id,
|
||||||
|
LeaseUse::Reported(&output(1_000)[..]),
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
(settled.billed_micro_usd, settled.overrun_micro_usd),
|
||||||
|
(15_000, 5_000)
|
||||||
|
);
|
||||||
|
assert_eq!(held(&service), 15_150);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn settlement_prices_each_model_at_its_own_rate_and_an_unknown_model_at_the_dearest() {
|
||||||
|
let store = TestStore::new();
|
||||||
|
let cheap = ModelPrice {
|
||||||
|
rates: Rates {
|
||||||
|
input: 1_000_000,
|
||||||
|
cache_write_short: 1_000_000,
|
||||||
|
cache_write_long: 1_000_000,
|
||||||
|
cache_read: 1_000_000,
|
||||||
|
output: 2_000_000,
|
||||||
|
},
|
||||||
|
..fixture_price()
|
||||||
|
};
|
||||||
|
let prices = PriceTable::from([
|
||||||
|
("gpt-5.4".to_string(), fixture_price()),
|
||||||
|
("cheap".to_string(), cheap),
|
||||||
|
]);
|
||||||
|
let mut service = TransportService::open(
|
||||||
|
&store.path,
|
||||||
|
"fixture-sse",
|
||||||
|
prices,
|
||||||
|
[CallerAdmission::new(
|
||||||
|
CALLER,
|
||||||
|
KEY,
|
||||||
|
1,
|
||||||
|
["gpt-5.4".to_string()],
|
||||||
|
1,
|
||||||
|
65_536,
|
||||||
|
32_768,
|
||||||
|
SpendCeiling(1_000_000),
|
||||||
|
SpendWindowMs(WINDOW_MS),
|
||||||
|
)
|
||||||
|
.unwrap()],
|
||||||
|
5_000,
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
let mixed: Vec<ModelUsage> = [("gpt-5.4", 10), ("cheap", 10), ("unheard-of", 10)]
|
||||||
|
.into_iter()
|
||||||
|
.flat_map(|(model, tokens)| used(model, tokens))
|
||||||
|
.collect();
|
||||||
|
let grant = granted(lease(&mut service, "run-1", 10_000));
|
||||||
|
// 10 tokens at 15, 10 at 2, and 10 at the dearest (15) per million.
|
||||||
|
assert_eq!(
|
||||||
|
billed(
|
||||||
|
service
|
||||||
|
.settle_lease(CALLER, &grant.lease_id, LeaseUse::Reported(&mixed[..]))
|
||||||
|
.unwrap()
|
||||||
|
),
|
||||||
|
150 + 20 + 150
|
||||||
|
);
|
||||||
|
// A run that reports no use pays the reservation, not nothing.
|
||||||
|
let grant = granted(lease(&mut service, "run-2", 10_000));
|
||||||
|
assert_eq!(
|
||||||
|
billed(
|
||||||
|
service
|
||||||
|
.settle_lease(CALLER, &grant.lease_id, LeaseUse::Reported(&[]))
|
||||||
|
.unwrap()
|
||||||
|
),
|
||||||
|
10_000
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn settling_an_empty_report_cannot_free_the_reservation() {
|
||||||
|
let store = TestStore::new();
|
||||||
|
let mut service = open(&store, 1_000);
|
||||||
|
let grant = granted(lease(&mut service, "run-1", 600));
|
||||||
|
assert_eq!(
|
||||||
|
billed(
|
||||||
|
service
|
||||||
|
.settle_lease(CALLER, &grant.lease_id, LeaseUse::Reported(&[]))
|
||||||
|
.unwrap()
|
||||||
|
),
|
||||||
|
600
|
||||||
|
);
|
||||||
|
assert_eq!(held(&service), 600);
|
||||||
|
assert!(matches!(
|
||||||
|
lease(&mut service, "run-2", 1_000),
|
||||||
|
LeaseAdmission::Refused(LeaseRefusal::SpendCeiling)
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_token_never_lent_bills_nothing_and_frees_the_window() {
|
||||||
|
let store = TestStore::new();
|
||||||
|
let mut service = open(&store, 1_000);
|
||||||
|
let grant = granted(lease(&mut service, "run-1", 600));
|
||||||
|
let settled = service
|
||||||
|
.settle_lease(CALLER, &grant.lease_id, LeaseUse::NeverLent)
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
(settled.billed_micro_usd, settled.overrun_micro_usd),
|
||||||
|
(0, 0)
|
||||||
|
);
|
||||||
|
assert_eq!(held(&service), 0);
|
||||||
|
granted(lease(&mut service, "run-2", 1_000));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_lease_settles_once() {
|
||||||
|
let store = TestStore::new();
|
||||||
|
let mut service = open(&store, 1_000_000);
|
||||||
|
let grant = granted(lease(&mut service, "run-1", 100));
|
||||||
|
service
|
||||||
|
.settle_lease(CALLER, &grant.lease_id, LeaseUse::Reservation)
|
||||||
|
.unwrap();
|
||||||
|
assert!(matches!(
|
||||||
|
service.settle_lease(CALLER, &grant.lease_id, LeaseUse::Reservation),
|
||||||
|
Err(ServiceError::NoActiveRequest)
|
||||||
|
));
|
||||||
|
assert_eq!(held(&service), 100);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_lease_left_open_by_a_crash_counts_at_its_reservation_after_restart() {
|
||||||
|
let store = TestStore::new();
|
||||||
|
let mut service = open(&store, 1_000_000);
|
||||||
|
granted(lease(&mut service, "run-1", 700));
|
||||||
|
drop(service);
|
||||||
|
let service = open(&store, 1_000_000);
|
||||||
|
assert_eq!(held(&service), 700);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_ledger_that_cannot_be_written_lends_nothing() {
|
||||||
|
let flaky = FlakyStore::new(1_000);
|
||||||
|
let allowance = flaky.allowance();
|
||||||
|
let mut service = open_with(flaky, 1_000_000);
|
||||||
|
allowance.store(0, Ordering::SeqCst);
|
||||||
|
assert!(
|
||||||
|
service
|
||||||
|
.begin_lease(CALLER, "run-1", 100, NOW + 1, NOW)
|
||||||
|
.is_err()
|
||||||
|
);
|
||||||
|
assert_eq!(held(&service), 0);
|
||||||
|
allowance.store(1_000, Ordering::SeqCst);
|
||||||
|
granted(lease(&mut service, "run-1", 100));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_failed_settle_write_frees_the_lease_and_keeps_the_reservation_counted() {
|
||||||
|
let flaky = FlakyStore::new(1_000);
|
||||||
|
let allowance = flaky.allowance();
|
||||||
|
let mut service = open_with(flaky, 1_000_000);
|
||||||
|
let grant = granted(lease(&mut service, "run-1", 100));
|
||||||
|
allowance.store(0, Ordering::SeqCst);
|
||||||
|
assert!(
|
||||||
|
service
|
||||||
|
.settle_lease(CALLER, &grant.lease_id, LeaseUse::Reported(&output(1)[..]))
|
||||||
|
.is_err()
|
||||||
|
);
|
||||||
|
assert_eq!(held(&service), 100);
|
||||||
|
allowance.store(1_000, Ordering::SeqCst);
|
||||||
|
granted(lease(&mut service, "run-2", 100));
|
||||||
|
}
|
||||||
|
|
||||||
|
fn begin_call(service: &mut TransportService, id: &str) -> TransportAdmission {
|
||||||
|
let mut call = request();
|
||||||
|
call.request_id = id.to_string();
|
||||||
|
let invocation =
|
||||||
|
TransportInvocation::new(CALLER, NOW + 1_000, [7; 32], MAX_COST, call).unwrap();
|
||||||
|
let sealed = crate::envelope::encrypt_message(
|
||||||
|
CALLER,
|
||||||
|
id,
|
||||||
|
crate::envelope::EnvelopeKind::Invocation,
|
||||||
|
&crate::contract::wire_encode(&invocation).unwrap(),
|
||||||
|
&security(KEY),
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
service.begin(&sealed, NOW, 1 << 20).unwrap()
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_lease_takes_no_request_slot_from_the_callers_calls() {
|
||||||
|
// The caller admits one concurrent call; an open lease must not use it.
|
||||||
|
let store = TestStore::new();
|
||||||
|
let mut service = open(&store, 1 << 40);
|
||||||
|
granted(lease(&mut service, "run-1", 100));
|
||||||
|
match begin_call(&mut service, "call-1") {
|
||||||
|
TransportAdmission::Execute(_) => {}
|
||||||
|
TransportAdmission::Reply(_) => panic!("the lease held the call's slot"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_call_in_flight_does_not_make_the_caller_busy_for_a_lease() {
|
||||||
|
let store = TestStore::new();
|
||||||
|
let mut service = open(&store, 1 << 40);
|
||||||
|
assert!(matches!(
|
||||||
|
begin_call(&mut service, "call-1"),
|
||||||
|
TransportAdmission::Execute(_)
|
||||||
|
));
|
||||||
|
granted(lease(&mut service, "run-1", 100));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn an_invocation_id_is_bounded_so_its_lease_id_fits_an_identity() {
|
||||||
|
let store = TestStore::new();
|
||||||
|
let mut service = open(&store, 1 << 40);
|
||||||
|
let longest = "r".repeat(MAX_IDENTITY_BYTES - LEASE_PREFIX.len());
|
||||||
|
let grant = granted(lease(&mut service, &longest, 100));
|
||||||
|
assert_eq!(grant.lease_id.len(), MAX_IDENTITY_BYTES);
|
||||||
|
service
|
||||||
|
.settle_lease(CALLER, &grant.lease_id, LeaseUse::Reservation)
|
||||||
|
.unwrap();
|
||||||
|
assert!(
|
||||||
|
service
|
||||||
|
.begin_lease(CALLER, &format!("{longest}r"), 100, NOW + 1, NOW)
|
||||||
|
.is_err()
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_settled_lease_leaves_its_charge_not_its_reservation_across_a_restart() {
|
||||||
|
let store = TestStore::new();
|
||||||
|
let mut service = open(&store, 1_000_000);
|
||||||
|
let grant = granted(lease(&mut service, "run-1", 10_000));
|
||||||
|
service
|
||||||
|
.settle_lease(CALLER, &grant.lease_id, LeaseUse::Reported(&output(10)[..]))
|
||||||
|
.unwrap();
|
||||||
|
drop(service);
|
||||||
|
let service = open(&store, 1_000_000);
|
||||||
|
assert_eq!(held(&service), 150);
|
||||||
|
}
|
||||||
|
}
|
||||||
Loading…
Reference in a new issue