Compare commits

...
Sign in to create a new pull request.
Author SHA1 Message Date
Meta Vens
549962cb8a quota meter: the weekly reading comes from the CLI's own rate-limit event; while a lease is live and the reading is stale the connector makes the smallest real call. The usage poll and its HTTP client go: the setup token can never reach that endpoint (squash of eureka-body/quota-meter-stream 7d4c1bf..c298d9d; Soul s4-s5)
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-11-03 13:12:31 +01:00
Meta Vens
7d4c1bf56e usage-meter-poll: gate fixes: the timeout is read per request, serve has a refusal test, the RFC 850 guard is pinned
The timeout field was read only by a test (dead_code under -D warnings), so
ask() applies it to each request. The serve wrapper is reached by a
subprocess refusal test (no network); the serve test fails inside 30 s
instead of hanging; "Wednesday 21-Oct-26 ..." (no comma) is refused.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
2026-11-01 03:30:55 +01:00
Meta Vens
8fb8b5a71d usage-meter-poll: the serve test asks a loopback fake, never the provider (finding serve-test-asks-real-provider)
serve_config is the public wrapper over a crate-private serve_config_with
that takes the usage client; the test runs it in-process against a fake and
asserts the fake received GET /api/oauth/usage with the bearer. The endpoint
is not a document field or env, so whoever writes the config cannot redirect
the token. The subprocess test that reached api.anthropic.com is deleted.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
2026-11-01 03:17:08 +01:00
Meta Vens
f242290ca5 usage-meter-poll: pin what the production client is built with (finding new-client-timeout-unpinned)
UsageClient carries its timeout and url; UsageClient::production() is what
start() builds, and the test asserts that client's url and 10 s timeout as
spec r2 says, replacing a 1..=60 s bound no rule names.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
2026-11-01 03:17:08 +01:00
Meta Vens
481428a7ce usage-meter-poll: Retry-After reads all three HTTP-date forms (finding retry-after-obsolete-date-forms)
RFC 9110 10.2.3 makes a recipient accept IMF-fixdate, RFC 850 and asctime;
the last two read as absent and the poll asked again at the floor. The
two-digit RFC 850 year more than 50 years ahead is the latest past year.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
2026-11-01 03:17:08 +01:00
Meta Vens
e17ef16d48 usage-meter-poll: a subprocess test that serve --config starts the poll
Spawns the claude-connector binary on a ready document and reads the
usage poll state line from its stderr; the connector must also stay up.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
2026-11-01 02:49:16 +01:00
Meta Vens
c1cc349b19 usage-meter-poll: the cell decides how a reading combines; a poll is as old as its question
QuotaCell::observe_at is the one decision for both sensors: a reading older
than the held one is dropped, and a rejection stands against a reading that
only reports utilization (status polled) until its window resets or a reading
that carries its own verdict replaces it. The poll stamps its answer with the
moment it asked. Also: Retry-After in HTTP-date form is honoured; a refused
token is not sent again until it reads differently (the poll thread no longer
exits); the client's timeout is a parameter so a silent endpoint is tested;
tests pin redirects(0), the two-window mapping, the slow-server race, and the
late poll against a rejection.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
2026-11-01 02:49:16 +01:00
Meta Vens
6982cbca7e usage-meter-poll: pin the log lines and the thread loop
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
2026-11-01 01:44:35 +01:00
Meta Vens
7f3a46313e usage-meter-poll: test helper no longer shadowed by its own binding
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
2026-11-01 01:37:29 +01:00
Meta Vens
9143d2c5e2 usage-meter-poll: serve_config starts the poll beside the leases
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
2026-11-01 01:35:51 +01:00
Meta Vens
79792eea85 usage-meter-poll: the usage poll, the QuotaCell's second sensor
usage_poll asks the subscription's usage endpoint when the seven_day
reading is 15 min old or older, maps five_hour and seven_day onto
PlanQuota with status polled, and honours Retry-After; failures wait,
log once per change of state, and never touch the cell.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
2026-11-01 01:35:50 +01:00
Meta Vens
c1802d85be lease v2: committed byte goldens, rustfmt
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-31 05:38:48 +01:00
Meta Vens
cfea6854ab lease tests: the fake unit record is global, keep it on invocation 1
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-31 05:30:41 +01:00
Meta Vens
9f10887169 lease v2: byte goldens, hold-loop and lease-helper tests
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-31 05:27:38 +01:00
Meta Vens
9db62027eb connector-usage-only: comment
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-31 05:08:19 +01:00
Meta Vens
7d9a6392a7 connector-usage-only: comments
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-31 04:59:02 +01:00
Meta Vens
a828602649 connector-usage-only: drop a stray doc comment
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-31 04:54:14 +01:00
Meta Vens
bb8909f2c2 connector-usage-only: regenerate v4 goldens and schemas, format, fix tests
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-31 04:51:05 +01:00
Meta Vens
28f2e5d2e7 connector-usage-only: test import
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-31 04:49:29 +01:00
Meta Vens
dcb5f109f7 connector-usage-only: store-failure tests, README front door
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-31 04:48:24 +01:00
Meta Vens
af5da3dfaa connector-usage-only: test fixes
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-31 04:44:45 +01:00
Meta Vens
702e903fdd connector-usage-only: fix serve test splice
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-31 04:44:01 +01:00
Meta Vens
06a47bb13f connector-usage-only: delete the dollar authorities (WIP, does not yet build)
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-31 04:43:18 +01:00
Meta Vens
d2b82a6c50 lease-persona-callers: the identity bound keeps room for the Persona sequence
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-31 01:59:39 +01:00
Meta Vens
47a4dcc6b4 lease-persona-callers: model-connector lease callers take the turn
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-31 01:58:12 +01:00
Meta Vens
1bf9fdb437 lease-persona-callers: rustfmt
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-31 01:57:21 +01:00
Meta Vens
13cbdaf23c lease-persona-callers: a lease caller names its turn kind; Persona and Self of one caller lease side by side
LeaseTurn {SelfRun, PersonaTurn} on LeaseCaller and LeaseCallerConfig; the unit derives from the turn; Busy is keyed (caller, turn); a Persona invocation leases in sequence as lease:<invocation>:<n>, one open at a time. config add-lease-caller and init take --turn self|persona.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-31 01:56:34 +01:00
Meta Vens
fdcfa063c7 connector-many-callers: one connector serves every Mind as its own caller; session handles are an HMAC under a connector secret over caller, session, instance, instructions, model and transcript prefix; config add-caller/add-lease-caller; one owner for distinct admissions (squash of eureka-body/connector-many-callers d1f25b5..a96b276; Soul s1-s2)
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-30 00:11:53 +01:00
Meta Vens
d1f25b5ae6 claude-self-run-lease: Self runs borrow the one Claude token over a unix socket bound to the peer's uid, its pidfd and its unit's InvocationID; the connector sets the reservation, settle bills reported per-model usage uncapped, no report bills the reservation (squash of eureka-body/claude-self-run-lease 480cf94..b5f26ef; Soul s1-s3)
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-29 21:54:03 +01:00
480cf941c8 claude-turn-schema-quota: schema-typed Persona replies via --json-schema, token on fd 3 never in the child env, per-window plan-quota capture, communal spend readout, not-logged-in bills zero (squash of eureka-body/claude-turn-schema-quota e867fad..429b344; Soul s1, s2)
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-29 20:18:35 +01:00
e867fad379 B4b1 claude-connector-binary: connector daemon over one CultCache config document; token from a systemd credential, pin before every spawn, private prompt dir, 401/403/429 door set, shared readiness for check and serve (squash of eureka-body/claude-connector-binary 200933f..67c7849; Soul s1, s2)
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-29 12:59:23 +01:00
59 changed files with 7883 additions and 2533 deletions

2
.gitattributes vendored
View file

