Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
3f87b19a2d | ||
|
|
f52eef1657 | ||
|
|
e80e522c10 | ||
| 9d0f4b400c | |||
| 2a65c4e7a2 | |||
| 42a157b049 | |||
| c47597b421 | |||
| 61cc59d9a0 | |||
|
|
0282244114 | ||
|
|
51982b7f63 | ||
|
|
254cf83d9d | ||
|
|
8331859537 | ||
|
|
6e87cf2e6c | ||
|
|
549962cb8a | ||
|
|
7d4c1bf56e | ||
|
|
8fb8b5a71d | ||
|
|
f242290ca5 | ||
|
|
481428a7ce | ||
|
|
e17ef16d48 | ||
|
|
c1cc349b19 | ||
|
|
6982cbca7e | ||
|
|
7f3a46313e | ||
|
|
9143d2c5e2 | ||
|
|
79792eea85 | ||
|
|
c1802d85be | ||
|
|
cfea6854ab | ||
|
|
9f10887169 | ||
|
|
9db62027eb | ||
|
|
7d9a6392a7 | ||
|
|
a828602649 | ||
|
|
bb8909f2c2 | ||
|
|
28f2e5d2e7 | ||
|
|
dcb5f109f7 | ||
|
|
af5da3dfaa | ||
|
|
702e903fdd | ||
|
|
06a47bb13f | ||
|
|
d2b82a6c50 | ||
|
|
47a4dcc6b4 | ||
|
|
1bf9fdb437 | ||
|
|
13cbdaf23c | ||
|
|
fdcfa063c7 | ||
|
|
d1f25b5ae6 | ||
| 480cf941c8 | |||
| e867fad379 | |||
| 200933f7cf | |||
|
|
e0c0832ae0 | ||
|
|
1d51e71ebf | ||
|
|
bf958fc327 | ||
|
|
ec0519d9ff |
63 changed files with 21374 additions and 15 deletions
2
.gitattributes
vendored
Normal file
2
.gitattributes
vendored
Normal file
|
|
@ -0,0 +1,2 @@
|
|||
schemas/*.json text eol=lf
|
||||
crates/model-connector/tests/fixtures/wire_v4/* -text
|
||||
1
.gitignore
vendored
Normal file
1
.gitignore
vendored
Normal file
|
|
@ -0,0 +1 @@
|
|||
/target
|
||||
1274
Cargo.lock
generated
Normal file
1274
Cargo.lock
generated
Normal file
File diff suppressed because it is too large
Load diff
9
Cargo.toml
Normal file
9
Cargo.toml
Normal file
|
|
@ -0,0 +1,9 @@
|
|||
[workspace]
|
||||
members = ["crates/*"]
|
||||
resolver = "3"
|
||||
|
||||
[workspace.package]
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
license = "AGPL-3.0-only"
|
||||
repository = "https://eureka.gamecult.org/forge/GameCult/ModelConnector"
|
||||
111
README.md
111
README.md
|
|
@ -1,26 +1,107 @@
|
|||
# ModelConnector
|
||||
|
||||
Status: planned. Nothing is implemented in this repository yet.
|
||||
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.
|
||||
|
||||
## What it will hold
|
||||
The connector meters nothing in dollars. It holds the subscription token, lends it, and
|
||||
publishes the plan's usage windows (`usage_readout`, `claude_usage_readout` on the lease socket, and `plan_quota` on every lease grant) for the
|
||||
minds to read and decide their own burn. The one usage-derived refusal is the provider's own
|
||||
(`ProviderRejected`), repeated until its reset.
|
||||
|
||||
The connector kit for GameCult minds:
|
||||
Status: the `model-connector` crate and the `claude-connector` backend crate and binary are live
|
||||
(`claude-connector serve|config|ask`); an OpenAI-compatible connector is not written.
|
||||
|
||||
- the connector contract, a Rust crate defining the typed boundary every model
|
||||
turn of a GameCult mind goes through;
|
||||
- connector daemons that implement the contract against a model provider,
|
||||
Claude first, then an OpenAI-compatible connector.
|
||||
## Entry points
|
||||
|
||||
The problem it solves: a mind should reach any model through one inspectable,
|
||||
typed contract, rather than each project wiring its own provider calls.
|
||||
Crate `crates/model-connector`:
|
||||
|
||||
## Where it comes from
|
||||
- `src/contract.rs` - the v4 wire contract (`gamecult.model.*.v4`): `ProviderRequest`, `TransportInvocation`,
|
||||
`TransportResult`, `TransportReceipt`, the typed `StopReason`, `FailureKind` and `Refusal`,
|
||||
`InputItem::ReasoningState` (a provider's opaque reasoning block, replayed by callers
|
||||
verbatim via `TransportResult::transcript_items`; the service refuses another provider's
|
||||
state in the in-flight tool exchange and drops it from earlier turns),
|
||||
and the canonical byte and digest functions that bind a result to its request. Every
|
||||
document is named MessagePack with unknown fields refused.
|
||||
- `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.
|
||||
- `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),
|
||||
in-flight and capacity refusals.
|
||||
- `src/serve.rs` (feature `daemon`) - `serve()`, the loopback accept loop, over the
|
||||
`ProviderBackend` trait. A connector binary implements `ProviderBackend` for its
|
||||
provider and calls `serve`; the kit names no provider.
|
||||
|
||||
[CodexConnector](https://github.com/GameCult/CodexConnector)
|
||||
(GameCult/CodexConnector) is the existing connector this kit is lifted from.
|
||||
The lift is the first piece of work here; until it lands there are no entry
|
||||
points to run.
|
||||
Crate `crates/claude-connector` (library and the `claude-connector` binary that wires it to `serve`):
|
||||
|
||||
## License
|
||||
- `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.
|
||||
- `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.
|
||||
|
||||
## Asking for one no-tools completion
|
||||
|
||||
`claude-connector ask` is a caller of a running connector for a Mind's Self run: the prompt is all
|
||||
of stdin, the answer text is stdout.
|
||||
|
||||
```
|
||||
claude-connector ask --caller-id ID --connection-key-file FILE --model MODEL (--system TEXT | --system-file FILE) [--endpoint HOST:PORT] [--effort low|medium|high|max] [--max-output-tokens N] [--timeout-secs N] < prompt > answer
|
||||
```
|
||||
|
||||
Defaults: endpoint `127.0.0.1:4110`, 8192 output tokens, 300 s. The key is read from the file (the
|
||||
same text the connector reads for the caller; never argv or environment). There is no tools flag,
|
||||
no working-directory flag and no other endpoint: the request is built by `ProviderRequest::new`
|
||||
(no tools) and the model child's cwd and environment stay the connector's. Any failure (refused,
|
||||
over a cap, model not allowed, wrong key or caller, timeout, provider failure, empty answer, empty
|
||||
stdin) exits nonzero with one `claude-connector: <cause>` line on stderr and nothing on stdout;
|
||||
the line never carries the key, the prompt or the provider's message. Each run uses a fresh request
|
||||
id. Code: `crates/claude-connector/src/ask.rs`; tests: `crates/claude-connector/tests/ask.rs`.
|
||||
|
||||
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
|
||||
a frame allows, and `MAX_IDENTITY_BYTES` (256) bounds caller ids, request ids and model names.
|
||||
`serve` refuses a frame/payload pair that violates them (`ServeError::FrameBound` states the frame
|
||||
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
|
||||
only: a v2 frame gets no answer.
|
||||
|
||||
## Verify
|
||||
|
||||
```
|
||||
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/published_schemas.rs` fails when `schemas/` differs from the types
|
||||
(`MODEL_CONNECTOR_REGENERATE_SCHEMAS=1` rewrites it).
|
||||
|
||||
## Provenance and license
|
||||
|
||||
AGPL-3.0-only. See `LICENSE`.
|
||||
|
||||
The contract, envelope, client, replay service and serve loop were lifted from
|
||||
[CodexConnector](https://github.com/GameCult/CodexConnector) at fd7cec4, formerly
|
||||
PolyForm Noncommercial. Its sole copyright holder relicensed the lifted material as
|
||||
AGPL-3.0-only for this repository.
|
||||
|
|
|
|||
32
crates/claude-connector/Cargo.toml
Normal file
32
crates/claude-connector/Cargo.toml
Normal file
|
|
@ -0,0 +1,32 @@
|
|||
[package]
|
||||
name = "claude-connector"
|
||||
description = "Claude backend for the connector kit: drives the unmodified claude binary in print stream-json mode"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
publish = false
|
||||
|
||||
[[bin]]
|
||||
name = "claude-connector"
|
||||
path = "src/main.rs"
|
||||
|
||||
[dependencies]
|
||||
anyhow = "1"
|
||||
chrono = { version = "0.4", default-features = false }
|
||||
hmac = "0.12"
|
||||
cultcache-rs = { git = "https://github.com/GameCult/CultLib.git", rev = "a4eae9056152335662f1115f97bf35fb00ab6d23", package = "cultcache-rs" }
|
||||
model-connector = { path = "../model-connector", features = ["daemon"] }
|
||||
rand_core = { version = "0.6", features = ["getrandom"] }
|
||||
rmp-serde = "1"
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
sha2 = "0.10"
|
||||
thiserror = "2"
|
||||
zeroize = "1"
|
||||
|
||||
[target.'cfg(unix)'.dependencies]
|
||||
libc = "0.2"
|
||||
|
||||
[dev-dependencies]
|
||||
rmpv = "1"
|
||||
336
crates/claude-connector/src/ask.rs
Normal file
336
crates/claude-connector/src/ask.rs
Normal file
|
|
@ -0,0 +1,336 @@
|
|||
//! `claude-connector ask`: one no-tools completion through a running connector.
|
||||
//!
|
||||
//! A thin caller over [`ConnectorClient`]. It builds a [`ProviderRequest`] with
|
||||
//! `ProviderRequest::new`, which carries no tools, and has no field, flag or
|
||||
//! path by which a caller could add one; the working directory and environment
|
||||
//! of the model child belong to the connector. Every failure is an
|
||||
//! [`AskError`] whose text names a cause and never echoes the key, the prompt
|
||||
//! or the instructions.
|
||||
|
||||
use std::net::SocketAddr;
|
||||
use std::path::Path;
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use model_connector::{
|
||||
ConnectorClient, ConnectorClientError, FailureKind, InputItem, ProviderRequest,
|
||||
ReasoningEffort, StopReason, TransportDisposition, TransportEventPayload, TransportInvocation,
|
||||
TransportOutcome, TransportResult, frame_bytes_for_payload, provider_request_sha256,
|
||||
};
|
||||
use rand_core::{OsRng, RngCore};
|
||||
use zeroize::Zeroizing;
|
||||
|
||||
/// The loopback endpoint a claude connector listens on.
|
||||
pub const DEFAULT_ENDPOINT: &str = "127.0.0.1:4110";
|
||||
/// The payload bound the connector's own `config init` defaults to.
|
||||
pub const PAYLOAD_BYTES: usize = 1 << 20;
|
||||
const MAX_KEY_BYTES: usize = 1024;
|
||||
pub const MAX_SYSTEM_BYTES: usize = 64 * 1024;
|
||||
|
||||
pub struct AskArgs {
|
||||
pub endpoint: SocketAddr,
|
||||
pub caller_id: String,
|
||||
pub key_file: std::path::PathBuf,
|
||||
pub model: String,
|
||||
pub instructions: String,
|
||||
pub effort: Option<ReasoningEffort>,
|
||||
pub max_output_tokens: u32,
|
||||
pub timeout: Duration,
|
||||
pub prompt: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum AskError {
|
||||
#[error("the prompt on stdin is empty")]
|
||||
EmptyPrompt,
|
||||
#[error("the prompt on stdin is larger than the {PAYLOAD_BYTES} byte payload bound")]
|
||||
PromptTooLarge,
|
||||
#[error("the instructions are empty")]
|
||||
EmptyInstructions,
|
||||
#[error("the connection key file cannot be read")]
|
||||
KeyUnreadable,
|
||||
#[error("the connection key file is not text")]
|
||||
KeyEncoding,
|
||||
#[error("the connection key file is not a regular file")]
|
||||
KeyNotRegular,
|
||||
#[error("the connection key file is readable by group or world (chmod 600)")]
|
||||
KeyExposed,
|
||||
#[error("the instructions file cannot be read as text from a regular file")]
|
||||
SystemUnreadable,
|
||||
#[error("the instructions file is larger than {MAX_SYSTEM_BYTES} bytes")]
|
||||
SystemTooLarge,
|
||||
#[error("standard input cannot be read")]
|
||||
StdinRead,
|
||||
#[error("standard input is not UTF-8 text")]
|
||||
StdinNotText,
|
||||
#[error("the answer cannot be written to standard output")]
|
||||
StdoutWrite,
|
||||
#[error("nothing is listening at the connector endpoint")]
|
||||
NotListening,
|
||||
#[error(
|
||||
"the connector closed the connection without answering (key, caller or frame not accepted)"
|
||||
)]
|
||||
ClosedWithoutAnswer,
|
||||
#[error("the connection key file is empty or longer than {MAX_KEY_BYTES} bytes")]
|
||||
KeyLength,
|
||||
#[error("the connection key is not usable as a key (surrounding whitespace?)")]
|
||||
KeyUnusable,
|
||||
#[error("the request is invalid for the connector contract")]
|
||||
Request,
|
||||
#[error("timed out waiting for the connector")]
|
||||
Timeout,
|
||||
#[error("the connector connection failed: {0}")]
|
||||
Connection(std::io::ErrorKind),
|
||||
#[error("the connector answer was not accepted (wrong key, caller or tampered frame): {0}")]
|
||||
Transport(model_connector::ServiceError),
|
||||
#[error("the request exceeded a frame bound")]
|
||||
FrameSize,
|
||||
#[error("the connector refused the request: {0:?}")]
|
||||
Refused(model_connector::Refusal),
|
||||
#[error("the provider failed: {0:?}")]
|
||||
Failed(FailureKind),
|
||||
#[error("the answer stopped as {0:?}, not a completed text")]
|
||||
NotCompleted(StopReason),
|
||||
#[error("the answer carried a tool call, which this command never requests")]
|
||||
ToolCall,
|
||||
#[error("the answer is empty")]
|
||||
EmptyAnswer,
|
||||
}
|
||||
|
||||
enum ReadFailure {
|
||||
Unreadable,
|
||||
NotRegular,
|
||||
TooLarge,
|
||||
Exposed,
|
||||
}
|
||||
|
||||
/// Opens `path` without blocking (a FIFO must not hang the open), requires the
|
||||
/// opened file to be a regular file (a symlink is followed and its target
|
||||
/// judged), optionally refuses one readable by group or world, and reads at
|
||||
/// most `cap` bytes, failing if there are more.
|
||||
fn read_regular(path: &Path, cap: usize, private: bool) -> Result<Zeroizing<Vec<u8>>, ReadFailure> {
|
||||
use std::io::Read;
|
||||
let mut options = std::fs::OpenOptions::new();
|
||||
options.read(true);
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::OpenOptionsExt;
|
||||
options.custom_flags(libc::O_NONBLOCK);
|
||||
}
|
||||
let file = options.open(path).map_err(|_| ReadFailure::Unreadable)?;
|
||||
let metadata = file.metadata().map_err(|_| ReadFailure::Unreadable)?;
|
||||
if !metadata.is_file() {
|
||||
return Err(ReadFailure::NotRegular);
|
||||
}
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
if private && metadata.permissions().mode() & 0o077 != 0 {
|
||||
return Err(ReadFailure::Exposed);
|
||||
}
|
||||
}
|
||||
let mut bytes = Zeroizing::new(Vec::new());
|
||||
file.take(cap as u64 + 1)
|
||||
.read_to_end(&mut bytes)
|
||||
.map_err(|_| ReadFailure::Unreadable)?;
|
||||
if bytes.len() > cap {
|
||||
return Err(ReadFailure::TooLarge);
|
||||
}
|
||||
Ok(bytes)
|
||||
}
|
||||
|
||||
/// Reads the key as the connector reads its own key file (the text without the
|
||||
/// trailing line ending). The bytes are zeroized on drop and never shown.
|
||||
pub fn read_key(path: &Path) -> Result<Zeroizing<String>, AskError> {
|
||||
let raw = read_regular(path, MAX_KEY_BYTES, true).map_err(|failure| match failure {
|
||||
ReadFailure::Unreadable => AskError::KeyUnreadable,
|
||||
ReadFailure::NotRegular => AskError::KeyNotRegular,
|
||||
ReadFailure::TooLarge => AskError::KeyLength,
|
||||
ReadFailure::Exposed => AskError::KeyExposed,
|
||||
})?;
|
||||
let text = std::str::from_utf8(raw.as_slice()).map_err(|_| AskError::KeyEncoding)?;
|
||||
let key = Zeroizing::new(text.trim_end_matches(['\r', '\n']).to_string());
|
||||
if key.is_empty() || key.len() > MAX_KEY_BYTES {
|
||||
return Err(AskError::KeyLength);
|
||||
}
|
||||
if key.trim() != key.as_str() {
|
||||
return Err(AskError::KeyUnusable);
|
||||
}
|
||||
Ok(key)
|
||||
}
|
||||
|
||||
/// The instructions file: a regular file of at most [`MAX_SYSTEM_BYTES`] of text.
|
||||
pub fn read_system_file(path: &Path) -> Result<String, AskError> {
|
||||
let raw = read_regular(path, MAX_SYSTEM_BYTES, false).map_err(|failure| match failure {
|
||||
ReadFailure::TooLarge => AskError::SystemTooLarge,
|
||||
_ => AskError::SystemUnreadable,
|
||||
})?;
|
||||
String::from_utf8(raw.to_vec()).map_err(|_| AskError::SystemUnreadable)
|
||||
}
|
||||
|
||||
fn fresh_id(prefix: &str) -> String {
|
||||
let mut nonce = [0_u8; 16];
|
||||
OsRng.fill_bytes(&mut nonce);
|
||||
let hex: String = nonce.iter().map(|b| format!("{b:02x}")).collect();
|
||||
format!("{prefix}-{hex}")
|
||||
}
|
||||
|
||||
/// The one invocation `ask` ever sends: instructions, one user text, no tools.
|
||||
pub fn build_invocation(args: &AskArgs) -> Result<TransportInvocation, AskError> {
|
||||
if args.prompt.trim().is_empty() {
|
||||
return Err(AskError::EmptyPrompt);
|
||||
}
|
||||
if args.prompt.len() > PAYLOAD_BYTES {
|
||||
return Err(AskError::PromptTooLarge);
|
||||
}
|
||||
if args.instructions.trim().is_empty() {
|
||||
return Err(AskError::EmptyInstructions);
|
||||
}
|
||||
let mut request = ProviderRequest::new(
|
||||
fresh_id("ask"),
|
||||
fresh_id("ask-conversation"),
|
||||
args.model.clone(),
|
||||
args.instructions.clone(),
|
||||
args.max_output_tokens,
|
||||
);
|
||||
request.input = vec![InputItem::UserText {
|
||||
text: args.prompt.clone(),
|
||||
}];
|
||||
request.reasoning_effort = args.effort;
|
||||
let digest = provider_request_sha256(&request).map_err(|_| AskError::Request)?;
|
||||
let now = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map_or(0, |d| d.as_millis() as u64);
|
||||
let invocation =
|
||||
TransportInvocation::new(args.caller_id.clone(), now + 120_000, digest, request)
|
||||
.map_err(|_| AskError::Request)?;
|
||||
invocation
|
||||
.validate(now, 300_000)
|
||||
.map_err(|_| AskError::Request)?;
|
||||
Ok(invocation)
|
||||
}
|
||||
|
||||
/// The answer text of a result, or the one cause it is not an answer.
|
||||
pub fn answer_text(result: &TransportResult) -> Result<String, AskError> {
|
||||
let (events, receipt) = match &result.disposition {
|
||||
TransportDisposition::Refused(reason) => return Err(AskError::Refused(reason.clone())),
|
||||
TransportDisposition::Transported { events, receipt } => (events, receipt),
|
||||
};
|
||||
match &receipt.outcome {
|
||||
TransportOutcome::Failed { kind, .. } => return Err(AskError::Failed(*kind)),
|
||||
TransportOutcome::Completed { stop_reason, .. } => {
|
||||
if !matches!(stop_reason, StopReason::EndTurn | StopReason::StopSequence) {
|
||||
return Err(AskError::NotCompleted(*stop_reason));
|
||||
}
|
||||
}
|
||||
}
|
||||
let mut text = String::new();
|
||||
for event in events {
|
||||
match &event.payload {
|
||||
TransportEventPayload::TextDelta { text: delta } => text.push_str(delta),
|
||||
TransportEventPayload::ToolCall { .. } => return Err(AskError::ToolCall),
|
||||
TransportEventPayload::ReasoningState { .. } => {}
|
||||
}
|
||||
}
|
||||
if text.trim().is_empty() {
|
||||
return Err(AskError::EmptyAnswer);
|
||||
}
|
||||
Ok(text)
|
||||
}
|
||||
|
||||
pub fn ask(args: &AskArgs) -> Result<String, AskError> {
|
||||
let invocation = build_invocation(args)?;
|
||||
let key = read_key(&args.key_file)?;
|
||||
let client = ConnectorClient::new(
|
||||
args.endpoint,
|
||||
key.as_str(),
|
||||
frame_bytes_for_payload(PAYLOAD_BYTES),
|
||||
Some(args.timeout),
|
||||
)
|
||||
.map_err(|error| match error {
|
||||
ConnectorClientError::Transport(_) => AskError::KeyUnusable,
|
||||
_ => AskError::Request,
|
||||
})?;
|
||||
let result = client.execute(&invocation).map_err(|error| match error {
|
||||
ConnectorClientError::Connection(io)
|
||||
if matches!(
|
||||
io.kind(),
|
||||
std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut
|
||||
) =>
|
||||
{
|
||||
AskError::Timeout
|
||||
}
|
||||
ConnectorClientError::Connection(io) => match io.kind() {
|
||||
std::io::ErrorKind::ConnectionRefused => AskError::NotListening,
|
||||
std::io::ErrorKind::UnexpectedEof
|
||||
| std::io::ErrorKind::ConnectionReset
|
||||
| std::io::ErrorKind::BrokenPipe => AskError::ClosedWithoutAnswer,
|
||||
kind => AskError::Connection(kind),
|
||||
},
|
||||
ConnectorClientError::Transport(error) => AskError::Transport(error),
|
||||
ConnectorClientError::FrameSize => AskError::FrameSize,
|
||||
ConnectorClientError::InvalidConfig | ConnectorClientError::RequestIdTooLong => {
|
||||
AskError::Request
|
||||
}
|
||||
})?;
|
||||
answer_text(&result)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::config::{ClaudeCodeConfig, ClaudeCodeLimits};
|
||||
use crate::handle::{SessionStore, plan};
|
||||
use model_connector::ToolDefinition;
|
||||
|
||||
struct NoSessions;
|
||||
impl SessionStore for NoSessions {
|
||||
fn exists(&self, _: &str) -> bool {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
fn args() -> AskArgs {
|
||||
AskArgs {
|
||||
endpoint: DEFAULT_ENDPOINT.parse().unwrap(),
|
||||
caller_id: "c".into(),
|
||||
key_file: "/nonexistent".into(),
|
||||
model: "m".into(),
|
||||
instructions: "i".into(),
|
||||
effort: None,
|
||||
max_output_tokens: 16,
|
||||
timeout: Duration::from_secs(1),
|
||||
prompt: "p".into(),
|
||||
}
|
||||
}
|
||||
|
||||
/// The claude backend's own planner accepts what `ask` builds and refuses
|
||||
/// the same request once a tool is forced into it.
|
||||
#[test]
|
||||
fn the_backend_planner_refuses_a_forced_tool_and_accepts_the_asked_request() {
|
||||
let config = ClaudeCodeConfig {
|
||||
claude_binary: "/claude".into(),
|
||||
claude_binary_sha256: [0; 32],
|
||||
work_dir: "/w".into(),
|
||||
config_dir: "/c".into(),
|
||||
prompt_dir: "/p".into(),
|
||||
token_credential: "t".into(),
|
||||
handle_secret_credential: "h".into(),
|
||||
instance: "i".into(),
|
||||
limits: ClaudeCodeLimits {
|
||||
deadline: Duration::from_secs(1),
|
||||
max_line_bytes: 1024,
|
||||
max_total_bytes: 4096,
|
||||
stderr_tail_bytes: 64,
|
||||
},
|
||||
};
|
||||
let mut request = build_invocation(&args()).unwrap().request;
|
||||
let secret = [7_u8; 32];
|
||||
assert!(plan(&request, &config, &secret, "c", &NoSessions).is_ok());
|
||||
request.tools.push(ToolDefinition {
|
||||
name: "t".into(),
|
||||
description: "d".into(),
|
||||
parameters_json: "{}".into(),
|
||||
});
|
||||
assert!(plan(&request, &config, &secret, "c", &NoSessions).is_err());
|
||||
}
|
||||
}
|
||||
474
crates/claude-connector/src/backend.rs
Normal file
474
crates/claude-connector/src/backend.rs
Normal file
|
|
@ -0,0 +1,474 @@
|
|||
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,
|
||||
};
|
||||
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::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
|
||||
/// `$CREDENTIALS_DIRECTORY/<name>`, read on every call so a rotated credential
|
||||
/// is picked up by the next spawn.
|
||||
pub struct CredentialTokenSource {
|
||||
directory: Option<PathBuf>,
|
||||
name: String,
|
||||
/// The configuration field `missing` names when the credential is absent.
|
||||
field: &'static str,
|
||||
}
|
||||
|
||||
impl CredentialTokenSource {
|
||||
/// The credential `name` under the directory systemd names in
|
||||
/// `CREDENTIALS_DIRECTORY`; with no such variable the source has no token.
|
||||
/// The variable names a directory, never a token.
|
||||
pub fn from_environment(name: impl Into<String>) -> Self {
|
||||
Self {
|
||||
directory: std::env::var_os("CREDENTIALS_DIRECTORY").map(PathBuf::from),
|
||||
name: name.into(),
|
||||
field: "token_credential",
|
||||
}
|
||||
}
|
||||
|
||||
/// The same source, naming `field` when its credential is absent.
|
||||
pub fn for_field(mut self, field: &'static str) -> Self {
|
||||
self.field = field;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn in_directory(directory: impl Into<PathBuf>, name: impl Into<String>) -> Self {
|
||||
Self {
|
||||
directory: Some(directory.into()),
|
||||
name: name.into(),
|
||||
field: "token_credential",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl TokenSource for CredentialTokenSource {
|
||||
fn read(&self) -> Option<Zeroizing<String>> {
|
||||
let directory = self.directory.as_ref()?;
|
||||
let mut token = Zeroizing::new(std::fs::read_to_string(directory.join(&self.name)).ok()?);
|
||||
let trimmed = token.trim_end_matches(['\r', '\n']).len();
|
||||
token.truncate(trimmed);
|
||||
(!token.is_empty()).then_some(token)
|
||||
}
|
||||
|
||||
/// No directory at all is the unit's environment lacking its credentials,
|
||||
/// which is not the same fault as a credential the directory lacks.
|
||||
fn missing(&self) -> &'static str {
|
||||
if self.directory.is_none() {
|
||||
"credentials_directory"
|
||||
} else {
|
||||
self.field
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Where the tail of a child's stderr goes. The backend never puts it in a
|
||||
/// result; the connector binary wires this to its log.
|
||||
pub trait StderrSink: Send + Sync {
|
||||
fn stderr_tail(&self, tail: &[u8]);
|
||||
}
|
||||
|
||||
/// Answers each call with one `claude` child process: never more than one per
|
||||
/// call, and a failed call is never run again.
|
||||
pub struct ClaudeCodeBackend {
|
||||
config: ClaudeCodeConfig,
|
||||
port: Box<dyn ProcessPort>,
|
||||
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 {
|
||||
/// The backend over real processes and the CLI's real session store.
|
||||
#[cfg(unix)]
|
||||
pub fn new(config: ClaudeCodeConfig, sink: Box<dyn StderrSink>) -> Self {
|
||||
let port = Box::new(crate::process::StdProcessPort::new(config.limits));
|
||||
let sessions = Box::new(crate::handle::FsSessionStore::new(
|
||||
config.config_dir.clone(),
|
||||
));
|
||||
let tokens = Box::new(CredentialTokenSource::from_environment(
|
||||
config.token_credential.clone(),
|
||||
));
|
||||
let handle_secret = Box::new(
|
||||
CredentialTokenSource::from_environment(config.handle_secret_credential.clone())
|
||||
.for_field("handle_secret_credential"),
|
||||
);
|
||||
Self::with_ports(config, port, sessions, sink, tokens, handle_secret)
|
||||
}
|
||||
|
||||
pub fn with_ports(
|
||||
config: ClaudeCodeConfig,
|
||||
port: Box<dyn ProcessPort>,
|
||||
sessions: Box<dyn SessionStore>,
|
||||
sink: Box<dyn StderrSink>,
|
||||
tokens: Box<dyn TokenSource>,
|
||||
handle_secret: Box<dyn TokenSource>,
|
||||
) -> Self {
|
||||
Self {
|
||||
config,
|
||||
port,
|
||||
sessions,
|
||||
sink,
|
||||
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:?}")),
|
||||
}
|
||||
}
|
||||
|
||||
/// The one pin check: `ready` and every spawn use it, and nothing caches
|
||||
/// its answer.
|
||||
fn binary_matches_pin(&self) -> bool {
|
||||
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(),
|
||||
) {
|
||||
Ok(plan) => plan,
|
||||
Err(error) => {
|
||||
return failed(
|
||||
FailureKind::InvalidRequest,
|
||||
error.to_string(),
|
||||
Some(zero_usage()),
|
||||
);
|
||||
}
|
||||
};
|
||||
let Some(token) = self.tokens.read() else {
|
||||
return failed(
|
||||
FailureKind::Authentication,
|
||||
"claude-code: the connector has no credential to give the child".to_string(),
|
||||
Some(zero_usage()),
|
||||
);
|
||||
};
|
||||
if !self.binary_matches_pin() {
|
||||
return failure(Failure::never_reached(
|
||||
FailureKind::Provider,
|
||||
"claude-code: the claude binary does not match its pin",
|
||||
));
|
||||
}
|
||||
let Ok(prompt) = PromptFile::create(&self.config.prompt_dir, &request.instructions) else {
|
||||
return failure(Failure::never_reached(
|
||||
FailureKind::Provider,
|
||||
"claude-code: the system prompt could not be staged",
|
||||
));
|
||||
};
|
||||
let argv = argv(&self.config, request, &plan, prompt.path(), ephemeral);
|
||||
let env = child_env(&self.config, request);
|
||||
let deadline = Instant::now() + self.config.limits.deadline;
|
||||
|
||||
let spawned = self.port.spawn(
|
||||
&argv,
|
||||
&env,
|
||||
&self.config.work_dir,
|
||||
stdin_bytes(plan.blocks()),
|
||||
&token,
|
||||
);
|
||||
let Ok(mut child) = spawned else {
|
||||
return failure(Failure::never_reached(
|
||||
FailureKind::Provider,
|
||||
"claude-code: the child could not be spawned",
|
||||
));
|
||||
};
|
||||
let finished = read_stream(
|
||||
child.as_mut(),
|
||||
deadline,
|
||||
&request.model,
|
||||
request.output_schema_json.as_deref(),
|
||||
&self.quota,
|
||||
);
|
||||
child.kill();
|
||||
child.wait();
|
||||
let tail = child.stderr_tail();
|
||||
if !tail.is_empty() {
|
||||
self.sink.stderr_tail(&tail);
|
||||
}
|
||||
|
||||
match finished {
|
||||
Ok(finished) => self.completed(request, caller, secret.as_bytes(), finished),
|
||||
Err(failure_) => failure(failure_),
|
||||
}
|
||||
}
|
||||
|
||||
fn completed(
|
||||
&self,
|
||||
request: &ProviderRequest,
|
||||
caller: &str,
|
||||
secret: &[u8],
|
||||
finished: Finished,
|
||||
) -> (Vec<TransportEvent>, TransportOutcome) {
|
||||
let mut payloads: Vec<TransportEventPayload> = finished
|
||||
.texts
|
||||
.iter()
|
||||
.map(|text| TransportEventPayload::TextDelta { text: text.clone() })
|
||||
.collect();
|
||||
|
||||
let mut prefix = request.input.clone();
|
||||
if !finished.texts.is_empty() {
|
||||
prefix.push(InputItem::AssistantText {
|
||||
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 events = payloads
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(sequence, payload)| TransportEvent {
|
||||
sequence: sequence as u64,
|
||||
payload,
|
||||
})
|
||||
.collect();
|
||||
(
|
||||
events,
|
||||
TransportOutcome::Completed {
|
||||
provider_response_id: finished.message_id,
|
||||
stop_reason: finished.stop_reason,
|
||||
usage: finished.usage,
|
||||
},
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl ProviderBackend for ClaudeCodeBackend {
|
||||
fn ready(&self) -> Result<(), &'static str> {
|
||||
readiness(
|
||||
&self.config,
|
||||
self.tokens.as_ref(),
|
||||
self.handle_secret.as_ref(),
|
||||
)
|
||||
}
|
||||
|
||||
fn provider(&self) -> &str {
|
||||
PROVIDER
|
||||
}
|
||||
|
||||
fn execute(&self, call: &ProviderCall) -> TransportResult {
|
||||
let (events, outcome) = self.run(call.request(), call.caller_runtime_id(), false);
|
||||
call.transported(events, outcome)
|
||||
}
|
||||
|
||||
fn plan_quota(&self) -> PlanQuotaSnapshot {
|
||||
self.quota.snapshot()
|
||||
}
|
||||
}
|
||||
|
||||
fn read_stream(
|
||||
child: &mut dyn RunningChild,
|
||||
deadline: Instant,
|
||||
model: &str,
|
||||
schema: Option<&str>,
|
||||
quota: &QuotaCell,
|
||||
) -> Result<Finished, Failure> {
|
||||
let mut state = StreamState::new(model, schema, quota);
|
||||
let finished = loop {
|
||||
match child.read_line(deadline) {
|
||||
StdoutRead::Line(line) => {
|
||||
if let Step::Done(done) = state.feed(&line) {
|
||||
break *done;
|
||||
}
|
||||
}
|
||||
StdoutRead::Eof => {
|
||||
return Err(Failure::cost_unknown(
|
||||
FailureKind::Provider,
|
||||
"claude-code: the stream ended without a result",
|
||||
));
|
||||
}
|
||||
StdoutRead::Timeout => {
|
||||
return Err(Failure::cost_unknown(
|
||||
FailureKind::Timeout,
|
||||
"claude-code: the deadline passed before a result",
|
||||
));
|
||||
}
|
||||
StdoutRead::Breach => {
|
||||
return Err(Failure::cost_unknown(
|
||||
FailureKind::Provider,
|
||||
"claude-code: the output passed its size cap",
|
||||
));
|
||||
}
|
||||
}
|
||||
};
|
||||
// The result is in; give the child until its deadline to close stdout.
|
||||
while matches!(child.read_line(deadline), StdoutRead::Line(_)) {}
|
||||
finished
|
||||
}
|
||||
|
||||
fn binary_matches_pin(config: &ClaudeCodeConfig) -> bool {
|
||||
let Ok(mut file) = File::open(&config.claude_binary) else {
|
||||
return false;
|
||||
};
|
||||
let mut hasher = Sha256::new();
|
||||
io::copy(&mut file, &mut hasher).is_ok()
|
||||
&& <[u8; 32]>::from(hasher.finalize()) == config.claude_binary_sha256
|
||||
}
|
||||
|
||||
/// Whether the backend can take work, naming the first config field that
|
||||
/// cannot be used. `config check` and `serve` both ask this.
|
||||
pub(crate) fn readiness(
|
||||
config: &ClaudeCodeConfig,
|
||||
tokens: &dyn TokenSource,
|
||||
handle_secret: &dyn TokenSource,
|
||||
) -> Result<(), &'static str> {
|
||||
if !config.work_dir.is_dir() {
|
||||
return Err("work_dir");
|
||||
}
|
||||
if !config.config_dir.is_dir() {
|
||||
return Err("config_dir");
|
||||
}
|
||||
if !config.prompt_dir.is_dir() {
|
||||
return Err("prompt_dir");
|
||||
}
|
||||
if !binary_matches_pin(config) {
|
||||
return Err("claude_binary_sha256");
|
||||
}
|
||||
// Last: outside the unit the credential is the one thing always absent,
|
||||
// and it must not hide the faults the document itself has.
|
||||
if usable_secret(handle_secret).is_none() {
|
||||
return Err(handle_secret.missing());
|
||||
}
|
||||
if tokens.read().is_none() {
|
||||
return Err(tokens.missing());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn zero_usage() -> Usage {
|
||||
Usage {
|
||||
input_tokens: 0,
|
||||
cache_read_tokens: 0,
|
||||
cache_write_tokens: 0,
|
||||
output_tokens: 0,
|
||||
reasoning_tokens: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn failed(
|
||||
kind: FailureKind,
|
||||
message: String,
|
||||
usage: Option<Usage>,
|
||||
) -> (Vec<TransportEvent>, TransportOutcome) {
|
||||
(
|
||||
Vec::new(),
|
||||
TransportOutcome::Failed {
|
||||
kind,
|
||||
message,
|
||||
usage,
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
fn failure(failure: Failure) -> (Vec<TransportEvent>, TransportOutcome) {
|
||||
failed(failure.kind, failure.message.to_string(), failure.usage)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
2792
crates/claude-connector/src/backend/tests.rs
Normal file
2792
crates/claude-connector/src/backend/tests.rs
Normal file
File diff suppressed because it is too large
Load diff
44
crates/claude-connector/src/config.rs
Normal file
44
crates/claude-connector/src/config.rs
Normal file
|
|
@ -0,0 +1,44 @@
|
|||
use std::path::PathBuf;
|
||||
use std::time::Duration;
|
||||
|
||||
/// The provider tag the kit puts on every receipt and on the reasoning state
|
||||
/// of this backend. It is distinct from any Messages-API backend's tag, so the
|
||||
/// kit treats each backend's state as foreign to the other.
|
||||
pub(crate) const PROVIDER: &str = "claude-code";
|
||||
|
||||
/// What the connector's binary hands the backend. Every path is absolute.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ClaudeCodeConfig {
|
||||
/// The `claude` executable.
|
||||
pub claude_binary: PathBuf,
|
||||
/// The SHA-256 the executable must have for [`ready`](model_connector::ProviderBackend::ready).
|
||||
pub claude_binary_sha256: [u8; 32],
|
||||
/// An empty directory the connector owns: the CLI writes its cwd, git state
|
||||
/// and platform into the system prompt.
|
||||
pub work_dir: PathBuf,
|
||||
/// Passed to the child as `CLAUDE_CONFIG_DIR` and `HOME`; it holds the
|
||||
/// CLI's session store.
|
||||
pub config_dir: PathBuf,
|
||||
/// A directory the connector owns (mode 0700), where each call's system
|
||||
/// prompt is staged for the child.
|
||||
pub prompt_dir: PathBuf,
|
||||
/// The name of the systemd credential that holds the child's token, under
|
||||
/// `$CREDENTIALS_DIRECTORY`.
|
||||
pub token_credential: String,
|
||||
/// The name of the systemd credential that holds the handle secret: the
|
||||
/// key every handle is authenticated with.
|
||||
pub handle_secret_credential: String,
|
||||
/// The connector instance id written into every handle.
|
||||
pub instance: String,
|
||||
pub limits: ClaudeCodeLimits,
|
||||
}
|
||||
|
||||
/// The bounds of one child process.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct ClaudeCodeLimits {
|
||||
/// The whole call, spawn to result.
|
||||
pub deadline: Duration,
|
||||
pub max_line_bytes: usize,
|
||||
pub max_total_bytes: usize,
|
||||
pub stderr_tail_bytes: usize,
|
||||
}
|
||||
830
crates/claude-connector/src/daemon.rs
Normal file
830
crates/claude-connector/src/daemon.rs
Normal file
|
|
@ -0,0 +1,830 @@
|
|||
//! The claude-connector's configuration document and its serve verb. One typed
|
||||
//! CultCache record owns every bound and path the backend and the serve loop
|
||||
//! run with; nothing in the connector's environment, flags or JSON can change
|
||||
//! them. Validation errors name a field and never echo a value.
|
||||
|
||||
use std::collections::BTreeSet;
|
||||
use std::net::{SocketAddr, TcpListener};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Duration;
|
||||
|
||||
use cultcache_rs::{CultCache, DatabaseEntry, SingleFileMessagePackBackingStore};
|
||||
use model_connector::{
|
||||
CallerAdmission, CallerConfig, LeaseTurn, MIN_PAYLOAD_BYTES, ProviderBackend, ServeError,
|
||||
ServeLimits, ServiceError, TransportService, load_admissions, serve_service,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::backend::readiness;
|
||||
#[cfg(target_os = "linux")]
|
||||
use crate::lease::{LeaseCaller, LeaseLimits, serve_leases, user_uid};
|
||||
#[cfg(target_os = "linux")]
|
||||
use crate::unit::SystemdUnits;
|
||||
#[cfg(unix)]
|
||||
use crate::meter;
|
||||
use crate::{
|
||||
ClaudeCodeBackend, ClaudeCodeConfig, ClaudeCodeLimits, CredentialTokenSource, StderrSink,
|
||||
};
|
||||
|
||||
const CONFIG_KEY: &str = "runtime";
|
||||
const MODEL_PREFIX: &str = "claude-";
|
||||
|
||||
/// A unix account that may lease the subscription token for a Self run or a
|
||||
/// Persona turn, the unit its process must run in (`mind-self@<unit_instance>`
|
||||
/// or `mind-persona@<unit_instance>`, by `turn`), and the configured caller
|
||||
/// whose name its leases carry.
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct LeaseCallerConfig {
|
||||
pub unix_user: String,
|
||||
pub caller_id: String,
|
||||
/// The kind of turn this account leases for: it fixes the unit, `mind-self@`
|
||||
/// or `mind-persona@`, its process must run in.
|
||||
pub turn: LeaseTurn,
|
||||
pub unit_instance: String,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, DatabaseEntry)]
|
||||
#[cultcache(type = "gamecult.model_connector.claude_connector_config.v2")]
|
||||
pub struct ClaudeConnectorConfig {
|
||||
/// A loopback socket address.
|
||||
#[cultcache(key = 0)]
|
||||
pub listen: String,
|
||||
#[cultcache(key = 1)]
|
||||
pub replay_store: PathBuf,
|
||||
#[cultcache(key = 2)]
|
||||
pub claude_binary: PathBuf,
|
||||
#[cultcache(key = 3)]
|
||||
pub claude_binary_sha256: [u8; 32],
|
||||
#[cultcache(key = 4)]
|
||||
pub work_dir: PathBuf,
|
||||
#[cultcache(key = 5)]
|
||||
pub config_dir: PathBuf,
|
||||
/// Where each call's system prompt is staged: a directory the connector
|
||||
/// owns, mode 0700.
|
||||
#[cultcache(key = 6)]
|
||||
pub prompt_dir: PathBuf,
|
||||
#[cultcache(key = 7)]
|
||||
pub instance: String,
|
||||
/// The name of the systemd credential that holds the child's token.
|
||||
#[cultcache(key = 8)]
|
||||
pub token_credential: String,
|
||||
#[cultcache(key = 9)]
|
||||
pub deadline_ms: u64,
|
||||
#[cultcache(key = 10)]
|
||||
pub max_line_bytes: usize,
|
||||
#[cultcache(key = 11)]
|
||||
pub max_total_bytes: usize,
|
||||
#[cultcache(key = 12)]
|
||||
pub stderr_tail_bytes: usize,
|
||||
#[cultcache(key = 13)]
|
||||
pub max_frame_bytes: usize,
|
||||
#[cultcache(key = 14)]
|
||||
pub max_connections: usize,
|
||||
#[cultcache(key = 15)]
|
||||
pub socket_timeout_ms: u64,
|
||||
#[cultcache(key = 16)]
|
||||
pub max_expiry_skew_ms: u64,
|
||||
/// Every caller of this connector: its own connection key and epoch, models,
|
||||
/// concurrency, payload and output bounds.
|
||||
/// Ids are unique. Handles are bound to the caller that was given them.
|
||||
#[cultcache(key = 17)]
|
||||
pub callers: Vec<CallerConfig>,
|
||||
/// Where the Self-run lease socket is bound (a systemd RuntimeDirectory),
|
||||
/// or none when this connector lends no token.
|
||||
#[cultcache(key = 19)]
|
||||
pub lease_socket: Option<PathBuf>,
|
||||
/// The accounts allowed to lease, each bound to a configured caller.
|
||||
/// Empty exactly when `lease_socket` is none.
|
||||
#[cultcache(key = 20)]
|
||||
pub lease_callers: Vec<LeaseCallerConfig>,
|
||||
/// The furthest a lease's expiry may lie from now.
|
||||
#[cultcache(key = 21)]
|
||||
pub lease_max_ms: u64,
|
||||
/// The name of the systemd credential that holds the handle secret (at
|
||||
/// least 32 bytes), the key every handle is authenticated with.
|
||||
#[cultcache(key = 22)]
|
||||
pub handle_secret_credential: String,
|
||||
}
|
||||
|
||||
/// A refusal. It names the field or stage that failed and nothing the
|
||||
/// operator or the file typed.
|
||||
#[derive(Debug, Error)]
|
||||
pub enum ConfigError {
|
||||
#[error("invalid configuration field {0}")]
|
||||
Invalid(&'static str),
|
||||
#[error("the configuration store could not be read as a claude connector document")]
|
||||
Store,
|
||||
#[error("the prompt directory is not a directory owned by this process with mode 0700")]
|
||||
PromptDir,
|
||||
#[error("the listener could not be bound")]
|
||||
Listen,
|
||||
/// Something the document names exists but cannot be used now.
|
||||
#[error("the configured {0} cannot be used")]
|
||||
Unusable(&'static str),
|
||||
/// The serve loop failed after it started.
|
||||
#[error("the serve loop failed at {0}")]
|
||||
Serving(&'static str),
|
||||
}
|
||||
|
||||
/// The kit's serve errors carry stored paths, ids and numbers; none of that
|
||||
/// text is shown. Each maps to the configuration field or stage it names.
|
||||
impl From<ServeError> for ConfigError {
|
||||
fn from(error: ServeError) -> Self {
|
||||
match error {
|
||||
ServeError::InvalidConfig(field) | ServeError::BackendUnready(field) => {
|
||||
Self::Unusable(field)
|
||||
}
|
||||
ServeError::FrameBound { .. } => Self::Invalid("max_frame_bytes"),
|
||||
ServeError::ConnectionKey(_) => Self::Unusable("connection_key_file"),
|
||||
ServeError::Service(error) => match error {
|
||||
ServiceError::ReplayStore(_) => Self::Unusable("replay_store"),
|
||||
ServiceError::InvalidAdmission
|
||||
| ServiceError::DuplicateCaller
|
||||
| ServiceError::SharedCallerKey => Self::Invalid("callers"),
|
||||
_ => Self::Serving("service"),
|
||||
},
|
||||
ServeError::ServicePoisoned => Self::Serving("service"),
|
||||
ServeError::Listen(_) => Self::Serving("listener"),
|
||||
ServeError::Thread(_) | ServeError::Connection(_) => Self::Serving("connection"),
|
||||
ServeError::FrameSize | ServeError::FrameEncoding => Self::Serving("frame"),
|
||||
ServeError::Clock => Self::Serving("clock"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ClaudeConnectorConfig {
|
||||
/// The one validation: `config init`, `config check` and `serve` all pass
|
||||
/// through it. Returns the listen address.
|
||||
pub fn validate(&self) -> Result<SocketAddr, ConfigError> {
|
||||
let listen = self
|
||||
.listen
|
||||
.parse::<SocketAddr>()
|
||||
.ok()
|
||||
.filter(|address| address.ip().is_loopback())
|
||||
.ok_or(ConfigError::Invalid("listen"))?;
|
||||
let paths = [
|
||||
("replay_store", &self.replay_store),
|
||||
("claude_binary", &self.claude_binary),
|
||||
("work_dir", &self.work_dir),
|
||||
("config_dir", &self.config_dir),
|
||||
("prompt_dir", &self.prompt_dir),
|
||||
];
|
||||
for (field, path) in paths {
|
||||
if !path.is_absolute() {
|
||||
return Err(ConfigError::Invalid(field));
|
||||
}
|
||||
}
|
||||
if self.claude_binary_sha256 == [0; 32] {
|
||||
return Err(ConfigError::Invalid("claude_binary_sha256"));
|
||||
}
|
||||
if !is_trimmed_nonempty(&self.instance) {
|
||||
return Err(ConfigError::Invalid("instance"));
|
||||
}
|
||||
if !is_credential_name(&self.token_credential) {
|
||||
return Err(ConfigError::Invalid("token_credential"));
|
||||
}
|
||||
if !is_credential_name(&self.handle_secret_credential)
|
||||
|| self.handle_secret_credential == self.token_credential
|
||||
{
|
||||
return Err(ConfigError::Invalid("handle_secret_credential"));
|
||||
}
|
||||
if self.deadline_ms == 0
|
||||
|| self.max_line_bytes == 0
|
||||
|| self.max_total_bytes < self.max_line_bytes
|
||||
|| self.stderr_tail_bytes == 0
|
||||
{
|
||||
return Err(ConfigError::Invalid("child limits"));
|
||||
}
|
||||
if self.max_frame_bytes == 0
|
||||
|| self.max_connections == 0
|
||||
|| self.socket_timeout_ms == 0
|
||||
|| self.max_expiry_skew_ms == 0
|
||||
{
|
||||
return Err(ConfigError::Invalid("serve limits"));
|
||||
}
|
||||
if self.callers.is_empty() {
|
||||
return Err(ConfigError::Invalid("callers"));
|
||||
}
|
||||
let mut ids = BTreeSet::new();
|
||||
for caller in &self.callers {
|
||||
if !is_trimmed_nonempty(&caller.caller_runtime_id) {
|
||||
return Err(ConfigError::Invalid("caller_runtime_id"));
|
||||
}
|
||||
if !ids.insert(caller.caller_runtime_id.as_str()) {
|
||||
return Err(ConfigError::Invalid("callers"));
|
||||
}
|
||||
if !caller.connection_key_file.is_absolute() {
|
||||
return Err(ConfigError::Invalid("connection_key_file"));
|
||||
}
|
||||
if caller.connection_key_epoch == 0 {
|
||||
return Err(ConfigError::Invalid("connection_key_epoch"));
|
||||
}
|
||||
if caller.max_concurrent_requests == 0
|
||||
|| caller.max_payload_bytes < MIN_PAYLOAD_BYTES
|
||||
|| caller.max_output_tokens == 0
|
||||
{
|
||||
return Err(ConfigError::Invalid("caller bounds"));
|
||||
}
|
||||
if caller.allowed_models.is_empty() {
|
||||
return Err(ConfigError::Invalid("allowed_models"));
|
||||
}
|
||||
let mut seen = BTreeSet::new();
|
||||
for model in &caller.allowed_models {
|
||||
if !has_model_id_form(model) || !seen.insert(model.as_str()) {
|
||||
return Err(ConfigError::Invalid("allowed_models"));
|
||||
}
|
||||
}
|
||||
}
|
||||
self.validate_leases()?;
|
||||
Ok(listen)
|
||||
}
|
||||
|
||||
fn validate_leases(&self) -> Result<(), ConfigError> {
|
||||
let Some(socket) = &self.lease_socket else {
|
||||
return if self.lease_callers.is_empty() {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(ConfigError::Invalid("lease_callers"))
|
||||
};
|
||||
};
|
||||
if !socket.is_absolute() {
|
||||
return Err(ConfigError::Invalid("lease_socket"));
|
||||
}
|
||||
if self.lease_max_ms == 0 {
|
||||
return Err(ConfigError::Invalid("lease_max_ms"));
|
||||
}
|
||||
let mut users = BTreeSet::new();
|
||||
if self.lease_callers.is_empty()
|
||||
|| !self.lease_callers.iter().all(|lease| {
|
||||
is_trimmed_nonempty(&lease.unix_user)
|
||||
&& self
|
||||
.callers
|
||||
.iter()
|
||||
.any(|caller| caller.caller_runtime_id == lease.caller_id)
|
||||
&& is_unit_instance(&lease.unit_instance)
|
||||
&& users.insert(lease.unix_user.as_str())
|
||||
})
|
||||
{
|
||||
return Err(ConfigError::Invalid("lease_callers"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn backend_config(&self) -> ClaudeCodeConfig {
|
||||
ClaudeCodeConfig {
|
||||
claude_binary: self.claude_binary.clone(),
|
||||
claude_binary_sha256: self.claude_binary_sha256,
|
||||
work_dir: self.work_dir.clone(),
|
||||
config_dir: self.config_dir.clone(),
|
||||
prompt_dir: self.prompt_dir.clone(),
|
||||
token_credential: self.token_credential.clone(),
|
||||
handle_secret_credential: self.handle_secret_credential.clone(),
|
||||
instance: self.instance.clone(),
|
||||
limits: ClaudeCodeLimits {
|
||||
deadline: Duration::from_millis(self.deadline_ms),
|
||||
max_line_bytes: self.max_line_bytes,
|
||||
max_total_bytes: self.max_total_bytes,
|
||||
stderr_tail_bytes: self.stderr_tail_bytes,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn serve_limits(&self) -> ServeLimits {
|
||||
ServeLimits {
|
||||
max_frame_bytes: self.max_frame_bytes,
|
||||
max_connections: self.max_connections,
|
||||
socket_timeout_ms: self.socket_timeout_ms,
|
||||
max_expiry_skew_ms: self.max_expiry_skew_ms,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The `<instance>` of a unit name, as plain as a systemd instance can be.
|
||||
fn is_unit_instance(text: &str) -> bool {
|
||||
!text.is_empty()
|
||||
&& text
|
||||
.bytes()
|
||||
.all(|b| b.is_ascii_alphanumeric() || matches!(b, b'-' | b'_'))
|
||||
}
|
||||
|
||||
fn is_trimmed_nonempty(text: &str) -> bool {
|
||||
!text.is_empty() && text.trim() == text
|
||||
}
|
||||
|
||||
/// The loose shape of a model id, only to catch typos:
|
||||
/// `claude-<family>-<major>`, then an optional `-<minor>` and an optional
|
||||
/// `-YYYYMMDD`, the family lowercase letters. No shape can tell an alias from
|
||||
/// a canonical id, so the known list is `allowed_models` itself; an alias that slips in fails its first
|
||||
/// call, when the stream's `init.model` differs from the request.
|
||||
fn has_model_id_form(model: &str) -> bool {
|
||||
let Some(rest) = model.strip_prefix(MODEL_PREFIX) else {
|
||||
return false;
|
||||
};
|
||||
let digits = |part: &str, most: usize| {
|
||||
(1..=most).contains(&part.len()) && part.bytes().all(|byte| byte.is_ascii_digit())
|
||||
};
|
||||
let date = |part: &str| part.len() == 8 && digits(part, 8);
|
||||
let mut parts = rest.split('-');
|
||||
let family = parts.next().unwrap_or("");
|
||||
let tail: Vec<&str> = parts.collect();
|
||||
let Some((major, after)) = tail.split_first() else {
|
||||
return false;
|
||||
};
|
||||
let after_ok = match after {
|
||||
[] => true,
|
||||
[one] => digits(one, 2) || date(one),
|
||||
[minor, last] => digits(minor, 2) && date(last),
|
||||
_ => false,
|
||||
};
|
||||
!family.is_empty()
|
||||
&& family.bytes().all(|byte| byte.is_ascii_lowercase())
|
||||
&& digits(major, 2)
|
||||
&& after_ok
|
||||
}
|
||||
|
||||
/// One file name under `$CREDENTIALS_DIRECTORY`, never a path.
|
||||
fn is_credential_name(name: &str) -> bool {
|
||||
!name.is_empty()
|
||||
&& name != "."
|
||||
&& name != ".."
|
||||
&& name
|
||||
.bytes()
|
||||
.all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'.' | b'_'))
|
||||
}
|
||||
|
||||
/// Reads the document without writing: a shared snapshot of exactly one
|
||||
/// record, with the exact key, type and schema id, in canonical positional
|
||||
/// MessagePack, then validated.
|
||||
pub fn load_config(path: &Path) -> Result<ClaudeConnectorConfig, ConfigError> {
|
||||
let store = SingleFileMessagePackBackingStore::new(path);
|
||||
let config = store
|
||||
.with_read_only_shared_snapshot(|entries| {
|
||||
let [envelope] = entries.as_slice() else {
|
||||
anyhow::bail!("not exactly one record");
|
||||
};
|
||||
if envelope.key != CONFIG_KEY
|
||||
|| envelope.r#type != ClaudeConnectorConfig::TYPE
|
||||
|| envelope.schema_id.as_deref() != Some(ClaudeConnectorConfig::TYPE)
|
||||
{
|
||||
anyhow::bail!("wrong typed envelope");
|
||||
}
|
||||
let config: ClaudeConnectorConfig = rmp_serde::from_slice(&envelope.payload)?;
|
||||
if rmp_serde::to_vec(&config)? != envelope.payload {
|
||||
anyhow::bail!("not canonical positional MessagePack");
|
||||
}
|
||||
Ok(config)
|
||||
})
|
||||
.map_err(|_| ConfigError::Store)?;
|
||||
config.validate()?;
|
||||
Ok(config)
|
||||
}
|
||||
|
||||
/// Validates, then writes the document as the store's only record.
|
||||
pub fn write_config(path: &Path, config: &ClaudeConnectorConfig) -> Result<(), ConfigError> {
|
||||
config.validate()?;
|
||||
let mut cache = CultCache::new();
|
||||
cache
|
||||
.register_entry_type::<ClaudeConnectorConfig>()
|
||||
.map_err(|_| ConfigError::Store)?;
|
||||
cache
|
||||
.add_generic_backing_store(SingleFileMessagePackBackingStore::new(path))
|
||||
.map_err(|_| ConfigError::Store)?;
|
||||
cache
|
||||
.pull_all_backing_stores()
|
||||
.map_err(|_| ConfigError::Store)?;
|
||||
cache
|
||||
.put(CONFIG_KEY, config)
|
||||
.map_err(|_| ConfigError::Store)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// The prompt directory holds every system prompt in flight: it must be a
|
||||
/// real directory (not a link) that this process owns and nobody else can
|
||||
/// enter.
|
||||
#[cfg(unix)]
|
||||
pub fn check_prompt_dir(dir: &Path) -> Result<(), ConfigError> {
|
||||
use std::os::unix::fs::{MetadataExt, PermissionsExt};
|
||||
let metadata = std::fs::symlink_metadata(dir).map_err(|_| ConfigError::PromptDir)?;
|
||||
// SAFETY: geteuid has no preconditions and cannot fail.
|
||||
let uid = unsafe { libc::geteuid() };
|
||||
if !metadata.is_dir()
|
||||
|| metadata.uid() != uid
|
||||
|| metadata.permissions().mode() & 0o7777 != 0o700
|
||||
{
|
||||
return Err(ConfigError::PromptDir);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// The child's stderr is the log's business by size only: its bytes are never
|
||||
/// written anywhere the connector controls.
|
||||
struct LogSink<W>(std::sync::Mutex<W>);
|
||||
|
||||
impl<W: std::io::Write + Send> StderrSink for LogSink<W> {
|
||||
fn stderr_tail(&self, tail: &[u8]) {
|
||||
if let Ok(mut log) = self.0.lock() {
|
||||
let _ = writeln!(log, "claude-connector: child stderr, {} bytes", tail.len());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Everything `serve` needs true before it binds, in one place so that
|
||||
/// `config check` passes exactly what `serve` will start: the prompt
|
||||
/// directory, the other directories, the credential, the pin, the connection
|
||||
/// key and the frame and connection bounds. Returns the admitted callers.
|
||||
fn readiness_of(config: &ClaudeConnectorConfig) -> Result<Vec<CallerAdmission>, ConfigError> {
|
||||
config.validate()?;
|
||||
#[cfg(unix)]
|
||||
check_prompt_dir(&config.prompt_dir)?;
|
||||
let admissions = load_admissions(&config.callers)?;
|
||||
config.serve_limits().validate(&admissions)?;
|
||||
let backend = config.backend_config();
|
||||
readiness(
|
||||
&backend,
|
||||
&CredentialTokenSource::from_environment(backend.token_credential.clone()),
|
||||
&CredentialTokenSource::from_environment(backend.handle_secret_credential.clone())
|
||||
.for_field("handle_secret_credential"),
|
||||
)
|
||||
.map_err(ConfigError::Unusable)?;
|
||||
if config.lease_socket.is_some() {
|
||||
lease_callers_of(config)?;
|
||||
}
|
||||
Ok(admissions)
|
||||
}
|
||||
|
||||
/// The accounts allowed to lease, resolved to the uids the kernel will report
|
||||
/// for their connections. An account the system does not have is not ready.
|
||||
#[cfg(target_os = "linux")]
|
||||
fn lease_callers_of(config: &ClaudeConnectorConfig) -> Result<Vec<LeaseCaller>, ConfigError> {
|
||||
config
|
||||
.lease_callers
|
||||
.iter()
|
||||
.map(|lease| {
|
||||
user_uid(&lease.unix_user)
|
||||
.map(|uid| LeaseCaller {
|
||||
uid,
|
||||
caller_id: lease.caller_id.clone(),
|
||||
turn: lease.turn,
|
||||
unit_instance: lease.unit_instance.clone(),
|
||||
})
|
||||
.ok_or(ConfigError::Unusable("lease_callers"))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Peer credentials identify a lease caller only on Linux.
|
||||
#[cfg(not(target_os = "linux"))]
|
||||
fn lease_callers_of(_: &ClaudeConnectorConfig) -> Result<(), ConfigError> {
|
||||
Err(ConfigError::Unusable("lease_socket"))
|
||||
}
|
||||
|
||||
/// Binds the lease socket (mode 0660; the unit's RuntimeDirectory and group
|
||||
/// decide who can reach it, and the kernel's peer uid decides who may lease)
|
||||
/// and serves it beside the RUDP loop. The lease listener is half of the
|
||||
/// connector: if it stops, so does the process.
|
||||
#[cfg(target_os = "linux")]
|
||||
fn start_leases(
|
||||
config: &ClaudeConnectorConfig,
|
||||
socket: &Path,
|
||||
service: Arc<Mutex<TransportService>>,
|
||||
quota: Arc<crate::QuotaCell>,
|
||||
) -> Result<(), ConfigError> {
|
||||
use std::os::unix::fs::{FileTypeExt, PermissionsExt};
|
||||
use std::os::unix::net::UnixListener;
|
||||
|
||||
let callers = lease_callers_of(config)?;
|
||||
// A socket a previous run left is replaced; anything else at the path is not ours.
|
||||
if std::fs::symlink_metadata(socket).is_ok_and(|found| found.file_type().is_socket()) {
|
||||
let _ = std::fs::remove_file(socket);
|
||||
}
|
||||
let listener = UnixListener::bind(socket).map_err(|_| ConfigError::Listen)?;
|
||||
std::fs::set_permissions(socket, std::fs::Permissions::from_mode(0o660))
|
||||
.map_err(|_| ConfigError::Listen)?;
|
||||
let limits = LeaseLimits {
|
||||
socket_timeout: Duration::from_millis(config.socket_timeout_ms),
|
||||
max_lease_ms: config.lease_max_ms,
|
||||
max_connections: config.max_connections,
|
||||
};
|
||||
let tokens = Arc::new(CredentialTokenSource::from_environment(
|
||||
config.token_credential.clone(),
|
||||
));
|
||||
std::thread::Builder::new()
|
||||
.name("claude-lease-listener".to_string())
|
||||
.spawn(move || {
|
||||
let _ = serve_leases(
|
||||
&listener,
|
||||
callers,
|
||||
limits,
|
||||
service,
|
||||
quota,
|
||||
tokens,
|
||||
Arc::new(SystemdUnits::system()),
|
||||
);
|
||||
eprintln!("claude-connector: the lease listener stopped");
|
||||
std::process::exit(1);
|
||||
})
|
||||
.map_err(|_| ConfigError::Serving("lease listener"))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Whether `serve` would start with this document and this environment.
|
||||
pub fn check_ready(config: &ClaudeConnectorConfig) -> Result<(), ConfigError> {
|
||||
readiness_of(config).map(drop)
|
||||
}
|
||||
|
||||
/// Loads the document and serves until the listener fails.
|
||||
#[cfg(unix)]
|
||||
pub fn serve_config(path: &Path) -> Result<(), ConfigError> {
|
||||
serve_config_with(path, |_| ())
|
||||
}
|
||||
|
||||
/// [`serve_config`] handing the quota cell to `started` once the meter runs:
|
||||
/// a test stands in for the swarm's lease requests only through here.
|
||||
#[cfg(unix)]
|
||||
fn serve_config_with(
|
||||
path: &Path,
|
||||
started: impl FnOnce(&Arc<crate::QuotaCell>),
|
||||
) -> Result<(), ConfigError> {
|
||||
let config = load_config(path)?;
|
||||
let admissions = readiness_of(&config)?;
|
||||
let listen = config.validate()?;
|
||||
let backend = Arc::new(ClaudeCodeBackend::new(
|
||||
config.backend_config(),
|
||||
Box::new(LogSink(std::sync::Mutex::new(std::io::stderr()))),
|
||||
));
|
||||
let listener = TcpListener::bind(listen).map_err(|_| ConfigError::Listen)?;
|
||||
let service = Arc::new(Mutex::new(
|
||||
TransportService::open(
|
||||
&config.replay_store,
|
||||
backend.provider(),
|
||||
admissions,
|
||||
config.max_expiry_skew_ms,
|
||||
)
|
||||
.map_err(|error| ConfigError::from(ServeError::Service(error)))?,
|
||||
));
|
||||
// The meter is running before the first lease can ask, so that lease can
|
||||
// wait for its first read.
|
||||
let quota = backend.quota();
|
||||
meter::start(quota.clone(), Box::new(backend.clone()))
|
||||
.map_err(|_| ConfigError::Serving("meter"))?;
|
||||
started("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)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use sha2::Digest;
|
||||
|
||||
use super::*;
|
||||
|
||||
/// A log that a test can read back.
|
||||
#[derive(Clone, Default)]
|
||||
struct Captured(Arc<Mutex<Vec<u8>>>);
|
||||
|
||||
impl std::io::Write for Captured {
|
||||
fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
|
||||
self.0.lock().unwrap().extend_from_slice(bytes);
|
||||
Ok(bytes.len())
|
||||
}
|
||||
fn flush(&mut self) -> std::io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// `serve` starts the meter beside the listener, wired to the backend's own
|
||||
/// cell: it makes no call while nobody asks, and once a lease is asked for
|
||||
/// it runs the claude binary and serve stays up. The binary is a script
|
||||
/// that replays a real captured call.
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn serve_config_starts_the_meter_and_stays_up() {
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
|
||||
const CALL: &str = include_str!("../tests/fixtures/claude-2.1.312-meter-call.jsonl");
|
||||
let dir =
|
||||
std::env::temp_dir().join(format!("claude-connector-serve-{}", std::process::id()));
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
let at = |name: &str| dir.join(name);
|
||||
for sub in ["work", "home", "prompts", "credentials"] {
|
||||
std::fs::create_dir_all(at(sub)).unwrap();
|
||||
}
|
||||
std::fs::set_permissions(at("prompts"), std::fs::Permissions::from_mode(0o700)).unwrap();
|
||||
std::fs::write(at("call.jsonl"), CALL).unwrap();
|
||||
let script = format!(
|
||||
"#!/bin/sh\ncat > /dev/null\necho spawned >> '{}'\ncat '{}'\n",
|
||||
at("spawned").display(),
|
||||
at("call.jsonl").display()
|
||||
);
|
||||
std::fs::write(at("claude"), &script).unwrap();
|
||||
std::fs::set_permissions(at("claude"), std::fs::Permissions::from_mode(0o755)).unwrap();
|
||||
std::fs::write(at("caller.key"), "k".repeat(48)).unwrap();
|
||||
std::fs::write(at("credentials/claude-token"), "sk-ant-oat01-fake-token").unwrap();
|
||||
std::fs::write(at("credentials/claude-handle-secret"), "s".repeat(48)).unwrap();
|
||||
let config = ClaudeConnectorConfig {
|
||||
listen: "127.0.0.1:0".to_string(),
|
||||
replay_store: at("replay.cc"),
|
||||
claude_binary: at("claude"),
|
||||
claude_binary_sha256: sha2::Sha256::digest(script.as_bytes()).into(),
|
||||
work_dir: at("work"),
|
||||
config_dir: at("home"),
|
||||
prompt_dir: at("prompts"),
|
||||
instance: "persona-a".to_string(),
|
||||
token_credential: "claude-token".to_string(),
|
||||
handle_secret_credential: "claude-handle-secret".to_string(),
|
||||
deadline_ms: 600_000,
|
||||
max_line_bytes: 4 << 20,
|
||||
max_total_bytes: 64 << 20,
|
||||
stderr_tail_bytes: 8192,
|
||||
max_frame_bytes: 2 << 20,
|
||||
max_connections: 8,
|
||||
socket_timeout_ms: 10_000,
|
||||
max_expiry_skew_ms: 300_000,
|
||||
callers: vec![CallerConfig {
|
||||
caller_runtime_id: "persona-a-caller".to_string(),
|
||||
connection_key_file: at("caller.key"),
|
||||
connection_key_epoch: 1,
|
||||
allowed_models: vec!["claude-haiku-5-5".to_string()],
|
||||
max_concurrent_requests: 4,
|
||||
max_payload_bytes: 1 << 20,
|
||||
max_output_tokens: 32_768,
|
||||
}],
|
||||
lease_socket: None,
|
||||
lease_callers: Vec::new(),
|
||||
lease_max_ms: 3_600_000,
|
||||
};
|
||||
write_config(&at("connector.cc"), &config).unwrap();
|
||||
// SAFETY: set once, to a directory only this test reads, before the
|
||||
// serve thread it starts.
|
||||
unsafe { std::env::set_var("CREDENTIALS_DIRECTORY", at("credentials")) };
|
||||
|
||||
let path = at("connector.cc");
|
||||
let lease = std::sync::Arc::new(std::sync::Mutex::new(None));
|
||||
let swarm = lease.clone();
|
||||
let served = std::thread::spawn(move || {
|
||||
serve_config_with(&path, |cell| {
|
||||
*swarm.lock().unwrap() = Some(cell.clone());
|
||||
})
|
||||
});
|
||||
|
||||
let deadline = std::time::Instant::now() + Duration::from_secs(30);
|
||||
let cell = loop {
|
||||
if let Some(cell) = lease.lock().unwrap().clone() {
|
||||
break cell;
|
||||
}
|
||||
assert!(std::time::Instant::now() < deadline, "serve never started");
|
||||
std::thread::sleep(Duration::from_millis(20));
|
||||
};
|
||||
std::thread::sleep(Duration::from_millis(500));
|
||||
assert!(!at("spawned").exists(), "a quiet swarm spends nothing");
|
||||
let _lease = cell.asking();
|
||||
|
||||
while !at("spawned").exists() {
|
||||
assert!(
|
||||
std::time::Instant::now() < deadline,
|
||||
"the meter never ran the claude binary"
|
||||
);
|
||||
std::thread::sleep(Duration::from_millis(50));
|
||||
}
|
||||
assert!(!served.is_finished(), "serve returned after starting the meter");
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serve_errors_name_a_field_or_stage_and_never_a_stored_value() {
|
||||
const CANARY: &str = "canary-path-8e2a";
|
||||
let io = || std::io::Error::other(CANARY);
|
||||
let cases: Vec<(ServeError, &str)> = vec![
|
||||
(ServeError::InvalidConfig("callers"), "callers"),
|
||||
(ServeError::BackendUnready("work_dir"), "work_dir"),
|
||||
(
|
||||
ServeError::FrameBound {
|
||||
max_frame_bytes: 4099,
|
||||
payload_bytes: 4096,
|
||||
needed_frame_bytes: 9999,
|
||||
},
|
||||
"max_frame_bytes",
|
||||
),
|
||||
(ServeError::ConnectionKey(io()), "connection_key_file"),
|
||||
(
|
||||
ServeError::Service(ServiceError::ReplayStore(format!("{CANARY}/replay.cc"))),
|
||||
"replay_store",
|
||||
),
|
||||
(
|
||||
ServeError::Service(ServiceError::InvalidAdmission),
|
||||
"callers",
|
||||
),
|
||||
(
|
||||
ServeError::Service(ServiceError::DuplicateCaller),
|
||||
"callers",
|
||||
),
|
||||
(
|
||||
ServeError::Service(ServiceError::SharedCallerKey),
|
||||
"callers",
|
||||
),
|
||||
(ServeError::Listen(io()), "listener"),
|
||||
];
|
||||
for (error, field) in cases {
|
||||
let shown = ConfigError::from(error);
|
||||
let text = format!("{shown} {shown:?}");
|
||||
assert!(text.contains(field), "{text}");
|
||||
for stored in [CANARY, "4099", "9999", "replay.cc"] {
|
||||
assert!(!text.contains(stored), "{text}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_log_gets_the_size_of_the_childs_stderr_and_never_its_bytes() {
|
||||
let log = Captured::default();
|
||||
let sink = LogSink(Mutex::new(log.clone()));
|
||||
sink.stderr_tail(b"canary-stderr-text, 22 bytes");
|
||||
let written = String::from_utf8(log.0.lock().unwrap().clone()).unwrap();
|
||||
assert_eq!(
|
||||
written,
|
||||
"claude-connector: child stderr, 28 bytes
|
||||
"
|
||||
);
|
||||
assert!(!written.contains("canary"));
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
mod prompt_dir {
|
||||
use std::os::unix::fs::{PermissionsExt, chown, symlink};
|
||||
use std::path::PathBuf;
|
||||
|
||||
use super::super::check_prompt_dir;
|
||||
|
||||
struct Dir(PathBuf);
|
||||
|
||||
impl Dir {
|
||||
fn new(name: &str) -> Self {
|
||||
let path = std::env::temp_dir()
|
||||
.join(format!("claude-prompt-dir-{}-{name}", std::process::id()));
|
||||
let _ = std::fs::remove_dir_all(&path);
|
||||
std::fs::create_dir_all(&path).unwrap();
|
||||
Self(path)
|
||||
}
|
||||
fn mode(&self, mode: u32) {
|
||||
std::fs::set_permissions(&self.0, std::fs::Permissions::from_mode(mode)).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Dir {
|
||||
fn drop(&mut self) {
|
||||
let _ = std::fs::set_permissions(&self.0, std::fs::Permissions::from_mode(0o700));
|
||||
let _ = std::fs::remove_dir_all(&self.0);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn only_mode_0700_passes() {
|
||||
let dir = Dir::new("modes");
|
||||
dir.mode(0o700);
|
||||
assert!(check_prompt_dir(&dir.0).is_ok());
|
||||
// Each of these leaves the owner able to enter, or adds a bit
|
||||
// beyond the owner's three; all are refused.
|
||||
for mode in [
|
||||
0o500, 0o300, 0o100, 0o600, 0o770, 0o707, 0o750, 0o755, 0o1700, 0o2700, 0o4700,
|
||||
] {
|
||||
dir.mode(mode);
|
||||
assert!(check_prompt_dir(&dir.0).is_err(), "{mode:o}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_link_to_a_good_directory_is_refused() {
|
||||
let dir = Dir::new("link-target");
|
||||
dir.mode(0o700);
|
||||
let link =
|
||||
std::env::temp_dir().join(format!("claude-prompt-dir-{}-link", std::process::id()));
|
||||
let _ = std::fs::remove_file(&link);
|
||||
symlink(&dir.0, &link).unwrap();
|
||||
let refused = check_prompt_dir(&link).is_err();
|
||||
std::fs::remove_file(&link).unwrap();
|
||||
assert!(refused);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_directory_owned_by_another_user_is_refused() {
|
||||
// SAFETY: geteuid has no preconditions and cannot fail.
|
||||
if unsafe { libc::geteuid() } != 0 {
|
||||
eprintln!("skipped: only root can give a directory away");
|
||||
return;
|
||||
}
|
||||
let dir = Dir::new("owner");
|
||||
dir.mode(0o700);
|
||||
assert!(check_prompt_dir(&dir.0).is_ok());
|
||||
chown(&dir.0, Some(1234), None).unwrap();
|
||||
assert!(check_prompt_dir(&dir.0).is_err());
|
||||
}
|
||||
}
|
||||
}
|
||||
308
crates/claude-connector/src/handle.rs
Normal file
308
crates/claude-connector/src/handle.rs
Normal file
|
|
@ -0,0 +1,308 @@
|
|||
//! The reasoning state of this backend and the decision it drives: resume the
|
||||
//! CLI session a handle names, or start a fresh one seeded with the caller's
|
||||
//! transcript. The session is a cache of the caller's transcript; the
|
||||
//! transcript is the truth.
|
||||
|
||||
use std::path::PathBuf;
|
||||
|
||||
use hmac::{Hmac, Mac};
|
||||
use model_connector::{InputItem, ProviderRequest};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sha2::Sha256;
|
||||
|
||||
use crate::config::{ClaudeCodeConfig, PROVIDER};
|
||||
|
||||
/// The state this backend hands a caller inside `ReasoningState`. A new
|
||||
/// variant is a new version.
|
||||
#[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.
|
||||
Ended {
|
||||
session_id: String,
|
||||
instance: String,
|
||||
prefix_sha256: [u8; 32],
|
||||
},
|
||||
}
|
||||
|
||||
impl ClaudeCodeHandle {
|
||||
pub(crate) fn encode(&self) -> Vec<u8> {
|
||||
rmp_serde::to_vec_named(self).expect("a handle of strings and bytes encodes")
|
||||
}
|
||||
|
||||
fn decode(state: &[u8]) -> Option<Self> {
|
||||
rmp_serde::from_slice(state).ok()
|
||||
}
|
||||
}
|
||||
|
||||
/// The smallest handle secret the connector accepts, in bytes.
|
||||
pub const MIN_SECRET_BYTES: usize = 32;
|
||||
|
||||
/// Everything a handle's tag binds: the minting caller, the instructions, the
|
||||
/// model, the session and instance the handle names, and the transcript items
|
||||
/// before the handle. Swapping any of them into a valid handle breaks the tag.
|
||||
pub(crate) struct Binding<'a> {
|
||||
pub(crate) caller: &'a str,
|
||||
pub(crate) instructions: &'a str,
|
||||
pub(crate) model: &'a str,
|
||||
pub(crate) session_id: &'a str,
|
||||
pub(crate) instance: &'a str,
|
||||
pub(crate) items: &'a [InputItem],
|
||||
}
|
||||
|
||||
fn prefix_mac(secret: &[u8], binding: &Binding<'_>) -> Hmac<Sha256> {
|
||||
let bytes = rmp_serde::to_vec_named(&(
|
||||
binding.caller,
|
||||
binding.instructions,
|
||||
binding.model,
|
||||
binding.session_id,
|
||||
binding.instance,
|
||||
binding.items,
|
||||
))
|
||||
.expect("strings and bytes encode");
|
||||
let mut mac = Hmac::<Sha256>::new_from_slice(secret).expect("an HMAC key has any length");
|
||||
mac.update(&bytes);
|
||||
mac
|
||||
}
|
||||
|
||||
/// The tag a handle carries: keyed named MessagePack of its [`Binding`].
|
||||
pub(crate) fn prefix_digest(secret: &[u8], binding: &Binding<'_>) -> [u8; 32] {
|
||||
prefix_mac(secret, binding).finalize().into_bytes().into()
|
||||
}
|
||||
|
||||
fn prefix_matches(secret: &[u8], binding: &Binding<'_>, tag: &[u8; 32]) -> bool {
|
||||
prefix_mac(secret, binding).verify_slice(tag).is_ok()
|
||||
}
|
||||
|
||||
/// Whether the CLI's session store still holds a session.
|
||||
pub trait SessionStore: Send + Sync {
|
||||
fn exists(&self, session_id: &str) -> bool;
|
||||
}
|
||||
|
||||
/// Looks for `<config_dir>/projects/*/<session_id>.jsonl`. A session id that
|
||||
/// is not a canonical UUID never names a path.
|
||||
pub struct FsSessionStore {
|
||||
config_dir: PathBuf,
|
||||
}
|
||||
|
||||
impl FsSessionStore {
|
||||
pub fn new(config_dir: PathBuf) -> Self {
|
||||
Self { config_dir }
|
||||
}
|
||||
}
|
||||
|
||||
fn is_canonical_uuid(id: &str) -> bool {
|
||||
id.len() == 36
|
||||
&& id.char_indices().all(|(at, c)| match at {
|
||||
8 | 13 | 18 | 23 => c == '-',
|
||||
_ => c.is_ascii_digit() || ('a'..='f').contains(&c),
|
||||
})
|
||||
}
|
||||
|
||||
impl SessionStore for FsSessionStore {
|
||||
fn exists(&self, session_id: &str) -> bool {
|
||||
if !is_canonical_uuid(session_id) {
|
||||
return false;
|
||||
}
|
||||
let Ok(projects) = std::fs::read_dir(self.config_dir.join("projects")) else {
|
||||
return false;
|
||||
};
|
||||
projects
|
||||
.flatten()
|
||||
.any(|project| project.path().join(format!("{session_id}.jsonl")).is_file())
|
||||
}
|
||||
}
|
||||
|
||||
/// Why a request cannot be answered by this backend; the text is the failure
|
||||
/// message and names what lifts it.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub(crate) enum PlanError {
|
||||
#[error("claude-code: tools are not supported by this backend yet")]
|
||||
Tools,
|
||||
#[error("claude-code: tool results are not supported by this backend yet")]
|
||||
ToolResult,
|
||||
#[error("claude-code: the request ends without user text")]
|
||||
NoUserText,
|
||||
#[error("claude-code: the request carries a reasoning state this backend cannot read")]
|
||||
UndecodableHandle,
|
||||
}
|
||||
|
||||
/// How a call reaches the CLI. `blocks` are the text blocks of the one user
|
||||
/// message the child is given.
|
||||
#[derive(Debug, PartialEq, Eq)]
|
||||
pub(crate) enum Plan {
|
||||
Resume {
|
||||
session_id: String,
|
||||
blocks: Vec<String>,
|
||||
},
|
||||
Seed {
|
||||
blocks: Vec<String>,
|
||||
},
|
||||
}
|
||||
|
||||
impl Plan {
|
||||
pub(crate) fn blocks(&self) -> &[String] {
|
||||
match self {
|
||||
Plan::Resume { blocks, .. } | Plan::Seed { blocks } => blocks,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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);
|
||||
}
|
||||
let input = &request.input;
|
||||
let run_start = input
|
||||
.iter()
|
||||
.rposition(|item| {
|
||||
!matches!(
|
||||
item,
|
||||
InputItem::UserText { .. } | InputItem::ToolResult { .. }
|
||||
)
|
||||
})
|
||||
.map_or(0, |at| at + 1);
|
||||
let run = &input[run_start..];
|
||||
if run
|
||||
.iter()
|
||||
.any(|item| matches!(item, InputItem::ToolResult { .. }))
|
||||
{
|
||||
return Err(PlanError::ToolResult);
|
||||
}
|
||||
let user_blocks: Vec<String> = run
|
||||
.iter()
|
||||
.filter_map(|item| match item {
|
||||
InputItem::UserText { text } => Some(text.clone()),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
if user_blocks.is_empty() {
|
||||
return Err(PlanError::NoUserText);
|
||||
}
|
||||
|
||||
let last_own = input.iter().rposition(
|
||||
|item| matches!(item, InputItem::ReasoningState { provider, .. } if provider == PROVIDER),
|
||||
);
|
||||
if let Some(at) = last_own {
|
||||
let InputItem::ReasoningState { state, .. } = &input[at] else {
|
||||
unreachable!("rposition matched a reasoning state");
|
||||
};
|
||||
let Some(ClaudeCodeHandle::Ended {
|
||||
session_id,
|
||||
instance,
|
||||
prefix_sha256,
|
||||
}) = ClaudeCodeHandle::decode(state)
|
||||
else {
|
||||
return Err(PlanError::UndecodableHandle);
|
||||
};
|
||||
let tail_is_user_text = input[at + 1..]
|
||||
.iter()
|
||||
.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,
|
||||
)
|
||||
&& sessions.exists(&session_id);
|
||||
if resumable {
|
||||
return Ok(Plan::Resume {
|
||||
session_id,
|
||||
blocks: user_blocks,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
let mut blocks = Vec::new();
|
||||
if let Some(transcript) = render_transcript(&input[..run_start]) {
|
||||
blocks.push(transcript);
|
||||
}
|
||||
blocks.extend(user_blocks);
|
||||
Ok(Plan::Seed { blocks })
|
||||
}
|
||||
|
||||
/// The earlier conversation as one deterministic text, or `None` when nothing
|
||||
/// of it is renderable. Reasoning states are omitted.
|
||||
fn render_transcript(items: &[InputItem]) -> Option<String> {
|
||||
let groups: Vec<String> = items
|
||||
.iter()
|
||||
.filter_map(|item| match item {
|
||||
InputItem::UserText { text } => Some(format!("[user]\n{text}")),
|
||||
InputItem::AssistantText { text } => Some(format!("[assistant]\n{text}")),
|
||||
InputItem::ToolCall {
|
||||
call_id,
|
||||
name,
|
||||
arguments,
|
||||
} => Some(format!("[tool_call name={name} id={call_id}]\n{arguments}")),
|
||||
InputItem::ToolResult { call_id, output } => {
|
||||
Some(format!("[tool_result id={call_id}]\n{output}"))
|
||||
}
|
||||
InputItem::ReasoningState { .. } => None,
|
||||
})
|
||||
.collect();
|
||||
if groups.is_empty() {
|
||||
return None;
|
||||
}
|
||||
Some(format!(
|
||||
"The conversation so far, replayed as text because the earlier session cannot be continued:\n\n{}",
|
||||
groups.join("\n\n")
|
||||
))
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct UserLine<'a> {
|
||||
#[serde(rename = "type")]
|
||||
kind: &'static str,
|
||||
message: UserMessage<'a>,
|
||||
parent_tool_use_id: Option<&'a str>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct UserMessage<'a> {
|
||||
role: &'static str,
|
||||
content: Vec<TextBlock<'a>>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct TextBlock<'a> {
|
||||
#[serde(rename = "type")]
|
||||
kind: &'static str,
|
||||
text: &'a str,
|
||||
}
|
||||
|
||||
/// The child's whole stdin: exactly one stream-json user line. Several lines
|
||||
/// would queue several turns.
|
||||
pub(crate) fn stdin_bytes(blocks: &[String]) -> Vec<u8> {
|
||||
let line = UserLine {
|
||||
kind: "user",
|
||||
message: UserMessage {
|
||||
role: "user",
|
||||
content: blocks
|
||||
.iter()
|
||||
.map(|text| TextBlock { kind: "text", text })
|
||||
.collect(),
|
||||
},
|
||||
parent_tool_use_id: None,
|
||||
};
|
||||
let mut bytes = serde_json::to_vec(&line).expect("typed strings encode");
|
||||
bytes.push(b'\n');
|
||||
bytes
|
||||
}
|
||||
215
crates/claude-connector/src/launch.rs
Normal file
215
crates/claude-connector/src/launch.rs
Normal file
|
|
@ -0,0 +1,215 @@
|
|||
//! The claude argv and the child environment: every flag and variable name
|
||||
//! this backend sends is a constant or built here, so a CLI upgrade is one
|
||||
//! place to re-probe.
|
||||
|
||||
use std::fmt;
|
||||
use std::fs::OpenOptions;
|
||||
use std::io::{self, Write};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
|
||||
use model_connector::{ProviderRequest, ReasoningEffort};
|
||||
|
||||
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;
|
||||
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)>);
|
||||
|
||||
impl ChildEnv {
|
||||
pub fn iter(&self) -> impl Iterator<Item = (&str, &str)> {
|
||||
self.0.iter().map(|(name, value)| (*name, value.as_str()))
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for ChildEnv {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_list()
|
||||
.entries(self.0.iter().map(|(name, _)| name))
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn child_env(config: &ClaudeCodeConfig, request: &ProviderRequest) -> 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),
|
||||
(
|
||||
"CLAUDE_CODE_MAX_OUTPUT_TOKENS",
|
||||
request.max_output_tokens.to_string(),
|
||||
),
|
||||
(TOKEN_FD_VAR, TOKEN_FD.to_string()),
|
||||
])
|
||||
}
|
||||
|
||||
/// The system prompt staged for one child: an exclusively created file only
|
||||
/// its owner can read, removed when this value drops, however the call ends.
|
||||
pub(crate) struct PromptFile(PathBuf);
|
||||
|
||||
static PROMPT_FILES: AtomicU64 = AtomicU64::new(0);
|
||||
|
||||
impl PromptFile {
|
||||
pub(crate) fn create(dir: &Path, instructions: &str) -> io::Result<Self> {
|
||||
let name = format!(
|
||||
"claude-connector-prompt-{}-{}",
|
||||
std::process::id(),
|
||||
PROMPT_FILES.fetch_add(1, Ordering::Relaxed)
|
||||
);
|
||||
Self::create_at(dir.join(name), instructions)
|
||||
}
|
||||
|
||||
/// Creates `path` and nothing that already exists there: a file, a link
|
||||
/// (dangling or not) or a directory at the name fails, and is left alone.
|
||||
fn create_at(path: PathBuf, instructions: &str) -> io::Result<Self> {
|
||||
let mut options = OpenOptions::new();
|
||||
options.write(true).create_new(true);
|
||||
#[cfg(unix)]
|
||||
std::os::unix::fs::OpenOptionsExt::mode(&mut options, 0o600);
|
||||
let mut file = options.open(&path)?;
|
||||
let owned = Self(path);
|
||||
file.write_all(instructions.as_bytes())?;
|
||||
Ok(owned)
|
||||
}
|
||||
|
||||
pub(crate) fn path(&self) -> &Path {
|
||||
&self.0
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for PromptFile {
|
||||
fn drop(&mut self) {
|
||||
let _ = std::fs::remove_file(&self.0);
|
||||
}
|
||||
}
|
||||
|
||||
fn effort(effort: ReasoningEffort) -> &'static str {
|
||||
match effort {
|
||||
ReasoningEffort::Low => "low",
|
||||
ReasoningEffort::Medium => "medium",
|
||||
ReasoningEffort::High => "high",
|
||||
ReasoningEffort::Max => "max",
|
||||
}
|
||||
}
|
||||
|
||||
/// The program and its arguments. The system prompt is replaced, by a file
|
||||
/// 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.
|
||||
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(),
|
||||
"-p",
|
||||
"--output-format",
|
||||
"stream-json",
|
||||
"--verbose",
|
||||
"--input-format",
|
||||
"stream-json",
|
||||
"--model",
|
||||
&request.model,
|
||||
"--system-prompt-file",
|
||||
&system_prompt_file.to_string_lossy(),
|
||||
"--tools",
|
||||
"",
|
||||
"--strict-mcp-config",
|
||||
"--setting-sources",
|
||||
"",
|
||||
"--disable-slash-commands",
|
||||
"--permission-prompts",
|
||||
"none",
|
||||
]
|
||||
.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(),
|
||||
session_id.clone(),
|
||||
"--fork-session".to_string(),
|
||||
]);
|
||||
}
|
||||
argv
|
||||
}
|
||||
|
||||
#[cfg(all(test, unix))]
|
||||
mod tests {
|
||||
use std::os::unix::fs::{PermissionsExt, symlink};
|
||||
|
||||
use super::*;
|
||||
|
||||
fn scratch(name: &str) -> PathBuf {
|
||||
let dir = std::env::temp_dir().join(format!(
|
||||
"claude-connector-launch-{name}-{}",
|
||||
std::process::id()
|
||||
));
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
std::fs::create_dir_all(&dir).unwrap();
|
||||
dir
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_prompt_is_created_private_and_removed_on_drop() {
|
||||
let dir = scratch("private");
|
||||
let prompt = PromptFile::create(&dir, "the prompt").unwrap();
|
||||
let path = prompt.path().to_path_buf();
|
||||
assert_eq!(path.parent(), Some(dir.as_path()));
|
||||
assert_eq!(std::fs::read_to_string(&path).unwrap(), "the prompt");
|
||||
let mode = std::fs::metadata(&path).unwrap().permissions().mode() & 0o777;
|
||||
assert_eq!(mode, 0o600);
|
||||
drop(prompt);
|
||||
assert!(!path.exists());
|
||||
std::fs::remove_dir_all(dir).unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_name_already_taken_is_never_written_or_removed() {
|
||||
let dir = scratch("taken");
|
||||
let victim = dir.join("victim");
|
||||
std::fs::write(&victim, "keep").unwrap();
|
||||
|
||||
let file = dir.join("file");
|
||||
std::fs::write(&file, "mine").unwrap();
|
||||
let link = dir.join("link");
|
||||
symlink(&victim, &link).unwrap();
|
||||
let dangling = dir.join("dangling");
|
||||
symlink(dir.join("nowhere"), &dangling).unwrap();
|
||||
|
||||
for taken in [&file, &link, &dangling] {
|
||||
assert!(PromptFile::create_at(taken.clone(), "the prompt").is_err());
|
||||
}
|
||||
assert_eq!(std::fs::read_to_string(&victim).unwrap(), "keep");
|
||||
assert_eq!(std::fs::read_to_string(&file).unwrap(), "mine");
|
||||
assert!(!dir.join("nowhere").exists());
|
||||
assert!(link.exists() && dangling.symlink_metadata().is_ok());
|
||||
std::fs::remove_dir_all(dir).unwrap();
|
||||
}
|
||||
}
|
||||
1006
crates/claude-connector/src/lease.rs
Normal file
1006
crates/claude-connector/src/lease.rs
Normal file
File diff suppressed because it is too large
Load diff
41
crates/claude-connector/src/lib.rs
Normal file
41
crates/claude-connector/src/lib.rs
Normal file
|
|
@ -0,0 +1,41 @@
|
|||
//! The Claude backend of the connector kit. [`ClaudeCodeBackend`] answers each
|
||||
//! call by driving the unmodified `claude` binary once, in print mode with
|
||||
//! stream-json in and out, with every tool, MCP server, settings source and
|
||||
//! plugin switched off. A binary wires it to [`model_connector::serve`].
|
||||
|
||||
pub mod ask;
|
||||
mod backend;
|
||||
mod config;
|
||||
mod daemon;
|
||||
mod handle;
|
||||
mod launch;
|
||||
#[cfg(target_os = "linux")]
|
||||
mod lease;
|
||||
mod meter;
|
||||
mod process;
|
||||
mod quota;
|
||||
mod stream;
|
||||
#[cfg(target_os = "linux")]
|
||||
mod unit;
|
||||
|
||||
pub use backend::{ClaudeCodeBackend, CredentialTokenSource, StderrSink, TokenSource};
|
||||
pub use config::{ClaudeCodeConfig, ClaudeCodeLimits};
|
||||
pub use daemon::{
|
||||
ClaudeConnectorConfig, ConfigError, LeaseCallerConfig, check_ready, load_config, write_config,
|
||||
};
|
||||
#[cfg(unix)]
|
||||
pub use daemon::{check_prompt_dir, serve_config};
|
||||
pub use handle::{FsSessionStore, MIN_SECRET_BYTES, SessionStore};
|
||||
pub use launch::ChildEnv;
|
||||
#[cfg(target_os = "linux")]
|
||||
pub use lease::{
|
||||
LeaseCaller, LeaseGranted, LeaseLimits, LeaseMessage, LeaseRefusalKind, LeaseRefused,
|
||||
LeaseRequest, LeaseUsageReadoutRequest, 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;
|
||||
460
crates/claude-connector/src/main.rs
Normal file
460
crates/claude-connector/src/main.rs
Normal file
|
|
@ -0,0 +1,460 @@
|
|||
//! The claude-connector binary.
|
||||
//!
|
||||
//! claude-connector serve --config PATH.cc
|
||||
//! claude-connector config init PATH.cc --flag value ...
|
||||
//! claude-connector config add-caller PATH.cc --caller-id ID --connection-key-file F ...
|
||||
//! claude-connector config add-lease-caller PATH.cc --lease-user U --caller-id ID ...
|
||||
//! claude-connector config check PATH.cc
|
||||
//! claude-connector ask --caller-id ID --connection-key-file F --model M (--system TEXT | --system-file F) < prompt
|
||||
//!
|
||||
//! `serve` takes its configuration from the typed document alone. The flags of
|
||||
//! `config init` exist only to author that document. No message here echoes a
|
||||
//! value that was typed or stored; refusals name a field or a flag.
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
use std::path::PathBuf;
|
||||
use std::process::ExitCode;
|
||||
|
||||
use claude_connector::ask::{self, AskArgs, AskError};
|
||||
use claude_connector::{
|
||||
ClaudeConnectorConfig, ConfigError, LeaseCallerConfig, MIN_SECRET_BYTES, check_ready,
|
||||
load_config, write_config,
|
||||
};
|
||||
use model_connector::{CallerConfig, LeaseTurn, frame_bytes_for_payload};
|
||||
|
||||
fn main() -> ExitCode {
|
||||
let args: Vec<String> = std::env::args().skip(1).collect();
|
||||
let args: Vec<&str> = args.iter().map(String::as_str).collect();
|
||||
let result = match args.as_slice() {
|
||||
["serve", "--config", path] => serve(path),
|
||||
["config", "init", path, flags @ ..] => init(path, flags),
|
||||
["config", "add-caller", path, flags @ ..] => add_caller(path, flags),
|
||||
["config", "add-lease-caller", path, flags @ ..] => add_lease_caller(path, flags),
|
||||
["config", "check", path] => check(path),
|
||||
["ask", flags @ ..] => run_ask(flags),
|
||||
_ => Err(Failure::Usage),
|
||||
};
|
||||
match result {
|
||||
Ok(()) => ExitCode::SUCCESS,
|
||||
Err(failure) => {
|
||||
eprintln!("claude-connector: {failure}");
|
||||
ExitCode::FAILURE
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
enum Failure {
|
||||
#[error(
|
||||
"usage: serve --config PATH | config init PATH --flag value ... | config add-caller PATH --flag value ... | config add-lease-caller PATH --flag value ... | config check PATH | ask --caller-id ID --connection-key-file FILE --model MODEL (--system TEXT | --system-file FILE) [--endpoint HOST:PORT] [--effort low|medium|high|max] [--max-output-tokens N] [--timeout-secs N] (prompt on stdin, answer on stdout) (check reads the credentials from the environment it runs in: run it in the unit's environment)"
|
||||
)]
|
||||
Usage,
|
||||
#[error(transparent)]
|
||||
Config(#[from] ConfigError),
|
||||
#[error("flag {0} is missing")]
|
||||
Missing(&'static str),
|
||||
#[error("flag {0} is malformed")]
|
||||
Malformed(&'static str),
|
||||
#[error("a flag is unknown, repeated or has no value")]
|
||||
Flags,
|
||||
#[error("{0}")]
|
||||
Ask(#[from] AskError),
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
fn serve(path: &str) -> Result<(), Failure> {
|
||||
Ok(claude_connector::serve_config(path.as_ref())?)
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
fn serve(_: &str) -> Result<(), Failure> {
|
||||
Err(Failure::Usage)
|
||||
}
|
||||
|
||||
fn check(path: &str) -> Result<(), Failure> {
|
||||
let config = load_config(path.as_ref())?;
|
||||
check_ready(&config)?;
|
||||
println!("claude connector configuration is valid and ready to serve");
|
||||
println!("callers: {}", config.callers.len());
|
||||
println!("lease_callers: {}", config.lease_callers.len());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
const ASK: [&str; 9] = [
|
||||
"--endpoint",
|
||||
"--caller-id",
|
||||
"--connection-key-file",
|
||||
"--model",
|
||||
"--system",
|
||||
"--system-file",
|
||||
"--effort",
|
||||
"--max-output-tokens",
|
||||
"--timeout-secs",
|
||||
];
|
||||
|
||||
/// One no-tools completion: the prompt is all of stdin, the answer is stdout.
|
||||
fn run_ask(args: &[&str]) -> Result<(), Failure> {
|
||||
use std::io::{Read, Write};
|
||||
|
||||
let flags = Flags::parse(args, &[&ASK], &[])?;
|
||||
let instructions = match (
|
||||
flags.single.get("--system"),
|
||||
flags.single.get("--system-file"),
|
||||
) {
|
||||
(Some(text), None) => text.clone(),
|
||||
(None, Some(path)) => ask::read_system_file(std::path::Path::new(path))?,
|
||||
_ => return Err(Failure::Missing("--system or --system-file (exactly one)")),
|
||||
};
|
||||
let effort = match flags.single.get("--effort").map(String::as_str) {
|
||||
None => None,
|
||||
Some("low") => Some(model_connector::ReasoningEffort::Low),
|
||||
Some("medium") => Some(model_connector::ReasoningEffort::Medium),
|
||||
Some("high") => Some(model_connector::ReasoningEffort::High),
|
||||
Some("max") => Some(model_connector::ReasoningEffort::Max),
|
||||
Some(_) => return Err(Failure::Malformed("--effort")),
|
||||
};
|
||||
let endpoint = flags
|
||||
.single
|
||||
.get("--endpoint")
|
||||
.map_or(ask::DEFAULT_ENDPOINT, String::as_str)
|
||||
.parse()
|
||||
.map_err(|_| Failure::Malformed("--endpoint"))?;
|
||||
let mut raw = Vec::new();
|
||||
std::io::stdin()
|
||||
.take(ask::PAYLOAD_BYTES as u64 + 1)
|
||||
.read_to_end(&mut raw)
|
||||
.map_err(|_| AskError::StdinRead)?;
|
||||
let prompt = String::from_utf8(raw).map_err(|_| AskError::StdinNotText)?;
|
||||
let answer = ask::ask(&AskArgs {
|
||||
endpoint,
|
||||
caller_id: flags.text("--caller-id")?,
|
||||
key_file: flags.path("--connection-key-file")?,
|
||||
model: flags.text("--model")?,
|
||||
instructions,
|
||||
effort,
|
||||
max_output_tokens: flags.number_or("--max-output-tokens", 8192)?,
|
||||
timeout: std::time::Duration::from_secs(flags.number_or("--timeout-secs", 300)?),
|
||||
prompt,
|
||||
})?;
|
||||
let mut out = std::io::stdout();
|
||||
out.write_all(answer.as_bytes())
|
||||
.and_then(|()| out.flush())
|
||||
.map_err(|_| AskError::StdoutWrite.into())
|
||||
}
|
||||
|
||||
/// The flags of one config verb: each single flag at most once, each list
|
||||
/// flag any number of times.
|
||||
struct Flags {
|
||||
single: BTreeMap<&'static str, String>,
|
||||
lists: BTreeMap<&'static str, Vec<String>>,
|
||||
}
|
||||
|
||||
/// The flags that describe one caller, shared by `config init` (which writes
|
||||
/// the first caller) and `config add-caller`.
|
||||
const CALLER: [&str; 5] = [
|
||||
"--caller-id",
|
||||
"--connection-key-file",
|
||||
"--connection-key-epoch",
|
||||
"--max-concurrent-requests",
|
||||
"--max-output-tokens",
|
||||
];
|
||||
|
||||
const INIT: [&str; 21] = [
|
||||
"--listen",
|
||||
"--replay-store",
|
||||
"--claude-binary",
|
||||
"--claude-sha256",
|
||||
"--work-dir",
|
||||
"--config-dir",
|
||||
"--prompt-dir",
|
||||
"--instance",
|
||||
"--token-credential",
|
||||
"--handle-secret-credential",
|
||||
"--handle-secret-file",
|
||||
"--deadline-ms",
|
||||
"--max-line-bytes",
|
||||
"--max-total-bytes",
|
||||
"--stderr-tail-bytes",
|
||||
"--max-payload-bytes",
|
||||
"--socket-timeout-ms",
|
||||
"--lease-socket",
|
||||
"--lease-max-ms",
|
||||
"--turn",
|
||||
"--lease-unit-instance",
|
||||
];
|
||||
|
||||
const LEASE_CALLER: [&str; 3] = ["--caller-id", "--turn", "--lease-unit-instance"];
|
||||
|
||||
impl Flags {
|
||||
fn parse(
|
||||
args: &[&str],
|
||||
singles: &[&[&'static str]],
|
||||
lists: &[&'static str],
|
||||
) -> Result<Self, Failure> {
|
||||
let mut flags = Flags {
|
||||
single: BTreeMap::new(),
|
||||
lists: BTreeMap::new(),
|
||||
};
|
||||
let mut args = args.iter();
|
||||
while let Some(&name) = args.next() {
|
||||
let Some(&value) = args.next() else {
|
||||
return Err(Failure::Flags);
|
||||
};
|
||||
if let Some(list) = lists.iter().find(|known| **known == name) {
|
||||
flags
|
||||
.lists
|
||||
.entry(*list)
|
||||
.or_default()
|
||||
.push(value.to_string());
|
||||
continue;
|
||||
}
|
||||
let Some(known) = singles
|
||||
.iter()
|
||||
.flat_map(|set| set.iter())
|
||||
.find(|k| **k == name)
|
||||
else {
|
||||
return Err(Failure::Flags);
|
||||
};
|
||||
if flags.single.insert(*known, value.to_string()).is_some() {
|
||||
return Err(Failure::Flags);
|
||||
}
|
||||
}
|
||||
Ok(flags)
|
||||
}
|
||||
|
||||
fn list(&self, name: &'static str) -> &[String] {
|
||||
self.lists.get(name).map_or(&[], Vec::as_slice)
|
||||
}
|
||||
|
||||
fn text(&self, name: &'static str) -> Result<String, Failure> {
|
||||
self.single.get(name).cloned().ok_or(Failure::Missing(name))
|
||||
}
|
||||
|
||||
fn path(&self, name: &'static str) -> Result<PathBuf, Failure> {
|
||||
self.text(name).map(PathBuf::from)
|
||||
}
|
||||
|
||||
fn number<T: std::str::FromStr>(&self, name: &'static str) -> Result<T, Failure> {
|
||||
self.text(name)?
|
||||
.parse()
|
||||
.map_err(|_| Failure::Malformed(name))
|
||||
}
|
||||
|
||||
/// The kind of turn a leasing account serves: `self` or `persona`.
|
||||
fn turn(&self) -> Result<LeaseTurn, Failure> {
|
||||
match self.text("--turn")?.as_str() {
|
||||
"self" => Ok(LeaseTurn::SelfRun),
|
||||
"persona" => Ok(LeaseTurn::PersonaTurn),
|
||||
_ => Err(Failure::Malformed("--turn")),
|
||||
}
|
||||
}
|
||||
|
||||
fn number_or<T: std::str::FromStr>(
|
||||
&self,
|
||||
name: &'static str,
|
||||
default: T,
|
||||
) -> Result<T, Failure> {
|
||||
match self.single.get(name) {
|
||||
None => Ok(default),
|
||||
Some(text) => text.parse().map_err(|_| Failure::Malformed(name)),
|
||||
}
|
||||
}
|
||||
|
||||
/// The caller these flags describe. Its payload bound comes from
|
||||
/// `--max-payload-bytes` when the verb takes one.
|
||||
fn caller(&self) -> Result<CallerConfig, Failure> {
|
||||
Ok(CallerConfig {
|
||||
caller_runtime_id: self.text("--caller-id")?,
|
||||
connection_key_file: self.path("--connection-key-file")?,
|
||||
connection_key_epoch: self.number("--connection-key-epoch")?,
|
||||
allowed_models: self.list("--allowed-model").to_vec(),
|
||||
max_concurrent_requests: self.number_or("--max-concurrent-requests", 4)?,
|
||||
max_payload_bytes: self.number_or("--max-payload-bytes", 1 << 20)?,
|
||||
max_output_tokens: self.number_or("--max-output-tokens", 32_768)?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn sha256(hex: &str) -> Option<[u8; 32]> {
|
||||
if hex.len() != 64 || !hex.is_ascii() {
|
||||
return None;
|
||||
}
|
||||
let mut digest = [0_u8; 32];
|
||||
for (byte, pair) in digest.iter_mut().zip(hex.as_bytes().chunks(2)) {
|
||||
*byte = u8::from_str_radix(std::str::from_utf8(pair).ok()?, 16).ok()?;
|
||||
}
|
||||
Some(digest)
|
||||
}
|
||||
|
||||
/// Writes 32 random bytes as 64 hex characters, mode 0600, unless the file is
|
||||
/// already there: the secret is never read back or shown, and a file that
|
||||
/// exists is the secret handles already depend on, so it is kept as it is when
|
||||
/// it is a regular file, mode 0600, owned by this user and at least
|
||||
/// `MIN_SECRET_BYTES` long, and refused untouched otherwise.
|
||||
#[cfg(unix)]
|
||||
fn ensure_secret_file(path: &std::path::Path) -> Result<(), Failure> {
|
||||
use std::io::{Read, Write};
|
||||
use std::os::unix::fs::OpenOptionsExt;
|
||||
|
||||
if !path.is_absolute() {
|
||||
return Err(Failure::Malformed("--handle-secret-file"));
|
||||
}
|
||||
let mut file = match std::fs::OpenOptions::new()
|
||||
.write(true)
|
||||
.create_new(true)
|
||||
.mode(0o600)
|
||||
.open(path)
|
||||
{
|
||||
Ok(file) => file,
|
||||
Err(_) => return existing_secret_is_usable(path),
|
||||
};
|
||||
let mut bytes = [0_u8; 32];
|
||||
std::fs::File::open("/dev/urandom")
|
||||
.and_then(|mut random| random.read_exact(&mut bytes))
|
||||
.map_err(|_| Failure::Malformed("--handle-secret-file"))?;
|
||||
let hex: String = bytes.iter().map(|byte| format!("{byte:02x}")).collect();
|
||||
file.write_all(hex.as_bytes())
|
||||
.map_err(|_| Failure::Malformed("--handle-secret-file"))
|
||||
}
|
||||
|
||||
/// An existing secret file is accepted only as the connector would use it: not
|
||||
/// a link, a regular file, readable by no one else, ours and long enough.
|
||||
#[cfg(unix)]
|
||||
fn existing_secret_is_usable(path: &std::path::Path) -> Result<(), Failure> {
|
||||
let metadata =
|
||||
std::fs::symlink_metadata(path).map_err(|_| Failure::Malformed("--handle-secret-file"))?;
|
||||
// SAFETY: geteuid has no preconditions and cannot fail.
|
||||
let uid = unsafe { libc::geteuid() };
|
||||
if secret_file_usable(&metadata, uid) {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(Failure::Malformed("--handle-secret-file"))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
fn secret_file_usable(metadata: &std::fs::Metadata, uid: u32) -> bool {
|
||||
use std::os::unix::fs::{MetadataExt, PermissionsExt};
|
||||
|
||||
metadata.is_file()
|
||||
&& metadata.permissions().mode() & 0o7777 == 0o600
|
||||
&& metadata.uid() == uid
|
||||
&& metadata.len() >= MIN_SECRET_BYTES as u64
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
fn ensure_secret_file(_: &std::path::Path) -> Result<(), Failure> {
|
||||
Err(Failure::Usage)
|
||||
}
|
||||
|
||||
fn init(path: &str, args: &[&str]) -> Result<(), Failure> {
|
||||
let flags = Flags::parse(
|
||||
args,
|
||||
&[&INIT, &CALLER],
|
||||
&["--allowed-model", "--lease-user"],
|
||||
)?;
|
||||
let first = flags.caller()?;
|
||||
let lease_callers = flags
|
||||
.list("--lease-user")
|
||||
.iter()
|
||||
.map(|unix_user| {
|
||||
Ok(LeaseCallerConfig {
|
||||
unix_user: unix_user.clone(),
|
||||
caller_id: first.caller_runtime_id.clone(),
|
||||
turn: flags.turn()?,
|
||||
unit_instance: flags.text("--lease-unit-instance")?,
|
||||
})
|
||||
})
|
||||
.collect::<Result<Vec<_>, Failure>>()?;
|
||||
let secret_file = flags.path("--handle-secret-file")?;
|
||||
let config = ClaudeConnectorConfig {
|
||||
listen: flags.text("--listen")?,
|
||||
replay_store: flags.path("--replay-store")?,
|
||||
claude_binary: flags.path("--claude-binary")?,
|
||||
claude_binary_sha256: sha256(&flags.text("--claude-sha256")?)
|
||||
.ok_or(Failure::Malformed("--claude-sha256"))?,
|
||||
work_dir: flags.path("--work-dir")?,
|
||||
config_dir: flags.path("--config-dir")?,
|
||||
prompt_dir: flags.path("--prompt-dir")?,
|
||||
instance: flags.text("--instance")?,
|
||||
token_credential: flags.text("--token-credential")?,
|
||||
deadline_ms: flags.number_or("--deadline-ms", 600_000)?,
|
||||
max_line_bytes: flags.number_or("--max-line-bytes", 4 << 20)?,
|
||||
max_total_bytes: flags.number_or("--max-total-bytes", 64 << 20)?,
|
||||
stderr_tail_bytes: flags.number_or("--stderr-tail-bytes", 8192)?,
|
||||
max_frame_bytes: frame_bytes_for_payload(first.max_payload_bytes),
|
||||
max_connections: first.max_concurrent_requests.max(8),
|
||||
socket_timeout_ms: flags.number_or("--socket-timeout-ms", 10_000)?,
|
||||
max_expiry_skew_ms: 300_000,
|
||||
callers: vec![first],
|
||||
lease_socket: flags.single.get("--lease-socket").map(PathBuf::from),
|
||||
lease_callers,
|
||||
lease_max_ms: flags.number_or("--lease-max-ms", 3_600_000)?,
|
||||
handle_secret_credential: flags.text("--handle-secret-credential")?,
|
||||
};
|
||||
// Validate before the secret file exists: a refused document leaves nothing.
|
||||
config.validate()?;
|
||||
ensure_secret_file(&secret_file)?;
|
||||
write_config(path.as_ref(), &config)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Appends one caller to an existing document. The serve limits grow to carry
|
||||
/// it: the frame to its payload bound, the connections to every caller's slots.
|
||||
fn add_caller(path: &str, args: &[&str]) -> Result<(), Failure> {
|
||||
let flags = Flags::parse(
|
||||
args,
|
||||
&[&CALLER, &["--max-payload-bytes"]],
|
||||
&["--allowed-model"],
|
||||
)?;
|
||||
let caller = flags.caller()?;
|
||||
let mut config = load_config(path.as_ref())?;
|
||||
config.max_frame_bytes = config
|
||||
.max_frame_bytes
|
||||
.max(frame_bytes_for_payload(caller.max_payload_bytes));
|
||||
config.callers.push(caller);
|
||||
let slots: usize = config
|
||||
.callers
|
||||
.iter()
|
||||
.map(|caller| caller.max_concurrent_requests)
|
||||
.sum();
|
||||
config.max_connections = config.max_connections.max(slots);
|
||||
write_config(path.as_ref(), &config)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Appends one leasing account to an existing document that lends a token.
|
||||
fn add_lease_caller(path: &str, args: &[&str]) -> Result<(), Failure> {
|
||||
let flags = Flags::parse(args, &[&LEASE_CALLER, &["--lease-user"]], &[])?;
|
||||
let mut config = load_config(path.as_ref())?;
|
||||
config.lease_callers.push(LeaseCallerConfig {
|
||||
unix_user: flags.text("--lease-user")?,
|
||||
caller_id: flags.text("--caller-id")?,
|
||||
turn: flags.turn()?,
|
||||
unit_instance: flags.text("--lease-unit-instance")?,
|
||||
});
|
||||
write_config(path.as_ref(), &config)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(all(test, unix))]
|
||||
mod tests {
|
||||
use std::os::unix::fs::{MetadataExt, PermissionsExt};
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn a_secret_file_owned_by_someone_else_is_not_usable() {
|
||||
let dir =
|
||||
std::env::temp_dir().join(format!("claude-connector-main-{}", std::process::id()));
|
||||
std::fs::create_dir_all(&dir).unwrap();
|
||||
let path = dir.join("handle.secret");
|
||||
std::fs::write(&path, "k".repeat(64)).unwrap();
|
||||
std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o600)).unwrap();
|
||||
let metadata = std::fs::symlink_metadata(&path).unwrap();
|
||||
let owner = metadata.uid();
|
||||
assert!(secret_file_usable(&metadata, owner));
|
||||
assert!(!secret_file_usable(&metadata, owner.wrapping_add(1)));
|
||||
std::fs::remove_dir_all(&dir).unwrap();
|
||||
}
|
||||
}
|
||||
651
crates/claude-connector/src/meter.rs
Normal file
651
crates/claude-connector/src/meter.rs
Normal file
|
|
@ -0,0 +1,651 @@
|
|||
//! 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::{PlanQuota, 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. A refresh is logged every time, with the reading
|
||||
/// it saw; the other states are logged once per change.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
enum State {
|
||||
/// The week window as the refresh left it; a window with no reading is
|
||||
/// `None`.
|
||||
Read {
|
||||
permille: Option<u32>,
|
||||
resets_at_unix_ms: u64,
|
||||
},
|
||||
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 {
|
||||
permille,
|
||||
resets_at_unix_ms,
|
||||
} => {
|
||||
let permille = permille.map_or("unknown".to_string(), |p| p.to_string());
|
||||
format!(
|
||||
"meter: weekly reading refreshed, permille {permille} resets_at_unix_ms {resets_at_unix_ms}"
|
||||
)
|
||||
}
|
||||
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 let Some(week) = self.week_refreshed_since(now_ms) {
|
||||
self.misses = 0;
|
||||
State::Read {
|
||||
permille: week.utilization_permille,
|
||||
resets_at_unix_ms: week.resets_at_unix_ms,
|
||||
}
|
||||
} 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)
|
||||
}
|
||||
|
||||
/// The week window if it was observed at or after `since_ms`, from one
|
||||
/// snapshot, so the logged value is the one that decided the refresh.
|
||||
fn week_refreshed_since(&self, since_ms: u64) -> Option<PlanQuota> {
|
||||
self.cell
|
||||
.snapshot()
|
||||
.windows
|
||||
.into_iter()
|
||||
.find(|quota| quota.window == WEEK && quota.observed_at_unix_ms >= since_ms)
|
||||
}
|
||||
|
||||
fn report(&mut self, state: State) {
|
||||
if matches!(state, State::Read { .. }) || 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 per refresh of the weekly reading and one per
|
||||
/// change of any other state. 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 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>,
|
||||
/// The used thousandths the reading reports; 60 when unset.
|
||||
permille: Option<u32>,
|
||||
/// The reset the reading reports; `RESET` when unset.
|
||||
resets_at: Option<u64>,
|
||||
}
|
||||
|
||||
/// 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 {
|
||||
let mut week = reading(&script.status);
|
||||
week[0].utilization_permille = Some(script.permille.unwrap_or(60));
|
||||
week[0].resets_at_unix_ms = script.resets_at.unwrap_or(RESET);
|
||||
self.cell.observe_at(at, week);
|
||||
}
|
||||
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,
|
||||
..Script::default()
|
||||
};
|
||||
}
|
||||
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()),
|
||||
..Script::default()
|
||||
};
|
||||
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![format!(
|
||||
"meter: weekly reading refreshed, permille 60 resets_at_unix_ms {RESET}"
|
||||
)]
|
||||
);
|
||||
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 a_refresh_logs_the_reading_and_an_unchanged_value_logs_again() {
|
||||
let mut rig = rig();
|
||||
rig.reads(NOW);
|
||||
rig.meter.tick(NOW);
|
||||
rig.reads(NOW + 15 * MIN);
|
||||
rig.meter.tick(NOW + 15 * MIN);
|
||||
let line =
|
||||
format!("meter: weekly reading refreshed, permille 60 resets_at_unix_ms {RESET}");
|
||||
assert_eq!(rig.lines(), vec![line.clone(), line]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_logged_value_follows_the_reading_and_an_expired_week_is_unknown() {
|
||||
let mut rig = rig();
|
||||
for (at, permille) in [(NOW, 60), (NOW + 15 * MIN, 75)] {
|
||||
*rig.script.lock().unwrap() = Script {
|
||||
reading_at: Some(at),
|
||||
status: "allowed".to_string(),
|
||||
permille: Some(permille),
|
||||
resets_at: Some(RESET + permille as u64),
|
||||
..Script::default()
|
||||
};
|
||||
rig.meter.tick(at);
|
||||
}
|
||||
let mut expired = self::rig();
|
||||
let now = unix_ms();
|
||||
*expired.script.lock().unwrap() = Script {
|
||||
reading_at: Some(now),
|
||||
status: "allowed".to_string(),
|
||||
resets_at: Some(1),
|
||||
..Script::default()
|
||||
};
|
||||
expired.meter.tick(now);
|
||||
assert_eq!(
|
||||
expired.lines(),
|
||||
vec!["meter: weekly reading refreshed, permille unknown resets_at_unix_ms 1"]
|
||||
);
|
||||
assert_eq!(
|
||||
rig.lines(),
|
||||
vec![
|
||||
format!(
|
||||
"meter: weekly reading refreshed, permille 60 resets_at_unix_ms {}",
|
||||
RESET + 60
|
||||
),
|
||||
format!(
|
||||
"meter: weekly reading refreshed, permille 75 resets_at_unix_ms {}",
|
||||
RESET + 75
|
||||
),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn idle_logs_once_per_change_however_often_it_repeats() {
|
||||
let mut rig = rig();
|
||||
rig.lease = None;
|
||||
for step in 0..3 {
|
||||
rig.meter.tick(NOW + 60 * MIN * (step + 1));
|
||||
}
|
||||
assert_eq!(rig.lines(), vec![State::Idle.line()]);
|
||||
}
|
||||
|
||||
#[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 {
|
||||
permille: None,
|
||||
resets_at_unix_ms: 7
|
||||
}
|
||||
.line(),
|
||||
"meter: weekly reading refreshed, permille unknown resets_at_unix_ms 7"
|
||||
);
|
||||
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,
|
||||
..Script::default()
|
||||
})),
|
||||
};
|
||||
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);
|
||||
}
|
||||
}
|
||||
362
crates/claude-connector/src/process.rs
Normal file
362
crates/claude-connector/src/process.rs
Normal file
|
|
@ -0,0 +1,362 @@
|
|||
//! The one OS seam: starting the child and reading its output.
|
||||
|
||||
use std::io;
|
||||
use std::path::Path;
|
||||
use std::time::Instant;
|
||||
|
||||
use crate::launch::ChildEnv;
|
||||
|
||||
/// What one read of the child's stdout produced.
|
||||
#[derive(Debug, PartialEq, Eq)]
|
||||
pub enum StdoutRead {
|
||||
/// One non-empty line, without its terminator.
|
||||
Line(Vec<u8>),
|
||||
/// The child closed its stdout.
|
||||
Eof,
|
||||
/// The deadline passed first.
|
||||
Timeout,
|
||||
/// A line or the total output passed its cap.
|
||||
Breach,
|
||||
}
|
||||
|
||||
/// 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.
|
||||
pub trait ProcessPort: Send + Sync {
|
||||
fn spawn(
|
||||
&self,
|
||||
argv: &[String],
|
||||
env: &ChildEnv,
|
||||
cwd: &Path,
|
||||
stdin: Vec<u8>,
|
||||
token: &str,
|
||||
) -> io::Result<Box<dyn RunningChild>>;
|
||||
}
|
||||
|
||||
/// A started child. The backend ends every child with `kill` then `wait`.
|
||||
pub trait RunningChild: Send {
|
||||
fn read_line(&mut self, deadline: Instant) -> StdoutRead;
|
||||
/// The last bytes the child wrote to stderr.
|
||||
fn stderr_tail(&self) -> Vec<u8>;
|
||||
/// Kills the child and everything it started.
|
||||
fn kill(&mut self);
|
||||
fn wait(&mut self);
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
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};
|
||||
use std::sync::mpsc::{Receiver, RecvTimeoutError, channel};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::thread;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
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(())
|
||||
}
|
||||
|
||||
/// Spawns each child in its own process group, so that killing the call
|
||||
/// kills what the child started too.
|
||||
pub struct StdProcessPort {
|
||||
limits: ClaudeCodeLimits,
|
||||
}
|
||||
|
||||
impl StdProcessPort {
|
||||
pub fn new(limits: ClaudeCodeLimits) -> Self {
|
||||
Self { limits }
|
||||
}
|
||||
}
|
||||
|
||||
impl ProcessPort for StdProcessPort {
|
||||
fn spawn(
|
||||
&self,
|
||||
argv: &[String],
|
||||
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))?;
|
||||
let mut command = Command::new(program);
|
||||
command
|
||||
.args(args)
|
||||
.env_clear()
|
||||
.envs(env.iter())
|
||||
.current_dir(cwd)
|
||||
.stdin(Stdio::piped())
|
||||
.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());
|
||||
|
||||
if let Some(mut pipe) = stdin_pipe {
|
||||
thread::spawn(move || {
|
||||
let _ = pipe.write_all(&stdin);
|
||||
});
|
||||
}
|
||||
|
||||
let (lines, receiver) = channel();
|
||||
if let Some(stdout) = stdout {
|
||||
let limits = self.limits;
|
||||
thread::spawn(move || {
|
||||
let mut reader = BufReader::new(stdout);
|
||||
let mut total = 0usize;
|
||||
loop {
|
||||
let mut line = Vec::new();
|
||||
let read = reader
|
||||
.by_ref()
|
||||
.take(limits.max_line_bytes as u64 + 1)
|
||||
.read_until(b'\n', &mut line);
|
||||
let Ok(read) = read else { break };
|
||||
if read == 0 {
|
||||
break;
|
||||
}
|
||||
total = total.saturating_add(read);
|
||||
let terminated = line.last() == Some(&b'\n');
|
||||
if terminated {
|
||||
line.pop();
|
||||
}
|
||||
if line.last() == Some(&b'\r') {
|
||||
line.pop();
|
||||
}
|
||||
if line.len() > limits.max_line_bytes || total > limits.max_total_bytes {
|
||||
let _ = lines.send(Piece::Breach);
|
||||
return;
|
||||
}
|
||||
if !line.is_empty() && lines.send(Piece::Line(line)).is_err() {
|
||||
return;
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
let tail = Arc::new(Mutex::new(Vec::new()));
|
||||
let stderr_reader = stderr.map(|mut stderr| {
|
||||
let (tail, keep) = (tail.clone(), self.limits.stderr_tail_bytes);
|
||||
thread::spawn(move || {
|
||||
let mut chunk = [0u8; 4096];
|
||||
while let Ok(read) = stderr.read(&mut chunk) {
|
||||
if read == 0 {
|
||||
break;
|
||||
}
|
||||
let mut held = tail.lock().unwrap_or_else(|e| e.into_inner());
|
||||
held.extend_from_slice(&chunk[..read]);
|
||||
let excess = held.len().saturating_sub(keep);
|
||||
held.drain(..excess);
|
||||
}
|
||||
})
|
||||
});
|
||||
|
||||
Ok(Box::new(StdChild {
|
||||
child,
|
||||
receiver,
|
||||
tail,
|
||||
stderr_reader,
|
||||
reaped: false,
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
enum Piece {
|
||||
Line(Vec<u8>),
|
||||
Breach,
|
||||
}
|
||||
|
||||
struct StdChild {
|
||||
child: Child,
|
||||
receiver: Receiver<Piece>,
|
||||
tail: Arc<Mutex<Vec<u8>>>,
|
||||
stderr_reader: Option<thread::JoinHandle<()>>,
|
||||
reaped: bool,
|
||||
}
|
||||
|
||||
impl RunningChild for StdChild {
|
||||
fn read_line(&mut self, deadline: Instant) -> StdoutRead {
|
||||
match self
|
||||
.receiver
|
||||
.recv_timeout(deadline.saturating_duration_since(Instant::now()))
|
||||
{
|
||||
Ok(Piece::Line(line)) => StdoutRead::Line(line),
|
||||
Ok(Piece::Breach) => StdoutRead::Breach,
|
||||
Err(RecvTimeoutError::Timeout) => StdoutRead::Timeout,
|
||||
Err(RecvTimeoutError::Disconnected) => StdoutRead::Eof,
|
||||
}
|
||||
}
|
||||
|
||||
fn stderr_tail(&self) -> Vec<u8> {
|
||||
self.tail.lock().unwrap_or_else(|e| e.into_inner()).clone()
|
||||
}
|
||||
|
||||
fn kill(&mut self) {
|
||||
// The group leader stays a zombie until `wait`, so its pid is
|
||||
// still ours to signal.
|
||||
if !self.reaped {
|
||||
// SAFETY: killpg takes a process group id and no pointers.
|
||||
unsafe {
|
||||
libc::killpg(self.child.id() as libc::pid_t, libc::SIGKILL);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn wait(&mut self) {
|
||||
let _ = self.child.wait();
|
||||
self.reaped = true;
|
||||
// Every writer to stderr died with the group, so the tail is
|
||||
// complete once its reader sees the end. A process that escaped
|
||||
// the group is not waited for.
|
||||
if let Some(reader) = self.stderr_reader.take() {
|
||||
for _ in 0..100 {
|
||||
if reader.is_finished() {
|
||||
let _ = reader.join();
|
||||
break;
|
||||
}
|
||||
thread::sleep(Duration::from_millis(2));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for StdChild {
|
||||
fn drop(&mut self) {
|
||||
self.kill();
|
||||
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());
|
||||
}
|
||||
}
|
||||
}
|
||||
405
crates/claude-connector/src/quota.rs
Normal file
405
crates/claude-connector/src/quota.rs
Normal file
|
|
@ -0,0 +1,405 @@
|
|||
//! The plan's own usage windows, as the subscription reports them. The
|
||||
//! subscription decides how much a plan may burn, and one sensor observes it:
|
||||
//! the CLI's `rate_limit_event` on the connector's own calls, which the
|
||||
//! `meter` makes when the week's reading goes stale. The latest observation of
|
||||
//! each window is kept here and published as the usage readout and lent with
|
||||
//! each lease.
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::{Arc, Condvar, Mutex};
|
||||
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use model_connector::{PlanQuota, PlanQuotaSnapshot};
|
||||
|
||||
/// The window the weekly pacing reads, and the one the meter keeps fresh.
|
||||
pub(crate) const WEEK: &str = "seven_day";
|
||||
/// The windows the pacing reads, and so the only ones the cell keeps.
|
||||
pub(crate) const PACED: [&str; 2] = ["five_hour", WEEK];
|
||||
|
||||
/// How often a lease waiting for its first reading looks at the cell again.
|
||||
const FIRST_READ_POLL: Duration = Duration::from_millis(50);
|
||||
|
||||
/// The latest observation of each plan window, shared by the backend that
|
||||
/// feeds it and whatever reads it out. Each window carries the time of its
|
||||
/// own observation: an event that names one window says nothing about the
|
||||
/// others.
|
||||
///
|
||||
/// It also keeps who is asking: a lease request is the swarm's only sign of
|
||||
/// life the connector sees, and the meter reads the plan only while someone
|
||||
/// asks.
|
||||
#[derive(Default)]
|
||||
pub struct QuotaCell {
|
||||
windows: Mutex<BTreeMap<String, PlanQuota>>,
|
||||
asks: Mutex<Asks>,
|
||||
asked: Condvar,
|
||||
metered: AtomicBool,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct Asks {
|
||||
/// Lease requests in flight or held.
|
||||
live: usize,
|
||||
/// Requests ever made; the meter sleeps until it changes.
|
||||
count: u64,
|
||||
/// When the latest request came, unix ms.
|
||||
last_ms: Option<u64>,
|
||||
}
|
||||
|
||||
/// A lease request while it is in flight or held. Dropping it is the lease
|
||||
/// ending, by whatever road.
|
||||
pub(crate) struct Asking(Arc<QuotaCell>);
|
||||
|
||||
impl Drop for Asking {
|
||||
fn drop(&mut self) {
|
||||
let mut asks = self.0.asks.lock().unwrap_or_else(|e| e.into_inner());
|
||||
asks.live = asks.live.saturating_sub(1);
|
||||
}
|
||||
}
|
||||
|
||||
impl QuotaCell {
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// Notes a lease request, now, and wakes the meter.
|
||||
pub(crate) fn asking(self: &Arc<Self>) -> Asking {
|
||||
self.asking_at(unix_ms())
|
||||
}
|
||||
|
||||
pub(crate) fn asking_at(self: &Arc<Self>, now: u64) -> Asking {
|
||||
{
|
||||
let mut asks = self.asks.lock().unwrap_or_else(|e| e.into_inner());
|
||||
asks.live += 1;
|
||||
asks.count += 1;
|
||||
asks.last_ms = Some(now);
|
||||
}
|
||||
self.asked.notify_all();
|
||||
Asking(self.clone())
|
||||
}
|
||||
|
||||
/// Whether a lease is held, or one was asked for less than `within` ago.
|
||||
pub(crate) fn wanted_at(&self, now: u64, within: Duration) -> bool {
|
||||
let asks = self.asks.lock().unwrap_or_else(|e| e.into_inner());
|
||||
asks.live > 0
|
||||
|| asks.last_ms.is_some_and(|last| {
|
||||
now < last.saturating_add(u64::try_from(within.as_millis()).unwrap_or(u64::MAX))
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn asks_so_far(&self) -> u64 {
|
||||
self.asks.lock().unwrap_or_else(|e| e.into_inner()).count
|
||||
}
|
||||
|
||||
/// Sleeps up to `limit`, or until a request beyond the first `seen` comes.
|
||||
pub(crate) fn sleep_until_asked(&self, seen: u64, limit: Duration) {
|
||||
let asks = self.asks.lock().unwrap_or_else(|e| e.into_inner());
|
||||
drop(
|
||||
self.asked
|
||||
.wait_timeout_while(asks, limit, |asks| asks.count == seen)
|
||||
.unwrap_or_else(|e| e.into_inner()),
|
||||
);
|
||||
}
|
||||
|
||||
/// Declares that a meter reads this cell, so a first lease may wait for it.
|
||||
pub(crate) fn attach_meter(&self) {
|
||||
self.metered.store(true, Ordering::SeqCst);
|
||||
}
|
||||
|
||||
/// The windows as they stand, but with no weekly reading yet (a restart
|
||||
/// leaves the cell empty) a metered cell is given up to `within` to get
|
||||
/// one, so a lease is not lent "unobserved" for want of a few seconds.
|
||||
pub(crate) fn snapshot_with_week(&self, within: Duration) -> PlanQuotaSnapshot {
|
||||
let deadline = Instant::now() + within;
|
||||
loop {
|
||||
let snapshot = self.snapshot();
|
||||
let read = snapshot
|
||||
.windows
|
||||
.iter()
|
||||
.any(|quota| quota.window == WEEK && quota.status != EXPIRED);
|
||||
let now = Instant::now();
|
||||
if read || !self.metered.load(Ordering::SeqCst) || now >= deadline {
|
||||
return snapshot;
|
||||
}
|
||||
std::thread::sleep(FIRST_READ_POLL.min(deadline - now));
|
||||
}
|
||||
}
|
||||
|
||||
/// Records the windows one event reported, stamping each with now. See
|
||||
/// [`QuotaCell::observe_at`] for how a reading combines with what is held.
|
||||
pub fn observe(&self, windows: impl IntoIterator<Item = PlanQuota>) {
|
||||
self.observe_at(unix_ms(), windows);
|
||||
}
|
||||
|
||||
/// Records readings taken at `at` (unix ms). A reading older than the
|
||||
/// held one is dropped, so a slow call cannot replace what a faster one
|
||||
/// saw meanwhile. Windows the readings did not mention stay as they were,
|
||||
/// stamp included.
|
||||
pub fn observe_at(&self, at: u64, windows: impl IntoIterator<Item = PlanQuota>) {
|
||||
let mut held = self.windows.lock().unwrap_or_else(|e| e.into_inner());
|
||||
for mut quota in windows {
|
||||
quota.observed_at_unix_ms = at;
|
||||
if held
|
||||
.get("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));
|
||||
}
|
||||
}
|
||||
445
crates/claude-connector/src/stream.rs
Normal file
445
crates/claude-connector/src/stream.rs
Normal file
|
|
@ -0,0 +1,445 @@
|
|||
//! Reading the CLI's stream-json output. The format is foreign and moves with
|
||||
//! the CLI, so every struct here ignores fields it does not name.
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use model_connector::{FailureKind, PlanQuota, StopReason, Usage};
|
||||
use serde::Deserialize;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::quota::{PACED, QuotaCell};
|
||||
|
||||
/// The tool the CLI gives the model to deliver an answer that must match a
|
||||
/// `--json-schema`. Its call is the structured output, not a tool exchange.
|
||||
const STRUCTURED_OUTPUT_TOOL: &str = "StructuredOutput";
|
||||
|
||||
/// A call that cannot complete, with the stage that failed. `usage` is `None`
|
||||
/// (unknown) unless a result event reported it or the
|
||||
/// call never reached the provider (zero).
|
||||
#[derive(Debug, PartialEq, Eq)]
|
||||
pub(crate) struct Failure {
|
||||
pub kind: FailureKind,
|
||||
pub message: &'static str,
|
||||
pub usage: Option<Usage>,
|
||||
}
|
||||
|
||||
impl Failure {
|
||||
/// A failure after the child ran: what the provider billed is unknown, so
|
||||
/// the usage is unknown.
|
||||
pub(crate) fn cost_unknown(kind: FailureKind, message: &'static str) -> Self {
|
||||
Self {
|
||||
kind,
|
||||
message,
|
||||
usage: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// A failure before anything was sent to the provider: it bills nothing.
|
||||
pub(crate) fn never_reached(kind: FailureKind, message: &'static str) -> Self {
|
||||
Self {
|
||||
kind,
|
||||
message,
|
||||
usage: Some(sum_usage(std::iter::empty())),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// What the stream said, once its result event arrived.
|
||||
#[derive(Debug, PartialEq, Eq)]
|
||||
pub(crate) struct Finished {
|
||||
pub session_id: String,
|
||||
/// The non-empty text blocks, in order.
|
||||
pub texts: Vec<String>,
|
||||
/// The id of the last assistant message.
|
||||
pub message_id: Option<String>,
|
||||
pub stop_reason: StopReason,
|
||||
pub usage: Usage,
|
||||
}
|
||||
|
||||
pub(crate) enum Step {
|
||||
More,
|
||||
Done(Box<Result<Finished, Failure>>),
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct Line {
|
||||
#[serde(rename = "type")]
|
||||
kind: String,
|
||||
subtype: Option<String>,
|
||||
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>,
|
||||
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>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct Message {
|
||||
id: Option<String>,
|
||||
content: Option<Vec<Block>>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct Block {
|
||||
#[serde(rename = "type")]
|
||||
kind: String,
|
||||
text: Option<String>,
|
||||
name: Option<String>,
|
||||
}
|
||||
|
||||
/// One model's entry. The four counts are required: an entry that lacks one
|
||||
/// is a shape this reader does not know, and unknown usage is never zero.
|
||||
#[derive(Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct ModelUsage {
|
||||
input_tokens: u64,
|
||||
output_tokens: u64,
|
||||
cache_read_input_tokens: u64,
|
||||
cache_creation_input_tokens: u64,
|
||||
thinking_tokens: Option<u64>,
|
||||
}
|
||||
|
||||
/// 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>,
|
||||
plan_rejected: bool,
|
||||
}
|
||||
|
||||
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());
|
||||
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,
|
||||
plan_rejected: false,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn feed(&mut self, raw: &[u8]) -> Step {
|
||||
let Ok(line) = serde_json::from_slice::<Line>(raw) else {
|
||||
return fail(
|
||||
FailureKind::Provider,
|
||||
"claude-code: a stream line did not parse",
|
||||
);
|
||||
};
|
||||
match line.kind.as_str() {
|
||||
"system" => match line.subtype.as_deref() {
|
||||
Some("init") => {
|
||||
if line.model.as_deref() != Some(self.model) {
|
||||
return fail(
|
||||
FailureKind::Provider,
|
||||
"claude-code: the session model differs from the request",
|
||||
);
|
||||
}
|
||||
self.session_id = line.session_id;
|
||||
Step::More
|
||||
}
|
||||
Some("compact_boundary") => fail(
|
||||
FailureKind::Provider,
|
||||
"claude-code: the session was compacted",
|
||||
),
|
||||
_ => Step::More,
|
||||
},
|
||||
"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));
|
||||
}
|
||||
Step::More
|
||||
}
|
||||
"result" => Step::Done(Box::new(self.result(line))),
|
||||
_ => Step::More,
|
||||
}
|
||||
}
|
||||
|
||||
fn assistant(&mut self, message: Option<Message>) -> Step {
|
||||
let Some(message) = message else {
|
||||
return Step::More;
|
||||
};
|
||||
if message.id.is_some() {
|
||||
self.message_id = message.id;
|
||||
}
|
||||
for block in message.content.unwrap_or_default() {
|
||||
match block.kind.as_str() {
|
||||
"text" => {
|
||||
if let Some(text) = block.text.filter(|text| !text.is_empty()) {
|
||||
self.texts.push(text);
|
||||
}
|
||||
}
|
||||
"tool_use"
|
||||
if self.structured && block.name.as_deref() == Some(STRUCTURED_OUTPUT_TOOL) => {
|
||||
}
|
||||
"tool_use" => {
|
||||
return fail(
|
||||
FailureKind::Provider,
|
||||
"claude-code: the model called a tool in a tool-less turn",
|
||||
);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Step::More
|
||||
}
|
||||
|
||||
fn result(&mut self, line: Line) -> Result<Finished, Failure> {
|
||||
let success = line.subtype.as_deref() == Some("success") && line.is_error == Some(false);
|
||||
let door_refusal = !success
|
||||
&& matches!(
|
||||
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,
|
||||
);
|
||||
if !success {
|
||||
return Err(Failure {
|
||||
kind: self.error_kind(&line.api_error_status, &line.result, &line.errors),
|
||||
message: "claude-code: the result reported an error",
|
||||
usage,
|
||||
});
|
||||
}
|
||||
let Some(usage) = usage else {
|
||||
return Err(Failure {
|
||||
kind: FailureKind::Provider,
|
||||
message: "claude-code: the result did not report usable usage",
|
||||
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,
|
||||
Some("model_context_window_exceeded") => StopReason::ContextWindow,
|
||||
_ => {
|
||||
return Err(Failure {
|
||||
kind: FailureKind::Provider,
|
||||
message: "claude-code: the result has an unknown stop reason",
|
||||
usage: Some(usage),
|
||||
});
|
||||
}
|
||||
};
|
||||
let Some(session_id) = self.session_id.take() else {
|
||||
return Err(Failure {
|
||||
kind: FailureKind::Provider,
|
||||
message: "claude-code: the stream never announced its session",
|
||||
usage: Some(usage),
|
||||
});
|
||||
};
|
||||
Ok(Finished {
|
||||
session_id,
|
||||
texts,
|
||||
message_id: self.message_id.take(),
|
||||
stop_reason,
|
||||
usage,
|
||||
})
|
||||
}
|
||||
|
||||
fn error_kind(
|
||||
&self,
|
||||
status: &Option<Value>,
|
||||
result: &Option<Value>,
|
||||
errors: &Option<Value>,
|
||||
) -> FailureKind {
|
||||
let status = status.as_ref().and_then(Value::as_u64);
|
||||
let names_rate_limit = [result, errors].into_iter().flatten().any(|value| {
|
||||
let text = value.to_string().to_lowercase();
|
||||
text.contains("rate limit") || text.contains("rate_limit")
|
||||
});
|
||||
match status {
|
||||
_ if self.plan_rejected || names_rate_limit || status == Some(429) => {
|
||||
FailureKind::RateLimited
|
||||
}
|
||||
Some(401 | 403) => FailureKind::Authentication,
|
||||
Some(400) => FailureKind::InvalidRequest,
|
||||
Some(500..=599) => FailureKind::Unavailable,
|
||||
_ => FailureKind::Provider,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 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()?;
|
||||
if models.is_empty() {
|
||||
return empty_is_zero.then(|| sum_usage(std::iter::empty()));
|
||||
}
|
||||
Some(sum_usage(models.into_values()))
|
||||
}
|
||||
|
||||
fn sum_usage(models: impl Iterator<Item = ModelUsage>) -> Usage {
|
||||
let mut usage = Usage {
|
||||
input_tokens: 0,
|
||||
cache_read_tokens: 0,
|
||||
cache_write_tokens: 0,
|
||||
output_tokens: 0,
|
||||
reasoning_tokens: None,
|
||||
};
|
||||
for model in models {
|
||||
usage.input_tokens = usage.input_tokens.saturating_add(model.input_tokens);
|
||||
usage.cache_read_tokens = usage
|
||||
.cache_read_tokens
|
||||
.saturating_add(model.cache_read_input_tokens);
|
||||
usage.cache_write_tokens = usage
|
||||
.cache_write_tokens
|
||||
.saturating_add(model.cache_creation_input_tokens);
|
||||
usage.output_tokens = usage.output_tokens.saturating_add(model.output_tokens);
|
||||
if let Some(thinking) = model.thinking_tokens {
|
||||
usage.reasoning_tokens =
|
||||
Some(usage.reasoning_tokens.unwrap_or(0).saturating_add(thinking));
|
||||
}
|
||||
}
|
||||
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}");
|
||||
}
|
||||
}
|
||||
}
|
||||
361
crates/claude-connector/src/unit.rs
Normal file
361
crates/claude-connector/src/unit.rs
Normal file
|
|
@ -0,0 +1,361 @@
|
|||
//! Which systemd unit invocation a process runs in, from the system's own
|
||||
//! record: the process's cgroup says the unit, and the service manager says
|
||||
//! that unit's current InvocationID. The lease listener asks this through the
|
||||
//! `PeerUnits` port; the answer never comes from the peer.
|
||||
//!
|
||||
//! This is the one place the connector asks systemd anything. It runs
|
||||
//! `systemctl show` for one read-only property of one unit and takes nothing
|
||||
//! from the output but a 32-digit hex id.
|
||||
|
||||
use std::fs;
|
||||
use std::io::Read;
|
||||
use std::os::fd::{AsRawFd, BorrowedFd};
|
||||
use std::path::PathBuf;
|
||||
use std::process::{Command, Stdio};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use crate::lease::{PeerUnit, PeerUnits};
|
||||
|
||||
/// Answers `PeerUnits` from `<proc>/<pid>/cgroup` and `systemctl show`.
|
||||
pub struct SystemdUnits {
|
||||
proc_root: PathBuf,
|
||||
systemctl: PathBuf,
|
||||
/// How long the manager has to answer before the lease is refused.
|
||||
deadline: Duration,
|
||||
}
|
||||
|
||||
impl SystemdUnits {
|
||||
/// The running system's `/proc` and `systemctl`.
|
||||
pub fn system() -> Self {
|
||||
Self {
|
||||
proc_root: PathBuf::from("/proc"),
|
||||
systemctl: PathBuf::from("/usr/bin/systemctl"),
|
||||
deadline: Duration::from_secs(5),
|
||||
}
|
||||
}
|
||||
|
||||
/// The pid a process handle names now, from the handle's own record. The
|
||||
/// handle holds the process's identity, so the number cannot have been
|
||||
/// recycled; a process that has exited reads as `-1`.
|
||||
fn pid_of(&self, process: BorrowedFd<'_>) -> Option<u32> {
|
||||
let info = fs::read_to_string(
|
||||
self.proc_root
|
||||
.join("self/fdinfo")
|
||||
.join(process.as_raw_fd().to_string()),
|
||||
)
|
||||
.ok()?;
|
||||
info.lines()
|
||||
.find_map(|line| line.strip_prefix("Pid:"))?
|
||||
.trim()
|
||||
.parse()
|
||||
.ok()
|
||||
}
|
||||
|
||||
fn unit_in_cgroup(&self, pid: u32) -> Option<String> {
|
||||
let cgroup = fs::read_to_string(self.proc_root.join(pid.to_string()).join("cgroup"));
|
||||
unit_of_cgroup(&cgroup.ok()?)
|
||||
}
|
||||
|
||||
/// `systemctl show` of `unit`'s InvocationID, killed and refused if it has
|
||||
/// not exited by the deadline.
|
||||
fn invocation_of(&self, unit: &str) -> Option<String> {
|
||||
let mut child = Command::new(&self.systemctl)
|
||||
.env_clear()
|
||||
.args(["show", "--property=InvocationID", "--value", "--", unit])
|
||||
.stdin(Stdio::null())
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::null())
|
||||
.spawn()
|
||||
.ok()?;
|
||||
let give_up = Instant::now() + self.deadline;
|
||||
let status = loop {
|
||||
match child.try_wait() {
|
||||
Ok(Some(status)) => break status,
|
||||
Ok(None) if Instant::now() < give_up => {
|
||||
std::thread::sleep(Duration::from_millis(10));
|
||||
}
|
||||
_ => {
|
||||
let _ = child.kill();
|
||||
let _ = child.wait();
|
||||
return None;
|
||||
}
|
||||
}
|
||||
};
|
||||
if !status.success() {
|
||||
return None;
|
||||
}
|
||||
// The id is 33 bytes; anything past a small bound is not an answer.
|
||||
let mut output = Vec::new();
|
||||
child
|
||||
.stdout
|
||||
.take()?
|
||||
.take(256)
|
||||
.read_to_end(&mut output)
|
||||
.ok()?;
|
||||
invocation_id_of(std::str::from_utf8(&output).ok()?)
|
||||
}
|
||||
}
|
||||
|
||||
/// The service unit a cgroup file puts its process in: the last component of
|
||||
/// the unified-hierarchy path, if it names a `.service`. A process in a
|
||||
/// sub-cgroup, a scope, a slice or the root is in no service unit this
|
||||
/// recognises.
|
||||
pub fn unit_of_cgroup(cgroup: &str) -> Option<String> {
|
||||
let path = cgroup.lines().find_map(|line| line.strip_prefix("0::"))?;
|
||||
let unit = path.rsplit('/').next()?;
|
||||
(unit.len() > ".service".len() && unit.ends_with(".service")).then(|| unit.to_string())
|
||||
}
|
||||
|
||||
/// An InvocationID as `systemctl show --value` prints it: 32 lowercase hex
|
||||
/// digits and a newline. Anything else (including the empty answer of a unit
|
||||
/// that is not running) is not an invocation.
|
||||
pub fn invocation_id_of(output: &str) -> Option<String> {
|
||||
let id = output.strip_suffix('\n').unwrap_or(output);
|
||||
(id.len() == 32 && id.bytes().all(|b| matches!(b, b'0'..=b'9' | b'a'..=b'f')))
|
||||
.then(|| id.to_string())
|
||||
}
|
||||
|
||||
impl PeerUnits for SystemdUnits {
|
||||
fn invocation_in(&self, process: BorrowedFd<'_>, unit: &str) -> Option<PeerUnit> {
|
||||
let pid = self.pid_of(process)?;
|
||||
// A process in any other unit is not asked about.
|
||||
if self.unit_in_cgroup(pid).as_deref() != Some(unit) {
|
||||
return None;
|
||||
}
|
||||
let invocation_id = self.invocation_of(unit)?;
|
||||
// The same process must still be in that unit now: one that left (or
|
||||
// died) while the manager was asked is not vouched for.
|
||||
(self.pid_of(process) == Some(pid) && self.unit_in_cgroup(pid).as_deref() == Some(unit))
|
||||
.then(|| PeerUnit {
|
||||
unit: unit.to_string(),
|
||||
invocation_id,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::os::fd::{AsFd, AsRawFd, OwnedFd};
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn a_service_unit_is_the_last_cgroup_component() {
|
||||
let service = "0::/system.slice/system-mind\\x2dself.slice/mind-self@huginn.service\n";
|
||||
assert_eq!(
|
||||
unit_of_cgroup(service).as_deref(),
|
||||
Some("mind-self@huginn.service")
|
||||
);
|
||||
// Several hierarchies listed: the unified one decides.
|
||||
let hybrid = "12:cpu:/foo\n0::/system.slice/mind-self@a.service\n";
|
||||
assert_eq!(
|
||||
unit_of_cgroup(hybrid).as_deref(),
|
||||
Some("mind-self@a.service")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_process_outside_a_service_unit_has_none() {
|
||||
for cgroup in [
|
||||
"0::/\n",
|
||||
"0::/user.slice/user-1000.slice/session-3.scope\n",
|
||||
"0::/user.slice/user-1000.slice/user@1000.service/app.slice/x.service/child\n",
|
||||
"0::/system.slice/mind-self@huginn.service/sub\n",
|
||||
"0::/system.slice/.service\n",
|
||||
"1:name=systemd:/system.slice/a.service\n",
|
||||
"",
|
||||
] {
|
||||
assert_eq!(unit_of_cgroup(cgroup), None, "{cgroup:?}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_invocation_id_is_exactly_32_lowercase_hex_digits() {
|
||||
let id = "0123456789abcdef0123456789abcdef";
|
||||
assert_eq!(invocation_id_of(&format!("{id}\n")).as_deref(), Some(id));
|
||||
assert_eq!(invocation_id_of(id).as_deref(), Some(id));
|
||||
for bad in [
|
||||
"",
|
||||
"\n",
|
||||
"0123456789ABCDEF0123456789ABCDEF\n",
|
||||
"0123456789abcdef0123456789abcde\n",
|
||||
"0123456789abcdef0123456789abcdef0\n",
|
||||
"0123456789abcdef0123456789abcdeg\n",
|
||||
"0123456789abcdef0123456789abcdef\n\n",
|
||||
] {
|
||||
assert_eq!(invocation_id_of(bad), None, "{bad:?}");
|
||||
}
|
||||
}
|
||||
|
||||
const ID: &str = "0123456789abcdef0123456789abcdef";
|
||||
const UNIT: &str = "mind-self@eureka.service";
|
||||
|
||||
/// A fake `/proc` with one process in `cgroup`, and a fake `systemctl`
|
||||
/// that records its arguments, runs `side_effect`, and exits `status`
|
||||
/// after printing `output`.
|
||||
struct Fake {
|
||||
dir: PathBuf,
|
||||
units: SystemdUnits,
|
||||
// A script being written while another test forks is "text file
|
||||
// busy" when it is run: fakes take turns.
|
||||
_turn: std::sync::MutexGuard<'static, ()>,
|
||||
}
|
||||
|
||||
static TURN: std::sync::Mutex<()> = std::sync::Mutex::new(());
|
||||
|
||||
impl Fake {
|
||||
fn new(name: &str, cgroup: &str, output: &str, status: i32, side_effect: &str) -> Self {
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
let turn = TURN
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
let dir = std::env::temp_dir().join(format!("unit-{}-{name}", std::process::id()));
|
||||
let _ = fs::remove_dir_all(&dir);
|
||||
let side_effect = side_effect.replace("{d}", &dir.display().to_string());
|
||||
fs::create_dir_all(dir.join("proc/42")).unwrap();
|
||||
fs::write(dir.join("proc/42/cgroup"), cgroup).unwrap();
|
||||
let script = dir.join("systemctl");
|
||||
fs::write(
|
||||
&script,
|
||||
format!(
|
||||
"#!/bin/sh\nprintf '%s' \"$*\" > {d}/args\n{side_effect}\nprintf '%s' '{output}'\nexit {status}\n",
|
||||
d = dir.display()
|
||||
),
|
||||
)
|
||||
.unwrap();
|
||||
fs::set_permissions(&script, fs::Permissions::from_mode(0o755)).unwrap();
|
||||
let units = SystemdUnits {
|
||||
proc_root: dir.join("proc"),
|
||||
systemctl: script,
|
||||
deadline: Duration::from_secs(5),
|
||||
};
|
||||
Self {
|
||||
dir,
|
||||
units,
|
||||
_turn: turn,
|
||||
}
|
||||
}
|
||||
|
||||
/// A process handle that names process `pid`, as the kernel's own
|
||||
/// record of the handle says (`-1` once the process has exited).
|
||||
fn process(&self, pid: i64) -> OwnedFd {
|
||||
let handle: OwnedFd = fs::File::open("/dev/null").unwrap().into();
|
||||
let info = self.dir.join("proc/self/fdinfo");
|
||||
fs::create_dir_all(&info).unwrap();
|
||||
fs::write(
|
||||
info.join(handle.as_raw_fd().to_string()),
|
||||
format!("pos:\t0\nflags:\t02000002\nPid:\t{pid}\nNSpid:\t{pid}\n"),
|
||||
)
|
||||
.unwrap();
|
||||
handle
|
||||
}
|
||||
|
||||
fn invocation_in(&self, pid: i64, unit: &str) -> Option<PeerUnit> {
|
||||
self.units.invocation_in(self.process(pid).as_fd(), unit)
|
||||
}
|
||||
|
||||
fn args(&self) -> String {
|
||||
fs::read_to_string(self.dir.join("args")).unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Fake {
|
||||
fn drop(&mut self) {
|
||||
let _ = fs::remove_dir_all(&self.dir);
|
||||
}
|
||||
}
|
||||
|
||||
fn eureka_cgroup() -> String {
|
||||
format!("0::/system.slice/{UNIT}\n")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_process_in_the_unit_is_that_unit_at_the_invocation_the_manager_reports() {
|
||||
let fake = Fake::new("ok", &eureka_cgroup(), &format!("{ID}\n"), 0, "");
|
||||
assert_eq!(
|
||||
fake.invocation_in(42, UNIT),
|
||||
Some(PeerUnit {
|
||||
unit: UNIT.to_string(),
|
||||
invocation_id: ID.to_string()
|
||||
})
|
||||
);
|
||||
// The manager is asked for exactly the invocation of that unit.
|
||||
assert_eq!(
|
||||
fake.args(),
|
||||
format!("show --property=InvocationID --value -- {UNIT}")
|
||||
);
|
||||
// A handle on a process that has exited, or on one with no cgroup
|
||||
// record, is in no unit.
|
||||
assert_eq!(fake.invocation_in(-1, UNIT), None);
|
||||
assert_eq!(fake.invocation_in(43, UNIT), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_process_outside_the_unit_is_not_asked_about() {
|
||||
let scope = Fake::new("scope", "0::/user.slice/session-3.scope\n", ID, 0, "");
|
||||
assert_eq!(scope.invocation_in(42, UNIT), None);
|
||||
assert!(!scope.dir.join("args").exists());
|
||||
drop(scope);
|
||||
// Another service, even a real one, is not the caller's unit: the
|
||||
// manager is never spawned for it.
|
||||
let other = Fake::new("other", "0::/system.slice/ssh.service\n", ID, 0, "");
|
||||
assert_eq!(other.invocation_in(42, UNIT), None);
|
||||
assert!(!other.dir.join("args").exists());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_manager_that_fails_or_names_no_invocation_vouches_for_nothing() {
|
||||
let failed = Fake::new("fail", &eureka_cgroup(), ID, 1, "");
|
||||
assert_eq!(failed.invocation_in(42, UNIT), None);
|
||||
drop(failed);
|
||||
let not_running = Fake::new("idle", &eureka_cgroup(), "\n", 0, "");
|
||||
assert_eq!(not_running.invocation_in(42, UNIT), None);
|
||||
let missing = SystemdUnits {
|
||||
proc_root: not_running.dir.join("proc"),
|
||||
systemctl: not_running.dir.join("no-such-systemctl"),
|
||||
deadline: Duration::from_secs(5),
|
||||
};
|
||||
assert_eq!(
|
||||
missing.invocation_in(not_running.process(42).as_fd(), UNIT),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_manager_that_does_not_answer_in_time_is_killed_and_the_lease_refused() {
|
||||
let mut fake = Fake::new("hung", &eureka_cgroup(), ID, 0, "exec sleep 30");
|
||||
fake.units.deadline = Duration::from_millis(200);
|
||||
let started = Instant::now();
|
||||
assert_eq!(fake.invocation_in(42, UNIT), None);
|
||||
assert!(started.elapsed() < Duration::from_secs(5));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_process_that_left_its_unit_while_the_manager_was_asked_is_not_vouched_for() {
|
||||
let fake = Fake::new(
|
||||
"moves",
|
||||
&eureka_cgroup(),
|
||||
ID,
|
||||
0,
|
||||
"printf '0::/system.slice/other.service\n' > {d}/proc/42/cgroup",
|
||||
);
|
||||
assert_eq!(fake.invocation_in(42, UNIT), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_process_that_exited_while_the_manager_was_asked_is_not_vouched_for() {
|
||||
let fake = Fake::new("dies", &eureka_cgroup(), ID, 0, "");
|
||||
let handle = fake.process(42);
|
||||
let side = fake
|
||||
.dir
|
||||
.join("proc/self/fdinfo")
|
||||
.join(handle.as_raw_fd().to_string());
|
||||
let mut script = fs::read_to_string(fake.dir.join("systemctl")).unwrap();
|
||||
script = script.replace(
|
||||
"exit 0",
|
||||
&format!("printf 'Pid:\t-1\n' > {}\nexit 0", side.display()),
|
||||
);
|
||||
fs::write(fake.dir.join("systemctl"), script).unwrap();
|
||||
assert_eq!(fake.units.invocation_in(handle.as_fd(), UNIT), None);
|
||||
}
|
||||
}
|
||||
769
crates/claude-connector/tests/ask.rs
Normal file
769
crates/claude-connector/tests/ask.rs
Normal file
|
|
@ -0,0 +1,769 @@
|
|||
//! `claude-connector ask` against a real serve loop on loopback with a scripted
|
||||
//! backend: the consumer side of the no-tools completion.
|
||||
|
||||
use std::net::{SocketAddr, TcpListener};
|
||||
use std::path::PathBuf;
|
||||
use std::process::{Command, Stdio};
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::thread;
|
||||
use std::time::Duration;
|
||||
|
||||
use claude_connector::ask::{AskArgs, AskError, ask, build_invocation};
|
||||
use model_connector::*;
|
||||
|
||||
const CALLER: &str = "ask-caller";
|
||||
const KEY: &str = "ask-test-key-0123456789-abcdefghijklmnop";
|
||||
const MODEL: &str = "claude-haiku-5-5";
|
||||
const FRAME: usize = 4 << 20;
|
||||
|
||||
/// What the scripted backend does with each call.
|
||||
#[derive(Clone)]
|
||||
enum Script {
|
||||
Text(&'static str),
|
||||
Empty,
|
||||
Failed(&'static str),
|
||||
Sleep(Duration),
|
||||
ToolCall,
|
||||
Stop(StopReason),
|
||||
}
|
||||
|
||||
struct Backend {
|
||||
script: Script,
|
||||
seen: Mutex<Vec<ProviderRequest>>,
|
||||
}
|
||||
|
||||
impl ProviderBackend for Backend {
|
||||
fn ready(&self) -> Result<(), &'static str> {
|
||||
Ok(())
|
||||
}
|
||||
fn provider(&self) -> &str {
|
||||
"claude-code"
|
||||
}
|
||||
fn execute(&self, call: &ProviderCall) -> TransportResult {
|
||||
self.seen.lock().unwrap().push(call.request().clone());
|
||||
let usage = Usage {
|
||||
input_tokens: 1,
|
||||
cache_read_tokens: 0,
|
||||
cache_write_tokens: 0,
|
||||
output_tokens: 1,
|
||||
reasoning_tokens: None,
|
||||
};
|
||||
let text = |t: &str| {
|
||||
vec![TransportEvent {
|
||||
sequence: 0,
|
||||
payload: TransportEventPayload::TextDelta { text: t.into() },
|
||||
}]
|
||||
};
|
||||
let stop = match &self.script {
|
||||
Script::Stop(reason) => *reason,
|
||||
_ => StopReason::EndTurn,
|
||||
};
|
||||
let completed = TransportOutcome::Completed {
|
||||
provider_response_id: None,
|
||||
stop_reason: stop,
|
||||
usage: usage.clone(),
|
||||
};
|
||||
match &self.script {
|
||||
Script::Text(t) => call.transported(text(t), completed),
|
||||
Script::Empty => call.transported(Vec::new(), completed),
|
||||
Script::Failed(message) => call.transported(
|
||||
Vec::new(),
|
||||
TransportOutcome::Failed {
|
||||
kind: FailureKind::Provider,
|
||||
message: message.to_string(),
|
||||
usage: None,
|
||||
},
|
||||
),
|
||||
Script::Stop(_) => call.transported(text("cut off"), completed),
|
||||
Script::ToolCall => call.transported(
|
||||
vec![TransportEvent {
|
||||
sequence: 0,
|
||||
payload: TransportEventPayload::ToolCall {
|
||||
call_id: "c".into(),
|
||||
name: "n".into(),
|
||||
arguments: "{}".into(),
|
||||
},
|
||||
}],
|
||||
completed,
|
||||
),
|
||||
Script::Sleep(d) => {
|
||||
thread::sleep(*d);
|
||||
call.transported(text("late"), completed)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn write_key(path: &std::path::Path, bytes: &[u8], mode: u32) {
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
std::fs::write(path, bytes).unwrap();
|
||||
std::fs::set_permissions(path, std::fs::Permissions::from_mode(mode)).unwrap();
|
||||
}
|
||||
|
||||
struct Rig {
|
||||
endpoint: SocketAddr,
|
||||
dir: PathBuf,
|
||||
key_file: PathBuf,
|
||||
backend: Arc<Backend>,
|
||||
}
|
||||
|
||||
fn rig(script: Script) -> Rig {
|
||||
static SEQ: AtomicUsize = AtomicUsize::new(0);
|
||||
let dir = std::env::temp_dir().join(format!(
|
||||
"claude-connector-ask-{}-{}",
|
||||
std::process::id(),
|
||||
SEQ.fetch_add(1, Ordering::Relaxed)
|
||||
));
|
||||
std::fs::create_dir_all(&dir).unwrap();
|
||||
let key_file = dir.join("caller.key");
|
||||
write_key(&key_file, format!("{KEY}\n").as_bytes(), 0o600);
|
||||
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
let endpoint = listener.local_addr().unwrap();
|
||||
let admission =
|
||||
CallerAdmission::new(CALLER, KEY, 1, [MODEL.to_string()], 4, 1 << 20, 4_096).unwrap();
|
||||
let limits = ServeLimits {
|
||||
max_frame_bytes: FRAME,
|
||||
max_connections: 8,
|
||||
socket_timeout_ms: 5_000,
|
||||
max_expiry_skew_ms: 300_000,
|
||||
};
|
||||
let backend = Arc::new(Backend {
|
||||
script,
|
||||
seen: Mutex::new(Vec::new()),
|
||||
});
|
||||
let served = backend.clone();
|
||||
let store = dir.join("replay.cc");
|
||||
thread::spawn(move || serve(listener, limits, vec![admission], &store, served));
|
||||
Rig {
|
||||
endpoint,
|
||||
dir,
|
||||
key_file,
|
||||
backend,
|
||||
}
|
||||
}
|
||||
|
||||
impl Rig {
|
||||
fn args(&self) -> AskArgs {
|
||||
AskArgs {
|
||||
endpoint: self.endpoint,
|
||||
caller_id: CALLER.into(),
|
||||
key_file: self.key_file.clone(),
|
||||
model: MODEL.into(),
|
||||
instructions: "Answer briefly.".into(),
|
||||
effort: Some(ReasoningEffort::Low),
|
||||
max_output_tokens: 1_024,
|
||||
timeout: Duration::from_secs(5),
|
||||
prompt: "SECRET-PROMPT-TEXT".into(),
|
||||
}
|
||||
}
|
||||
fn seen(&self) -> Vec<ProviderRequest> {
|
||||
self.backend.seen.lock().unwrap().clone()
|
||||
}
|
||||
}
|
||||
|
||||
fn message(error: &AskError) -> String {
|
||||
error.to_string()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn success_returns_the_answer_text_and_the_backend_saw_a_toolless_request() {
|
||||
let rig = rig(Script::Text("four"));
|
||||
assert_eq!(ask(&rig.args()).unwrap(), "four");
|
||||
let seen = rig.seen();
|
||||
assert_eq!(seen.len(), 1);
|
||||
assert!(seen[0].tools.is_empty());
|
||||
assert_eq!(seen[0].model, MODEL);
|
||||
assert_eq!(seen[0].instructions, "Answer briefly.");
|
||||
assert_eq!(seen[0].reasoning_effort, Some(ReasoningEffort::Low));
|
||||
assert_eq!(seen[0].max_output_tokens, 1_024);
|
||||
assert_eq!(
|
||||
seen[0].input,
|
||||
vec![InputItem::UserText {
|
||||
text: "SECRET-PROMPT-TEXT".into()
|
||||
}]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn each_run_uses_a_fresh_request_id_so_replay_never_answers_a_new_question() {
|
||||
let rig = rig(Script::Text("again"));
|
||||
ask(&rig.args()).unwrap();
|
||||
ask(&rig.args()).unwrap();
|
||||
let seen = rig.seen();
|
||||
assert_eq!(seen.len(), 2, "the provider ran once per run");
|
||||
assert_ne!(seen[0].request_id, seen[1].request_id);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_cli_cannot_construct_a_request_with_tools() {
|
||||
let rig = rig(Script::Text("x"));
|
||||
let invocation = build_invocation(&rig.args()).unwrap();
|
||||
assert!(invocation.request.tools.is_empty());
|
||||
assert_eq!(invocation.request.tool_choice, ToolChoice::Auto);
|
||||
assert!(
|
||||
invocation
|
||||
.request
|
||||
.input
|
||||
.iter()
|
||||
.all(|item| matches!(item, InputItem::UserText { .. }))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_model_outside_allowed_models_is_refused_loudly() {
|
||||
let rig = rig(Script::Text("x"));
|
||||
let mut args = rig.args();
|
||||
args.model = "claude-opus-5".into();
|
||||
let error = ask(&args).unwrap_err();
|
||||
assert!(
|
||||
matches!(error, AskError::Refused(Refusal::ModelNotAllowed)),
|
||||
"{error}"
|
||||
);
|
||||
assert!(rig.seen().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_output_bound_over_the_callers_cap_is_refused() {
|
||||
let rig = rig(Script::Text("x"));
|
||||
let mut args = rig.args();
|
||||
args.max_output_tokens = 4_097;
|
||||
assert!(matches!(
|
||||
ask(&args).unwrap_err(),
|
||||
AskError::Refused(Refusal::OutputBound)
|
||||
));
|
||||
assert!(rig.seen().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_wrong_key_or_caller_is_an_error_and_never_reaches_the_provider() {
|
||||
let rig = rig(Script::Text("x"));
|
||||
let mut wrong_key = rig.args();
|
||||
let other = rig.dir.join("other.key");
|
||||
write_key(&other, b"some-other-connection-key-0123456789", 0o600);
|
||||
wrong_key.key_file = other;
|
||||
let error = ask(&wrong_key).unwrap_err();
|
||||
assert!(!message(&error).contains("some-other"), "{error}");
|
||||
let mut wrong_caller = rig.args();
|
||||
wrong_caller.caller_id = "someone-else".into();
|
||||
ask(&wrong_caller).unwrap_err();
|
||||
assert!(rig.seen().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_unusable_key_file_errors_without_echoing_its_contents() {
|
||||
let rig = rig(Script::Text("x"));
|
||||
let cases: Vec<(&str, Vec<u8>)> = vec![
|
||||
("empty", b"\n".to_vec()),
|
||||
("long", vec![b'k'; 2_000]),
|
||||
("binary", vec![0xff, 0xfe, 0xfd, 0x41]),
|
||||
("padded", b" spaced-secret-value-padded ".to_vec()),
|
||||
];
|
||||
for (name, bytes) in cases {
|
||||
let path = rig.dir.join(format!("{name}.key"));
|
||||
write_key(&path, &bytes, 0o600);
|
||||
let mut args = rig.args();
|
||||
args.key_file = path;
|
||||
let text = message(&ask(&args).unwrap_err());
|
||||
assert!(!text.contains("spaced-secret"), "{name}: {text}");
|
||||
assert!(!text.contains("kkkk"), "{name}: {text}");
|
||||
}
|
||||
let mut missing = rig.args();
|
||||
missing.key_file = rig.dir.join("absent.key");
|
||||
assert!(matches!(
|
||||
ask(&missing).unwrap_err(),
|
||||
AskError::KeyUnreadable
|
||||
));
|
||||
let mut directory = rig.args();
|
||||
directory.key_file = rig.dir.clone();
|
||||
assert!(matches!(
|
||||
ask(&directory).unwrap_err(),
|
||||
AskError::KeyNotRegular
|
||||
));
|
||||
assert!(rig.seen().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_empty_answer_is_an_error() {
|
||||
for script in [Script::Empty, Script::Text(" \n")] {
|
||||
let rig = rig(script);
|
||||
assert!(matches!(
|
||||
ask(&rig.args()).unwrap_err(),
|
||||
AskError::EmptyAnswer
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_empty_prompt_is_an_error_before_any_connection() {
|
||||
let rig = rig(Script::Text("x"));
|
||||
let mut args = rig.args();
|
||||
args.prompt = " \n\t".into();
|
||||
assert!(matches!(ask(&args).unwrap_err(), AskError::EmptyPrompt));
|
||||
assert!(rig.seen().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_provider_failure_names_its_kind_and_not_its_message() {
|
||||
let rig = rig(Script::Failed("stderr tail with SECRET-PROMPT-TEXT"));
|
||||
let error = ask(&rig.args()).unwrap_err();
|
||||
assert!(matches!(error, AskError::Failed(FailureKind::Provider)));
|
||||
assert!(!message(&error).contains("SECRET"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_connector_that_does_not_answer_in_time_is_a_timeout() {
|
||||
let rig = rig(Script::Sleep(Duration::from_secs(3)));
|
||||
let mut args = rig.args();
|
||||
args.timeout = Duration::from_millis(300);
|
||||
assert!(matches!(ask(&args).unwrap_err(), AskError::Timeout));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_dead_endpoint_is_a_connection_error() {
|
||||
let rig = rig(Script::Text("x"));
|
||||
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
let mut args = rig.args();
|
||||
args.endpoint = listener.local_addr().unwrap();
|
||||
drop(listener);
|
||||
assert!(matches!(ask(&args).unwrap_err(), AskError::NotListening));
|
||||
}
|
||||
|
||||
fn run_binary(rig: &Rig, stdin: &str, extra: &[&str]) -> std::process::Output {
|
||||
run_binary_with_key(rig, &rig.key_file, stdin, extra)
|
||||
}
|
||||
|
||||
fn run_binary_with_key(
|
||||
rig: &Rig,
|
||||
key_file: &std::path::Path,
|
||||
stdin: &str,
|
||||
extra: &[&str],
|
||||
) -> std::process::Output {
|
||||
let mut child = Command::new(env!("CARGO_BIN_EXE_claude-connector"))
|
||||
.args(["ask", "--endpoint", &rig.endpoint.to_string()])
|
||||
.args([
|
||||
"--caller-id",
|
||||
CALLER,
|
||||
"--model",
|
||||
MODEL,
|
||||
"--system",
|
||||
"Be brief.",
|
||||
])
|
||||
.args(["--connection-key-file", key_file.to_str().unwrap()])
|
||||
.args(extra)
|
||||
.stdin(Stdio::piped())
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped())
|
||||
.spawn()
|
||||
.unwrap();
|
||||
{
|
||||
use std::io::Write;
|
||||
let _ = child.stdin.take().unwrap().write_all(stdin.as_bytes());
|
||||
}
|
||||
child.wait_with_output().unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_toolcall_in_the_answer_is_an_error() {
|
||||
let rig = rig(Script::ToolCall);
|
||||
assert!(matches!(ask(&rig.args()).unwrap_err(), AskError::ToolCall));
|
||||
}
|
||||
|
||||
/// The rule: only a turn that ended (end of turn or a stop sequence) is an
|
||||
/// answer; a truncated, refused or context-limited turn is an error even when
|
||||
/// it carries text.
|
||||
#[test]
|
||||
fn a_turn_that_did_not_end_is_an_error() {
|
||||
for reason in [
|
||||
StopReason::MaxOutputTokens,
|
||||
StopReason::Refusal,
|
||||
StopReason::ContextWindow,
|
||||
StopReason::ToolUse,
|
||||
] {
|
||||
let rig = rig(Script::Stop(reason));
|
||||
assert!(
|
||||
matches!(ask(&rig.args()).unwrap_err(), AskError::NotCompleted(r) if r == reason),
|
||||
"{reason:?}"
|
||||
);
|
||||
}
|
||||
let rig = rig(Script::Stop(StopReason::StopSequence));
|
||||
assert_eq!(ask(&rig.args()).unwrap(), "cut off");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_fifo_key_file_errors_promptly_instead_of_hanging() {
|
||||
let rig = rig(Script::Text("x"));
|
||||
let fifo = rig.dir.join("fifo.key");
|
||||
let path = std::ffi::CString::new(fifo.to_str().unwrap()).unwrap();
|
||||
assert_eq!(unsafe { libc::mkfifo(path.as_ptr(), 0o600) }, 0);
|
||||
let mut args = rig.args();
|
||||
args.key_file = fifo;
|
||||
let (tx, rx) = std::sync::mpsc::channel();
|
||||
thread::spawn(move || tx.send(ask(&args)).unwrap());
|
||||
let error = rx
|
||||
.recv_timeout(Duration::from_secs(3))
|
||||
.expect("reading a FIFO key file must not hang")
|
||||
.unwrap_err();
|
||||
assert!(matches!(error, AskError::KeyNotRegular), "{error}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_key_file_readable_by_group_or_world_is_refused_and_a_private_one_accepted() {
|
||||
let rig = rig(Script::Text("fine"));
|
||||
for mode in [0o644, 0o640, 0o604] {
|
||||
let path = rig.dir.join(format!("mode-{mode:o}.key"));
|
||||
write_key(&path, KEY.as_bytes(), mode);
|
||||
let mut args = rig.args();
|
||||
args.key_file = path;
|
||||
let error = ask(&args).unwrap_err();
|
||||
assert!(matches!(error, AskError::KeyExposed), "{mode:o}: {error}");
|
||||
assert!(!message(&error).contains(KEY));
|
||||
}
|
||||
assert!(rig.seen().is_empty());
|
||||
assert_eq!(ask(&rig.args()).unwrap(), "fine");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_symlinked_key_is_judged_by_its_target() {
|
||||
let rig = rig(Script::Text("linked"));
|
||||
let link = rig.dir.join("link.key");
|
||||
std::os::unix::fs::symlink(&rig.key_file, &link).unwrap();
|
||||
let mut args = rig.args();
|
||||
args.key_file = link;
|
||||
assert_eq!(ask(&args).unwrap(), "linked");
|
||||
let dir_link = rig.dir.join("dir.key");
|
||||
std::os::unix::fs::symlink(&rig.dir, &dir_link).unwrap();
|
||||
args.key_file = dir_link;
|
||||
assert!(matches!(ask(&args).unwrap_err(), AskError::KeyNotRegular));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_key_read_is_bounded_at_1024_bytes() {
|
||||
let rig = rig(Script::Text("x"));
|
||||
let path = rig.dir.join("big.key");
|
||||
write_key(&path, &vec![b'k'; 1_025], 0o600);
|
||||
assert!(matches!(
|
||||
claude_connector::ask::read_key(&path).unwrap_err(),
|
||||
AskError::KeyLength
|
||||
));
|
||||
write_key(&path, &vec![b'k'; 1_024], 0o600);
|
||||
assert!(claude_connector::ask::read_key(&path).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_instructions_file_is_capped_and_must_be_a_regular_file() {
|
||||
use claude_connector::ask::{MAX_SYSTEM_BYTES, read_system_file};
|
||||
let rig = rig(Script::Text("x"));
|
||||
let path = rig.dir.join("system.txt");
|
||||
std::fs::write(&path, vec![b's'; MAX_SYSTEM_BYTES]).unwrap();
|
||||
assert_eq!(read_system_file(&path).unwrap().len(), MAX_SYSTEM_BYTES);
|
||||
std::fs::write(&path, vec![b's'; MAX_SYSTEM_BYTES + 1]).unwrap();
|
||||
assert!(matches!(
|
||||
read_system_file(&path).unwrap_err(),
|
||||
AskError::SystemTooLarge
|
||||
));
|
||||
assert!(matches!(
|
||||
read_system_file(&rig.dir).unwrap_err(),
|
||||
AskError::SystemUnreadable
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_binary_never_prints_key_bytes_for_a_bad_key_file() {
|
||||
let rig = rig(Script::Text("x"));
|
||||
let cases: [(&str, &[u8], u32); 3] = [
|
||||
("exposed", b"exposed-key-bytes-0123456789", 0o644),
|
||||
("padded", b" padded-key-bytes-0123456789 ", 0o600),
|
||||
("wrong", b"wrong-but-wellformed-key-0123456789", 0o600),
|
||||
];
|
||||
for (name, bytes, mode) in cases {
|
||||
let path = rig.dir.join(format!("bin-{name}.key"));
|
||||
write_key(&path, bytes, mode);
|
||||
let out = run_binary_with_key(&rig, &path, "q", &[]);
|
||||
assert!(!out.status.success(), "{name}");
|
||||
let shown = [out.stdout.clone(), out.stderr.clone()].concat();
|
||||
let shown = String::from_utf8_lossy(&shown);
|
||||
assert!(
|
||||
!shown.contains("key-bytes") && !shown.contains("wellformed"),
|
||||
"{name}: {shown}"
|
||||
);
|
||||
assert!(out.stdout.is_empty());
|
||||
}
|
||||
let out = run_binary(&rig, "q", &["--system-file", "/nonexistent"]);
|
||||
assert!(!out.status.success());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stdin_and_stdout_failures_have_their_own_wording() {
|
||||
let rig = rig(Script::Text("x"));
|
||||
let mut child = Command::new(env!("CARGO_BIN_EXE_claude-connector"))
|
||||
.args(["ask", "--endpoint", &rig.endpoint.to_string()])
|
||||
.args(["--caller-id", CALLER, "--model", MODEL, "--system", "s"])
|
||||
.args(["--connection-key-file", rig.key_file.to_str().unwrap()])
|
||||
.stdin(Stdio::piped())
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped())
|
||||
.spawn()
|
||||
.unwrap();
|
||||
{
|
||||
use std::io::Write;
|
||||
child
|
||||
.stdin
|
||||
.take()
|
||||
.unwrap()
|
||||
.write_all(&[0xff, 0xfe])
|
||||
.unwrap();
|
||||
}
|
||||
let out = child.wait_with_output().unwrap();
|
||||
let err = String::from_utf8(out.stderr).unwrap();
|
||||
assert!(err.contains("not UTF-8"), "{err}");
|
||||
assert!(!err.contains("flag"), "{err}");
|
||||
}
|
||||
|
||||
/// What the wire returns for a tools-bearing request that gets past the CLI's
|
||||
/// own construction (here forced by hand) when the real claude backend sits
|
||||
/// behind the real serve loop: the service admits it, the backend's planner
|
||||
/// refuses it, the caller receives a transported `Failed(InvalidRequest)`
|
||||
/// result, and the claude child is never started.
|
||||
#[test]
|
||||
fn a_forced_tools_request_reaches_the_planner_is_refused_and_no_child_starts() {
|
||||
use claude_connector::{
|
||||
ChildEnv, ClaudeCodeBackend, ClaudeCodeConfig, ClaudeCodeLimits, ProcessPort, RunningChild,
|
||||
SessionStore, StderrSink, TokenSource,
|
||||
};
|
||||
use sha2::{Digest, Sha256};
|
||||
use zeroize::Zeroizing;
|
||||
|
||||
struct CountingPort(Arc<AtomicUsize>);
|
||||
impl ProcessPort for CountingPort {
|
||||
fn spawn(
|
||||
&self,
|
||||
_: &[String],
|
||||
_: &ChildEnv,
|
||||
_: &std::path::Path,
|
||||
_: Vec<u8>,
|
||||
_: &str,
|
||||
) -> std::io::Result<Box<dyn RunningChild>> {
|
||||
self.0.fetch_add(1, Ordering::SeqCst);
|
||||
Err(std::io::Error::other("no child in this test"))
|
||||
}
|
||||
}
|
||||
struct Fixed(&'static str);
|
||||
impl TokenSource for Fixed {
|
||||
fn read(&self) -> Option<Zeroizing<String>> {
|
||||
Some(Zeroizing::new(self.0.to_string()))
|
||||
}
|
||||
}
|
||||
struct NoSessions;
|
||||
impl SessionStore for NoSessions {
|
||||
fn exists(&self, _: &str) -> bool {
|
||||
false
|
||||
}
|
||||
}
|
||||
struct Quiet;
|
||||
impl StderrSink for Quiet {
|
||||
fn stderr_tail(&self, _: &[u8]) {}
|
||||
}
|
||||
|
||||
let dir =
|
||||
std::env::temp_dir().join(format!("claude-connector-ask-tools-{}", std::process::id()));
|
||||
for sub in ["work", "config", "prompt"] {
|
||||
std::fs::create_dir_all(dir.join(sub)).unwrap();
|
||||
}
|
||||
let binary = dir.join("claude");
|
||||
std::fs::write(&binary, b"pretend claude").unwrap();
|
||||
let spawns = Arc::new(AtomicUsize::new(0));
|
||||
let backend = ClaudeCodeBackend::with_ports(
|
||||
ClaudeCodeConfig {
|
||||
claude_binary: binary,
|
||||
claude_binary_sha256: Sha256::digest(b"pretend claude").into(),
|
||||
work_dir: dir.join("work"),
|
||||
config_dir: dir.join("config"),
|
||||
prompt_dir: dir.join("prompt"),
|
||||
token_credential: "t".into(),
|
||||
handle_secret_credential: "h".into(),
|
||||
instance: "i".into(),
|
||||
limits: ClaudeCodeLimits {
|
||||
deadline: Duration::from_secs(5),
|
||||
max_line_bytes: 1 << 16,
|
||||
max_total_bytes: 1 << 20,
|
||||
stderr_tail_bytes: 64,
|
||||
},
|
||||
},
|
||||
Box::new(CountingPort(spawns.clone())),
|
||||
Box::new(NoSessions),
|
||||
Box::new(Quiet),
|
||||
Box::new(Fixed("token")),
|
||||
Box::new(Fixed(
|
||||
"handle-secret-0000000000000000000000000000000000000000",
|
||||
)),
|
||||
);
|
||||
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
let endpoint = listener.local_addr().unwrap();
|
||||
let admission =
|
||||
CallerAdmission::new(CALLER, KEY, 1, [MODEL.to_string()], 4, 1 << 20, 4_096).unwrap();
|
||||
let limits = ServeLimits {
|
||||
max_frame_bytes: FRAME,
|
||||
max_connections: 8,
|
||||
socket_timeout_ms: 5_000,
|
||||
max_expiry_skew_ms: 300_000,
|
||||
};
|
||||
let store = dir.join("replay.cc");
|
||||
thread::spawn(move || serve(listener, limits, vec![admission], &store, Arc::new(backend)));
|
||||
|
||||
let mut args = AskArgs {
|
||||
endpoint,
|
||||
caller_id: CALLER.into(),
|
||||
key_file: dir.join("unused"),
|
||||
model: MODEL.into(),
|
||||
instructions: "i".into(),
|
||||
effort: None,
|
||||
max_output_tokens: 64,
|
||||
timeout: Duration::from_secs(5),
|
||||
prompt: "p".into(),
|
||||
};
|
||||
let mut forced = build_invocation(&args).unwrap().request;
|
||||
forced.tools.push(ToolDefinition {
|
||||
name: "lookup".into(),
|
||||
description: "d".into(),
|
||||
parameters_json: "{}".into(),
|
||||
});
|
||||
let now = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_millis() as u64;
|
||||
let invocation = TransportInvocation::new(
|
||||
CALLER,
|
||||
now + 60_000,
|
||||
provider_request_sha256(&forced).unwrap(),
|
||||
forced,
|
||||
)
|
||||
.unwrap();
|
||||
let client = ConnectorClient::new(endpoint, KEY, FRAME, Some(Duration::from_secs(5))).unwrap();
|
||||
let result = client.execute(&invocation).unwrap();
|
||||
let TransportDisposition::Transported { receipt, .. } = &result.disposition else {
|
||||
panic!(
|
||||
"expected the planner's failed result, got {:?}",
|
||||
result.disposition
|
||||
);
|
||||
};
|
||||
assert!(matches!(
|
||||
receipt.outcome,
|
||||
TransportOutcome::Failed {
|
||||
kind: FailureKind::InvalidRequest,
|
||||
..
|
||||
}
|
||||
));
|
||||
assert_eq!(
|
||||
spawns.load(Ordering::SeqCst),
|
||||
0,
|
||||
"the claude child must never start"
|
||||
);
|
||||
// The same CLI path with no tools would have reached the spawn.
|
||||
args.key_file = dir.join("unused");
|
||||
let clean = client.execute(&build_invocation(&args).unwrap()).unwrap();
|
||||
assert!(matches!(
|
||||
clean.disposition,
|
||||
TransportDisposition::Transported { .. }
|
||||
));
|
||||
assert_eq!(spawns.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_binary_prints_the_answer_on_success() {
|
||||
let rig = rig(Script::Text("pong"));
|
||||
let out = run_binary(
|
||||
&rig,
|
||||
"ping",
|
||||
&["--effort", "low", "--max-output-tokens", "1024"],
|
||||
);
|
||||
assert!(
|
||||
out.status.success(),
|
||||
"{}",
|
||||
String::from_utf8_lossy(&out.stderr)
|
||||
);
|
||||
assert_eq!(String::from_utf8(out.stdout).unwrap(), "pong");
|
||||
assert!(out.stderr.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_binary_fails_loudly_with_one_stderr_line_and_empty_stdout() {
|
||||
let rig = rig(Script::Text("x"));
|
||||
let out = run_binary(&rig, "", &[]);
|
||||
assert!(!out.status.success());
|
||||
assert!(out.stdout.is_empty());
|
||||
let err = String::from_utf8(out.stderr).unwrap();
|
||||
assert_eq!(err.trim_end().lines().count(), 1, "{err}");
|
||||
assert!(err.starts_with("claude-connector: "));
|
||||
assert!(!err.contains(KEY));
|
||||
// Tools, a working directory and a second endpoint are not flags.
|
||||
for flag in ["--tools", "--cwd", "--work-dir"] {
|
||||
let out = run_binary(&rig, "q", &[flag, "x"]);
|
||||
assert!(!out.status.success());
|
||||
assert!(out.stdout.is_empty());
|
||||
}
|
||||
assert!(rig.seen().is_empty());
|
||||
}
|
||||
|
||||
/// Runs the binary with `--system-file` instead of `--system`, killing it if it
|
||||
/// has not exited within five seconds (a hang is a failure, not a stuck suite).
|
||||
fn run_with_system_file(rig: &Rig, file: &std::path::Path) -> std::process::Output {
|
||||
use std::io::Write;
|
||||
let mut child = Command::new(env!("CARGO_BIN_EXE_claude-connector"))
|
||||
.args(["ask", "--endpoint", &rig.endpoint.to_string()])
|
||||
.args(["--caller-id", CALLER, "--model", MODEL])
|
||||
.args(["--max-output-tokens", "1024"])
|
||||
.args(["--connection-key-file", rig.key_file.to_str().unwrap()])
|
||||
.args(["--system-file", file.to_str().unwrap()])
|
||||
.stdin(Stdio::piped())
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped())
|
||||
.spawn()
|
||||
.unwrap();
|
||||
let _ = child.stdin.take().unwrap().write_all(b"question");
|
||||
let started = std::time::Instant::now();
|
||||
while child.try_wait().unwrap().is_none() {
|
||||
if started.elapsed() > Duration::from_secs(5) {
|
||||
child.kill().unwrap();
|
||||
let _ = child.wait();
|
||||
panic!("the binary hung on its --system-file");
|
||||
}
|
||||
thread::sleep(Duration::from_millis(20));
|
||||
}
|
||||
child.wait_with_output().unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_binary_bounds_and_guards_its_system_file() {
|
||||
use claude_connector::ask::MAX_SYSTEM_BYTES;
|
||||
let rig = rig(Script::Text("accepted"));
|
||||
|
||||
let fifo = rig.dir.join("system.fifo");
|
||||
let path = std::ffi::CString::new(fifo.to_str().unwrap()).unwrap();
|
||||
assert_eq!(unsafe { libc::mkfifo(path.as_ptr(), 0o600) }, 0);
|
||||
let out = run_with_system_file(&rig, &fifo);
|
||||
assert!(!out.status.success());
|
||||
assert!(out.stdout.is_empty());
|
||||
|
||||
let zero = run_with_system_file(&rig, std::path::Path::new("/dev/zero"));
|
||||
assert!(!zero.status.success());
|
||||
assert!(zero.stdout.is_empty());
|
||||
|
||||
let file = rig.dir.join("system.txt");
|
||||
std::fs::write(&file, vec![b's'; MAX_SYSTEM_BYTES + 1]).unwrap();
|
||||
let over = run_with_system_file(&rig, &file);
|
||||
assert!(!over.status.success());
|
||||
assert!(over.stdout.is_empty());
|
||||
assert!(String::from_utf8_lossy(&over.stderr).contains("larger than"));
|
||||
assert!(rig.seen().is_empty());
|
||||
|
||||
std::fs::write(&file, vec![b's'; MAX_SYSTEM_BYTES]).unwrap();
|
||||
let exact = run_with_system_file(&rig, &file);
|
||||
assert!(
|
||||
exact.status.success(),
|
||||
"{}",
|
||||
String::from_utf8_lossy(&exact.stderr)
|
||||
);
|
||||
assert_eq!(String::from_utf8(exact.stdout).unwrap(), "accepted");
|
||||
assert_eq!(rig.seen()[0].instructions.len(), MAX_SYSTEM_BYTES);
|
||||
}
|
||||
297
crates/claude-connector/tests/config_document.rs
Normal file
297
crates/claude-connector/tests/config_document.rs
Normal file
|
|
@ -0,0 +1,297 @@
|
|||
//! The claude-connector's configuration document and binary: what the
|
||||
//! document accepts, what it refuses, and that a document of the dollar era
|
||||
//! fails closed.
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::process::{Command, Output};
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
|
||||
use claude_connector::{
|
||||
ClaudeConnectorConfig, ConfigError, LeaseCallerConfig, load_config, write_config,
|
||||
};
|
||||
use cultcache_rs::{CultCache, DatabaseEntry, SingleFileMessagePackBackingStore};
|
||||
use model_connector::{CallerConfig, LeaseTurn};
|
||||
|
||||
const OPUS: &str = "claude-opus-5-1";
|
||||
const HAIKU: &str = "claude-haiku-5-5";
|
||||
|
||||
static SEQUENCE: AtomicUsize = AtomicUsize::new(0);
|
||||
|
||||
struct Scratch(PathBuf);
|
||||
|
||||
impl Scratch {
|
||||
fn new() -> Self {
|
||||
let dir = std::env::temp_dir().join(format!(
|
||||
"claude-connector-config-{}-{}",
|
||||
std::process::id(),
|
||||
SEQUENCE.fetch_add(1, Ordering::SeqCst)
|
||||
));
|
||||
std::fs::create_dir_all(&dir).unwrap();
|
||||
Self(dir)
|
||||
}
|
||||
fn path(&self, name: &str) -> PathBuf {
|
||||
self.0.join(name)
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Scratch {
|
||||
fn drop(&mut self) {
|
||||
let _ = std::fs::remove_dir_all(&self.0);
|
||||
}
|
||||
}
|
||||
|
||||
fn valid() -> ClaudeConnectorConfig {
|
||||
ClaudeConnectorConfig {
|
||||
listen: "127.0.0.1:7443".to_string(),
|
||||
replay_store: "/var/lib/claude-connector/replay.cc".into(),
|
||||
claude_binary: "/opt/claude/bin/claude".into(),
|
||||
claude_binary_sha256: [7; 32],
|
||||
work_dir: "/var/lib/claude-connector/work".into(),
|
||||
config_dir: "/var/lib/claude-connector/home".into(),
|
||||
prompt_dir: "/var/lib/claude-connector/prompts".into(),
|
||||
instance: "persona-a".to_string(),
|
||||
token_credential: "claude-token".to_string(),
|
||||
handle_secret_credential: "claude-handle-secret".to_string(),
|
||||
deadline_ms: 600_000,
|
||||
max_line_bytes: 4 << 20,
|
||||
max_total_bytes: 64 << 20,
|
||||
stderr_tail_bytes: 8192,
|
||||
max_frame_bytes: 2 << 20,
|
||||
max_connections: 8,
|
||||
socket_timeout_ms: 10_000,
|
||||
max_expiry_skew_ms: 300_000,
|
||||
callers: vec![CallerConfig {
|
||||
caller_runtime_id: "persona-a-caller".to_string(),
|
||||
connection_key_file: "/etc/claude-connector/caller.key".into(),
|
||||
connection_key_epoch: 1,
|
||||
allowed_models: vec![HAIKU.to_string(), OPUS.to_string()],
|
||||
max_concurrent_requests: 4,
|
||||
max_payload_bytes: 1 << 20,
|
||||
max_output_tokens: 32_768,
|
||||
}],
|
||||
lease_socket: Some("/run/claude-connector/lease.sock".into()),
|
||||
lease_callers: vec![LeaseCallerConfig {
|
||||
unix_user: "mind-eureka-self".to_string(),
|
||||
caller_id: "persona-a-caller".to_string(),
|
||||
turn: LeaseTurn::SelfRun,
|
||||
unit_instance: "eureka".to_string(),
|
||||
}],
|
||||
lease_max_ms: 3_600_000,
|
||||
}
|
||||
}
|
||||
|
||||
fn refused(config: &ClaudeConnectorConfig) -> &'static str {
|
||||
match config.validate() {
|
||||
Err(ConfigError::Invalid(field)) => field,
|
||||
other => panic!("expected a refusal naming a field, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn config_round_trips_through_cultcache() {
|
||||
let scratch = Scratch::new();
|
||||
let path = scratch.path("connector.cc");
|
||||
let config = valid();
|
||||
write_config(&path, &config).unwrap();
|
||||
assert_eq!(load_config(&path).unwrap(), config);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_valid_document_validates() {
|
||||
assert!(valid().validate().is_ok());
|
||||
}
|
||||
|
||||
/// A document of the dollar era: the first type id, carrying a price table.
|
||||
#[derive(Clone, Debug, PartialEq, Eq, DatabaseEntry)]
|
||||
#[cultcache(type = "gamecult.model_connector.claude_connector_config.v1")]
|
||||
struct DollarEraConfig {
|
||||
#[cultcache(key = 0)]
|
||||
listen: String,
|
||||
#[cultcache(key = 18)]
|
||||
prices: BTreeMap<String, u64>,
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_dollar_era_config_is_refused_at_load() {
|
||||
let scratch = Scratch::new();
|
||||
let path = scratch.path("old.cc");
|
||||
{
|
||||
let mut cache = CultCache::new();
|
||||
cache.register_entry_type::<DollarEraConfig>().unwrap();
|
||||
cache
|
||||
.add_generic_backing_store(SingleFileMessagePackBackingStore::new(&path))
|
||||
.unwrap();
|
||||
cache
|
||||
.put(
|
||||
"runtime".to_string(),
|
||||
&DollarEraConfig {
|
||||
listen: "127.0.0.1:7443".to_string(),
|
||||
prices: BTreeMap::from([(OPUS.to_string(), 5_000_000)]),
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
assert!(matches!(load_config(&path), Err(ConfigError::Store)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_known_models_are_the_allowed_models_by_form_alone() {
|
||||
for good in [
|
||||
"claude-opus-5-1",
|
||||
"claude-haiku-5-5",
|
||||
"claude-opus-5",
|
||||
"claude-opus-5-1-20260101",
|
||||
] {
|
||||
let mut config = valid();
|
||||
config.callers[0].allowed_models = vec![good.to_string()];
|
||||
assert!(config.validate().is_ok(), "{good}");
|
||||
}
|
||||
for bad in [
|
||||
"sonnet",
|
||||
"claude-",
|
||||
"claude-sonnet-latest",
|
||||
"claude-Opus-5-1",
|
||||
"claude-opus-5-1-2026",
|
||||
"claude-opus-5-1-20260101-x",
|
||||
] {
|
||||
let mut config = valid();
|
||||
config.callers[0].allowed_models = vec![bad.to_string()];
|
||||
assert_eq!(refused(&config), "allowed_models", "{bad}");
|
||||
}
|
||||
let mut twice = valid();
|
||||
twice.callers[0].allowed_models = vec![OPUS.to_string(), OPUS.to_string()];
|
||||
assert_eq!(refused(&twice), "allowed_models");
|
||||
let mut none = valid();
|
||||
none.callers[0].allowed_models.clear();
|
||||
assert_eq!(refused(&none), "allowed_models");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn caller_bounds_are_required_nonzero() {
|
||||
let edits: [fn(&mut ClaudeConnectorConfig); 3] = [
|
||||
|c| c.callers[0].max_concurrent_requests = 0,
|
||||
|c| c.callers[0].max_output_tokens = 0,
|
||||
|c| c.callers[0].max_payload_bytes = 0,
|
||||
];
|
||||
for edit in edits {
|
||||
let mut config = valid();
|
||||
edit(&mut config);
|
||||
assert_eq!(refused(&config), "caller bounds");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_lease_caller_names_a_configured_caller_and_a_socket_exists_exactly_when_it_leases() {
|
||||
let mut stranger = valid();
|
||||
stranger.lease_callers[0].caller_id = "nobody".to_string();
|
||||
assert_eq!(refused(&stranger), "lease_callers");
|
||||
|
||||
let mut no_callers = valid();
|
||||
no_callers.lease_callers.clear();
|
||||
assert_eq!(refused(&no_callers), "lease_callers");
|
||||
|
||||
let mut no_socket = valid();
|
||||
no_socket.lease_socket = None;
|
||||
assert_eq!(refused(&no_socket), "lease_callers");
|
||||
no_socket.lease_callers.clear();
|
||||
assert!(no_socket.validate().is_ok());
|
||||
|
||||
let mut twice = valid();
|
||||
twice.lease_callers.push(twice.lease_callers[0].clone());
|
||||
assert_eq!(refused(&twice), "lease_callers");
|
||||
}
|
||||
|
||||
fn run(args: &[&str]) -> Output {
|
||||
Command::new(env!("CARGO_BIN_EXE_claude-connector"))
|
||||
.args(args)
|
||||
.env_clear()
|
||||
.output()
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn init_flags(scratch: &Scratch, extra: &[&str]) -> Vec<String> {
|
||||
let path = |name: &str| scratch.path(name).to_string_lossy().into_owned();
|
||||
let mut flags: Vec<String> = [
|
||||
"--listen",
|
||||
"127.0.0.1:0",
|
||||
"--replay-store",
|
||||
&path("replay.cc"),
|
||||
"--claude-binary",
|
||||
&path("claude"),
|
||||
"--claude-sha256",
|
||||
&"ab".repeat(32),
|
||||
"--work-dir",
|
||||
&path("work"),
|
||||
"--config-dir",
|
||||
&path("home"),
|
||||
"--prompt-dir",
|
||||
&path("prompts"),
|
||||
"--instance",
|
||||
"persona-a",
|
||||
"--token-credential",
|
||||
"claude-token",
|
||||
"--handle-secret-credential",
|
||||
"claude-handle-secret",
|
||||
"--handle-secret-file",
|
||||
&path("handle.secret"),
|
||||
"--caller-id",
|
||||
"persona-a-caller",
|
||||
"--connection-key-file",
|
||||
&path("caller.key"),
|
||||
"--connection-key-epoch",
|
||||
"1",
|
||||
"--allowed-model",
|
||||
HAIKU,
|
||||
]
|
||||
.map(String::from)
|
||||
.to_vec();
|
||||
flags.extend(extra.iter().map(|flag| flag.to_string()));
|
||||
flags
|
||||
}
|
||||
|
||||
fn init(scratch: &Scratch, flags: &[String]) -> Output {
|
||||
let config = scratch.path("connector.cc");
|
||||
let mut args = vec!["config", "init", config.to_str().unwrap()];
|
||||
args.extend(flags.iter().map(String::as_str));
|
||||
run(&args)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn config_init_writes_the_flags_into_the_document() {
|
||||
let scratch = Scratch::new();
|
||||
let made = init(&scratch, &init_flags(&scratch, &[]));
|
||||
assert!(made.status.success(), "{:?}", made);
|
||||
let config = load_config(&scratch.path("connector.cc")).unwrap();
|
||||
assert_eq!(config.callers[0].allowed_models, [HAIKU]);
|
||||
assert_eq!(config.claude_binary_sha256, [0xab; 32]);
|
||||
assert!(config.lease_callers.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_dollar_flags_are_not_flags() {
|
||||
let dollar_flags: [&[&str]; 4] = [
|
||||
&["--spend-ceiling-micro-usd", "50000000"],
|
||||
&["--spend-window-ms", "3600000"],
|
||||
&["--price", "claude-haiku-5-5:1:1:1:1:1"],
|
||||
&["--lease-reservation-micro-usd", "5000000"],
|
||||
];
|
||||
for extra in dollar_flags {
|
||||
let scratch = Scratch::new();
|
||||
let made = init(&scratch, &init_flags(&scratch, extra));
|
||||
assert!(!made.status.success(), "{extra:?} was accepted");
|
||||
assert!(
|
||||
!Path::new(&scratch.path("connector.cc")).exists(),
|
||||
"{extra:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serve_refuses_a_document_that_is_not_there() {
|
||||
let scratch = Scratch::new();
|
||||
let missing = scratch.path("missing.cc");
|
||||
let served = run(&["serve", "--config", missing.to_str().unwrap()]);
|
||||
assert!(!served.status.success(), "{served:?}");
|
||||
assert!(!String::from_utf8_lossy(&served.stderr).is_empty());
|
||||
}
|
||||
4
crates/claude-connector/tests/fixtures/claude-2.1.312-meter-call.jsonl
vendored
Normal file
4
crates/claude-connector/tests/fixtures/claude-2.1.312-meter-call.jsonl
vendored
Normal file
|
|
@ -0,0 +1,4 @@
|
|||
{"type":"system","subtype":"init","cwd":"/tmp/meter-probe.Ye68uq/work","session_id":"1c3d29e9-8277-4ee0-86df-efc86d994f01","tools":[],"mcp_servers":[],"model":"claude-haiku-5-5","permissionMode":"bypassPermissions","slash_commands":[],"apiKeySource":"none","claude_code_version":"2.1.312","output_style":"default","agents":["claude","Explore","general-purpose","Plan","statusline-setup"],"skills":[],"plugins":[{"name":"cc-plugin-sec-default","path":"builtin","source":"cc-plugin-sec-default@builtin"},{"name":"cc-plugin-agents-md","path":"builtin","source":"cc-plugin-agents-md@builtin"},{"name":"cc-plugin-telemetry","path":"builtin","source":"cc-plugin-telemetry@builtin"},{"name":"cc-plugin-plugin-authoring","path":"builtin","source":"cc-plugin-plugin-authoring@builtin"}],"capabilities":["interrupt_receipt_v1","interrupt_cancel_queued_v1","interrupt_send_now_v1","interrupt_requeue_taken_v1","msg_lifecycle_v1","request_marker_lists_v1","ask_user_question_unanswered_v1","lifecycle_wakeup_source_v1","sdk_mcp_tools_list_changed","sdk_mcp_manifests","mcp_read_resource_v1","mcp_tool_ui_meta_v1","ui_surface_v1","mcp_call_result_not_for_model_v1"],"analytics_disabled":false,"product_feedback_disabled":false,"uuid":"8ad7f8d7-2831-43f2-a221-8af6323dcbd9","memory_paths":{"auto":"/tmp/meter-probe.Ye68uq/home/projects/-tmp-meter-probe-Ye68uq-work/memory/"},"original_cwd":"/tmp/meter-probe.Ye68uq/work","additional_directories":[],"messaging_socket_path":"/tmp/cc-socks/2838104.sock","fast_mode_state":"off","fast_mode_disabled_reason":"sdk_opt_in_required","per_turn_effort_active":true,"view_mode":"default"}
|
||||
{"type":"assistant","message":{"model":"claude-haiku-5-5","id":"msg_011Cgfu1sQvRAVmjwWNHc58s","type":"message","role":"assistant","content":[{"type":"text","text":"ok"}],"container":null,"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":384,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"cache_creation":{"ephemeral_5m_input_tokens":0,"ephemeral_1h_input_tokens":0},"output_tokens":4,"service_tier":"standard","inference_geo":"global","speed":"standard"},"input_transformations":[],"diagnostics":null,"context_management":null},"parent_tool_use_id":null,"session_id":"1c3d29e9-8277-4ee0-86df-efc86d994f01","uuid":"f80db687-b098-415c-ac44-51c40a6cb63a","timestamp":"2026-11-03T11:14:35.024Z","request_id":"req_011Cgfu1sCXBsDrfz8a1LoMd"}
|
||||
{"type":"rate_limit_event","rate_limit_info":{"status":"allowed","resetsAt":1794261600,"rateLimitType":"seven_day","overageStatus":"rejected","overageDisabledReason":"org_level_disabled","isUsingOverage":false,"unifiedWindows":{"seven_day":{"utilization":0.06,"resetsAt":1794261600,"status":"allowed"}}},"uuid":"acb3b364-a218-4965-b967-f73028254f8b","session_id":"1c3d29e9-8277-4ee0-86df-efc86d994f01"}
|
||||
{"duration_api_ms":995,"stop_reason":"end_turn","session_id":"1c3d29e9-8277-4ee0-86df-efc86d994f01","total_cost_usd":0.000040400000000000006,"usage":{"input_tokens":384,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"output_tokens":4,"output_tokens_details":{"thinking_tokens":0},"server_tool_use":{"web_search_requests":0,"web_fetch_requests":0},"service_tier":"standard","cache_creation":{"ephemeral_1h_input_tokens":0,"ephemeral_5m_input_tokens":0},"inference_geo":"global","iterations":[{"input_tokens":384,"output_tokens":4,"cache_read_input_tokens":0,"cache_creation_input_tokens":0,"cache_creation":{"ephemeral_5m_input_tokens":0,"ephemeral_1h_input_tokens":0},"type":"message"}],"speed":"standard","fallback_credit":null},"modelUsage":{"claude-haiku-5-5":{"inputTokens":384,"outputTokens":4,"cacheReadInputTokens":0,"cacheCreationInputTokens":0,"webSearchRequests":0,"costUSD":0.000040400000000000006,"contextWindow":1000000,"maxOutputTokens":128000,"thinkingTokens":0,"cacheCreationEphemeral5mInputTokens":0,"cacheCreationEphemeral1hInputTokens":0,"canonicalModel":"claude-haiku-5-5","provider":"firstParty","costBasis":"list"}},"permission_denials":[],"terminal_reason":"completed","fast_mode_state":"off","fast_mode_disabled_reason":"sdk_opt_in_required","subagent_stats":{"spawned":0,"requested":{"background":0,"foreground":0,"unset":0},"started_in_background":0,"max_depth":0,"spawned_by_subagents":0,"completed":0,"failed":0,"killed":{"parent":0,"user":0,"system":0},"refused":{"depth_limit":0,"concurrency_limit":0,"budget":0},"by_type":{}},"safety_stops":0,"is_error":false,"num_turns":1,"subtype":"success","api_error_status":null,"result":"ok","ttft_ms":1210,"type":"result","duration_ms":1228,"uuid":"d0b1b31b-1428-45aa-b90e-d35903c9b91b","ttft_stream_ms":1209,"time_to_request_ms":232,"first_content_frame_ms":1209,"queued_turn_count":0,"result_index":0,"model_usage_since_last_result":{"claude-haiku-5-5":{"inputTokens":384,"outputTokens":4,"cacheReadInputTokens":0,"cacheCreationInputTokens":0,"webSearchRequests":0,"costUSD":0.000040400000000000006,"contextWindow":1000000,"maxOutputTokens":128000,"thinkingTokens":0,"cacheCreationEphemeral5mInputTokens":0,"cacheCreationEphemeral1hInputTokens":0,"canonicalModel":"claude-haiku-5-5","provider":"firstParty","costBasis":"list"}}}
|
||||
20
crates/claude-connector/tests/fixtures/lease_v2/README.md
vendored
Normal file
20
crates/claude-connector/tests/fixtures/lease_v2/README.md
vendored
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
# 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.
|
||||
|
||||
- `usage_readout_request.bin`: `claude_usage_readout_request`, an empty map; opens no lease.
|
||||
- `usage_readout.bin`: `claude_usage_readout`, the kit's `UsageReadout` (`schema_id`, `at_unix_ms`,
|
||||
`plan_quota` with `five_hour` and `seven_day` windows); no fd travels with it.
|
||||
|
||||
Huginn's lease client vendors these files from this repository and tests against the same
|
||||
bytes. Generator: `tests/lease_wire_golden.rs`; run it with `LEASE_V2_REGENERATE=1`, then
|
||||
review the diff: a changed byte is a wire change.
|
||||
BIN
crates/claude-connector/tests/fixtures/lease_v2/grant.bin
vendored
Normal file
BIN
crates/claude-connector/tests/fixtures/lease_v2/grant.bin
vendored
Normal file
Binary file not shown.
BIN
crates/claude-connector/tests/fixtures/lease_v2/refusal_busy.bin
vendored
Normal file
BIN
crates/claude-connector/tests/fixtures/lease_v2/refusal_busy.bin
vendored
Normal file
Binary file not shown.
BIN
crates/claude-connector/tests/fixtures/lease_v2/refusal_expired.bin
vendored
Normal file
BIN
crates/claude-connector/tests/fixtures/lease_v2/refusal_expired.bin
vendored
Normal file
Binary file not shown.
BIN
crates/claude-connector/tests/fixtures/lease_v2/refusal_provider_rejected.bin
vendored
Normal file
BIN
crates/claude-connector/tests/fixtures/lease_v2/refusal_provider_rejected.bin
vendored
Normal file
Binary file not shown.
BIN
crates/claude-connector/tests/fixtures/lease_v2/refusal_unknown_caller.bin
vendored
Normal file
BIN
crates/claude-connector/tests/fixtures/lease_v2/refusal_unknown_caller.bin
vendored
Normal file
Binary file not shown.
BIN
crates/claude-connector/tests/fixtures/lease_v2/request.bin
vendored
Normal file
BIN
crates/claude-connector/tests/fixtures/lease_v2/request.bin
vendored
Normal file
Binary file not shown.
BIN
crates/claude-connector/tests/fixtures/lease_v2/usage_readout.bin
vendored
Normal file
BIN
crates/claude-connector/tests/fixtures/lease_v2/usage_readout.bin
vendored
Normal file
Binary file not shown.
BIN
crates/claude-connector/tests/fixtures/lease_v2/usage_readout_request.bin
vendored
Normal file
BIN
crates/claude-connector/tests/fixtures/lease_v2/usage_readout_request.bin
vendored
Normal file
Binary file not shown.
728
crates/claude-connector/tests/lease.rs
Normal file
728
crates/claude-connector/tests/lease.rs
Normal file
|
|
@ -0,0 +1,728 @@
|
|||
//! 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, LeaseUsageReadoutRequest, PeerUnit, PeerUnits, QuotaCell, TokenSource,
|
||||
decode_frame, encode_frame, serve_leases,
|
||||
};
|
||||
use model_connector::{CallerAdmission, LeaseTurn, PlanQuota, TransportService, UsageReadout};
|
||||
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);
|
||||
}
|
||||
|
||||
/// What a usage readout request came back with.
|
||||
fn read_usage(harness: &Harness) -> Answer {
|
||||
let mut stream = UnixStream::connect(&harness.socket).unwrap();
|
||||
stream
|
||||
.set_read_timeout(Some(Duration::from_secs(10)))
|
||||
.unwrap();
|
||||
stream
|
||||
.write_all(&encode_frame(&LeaseMessage::ClaudeUsageReadoutRequest(
|
||||
LeaseUsageReadoutRequest {},
|
||||
)))
|
||||
.unwrap();
|
||||
let (raw, fds) = receive(&stream);
|
||||
let message = decode_frame(&raw);
|
||||
Answer {
|
||||
_stream: stream,
|
||||
raw,
|
||||
message,
|
||||
fds,
|
||||
}
|
||||
}
|
||||
|
||||
fn usage_of(answer: &Answer) -> UsageReadout {
|
||||
match &answer.message {
|
||||
Some(LeaseMessage::ClaudeUsageReadout(readout)) => readout.clone(),
|
||||
other => panic!("expected a usage readout, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_readout_answers_the_cells_windows_and_lends_nothing() {
|
||||
let harness = start(setup());
|
||||
let reset = now_ms() + 3_600_000;
|
||||
harness.quota.observe([
|
||||
window("five_hour", 420, "allowed", reset),
|
||||
window("seven_day", 999, "allowed_warning", reset + 86_400_000),
|
||||
]);
|
||||
let before = now_ms();
|
||||
let answer = read_usage(&harness);
|
||||
let readout = usage_of(&answer);
|
||||
assert!(answer.fds.is_empty(), "a readout lends no descriptor");
|
||||
assert_eq!(readout.plan_quota, harness.quota.snapshot().windows);
|
||||
assert_eq!(readout.plan_quota.len(), 2);
|
||||
assert!(readout.at_unix_ms >= before && readout.at_unix_ms <= now_ms());
|
||||
assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 0);
|
||||
// No lease was opened: the next request is granted.
|
||||
assert!(granted_ok(&ask(&harness, 60_000)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_peer_outside_the_lease_callers_still_reads_usage() {
|
||||
let harness = start(Setup {
|
||||
uid: own_uid().wrapping_add(1),
|
||||
..setup()
|
||||
});
|
||||
harness
|
||||
.quota
|
||||
.observe([window("five_hour", 420, "allowed", now_ms() + 3_600_000)]);
|
||||
let refused = ask(&harness, 60_000);
|
||||
assert_eq!(refusal_of(&refused), Some(LeaseRefusalKind::UnknownCaller));
|
||||
let readout = usage_of(&read_usage(&harness));
|
||||
assert_eq!(readout.plan_quota.len(), 1);
|
||||
assert_eq!(readout.plan_quota[0].utilization_permille, Some(420));
|
||||
assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_readout_beside_a_held_lease_leaves_it_held() {
|
||||
let harness = start(setup());
|
||||
let held = ask(&harness, 60_000);
|
||||
assert!(granted_ok(&held));
|
||||
assert_eq!(usage_of(&read_usage(&harness)).plan_quota.len(), 0);
|
||||
harness.enter(2);
|
||||
let second = ask(&harness, 60_000);
|
||||
assert_eq!(refusal_of(&second), Some(LeaseRefusalKind::Busy));
|
||||
assert!(granted_ok(&held));
|
||||
}
|
||||
134
crates/claude-connector/tests/lease_wire_golden.rs
Normal file
134
crates/claude-connector/tests/lease_wire_golden.rs
Normal file
|
|
@ -0,0 +1,134 @@
|
|||
//! 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,
|
||||
LeaseUsageReadoutRequest, decode_frame, encode_frame,
|
||||
};
|
||||
use model_connector::{PlanQuota, PlanQuotaSnapshot, UsageReadout};
|
||||
|
||||
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),
|
||||
],
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_usage_readout_frames_are_pinned() {
|
||||
let window = |name: &str, permille, status: &str, resets| PlanQuota {
|
||||
window: name.to_string(),
|
||||
utilization_permille: Some(permille),
|
||||
resets_at_unix_ms: resets,
|
||||
status: status.to_string(),
|
||||
observed_at_unix_ms: 1_793_300_000_000,
|
||||
};
|
||||
let request = LeaseMessage::ClaudeUsageReadoutRequest(LeaseUsageReadoutRequest {});
|
||||
let answer = LeaseMessage::ClaudeUsageReadout(UsageReadout::new(
|
||||
1_793_300_100_000,
|
||||
PlanQuotaSnapshot {
|
||||
windows: vec![
|
||||
window("five_hour", 420, "allowed", 1_793_307_000_000),
|
||||
window("seven_day", 999, "allowed_warning", 1_793_800_000_000),
|
||||
],
|
||||
},
|
||||
));
|
||||
assert_frame("usage_readout_request.bin", &request);
|
||||
assert_frame("usage_readout.bin", &answer);
|
||||
for (message, spelled) in [
|
||||
(&request, "claude_usage_readout_request"),
|
||||
(&answer, "claude_usage_readout"),
|
||||
] {
|
||||
let bytes = encode_frame(message);
|
||||
let value: rmpv::Value = rmp_serde::from_slice(&bytes[4..]).unwrap();
|
||||
let rmpv::Value::Map(outer) = value else {
|
||||
panic!("a frame is a map")
|
||||
};
|
||||
assert_eq!(outer[0].0.as_str(), Some("kind"));
|
||||
assert_eq!(outer[0].1.as_str(), Some(spelled));
|
||||
}
|
||||
}
|
||||
31
crates/model-connector/Cargo.toml
Normal file
31
crates/model-connector/Cargo.toml
Normal file
|
|
@ -0,0 +1,31 @@
|
|||
[package]
|
||||
name = "model-connector"
|
||||
description = "Provider-neutral connector contract, envelope, client, replay service and serve loop for GameCult minds"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
publish = false
|
||||
|
||||
[features]
|
||||
default = []
|
||||
daemon = ["dep:cultcache-rs"]
|
||||
|
||||
[dependencies]
|
||||
aes-gcm = "0.10"
|
||||
base64 = "0.22"
|
||||
cultnet-rs = { git = "https://github.com/GameCult/CultLib.git", rev = "a4eae9056152335662f1115f97bf35fb00ab6d23", package = "cultnet-rs" }
|
||||
cultcache-rs = { git = "https://github.com/GameCult/CultLib.git", rev = "a4eae9056152335662f1115f97bf35fb00ab6d23", package = "cultcache-rs", optional = true }
|
||||
rand_core = { version = "0.6", features = ["getrandom"] }
|
||||
rmp-serde = "1"
|
||||
rmpv = "1"
|
||||
schemars = "1"
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_bytes = "0.11"
|
||||
serde_json = "1"
|
||||
sha2 = "0.10"
|
||||
thiserror = "2"
|
||||
zeroize = "1"
|
||||
|
||||
[dev-dependencies]
|
||||
anyhow = "1"
|
||||
311
crates/model-connector/src/client.rs
Normal file
311
crates/model-connector/src/client.rs
Normal file
|
|
@ -0,0 +1,311 @@
|
|||
use std::net::{SocketAddr, TcpStream};
|
||||
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,
|
||||
};
|
||||
use crate::readout::UsageReadout;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ConnectorClient {
|
||||
endpoint: SocketAddr,
|
||||
security: TransportKey,
|
||||
max_frame_bytes: usize,
|
||||
response_timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
impl ConnectorClient {
|
||||
pub fn new(
|
||||
endpoint: SocketAddr,
|
||||
connection_key: impl Into<String>,
|
||||
max_frame_bytes: usize,
|
||||
response_timeout: Option<Duration>,
|
||||
) -> Result<Self, ConnectorClientError> {
|
||||
if endpoint.port() == 0
|
||||
|| endpoint.ip().is_unspecified()
|
||||
|| endpoint.ip().is_multicast()
|
||||
|| endpoint.ip() == std::net::Ipv4Addr::BROADCAST
|
||||
|| !(4096..=u32::MAX as usize).contains(&max_frame_bytes)
|
||||
|| response_timeout.is_some_and(|timeout| timeout.is_zero())
|
||||
{
|
||||
return Err(ConnectorClientError::InvalidConfig);
|
||||
}
|
||||
let connection_key = Zeroizing::new(connection_key.into());
|
||||
let security = TransportKey::from_connection_secret(connection_key.as_str())?;
|
||||
Ok(Self {
|
||||
endpoint,
|
||||
security,
|
||||
max_frame_bytes,
|
||||
response_timeout,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn execute(
|
||||
&self,
|
||||
invocation: &TransportInvocation,
|
||||
) -> Result<TransportResult, ConnectorClientError> {
|
||||
if invocation.request_id().len() > MAX_IDENTITY_BYTES {
|
||||
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))
|
||||
.min(Duration::from_secs(10));
|
||||
let mut stream = TcpStream::connect_timeout(&self.endpoint, connect_timeout)
|
||||
.map_err(ConnectorClientError::Connection)?;
|
||||
stream
|
||||
.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)
|
||||
.map_err(client_frame_error)?;
|
||||
read_transport_frame(&mut stream, self.max_frame_bytes).map_err(client_frame_error)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod client_endpoint_tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn client_admits_explicit_unicast_and_refuses_non_destinations() {
|
||||
for endpoint in ["127.0.0.1:4103", "172.19.0.1:4103", "[::1]:4103"] {
|
||||
assert!(
|
||||
ConnectorClient::new(
|
||||
endpoint.parse().unwrap(),
|
||||
"bounded-client-endpoint-key",
|
||||
64 * 1024,
|
||||
Some(Duration::from_secs(1)),
|
||||
)
|
||||
.is_ok(),
|
||||
"{endpoint} should be an admissible explicit destination"
|
||||
);
|
||||
}
|
||||
for endpoint in [
|
||||
"0.0.0.0:4103",
|
||||
"255.255.255.255:4103",
|
||||
"224.0.0.1:4103",
|
||||
"[::]:4103",
|
||||
"[ff02::1]:4103",
|
||||
"127.0.0.1:0",
|
||||
] {
|
||||
assert!(matches!(
|
||||
ConnectorClient::new(
|
||||
endpoint.parse().unwrap(),
|
||||
"bounded-client-endpoint-key",
|
||||
64 * 1024,
|
||||
Some(Duration::from_secs(1)),
|
||||
),
|
||||
Err(ConnectorClientError::InvalidConfig)
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum ConnectorClientError {
|
||||
#[error("invalid connector client configuration")]
|
||||
InvalidConfig,
|
||||
#[error("connector client connection failed")]
|
||||
Connection(#[source] std::io::Error),
|
||||
#[error("connector frame exceeded its bound")]
|
||||
FrameSize,
|
||||
#[error("request id is longer than the connector identity bound")]
|
||||
RequestIdTooLong,
|
||||
#[error(transparent)]
|
||||
Transport(#[from] ServiceError),
|
||||
}
|
||||
|
||||
fn client_frame_error(error: TransportFrameError) -> ConnectorClientError {
|
||||
match error {
|
||||
TransportFrameError::Connection(error) => ConnectorClientError::Connection(error),
|
||||
TransportFrameError::Size => ConnectorClientError::FrameSize,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod fixture_tests {
|
||||
use super::*;
|
||||
use crate::fixtures::*;
|
||||
use crate::{contract::*, envelope::*};
|
||||
|
||||
#[test]
|
||||
fn shared_client_transports_one_exact_encrypted_request_and_result() {
|
||||
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
let endpoint = listener.local_addr().unwrap();
|
||||
let invocation = invocation();
|
||||
let expected = invocation.clone();
|
||||
let security = security("epiphany-distinct-test-key");
|
||||
let server = std::thread::spawn(move || {
|
||||
let (mut stream, _) = listener.accept().unwrap();
|
||||
let request = read_transport_frame(&mut stream, 64 * 1024).unwrap();
|
||||
let envelope = decode_invocation_frame(&request).unwrap();
|
||||
assert_eq!(decrypt_invocation(&envelope, &security).unwrap(), expected);
|
||||
let response = encrypt_result(
|
||||
&TransportResult::refused(&expected, Refusal::ModelNotAllowed),
|
||||
&security,
|
||||
)
|
||||
.unwrap();
|
||||
write_transport_frame(
|
||||
&mut stream,
|
||||
&encode_result_frame(&response).unwrap(),
|
||||
64 * 1024,
|
||||
)
|
||||
.unwrap();
|
||||
});
|
||||
|
||||
let client = ConnectorClient::new(
|
||||
endpoint,
|
||||
"epiphany-distinct-test-key",
|
||||
64 * 1024,
|
||||
Some(Duration::from_secs(1)),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
client.execute(&invocation).unwrap().disposition,
|
||||
TransportDisposition::Refused(Refusal::ModelNotAllowed)
|
||||
);
|
||||
server.join().unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_over_long_request_id_is_refused_locally_before_any_byte_is_sent() {
|
||||
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
listener.set_nonblocking(true).unwrap();
|
||||
let client = ConnectorClient::new(
|
||||
listener.local_addr().unwrap(),
|
||||
"epiphany-distinct-test-key",
|
||||
64 * 1024,
|
||||
Some(Duration::from_secs(1)),
|
||||
)
|
||||
.unwrap();
|
||||
let mut long = invocation();
|
||||
long.request.request_id = "r".repeat(MAX_IDENTITY_BYTES + 1);
|
||||
assert!(matches!(
|
||||
client.execute(&long),
|
||||
Err(ConnectorClientError::RequestIdTooLong)
|
||||
));
|
||||
assert_eq!(
|
||||
listener.accept().unwrap_err().kind(),
|
||||
std::io::ErrorKind::WouldBlock,
|
||||
"the client must not even connect"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_client_measures_a_request_id_in_bytes() {
|
||||
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
listener.set_nonblocking(true).unwrap();
|
||||
let client = ConnectorClient::new(
|
||||
listener.local_addr().unwrap(),
|
||||
"epiphany-distinct-test-key",
|
||||
64 * 1024,
|
||||
Some(Duration::from_secs(1)),
|
||||
)
|
||||
.unwrap();
|
||||
// 86 characters, 258 bytes: under the bound by characters, over by bytes.
|
||||
let id = "\u{20AC}".repeat(86);
|
||||
assert_eq!((id.chars().count(), id.len()), (86, 258));
|
||||
let mut long = invocation();
|
||||
long.request.request_id = id;
|
||||
assert!(matches!(
|
||||
client.execute(&long),
|
||||
Err(ConnectorClientError::RequestIdTooLong)
|
||||
));
|
||||
assert_eq!(
|
||||
listener.accept().unwrap_err().kind(),
|
||||
std::io::ErrorKind::WouldBlock,
|
||||
"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();
|
||||
}
|
||||
}
|
||||
}
|
||||
1188
crates/model-connector/src/contract.rs
Normal file
1188
crates/model-connector/src/contract.rs
Normal file
File diff suppressed because it is too large
Load diff
1333
crates/model-connector/src/envelope.rs
Normal file
1333
crates/model-connector/src/envelope.rs
Normal file
File diff suppressed because it is too large
Load diff
205
crates/model-connector/src/fixtures.rs
Normal file
205
crates/model-connector/src/fixtures.rs
Normal file
|
|
@ -0,0 +1,205 @@
|
|||
//! Shared fixtures for the carried contract, envelope, client and service tests.
|
||||
use crate::contract::*;
|
||||
|
||||
#[cfg(feature = "daemon")]
|
||||
use crate::service::*;
|
||||
|
||||
#[cfg(feature = "daemon")]
|
||||
use std::collections::BTreeMap;
|
||||
#[cfg(feature = "daemon")]
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
#[cfg(feature = "daemon")]
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
#[cfg(feature = "daemon")]
|
||||
use cultcache_rs::{CacheBackingStore, CultCacheEnvelope, PushAllOptions};
|
||||
|
||||
#[cfg(feature = "daemon")]
|
||||
static TEST_STORE_SEQUENCE: AtomicUsize = AtomicUsize::new(0);
|
||||
|
||||
#[cfg(feature = "daemon")]
|
||||
pub(crate) struct TestStore {
|
||||
pub(crate) root: std::path::PathBuf,
|
||||
pub(crate) path: std::path::PathBuf,
|
||||
}
|
||||
|
||||
#[cfg(feature = "daemon")]
|
||||
impl TestStore {
|
||||
pub(crate) fn new() -> Self {
|
||||
let sequence = TEST_STORE_SEQUENCE.fetch_add(1, Ordering::Relaxed);
|
||||
let root = std::env::temp_dir().join(format!(
|
||||
"model-connector-replay-{}-{sequence}",
|
||||
std::process::id()
|
||||
));
|
||||
std::fs::create_dir_all(&root).unwrap();
|
||||
let path = root.join("replay.cc");
|
||||
Self { root, path }
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "daemon")]
|
||||
impl Drop for TestStore {
|
||||
fn drop(&mut self) {
|
||||
let _ = std::fs::remove_dir_all(&self.root);
|
||||
}
|
||||
}
|
||||
|
||||
/// An in-memory ledger that accepts a fixed number of writes and then refuses
|
||||
/// every one, so a test can make the replay store fail where it chooses. A
|
||||
/// store reopened from it sees what was written, like a restart.
|
||||
#[cfg(feature = "daemon")]
|
||||
pub(crate) struct FlakyStore {
|
||||
held: Arc<Mutex<BTreeMap<String, CultCacheEnvelope>>>,
|
||||
pushes_allowed: Arc<AtomicUsize>,
|
||||
}
|
||||
|
||||
#[cfg(feature = "daemon")]
|
||||
impl FlakyStore {
|
||||
pub(crate) fn new(pushes_allowed: usize) -> Self {
|
||||
Self {
|
||||
held: Arc::default(),
|
||||
pushes_allowed: Arc::new(AtomicUsize::new(pushes_allowed)),
|
||||
}
|
||||
}
|
||||
|
||||
/// The remaining writes this store accepts, shared with it after it has
|
||||
/// been handed to a service.
|
||||
pub(crate) fn allowance(&self) -> Arc<AtomicUsize> {
|
||||
self.pushes_allowed.clone()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "daemon")]
|
||||
impl CacheBackingStore for FlakyStore {
|
||||
fn pull_all(&self) -> anyhow::Result<Vec<CultCacheEnvelope>> {
|
||||
Ok(self.held.lock().unwrap().values().cloned().collect())
|
||||
}
|
||||
|
||||
fn push(&mut self, entry: &CultCacheEnvelope) -> anyhow::Result<()> {
|
||||
let allowed = self.pushes_allowed.load(Ordering::SeqCst);
|
||||
if allowed == 0 {
|
||||
return Err(anyhow::anyhow!("ledger write refused"));
|
||||
}
|
||||
self.pushes_allowed.store(allowed - 1, Ordering::SeqCst);
|
||||
self.held
|
||||
.lock()
|
||||
.unwrap()
|
||||
.insert(entry.key.clone(), entry.clone());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn delete(&mut self, entry: &CultCacheEnvelope) -> anyhow::Result<()> {
|
||||
self.held.lock().unwrap().remove(&entry.key);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn push_all(
|
||||
&mut self,
|
||||
entries: &[CultCacheEnvelope],
|
||||
_options: PushAllOptions,
|
||||
) -> anyhow::Result<()> {
|
||||
*self.held.lock().unwrap() = entries
|
||||
.iter()
|
||||
.map(|entry| (entry.key.clone(), entry.clone()))
|
||||
.collect();
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn request() -> ProviderRequest {
|
||||
let mut request = ProviderRequest::new(
|
||||
"request-1",
|
||||
"conversation-1",
|
||||
"gpt-5.4",
|
||||
"Return the typed result.",
|
||||
4_096,
|
||||
);
|
||||
request.input.push(InputItem::UserText {
|
||||
text: "Projected state".to_string(),
|
||||
});
|
||||
request.tools.push(ToolDefinition {
|
||||
name: "read_source".to_string(),
|
||||
description: "Read one admitted source.".to_string(),
|
||||
parameters_json: r#"{"type":"object"}"#.to_string(),
|
||||
});
|
||||
request
|
||||
}
|
||||
|
||||
pub(crate) fn invocation() -> TransportInvocation {
|
||||
TransportInvocation::new("epiphany-yggdrasil", 2_000, [7; 32], request()).unwrap()
|
||||
}
|
||||
|
||||
pub(crate) fn receipt(invocation: &TransportInvocation) -> TransportReceipt {
|
||||
TransportReceipt {
|
||||
schema_id: RECEIPT_SCHEMA_ID.to_string(),
|
||||
request_id: invocation.request_id().to_string(),
|
||||
caller_runtime_id: invocation.caller_runtime_id.clone(),
|
||||
native_request_sha256: invocation.native_request_sha256,
|
||||
provider_request_sha256: invocation.provider_request_sha256,
|
||||
model: invocation.request.model.clone(),
|
||||
provider: "fixture-sse".to_string(),
|
||||
foreign_reasoning_dropped: 0,
|
||||
outcome: TransportOutcome::Completed {
|
||||
provider_response_id: Some("response-1".to_string()),
|
||||
stop_reason: StopReason::EndTurn,
|
||||
usage: Usage {
|
||||
input_tokens: 10,
|
||||
cache_read_tokens: 0,
|
||||
cache_write_tokens: 0,
|
||||
output_tokens: 3,
|
||||
reasoning_tokens: Some(1),
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "daemon")]
|
||||
pub(crate) fn caller(
|
||||
caller_runtime_id: &str,
|
||||
key: &str,
|
||||
max_concurrent_requests: usize,
|
||||
) -> CallerAdmission {
|
||||
CallerAdmission::new(
|
||||
caller_runtime_id,
|
||||
key,
|
||||
1,
|
||||
["gpt-5.4".to_string()],
|
||||
max_concurrent_requests,
|
||||
64 * 1024,
|
||||
32_768,
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[cfg(feature = "daemon")]
|
||||
pub(crate) fn open_service(store: &TestStore, max_concurrent_requests: usize) -> TransportService {
|
||||
TransportService::open(
|
||||
&store.path,
|
||||
"fixture-sse",
|
||||
[
|
||||
caller(
|
||||
"epiphany-yggdrasil",
|
||||
"epiphany-distinct-test-key",
|
||||
max_concurrent_requests,
|
||||
),
|
||||
caller(
|
||||
"ghostlight-yggdrasil",
|
||||
"ghostlight-distinct-test-key",
|
||||
max_concurrent_requests,
|
||||
),
|
||||
],
|
||||
5_000,
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
pub(crate) fn security(key: &str) -> TransportKey {
|
||||
TransportKey::from_connection_secret(key).unwrap()
|
||||
}
|
||||
|
||||
#[cfg(feature = "daemon")]
|
||||
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()
|
||||
}
|
||||
22
crates/model-connector/src/lib.rs
Normal file
22
crates/model-connector/src/lib.rs
Normal file
|
|
@ -0,0 +1,22 @@
|
|||
mod client;
|
||||
mod contract;
|
||||
mod envelope;
|
||||
mod readout;
|
||||
|
||||
#[cfg(test)]
|
||||
mod fixtures;
|
||||
|
||||
#[cfg(feature = "daemon")]
|
||||
mod serve;
|
||||
#[cfg(feature = "daemon")]
|
||||
mod service;
|
||||
|
||||
pub use client::*;
|
||||
pub use contract::*;
|
||||
pub use envelope::*;
|
||||
pub use readout::*;
|
||||
|
||||
#[cfg(feature = "daemon")]
|
||||
pub use serve::*;
|
||||
#[cfg(feature = "daemon")]
|
||||
pub use service::*;
|
||||
106
crates/model-connector/src/readout.rs
Normal file
106
crates/model-connector/src/readout.rs
Normal file
|
|
@ -0,0 +1,106 @@
|
|||
//! The usage readout: the subscription plan's own usage windows as the backend
|
||||
//! last saw them. The connector meters nothing in dollars; it publishes what
|
||||
//! the provider reports, and the minds read it and decide their own burn.
|
||||
|
||||
use schemars::JsonSchema;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::contract::ContractError;
|
||||
|
||||
pub const USAGE_READOUT_SCHEMA_ID: &str = "gamecult.model.usage_readout.v1";
|
||||
|
||||
/// One window of the subscription plan, as the provider reported it.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct PlanQuota {
|
||||
pub window: String,
|
||||
/// How much of the window was used when it was observed, in thousandths;
|
||||
/// `None` once the window has reset (`status` is then `expired`), because
|
||||
/// the old reading says nothing about the new window.
|
||||
pub utilization_permille: Option<u32>,
|
||||
pub resets_at_unix_ms: u64,
|
||||
pub status: String,
|
||||
/// When this window was last reported. Its age in a readout is the
|
||||
/// readout's `at_unix_ms` minus this.
|
||||
pub observed_at_unix_ms: u64,
|
||||
}
|
||||
|
||||
/// What a backend last saw of the plan's windows, each stamped with its own
|
||||
/// observation time. A backend with no plan to report returns the default.
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct PlanQuotaSnapshot {
|
||||
pub windows: Vec<PlanQuota>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct UsageReadout {
|
||||
pub schema_id: String,
|
||||
pub at_unix_ms: u64,
|
||||
pub plan_quota: Vec<PlanQuota>,
|
||||
}
|
||||
|
||||
impl UsageReadout {
|
||||
pub fn new(at_unix_ms: u64, plan: PlanQuotaSnapshot) -> Self {
|
||||
Self {
|
||||
schema_id: USAGE_READOUT_SCHEMA_ID.to_string(),
|
||||
at_unix_ms,
|
||||
plan_quota: plan.windows,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn validate(&self) -> Result<(), ContractError> {
|
||||
if self.schema_id != USAGE_READOUT_SCHEMA_ID {
|
||||
return Err(ContractError::Schema);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::contract::{wire_decode, wire_encode};
|
||||
|
||||
fn readout() -> UsageReadout {
|
||||
UsageReadout::new(
|
||||
1_000,
|
||||
PlanQuotaSnapshot {
|
||||
windows: vec![PlanQuota {
|
||||
window: "five_hour".into(),
|
||||
utilization_permille: Some(80),
|
||||
resets_at_unix_ms: 1_793_307_000_000,
|
||||
status: "allowed".into(),
|
||||
observed_at_unix_ms: 999,
|
||||
}],
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_readout_is_a_named_map_that_decodes_to_itself() {
|
||||
let bytes = wire_encode(&readout()).unwrap();
|
||||
let value: rmpv::Value = rmp_serde::from_slice(&bytes).unwrap();
|
||||
let rmpv::Value::Map(fields) = value else {
|
||||
panic!("a readout is a map, not a positional array");
|
||||
};
|
||||
let names: Vec<&str> = fields.iter().filter_map(|(k, _)| k.as_str()).collect();
|
||||
assert_eq!(names, ["schema_id", "at_unix_ms", "plan_quota"]);
|
||||
assert_eq!(wire_decode::<UsageReadout>(&bytes).unwrap(), readout());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_unobserved_plan_is_none_not_zero() {
|
||||
let bare = UsageReadout::new(1, PlanQuotaSnapshot::default());
|
||||
let decoded: UsageReadout = wire_decode(&wire_encode(&bare).unwrap()).unwrap();
|
||||
assert!(decoded.plan_quota.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_readout_of_another_schema_is_refused() {
|
||||
let mut other = readout();
|
||||
other.schema_id = "gamecult.model.usage_readout.v2".into();
|
||||
assert!(other.validate().is_err());
|
||||
assert!(readout().validate().is_ok());
|
||||
}
|
||||
}
|
||||
1449
crates/model-connector/src/serve.rs
Normal file
1449
crates/model-connector/src/serve.rs
Normal file
File diff suppressed because it is too large
Load diff
2183
crates/model-connector/src/service.rs
Normal file
2183
crates/model-connector/src/service.rs
Normal file
File diff suppressed because it is too large
Load diff
429
crates/model-connector/src/service/lease.rs
Normal file
429
crates/model-connector/src/service/lease.rs
Normal file
|
|
@ -0,0 +1,429 @@
|
|||
//! A credential lease: a Self run's or a Persona turn's hold on the token the
|
||||
//! connector lends. The lease is memory bound to its holder's open connection,
|
||||
//! not a record and not a charge: it decides `Busy` while it is held and
|
||||
//! nothing else. It ends when its holder lets go (`end_lease`), when the unit
|
||||
//! invocation it served ends, or when its expiry passes. Nothing is billed and
|
||||
//! nothing is reserved. The only usage-derived refusal is the provider's own
|
||||
//! (`ProviderRejected`), repeated with the provider's own expiry.
|
||||
|
||||
use super::*;
|
||||
|
||||
/// The request-id namespace of leases.
|
||||
const LEASE_PREFIX: &str = "lease:";
|
||||
|
||||
/// What a Persona turn's lease id adds to its invocation id, so a Persona
|
||||
/// lease and a Self lease of one caller and one invocation never share an id.
|
||||
const PERSONA_SUFFIX: &str = ":persona";
|
||||
|
||||
/// Which kind of turn a lease serves. It is a property of the lease caller's
|
||||
/// configured row, never of a request. A caller holds one Self lease at a
|
||||
/// time; a Persona turn holds one lease per live invocation, one at a time.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub enum LeaseTurn {
|
||||
SelfRun,
|
||||
PersonaTurn,
|
||||
}
|
||||
|
||||
/// What an admitted lease holds.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct LeaseGrant {
|
||||
pub lease_id: String,
|
||||
}
|
||||
|
||||
/// Why no lease was granted.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum LeaseRefusal {
|
||||
UnknownCaller,
|
||||
/// A Self run: the caller holds a Self lease now. A Persona turn: this
|
||||
/// invocation holds a lease now.
|
||||
Busy,
|
||||
/// The provider is refusing right now and has said until when.
|
||||
ProviderRejected,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum LeaseAdmission {
|
||||
Granted(LeaseGrant),
|
||||
Refused(LeaseRefusal),
|
||||
}
|
||||
|
||||
impl TransportService {
|
||||
/// Grants `caller_runtime_id` a lease for the unit invocation
|
||||
/// `invocation_id` until `expires_at_unix_ms`. A Self run's lease is
|
||||
/// `lease:<invocation>` and one is open per caller; a Persona turn's is
|
||||
/// `lease:<invocation>:persona`, one open per invocation; neither blocks
|
||||
/// the other. A lease takes no request slot. `plan` is the backend's plan
|
||||
/// quota as it stands: a provider that is refusing right now refuses the
|
||||
/// lease. A lease past its expiry no longer holds anything.
|
||||
pub fn begin_lease(
|
||||
&mut self,
|
||||
caller_runtime_id: &str,
|
||||
turn: LeaseTurn,
|
||||
invocation_id: &str,
|
||||
expires_at_unix_ms: u64,
|
||||
now_unix_ms: u64,
|
||||
plan: &PlanQuotaSnapshot,
|
||||
) -> Result<LeaseAdmission, ServiceError> {
|
||||
if !self.callers.contains_key(caller_runtime_id) {
|
||||
return Ok(LeaseAdmission::Refused(LeaseRefusal::UnknownCaller));
|
||||
}
|
||||
require_id(invocation_id, "invocation_id")?;
|
||||
let suffix = match turn {
|
||||
LeaseTurn::SelfRun => "",
|
||||
LeaseTurn::PersonaTurn => PERSONA_SUFFIX,
|
||||
};
|
||||
let lease_id = format!("{LEASE_PREFIX}{invocation_id}{suffix}");
|
||||
if lease_id.len() > MAX_IDENTITY_BYTES {
|
||||
return Err(ContractError::Invalid("invocation_id").into());
|
||||
}
|
||||
self.leases
|
||||
.retain(|_, lease| lease.expires_at_unix_ms > now_unix_ms);
|
||||
let identity = RequestIdentity {
|
||||
caller_runtime_id: caller_runtime_id.to_string(),
|
||||
request_id: lease_id.clone(),
|
||||
};
|
||||
let busy = match turn {
|
||||
LeaseTurn::SelfRun => self.leases.iter().any(|(held, lease)| {
|
||||
lease.turn == LeaseTurn::SelfRun && held.caller_runtime_id == caller_runtime_id
|
||||
}),
|
||||
LeaseTurn::PersonaTurn => self.leases.contains_key(&identity),
|
||||
};
|
||||
if busy {
|
||||
return Ok(LeaseAdmission::Refused(LeaseRefusal::Busy));
|
||||
}
|
||||
if provider_rejects(plan, now_unix_ms) {
|
||||
return Ok(LeaseAdmission::Refused(LeaseRefusal::ProviderRejected));
|
||||
}
|
||||
self.leases.insert(
|
||||
identity,
|
||||
LiveLease {
|
||||
turn,
|
||||
expires_at_unix_ms,
|
||||
},
|
||||
);
|
||||
Ok(LeaseAdmission::Granted(LeaseGrant { lease_id }))
|
||||
}
|
||||
|
||||
/// Ends a lease its holder let go of. Ending one that is not held is a
|
||||
/// no-op: a holder may close after its expiry already freed the lease.
|
||||
pub fn end_lease(&mut self, caller_runtime_id: &str, lease_id: &str) {
|
||||
self.leases.remove(&RequestIdentity {
|
||||
caller_runtime_id: caller_runtime_id.to_string(),
|
||||
request_id: lease_id.to_string(),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::fixtures::*;
|
||||
use crate::readout::PlanQuota;
|
||||
|
||||
const CALLER: &str = "eureka";
|
||||
const KEY: &str = "lease-distinct-test-key";
|
||||
const NOW: u64 = 7_200_500;
|
||||
|
||||
fn open(store: &TestStore) -> TransportService {
|
||||
TransportService::open(
|
||||
&store.path,
|
||||
"fixture-sse",
|
||||
[
|
||||
CallerAdmission::new(CALLER, KEY, 1, ["gpt-5.4".to_string()], 1, 65_536, 32_768)
|
||||
.unwrap(),
|
||||
],
|
||||
5_000,
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn none() -> PlanQuotaSnapshot {
|
||||
PlanQuotaSnapshot::default()
|
||||
}
|
||||
|
||||
fn window(permille: u32, status: &str, resets_at: u64) -> PlanQuotaSnapshot {
|
||||
PlanQuotaSnapshot {
|
||||
windows: vec![PlanQuota {
|
||||
window: "five_hour".to_string(),
|
||||
utilization_permille: Some(permille),
|
||||
resets_at_unix_ms: resets_at,
|
||||
status: status.to_string(),
|
||||
observed_at_unix_ms: NOW,
|
||||
}],
|
||||
}
|
||||
}
|
||||
|
||||
fn lease_in(
|
||||
service: &mut TransportService,
|
||||
turn: LeaseTurn,
|
||||
run: &str,
|
||||
plan: &PlanQuotaSnapshot,
|
||||
) -> LeaseAdmission {
|
||||
service
|
||||
.begin_lease(CALLER, turn, run, NOW + 60_000, NOW, plan)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn lease(service: &mut TransportService, run: &str) -> LeaseAdmission {
|
||||
lease_in(service, LeaseTurn::SelfRun, run, &none())
|
||||
}
|
||||
|
||||
fn granted(admission: LeaseAdmission) -> LeaseGrant {
|
||||
match admission {
|
||||
LeaseAdmission::Granted(grant) => grant,
|
||||
LeaseAdmission::Refused(why) => panic!("refused {why:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_lease_is_named_by_its_invocation() {
|
||||
let store = TestStore::new();
|
||||
let mut service = open(&store);
|
||||
assert_eq!(
|
||||
granted(lease(&mut service, "run-1")).lease_id,
|
||||
"lease:run-1"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_caller_holds_one_self_lease_at_a_time() {
|
||||
let store = TestStore::new();
|
||||
let mut service = open(&store);
|
||||
granted(lease(&mut service, "run-1"));
|
||||
assert_eq!(
|
||||
lease(&mut service, "run-2"),
|
||||
LeaseAdmission::Refused(LeaseRefusal::Busy)
|
||||
);
|
||||
assert_eq!(
|
||||
lease(&mut service, "run-1"),
|
||||
LeaseAdmission::Refused(LeaseRefusal::Busy)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_lease_ends_when_its_holder_closes() {
|
||||
let store = TestStore::new();
|
||||
let mut service = open(&store);
|
||||
let grant = granted(lease(&mut service, "run-1"));
|
||||
service.end_lease(CALLER, &grant.lease_id);
|
||||
// The same invocation may lease again: nothing remembers a closed
|
||||
// lease, and nothing was billed for it.
|
||||
granted(lease(&mut service, "run-1"));
|
||||
service.end_lease(CALLER, "lease:not-held");
|
||||
assert_eq!(
|
||||
lease(&mut service, "run-2"),
|
||||
LeaseAdmission::Refused(LeaseRefusal::Busy),
|
||||
"ending a lease that is not held frees nothing else"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_lease_ends_at_its_expiry_without_its_holder() {
|
||||
let store = TestStore::new();
|
||||
let mut service = open(&store);
|
||||
granted(
|
||||
service
|
||||
.begin_lease(CALLER, LeaseTurn::SelfRun, "run-1", NOW + 10, NOW, &none())
|
||||
.unwrap(),
|
||||
);
|
||||
assert_eq!(
|
||||
service
|
||||
.begin_lease(
|
||||
CALLER,
|
||||
LeaseTurn::SelfRun,
|
||||
"run-2",
|
||||
NOW + 99,
|
||||
NOW + 9,
|
||||
&none()
|
||||
)
|
||||
.unwrap(),
|
||||
LeaseAdmission::Refused(LeaseRefusal::Busy),
|
||||
"one millisecond before its expiry the lease still holds"
|
||||
);
|
||||
granted(
|
||||
service
|
||||
.begin_lease(
|
||||
CALLER,
|
||||
LeaseTurn::SelfRun,
|
||||
"run-2",
|
||||
NOW + 99,
|
||||
NOW + 10,
|
||||
&none(),
|
||||
)
|
||||
.unwrap(),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_stranger_is_refused() {
|
||||
let store = TestStore::new();
|
||||
let mut service = open(&store);
|
||||
assert_eq!(
|
||||
service
|
||||
.begin_lease(
|
||||
"stranger",
|
||||
LeaseTurn::SelfRun,
|
||||
"run-1",
|
||||
NOW + 1,
|
||||
NOW,
|
||||
&none()
|
||||
)
|
||||
.unwrap(),
|
||||
LeaseAdmission::Refused(LeaseRefusal::UnknownCaller)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_persona_holds_one_lease_per_invocation_and_never_blocks_the_self_lease() {
|
||||
let store = TestStore::new();
|
||||
let mut service = open(&store);
|
||||
let persona = |service: &mut TransportService, run: &str| {
|
||||
lease_in(service, LeaseTurn::PersonaTurn, run, &none())
|
||||
};
|
||||
let first = granted(persona(&mut service, "turn-1"));
|
||||
assert_eq!(first.lease_id, "lease:turn-1:persona");
|
||||
assert_eq!(
|
||||
persona(&mut service, "turn-1"),
|
||||
LeaseAdmission::Refused(LeaseRefusal::Busy)
|
||||
);
|
||||
granted(persona(&mut service, "turn-2"));
|
||||
granted(lease(&mut service, "turn-1"));
|
||||
service.end_lease(CALLER, &first.lease_id);
|
||||
granted(persona(&mut service, "turn-1"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_rejected_window_refuses_invoke_and_lease_until_its_reset() {
|
||||
let store = TestStore::new();
|
||||
let mut service = open(&store);
|
||||
let rejected = window(1_000, "rejected", NOW + 5_000);
|
||||
assert_eq!(
|
||||
lease_in(&mut service, LeaseTurn::SelfRun, "run-1", &rejected),
|
||||
LeaseAdmission::Refused(LeaseRefusal::ProviderRejected)
|
||||
);
|
||||
let call = |service: &mut TransportService, now: u64, plan: &PlanQuotaSnapshot| {
|
||||
let mut asked = request();
|
||||
asked.request_id = format!("call-{now}");
|
||||
let invocation = TransportInvocation::new(CALLER, now + 1_000, [7; 32], asked).unwrap();
|
||||
let sealed = crate::envelope::encrypt_message(
|
||||
CALLER,
|
||||
invocation.request_id(),
|
||||
crate::envelope::EnvelopeKind::Invocation,
|
||||
&crate::contract::wire_encode(&invocation).unwrap(),
|
||||
&security(KEY),
|
||||
)
|
||||
.unwrap();
|
||||
let TransportAdmission::Reply(reply) =
|
||||
service.begin(&sealed, now, plan, 1 << 20).unwrap()
|
||||
else {
|
||||
return None;
|
||||
};
|
||||
let result =
|
||||
crate::envelope::decrypt_result(&reply, &security(KEY), &invocation).unwrap();
|
||||
match result.disposition {
|
||||
TransportDisposition::Refused(refusal) => Some(refusal),
|
||||
other => panic!("unexpected {other:?}"),
|
||||
}
|
||||
};
|
||||
assert_eq!(
|
||||
call(&mut service, NOW, &rejected),
|
||||
Some(Refusal::ProviderRejected)
|
||||
);
|
||||
assert_eq!(
|
||||
call(&mut service, NOW + 4_999, &rejected),
|
||||
Some(Refusal::ProviderRejected)
|
||||
);
|
||||
assert_eq!(call(&mut service, NOW + 5_000, &rejected), None);
|
||||
assert!(matches!(
|
||||
service
|
||||
.begin_lease(
|
||||
CALLER,
|
||||
LeaseTurn::SelfRun,
|
||||
"run-1",
|
||||
NOW + 9_000,
|
||||
NOW + 5_000,
|
||||
&rejected
|
||||
)
|
||||
.unwrap(),
|
||||
LeaseAdmission::Granted(_)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn high_utilization_refuses_nothing() {
|
||||
let store = TestStore::new();
|
||||
let mut service = open(&store);
|
||||
let full = window(999, "allowed", NOW + 5_000);
|
||||
granted(lease_in(&mut service, LeaseTurn::SelfRun, "run-1", &full));
|
||||
let warned = window(1_000, "allowed_warning", NOW + 5_000);
|
||||
granted(lease_in(
|
||||
&mut service,
|
||||
LeaseTurn::PersonaTurn,
|
||||
"turn-1",
|
||||
&warned,
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_invocation_id_is_bounded_so_its_lease_id_fits_an_identity() {
|
||||
let store = TestStore::new();
|
||||
let mut service = open(&store);
|
||||
let longest = "r".repeat(MAX_IDENTITY_BYTES - LEASE_PREFIX.len());
|
||||
let grant = granted(lease(&mut service, &longest));
|
||||
assert_eq!(grant.lease_id.len(), MAX_IDENTITY_BYTES);
|
||||
assert!(lease_try(&mut service, LeaseTurn::SelfRun, &format!("{longest}r")).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_persona_invocation_id_leaves_room_for_its_suffix() {
|
||||
let store = TestStore::new();
|
||||
let mut service = open(&store);
|
||||
let longest = "r".repeat(MAX_IDENTITY_BYTES - LEASE_PREFIX.len() - PERSONA_SUFFIX.len());
|
||||
let grant = granted(lease_in(
|
||||
&mut service,
|
||||
LeaseTurn::PersonaTurn,
|
||||
&longest,
|
||||
&none(),
|
||||
));
|
||||
assert_eq!(grant.lease_id.len(), MAX_IDENTITY_BYTES);
|
||||
assert!(lease_try(&mut service, LeaseTurn::PersonaTurn, &format!("{longest}r")).is_err());
|
||||
}
|
||||
|
||||
fn lease_try(
|
||||
service: &mut TransportService,
|
||||
turn: LeaseTurn,
|
||||
run: &str,
|
||||
) -> Result<LeaseAdmission, ServiceError> {
|
||||
service.begin_lease(CALLER, turn, run, NOW + 60_000, NOW, &none())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_lease_takes_no_request_slot_from_the_callers_calls() {
|
||||
// The caller admits one concurrent call; an open lease must not use it.
|
||||
let store = TestStore::new();
|
||||
let mut service = open(&store);
|
||||
granted(lease(&mut service, "run-1"));
|
||||
let mut call = request();
|
||||
call.request_id = "call-1".to_string();
|
||||
let invocation = TransportInvocation::new(CALLER, NOW + 1_000, [7; 32], call).unwrap();
|
||||
let sealed = crate::envelope::encrypt_message(
|
||||
CALLER,
|
||||
"call-1",
|
||||
crate::envelope::EnvelopeKind::Invocation,
|
||||
&crate::contract::wire_encode(&invocation).unwrap(),
|
||||
&security(KEY),
|
||||
)
|
||||
.unwrap();
|
||||
assert!(matches!(
|
||||
service.begin(&sealed, NOW, &none(), 1 << 20).unwrap(),
|
||||
TransportAdmission::Execute(_)
|
||||
));
|
||||
// And a call in flight does not make the caller busy for a lease.
|
||||
granted(lease_in(
|
||||
&mut service,
|
||||
LeaseTurn::PersonaTurn,
|
||||
"turn-1",
|
||||
&none(),
|
||||
));
|
||||
}
|
||||
}
|
||||
16
crates/model-connector/tests/fixtures/wire_v4/README.md
vendored
Normal file
16
crates/model-connector/tests/fixtures/wire_v4/README.md
vendored
Normal file
|
|
@ -0,0 +1,16 @@
|
|||
# wire_v4 fixtures
|
||||
|
||||
The named-MessagePack bytes of the v4 documents and one sealed carriage frame,
|
||||
produced by the kit at the revision that added them. They pin drift; the types in
|
||||
`src/contract.rs` own the shape, and `schemas/` publishes it.
|
||||
|
||||
- `provider_request.mp`, `transport_invocation.mp`, `transport_result_transported.mp`,
|
||||
`transport_result_refused.mp`, `transport_receipt_failed.mp`, `transport_envelope.mp`,
|
||||
`input_reasoning_state.mp` (an `InputItem::ReasoningState`, its state as MessagePack `bin`):
|
||||
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`.
|
||||
|
||||
Generator: `tests/wire_golden_v4.rs` itself. Run it with `WIRE_V4_REGENERATE=1` to rewrite
|
||||
the files from the kit, then review the diff: a changed byte is a wire change.
|
||||
1
crates/model-connector/tests/fixtures/wire_v4/carriage_invocation.frame
vendored
Normal file
1
crates/model-connector/tests/fixtures/wire_v4/carriage_invocation.frame
vendored
Normal file
|
|
@ -0,0 +1 @@
|
|||
<EFBFBD>messageIdョgolden-requestゥoperationヲinvokeァpayloadレ(hqlzY2hlbWFfaWTZJGdhbWVjdWx0Lm1vZGVsLnRyYW5zcG9ydF9lbnZlbG9wZS52NLFjYWxsZXJfcnVudGltZV9pZK1nb2xkZW4tY2FsbGVyqnJlcXVlc3RfaWSuZ29sZGVuLXJlcXVlc3SsbWVzc2FnZV9raW5kqkludm9jYXRpb26lbm9uY2WcXSbMzU3MhsyCzLAtzKLMhcy8zMSqY2lwaGVydGV4dMUDMY816tn0FGQ9jqxE+KyDVcxmhufaQI7vzEXO8NXTQPOYpRuMcRDM/j/QzNF2Lb3Fak0G+hHRKrXDxJ3MeKpoELNYOEyxbKfxn88mjAyyc7BVPM7+0Bk9ye3m5oBF6hoyCgOHSTwESIfq0bim+9oQh1IDp+HdFjztK/XjM731GRzKJmGgxsKqQsz42zq8y9mjAXl3V1RR9ssdKvk2k5U56/3liOeHDU96kOa7IlVCBm6qnn06hg2DTtnBF4xGj6PihHj4KufVkAcuQhRCZWOKsploToXo7wy+jcQDrnhHIUCMN6sk3P9YJh1ejuleDM9rp3RYiXcnzRwMNXK1AGpNraTgWn91d6L/1Ff+gaQBgPxSjW6FATshJf9Iqk9pEhT8xeCQMmNj40CdlVFYNz/TsRJYW45BXm8g83Jnh3Wj0Z90WJ80bwwgZbjPWx6wCOCID6Jo8wiubF6411ygi+GUHTR9rZFhrsGo6TSmsbvIAuabUJcskqKHW7tqVkfpz5JdjyeSSDvKcZ3PQyVne7Lyo8nzNh0g06a7CEUvq9kvr4zfTCfwk/1YVSrZ+7U8ePGj27V3D8+BH/pafiLR5FxxBKWBTYe7n5irVP5gLw4WlzOrSriXpolmaPWujCm3apml45mhlYARdfJJbpQxnQGb6cI3D/vAxxI50BPLB82JqCjNFex9xwfP5jSRDYah9wLm4yVymBqZ8Wvn0maLz4OvF2EonJBwEtN73UBULkdeX6g9GTiu2MKtCjSMm33SkwiHfvDdwMDS11sINK+YQLF/3HobbnAk2n963JahlddNRg79T96GksjN0br8u1PNMC+/uBUHEQvUBfRnR4tKp+SeLF6zMn228orgHgWr7sxilMwTOP1DgXpbEE19M0P8XC0CA/pTeHr4eFHzakPLotxF8avQfPKAXbXWgNGK6uWj1NGPshD6EN3pLvePOc4uObqHMKz/dg52/7m1VTHhYMivSRenG6H/BrWfUXQvcgCPalg7C4KCPYINb3nnRxP0NIOBP+KJ4S2haU94ZNHVr2HB038xNzUm2HrONVBsxuuTckYDmHfCa6FmvVJv4she5wyuPHk=ッpayloadEncodingイmessagepack-base64ュpayloadSchemaル$gamecult.model.transport_envelope.v4ュschemaVersionシcultnet.operation_request.v0ゥserviceIdクgamecult.model.connector
|
||||
1
crates/model-connector/tests/fixtures/wire_v4/input_reasoning_state.mp
vendored
Normal file
1
crates/model-connector/tests/fixtures/wire_v4/input_reasoning_state.mp
vendored
Normal file
|
|
@ -0,0 +1 @@
|
|||
<EFBFBD>ReasoningState<EFBFBD>provider淮nthropic-messages叮odel畚laude-opus-4-5叫tate<74><65><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>﹜ㄓ它夾帚型陋秣捲陷絮溢劃遞蝨螃謝藥齪ヌ圴佮迓玿旂衲欶趹欹詘棰葮摵蜠樉賥濋錎膼瀔嚦黀蘜蠲╭帉爎<E5B889>
|
||||
BIN
crates/model-connector/tests/fixtures/wire_v4/provider_request.mp
vendored
Normal file
BIN
crates/model-connector/tests/fixtures/wire_v4/provider_request.mp
vendored
Normal file
Binary file not shown.
BIN
crates/model-connector/tests/fixtures/wire_v4/transport_envelope.mp
vendored
Normal file
BIN
crates/model-connector/tests/fixtures/wire_v4/transport_envelope.mp
vendored
Normal file
Binary file not shown.
BIN
crates/model-connector/tests/fixtures/wire_v4/transport_invocation.mp
vendored
Normal file
BIN
crates/model-connector/tests/fixtures/wire_v4/transport_invocation.mp
vendored
Normal file
Binary file not shown.
BIN
crates/model-connector/tests/fixtures/wire_v4/transport_receipt_failed.mp
vendored
Normal file
BIN
crates/model-connector/tests/fixtures/wire_v4/transport_receipt_failed.mp
vendored
Normal file
Binary file not shown.
1
crates/model-connector/tests/fixtures/wire_v4/transport_result_refused.mp
vendored
Normal file
1
crates/model-connector/tests/fixtures/wire_v4/transport_result_refused.mp
vendored
Normal file
|
|
@ -0,0 +1 @@
|
|||
„©schema_idÙ"gamecult.model.transport_result.v4ªrequest_id®golden-request±caller_runtime_idgolden-caller«disposition<6F>§Refused«OutputBound
|
||||
BIN
crates/model-connector/tests/fixtures/wire_v4/transport_result_transported.mp
vendored
Normal file
BIN
crates/model-connector/tests/fixtures/wire_v4/transport_result_transported.mp
vendored
Normal file
Binary file not shown.
117
crates/model-connector/tests/published_schemas.rs
Normal file
117
crates/model-connector/tests/published_schemas.rs
Normal file
|
|
@ -0,0 +1,117 @@
|
|||
//! `schemas/` at the repository root publishes one JSON Schema per v4 document,
|
||||
//! generated from the types. This test regenerates them and fails on any
|
||||
//! difference; run it with `MODEL_CONNECTOR_REGENERATE_SCHEMAS=1` to rewrite the
|
||||
//! files from the types.
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use model_connector::*;
|
||||
use schemars::{JsonSchema, schema_for};
|
||||
|
||||
const REGENERATE: &str = "MODEL_CONNECTOR_REGENERATE_SCHEMAS";
|
||||
|
||||
fn render<T: JsonSchema>() -> String {
|
||||
let mut text = serde_json::to_string_pretty(&schema_for!(T)).unwrap();
|
||||
text.push('\n');
|
||||
text
|
||||
}
|
||||
|
||||
fn schemas_dir() -> PathBuf {
|
||||
Path::new(env!("CARGO_MANIFEST_DIR")).join("../../schemas")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn published_schemas_are_the_types() {
|
||||
let documents = [
|
||||
("provider_request", render::<ProviderRequest>()),
|
||||
("transport_invocation", render::<TransportInvocation>()),
|
||||
("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() {
|
||||
std::fs::create_dir_all(&dir).unwrap();
|
||||
for (name, text) in &documents {
|
||||
std::fs::write(dir.join(format!("{name}.json")), text).unwrap();
|
||||
}
|
||||
}
|
||||
for (name, text) in &documents {
|
||||
let committed = std::fs::read_to_string(dir.join(format!("{name}.json")))
|
||||
.unwrap_or_else(|_| panic!("schemas/{name}.json is missing"));
|
||||
assert_eq!(&committed, text, "schemas/{name}.json is not the type's");
|
||||
}
|
||||
let mut on_disk: Vec<String> = std::fs::read_dir(&dir)
|
||||
.unwrap()
|
||||
.map(|entry| entry.unwrap().file_name().into_string().unwrap())
|
||||
.collect();
|
||||
on_disk.sort();
|
||||
let mut expected: Vec<String> = documents
|
||||
.iter()
|
||||
.map(|(name, _)| format!("{name}.json"))
|
||||
.collect();
|
||||
expected.sort();
|
||||
assert_eq!(
|
||||
on_disk, expected,
|
||||
"schemas/ holds exactly the six documents"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_published_schema_refuses_an_unknown_field() {
|
||||
let schema = serde_json::to_value(schema_for!(ProviderRequest)).unwrap();
|
||||
assert_eq!(schema["additionalProperties"], serde_json::json!(false));
|
||||
let required = schema["required"].as_array().unwrap();
|
||||
assert!(required.contains(&serde_json::json!("max_output_tokens")));
|
||||
}
|
||||
|
||||
/// The schema and the wire agree on the ciphertext: the schema says binary, not
|
||||
/// an array of integers, and the encoder writes MessagePack `bin`.
|
||||
#[test]
|
||||
fn the_published_ciphertext_is_what_the_wire_sends() {
|
||||
let schema = serde_json::to_value(schema_for!(TransportEnvelope)).unwrap();
|
||||
let ciphertext = &schema["properties"]["ciphertext"];
|
||||
assert_eq!(ciphertext["type"], "string");
|
||||
assert_eq!(ciphertext["format"], "binary");
|
||||
assert!(ciphertext.get("items").is_none());
|
||||
let published = std::fs::read_to_string(schemas_dir().join("transport_envelope.json")).unwrap();
|
||||
let published: serde_json::Value = serde_json::from_str(&published).unwrap();
|
||||
assert_eq!(&published["properties"]["ciphertext"], ciphertext);
|
||||
|
||||
let envelope = TransportEnvelope {
|
||||
schema_id: "gamecult.model.transport_envelope.v4".to_string(),
|
||||
caller_runtime_id: "caller".to_string(),
|
||||
request_id: "request".to_string(),
|
||||
message_kind: EnvelopeKind::Invocation,
|
||||
nonce: [1; 12],
|
||||
ciphertext: vec![9; 40],
|
||||
};
|
||||
let sent =
|
||||
rmpv::decode::read_value(&mut rmp_serde::to_vec_named(&envelope).unwrap().as_slice())
|
||||
.unwrap();
|
||||
let rmpv::Value::Map(entries) = sent else {
|
||||
panic!("the envelope is a named map");
|
||||
};
|
||||
let (_, wire) = entries
|
||||
.iter()
|
||||
.find(|(key, _)| key.as_str() == Some("ciphertext"))
|
||||
.unwrap();
|
||||
assert!(matches!(wire, rmpv::Value::Binary(bytes) if bytes.len() == 40));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_refusal_schema_is_one_flat_string_enum() {
|
||||
let schema: serde_json::Value = serde_json::from_str(&render::<Refusal>()).unwrap();
|
||||
assert!(schema.get("oneOf").is_none(), "{schema}");
|
||||
assert_eq!(schema["type"], "string");
|
||||
let variants: Vec<&str> = schema["enum"]
|
||||
.as_array()
|
||||
.expect("a closed set of strings")
|
||||
.iter()
|
||||
.map(|value| value.as_str().unwrap())
|
||||
.collect();
|
||||
assert_eq!(variants.len(), 13);
|
||||
assert!(variants.contains(&"ForeignReasoningState"));
|
||||
assert!(variants.contains(&"ProviderRejected"));
|
||||
}
|
||||
349
crates/model-connector/tests/reasoning_state.rs
Normal file
349
crates/model-connector/tests/reasoning_state.rs
Normal file
|
|
@ -0,0 +1,349 @@
|
|||
//! The reasoning-state rules through a live serve loop and the public surface
|
||||
//! only: what a backend is shown, what a caller gets back and what it replays.
|
||||
#![cfg(feature = "daemon")]
|
||||
|
||||
use std::net::TcpListener;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::thread;
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use model_connector::*;
|
||||
|
||||
const CALLER: &str = "reasoning-caller";
|
||||
const KEY: &str = "reasoning-distinct-test-key";
|
||||
const PROVIDER: &str = "anthropic-messages";
|
||||
const FRAME: usize = 256 * 1024;
|
||||
|
||||
/// Answers every call with fixed events and records what it was shown.
|
||||
struct RecordingBackend {
|
||||
events: Vec<TransportEventPayload>,
|
||||
executions: AtomicUsize,
|
||||
seen: Mutex<Vec<ProviderRequest>>,
|
||||
}
|
||||
|
||||
impl RecordingBackend {
|
||||
fn new(events: Vec<TransportEventPayload>) -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
events,
|
||||
executions: AtomicUsize::new(0),
|
||||
seen: Mutex::new(Vec::new()),
|
||||
})
|
||||
}
|
||||
|
||||
fn seen(&self) -> Vec<ProviderRequest> {
|
||||
self.seen.lock().unwrap().clone()
|
||||
}
|
||||
}
|
||||
|
||||
impl ProviderBackend for RecordingBackend {
|
||||
fn ready(&self) -> Result<(), &'static str> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn provider(&self) -> &str {
|
||||
PROVIDER
|
||||
}
|
||||
|
||||
fn execute(&self, call: &ProviderCall) -> TransportResult {
|
||||
self.executions.fetch_add(1, Ordering::SeqCst);
|
||||
self.seen.lock().unwrap().push(call.request().clone());
|
||||
call.transported(
|
||||
self.events
|
||||
.iter()
|
||||
.cloned()
|
||||
.enumerate()
|
||||
.map(|(sequence, payload)| TransportEvent {
|
||||
sequence: sequence as u64,
|
||||
payload,
|
||||
})
|
||||
.collect(),
|
||||
TransportOutcome::Completed {
|
||||
provider_response_id: Some("msg-1".to_string()),
|
||||
stop_reason: StopReason::ToolUse,
|
||||
usage: Usage {
|
||||
input_tokens: 1,
|
||||
cache_read_tokens: 0,
|
||||
cache_write_tokens: 0,
|
||||
output_tokens: 1,
|
||||
reasoning_tokens: None,
|
||||
},
|
||||
},
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
fn now_ms() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_millis() as u64
|
||||
}
|
||||
|
||||
/// A running serve loop and a client for it.
|
||||
fn start(backend: Arc<RecordingBackend>) -> ConnectorClient {
|
||||
start_with(backend, 4)
|
||||
}
|
||||
|
||||
/// As `start`, with the caller's concurrent-request limit.
|
||||
fn start_with(backend: Arc<dyn ProviderBackend>, max_concurrent: usize) -> ConnectorClient {
|
||||
static SEQUENCE: AtomicUsize = AtomicUsize::new(0);
|
||||
let root: PathBuf = std::env::temp_dir().join(format!(
|
||||
"model-connector-reasoning-{}-{}",
|
||||
std::process::id(),
|
||||
SEQUENCE.fetch_add(1, Ordering::Relaxed)
|
||||
));
|
||||
std::fs::create_dir_all(&root).unwrap();
|
||||
let store = root.join("replay.cc");
|
||||
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
let endpoint = listener.local_addr().unwrap();
|
||||
let admission = CallerAdmission::new(
|
||||
CALLER,
|
||||
KEY,
|
||||
1,
|
||||
["claude-opus-4-5".to_string()],
|
||||
max_concurrent,
|
||||
64 * 1024,
|
||||
32_768,
|
||||
)
|
||||
.unwrap();
|
||||
let limits = ServeLimits {
|
||||
max_frame_bytes: FRAME,
|
||||
max_connections: 8,
|
||||
socket_timeout_ms: 5_000,
|
||||
max_expiry_skew_ms: 300_000,
|
||||
};
|
||||
thread::spawn(move || serve(listener, limits, vec![admission], &store, backend));
|
||||
ConnectorClient::new(endpoint, KEY, FRAME, Some(Duration::from_secs(5))).unwrap()
|
||||
}
|
||||
|
||||
fn invocation(request_id: &str, input: Vec<InputItem>) -> TransportInvocation {
|
||||
let mut request = ProviderRequest::new(
|
||||
request_id,
|
||||
"conversation-1",
|
||||
"claude-opus-4-5",
|
||||
"Return the typed result.",
|
||||
4_096,
|
||||
);
|
||||
request.input = input;
|
||||
TransportInvocation::new(CALLER, now_ms() + 60_000, [7; 32], request).unwrap()
|
||||
}
|
||||
|
||||
fn user(text: &str) -> InputItem {
|
||||
InputItem::UserText {
|
||||
text: text.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn state(provider: &str, byte: u8) -> InputItem {
|
||||
InputItem::ReasoningState {
|
||||
provider: provider.to_string(),
|
||||
model: "claude-opus-4-5".to_string(),
|
||||
state: vec![0x80 | byte, byte],
|
||||
}
|
||||
}
|
||||
|
||||
fn receipt_of(result: &TransportResult) -> &TransportReceipt {
|
||||
let TransportDisposition::Transported { receipt, .. } = &result.disposition else {
|
||||
panic!(
|
||||
"expected a transported result, got {:?}",
|
||||
result.disposition
|
||||
);
|
||||
};
|
||||
receipt
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn foreign_reasoning_in_flight_is_refused_and_never_executed() {
|
||||
let backend = RecordingBackend::new(vec![TransportEventPayload::TextDelta {
|
||||
text: "done".to_string(),
|
||||
}]);
|
||||
let client = start(backend.clone());
|
||||
let result = client
|
||||
.execute(&invocation(
|
||||
"foreign-in-flight",
|
||||
vec![
|
||||
user("question"),
|
||||
InputItem::ToolCall {
|
||||
call_id: "call-1".to_string(),
|
||||
name: "read_source".to_string(),
|
||||
arguments: "{}".to_string(),
|
||||
},
|
||||
state("another-provider", 1),
|
||||
InputItem::ToolResult {
|
||||
call_id: "call-1".to_string(),
|
||||
output: "evidence".to_string(),
|
||||
},
|
||||
],
|
||||
))
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
result.disposition,
|
||||
TransportDisposition::Refused(Refusal::ForeignReasoningState)
|
||||
);
|
||||
assert_eq!(backend.executions.load(Ordering::SeqCst), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_backend_never_sees_foreign_reasoning_state() {
|
||||
let backend = RecordingBackend::new(vec![TransportEventPayload::TextDelta {
|
||||
text: "done".to_string(),
|
||||
}]);
|
||||
let client = start(backend.clone());
|
||||
let own = state(PROVIDER, 2);
|
||||
let result = client
|
||||
.execute(&invocation(
|
||||
"foreign-earlier",
|
||||
vec![
|
||||
state("another-provider", 1),
|
||||
user("first"),
|
||||
own.clone(),
|
||||
state("a-third-provider", 3),
|
||||
user("second"),
|
||||
],
|
||||
))
|
||||
.unwrap();
|
||||
assert_eq!(receipt_of(&result).foreign_reasoning_dropped, 2);
|
||||
assert_eq!(receipt_of(&result).provider, PROVIDER);
|
||||
assert_eq!(
|
||||
backend.seen()[0].input,
|
||||
vec![user("first"), own, user("second")]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn replayed_transcript_reaches_the_same_provider_verbatim() {
|
||||
let thinking = vec![0x80, 0xfe, 0x00, 0xff];
|
||||
let backend = RecordingBackend::new(vec![
|
||||
TransportEventPayload::TextDelta {
|
||||
text: "let me ".to_string(),
|
||||
},
|
||||
TransportEventPayload::TextDelta {
|
||||
text: "look".to_string(),
|
||||
},
|
||||
TransportEventPayload::ReasoningState {
|
||||
state: thinking.clone(),
|
||||
},
|
||||
TransportEventPayload::ToolCall {
|
||||
call_id: "call-1".to_string(),
|
||||
name: "read_source".to_string(),
|
||||
arguments: "{}".to_string(),
|
||||
},
|
||||
]);
|
||||
let client = start(backend.clone());
|
||||
let first = client
|
||||
.execute(&invocation("turn-1", vec![user("question")]))
|
||||
.unwrap();
|
||||
let added = first.transcript_items();
|
||||
assert_eq!(
|
||||
added,
|
||||
vec![
|
||||
InputItem::AssistantText {
|
||||
text: "let me look".to_string()
|
||||
},
|
||||
InputItem::ReasoningState {
|
||||
provider: PROVIDER.to_string(),
|
||||
model: "claude-opus-4-5".to_string(),
|
||||
state: thinking,
|
||||
},
|
||||
InputItem::ToolCall {
|
||||
call_id: "call-1".to_string(),
|
||||
name: "read_source".to_string(),
|
||||
arguments: "{}".to_string(),
|
||||
},
|
||||
]
|
||||
);
|
||||
|
||||
// The caller appends the items untouched and answers the tool call; the
|
||||
// next call is in flight on the same provider and is admitted as it is.
|
||||
let mut next = vec![user("question")];
|
||||
next.extend(added.clone());
|
||||
next.push(InputItem::ToolResult {
|
||||
call_id: "call-1".to_string(),
|
||||
output: "evidence".to_string(),
|
||||
});
|
||||
let second = client.execute(&invocation("turn-2", next.clone())).unwrap();
|
||||
assert_eq!(receipt_of(&second).foreign_reasoning_dropped, 0);
|
||||
assert_eq!(backend.seen()[1].input, next);
|
||||
}
|
||||
|
||||
/// A backend that answers with a receipt it has edited after the service
|
||||
/// derived it.
|
||||
struct TamperingBackend {
|
||||
edit: fn(&mut TransportReceipt),
|
||||
}
|
||||
|
||||
impl ProviderBackend for TamperingBackend {
|
||||
fn ready(&self) -> Result<(), &'static str> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn provider(&self) -> &str {
|
||||
PROVIDER
|
||||
}
|
||||
|
||||
fn execute(&self, call: &ProviderCall) -> TransportResult {
|
||||
let mut result = call.transported(
|
||||
vec![TransportEvent {
|
||||
sequence: 0,
|
||||
payload: TransportEventPayload::TextDelta {
|
||||
text: "done".to_string(),
|
||||
},
|
||||
}],
|
||||
TransportOutcome::Completed {
|
||||
provider_response_id: None,
|
||||
stop_reason: StopReason::EndTurn,
|
||||
usage: Usage {
|
||||
input_tokens: 1,
|
||||
cache_read_tokens: 0,
|
||||
cache_write_tokens: 0,
|
||||
output_tokens: 1,
|
||||
reasoning_tokens: None,
|
||||
},
|
||||
},
|
||||
);
|
||||
let TransportDisposition::Transported { receipt, .. } = &mut result.disposition else {
|
||||
unreachable!();
|
||||
};
|
||||
(self.edit)(receipt);
|
||||
result
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
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(),
|
||||
];
|
||||
for edit in edits {
|
||||
// One concurrent request: a wedged slot would refuse the next id.
|
||||
let client = start_with(Arc::new(TamperingBackend { edit }), 1);
|
||||
let first = invocation("tampered", vec![user("question")]);
|
||||
let answer = client.execute(&first).unwrap();
|
||||
let TransportDisposition::Transported { events, receipt } = &answer.disposition else {
|
||||
panic!("expected a typed failure, got {:?}", answer.disposition);
|
||||
};
|
||||
assert!(events.is_empty());
|
||||
assert!(matches!(
|
||||
receipt.outcome,
|
||||
TransportOutcome::Failed {
|
||||
kind: FailureKind::Provider,
|
||||
..
|
||||
}
|
||||
));
|
||||
// The same id replays that failure rather than conflicting.
|
||||
assert_eq!(client.execute(&first).unwrap(), answer);
|
||||
// The next id is not refused for capacity.
|
||||
let next = client
|
||||
.execute(&invocation("tampered-next", vec![user("again")]))
|
||||
.unwrap();
|
||||
assert!(
|
||||
!matches!(next.disposition, TransportDisposition::Refused(_)),
|
||||
"{:?}",
|
||||
next.disposition
|
||||
);
|
||||
}
|
||||
}
|
||||
251
crates/model-connector/tests/wire_golden_v4.rs
Normal file
251
crates/model-connector/tests/wire_golden_v4.rs
Normal file
|
|
@ -0,0 +1,251 @@
|
|||
//! Pins the bytes of the v4 wire so a drift is a failing test. The fixtures in
|
||||
//! `fixtures/wire_v4` are produced by this file (see its README); they pin the
|
||||
//! kit against itself and are not a second contract: the types in
|
||||
//! `model_connector::contract` own the shape.
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use base64::Engine;
|
||||
use base64::engine::general_purpose::STANDARD;
|
||||
use cultnet_rs::{
|
||||
CultNetMessage, CultNetWireContract, decode_cultnet_message_from_slice,
|
||||
encode_cultnet_message_to_vec,
|
||||
};
|
||||
use model_connector::*;
|
||||
use serde::Serialize;
|
||||
use serde::de::DeserializeOwned;
|
||||
|
||||
const KEY: &str = "golden-distinct-test-key";
|
||||
const CALLER: &str = "golden-caller";
|
||||
const REGENERATE: &str = "WIRE_V4_REGENERATE";
|
||||
|
||||
fn fixtures() -> PathBuf {
|
||||
Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/wire_v4")
|
||||
}
|
||||
|
||||
fn named<T: Serialize>(value: &T) -> Vec<u8> {
|
||||
rmp_serde::to_vec_named(value).unwrap()
|
||||
}
|
||||
|
||||
/// 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
|
||||
/// from the value first.
|
||||
fn assert_wire<T: Serialize + DeserializeOwned + PartialEq + std::fmt::Debug>(
|
||||
name: &str,
|
||||
value: &T,
|
||||
) {
|
||||
let path = fixtures().join(name);
|
||||
if std::env::var_os(REGENERATE).is_some() {
|
||||
std::fs::create_dir_all(fixtures()).unwrap();
|
||||
std::fs::write(&path, named(value)).unwrap();
|
||||
}
|
||||
let golden = std::fs::read(&path).unwrap_or_else(|_| panic!("{name}: fixture is missing"));
|
||||
assert_eq!(named(value), golden, "{name}: serialized bytes");
|
||||
let decoded: T = rmp_serde::from_slice(&golden).unwrap();
|
||||
assert_eq!(&decoded, value, "{name}: decoded value");
|
||||
}
|
||||
|
||||
fn request() -> ProviderRequest {
|
||||
let mut request = ProviderRequest::new(
|
||||
"golden-request",
|
||||
"golden-conversation",
|
||||
"claude-opus-4-5",
|
||||
"Return the typed result.",
|
||||
8_192,
|
||||
);
|
||||
request.input = vec![
|
||||
InputItem::UserText {
|
||||
text: "Projected state".to_string(),
|
||||
},
|
||||
InputItem::ToolCall {
|
||||
call_id: "call-1".to_string(),
|
||||
name: "read_source".to_string(),
|
||||
arguments: "{}".to_string(),
|
||||
},
|
||||
InputItem::ToolResult {
|
||||
call_id: "call-1".to_string(),
|
||||
output: "evidence".to_string(),
|
||||
},
|
||||
];
|
||||
request.reasoning_effort = Some(ReasoningEffort::High);
|
||||
request.cache_retention = Some(CacheRetention::Short);
|
||||
request.tools = vec![ToolDefinition {
|
||||
name: "read_source".to_string(),
|
||||
description: "Read one admitted source.".to_string(),
|
||||
parameters_json: r#"{"type":"object"}"#.to_string(),
|
||||
}];
|
||||
request.tool_choice = ToolChoice::Required;
|
||||
request.parallel_tool_calls = true;
|
||||
request.output_schema_json = Some(r#"{"type":"object"}"#.to_string());
|
||||
request
|
||||
}
|
||||
|
||||
fn invocation() -> TransportInvocation {
|
||||
TransportInvocation::new(CALLER, 1_800_000_000_000, [7; 32], request()).unwrap()
|
||||
}
|
||||
|
||||
fn receipt(outcome: TransportOutcome) -> TransportReceipt {
|
||||
let invocation = invocation();
|
||||
TransportReceipt {
|
||||
schema_id: RECEIPT_SCHEMA_ID.to_string(),
|
||||
request_id: invocation.request_id().to_string(),
|
||||
caller_runtime_id: CALLER.to_string(),
|
||||
native_request_sha256: invocation.native_request_sha256,
|
||||
provider_request_sha256: invocation.provider_request_sha256,
|
||||
model: invocation.request.model.clone(),
|
||||
provider: "anthropic-messages".to_string(),
|
||||
foreign_reasoning_dropped: 2,
|
||||
outcome,
|
||||
}
|
||||
}
|
||||
|
||||
fn completed() -> TransportOutcome {
|
||||
TransportOutcome::Completed {
|
||||
provider_response_id: Some("msg-1".to_string()),
|
||||
stop_reason: StopReason::ToolUse,
|
||||
usage: Usage {
|
||||
input_tokens: 120,
|
||||
cache_read_tokens: 40,
|
||||
cache_write_tokens: 8,
|
||||
output_tokens: 33,
|
||||
reasoning_tokens: Some(12),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn wire_golden_v4() {
|
||||
let invocation = invocation();
|
||||
assert_wire("provider_request.mp", &invocation.request);
|
||||
assert_wire("transport_invocation.mp", &invocation);
|
||||
|
||||
// The reasoning bytes are all >= 0x80 so a byte-array encoding cannot pass
|
||||
// for MessagePack bin.
|
||||
let reasoning = || (0x80..=0xff_u8).collect::<Vec<u8>>();
|
||||
assert_wire(
|
||||
"input_reasoning_state.mp",
|
||||
&InputItem::ReasoningState {
|
||||
provider: "anthropic-messages".to_string(),
|
||||
model: "claude-opus-4-5".to_string(),
|
||||
state: reasoning(),
|
||||
},
|
||||
);
|
||||
|
||||
let transported = TransportResult {
|
||||
schema_id: RESULT_SCHEMA_ID.to_string(),
|
||||
request_id: invocation.request_id().to_string(),
|
||||
caller_runtime_id: CALLER.to_string(),
|
||||
disposition: TransportDisposition::Transported {
|
||||
events: vec![
|
||||
TransportEvent {
|
||||
sequence: 0,
|
||||
payload: TransportEventPayload::TextDelta {
|
||||
text: "thinking".to_string(),
|
||||
},
|
||||
},
|
||||
TransportEvent {
|
||||
sequence: 1,
|
||||
payload: TransportEventPayload::ReasoningState { state: reasoning() },
|
||||
},
|
||||
TransportEvent {
|
||||
sequence: 2,
|
||||
payload: TransportEventPayload::ToolCall {
|
||||
call_id: "call-2".to_string(),
|
||||
name: "read_source".to_string(),
|
||||
arguments: "{}".to_string(),
|
||||
},
|
||||
},
|
||||
],
|
||||
receipt: Box::new(receipt(completed())),
|
||||
},
|
||||
};
|
||||
assert_eq!(transported.validate_against(&invocation), Ok(()));
|
||||
assert_wire("transport_result_transported.mp", &transported);
|
||||
|
||||
let refused = TransportResult::refused(&invocation, Refusal::OutputBound);
|
||||
assert_wire("transport_result_refused.mp", &refused);
|
||||
|
||||
let failed = receipt(TransportOutcome::Failed {
|
||||
kind: FailureKind::RateLimited,
|
||||
message: "provider asked to slow down".to_string(),
|
||||
usage: Some(Usage {
|
||||
input_tokens: 120,
|
||||
cache_read_tokens: 0,
|
||||
cache_write_tokens: 0,
|
||||
output_tokens: 0,
|
||||
reasoning_tokens: None,
|
||||
}),
|
||||
});
|
||||
assert_eq!(failed.validate_against(&invocation), Ok(()));
|
||||
assert_wire("transport_receipt_failed.mp", &failed);
|
||||
|
||||
let envelope = TransportEnvelope {
|
||||
schema_id: ENVELOPE_SCHEMA_ID.to_string(),
|
||||
caller_runtime_id: CALLER.to_string(),
|
||||
request_id: "golden-request".to_string(),
|
||||
message_kind: EnvelopeKind::Invocation,
|
||||
nonce: [0x80, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 0xff],
|
||||
ciphertext: (0..40_u8).map(|byte| byte.wrapping_mul(7)).collect(),
|
||||
};
|
||||
assert_wire("transport_envelope.mp", &envelope);
|
||||
}
|
||||
|
||||
/// The committed carriage frame is a `cultnet.operation_request.v0` for the
|
||||
/// connector, read here with CultNet itself and not the kit's own carriage code.
|
||||
/// Its payload is the base64 of a sealed envelope that decrypts under the
|
||||
/// fixed key to the golden invocation.
|
||||
#[test]
|
||||
fn the_golden_carriage_frame_decrypts_and_validates_under_the_fixed_key() {
|
||||
let path = fixtures().join("carriage_invocation.frame");
|
||||
let key = TransportKey::from_connection_secret(KEY).unwrap();
|
||||
let invocation = invocation();
|
||||
if std::env::var_os(REGENERATE).is_some() {
|
||||
let envelope = encrypt_invocation(&invocation, &key).unwrap();
|
||||
let message = CultNetMessage::OperationRequest {
|
||||
message_id: envelope.request_id.clone(),
|
||||
service_id: CONNECTOR_SERVICE_ID.to_string(),
|
||||
operation: INVOKE_OPERATION.to_string(),
|
||||
payload_schema: ENVELOPE_SCHEMA_ID.to_string(),
|
||||
payload_encoding: "messagepack-base64".to_string(),
|
||||
payload: STANDARD.encode(named(&envelope)),
|
||||
source_runtime_id: None,
|
||||
target_runtime_id: None,
|
||||
};
|
||||
std::fs::write(
|
||||
&path,
|
||||
encode_cultnet_message_to_vec(&message, CultNetWireContract::CultNetSchemaV0).unwrap(),
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
let frame = std::fs::read(&path).expect("carriage_invocation.frame is missing");
|
||||
let CultNetMessage::OperationRequest {
|
||||
message_id,
|
||||
service_id,
|
||||
operation,
|
||||
payload_schema,
|
||||
payload,
|
||||
..
|
||||
} = decode_cultnet_message_from_slice(&frame, CultNetWireContract::CultNetSchemaV0).unwrap()
|
||||
else {
|
||||
panic!("the golden frame is a CultNet operation request");
|
||||
};
|
||||
assert_eq!(service_id, "gamecult.model.connector");
|
||||
assert_eq!(operation, "invoke");
|
||||
assert_eq!(payload_schema, "gamecult.model.transport_envelope.v4");
|
||||
assert_eq!(message_id, invocation.request_id());
|
||||
let envelope: TransportEnvelope =
|
||||
rmp_serde::from_slice(&STANDARD.decode(payload).unwrap()).unwrap();
|
||||
assert_eq!(envelope.request_id, message_id);
|
||||
let decrypted = decrypt_invocation(&envelope, &key).unwrap();
|
||||
assert_eq!(decrypted, invocation);
|
||||
assert_eq!(decrypted.validate(1_799_999_999_000, 5_000), Ok(()));
|
||||
}
|
||||
|
||||
/// The reasoning bytes travel as MessagePack `bin` (0xc4 length marker), never
|
||||
/// as an array of integers.
|
||||
#[test]
|
||||
fn reasoning_state_is_messagepack_bin() {
|
||||
let golden = std::fs::read(fixtures().join("input_reasoning_state.mp")).unwrap();
|
||||
let marker = golden.windows(3).any(|window| window == [0xc4, 0x80, 0x80]);
|
||||
assert!(marker, "state must be a bin 8 of 128 bytes starting 0x80");
|
||||
}
|
||||
261
schemas/provider_request.json
Normal file
261
schemas/provider_request.json
Normal file
|
|
@ -0,0 +1,261 @@
|
|||
{
|
||||
"$schema": "https://json-schema.org/draft/2020-12/schema",
|
||||
"title": "ProviderRequest",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"cache_retention": {
|
||||
"anyOf": [
|
||||
{
|
||||
"$ref": "#/$defs/CacheRetention"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
]
|
||||
},
|
||||
"conversation_id": {
|
||||
"type": "string"
|
||||
},
|
||||
"input": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"$ref": "#/$defs/InputItem"
|
||||
}
|
||||
},
|
||||
"instructions": {
|
||||
"type": "string"
|
||||
},
|
||||
"max_output_tokens": {
|
||||
"type": "integer",
|
||||
"format": "uint32",
|
||||
"minimum": 0
|
||||
},
|
||||
"model": {
|
||||
"type": "string"
|
||||
},
|
||||
"output_schema_json": {
|
||||
"type": [
|
||||
"string",
|
||||
"null"
|
||||
]
|
||||
},
|
||||
"parallel_tool_calls": {
|
||||
"type": "boolean"
|
||||
},
|
||||
"reasoning_effort": {
|
||||
"anyOf": [
|
||||
{
|
||||
"$ref": "#/$defs/ReasoningEffort"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
]
|
||||
},
|
||||
"request_id": {
|
||||
"type": "string"
|
||||
},
|
||||
"schema_id": {
|
||||
"type": "string"
|
||||
},
|
||||
"tool_choice": {
|
||||
"$ref": "#/$defs/ToolChoice"
|
||||
},
|
||||
"tools": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"$ref": "#/$defs/ToolDefinition"
|
||||
}
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"schema_id",
|
||||
"request_id",
|
||||
"conversation_id",
|
||||
"model",
|
||||
"instructions",
|
||||
"input",
|
||||
"tools",
|
||||
"tool_choice",
|
||||
"parallel_tool_calls",
|
||||
"max_output_tokens"
|
||||
],
|
||||
"$defs": {
|
||||
"CacheRetention": {
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"short",
|
||||
"long"
|
||||
]
|
||||
},
|
||||
"InputItem": {
|
||||
"oneOf": [
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"UserText": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"text": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"text"
|
||||
]
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"UserText"
|
||||
]
|
||||
},
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"AssistantText": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"text": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"text"
|
||||
]
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"AssistantText"
|
||||
]
|
||||
},
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"ToolCall": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"arguments": {
|
||||
"type": "string"
|
||||
},
|
||||
"call_id": {
|
||||
"type": "string"
|
||||
},
|
||||
"name": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"call_id",
|
||||
"name",
|
||||
"arguments"
|
||||
]
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"ToolCall"
|
||||
]
|
||||
},
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"ToolResult": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"call_id": {
|
||||
"type": "string"
|
||||
},
|
||||
"output": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"call_id",
|
||||
"output"
|
||||
]
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"ToolResult"
|
||||
]
|
||||
},
|
||||
{
|
||||
"description": "One provider reasoning block, opaque to every caller: `state` is the\nconnector's own encoding of the whole block. Callers replay it\nverbatim in transcript order and never inspect or branch on it; only\nthe connector of `provider` reads it, and the service drops or refuses\nthe state of any other provider before a backend sees the request.",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"ReasoningState": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"model": {
|
||||
"type": "string"
|
||||
},
|
||||
"provider": {
|
||||
"type": "string"
|
||||
},
|
||||
"state": {
|
||||
"description": "MessagePack bin: raw bytes, never an array of integers",
|
||||
"type": "string",
|
||||
"format": "binary"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"provider",
|
||||
"model",
|
||||
"state"
|
||||
]
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"ReasoningState"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
"ReasoningEffort": {
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"low",
|
||||
"medium",
|
||||
"high",
|
||||
"max"
|
||||
]
|
||||
},
|
||||
"ToolChoice": {
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"auto",
|
||||
"required"
|
||||
]
|
||||
},
|
||||
"ToolDefinition": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"description": {
|
||||
"type": "string"
|
||||
},
|
||||
"name": {
|
||||
"type": "string"
|
||||
},
|
||||
"parameters_json": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"name",
|
||||
"description",
|
||||
"parameters_json"
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
55
schemas/transport_envelope.json
Normal file
55
schemas/transport_envelope.json
Normal file
|
|
@ -0,0 +1,55 @@
|
|||
{
|
||||
"$schema": "https://json-schema.org/draft/2020-12/schema",
|
||||
"title": "TransportEnvelope",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"caller_runtime_id": {
|
||||
"type": "string"
|
||||
},
|
||||
"ciphertext": {
|
||||
"description": "MessagePack bin: raw bytes, never an array of integers",
|
||||
"type": "string",
|
||||
"format": "binary"
|
||||
},
|
||||
"message_kind": {
|
||||
"$ref": "#/$defs/EnvelopeKind"
|
||||
},
|
||||
"nonce": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "integer",
|
||||
"format": "uint8",
|
||||
"maximum": 255,
|
||||
"minimum": 0
|
||||
},
|
||||
"maxItems": 12,
|
||||
"minItems": 12
|
||||
},
|
||||
"request_id": {
|
||||
"type": "string"
|
||||
},
|
||||
"schema_id": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"schema_id",
|
||||
"caller_runtime_id",
|
||||
"request_id",
|
||||
"message_kind",
|
||||
"nonce",
|
||||
"ciphertext"
|
||||
],
|
||||
"$defs": {
|
||||
"EnvelopeKind": {
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"Invocation",
|
||||
"Result",
|
||||
"UsageReadoutRequest",
|
||||
"UsageReadoutResult"
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
311
schemas/transport_invocation.json
Normal file
311
schemas/transport_invocation.json
Normal file
|
|
@ -0,0 +1,311 @@
|
|||
{
|
||||
"$schema": "https://json-schema.org/draft/2020-12/schema",
|
||||
"title": "TransportInvocation",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"caller_runtime_id": {
|
||||
"type": "string"
|
||||
},
|
||||
"expires_at_unix_ms": {
|
||||
"type": "integer",
|
||||
"format": "uint64",
|
||||
"minimum": 0
|
||||
},
|
||||
"native_request_sha256": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "integer",
|
||||
"format": "uint8",
|
||||
"maximum": 255,
|
||||
"minimum": 0
|
||||
},
|
||||
"maxItems": 32,
|
||||
"minItems": 32
|
||||
},
|
||||
"provider_request_sha256": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "integer",
|
||||
"format": "uint8",
|
||||
"maximum": 255,
|
||||
"minimum": 0
|
||||
},
|
||||
"maxItems": 32,
|
||||
"minItems": 32
|
||||
},
|
||||
"request": {
|
||||
"$ref": "#/$defs/ProviderRequest"
|
||||
},
|
||||
"schema_id": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"schema_id",
|
||||
"caller_runtime_id",
|
||||
"expires_at_unix_ms",
|
||||
"native_request_sha256",
|
||||
"provider_request_sha256",
|
||||
"request"
|
||||
],
|
||||
"$defs": {
|
||||
"CacheRetention": {
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"short",
|
||||
"long"
|
||||
]
|
||||
},
|
||||
"InputItem": {
|
||||
"oneOf": [
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"UserText": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"text": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"text"
|
||||
]
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"UserText"
|
||||
]
|
||||
},
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"AssistantText": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"text": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"text"
|
||||
]
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"AssistantText"
|
||||
]
|
||||
},
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"ToolCall": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"arguments": {
|
||||
"type": "string"
|
||||
},
|
||||
"call_id": {
|
||||
"type": "string"
|
||||
},
|
||||
"name": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"call_id",
|
||||
"name",
|
||||
"arguments"
|
||||
]
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"ToolCall"
|
||||
]
|
||||
},
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"ToolResult": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"call_id": {
|
||||
"type": "string"
|
||||
},
|
||||
"output": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"call_id",
|
||||
"output"
|
||||
]
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"ToolResult"
|
||||
]
|
||||
},
|
||||
{
|
||||
"description": "One provider reasoning block, opaque to every caller: `state` is the\nconnector's own encoding of the whole block. Callers replay it\nverbatim in transcript order and never inspect or branch on it; only\nthe connector of `provider` reads it, and the service drops or refuses\nthe state of any other provider before a backend sees the request.",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"ReasoningState": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"model": {
|
||||
"type": "string"
|
||||
},
|
||||
"provider": {
|
||||
"type": "string"
|
||||
},
|
||||
"state": {
|
||||
"description": "MessagePack bin: raw bytes, never an array of integers",
|
||||
"type": "string",
|
||||
"format": "binary"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"provider",
|
||||
"model",
|
||||
"state"
|
||||
]
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"ReasoningState"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
"ProviderRequest": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"cache_retention": {
|
||||
"anyOf": [
|
||||
{
|
||||
"$ref": "#/$defs/CacheRetention"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
]
|
||||
},
|
||||
"conversation_id": {
|
||||
"type": "string"
|
||||
},
|
||||
"input": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"$ref": "#/$defs/InputItem"
|
||||
}
|
||||
},
|
||||
"instructions": {
|
||||
"type": "string"
|
||||
},
|
||||
"max_output_tokens": {
|
||||
"type": "integer",
|
||||
"format": "uint32",
|
||||
"minimum": 0
|
||||
},
|
||||
"model": {
|
||||
"type": "string"
|
||||
},
|
||||
"output_schema_json": {
|
||||
"type": [
|
||||
"string",
|
||||
"null"
|
||||
]
|
||||
},
|
||||
"parallel_tool_calls": {
|
||||
"type": "boolean"
|
||||
},
|
||||
"reasoning_effort": {
|
||||
"anyOf": [
|
||||
{
|
||||
"$ref": "#/$defs/ReasoningEffort"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
]
|
||||
},
|
||||
"request_id": {
|
||||
"type": "string"
|
||||
},
|
||||
"schema_id": {
|
||||
"type": "string"
|
||||
},
|
||||
"tool_choice": {
|
||||
"$ref": "#/$defs/ToolChoice"
|
||||
},
|
||||
"tools": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"$ref": "#/$defs/ToolDefinition"
|
||||
}
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"schema_id",
|
||||
"request_id",
|
||||
"conversation_id",
|
||||
"model",
|
||||
"instructions",
|
||||
"input",
|
||||
"tools",
|
||||
"tool_choice",
|
||||
"parallel_tool_calls",
|
||||
"max_output_tokens"
|
||||
]
|
||||
},
|
||||
"ReasoningEffort": {
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"low",
|
||||
"medium",
|
||||
"high",
|
||||
"max"
|
||||
]
|
||||
},
|
||||
"ToolChoice": {
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"auto",
|
||||
"required"
|
||||
]
|
||||
},
|
||||
"ToolDefinition": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"description": {
|
||||
"type": "string"
|
||||
},
|
||||
"name": {
|
||||
"type": "string"
|
||||
},
|
||||
"parameters_json": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"name",
|
||||
"description",
|
||||
"parameters_json"
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
204
schemas/transport_receipt.json
Normal file
204
schemas/transport_receipt.json
Normal file
|
|
@ -0,0 +1,204 @@
|
|||
{
|
||||
"$schema": "https://json-schema.org/draft/2020-12/schema",
|
||||
"title": "TransportReceipt",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"caller_runtime_id": {
|
||||
"type": "string"
|
||||
},
|
||||
"foreign_reasoning_dropped": {
|
||||
"description": "How many foreign `ReasoningState` items the service removed from the\nrequest before the backend saw it.",
|
||||
"type": "integer",
|
||||
"format": "uint32",
|
||||
"minimum": 0
|
||||
},
|
||||
"model": {
|
||||
"type": "string"
|
||||
},
|
||||
"native_request_sha256": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "integer",
|
||||
"format": "uint8",
|
||||
"maximum": 255,
|
||||
"minimum": 0
|
||||
},
|
||||
"maxItems": 32,
|
||||
"minItems": 32
|
||||
},
|
||||
"outcome": {
|
||||
"$ref": "#/$defs/TransportOutcome"
|
||||
},
|
||||
"provider": {
|
||||
"type": "string"
|
||||
},
|
||||
"provider_request_sha256": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "integer",
|
||||
"format": "uint8",
|
||||
"maximum": 255,
|
||||
"minimum": 0
|
||||
},
|
||||
"maxItems": 32,
|
||||
"minItems": 32
|
||||
},
|
||||
"request_id": {
|
||||
"type": "string"
|
||||
},
|
||||
"schema_id": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"schema_id",
|
||||
"request_id",
|
||||
"caller_runtime_id",
|
||||
"native_request_sha256",
|
||||
"provider_request_sha256",
|
||||
"model",
|
||||
"provider",
|
||||
"foreign_reasoning_dropped",
|
||||
"outcome"
|
||||
],
|
||||
"$defs": {
|
||||
"FailureKind": {
|
||||
"description": "What kind of failure a provider reported. `ResultSize` is the connector's own\nstand-in for a result whose sealed frame cannot be carried.",
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"result_size",
|
||||
"invalid_request",
|
||||
"authentication",
|
||||
"rate_limited",
|
||||
"unavailable",
|
||||
"timeout",
|
||||
"provider"
|
||||
]
|
||||
},
|
||||
"StopReason": {
|
||||
"description": "Why a provider stopped producing output.",
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"end_turn",
|
||||
"tool_use",
|
||||
"max_output_tokens",
|
||||
"stop_sequence",
|
||||
"refusal",
|
||||
"context_window"
|
||||
]
|
||||
},
|
||||
"TransportOutcome": {
|
||||
"oneOf": [
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"Completed": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"provider_response_id": {
|
||||
"type": [
|
||||
"string",
|
||||
"null"
|
||||
]
|
||||
},
|
||||
"stop_reason": {
|
||||
"$ref": "#/$defs/StopReason"
|
||||
},
|
||||
"usage": {
|
||||
"$ref": "#/$defs/Usage"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"stop_reason",
|
||||
"usage"
|
||||
]
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"Completed"
|
||||
]
|
||||
},
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"Failed": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"kind": {
|
||||
"$ref": "#/$defs/FailureKind"
|
||||
},
|
||||
"message": {
|
||||
"type": "string"
|
||||
},
|
||||
"usage": {
|
||||
"description": "What the provider billed for the failed request: `Some` when it\nreported usage (`Some` of zeros when it billed nothing), `None`\nwhen unknown.",
|
||||
"anyOf": [
|
||||
{
|
||||
"$ref": "#/$defs/Usage"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"kind",
|
||||
"message"
|
||||
]
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"Failed"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
"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.",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"cache_read_tokens": {
|
||||
"type": "integer",
|
||||
"format": "uint64",
|
||||
"minimum": 0
|
||||
},
|
||||
"cache_write_tokens": {
|
||||
"type": "integer",
|
||||
"format": "uint64",
|
||||
"minimum": 0
|
||||
},
|
||||
"input_tokens": {
|
||||
"type": "integer",
|
||||
"format": "uint64",
|
||||
"minimum": 0
|
||||
},
|
||||
"output_tokens": {
|
||||
"type": "integer",
|
||||
"format": "uint64",
|
||||
"minimum": 0
|
||||
},
|
||||
"reasoning_tokens": {
|
||||
"type": [
|
||||
"integer",
|
||||
"null"
|
||||
],
|
||||
"format": "uint64",
|
||||
"minimum": 0
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"input_tokens",
|
||||
"cache_read_tokens",
|
||||
"cache_write_tokens",
|
||||
"output_tokens"
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
386
schemas/transport_result.json
Normal file
386
schemas/transport_result.json
Normal file
|
|
@ -0,0 +1,386 @@
|
|||
{
|
||||
"$schema": "https://json-schema.org/draft/2020-12/schema",
|
||||
"title": "TransportResult",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"caller_runtime_id": {
|
||||
"type": "string"
|
||||
},
|
||||
"disposition": {
|
||||
"$ref": "#/$defs/TransportDisposition"
|
||||
},
|
||||
"request_id": {
|
||||
"type": "string"
|
||||
},
|
||||
"schema_id": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"schema_id",
|
||||
"request_id",
|
||||
"caller_runtime_id",
|
||||
"disposition"
|
||||
],
|
||||
"$defs": {
|
||||
"FailureKind": {
|
||||
"description": "What kind of failure a provider reported. `ResultSize` is the connector's own\nstand-in for a result whose sealed frame cannot be carried.",
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"result_size",
|
||||
"invalid_request",
|
||||
"authentication",
|
||||
"rate_limited",
|
||||
"unavailable",
|
||||
"timeout",
|
||||
"provider"
|
||||
]
|
||||
},
|
||||
"Refusal": {
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"Expired",
|
||||
"IdentitySubstitution",
|
||||
"ProviderDigestSubstitution",
|
||||
"ModelNotAllowed",
|
||||
"OutputBound",
|
||||
"PayloadBound",
|
||||
"Capacity",
|
||||
"InFlight",
|
||||
"Indeterminate",
|
||||
"ReplayConflict",
|
||||
"Malformed",
|
||||
"ForeignReasoningState",
|
||||
"ProviderRejected"
|
||||
]
|
||||
},
|
||||
"StopReason": {
|
||||
"description": "Why a provider stopped producing output.",
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"end_turn",
|
||||
"tool_use",
|
||||
"max_output_tokens",
|
||||
"stop_sequence",
|
||||
"refusal",
|
||||
"context_window"
|
||||
]
|
||||
},
|
||||
"TransportDisposition": {
|
||||
"oneOf": [
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"Refused": {
|
||||
"$ref": "#/$defs/Refusal"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"Refused"
|
||||
]
|
||||
},
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"Transported": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"events": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"$ref": "#/$defs/TransportEvent"
|
||||
}
|
||||
},
|
||||
"receipt": {
|
||||
"$ref": "#/$defs/TransportReceipt"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"events",
|
||||
"receipt"
|
||||
]
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"Transported"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
"TransportEvent": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"payload": {
|
||||
"$ref": "#/$defs/TransportEventPayload"
|
||||
},
|
||||
"sequence": {
|
||||
"type": "integer",
|
||||
"format": "uint64",
|
||||
"minimum": 0
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"sequence",
|
||||
"payload"
|
||||
]
|
||||
},
|
||||
"TransportEventPayload": {
|
||||
"oneOf": [
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"TextDelta": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"text": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"text"
|
||||
]
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"TextDelta"
|
||||
]
|
||||
},
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"ToolCall": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"arguments": {
|
||||
"type": "string"
|
||||
},
|
||||
"call_id": {
|
||||
"type": "string"
|
||||
},
|
||||
"name": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"call_id",
|
||||
"name",
|
||||
"arguments"
|
||||
]
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"ToolCall"
|
||||
]
|
||||
},
|
||||
{
|
||||
"description": "One reasoning block of the provider that answered. It carries no\nprovider or model: the receipt names them, so a backend cannot tag its\nown state.",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"ReasoningState": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"state": {
|
||||
"description": "MessagePack bin: raw bytes, never an array of integers",
|
||||
"type": "string",
|
||||
"format": "binary"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"state"
|
||||
]
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"ReasoningState"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
"TransportOutcome": {
|
||||
"oneOf": [
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"Completed": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"provider_response_id": {
|
||||
"type": [
|
||||
"string",
|
||||
"null"
|
||||
]
|
||||
},
|
||||
"stop_reason": {
|
||||
"$ref": "#/$defs/StopReason"
|
||||
},
|
||||
"usage": {
|
||||
"$ref": "#/$defs/Usage"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"stop_reason",
|
||||
"usage"
|
||||
]
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"Completed"
|
||||
]
|
||||
},
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"Failed": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"kind": {
|
||||
"$ref": "#/$defs/FailureKind"
|
||||
},
|
||||
"message": {
|
||||
"type": "string"
|
||||
},
|
||||
"usage": {
|
||||
"description": "What the provider billed for the failed request: `Some` when it\nreported usage (`Some` of zeros when it billed nothing), `None`\nwhen unknown.",
|
||||
"anyOf": [
|
||||
{
|
||||
"$ref": "#/$defs/Usage"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"kind",
|
||||
"message"
|
||||
]
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"Failed"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
"TransportReceipt": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"caller_runtime_id": {
|
||||
"type": "string"
|
||||
},
|
||||
"foreign_reasoning_dropped": {
|
||||
"description": "How many foreign `ReasoningState` items the service removed from the\nrequest before the backend saw it.",
|
||||
"type": "integer",
|
||||
"format": "uint32",
|
||||
"minimum": 0
|
||||
},
|
||||
"model": {
|
||||
"type": "string"
|
||||
},
|
||||
"native_request_sha256": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "integer",
|
||||
"format": "uint8",
|
||||
"maximum": 255,
|
||||
"minimum": 0
|
||||
},
|
||||
"maxItems": 32,
|
||||
"minItems": 32
|
||||
},
|
||||
"outcome": {
|
||||
"$ref": "#/$defs/TransportOutcome"
|
||||
},
|
||||
"provider": {
|
||||
"type": "string"
|
||||
},
|
||||
"provider_request_sha256": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "integer",
|
||||
"format": "uint8",
|
||||
"maximum": 255,
|
||||
"minimum": 0
|
||||
},
|
||||
"maxItems": 32,
|
||||
"minItems": 32
|
||||
},
|
||||
"request_id": {
|
||||
"type": "string"
|
||||
},
|
||||
"schema_id": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"schema_id",
|
||||
"request_id",
|
||||
"caller_runtime_id",
|
||||
"native_request_sha256",
|
||||
"provider_request_sha256",
|
||||
"model",
|
||||
"provider",
|
||||
"foreign_reasoning_dropped",
|
||||
"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.",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"cache_read_tokens": {
|
||||
"type": "integer",
|
||||
"format": "uint64",
|
||||
"minimum": 0
|
||||
},
|
||||
"cache_write_tokens": {
|
||||
"type": "integer",
|
||||
"format": "uint64",
|
||||
"minimum": 0
|
||||
},
|
||||
"input_tokens": {
|
||||
"type": "integer",
|
||||
"format": "uint64",
|
||||
"minimum": 0
|
||||
},
|
||||
"output_tokens": {
|
||||
"type": "integer",
|
||||
"format": "uint64",
|
||||
"minimum": 0
|
||||
},
|
||||
"reasoning_tokens": {
|
||||
"type": [
|
||||
"integer",
|
||||
"null"
|
||||
],
|
||||
"format": "uint64",
|
||||
"minimum": 0
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"input_tokens",
|
||||
"cache_read_tokens",
|
||||
"cache_write_tokens",
|
||||
"output_tokens"
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
68
schemas/usage_readout.json
Normal file
68
schemas/usage_readout.json
Normal file
|
|
@ -0,0 +1,68 @@
|
|||
{
|
||||
"$schema": "https://json-schema.org/draft/2020-12/schema",
|
||||
"title": "UsageReadout",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"at_unix_ms": {
|
||||
"type": "integer",
|
||||
"format": "uint64",
|
||||
"minimum": 0
|
||||
},
|
||||
"plan_quota": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"$ref": "#/$defs/PlanQuota"
|
||||
}
|
||||
},
|
||||
"schema_id": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"schema_id",
|
||||
"at_unix_ms",
|
||||
"plan_quota"
|
||||
],
|
||||
"$defs": {
|
||||
"PlanQuota": {
|
||||
"description": "One window of the subscription plan, as the provider reported it.",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"observed_at_unix_ms": {
|
||||
"description": "When this window was last reported. Its age in a readout is the\nreadout's `at_unix_ms` minus this.",
|
||||
"type": "integer",
|
||||
"format": "uint64",
|
||||
"minimum": 0
|
||||
},
|
||||
"resets_at_unix_ms": {
|
||||
"type": "integer",
|
||||
"format": "uint64",
|
||||
"minimum": 0
|
||||
},
|
||||
"status": {
|
||||
"type": "string"
|
||||
},
|
||||
"utilization_permille": {
|
||||
"description": "How much of the window was used when it was observed, in thousandths;\n`None` once the window has reset (`status` is then `expired`), because\nthe old reading says nothing about the new window.",
|
||||
"type": [
|
||||
"integer",
|
||||
"null"
|
||||
],
|
||||
"format": "uint32",
|
||||
"minimum": 0
|
||||
},
|
||||
"window": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"window",
|
||||
"resets_at_unix_ms",
|
||||
"status",
|
||||
"observed_at_unix_ms"
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue