diff --git a/README.md b/README.md index 74bd5b5..7ffd262 100644 --- a/README.md +++ b/README.md @@ -5,7 +5,7 @@ between a caller that needs a model turn and a daemon that holds the provider credential, so no project wires its own provider calls. The connector meters nothing in dollars. It holds the subscription token, lends it, and -publishes the plan's usage windows (`usage_readout`, and `plan_quota` on every lease grant) for the +publishes the plan's usage windows (`usage_readout`, `claude_usage_readout` on the lease socket, and `plan_quota` on every lease grant) for the minds to read and decide their own burn. The one usage-derived refusal is the provider's own (`ProviderRejected`), repeated until its reset. diff --git a/crates/claude-connector/src/lease.rs b/crates/claude-connector/src/lease.rs index 7e503aa..a049bba 100644 --- a/crates/claude-connector/src/lease.rs +++ b/crates/claude-connector/src/lease.rs @@ -10,6 +10,10 @@ //! provider's own, repeated until its own reset. Nothing the run says about //! itself decides who it is or what the plan allows. //! +//! The same socket also answers a usage readout to any peer that can reach it +//! (the claude-lease group): no lease, no caller check, no token read and no +//! wake of the meter. Usage is state every Mind reads. +//! //! 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` @@ -27,7 +31,7 @@ use std::thread; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use model_connector::{ - LeaseAdmission, LeaseGrant, LeaseRefusal, LeaseTurn, PlanQuota, TransportService, + LeaseAdmission, LeaseGrant, LeaseRefusal, LeaseTurn, PlanQuota, TransportService, UsageReadout, }; use serde::{Deserialize, Serialize}; @@ -48,6 +52,10 @@ pub enum LeaseMessage { /// holds the token, its write end closed. ClaudeLeaseGrant(LeaseGranted), ClaudeLeaseRefusal(LeaseRefused), + /// Asks for the plan's usage windows. Opens no lease. + ClaudeUsageReadoutRequest(LeaseUsageReadoutRequest), + /// The answer to a readout request: the cell's windows as of now, no fd. + ClaudeUsageReadout(UsageReadout), } /// A request names only when the lease ends. Who asks is the connector's to know. @@ -56,6 +64,10 @@ pub struct LeaseRequest { pub expires_at_unix_ms: u64, } +/// A readout request names nothing: the answer is the same for every asker. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct LeaseUsageReadoutRequest {} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct LeaseGranted { pub lease_id: String, @@ -383,10 +395,22 @@ struct Lender { quota: Arc, tokens: Arc, units: Arc, + /// Connections holding or waiting on a lease, at most `max_connections`. + leases: Arc, + /// Connections being answered a readout, at most `READOUT_CONNECTIONS`. + readouts: Arc, } -/// Serves leases on `listener` until it fails. Each connection is one lease: -/// request, then grant or refusal, then (after a grant) the hold. A caller is +/// How many usage readouts are served at once, beside the leases. A readout +/// holds nothing, so a few suffice, and a flood of them cannot take a lease's +/// slot: the accept loop admits `max_connections` plus this many. +const READOUT_CONNECTIONS: usize = 4; + +/// Serves leases on `listener` until it fails. Each connection is one lease or +/// one usage readout. A lease: request, then grant or refusal, then (after a +/// grant) the hold. A readout: request, then one frame, then close; its reach +/// is the socket's group, and it reads the plan's windows without waking the +/// meter. A caller is /// the account whose uid the kernel reports for the connection, and only a /// process in that caller's unit may lease, once per unit invocation for a Self /// run and once per turn for a Persona; no @@ -407,11 +431,14 @@ pub fn serve_leases( quota, tokens, units, + leases: Arc::new(AtomicUsize::new(0)), + readouts: Arc::new(AtomicUsize::new(0)), }); let active = Arc::new(AtomicUsize::new(0)); loop { let (stream, _) = listener.accept()?; - let Some(permit) = Permit::acquire(&active, limits.max_connections) else { + let admitted = limits.max_connections.saturating_add(READOUT_CONNECTIONS); + let Some(permit) = Permit::acquire(&active, admitted) else { continue; }; let lender = lender.clone(); @@ -421,21 +448,45 @@ pub fn serve_leases( let _permit = permit; // A connection that fails is dropped; the lease it may have // held is ended inside `lend`. - let _ = lender.lend(stream); + let _ = lender.serve_one(stream); })?; } } impl Lender { - fn lend(&self, mut stream: UnixStream) -> io::Result<()> { + fn serve_one(&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)); - }; + match read_message(&mut stream, request_deadline)? { + LeaseMessage::ClaudeLeaseRequest(request) => { + // A lease over the limit is dropped unanswered; readouts + // never count against it. + let Some(_slot) = Permit::acquire(&self.leases, self.limits.max_connections) else { + return Ok(()); + }; + let peer = peer_credentials(&stream)?; + self.lend(stream, peer, request) + } + LeaseMessage::ClaudeUsageReadoutRequest(_) => { + let Some(_slot) = Permit::acquire(&self.readouts, READOUT_CONNECTIONS) else { + return refuse( + &mut stream, + LeaseRefusalKind::Busy, + "too many usage readouts at once", + ); + }; + stream.write_all(&encode_frame(&readout(&self.quota, unix_ms()))) + } + _ => Err(io::Error::from(io::ErrorKind::InvalidData)), + } + } + + fn lend( + &self, + mut stream: UnixStream, + peer: PeerCredentials, + request: LeaseRequest, + ) -> io::Result<()> { let Some(caller) = self.callers.iter().find(|caller| caller.uid == peer.uid) else { return refuse( &mut stream, @@ -609,6 +660,12 @@ impl Lender { } } +/// The plan's windows as the cell holds them at `now_ms`. It takes the cell by +/// reference so it cannot note a request, and never waits for a first read. +fn readout(quota: &QuotaCell, now_ms: u64) -> LeaseMessage { + LeaseMessage::ClaudeUsageReadout(UsageReadout::new(now_ms, quota.snapshot())) +} + /// The longest a lease waits for the meter's first read of the week. A call /// takes a few seconds; the holder's own wait for the grant is far longer. const FIRST_READ_WAIT: Duration = Duration::from_secs(10); @@ -629,6 +686,210 @@ fn refuse(stream: &mut UnixStream, kind: LeaseRefusalKind, detail: &str) -> io:: mod tests { use super::*; + struct NoTokens; + impl TokenSource for NoTokens { + fn read(&self) -> Option> { + panic!("a readout read the token") + } + } + struct NoUnits; + impl PeerUnits for NoUnits { + fn invocation_in(&self, _: BorrowedFd<'_>, _: &str) -> Option { + panic!("a readout looked up the peer's unit") + } + } + + fn lender( + quota: &Arc, + dir: &std::path::Path, + max_connections: usize, + ) -> (Arc>, Lender) { + let service = Arc::new(Mutex::new( + TransportService::open( + &dir.join("replay.cc"), + "claude-code", + [model_connector::CallerAdmission::new( + "lease-test-caller", + "lease-test-connection-key", + 1, + ["model".to_string()], + 1, + 65_536, + 32_768, + ) + .unwrap()], + 300_000, + ) + .unwrap(), + )); + let lender = Lender { + callers: Vec::new(), + limits: LeaseLimits { + socket_timeout: Duration::from_secs(2), + max_lease_ms: 3_600_000, + max_connections, + }, + service: service.clone(), + quota: quota.clone(), + tokens: Arc::new(NoTokens), + units: Arc::new(NoUnits), + leases: Arc::new(AtomicUsize::new(0)), + readouts: Arc::new(AtomicUsize::new(0)), + }; + (service, lender) + } + + /// Serves `lender`'s parts on a socket in `dir`; the path comes back. + fn serve(lender: Lender, dir: &std::path::Path) -> std::path::PathBuf { + let socket = dir.join("lease.sock"); + let listener = UnixListener::bind(&socket).unwrap(); + let (service, quota) = (lender.service.clone(), lender.quota.clone()); + thread::spawn(move || { + let _ = serve_leases( + &listener, + lender.callers, + lender.limits, + service, + quota, + lender.tokens, + lender.units, + ); + }); + socket + } + + fn ask(socket: &std::path::Path, message: &LeaseMessage) -> Option { + let mut stream = UnixStream::connect(socket).ok()?; + stream + .set_read_timeout(Some(Duration::from_secs(3))) + .unwrap(); + stream.write_all(&encode_frame(message)).ok()?; + let mut raw = Vec::new(); + let _ = stream.read_to_end(&mut raw); + decode_frame(&raw) + } + + fn readout_request() -> LeaseMessage { + LeaseMessage::ClaudeUsageReadoutRequest(LeaseUsageReadoutRequest {}) + } + + fn scratch(name: &str) -> std::path::PathBuf { + let dir = std::env::temp_dir().join(format!("{name}-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + dir + } + + #[test] + fn a_socket_readout_never_asks_the_meter_waits_for_a_week_or_takes_the_service_lock() { + let dir = scratch("readout-arm"); + let quota = Arc::new(QuotaCell::new()); + quota.attach_meter(); + let (service, lender) = lender("a, &dir, 8); + let socket = serve(lender, &dir); + // Held throughout: a readout that takes it hangs past the read timeout. + let guard = service.lock().unwrap(); + for _ in 0..20 { + let started = Instant::now(); + let Some(LeaseMessage::ClaudeUsageReadout(answer)) = ask(&socket, &readout_request()) + else { + panic!("no readout"); + }; + assert!(started.elapsed() < Duration::from_secs(1)); + assert!(answer.plan_quota.is_empty()); + } + drop(guard); + assert_eq!(quota.asks_so_far(), 0); + assert!(!quota.wanted_at(unix_ms(), Duration::from_secs(15 * 60))); + let _ = std::fs::remove_dir_all(&dir); + } + + #[test] + fn a_flood_of_readouts_at_their_limit_leaves_every_lease_request_answered() { + let dir = scratch("readout-flood"); + let quota = Arc::new(QuotaCell::new()); + // Fewer lease slots than flooding readers: the readers must not need them. + let (_service, lender) = lender("a, &dir, 2); + let socket = serve(lender, &dir); + let stop = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let flood: Vec<_> = (0..READOUT_CONNECTIONS) + .map(|_| { + let (socket, stop) = (socket.clone(), stop.clone()); + thread::spawn(move || { + while !stop.load(Ordering::Acquire) { + let _ = ask(&socket, &readout_request()); + } + }) + }) + .collect(); + // No caller is configured, so each lease is answered UnknownCaller: + // an answer is all that is asked. + let lease = LeaseMessage::ClaudeLeaseRequest(LeaseRequest { + expires_at_unix_ms: unix_ms() + 60_000, + }); + let unanswered = (0..40) + .filter(|_| { + !matches!( + ask(&socket, &lease), + Some(LeaseMessage::ClaudeLeaseRefusal(LeaseRefused { + kind: LeaseRefusalKind::UnknownCaller, + .. + })) + ) + }) + .count(); + stop.store(true, Ordering::Release); + for client in flood { + client.join().unwrap(); + } + assert_eq!(unanswered, 0); + let _ = std::fs::remove_dir_all(&dir); + } + + #[test] + fn a_readout_over_the_readout_limit_is_refused_busy_not_dropped() { + let dir = scratch("readout-limit"); + let quota = Arc::new(QuotaCell::new()); + let (_service, lender) = lender("a, &dir, 8); + let held: Vec<_> = (0..READOUT_CONNECTIONS) + .map(|_| Permit::acquire(&lender.readouts, READOUT_CONNECTIONS).unwrap()) + .collect(); + let (mut client, server) = UnixStream::pair().unwrap(); + client.write_all(&encode_frame(&readout_request())).unwrap(); + lender.serve_one(server).unwrap(); + let mut raw = Vec::new(); + client.read_to_end(&mut raw).unwrap(); + let Some(LeaseMessage::ClaudeLeaseRefusal(refusal)) = decode_frame(&raw) else { + panic!("no refusal"); + }; + assert_eq!(refusal.kind, LeaseRefusalKind::Busy); + drop(held); + let _ = std::fs::remove_dir_all(&dir); + } + + #[test] + fn a_readout_of_an_empty_metered_cell_does_not_wait() { + let quota = QuotaCell::new(); + quota.attach_meter(); + let started = Instant::now(); + let LeaseMessage::ClaudeUsageReadout(answer) = readout("a, 77) else { + panic!("a readout answers with a readout"); + }; + assert!(started.elapsed() < Duration::from_secs(1)); + assert!(answer.plan_quota.is_empty()); + assert_eq!(answer.at_unix_ms, 77); + } + + #[test] + fn a_readout_is_no_sign_of_life() { + let quota = Arc::new(QuotaCell::new()); + let within = Duration::from_secs(15 * 60); + let now = unix_ms(); + let _ = readout("a, now); + assert!(!quota.wanted_at(now, within)); + let _asking = quota.asking(); + assert!(quota.wanted_at(now, within)); + } + /// A refusal frame whose body is exactly `body_bytes` long. fn frame_of(body_bytes: usize) -> Vec { let frame_with = |pad: usize| { diff --git a/crates/claude-connector/src/lib.rs b/crates/claude-connector/src/lib.rs index 671ee36..88de37b 100644 --- a/crates/claude-connector/src/lib.rs +++ b/crates/claude-connector/src/lib.rs @@ -29,8 +29,8 @@ pub use launch::ChildEnv; #[cfg(target_os = "linux")] pub use lease::{ LeaseCaller, LeaseGranted, LeaseLimits, LeaseMessage, LeaseRefusalKind, LeaseRefused, - LeaseRequest, MAX_LEASE_FRAME_BYTES, PeerCredentials, PeerUnit, PeerUnits, decode_frame, - encode_frame, peer_credentials, serve_leases, user_uid, + LeaseRequest, LeaseUsageReadoutRequest, MAX_LEASE_FRAME_BYTES, PeerCredentials, PeerUnit, + PeerUnits, decode_frame, encode_frame, peer_credentials, serve_leases, user_uid, }; #[cfg(unix)] pub use process::StdProcessPort; diff --git a/crates/claude-connector/tests/fixtures/lease_v2/README.md b/crates/claude-connector/tests/fixtures/lease_v2/README.md index 3cc9766..4204337 100644 --- a/crates/claude-connector/tests/fixtures/lease_v2/README.md +++ b/crates/claude-connector/tests/fixtures/lease_v2/README.md @@ -11,6 +11,10 @@ big-endian length, then the named-MessagePack map `{kind, value}`). `seven_day` windows (utilization_permille, resets_at_unix_ms, status, observed_at_unix_ms). On the wire the grant travels with one SCM_RIGHTS fd; the file holds the frame only. +- `usage_readout_request.bin`: `claude_usage_readout_request`, an empty map; opens no lease. +- `usage_readout.bin`: `claude_usage_readout`, the kit's `UsageReadout` (`schema_id`, `at_unix_ms`, + `plan_quota` with `five_hour` and `seven_day` windows); no fd travels with it. + Huginn's lease client vendors these files from this repository and tests against the same bytes. Generator: `tests/lease_wire_golden.rs`; run it with `LEASE_V2_REGENERATE=1`, then review the diff: a changed byte is a wire change. diff --git a/crates/claude-connector/tests/fixtures/lease_v2/usage_readout.bin b/crates/claude-connector/tests/fixtures/lease_v2/usage_readout.bin new file mode 100644 index 0000000..3fe12c8 Binary files /dev/null and b/crates/claude-connector/tests/fixtures/lease_v2/usage_readout.bin differ diff --git a/crates/claude-connector/tests/fixtures/lease_v2/usage_readout_request.bin b/crates/claude-connector/tests/fixtures/lease_v2/usage_readout_request.bin new file mode 100644 index 0000000..f3fae6c Binary files /dev/null and b/crates/claude-connector/tests/fixtures/lease_v2/usage_readout_request.bin differ diff --git a/crates/claude-connector/tests/lease.rs b/crates/claude-connector/tests/lease.rs index 77c8295..9c6f27d 100644 --- a/crates/claude-connector/tests/lease.rs +++ b/crates/claude-connector/tests/lease.rs @@ -15,10 +15,10 @@ use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use claude_connector::{ LeaseCaller, LeaseGranted, LeaseLimits, LeaseMessage, LeaseRefusalKind, LeaseRefused, - LeaseRequest, PeerUnit, PeerUnits, QuotaCell, TokenSource, decode_frame, encode_frame, - serve_leases, + LeaseRequest, LeaseUsageReadoutRequest, PeerUnit, PeerUnits, QuotaCell, TokenSource, + decode_frame, encode_frame, serve_leases, }; -use model_connector::{CallerAdmission, LeaseTurn, PlanQuota, TransportService}; +use model_connector::{CallerAdmission, LeaseTurn, PlanQuota, TransportService, UsageReadout}; use zeroize::Zeroizing; const CALLER: &str = "eureka"; @@ -649,3 +649,80 @@ fn a_request_trickled_past_the_socket_timeout_is_dropped_without_a_lease() { assert!(started.elapsed() < Duration::from_secs(6)); assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 0); } + +/// What a usage readout request came back with. +fn read_usage(harness: &Harness) -> Answer { + let mut stream = UnixStream::connect(&harness.socket).unwrap(); + stream + .set_read_timeout(Some(Duration::from_secs(10))) + .unwrap(); + stream + .write_all(&encode_frame(&LeaseMessage::ClaudeUsageReadoutRequest( + LeaseUsageReadoutRequest {}, + ))) + .unwrap(); + let (raw, fds) = receive(&stream); + let message = decode_frame(&raw); + Answer { + _stream: stream, + raw, + message, + fds, + } +} + +fn usage_of(answer: &Answer) -> UsageReadout { + match &answer.message { + Some(LeaseMessage::ClaudeUsageReadout(readout)) => readout.clone(), + other => panic!("expected a usage readout, got {other:?}"), + } +} + +#[test] +fn a_readout_answers_the_cells_windows_and_lends_nothing() { + let harness = start(setup()); + let reset = now_ms() + 3_600_000; + harness.quota.observe([ + window("five_hour", 420, "allowed", reset), + window("seven_day", 999, "allowed_warning", reset + 86_400_000), + ]); + let before = now_ms(); + let answer = read_usage(&harness); + let readout = usage_of(&answer); + assert!(answer.fds.is_empty(), "a readout lends no descriptor"); + assert_eq!(readout.plan_quota, harness.quota.snapshot().windows); + assert_eq!(readout.plan_quota.len(), 2); + assert!(readout.at_unix_ms >= before && readout.at_unix_ms <= now_ms()); + assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 0); + // No lease was opened: the next request is granted. + assert!(granted_ok(&ask(&harness, 60_000))); +} + +#[test] +fn a_peer_outside_the_lease_callers_still_reads_usage() { + let harness = start(Setup { + uid: own_uid().wrapping_add(1), + ..setup() + }); + harness + .quota + .observe([window("five_hour", 420, "allowed", now_ms() + 3_600_000)]); + let refused = ask(&harness, 60_000); + assert_eq!(refusal_of(&refused), Some(LeaseRefusalKind::UnknownCaller)); + let readout = usage_of(&read_usage(&harness)); + assert_eq!(readout.plan_quota.len(), 1); + assert_eq!(readout.plan_quota[0].utilization_permille, Some(420)); + assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 0); +} + +#[test] +fn a_readout_beside_a_held_lease_leaves_it_held() { + let harness = start(setup()); + let held = ask(&harness, 60_000); + assert!(granted_ok(&held)); + assert_eq!(usage_of(&read_usage(&harness)).plan_quota.len(), 0); + harness.enter(2); + let second = ask(&harness, 60_000); + assert_eq!(refusal_of(&second), Some(LeaseRefusalKind::Busy)); + assert!(granted_ok(&held)); +} diff --git a/crates/claude-connector/tests/lease_wire_golden.rs b/crates/claude-connector/tests/lease_wire_golden.rs index 7b9cea2..5b51f8a 100644 --- a/crates/claude-connector/tests/lease_wire_golden.rs +++ b/crates/claude-connector/tests/lease_wire_golden.rs @@ -8,10 +8,10 @@ use std::path::{Path, PathBuf}; use claude_connector::{ - LeaseGranted, LeaseMessage, LeaseRefusalKind, LeaseRefused, LeaseRequest, decode_frame, - encode_frame, + LeaseGranted, LeaseMessage, LeaseRefusalKind, LeaseRefused, LeaseRequest, + LeaseUsageReadoutRequest, decode_frame, encode_frame, }; -use model_connector::PlanQuota; +use model_connector::{PlanQuota, PlanQuotaSnapshot, UsageReadout}; const REGENERATE: &str = "LEASE_V2_REGENERATE"; @@ -97,3 +97,38 @@ fn a_grant_with_both_usage_windows_is_pinned() { }), ); } + +#[test] +fn the_usage_readout_frames_are_pinned() { + let window = |name: &str, permille, status: &str, resets| PlanQuota { + window: name.to_string(), + utilization_permille: Some(permille), + resets_at_unix_ms: resets, + status: status.to_string(), + observed_at_unix_ms: 1_793_300_000_000, + }; + let request = LeaseMessage::ClaudeUsageReadoutRequest(LeaseUsageReadoutRequest {}); + let answer = LeaseMessage::ClaudeUsageReadout(UsageReadout::new( + 1_793_300_100_000, + PlanQuotaSnapshot { + windows: vec![ + window("five_hour", 420, "allowed", 1_793_307_000_000), + window("seven_day", 999, "allowed_warning", 1_793_800_000_000), + ], + }, + )); + assert_frame("usage_readout_request.bin", &request); + assert_frame("usage_readout.bin", &answer); + for (message, spelled) in [ + (&request, "claude_usage_readout_request"), + (&answer, "claude_usage_readout"), + ] { + let bytes = encode_frame(message); + let value: rmpv::Value = rmp_serde::from_slice(&bytes[4..]).unwrap(); + let rmpv::Value::Map(outer) = value else { + panic!("a frame is a map") + }; + assert_eq!(outer[0].0.as_str(), Some("kind")); + assert_eq!(outer[0].1.as_str(), Some(spelled)); + } +}