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

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)

Browse source 
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Meta Vens 2026-10-29 21:54:03 +01:00
commit d1f25b5ae6
15 changed files with 3421 additions and 28 deletions
Download patch file Download diff file Expand all files Collapse all files

171
crates/claude-connector/src/daemon.rs
View file

@ -6,17 +6,22 @@
use std::collections::{BTreeMap, BTreeSet};
use std::net::{SocketAddr, TcpListener};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use cultcache_rs::{CultCache, DatabaseEntry, SingleFileMessagePackBackingStore};
use model_connector::{
CallerAdmission, CallerConfig, MIN_PAYLOAD_BYTES, ModelPrice, ServeError, ServeLimits,
ServiceError, load_admissions, serve,
CallerAdmission, CallerConfig, MIN_PAYLOAD_BYTES, ModelPrice, ProviderBackend, ServeError,
ServeLimits, ServiceError, TransportService, load_admissions, serve_service,
};
use serde::{Deserialize, Serialize};
use thiserror::Error;
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::{
ClaudeCodeBackend, ClaudeCodeConfig, ClaudeCodeLimits, CredentialTokenSource, StderrSink,
};
@ -24,6 +29,18 @@ use crate::{
const CONFIG_KEY: &str = "runtime";
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)]
#[cultcache(type = "gamecult.model_connector.claude_connector_config.v1")]
pub struct ClaudeConnectorConfig {
@ -72,6 +89,17 @@ pub struct ClaudeConnectorConfig {
/// ids exactly as the CLI reports them in its init event.
#[cultcache(key = 18)]
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
@ -205,9 +233,39 @@ impl ClaudeConnectorConfig {
{
return Err(ConfigError::Invalid("prices"));
}
self.validate_leases(&caller.caller_runtime_id, caller.spend_ceiling_micro_usd)?;
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 {
ClaudeCodeConfig {
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 {
!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()),
)
.map_err(ConfigError::Unusable)?;
if config.lease_socket.is_some() {
lease_callers_of(config)?;
}
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.
pub fn check_ready(config: &ClaudeConnectorConfig) -> Result<(), ConfigError> {
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()))),
));
let listener = TcpListener::bind(listen).map_err(|_| ConfigError::Listen)?;
serve(
listener,
config.serve_limits(),
admissions,
config.prices.clone(),
&config.replay_store,
backend,
)?;
let service = Arc::new(Mutex::new(
TransportService::open(
&config.replay_store,
backend.provider(),
config.prices.clone(),
admissions,
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(())
}

729
crates/claude-connector/src/lease.rs Normal file
View 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();
}
}

16
crates/claude-connector/src/lib.rs
View file

@ -8,18 +8,32 @@ mod config;
mod daemon;
mod handle;
mod launch;
#[cfg(target_os = "linux")]
mod lease;
mod process;
mod quota;
mod stream;
#[cfg(target_os = "linux")]
mod unit;
pub use backend::{ClaudeCodeBackend, CredentialTokenSource, StderrSink, TokenSource};
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)]
pub use daemon::{check_prompt_dir, serve_config};
pub use handle::{FsSessionStore, SessionStore};
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)]
pub use process::StdProcessPort;
pub use process::{ProcessPort, RunningChild, StdoutRead};
pub use quota::QuotaCell;
#[cfg(target_os = "linux")]
pub use unit::SystemdUnits;

29
crates/claude-connector/src/main.rs
View file

@ -13,7 +13,7 @@ use std::path::PathBuf;
use std::process::ExitCode;
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};
@ -76,9 +76,10 @@ struct Flags {
single: BTreeMap<&'static str, String>,
models: Vec<String>,
prices: Vec<String>,
lease_users: Vec<String>,
}
const SINGLE: [&str; 22] = [
const SINGLE: [&str; 26] = [
"--listen",
"--replay-store",
"--claude-binary",
@ -101,6 +102,10 @@ const SINGLE: [&str; 22] = [
"--max-payload-bytes",
"--max-output-tokens",
"--socket-timeout-ms",
"--lease-socket",
"--lease-max-ms",
"--lease-unit-instance",
"--lease-reservation-micro-usd",
];
impl Flags {
@ -109,6 +114,7 @@ impl Flags {
single: BTreeMap::new(),
models: Vec::new(),
prices: Vec::new(),
lease_users: Vec::new(),
};
let mut args = args.iter();
while let Some(&name) = args.next() {
@ -118,6 +124,7 @@ impl Flags {
match name {
"--allowed-model" => flags.models.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 {
return Err(Failure::Flags);
@ -206,6 +213,19 @@ fn init(path: &str, args: &[&str]) -> Result<(), Failure> {
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 {
listen: flags.text("--listen")?,
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)?,
max_expiry_skew_ms: 300_000,
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_epoch: flags.number("--connection-key-epoch")?,
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")?,
}],
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)?;
Ok(())

361
crates/claude-connector/src/unit.rs Normal file
View 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);
}
}

219
crates/claude-connector/tests/config_document.rs
View file

@ -8,7 +8,9 @@ use std::process::{Command, Output};
use std::sync::atomic::{AtomicUsize, Ordering};
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::{
CacheBackingStore, CultCache, CultCacheEnvelope, DatabaseEntry,
SingleFileMessagePackBackingStore,
@ -92,6 +94,14 @@ fn valid() -> ClaudeConnectorConfig {
spend_window_ms: 3_600_000,
}],
prices: BTreeMap::from([(HAIKU.to_string(), price()), (OPUS.to_string(), price())]),
lease_socket: Some("/run/claude-connector/lease.sock".into()),
lease_callers: vec![LeaseCallerConfig {
unix_user: "mind-eureka-self".to_string(),
caller_id: "persona-a-caller".to_string(),
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(&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
View file

Binary file not shown.

BIN
crates/claude-connector/tests/fixtures/lease_v1/refusal.bin vendored Normal file
View file

Binary file not shown.

BIN
crates/claude-connector/tests/fixtures/lease_v1/request.bin vendored Normal file
View file

Binary file not shown.

BIN
crates/claude-connector/tests/fixtures/lease_v1/settle.bin vendored Normal file
View file

Binary file not shown.

BIN
crates/claude-connector/tests/fixtures/lease_v1/settled.bin vendored Normal file
View file

Binary file not shown.

1166
crates/claude-connector/tests/lease.rs Normal file
View file

File diff suppressed because it is too large Load diff

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

@ -171,15 +171,7 @@ pub fn serve(
backend: Arc<dyn ProviderBackend>,
) -> Result<(), ServeError> {
limits.validate(&admissions)?;
if !listener
.local_addr()
.map_err(ServeError::Listen)?
.ip()
.is_loopback()
{
return Err(ServeError::InvalidConfig("listener address"));
}
backend.ready().map_err(ServeError::BackendUnready)?;
admit_listener(&listener, backend.as_ref())?;
let service = Arc::new(Mutex::new(TransportService::open(
replay_store,
backend.provider(),
@ -190,6 +182,34 @@ pub fn serve(
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
/// peer address that a loopback listener never reports.
fn accept_loop(

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

@ -317,6 +317,9 @@ fn prune_to_latest(
struct LiveRequest {
window: u64,
reserved_micro_usd: u64,
/// A Self run's credential lease, not a provider call: it holds its
/// reservation in the window like a call does but takes no request slot.
lease: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
@ -338,6 +341,12 @@ enum ReplayState {
charged_micro_usd: u64,
overrun_micro_usd: u64,
},
/// A credential lease that settled: it has no answer to replay, only the
/// charge it left in its window and how far that passed the reservation.
LeaseSettled {
charged_micro_usd: u64,
overrun_micro_usd: u64,
},
}
#[derive(Clone, Debug, PartialEq, Eq, DatabaseEntry)]
@ -586,6 +595,9 @@ impl TransportService {
ReplayState::Active { reserved_micro_usd } => *reserved_micro_usd,
ReplayState::Completed {
charged_micro_usd, ..
}
| ReplayState::LeaseSettled {
charged_micro_usd, ..
} => *charged_micro_usd,
};
add_spend(
@ -787,6 +799,9 @@ impl TransportService {
.map_err(replay_store_error)?;
Ok(TransportAdmission::Reply(stand_in))
}
ReplayState::LeaseSettled { .. } => {
self.reply_refusal(&invocation, Refusal::ReplayConflict, &security)
}
ReplayState::Active { .. } => {
let refusal = if self.live.contains_key(&identity) {
Refusal::InFlight
@ -810,8 +825,10 @@ impl TransportService {
};
let active_for_caller = self
.live
.keys()
.filter(|key| key.caller_runtime_id == invocation.caller_runtime_id)
.iter()
.filter(|(key, request)| {
!request.lease && key.caller_runtime_id == invocation.caller_runtime_id
})
.count();
if active_for_caller >= max_concurrent_requests {
return self.reply_refusal(&invocation, Refusal::Capacity, &security);
@ -873,6 +890,7 @@ impl TransportService {
LiveRequest {
window,
reserved_micro_usd,
lease: false,
},
);
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())
}
mod lease;
pub use lease::{LeaseAdmission, LeaseGrant, LeaseRefusal, LeaseSettlement, LeaseUse, ModelUsage};
fn contract_refusal(error: &ContractError) -> Refusal {
match error {
ContractError::Expiry => Refusal::Expired,
@ -2525,7 +2546,9 @@ mod tests {
ReplayState::Completed {
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
View 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);
}
}