@ -1,2 +1,2 @@
schemas/*.json text eol=lf
crates/model-connector/tests/fixtures/wire_v3/* -text
crates/model-connector/tests/fixtures/wire_v4/* -text

5
Cargo.lock generated
View file

@ -138,9 +138,14 @@ dependencies = [
name = "claude-connector"
version = "0.1.0"
dependencies = [
"anyhow",
"chrono",
"cultcache-rs",
"hmac",
"libc",
"model-connector",
"rmp-serde",
"rmpv",
"serde",
"serde_json",
"sha2",

View file

@ -4,6 +4,11 @@ The connector kit for GameCult minds: one typed, encrypted, replay-safe contract
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
minds to read and decide their own burn. The one usage-derived refusal is the provider's own
(`ProviderRejected`), repeated until its reset.
Status: the `model-connector` crate and the `claude-connector` backend crate are live. No
connector binary exists yet (the Claude connector binary is first, then an OpenAI-compatible one).
@ -11,7 +16,7 @@ connector binary exists yet (the Claude connector binary is first, then an OpenA
Crate `crates/model-connector`:
- `src/contract.rs` - the v3 wire contract (`gamecult.model.*.v3`): `ProviderRequest`, `TransportInvocation`,
- `src/contract.rs` - the v4 wire contract (`gamecult.model.*.v4`): `ProviderRequest`, `TransportInvocation`,
`TransportResult`, `TransportReceipt`, the typed `StopReason`, `FailureKind` and `Refusal`,
`InputItem::ReasoningState` (a provider's opaque reasoning block, replayed by callers
verbatim via `TransportResult::transcript_items`; the service refuses another provider's
@ -21,9 +26,12 @@ Crate `crates/model-connector`:
- `src/envelope.rs` - the AES-256-GCM `TransportEnvelope` (binary ciphertext), its
identity-binding AAD, the length-prefixed framing, and the CultNet carriage: each frame is a
`cultnet.operation_request.v0` / `operation_response.v0` for service
`gamecult.model.connector`, operation `invoke`, its `message_id` the request id and its
payload the base64 of the sealed envelope.
- `schemas/` - one JSON Schema per v3 document, generated from the types.
`gamecult.model.connector`, operation `invoke` (or `usage_readout`), its `message_id` the
request id and its payload the base64 of the sealed envelope. `usage_readout` answers a
`UsageReadout` (`gamecult.model.usage_readout.v1`): the backend's plan quota windows, each with
utilization, reset, status and age;
`ConnectorClient::usage_readout` asks.
- `schemas/` - one JSON Schema per v4 document, generated from the types.
- `src/client.rs` - `ConnectorClient`, the one-request caller side.
- `src/service.rs` (feature `daemon`) - `TransportService`: caller admission, durable
replay (a retried request is answered byte-identically and the provider runs once),
@ -37,13 +45,16 @@ Crate `crates/claude-connector` (library only; the binary that wires it to `serv
- `ClaudeCodeBackend` implements `ProviderBackend` for provider `claude-code` by driving the
unmodified `claude` binary, one child per call, in print stream-json mode with no tools, no MCP
servers, no settings sources and the system prompt replaced. Tool-less turns only: a request
with tools, an output schema or a final tool result is refused before any spawn.
- Reasoning state is a session handle (`ClaudeCodeHandle`); a handle that cannot be proven to
match the request re-seeds from the caller's transcript, and every resume forks.
with tools or a final tool result is refused before any spawn. A request's output schema goes
to the CLI as `--json-schema` and the result's structured output is the one assistant text.
- Reasoning state is a session handle (`ClaudeCodeHandle`) authenticated with a connector secret
and bound to the caller it was minted for; a handle that cannot be proven to match the request
and its caller re-seeds from the caller's transcript, and every resume forks.
- `ProcessPort` is the one OS seam (`StdProcessPort`, unix, kills the child's whole process
group); `SessionStore` and `StderrSink` are the other ports. `ClaudeCodeConfig` is what the
binary hands in; the credential is read by name from the connector's own environment
(`CLAUDE_CODE_OAUTH_TOKEN`) and copied only into the child.
binary hands in; the credential is read from a systemd credential and reaches the child on a
pipe (`CLAUDE_CODE_OAUTH_TOKEN_FILE_DESCRIPTOR`), never in an environment. `QuotaCell` keeps the
plan windows the CLI's `rate_limit_event`s report, for the usage readout and the lease grant.
Frame arithmetic is public in `src/envelope.rs`, measured on the carriage frame: `frame_bytes_for_payload(n)` is the frame a
caller payload bound of `n` needs, `payload_bytes_for_frame(frame)` is the largest payload bound
@ -52,7 +63,7 @@ a frame allows, and `MAX_IDENTITY_BYTES` (256) bounds caller ids, request ids an
needed); the smallest legal frame is `frame_bytes_for_payload(MIN_PAYLOAD_BYTES)`. A request id past
the bound is refused on the wire as `Refusal::Malformed`, the reply echoing the id (the connection is dropped only if that reply cannot fit the frame); `ConnectorClient` refuses such an id locally with `RequestIdTooLong`.
The default feature set compiles contract, envelope and client only. The kit speaks v3
The default feature set compiles contract, envelope and client only. The kit speaks v4
only: a v2 frame gets no answer.
## Verify
@ -62,9 +73,9 @@ cargo test --locked --workspace
cargo test --locked --workspace --features model-connector/daemon
```
`crates/model-connector/tests/wire_golden_v3.rs` pins the v3 documents and one sealed carriage
frame to the fixtures in `crates/model-connector/tests/fixtures/wire_v3` (regenerate with
`WIRE_V3_REGENERATE=1`; see the fixtures' README), and
`crates/model-connector/tests/wire_golden_v4.rs` pins the v4 documents and one sealed carriage
frame to the fixtures in `crates/model-connector/tests/fixtures/wire_v4` (regenerate with
`WIRE_V4_REGENERATE=1`; see the fixtures' README), and
`crates/model-connector/tests/published_schemas.rs` fails when `schemas/` differs from the types
(`MODEL_CONNECTOR_REGENERATE_SCHEMAS=1` rewrites it).

View file

@ -7,7 +7,15 @@ license.workspace = true
repository.workspace = true
publish = false
[[bin]]
name = "claude-connector"
path = "src/main.rs"
[dependencies]
anyhow = "1"
chrono = { version = "0.4", default-features = false }
hmac = "0.12"
cultcache-rs = { git = "https://github.com/GameCult/CultLib.git", rev = "a4eae9056152335662f1115f97bf35fb00ab6d23", package = "cultcache-rs" }
model-connector = { path = "../model-connector", features = ["daemon"] }
rmp-serde = "1"
serde = { version = "1", features = ["derive"] }
@ -18,3 +26,6 @@ zeroize = "1"
[target.'cfg(unix)'.dependencies]
libc = "0.2"
[dev-dependencies]
rmpv = "1"

View file

@ -1,22 +1,99 @@
use std::fs::File;
use std::io;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Instant;
use model_connector::{
FailureKind, InputItem, ProviderBackend, ProviderCall, ProviderRequest, TransportEvent,
TransportEventPayload, TransportOutcome, TransportResult, Usage,
FailureKind, InputItem, PlanQuotaSnapshot, ProviderBackend, ProviderCall, ProviderRequest,
TransportEvent, TransportEventPayload, TransportOutcome, TransportResult, Usage,
};
use sha2::{Digest, Sha256};
use zeroize::Zeroizing;
use crate::config::{ClaudeCodeConfig, PROVIDER};
use crate::handle::{ClaudeCodeHandle, SessionStore, plan, prefix_digest, stdin_bytes};
use crate::launch::{PromptFile, TOKEN_VAR, argv, child_env};
use crate::handle::{
Binding, ClaudeCodeHandle, MIN_SECRET_BYTES, SessionStore, plan, prefix_digest, stdin_bytes,
};
use crate::launch::{PromptFile, argv, child_env};
use crate::process::{ProcessPort, RunningChild, StdoutRead};
use crate::quota::QuotaCell;
use crate::stream::{Failure, Finished, Step, StreamState};
/// How the backend reads its own environment.
type EnvLookup = Box<dyn Fn(&str) -> Option<String> + Send + Sync>;
/// The meter's call: the cheapest model, a trivial prompt, and an answer cap
/// that only bounds a runaway.
const METER_MODEL: &str = "claude-haiku-5-5";
const METER_CALLER: &str = "meter";
const METER_MAX_OUTPUT_TOKENS: u32 = 16;
/// Where the child's credential comes from, read afresh at every use. The
/// connector's own environment is never a source (ruling
/// claude-connector-token-credential).
pub trait TokenSource: Send + Sync {
/// The token, or `None` when there is none or it is empty.
fn read(&self) -> Option<Zeroizing<String>>;
/// What `config check` names when `read` has no token.
fn missing(&self) -> &'static str {
"token_credential"
}
}
/// A token in a file named by a systemd credential: the file
/// `$CREDENTIALS_DIRECTORY/<name>`, read on every call so a rotated credential
/// is picked up by the next spawn.
pub struct CredentialTokenSource {
directory: Option<PathBuf>,
name: String,
/// The configuration field `missing` names when the credential is absent.
field: &'static str,
}
impl CredentialTokenSource {
/// The credential `name` under the directory systemd names in
/// `CREDENTIALS_DIRECTORY`; with no such variable the source has no token.
/// The variable names a directory, never a token.
pub fn from_environment(name: impl Into<String>) -> Self {
Self {
directory: std::env::var_os("CREDENTIALS_DIRECTORY").map(PathBuf::from),
name: name.into(),
field: "token_credential",
}
}
/// The same source, naming `field` when its credential is absent.
pub fn for_field(mut self, field: &'static str) -> Self {
self.field = field;
self
}
pub fn in_directory(directory: impl Into<PathBuf>, name: impl Into<String>) -> Self {
Self {
directory: Some(directory.into()),
name: name.into(),
field: "token_credential",
}
}
}
impl TokenSource for CredentialTokenSource {
fn read(&self) -> Option<Zeroizing<String>> {
let directory = self.directory.as_ref()?;
let mut token = Zeroizing::new(std::fs::read_to_string(directory.join(&self.name)).ok()?);
let trimmed = token.trim_end_matches(['\r', '\n']).len();
token.truncate(trimmed);
(!token.is_empty()).then_some(token)
}
/// No directory at all is the unit's environment lacking its credentials,
/// which is not the same fault as a credential the directory lacks.
fn missing(&self) -> &'static str {
if self.directory.is_none() {
"credentials_directory"
} else {
self.field
}
}
}
/// Where the tail of a child's stderr goes. The backend never puts it in a
/// result; the connector binary wires this to its log.
@ -31,7 +108,16 @@ pub struct ClaudeCodeBackend {
port: Box<dyn ProcessPort>,
sessions: Box<dyn SessionStore>,
sink: Box<dyn StderrSink>,
env_lookup: EnvLookup,
tokens: Box<dyn TokenSource>,
handle_secret: Box<dyn TokenSource>,
quota: Arc<QuotaCell>,
}
/// The handle secret, or `None` when it is absent or shorter than a key.
fn usable_secret(source: &dyn TokenSource) -> Option<Zeroizing<String>> {
source
.read()
.filter(|secret| secret.len() >= MIN_SECRET_BYTES)
}
impl ClaudeCodeBackend {
@ -42,7 +128,14 @@ impl ClaudeCodeBackend {
let sessions = Box::new(crate::handle::FsSessionStore::new(
config.config_dir.clone(),
));
Self::with_ports(config, port, sessions, sink)
let tokens = Box::new(CredentialTokenSource::from_environment(
config.token_credential.clone(),
));
let handle_secret = Box::new(
CredentialTokenSource::from_environment(config.handle_secret_credential.clone())
.for_field("handle_secret_credential"),
);
Self::with_ports(config, port, sessions, sink, tokens, handle_secret)
}
pub fn with_ports(
@ -50,33 +143,74 @@ impl ClaudeCodeBackend {
port: Box<dyn ProcessPort>,
sessions: Box<dyn SessionStore>,
sink: Box<dyn StderrSink>,
tokens: Box<dyn TokenSource>,
handle_secret: Box<dyn TokenSource>,
) -> Self {
Self {
config,
port,
sessions,
sink,
env_lookup: Box::new(|name| std::env::var(name).ok()),
tokens,
handle_secret,
quota: Arc::new(QuotaCell::new()),
}
}
fn token(&self) -> Option<Zeroizing<String>> {
(self.env_lookup)(TOKEN_VAR)
.filter(|value| !value.is_empty())
.map(Zeroizing::new)
/// The plan windows every stream this backend reads feeds.
pub fn quota(&self) -> Arc<QuotaCell> {
self.quota.clone()
}
/// The meter's call: the smallest real one, through the same child,
/// stream and cell as any other. The stream's `rate_limit_event` is what
/// the caller wants; this returns what the call billed, or the failure's
/// kind (never its message).
pub(crate) fn read_meter(&self) -> Result<Usage, String> {
let mut request = ProviderRequest::new(
METER_CALLER,
METER_CALLER,
METER_MODEL,
"Reply with the single word ok.",
METER_MAX_OUTPUT_TOKENS,
);
request.input = vec![InputItem::UserText {
text: "ok?".to_string(),
}];
match self.run(&request, METER_CALLER, true).1 {
TransportOutcome::Completed { usage, .. } => Ok(usage),
TransportOutcome::Failed { kind, .. } => Err(format!("{kind:?}")),
}
}
/// The one pin check: `ready` and every spawn use it, and nothing caches
/// its answer.
fn binary_matches_pin(&self) -> bool {
let Ok(mut file) = File::open(&self.config.claude_binary) else {
return false;
};
let mut hasher = Sha256::new();
io::copy(&mut file, &mut hasher).is_ok()
&& <[u8; 32]>::from(hasher.finalize()) == self.config.claude_binary_sha256
binary_matches_pin(&self.config)
}
fn run(&self, request: &ProviderRequest) -> (Vec<TransportEvent>, TransportOutcome) {
let plan = match plan(request, &self.config, self.sessions.as_ref()) {
/// `ephemeral` leaves no transcript behind: the meter keeps no handle and
/// resumes nothing, every other call is resumed by its session.
fn run(
&self,
request: &ProviderRequest,
caller: &str,
ephemeral: bool,
) -> (Vec<TransportEvent>, TransportOutcome) {
let Some(secret) = usable_secret(self.handle_secret.as_ref()) else {
return failed(
FailureKind::Authentication,
"claude-code: the connector has no handle secret".to_string(),
Some(zero_usage()),
);
};
let plan = match plan(
request,
&self.config,
secret.as_bytes(),
caller,
self.sessions.as_ref(),
) {
Ok(plan) => plan,
Err(error) => {
return failed(
@ -86,21 +220,27 @@ impl ClaudeCodeBackend {
);
}
};
let Some(token) = self.token() else {
let Some(token) = self.tokens.read() else {
return failed(
FailureKind::Authentication,
"claude-code: the connector has no credential to give the child".to_string(),
Some(zero_usage()),
);
};
let Ok(prompt) = PromptFile::create(&request.instructions) else {
return failure(Failure::unbilled(
if !self.binary_matches_pin() {
return failure(Failure::never_reached(
FailureKind::Provider,
"claude-code: the claude binary does not match its pin",
));
}
let Ok(prompt) = PromptFile::create(&self.config.prompt_dir, &request.instructions) else {
return failure(Failure::never_reached(
FailureKind::Provider,
"claude-code: the system prompt could not be staged",
));
};
let argv = argv(&self.config, request, &plan, prompt.path());
let env = child_env(&self.config, request, token);
let argv = argv(&self.config, request, &plan, prompt.path(), ephemeral);
let env = child_env(&self.config, request);
let deadline = Instant::now() + self.config.limits.deadline;
let spawned = self.port.spawn(
@ -108,14 +248,21 @@ impl ClaudeCodeBackend {
&env,
&self.config.work_dir,
stdin_bytes(plan.blocks()),
&token,
);
let Ok(mut child) = spawned else {
return failure(Failure::unbilled(
return failure(Failure::never_reached(
FailureKind::Provider,
"claude-code: the child could not be spawned",
));
};
let finished = read_stream(child.as_mut(), deadline, &request.model);
let finished = read_stream(
child.as_mut(),
deadline,
&request.model,
request.output_schema_json.as_deref(),
&self.quota,
);
child.kill();
child.wait();
let tail = child.stderr_tail();
@ -124,7 +271,7 @@ impl ClaudeCodeBackend {
}
match finished {
Ok(finished) => self.completed(request, finished),
Ok(finished) => self.completed(request, caller, secret.as_bytes(), finished),
Err(failure_) => failure(failure_),
}
}
@ -132,6 +279,8 @@ impl ClaudeCodeBackend {
fn completed(
&self,
request: &ProviderRequest,
caller: &str,
secret: &[u8],
finished: Finished,
) -> (Vec<TransportEvent>, TransportOutcome) {
let mut payloads: Vec<TransportEventPayload> = finished
@ -146,14 +295,30 @@ impl ClaudeCodeBackend {
text: finished.texts.concat(),
});
}
let handle = ClaudeCodeHandle::Ended {
session_id: finished.session_id,
instance: self.config.instance.clone(),
prefix_sha256: prefix_digest(&request.instructions, &request.model, &prefix),
};
payloads.push(TransportEventPayload::ReasoningState {
state: handle.encode(),
});
// A schema turn leaves the CLI's StructuredOutput tool exchange in its
// session, and a later turn without the schema runs with no tools: it
// must not resume that session. So a schema turn hands out no handle,
// and whatever follows it is seeded from the caller's transcript.
if request.output_schema_json.is_none() {
let handle = ClaudeCodeHandle::Ended {
session_id: finished.session_id.clone(),
instance: self.config.instance.clone(),
prefix_sha256: prefix_digest(
secret,
&Binding {
caller,
instructions: &request.instructions,
model: &request.model,
session_id: &finished.session_id,
instance: &self.config.instance,
items: &prefix,
},
),
};
payloads.push(TransportEventPayload::ReasoningState {
state: handle.encode(),
});
}
let events = payloads
.into_iter()
@ -175,11 +340,12 @@ impl ClaudeCodeBackend {
}
impl ProviderBackend for ClaudeCodeBackend {
fn ready(&self) -> bool {
self.config.work_dir.is_dir()
&& self.config.config_dir.is_dir()
&& self.token().is_some()
&& self.binary_matches_pin()
fn ready(&self) -> Result<(), &'static str> {
readiness(
&self.config,
self.tokens.as_ref(),
self.handle_secret.as_ref(),
)
}
fn provider(&self) -> &str {
@ -187,17 +353,23 @@ impl ProviderBackend for ClaudeCodeBackend {
}
fn execute(&self, call: &ProviderCall) -> TransportResult {
let (events, outcome) = self.run(call.request());
let (events, outcome) = self.run(call.request(), call.caller_runtime_id(), false);
call.transported(events, outcome)
}
fn plan_quota(&self) -> PlanQuotaSnapshot {
self.quota.snapshot()
}
}
fn read_stream(
child: &mut dyn RunningChild,
deadline: Instant,
model: &str,
schema: Option<&str>,
quota: &QuotaCell,
) -> Result<Finished, Failure> {
let mut state = StreamState::new(model);
let mut state = StreamState::new(model, schema, quota);
let finished = loop {
match child.read_line(deadline) {
StdoutRead::Line(line) => {
@ -206,19 +378,19 @@ fn read_stream(
}
}
StdoutRead::Eof => {
return Err(Failure::unbilled(
return Err(Failure::cost_unknown(
FailureKind::Provider,
"claude-code: the stream ended without a result",
));
}
StdoutRead::Timeout => {
return Err(Failure::unbilled(
return Err(Failure::cost_unknown(
FailureKind::Timeout,
"claude-code: the deadline passed before a result",
));
}
StdoutRead::Breach => {
return Err(Failure::unbilled(
return Err(Failure::cost_unknown(
FailureKind::Provider,
"claude-code: the output passed its size cap",
));
@ -230,6 +402,45 @@ fn read_stream(
finished
}
fn binary_matches_pin(config: &ClaudeCodeConfig) -> bool {
let Ok(mut file) = File::open(&config.claude_binary) else {
return false;
};
let mut hasher = Sha256::new();
io::copy(&mut file, &mut hasher).is_ok()
&& <[u8; 32]>::from(hasher.finalize()) == config.claude_binary_sha256
}
/// Whether the backend can take work, naming the first config field that
/// cannot be used. `config check` and `serve` both ask this.
pub(crate) fn readiness(
config: &ClaudeCodeConfig,
tokens: &dyn TokenSource,
handle_secret: &dyn TokenSource,
) -> Result<(), &'static str> {
if !config.work_dir.is_dir() {
return Err("work_dir");
}
if !config.config_dir.is_dir() {
return Err("config_dir");
}
if !config.prompt_dir.is_dir() {
return Err("prompt_dir");
}
if !binary_matches_pin(config) {
return Err("claude_binary_sha256");
}
// Last: outside the unit the credential is the one thing always absent,
// and it must not hide the faults the document itself has.
if usable_secret(handle_secret).is_none() {
return Err(handle_secret.missing());
}
if tokens.read().is_none() {
return Err(tokens.missing());
}
Ok(())
}
fn zero_usage() -> Usage {
Usage {
input_tokens: 0,

File diff suppressed because it is too large Load diff

View file

@ -19,6 +19,15 @@ pub struct ClaudeCodeConfig {
/// Passed to the child as `CLAUDE_CONFIG_DIR` and `HOME`; it holds the
/// CLI's session store.
pub config_dir: PathBuf,
/// A directory the connector owns (mode 0700), where each call's system
/// prompt is staged for the child.
pub prompt_dir: PathBuf,
/// The name of the systemd credential that holds the child's token, under
/// `$CREDENTIALS_DIRECTORY`.
pub token_credential: String,
/// The name of the systemd credential that holds the handle secret: the
/// key every handle is authenticated with.
pub handle_secret_credential: String,
/// The connector instance id written into every handle.
pub instance: String,
pub limits: ClaudeCodeLimits,

View file

@ -0,0 +1,830 @@
//! The claude-connector's configuration document and its serve verb. One typed
//! CultCache record owns every bound and path the backend and the serve loop
//! run with; nothing in the connector's environment, flags or JSON can change
//! them. Validation errors name a field and never echo a value.
use std::collections::BTreeSet;
use std::net::{SocketAddr, TcpListener};
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use cultcache_rs::{CultCache, DatabaseEntry, SingleFileMessagePackBackingStore};
use model_connector::{
CallerAdmission, CallerConfig, LeaseTurn, MIN_PAYLOAD_BYTES, 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;
#[cfg(unix)]
use crate::meter;
use crate::{
ClaudeCodeBackend, ClaudeCodeConfig, ClaudeCodeLimits, CredentialTokenSource, StderrSink,
};
const CONFIG_KEY: &str = "runtime";
const MODEL_PREFIX: &str = "claude-";
/// A unix account that may lease the subscription token for a Self run or a
/// Persona turn, the unit its process must run in (`mind-self@<unit_instance>`
/// or `mind-persona@<unit_instance>`, by `turn`), and the configured caller
/// whose name its leases carry.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct LeaseCallerConfig {
pub unix_user: String,
pub caller_id: String,
/// The kind of turn this account leases for: it fixes the unit, `mind-self@`
/// or `mind-persona@`, its process must run in.
pub turn: LeaseTurn,
pub unit_instance: String,
}
#[derive(Clone, Debug, PartialEq, Eq, DatabaseEntry)]
#[cultcache(type = "gamecult.model_connector.claude_connector_config.v2")]
pub struct ClaudeConnectorConfig {
/// A loopback socket address.
#[cultcache(key = 0)]
pub listen: String,
#[cultcache(key = 1)]
pub replay_store: PathBuf,
#[cultcache(key = 2)]
pub claude_binary: PathBuf,
#[cultcache(key = 3)]
pub claude_binary_sha256: [u8; 32],
#[cultcache(key = 4)]
pub work_dir: PathBuf,
#[cultcache(key = 5)]
pub config_dir: PathBuf,
/// Where each call's system prompt is staged: a directory the connector
/// owns, mode 0700.
#[cultcache(key = 6)]
pub prompt_dir: PathBuf,
#[cultcache(key = 7)]
pub instance: String,
/// The name of the systemd credential that holds the child's token.
#[cultcache(key = 8)]
pub token_credential: String,
#[cultcache(key = 9)]
pub deadline_ms: u64,
#[cultcache(key = 10)]
pub max_line_bytes: usize,
#[cultcache(key = 11)]
pub max_total_bytes: usize,
#[cultcache(key = 12)]
pub stderr_tail_bytes: usize,
#[cultcache(key = 13)]
pub max_frame_bytes: usize,
#[cultcache(key = 14)]
pub max_connections: usize,
#[cultcache(key = 15)]
pub socket_timeout_ms: u64,
#[cultcache(key = 16)]
pub max_expiry_skew_ms: u64,
/// Every caller of this connector: its own connection key and epoch, models,
/// concurrency, payload and output bounds.
/// Ids are unique. Handles are bound to the caller that was given them.
#[cultcache(key = 17)]
pub callers: Vec<CallerConfig>,
/// 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,
/// The name of the systemd credential that holds the handle secret (at
/// least 32 bytes), the key every handle is authenticated with.
#[cultcache(key = 22)]
pub handle_secret_credential: String,
}
/// A refusal. It names the field or stage that failed and nothing the
/// operator or the file typed.
#[derive(Debug, Error)]
pub enum ConfigError {
#[error("invalid configuration field {0}")]
Invalid(&'static str),
#[error("the configuration store could not be read as a claude connector document")]
Store,
#[error("the prompt directory is not a directory owned by this process with mode 0700")]
PromptDir,
#[error("the listener could not be bound")]
Listen,
/// Something the document names exists but cannot be used now.
#[error("the configured {0} cannot be used")]
Unusable(&'static str),
/// The serve loop failed after it started.
#[error("the serve loop failed at {0}")]
Serving(&'static str),
}
/// The kit's serve errors carry stored paths, ids and numbers; none of that
/// text is shown. Each maps to the configuration field or stage it names.
impl From<ServeError> for ConfigError {
fn from(error: ServeError) -> Self {
match error {
ServeError::InvalidConfig(field) | ServeError::BackendUnready(field) => {
Self::Unusable(field)
}
ServeError::FrameBound { .. } => Self::Invalid("max_frame_bytes"),
ServeError::ConnectionKey(_) => Self::Unusable("connection_key_file"),
ServeError::Service(error) => match error {
ServiceError::ReplayStore(_) => Self::Unusable("replay_store"),
ServiceError::InvalidAdmission
| ServiceError::DuplicateCaller
| ServiceError::SharedCallerKey => Self::Invalid("callers"),
_ => Self::Serving("service"),
},
ServeError::ServicePoisoned => Self::Serving("service"),
ServeError::Listen(_) => Self::Serving("listener"),
ServeError::Thread(_) | ServeError::Connection(_) => Self::Serving("connection"),
ServeError::FrameSize | ServeError::FrameEncoding => Self::Serving("frame"),
ServeError::Clock => Self::Serving("clock"),
}
}
}
impl ClaudeConnectorConfig {
/// The one validation: `config init`, `config check` and `serve` all pass
/// through it. Returns the listen address.
pub fn validate(&self) -> Result<SocketAddr, ConfigError> {
let listen = self
.listen
.parse::<SocketAddr>()
.ok()
.filter(|address| address.ip().is_loopback())
.ok_or(ConfigError::Invalid("listen"))?;
let paths = [
("replay_store", &self.replay_store),
("claude_binary", &self.claude_binary),
("work_dir", &self.work_dir),
("config_dir", &self.config_dir),
("prompt_dir", &self.prompt_dir),
];
for (field, path) in paths {
if !path.is_absolute() {
return Err(ConfigError::Invalid(field));
}
}
if self.claude_binary_sha256 == [0; 32] {
return Err(ConfigError::Invalid("claude_binary_sha256"));
}
if !is_trimmed_nonempty(&self.instance) {
return Err(ConfigError::Invalid("instance"));
}
if !is_credential_name(&self.token_credential) {
return Err(ConfigError::Invalid("token_credential"));
}
if !is_credential_name(&self.handle_secret_credential)
|| self.handle_secret_credential == self.token_credential
{
return Err(ConfigError::Invalid("handle_secret_credential"));
}
if self.deadline_ms == 0
|| self.max_line_bytes == 0
|| self.max_total_bytes < self.max_line_bytes
|| self.stderr_tail_bytes == 0
{
return Err(ConfigError::Invalid("child limits"));
}
if self.max_frame_bytes == 0
|| self.max_connections == 0
|| self.socket_timeout_ms == 0
|| self.max_expiry_skew_ms == 0
{
return Err(ConfigError::Invalid("serve limits"));
}
if self.callers.is_empty() {
return Err(ConfigError::Invalid("callers"));
}
let mut ids = BTreeSet::new();
for caller in &self.callers {
if !is_trimmed_nonempty(&caller.caller_runtime_id) {
return Err(ConfigError::Invalid("caller_runtime_id"));
}
if !ids.insert(caller.caller_runtime_id.as_str()) {
return Err(ConfigError::Invalid("callers"));
}
if !caller.connection_key_file.is_absolute() {
return Err(ConfigError::Invalid("connection_key_file"));
}
if caller.connection_key_epoch == 0 {
return Err(ConfigError::Invalid("connection_key_epoch"));
}
if caller.max_concurrent_requests == 0
|| caller.max_payload_bytes < MIN_PAYLOAD_BYTES
|| caller.max_output_tokens == 0
{
return Err(ConfigError::Invalid("caller bounds"));
}
if caller.allowed_models.is_empty() {
return Err(ConfigError::Invalid("allowed_models"));
}
let mut seen = BTreeSet::new();
for model in &caller.allowed_models {
if !has_model_id_form(model) || !seen.insert(model.as_str()) {
return Err(ConfigError::Invalid("allowed_models"));
}
}
}
self.validate_leases()?;
Ok(listen)
}
fn validate_leases(&self) -> 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)
&& self
.callers
.iter()
.any(|caller| caller.caller_runtime_id == lease.caller_id)
&& is_unit_instance(&lease.unit_instance)
&& 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(),
claude_binary_sha256: self.claude_binary_sha256,
work_dir: self.work_dir.clone(),
config_dir: self.config_dir.clone(),
prompt_dir: self.prompt_dir.clone(),
token_credential: self.token_credential.clone(),
handle_secret_credential: self.handle_secret_credential.clone(),
instance: self.instance.clone(),
limits: ClaudeCodeLimits {
deadline: Duration::from_millis(self.deadline_ms),
max_line_bytes: self.max_line_bytes,
max_total_bytes: self.max_total_bytes,
stderr_tail_bytes: self.stderr_tail_bytes,
},
}
}
fn serve_limits(&self) -> ServeLimits {
ServeLimits {
max_frame_bytes: self.max_frame_bytes,
max_connections: self.max_connections,
socket_timeout_ms: self.socket_timeout_ms,
max_expiry_skew_ms: self.max_expiry_skew_ms,
}
}
}
/// 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
}
/// The loose shape of a model id, only to catch typos:
/// `claude-<family>-<major>`, then an optional `-<minor>` and an optional
/// `-YYYYMMDD`, the family lowercase letters. No shape can tell an alias from
/// a canonical id, so the known list is `allowed_models` itself; an alias that slips in fails its first
/// call, when the stream's `init.model` differs from the request.
fn has_model_id_form(model: &str) -> bool {
let Some(rest) = model.strip_prefix(MODEL_PREFIX) else {
return false;
};
let digits = |part: &str, most: usize| {
(1..=most).contains(&part.len()) && part.bytes().all(|byte| byte.is_ascii_digit())
};
let date = |part: &str| part.len() == 8 && digits(part, 8);
let mut parts = rest.split('-');
let family = parts.next().unwrap_or("");
let tail: Vec<&str> = parts.collect();
let Some((major, after)) = tail.split_first() else {
return false;
};
let after_ok = match after {
[] => true,
[one] => digits(one, 2) || date(one),
[minor, last] => digits(minor, 2) && date(last),
_ => false,
};
!family.is_empty()
&& family.bytes().all(|byte| byte.is_ascii_lowercase())
&& digits(major, 2)
&& after_ok
}
/// One file name under `$CREDENTIALS_DIRECTORY`, never a path.
fn is_credential_name(name: &str) -> bool {
!name.is_empty()
&& name != "."
&& name != ".."
&& name
.bytes()
.all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'.' | b'_'))
}
/// Reads the document without writing: a shared snapshot of exactly one
/// record, with the exact key, type and schema id, in canonical positional
/// MessagePack, then validated.
pub fn load_config(path: &Path) -> Result<ClaudeConnectorConfig, ConfigError> {
let store = SingleFileMessagePackBackingStore::new(path);
let config = store
.with_read_only_shared_snapshot(|entries| {
let [envelope] = entries.as_slice() else {
anyhow::bail!("not exactly one record");
};
if envelope.key != CONFIG_KEY
|| envelope.r#type != ClaudeConnectorConfig::TYPE
|| envelope.schema_id.as_deref() != Some(ClaudeConnectorConfig::TYPE)
{
anyhow::bail!("wrong typed envelope");
}
let config: ClaudeConnectorConfig = rmp_serde::from_slice(&envelope.payload)?;
if rmp_serde::to_vec(&config)? != envelope.payload {
anyhow::bail!("not canonical positional MessagePack");
}
Ok(config)
})
.map_err(|_| ConfigError::Store)?;
config.validate()?;
Ok(config)
}
/// Validates, then writes the document as the store's only record.
pub fn write_config(path: &Path, config: &ClaudeConnectorConfig) -> Result<(), ConfigError> {
config.validate()?;
let mut cache = CultCache::new();
cache
.register_entry_type::<ClaudeConnectorConfig>()
.map_err(|_| ConfigError::Store)?;
cache
.add_generic_backing_store(SingleFileMessagePackBackingStore::new(path))
.map_err(|_| ConfigError::Store)?;
cache
.pull_all_backing_stores()
.map_err(|_| ConfigError::Store)?;
cache
.put(CONFIG_KEY, config)
.map_err(|_| ConfigError::Store)?;
Ok(())
}
/// The prompt directory holds every system prompt in flight: it must be a
/// real directory (not a link) that this process owns and nobody else can
/// enter.
#[cfg(unix)]
pub fn check_prompt_dir(dir: &Path) -> Result<(), ConfigError> {
use std::os::unix::fs::{MetadataExt, PermissionsExt};
let metadata = std::fs::symlink_metadata(dir).map_err(|_| ConfigError::PromptDir)?;
// SAFETY: geteuid has no preconditions and cannot fail.
let uid = unsafe { libc::geteuid() };
if !metadata.is_dir()
|| metadata.uid() != uid
|| metadata.permissions().mode() & 0o7777 != 0o700
{
return Err(ConfigError::PromptDir);
}
Ok(())
}
/// The child's stderr is the log's business by size only: its bytes are never
/// written anywhere the connector controls.
struct LogSink<W>(std::sync::Mutex<W>);
impl<W: std::io::Write + Send> StderrSink for LogSink<W> {
fn stderr_tail(&self, tail: &[u8]) {
if let Ok(mut log) = self.0.lock() {
let _ = writeln!(log, "claude-connector: child stderr, {} bytes", tail.len());
}
}
}
/// Everything `serve` needs true before it binds, in one place so that
/// `config check` passes exactly what `serve` will start: the prompt
/// directory, the other directories, the credential, the pin, the connection
/// key and the frame and connection bounds. Returns the admitted callers.
fn readiness_of(config: &ClaudeConnectorConfig) -> Result<Vec<CallerAdmission>, ConfigError> {
config.validate()?;
#[cfg(unix)]
check_prompt_dir(&config.prompt_dir)?;
let admissions = load_admissions(&config.callers)?;
config.serve_limits().validate(&admissions)?;
let backend = config.backend_config();
readiness(
&backend,
&CredentialTokenSource::from_environment(backend.token_credential.clone()),
&CredentialTokenSource::from_environment(backend.handle_secret_credential.clone())
.for_field("handle_secret_credential"),
)
.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(),
turn: lease.turn,
unit_instance: lease.unit_instance.clone(),
})
.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>>,
quota: Arc<crate::QuotaCell>,
) -> 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,
quota,
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)
}
/// Loads the document and serves until the listener fails.
#[cfg(unix)]
pub fn serve_config(path: &Path) -> Result<(), ConfigError> {
serve_config_with(path, |_| ())
}
/// [`serve_config`] handing the quota cell to `started` once the meter runs:
/// a test stands in for the swarm's lease requests only through here.
#[cfg(unix)]
fn serve_config_with(
path: &Path,
started: impl FnOnce(&Arc<crate::QuotaCell>),
) -> Result<(), ConfigError> {
let config = load_config(path)?;
let admissions = readiness_of(&config)?;
let listen = config.validate()?;
let backend = Arc::new(ClaudeCodeBackend::new(
config.backend_config(),
Box::new(LogSink(std::sync::Mutex::new(std::io::stderr()))),
));
let listener = TcpListener::bind(listen).map_err(|_| ConfigError::Listen)?;
let service = Arc::new(Mutex::new(
TransportService::open(
&config.replay_store,
backend.provider(),
admissions,
config.max_expiry_skew_ms,
)
.map_err(|error| ConfigError::from(ServeError::Service(error)))?,
));
// The meter is running before the first lease can ask, so that lease can
// wait for its first read.
let quota = backend.quota();
meter::start(quota.clone(), Box::new(backend.clone()))
.map_err(|_| ConfigError::Serving("meter"))?;
started(&quota);
#[cfg(target_os = "linux")]
if let Some(socket) = &config.lease_socket {
start_leases(&config, socket, service.clone(), quota.clone())?;
}
serve_service(listener, config.serve_limits(), service, backend)?;
Ok(())
}
#[cfg(test)]
mod tests {
use std::sync::{Arc, Mutex};
use sha2::Digest;
use super::*;
/// A log that a test can read back.
#[derive(Clone, Default)]
struct Captured(Arc<Mutex<Vec<u8>>>);
impl std::io::Write for Captured {
fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
self.0.lock().unwrap().extend_from_slice(bytes);
Ok(bytes.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
/// `serve` starts the meter beside the listener, wired to the backend's own
/// cell: it makes no call while nobody asks, and once a lease is asked for
/// it runs the claude binary and serve stays up. The binary is a script
/// that replays a real captured call.
#[cfg(unix)]
#[test]
fn serve_config_starts_the_meter_and_stays_up() {
use std::os::unix::fs::PermissionsExt;
const CALL: &str = include_str!("../tests/fixtures/claude-2.1.312-meter-call.jsonl");
let dir =
std::env::temp_dir().join(format!("claude-connector-serve-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
let at = |name: &str| dir.join(name);
for sub in ["work", "home", "prompts", "credentials"] {
std::fs::create_dir_all(at(sub)).unwrap();
}
std::fs::set_permissions(at("prompts"), std::fs::Permissions::from_mode(0o700)).unwrap();
std::fs::write(at("call.jsonl"), CALL).unwrap();
let script = format!(
"#!/bin/sh\ncat > /dev/null\necho spawned >> '{}'\ncat '{}'\n",
at("spawned").display(),
at("call.jsonl").display()
);
std::fs::write(at("claude"), &script).unwrap();
std::fs::set_permissions(at("claude"), std::fs::Permissions::from_mode(0o755)).unwrap();
std::fs::write(at("caller.key"), "k".repeat(48)).unwrap();
std::fs::write(at("credentials/claude-token"), "sk-ant-oat01-fake-token").unwrap();
std::fs::write(at("credentials/claude-handle-secret"), "s".repeat(48)).unwrap();
let config = ClaudeConnectorConfig {
listen: "127.0.0.1:0".to_string(),
replay_store: at("replay.cc"),
claude_binary: at("claude"),
claude_binary_sha256: sha2::Sha256::digest(script.as_bytes()).into(),
work_dir: at("work"),
config_dir: at("home"),
prompt_dir: at("prompts"),
instance: "persona-a".to_string(),
token_credential: "claude-token".to_string(),
handle_secret_credential: "claude-handle-secret".to_string(),
deadline_ms: 600_000,
max_line_bytes: 4 << 20,
max_total_bytes: 64 << 20,
stderr_tail_bytes: 8192,
max_frame_bytes: 2 << 20,
max_connections: 8,
socket_timeout_ms: 10_000,
max_expiry_skew_ms: 300_000,
callers: vec![CallerConfig {
caller_runtime_id: "persona-a-caller".to_string(),
connection_key_file: at("caller.key"),
connection_key_epoch: 1,
allowed_models: vec!["claude-haiku-5-5".to_string()],
max_concurrent_requests: 4,
max_payload_bytes: 1 << 20,
max_output_tokens: 32_768,
}],
lease_socket: None,
lease_callers: Vec::new(),
lease_max_ms: 3_600_000,
};
write_config(&at("connector.cc"), &config).unwrap();
// SAFETY: set once, to a directory only this test reads, before the
// serve thread it starts.
unsafe { std::env::set_var("CREDENTIALS_DIRECTORY", at("credentials")) };
let path = at("connector.cc");
let lease = std::sync::Arc::new(std::sync::Mutex::new(None));
let swarm = lease.clone();
let served = std::thread::spawn(move || {
serve_config_with(&path, |cell| {
*swarm.lock().unwrap() = Some(cell.clone());
})
});
let deadline = std::time::Instant::now() + Duration::from_secs(30);
let cell = loop {
if let Some(cell) = lease.lock().unwrap().clone() {
break cell;
}
assert!(std::time::Instant::now() < deadline, "serve never started");
std::thread::sleep(Duration::from_millis(20));
};
std::thread::sleep(Duration::from_millis(500));
assert!(!at("spawned").exists(), "a quiet swarm spends nothing");
let _lease = cell.asking();
while !at("spawned").exists() {
assert!(
std::time::Instant::now() < deadline,
"the meter never ran the claude binary"
);
std::thread::sleep(Duration::from_millis(50));
}
assert!(!served.is_finished(), "serve returned after starting the meter");
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn serve_errors_name_a_field_or_stage_and_never_a_stored_value() {
const CANARY: &str = "canary-path-8e2a";
let io = || std::io::Error::other(CANARY);
let cases: Vec<(ServeError, &str)> = vec![
(ServeError::InvalidConfig("callers"), "callers"),
(ServeError::BackendUnready("work_dir"), "work_dir"),
(
ServeError::FrameBound {
max_frame_bytes: 4099,
payload_bytes: 4096,
needed_frame_bytes: 9999,
},
"max_frame_bytes",
),
(ServeError::ConnectionKey(io()), "connection_key_file"),
(
ServeError::Service(ServiceError::ReplayStore(format!("{CANARY}/replay.cc"))),
"replay_store",
),
(
ServeError::Service(ServiceError::InvalidAdmission),
"callers",
),
(
ServeError::Service(ServiceError::DuplicateCaller),
"callers",
),
(
ServeError::Service(ServiceError::SharedCallerKey),
"callers",
),
(ServeError::Listen(io()), "listener"),
];
for (error, field) in cases {
let shown = ConfigError::from(error);
let text = format!("{shown} {shown:?}");
assert!(text.contains(field), "{text}");
for stored in [CANARY, "4099", "9999", "replay.cc"] {
assert!(!text.contains(stored), "{text}");
}
}
}
#[test]
fn the_log_gets_the_size_of_the_childs_stderr_and_never_its_bytes() {
let log = Captured::default();
let sink = LogSink(Mutex::new(log.clone()));
sink.stderr_tail(b"canary-stderr-text, 22 bytes");
let written = String::from_utf8(log.0.lock().unwrap().clone()).unwrap();
assert_eq!(
written,
"claude-connector: child stderr, 28 bytes
"
);
assert!(!written.contains("canary"));
}
#[cfg(unix)]
mod prompt_dir {
use std::os::unix::fs::{PermissionsExt, chown, symlink};
use std::path::PathBuf;
use super::super::check_prompt_dir;
struct Dir(PathBuf);
impl Dir {
fn new(name: &str) -> Self {
let path = std::env::temp_dir()
.join(format!("claude-prompt-dir-{}-{name}", std::process::id()));
let _ = std::fs::remove_dir_all(&path);
std::fs::create_dir_all(&path).unwrap();
Self(path)
}
fn mode(&self, mode: u32) {
std::fs::set_permissions(&self.0, std::fs::Permissions::from_mode(mode)).unwrap();
}
}
impl Drop for Dir {
fn drop(&mut self) {
let _ = std::fs::set_permissions(&self.0, std::fs::Permissions::from_mode(0o700));
let _ = std::fs::remove_dir_all(&self.0);
}
}
#[test]
fn only_mode_0700_passes() {
let dir = Dir::new("modes");
dir.mode(0o700);
assert!(check_prompt_dir(&dir.0).is_ok());
// Each of these leaves the owner able to enter, or adds a bit
// beyond the owner's three; all are refused.
for mode in [
0o500, 0o300, 0o100, 0o600, 0o770, 0o707, 0o750, 0o755, 0o1700, 0o2700, 0o4700,
] {
dir.mode(mode);
assert!(check_prompt_dir(&dir.0).is_err(), "{mode:o}");
}
}
#[test]
fn a_link_to_a_good_directory_is_refused() {
let dir = Dir::new("link-target");
dir.mode(0o700);
let link =
std::env::temp_dir().join(format!("claude-prompt-dir-{}-link", std::process::id()));
let _ = std::fs::remove_file(&link);
symlink(&dir.0, &link).unwrap();
let refused = check_prompt_dir(&link).is_err();
std::fs::remove_file(&link).unwrap();
assert!(refused);
}
#[test]
fn a_directory_owned_by_another_user_is_refused() {
// SAFETY: geteuid has no preconditions and cannot fail.
if unsafe { libc::geteuid() } != 0 {
eprintln!("skipped: only root can give a directory away");
return;
}
let dir = Dir::new("owner");
dir.mode(0o700);
assert!(check_prompt_dir(&dir.0).is_ok());
chown(&dir.0, Some(1234), None).unwrap();
assert!(check_prompt_dir(&dir.0).is_err());
}
}
}

View file

@ -5,9 +5,10 @@
use std::path::PathBuf;
use hmac::{Hmac, Mac};
use model_connector::{InputItem, ProviderRequest};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use sha2::Sha256;
use crate::config::{ClaudeCodeConfig, PROVIDER};
@ -16,9 +17,11 @@ use crate::config::{ClaudeCodeConfig, PROVIDER};
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub(crate) enum ClaudeCodeHandle {
/// The session a finished turn left behind. `prefix_sha256` is the digest
/// of the instructions, the model and every transcript item that precedes
/// this handle.
/// The session a finished turn left behind. `prefix_sha256` is the
/// HMAC-SHA-256, keyed with the connector's handle secret, of the caller
/// that minted the handle, the instructions, the model, this session and
/// instance, and every transcript item that precedes this handle. Another caller, or anyone
/// without the secret, cannot produce or satisfy it.
Ended {
session_id: String,
instance: String,
@ -36,12 +39,43 @@ impl ClaudeCodeHandle {
}
}
/// The digest a handle binds: named MessagePack of the instructions, the model
/// and the transcript items before the handle.
pub(crate) fn prefix_digest(instructions: &str, model: &str, items: &[InputItem]) -> [u8; 32] {
let bytes =
rmp_serde::to_vec_named(&(instructions, model, items)).expect("strings and bytes encode");
Sha256::digest(bytes).into()
/// The smallest handle secret the connector accepts, in bytes.
pub const MIN_SECRET_BYTES: usize = 32;
/// Everything a handle's tag binds: the minting caller, the instructions, the
/// model, the session and instance the handle names, and the transcript items
/// before the handle. Swapping any of them into a valid handle breaks the tag.
pub(crate) struct Binding<'a> {
pub(crate) caller: &'a str,
pub(crate) instructions: &'a str,
pub(crate) model: &'a str,
pub(crate) session_id: &'a str,
pub(crate) instance: &'a str,
pub(crate) items: &'a [InputItem],
}
fn prefix_mac(secret: &[u8], binding: &Binding<'_>) -> Hmac<Sha256> {
let bytes = rmp_serde::to_vec_named(&(
binding.caller,
binding.instructions,
binding.model,
binding.session_id,
binding.instance,
binding.items,
))
.expect("strings and bytes encode");
let mut mac = Hmac::<Sha256>::new_from_slice(secret).expect("an HMAC key has any length");
mac.update(&bytes);
mac
}
/// The tag a handle carries: keyed named MessagePack of its [`Binding`].
pub(crate) fn prefix_digest(secret: &[u8], binding: &Binding<'_>) -> [u8; 32] {
prefix_mac(secret, binding).finalize().into_bytes().into()
}
fn prefix_matches(secret: &[u8], binding: &Binding<'_>, tag: &[u8; 32]) -> bool {
prefix_mac(secret, binding).verify_slice(tag).is_ok()
}
/// Whether the CLI's session store still holds a session.
@ -89,8 +123,6 @@ impl SessionStore for FsSessionStore {
pub(crate) enum PlanError {
#[error("claude-code: tools are not supported by this backend yet")]
Tools,
#[error("claude-code: structured output is not supported by this backend yet")]
OutputSchema,
#[error("claude-code: tool results are not supported by this backend yet")]
ToolResult,
#[error("claude-code: the request ends without user text")]
@ -123,14 +155,13 @@ impl Plan {
pub(crate) fn plan(
request: &ProviderRequest,
config: &ClaudeCodeConfig,
secret: &[u8],
caller: &str,
sessions: &dyn SessionStore,
) -> Result<Plan, PlanError> {
if !request.tools.is_empty() {
return Err(PlanError::Tools);
}
if request.output_schema_json.is_some() {
return Err(PlanError::OutputSchema);
}
let input = &request.input;
let run_start = input
.iter()
@ -179,7 +210,18 @@ pub(crate) fn plan(
.all(|item| matches!(item, InputItem::UserText { .. }));
let resumable = instance == config.instance
&& tail_is_user_text
&& prefix_sha256 == prefix_digest(&request.instructions, &request.model, &input[..at])
&& prefix_matches(
secret,
&Binding {
caller,
instructions: &request.instructions,
model: &request.model,
session_id: &session_id,
instance: &instance,
items: &input[..at],
},
&prefix_sha256,
)
&& sessions.exists(&session_id);
if resumable {
return Ok(Plan::Resume {

View file

@ -9,19 +9,22 @@ use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use model_connector::{ProviderRequest, ReasoningEffort};
use zeroize::Zeroizing;
use crate::config::ClaudeCodeConfig;
use crate::handle::Plan;
/// The one credential variable, copied by name from the connector's own
/// environment into the child and nowhere else.
pub(crate) const TOKEN_VAR: &str = "CLAUDE_CODE_OAUTH_TOKEN";
/// The variable that tells the CLI which file descriptor holds its credential.
/// The token itself never enters the environment: the process port writes it
/// into a pipe and hands the child the read end at [`TOKEN_FD`].
pub(crate) const TOKEN_FD_VAR: &str = "CLAUDE_CODE_OAUTH_TOKEN_FILE_DESCRIPTOR";
/// The descriptor number the child finds the token pipe at: the first one
/// after stdin, stdout and stderr.
pub(crate) const TOKEN_FD: i32 = 3;
const CHILD_PATH: &str = "/usr/local/bin:/usr/bin:/bin";
/// The environment of the child: exactly these names and nothing inherited.
/// Values are never formatted; `Debug` lists names only.
pub struct ChildEnv(Vec<(&'static str, Zeroizing<String>)>);
/// `Debug` lists names only.
pub struct ChildEnv(Vec<(&'static str, String)>);
impl ChildEnv {
pub fn iter(&self) -> impl Iterator<Item = (&str, &str)> {
@ -37,21 +40,17 @@ impl fmt::Debug for ChildEnv {
}
}
pub(crate) fn child_env(
config: &ClaudeCodeConfig,
request: &ProviderRequest,
token: Zeroizing<String>,
) -> ChildEnv {
pub(crate) fn child_env(config: &ClaudeCodeConfig, request: &ProviderRequest) -> ChildEnv {
let config_dir = config.config_dir.to_string_lossy().into_owned();
ChildEnv(vec![
("PATH", Zeroizing::new(CHILD_PATH.to_string())),
("HOME", Zeroizing::new(config_dir.clone())),
("CLAUDE_CONFIG_DIR", Zeroizing::new(config_dir)),
("PATH", CHILD_PATH.to_string()),
("HOME", config_dir.clone()),
("CLAUDE_CONFIG_DIR", config_dir),
(
"CLAUDE_CODE_MAX_OUTPUT_TOKENS",
Zeroizing::new(request.max_output_tokens.to_string()),
request.max_output_tokens.to_string(),
),
(TOKEN_VAR, token),
(TOKEN_FD_VAR, TOKEN_FD.to_string()),
])
}
@ -62,13 +61,18 @@ pub(crate) struct PromptFile(PathBuf);
static PROMPT_FILES: AtomicU64 = AtomicU64::new(0);
impl PromptFile {
pub(crate) fn create(instructions: &str) -> io::Result<Self> {
pub(crate) fn create(dir: &Path, instructions: &str) -> io::Result<Self> {
let name = format!(
"claude-connector-prompt-{}-{}",
std::process::id(),
PROMPT_FILES.fetch_add(1, Ordering::Relaxed)
);
let path = std::env::temp_dir().join(name);
Self::create_at(dir.join(name), instructions)
}
/// Creates `path` and nothing that already exists there: a file, a link
/// (dangling or not) or a directory at the name fails, and is left alone.
fn create_at(path: PathBuf, instructions: &str) -> io::Result<Self> {
let mut options = OpenOptions::new();
options.write(true).create_new(true);
#[cfg(unix)]
@ -103,12 +107,15 @@ fn effort(effort: ReasoningEffort) -> &'static str {
/// rather than an argument because one argument is capped near 128 KiB
/// (E2BIG) and instructions are not, the child has
/// no tools, no MCP servers, no settings, no slash commands and nobody to ask
/// for permission, and a resume always forks.
/// for permission, and a resume always forks. An ephemeral call (the
/// meter's) keeps no session on disk. A request's output schema goes to the
/// CLI as `--json-schema`; without one the flag is absent.
pub(crate) fn argv(
config: &ClaudeCodeConfig,
request: &ProviderRequest,
plan: &Plan,
system_prompt_file: &Path,
ephemeral: bool,
) -> Vec<String> {
let mut argv: Vec<String> = [
config.claude_binary.to_string_lossy().as_ref(),
@ -134,9 +141,15 @@ pub(crate) fn argv(
.into_iter()
.map(str::to_string)
.collect();
if ephemeral {
argv.push("--no-session-persistence".to_string());
}
if let Some(level) = request.reasoning_effort {
argv.extend(["--effort".to_string(), effort(level).to_string()]);
}
if let Some(schema) = &request.output_schema_json {
argv.extend(["--json-schema".to_string(), schema.clone()]);
}
if let Plan::Resume { session_id, .. } = plan {
argv.extend([
"--resume".to_string(),
@ -146,3 +159,57 @@ pub(crate) fn argv(
}
argv
}
#[cfg(all(test, unix))]
mod tests {
use std::os::unix::fs::{PermissionsExt, symlink};
use super::*;
fn scratch(name: &str) -> PathBuf {
let dir = std::env::temp_dir().join(format!(
"claude-connector-launch-{name}-{}",
std::process::id()
));
let _ = std::fs::remove_dir_all(&dir);
std::fs::create_dir_all(&dir).unwrap();
dir
}
#[test]
fn a_prompt_is_created_private_and_removed_on_drop() {
let dir = scratch("private");
let prompt = PromptFile::create(&dir, "the prompt").unwrap();
let path = prompt.path().to_path_buf();
assert_eq!(path.parent(), Some(dir.as_path()));
assert_eq!(std::fs::read_to_string(&path).unwrap(), "the prompt");
let mode = std::fs::metadata(&path).unwrap().permissions().mode() & 0o777;
assert_eq!(mode, 0o600);
drop(prompt);
assert!(!path.exists());
std::fs::remove_dir_all(dir).unwrap();
}
#[test]
fn a_name_already_taken_is_never_written_or_removed() {
let dir = scratch("taken");
let victim = dir.join("victim");
std::fs::write(&victim, "keep").unwrap();
let file = dir.join("file");
std::fs::write(&file, "mine").unwrap();
let link = dir.join("link");
symlink(&victim, &link).unwrap();
let dangling = dir.join("dangling");
symlink(dir.join("nowhere"), &dangling).unwrap();
for taken in [&file, &link, &dangling] {
assert!(PromptFile::create_at(taken.clone(), "the prompt").is_err());
}
assert_eq!(std::fs::read_to_string(&victim).unwrap(), "keep");
assert_eq!(std::fs::read_to_string(&file).unwrap(), "mine");
assert!(!dir.join("nowhere").exists());
assert!(link.exists() && dangling.symlink_metadata().is_ok());
std::fs::remove_dir_all(dir).unwrap();
}
}

View file

@ -0,0 +1,745 @@
//! The Self-run credential lease. The connector is the only holder of the
//! subscription token and the one owner of Self-run admission (rulings
//! self-runs-through-connector and self-run-credential-lease). 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), and hands
//! back one pipe fd that holds the token, with the plan's usage windows. The
//! lease lasts while the holder keeps the socket open: it ends when the holder
//! closes it, when its unit invocation ends, or when its expiry passes. Nothing
//! is reserved and nothing is billed; the only usage-derived refusal is the
//! provider's own, repeated until its own reset. Nothing the run says about
//! itself decides who it is 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, LeaseTurn, PlanQuota, TransportService,
};
use serde::{Deserialize, Serialize};
use crate::{QuotaCell, 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),
}
/// A request names only when the lease ends. Who asks is 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,
/// The plan's usage windows as the connector last saw them, so the run
/// can read its own burn.
pub plan_quota: Vec<PlanQuota>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum LeaseRefusalKind {
UnknownCaller,
Expired,
Busy,
ProviderRejected,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct LeaseRefused {
pub kind: LeaseRefusalKind,
pub detail: String,
}
/// 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 unit its process must run in (the
/// Self or the Persona unit of its instance, by `turn`), and the
/// caller whose name its leases carry.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LeaseCaller {
pub uid: u32,
pub caller_id: String,
/// Which kind of turn this account leases for; it picks the unit.
pub turn: LeaseTurn,
/// The `<instance>` of `mind-self@<instance>.service` or
/// `mind-persona@<instance>.service`.
pub unit_instance: String,
}
impl LeaseCaller {
/// The one unit a lease to this caller may be requested from.
pub fn unit(&self) -> String {
match self.turn {
LeaseTurn::SelfRun => format!("mind-self@{}.service", self.unit_instance),
LeaseTurn::PersonaTurn => format!("mind-persona@{}.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>>,
quota: Arc<QuotaCell>,
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) the hold. 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
/// field of any message names either.
pub fn serve_leases(
listener: &UnixListener,
callers: Vec<LeaseCaller>,
limits: LeaseLimits,
service: Arc<Mutex<TransportService>>,
quota: Arc<QuotaCell>,
tokens: Arc<dyn TokenSource>,
units: Arc<dyn PeerUnits>,
) -> io::Result<()> {
let lender = Arc::new(Lender {
callers,
limits,
service,
quota,
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 ended 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 lease is claimed before the token is read, let alone lent. The
// plan the refusal reads is the one this process last observed, but a
// cell with no weekly reading (a restart empties it) gets the meter's
// first read, bounded, before the lease is lent "unobserved" for its
// whole life. The request is the swarm's sign of life, however it
// ends: the meter reads while one is held or lately asked for.
let _asking = self.quota.asking();
let plan = self.quota.snapshot_with_week(FIRST_READ_WAIT);
let now = unix_ms();
let admission = self
.service
.lock()
.map_err(|_| io::Error::other("service lock"))?
.begin_lease(
&caller.caller_id,
caller.turn,
&invocation.invocation_id,
request.expires_at_unix_ms,
now,
&plan,
)
.map_err(|_| io::Error::other("lease"))?;
let grant = match admission {
LeaseAdmission::Granted(grant) => grant,
LeaseAdmission::Refused(why) => {
let (kind, detail) = match why {
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",
),
LeaseRefusal::ProviderRejected => (
LeaseRefusalKind::ProviderRejected,
"the provider is refusing until its own reset",
),
};
return refuse(&mut stream, kind, detail);
}
};
self.hold(
stream,
caller,
&grant,
&invocation,
peer.process,
request.expires_at_unix_ms,
plan.windows,
)
}
/// Everything after a lease exists: lend the token, then keep the lease
/// while its holder keeps the socket and its unit invocation, and end it
/// exactly once whatever happens.
// One argument per fact the held lease binds.
#[allow(clippy::too_many_arguments)]
fn hold(
&self,
mut stream: UnixStream,
caller: &LeaseCaller,
grant: &LeaseGrant,
invocation: &PeerUnit,
process: Option<OwnedFd>,
expires_at_unix_ms: u64,
plan_quota: Vec<PlanQuota>,
) -> io::Result<()> {
let held = self.keep(
&mut stream,
caller,
grant,
invocation,
process.as_ref(),
expires_at_unix_ms,
plan_quota,
);
self.service
.lock()
.map_err(|_| io::Error::other("service lock"))?
.end_lease(&caller.caller_id, &grant.lease_id);
held
}
// One argument per fact the held lease binds.
#[allow(clippy::too_many_arguments)]
fn keep(
&self,
stream: &mut UnixStream,
caller: &LeaseCaller,
grant: &LeaseGrant,
invocation: &PeerUnit,
process: Option<&OwnedFd>,
expires_at_unix_ms: u64,
plan_quota: Vec<PlanQuota>,
) -> io::Result<()> {
// A token that cannot be read or piped was never lent: the lease ends.
let Some(read_end) = self
.tokens
.read()
.and_then(|token| token_pipe(token.as_bytes()).ok())
else {
return Ok(());
};
let frame = encode_frame(&LeaseMessage::ClaudeLeaseGrant(LeaseGranted {
lease_id: grant.lease_id.clone(),
plan_quota,
}));
let sent = send_with_fd(stream, &frame, read_end.as_raw_fd());
drop(read_end);
sent?;
// The holder says nothing more. The lease ends when it closes the
// socket, when its unit invocation is no longer the one that leased,
// or at the expiry.
let mut byte = [0_u8; 1];
loop {
let left = expires_at_unix_ms.saturating_sub(unix_ms());
if left == 0 {
return Ok(());
}
stream.set_read_timeout(Some(Duration::from_millis(left.min(HOLD_POLL_MS))))?;
match stream.read(&mut byte) {
Ok(0) => return Ok(()),
Ok(_) => {}
Err(error)
if matches!(
error.kind(),
io::ErrorKind::WouldBlock | io::ErrorKind::TimedOut
) => {}
Err(error) if error.kind() == io::ErrorKind::Interrupted => {}
Err(_) => return Ok(()),
}
let still = process
.and_then(|process| self.units.invocation_in(process.as_fd(), &caller.unit()))
.is_some_and(|found| found == *invocation);
if !still {
return Ok(());
}
}
}
}
/// 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);
/// How often a held lease looks again at its holder's unit invocation.
const HOLD_POLL_MS: u64 = 500;
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 refusal frame whose body is exactly `body_bytes` long.
fn frame_of(body_bytes: usize) -> Vec<u8> {
let frame_with = |pad: usize| {
encode_frame(&LeaseMessage::ClaudeLeaseRefusal(LeaseRefused {
kind: LeaseRefusalKind::Busy,
detail: "x".repeat(pad),
}))
};
// 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 = 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 = 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 = frame_of(MAX_LEASE_FRAME_BYTES + 1);
assert!(within(&past_bound).is_err());
let at_bound = 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::ClaudeLeaseRefusal(LeaseRefused {
kind: LeaseRefusalKind::Busy,
detail: "lease:x".to_string(),
}));
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();
}
}

View file

@ -5,15 +5,36 @@
mod backend;
mod config;
mod daemon;
mod handle;
mod launch;
#[cfg(target_os = "linux")]
mod lease;
mod meter;
mod process;
mod quota;
mod stream;
#[cfg(target_os = "linux")]
mod unit;
pub use backend::{ClaudeCodeBackend, StderrSink};
pub use backend::{ClaudeCodeBackend, CredentialTokenSource, StderrSink, TokenSource};
pub use config::{ClaudeCodeConfig, ClaudeCodeLimits};
pub use handle::{FsSessionStore, SessionStore};
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, MIN_SECRET_BYTES, SessionStore};
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,
};
#[cfg(unix)]
pub use process::StdProcessPort;
pub use process::{ProcessPort, RunningChild, StdoutRead};
pub use quota::QuotaCell;
#[cfg(target_os = "linux")]
pub use unit::SystemdUnits;

View file

@ -0,0 +1,393 @@
//! The claude-connector binary.
//!
//! claude-connector serve --config PATH.cc
//! claude-connector config init PATH.cc --flag value ...
//! claude-connector config add-caller PATH.cc --caller-id ID --connection-key-file F ...
//! claude-connector config add-lease-caller PATH.cc --lease-user U --caller-id ID ...
//! claude-connector config check PATH.cc
//!
//! `serve` takes its configuration from the typed document alone. The flags of
//! `config init` exist only to author that document. No message here echoes a
//! value that was typed or stored; refusals name a field or a flag.
use std::collections::BTreeMap;
use std::path::PathBuf;
use std::process::ExitCode;
use claude_connector::{
ClaudeConnectorConfig, ConfigError, LeaseCallerConfig, MIN_SECRET_BYTES, check_ready,
load_config, write_config,
};
use model_connector::{CallerConfig, LeaseTurn, frame_bytes_for_payload};
fn main() -> ExitCode {
let args: Vec<String> = std::env::args().skip(1).collect();
let args: Vec<&str> = args.iter().map(String::as_str).collect();
let result = match args.as_slice() {
["serve", "--config", path] => serve(path),
["config", "init", path, flags @ ..] => init(path, flags),
["config", "add-caller", path, flags @ ..] => add_caller(path, flags),
["config", "add-lease-caller", path, flags @ ..] => add_lease_caller(path, flags),
["config", "check", path] => check(path),
_ => Err(Failure::Usage),
};
match result {
Ok(()) => ExitCode::SUCCESS,
Err(failure) => {
eprintln!("claude-connector: {failure}");
ExitCode::FAILURE
}
}
}
#[derive(Debug, thiserror::Error)]
enum Failure {
#[error(
"usage: serve --config PATH | config init PATH --flag value ... | config add-caller PATH --flag value ... | config add-lease-caller PATH --flag value ... | config check PATH (check reads the credentials from the environment it runs in: run it in the unit's environment)"
)]
Usage,
#[error(transparent)]
Config(#[from] ConfigError),
#[error("flag {0} is missing")]
Missing(&'static str),
#[error("flag {0} is malformed")]
Malformed(&'static str),
#[error("a flag is unknown, repeated or has no value")]
Flags,
}
#[cfg(unix)]
fn serve(path: &str) -> Result<(), Failure> {
Ok(claude_connector::serve_config(path.as_ref())?)
}
#[cfg(not(unix))]
fn serve(_: &str) -> Result<(), Failure> {
Err(Failure::Usage)
}
fn check(path: &str) -> Result<(), Failure> {
let config = load_config(path.as_ref())?;
check_ready(&config)?;
println!("claude connector configuration is valid and ready to serve");
println!("callers: {}", config.callers.len());
println!("lease_callers: {}", config.lease_callers.len());
Ok(())
}
/// The flags of one config verb: each single flag at most once, each list
/// flag any number of times.
struct Flags {
single: BTreeMap<&'static str, String>,
lists: BTreeMap<&'static str, Vec<String>>,
}
/// The flags that describe one caller, shared by `config init` (which writes
/// the first caller) and `config add-caller`.
const CALLER: [&str; 5] = [
"--caller-id",
"--connection-key-file",
"--connection-key-epoch",
"--max-concurrent-requests",
"--max-output-tokens",
];
const INIT: [&str; 21] = [
"--listen",
"--replay-store",
"--claude-binary",
"--claude-sha256",
"--work-dir",
"--config-dir",
"--prompt-dir",
"--instance",
"--token-credential",
"--handle-secret-credential",
"--handle-secret-file",
"--deadline-ms",
"--max-line-bytes",
"--max-total-bytes",
"--stderr-tail-bytes",
"--max-payload-bytes",
"--socket-timeout-ms",
"--lease-socket",
"--lease-max-ms",
"--turn",
"--lease-unit-instance",
];
const LEASE_CALLER: [&str; 3] = ["--caller-id", "--turn", "--lease-unit-instance"];
impl Flags {
fn parse(
args: &[&str],
singles: &[&[&'static str]],
lists: &[&'static str],
) -> Result<Self, Failure> {
let mut flags = Flags {
single: BTreeMap::new(),
lists: BTreeMap::new(),
};
let mut args = args.iter();
while let Some(&name) = args.next() {
let Some(&value) = args.next() else {
return Err(Failure::Flags);
};
if let Some(list) = lists.iter().find(|known| **known == name) {
flags
.lists
.entry(*list)
.or_default()
.push(value.to_string());
continue;
}
let Some(known) = singles
.iter()
.flat_map(|set| set.iter())
.find(|k| **k == name)
else {
return Err(Failure::Flags);
};
if flags.single.insert(*known, value.to_string()).is_some() {
return Err(Failure::Flags);
}
}
Ok(flags)
}
fn list(&self, name: &'static str) -> &[String] {
self.lists.get(name).map_or(&[], Vec::as_slice)
}
fn text(&self, name: &'static str) -> Result<String, Failure> {
self.single.get(name).cloned().ok_or(Failure::Missing(name))
}
fn path(&self, name: &'static str) -> Result<PathBuf, Failure> {
self.text(name).map(PathBuf::from)
}
fn number<T: std::str::FromStr>(&self, name: &'static str) -> Result<T, Failure> {
self.text(name)?
.parse()
.map_err(|_| Failure::Malformed(name))
}
/// The kind of turn a leasing account serves: `self` or `persona`.
fn turn(&self) -> Result<LeaseTurn, Failure> {
match self.text("--turn")?.as_str() {
"self" => Ok(LeaseTurn::SelfRun),
"persona" => Ok(LeaseTurn::PersonaTurn),
_ => Err(Failure::Malformed("--turn")),
}
}
fn number_or<T: std::str::FromStr>(
&self,
name: &'static str,
default: T,
) -> Result<T, Failure> {
match self.single.get(name) {
None => Ok(default),
Some(text) => text.parse().map_err(|_| Failure::Malformed(name)),
}
}
/// The caller these flags describe. Its payload bound comes from
/// `--max-payload-bytes` when the verb takes one.
fn caller(&self) -> Result<CallerConfig, Failure> {
Ok(CallerConfig {
caller_runtime_id: self.text("--caller-id")?,
connection_key_file: self.path("--connection-key-file")?,
connection_key_epoch: self.number("--connection-key-epoch")?,
allowed_models: self.list("--allowed-model").to_vec(),
max_concurrent_requests: self.number_or("--max-concurrent-requests", 4)?,
max_payload_bytes: self.number_or("--max-payload-bytes", 1 << 20)?,
max_output_tokens: self.number_or("--max-output-tokens", 32_768)?,
})
}
}
fn sha256(hex: &str) -> Option<[u8; 32]> {
if hex.len() != 64 || !hex.is_ascii() {
return None;
}
let mut digest = [0_u8; 32];
for (byte, pair) in digest.iter_mut().zip(hex.as_bytes().chunks(2)) {
*byte = u8::from_str_radix(std::str::from_utf8(pair).ok()?, 16).ok()?;
}
Some(digest)
}
/// Writes 32 random bytes as 64 hex characters, mode 0600, unless the file is
/// already there: the secret is never read back or shown, and a file that
/// exists is the secret handles already depend on, so it is kept as it is when
/// it is a regular file, mode 0600, owned by this user and at least
/// `MIN_SECRET_BYTES` long, and refused untouched otherwise.
#[cfg(unix)]
fn ensure_secret_file(path: &std::path::Path) -> Result<(), Failure> {
use std::io::{Read, Write};
use std::os::unix::fs::OpenOptionsExt;
if !path.is_absolute() {
return Err(Failure::Malformed("--handle-secret-file"));
}
let mut file = match std::fs::OpenOptions::new()
.write(true)
.create_new(true)
.mode(0o600)
.open(path)
{
Ok(file) => file,
Err(_) => return existing_secret_is_usable(path),
};
let mut bytes = [0_u8; 32];
std::fs::File::open("/dev/urandom")
.and_then(|mut random| random.read_exact(&mut bytes))
.map_err(|_| Failure::Malformed("--handle-secret-file"))?;
let hex: String = bytes.iter().map(|byte| format!("{byte:02x}")).collect();
file.write_all(hex.as_bytes())
.map_err(|_| Failure::Malformed("--handle-secret-file"))
}
/// An existing secret file is accepted only as the connector would use it: not
/// a link, a regular file, readable by no one else, ours and long enough.
#[cfg(unix)]
fn existing_secret_is_usable(path: &std::path::Path) -> Result<(), Failure> {
let metadata =
std::fs::symlink_metadata(path).map_err(|_| Failure::Malformed("--handle-secret-file"))?;
// SAFETY: geteuid has no preconditions and cannot fail.
let uid = unsafe { libc::geteuid() };
if secret_file_usable(&metadata, uid) {
Ok(())
} else {
Err(Failure::Malformed("--handle-secret-file"))
}
}
#[cfg(unix)]
fn secret_file_usable(metadata: &std::fs::Metadata, uid: u32) -> bool {
use std::os::unix::fs::{MetadataExt, PermissionsExt};
metadata.is_file()
&& metadata.permissions().mode() & 0o7777 == 0o600
&& metadata.uid() == uid
&& metadata.len() >= MIN_SECRET_BYTES as u64
}
#[cfg(not(unix))]
fn ensure_secret_file(_: &std::path::Path) -> Result<(), Failure> {
Err(Failure::Usage)
}
fn init(path: &str, args: &[&str]) -> Result<(), Failure> {
let flags = Flags::parse(
args,
&[&INIT, &CALLER],
&["--allowed-model", "--lease-user"],
)?;
let first = flags.caller()?;
let lease_callers = flags
.list("--lease-user")
.iter()
.map(|unix_user| {
Ok(LeaseCallerConfig {
unix_user: unix_user.clone(),
caller_id: first.caller_runtime_id.clone(),
turn: flags.turn()?,
unit_instance: flags.text("--lease-unit-instance")?,
})
})
.collect::<Result<Vec<_>, Failure>>()?;
let secret_file = flags.path("--handle-secret-file")?;
let config = ClaudeConnectorConfig {
listen: flags.text("--listen")?,
replay_store: flags.path("--replay-store")?,
claude_binary: flags.path("--claude-binary")?,
claude_binary_sha256: sha256(&flags.text("--claude-sha256")?)
.ok_or(Failure::Malformed("--claude-sha256"))?,
work_dir: flags.path("--work-dir")?,
config_dir: flags.path("--config-dir")?,
prompt_dir: flags.path("--prompt-dir")?,
instance: flags.text("--instance")?,
token_credential: flags.text("--token-credential")?,
deadline_ms: flags.number_or("--deadline-ms", 600_000)?,
max_line_bytes: flags.number_or("--max-line-bytes", 4 << 20)?,
max_total_bytes: flags.number_or("--max-total-bytes", 64 << 20)?,
stderr_tail_bytes: flags.number_or("--stderr-tail-bytes", 8192)?,
max_frame_bytes: frame_bytes_for_payload(first.max_payload_bytes),
max_connections: first.max_concurrent_requests.max(8),
socket_timeout_ms: flags.number_or("--socket-timeout-ms", 10_000)?,
max_expiry_skew_ms: 300_000,
callers: vec![first],
lease_socket: flags.single.get("--lease-socket").map(PathBuf::from),
lease_callers,
lease_max_ms: flags.number_or("--lease-max-ms", 3_600_000)?,
handle_secret_credential: flags.text("--handle-secret-credential")?,
};
// Validate before the secret file exists: a refused document leaves nothing.
config.validate()?;
ensure_secret_file(&secret_file)?;
write_config(path.as_ref(), &config)?;
Ok(())
}
/// Appends one caller to an existing document. The serve limits grow to carry
/// it: the frame to its payload bound, the connections to every caller's slots.
fn add_caller(path: &str, args: &[&str]) -> Result<(), Failure> {
let flags = Flags::parse(
args,
&[&CALLER, &["--max-payload-bytes"]],
&["--allowed-model"],
)?;
let caller = flags.caller()?;
let mut config = load_config(path.as_ref())?;
config.max_frame_bytes = config
.max_frame_bytes
.max(frame_bytes_for_payload(caller.max_payload_bytes));
config.callers.push(caller);
let slots: usize = config
.callers
.iter()
.map(|caller| caller.max_concurrent_requests)
.sum();
config.max_connections = config.max_connections.max(slots);
write_config(path.as_ref(), &config)?;
Ok(())
}
/// Appends one leasing account to an existing document that lends a token.
fn add_lease_caller(path: &str, args: &[&str]) -> Result<(), Failure> {
let flags = Flags::parse(args, &[&LEASE_CALLER, &["--lease-user"]], &[])?;
let mut config = load_config(path.as_ref())?;
config.lease_callers.push(LeaseCallerConfig {
unix_user: flags.text("--lease-user")?,
caller_id: flags.text("--caller-id")?,
turn: flags.turn()?,
unit_instance: flags.text("--lease-unit-instance")?,
});
write_config(path.as_ref(), &config)?;
Ok(())
}
#[cfg(all(test, unix))]
mod tests {
use std::os::unix::fs::{MetadataExt, PermissionsExt};
use super::*;
#[test]
fn a_secret_file_owned_by_someone_else_is_not_usable() {
let dir =
std::env::temp_dir().join(format!("claude-connector-main-{}", std::process::id()));
std::fs::create_dir_all(&dir).unwrap();
let path = dir.join("handle.secret");
std::fs::write(&path, "k".repeat(64)).unwrap();
std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o600)).unwrap();
let metadata = std::fs::symlink_metadata(&path).unwrap();
let owner = metadata.uid();
assert!(secret_file_usable(&metadata, owner));
assert!(!secret_file_usable(&metadata, owner.wrapping_add(1)));
std::fs::remove_dir_all(&dir).unwrap();
}
}

View file

@ -0,0 +1,543 @@
//! The meter: keeps the plan's weekly reading fresh while the swarm is
//! working. The subscription reports its windows only inside the CLI's
//! `rate_limit_event`, on a call the CLI makes, and Personas make theirs under
//! leases that report nothing back. So when the `seven_day` reading in the
//! [`QuotaCell`] is `STALE` old, the meter makes the smallest real call
//! through the backend's own path (one `claude -p` on the cheapest model with
//! a trivial prompt) and the stream fills the cell exactly as it does for any
//! call. The meter holds no credential, speaks no protocol and writes nothing
//! itself.
//!
//! The meter is driven by the connector's own lease state. It reads only
//! while a lease is held or one was asked for less than `STALE` ago: a held
//! swarm brake stops the units, they ask for nothing, and a braked or quiet
//! swarm spends nothing. A call that fails, or completes without refreshing
//! the reading, is not retried at the staleness period: the wait doubles per
//! consecutive miss, up to `CEILING`.
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use model_connector::Usage;
use crate::backend::ClaudeCodeBackend;
use crate::quota::{EXPIRED, QuotaCell, WEEK};
/// A reading this old, or older, is stale; also the least gap between calls.
const STALE: Duration = Duration::from_secs(15 * 60);
/// The longest any backoff may wait; also how long an idle meter sleeps
/// before looking again, since a lease request wakes it.
const CEILING: Duration = Duration::from_secs(6 * 60 * 60);
/// The smallest real call, as the backend makes it.
pub(crate) trait Probe: Send {
/// What the call billed, or the failure's kind. Never a message: a
/// provider's words do not enter the meter's log.
fn call(&self) -> Result<Usage, String>;
}
impl Probe for Arc<ClaudeCodeBackend> {
fn call(&self) -> Result<Usage, String> {
self.read_meter()
}
}
/// Why the meter is waiting; logged once per change.
#[derive(Debug, Clone, PartialEq, Eq)]
enum State {
Read,
Idle,
/// The call failed; the kind is the failure's kind, never its message.
Failed(String),
Unread,
}
impl State {
fn line(&self) -> String {
match self {
Self::Read => "meter: weekly reading refreshed".to_string(),
Self::Idle => "meter: no lease asked lately, making no call".to_string(),
Self::Failed(kind) => format!("meter: the call failed ({kind}), backing off"),
Self::Unread => {
"meter: the call completed without a weekly reading, backing off".to_string()
}
}
}
}
struct Meter {
cell: Arc<QuotaCell>,
probe: Box<dyn Probe>,
log: Box<dyn FnMut(&str) + Send>,
/// No call before this instant, unix ms.
not_before_ms: u64,
/// Consecutive calls that left the reading unrefreshed.
misses: u32,
last_state: Option<State>,
}
impl Meter {
fn new(cell: Arc<QuotaCell>, probe: Box<dyn Probe>, log: Box<dyn FnMut(&str) + Send>) -> Self {
Self {
cell,
probe,
log,
not_before_ms: 0,
misses: 0,
last_state: None,
}
}
/// Makes the call if one is due at `now_ms`; returns how long to wait
/// before looking again. Nothing is due while no lease is held or lately
/// asked for, nor before the last miss's backoff has passed, nor while the
/// weekly reading is younger than `STALE`.
fn tick(&mut self, now_ms: u64) -> Duration {
if !self.cell.wanted_at(now_ms, STALE) {
self.report(State::Idle);
return CEILING;
}
let due_at = self.due_at_ms();
if now_ms < due_at {
return Duration::from_millis(due_at - now_ms);
}
// A failed call may still have refreshed the reading: a rejected turn
// reports its windows and then errors. Only the cell says whether the
// week was read.
let called = self.probe.call();
let state = if self.refreshed_since(now_ms) {
self.misses = 0;
State::Read
} else {
self.misses = self.misses.saturating_add(1);
match called {
Ok(_) => State::Unread,
Err(kind) => State::Failed(kind),
}
};
self.report(state);
self.not_before_ms = now_ms.saturating_add(duration_ms(backoff(self.misses)));
Duration::from_millis(self.due_at_ms().saturating_sub(now_ms))
}
/// The earliest instant the next call may happen. A window past its reset
/// described a week that is over, so it is due at once.
fn due_at_ms(&self) -> u64 {
let stale_at = self
.cell
.snapshot()
.windows
.iter()
.find(|quota| quota.window == WEEK)
.map_or(0, |quota| {
if quota.status == EXPIRED {
0
} else {
quota.observed_at_unix_ms.saturating_add(duration_ms(STALE))
}
});
stale_at.max(self.not_before_ms)
}
fn refreshed_since(&self, since_ms: u64) -> bool {
self.cell
.snapshot()
.windows
.iter()
.any(|quota| quota.window == WEEK && quota.observed_at_unix_ms >= since_ms)
}
fn report(&mut self, state: State) {
if self.last_state.as_ref() != Some(&state) {
(self.log)(&state.line());
self.last_state = Some(state);
}
}
fn run(mut self) {
loop {
let seen = self.cell.asks_so_far();
let wait = self.tick(unix_ms());
self.cell.sleep_until_asked(seen, wait);
}
}
}
/// The wait after `misses` consecutive unrefreshed calls: the staleness
/// period, doubling per miss beyond the first, up to the ceiling.
fn backoff(misses: u32) -> Duration {
let doublings = misses.saturating_sub(1).min(16);
STALE.saturating_mul(1 << doublings).min(CEILING)
}
fn duration_ms(duration: Duration) -> u64 {
u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
}
fn unix_ms() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map_or(0, duration_ms)
}
/// Starts the meter thread. It writes only through the calls it makes, and
/// logs one line to stderr when its state changes. From here a lease asked of
/// a cell with no weekly reading waits for the meter's first read.
pub(crate) fn start(
cell: Arc<QuotaCell>,
probe: Box<dyn Probe>,
) -> std::io::Result<std::thread::JoinHandle<()>> {
cell.attach_meter();
let meter = Meter::new(
cell,
probe,
Box::new(|line| eprintln!("claude-connector: {line}")),
);
std::thread::Builder::new()
.name("claude-meter".to_string())
.spawn(move || meter.run())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::quota::Asking;
use model_connector::PlanQuota;
use std::sync::Mutex;
use std::sync::atomic::{AtomicUsize, Ordering};
const MIN: u64 = 60_000;
const NOW: u64 = 1_800_000_000_000;
const RESET: u64 = 4_102_444_800_000;
fn reading(status: &str) -> Vec<PlanQuota> {
vec![PlanQuota {
window: WEEK.to_string(),
utilization_permille: Some(60),
resets_at_unix_ms: RESET,
status: status.to_string(),
observed_at_unix_ms: 0,
}]
}
/// What the next call does: the instant its stream stamps a weekly
/// reading (if any) and the kind it then fails with (if any).
#[derive(Clone, Default)]
struct Script {
reading_at: Option<u64>,
status: String,
fails_with: Option<String>,
}
/// A call that answers as scripted and feeds the cell the way the stream
/// does, before it fails, as a rejected turn does.
struct Calls {
cell: Arc<QuotaCell>,
made: Arc<AtomicUsize>,
script: Arc<Mutex<Script>>,
}
impl Probe for Calls {
fn call(&self) -> Result<Usage, String> {
self.made.fetch_add(1, Ordering::SeqCst);
let script = self.script.lock().unwrap().clone();
if let Some(at) = script.reading_at {
self.cell.observe_at(at, reading(&script.status));
}
if let Some(kind) = script.fails_with {
return Err(kind);
}
Ok(Usage {
input_tokens: 385,
cache_read_tokens: 0,
cache_write_tokens: 0,
output_tokens: 4,
reasoning_tokens: None,
})
}
}
struct Rig {
meter: Meter,
cell: Arc<QuotaCell>,
made: Arc<AtomicUsize>,
script: Arc<Mutex<Script>>,
log: Arc<Mutex<Vec<String>>>,
/// The lease the swarm holds until a test lets go of it.
lease: Option<Asking>,
}
fn rig() -> Rig {
let cell = Arc::new(QuotaCell::new());
let made = Arc::new(AtomicUsize::new(0));
let script = Arc::new(Mutex::new(Script::default()));
let log = Arc::new(Mutex::new(Vec::new()));
let sink = log.clone();
let meter = Meter::new(
cell.clone(),
Box::new(Calls {
cell: cell.clone(),
made: made.clone(),
script: script.clone(),
}),
Box::new(move |line| sink.lock().unwrap().push(line.to_string())),
);
let lease = Some(cell.asking_at(NOW));
Rig {
meter,
cell,
made,
script,
log,
lease,
}
}
impl Rig {
fn reads(&self, at: u64) {
*self.script.lock().unwrap() = Script {
reading_at: Some(at),
status: "allowed".to_string(),
fails_with: None,
};
}
fn reads_nothing(&self) {
*self.script.lock().unwrap() = Script::default();
}
fn fails(&self, kind: &str) {
*self.script.lock().unwrap() = Script {
fails_with: Some(kind.to_string()),
..Script::default()
};
}
fn calls(&self) -> usize {
self.made.load(Ordering::SeqCst)
}
fn lines(&self) -> Vec<String> {
self.log.lock().unwrap().clone()
}
}
#[test]
fn an_unread_week_triggers_one_call_and_the_fresh_reading_ends_the_asking() {
let mut rig = rig();
rig.reads(NOW);
assert_eq!(rig.meter.tick(NOW), STALE, "fresh for the staleness period");
assert_eq!(rig.calls(), 1);
assert_eq!(
rig.meter.tick(NOW + 5 * MIN),
Duration::from_millis(10 * MIN)
);
assert_eq!(rig.calls(), 1, "a fresh reading makes no call");
}
#[test]
fn a_stale_reading_triggers_a_call_at_the_period_and_not_before() {
let mut rig = rig();
rig.cell.observe_at(NOW, reading("allowed"));
rig.reads(NOW + 15 * MIN);
rig.meter.tick(NOW + 15 * MIN - 1);
assert_eq!(rig.calls(), 0, "one millisecond short of stale");
rig.meter.tick(NOW + 15 * MIN);
assert_eq!(rig.calls(), 1);
}
#[test]
fn a_window_past_its_reset_is_read_at_once_not_after_the_period() {
let mut rig = rig();
let now = unix_ms();
let mut over = reading("allowed").remove(0);
over.resets_at_unix_ms = 1;
rig.cell.observe_at(now, vec![over]);
rig.reads(now);
rig.meter.tick(now);
assert_eq!(rig.calls(), 1);
}
#[test]
fn a_swarm_that_asks_for_nothing_spends_nothing_however_stale_the_reading() {
let mut rig = rig();
drop(rig.lease.take());
for step in 0..10 {
let at = NOW + 15 * MIN + step * 60 * MIN;
assert_eq!(rig.meter.tick(at), CEILING);
}
assert_eq!(rig.calls(), 0);
assert_eq!(rig.lines(), vec![State::Idle.line()]);
}
#[test]
fn a_lease_asked_for_wakes_the_meter_and_the_empty_cell_is_read() {
let mut rig = rig();
drop(rig.lease.take());
rig.reads(NOW + 180 * MIN);
rig.meter.tick(NOW + 180 * MIN);
assert_eq!(rig.calls(), 0, "asked three hours ago and let go");
let _later = rig.cell.asking_at(NOW + 180 * MIN);
rig.meter.tick(NOW + 180 * MIN);
assert_eq!(rig.calls(), 1);
}
#[test]
fn a_lease_let_go_is_still_wanted_for_exactly_one_staleness_period() {
let mut rig = rig();
drop(rig.lease.take());
rig.reads(NOW + 15 * MIN - 1);
rig.meter.tick(NOW + 15 * MIN - 1);
assert_eq!(rig.calls(), 1, "one millisecond inside the period");
rig.reads(NOW + 30 * MIN);
rig.meter.tick(NOW + 30 * MIN);
assert_eq!(rig.calls(), 1, "a period after the last ask: unwanted");
}
#[test]
fn a_lease_held_for_four_hours_keeps_the_week_read_throughout() {
let mut rig = rig();
for step in 0..=16 {
let at = NOW + step * 15 * MIN;
rig.reads(at);
rig.meter.tick(at);
}
assert_eq!(rig.calls(), 17, "one read per period, the lease never lapsing");
}
#[test]
fn a_failed_call_leaves_the_cell_and_backs_off_instead_of_hammering() {
let mut rig = rig();
rig.cell.observe_at(NOW, reading("allowed"));
let before = rig.cell.snapshot();
rig.fails("Provider");
let mut at = NOW + 15 * MIN;
let mut waits = Vec::new();
for _ in 0..4 {
waits.push(rig.meter.tick(at));
at += duration_ms(*waits.last().unwrap());
}
assert_eq!(rig.calls(), 4);
assert_eq!(rig.cell.snapshot(), before, "the cell is untouched");
assert_eq!(
waits,
[15, 30, 60, 120].map(|m| Duration::from_secs(m * 60)),
"the wait doubles per consecutive miss"
);
rig.meter.tick(at - MIN);
assert_eq!(rig.calls(), 4, "nothing is asked before the backoff ends");
assert_eq!(
rig.lines(),
vec!["meter: the call failed (Provider), backing off"]
);
}
#[test]
fn the_failure_kind_is_logged_again_when_it_changes_and_not_when_it_repeats() {
let mut rig = rig();
rig.fails("Authentication");
rig.meter.tick(NOW);
rig.meter.tick(NOW + 15 * MIN);
rig.fails("InvalidRequest");
rig.meter.tick(NOW + 45 * MIN);
assert_eq!(
rig.lines(),
vec![
"meter: the call failed (Authentication), backing off",
"meter: the call failed (InvalidRequest), backing off",
]
);
}
#[test]
fn the_backoff_stops_at_the_ceiling_and_one_good_reading_resets_it() {
assert_eq!(backoff(0), STALE);
assert_eq!(backoff(1), STALE);
assert_eq!(backoff(u32::MAX), CEILING);
let mut rig = rig();
rig.fails("Provider");
rig.meter.tick(NOW);
rig.meter.tick(NOW + 15 * MIN);
assert_eq!(rig.meter.misses, 2);
rig.reads(NOW + 45 * MIN);
rig.meter.tick(NOW + 45 * MIN);
assert_eq!(rig.meter.misses, 0);
}
#[test]
fn a_call_that_completes_without_a_weekly_reading_is_a_miss_too() {
let mut rig = rig();
rig.reads_nothing();
assert_eq!(rig.meter.tick(NOW), STALE);
assert_eq!(rig.meter.tick(NOW + 15 * MIN), 2 * STALE);
assert_eq!(rig.calls(), 2);
assert_eq!(rig.lines(), vec![State::Unread.line()]);
}
#[test]
fn a_rejected_turn_that_refreshed_the_reading_and_then_errored_is_no_miss() {
let mut rig = rig();
*rig.script.lock().unwrap() = Script {
reading_at: Some(NOW),
status: "rejected".to_string(),
fails_with: Some("RateLimited".to_string()),
};
assert_eq!(rig.meter.tick(NOW), STALE, "the next read is a period away");
assert_eq!(rig.meter.misses, 0);
assert_eq!(rig.lines(), vec![State::Read.line()]);
rig.reads(NOW + 15 * MIN);
rig.meter.tick(NOW + 15 * MIN);
assert_eq!(rig.calls(), 2, "a rejected reading is kept fresh");
}
#[test]
fn a_reading_from_before_the_call_does_not_count_as_its_answer() {
let mut rig = rig();
// A slow concurrent sensor stamped earlier than this call began.
rig.reads(NOW - 1);
rig.meter.tick(NOW);
assert_eq!(rig.meter.misses, 1);
}
#[test]
fn the_log_says_what_the_meter_is_doing_in_words() {
assert_eq!(
State::Idle.line(),
"meter: no lease asked lately, making no call"
);
assert_eq!(State::Read.line(), "meter: weekly reading refreshed");
assert_eq!(
State::Unread.line(),
"meter: the call completed without a weekly reading, backing off"
);
}
#[test]
fn the_clock_is_the_system_clock_in_unix_milliseconds() {
let before = SystemTime::now().duration_since(UNIX_EPOCH).unwrap();
let read = unix_ms();
let after = SystemTime::now().duration_since(UNIX_EPOCH).unwrap();
assert!((before.as_millis() as u64..=after.as_millis() as u64).contains(&read));
}
#[test]
fn a_started_meter_sleeps_until_a_lease_asks_and_then_reads() {
let cell = Arc::new(QuotaCell::new());
let made = Arc::new(AtomicUsize::new(0));
let probe = Calls {
cell: cell.clone(),
made: made.clone(),
script: Arc::new(Mutex::new(Script {
reading_at: Some(unix_ms() + 60 * MIN),
status: "allowed".to_string(),
fails_with: None,
})),
};
start(cell.clone(), Box::new(probe)).unwrap();
std::thread::sleep(Duration::from_millis(300));
assert_eq!(made.load(Ordering::SeqCst), 0, "no lease, no call");
let _lease = cell.asking();
let deadline = std::time::Instant::now() + Duration::from_secs(10);
while made.load(Ordering::SeqCst) == 0 {
assert!(std::time::Instant::now() < deadline, "the ask woke nobody");
std::thread::sleep(Duration::from_millis(20));
}
assert_eq!(cell.snapshot_with_week(Duration::ZERO).windows.len(), 1);
}
}

View file

@ -20,7 +20,9 @@ pub enum StdoutRead {
}
/// Starts a `claude` child. `argv[0]` is the program; the environment is
/// exactly `env`; `stdin` is written and closed.
/// exactly `env`; `stdin` is written and closed. `token` is the child's
/// credential: it reaches the child on a pipe at `launch::TOKEN_FD`, once, and
/// is in no environment.
pub trait ProcessPort: Send + Sync {
fn spawn(
&self,
@ -28,6 +30,7 @@ pub trait ProcessPort: Send + Sync {
env: &ChildEnv,
cwd: &Path,
stdin: Vec<u8>,
token: &str,
) -> io::Result<Box<dyn RunningChild>>;
}
@ -47,6 +50,7 @@ pub use unix::StdProcessPort;
#[cfg(unix)]
mod unix {
use std::io::{self, BufRead, BufReader, Read, Write};
use std::os::fd::{AsRawFd, RawFd};
use std::os::unix::process::CommandExt;
use std::path::Path;
use std::process::{Child, Command, Stdio};
@ -57,7 +61,40 @@ mod unix {
use super::{ProcessPort, RunningChild, StdoutRead};
use crate::config::ClaudeCodeLimits;
use crate::launch::ChildEnv;
use crate::launch::{ChildEnv, TOKEN_FD};
/// A pipe holds a token without blocking the writer; a token is a few
/// hundred bytes, so anything near the pipe's capacity is not one.
const MAX_TOKEN_BYTES: usize = 4096;
fn cvt(result: libc::c_int) -> io::Result<libc::c_int> {
if result < 0 {
Err(io::Error::last_os_error())
} else {
Ok(result)
}
}
/// Makes `source` also descriptor `target`, open across exec. Runs between
/// fork and exec in the child, so it calls nothing but `dup2` and `fcntl`.
fn place_fd(source: RawFd, target: RawFd) -> io::Result<()> {
// SAFETY: dup2 and fcntl take integers only and are async-signal-safe.
unsafe {
if source == target {
// dup2 onto itself changes nothing, so the close-on-exec flag
// the pipe was created with must be cleared by hand.
let flags = cvt(libc::fcntl(target, libc::F_GETFD))?;
cvt(libc::fcntl(
target,
libc::F_SETFD,
flags & !libc::FD_CLOEXEC,
))?;
} else {
cvt(libc::dup2(source, target))?;
}
}
Ok(())
}
/// Spawns each child in its own process group, so that killing the call
/// kills what the child started too.
@ -78,7 +115,11 @@ mod unix {
env: &ChildEnv,
cwd: &Path,
stdin: Vec<u8>,
token: &str,
) -> io::Result<Box<dyn RunningChild>> {
if token.len() > MAX_TOKEN_BYTES {
return Err(io::Error::from(io::ErrorKind::InvalidInput));
}
let (program, args) = argv
.split_first()
.ok_or_else(|| io::Error::from(io::ErrorKind::InvalidInput))?;
@ -92,7 +133,17 @@ mod unix {
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.process_group(0);
// The write end closes here, so the child sees the token and then
// the end of the pipe. Both ends are close-on-exec; only the copy
// `place_fd` makes reaches the child.
let (reader, mut writer) = std::io::pipe()?;
writer.write_all(token.as_bytes())?;
drop(writer);
let source = reader.as_raw_fd();
// SAFETY: the closure calls only async-signal-safe functions.
unsafe { command.pre_exec(move || place_fd(source, TOKEN_FD)) };
let mut child = command.spawn()?;
drop(reader);
let (stdin_pipe, stdout, stderr) =
(child.stdin.take(), child.stdout.take(), child.stderr.take());
@ -229,4 +280,83 @@ mod unix {
self.wait();
}
}
#[cfg(test)]
mod tests {
use super::*;
fn cloexec(fd: RawFd) -> bool {
// SAFETY: fcntl on an integer descriptor.
let flags = unsafe { libc::fcntl(fd, libc::F_GETFD) };
assert!(flags >= 0, "descriptor {fd} is not open");
flags & libc::FD_CLOEXEC != 0
}
/// A descriptor at or above 300 on the same pipe, so a test never
/// touches a low number another thread may be using.
fn high_dup(source: RawFd, close_on_exec: bool) -> RawFd {
let command = if close_on_exec {
libc::F_DUPFD_CLOEXEC
} else {
libc::F_DUPFD
};
// SAFETY: fcntl on an integer descriptor.
let fd = unsafe { libc::fcntl(source, command, 300) };
assert!(fd >= 300);
fd
}
/// The three tests below name descriptor numbers by value, and a number
/// one releases is the lowest free one for the next `high_dup`. They
/// hold this for their whole body, so none sees a number another
/// freed or places over one another owns.
static FDS: std::sync::Mutex<()> = std::sync::Mutex::new(());
fn fds() -> std::sync::MutexGuard<'static, ()> {
FDS.lock().unwrap_or_else(|e| e.into_inner())
}
fn close(fd: RawFd) {
// SAFETY: closing a descriptor the test owns.
unsafe { libc::close(fd) };
}
#[test]
fn placing_a_descriptor_on_itself_clears_close_on_exec() {
let _held = fds();
let (reader, _writer) = std::io::pipe().unwrap();
let fd = high_dup(reader.as_raw_fd(), true);
assert!(cloexec(fd));
place_fd(fd, fd).unwrap();
assert!(!cloexec(fd), "it must survive exec");
// One that already survives is left alone, not refused.
place_fd(fd, fd).unwrap();
assert!(!cloexec(fd));
close(fd);
}
#[test]
fn placing_a_descriptor_elsewhere_copies_it_open_across_exec() {
let _held = fds();
let (reader, _writer) = std::io::pipe().unwrap();
let source = high_dup(reader.as_raw_fd(), true);
let target = high_dup(reader.as_raw_fd(), true);
assert!(cloexec(target));
place_fd(source, target).unwrap();
assert!(!cloexec(target), "the copy survives exec");
assert!(cloexec(source), "the source does not");
close(source);
close(target);
}
#[test]
fn a_descriptor_that_is_not_open_is_refused_either_way() {
let _held = fds();
let (reader, _writer) = std::io::pipe().unwrap();
let closed = high_dup(reader.as_raw_fd(), true);
close(closed);
assert!(place_fd(closed, closed).is_err());
assert!(place_fd(closed, closed + 1).is_err());
}
}
}

View file

@ -0,0 +1,405 @@
//! The plan's own usage windows, as the subscription reports them. The
//! subscription decides how much a plan may burn, and one sensor observes it:
//! the CLI's `rate_limit_event` on the connector's own calls, which the
//! `meter` makes when the week's reading goes stale. The latest observation of
//! each window is kept here and published as the usage readout and lent with
//! each lease.
use std::collections::BTreeMap;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Condvar, Mutex};
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use model_connector::{PlanQuota, PlanQuotaSnapshot};
/// The window the weekly pacing reads, and the one the meter keeps fresh.
pub(crate) const WEEK: &str = "seven_day";
/// The windows the pacing reads, and so the only ones the cell keeps.
pub(crate) const PACED: [&str; 2] = ["five_hour", WEEK];
/// How often a lease waiting for its first reading looks at the cell again.
const FIRST_READ_POLL: Duration = Duration::from_millis(50);
/// The latest observation of each plan window, shared by the backend that
/// feeds it and whatever reads it out. Each window carries the time of its
/// own observation: an event that names one window says nothing about the
/// others.
///
/// It also keeps who is asking: a lease request is the swarm's only sign of
/// life the connector sees, and the meter reads the plan only while someone
/// asks.
#[derive(Default)]
pub struct QuotaCell {
windows: Mutex<BTreeMap<String, PlanQuota>>,
asks: Mutex<Asks>,
asked: Condvar,
metered: AtomicBool,
}
#[derive(Default)]
struct Asks {
/// Lease requests in flight or held.
live: usize,
/// Requests ever made; the meter sleeps until it changes.
count: u64,
/// When the latest request came, unix ms.
last_ms: Option<u64>,
}
/// A lease request while it is in flight or held. Dropping it is the lease
/// ending, by whatever road.
pub(crate) struct Asking(Arc<QuotaCell>);
impl Drop for Asking {
fn drop(&mut self) {
let mut asks = self.0.asks.lock().unwrap_or_else(|e| e.into_inner());
asks.live = asks.live.saturating_sub(1);
}
}
impl QuotaCell {
pub fn new() -> Self {
Self::default()
}
/// Notes a lease request, now, and wakes the meter.
pub(crate) fn asking(self: &Arc<Self>) -> Asking {
self.asking_at(unix_ms())
}
pub(crate) fn asking_at(self: &Arc<Self>, now: u64) -> Asking {
{
let mut asks = self.asks.lock().unwrap_or_else(|e| e.into_inner());
asks.live += 1;
asks.count += 1;
asks.last_ms = Some(now);
}
self.asked.notify_all();
Asking(self.clone())
}
/// Whether a lease is held, or one was asked for less than `within` ago.
pub(crate) fn wanted_at(&self, now: u64, within: Duration) -> bool {
let asks = self.asks.lock().unwrap_or_else(|e| e.into_inner());
asks.live > 0
|| asks.last_ms.is_some_and(|last| {
now < last.saturating_add(u64::try_from(within.as_millis()).unwrap_or(u64::MAX))
})
}
pub(crate) fn asks_so_far(&self) -> u64 {
self.asks.lock().unwrap_or_else(|e| e.into_inner()).count
}
/// Sleeps up to `limit`, or until a request beyond the first `seen` comes.
pub(crate) fn sleep_until_asked(&self, seen: u64, limit: Duration) {
let asks = self.asks.lock().unwrap_or_else(|e| e.into_inner());
drop(
self.asked
.wait_timeout_while(asks, limit, |asks| asks.count == seen)
.unwrap_or_else(|e| e.into_inner()),
);
}
/// Declares that a meter reads this cell, so a first lease may wait for it.
pub(crate) fn attach_meter(&self) {
self.metered.store(true, Ordering::SeqCst);
}
/// The windows as they stand, but with no weekly reading yet (a restart
/// leaves the cell empty) a metered cell is given up to `within` to get
/// one, so a lease is not lent "unobserved" for want of a few seconds.
pub(crate) fn snapshot_with_week(&self, within: Duration) -> PlanQuotaSnapshot {
let deadline = Instant::now() + within;
loop {
let snapshot = self.snapshot();
let read = snapshot
.windows
.iter()
.any(|quota| quota.window == WEEK && quota.status != EXPIRED);
let now = Instant::now();
if read || !self.metered.load(Ordering::SeqCst) || now >= deadline {
return snapshot;
}
std::thread::sleep(FIRST_READ_POLL.min(deadline - now));
}
}
/// Records the windows one event reported, stamping each with now. See
/// [`QuotaCell::observe_at`] for how a reading combines with what is held.
pub fn observe(&self, windows: impl IntoIterator<Item = PlanQuota>) {
self.observe_at(unix_ms(), windows);
}
/// Records readings taken at `at` (unix ms). A reading older than the
/// held one is dropped, so a slow call cannot replace what a faster one
/// saw meanwhile. Windows the readings did not mention stay as they were,
/// stamp included.
pub fn observe_at(&self, at: u64, windows: impl IntoIterator<Item = PlanQuota>) {
let mut held = self.windows.lock().unwrap_or_else(|e| e.into_inner());
for mut quota in windows {
quota.observed_at_unix_ms = at;
if held
.get(&quota.window)
.is_none_or(|current| quota.observed_at_unix_ms >= current.observed_at_unix_ms)
{
held.insert(quota.window.clone(), quota);
}
}
}
pub fn snapshot(&self) -> PlanQuotaSnapshot {
self.snapshot_at(unix_ms())
}
/// The windows as of `now`. One whose reset time has passed is reported
/// `expired` with no utilization: the reading described a window that no
/// longer exists.
fn snapshot_at(&self, now: u64) -> PlanQuotaSnapshot {
let held = self.windows.lock().unwrap_or_else(|e| e.into_inner());
PlanQuotaSnapshot {
windows: held
.values()
.cloned()
.map(|mut quota| {
if now >= quota.resets_at_unix_ms {
quota.utilization_permille = None;
quota.status = EXPIRED.to_string();
}
quota
})
.collect(),
}
}
}
/// The status of a window whose reset has passed since it was observed.
pub(crate) const EXPIRED: &str = "expired";
fn unix_ms() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map_or(0, |elapsed| {
u64::try_from(elapsed.as_millis()).unwrap_or(u64::MAX)
})
}
#[cfg(test)]
mod tests {
use super::*;
const RESET: u64 = 4_102_444_800_000;
fn quota(window: &str, permille: u32, observed_at: u64) -> PlanQuota {
PlanQuota {
window: window.to_string(),
utilization_permille: Some(permille),
resets_at_unix_ms: RESET,
status: "allowed".to_string(),
observed_at_unix_ms: observed_at,
}
}
#[test]
fn an_empty_cell_has_observed_nothing() {
assert!(QuotaCell::new().snapshot().windows.is_empty());
}
#[test]
fn each_window_keeps_its_own_latest_observation_and_its_own_time() {
let cell = QuotaCell::new();
cell.observe_at(100, [quota("five_hour", 80, 0), quota("seven_day", 300, 0)]);
cell.observe_at(200, [quota("five_hour", 120, 0)]);
assert_eq!(
cell.snapshot_at(250).windows,
vec![quota("five_hour", 120, 200), quota("seven_day", 300, 100)],
"the window the later event did not name keeps its older stamp"
);
}
#[test]
fn an_event_with_no_windows_changes_nothing() {
let cell = QuotaCell::new();
cell.observe_at(100, [quota("five_hour", 80, 0)]);
cell.observe_at(900, Vec::new());
assert_eq!(
cell.snapshot_at(950).windows,
vec![quota("five_hour", 80, 100)]
);
}
#[test]
fn a_window_past_its_reset_is_expired_not_old_utilization() {
let cell = QuotaCell::new();
cell.observe_at(100, [quota("five_hour", 80, 0)]);
let live = cell.snapshot_at(RESET - 1).windows;
assert_eq!(live, vec![quota("five_hour", 80, 100)]);
let expired = cell.snapshot_at(RESET).windows;
assert_eq!(expired.len(), 1);
assert_eq!(expired[0].status, "expired");
assert_eq!(expired[0].utilization_permille, None);
assert_eq!(
expired[0].observed_at_unix_ms, 100,
"its age stays readable"
);
assert_eq!(expired[0].resets_at_unix_ms, RESET);
}
fn with_status(mut quota: PlanQuota, status: &str) -> PlanQuota {
quota.status = status.to_string();
quota
}
fn held(cell: &QuotaCell) -> (String, Option<u32>, u64) {
let window = cell.snapshot_at(250).windows.remove(0);
(
window.status,
window.utilization_permille,
window.observed_at_unix_ms,
)
}
#[test]
fn a_reading_older_than_the_held_one_is_dropped_and_an_equal_one_replaces() {
let cell = QuotaCell::new();
cell.observe_at(200, [quota("five_hour", 80, 0)]);
cell.observe_at(100, [quota("five_hour", 10, 0)]);
assert_eq!(held(&cell), ("allowed".to_string(), Some(80), 200));
cell.observe_at(200, [quota("five_hour", 90, 0)]);
assert_eq!(held(&cell), ("allowed".to_string(), Some(90), 200));
}
#[test]
fn a_newer_reading_lifts_a_rejection() {
let cell = QuotaCell::new();
cell.observe_at(100, [with_status(quota("five_hour", 1000, 0), "rejected")]);
cell.observe_at(200, [quota("five_hour", 300, 0)]);
assert_eq!(held(&cell), ("allowed".to_string(), Some(300), 200));
}
#[test]
fn the_live_clock_stamps_and_expires() {
let cell = QuotaCell::new();
let before = unix_ms();
let mut past = quota("old", 10, 0);
past.resets_at_unix_ms = 1;
cell.observe([past, quota("five_hour", 80, 0)]);
let after = unix_ms();
let windows = cell.snapshot().windows;
let stamp = windows[0].observed_at_unix_ms;
assert!((before..=after).contains(&stamp));
assert_eq!(windows[0].window, "five_hour");
assert_eq!(windows[0].utilization_permille, Some(80));
assert_eq!(windows[1].status, "expired");
}
const HOUR: u64 = 3_600_000;
#[test]
fn nobody_wants_the_plan_before_anyone_asks() {
let cell = Arc::new(QuotaCell::new());
assert!(!cell.wanted_at(0, Duration::from_secs(900)));
}
#[test]
fn a_lease_in_flight_or_held_is_wanted_for_as_long_as_it_lives() {
let cell = Arc::new(QuotaCell::new());
let lease = cell.asking_at(1_000);
let other = cell.asking_at(2_000);
let day = Duration::from_secs(86_400);
assert!(cell.wanted_at(1_000 + 4 * HOUR, Duration::from_secs(1)));
drop(lease);
assert!(cell.wanted_at(1_000 + 4 * HOUR, Duration::from_secs(1)), "one lease still lives");
drop(other);
assert!(!cell.wanted_at(1_000 + 4 * HOUR, Duration::from_secs(1)));
assert!(cell.wanted_at(2_000 + 4 * HOUR, day), "within the window of the last ask");
}
#[test]
fn a_lease_let_go_is_wanted_until_exactly_the_period_after_its_ask() {
let cell = Arc::new(QuotaCell::new());
drop(cell.asking_at(10_000));
let period = Duration::from_millis(900);
assert!(cell.wanted_at(10_899, period));
assert!(!cell.wanted_at(10_900, period));
}
#[test]
fn each_ask_is_counted_and_wakes_a_sleeper_at_once() {
let cell = Arc::new(QuotaCell::new());
assert_eq!(cell.asks_so_far(), 0);
let sleeper = {
let cell = cell.clone();
std::thread::spawn(move || {
let started = Instant::now();
cell.sleep_until_asked(0, Duration::from_secs(30));
started.elapsed()
})
};
std::thread::sleep(Duration::from_millis(100));
drop(cell.asking());
assert!(sleeper.join().unwrap() < Duration::from_secs(10), "woken, not timed out");
assert_eq!(cell.asks_so_far(), 1);
// An ask that came after `seen` was read ends the sleep before it starts.
let started = Instant::now();
cell.sleep_until_asked(0, Duration::from_secs(30));
assert!(started.elapsed() < Duration::from_secs(10));
// Nothing new: the sleep is the whole limit.
let started = Instant::now();
cell.sleep_until_asked(1, Duration::from_millis(200));
assert!(started.elapsed() >= Duration::from_millis(200));
}
#[test]
fn a_metered_cell_with_no_week_gives_a_first_lease_the_reading_when_it_lands() {
let cell = Arc::new(QuotaCell::new());
cell.attach_meter();
let feeder = {
let cell = cell.clone();
std::thread::spawn(move || {
std::thread::sleep(Duration::from_millis(200));
cell.observe([quota(WEEK, 60, 0)]);
})
};
let plan = cell.snapshot_with_week(Duration::from_secs(20));
feeder.join().unwrap();
assert_eq!(plan.windows.len(), 1, "the lease waited for the week");
assert_eq!(plan.windows[0].window, WEEK);
}
#[test]
fn the_wait_for_the_first_reading_is_bounded() {
let cell = Arc::new(QuotaCell::new());
cell.attach_meter();
let started = Instant::now();
let plan = cell.snapshot_with_week(Duration::from_millis(300));
let waited = started.elapsed();
assert!(plan.windows.is_empty(), "no reading came: unobserved, not hung");
assert!(waited >= Duration::from_millis(300), "{waited:?}");
assert!(waited < Duration::from_secs(10), "{waited:?}");
}
#[test]
fn a_cell_nobody_meters_or_that_holds_the_week_is_never_waited_on() {
let bare = QuotaCell::new();
let started = Instant::now();
assert!(bare.snapshot_with_week(Duration::from_secs(20)).windows.is_empty());
let held = QuotaCell::new();
held.attach_meter();
held.observe([quota(WEEK, 60, 0)]);
assert_eq!(held.snapshot_with_week(Duration::from_secs(20)).windows.len(), 1);
assert!(started.elapsed() < Duration::from_secs(5));
}
#[test]
fn a_week_past_its_reset_is_waited_on_like_no_week() {
let cell = QuotaCell::new();
cell.attach_meter();
let mut over = quota(WEEK, 60, 0);
over.resets_at_unix_ms = 1;
cell.observe([over]);
let started = Instant::now();
cell.snapshot_with_week(Duration::from_millis(300));
assert!(started.elapsed() >= Duration::from_millis(300));
}
}

View file

@ -3,12 +3,19 @@
use std::collections::BTreeMap;
use model_connector::{FailureKind, StopReason, Usage};
use model_connector::{FailureKind, PlanQuota, StopReason, Usage};
use serde::Deserialize;
use serde_json::Value;
use crate::quota::{PACED, QuotaCell};
/// The tool the CLI gives the model to deliver an answer that must match a
/// `--json-schema`. Its call is the structured output, not a tool exchange.
const STRUCTURED_OUTPUT_TOOL: &str = "StructuredOutput";
/// A call that cannot complete, with the stage that failed. `usage` is `None`
/// unless a result event reported it.
/// (unknown) unless a result event reported it or the
/// call never reached the provider (zero).
#[derive(Debug, PartialEq, Eq)]
pub(crate) struct Failure {
pub kind: FailureKind,
@ -17,13 +24,24 @@ pub(crate) struct Failure {
}
impl Failure {
pub(crate) fn unbilled(kind: FailureKind, message: &'static str) -> Self {
/// A failure after the child ran: what the provider billed is unknown, so
/// the usage is unknown.
pub(crate) fn cost_unknown(kind: FailureKind, message: &'static str) -> Self {
Self {
kind,
message,
usage: None,
}
}
/// A failure before anything was sent to the provider: it bills nothing.
pub(crate) fn never_reached(kind: FailureKind, message: &'static str) -> Self {
Self {
kind,
message,
usage: Some(sum_usage(std::iter::empty())),
}
}
}
/// What the stream said, once its result event arrived.
@ -51,12 +69,18 @@ struct Line {
session_id: Option<String>,
model: Option<String>,
message: Option<Message>,
rate_limit_info: Option<RateLimitInfo>,
/// Read by hand: the shape moves with the CLI, and a field that changes
/// type must cost the quota reading, never the call.
rate_limit_info: Option<Value>,
is_error: Option<bool>,
stop_reason: Option<String>,
api_error_status: Option<Value>,
/// Milliseconds spent in API calls; zero means no call was made.
duration_api_ms: Option<u64>,
result: Option<Value>,
errors: Option<Value>,
/// The answer to a `--json-schema` request, an object.
structured_output: Option<Value>,
/// What this turn alone consumed. `modelUsage` is deliberately not read:
/// on a resumed session it carries every earlier turn too.
model_usage_since_last_result: Option<Value>,
@ -73,11 +97,7 @@ struct Block {
#[serde(rename = "type")]
kind: String,
text: Option<String>,
}
#[derive(Deserialize)]
struct RateLimitInfo {
status: Option<String>,
name: Option<String>,
}
/// One model's entry. The four counts are required: an entry that lacks one
@ -95,6 +115,13 @@ struct ModelUsage {
/// One call's stream, line by line.
pub(crate) struct StreamState<'a> {
model: &'a str,
/// Whether the request carried an output schema, so the answer is the
/// result's structured output.
structured: bool,
/// Whether that schema's top type admits an object (absent counts): the
/// structured output must be an object, never a bare scalar or array.
schema_top_is_object: bool,
quota: &'a QuotaCell,
session_id: Option<String>,
texts: Vec<String>,
message_id: Option<String>,
@ -102,9 +129,16 @@ pub(crate) struct StreamState<'a> {
}
impl<'a> StreamState<'a> {
pub(crate) fn new(model: &'a str) -> Self {
/// `schema` is the request's output schema text, when it carried one.
pub(crate) fn new(model: &'a str, schema: Option<&str>, quota: &'a QuotaCell) -> Self {
let top = schema
.and_then(|text| serde_json::from_str::<Value>(text).ok())
.and_then(|schema| schema.get("type").cloned());
Self {
model,
structured: schema.is_some(),
schema_top_is_object: top.is_none_or(|top| top.as_str() == Some("object")),
quota,
session_id: None,
texts: Vec::new(),
message_id: None,
@ -139,10 +173,11 @@ impl<'a> StreamState<'a> {
},
"assistant" => self.assistant(line.message),
"rate_limit_event" => {
self.plan_rejected |= line
.rate_limit_info
.and_then(|info| info.status)
.is_some_and(|status| status == "rejected");
if let Some(info) = &line.rate_limit_info {
self.plan_rejected |=
info.get("status").and_then(Value::as_str) == Some("rejected");
self.quota.observe(plan_windows(info));
}
Step::More
}
"result" => Step::Done(Box::new(self.result(line))),
@ -164,6 +199,9 @@ impl<'a> StreamState<'a> {
self.texts.push(text);
}
}
"tool_use"
if self.structured && block.name.as_deref() == Some(STRUCTURED_OUTPUT_TOOL) => {
}
"tool_use" => {
return fail(
FailureKind::Provider,
@ -181,9 +219,18 @@ impl<'a> StreamState<'a> {
let door_refusal = !success
&& matches!(
line.api_error_status.as_ref().and_then(Value::as_u64),
Some(401 | 403)
Some(401 | 403 | 429)
);
let usage = read_usage(line.model_usage_since_last_result, door_refusal);
// The CLI found no credential and never called the API: an error with
// no status and no API time at all.
let no_call = !success
&& line.duration_api_ms == Some(0)
&& line.api_error_status.as_ref().is_none_or(Value::is_null);
let usage = read_usage(
line.model_usage_since_last_result,
door_refusal || no_call,
no_call,
);
if !success {
return Err(Failure {
kind: self.error_kind(&line.api_error_status, &line.result, &line.errors),
@ -198,8 +245,26 @@ impl<'a> StreamState<'a> {
usage: None,
});
};
let mut texts = std::mem::take(&mut self.texts);
if self.structured {
// The answer is the structured output alone; any prose around the
// tool call is not part of it.
let Some(output) = line
.structured_output
.filter(|value| value.is_object() && self.schema_top_is_object)
else {
return Err(Failure {
kind: FailureKind::Provider,
message: "claude-code: a schema was sent and the result carried no structured output object",
usage: Some(usage),
});
};
texts = vec![output.to_string()];
}
let stop_reason = match line.stop_reason.as_deref() {
Some("end_turn") => StopReason::EndTurn,
// The structured-output tool call is how such a turn ends.
Some("tool_use") if self.structured => StopReason::EndTurn,
Some("max_tokens") => StopReason::MaxOutputTokens,
Some("stop_sequence") => StopReason::StopSequence,
Some("refusal") => StopReason::Refusal,
@ -221,7 +286,7 @@ impl<'a> StreamState<'a> {
};
Ok(Finished {
session_id,
texts: std::mem::take(&mut self.texts),
texts,
message_id: self.message_id.take(),
stop_reason,
usage,
@ -251,20 +316,58 @@ impl<'a> StreamState<'a> {
}
}
/// The plan windows one `rate_limit_event` reports under
/// `rate_limit_info.unifiedWindows`: utilization as a fraction of the window,
/// reset as unix seconds. A window missing any of its three fields is skipped;
/// an unknown reading is never zero.
/// Only the windows pacing reads are taken ([`PACED`]): the event also names
/// per-model buckets, and one of those reading `rejected` must not close every
/// lease.
fn plan_windows(info: &Value) -> Vec<PlanQuota> {
let Some(windows) = info.get("unifiedWindows").and_then(Value::as_object) else {
return Vec::new();
};
windows
.iter()
.filter(|(window, _)| PACED.contains(&window.as_str()))
.filter_map(|(window, reading)| {
let utilization = reading.get("utilization")?.as_f64()?;
let resets_at = reading.get("resetsAt")?.as_u64()?;
let status = reading.get("status")?.as_str()?;
if !utilization.is_finite() || utilization < 0.0 {
return None;
}
Some(PlanQuota {
window: window.clone(),
// `as` saturates a float at the integer's bounds.
utilization_permille: Some((utilization * 1000.0).round() as u32),
resets_at_unix_ms: resets_at.saturating_mul(1000),
status: status.to_string(),
// Stamped by the cell that keeps it.
observed_at_unix_ms: 0,
})
})
.collect()
}
fn fail(kind: FailureKind, message: &'static str) -> Step {
Step::Done(Box::new(Err(Failure::unbilled(kind, message))))
Step::Done(Box::new(Err(Failure::cost_unknown(kind, message))))
}
/// What this turn consumed, summed over every model it used (harness side
/// calls included), or `None` when the result does not say: the field absent,
/// of another shape, or empty. An empty report is believed only for a refusal
/// at the door (401 or 403), which the provider cannot have billed; elsewhere
/// the CLI also zeroes usage on crashes after real work. `None` is billed the
/// reservation by the kit.
fn read_usage(reported: Option<Value>, door_refusal: bool) -> Option<Usage> {
let models: BTreeMap<String, ModelUsage> = serde_json::from_value(reported?).ok()?;
/// of another shape, or empty. An empty report is believed only where the
/// provider cannot have billed (`empty_is_zero`: a refusal at the door, 401,
/// 403 or 429, or a call never made); elsewhere the CLI also zeroes usage on
/// crashes after real work. A call never made (`absent_is_zero`) may also
/// omit the field. `None` is unknown usage.
fn read_usage(reported: Option<Value>, empty_is_zero: bool, absent_is_zero: bool) -> Option<Usage> {
let Some(reported) = reported else {
return absent_is_zero.then(|| sum_usage(std::iter::empty()));
};
let models: BTreeMap<String, ModelUsage> = serde_json::from_value(reported).ok()?;
if models.is_empty() {
return door_refusal.then(|| sum_usage(std::iter::empty()));
return empty_is_zero.then(|| sum_usage(std::iter::empty()));
}
Some(sum_usage(models.into_values()))
}
@ -293,3 +396,50 @@ fn sum_usage(models: impl Iterator<Item = ModelUsage>) -> Usage {
}
usage
}
#[cfg(test)]
mod tests {
use super::*;
fn windows(body: &str) -> Vec<(String, Option<u32>, u64)> {
let info: Value = serde_json::from_str(&format!(r#"{{"unifiedWindows":{body}}}"#)).unwrap();
plan_windows(&info)
.into_iter()
.map(|q| (q.window, q.utilization_permille, q.resets_at_unix_ms))
.collect()
}
#[test]
fn only_the_five_hour_and_seven_day_windows_are_taken() {
let taken = windows(
r#"{"five_hour":{"utilization":0.5,"resetsAt":2,"status":"allowed"},
"seven_day":{"utilization":0.25,"resetsAt":3,"status":"allowed"},
"seven_day_opus":{"utilization":0.1,"resetsAt":4,"status":"allowed"},
"seven_day_overage_included":{"utilization":1,"resetsAt":5,"status":"rejected"}}"#,
);
assert_eq!(
taken,
vec![
("five_hour".to_string(), Some(500), 2000),
("seven_day".to_string(), Some(250), 3000),
]
);
}
#[test]
fn a_window_missing_a_field_or_of_the_wrong_type_costs_only_itself() {
// Unknown is never zero: such a window is skipped, not read as empty.
for broken in [
r#"{"resetsAt":2,"status":"allowed"}"#,
r#"{"utilization":0.5,"status":"allowed"}"#,
r#"{"utilization":0.5,"resetsAt":2}"#,
r#"{"utilization":"high","resetsAt":2,"status":"allowed"}"#,
r#"{"utilization":-0.5,"resetsAt":2,"status":"allowed"}"#,
] {
let both = windows(&format!(
r#"{{"five_hour":{broken},"seven_day":{{"utilization":0,"resetsAt":3,"status":"allowed"}}}}"#
));
assert_eq!(both, vec![("seven_day".to_string(), Some(0), 3000)], "{broken}");
}
}
}

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);
}
}

View file

@ -0,0 +1,297 @@
//! The claude-connector's configuration document and binary: what the
//! document accepts, what it refuses, and that a document of the dollar era
//! fails closed.
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::process::{Command, Output};
use std::sync::atomic::{AtomicUsize, Ordering};
use claude_connector::{
ClaudeConnectorConfig, ConfigError, LeaseCallerConfig, load_config, write_config,
};
use cultcache_rs::{CultCache, DatabaseEntry, SingleFileMessagePackBackingStore};
use model_connector::{CallerConfig, LeaseTurn};
const OPUS: &str = "claude-opus-5-1";
const HAIKU: &str = "claude-haiku-5-5";
static SEQUENCE: AtomicUsize = AtomicUsize::new(0);
struct Scratch(PathBuf);
impl Scratch {
fn new() -> Self {
let dir = std::env::temp_dir().join(format!(
"claude-connector-config-{}-{}",
std::process::id(),
SEQUENCE.fetch_add(1, Ordering::SeqCst)
));
std::fs::create_dir_all(&dir).unwrap();
Self(dir)
}
fn path(&self, name: &str) -> PathBuf {
self.0.join(name)
}
}
impl Drop for Scratch {
fn drop(&mut self) {
let _ = std::fs::remove_dir_all(&self.0);
}
}
fn valid() -> ClaudeConnectorConfig {
ClaudeConnectorConfig {
listen: "127.0.0.1:7443".to_string(),
replay_store: "/var/lib/claude-connector/replay.cc".into(),
claude_binary: "/opt/claude/bin/claude".into(),
claude_binary_sha256: [7; 32],
work_dir: "/var/lib/claude-connector/work".into(),
config_dir: "/var/lib/claude-connector/home".into(),
prompt_dir: "/var/lib/claude-connector/prompts".into(),
instance: "persona-a".to_string(),
token_credential: "claude-token".to_string(),
handle_secret_credential: "claude-handle-secret".to_string(),
deadline_ms: 600_000,
max_line_bytes: 4 << 20,
max_total_bytes: 64 << 20,
stderr_tail_bytes: 8192,
max_frame_bytes: 2 << 20,
max_connections: 8,
socket_timeout_ms: 10_000,
max_expiry_skew_ms: 300_000,
callers: vec![CallerConfig {
caller_runtime_id: "persona-a-caller".to_string(),
connection_key_file: "/etc/claude-connector/caller.key".into(),
connection_key_epoch: 1,
allowed_models: vec![HAIKU.to_string(), OPUS.to_string()],
max_concurrent_requests: 4,
max_payload_bytes: 1 << 20,
max_output_tokens: 32_768,
}],
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(),
turn: LeaseTurn::SelfRun,
unit_instance: "eureka".to_string(),
}],
lease_max_ms: 3_600_000,
}
}
fn refused(config: &ClaudeConnectorConfig) -> &'static str {
match config.validate() {
Err(ConfigError::Invalid(field)) => field,
other => panic!("expected a refusal naming a field, got {other:?}"),
}
}
#[test]
fn config_round_trips_through_cultcache() {
let scratch = Scratch::new();
let path = scratch.path("connector.cc");
let config = valid();
write_config(&path, &config).unwrap();
assert_eq!(load_config(&path).unwrap(), config);
}
#[test]
fn a_valid_document_validates() {
assert!(valid().validate().is_ok());
}
/// A document of the dollar era: the first type id, carrying a price table.
#[derive(Clone, Debug, PartialEq, Eq, DatabaseEntry)]
#[cultcache(type = "gamecult.model_connector.claude_connector_config.v1")]
struct DollarEraConfig {
#[cultcache(key = 0)]
listen: String,
#[cultcache(key = 18)]
prices: BTreeMap<String, u64>,
}
#[test]
fn a_dollar_era_config_is_refused_at_load() {
let scratch = Scratch::new();
let path = scratch.path("old.cc");
{
let mut cache = CultCache::new();
cache.register_entry_type::<DollarEraConfig>().unwrap();
cache
.add_generic_backing_store(SingleFileMessagePackBackingStore::new(&path))
.unwrap();
cache
.put(
"runtime".to_string(),
&DollarEraConfig {
listen: "127.0.0.1:7443".to_string(),
prices: BTreeMap::from([(OPUS.to_string(), 5_000_000)]),
},
)
.unwrap();
}
assert!(matches!(load_config(&path), Err(ConfigError::Store)));
}
#[test]
fn the_known_models_are_the_allowed_models_by_form_alone() {
for good in [
"claude-opus-5-1",
"claude-haiku-5-5",
"claude-opus-5",
"claude-opus-5-1-20260101",
] {
let mut config = valid();
config.callers[0].allowed_models = vec![good.to_string()];
assert!(config.validate().is_ok(), "{good}");
}
for bad in [
"sonnet",
"claude-",
"claude-sonnet-latest",
"claude-Opus-5-1",
"claude-opus-5-1-2026",
"claude-opus-5-1-20260101-x",
] {
let mut config = valid();
config.callers[0].allowed_models = vec![bad.to_string()];
assert_eq!(refused(&config), "allowed_models", "{bad}");
}
let mut twice = valid();
twice.callers[0].allowed_models = vec![OPUS.to_string(), OPUS.to_string()];
assert_eq!(refused(&twice), "allowed_models");
let mut none = valid();
none.callers[0].allowed_models.clear();
assert_eq!(refused(&none), "allowed_models");
}
#[test]
fn caller_bounds_are_required_nonzero() {
let edits: [fn(&mut ClaudeConnectorConfig); 3] = [
|c| c.callers[0].max_concurrent_requests = 0,
|c| c.callers[0].max_output_tokens = 0,
|c| c.callers[0].max_payload_bytes = 0,
];
for edit in edits {
let mut config = valid();
edit(&mut config);
assert_eq!(refused(&config), "caller bounds");
}
}
#[test]
fn a_lease_caller_names_a_configured_caller_and_a_socket_exists_exactly_when_it_leases() {
let mut stranger = valid();
stranger.lease_callers[0].caller_id = "nobody".to_string();
assert_eq!(refused(&stranger), "lease_callers");
let mut no_callers = valid();
no_callers.lease_callers.clear();
assert_eq!(refused(&no_callers), "lease_callers");
let mut no_socket = valid();
no_socket.lease_socket = None;
assert_eq!(refused(&no_socket), "lease_callers");
no_socket.lease_callers.clear();
assert!(no_socket.validate().is_ok());
let mut twice = valid();
twice.lease_callers.push(twice.lease_callers[0].clone());
assert_eq!(refused(&twice), "lease_callers");
}
fn run(args: &[&str]) -> Output {
Command::new(env!("CARGO_BIN_EXE_claude-connector"))
.args(args)
.env_clear()
.output()
.unwrap()
}
fn init_flags(scratch: &Scratch, extra: &[&str]) -> Vec<String> {
let path = |name: &str| scratch.path(name).to_string_lossy().into_owned();
let mut flags: Vec<String> = [
"--listen",
"127.0.0.1:0",
"--replay-store",
&path("replay.cc"),
"--claude-binary",
&path("claude"),
"--claude-sha256",
&"ab".repeat(32),
"--work-dir",
&path("work"),
"--config-dir",
&path("home"),
"--prompt-dir",
&path("prompts"),
"--instance",
"persona-a",
"--token-credential",
"claude-token",
"--handle-secret-credential",
"claude-handle-secret",
"--handle-secret-file",
&path("handle.secret"),
"--caller-id",
"persona-a-caller",
"--connection-key-file",
&path("caller.key"),
"--connection-key-epoch",
"1",
"--allowed-model",
HAIKU,
]
.map(String::from)
.to_vec();
flags.extend(extra.iter().map(|flag| flag.to_string()));
flags
}
fn init(scratch: &Scratch, flags: &[String]) -> Output {
let config = scratch.path("connector.cc");
let mut args = vec!["config", "init", config.to_str().unwrap()];
args.extend(flags.iter().map(String::as_str));
run(&args)
}
#[test]
fn config_init_writes_the_flags_into_the_document() {
let scratch = Scratch::new();
let made = init(&scratch, &init_flags(&scratch, &[]));
assert!(made.status.success(), "{:?}", made);
let config = load_config(&scratch.path("connector.cc")).unwrap();
assert_eq!(config.callers[0].allowed_models, [HAIKU]);
assert_eq!(config.claude_binary_sha256, [0xab; 32]);
assert!(config.lease_callers.is_empty());
}
#[test]
fn the_dollar_flags_are_not_flags() {
let dollar_flags: [&[&str]; 4] = [
&["--spend-ceiling-micro-usd", "50000000"],
&["--spend-window-ms", "3600000"],
&["--price", "claude-haiku-5-5:1:1:1:1:1"],
&["--lease-reservation-micro-usd", "5000000"],
];
for extra in dollar_flags {
let scratch = Scratch::new();
let made = init(&scratch, &init_flags(&scratch, extra));
assert!(!made.status.success(), "{extra:?} was accepted");
assert!(
!Path::new(&scratch.path("connector.cc")).exists(),
"{extra:?}"
);
}
}
#[test]
fn serve_refuses_a_document_that_is_not_there() {
let scratch = Scratch::new();
let missing = scratch.path("missing.cc");
let served = run(&["serve", "--config", missing.to_str().unwrap()]);
assert!(!served.status.success(), "{served:?}");
assert!(!String::from_utf8_lossy(&served.stderr).is_empty());
}

View file

@ -0,0 +1,4 @@
{"type":"system","subtype":"init","cwd":"/tmp/meter-probe.Ye68uq/work","session_id":"1c3d29e9-8277-4ee0-86df-efc86d994f01","tools":[],"mcp_servers":[],"model":"claude-haiku-5-5","permissionMode":"bypassPermissions","slash_commands":[],"apiKeySource":"none","claude_code_version":"2.1.312","output_style":"default","agents":["claude","Explore","general-purpose","Plan","statusline-setup"],"skills":[],"plugins":[{"name":"cc-plugin-sec-default","path":"builtin","source":"cc-plugin-sec-default@builtin"},{"name":"cc-plugin-agents-md","path":"builtin","source":"cc-plugin-agents-md@builtin"},{"name":"cc-plugin-telemetry","path":"builtin","source":"cc-plugin-telemetry@builtin"},{"name":"cc-plugin-plugin-authoring","path":"builtin","source":"cc-plugin-plugin-authoring@builtin"}],"capabilities":["interrupt_receipt_v1","interrupt_cancel_queued_v1","interrupt_send_now_v1","interrupt_requeue_taken_v1","msg_lifecycle_v1","request_marker_lists_v1","ask_user_question_unanswered_v1","lifecycle_wakeup_source_v1","sdk_mcp_tools_list_changed","sdk_mcp_manifests","mcp_read_resource_v1","mcp_tool_ui_meta_v1","ui_surface_v1","mcp_call_result_not_for_model_v1"],"analytics_disabled":false,"product_feedback_disabled":false,"uuid":"8ad7f8d7-2831-43f2-a221-8af6323dcbd9","memory_paths":{"auto":"/tmp/meter-probe.Ye68uq/home/projects/-tmp-meter-probe-Ye68uq-work/memory/"},"original_cwd":"/tmp/meter-probe.Ye68uq/work","additional_directories":[],"messaging_socket_path":"/tmp/cc-socks/2838104.sock","fast_mode_state":"off","fast_mode_disabled_reason":"sdk_opt_in_required","per_turn_effort_active":true,"view_mode":"default"}
{"type":"assistant","message":{"model":"claude-haiku-5-5","id":"msg_011Cgfu1sQvRAVmjwWNHc58s","type":"message","role":"assistant","content":[{"type":"text","text":"ok"}],"container":null,"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":384,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"cache_creation":{"ephemeral_5m_input_tokens":0,"ephemeral_1h_input_tokens":0},"output_tokens":4,"service_tier":"standard","inference_geo":"global","speed":"standard"},"input_transformations":[],"diagnostics":null,"context_management":null},"parent_tool_use_id":null,"session_id":"1c3d29e9-8277-4ee0-86df-efc86d994f01","uuid":"f80db687-b098-415c-ac44-51c40a6cb63a","timestamp":"2026-11-03T11:14:35.024Z","request_id":"req_011Cgfu1sCXBsDrfz8a1LoMd"}
{"type":"rate_limit_event","rate_limit_info":{"status":"allowed","resetsAt":1794261600,"rateLimitType":"seven_day","overageStatus":"rejected","overageDisabledReason":"org_level_disabled","isUsingOverage":false,"unifiedWindows":{"seven_day":{"utilization":0.06,"resetsAt":1794261600,"status":"allowed"}}},"uuid":"acb3b364-a218-4965-b967-f73028254f8b","session_id":"1c3d29e9-8277-4ee0-86df-efc86d994f01"}
{"duration_api_ms":995,"stop_reason":"end_turn","session_id":"1c3d29e9-8277-4ee0-86df-efc86d994f01","total_cost_usd":0.000040400000000000006,"usage":{"input_tokens":384,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"output_tokens":4,"output_tokens_details":{"thinking_tokens":0},"server_tool_use":{"web_search_requests":0,"web_fetch_requests":0},"service_tier":"standard","cache_creation":{"ephemeral_1h_input_tokens":0,"ephemeral_5m_input_tokens":0},"inference_geo":"global","iterations":[{"input_tokens":384,"output_tokens":4,"cache_read_input_tokens":0,"cache_creation_input_tokens":0,"cache_creation":{"ephemeral_5m_input_tokens":0,"ephemeral_1h_input_tokens":0},"type":"message"}],"speed":"standard","fallback_credit":null},"modelUsage":{"claude-haiku-5-5":{"inputTokens":384,"outputTokens":4,"cacheReadInputTokens":0,"cacheCreationInputTokens":0,"webSearchRequests":0,"costUSD":0.000040400000000000006,"contextWindow":1000000,"maxOutputTokens":128000,"thinkingTokens":0,"cacheCreationEphemeral5mInputTokens":0,"cacheCreationEphemeral1hInputTokens":0,"canonicalModel":"claude-haiku-5-5","provider":"firstParty","costBasis":"list"}},"permission_denials":[],"terminal_reason":"completed","fast_mode_state":"off","fast_mode_disabled_reason":"sdk_opt_in_required","subagent_stats":{"spawned":0,"requested":{"background":0,"foreground":0,"unset":0},"started_in_background":0,"max_depth":0,"spawned_by_subagents":0,"completed":0,"failed":0,"killed":{"parent":0,"user":0,"system":0},"refused":{"depth_limit":0,"concurrency_limit":0,"budget":0},"by_type":{}},"safety_stops":0,"is_error":false,"num_turns":1,"subtype":"success","api_error_status":null,"result":"ok","ttft_ms":1210,"type":"result","duration_ms":1228,"uuid":"d0b1b31b-1428-45aa-b90e-d35903c9b91b","ttft_stream_ms":1209,"time_to_request_ms":232,"first_content_frame_ms":1209,"queued_turn_count":0,"result_index":0,"model_usage_since_last_result":{"claude-haiku-5-5":{"inputTokens":384,"outputTokens":4,"cacheReadInputTokens":0,"cacheCreationInputTokens":0,"webSearchRequests":0,"costUSD":0.000040400000000000006,"contextWindow":1000000,"maxOutputTokens":128000,"thinkingTokens":0,"cacheCreationEphemeral5mInputTokens":0,"cacheCreationEphemeral1hInputTokens":0,"canonicalModel":"claude-haiku-5-5","provider":"firstParty","costBasis":"list"}}}

View file

@ -0,0 +1,16 @@
# lease_v2 fixtures
The exact bytes of the claude lease protocol v2: each file is one whole frame (a u32
big-endian length, then the named-MessagePack map `{kind, value}`).
- `request.bin`: `claude_lease_request`.
- `refusal_unknown_caller.bin`, `refusal_expired.bin`, `refusal_busy.bin`,
`refusal_provider_rejected.bin`: `claude_lease_refusal`; `kind` is the variant name as a
PascalCase string.
- `grant.bin`: `claude_lease_grant` with `lease_id` and `plan_quota` carrying `five_hour` and
`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.
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.

Binary file not shown.

Binary file not shown.

Binary file not shown.

Binary file not shown.

View file

@ -0,0 +1,651 @@
//! The Self-run credential lease over a real unix socket: who may lease, where
//! the token goes, what ends a lease, and the one usage-derived refusal, the
//! provider's own. Nothing is reserved and nothing is billed. The unit a peer
//! runs in is the `PeerUnits` port, answered here by a fake.
#![cfg(target_os = "linux")]
use std::fs::File;
use std::io::{Read, Write};
use std::os::fd::{AsRawFd, BorrowedFd, FromRawFd, OwnedFd, RawFd};
use std::os::unix::net::{UnixListener, UnixStream};
use std::path::PathBuf;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
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,
};
use model_connector::{CallerAdmission, LeaseTurn, PlanQuota, TransportService};
use zeroize::Zeroizing;
const CALLER: &str = "eureka";
const MODEL: &str = "claude-opus-5-1";
const TOKEN: &str = "sk-canary-lease-4c1f9a";
const INSTANCE: &str = "eureka";
const UNIT: &str = "mind-self@eureka.service";
static SEQUENCE: AtomicUsize = AtomicUsize::new(0);
fn now_ms() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_millis() as u64
}
fn own_uid() -> u32 {
// SAFETY: geteuid has no preconditions and cannot fail.
unsafe { libc::geteuid() }
}
/// The 32-digit invocation id of the test's `n`th unit invocation.
fn invocation(n: u64) -> String {
format!("{n:032x}")
}
type Probe = Box<dyn Fn() + Send + Sync>;
/// A token source that counts its reads and can run a probe at the moment of
/// each one.
struct FakeTokens {
token: Option<&'static str>,
reads: AtomicUsize,
probe: Mutex<Option<Probe>>,
}
impl TokenSource for FakeTokens {
fn read(&self) -> Option<Zeroizing<String>> {
self.reads.fetch_add(1, Ordering::SeqCst);
if let Some(probe) = self.probe.lock().unwrap().as_ref() {
probe();
}
self.token.map(|token| Zeroizing::new(token.to_string()))
}
}
/// The system's record of where every peer runs, as the test sets it, and the
/// pid each handle it was asked about names (by the kernel's own record).
struct FakeUnits {
peer: Mutex<Option<PeerUnit>>,
asked_about: Mutex<Vec<(String, String)>>,
}
impl PeerUnits for FakeUnits {
fn invocation_in(&self, process: BorrowedFd<'_>, unit: &str) -> Option<PeerUnit> {
let info =
std::fs::read_to_string(format!("/proc/self/fdinfo/{}", process.as_raw_fd())).unwrap();
let pid = info
.lines()
.find_map(|line| line.strip_prefix("Pid:"))
.unwrap()
.trim()
.to_string();
self.asked_about
.lock()
.unwrap()
.push((pid, unit.to_string()));
self.peer.lock().unwrap().clone()
}
}
struct Harness {
dir: PathBuf,
socket: PathBuf,
_service: Arc<Mutex<TransportService>>,
quota: Arc<QuotaCell>,
tokens: Arc<FakeTokens>,
/// The fake behind the listener, when it has one.
units: Option<Arc<FakeUnits>>,
}
impl Harness {
/// The peers connect from invocation `n` of the Self unit.
fn enter(&self, n: u64) {
self.run_in(UNIT, n);
}
fn run_in(&self, unit: &str, n: u64) {
*self.units.as_ref().unwrap().peer.lock().unwrap() = Some(PeerUnit {
unit: unit.to_string(),
invocation_id: invocation(n),
});
}
}
impl Drop for Harness {
fn drop(&mut self) {
let _ = std::fs::remove_dir_all(&self.dir);
}
}
fn receive(stream: &UnixStream) -> (Vec<u8>, Vec<OwnedFd>) {
use std::os::fd::AsRawFd;
let mut buffer = vec![0_u8; 70_000];
let mut control = vec![0_u64; 32];
let mut vector = libc::iovec {
iov_base: buffer.as_mut_ptr().cast(),
iov_len: buffer.len(),
};
// SAFETY: the buffers outlive the call and the lengths describe them; the
// descriptors read out of the control messages were just installed in this
// process by the kernel and are owned by nothing else.
unsafe {
let mut message: libc::msghdr = std::mem::zeroed();
message.msg_iov = &mut vector;
message.msg_iovlen = 1;
message.msg_control = control.as_mut_ptr().cast();
message.msg_controllen = (control.len() * 8) as _;
let received = libc::recvmsg(stream.as_raw_fd(), &mut message, libc::MSG_CMSG_CLOEXEC);
if received < 0 {
// A listener that drops a connection with our request unread
// resets it: that is end of file for this test's purposes.
let error = std::io::Error::last_os_error();
assert_eq!(error.kind(), std::io::ErrorKind::ConnectionReset, "{error}");
return (Vec::new(), Vec::new());
}
let mut fds = Vec::new();
let mut header = libc::CMSG_FIRSTHDR(&message);
while !header.is_null() {
if (*header).cmsg_level == libc::SOL_SOCKET && (*header).cmsg_type == libc::SCM_RIGHTS {
let bytes = (*header).cmsg_len as usize - libc::CMSG_LEN(0) as usize;
let data = libc::CMSG_DATA(header).cast::<RawFd>();
for index in 0..bytes / std::mem::size_of::<RawFd>() {
fds.push(OwnedFd::from_raw_fd(data.add(index).read_unaligned()));
}
}
header = libc::CMSG_NXTHDR(&message, header);
}
buffer.truncate(received as usize);
(buffer, fds)
}
}
fn request(ttl_ms: u64) -> LeaseMessage {
LeaseMessage::ClaudeLeaseRequest(LeaseRequest {
expires_at_unix_ms: now_ms() + ttl_ms,
})
}
/// What one request to the listener came back with.
struct Answer {
_stream: UnixStream,
raw: Vec<u8>,
message: Option<LeaseMessage>,
fds: Vec<OwnedFd>,
}
fn ask(harness: &Harness, ttl_ms: u64) -> Answer {
ask_at(&harness.socket, ttl_ms)
}
fn ask_at(socket: &std::path::Path, ttl_ms: u64) -> Answer {
let mut stream = UnixStream::connect(socket).unwrap();
stream
.set_read_timeout(Some(Duration::from_secs(10)))
.unwrap();
// A listener at its limit has already dropped us: the write may fail.
let _ = stream.write_all(&encode_frame(&request(ttl_ms)));
let (raw, fds) = receive(&stream);
let message = decode_frame(&raw);
Answer {
_stream: stream,
raw,
message,
fds,
}
}
fn refusal_of(answer: &Answer) -> Option<LeaseRefusalKind> {
match &answer.message {
Some(LeaseMessage::ClaudeLeaseRefusal(LeaseRefused { kind, .. })) => Some(*kind),
_ => None,
}
}
fn granted(answer: &Answer) -> LeaseGranted {
match &answer.message {
Some(LeaseMessage::ClaudeLeaseGrant(grant)) => grant.clone(),
other => panic!("expected a grant, got {other:?}"),
}
}
fn granted_ok(answer: &Answer) -> bool {
matches!(answer.message, Some(LeaseMessage::ClaudeLeaseGrant(_))) && answer.fds.len() == 1
}
fn contains(haystack: &[u8], needle: &str) -> bool {
haystack
.windows(needle.len())
.any(|window| window == needle.as_bytes())
}
/// What a harness is built from.
struct Setup {
uid: u32,
turn: LeaseTurn,
token: Option<&'static str>,
max_connections: usize,
}
fn setup() -> Setup {
Setup {
uid: own_uid(),
turn: LeaseTurn::SelfRun,
token: Some(TOKEN),
max_connections: 8,
}
}
/// A listener serving `uid` as the caller `eureka`. Peers start in invocation
/// 1 of the Self unit.
fn start(setup: Setup) -> Harness {
let units = Arc::new(FakeUnits {
peer: Mutex::new(None),
asked_about: Mutex::new(Vec::new()),
});
let dir = std::env::temp_dir().join(format!(
"claude-lease-{}-{}",
std::process::id(),
SEQUENCE.fetch_add(1, Ordering::SeqCst)
));
std::fs::create_dir_all(&dir).unwrap();
let socket = dir.join("lease.sock");
let service = Arc::new(Mutex::new(
TransportService::open(
&dir.join("replay.cc"),
"claude-code",
[CallerAdmission::new(
CALLER,
"lease-test-connection-key",
1,
[MODEL.to_string()],
1,
65_536,
32_768,
)
.unwrap()],
300_000,
)
.unwrap(),
));
let tokens = Arc::new(FakeTokens {
token: setup.token,
reads: AtomicUsize::new(0),
probe: Mutex::new(None),
});
let quota = Arc::new(QuotaCell::new());
let listener = UnixListener::bind(&socket).unwrap();
let serving = (
service.clone(),
quota.clone(),
tokens.clone(),
units.clone(),
);
let caller = LeaseCaller {
uid: setup.uid,
caller_id: CALLER.to_string(),
turn: setup.turn,
unit_instance: INSTANCE.to_string(),
};
std::thread::spawn(move || {
let _ = serve_leases(
&listener,
vec![caller],
LeaseLimits {
socket_timeout: Duration::from_secs(2),
max_lease_ms: 3_600_000,
max_connections: setup.max_connections,
},
serving.0,
serving.1,
serving.2,
serving.3,
);
});
let harness = Harness {
dir,
socket,
_service: service,
quota,
tokens,
units: Some(units),
};
harness.enter(1);
harness
}
fn window(name: &str, permille: u32, status: &str, resets_at_unix_ms: u64) -> PlanQuota {
PlanQuota {
window: name.to_string(),
utilization_permille: Some(permille),
resets_at_unix_ms,
status: status.to_string(),
observed_at_unix_ms: 0,
}
}
/// Asks as fresh unit invocations until a lease is granted, or panics.
fn ask_until_granted(harness: &Harness, first_invocation: u64) -> Answer {
let deadline = Instant::now() + Duration::from_secs(8);
let mut n = first_invocation;
loop {
harness.enter(n);
let answer = ask(harness, 60_000);
if granted_ok(&answer) {
return answer;
}
assert!(Instant::now() < deadline, "no lease was ever granted");
n += 1;
std::thread::sleep(Duration::from_millis(100));
}
}
#[test]
fn a_grant_passes_one_fd_holding_the_token_then_end_of_file() {
let harness = start(setup());
let mut answer = ask(&harness, 60_000);
assert_eq!(
granted(&answer).lease_id,
format!("lease:{}", invocation(1))
);
assert_eq!(answer.fds.len(), 1);
let mut read = Vec::new();
File::from(answer.fds.remove(0))
.read_to_end(&mut read)
.unwrap();
assert_eq!(read, TOKEN.as_bytes());
assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 1);
assert!(
!contains(&answer.raw, TOKEN),
"the token is in the pipe only"
);
}
#[test]
fn a_grant_carries_the_plan_usage_as_the_connector_last_saw_it() {
let harness = start(setup());
let reset = now_ms() + 3_600_000;
harness
.quota
.observe([window("five_hour", 420, "allowed", reset)]);
let answer = ask(&harness, 60_000);
let quota = granted(&answer).plan_quota;
assert_eq!(quota.len(), 1);
assert_eq!(quota[0].window, "five_hour");
assert_eq!(quota[0].utilization_permille, Some(420));
assert_eq!(quota[0].status, "allowed");
assert_eq!(quota[0].resets_at_unix_ms, reset);
}
#[test]
fn a_lease_ends_when_its_holder_closes() {
let harness = start(setup());
let first = ask(&harness, 60_000);
assert!(granted_ok(&first));
harness.enter(2);
let second = ask(&harness, 60_000);
assert_eq!(refusal_of(&second), Some(LeaseRefusalKind::Busy));
assert!(second.fds.is_empty());
drop(first);
let third = ask_until_granted(&harness, 3);
assert!(granted_ok(&third));
}
#[test]
fn a_lease_ends_when_its_unit_invocation_ends() {
let harness = start(setup());
let first = ask(&harness, 60_000);
assert!(granted_ok(&first));
// The invocation is over: the system no longer places the peer in it. The
// holder still has its socket open, and the lease does not outlive it.
*harness.units.as_ref().unwrap().peer.lock().unwrap() = None;
std::thread::sleep(Duration::from_millis(1_200));
harness.enter(2);
assert!(granted_ok(&ask(&harness, 60_000)));
drop(first);
}
#[test]
fn a_lease_ends_at_its_expiry_though_its_holder_is_still_there() {
let harness = start(setup());
let first = ask(&harness, 700);
assert!(granted_ok(&first));
let next = ask_until_granted(&harness, 2);
assert!(granted_ok(&next));
drop(first);
}
#[test]
fn a_rejected_window_refuses_the_lease_until_its_reset() {
let harness = start(setup());
let reset = now_ms() + 1_500;
harness
.quota
.observe([window("five_hour", 1_000, "rejected", reset)]);
let refused = ask(&harness, 60_000);
assert_eq!(
refusal_of(&refused),
Some(LeaseRefusalKind::ProviderRejected)
);
assert!(refused.fds.is_empty());
assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 0);
// At the provider's own reset the refusal lifts, with no one lifting it.
let lifted = ask_until_granted(&harness, 2);
assert!(granted_ok(&lifted));
assert!(now_ms() >= reset);
}
#[test]
fn high_utilization_refuses_nothing() {
let harness = start(setup());
let reset = now_ms() + 3_600_000;
harness.quota.observe([
window("five_hour", 999, "allowed", reset),
window("seven_day", 1_000, "allowed_warning", reset),
]);
assert!(granted_ok(&ask(&harness, 60_000)));
}
#[test]
fn a_peer_uid_outside_the_lease_callers_is_refused_and_nothing_is_read() {
let harness = start(Setup {
uid: own_uid().wrapping_add(1),
..setup()
});
let answer = ask(&harness, 60_000);
assert_eq!(refusal_of(&answer), Some(LeaseRefusalKind::UnknownCaller));
assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 0);
}
#[test]
fn the_configured_uid_outside_the_self_unit_is_refused_and_nothing_is_read() {
let harness = start(setup());
harness.run_in("mind-persona@eureka.service", 1);
let answer = ask(&harness, 60_000);
assert_eq!(refusal_of(&answer), Some(LeaseRefusalKind::UnknownCaller));
assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 0);
}
#[test]
fn an_expiry_in_the_past_or_past_the_bound_is_refused() {
let harness = start(setup());
let past = ask(&harness, 0);
assert_eq!(refusal_of(&past), Some(LeaseRefusalKind::Expired));
let far = ask(&harness, 3_600_001 + 60_000);
assert_eq!(refusal_of(&far), Some(LeaseRefusalKind::Expired));
assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 0);
}
#[test]
fn a_token_that_cannot_be_read_is_never_lent_and_frees_the_lease() {
let harness = start(Setup {
token: None,
..setup()
});
let first = ask(&harness, 60_000);
assert!(first.raw.is_empty() && first.fds.is_empty());
harness.enter(2);
let again = ask(&harness, 60_000);
assert!(again.message.is_none() && again.fds.is_empty(), "not Busy");
assert!(harness.tokens.reads.load(Ordering::SeqCst) >= 2);
}
#[test]
fn the_protocol_has_no_settle_frame() {
let grant = encode_frame(&LeaseMessage::ClaudeLeaseGrant(LeaseGranted {
lease_id: "lease:x".to_string(),
plan_quota: Vec::new(),
}));
let value: rmpv::Value = rmp_serde::from_slice(&grant[4..]).unwrap();
assert_eq!(
value,
rmpv::Value::Map(vec![
("kind".into(), "claude_lease_grant".into()),
(
"value".into(),
rmpv::Value::Map(vec![
("lease_id".into(), "lease:x".into()),
("plan_quota".into(), rmpv::Value::Array(Vec::new())),
])
),
])
);
// A frame of the old kind is no message at all.
let old = rmpv::Value::Map(vec![
("kind".into(), "claude_lease_settle".into()),
(
"value".into(),
rmpv::Value::Map(vec![
("lease_id".into(), "lease:x".into()),
("usage".into(), rmpv::Value::Array(Vec::new())),
]),
),
]);
let mut body = Vec::new();
rmpv::encode::write_value(&mut body, &old).unwrap();
let mut frame = (body.len() as u32).to_be_bytes().to_vec();
frame.extend(body);
assert_eq!(decode_frame(&frame), None);
}
#[test]
fn connections_past_the_limit_are_dropped_and_finished_ones_free_their_place() {
let harness = start(Setup {
max_connections: 1,
..setup()
});
let first = ask(&harness, 60_000);
assert!(granted_ok(&first));
let second = ask(&harness, 60_000);
assert!(second.message.is_none() && second.fds.is_empty());
drop(first);
assert!(granted_ok(&ask_until_granted(&harness, 2)));
}
#[test]
fn a_lease_survives_many_polls_while_its_holder_and_invocation_stay() {
let harness = start(setup());
let mut first = ask(&harness, 60_000);
assert!(granted_ok(&first));
// Three polls later the lease still holds: a second invocation is Busy.
std::thread::sleep(Duration::from_millis(1_700));
harness.enter(2);
let second = ask(&harness, 60_000);
assert_eq!(refusal_of(&second), Some(LeaseRefusalKind::Busy));
// The system still places the holder in invocation 1.
harness.enter(1);
// The holder talking does not end it either.
first._stream.write_all(b"x").unwrap();
std::thread::sleep(Duration::from_millis(1_200));
assert_eq!(
refusal_of(&ask(&harness, 60_000)),
Some(LeaseRefusalKind::Busy)
);
}
#[test]
fn a_lease_ends_when_the_unit_runs_a_different_invocation() {
let harness = start(setup());
let first = ask(&harness, 60_000);
assert!(granted_ok(&first));
// The same unit, a new invocation: the one that leased is over.
harness.enter(2);
std::thread::sleep(Duration::from_millis(1_200));
assert!(granted_ok(&ask(&harness, 60_000)));
drop(first);
}
#[test]
fn a_lease_asks_the_system_about_its_holders_process_for_its_own_unit() {
let harness = start(setup());
let first = ask(&harness, 60_000);
assert!(granted_ok(&first));
std::thread::sleep(Duration::from_millis(1_200));
let asked = harness
.units
.as_ref()
.unwrap()
.asked_about
.lock()
.unwrap()
.clone();
assert!(asked.len() >= 3, "the hold polled: {asked:?}");
assert!(asked.iter().all(|(_, unit)| unit == UNIT));
}
#[test]
fn an_account_resolves_to_its_uid_and_a_missing_one_to_none() {
assert_eq!(claude_connector::user_uid("root"), Some(0));
assert_eq!(claude_connector::user_uid("no-such-account-4c1f9a"), None);
assert_eq!(claude_connector::user_uid("bad\0name"), None);
}
#[test]
fn a_token_up_to_one_atomic_pipe_write_is_lent_and_a_longer_one_is_not() {
let exact: &'static str = Box::leak("a".repeat(4096).into_boxed_str());
let harness = start(Setup {
token: Some(exact),
..setup()
});
let mut answer = ask(&harness, 60_000);
assert!(granted_ok(&answer));
let mut read = Vec::new();
File::from(answer.fds.remove(0))
.read_to_end(&mut read)
.unwrap();
assert_eq!(read.len(), 4096);
let long: &'static str = Box::leak("a".repeat(4097).into_boxed_str());
let harness = start(Setup {
token: Some(long),
..setup()
});
let refused = ask(&harness, 60_000);
assert!(refused.message.is_none() && refused.fds.is_empty());
}
#[test]
fn a_request_trickled_past_the_socket_timeout_is_dropped_without_a_lease() {
let harness = start(setup());
let mut stream = UnixStream::connect(&harness.socket).unwrap();
stream
.set_read_timeout(Some(Duration::from_secs(8)))
.unwrap();
stream.write_all(&[0, 0]).unwrap();
let started = Instant::now();
let mut byte = [0_u8; 1];
let outcome = stream.read(&mut byte);
// The listener gave up at its own two-second bound and closed the
// connection: end of file or a reset, never our own read timing out.
assert!(
matches!(outcome, Ok(0))
|| outcome
.as_ref()
.is_err_and(|e| e.kind() == std::io::ErrorKind::ConnectionReset),
"{outcome:?}"
);
assert!(started.elapsed() < Duration::from_secs(6));
assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 0);
}

View file

@ -0,0 +1,99 @@
//! Pins the bytes of the lease protocol v2: every frame a lease client sends
//! or receives. The fixtures in `fixtures/lease_v2` are the contract's bytes
//! for other implementations (Huginn's lease client vendors them and tests
//! against the same bytes); the types in `claude_connector::lease` own the
//! shape. Regenerate with `LEASE_V2_REGENERATE=1`, then review the diff.
#![cfg(target_os = "linux")]
use std::path::{Path, PathBuf};
use claude_connector::{
LeaseGranted, LeaseMessage, LeaseRefusalKind, LeaseRefused, LeaseRequest, decode_frame,
encode_frame,
};
use model_connector::PlanQuota;
const REGENERATE: &str = "LEASE_V2_REGENERATE";
fn fixtures() -> PathBuf {
Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/lease_v2")
}
fn assert_frame(name: &str, message: &LeaseMessage) {
let path = fixtures().join(name);
let bytes = encode_frame(message);
if std::env::var_os(REGENERATE).is_some() {
std::fs::create_dir_all(fixtures()).unwrap();
std::fs::write(&path, &bytes).unwrap();
}
let committed = std::fs::read(&path).unwrap_or_else(|_| panic!("{name} is missing"));
assert_eq!(committed, bytes, "{name} is not what the types encode");
assert_eq!(decode_frame(&committed).as_ref(), Some(message), "{name}");
}
#[test]
fn the_request_frame_is_pinned() {
assert_frame(
"request.bin",
&LeaseMessage::ClaudeLeaseRequest(LeaseRequest {
expires_at_unix_ms: 1_793_307_000_000,
}),
);
}
#[test]
fn every_refusal_frame_is_pinned_with_its_kind_as_a_pascal_case_string() {
let kinds = [
(
"refusal_unknown_caller.bin",
LeaseRefusalKind::UnknownCaller,
"UnknownCaller",
),
("refusal_expired.bin", LeaseRefusalKind::Expired, "Expired"),
("refusal_busy.bin", LeaseRefusalKind::Busy, "Busy"),
(
"refusal_provider_rejected.bin",
LeaseRefusalKind::ProviderRejected,
"ProviderRejected",
),
];
for (name, kind, spelled) in kinds {
let message = LeaseMessage::ClaudeLeaseRefusal(LeaseRefused {
kind,
detail: "the reason".to_string(),
});
assert_frame(name, &message);
// The kind travels as the variant's name, which other runtimes read.
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")
};
let rmpv::Value::Map(inner) = &outer[1].1 else {
panic!("the value is a map")
};
assert_eq!(inner[0].0.as_str(), Some("kind"));
assert_eq!(inner[0].1.as_str(), Some(spelled));
}
}
#[test]
fn a_grant_with_both_usage_windows_is_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,
};
assert_frame(
"grant.bin",
&LeaseMessage::ClaudeLeaseGrant(LeaseGranted {
lease_id: "lease:0123456789abcdef0123456789abcdef".to_string(),
plan_quota: vec![
window("five_hour", 420, "allowed", 1_793_307_000_000),
window("seven_day", 999, "allowed_warning", 1_793_800_000_000),
],
}),
);
}

View file

@ -4,11 +4,16 @@ use std::time::Duration;
use thiserror::Error;
use zeroize::Zeroizing;
use rand_core::{OsRng, RngCore};
use crate::contract::{TransportInvocation, TransportKey, TransportResult};
use crate::envelope::{
MAX_IDENTITY_BYTES, ServiceError, TransportFrameError, decode_result_frame, decrypt_result,
encode_invocation_frame, encrypt_invocation, read_transport_frame, write_transport_frame,
MAX_IDENTITY_BYTES, ServiceError, TransportFrameError, decode_result_frame,
decode_usage_readout_response_frame, decrypt_result, decrypt_usage_readout,
encode_invocation_frame, encode_usage_readout_request_frame, encrypt_invocation,
encrypt_usage_readout_request, read_transport_frame, write_transport_frame,
};
use crate::readout::UsageReadout;
#[derive(Clone)]
pub struct ConnectorClient {
@ -52,6 +57,34 @@ impl ConnectorClient {
return Err(ConnectorClientError::RequestIdTooLong);
}
let request = encode_invocation_frame(&encrypt_invocation(invocation, &self.security)?)?;
let response = self.exchange(&request)?;
let envelope = decode_result_frame(&response)?;
decrypt_result(&envelope, &self.security, invocation).map_err(Into::into)
}
/// Reads the plan's usage windows as `caller_runtime_id`,
/// which must be a caller whose connection key this client holds.
pub fn usage_readout(
&self,
caller_runtime_id: &str,
) -> Result<UsageReadout, ConnectorClientError> {
let mut nonce = [0_u8; 16];
OsRng.fill_bytes(&mut nonce);
let request_id = format!(
"usage-readout-{}",
nonce.iter().map(|b| format!("{b:02x}")).collect::<String>()
);
let sealed = encrypt_usage_readout_request(caller_runtime_id, &request_id, &self.security)?;
let response = self.exchange(&encode_usage_readout_request_frame(&sealed)?)?;
let envelope = decode_usage_readout_response_frame(&response)?;
if envelope.caller_runtime_id != caller_runtime_id || envelope.request_id != request_id {
return Err(ServiceError::OuterIdentity.into());
}
decrypt_usage_readout(&envelope, &self.security).map_err(Into::into)
}
/// Sends one frame on a fresh connection and reads the one answer.
fn exchange(&self, request: &[u8]) -> Result<Vec<u8>, ConnectorClientError> {
let connect_timeout = self
.response_timeout
.unwrap_or(Duration::from_secs(10))
@ -62,12 +95,9 @@ impl ConnectorClient {
.set_read_timeout(self.response_timeout)
.and_then(|()| stream.set_write_timeout(self.response_timeout))
.map_err(ConnectorClientError::Connection)?;
write_transport_frame(&mut stream, &request, self.max_frame_bytes)
write_transport_frame(&mut stream, request, self.max_frame_bytes)
.map_err(client_frame_error)?;
let response =
read_transport_frame(&mut stream, self.max_frame_bytes).map_err(client_frame_error)?;
let envelope = decode_result_frame(&response)?;
decrypt_result(&envelope, &self.security, invocation).map_err(Into::into)
read_transport_frame(&mut stream, self.max_frame_bytes).map_err(client_frame_error)
}
}
@ -226,4 +256,56 @@ mod fixture_tests {
"the client must not even connect"
);
}
#[cfg(feature = "daemon")]
#[test]
fn a_readout_is_accepted_only_for_the_identity_that_asked() {
use crate::readout::PlanQuotaSnapshot;
// (caller the answer is sealed for, request id override, accepted)
let cases = [
("caller-1", None, true),
("caller-2", None, false),
("caller-1", Some("some-other-request"), false),
];
for (answer_caller, answer_request, accepted) in cases {
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let endpoint = listener.local_addr().unwrap();
let key = security("readout-client-key");
let server = std::thread::spawn(move || {
let (mut stream, _) = listener.accept().unwrap();
let request = read_transport_frame(&mut stream, 64 * 1024).unwrap();
let RequestFrame::UsageReadout(asked) = decode_request_frame(&request).unwrap()
else {
panic!("the client sent something other than a readout request");
};
let request_id = answer_request.unwrap_or(&asked.request_id);
let readout = UsageReadout::new(1, PlanQuotaSnapshot::default());
let sealed =
encrypt_usage_readout(answer_caller, request_id, &readout, &key).unwrap();
write_transport_frame(
&mut stream,
&encode_usage_readout_response_frame(&sealed).unwrap(),
64 * 1024,
)
.unwrap();
});
let client = ConnectorClient::new(
endpoint,
"readout-client-key",
64 * 1024,
Some(Duration::from_secs(1)),
)
.unwrap();
let answer = client.usage_readout("caller-1");
if accepted {
assert_eq!(answer.unwrap().at_unix_ms, 1);
} else {
assert!(matches!(
answer,
Err(ConnectorClientError::Transport(ServiceError::OuterIdentity))
));
}
server.join().unwrap();
}
}
}

View file

@ -9,11 +9,11 @@ use zeroize::Zeroize;
use crate::envelope::{ServiceError, binary_schema};
pub const PROVIDER_REQUEST_SCHEMA_ID: &str = "gamecult.model.provider_request.v3";
pub const INVOCATION_SCHEMA_ID: &str = "gamecult.model.transport_invocation.v3";
pub const RESULT_SCHEMA_ID: &str = "gamecult.model.transport_result.v3";
pub const RECEIPT_SCHEMA_ID: &str = "gamecult.model.transport_receipt.v3";
pub const ENVELOPE_SCHEMA_ID: &str = "gamecult.model.transport_envelope.v3";
pub const PROVIDER_REQUEST_SCHEMA_ID: &str = "gamecult.model.provider_request.v4";
pub const INVOCATION_SCHEMA_ID: &str = "gamecult.model.transport_invocation.v4";
pub const RESULT_SCHEMA_ID: &str = "gamecult.model.transport_result.v4";
pub const RECEIPT_SCHEMA_ID: &str = "gamecult.model.transport_receipt.v4";
pub const ENVELOPE_SCHEMA_ID: &str = "gamecult.model.transport_envelope.v4";
#[derive(Clone, PartialEq, Eq)]
pub struct TransportKey(pub(crate) [u8; 32]);
@ -190,10 +190,10 @@ pub enum FailureKind {
Provider,
}
/// Token accounting for one request, as the backend reports it; the service
/// prices it. Cache reads and writes are counted apart from plain input.
/// Token accounting for one request, as the backend reports it. Cache reads
/// and writes are counted apart from plain input.
/// `output_tokens` includes reasoning tokens: `reasoning_tokens` is a detail
/// (absent when the provider does not report it) and is never priced apart.
/// (absent when the provider does not report it) and is never counted apart.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct Usage {
@ -252,10 +252,6 @@ pub struct TransportInvocation {
pub expires_at_unix_ms: u64,
pub native_request_sha256: [u8; 32],
pub provider_request_sha256: [u8; 32],
/// The most this request may cost, in micro-USD: the caller's own bound
/// on it. Required nonzero. It is part of the invocation's identity, so a
/// changed budget under a recorded request id is a `ReplayConflict`.
pub max_cost_micro_usd: u64,
pub request: ProviderRequest,
}
@ -264,7 +260,6 @@ impl TransportInvocation {
caller_runtime_id: impl Into<String>,
expires_at_unix_ms: u64,
native_request_sha256: [u8; 32],
max_cost_micro_usd: u64,
request: ProviderRequest,
) -> Result<Self, ContractError> {
let provider_request_sha256 = provider_request_sha256(&request)?;
@ -274,7 +269,6 @@ impl TransportInvocation {
expires_at_unix_ms,
native_request_sha256,
provider_request_sha256,
max_cost_micro_usd,
request,
})
}
@ -287,9 +281,6 @@ impl TransportInvocation {
if self.native_request_sha256 == [0; 32] {
return Err(ContractError::Invalid("native_request_sha256"));
}
if self.max_cost_micro_usd == 0 {
return Err(ContractError::Invalid("max_cost_micro_usd"));
}
self.request.validate()?;
if self.provider_request_sha256 != provider_request_sha256(&self.request)? {
return Err(ContractError::ProviderDigest);
@ -485,15 +476,6 @@ pub struct TransportReceipt {
/// How many foreign `ReasoningState` items the service removed from the
/// request before the backend saw it.
pub foreign_reasoning_dropped: u32,
/// What the request cost, in micro-USD. The service derives it from the
/// outcome's usage and its price table; a backend never writes it.
pub cost_micro_usd: u64,
/// How far `cost_micro_usd` exceeds what the service reserved for the
/// request: zero unless the backend reported usage above the worst case
/// the request allowed, which is a breach of the backend's contract. The
/// cost is still the true one and the caller's window is charged all of
/// it. The service derives this; a backend never writes it.
pub overrun_micro_usd: u64,
pub outcome: TransportOutcome,
}
@ -558,12 +540,6 @@ pub enum Refusal {
ModelNotAllowed,
OutputBound,
PayloadBound,
// The worst case of the request exceeds the invocation's own
// `max_cost_micro_usd`. Plain comments: a doc comment would split the
// published schema into a union.
CostBound,
// The caller's spend window cannot hold the request's reservation.
SpendCeiling,
Capacity,
InFlight,
Indeterminate,
@ -573,15 +549,12 @@ pub enum Refusal {
// exchange, where dropping it would break the provider's own rules. A plain
// comment: a doc comment would split the published schema into a union.
ForeignReasoningState,
// The caller's newest admission window starts after the clock's reading,
// and the request does not fit what that window holds: the host's clock
// is behind the window and the caller stays refused until it catches up.
// A plain comment: a doc comment would split the published schema into a
// union.
ClockBehindAdmissionWindow,
// The provider itself is refusing right now: a plan window it reported
// `rejected` has not reached its own reset. Plain comments, as above.
ProviderRejected,
}
/// The one wire codec of every v3 document: named MessagePack (struct fields
/// The one wire codec of every v4 document: named MessagePack (struct fields
/// are keyed by name), with unknown fields refused by the types themselves.
pub(crate) fn wire_encode<T: Serialize + ?Sized>(document: &T) -> Result<Vec<u8>, ContractError> {
rmp_serde::to_vec_named(document).map_err(|_| ContractError::Encoding)
@ -787,28 +760,23 @@ mod tests {
}
#[test]
fn max_cost_zero_is_malformed() {
let mut invocation = invocation();
invocation.max_cost_micro_usd = 0;
assert_eq!(
invocation.validate(1_000, 2_000),
Err(ContractError::Invalid("max_cost_micro_usd"))
);
invocation.max_cost_micro_usd = 1;
assert_eq!(invocation.validate(1_000, 2_000), Ok(()));
}
#[test]
fn max_cost_is_part_of_the_invocations_identity() {
fn the_invocations_identity_covers_its_request() {
let one = invocation();
let mut other = one.clone();
other.max_cost_micro_usd += 1;
let mut request = one.request.clone();
request.max_output_tokens += 1;
let other = TransportInvocation::new(
one.caller_runtime_id.clone(),
one.expires_at_unix_ms,
one.native_request_sha256,
request,
)
.unwrap();
assert_ne!(
canonical_invocation_bytes(&one).unwrap(),
canonical_invocation_bytes(&other).unwrap()
);
let decoded: TransportInvocation = wire_decode(&wire_encode(&one).unwrap()).unwrap();
assert_eq!(decoded.max_cost_micro_usd, one.max_cost_micro_usd);
assert_eq!(decoded, one);
}
#[test]
@ -869,7 +837,7 @@ mod tests {
}
#[test]
fn v3_documents_carry_their_v3_ids() {
fn v4_documents_carry_their_v4_ids() {
assert_eq!(
[
PROVIDER_REQUEST_SCHEMA_ID,
@ -879,11 +847,11 @@ mod tests {
ENVELOPE_SCHEMA_ID
],
[
"gamecult.model.provider_request.v3",
"gamecult.model.transport_invocation.v3",
"gamecult.model.transport_result.v3",
"gamecult.model.transport_receipt.v3",
"gamecult.model.transport_envelope.v3"
"gamecult.model.provider_request.v4",
"gamecult.model.transport_invocation.v4",
"gamecult.model.transport_result.v4",
"gamecult.model.transport_receipt.v4",
"gamecult.model.transport_envelope.v4"
]
);
let invocation = invocation();
@ -922,7 +890,7 @@ mod tests {
let bytes = wire_encode(document).unwrap();
let mut value = rmpv::decode::read_value(&mut bytes.as_slice()).unwrap();
let rmpv::Value::Map(entries) = &mut value else {
panic!("a v3 document is a named map");
panic!("a v4 document is a named map");
};
entries.push((rmpv::Value::from("unexpected_field"), rmpv::Value::from(1)));
let mut out = Vec::new();

View file

@ -17,11 +17,14 @@ use crate::contract::{
ContractError, ENVELOPE_SCHEMA_ID, TransportInvocation, TransportKey, TransportResult,
canonical_invocation_bytes, require_id, wire_decode, wire_encode,
};
use crate::readout::UsageReadout;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
pub enum EnvelopeKind {
Invocation,
Result,
UsageReadoutRequest,
UsageReadoutResult,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
@ -49,8 +52,10 @@ pub(crate) fn binary_schema(_: &mut schemars::SchemaGenerator) -> schemars::Sche
/// The CultNet service every frame of this connector addresses.
pub const CONNECTOR_SERVICE_ID: &str = "gamecult.model.connector";
/// The one operation of that service: run one sealed invocation.
/// The operation that runs one sealed invocation.
pub const INVOKE_OPERATION: &str = "invoke";
/// The operation that reads the plan's usage.
pub const USAGE_READOUT_OPERATION: &str = "usage_readout";
const PAYLOAD_ENCODING: &str = "messagepack-base64";
const OK_STATUS: &str = "ok";
@ -69,16 +74,39 @@ fn seal_payload(envelope: &TransportEnvelope) -> Result<String, ServiceError> {
Ok(STANDARD.encode(wire_encode(envelope)?))
}
/// The CultNet frame an invocation travels in: a `cultnet.operation_request.v0`
/// for [`CONNECTOR_SERVICE_ID`] and [`INVOKE_OPERATION`], its `message_id` the
/// The two operations of the connector service.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum Operation {
Invoke,
UsageReadout,
}
impl Operation {
fn name(self) -> &'static str {
match self {
Operation::Invoke => INVOKE_OPERATION,
Operation::UsageReadout => USAGE_READOUT_OPERATION,
}
}
pub(crate) fn from_name(name: &str) -> Option<Self> {
[Operation::Invoke, Operation::UsageReadout]
.into_iter()
.find(|operation| operation.name() == name)
}
}
/// The CultNet frame a request travels in: a `cultnet.operation_request.v0`
/// for [`CONNECTOR_SERVICE_ID`] and the named operation, its `message_id` the
/// request id and its payload the base64 of the named sealed envelope.
pub(crate) fn encode_invocation_frame(
fn request_frame(
operation: Operation,
envelope: &TransportEnvelope,
) -> Result<Vec<u8>, ServiceError> {
carriage_bytes(&CultNetMessage::OperationRequest {
message_id: envelope.request_id.clone(),
service_id: CONNECTOR_SERVICE_ID.to_string(),
operation: INVOKE_OPERATION.to_string(),
operation: operation.name().to_string(),
payload_schema: ENVELOPE_SCHEMA_ID.to_string(),
payload_encoding: PAYLOAD_ENCODING.to_string(),
payload: seal_payload(envelope)?,
@ -87,14 +115,17 @@ pub(crate) fn encode_invocation_frame(
})
}
/// The CultNet frame a result travels in: a `cultnet.operation_response.v0`
/// with status `ok`, otherwise shaped like the invocation's.
/// The CultNet frame a response travels in: a `cultnet.operation_response.v0`
/// with status `ok`, otherwise shaped like the request's.
#[cfg(any(feature = "daemon", test))]
pub(crate) fn encode_result_frame(envelope: &TransportEnvelope) -> Result<Vec<u8>, ServiceError> {
fn response_frame(
operation: Operation,
envelope: &TransportEnvelope,
) -> Result<Vec<u8>, ServiceError> {
carriage_bytes(&CultNetMessage::OperationResponse {
message_id: envelope.request_id.clone(),
service_id: CONNECTOR_SERVICE_ID.to_string(),
operation: INVOKE_OPERATION.to_string(),
operation: operation.name().to_string(),
status: OK_STATUS.to_string(),
payload_schema: ENVELOPE_SCHEMA_ID.to_string(),
payload_encoding: PAYLOAD_ENCODING.to_string(),
@ -104,20 +135,43 @@ pub(crate) fn encode_result_frame(envelope: &TransportEnvelope) -> Result<Vec<u8
})
}
pub(crate) fn encode_invocation_frame(
envelope: &TransportEnvelope,
) -> Result<Vec<u8>, ServiceError> {
request_frame(Operation::Invoke, envelope)
}
pub(crate) fn encode_usage_readout_request_frame(
envelope: &TransportEnvelope,
) -> Result<Vec<u8>, ServiceError> {
request_frame(Operation::UsageReadout, envelope)
}
#[cfg(any(feature = "daemon", test))]
pub(crate) fn encode_result_frame(envelope: &TransportEnvelope) -> Result<Vec<u8>, ServiceError> {
response_frame(Operation::Invoke, envelope)
}
#[cfg(feature = "daemon")]
pub(crate) fn encode_usage_readout_response_frame(
envelope: &TransportEnvelope,
) -> Result<Vec<u8>, ServiceError> {
response_frame(Operation::UsageReadout, envelope)
}
/// Opens the carriage of one frame. Everything the carriage says is checked
/// here, before any decryption: the service, the operation, the payload schema,
/// and that the message id is the sealed envelope's request id.
/// here, before any decryption: the service, that the operation is one this
/// service has, the payload schema, and that the message id is the sealed
/// envelope's request id.
fn open_carriage(
message_id: &str,
service_id: &str,
operation: &str,
payload_schema: &str,
payload: &str,
) -> Result<TransportEnvelope, ServiceError> {
if service_id != CONNECTOR_SERVICE_ID
|| operation != INVOKE_OPERATION
|| payload_schema != ENVELOPE_SCHEMA_ID
{
) -> Result<(Operation, TransportEnvelope), ServiceError> {
let operation = Operation::from_name(operation).ok_or(ServiceError::Carriage)?;
if service_id != CONNECTOR_SERVICE_ID || payload_schema != ENVELOPE_SCHEMA_ID {
return Err(ServiceError::Carriage);
}
let bytes = STANDARD
@ -127,7 +181,7 @@ fn open_carriage(
if envelope.request_id != message_id {
return Err(ServiceError::OuterIdentity);
}
Ok(envelope)
Ok((operation, envelope))
}
fn read_carriage(bytes: &[u8]) -> Result<CultNetMessage, ServiceError> {
@ -135,8 +189,15 @@ fn read_carriage(bytes: &[u8]) -> Result<CultNetMessage, ServiceError> {
.map_err(|_| ServiceError::Carriage)
}
/// A request frame, by the operation it names.
#[cfg(any(feature = "daemon", test))]
pub(crate) fn decode_invocation_frame(bytes: &[u8]) -> Result<TransportEnvelope, ServiceError> {
pub(crate) enum RequestFrame {
Invoke(TransportEnvelope),
UsageReadout(TransportEnvelope),
}
#[cfg(any(feature = "daemon", test))]
pub(crate) fn decode_request_frame(bytes: &[u8]) -> Result<RequestFrame, ServiceError> {
match read_carriage(bytes)? {
CultNetMessage::OperationRequest {
message_id,
@ -145,18 +206,34 @@ pub(crate) fn decode_invocation_frame(bytes: &[u8]) -> Result<TransportEnvelope,
payload_schema,
payload,
..
} => open_carriage(
&message_id,
&service_id,
&operation,
&payload_schema,
&payload,
),
} => {
match open_carriage(
&message_id,
&service_id,
&operation,
&payload_schema,
&payload,
)? {
(Operation::Invoke, envelope) => Ok(RequestFrame::Invoke(envelope)),
(Operation::UsageReadout, envelope) => Ok(RequestFrame::UsageReadout(envelope)),
}
}
_ => Err(ServiceError::Carriage),
}
}
pub(crate) fn decode_result_frame(bytes: &[u8]) -> Result<TransportEnvelope, ServiceError> {
#[cfg(test)]
pub(crate) fn decode_invocation_frame(bytes: &[u8]) -> Result<TransportEnvelope, ServiceError> {
match decode_request_frame(bytes)? {
RequestFrame::Invoke(envelope) => Ok(envelope),
RequestFrame::UsageReadout(_) => Err(ServiceError::Carriage),
}
}
fn decode_response_frame(
bytes: &[u8],
expected: Operation,
) -> Result<TransportEnvelope, ServiceError> {
match read_carriage(bytes)? {
CultNetMessage::OperationResponse {
message_id,
@ -166,17 +243,32 @@ pub(crate) fn decode_result_frame(bytes: &[u8]) -> Result<TransportEnvelope, Ser
payload_schema,
payload,
..
} if status == OK_STATUS => open_carriage(
&message_id,
&service_id,
&operation,
&payload_schema,
&payload,
),
} if status == OK_STATUS => {
match open_carriage(
&message_id,
&service_id,
&operation,
&payload_schema,
&payload,
)? {
(found, envelope) if found == expected => Ok(envelope),
_ => Err(ServiceError::Carriage),
}
}
_ => Err(ServiceError::Carriage),
}
}
pub(crate) fn decode_result_frame(bytes: &[u8]) -> Result<TransportEnvelope, ServiceError> {
decode_response_frame(bytes, Operation::Invoke)
}
pub(crate) fn decode_usage_readout_response_frame(
bytes: &[u8],
) -> Result<TransportEnvelope, ServiceError> {
decode_response_frame(bytes, Operation::UsageReadout)
}
/// The most bytes one frame holding a sealed envelope of `ciphertext_len`
/// ciphertext bytes can occupy, with ids at their longest: the frame a caller
/// payload bound (`CallerAdmission`'s `max_payload_bytes`, a bound on
@ -324,6 +416,59 @@ pub fn decrypt_result(
Ok(result)
}
/// A usage readout request carries no content: the sealed envelope's
/// authenticated identity and kind are the whole request.
pub(crate) fn encrypt_usage_readout_request(
caller_runtime_id: &str,
request_id: &str,
security: &TransportKey,
) -> Result<TransportEnvelope, ServiceError> {
encrypt_message(
caller_runtime_id,
request_id,
EnvelopeKind::UsageReadoutRequest,
&[],
security,
)
}
#[cfg(any(feature = "daemon", test))]
pub(crate) fn decrypt_usage_readout_request(
envelope: &TransportEnvelope,
security: &TransportKey,
) -> Result<(), ServiceError> {
validate_envelope(envelope, EnvelopeKind::UsageReadoutRequest)?;
decrypt_message(envelope, security).map(drop)
}
#[cfg(feature = "daemon")]
pub(crate) fn encrypt_usage_readout(
caller_runtime_id: &str,
request_id: &str,
readout: &UsageReadout,
security: &TransportKey,
) -> Result<TransportEnvelope, ServiceError> {
let plaintext = wire_encode(readout).map_err(|_| ServiceError::Encoding)?;
encrypt_message(
caller_runtime_id,
request_id,
EnvelopeKind::UsageReadoutResult,
&plaintext,
security,
)
}
pub(crate) fn decrypt_usage_readout(
envelope: &TransportEnvelope,
security: &TransportKey,
) -> Result<UsageReadout, ServiceError> {
validate_envelope(envelope, EnvelopeKind::UsageReadoutResult)?;
let plaintext = decrypt_message(envelope, security)?;
let readout: UsageReadout = wire_decode(&plaintext).map_err(|_| ServiceError::Encoding)?;
readout.validate()?;
Ok(readout)
}
#[derive(Debug)]
pub(crate) enum TransportFrameError {
Connection(std::io::Error),
@ -462,8 +607,6 @@ pub enum ServiceError {
SharedCallerKey,
#[error("connector replay store failed: {0}")]
ReplayStore(String),
#[error("an allowed model has no price: {0}")]
UnpricedModel(String),
#[error("connector replay record is malformed")]
InvalidReplayRecord,
#[error("caller key changed while durable replay records still exist")]
@ -927,7 +1070,7 @@ mod tests {
fn golden_frame() -> Vec<u8> {
std::fs::read(
std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
.join("tests/fixtures/wire_v3/carriage_invocation.frame"),
.join("tests/fixtures/wire_v4/carriage_invocation.frame"),
)
.expect("carriage_invocation.frame is committed")
}
@ -964,7 +1107,7 @@ mod tests {
assert_eq!(text(&request, "operation"), "invoke");
assert_eq!(
text(&request, "payloadSchema"),
"gamecult.model.transport_envelope.v3"
"gamecult.model.transport_envelope.v4"
);
assert_eq!(text(&request, "payloadEncoding"), "messagepack-base64");
assert_eq!(text(&request, "messageId"), envelope.request_id);
@ -978,7 +1121,7 @@ mod tests {
assert_eq!(text(&response, "status"), "ok");
assert_eq!(
text(&response, "payloadSchema"),
"gamecult.model.transport_envelope.v3"
"gamecult.model.transport_envelope.v4"
);
assert_eq!(text(&response, "payloadEncoding"), "messagepack-base64");
assert_eq!(text(&response, "messageId"), envelope.request_id);
@ -1138,4 +1281,53 @@ mod tests {
);
}
}
#[test]
fn a_readout_request_is_not_an_invocation_and_names_its_own_operation() {
let key = TransportKey::from_connection_secret("readout-test-key").unwrap();
let sealed = encrypt_usage_readout_request("caller-1", "readout-1", &key).unwrap();
let frame = encode_usage_readout_request_frame(&sealed).unwrap();
assert!(matches!(
decode_request_frame(&frame),
Ok(RequestFrame::UsageReadout(found)) if found == sealed
));
assert_eq!(decode_invocation_frame(&frame), Err(ServiceError::Carriage));
// Sealed as one kind, the authenticated data refuses it as another.
assert!(decrypt_usage_readout_request(&sealed, &key).is_ok());
assert_eq!(
decrypt_usage_readout(&sealed, &key).unwrap_err(),
ServiceError::Envelope
);
let mut relabelled = sealed.clone();
relabelled.message_kind = EnvelopeKind::UsageReadoutResult;
assert_eq!(
decrypt_usage_readout(&relabelled, &key).unwrap_err(),
ServiceError::Encryption
);
}
#[cfg(feature = "daemon")]
#[test]
fn a_response_is_read_as_its_own_operation_only() {
use crate::readout::PlanQuotaSnapshot;
let key = TransportKey::from_connection_secret("readout-test-key").unwrap();
let readout = UsageReadout::new(1, PlanQuotaSnapshot::default());
let sealed = encrypt_usage_readout("caller-1", "readout-1", &readout, &key).unwrap();
let readout_frame = encode_usage_readout_response_frame(&sealed).unwrap();
let invoke_frame = encode_result_frame(&sealed).unwrap();
assert_eq!(
decode_usage_readout_response_frame(&readout_frame).unwrap(),
sealed
);
assert_eq!(decode_result_frame(&invoke_frame).unwrap(), sealed);
assert_eq!(
decode_result_frame(&readout_frame),
Err(ServiceError::Carriage)
);
assert_eq!(
decode_usage_readout_response_frame(&invoke_frame),
Err(ServiceError::Carriage)
);
assert_eq!(decrypt_usage_readout(&sealed, &key).unwrap(), readout);
}
}

View file

@ -67,14 +67,6 @@ impl FlakyStore {
pub(crate) fn allowance(&self) -> Arc<AtomicUsize> {
self.pushes_allowed.clone()
}
/// A store over the same records that accepts `pushes_allowed` more writes.
pub(crate) fn reopened(&self, pushes_allowed: usize) -> Self {
Self {
held: self.held.clone(),
pushes_allowed: Arc::new(AtomicUsize::new(pushes_allowed)),
}
}
}
#[cfg(feature = "daemon")]
@ -133,16 +125,8 @@ pub(crate) fn request() -> ProviderRequest {
request
}
/// Generous next to the fixture request's worst case (about 64,000), so it
/// binds only in the tests that are about spend.
pub(crate) const MAX_COST: u64 = 10_000_000;
/// What the fixture usage (10 input and 3 output tokens) costs under
/// [`fixture_price`]: 10 x 3,000,000 + 3 x 15,000,000 = 75,000,000 / 1e6.
pub(crate) const FIXTURE_USAGE_COST: u64 = 75;
pub(crate) fn invocation() -> TransportInvocation {
TransportInvocation::new("epiphany-yggdrasil", 2_000, [7; 32], MAX_COST, request()).unwrap()
TransportInvocation::new("epiphany-yggdrasil", 2_000, [7; 32], request()).unwrap()
}
pub(crate) fn receipt(invocation: &TransportInvocation) -> TransportReceipt {
@ -155,8 +139,6 @@ pub(crate) fn receipt(invocation: &TransportInvocation) -> TransportReceipt {
model: invocation.request.model.clone(),
provider: "fixture-sse".to_string(),
foreign_reasoning_dropped: 0,
cost_micro_usd: FIXTURE_USAGE_COST,
overrun_micro_usd: 0,
outcome: TransportOutcome::Completed {
provider_response_id: Some("response-1".to_string()),
stop_reason: StopReason::EndTurn,
@ -170,33 +152,6 @@ pub(crate) fn receipt(invocation: &TransportInvocation) -> TransportReceipt {
},
}
}
/// A price with a distinct nonzero rate per class, so a swapped rate changes a cost.
#[cfg(feature = "daemon")]
pub(crate) fn fixture_price() -> ModelPrice {
ModelPrice {
rates: Rates {
input: 3_000_000,
cache_write_short: 3_750_000,
cache_write_long: 6_000_000,
cache_read: 300_000,
output: 15_000_000,
},
long_context: None,
overhead_tokens: 100,
}
}
#[cfg(feature = "daemon")]
pub(crate) fn fixture_prices() -> PriceTable {
PriceTable::from([("gpt-5.4".to_string(), fixture_price())])
}
/// Far above anything a fixture request reserves: the ceiling binds only in
/// the tests that are about it.
#[cfg(feature = "daemon")]
pub(crate) const CEILING: u64 = 1 << 40;
#[cfg(feature = "daemon")]
pub(crate) const WINDOW_MS: u64 = 3_600_000;
#[cfg(feature = "daemon")]
pub(crate) fn caller(
@ -212,8 +167,6 @@ pub(crate) fn caller(
max_concurrent_requests,
64 * 1024,
32_768,
SpendCeiling(CEILING),
SpendWindowMs(WINDOW_MS),
)
.unwrap()
}
@ -223,7 +176,6 @@ pub(crate) fn open_service(store: &TestStore, max_concurrent_requests: usize) ->
TransportService::open(
&store.path,
"fixture-sse",
fixture_prices(),
[
caller(
"epiphany-yggdrasil",
@ -249,5 +201,5 @@ pub(crate) fn security(key: &str) -> TransportKey {
pub(crate) fn invocation_for(caller_runtime_id: &str, request_id: &str) -> TransportInvocation {
let mut request = request();
request.request_id = request_id.to_string();
TransportInvocation::new(caller_runtime_id, 2_000, [7; 32], MAX_COST, request).unwrap()
TransportInvocation::new(caller_runtime_id, 2_000, [7; 32], request).unwrap()
}

View file

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

View file

@ -0,0 +1,106 @@
//! The usage readout: the subscription plan's own usage windows as the backend
//! last saw them. The connector meters nothing in dollars; it publishes what
//! the provider reports, and the minds read it and decide their own burn.
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use crate::contract::ContractError;
pub const USAGE_READOUT_SCHEMA_ID: &str = "gamecult.model.usage_readout.v1";
/// One window of the subscription plan, as the provider reported it.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct PlanQuota {
pub window: String,
/// How much of the window was used when it was observed, in thousandths;
/// `None` once the window has reset (`status` is then `expired`), because
/// the old reading says nothing about the new window.
pub utilization_permille: Option<u32>,
pub resets_at_unix_ms: u64,
pub status: String,
/// When this window was last reported. Its age in a readout is the
/// readout's `at_unix_ms` minus this.
pub observed_at_unix_ms: u64,
}
/// What a backend last saw of the plan's windows, each stamped with its own
/// observation time. A backend with no plan to report returns the default.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct PlanQuotaSnapshot {
pub windows: Vec<PlanQuota>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct UsageReadout {
pub schema_id: String,
pub at_unix_ms: u64,
pub plan_quota: Vec<PlanQuota>,
}
impl UsageReadout {
pub fn new(at_unix_ms: u64, plan: PlanQuotaSnapshot) -> Self {
Self {
schema_id: USAGE_READOUT_SCHEMA_ID.to_string(),
at_unix_ms,
plan_quota: plan.windows,
}
}
pub(crate) fn validate(&self) -> Result<(), ContractError> {
if self.schema_id != USAGE_READOUT_SCHEMA_ID {
return Err(ContractError::Schema);
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::contract::{wire_decode, wire_encode};
fn readout() -> UsageReadout {
UsageReadout::new(
1_000,
PlanQuotaSnapshot {
windows: vec![PlanQuota {
window: "five_hour".into(),
utilization_permille: Some(80),
resets_at_unix_ms: 1_793_307_000_000,
status: "allowed".into(),
observed_at_unix_ms: 999,
}],
},
)
}
#[test]
fn a_readout_is_a_named_map_that_decodes_to_itself() {
let bytes = wire_encode(&readout()).unwrap();
let value: rmpv::Value = rmp_serde::from_slice(&bytes).unwrap();
let rmpv::Value::Map(fields) = value else {
panic!("a readout is a map, not a positional array");
};
let names: Vec<&str> = fields.iter().filter_map(|(k, _)| k.as_str()).collect();
assert_eq!(names, ["schema_id", "at_unix_ms", "plan_quota"]);
assert_eq!(wire_decode::<UsageReadout>(&bytes).unwrap(), readout());
}
#[test]
fn an_unobserved_plan_is_none_not_zero() {
let bare = UsageReadout::new(1, PlanQuotaSnapshot::default());
let decoded: UsageReadout = wire_decode(&wire_encode(&bare).unwrap()).unwrap();
assert!(decoded.plan_quota.is_empty());
}
#[test]
fn a_readout_of_another_schema_is_refused() {
let mut other = readout();
other.schema_id = "gamecult.model.usage_readout.v2".into();
assert!(other.validate().is_err());
assert!(readout().validate().is_ok());
}
}

View file

@ -12,25 +12,33 @@ use zeroize::Zeroizing;
use crate::contract::TransportResult;
use crate::envelope::{
ServiceError, TransportFrameError, decode_invocation_frame, encode_result_frame,
frame_bytes_for_payload, read_transport_frame, write_transport_frame,
RequestFrame, ServiceError, TransportFrameError, decode_request_frame, encode_result_frame,
encode_usage_readout_response_frame, frame_bytes_for_payload, read_transport_frame,
write_transport_frame,
};
use crate::readout::PlanQuotaSnapshot;
use crate::service::{
CallerAdmission, PriceTable, ProviderCall, SpendCeiling, SpendWindowMs, TransportAdmission,
TransportService,
CallerAdmission, ProviderCall, TransportAdmission, TransportService, distinct_admissions,
};
/// The only provider seam of the connector kit. A connector binary implements
/// it over one model provider; the kit names no provider.
pub trait ProviderBackend: Send + Sync {
/// Whether the provider can take work. `serve` consults it once, before it
/// opens the replay store or accepts a connection; the request path does
/// not consult it.
fn ready(&self) -> bool;
/// Whether the provider can take work; `Err` names the backend field that
/// is not ready, never a value. `serve` consults it once, before it opens
/// the replay store or accepts a connection; the request path does not
/// consult it.
fn ready(&self) -> Result<(), &'static str>;
/// The tag of the provider this backend fronts. The service owns it from
/// `serve` on: it names the provider on every receipt and decides which
/// reasoning state is foreign.
fn provider(&self) -> &str;
/// What the provider's own plan allows and has used, for the usage
/// readout and the provider's own refusals. A provider with no plan to report
/// keeps the default.
fn plan_quota(&self) -> PlanQuotaSnapshot {
PlanQuotaSnapshot::default()
}
/// Executes one admitted call. The service has already bound it to its
/// caller and recorded its replay claim. The call cannot show a foreign
/// reasoning item and builds the result's receipt itself.
@ -59,7 +67,7 @@ impl ServeLimits {
/// largest whose sealed frame, at its worst-case encoding, fits the frame.
/// An admitted payload therefore always fits, and `complete` is given the
/// same frame bound to fit every result in.
fn validate(&self, admissions: &[CallerAdmission]) -> Result<(), ServeError> {
pub fn validate(&self, admissions: &[CallerAdmission]) -> Result<(), ServeError> {
if self.max_frame_bytes == 0
|| self.max_frame_bytes > u32::MAX as usize
|| self.max_connections == 0
@ -79,6 +87,7 @@ impl ServeLimits {
if admissions.is_empty() {
return Err(ServeError::InvalidConfig("callers"));
}
distinct_admissions(admissions)?;
let slots = admissions
.iter()
.try_fold(0_usize, |total, caller| {
@ -114,12 +123,6 @@ pub struct CallerConfig {
pub max_concurrent_requests: usize,
pub max_payload_bytes: usize,
pub max_output_tokens: u32,
/// The most the caller may spend in one window, in micro-USD. Required
/// nonzero. A fuse set once, generously above what the dial allows.
pub spend_ceiling_micro_usd: u64,
/// The length of the caller's spend window, in milliseconds. Windows are
/// fixed and aligned to the unix epoch. Required nonzero.
pub spend_window_ms: u64,
}
pub fn load_admissions(callers: &[CallerConfig]) -> Result<Vec<CallerAdmission>, ServeError> {
@ -140,8 +143,6 @@ pub fn load_admissions(callers: &[CallerConfig]) -> Result<Vec<CallerAdmission>,
caller.max_concurrent_requests,
caller.max_payload_bytes,
caller.max_output_tokens,
SpendCeiling(caller.spend_ceiling_micro_usd),
SpendWindowMs(caller.spend_window_ms),
)
.map_err(ServeError::Service)?,
);
@ -150,19 +151,43 @@ pub fn load_admissions(callers: &[CallerConfig]) -> Result<Vec<CallerAdmission>,
}
/// Serves admitted callers on a loopback listener until the listener fails.
/// `prices` is the connector binary's config document's table of what every
/// allowed model costs (B4 owns that document).
/// The caller binds; a listener that is not bound to loopback is refused, and
/// peers that are not loopback are dropped.
pub fn serve(
listener: TcpListener,
limits: ServeLimits,
admissions: Vec<CallerAdmission>,
prices: PriceTable,
replay_store: &Path,
backend: Arc<dyn ProviderBackend>,
) -> Result<(), ServeError> {
limits.validate(&admissions)?;
admit_listener(&listener, backend.as_ref())?;
let service = Arc::new(Mutex::new(TransportService::open(
replay_store,
backend.provider(),
admissions,
limits.max_expiry_skew_ms,
)?));
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 service (the Self-run credential lease) shares its
/// replay store and plan. 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 store: 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)?
@ -171,17 +196,7 @@ pub fn serve(
{
return Err(ServeError::InvalidConfig("listener address"));
}
if !backend.ready() {
return Err(ServeError::BackendUnready);
}
let service = Arc::new(Mutex::new(TransportService::open(
replay_store,
backend.provider(),
prices,
admissions,
limits.max_expiry_skew_ms,
)?));
accept_loop(|| listener.accept(), limits, service, backend)
backend.ready().map_err(ServeError::BackendUnready)
}
/// The accept loop over any source of connections, so a test can present a
@ -233,21 +248,37 @@ fn serve_connection(
.and_then(|()| stream.set_write_timeout(Some(timeout)))
.map_err(ServeError::Connection)?;
let request = read_transport_frame(&mut stream, limits.max_frame_bytes).map_err(frame_error)?;
let response = execute_inference_frame(&request, service, backend, limits)?;
let response = answer_frame(&request, service, backend, limits)?;
write_transport_frame(&mut stream, &response, limits.max_frame_bytes).map_err(frame_error)
}
fn execute_inference_frame(
fn answer_frame(
request: &[u8],
service: &Mutex<TransportService>,
backend: &dyn ProviderBackend,
limits: &ServeLimits,
) -> Result<Vec<u8>, ServeError> {
let envelope = decode_invocation_frame(request).map_err(|_| ServeError::FrameEncoding)?;
let envelope = match decode_request_frame(request).map_err(|_| ServeError::FrameEncoding)? {
RequestFrame::Invoke(envelope) => envelope,
RequestFrame::UsageReadout(envelope) => {
let plan = backend.plan_quota();
let response = service
.lock()
.map_err(|_| ServeError::ServicePoisoned)?
.usage_readout(&envelope, unix_ms()?, plan, limits.max_frame_bytes)?;
return encode_usage_readout_response_frame(&response)
.map_err(|_| ServeError::FrameEncoding);
}
};
let admission = service
.lock()
.map_err(|_| ServeError::ServicePoisoned)?
.begin(&envelope, unix_ms()?, limits.max_frame_bytes)?;
.begin(
&envelope,
unix_ms()?,
&backend.plan_quota(),
limits.max_frame_bytes,
)?;
let response = match admission {
TransportAdmission::Reply(response) => response,
TransportAdmission::Execute(claim) => {
@ -322,8 +353,8 @@ pub enum ServeError {
ConnectionKey(#[source] std::io::Error),
#[error(transparent)]
Service(#[from] ServiceError),
#[error("provider backend is not ready")]
BackendUnready,
#[error("provider backend is not ready: {0}")]
BackendUnready(&'static str),
#[error("connector service lock was poisoned")]
ServicePoisoned,
#[error("connector listener failed")]
@ -351,10 +382,7 @@ mod tests {
FailureKind, Refusal, TransportDisposition, TransportEvent, TransportEventPayload,
TransportInvocation, TransportOutcome, TransportResult,
};
use crate::fixtures::{
CEILING, MAX_COST, TestStore, WINDOW_MS, caller, fixture_prices, invocation_for, receipt,
request,
};
use crate::fixtures::{TestStore, caller, invocation_for, receipt, request};
const CALLER: &str = "epiphany-yggdrasil";
const KEY: &str = "epiphany-distinct-test-key";
@ -379,9 +407,9 @@ mod tests {
}
impl ProviderBackend for CountingBackend {
fn ready(&self) -> bool {
fn ready(&self) -> Result<(), &'static str> {
self.ready_calls.fetch_add(1, Ordering::SeqCst);
self.ready
self.ready.then_some(()).ok_or("fixture")
}
fn provider(&self) -> &str {
@ -430,14 +458,7 @@ mod tests {
let (sender, receiver) = std::sync::mpsc::channel();
let path = path.to_path_buf();
thread::spawn(move || {
let _ = sender.send(serve(
listener,
limits,
admissions,
fixture_prices(),
&path,
backend,
));
let _ = sender.send(serve(listener, limits, admissions, &path, backend));
});
receiver.recv_timeout(Duration::from_secs(2)).ok()
}
@ -463,16 +484,7 @@ mod tests {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let endpoint = listener.local_addr().unwrap();
let path = store.path.clone();
thread::spawn(move || {
serve(
listener,
limits(),
admissions,
fixture_prices(),
&path,
backend,
)
});
thread::spawn(move || serve(listener, limits(), admissions, &path, backend));
let client =
ConnectorClient::new(endpoint, KEY, FRAME, Some(Duration::from_secs(5))).unwrap();
(client, endpoint)
@ -490,35 +502,128 @@ mod tests {
max_concurrent_requests: *capacity,
max_payload_bytes: 1_048_576,
max_output_tokens: 16_384,
spend_ceiling_micro_usd: CEILING,
spend_window_ms: WINDOW_MS,
})
.collect()
}
/// A counting backend that reports a plan, as a provider with one would.
struct PlanBackend {
inner: Arc<CountingBackend>,
plan: PlanQuotaSnapshot,
}
impl ProviderBackend for PlanBackend {
fn ready(&self) -> Result<(), &'static str> {
self.inner.ready()
}
fn provider(&self) -> &str {
self.inner.provider()
}
fn plan_quota(&self) -> PlanQuotaSnapshot {
self.plan.clone()
}
fn execute(&self, call: &ProviderCall) -> TransportResult {
self.inner.execute(call)
}
}
#[test]
fn serve_refuses_a_request_over_its_max_cost_without_executing_it() {
fn serve_reads_out_plan_usage_without_executing_anything() {
let store = TestStore::new();
let backend = CountingBackend::new(true);
let (client, _) = start(backend.clone(), &store);
let mut invocation = fresh("request-costly");
invocation.max_cost_micro_usd = 1;
let refused = client.execute(&invocation).unwrap();
let inner = CountingBackend::new(true);
let plan = PlanQuotaSnapshot {
windows: vec![crate::readout::PlanQuota {
window: "five_hour".into(),
utilization_permille: Some(80),
resets_at_unix_ms: 1_793_307_000_000,
status: "allowed".into(),
observed_at_unix_ms: 1_793_300_000_000,
}],
};
let backend = Arc::new(PlanBackend {
inner: inner.clone(),
plan: plan.clone(),
});
let (client, endpoint) = start_with(backend, &store, vec![caller(CALLER, KEY, 4)]);
client.execute(&fresh("request-1")).unwrap();
assert_eq!(inner.executions.load(Ordering::SeqCst), 1);
let before = now_ms();
let readout = client.usage_readout(CALLER).unwrap();
assert!(readout.at_unix_ms >= before && readout.at_unix_ms <= now_ms());
assert_eq!(readout.plan_quota, plan.windows);
assert_eq!(
inner.executions.load(Ordering::SeqCst),
1,
"a readout runs nothing"
);
// A client without the caller's key is dropped, not answered.
let stranger = ConnectorClient::new(
endpoint,
"some-other-key",
FRAME,
Some(Duration::from_secs(5)),
)
.unwrap();
assert!(stranger.usage_readout(CALLER).is_err());
}
#[test]
fn a_backend_with_no_plan_reads_out_none() {
let store = TestStore::new();
let (client, _) = start(CountingBackend::new(true), &store);
let readout = client.usage_readout(CALLER).unwrap();
assert!(readout.plan_quota.is_empty());
}
fn window_plan(permille: u32, status: &str, resets_at_unix_ms: u64) -> PlanQuotaSnapshot {
PlanQuotaSnapshot {
windows: vec![crate::readout::PlanQuota {
window: "five_hour".into(),
utilization_permille: Some(permille),
resets_at_unix_ms,
status: status.into(),
observed_at_unix_ms: 1,
}],
}
}
#[test]
fn serve_repeats_the_providers_rejection_and_executes_nothing_while_it_stands() {
let store = TestStore::new();
let inner = CountingBackend::new(true);
let backend = Arc::new(PlanBackend {
inner: inner.clone(),
plan: window_plan(1_000, "rejected", now_ms() + 3_600_000),
});
let (client, _) = start_with(backend, &store, vec![caller(CALLER, KEY, 4)]);
let refused = client.execute(&fresh("request-rejected")).unwrap();
assert_eq!(
refused.disposition,
TransportDisposition::Refused(Refusal::CostBound)
TransportDisposition::Refused(Refusal::ProviderRejected)
);
assert_eq!(backend.executions.load(Ordering::SeqCst), 0);
assert_eq!(inner.executions.load(Ordering::SeqCst), 0);
}
// The refusal was not recorded: the same id with a budget that holds
// the worst case is admitted and executed.
invocation.max_cost_micro_usd = MAX_COST;
let admitted = client.execute(&invocation).unwrap();
#[test]
fn serve_admits_a_full_window_the_provider_still_allows() {
let store = TestStore::new();
let inner = CountingBackend::new(true);
let backend = Arc::new(PlanBackend {
inner: inner.clone(),
plan: window_plan(999, "allowed", now_ms() + 3_600_000),
});
let (client, _) = start_with(backend, &store, vec![caller(CALLER, KEY, 4)]);
let admitted = client.execute(&fresh("request-full")).unwrap();
assert!(matches!(
admitted.disposition,
TransportDisposition::Transported { .. }
));
assert_eq!(backend.executions.load(Ordering::SeqCst), 1);
assert_eq!(inner.executions.load(Ordering::SeqCst), 1);
}
#[test]
@ -557,7 +662,6 @@ mod tests {
CALLER,
invocation.expires_at_unix_ms,
[7; 32],
MAX_COST,
conflicting_request,
)
.unwrap();
@ -642,7 +746,6 @@ mod tests {
let mut service = TransportService::open(
&store.path,
"fixture-sse",
fixture_prices(),
vec![caller(CALLER, KEY, 4)],
300_000,
)
@ -650,8 +753,9 @@ mod tests {
let envelope =
crate::envelope::encrypt_invocation(&invocation, &crate::fixtures::security(KEY))
.unwrap();
let TransportAdmission::Execute(claim) =
service.begin(&envelope, now_ms(), 1 << 20).unwrap()
let TransportAdmission::Execute(claim) = service
.begin(&envelope, now_ms(), &PlanQuotaSnapshot::default(), 1 << 20)
.unwrap()
else {
panic!("a fresh request must execute");
};
@ -674,7 +778,6 @@ mod tests {
listener,
narrowed,
vec![caller(CALLER, KEY, 4)],
fixture_prices(),
&path,
served,
)
@ -710,7 +813,10 @@ mod tests {
&store.path,
backend.clone(),
);
assert!(matches!(result, Some(Err(ServeError::BackendUnready))));
assert!(matches!(
result,
Some(Err(ServeError::BackendUnready("fixture")))
));
assert_eq!(backend.executions.load(Ordering::SeqCst), 0);
assert!(
!store.path.exists(),
@ -757,13 +863,7 @@ mod tests {
fn the_smallest_frame_also_carries_the_failure_that_stands_in_for_a_result() {
let long = "x".repeat(crate::envelope::MAX_IDENTITY_BYTES);
let plaintext = crate::contract::wire_encode(&crate::service::result_size_failure(
&long,
&long,
&long,
[0xff; 32],
[0xff; 32],
0,
u64::MAX,
&long, &long, &long, [0xff; 32], [0xff; 32], 0,
))
.unwrap()
.len();
@ -829,6 +929,26 @@ mod tests {
));
}
#[test]
fn the_limits_refuse_callers_that_share_a_key_or_an_id() {
let limits = limits();
let shared_key = [
caller("one", "shared-key", 1),
caller("two", "shared-key", 1),
];
assert!(matches!(
limits.validate(&shared_key),
Err(ServeError::Service(ServiceError::SharedCallerKey))
));
let shared_id = [caller("one", "key-a", 1), caller("one", "key-b", 1)];
assert!(matches!(
limits.validate(&shared_id),
Err(ServeError::Service(ServiceError::DuplicateCaller))
));
let distinct = [caller("one", "key-a", 1), caller("two", "key-b", 1)];
assert!(limits.validate(&distinct).is_ok());
}
#[test]
fn only_loopback_peers_are_admitted() {
assert!(peer_admitted(IpAddr::from([127, 0, 0, 1])));
@ -851,18 +971,7 @@ mod tests {
}
fn admission(id: &str, key: &str, slots: usize, payload: usize) -> CallerAdmission {
CallerAdmission::new(
id,
key,
1,
["gpt-5.4".to_string()],
slots,
payload,
32_768,
SpendCeiling(CEILING),
SpendWindowMs(WINDOW_MS),
)
.unwrap()
CallerAdmission::new(id, key, 1, ["gpt-5.4".to_string()], slots, payload, 32_768).unwrap()
}
/// `Some(field)` when serve refused the configuration, `None` while it serves.
@ -968,7 +1077,7 @@ mod tests {
..limits()
};
let admissions = vec![admission(CALLER, KEY, 4, bound)];
thread::spawn(move || serve(listener, tight, admissions, fixture_prices(), &path, served));
thread::spawn(move || serve(listener, tight, admissions, &path, served));
let client =
ConnectorClient::new(endpoint, KEY, frame, Some(Duration::from_secs(5))).unwrap();
let reply = client.execute(&invocation).unwrap();
@ -1016,8 +1125,8 @@ mod tests {
}
impl ProviderBackend for SizedBackend {
fn ready(&self) -> bool {
true
fn ready(&self) -> Result<(), &'static str> {
Ok(())
}
fn provider(&self) -> &str {
@ -1083,7 +1192,6 @@ mod tests {
let mut service = TransportService::open(
&store.path,
"fixture-sse",
fixture_prices(),
vec![caller(CALLER, KEY, 4)],
300_000,
)
@ -1093,8 +1201,9 @@ mod tests {
&crate::fixtures::security(KEY),
)
.unwrap();
let TransportAdmission::Execute(claim) =
service.begin(&envelope, now_ms(), 1 << 20).unwrap()
let TransportAdmission::Execute(claim) = service
.begin(&envelope, now_ms(), &PlanQuotaSnapshot::default(), 1 << 20)
.unwrap()
else {
panic!("a fresh request must execute");
};
@ -1186,7 +1295,6 @@ mod tests {
TransportService::open(
&store.path,
"fixture-sse",
fixture_prices(),
vec![caller(CALLER, KEY, 4)],
300_000,
)
@ -1244,16 +1352,7 @@ mod tests {
let backend = CountingBackend::new(true);
let path = store.path.clone();
let served = backend.clone();
thread::spawn(move || {
serve(
listener,
limits(),
admissions,
fixture_prices(),
&path,
served,
)
});
thread::spawn(move || serve(listener, limits(), admissions, &path, served));
let client =
ConnectorClient::new(endpoint, KEY, FRAME, Some(Duration::from_secs(5))).unwrap();
client.execute(&fresh("request-key")).unwrap();

File diff suppressed because it is too large Load diff

View file

@ -0,0 +1,429 @@
//! A credential lease: a Self run's or a Persona turn's hold on the token the
//! connector lends. The lease is memory bound to its holder's open connection,
//! not a record and not a charge: it decides `Busy` while it is held and
//! nothing else. It ends when its holder lets go (`end_lease`), when the unit
//! invocation it served ends, or when its expiry passes. Nothing is billed and
//! nothing is reserved. The only usage-derived refusal is the provider's own
//! (`ProviderRejected`), repeated with the provider's own expiry.
use super::*;
/// The request-id namespace of leases.
const LEASE_PREFIX: &str = "lease:";
/// What a Persona turn's lease id adds to its invocation id, so a Persona
/// lease and a Self lease of one caller and one invocation never share an id.
const PERSONA_SUFFIX: &str = ":persona";
/// Which kind of turn a lease serves. It is a property of the lease caller's
/// configured row, never of a request. A caller holds one Self lease at a
/// time; a Persona turn holds one lease per live invocation, one at a time.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum LeaseTurn {
SelfRun,
PersonaTurn,
}
/// What an admitted lease holds.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LeaseGrant {
pub lease_id: String,
}
/// Why no lease was granted.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LeaseRefusal {
UnknownCaller,
/// A Self run: the caller holds a Self lease now. A Persona turn: this
/// invocation holds a lease now.
Busy,
/// The provider is refusing right now and has said until when.
ProviderRejected,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum LeaseAdmission {
Granted(LeaseGrant),
Refused(LeaseRefusal),
}
impl TransportService {
/// Grants `caller_runtime_id` a lease for the unit invocation
/// `invocation_id` until `expires_at_unix_ms`. A Self run's lease is
/// `lease:<invocation>` and one is open per caller; a Persona turn's is
/// `lease:<invocation>:persona`, one open per invocation; neither blocks
/// the other. A lease takes no request slot. `plan` is the backend's plan
/// quota as it stands: a provider that is refusing right now refuses the
/// lease. A lease past its expiry no longer holds anything.
pub fn begin_lease(
&mut self,
caller_runtime_id: &str,
turn: LeaseTurn,
invocation_id: &str,
expires_at_unix_ms: u64,
now_unix_ms: u64,
plan: &PlanQuotaSnapshot,
) -> Result<LeaseAdmission, ServiceError> {
if !self.callers.contains_key(caller_runtime_id) {
return Ok(LeaseAdmission::Refused(LeaseRefusal::UnknownCaller));
}
require_id(invocation_id, "invocation_id")?;
let suffix = match turn {
LeaseTurn::SelfRun => "",
LeaseTurn::PersonaTurn => PERSONA_SUFFIX,
};
let lease_id = format!("{LEASE_PREFIX}{invocation_id}{suffix}");
if lease_id.len() > MAX_IDENTITY_BYTES {
return Err(ContractError::Invalid("invocation_id").into());
}
self.leases
.retain(|_, lease| lease.expires_at_unix_ms > now_unix_ms);
let identity = RequestIdentity {
caller_runtime_id: caller_runtime_id.to_string(),
request_id: lease_id.clone(),
};
let busy = match turn {
LeaseTurn::SelfRun => self.leases.iter().any(|(held, lease)| {
lease.turn == LeaseTurn::SelfRun && held.caller_runtime_id == caller_runtime_id
}),
LeaseTurn::PersonaTurn => self.leases.contains_key(&identity),
};
if busy {
return Ok(LeaseAdmission::Refused(LeaseRefusal::Busy));
}
if provider_rejects(plan, now_unix_ms) {
return Ok(LeaseAdmission::Refused(LeaseRefusal::ProviderRejected));
}
self.leases.insert(
identity,
LiveLease {
turn,
expires_at_unix_ms,
},
);
Ok(LeaseAdmission::Granted(LeaseGrant { lease_id }))
}
/// Ends a lease its holder let go of. Ending one that is not held is a
/// no-op: a holder may close after its expiry already freed the lease.
pub fn end_lease(&mut self, caller_runtime_id: &str, lease_id: &str) {
self.leases.remove(&RequestIdentity {
caller_runtime_id: caller_runtime_id.to_string(),
request_id: lease_id.to_string(),
});
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::fixtures::*;
use crate::readout::PlanQuota;
const CALLER: &str = "eureka";
const KEY: &str = "lease-distinct-test-key";
const NOW: u64 = 7_200_500;
fn open(store: &TestStore) -> TransportService {
TransportService::open(
&store.path,
"fixture-sse",
[
CallerAdmission::new(CALLER, KEY, 1, ["gpt-5.4".to_string()], 1, 65_536, 32_768)
.unwrap(),
],
5_000,
)
.unwrap()
}
fn none() -> PlanQuotaSnapshot {
PlanQuotaSnapshot::default()
}
fn window(permille: u32, status: &str, resets_at: u64) -> PlanQuotaSnapshot {
PlanQuotaSnapshot {
windows: vec![PlanQuota {
window: "five_hour".to_string(),
utilization_permille: Some(permille),
resets_at_unix_ms: resets_at,
status: status.to_string(),
observed_at_unix_ms: NOW,
}],
}
}
fn lease_in(
service: &mut TransportService,
turn: LeaseTurn,
run: &str,
plan: &PlanQuotaSnapshot,
) -> LeaseAdmission {
service
.begin_lease(CALLER, turn, run, NOW + 60_000, NOW, plan)
.unwrap()
}
fn lease(service: &mut TransportService, run: &str) -> LeaseAdmission {
lease_in(service, LeaseTurn::SelfRun, run, &none())
}
fn granted(admission: LeaseAdmission) -> LeaseGrant {
match admission {
LeaseAdmission::Granted(grant) => grant,
LeaseAdmission::Refused(why) => panic!("refused {why:?}"),
}
}
#[test]
fn a_lease_is_named_by_its_invocation() {
let store = TestStore::new();
let mut service = open(&store);
assert_eq!(
granted(lease(&mut service, "run-1")).lease_id,
"lease:run-1"
);
}
#[test]
fn a_caller_holds_one_self_lease_at_a_time() {
let store = TestStore::new();
let mut service = open(&store);
granted(lease(&mut service, "run-1"));
assert_eq!(
lease(&mut service, "run-2"),
LeaseAdmission::Refused(LeaseRefusal::Busy)
);
assert_eq!(
lease(&mut service, "run-1"),
LeaseAdmission::Refused(LeaseRefusal::Busy)
);
}
#[test]
fn a_lease_ends_when_its_holder_closes() {
let store = TestStore::new();
let mut service = open(&store);
let grant = granted(lease(&mut service, "run-1"));
service.end_lease(CALLER, &grant.lease_id);
// The same invocation may lease again: nothing remembers a closed
// lease, and nothing was billed for it.
granted(lease(&mut service, "run-1"));
service.end_lease(CALLER, "lease:not-held");
assert_eq!(
lease(&mut service, "run-2"),
LeaseAdmission::Refused(LeaseRefusal::Busy),
"ending a lease that is not held frees nothing else"
);
}
#[test]
fn a_lease_ends_at_its_expiry_without_its_holder() {
let store = TestStore::new();
let mut service = open(&store);
granted(
service
.begin_lease(CALLER, LeaseTurn::SelfRun, "run-1", NOW + 10, NOW, &none())
.unwrap(),
);
assert_eq!(
service
.begin_lease(
CALLER,
LeaseTurn::SelfRun,
"run-2",
NOW + 99,
NOW + 9,
&none()
)
.unwrap(),
LeaseAdmission::Refused(LeaseRefusal::Busy),
"one millisecond before its expiry the lease still holds"
);
granted(
service
.begin_lease(
CALLER,
LeaseTurn::SelfRun,
"run-2",
NOW + 99,
NOW + 10,
&none(),
)
.unwrap(),
);
}
#[test]
fn a_stranger_is_refused() {
let store = TestStore::new();
let mut service = open(&store);
assert_eq!(
service
.begin_lease(
"stranger",
LeaseTurn::SelfRun,
"run-1",
NOW + 1,
NOW,
&none()
)
.unwrap(),
LeaseAdmission::Refused(LeaseRefusal::UnknownCaller)
);
}
#[test]
fn a_persona_holds_one_lease_per_invocation_and_never_blocks_the_self_lease() {
let store = TestStore::new();
let mut service = open(&store);
let persona = |service: &mut TransportService, run: &str| {
lease_in(service, LeaseTurn::PersonaTurn, run, &none())
};
let first = granted(persona(&mut service, "turn-1"));
assert_eq!(first.lease_id, "lease:turn-1:persona");
assert_eq!(
persona(&mut service, "turn-1"),
LeaseAdmission::Refused(LeaseRefusal::Busy)
);
granted(persona(&mut service, "turn-2"));
granted(lease(&mut service, "turn-1"));
service.end_lease(CALLER, &first.lease_id);
granted(persona(&mut service, "turn-1"));
}
#[test]
fn a_rejected_window_refuses_invoke_and_lease_until_its_reset() {
let store = TestStore::new();
let mut service = open(&store);
let rejected = window(1_000, "rejected", NOW + 5_000);
assert_eq!(
lease_in(&mut service, LeaseTurn::SelfRun, "run-1", &rejected),
LeaseAdmission::Refused(LeaseRefusal::ProviderRejected)
);
let call = |service: &mut TransportService, now: u64, plan: &PlanQuotaSnapshot| {
let mut asked = request();
asked.request_id = format!("call-{now}");
let invocation = TransportInvocation::new(CALLER, now + 1_000, [7; 32], asked).unwrap();
let sealed = crate::envelope::encrypt_message(
CALLER,
invocation.request_id(),
crate::envelope::EnvelopeKind::Invocation,
&crate::contract::wire_encode(&invocation).unwrap(),
&security(KEY),
)
.unwrap();
let TransportAdmission::Reply(reply) =
service.begin(&sealed, now, plan, 1 << 20).unwrap()
else {
return None;
};
let result =
crate::envelope::decrypt_result(&reply, &security(KEY), &invocation).unwrap();
match result.disposition {
TransportDisposition::Refused(refusal) => Some(refusal),
other => panic!("unexpected {other:?}"),
}
};
assert_eq!(
call(&mut service, NOW, &rejected),
Some(Refusal::ProviderRejected)
);
assert_eq!(
call(&mut service, NOW + 4_999, &rejected),
Some(Refusal::ProviderRejected)
);
assert_eq!(call(&mut service, NOW + 5_000, &rejected), None);
assert!(matches!(
service
.begin_lease(
CALLER,
LeaseTurn::SelfRun,
"run-1",
NOW + 9_000,
NOW + 5_000,
&rejected
)
.unwrap(),
LeaseAdmission::Granted(_)
));
}
#[test]
fn high_utilization_refuses_nothing() {
let store = TestStore::new();
let mut service = open(&store);
let full = window(999, "allowed", NOW + 5_000);
granted(lease_in(&mut service, LeaseTurn::SelfRun, "run-1", &full));
let warned = window(1_000, "allowed_warning", NOW + 5_000);
granted(lease_in(
&mut service,
LeaseTurn::PersonaTurn,
"turn-1",
&warned,
));
}
#[test]
fn an_invocation_id_is_bounded_so_its_lease_id_fits_an_identity() {
let store = TestStore::new();
let mut service = open(&store);
let longest = "r".repeat(MAX_IDENTITY_BYTES - LEASE_PREFIX.len());
let grant = granted(lease(&mut service, &longest));
assert_eq!(grant.lease_id.len(), MAX_IDENTITY_BYTES);
assert!(lease_try(&mut service, LeaseTurn::SelfRun, &format!("{longest}r")).is_err());
}
#[test]
fn a_persona_invocation_id_leaves_room_for_its_suffix() {
let store = TestStore::new();
let mut service = open(&store);
let longest = "r".repeat(MAX_IDENTITY_BYTES - LEASE_PREFIX.len() - PERSONA_SUFFIX.len());
let grant = granted(lease_in(
&mut service,
LeaseTurn::PersonaTurn,
&longest,
&none(),
));
assert_eq!(grant.lease_id.len(), MAX_IDENTITY_BYTES);
assert!(lease_try(&mut service, LeaseTurn::PersonaTurn, &format!("{longest}r")).is_err());
}
fn lease_try(
service: &mut TransportService,
turn: LeaseTurn,
run: &str,
) -> Result<LeaseAdmission, ServiceError> {
service.begin_lease(CALLER, turn, run, NOW + 60_000, NOW, &none())
}
#[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);
granted(lease(&mut service, "run-1"));
let mut call = request();
call.request_id = "call-1".to_string();
let invocation = TransportInvocation::new(CALLER, NOW + 1_000, [7; 32], call).unwrap();
let sealed = crate::envelope::encrypt_message(
CALLER,
"call-1",
crate::envelope::EnvelopeKind::Invocation,
&crate::contract::wire_encode(&invocation).unwrap(),
&security(KEY),
)
.unwrap();
assert!(matches!(
service.begin(&sealed, NOW, &none(), 1 << 20).unwrap(),
TransportAdmission::Execute(_)
));
// And a call in flight does not make the caller busy for a lease.
granted(lease_in(
&mut service,
LeaseTurn::PersonaTurn,
"turn-1",
&none(),
));
}
}

View file

@ -1 +0,0 @@
<EFBFBD>messageIdョgolden-requestゥoperationヲinvokeァpayloadレDhqlzY2hlbWFfaWTZJGdhbWVjdWx0Lm1vZGVsLnRyYW5zcG9ydF9lbnZlbG9wZS52M7FjYWxsZXJfcnVudGltZV9pZK1nb2xkZW4tY2FsbGVyqnJlcXVlc3RfaWSuZ29sZGVuLXJlcXVlc3SsbWVzc2FnZV9raW5kqkludm9jYXRpb26lbm9uY2WczO7MxRrMxw7MhTh7SMzwzNB2qmNpcGhlcnRleHTFA0d65EskLG6kIwJ/G+XDK03gZOGtLFi6vZHJbI+/KpRLwnLT0TPF7RQa0KD98CrW+VwkKY3sRa4+Z9+/5OdPq1YmDFwbxvMVjZsFjXXw84tzvuLKHsb/3kQSEvAlFE4K6FXHvznLcMo4CnGUzoXwNGALf+xDTqFG0aFbdG4xGynV/kttsL6f34F4+l6xiXMUy87Q8/1UyOK/YaVxMMmlT3tdTllhq8YU8Dvfq0KuMaj/f9Anya8A5bJwlxafZPtc53RoH1/oNjLdaDqvwrTEnTB9TvKlnEOxKyyKhmmwGaC03qqPJXRaj8cDQO6EHDbML4J0zMN8TvssQISd5UCeVMkvDGjmvEOmRUoPk1/SnSrIHdha4+6rU4lUSumloLmlx5vv++Mq4AqzOdd/vPB6Ov/AoF32X27o4RI6H3xCKKHjKC4m3+cVjXCrOcGojuRf9yeDFkgbiPaC31hmSc8rAm6HO/NcaIHvx+G+6AaurAEhDLdF6DdrXb4xwjuv6yWQKihNvLXpnwmhRFkSeD5lS6FHW+CoOIdq/RBnFFyDrboTPWFoxQYIrBuqFoHDOajIz7dXYif0sf9sB9JBrDhlcbTpiGcyURfzDhW55NOPqKFGi6ropaJ1/SyLP3f7E57kA22aQ2Op7o36Tu06Pnzo2wI7UsL/Qz27F+AGAukQCvTJyjxARiGr3TRpChjVXx6GXhLrG1T9P+Mrhyrt9vwwhboN8qz1uvDZwKhGWBKXPXkxluq/NY1DwB3JTCCZfEwukWjMZbqiINLEKvW9BStcBR1kfXcK6sQvn/F0g9r81e9brbe+w3U24GERnHwgruOqytTeax/xPbQ7uISV48g5L8DAYf788KhUXwI43QDHNBdBUyo0rCsgoWJKeO05fcAgKgYYGsyLtMW10Vstw3q4POyn3o7rrBz4S2cnl2iZgr0ixg/KgF0V+loEVuj/6rYrcW0SLCHwV1+w1ClFMn0ZPsrz28EoZMA617G0Rrch66cS5uI9/Zf/wLX3z6n7xoPLNXk51ZsLSZ2bqEQkcBJbW/dT5MdH6anWVogRDvIW1coUbkIq9l+pkQSuelWsfakX4UQwYuUD0UwNbqxHiGwlEFj0f1bdhEj26Q==ッpayloadEncodingイmessagepack-base64ュpayloadSchemaル$gamecult.model.transport_envelope.v3ュschemaVersionシcultnet.operation_request.v0ゥserviceIdクgamecult.model.connector

View file

@ -1,6 +1,6 @@
# wire_v3 fixtures
# wire_v4 fixtures
The named-MessagePack bytes of the v3 documents and one sealed carriage frame,
The named-MessagePack bytes of the v4 documents and one sealed carriage frame,
produced by the kit at the revision that added them. They pin drift; the types in
`src/contract.rs` own the shape, and `schemas/` publishes it.
@ -10,7 +10,7 @@ produced by the kit at the revision that added them. They pin drift; the types i
the plaintext documents.
- `carriage_invocation.frame`: a `cultnet.operation_request.v0` for service
`gamecult.model.connector`, operation `invoke`, whose payload is the base64 of an envelope
sealed under the fixed test key in `tests/wire_golden_v3.rs`.
sealed under the fixed test key in `tests/wire_golden_v4.rs`.
Generator: `tests/wire_golden_v3.rs` itself. Run it with `WIRE_V3_REGENERATE=1` to rewrite
Generator: `tests/wire_golden_v4.rs` itself. Run it with `WIRE_V4_REGENERATE=1` to rewrite
the files from the kit, then review the diff: a changed byte is a wire change.

View file

@ -0,0 +1 @@
<EFBFBD>messageIdョgolden-requestゥoperationヲinvokeァpayloadレ(hqlzY2hlbWFfaWTZJGdhbWVjdWx0Lm1vZGVsLnRyYW5zcG9ydF9lbnZlbG9wZS52NLFjYWxsZXJfcnVudGltZV9pZK1nb2xkZW4tY2FsbGVyqnJlcXVlc3RfaWSuZ29sZGVuLXJlcXVlc3SsbWVzc2FnZV9raW5kqkludm9jYXRpb26lbm9uY2WcXSbMzU3MhsyCzLAtzKLMhcy8zMSqY2lwaGVydGV4dMUDMY816tn0FGQ9jqxE+KyDVcxmhufaQI7vzEXO8NXTQPOYpRuMcRDM/j/QzNF2Lb3Fak0G+hHRKrXDxJ3MeKpoELNYOEyxbKfxn88mjAyyc7BVPM7+0Bk9ye3m5oBF6hoyCgOHSTwESIfq0bim+9oQh1IDp+HdFjztK/XjM731GRzKJmGgxsKqQsz42zq8y9mjAXl3V1RR9ssdKvk2k5U56/3liOeHDU96kOa7IlVCBm6qnn06hg2DTtnBF4xGj6PihHj4KufVkAcuQhRCZWOKsploToXo7wy+jcQDrnhHIUCMN6sk3P9YJh1ejuleDM9rp3RYiXcnzRwMNXK1AGpNraTgWn91d6L/1Ff+gaQBgPxSjW6FATshJf9Iqk9pEhT8xeCQMmNj40CdlVFYNz/TsRJYW45BXm8g83Jnh3Wj0Z90WJ80bwwgZbjPWx6wCOCID6Jo8wiubF6411ygi+GUHTR9rZFhrsGo6TSmsbvIAuabUJcskqKHW7tqVkfpz5JdjyeSSDvKcZ3PQyVne7Lyo8nzNh0g06a7CEUvq9kvr4zfTCfwk/1YVSrZ+7U8ePGj27V3D8+BH/pafiLR5FxxBKWBTYe7n5irVP5gLw4WlzOrSriXpolmaPWujCm3apml45mhlYARdfJJbpQxnQGb6cI3D/vAxxI50BPLB82JqCjNFex9xwfP5jSRDYah9wLm4yVymBqZ8Wvn0maLz4OvF2EonJBwEtN73UBULkdeX6g9GTiu2MKtCjSMm33SkwiHfvDdwMDS11sINK+YQLF/3HobbnAk2n963JahlddNRg79T96GksjN0br8u1PNMC+/uBUHEQvUBfRnR4tKp+SeLF6zMn228orgHgWr7sxilMwTOP1DgXpbEE19M0P8XC0CA/pTeHr4eFHzakPLotxF8avQfPKAXbXWgNGK6uWj1NGPshD6EN3pLvePOc4uObqHMKz/dg52/7m1VTHhYMivSRenG6H/BrWfUXQvcgCPalg7C4KCPYINb3nnRxP0NIOBP+KJ4S2haU94ZNHVr2HB038xNzUm2HrONVBsxuuTckYDmHfCa6FmvVJv4she5wyuPHk=ッpayloadEncodingイmessagepack-base64ュpayloadSchemaル$gamecult.model.transport_envelope.v4ュschemaVersionシcultnet.operation_request.v0ゥserviceIdクgamecult.model.connector

View file

@ -1 +1 @@
„©schema_idÙ"gamecult.model.transport_result.v3ªrequest_id®golden-request±caller_runtime_id­golden-caller«disposition<6F>§Refused«OutputBound
„©schema_idÙ"gamecult.model.transport_result.v4ªrequest_id®golden-request±caller_runtime_id­golden-caller«disposition<6F>§Refused«OutputBound

View file

@ -1,4 +1,4 @@
//! `schemas/` at the repository root publishes one JSON Schema per v3 document,
//! `schemas/` at the repository root publishes one JSON Schema per v4 document,
//! generated from the types. This test regenerates them and fails on any
//! difference; run it with `MODEL_CONNECTOR_REGENERATE_SCHEMAS=1` to rewrite the
//! files from the types.
@ -28,6 +28,7 @@ fn published_schemas_are_the_types() {
("transport_result", render::<TransportResult>()),
("transport_receipt", render::<TransportReceipt>()),
("transport_envelope", render::<TransportEnvelope>()),
("usage_readout", render::<UsageReadout>()),
];
let dir = schemas_dir();
if std::env::var_os(REGENERATE).is_some() {
@ -53,7 +54,7 @@ fn published_schemas_are_the_types() {
expected.sort();
assert_eq!(
on_disk, expected,
"schemas/ holds exactly the five documents"
"schemas/ holds exactly the six documents"
);
}
@ -79,7 +80,7 @@ fn the_published_ciphertext_is_what_the_wire_sends() {
assert_eq!(&published["properties"]["ciphertext"], ciphertext);
let envelope = TransportEnvelope {
schema_id: "gamecult.model.transport_envelope.v3".to_string(),
schema_id: "gamecult.model.transport_envelope.v4".to_string(),
caller_runtime_id: "caller".to_string(),
request_id: "request".to_string(),
message_kind: EnvelopeKind::Invocation,
@ -110,9 +111,7 @@ fn the_refusal_schema_is_one_flat_string_enum() {
.iter()
.map(|value| value.as_str().unwrap())
.collect();
assert_eq!(variants.len(), 15);
assert!(variants.contains(&"ClockBehindAdmissionWindow"));
assert_eq!(variants.len(), 13);
assert!(variants.contains(&"ForeignReasoningState"));
assert!(variants.contains(&"CostBound"));
assert!(variants.contains(&"SpendCeiling"));
assert!(variants.contains(&"ProviderRejected"));
}

View file

@ -38,8 +38,8 @@ impl RecordingBackend {
}
impl ProviderBackend for RecordingBackend {
fn ready(&self) -> bool {
true
fn ready(&self) -> Result<(), &'static str> {
Ok(())
}
fn provider(&self) -> &str {
@ -106,8 +106,6 @@ fn start_with(backend: Arc<dyn ProviderBackend>, max_concurrent: usize) -> Conne
max_concurrent,
64 * 1024,
32_768,
SpendCeiling(1 << 40),
SpendWindowMs(3_600_000),
)
.unwrap();
let limits = ServeLimits {
@ -116,21 +114,7 @@ fn start_with(backend: Arc<dyn ProviderBackend>, max_concurrent: usize) -> Conne
socket_timeout_ms: 5_000,
max_expiry_skew_ms: 300_000,
};
let prices = PriceTable::from([(
"claude-opus-4-5".to_string(),
ModelPrice {
rates: Rates {
input: 5_000_000,
cache_write_short: 6_250_000,
cache_write_long: 10_000_000,
cache_read: 500_000,
output: 25_000_000,
},
long_context: None,
overhead_tokens: 100,
},
)]);
thread::spawn(move || serve(listener, limits, vec![admission], prices, &store, backend));
thread::spawn(move || serve(listener, limits, vec![admission], &store, backend));
ConnectorClient::new(endpoint, KEY, FRAME, Some(Duration::from_secs(5))).unwrap()
}
@ -143,7 +127,7 @@ fn invocation(request_id: &str, input: Vec<InputItem>) -> TransportInvocation {
4_096,
);
request.input = input;
TransportInvocation::new(CALLER, now_ms() + 60_000, [7; 32], 10_000_000, request).unwrap()
TransportInvocation::new(CALLER, now_ms() + 60_000, [7; 32], request).unwrap()
}
fn user(text: &str) -> InputItem {
@ -291,8 +275,8 @@ struct TamperingBackend {
}
impl ProviderBackend for TamperingBackend {
fn ready(&self) -> bool {
true
fn ready(&self) -> Result<(), &'static str> {
Ok(())
}
fn provider(&self) -> &str {
@ -332,7 +316,7 @@ fn a_refused_completion_reaches_the_caller_as_a_typed_failure_and_frees_the_slot
let edits: [fn(&mut TransportReceipt); 3] = [
|receipt| receipt.provider = "another-provider".to_string(),
|receipt| receipt.foreign_reasoning_dropped += 1,
|receipt| receipt.cost_micro_usd += 1,
|receipt| receipt.request_id = "another-request".to_string(),
];
for edit in edits {
// One concurrent request: a wedged slot would refuse the next id.

View file

@ -1,5 +1,5 @@
//! Pins the bytes of the v3 wire so a drift is a failing test. The fixtures in
//! `fixtures/wire_v3` are produced by this file (see its README); they pin the
//! Pins the bytes of the v4 wire so a drift is a failing test. The fixtures in
//! `fixtures/wire_v4` are produced by this file (see its README); they pin the
//! kit against itself and are not a second contract: the types in
//! `model_connector::contract` own the shape.
@ -17,10 +17,10 @@ use serde::de::DeserializeOwned;
const KEY: &str = "golden-distinct-test-key";
const CALLER: &str = "golden-caller";
const REGENERATE: &str = "WIRE_V3_REGENERATE";
const REGENERATE: &str = "WIRE_V4_REGENERATE";
fn fixtures() -> PathBuf {
Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/wire_v3")
Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/wire_v4")
}
fn named<T: Serialize>(value: &T) -> Vec<u8> {
@ -28,7 +28,7 @@ fn named<T: Serialize>(value: &T) -> Vec<u8> {
}
/// The value serializes to the committed bytes and the committed bytes decode
/// back to it. With `WIRE_V3_REGENERATE` set, the committed bytes are rewritten
/// back to it. With `WIRE_V4_REGENERATE` set, the committed bytes are rewritten
/// from the value first.
fn assert_wire<T: Serialize + DeserializeOwned + PartialEq + std::fmt::Debug>(
name: &str,
@ -81,7 +81,7 @@ fn request() -> ProviderRequest {
}
fn invocation() -> TransportInvocation {
TransportInvocation::new(CALLER, 1_800_000_000_000, [7; 32], 5_000_000, request()).unwrap()
TransportInvocation::new(CALLER, 1_800_000_000_000, [7; 32], request()).unwrap()
}
fn receipt(outcome: TransportOutcome) -> TransportReceipt {
@ -95,8 +95,6 @@ fn receipt(outcome: TransportOutcome) -> TransportReceipt {
model: invocation.request.model.clone(),
provider: "anthropic-messages".to_string(),
foreign_reasoning_dropped: 2,
cost_micro_usd: 1_234_567,
overrun_micro_usd: 89,
outcome,
}
}
@ -116,7 +114,7 @@ fn completed() -> TransportOutcome {
}
#[test]
fn wire_golden_v3() {
fn wire_golden_v4() {
let invocation = invocation();
assert_wire("provider_request.mp", &invocation.request);
assert_wire("transport_invocation.mp", &invocation);
@ -233,7 +231,7 @@ fn the_golden_carriage_frame_decrypts_and_validates_under_the_fixed_key() {
};
assert_eq!(service_id, "gamecult.model.connector");
assert_eq!(operation, "invoke");
assert_eq!(payload_schema, "gamecult.model.transport_envelope.v3");
assert_eq!(payload_schema, "gamecult.model.transport_envelope.v4");
assert_eq!(message_id, invocation.request_id());
let envelope: TransportEnvelope =
rmp_serde::from_slice(&STANDARD.decode(payload).unwrap()).unwrap();

View file

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

View file

@ -11,12 +11,6 @@
"format": "uint64",
"minimum": 0
},
"max_cost_micro_usd": {
"description": "The most this request may cost, in micro-USD: the caller's own bound\non it. Required nonzero. It is part of the invocation's identity, so a\nchanged budget under a recorded request id is a `ReplayConflict`.",
"type": "integer",
"format": "uint64",
"minimum": 0
},
"native_request_sha256": {
"type": "array",
"items": {
@ -53,7 +47,6 @@
"expires_at_unix_ms",
"native_request_sha256",
"provider_request_sha256",
"max_cost_micro_usd",
"request"
],
"$defs": {

View file

@ -6,12 +6,6 @@
"caller_runtime_id": {
"type": "string"
},
"cost_micro_usd": {
"description": "What the request cost, in micro-USD. The service derives it from the\noutcome's usage and its price table; a backend never writes it.",
"type": "integer",
"format": "uint64",
"minimum": 0
},
"foreign_reasoning_dropped": {
"description": "How many foreign `ReasoningState` items the service removed from the\nrequest before the backend saw it.",
"type": "integer",
@ -35,12 +29,6 @@
"outcome": {
"$ref": "#/$defs/TransportOutcome"
},
"overrun_micro_usd": {
"description": "How far `cost_micro_usd` exceeds what the service reserved for the\nrequest: zero unless the backend reported usage above the worst case\nthe request allowed, which is a breach of the backend's contract. The\ncost is still the true one and the caller's window is charged all of\nit. The service derives this; a backend never writes it.",
"type": "integer",
"format": "uint64",
"minimum": 0
},
"provider": {
"type": "string"
},
@ -72,8 +60,6 @@
"model",
"provider",
"foreign_reasoning_dropped",
"cost_micro_usd",
"overrun_micro_usd",
"outcome"
],
"$defs": {
@ -174,7 +160,7 @@
]
},
"Usage": {
"description": "Token accounting for one request, as the backend reports it; the service\nprices it. Cache reads and writes are counted apart from plain input.\n`output_tokens` includes reasoning tokens: `reasoning_tokens` is a detail\n(absent when the provider does not report it) and is never priced apart.",
"description": "Token accounting for one request, as the backend reports it. Cache reads\nand writes are counted apart from plain input.\n`output_tokens` includes reasoning tokens: `reasoning_tokens` is a detail\n(absent when the provider does not report it) and is never counted apart.",
"type": "object",
"properties": {
"cache_read_tokens": {

View file

@ -46,15 +46,13 @@
"ModelNotAllowed",
"OutputBound",
"PayloadBound",
"CostBound",
"SpendCeiling",
"Capacity",
"InFlight",
"Indeterminate",
"ReplayConflict",
"Malformed",
"ForeignReasoningState",
"ClockBehindAdmissionWindow"
"ProviderRejected"
]
},
"StopReason": {
@ -286,12 +284,6 @@
"caller_runtime_id": {
"type": "string"
},
"cost_micro_usd": {
"description": "What the request cost, in micro-USD. The service derives it from the\noutcome's usage and its price table; a backend never writes it.",
"type": "integer",
"format": "uint64",
"minimum": 0
},
"foreign_reasoning_dropped": {
"description": "How many foreign `ReasoningState` items the service removed from the\nrequest before the backend saw it.",
"type": "integer",
@ -315,12 +307,6 @@
"outcome": {
"$ref": "#/$defs/TransportOutcome"
},
"overrun_micro_usd": {
"description": "How far `cost_micro_usd` exceeds what the service reserved for the\nrequest: zero unless the backend reported usage above the worst case\nthe request allowed, which is a breach of the backend's contract. The\ncost is still the true one and the caller's window is charged all of\nit. The service derives this; a backend never writes it.",
"type": "integer",
"format": "uint64",
"minimum": 0
},
"provider": {
"type": "string"
},
@ -352,13 +338,11 @@
"model",
"provider",
"foreign_reasoning_dropped",
"cost_micro_usd",
"overrun_micro_usd",
"outcome"
]
},
"Usage": {
"description": "Token accounting for one request, as the backend reports it; the service\nprices it. Cache reads and writes are counted apart from plain input.\n`output_tokens` includes reasoning tokens: `reasoning_tokens` is a detail\n(absent when the provider does not report it) and is never priced apart.",
"description": "Token accounting for one request, as the backend reports it. Cache reads\nand writes are counted apart from plain input.\n`output_tokens` includes reasoning tokens: `reasoning_tokens` is a detail\n(absent when the provider does not report it) and is never counted apart.",
"type": "object",
"properties": {
"cache_read_tokens": {

View file

@ -0,0 +1,68 @@
{
"$schema": "https://json-schema.org/draft/2020-12/schema",
"title": "UsageReadout",
"type": "object",
"properties": {
"at_unix_ms": {
"type": "integer",
"format": "uint64",
"minimum": 0
},
"plan_quota": {
"type": "array",
"items": {
"$ref": "#/$defs/PlanQuota"
}
},
"schema_id": {
"type": "string"
}
},
"additionalProperties": false,
"required": [
"schema_id",
"at_unix_ms",
"plan_quota"
],
"$defs": {
"PlanQuota": {
"description": "One window of the subscription plan, as the provider reported it.",
"type": "object",
"properties": {
"observed_at_unix_ms": {
"description": "When this window was last reported. Its age in a readout is the\nreadout's `at_unix_ms` minus this.",
"type": "integer",
"format": "uint64",
"minimum": 0
},
"resets_at_unix_ms": {
"type": "integer",
"format": "uint64",
"minimum": 0
},
"status": {
"type": "string"
},
"utilization_permille": {
"description": "How much of the window was used when it was observed, in thousandths;\n`None` once the window has reset (`status` is then `expired`), because\nthe old reading says nothing about the new window.",
"type": [
"integer",
"null"
],
"format": "uint32",
"minimum": 0
},
"window": {
"type": "string"
}
},
"additionalProperties": false,
"required": [
"window",
"resets_at_unix_ms",
"status",
"observed_at_unix_ms"
]
}
}
}