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..1749f6f 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(UsageReadoutRequest), + /// 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 UsageReadoutRequest {} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct LeaseGranted { pub lease_id: String, @@ -385,8 +397,11 @@ struct Lender { units: 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 +/// 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 @@ -421,21 +436,33 @@ 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) => { + let peer = peer_credentials(&stream)?; + self.lend(stream, peer, request) + } + LeaseMessage::ClaudeUsageReadoutRequest(_) => { + 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 +636,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 +662,30 @@ fn refuse(stream: &mut UnixStream, kind: LeaseRefusalKind, detail: &str) -> io:: mod tests { use super::*; + #[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..52f957f 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, MAX_LEASE_FRAME_BYTES, PeerCredentials, PeerUnit, PeerUnits, UsageReadoutRequest, + 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/lease.rs b/crates/claude-connector/tests/lease.rs index 77c8295..5115e8a 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, PeerUnit, PeerUnits, QuotaCell, TokenSource, UsageReadoutRequest, 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,90 @@ 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( + UsageReadoutRequest {}, + ))) + .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)); +} + +#[test] +fn a_readout_does_not_wait_for_the_meter() { + let harness = start(setup()); + harness.quota.attach_meter(); + let started = Instant::now(); + let readout = usage_of(&read_usage(&harness)); + assert!(started.elapsed() < Duration::from_secs(5)); + assert!(readout.plan_quota.is_empty()); +} diff --git a/crates/claude-connector/tests/lease_wire_golden.rs b/crates/claude-connector/tests/lease_wire_golden.rs index 7b9cea2..c713889 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, UsageReadoutRequest, + 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(UsageReadoutRequest {}); + 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)); + } +}