Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
67c7849cb6 | ||
|
|
85ed36d292 | ||
|
|
2f2f7a9db5 | ||
|
|
83c93b2bb8 | ||
|
|
b44c5a3b75 | ||
|
|
9a062f0846 |
59 changed files with 3328 additions and 6539 deletions
2
.gitattributes
vendored
2
.gitattributes
vendored
|
|
@ -1,2 +1,2 @@
|
|||
schemas/*.json text eol=lf
|
||||
crates/model-connector/tests/fixtures/wire_v4/* -text
|
||||
crates/model-connector/tests/fixtures/wire_v3/* -text
|
||||
|
|
|
|||
3
Cargo.lock
generated
3
Cargo.lock
generated
|
|
@ -139,13 +139,10 @@ name = "claude-connector"
|
|||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"chrono",
|
||||
"cultcache-rs",
|
||||
"hmac",
|
||||
"libc",
|
||||
"model-connector",
|
||||
"rmp-serde",
|
||||
"rmpv",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"sha2",
|
||||
|
|
|
|||
37
README.md
37
README.md
|
|
@ -4,11 +4,6 @@ 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).
|
||||
|
||||
|
|
@ -16,7 +11,7 @@ connector binary exists yet (the Claude connector binary is first, then an OpenA
|
|||
|
||||
Crate `crates/model-connector`:
|
||||
|
||||
- `src/contract.rs` - the v4 wire contract (`gamecult.model.*.v4`): `ProviderRequest`, `TransportInvocation`,
|
||||
- `src/contract.rs` - the v3 wire contract (`gamecult.model.*.v3`): `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
|
||||
|
|
@ -26,12 +21,9 @@ 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` (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.
|
||||
`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.
|
||||
- `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),
|
||||
|
|
@ -45,16 +37,13 @@ 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 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.
|
||||
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.
|
||||
- `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 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.
|
||||
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.
|
||||
|
||||
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
|
||||
|
|
@ -63,7 +52,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 v4
|
||||
The default feature set compiles contract, envelope and client only. The kit speaks v3
|
||||
only: a v2 frame gets no answer.
|
||||
|
||||
## Verify
|
||||
|
|
@ -73,9 +62,9 @@ cargo test --locked --workspace
|
|||
cargo test --locked --workspace --features model-connector/daemon
|
||||
```
|
||||
|
||||
`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/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/published_schemas.rs` fails when `schemas/` differs from the types
|
||||
(`MODEL_CONNECTOR_REGENERATE_SCHEMAS=1` rewrites it).
|
||||
|
||||
|
|
|
|||
|
|
@ -13,8 +13,6 @@ 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"
|
||||
|
|
@ -26,6 +24,3 @@ zeroize = "1"
|
|||
|
||||
[target.'cfg(unix)'.dependencies]
|
||||
libc = "0.2"
|
||||
|
||||
[dev-dependencies]
|
||||
rmpv = "1"
|
||||
|
|
|
|||
|
|
@ -1,41 +1,27 @@
|
|||
use std::fs::File;
|
||||
use std::io;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
|
||||
use model_connector::{
|
||||
FailureKind, InputItem, PlanQuotaSnapshot, ProviderBackend, ProviderCall, ProviderRequest,
|
||||
TransportEvent, TransportEventPayload, TransportOutcome, TransportResult, Usage,
|
||||
FailureKind, InputItem, ProviderBackend, ProviderCall, ProviderRequest, TransportEvent,
|
||||
TransportEventPayload, TransportOutcome, TransportResult, Usage,
|
||||
};
|
||||
use sha2::{Digest, Sha256};
|
||||
use zeroize::Zeroizing;
|
||||
|
||||
use crate::config::{ClaudeCodeConfig, PROVIDER};
|
||||
use crate::handle::{
|
||||
Binding, ClaudeCodeHandle, MIN_SECRET_BYTES, SessionStore, plan, prefix_digest, stdin_bytes,
|
||||
};
|
||||
use crate::handle::{ClaudeCodeHandle, 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};
|
||||
|
||||
/// 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
|
||||
|
|
@ -44,8 +30,6 @@ pub trait TokenSource: Send + Sync {
|
|||
pub struct CredentialTokenSource {
|
||||
directory: Option<PathBuf>,
|
||||
name: String,
|
||||
/// The configuration field `missing` names when the credential is absent.
|
||||
field: &'static str,
|
||||
}
|
||||
|
||||
impl CredentialTokenSource {
|
||||
|
|
@ -56,21 +40,13 @@ impl CredentialTokenSource {
|
|||
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",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -83,16 +59,6 @@ impl TokenSource for CredentialTokenSource {
|
|||
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
|
||||
|
|
@ -109,15 +75,6 @@ pub struct ClaudeCodeBackend {
|
|||
sessions: Box<dyn SessionStore>,
|
||||
sink: Box<dyn StderrSink>,
|
||||
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 {
|
||||
|
|
@ -131,11 +88,7 @@ impl ClaudeCodeBackend {
|
|||
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)
|
||||
Self::with_ports(config, port, sessions, sink, tokens)
|
||||
}
|
||||
|
||||
pub fn with_ports(
|
||||
|
|
@ -144,7 +97,6 @@ impl ClaudeCodeBackend {
|
|||
sessions: Box<dyn SessionStore>,
|
||||
sink: Box<dyn StderrSink>,
|
||||
tokens: Box<dyn TokenSource>,
|
||||
handle_secret: Box<dyn TokenSource>,
|
||||
) -> Self {
|
||||
Self {
|
||||
config,
|
||||
|
|
@ -152,34 +104,6 @@ impl ClaudeCodeBackend {
|
|||
sessions,
|
||||
sink,
|
||||
tokens,
|
||||
handle_secret,
|
||||
quota: Arc::new(QuotaCell::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:?}")),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -189,28 +113,8 @@ impl ClaudeCodeBackend {
|
|||
binary_matches_pin(&self.config)
|
||||
}
|
||||
|
||||
/// `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(),
|
||||
) {
|
||||
fn run(&self, request: &ProviderRequest) -> (Vec<TransportEvent>, TransportOutcome) {
|
||||
let plan = match plan(request, &self.config, self.sessions.as_ref()) {
|
||||
Ok(plan) => plan,
|
||||
Err(error) => {
|
||||
return failed(
|
||||
|
|
@ -239,8 +143,8 @@ impl ClaudeCodeBackend {
|
|||
"claude-code: the system prompt could not be staged",
|
||||
));
|
||||
};
|
||||
let argv = argv(&self.config, request, &plan, prompt.path(), ephemeral);
|
||||
let env = child_env(&self.config, request);
|
||||
let argv = argv(&self.config, request, &plan, prompt.path());
|
||||
let env = child_env(&self.config, request, token);
|
||||
let deadline = Instant::now() + self.config.limits.deadline;
|
||||
|
||||
let spawned = self.port.spawn(
|
||||
|
|
@ -248,7 +152,6 @@ impl ClaudeCodeBackend {
|
|||
&env,
|
||||
&self.config.work_dir,
|
||||
stdin_bytes(plan.blocks()),
|
||||
&token,
|
||||
);
|
||||
let Ok(mut child) = spawned else {
|
||||
return failure(Failure::never_reached(
|
||||
|
|
@ -256,13 +159,7 @@ impl ClaudeCodeBackend {
|
|||
"claude-code: the child could not be spawned",
|
||||
));
|
||||
};
|
||||
let finished = read_stream(
|
||||
child.as_mut(),
|
||||
deadline,
|
||||
&request.model,
|
||||
request.output_schema_json.as_deref(),
|
||||
&self.quota,
|
||||
);
|
||||
let finished = read_stream(child.as_mut(), deadline, &request.model);
|
||||
child.kill();
|
||||
child.wait();
|
||||
let tail = child.stderr_tail();
|
||||
|
|
@ -271,7 +168,7 @@ impl ClaudeCodeBackend {
|
|||
}
|
||||
|
||||
match finished {
|
||||
Ok(finished) => self.completed(request, caller, secret.as_bytes(), finished),
|
||||
Ok(finished) => self.completed(request, finished),
|
||||
Err(failure_) => failure(failure_),
|
||||
}
|
||||
}
|
||||
|
|
@ -279,8 +176,6 @@ impl ClaudeCodeBackend {
|
|||
fn completed(
|
||||
&self,
|
||||
request: &ProviderRequest,
|
||||
caller: &str,
|
||||
secret: &[u8],
|
||||
finished: Finished,
|
||||
) -> (Vec<TransportEvent>, TransportOutcome) {
|
||||
let mut payloads: Vec<TransportEventPayload> = finished
|
||||
|
|
@ -295,30 +190,14 @@ impl ClaudeCodeBackend {
|
|||
text: finished.texts.concat(),
|
||||
});
|
||||
}
|
||||
// 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 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(),
|
||||
});
|
||||
|
||||
let events = payloads
|
||||
.into_iter()
|
||||
|
|
@ -341,11 +220,7 @@ impl ClaudeCodeBackend {
|
|||
|
||||
impl ProviderBackend for ClaudeCodeBackend {
|
||||
fn ready(&self) -> Result<(), &'static str> {
|
||||
readiness(
|
||||
&self.config,
|
||||
self.tokens.as_ref(),
|
||||
self.handle_secret.as_ref(),
|
||||
)
|
||||
readiness(&self.config, self.tokens.as_ref())
|
||||
}
|
||||
|
||||
fn provider(&self) -> &str {
|
||||
|
|
@ -353,23 +228,17 @@ impl ProviderBackend for ClaudeCodeBackend {
|
|||
}
|
||||
|
||||
fn execute(&self, call: &ProviderCall) -> TransportResult {
|
||||
let (events, outcome) = self.run(call.request(), call.caller_runtime_id(), false);
|
||||
let (events, outcome) = self.run(call.request());
|
||||
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, schema, quota);
|
||||
let mut state = StreamState::new(model);
|
||||
let finished = loop {
|
||||
match child.read_line(deadline) {
|
||||
StdoutRead::Line(line) => {
|
||||
|
|
@ -416,7 +285,6 @@ fn binary_matches_pin(config: &ClaudeCodeConfig) -> bool {
|
|||
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");
|
||||
|
|
@ -427,17 +295,12 @@ pub(crate) fn readiness(
|
|||
if !config.prompt_dir.is_dir() {
|
||||
return Err("prompt_dir");
|
||||
}
|
||||
if tokens.read().is_none() {
|
||||
return Err("token_credential");
|
||||
}
|
||||
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(())
|
||||
}
|
||||
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -25,9 +25,6 @@ pub struct ClaudeCodeConfig {
|
|||
/// 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,
|
||||
|
|
|
|||
|
|
@ -3,27 +3,20 @@
|
|||
//! 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::collections::{BTreeMap, BTreeSet};
|
||||
use std::net::{SocketAddr, TcpListener};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::sync::Arc;
|
||||
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,
|
||||
CallerAdmission, CallerConfig, MIN_PAYLOAD_BYTES, ModelPrice, ServeError, ServeLimits,
|
||||
ServiceError, load_admissions, serve,
|
||||
};
|
||||
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,
|
||||
};
|
||||
|
|
@ -31,22 +24,8 @@ use crate::{
|
|||
const CONFIG_KEY: &str = "runtime";
|
||||
const MODEL_PREFIX: &str = "claude-";
|
||||
|
||||
/// A unix account that may lease the subscription token for a Self run 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")]
|
||||
#[cultcache(type = "gamecult.model_connector.claude_connector_config.v1")]
|
||||
pub struct ClaudeConnectorConfig {
|
||||
/// A loopback socket address.
|
||||
#[cultcache(key = 0)]
|
||||
|
|
@ -86,26 +65,12 @@ pub struct ClaudeConnectorConfig {
|
|||
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.
|
||||
/// Exactly one caller until claude-handle-caller-binding lands.
|
||||
#[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,
|
||||
/// What each allowed model costs; keys are canonical model ids.
|
||||
#[cultcache(key = 18)]
|
||||
pub prices: BTreeMap<String, ModelPrice>,
|
||||
}
|
||||
|
||||
/// A refusal. It names the field or stage that failed and nothing the
|
||||
|
|
@ -140,6 +105,7 @@ impl From<ServeError> for ConfigError {
|
|||
ServeError::ConnectionKey(_) => Self::Unusable("connection_key_file"),
|
||||
ServeError::Service(error) => match error {
|
||||
ServiceError::ReplayStore(_) => Self::Unusable("replay_store"),
|
||||
ServiceError::UnpricedModel(_) => Self::Invalid("prices"),
|
||||
ServiceError::InvalidAdmission
|
||||
| ServiceError::DuplicateCaller
|
||||
| ServiceError::SharedCallerKey => Self::Invalid("callers"),
|
||||
|
|
@ -185,11 +151,6 @@ impl ClaudeConnectorConfig {
|
|||
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
|
||||
|
|
@ -204,72 +165,46 @@ impl ClaudeConnectorConfig {
|
|||
{
|
||||
return Err(ConfigError::Invalid("serve limits"));
|
||||
}
|
||||
if self.callers.is_empty() {
|
||||
let [caller] = self.callers.as_slice() else {
|
||||
return Err(ConfigError::Invalid("callers"));
|
||||
};
|
||||
if !is_trimmed_nonempty(&caller.caller_runtime_id) {
|
||||
return Err(ConfigError::Invalid("caller_runtime_id"));
|
||||
}
|
||||
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
|
||||
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
|
||||
|| caller.spend_ceiling_micro_usd == 0
|
||||
|| caller.spend_window_ms == 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 !is_canonical_model(model)
|
||||
|| !seen.insert(model.as_str())
|
||||
|| !self.prices.contains_key(model)
|
||||
{
|
||||
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())
|
||||
})
|
||||
if !self
|
||||
.prices
|
||||
.iter()
|
||||
.all(|(model, price)| is_canonical_model(model) && is_fully_priced(price))
|
||||
{
|
||||
return Err(ConfigError::Invalid("lease_callers"));
|
||||
return Err(ConfigError::Invalid("prices"));
|
||||
}
|
||||
Ok(())
|
||||
Ok(listen)
|
||||
}
|
||||
|
||||
fn backend_config(&self) -> ClaudeCodeConfig {
|
||||
|
|
@ -280,7 +215,6 @@ impl ClaudeConnectorConfig {
|
|||
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),
|
||||
|
|
@ -301,47 +235,52 @@ impl ClaudeConnectorConfig {
|
|||
}
|
||||
}
|
||||
|
||||
/// The `<instance>` of a unit name, as plain as a systemd instance can be.
|
||||
fn is_unit_instance(text: &str) -> bool {
|
||||
!text.is_empty()
|
||||
&& text
|
||||
.bytes()
|
||||
.all(|b| b.is_ascii_alphanumeric() || matches!(b, b'-' | b'_'))
|
||||
}
|
||||
|
||||
fn is_trimmed_nonempty(text: &str) -> bool {
|
||||
!text.is_empty() && text.trim() == text
|
||||
}
|
||||
|
||||
/// 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 {
|
||||
/// A model id the CLI echoes back unchanged: `claude-<family>-<major>-<minor>`
|
||||
/// with an optional `-YYYYMMDD` suffix, the family lowercase letters. Aliases
|
||||
/// and bare prefixes are refused because the stream's `init.model` must equal
|
||||
/// the requested model.
|
||||
fn is_canonical_model(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,
|
||||
};
|
||||
let major = parts.next().unwrap_or("");
|
||||
let minor = parts.next().unwrap_or("");
|
||||
let date = parts.next();
|
||||
!family.is_empty()
|
||||
&& family.bytes().all(|byte| byte.is_ascii_lowercase())
|
||||
&& digits(major, 2)
|
||||
&& after_ok
|
||||
&& digits(minor, 2)
|
||||
&& date.is_none_or(|date| date.len() == 8 && digits(date, 8))
|
||||
&& parts.next().is_none()
|
||||
}
|
||||
|
||||
/// A price is a fuse only if every component it bills is nonzero: a zero rate
|
||||
/// lets a call cost nothing against the caller's ceiling.
|
||||
fn is_fully_priced(price: &ModelPrice) -> bool {
|
||||
let long = price.long_context.map(|long| long.rates);
|
||||
[Some(price.rates), long]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.all(|rates| {
|
||||
[
|
||||
rates.input,
|
||||
rates.cache_write_short,
|
||||
rates.cache_write_long,
|
||||
rates.cache_read,
|
||||
rates.output,
|
||||
]
|
||||
.iter()
|
||||
.all(|rate| *rate > 0)
|
||||
})
|
||||
}
|
||||
|
||||
/// One file name under `$CREDENTIALS_DIRECTORY`, never a path.
|
||||
|
|
@ -444,91 +383,11 @@ fn readiness_of(config: &ClaudeConnectorConfig) -> Result<Vec<CallerAdmission>,
|
|||
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)
|
||||
|
|
@ -537,16 +396,6 @@ pub fn check_ready(config: &ClaudeConnectorConfig) -> Result<(), ConfigError> {
|
|||
/// 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()?;
|
||||
|
|
@ -555,26 +404,14 @@ fn serve_config_with(
|
|||
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("a);
|
||||
#[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)?;
|
||||
serve(
|
||||
listener,
|
||||
config.serve_limits(),
|
||||
admissions,
|
||||
config.prices.clone(),
|
||||
&config.replay_store,
|
||||
backend,
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
|
@ -582,8 +419,6 @@ fn serve_config_with(
|
|||
mod tests {
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use sha2::Digest;
|
||||
|
||||
use super::*;
|
||||
|
||||
/// A log that a test can read back.
|
||||
|
|
@ -600,104 +435,6 @@ mod tests {
|
|||
}
|
||||
}
|
||||
|
||||
/// `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";
|
||||
|
|
@ -718,6 +455,10 @@ mod tests {
|
|||
ServeError::Service(ServiceError::ReplayStore(format!("{CANARY}/replay.cc"))),
|
||||
"replay_store",
|
||||
),
|
||||
(
|
||||
ServeError::Service(ServiceError::UnpricedModel(CANARY.to_string())),
|
||||
"prices",
|
||||
),
|
||||
(
|
||||
ServeError::Service(ServiceError::InvalidAdmission),
|
||||
"callers",
|
||||
|
|
|
|||
|
|
@ -5,10 +5,9 @@
|
|||
|
||||
use std::path::PathBuf;
|
||||
|
||||
use hmac::{Hmac, Mac};
|
||||
use model_connector::{InputItem, ProviderRequest};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sha2::Sha256;
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use crate::config::{ClaudeCodeConfig, PROVIDER};
|
||||
|
||||
|
|
@ -17,11 +16,9 @@ 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
|
||||
/// 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.
|
||||
/// 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.
|
||||
Ended {
|
||||
session_id: String,
|
||||
instance: String,
|
||||
|
|
@ -39,43 +36,12 @@ impl ClaudeCodeHandle {
|
|||
}
|
||||
}
|
||||
|
||||
/// 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()
|
||||
/// 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()
|
||||
}
|
||||
|
||||
/// Whether the CLI's session store still holds a session.
|
||||
|
|
@ -123,6 +89,8 @@ 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")]
|
||||
|
|
@ -155,13 +123,14 @@ 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()
|
||||
|
|
@ -210,18 +179,7 @@ pub(crate) fn plan(
|
|||
.all(|item| matches!(item, InputItem::UserText { .. }));
|
||||
let resumable = instance == config.instance
|
||||
&& tail_is_user_text
|
||||
&& prefix_matches(
|
||||
secret,
|
||||
&Binding {
|
||||
caller,
|
||||
instructions: &request.instructions,
|
||||
model: &request.model,
|
||||
session_id: &session_id,
|
||||
instance: &instance,
|
||||
items: &input[..at],
|
||||
},
|
||||
&prefix_sha256,
|
||||
)
|
||||
&& prefix_sha256 == prefix_digest(&request.instructions, &request.model, &input[..at])
|
||||
&& sessions.exists(&session_id);
|
||||
if resumable {
|
||||
return Ok(Plan::Resume {
|
||||
|
|
|
|||
|
|
@ -9,22 +9,19 @@ 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 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;
|
||||
/// The one credential variable: the child's environment is the only
|
||||
/// environment the token ever enters.
|
||||
pub(crate) const TOKEN_VAR: &str = "CLAUDE_CODE_OAUTH_TOKEN";
|
||||
const CHILD_PATH: &str = "/usr/local/bin:/usr/bin:/bin";
|
||||
|
||||
/// The environment of the child: exactly these names and nothing inherited.
|
||||
/// `Debug` lists names only.
|
||||
pub struct ChildEnv(Vec<(&'static str, String)>);
|
||||
/// Values are never formatted; `Debug` lists names only.
|
||||
pub struct ChildEnv(Vec<(&'static str, Zeroizing<String>)>);
|
||||
|
||||
impl ChildEnv {
|
||||
pub fn iter(&self) -> impl Iterator<Item = (&str, &str)> {
|
||||
|
|
@ -40,17 +37,21 @@ impl fmt::Debug for ChildEnv {
|
|||
}
|
||||
}
|
||||
|
||||
pub(crate) fn child_env(config: &ClaudeCodeConfig, request: &ProviderRequest) -> ChildEnv {
|
||||
pub(crate) fn child_env(
|
||||
config: &ClaudeCodeConfig,
|
||||
request: &ProviderRequest,
|
||||
token: Zeroizing<String>,
|
||||
) -> ChildEnv {
|
||||
let config_dir = config.config_dir.to_string_lossy().into_owned();
|
||||
ChildEnv(vec![
|
||||
("PATH", CHILD_PATH.to_string()),
|
||||
("HOME", config_dir.clone()),
|
||||
("CLAUDE_CONFIG_DIR", config_dir),
|
||||
("PATH", Zeroizing::new(CHILD_PATH.to_string())),
|
||||
("HOME", Zeroizing::new(config_dir.clone())),
|
||||
("CLAUDE_CONFIG_DIR", Zeroizing::new(config_dir)),
|
||||
(
|
||||
"CLAUDE_CODE_MAX_OUTPUT_TOKENS",
|
||||
request.max_output_tokens.to_string(),
|
||||
Zeroizing::new(request.max_output_tokens.to_string()),
|
||||
),
|
||||
(TOKEN_FD_VAR, TOKEN_FD.to_string()),
|
||||
(TOKEN_VAR, token),
|
||||
])
|
||||
}
|
||||
|
||||
|
|
@ -107,15 +108,12 @@ 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. 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.
|
||||
/// for permission, and a resume always forks.
|
||||
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(),
|
||||
|
|
@ -141,15 +139,9 @@ 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(),
|
||||
|
|
|
|||
|
|
@ -1,745 +0,0 @@
|
|||
//! 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();
|
||||
}
|
||||
}
|
||||
|
|
@ -8,33 +8,16 @@ 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, CredentialTokenSource, StderrSink, TokenSource};
|
||||
pub use config::{ClaudeCodeConfig, ClaudeCodeLimits};
|
||||
pub use daemon::{
|
||||
ClaudeConnectorConfig, ConfigError, LeaseCallerConfig, check_ready, load_config, write_config,
|
||||
};
|
||||
pub use daemon::{ClaudeConnectorConfig, ConfigError, 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 handle::{FsSessionStore, 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;
|
||||
|
|
|
|||
|
|
@ -2,8 +2,6 @@
|
|||
//!
|
||||
//! 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
|
||||
|
|
@ -15,10 +13,9 @@ use std::path::PathBuf;
|
|||
use std::process::ExitCode;
|
||||
|
||||
use claude_connector::{
|
||||
ClaudeConnectorConfig, ConfigError, LeaseCallerConfig, MIN_SECRET_BYTES, check_ready,
|
||||
load_config, write_config,
|
||||
ClaudeConnectorConfig, ConfigError, check_ready, load_config, write_config,
|
||||
};
|
||||
use model_connector::{CallerConfig, LeaseTurn, frame_bytes_for_payload};
|
||||
use model_connector::{CallerConfig, ModelPrice, Rates, frame_bytes_for_payload};
|
||||
|
||||
fn main() -> ExitCode {
|
||||
let args: Vec<String> = std::env::args().skip(1).collect();
|
||||
|
|
@ -26,8 +23,6 @@ fn main() -> ExitCode {
|
|||
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),
|
||||
};
|
||||
|
|
@ -42,9 +37,7 @@ fn main() -> ExitCode {
|
|||
|
||||
#[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)"
|
||||
)]
|
||||
#[error("usage: serve --config PATH | config init PATH --flag value ... | config check PATH")]
|
||||
Usage,
|
||||
#[error(transparent)]
|
||||
Config(#[from] ConfigError),
|
||||
|
|
@ -71,28 +64,19 @@ fn check(path: &str) -> Result<(), Failure> {
|
|||
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());
|
||||
println!("allowed_models: {}", config.callers[0].allowed_models.len());
|
||||
println!("prices: {}", config.prices.len());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// The flags of one config verb: each single flag at most once, each list
|
||||
/// flag any number of times.
|
||||
/// The flags of `config init`, each given at most once except the two lists.
|
||||
struct Flags {
|
||||
single: BTreeMap<&'static str, String>,
|
||||
lists: BTreeMap<&'static str, Vec<String>>,
|
||||
models: Vec<String>,
|
||||
prices: 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] = [
|
||||
const SINGLE: [&str; 22] = [
|
||||
"--listen",
|
||||
"--replay-store",
|
||||
"--claude-binary",
|
||||
|
|
@ -102,63 +86,49 @@ const INIT: [&str; 21] = [
|
|||
"--prompt-dir",
|
||||
"--instance",
|
||||
"--token-credential",
|
||||
"--handle-secret-credential",
|
||||
"--handle-secret-file",
|
||||
"--caller-id",
|
||||
"--connection-key-file",
|
||||
"--connection-key-epoch",
|
||||
"--spend-ceiling-micro-usd",
|
||||
"--spend-window-ms",
|
||||
"--deadline-ms",
|
||||
"--max-line-bytes",
|
||||
"--max-total-bytes",
|
||||
"--stderr-tail-bytes",
|
||||
"--max-concurrent-requests",
|
||||
"--max-payload-bytes",
|
||||
"--max-output-tokens",
|
||||
"--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> {
|
||||
fn parse(args: &[&str]) -> Result<Self, Failure> {
|
||||
let mut flags = Flags {
|
||||
single: BTreeMap::new(),
|
||||
lists: BTreeMap::new(),
|
||||
models: Vec::new(),
|
||||
prices: Vec::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);
|
||||
match name {
|
||||
"--allowed-model" => flags.models.push(value.to_string()),
|
||||
"--price" => flags.prices.push(value.to_string()),
|
||||
_ => {
|
||||
let Some(known) = SINGLE.iter().find(|known| **known == 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))
|
||||
}
|
||||
|
|
@ -173,15 +143,6 @@ impl Flags {
|
|||
.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,
|
||||
|
|
@ -192,20 +153,6 @@ impl Flags {
|
|||
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]> {
|
||||
|
|
@ -219,87 +166,44 @@ fn sha256(hex: &str) -> Option<[u8; 32]> {
|
|||
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"));
|
||||
/// MODEL:INPUT:OUTPUT:CACHE_WRITE:CACHE_READ:OVERHEAD, rates in micro-USD per
|
||||
/// million tokens. The one cache-write rate prices both retentions.
|
||||
fn price(text: &str) -> Option<(String, ModelPrice)> {
|
||||
let mut parts = text.split(':');
|
||||
let model = parts.next()?.to_string();
|
||||
let mut next = || parts.next()?.parse::<u64>().ok();
|
||||
let (input, output, cache_write, cache_read, overhead_tokens) =
|
||||
(next()?, next()?, next()?, next()?, next()?);
|
||||
if parts.next().is_some() {
|
||||
return None;
|
||||
}
|
||||
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)
|
||||
Some((
|
||||
model,
|
||||
ModelPrice {
|
||||
rates: Rates {
|
||||
input,
|
||||
cache_write_short: cache_write,
|
||||
cache_write_long: cache_write,
|
||||
cache_read,
|
||||
output,
|
||||
},
|
||||
long_context: None,
|
||||
overhead_tokens,
|
||||
},
|
||||
))
|
||||
}
|
||||
|
||||
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 flags = Flags::parse(args)?;
|
||||
let max_payload_bytes: usize = flags.number_or("--max-payload-bytes", 1 << 20)?;
|
||||
let max_concurrent_requests: usize = flags.number_or("--max-concurrent-requests", 4)?;
|
||||
let mut prices = BTreeMap::new();
|
||||
for text in &flags.prices {
|
||||
let (model, price) = price(text).ok_or(Failure::Malformed("--price"))?;
|
||||
if prices.insert(model, price).is_some() {
|
||||
return Err(Failure::Malformed("--price"));
|
||||
}
|
||||
}
|
||||
let config = ClaudeConnectorConfig {
|
||||
listen: flags.text("--listen")?,
|
||||
replay_store: flags.path("--replay-store")?,
|
||||
|
|
@ -315,79 +219,23 @@ fn init(path: &str, args: &[&str]) -> Result<(), Failure> {
|
|||
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),
|
||||
max_frame_bytes: frame_bytes_for_payload(max_payload_bytes),
|
||||
max_connections: 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")?,
|
||||
callers: vec![CallerConfig {
|
||||
caller_runtime_id: flags.text("--caller-id")?,
|
||||
connection_key_file: flags.path("--connection-key-file")?,
|
||||
connection_key_epoch: flags.number("--connection-key-epoch")?,
|
||||
allowed_models: flags.models.clone(),
|
||||
max_concurrent_requests,
|
||||
max_payload_bytes,
|
||||
max_output_tokens: flags.number_or("--max-output-tokens", 32_768)?,
|
||||
spend_ceiling_micro_usd: flags.number("--spend-ceiling-micro-usd")?,
|
||||
spend_window_ms: flags.number("--spend-window-ms")?,
|
||||
}],
|
||||
prices,
|
||||
};
|
||||
// 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();
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,543 +0,0 @@
|
|||
//! 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);
|
||||
}
|
||||
}
|
||||
|
|
@ -20,9 +20,7 @@ pub enum StdoutRead {
|
|||
}
|
||||
|
||||
/// Starts a `claude` child. `argv[0]` is the program; the environment is
|
||||
/// 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.
|
||||
/// exactly `env`; `stdin` is written and closed.
|
||||
pub trait ProcessPort: Send + Sync {
|
||||
fn spawn(
|
||||
&self,
|
||||
|
|
@ -30,7 +28,6 @@ pub trait ProcessPort: Send + Sync {
|
|||
env: &ChildEnv,
|
||||
cwd: &Path,
|
||||
stdin: Vec<u8>,
|
||||
token: &str,
|
||||
) -> io::Result<Box<dyn RunningChild>>;
|
||||
}
|
||||
|
||||
|
|
@ -50,7 +47,6 @@ 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};
|
||||
|
|
@ -61,40 +57,7 @@ mod unix {
|
|||
|
||||
use super::{ProcessPort, RunningChild, StdoutRead};
|
||||
use crate::config::ClaudeCodeLimits;
|
||||
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(())
|
||||
}
|
||||
use crate::launch::ChildEnv;
|
||||
|
||||
/// Spawns each child in its own process group, so that killing the call
|
||||
/// kills what the child started too.
|
||||
|
|
@ -115,11 +78,7 @@ 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))?;
|
||||
|
|
@ -133,17 +92,7 @@ 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());
|
||||
|
||||
|
|
@ -280,83 +229,4 @@ 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());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,405 +0,0 @@
|
|||
//! 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("a.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));
|
||||
}
|
||||
}
|
||||
|
|
@ -3,18 +3,12 @@
|
|||
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use model_connector::{FailureKind, PlanQuota, StopReason, Usage};
|
||||
use model_connector::{FailureKind, 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`
|
||||
/// (unknown) unless a result event reported it or the
|
||||
/// (the kit bills its reservation) unless a result event reported it or the
|
||||
/// call never reached the provider (zero).
|
||||
#[derive(Debug, PartialEq, Eq)]
|
||||
pub(crate) struct Failure {
|
||||
|
|
@ -25,7 +19,7 @@ pub(crate) struct Failure {
|
|||
|
||||
impl Failure {
|
||||
/// A failure after the child ran: what the provider billed is unknown, so
|
||||
/// the usage is unknown.
|
||||
/// the kit bills its reservation.
|
||||
pub(crate) fn cost_unknown(kind: FailureKind, message: &'static str) -> Self {
|
||||
Self {
|
||||
kind,
|
||||
|
|
@ -69,18 +63,12 @@ struct Line {
|
|||
session_id: Option<String>,
|
||||
model: Option<String>,
|
||||
message: Option<Message>,
|
||||
/// 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>,
|
||||
rate_limit_info: Option<RateLimitInfo>,
|
||||
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>,
|
||||
|
|
@ -97,7 +85,11 @@ struct Block {
|
|||
#[serde(rename = "type")]
|
||||
kind: String,
|
||||
text: Option<String>,
|
||||
name: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct RateLimitInfo {
|
||||
status: Option<String>,
|
||||
}
|
||||
|
||||
/// One model's entry. The four counts are required: an entry that lacks one
|
||||
|
|
@ -115,13 +107,6 @@ 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>,
|
||||
|
|
@ -129,16 +114,9 @@ pub(crate) struct StreamState<'a> {
|
|||
}
|
||||
|
||||
impl<'a> StreamState<'a> {
|
||||
/// `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());
|
||||
pub(crate) fn new(model: &'a str) -> Self {
|
||||
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,
|
||||
|
|
@ -173,11 +151,10 @@ impl<'a> StreamState<'a> {
|
|||
},
|
||||
"assistant" => self.assistant(line.message),
|
||||
"rate_limit_event" => {
|
||||
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));
|
||||
}
|
||||
self.plan_rejected |= line
|
||||
.rate_limit_info
|
||||
.and_then(|info| info.status)
|
||||
.is_some_and(|status| status == "rejected");
|
||||
Step::More
|
||||
}
|
||||
"result" => Step::Done(Box::new(self.result(line))),
|
||||
|
|
@ -199,9 +176,6 @@ 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,
|
||||
|
|
@ -221,16 +195,7 @@ impl<'a> StreamState<'a> {
|
|||
line.api_error_status.as_ref().and_then(Value::as_u64),
|
||||
Some(401 | 403 | 429)
|
||||
);
|
||||
// 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,
|
||||
);
|
||||
let usage = read_usage(line.model_usage_since_last_result, door_refusal);
|
||||
if !success {
|
||||
return Err(Failure {
|
||||
kind: self.error_kind(&line.api_error_status, &line.result, &line.errors),
|
||||
|
|
@ -245,26 +210,8 @@ 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,
|
||||
|
|
@ -286,7 +233,7 @@ impl<'a> StreamState<'a> {
|
|||
};
|
||||
Ok(Finished {
|
||||
session_id,
|
||||
texts,
|
||||
texts: std::mem::take(&mut self.texts),
|
||||
message_id: self.message_id.take(),
|
||||
stop_reason,
|
||||
usage,
|
||||
|
|
@ -316,58 +263,20 @@ 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::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 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()?;
|
||||
/// of another shape, or empty. An empty report is believed only for a refusal
|
||||
/// at the door (401, 403 or 429), 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()?;
|
||||
if models.is_empty() {
|
||||
return empty_is_zero.then(|| sum_usage(std::iter::empty()));
|
||||
return door_refusal.then(|| sum_usage(std::iter::empty()));
|
||||
}
|
||||
Some(sum_usage(models.into_values()))
|
||||
}
|
||||
|
|
@ -396,50 +305,3 @@ 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}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,361 +0,0 @@
|
|||
//! 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);
|
||||
}
|
||||
}
|
||||
|
|
@ -1,21 +1,26 @@
|
|||
//! 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.
|
||||
//! document accepts, what it refuses, and that no refusal or output echoes a
|
||||
//! byte of a value that was typed, stored or read from a credential.
|
||||
|
||||
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 claude_connector::{ClaudeConnectorConfig, ConfigError, load_config, write_config};
|
||||
use cultcache_rs::{
|
||||
CacheBackingStore, CultCache, CultCacheEnvelope, DatabaseEntry,
|
||||
SingleFileMessagePackBackingStore,
|
||||
};
|
||||
use cultcache_rs::{CultCache, DatabaseEntry, SingleFileMessagePackBackingStore};
|
||||
use model_connector::{CallerConfig, LeaseTurn};
|
||||
use model_connector::{CallerConfig, MIN_PAYLOAD_BYTES, ModelPrice, Rates};
|
||||
|
||||
const CANARY: &str = "canary-9f3b7c1d";
|
||||
const OPUS: &str = "claude-opus-5-1";
|
||||
const HAIKU: &str = "claude-haiku-5-5";
|
||||
|
||||
/// An edit that should make the valid document refuse.
|
||||
type Change = Box<dyn Fn(&mut ClaudeConnectorConfig)>;
|
||||
|
||||
static SEQUENCE: AtomicUsize = AtomicUsize::new(0);
|
||||
|
||||
struct Scratch(PathBuf);
|
||||
|
|
@ -41,6 +46,20 @@ impl Drop for Scratch {
|
|||
}
|
||||
}
|
||||
|
||||
fn price() -> ModelPrice {
|
||||
ModelPrice {
|
||||
rates: Rates {
|
||||
input: 3_000_000,
|
||||
cache_write_short: 3_750_000,
|
||||
cache_write_long: 3_750_000,
|
||||
cache_read: 300_000,
|
||||
output: 15_000_000,
|
||||
},
|
||||
long_context: None,
|
||||
overhead_tokens: 100,
|
||||
}
|
||||
}
|
||||
|
||||
fn valid() -> ClaudeConnectorConfig {
|
||||
ClaudeConnectorConfig {
|
||||
listen: "127.0.0.1:7443".to_string(),
|
||||
|
|
@ -52,7 +71,6 @@ fn valid() -> ClaudeConnectorConfig {
|
|||
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,
|
||||
|
|
@ -69,15 +87,10 @@ fn valid() -> ClaudeConnectorConfig {
|
|||
max_concurrent_requests: 4,
|
||||
max_payload_bytes: 1 << 20,
|
||||
max_output_tokens: 32_768,
|
||||
spend_ceiling_micro_usd: 50_000_000,
|
||||
spend_window_ms: 3_600_000,
|
||||
}],
|
||||
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,
|
||||
prices: BTreeMap::from([(HAIKU.to_string(), price()), (OPUS.to_string(), price())]),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -98,119 +111,456 @@ fn config_round_trips_through_cultcache() {
|
|||
}
|
||||
|
||||
#[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() {
|
||||
fn a_second_record_or_the_wrong_type_is_refused() {
|
||||
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();
|
||||
}
|
||||
let path = scratch.path("two.cc");
|
||||
let config = valid();
|
||||
write_config(&path, &config).unwrap();
|
||||
let mut cache = CultCache::new();
|
||||
cache
|
||||
.register_entry_type::<ClaudeConnectorConfig>()
|
||||
.unwrap();
|
||||
cache
|
||||
.add_generic_backing_store(SingleFileMessagePackBackingStore::new(&path))
|
||||
.unwrap();
|
||||
cache.pull_all_backing_stores().unwrap();
|
||||
cache.put("extra", &config).unwrap();
|
||||
assert!(matches!(load_config(&path), Err(ConfigError::Store)));
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, DatabaseEntry)]
|
||||
#[cultcache(type = "gamecult.other.document.v1")]
|
||||
struct Other {
|
||||
#[cultcache(key = 0)]
|
||||
value: u32,
|
||||
}
|
||||
let other = scratch.path("other.cc");
|
||||
let mut cache = CultCache::new();
|
||||
cache.register_entry_type::<Other>().unwrap();
|
||||
cache
|
||||
.add_generic_backing_store(SingleFileMessagePackBackingStore::new(&other))
|
||||
.unwrap();
|
||||
cache.pull_all_backing_stores().unwrap();
|
||||
cache.put("runtime", &Other { value: 1 }).unwrap();
|
||||
assert!(matches!(load_config(&other), Err(ConfigError::Store)));
|
||||
|
||||
assert!(matches!(
|
||||
load_config(&scratch.path("absent.cc")),
|
||||
Err(ConfigError::Store)
|
||||
));
|
||||
}
|
||||
|
||||
fn store_raw(path: &Path, key: &str, ty: &str, schema: Option<&str>, payload: Vec<u8>) {
|
||||
let mut store = SingleFileMessagePackBackingStore::new(path);
|
||||
store
|
||||
.push(&CultCacheEnvelope {
|
||||
key: key.to_string(),
|
||||
r#type: ty.to_string(),
|
||||
payload,
|
||||
stored_at: "2026-10-29T00:00:00Z".to_string(),
|
||||
schema_id: schema.map(str::to_string),
|
||||
})
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[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");
|
||||
}
|
||||
fn the_envelope_must_be_exactly_the_documents_key_type_schema_and_encoding() {
|
||||
let scratch = Scratch::new();
|
||||
let canonical = rmp_serde::to_vec(&valid()).unwrap();
|
||||
let ty = ClaudeConnectorConfig::TYPE;
|
||||
let mut trailing = canonical.clone();
|
||||
trailing.push(0xc0);
|
||||
|
||||
#[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,
|
||||
let cases = vec![
|
||||
("wrong key", "other", ty, Some(ty), canonical.clone()),
|
||||
(
|
||||
"wrong type",
|
||||
"runtime",
|
||||
"gamecult.other.v1",
|
||||
Some(ty),
|
||||
canonical.clone(),
|
||||
),
|
||||
(
|
||||
"wrong schema",
|
||||
"runtime",
|
||||
ty,
|
||||
Some("gamecult.other.v1"),
|
||||
canonical.clone(),
|
||||
),
|
||||
("not canonical", "runtime", ty, Some(ty), trailing),
|
||||
];
|
||||
for edit in edits {
|
||||
for (name, key, ty, schema, payload) in cases {
|
||||
let path = scratch.path(&format!("{}.cc", name.replace(' ', "-")));
|
||||
store_raw(&path, key, ty, schema, payload);
|
||||
assert!(
|
||||
matches!(load_config(&path), Err(ConfigError::Store)),
|
||||
"{name}"
|
||||
);
|
||||
}
|
||||
|
||||
// The same bytes with every part right load.
|
||||
let path = scratch.path("right.cc");
|
||||
store_raw(&path, "runtime", ty, Some(ty), canonical);
|
||||
assert_eq!(load_config(&path).unwrap(), valid());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bounds_sit_exactly_where_the_connector_needs_them() {
|
||||
let mut config = valid();
|
||||
config.max_total_bytes = config.max_line_bytes;
|
||||
config.validate().unwrap();
|
||||
config.callers[0].max_payload_bytes = MIN_PAYLOAD_BYTES;
|
||||
config.validate().unwrap();
|
||||
config.callers[0].max_payload_bytes = MIN_PAYLOAD_BYTES - 1;
|
||||
assert_eq!(refused(&config), "caller bounds");
|
||||
|
||||
for set in [
|
||||
(|c: &mut ClaudeConnectorConfig| c.max_connections = 0) as fn(&mut ClaudeConnectorConfig),
|
||||
|c| c.socket_timeout_ms = 0,
|
||||
|c| c.max_expiry_skew_ms = 0,
|
||||
|c| c.max_frame_bytes = 0,
|
||||
] {
|
||||
let mut config = valid();
|
||||
edit(&mut config);
|
||||
assert_eq!(refused(&config), "caller bounds");
|
||||
set(&mut config);
|
||||
assert_eq!(refused(&config), "serve limits");
|
||||
}
|
||||
let mut config = valid();
|
||||
config.prices.insert("claude-".to_string(), price());
|
||||
assert_eq!(refused(&config), "prices");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_write_refuses_an_invalid_document_and_leaves_no_file() {
|
||||
let scratch = Scratch::new();
|
||||
let path = scratch.path("bad.cc");
|
||||
let mut config = valid();
|
||||
config.callers[0].allowed_models = vec!["opus".to_string()];
|
||||
assert!(write_config(&path, &config).is_err());
|
||||
assert!(!path.exists());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn config_refuses_aliases_unpriced_models_and_second_callers() {
|
||||
let valid_config = valid();
|
||||
assert_eq!(valid_config.validate().unwrap().port(), 7443);
|
||||
|
||||
let cases: Vec<(&str, Change)> = vec![
|
||||
(
|
||||
"alias allowed",
|
||||
Box::new(|c| {
|
||||
c.callers[0].allowed_models = vec!["opus".to_string()];
|
||||
}),
|
||||
),
|
||||
(
|
||||
"alias price key",
|
||||
Box::new(|c| {
|
||||
c.prices.insert("sonnet".to_string(), price());
|
||||
}),
|
||||
),
|
||||
(
|
||||
"bare prefix",
|
||||
Box::new(|c| {
|
||||
c.callers[0].allowed_models = vec!["claude-".to_string()];
|
||||
}),
|
||||
),
|
||||
(
|
||||
"family only",
|
||||
Box::new(|c| {
|
||||
c.callers[0].allowed_models = vec!["claude-haiku".to_string()];
|
||||
c.prices.insert("claude-haiku".to_string(), price());
|
||||
}),
|
||||
),
|
||||
(
|
||||
"latest alias",
|
||||
Box::new(|c| {
|
||||
c.callers[0].allowed_models = vec!["claude-sonnet-latest".to_string()];
|
||||
c.prices.insert("claude-sonnet-latest".to_string(), price());
|
||||
}),
|
||||
),
|
||||
(
|
||||
"no minor",
|
||||
Box::new(|c| {
|
||||
c.callers[0].allowed_models = vec!["claude-opus-5".to_string()];
|
||||
c.prices.insert("claude-opus-5".to_string(), price());
|
||||
}),
|
||||
),
|
||||
(
|
||||
"uppercase family",
|
||||
Box::new(|c| {
|
||||
c.callers[0].allowed_models = vec!["claude-Opus-5-1".to_string()];
|
||||
c.prices.insert("claude-Opus-5-1".to_string(), price());
|
||||
}),
|
||||
),
|
||||
(
|
||||
"short date",
|
||||
Box::new(|c| {
|
||||
c.callers[0].allowed_models = vec!["claude-opus-5-1-2026".to_string()];
|
||||
c.prices.insert("claude-opus-5-1-2026".to_string(), price());
|
||||
}),
|
||||
),
|
||||
(
|
||||
"trailing part",
|
||||
Box::new(|c| {
|
||||
c.callers[0].allowed_models = vec!["claude-opus-5-1-20260101-x".to_string()];
|
||||
c.prices
|
||||
.insert("claude-opus-5-1-20260101-x".to_string(), price());
|
||||
}),
|
||||
),
|
||||
(
|
||||
"unpriced allowed",
|
||||
Box::new(|c| {
|
||||
c.prices.remove(OPUS);
|
||||
}),
|
||||
),
|
||||
(
|
||||
"duplicate allowed",
|
||||
Box::new(|c| {
|
||||
c.callers[0].allowed_models = vec![HAIKU.to_string(), HAIKU.to_string()];
|
||||
}),
|
||||
),
|
||||
(
|
||||
"no allowed",
|
||||
Box::new(|c| c.callers[0].allowed_models.clear()),
|
||||
),
|
||||
(
|
||||
"second caller",
|
||||
Box::new(|c| {
|
||||
let mut second = c.callers[0].clone();
|
||||
second.caller_runtime_id = "other".to_string();
|
||||
c.callers.push(second);
|
||||
}),
|
||||
),
|
||||
("no caller", Box::new(|c| c.callers.clear())),
|
||||
];
|
||||
let fields = [
|
||||
"allowed_models",
|
||||
"prices",
|
||||
"allowed_models",
|
||||
"allowed_models",
|
||||
"allowed_models",
|
||||
"allowed_models",
|
||||
"allowed_models",
|
||||
"allowed_models",
|
||||
"allowed_models",
|
||||
"allowed_models",
|
||||
"allowed_models",
|
||||
"allowed_models",
|
||||
"callers",
|
||||
"callers",
|
||||
];
|
||||
for ((name, change), field) in cases.iter().zip(fields) {
|
||||
let mut config = valid();
|
||||
change(&mut config);
|
||||
assert_eq!(refused(&config), field, "{name}");
|
||||
}
|
||||
}
|
||||
|
||||
#[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 a_dated_canonical_id_is_allowed() {
|
||||
let dated = "claude-opus-5-1-20260101".to_string();
|
||||
let mut config = valid();
|
||||
config.callers[0].allowed_models = vec![dated.clone()];
|
||||
config.prices.insert(dated, price());
|
||||
config.validate().unwrap();
|
||||
}
|
||||
|
||||
fn run(args: &[&str]) -> Output {
|
||||
Command::new(env!("CARGO_BIN_EXE_claude-connector"))
|
||||
.args(args)
|
||||
.env_clear()
|
||||
.output()
|
||||
.unwrap()
|
||||
#[test]
|
||||
fn a_price_with_a_zero_billed_component_is_refused() {
|
||||
type Zero = fn(&mut Rates);
|
||||
let zeroed: Vec<(&str, Zero)> = vec![
|
||||
("input", |r| r.input = 0),
|
||||
("cache_write_short", |r| r.cache_write_short = 0),
|
||||
("cache_write_long", |r| r.cache_write_long = 0),
|
||||
("cache_read", |r| r.cache_read = 0),
|
||||
("output", |r| r.output = 0),
|
||||
];
|
||||
for (name, zero) in &zeroed {
|
||||
let mut config = valid();
|
||||
zero(&mut config.prices.get_mut(HAIKU).unwrap().rates);
|
||||
assert_eq!(refused(&config), "prices", "{name}");
|
||||
|
||||
let mut config = valid();
|
||||
let mut long = price().rates;
|
||||
zero(&mut long);
|
||||
config.prices.get_mut(HAIKU).unwrap().long_context = Some(model_connector::LongContext {
|
||||
above_prompt_tokens: 200_000,
|
||||
rates: long,
|
||||
});
|
||||
assert_eq!(refused(&config), "prices", "long context {name}");
|
||||
}
|
||||
// A zero overhead is not a rate.
|
||||
let mut config = valid();
|
||||
config.prices.get_mut(HAIKU).unwrap().overhead_tokens = 0;
|
||||
config.validate().unwrap();
|
||||
}
|
||||
|
||||
fn init_flags(scratch: &Scratch, extra: &[&str]) -> Vec<String> {
|
||||
#[test]
|
||||
fn config_refuses_exposed_listeners_relative_paths_and_empty_bounds() {
|
||||
let cases: Vec<(&str, &str, Change)> = vec![
|
||||
(
|
||||
"wildcard listen",
|
||||
"listen",
|
||||
Box::new(|c| c.listen = "0.0.0.0:7443".into()),
|
||||
),
|
||||
(
|
||||
"lan listen",
|
||||
"listen",
|
||||
Box::new(|c| c.listen = "192.168.1.9:7443".into()),
|
||||
),
|
||||
(
|
||||
"hostname listen",
|
||||
"listen",
|
||||
Box::new(|c| c.listen = "localhost:7443".into()),
|
||||
),
|
||||
(
|
||||
"replay",
|
||||
"replay_store",
|
||||
Box::new(|c| c.replay_store = "replay.cc".into()),
|
||||
),
|
||||
(
|
||||
"binary",
|
||||
"claude_binary",
|
||||
Box::new(|c| c.claude_binary = "claude".into()),
|
||||
),
|
||||
("work", "work_dir", Box::new(|c| c.work_dir = "work".into())),
|
||||
(
|
||||
"config",
|
||||
"config_dir",
|
||||
Box::new(|c| c.config_dir = "cfg".into()),
|
||||
),
|
||||
(
|
||||
"prompts",
|
||||
"prompt_dir",
|
||||
Box::new(|c| c.prompt_dir = "prompts".into()),
|
||||
),
|
||||
(
|
||||
"key file",
|
||||
"connection_key_file",
|
||||
Box::new(|c| {
|
||||
c.callers[0].connection_key_file = "caller.key".into();
|
||||
}),
|
||||
),
|
||||
(
|
||||
"zero pin",
|
||||
"claude_binary_sha256",
|
||||
Box::new(|c| c.claude_binary_sha256 = [0; 32]),
|
||||
),
|
||||
(
|
||||
"blank instance",
|
||||
"instance",
|
||||
Box::new(|c| c.instance = " ".into()),
|
||||
),
|
||||
(
|
||||
"credential path",
|
||||
"token_credential",
|
||||
Box::new(|c| c.token_credential = "a/b".into()),
|
||||
),
|
||||
(
|
||||
"credential parent",
|
||||
"token_credential",
|
||||
Box::new(|c| c.token_credential = "..".into()),
|
||||
),
|
||||
(
|
||||
"credential empty",
|
||||
"token_credential",
|
||||
Box::new(|c| c.token_credential = String::new()),
|
||||
),
|
||||
(
|
||||
"zero deadline",
|
||||
"child limits",
|
||||
Box::new(|c| c.deadline_ms = 0),
|
||||
),
|
||||
(
|
||||
"total below line",
|
||||
"child limits",
|
||||
Box::new(|c| c.max_total_bytes = c.max_line_bytes - 1),
|
||||
),
|
||||
(
|
||||
"zero frame",
|
||||
"serve limits",
|
||||
Box::new(|c| c.max_frame_bytes = 0),
|
||||
),
|
||||
(
|
||||
"zero epoch",
|
||||
"connection_key_epoch",
|
||||
Box::new(|c| c.callers[0].connection_key_epoch = 0),
|
||||
),
|
||||
(
|
||||
"tiny payload",
|
||||
"caller bounds",
|
||||
Box::new(|c| c.callers[0].max_payload_bytes = 16),
|
||||
),
|
||||
(
|
||||
"no ceiling",
|
||||
"caller bounds",
|
||||
Box::new(|c| c.callers[0].spend_ceiling_micro_usd = 0),
|
||||
),
|
||||
];
|
||||
for (name, field, change) in cases {
|
||||
let mut config = valid();
|
||||
change(&mut config);
|
||||
assert_eq!(refused(&config), field, "{name}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_refusal_echoes_a_value() {
|
||||
let canary_path: PathBuf = format!("{CANARY}/relative").into();
|
||||
let canaried: Vec<Change> = vec![
|
||||
Box::new(|c| c.listen = format!("{CANARY}:1")),
|
||||
Box::new(|c| c.instance = format!(" {CANARY} ")),
|
||||
Box::new(|c| c.token_credential = format!("{CANARY}/x")),
|
||||
Box::new(|c| c.callers[0].allowed_models = vec![format!("{CANARY}-model")]),
|
||||
Box::new(|c| {
|
||||
c.prices.insert(format!("{CANARY}-alias"), price());
|
||||
}),
|
||||
Box::new(|c| c.callers[0].caller_runtime_id = format!(" {CANARY}")),
|
||||
Box::new({
|
||||
let path = canary_path.clone();
|
||||
move |c| c.replay_store = path.clone()
|
||||
}),
|
||||
Box::new({
|
||||
let path = canary_path.clone();
|
||||
move |c| c.callers[0].connection_key_file = path.clone()
|
||||
}),
|
||||
Box::new(|c| {
|
||||
let mut second = c.callers[0].clone();
|
||||
second.caller_runtime_id = CANARY.to_string();
|
||||
c.callers.push(second);
|
||||
}),
|
||||
];
|
||||
for change in canaried {
|
||||
let mut config = valid();
|
||||
change(&mut config);
|
||||
let error = config.validate().unwrap_err();
|
||||
for shown in [error.to_string(), format!("{error:?}")] {
|
||||
assert!(!shown.contains(CANARY), "{shown}");
|
||||
}
|
||||
let scratch = Scratch::new();
|
||||
let write = write_config(&scratch.path("x.cc"), &config).unwrap_err();
|
||||
assert!(!format!("{write} {write:?}").contains(CANARY));
|
||||
}
|
||||
}
|
||||
|
||||
// ---- the binary ----
|
||||
|
||||
fn run(args: &[&str], envs: &[(&str, &Path)]) -> Output {
|
||||
let mut command = Command::new(env!("CARGO_BIN_EXE_claude-connector"));
|
||||
command.args(args).env_clear();
|
||||
for (name, value) in envs {
|
||||
command.env(name, value);
|
||||
}
|
||||
command.output().unwrap()
|
||||
}
|
||||
|
||||
fn shown(output: &Output) -> String {
|
||||
format!(
|
||||
"{}{}",
|
||||
String::from_utf8_lossy(&output.stdout),
|
||||
String::from_utf8_lossy(&output.stderr)
|
||||
)
|
||||
}
|
||||
|
||||
fn init_flags<'a>(scratch: &'a Scratch, extra: &[&'a str]) -> Vec<String> {
|
||||
let path = |name: &str| scratch.path(name).to_string_lossy().into_owned();
|
||||
let mut flags: Vec<String> = [
|
||||
"--listen",
|
||||
|
|
@ -231,18 +581,20 @@ fn init_flags(scratch: &Scratch, extra: &[&str]) -> Vec<String> {
|
|||
"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",
|
||||
"--spend-ceiling-micro-usd",
|
||||
"50000000",
|
||||
"--spend-window-ms",
|
||||
"3600000",
|
||||
"--allowed-model",
|
||||
HAIKU,
|
||||
"--price",
|
||||
&format!("{HAIKU}:1000000:5000000:1250000:100000:100"),
|
||||
]
|
||||
.map(String::from)
|
||||
.to_vec();
|
||||
|
|
@ -254,44 +606,297 @@ 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)
|
||||
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);
|
||||
assert!(made.status.success(), "{}", shown(&made));
|
||||
let config = load_config(&scratch.path("connector.cc")).unwrap();
|
||||
assert_eq!(config.callers[0].allowed_models, [HAIKU]);
|
||||
assert_eq!(config.prices[HAIKU].rates.output, 5_000_000);
|
||||
assert_eq!(config.prices[HAIKU].rates.cache_write_long, 1_250_000);
|
||||
assert_eq!(config.prices[HAIKU].overhead_tokens, 100);
|
||||
assert_eq!(config.claude_binary_sha256, [0xab; 32]);
|
||||
assert!(config.lease_callers.is_empty());
|
||||
}
|
||||
|
||||
/// A scratch directory in which `serve` would start: every directory, the
|
||||
/// pinned binary, the connection key and the credential exist. Returns the
|
||||
/// `config init` flags and the credential directory.
|
||||
#[cfg(unix)]
|
||||
fn prepare(scratch: &Scratch) -> (Vec<String>, PathBuf) {
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
|
||||
for name in ["work", "home", "prompts", "credentials"] {
|
||||
std::fs::create_dir_all(scratch.path(name)).unwrap();
|
||||
}
|
||||
std::fs::set_permissions(
|
||||
scratch.path("prompts"),
|
||||
std::fs::Permissions::from_mode(0o700),
|
||||
)
|
||||
.unwrap();
|
||||
std::fs::write(scratch.path("credentials/claude-token"), CANARY).unwrap();
|
||||
std::fs::write(scratch.path("caller.key"), "caller-secret").unwrap();
|
||||
std::fs::write(scratch.path("claude"), b"pinned claude").unwrap();
|
||||
let sha: String = Sha256::digest(b"pinned claude")
|
||||
.iter()
|
||||
.map(|byte| format!("{byte:02x}"))
|
||||
.collect();
|
||||
let mut flags = init_flags(scratch, &[]);
|
||||
let at = flags
|
||||
.iter()
|
||||
.position(|flag| flag == "--claude-sha256")
|
||||
.unwrap();
|
||||
flags[at + 1] = sha;
|
||||
(flags, scratch.path("credentials"))
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
fn check_and_serve(config: &Path, credentials: Option<&Path>) -> (Output, Output) {
|
||||
let envs: Vec<(&str, &Path)> = credentials
|
||||
.map(|dir| ("CREDENTIALS_DIRECTORY", dir))
|
||||
.into_iter()
|
||||
.collect();
|
||||
let path = config.to_str().unwrap();
|
||||
(
|
||||
run(&["config", "check", path], &envs),
|
||||
run(&["serve", "--config", path], &envs),
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn check_passes_a_document_serve_can_start_and_reports_counts_only() {
|
||||
let scratch = Scratch::new();
|
||||
let (flags, credentials) = prepare(&scratch);
|
||||
let made = init(&scratch, &flags);
|
||||
assert!(made.status.success(), "{}", shown(&made));
|
||||
let checked = run(
|
||||
&[
|
||||
"config",
|
||||
"check",
|
||||
scratch.path("connector.cc").to_str().unwrap(),
|
||||
],
|
||||
&[("CREDENTIALS_DIRECTORY", &credentials)],
|
||||
);
|
||||
assert!(checked.status.success(), "{}", shown(&checked));
|
||||
let text = shown(&checked);
|
||||
assert!(text.contains("callers: 1") && text.contains("allowed_models: 1"));
|
||||
assert!(text.contains("prices: 1"));
|
||||
for secret in ["persona-a", HAIKU, CANARY, "caller-secret"] {
|
||||
assert!(!text.contains(secret), "{text}");
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn check_refuses_exactly_what_serve_refuses_and_names_the_field() {
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
|
||||
type Break = fn(&Scratch, &Path);
|
||||
let cases: Vec<(&str, &str, Break)> = vec![
|
||||
("open prompt dir", "prompt directory", |s, _| {
|
||||
std::fs::set_permissions(s.path("prompts"), std::fs::Permissions::from_mode(0o755))
|
||||
.unwrap();
|
||||
}),
|
||||
("no key file", "connection_key_file", |s, _| {
|
||||
std::fs::remove_file(s.path("caller.key")).unwrap();
|
||||
}),
|
||||
("binary off its pin", "claude_binary_sha256", |s, _| {
|
||||
std::fs::write(s.path("claude"), b"other bytes").unwrap();
|
||||
}),
|
||||
("no credential", "token_credential", |_, credentials| {
|
||||
std::fs::remove_file(credentials.join("claude-token")).unwrap();
|
||||
}),
|
||||
("no work dir", "work_dir", |s, _| {
|
||||
std::fs::remove_dir(s.path("work")).unwrap();
|
||||
}),
|
||||
("no config dir", "config_dir", |s, _| {
|
||||
std::fs::remove_dir(s.path("home")).unwrap();
|
||||
}),
|
||||
("frame too small", "max_frame_bytes", |s, _| {
|
||||
let path = s.path("connector.cc");
|
||||
let mut config = load_config(&path).unwrap();
|
||||
config.max_frame_bytes = 4099;
|
||||
write_config(&path, &config).unwrap();
|
||||
}),
|
||||
("too few connections", "connection capacity", |s, _| {
|
||||
let path = s.path("connector.cc");
|
||||
let mut config = load_config(&path).unwrap();
|
||||
config.max_connections = 1;
|
||||
write_config(&path, &config).unwrap();
|
||||
}),
|
||||
];
|
||||
for (name, field, break_it) in cases {
|
||||
let scratch = Scratch::new();
|
||||
let (flags, credentials) = prepare(&scratch);
|
||||
let made = init(&scratch, &flags);
|
||||
assert!(made.status.success(), "{name}: {}", shown(&made));
|
||||
break_it(&scratch, &credentials);
|
||||
let (checked, served) = check_and_serve(&scratch.path("connector.cc"), Some(&credentials));
|
||||
assert!(!checked.status.success(), "{name} check passed");
|
||||
assert!(!served.status.success(), "{name} serve started");
|
||||
assert!(
|
||||
shown(&checked).contains(field),
|
||||
"{name}: {}",
|
||||
shown(&checked)
|
||||
);
|
||||
assert_eq!(shown(&checked), shown(&served), "{name}");
|
||||
for secret in [CANARY, "caller-secret", "other bytes"] {
|
||||
assert!(!shown(&checked).contains(secret), "{name}");
|
||||
}
|
||||
}
|
||||
|
||||
// No credential directory at all is the same refusal.
|
||||
let scratch = Scratch::new();
|
||||
let (flags, _) = prepare(&scratch);
|
||||
assert!(init(&scratch, &flags).status.success());
|
||||
let (checked, served) = check_and_serve(&scratch.path("connector.cc"), None);
|
||||
assert!(!checked.status.success() && !served.status.success());
|
||||
assert!(shown(&checked).contains("token_credential"));
|
||||
assert_eq!(shown(&checked), shown(&served));
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn serve_names_a_corrupt_replay_store_without_its_path() {
|
||||
let scratch = Scratch::new();
|
||||
let (mut flags, credentials) = prepare(&scratch);
|
||||
let store = scratch.path(&format!("{CANARY}-store/replay.cc"));
|
||||
std::fs::create_dir_all(store.parent().unwrap()).unwrap();
|
||||
std::fs::write(&store, b"not a cultcache store").unwrap();
|
||||
let at = flags
|
||||
.iter()
|
||||
.position(|flag| flag == "--replay-store")
|
||||
.unwrap();
|
||||
flags[at + 1] = store.to_string_lossy().into_owned();
|
||||
assert!(init(&scratch, &flags).status.success());
|
||||
let (_, served) = check_and_serve(&scratch.path("connector.cc"), Some(&credentials));
|
||||
assert!(!served.status.success());
|
||||
let text = shown(&served);
|
||||
assert!(text.contains("replay_store"), "{text}");
|
||||
assert!(
|
||||
!text.contains(CANARY) && !text.contains("replay.cc"),
|
||||
"{text}"
|
||||
);
|
||||
}
|
||||
|
||||
#[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"],
|
||||
fn the_binary_refuses_without_echoing_what_it_was_given() {
|
||||
let spaced = format!(" {CANARY}");
|
||||
// Each pair replaces the value of an existing flag; the others are appended.
|
||||
let replaced: Vec<(&str, &str)> = vec![
|
||||
("--listen", CANARY),
|
||||
("--claude-sha256", CANARY),
|
||||
("--claude-sha256", "abcd"),
|
||||
("--price", CANARY),
|
||||
("--instance", &spaced),
|
||||
("--connection-key-epoch", CANARY),
|
||||
("--allowed-model", CANARY),
|
||||
];
|
||||
for extra in dollar_flags {
|
||||
let appended: Vec<Vec<&str>> = vec![
|
||||
vec![CANARY, CANARY],
|
||||
vec!["--listen"],
|
||||
vec!["--listen", "127.0.0.1:1"],
|
||||
];
|
||||
let mut attempts: Vec<Vec<String>> = Vec::new();
|
||||
for (name, value) in replaced {
|
||||
let scratch = Scratch::new();
|
||||
let made = init(&scratch, &init_flags(&scratch, extra));
|
||||
assert!(!made.status.success(), "{extra:?} was accepted");
|
||||
let mut flags = init_flags(&scratch, &[]);
|
||||
let at = flags.iter().position(|flag| flag == name).unwrap();
|
||||
flags[at + 1] = value.to_string();
|
||||
attempts.push(flags);
|
||||
}
|
||||
for extra in appended {
|
||||
let scratch = Scratch::new();
|
||||
attempts.push(init_flags(&scratch, &extra));
|
||||
}
|
||||
for flags in attempts {
|
||||
let scratch = Scratch::new();
|
||||
let made = init(&scratch, &flags);
|
||||
assert!(!made.status.success(), "{flags:?}");
|
||||
assert!(!shown(&made).contains(CANARY), "{}", shown(&made));
|
||||
assert!(!scratch.path("connector.cc").exists());
|
||||
}
|
||||
|
||||
// A missing flag is refused by name.
|
||||
let scratch = Scratch::new();
|
||||
let mut flags = init_flags(&scratch, &[]);
|
||||
let at = flags.iter().position(|f| f == "--instance").unwrap();
|
||||
flags.drain(at..at + 2);
|
||||
let made = init(&scratch, &flags);
|
||||
assert!(!made.status.success());
|
||||
assert!(shown(&made).contains("--instance"));
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn serve_refuses_an_unsafe_prompt_dir_and_prints_no_credential() {
|
||||
use std::os::unix::fs::{PermissionsExt, symlink};
|
||||
|
||||
let scratch = Scratch::new();
|
||||
let credentials = scratch.path("credentials");
|
||||
std::fs::create_dir_all(&credentials).unwrap();
|
||||
std::fs::write(credentials.join("claude-token"), CANARY).unwrap();
|
||||
std::fs::write(scratch.path("caller.key"), "caller-secret").unwrap();
|
||||
let made = init(&scratch, &init_flags(&scratch, &[]));
|
||||
assert!(made.status.success(), "{}", shown(&made));
|
||||
let config = scratch.path("connector.cc");
|
||||
let serve = |config: &Path| {
|
||||
run(
|
||||
&["serve", "--config", config.to_str().unwrap()],
|
||||
&[("CREDENTIALS_DIRECTORY", &credentials)],
|
||||
)
|
||||
};
|
||||
|
||||
// No prompt directory, then one others can enter, then a link to a good one.
|
||||
let prompts = scratch.path("prompts");
|
||||
let refused = serve(&config);
|
||||
assert!(!refused.status.success());
|
||||
assert!(
|
||||
shown(&refused).contains("prompt directory"),
|
||||
"{}",
|
||||
shown(&refused)
|
||||
);
|
||||
std::fs::create_dir_all(&prompts).unwrap();
|
||||
std::fs::set_permissions(&prompts, std::fs::Permissions::from_mode(0o755)).unwrap();
|
||||
let open = serve(&config);
|
||||
assert!(
|
||||
shown(&open).contains("prompt directory"),
|
||||
"{}",
|
||||
shown(&open)
|
||||
);
|
||||
assert!(claude_connector::check_prompt_dir(&prompts).is_err());
|
||||
std::fs::set_permissions(&prompts, std::fs::Permissions::from_mode(0o700)).unwrap();
|
||||
assert!(claude_connector::check_prompt_dir(&prompts).is_ok());
|
||||
let linked = scratch.path("linked");
|
||||
symlink(&prompts, &linked).unwrap();
|
||||
assert!(claude_connector::check_prompt_dir(&linked).is_err());
|
||||
|
||||
// A safe directory but no claude binary: the backend is not ready, and
|
||||
// nothing the credential holds is printed.
|
||||
let not_ready = serve(&config);
|
||||
assert!(!not_ready.status.success());
|
||||
for output in [&refused, &open, ¬_ready] {
|
||||
let text = shown(output);
|
||||
assert!(
|
||||
!Path::new(&scratch.path("connector.cc")).exists(),
|
||||
"{extra:?}"
|
||||
!text.contains(CANARY) && !text.contains("caller-secret"),
|
||||
"{text}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn serve_refuses_a_document_that_is_not_there() {
|
||||
fn a_private_regular_file_is_not_a_prompt_dir() {
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
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());
|
||||
let file = scratch.path("file");
|
||||
std::fs::write(&file, "x").unwrap();
|
||||
std::fs::set_permissions(&file, std::fs::Permissions::from_mode(0o700)).unwrap();
|
||||
assert!(claude_connector::check_prompt_dir(&file).is_err());
|
||||
assert!(claude_connector::check_prompt_dir(&scratch.path("absent")).is_err());
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,4 +0,0 @@
|
|||
{"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"}}}
|
||||
|
|
@ -1,16 +0,0 @@
|
|||
# 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.
Binary file not shown.
Binary file not shown.
|
|
@ -1,651 +0,0 @@
|
|||
//! 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);
|
||||
}
|
||||
|
|
@ -1,99 +0,0 @@
|
|||
//! 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),
|
||||
],
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
|
@ -4,16 +4,11 @@ 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,
|
||||
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,
|
||||
MAX_IDENTITY_BYTES, ServiceError, TransportFrameError, decode_result_frame, decrypt_result,
|
||||
encode_invocation_frame, encrypt_invocation, read_transport_frame, write_transport_frame,
|
||||
};
|
||||
use crate::readout::UsageReadout;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ConnectorClient {
|
||||
|
|
@ -57,34 +52,6 @@ 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))
|
||||
|
|
@ -95,9 +62,12 @@ 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)?;
|
||||
read_transport_frame(&mut stream, 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)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -256,56 +226,4 @@ 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();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -9,11 +9,11 @@ use zeroize::Zeroize;
|
|||
|
||||
use crate::envelope::{ServiceError, binary_schema};
|
||||
|
||||
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";
|
||||
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";
|
||||
|
||||
#[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. Cache reads
|
||||
/// and writes are counted apart from plain input.
|
||||
/// Token accounting for one request, as the backend reports it; the service
|
||||
/// prices 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 counted apart.
|
||||
/// (absent when the provider does not report it) and is never priced apart.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct Usage {
|
||||
|
|
@ -252,6 +252,10 @@ 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,
|
||||
}
|
||||
|
||||
|
|
@ -260,6 +264,7 @@ 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)?;
|
||||
|
|
@ -269,6 +274,7 @@ impl TransportInvocation {
|
|||
expires_at_unix_ms,
|
||||
native_request_sha256,
|
||||
provider_request_sha256,
|
||||
max_cost_micro_usd,
|
||||
request,
|
||||
})
|
||||
}
|
||||
|
|
@ -281,6 +287,9 @@ 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);
|
||||
|
|
@ -476,6 +485,15 @@ 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,
|
||||
}
|
||||
|
||||
|
|
@ -540,6 +558,12 @@ 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,
|
||||
|
|
@ -549,12 +573,15 @@ 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 provider itself is refusing right now: a plan window it reported
|
||||
// `rejected` has not reached its own reset. Plain comments, as above.
|
||||
ProviderRejected,
|
||||
// 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 one wire codec of every v4 document: named MessagePack (struct fields
|
||||
/// The one wire codec of every v3 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)
|
||||
|
|
@ -760,23 +787,28 @@ mod tests {
|
|||
}
|
||||
|
||||
#[test]
|
||||
fn the_invocations_identity_covers_its_request() {
|
||||
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() {
|
||||
let one = invocation();
|
||||
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();
|
||||
let mut other = one.clone();
|
||||
other.max_cost_micro_usd += 1;
|
||||
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, one);
|
||||
assert_eq!(decoded.max_cost_micro_usd, one.max_cost_micro_usd);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
|
@ -837,7 +869,7 @@ mod tests {
|
|||
}
|
||||
|
||||
#[test]
|
||||
fn v4_documents_carry_their_v4_ids() {
|
||||
fn v3_documents_carry_their_v3_ids() {
|
||||
assert_eq!(
|
||||
[
|
||||
PROVIDER_REQUEST_SCHEMA_ID,
|
||||
|
|
@ -847,11 +879,11 @@ mod tests {
|
|||
ENVELOPE_SCHEMA_ID
|
||||
],
|
||||
[
|
||||
"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"
|
||||
"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"
|
||||
]
|
||||
);
|
||||
let invocation = invocation();
|
||||
|
|
@ -890,7 +922,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 v4 document is a named map");
|
||||
panic!("a v3 document is a named map");
|
||||
};
|
||||
entries.push((rmpv::Value::from("unexpected_field"), rmpv::Value::from(1)));
|
||||
let mut out = Vec::new();
|
||||
|
|
|
|||
|
|
@ -17,14 +17,11 @@ 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)]
|
||||
|
|
@ -52,10 +49,8 @@ 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 operation that runs one sealed invocation.
|
||||
/// The one operation of that service: run 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";
|
||||
|
|
@ -74,39 +69,16 @@ fn seal_payload(envelope: &TransportEnvelope) -> Result<String, ServiceError> {
|
|||
Ok(STANDARD.encode(wire_encode(envelope)?))
|
||||
}
|
||||
|
||||
/// 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
|
||||
/// The CultNet frame an invocation travels in: a `cultnet.operation_request.v0`
|
||||
/// for [`CONNECTOR_SERVICE_ID`] and [`INVOKE_OPERATION`], its `message_id` the
|
||||
/// request id and its payload the base64 of the named sealed envelope.
|
||||
fn request_frame(
|
||||
operation: Operation,
|
||||
pub(crate) fn encode_invocation_frame(
|
||||
envelope: &TransportEnvelope,
|
||||
) -> Result<Vec<u8>, ServiceError> {
|
||||
carriage_bytes(&CultNetMessage::OperationRequest {
|
||||
message_id: envelope.request_id.clone(),
|
||||
service_id: CONNECTOR_SERVICE_ID.to_string(),
|
||||
operation: operation.name().to_string(),
|
||||
operation: INVOKE_OPERATION.to_string(),
|
||||
payload_schema: ENVELOPE_SCHEMA_ID.to_string(),
|
||||
payload_encoding: PAYLOAD_ENCODING.to_string(),
|
||||
payload: seal_payload(envelope)?,
|
||||
|
|
@ -115,17 +87,14 @@ fn request_frame(
|
|||
})
|
||||
}
|
||||
|
||||
/// The CultNet frame a response travels in: a `cultnet.operation_response.v0`
|
||||
/// with status `ok`, otherwise shaped like the request's.
|
||||
/// The CultNet frame a result travels in: a `cultnet.operation_response.v0`
|
||||
/// with status `ok`, otherwise shaped like the invocation's.
|
||||
#[cfg(any(feature = "daemon", test))]
|
||||
fn response_frame(
|
||||
operation: Operation,
|
||||
envelope: &TransportEnvelope,
|
||||
) -> Result<Vec<u8>, ServiceError> {
|
||||
pub(crate) fn encode_result_frame(envelope: &TransportEnvelope) -> Result<Vec<u8>, ServiceError> {
|
||||
carriage_bytes(&CultNetMessage::OperationResponse {
|
||||
message_id: envelope.request_id.clone(),
|
||||
service_id: CONNECTOR_SERVICE_ID.to_string(),
|
||||
operation: operation.name().to_string(),
|
||||
operation: INVOKE_OPERATION.to_string(),
|
||||
status: OK_STATUS.to_string(),
|
||||
payload_schema: ENVELOPE_SCHEMA_ID.to_string(),
|
||||
payload_encoding: PAYLOAD_ENCODING.to_string(),
|
||||
|
|
@ -135,43 +104,20 @@ fn response_frame(
|
|||
})
|
||||
}
|
||||
|
||||
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, that the operation is one this
|
||||
/// service has, the payload schema, and that the message id is the sealed
|
||||
/// envelope's request id.
|
||||
/// here, before any decryption: the service, the operation, 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<(Operation, TransportEnvelope), ServiceError> {
|
||||
let operation = Operation::from_name(operation).ok_or(ServiceError::Carriage)?;
|
||||
if service_id != CONNECTOR_SERVICE_ID || payload_schema != ENVELOPE_SCHEMA_ID {
|
||||
) -> Result<TransportEnvelope, ServiceError> {
|
||||
if service_id != CONNECTOR_SERVICE_ID
|
||||
|| operation != INVOKE_OPERATION
|
||||
|| payload_schema != ENVELOPE_SCHEMA_ID
|
||||
{
|
||||
return Err(ServiceError::Carriage);
|
||||
}
|
||||
let bytes = STANDARD
|
||||
|
|
@ -181,7 +127,7 @@ fn open_carriage(
|
|||
if envelope.request_id != message_id {
|
||||
return Err(ServiceError::OuterIdentity);
|
||||
}
|
||||
Ok((operation, envelope))
|
||||
Ok(envelope)
|
||||
}
|
||||
|
||||
fn read_carriage(bytes: &[u8]) -> Result<CultNetMessage, ServiceError> {
|
||||
|
|
@ -189,15 +135,8 @@ 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) enum RequestFrame {
|
||||
Invoke(TransportEnvelope),
|
||||
UsageReadout(TransportEnvelope),
|
||||
}
|
||||
|
||||
#[cfg(any(feature = "daemon", test))]
|
||||
pub(crate) fn decode_request_frame(bytes: &[u8]) -> Result<RequestFrame, ServiceError> {
|
||||
pub(crate) fn decode_invocation_frame(bytes: &[u8]) -> Result<TransportEnvelope, ServiceError> {
|
||||
match read_carriage(bytes)? {
|
||||
CultNetMessage::OperationRequest {
|
||||
message_id,
|
||||
|
|
@ -206,34 +145,18 @@ pub(crate) fn decode_request_frame(bytes: &[u8]) -> Result<RequestFrame, Service
|
|||
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)),
|
||||
}
|
||||
}
|
||||
} => open_carriage(
|
||||
&message_id,
|
||||
&service_id,
|
||||
&operation,
|
||||
&payload_schema,
|
||||
&payload,
|
||||
),
|
||||
_ => Err(ServiceError::Carriage),
|
||||
}
|
||||
}
|
||||
|
||||
#[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> {
|
||||
pub(crate) fn decode_result_frame(bytes: &[u8]) -> Result<TransportEnvelope, ServiceError> {
|
||||
match read_carriage(bytes)? {
|
||||
CultNetMessage::OperationResponse {
|
||||
message_id,
|
||||
|
|
@ -243,32 +166,17 @@ fn decode_response_frame(
|
|||
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),
|
||||
}
|
||||
}
|
||||
} if status == OK_STATUS => open_carriage(
|
||||
&message_id,
|
||||
&service_id,
|
||||
&operation,
|
||||
&payload_schema,
|
||||
&payload,
|
||||
),
|
||||
_ => 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
|
||||
|
|
@ -416,59 +324,6 @@ 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),
|
||||
|
|
@ -607,6 +462,8 @@ 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")]
|
||||
|
|
@ -1070,7 +927,7 @@ mod tests {
|
|||
fn golden_frame() -> Vec<u8> {
|
||||
std::fs::read(
|
||||
std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
|
||||
.join("tests/fixtures/wire_v4/carriage_invocation.frame"),
|
||||
.join("tests/fixtures/wire_v3/carriage_invocation.frame"),
|
||||
)
|
||||
.expect("carriage_invocation.frame is committed")
|
||||
}
|
||||
|
|
@ -1107,7 +964,7 @@ mod tests {
|
|||
assert_eq!(text(&request, "operation"), "invoke");
|
||||
assert_eq!(
|
||||
text(&request, "payloadSchema"),
|
||||
"gamecult.model.transport_envelope.v4"
|
||||
"gamecult.model.transport_envelope.v3"
|
||||
);
|
||||
assert_eq!(text(&request, "payloadEncoding"), "messagepack-base64");
|
||||
assert_eq!(text(&request, "messageId"), envelope.request_id);
|
||||
|
|
@ -1121,7 +978,7 @@ mod tests {
|
|||
assert_eq!(text(&response, "status"), "ok");
|
||||
assert_eq!(
|
||||
text(&response, "payloadSchema"),
|
||||
"gamecult.model.transport_envelope.v4"
|
||||
"gamecult.model.transport_envelope.v3"
|
||||
);
|
||||
assert_eq!(text(&response, "payloadEncoding"), "messagepack-base64");
|
||||
assert_eq!(text(&response, "messageId"), envelope.request_id);
|
||||
|
|
@ -1281,53 +1138,4 @@ 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);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -67,6 +67,14 @@ 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")]
|
||||
|
|
@ -125,8 +133,16 @@ 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], request()).unwrap()
|
||||
TransportInvocation::new("epiphany-yggdrasil", 2_000, [7; 32], MAX_COST, request()).unwrap()
|
||||
}
|
||||
|
||||
pub(crate) fn receipt(invocation: &TransportInvocation) -> TransportReceipt {
|
||||
|
|
@ -139,6 +155,8 @@ 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,
|
||||
|
|
@ -152,6 +170,33 @@ 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(
|
||||
|
|
@ -167,6 +212,8 @@ pub(crate) fn caller(
|
|||
max_concurrent_requests,
|
||||
64 * 1024,
|
||||
32_768,
|
||||
SpendCeiling(CEILING),
|
||||
SpendWindowMs(WINDOW_MS),
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
|
|
@ -176,6 +223,7 @@ pub(crate) fn open_service(store: &TestStore, max_concurrent_requests: usize) ->
|
|||
TransportService::open(
|
||||
&store.path,
|
||||
"fixture-sse",
|
||||
fixture_prices(),
|
||||
[
|
||||
caller(
|
||||
"epiphany-yggdrasil",
|
||||
|
|
@ -201,5 +249,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], request).unwrap()
|
||||
TransportInvocation::new(caller_runtime_id, 2_000, [7; 32], MAX_COST, request).unwrap()
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,7 +1,6 @@
|
|||
mod client;
|
||||
mod contract;
|
||||
mod envelope;
|
||||
mod readout;
|
||||
|
||||
#[cfg(test)]
|
||||
mod fixtures;
|
||||
|
|
@ -14,7 +13,6 @@ mod service;
|
|||
pub use client::*;
|
||||
pub use contract::*;
|
||||
pub use envelope::*;
|
||||
pub use readout::*;
|
||||
|
||||
#[cfg(feature = "daemon")]
|
||||
pub use serve::*;
|
||||
|
|
|
|||
|
|
@ -1,106 +0,0 @@
|
|||
//! 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());
|
||||
}
|
||||
}
|
||||
|
|
@ -12,13 +12,12 @@ use zeroize::Zeroizing;
|
|||
|
||||
use crate::contract::TransportResult;
|
||||
use crate::envelope::{
|
||||
RequestFrame, ServiceError, TransportFrameError, decode_request_frame, encode_result_frame,
|
||||
encode_usage_readout_response_frame, frame_bytes_for_payload, read_transport_frame,
|
||||
write_transport_frame,
|
||||
ServiceError, TransportFrameError, decode_invocation_frame, encode_result_frame,
|
||||
frame_bytes_for_payload, read_transport_frame, write_transport_frame,
|
||||
};
|
||||
use crate::readout::PlanQuotaSnapshot;
|
||||
use crate::service::{
|
||||
CallerAdmission, ProviderCall, TransportAdmission, TransportService, distinct_admissions,
|
||||
CallerAdmission, PriceTable, ProviderCall, SpendCeiling, SpendWindowMs, TransportAdmission,
|
||||
TransportService,
|
||||
};
|
||||
|
||||
/// The only provider seam of the connector kit. A connector binary implements
|
||||
|
|
@ -33,12 +32,6 @@ pub trait ProviderBackend: Send + Sync {
|
|||
/// `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.
|
||||
|
|
@ -87,7 +80,6 @@ impl ServeLimits {
|
|||
if admissions.is_empty() {
|
||||
return Err(ServeError::InvalidConfig("callers"));
|
||||
}
|
||||
distinct_admissions(admissions)?;
|
||||
let slots = admissions
|
||||
.iter()
|
||||
.try_fold(0_usize, |total, caller| {
|
||||
|
|
@ -123,6 +115,12 @@ 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> {
|
||||
|
|
@ -143,6 +141,8 @@ 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)?,
|
||||
);
|
||||
|
|
@ -151,43 +151,19 @@ 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)?
|
||||
|
|
@ -196,7 +172,15 @@ fn admit_listener(listener: &TcpListener, backend: &dyn ProviderBackend) -> Resu
|
|||
{
|
||||
return Err(ServeError::InvalidConfig("listener address"));
|
||||
}
|
||||
backend.ready().map_err(ServeError::BackendUnready)
|
||||
backend.ready().map_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)
|
||||
}
|
||||
|
||||
/// The accept loop over any source of connections, so a test can present a
|
||||
|
|
@ -248,37 +232,21 @@ 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 = answer_frame(&request, service, backend, limits)?;
|
||||
let response = execute_inference_frame(&request, service, backend, limits)?;
|
||||
write_transport_frame(&mut stream, &response, limits.max_frame_bytes).map_err(frame_error)
|
||||
}
|
||||
|
||||
fn answer_frame(
|
||||
fn execute_inference_frame(
|
||||
request: &[u8],
|
||||
service: &Mutex<TransportService>,
|
||||
backend: &dyn ProviderBackend,
|
||||
limits: &ServeLimits,
|
||||
) -> Result<Vec<u8>, ServeError> {
|
||||
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 envelope = decode_invocation_frame(request).map_err(|_| ServeError::FrameEncoding)?;
|
||||
let admission = service
|
||||
.lock()
|
||||
.map_err(|_| ServeError::ServicePoisoned)?
|
||||
.begin(
|
||||
&envelope,
|
||||
unix_ms()?,
|
||||
&backend.plan_quota(),
|
||||
limits.max_frame_bytes,
|
||||
)?;
|
||||
.begin(&envelope, unix_ms()?, limits.max_frame_bytes)?;
|
||||
let response = match admission {
|
||||
TransportAdmission::Reply(response) => response,
|
||||
TransportAdmission::Execute(claim) => {
|
||||
|
|
@ -382,7 +350,10 @@ mod tests {
|
|||
FailureKind, Refusal, TransportDisposition, TransportEvent, TransportEventPayload,
|
||||
TransportInvocation, TransportOutcome, TransportResult,
|
||||
};
|
||||
use crate::fixtures::{TestStore, caller, invocation_for, receipt, request};
|
||||
use crate::fixtures::{
|
||||
CEILING, MAX_COST, TestStore, WINDOW_MS, caller, fixture_prices, invocation_for, receipt,
|
||||
request,
|
||||
};
|
||||
|
||||
const CALLER: &str = "epiphany-yggdrasil";
|
||||
const KEY: &str = "epiphany-distinct-test-key";
|
||||
|
|
@ -458,7 +429,14 @@ 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, &path, backend));
|
||||
let _ = sender.send(serve(
|
||||
listener,
|
||||
limits,
|
||||
admissions,
|
||||
fixture_prices(),
|
||||
&path,
|
||||
backend,
|
||||
));
|
||||
});
|
||||
receiver.recv_timeout(Duration::from_secs(2)).ok()
|
||||
}
|
||||
|
|
@ -484,7 +462,16 @@ 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, &path, backend));
|
||||
thread::spawn(move || {
|
||||
serve(
|
||||
listener,
|
||||
limits(),
|
||||
admissions,
|
||||
fixture_prices(),
|
||||
&path,
|
||||
backend,
|
||||
)
|
||||
});
|
||||
let client =
|
||||
ConnectorClient::new(endpoint, KEY, FRAME, Some(Duration::from_secs(5))).unwrap();
|
||||
(client, endpoint)
|
||||
|
|
@ -502,128 +489,35 @@ 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_reads_out_plan_usage_without_executing_anything() {
|
||||
fn serve_refuses_a_request_over_its_max_cost_without_executing_it() {
|
||||
let store = TestStore::new();
|
||||
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();
|
||||
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();
|
||||
assert_eq!(
|
||||
refused.disposition,
|
||||
TransportDisposition::Refused(Refusal::ProviderRejected)
|
||||
TransportDisposition::Refused(Refusal::CostBound)
|
||||
);
|
||||
assert_eq!(inner.executions.load(Ordering::SeqCst), 0);
|
||||
}
|
||||
assert_eq!(backend.executions.load(Ordering::SeqCst), 0);
|
||||
|
||||
#[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();
|
||||
// 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();
|
||||
assert!(matches!(
|
||||
admitted.disposition,
|
||||
TransportDisposition::Transported { .. }
|
||||
));
|
||||
assert_eq!(inner.executions.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(backend.executions.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
|
@ -662,6 +556,7 @@ mod tests {
|
|||
CALLER,
|
||||
invocation.expires_at_unix_ms,
|
||||
[7; 32],
|
||||
MAX_COST,
|
||||
conflicting_request,
|
||||
)
|
||||
.unwrap();
|
||||
|
|
@ -746,6 +641,7 @@ mod tests {
|
|||
let mut service = TransportService::open(
|
||||
&store.path,
|
||||
"fixture-sse",
|
||||
fixture_prices(),
|
||||
vec![caller(CALLER, KEY, 4)],
|
||||
300_000,
|
||||
)
|
||||
|
|
@ -753,9 +649,8 @@ mod tests {
|
|||
let envelope =
|
||||
crate::envelope::encrypt_invocation(&invocation, &crate::fixtures::security(KEY))
|
||||
.unwrap();
|
||||
let TransportAdmission::Execute(claim) = service
|
||||
.begin(&envelope, now_ms(), &PlanQuotaSnapshot::default(), 1 << 20)
|
||||
.unwrap()
|
||||
let TransportAdmission::Execute(claim) =
|
||||
service.begin(&envelope, now_ms(), 1 << 20).unwrap()
|
||||
else {
|
||||
panic!("a fresh request must execute");
|
||||
};
|
||||
|
|
@ -778,6 +673,7 @@ mod tests {
|
|||
listener,
|
||||
narrowed,
|
||||
vec![caller(CALLER, KEY, 4)],
|
||||
fixture_prices(),
|
||||
&path,
|
||||
served,
|
||||
)
|
||||
|
|
@ -863,7 +759,13 @@ 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,
|
||||
&long,
|
||||
&long,
|
||||
&long,
|
||||
[0xff; 32],
|
||||
[0xff; 32],
|
||||
0,
|
||||
u64::MAX,
|
||||
))
|
||||
.unwrap()
|
||||
.len();
|
||||
|
|
@ -929,26 +831,6 @@ 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])));
|
||||
|
|
@ -971,7 +853,18 @@ 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).unwrap()
|
||||
CallerAdmission::new(
|
||||
id,
|
||||
key,
|
||||
1,
|
||||
["gpt-5.4".to_string()],
|
||||
slots,
|
||||
payload,
|
||||
32_768,
|
||||
SpendCeiling(CEILING),
|
||||
SpendWindowMs(WINDOW_MS),
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
/// `Some(field)` when serve refused the configuration, `None` while it serves.
|
||||
|
|
@ -1077,7 +970,7 @@ mod tests {
|
|||
..limits()
|
||||
};
|
||||
let admissions = vec![admission(CALLER, KEY, 4, bound)];
|
||||
thread::spawn(move || serve(listener, tight, admissions, &path, served));
|
||||
thread::spawn(move || serve(listener, tight, admissions, fixture_prices(), &path, served));
|
||||
let client =
|
||||
ConnectorClient::new(endpoint, KEY, frame, Some(Duration::from_secs(5))).unwrap();
|
||||
let reply = client.execute(&invocation).unwrap();
|
||||
|
|
@ -1192,6 +1085,7 @@ mod tests {
|
|||
let mut service = TransportService::open(
|
||||
&store.path,
|
||||
"fixture-sse",
|
||||
fixture_prices(),
|
||||
vec![caller(CALLER, KEY, 4)],
|
||||
300_000,
|
||||
)
|
||||
|
|
@ -1201,9 +1095,8 @@ mod tests {
|
|||
&crate::fixtures::security(KEY),
|
||||
)
|
||||
.unwrap();
|
||||
let TransportAdmission::Execute(claim) = service
|
||||
.begin(&envelope, now_ms(), &PlanQuotaSnapshot::default(), 1 << 20)
|
||||
.unwrap()
|
||||
let TransportAdmission::Execute(claim) =
|
||||
service.begin(&envelope, now_ms(), 1 << 20).unwrap()
|
||||
else {
|
||||
panic!("a fresh request must execute");
|
||||
};
|
||||
|
|
@ -1295,6 +1188,7 @@ mod tests {
|
|||
TransportService::open(
|
||||
&store.path,
|
||||
"fixture-sse",
|
||||
fixture_prices(),
|
||||
vec![caller(CALLER, KEY, 4)],
|
||||
300_000,
|
||||
)
|
||||
|
|
@ -1352,7 +1246,16 @@ mod tests {
|
|||
let backend = CountingBackend::new(true);
|
||||
let path = store.path.clone();
|
||||
let served = backend.clone();
|
||||
thread::spawn(move || serve(listener, limits(), admissions, &path, served));
|
||||
thread::spawn(move || {
|
||||
serve(
|
||||
listener,
|
||||
limits(),
|
||||
admissions,
|
||||
fixture_prices(),
|
||||
&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
|
|
@ -1,429 +0,0 @@
|
|||
//! 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(),
|
||||
));
|
||||
}
|
||||
}
|
||||
|
|
@ -1,6 +1,6 @@
|
|||
# wire_v4 fixtures
|
||||
# wire_v3 fixtures
|
||||
|
||||
The named-MessagePack bytes of the v4 documents and one sealed carriage frame,
|
||||
The named-MessagePack bytes of the v3 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_v4.rs`.
|
||||
sealed under the fixed test key in `tests/wire_golden_v3.rs`.
|
||||
|
||||
Generator: `tests/wire_golden_v4.rs` itself. Run it with `WIRE_V4_REGENERATE=1` to rewrite
|
||||
Generator: `tests/wire_golden_v3.rs` itself. Run it with `WIRE_V3_REGENERATE=1` to rewrite
|
||||
the files from the kit, then review the diff: a changed byte is a wire change.
|
||||
1
crates/model-connector/tests/fixtures/wire_v3/carriage_invocation.frame
vendored
Normal file
1
crates/model-connector/tests/fixtures/wire_v3/carriage_invocation.frame
vendored
Normal file
|
|
@ -0,0 +1 @@
|
|||
<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
|
||||
Binary file not shown.
Binary file not shown.
BIN
crates/model-connector/tests/fixtures/wire_v3/transport_invocation.mp
vendored
Normal file
BIN
crates/model-connector/tests/fixtures/wire_v3/transport_invocation.mp
vendored
Normal file
Binary file not shown.
BIN
crates/model-connector/tests/fixtures/wire_v3/transport_receipt_failed.mp
vendored
Normal file
BIN
crates/model-connector/tests/fixtures/wire_v3/transport_receipt_failed.mp
vendored
Normal file
Binary file not shown.
|
|
@ -1 +1 @@
|
|||
„©schema_idÙ"gamecult.model.transport_result.v4ªrequest_id®golden-request±caller_runtime_idgolden-caller«disposition<6F>§Refused«OutputBound
|
||||
„©schema_idÙ"gamecult.model.transport_result.v3ªrequest_id®golden-request±caller_runtime_idgolden-caller«disposition<6F>§Refused«OutputBound
|
||||
BIN
crates/model-connector/tests/fixtures/wire_v3/transport_result_transported.mp
vendored
Normal file
BIN
crates/model-connector/tests/fixtures/wire_v3/transport_result_transported.mp
vendored
Normal file
Binary file not shown.
|
|
@ -1 +0,0 @@
|
|||
<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
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
|
|
@ -1,4 +1,4 @@
|
|||
//! `schemas/` at the repository root publishes one JSON Schema per v4 document,
|
||||
//! `schemas/` at the repository root publishes one JSON Schema per v3 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,7 +28,6 @@ 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() {
|
||||
|
|
@ -54,7 +53,7 @@ fn published_schemas_are_the_types() {
|
|||
expected.sort();
|
||||
assert_eq!(
|
||||
on_disk, expected,
|
||||
"schemas/ holds exactly the six documents"
|
||||
"schemas/ holds exactly the five documents"
|
||||
);
|
||||
}
|
||||
|
||||
|
|
@ -80,7 +79,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.v4".to_string(),
|
||||
schema_id: "gamecult.model.transport_envelope.v3".to_string(),
|
||||
caller_runtime_id: "caller".to_string(),
|
||||
request_id: "request".to_string(),
|
||||
message_kind: EnvelopeKind::Invocation,
|
||||
|
|
@ -111,7 +110,9 @@ fn the_refusal_schema_is_one_flat_string_enum() {
|
|||
.iter()
|
||||
.map(|value| value.as_str().unwrap())
|
||||
.collect();
|
||||
assert_eq!(variants.len(), 13);
|
||||
assert_eq!(variants.len(), 15);
|
||||
assert!(variants.contains(&"ClockBehindAdmissionWindow"));
|
||||
assert!(variants.contains(&"ForeignReasoningState"));
|
||||
assert!(variants.contains(&"ProviderRejected"));
|
||||
assert!(variants.contains(&"CostBound"));
|
||||
assert!(variants.contains(&"SpendCeiling"));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -106,6 +106,8 @@ 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 {
|
||||
|
|
@ -114,7 +116,21 @@ fn start_with(backend: Arc<dyn ProviderBackend>, max_concurrent: usize) -> Conne
|
|||
socket_timeout_ms: 5_000,
|
||||
max_expiry_skew_ms: 300_000,
|
||||
};
|
||||
thread::spawn(move || serve(listener, limits, vec![admission], &store, backend));
|
||||
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));
|
||||
ConnectorClient::new(endpoint, KEY, FRAME, Some(Duration::from_secs(5))).unwrap()
|
||||
}
|
||||
|
||||
|
|
@ -127,7 +143,7 @@ fn invocation(request_id: &str, input: Vec<InputItem>) -> TransportInvocation {
|
|||
4_096,
|
||||
);
|
||||
request.input = input;
|
||||
TransportInvocation::new(CALLER, now_ms() + 60_000, [7; 32], request).unwrap()
|
||||
TransportInvocation::new(CALLER, now_ms() + 60_000, [7; 32], 10_000_000, request).unwrap()
|
||||
}
|
||||
|
||||
fn user(text: &str) -> InputItem {
|
||||
|
|
@ -316,7 +332,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.request_id = "another-request".to_string(),
|
||||
|receipt| receipt.cost_micro_usd += 1,
|
||||
];
|
||||
for edit in edits {
|
||||
// One concurrent request: a wedged slot would refuse the next id.
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
//! 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
|
||||
//! 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
|
||||
//! 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_V4_REGENERATE";
|
||||
const REGENERATE: &str = "WIRE_V3_REGENERATE";
|
||||
|
||||
fn fixtures() -> PathBuf {
|
||||
Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/wire_v4")
|
||||
Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/wire_v3")
|
||||
}
|
||||
|
||||
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_V4_REGENERATE` set, the committed bytes are rewritten
|
||||
/// back to it. With `WIRE_V3_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], request()).unwrap()
|
||||
TransportInvocation::new(CALLER, 1_800_000_000_000, [7; 32], 5_000_000, request()).unwrap()
|
||||
}
|
||||
|
||||
fn receipt(outcome: TransportOutcome) -> TransportReceipt {
|
||||
|
|
@ -95,6 +95,8 @@ 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,
|
||||
}
|
||||
}
|
||||
|
|
@ -114,7 +116,7 @@ fn completed() -> TransportOutcome {
|
|||
}
|
||||
|
||||
#[test]
|
||||
fn wire_golden_v4() {
|
||||
fn wire_golden_v3() {
|
||||
let invocation = invocation();
|
||||
assert_wire("provider_request.mp", &invocation.request);
|
||||
assert_wire("transport_invocation.mp", &invocation);
|
||||
|
|
@ -231,7 +233,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.v4");
|
||||
assert_eq!(payload_schema, "gamecult.model.transport_envelope.v3");
|
||||
assert_eq!(message_id, invocation.request_id());
|
||||
let envelope: TransportEnvelope =
|
||||
rmp_serde::from_slice(&STANDARD.decode(payload).unwrap()).unwrap();
|
||||
|
|
@ -46,9 +46,7 @@
|
|||
"type": "string",
|
||||
"enum": [
|
||||
"Invocation",
|
||||
"Result",
|
||||
"UsageReadoutRequest",
|
||||
"UsageReadoutResult"
|
||||
"Result"
|
||||
]
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -11,6 +11,12 @@
|
|||
"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": {
|
||||
|
|
@ -47,6 +53,7 @@
|
|||
"expires_at_unix_ms",
|
||||
"native_request_sha256",
|
||||
"provider_request_sha256",
|
||||
"max_cost_micro_usd",
|
||||
"request"
|
||||
],
|
||||
"$defs": {
|
||||
|
|
|
|||
|
|
@ -6,6 +6,12 @@
|
|||
"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",
|
||||
|
|
@ -29,6 +35,12 @@
|
|||
"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"
|
||||
},
|
||||
|
|
@ -60,6 +72,8 @@
|
|||
"model",
|
||||
"provider",
|
||||
"foreign_reasoning_dropped",
|
||||
"cost_micro_usd",
|
||||
"overrun_micro_usd",
|
||||
"outcome"
|
||||
],
|
||||
"$defs": {
|
||||
|
|
@ -160,7 +174,7 @@
|
|||
]
|
||||
},
|
||||
"Usage": {
|
||||
"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.",
|
||||
"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.",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"cache_read_tokens": {
|
||||
|
|
|
|||
|
|
@ -46,13 +46,15 @@
|
|||
"ModelNotAllowed",
|
||||
"OutputBound",
|
||||
"PayloadBound",
|
||||
"CostBound",
|
||||
"SpendCeiling",
|
||||
"Capacity",
|
||||
"InFlight",
|
||||
"Indeterminate",
|
||||
"ReplayConflict",
|
||||
"Malformed",
|
||||
"ForeignReasoningState",
|
||||
"ProviderRejected"
|
||||
"ClockBehindAdmissionWindow"
|
||||
]
|
||||
},
|
||||
"StopReason": {
|
||||
|
|
@ -284,6 +286,12 @@
|
|||
"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",
|
||||
|
|
@ -307,6 +315,12 @@
|
|||
"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"
|
||||
},
|
||||
|
|
@ -338,11 +352,13 @@
|
|||
"model",
|
||||
"provider",
|
||||
"foreign_reasoning_dropped",
|
||||
"cost_micro_usd",
|
||||
"overrun_micro_usd",
|
||||
"outcome"
|
||||
]
|
||||
},
|
||||
"Usage": {
|
||||
"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.",
|
||||
"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.",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"cache_read_tokens": {
|
||||
|
|
|
|||
|
|
@ -1,68 +0,0 @@
|
|||
{
|
||||
"$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"
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue