GameCult/ModelConnector
Watch
9
0
Fork
You've already forked ModelConnector
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

claude-turn-schema-quota: schema-typed Persona replies via --json-schema, token on fd 3 never in the child env, per-window plan-quota capture, communal spend readout, not-logged-in bills zero (squash of eureka-body/claude-turn-schema-quota e867fad..429b344; Soul s1, s2)

Browse source 
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Eureka 2026-10-29 20:18:35 +01:00
commit 480cf941c8
21 changed files with 2021 additions and 187 deletions
Download patch file Download diff file Expand all files Collapse all files

15
README.md
View file

@ -21,8 +21,11 @@ Crate `crates/model-connector`:
- `src/envelope.rs` - the AES-256-GCM `TransportEnvelope` (binary ciphertext), its
identity-binding AAD, the length-prefixed framing, and the CultNet carriage: each frame is a
`cultnet.operation_request.v0` / `operation_response.v0` for service
`gamecult.model.connector`, operation `invoke`, its `message_id` the request id and its
payload the base64 of the sealed envelope.
`gamecult.model.connector`, operation `invoke` (or `spend_readout`), its `message_id` the
request id and its payload the base64 of the sealed envelope. `spend_readout` answers a
`SpendReadout` (`gamecult.model.spend_readout.v1`): each caller's spent and reserved
micro-USD against its ceiling in its current window, and the backend's plan quota;
`ConnectorClient::spend_readout` asks.
- `schemas/` - one JSON Schema per v3 document, generated from the types.
- `src/client.rs` - `ConnectorClient`, the one-request caller side.
- `src/service.rs` (feature `daemon`) - `TransportService`: caller admission, durable
@ -37,13 +40,15 @@ Crate `crates/claude-connector` (library only; the binary that wires it to `serv
- `ClaudeCodeBackend` implements `ProviderBackend` for provider `claude-code` by driving the
unmodified `claude` binary, one child per call, in print stream-json mode with no tools, no MCP
servers, no settings sources and the system prompt replaced. Tool-less turns only: a request
with tools, an output schema or a final tool result is refused before any spawn.
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`); a handle that cannot be proven to
match the request re-seeds from the caller's transcript, and every resume forks.
- `ProcessPort` is the one OS seam (`StdProcessPort`, unix, kills the child's whole process
group); `SessionStore` and `StderrSink` are the other ports. `ClaudeCodeConfig` is what the
binary hands in; the credential is read by name from the connector's own environment
(`CLAUDE_CODE_OAUTH_TOKEN`) and copied only into the child.
binary hands in; the credential is read from a systemd credential and reaches the child on a
pipe (`CLAUDE_CODE_OAUTH_TOKEN_FILE_DESCRIPTOR`), never in an environment. `QuotaCell` keeps the
plan windows the CLI's `rate_limit_event`s report, for the spend readout.
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

76
crates/claude-connector/src/backend.rs
View file

@ -1,11 +1,12 @@
use std::fs::File;
use std::io;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Instant;
use model_connector::{
FailureKind, InputItem, ProviderBackend, ProviderCall, ProviderRequest, TransportEvent,
TransportEventPayload, TransportOutcome, TransportResult, Usage,
FailureKind, InputItem, PlanQuotaSnapshot, ProviderBackend, ProviderCall, ProviderRequest,
TransportEvent, TransportEventPayload, TransportOutcome, TransportResult, Usage,
};
use sha2::{Digest, Sha256};
use zeroize::Zeroizing;
@ -14,6 +15,7 @@ use crate::config::{ClaudeCodeConfig, PROVIDER};
use crate::handle::{ClaudeCodeHandle, SessionStore, plan, prefix_digest, stdin_bytes};
use crate::launch::{PromptFile, argv, child_env};
use crate::process::{ProcessPort, RunningChild, StdoutRead};
use crate::quota::QuotaCell;
use crate::stream::{Failure, Finished, Step, StreamState};
/// Where the child's credential comes from, read afresh at every use. The
@ -22,6 +24,10 @@ use crate::stream::{Failure, Finished, Step, StreamState};
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
@ -59,6 +65,16 @@ impl TokenSource for CredentialTokenSource {
token.truncate(trimmed);
(!token.is_empty()).then_some(token)
}
/// No directory at all is the unit's environment lacking its credentials,
/// which is not the same fault as a credential the directory lacks.
fn missing(&self) -> &'static str {
if self.directory.is_none() {
"credentials_directory"
} else {
"token_credential"
}
}
}
/// Where the tail of a child's stderr goes. The backend never puts it in a
@ -75,6 +91,7 @@ pub struct ClaudeCodeBackend {
sessions: Box<dyn SessionStore>,
sink: Box<dyn StderrSink>,
tokens: Box<dyn TokenSource>,
quota: Arc<QuotaCell>,
}
impl ClaudeCodeBackend {
@ -104,9 +121,15 @@ impl ClaudeCodeBackend {
sessions,
sink,
tokens,
quota: Arc::new(QuotaCell::new()),
}
}
/// The plan windows every stream this backend reads feeds.
pub fn quota(&self) -> Arc<QuotaCell> {
self.quota.clone()
}
/// The one pin check: `ready` and every spawn use it, and nothing caches
/// its answer.
fn binary_matches_pin(&self) -> bool {
@ -144,7 +167,7 @@ impl ClaudeCodeBackend {
));
};
let argv = argv(&self.config, request, &plan, prompt.path());
let env = child_env(&self.config, request, token);
let env = child_env(&self.config, request);
let deadline = Instant::now() + self.config.limits.deadline;
let spawned = self.port.spawn(
@ -152,6 +175,7 @@ impl ClaudeCodeBackend {
&env,
&self.config.work_dir,
stdin_bytes(plan.blocks()),
&token,
);
let Ok(mut child) = spawned else {
return failure(Failure::never_reached(
@ -159,7 +183,13 @@ impl ClaudeCodeBackend {
"claude-code: the child could not be spawned",
));
};
let finished = read_stream(child.as_mut(), deadline, &request.model);
let finished = read_stream(
child.as_mut(),
deadline,
&request.model,
request.output_schema_json.as_deref(),
&self.quota,
);
child.kill();
child.wait();
let tail = child.stderr_tail();
@ -190,14 +220,20 @@ impl ClaudeCodeBackend {
text: finished.texts.concat(),
});
}
let handle = ClaudeCodeHandle::Ended {
session_id: finished.session_id,
instance: self.config.instance.clone(),
prefix_sha256: prefix_digest(&request.instructions, &request.model, &prefix),
};
payloads.push(TransportEventPayload::ReasoningState {
state: handle.encode(),
});
// A schema turn leaves the CLI's StructuredOutput tool exchange in its
// session, and a later turn without the schema runs with no tools: it
// must not resume that session. So a schema turn hands out no handle,
// and whatever follows it is seeded from the caller's transcript.
if request.output_schema_json.is_none() {
let handle = ClaudeCodeHandle::Ended {
session_id: finished.session_id,
instance: self.config.instance.clone(),
prefix_sha256: prefix_digest(&request.instructions, &request.model, &prefix),
};
payloads.push(TransportEventPayload::ReasoningState {
state: handle.encode(),
});
}
let events = payloads
.into_iter()
@ -231,14 +267,20 @@ impl ProviderBackend for ClaudeCodeBackend {
let (events, outcome) = self.run(call.request());
call.transported(events, outcome)
}
fn plan_quota(&self) -> PlanQuotaSnapshot {
self.quota.snapshot()
}
}
fn read_stream(
child: &mut dyn RunningChild,
deadline: Instant,
model: &str,
schema: Option<&str>,
quota: &QuotaCell,
) -> Result<Finished, Failure> {
let mut state = StreamState::new(model);
let mut state = StreamState::new(model, schema, quota);
let finished = loop {
match child.read_line(deadline) {
StdoutRead::Line(line) => {
@ -295,12 +337,14 @@ pub(crate) fn readiness(
if !config.prompt_dir.is_dir() {
return Err("prompt_dir");
}
if tokens.read().is_none() {
return Err("token_credential");
}
if !binary_matches_pin(config) {
return Err("claude_binary_sha256");
}
// Last: outside the unit the credential is the one thing always absent,
// and it must not hide the faults the document itself has.
if tokens.read().is_none() {
return Err(tokens.missing());
}
Ok(())
}

523
crates/claude-connector/src/backend/tests.rs
View file

@ -6,16 +6,19 @@ use std::time::Duration;
use model_connector::{
CallerAdmission, ModelPrice, PriceTable, Rates, ReasoningEffort, SpendCeiling, SpendWindowMs,
ToolDefinition, TransportAdmission, TransportDisposition, TransportInvocation, TransportKey,
TransportReceipt, TransportService, decrypt_result, encrypt_invocation,
StopReason, ToolDefinition, TransportAdmission, TransportDisposition, TransportInvocation,
TransportKey, TransportReceipt, TransportService, decrypt_result, encrypt_invocation,
};
use super::*;
use crate::config::ClaudeCodeLimits;
use crate::handle::{FsSessionStore, SessionStore};
use crate::launch::{ChildEnv, TOKEN_VAR};
use crate::launch::ChildEnv;
const CANARY_TOKEN: &str = "canary-oauth-token-value";
/// The variable the connector's own environment must never hand on.
const TOKEN_ENV: &str = "CLAUDE_CODE_OAUTH_TOKEN";
const TOKEN_FD_ENV: &str = "CLAUDE_CODE_OAUTH_TOKEN_FILE_DESCRIPTOR";
const MODEL: &str = "claude-haiku-5-5";
const INSTANCE: &str = "persona-a";
const SESSION: &str = "11111111-2222-3333-4444-555555555555";
@ -62,11 +65,71 @@ const USAGE: &str = r#"{"claude-haiku-5-5":{"inputTokens":2,"outputTokens":4,"ca
const CUMULATIVE: &str = r#"{"claude-haiku-5-5":{"inputTokens":900000,"outputTokens":900000,"cacheReadInputTokens":900000,"cacheCreationInputTokens":900000,"thinkingTokens":900000,"costUSD":99.0}}"#;
fn result_line(stop_reason: &str, since_last: &str) -> String {
result_line_with(stop_reason, since_last, "")
}
/// A success result with `extra` (a leading comma and fields) appended.
fn result_line_with(stop_reason: &str, since_last: &str, extra: &str) -> String {
format!(
r#"{{"type":"result","subtype":"success","is_error":false,"stop_reason":"{stop_reason}","session_id":"{SESSION}","total_cost_usd":99.0,"modelUsage":{CUMULATIVE},"model_usage_since_last_result":{since_last}}}"#
r#"{{"type":"result","subtype":"success","is_error":false,"stop_reason":"{stop_reason}","session_id":"{SESSION}","total_cost_usd":99.0,"modelUsage":{CUMULATIVE},"model_usage_since_last_result":{since_last}{extra}}}"#
)
}
// The shapes `claude -p --json-schema` produced on the pinned CLI (2.1.312):
// the answer is a call of the CLI's own StructuredOutput tool, the result
// carries it as `structured_output` with stop_reason `tool_use`, and a
// `rate_limit_event` reports the plan's windows under `unifiedWindows`.
const SCHEMA: &str = r#"{"type":"object","properties":{"answer":{"type":"string"},"n":{"type":"integer"}},"required":["answer","n"],"additionalProperties":false}"#;
const STRUCTURED: &str = r#"{"answer":"say hello","n":0}"#;
fn structured_call_line(name: &str) -> String {
format!(
r#"{{"type":"assistant","message":{{"id":"msg_s","role":"assistant","content":[{{"type":"tool_use","id":"toolu_01","name":"{name}","input":{STRUCTURED},"caller":{{"type":"direct"}}}}],"stop_reason":null}},"session_id":"{SESSION}"}}"#
)
}
fn structured_tool_result_line() -> String {
format!(
r#"{{"type":"user","message":{{"role":"user","content":[{{"tool_use_id":"toolu_01","type":"tool_result","content":"Structured output provided successfully"}}]}},"parent_tool_use_id":null,"session_id":"{SESSION}"}}"#
)
}
fn plan_rate_limit_line() -> String {
r#"{"type":"rate_limit_event","rate_limit_info":{"status":"allowed","resetsAt":4102444800,"rateLimitType":"five_hour","overageStatus":"rejected","isUsingOverage":false,"unifiedWindows":{"five_hour":{"utilization":0.08,"resetsAt":4102444800,"status":"allowed"},"seven_day":{"utilization":0.3,"resetsAt":4102531200,"status":"allowed"}}}}"#.to_string()
}
fn structured_result_line() -> String {
result_line_with(
"tool_use",
USAGE,
&format!(r#","structured_output":{STRUCTURED},"result":"ignored""#),
)
}
fn structured_script() -> Script {
Script::lines(vec![
init_line(MODEL, SESSION),
structured_call_line("StructuredOutput"),
structured_tool_result_line(),
plan_rate_limit_line(),
structured_result_line(),
])
}
fn schema_request() -> ProviderRequest {
let mut request = request(vec![user("hello")]);
request.output_schema_json = Some(SCHEMA.to_string());
request
}
fn a_tool() -> ToolDefinition {
ToolDefinition {
name: "lookup".into(),
description: "d".into(),
parameters_json: "{}".into(),
}
}
fn error_result_line(extra: &str) -> String {
format!(
r#"{{"type":"result","subtype":"error_during_execution","is_error":true,"session_id":"{SESSION}","total_cost_usd":99.0,"modelUsage":{CUMULATIVE},"model_usage_since_last_result":{USAGE}{extra}}}"#
@ -146,6 +209,8 @@ struct Seen {
env: Vec<(String, String)>,
cwd: PathBuf,
stdin: Vec<u8>,
/// The token the backend handed the port, which a real port puts on a pipe.
token: String,
/// The file `--system-prompt-file` named, as it was when the child started.
prompt_file: Option<(PathBuf, String)>,
/// Every deadline the backend passed to a read of this child.
@ -173,6 +238,7 @@ impl ProcessPort for ScriptedPort {
env: &ChildEnv,
cwd: &Path,
stdin: Vec<u8>,
token: &str,
) -> io::Result<Box<dyn RunningChild>> {
self.spawns.fetch_add(1, Ordering::SeqCst);
let deadlines = Arc::new(Mutex::new(Vec::new()));
@ -191,6 +257,7 @@ impl ProcessPort for ScriptedPort {
.collect(),
cwd: cwd.to_path_buf(),
stdin,
token: token.to_string(),
};
let script = self
.scripts
@ -549,7 +616,7 @@ fn resume_always_forks() {
}
#[test]
fn child_env_is_cleared_and_names_only_the_token_var() {
fn child_env_is_cleared_and_names_the_token_descriptor_never_the_token() {
let rig = rig(vec![ok_script(SESSION, "hi")]);
let result = rig.run(request(vec![user("hello")]));
let seen = rig.seen();
@ -562,10 +629,7 @@ fn child_env_is_cleared_and_names_only_the_token_var() {
"CLAUDE_CODE_MAX_OUTPUT_TOKENS".to_string(),
"1024".to_string()
),
(
"CLAUDE_CODE_OAUTH_TOKEN".to_string(),
CANARY_TOKEN.to_string()
),
(TOKEN_FD_ENV.to_string(), "3".to_string()),
("CLAUDE_CONFIG_DIR".to_string(), "/cfg".to_string()),
("HOME".to_string(), "/cfg".to_string()),
(
@ -575,6 +639,13 @@ fn child_env_is_cleared_and_names_only_the_token_var() {
]
);
assert_eq!(seen.cwd, PathBuf::from("/work"));
// The credential travels beside the environment, to the port, alone.
assert_eq!(seen.token, CANARY_TOKEN);
assert!(
seen.env
.iter()
.all(|(name, value)| name != TOKEN_ENV && value != CANARY_TOKEN)
);
assert!(!format!("{result:?}").contains(CANARY_TOKEN));
}
@ -600,11 +671,10 @@ fn the_token_never_reaches_a_message_an_event_or_a_debug_form() {
let env = child_env(
&config(rig.dir.path(), INSTANCE, limits()),
&request(vec![user("x")]),
Zeroizing::new(CANARY_TOKEN.to_string()),
);
let shown = format!("{env:?}");
assert!(!shown.contains(CANARY_TOKEN));
assert!(shown.contains("CLAUDE_CODE_OAUTH_TOKEN"));
assert!(shown.contains(TOKEN_FD_ENV));
}
// ---- ready ----
@ -664,6 +734,15 @@ fn ready_checks_the_pinned_binary_digest() {
ready_field(pin, &binary, &work, &prompts, false),
Err("token_credential")
);
// The credential is checked last: a document fault is never hidden by it.
assert_eq!(
ready_field([0; 32], &binary, &work, &prompts, false),
Err("claude_binary_sha256")
);
assert_eq!(
ready_field(pin, &binary, &dir.path().join("absent"), &prompts, false),
Err("work_dir")
);
assert!(ready(pin, &binary, &work, &prompts, true));
assert!(!ready([0; 32], &binary, &work, &prompts, true), "other pin");
@ -692,7 +771,7 @@ fn token_comes_from_the_credential_file_not_the_environment() {
static CANARY_PARENT: std::sync::Once = std::sync::Once::new();
CANARY_PARENT.call_once(|| {
// SAFETY: set once to a value no code path reads; the test proves none does.
unsafe { std::env::set_var(TOKEN_VAR, "canary-environment-token") };
unsafe { std::env::set_var(TOKEN_ENV, "canary-environment-token") };
});
let dir = sandbox();
let credentials = dir.path().join("credentials");
@ -730,7 +809,7 @@ fn token_comes_from_the_credential_file_not_the_environment() {
assert!(seen.lock().unwrap().argv.is_empty(), "no child was spawned");
// The file is read afresh at the call, its trailing newline trimmed, and
// the child's environment carries the file's value.
// the port is handed the file's value, which is in no environment.
std::fs::write(
credentials.join("claude-token"),
format!("{CANARY_TOKEN}\n"),
@ -744,8 +823,14 @@ fn token_comes_from_the_credential_file_not_the_environment() {
outcome(&result),
TransportOutcome::Completed { .. }
));
let env = seen.lock().unwrap().env.clone();
assert!(env.contains(&(TOKEN_VAR.to_string(), CANARY_TOKEN.to_string())));
let given = seen.lock().unwrap().clone();
assert_eq!(given.token, CANARY_TOKEN);
assert!(
given
.env
.iter()
.all(|(name, value)| name != TOKEN_ENV && !value.contains("canary-environment-token"))
);
// An empty credential is no token.
std::fs::write(credentials.join("claude-token"), "\n").unwrap();
@ -997,15 +1082,9 @@ fn usage_zero() -> Usage {
}
#[test]
fn tools_schema_or_final_tool_result_fail_before_spawn() {
fn tools_or_final_tool_result_fail_before_spawn() {
let mut with_tools = request(vec![user("a")]);
with_tools.tools.push(ToolDefinition {
name: "lookup".into(),
description: "d".into(),
parameters_json: "{}".into(),
});
let mut with_schema = request(vec![user("a")]);
with_schema.output_schema_json = Some("{}".into());
with_tools.tools.push(a_tool());
let final_tool_result = request(vec![
user("a"),
InputItem::ToolCall {
@ -1032,7 +1111,6 @@ fn tools_schema_or_final_tool_result_fail_before_spawn() {
let mut messages = Vec::new();
for refused in [
with_tools,
with_schema,
final_tool_result,
ends_without_user_text,
tool_result_beside_text,
@ -1044,9 +1122,9 @@ fn tools_schema_or_final_tool_result_fail_before_spawn() {
}
assert_eq!(rig.spawns(), 0);
// Each refusal names what lifts it.
assert_eq!(messages[2], messages[4]);
assert_eq!(messages[1], messages[3]);
let distinct: std::collections::HashSet<_> = messages.iter().collect();
assert_eq!(distinct.len(), 4, "{messages:?}");
assert_eq!(distinct.len(), 3, "{messages:?}");
}
// ---- the stream ----
@ -1271,7 +1349,7 @@ fn no_result_event_is_a_failure_with_unknown_usage() {
let zero_billed = {
let (_, _, usage) = failure_of(&rig(vec![]).run({
let mut bad = request(vec![user("a")]);
bad.output_schema_json = Some("{}".into());
bad.tools.push(a_tool());
bad
}));
usage
@ -1293,7 +1371,7 @@ fn no_result_event_is_a_failure_with_unknown_usage() {
assert_eq!(
parts(&rig(vec![]).run({
let mut bad = request(vec![user("a")]);
bad.output_schema_json = Some("{}".into());
bad.tools.push(a_tool());
bad
}))
.1
@ -1374,6 +1452,7 @@ fn a_port_that_cannot_spawn_is_a_provider_failure() {
_: &ChildEnv,
_: &Path,
_: Vec<u8>,
_: &str,
) -> io::Result<Box<dyn RunningChild>> {
Err(io::Error::from(io::ErrorKind::NotFound))
}
@ -1742,6 +1821,199 @@ impl Drop for TempDir {
}
}
// ---- structured output and the plan quota ----
#[test]
fn schema_request_lowers_to_json_schema() {
let rig = rig(vec![structured_script(), ok_script(SESSION, "hi")]);
rig.run(schema_request());
let argv = rig.seen().argv;
assert_eq!(arg_after(&argv, "--json-schema"), Some(SCHEMA));
assert_eq!(argv.iter().filter(|a| *a == "--json-schema").count(), 1);
// The tool-less hygiene stays with it.
assert_eq!(arg_after(&argv, "--tools"), Some(""));
assert!(argv.iter().any(|a| a == "--strict-mcp-config"));
assert!(!argv.iter().any(|a| a == "--mcp-config"));
rig.run(request(vec![user("hello")]));
assert!(!rig.seen().argv.iter().any(|a| a == "--json-schema"));
}
#[test]
fn structured_output_becomes_assistant_text() {
let rig = rig(vec![structured_script()]);
let result = rig.run(schema_request());
let TransportOutcome::Completed {
stop_reason, usage, ..
} = outcome(&result)
else {
panic!("expected completion, got {:?}", outcome(&result));
};
assert_eq!(*stop_reason, StopReason::EndTurn);
assert_eq!(*usage, usage_cost_fixture());
let (events, _) = parts(&result);
let texts: Vec<&str> = events
.iter()
.filter_map(|event| match &event.payload {
TransportEventPayload::TextDelta { text } => Some(text.as_str()),
_ => None,
})
.collect();
assert_eq!(texts, [STRUCTURED]);
let items = result.transcript_items();
assert_eq!(items[0], assistant(STRUCTURED));
// The answer is JSON the caller decodes with serde: it is the object.
let decoded: serde_json::Value = serde_json::from_str(STRUCTURED).unwrap();
assert_eq!(decoded["answer"], "say hello");
}
fn usage_cost_fixture() -> Usage {
usage(2, 7, 1043, 4, Some(3))
}
#[test]
fn a_schema_turn_without_structured_output_is_a_provider_failure_billed_as_reported() {
for line in [
result_line("tool_use", USAGE),
result_line_with("end_turn", USAGE, r#","structured_output":null"#),
] {
let rig = rig(vec![Script::lines(vec![
init_line(MODEL, SESSION),
text_line("msg_1", "prose, not the schema"),
line,
])]);
let result = rig.run(schema_request());
let (kind, _, usage) = failure_of(&result);
assert_eq!(kind, FailureKind::Provider);
assert_eq!(usage, Some(usage_cost_fixture()), "the turn was billed");
assert!(result.transcript_items().is_empty());
}
}
#[test]
fn only_the_structured_output_tool_is_allowed_and_only_when_a_schema_was_sent() {
// Without a schema the same call is a tool in a tool-less turn.
let first = rig(vec![structured_script()]);
let (kind, message, _) = failure_of(&first.run(request(vec![user("hello")])));
assert_eq!(kind, FailureKind::Provider);
assert!(message.contains("tool"), "{message}");
// With a schema, any other tool is still refused.
let second = rig(vec![Script::lines(vec![
init_line(MODEL, SESSION),
structured_call_line("Bash"),
structured_result_line(),
])]);
let (kind, message, _) = failure_of(&second.run(schema_request()));
assert_eq!(kind, FailureKind::Provider);
assert!(message.contains("tool"), "{message}");
// `tool_use` is a stop reason only for a schema turn.
let third = rig(vec![Script::lines(vec![
init_line(MODEL, SESSION),
text_line("msg_1", "hi"),
result_line("tool_use", USAGE),
])]);
let (kind, message, _) = failure_of(&third.run(request(vec![user("hello")])));
assert_eq!(kind, FailureKind::Provider);
assert!(message.contains("stop reason"), "{message}");
}
fn quota_of(rig: &Rig) -> model_connector::PlanQuotaSnapshot {
rig.backend.plan_quota()
}
/// The windows with their observation stamps zeroed, for comparing readings.
fn readings(rig: &Rig) -> Vec<model_connector::PlanQuota> {
let mut windows = quota_of(rig).windows;
for window in &mut windows {
window.observed_at_unix_ms = 0;
}
windows
}
fn window(name: &str, permille: u32, resets_at_unix_ms: u64) -> model_connector::PlanQuota {
model_connector::PlanQuota {
window: name.to_string(),
utilization_permille: Some(permille),
resets_at_unix_ms,
status: "allowed".to_string(),
observed_at_unix_ms: 0,
}
}
#[test]
fn rate_limit_event_updates_quota_cell() {
let rig = rig(vec![structured_script()]);
assert!(quota_of(&rig).windows.is_empty());
let before = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_millis() as u64;
rig.run(schema_request());
let snapshot = quota_of(&rig);
assert_eq!(
readings(&rig),
vec![
window("five_hour", 80, 4_102_444_800_000),
window("seven_day", 300, 4_102_531_200_000),
]
);
let after = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_millis() as u64;
for each in &snapshot.windows {
assert!((before..=after).contains(&each.observed_at_unix_ms));
}
// The cell is the one the backend hands out.
assert_eq!(rig.backend.quota().snapshot(), snapshot);
}
#[test]
fn a_rejected_turn_still_reports_its_windows_and_a_bad_window_costs_only_itself() {
let event = r#"{"type":"rate_limit_event","rate_limit_info":{"status":"rejected","unifiedWindows":{"five_hour":{"utilization":0.9996,"resetsAt":4102444800,"status":"rejected"},"seven_day":{"resetsAt":4102531200,"status":"allowed"},"idle":{"utilization":0,"resetsAt":4102444800,"status":"allowed"},"odd":{"utilization":"high","resetsAt":1,"status":"allowed"},"negative":{"utilization":-0.5,"resetsAt":1,"status":"allowed"}}}}"#;
let rig = rig(vec![Script::lines(vec![
init_line(MODEL, SESSION),
event.to_string(),
error_result_line(r#","api_error_status":429"#),
])]);
let (kind, _, _) = failure_of(&rig.run(request(vec![user("hello")])));
assert_eq!(kind, FailureKind::RateLimited);
let mut rejected = window("five_hour", 1000, 4_102_444_800_000);
rejected.status = "rejected".to_string();
// Unknown is never zero: a window missing a field or of another type is
// skipped, not read as empty.
// A window at exactly zero use is a reading, and is kept.
assert_eq!(
readings(&rig),
vec![rejected, window("idle", 0, 4_102_444_800_000)]
);
}
#[test]
fn a_later_event_replaces_only_the_windows_it_names() {
let only_five = r#"{"type":"rate_limit_event","rate_limit_info":{"status":"allowed","unifiedWindows":{"five_hour":{"utilization":0.5,"resetsAt":4102448400,"status":"allowed"}}}}"#;
let rig = rig(vec![
structured_script(),
Script::lines(vec![
init_line(MODEL, SESSION),
only_five.to_string(),
text_line("msg_1", "hi"),
result_line("end_turn", USAGE),
]),
]);
rig.run(schema_request());
rig.run(request(vec![user("again")]));
assert_eq!(
readings(&rig),
vec![
window("five_hour", 500, 4_102_448_400_000),
window("seven_day", 300, 4_102_531_200_000),
]
);
}
#[cfg(unix)]
mod real {
use super::*;
@ -1892,11 +2164,7 @@ mod real {
);
let fake = fake(&body, limits());
let config = config(fake.dir.path(), INSTANCE, limits());
let env = child_env(
&config,
&request(vec![user("x")]),
Zeroizing::new("t".to_string()),
);
let env = child_env(&config, &request(vec![user("x")]));
let argv = vec![
fake.dir
.path()
@ -1906,7 +2174,7 @@ mod real {
];
let port = crate::process::StdProcessPort::new(limits());
let child = port
.spawn(&argv, &env, fake.dir.path(), Vec::new())
.spawn(&argv, &env, fake.dir.path(), Vec::new(), "t")
.unwrap();
let (leader_pid, grandchild_pid) = (pid_in(&leader), pid_in(&grandchild));
drop(child);
@ -1966,13 +2234,14 @@ mod real {
names,
[
"CLAUDE_CODE_MAX_OUTPUT_TOKENS",
"CLAUDE_CODE_OAUTH_TOKEN",
TOKEN_FD_ENV,
"CLAUDE_CONFIG_DIR",
"HOME",
"PATH"
]
);
assert!(env.contains(&format!("CLAUDE_CODE_OAUTH_TOKEN={CANARY_TOKEN}")));
assert!(env.contains(&format!("{TOKEN_FD_ENV}=3")));
assert!(!env.contains(CANARY_TOKEN));
let cwd = std::fs::read_to_string(probe.path().join("pwd")).unwrap();
assert_eq!(cwd.trim(), fake.dir.path().join("work").to_str().unwrap());
let stdin = std::fs::read_to_string(probe.path().join("stdin")).unwrap();
@ -1980,6 +2249,59 @@ mod real {
assert!(stdin.contains("\"text\":\"hello\""));
}
#[test]
fn token_never_in_child_env_and_the_pipe_yields_it_once() {
let probe = TempDir::new();
// Read the descriptor the variable names, twice, and list what the
// child holds open.
let body = format!(
"fd=\"${TOKEN_FD_ENV}\"\neval \"cat <&$fd\" > '{d}/first'\neval \"cat <&$fd\" > '{d}/second'\nenv > '{d}/env'\nls /proc/$$/fd > '{d}/fds'\n{lines}",
d = probe.path().display(),
lines = printf_lines(&good_lines())
);
let fake = fake(&body, limits());
let result = fake.run(request(vec![user("hello")]));
assert!(matches!(
outcome(&result),
TransportOutcome::Completed { .. }
));
let read = |name: &str| std::fs::read_to_string(probe.path().join(name)).unwrap();
assert_eq!(read("first"), CANARY_TOKEN, "the pipe holds the token");
assert_eq!(read("second"), "", "and nothing the second time");
let env = read("env");
assert!(!env.contains(CANARY_TOKEN) && !env.contains(&format!("{TOKEN_ENV}=")));
assert!(env.contains(&format!("{TOKEN_FD_ENV}=3")));
// 0 1 2, the pipe at 3, and the shell's own working descriptors:
// no descriptor of the connector's other pipes survives into the child.
let held: Vec<u32> = read("fds")
.lines()
.filter_map(|l| l.trim().parse().ok())
.collect();
assert!(held.contains(&3));
}
#[test]
fn a_token_too_large_for_a_pipe_is_refused_at_its_bound_not_hung() {
let fake = fake("exit 0", limits());
let config = config(fake.dir.path(), INSTANCE, limits());
let env = child_env(&config, &request(vec![user("x")]));
let port = crate::process::StdProcessPort::new(limits());
let argv = vec![
fake.dir
.path()
.join("claude")
.to_string_lossy()
.into_owned(),
];
let spawn = |token: &str| port.spawn(&argv, &env, fake.dir.path(), Vec::new(), token);
let at_bound = spawn(&"t".repeat(4096)).expect("4096 bytes fit the bound");
drop(at_bound);
assert_eq!(
spawn(&"t".repeat(4097)).err().map(|e| e.kind()),
Some(io::ErrorKind::InvalidInput)
);
}
#[test]
fn the_output_caps_hold_at_their_exact_sizes() {
let lines = good_lines();
@ -2109,3 +2431,132 @@ mod real {
));
}
}
// ---- a call that never reached the provider ----
/// The shape claude 2.1.312 gave with no usable credential: an error result
/// with no status, no API time and no usage rows. `extra` replaces the
/// per-turn usage field.
fn not_logged_in_line(status: &str, api_ms: u64, extra: &str) -> String {
format!(
r#"{{"type":"result","subtype":"success","is_error":true,{status}"duration_ms":118,"duration_api_ms":{api_ms},"stop_reason":"stop_sequence","result":"Not logged in","total_cost_usd":0,"modelUsage":{{}}{extra}}}"#
)
}
#[test]
fn a_result_that_made_no_api_call_bills_nothing() {
let reservation = reservation();
let empty = r#","model_usage_since_last_result":{}"#;
let cases = [
// (status field, api ms, usage field, billed nothing)
(r#""api_error_status":null,"#, 0, empty, true),
("", 0, empty, true),
("", 0, "", true),
// An API call was made, or the status names a server fault, or the
// CLI reported rows: the reservation stands.
("", 2300, empty, false),
(r#""api_error_status":500,"#, 0, empty, false),
(
r#""api_error_status":null,"#,
0,
r#","model_usage_since_last_result":"?""#,
false,
),
];
for (status, api_ms, extra, free) in cases {
let line = not_logged_in_line(status, api_ms, extra);
let result = rig(vec![Script::lines(vec![
init_line(MODEL, SESSION),
line.clone(),
])])
.run(request(vec![user("hi")]));
let (_, _, usage) = failure_of(&result);
if free {
assert_eq!(usage, Some(usage_zero()), "{line}");
assert_eq!(parts(&result).1.cost_micro_usd, 0, "{line}");
} else {
assert_eq!(usage, None, "{line}");
assert_eq!(parts(&result).1.cost_micro_usd, reservation, "{line}");
}
}
}
#[test]
fn a_no_call_result_that_reports_rows_is_billed_those_rows() {
let line = not_logged_in_line(
"",
0,
&format!(r#","model_usage_since_last_result":{USAGE}"#),
);
let result = rig(vec![Script::lines(vec![init_line(MODEL, SESSION), line])])
.run(request(vec![user("hi")]));
let (_, _, billed) = failure_of(&result);
assert_eq!(billed, Some(usage(2, 7, 1043, 4, Some(3))));
}
// ---- the structured output is an object ----
#[test]
fn a_structured_output_that_is_not_an_object_fails_as_provider_with_usage() {
for output in ["\"text\"", "[1]", "42", "true", "null"] {
let result = rig(vec![Script::lines(vec![
init_line(MODEL, SESSION),
structured_call_line("StructuredOutput"),
result_line_with(
"tool_use",
USAGE,
&format!(r#","structured_output":{output}"#),
),
])])
.run(schema_request());
let (kind, message, usage) = failure_of(&result);
assert_eq!(kind, FailureKind::Provider, "{output}");
assert!(message.contains("structured output"), "{message}");
assert!(usage.is_some(), "{output}: the turn ran and is billed");
}
}
#[test]
fn an_object_answer_to_a_schema_whose_top_type_is_not_object_fails() {
for (schema, accepted) in [
(r#"{"type":"array"}"#, false),
(r#"{"type":"string"}"#, false),
(r#"{"type":["object","null"]}"#, false),
(r#"{"type":"object"}"#, true),
(r#"{"properties":{}}"#, true),
] {
let mut asked = request(vec![user("hello")]);
asked.output_schema_json = Some(schema.to_string());
let result = rig(vec![structured_script()]).run(asked);
assert_eq!(
matches!(outcome(&result), TransportOutcome::Completed { .. }),
accepted,
"{schema}"
);
}
}
// ---- a schema turn does not hand out a session to resume ----
#[test]
fn a_turn_after_a_schema_turn_is_seeded_not_resumed() {
let rig = rig(vec![structured_script(), ok_script(SESSION, "next")]);
let first = rig.run(schema_request());
assert!(
!parts(&first)
.0
.iter()
.any(|e| matches!(e.payload, TransportEventPayload::ReasoningState { .. })),
"a schema turn leaves no handle"
);
rig.run(request(vec![
user("hello"),
InputItem::AssistantText {
text: STRUCTURED.to_string(),
},
user("and now plainly"),
]));
let argv = rig.seen().argv;
assert!(!argv.iter().any(|a| a == "--resume"), "{argv:?}");
assert!(!argv.iter().any(|a| a == "--json-schema"));
}

37
crates/claude-connector/src/daemon.rs
View file

@ -68,7 +68,8 @@ pub struct ClaudeConnectorConfig {
/// Exactly one caller until claude-handle-caller-binding lands.
#[cultcache(key = 17)]
pub callers: Vec<CallerConfig>,
/// What each allowed model costs; keys are canonical model ids.
/// What each allowed model costs. The keys are the known list of models:
/// ids exactly as the CLI reports them in its init event.
#[cultcache(key = 18)]
pub prices: BTreeMap<String, ModelPrice>,
}
@ -190,7 +191,7 @@ impl ClaudeConnectorConfig {
}
let mut seen = BTreeSet::new();
for model in &caller.allowed_models {
if !is_canonical_model(model)
if !has_model_id_form(model)
|| !seen.insert(model.as_str())
|| !self.prices.contains_key(model)
{
@ -200,7 +201,7 @@ impl ClaudeConnectorConfig {
if !self
.prices
.iter()
.all(|(model, price)| is_canonical_model(model) && is_fully_priced(price))
.all(|(model, price)| has_model_id_form(model) && is_fully_priced(price))
{
return Err(ConfigError::Invalid("prices"));
}
@ -239,28 +240,36 @@ fn is_trimmed_nonempty(text: &str) -> bool {
!text.is_empty() && text.trim() == text
}
/// A model id the CLI echoes back unchanged: `claude-<family>-<major>-<minor>`
/// with an optional `-YYYYMMDD` suffix, the family lowercase letters. Aliases
/// and bare prefixes are refused because the stream's `init.model` must equal
/// the requested model.
fn is_canonical_model(model: &str) -> bool {
/// 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 the price table's keys (ruling
/// claude-model-ids-from-price-table); 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 major = parts.next().unwrap_or("");
let minor = parts.next().unwrap_or("");
let date = parts.next();
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)
&& digits(minor, 2)
&& date.is_none_or(|date| date.len() == 8 && digits(date, 8))
&& parts.next().is_none()
&& after_ok
}
/// A price is a fuse only if every component it bills is nonzero: a zero rate

5
crates/claude-connector/src/handle.rs
View file

@ -89,8 +89,6 @@ impl SessionStore for FsSessionStore {
pub(crate) enum PlanError {
#[error("claude-code: tools are not supported by this backend yet")]
Tools,
#[error("claude-code: structured output is not supported by this backend yet")]
OutputSchema,
#[error("claude-code: tool results are not supported by this backend yet")]
ToolResult,
#[error("claude-code: the request ends without user text")]
@ -128,9 +126,6 @@ pub(crate) fn plan(
if !request.tools.is_empty() {
return Err(PlanError::Tools);
}
if request.output_schema_json.is_some() {
return Err(PlanError::OutputSchema);
}
let input = &request.input;
let run_start = input
.iter()

37
crates/claude-connector/src/launch.rs
View file

@ -9,19 +9,22 @@ use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use model_connector::{ProviderRequest, ReasoningEffort};
use zeroize::Zeroizing;
use crate::config::ClaudeCodeConfig;
use crate::handle::Plan;
/// The one credential variable: the child's environment is the only
/// environment the token ever enters.
pub(crate) const TOKEN_VAR: &str = "CLAUDE_CODE_OAUTH_TOKEN";
/// The variable that tells the CLI which file descriptor holds its credential.
/// The token itself never enters the environment: the process port writes it
/// into a pipe and hands the child the read end at [`TOKEN_FD`].
pub(crate) const TOKEN_FD_VAR: &str = "CLAUDE_CODE_OAUTH_TOKEN_FILE_DESCRIPTOR";
/// The descriptor number the child finds the token pipe at: the first one
/// after stdin, stdout and stderr.
pub(crate) const TOKEN_FD: i32 = 3;
const CHILD_PATH: &str = "/usr/local/bin:/usr/bin:/bin";
/// The environment of the child: exactly these names and nothing inherited.
/// Values are never formatted; `Debug` lists names only.
pub struct ChildEnv(Vec<(&'static str, Zeroizing<String>)>);
/// `Debug` lists names only.
pub struct ChildEnv(Vec<(&'static str, String)>);
impl ChildEnv {
pub fn iter(&self) -> impl Iterator<Item = (&str, &str)> {
@ -37,21 +40,17 @@ impl fmt::Debug for ChildEnv {
}
}
pub(crate) fn child_env(
config: &ClaudeCodeConfig,
request: &ProviderRequest,
token: Zeroizing<String>,
) -> ChildEnv {
pub(crate) fn child_env(config: &ClaudeCodeConfig, request: &ProviderRequest) -> ChildEnv {
let config_dir = config.config_dir.to_string_lossy().into_owned();
ChildEnv(vec![
("PATH", Zeroizing::new(CHILD_PATH.to_string())),
("HOME", Zeroizing::new(config_dir.clone())),
("CLAUDE_CONFIG_DIR", Zeroizing::new(config_dir)),
("PATH", CHILD_PATH.to_string()),
("HOME", config_dir.clone()),
("CLAUDE_CONFIG_DIR", config_dir),
(
"CLAUDE_CODE_MAX_OUTPUT_TOKENS",
Zeroizing::new(request.max_output_tokens.to_string()),
request.max_output_tokens.to_string(),
),
(TOKEN_VAR, token),
(TOKEN_FD_VAR, TOKEN_FD.to_string()),
])
}
@ -108,7 +107,8 @@ fn effort(effort: ReasoningEffort) -> &'static str {
/// rather than an argument because one argument is capped near 128 KiB
/// (E2BIG) and instructions are not, the child has
/// no tools, no MCP servers, no settings, no slash commands and nobody to ask
/// for permission, and a resume always forks.
/// for permission, and a resume always forks. 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,
@ -142,6 +142,9 @@ pub(crate) fn argv(
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(),

2
crates/claude-connector/src/lib.rs
View file

@ -9,6 +9,7 @@ mod daemon;
mod handle;
mod launch;
mod process;
mod quota;
mod stream;
pub use backend::{ClaudeCodeBackend, CredentialTokenSource, StderrSink, TokenSource};
@ -21,3 +22,4 @@ pub use launch::ChildEnv;
#[cfg(unix)]
pub use process::StdProcessPort;
pub use process::{ProcessPort, RunningChild, StdoutRead};
pub use quota::QuotaCell;

4
crates/claude-connector/src/main.rs
View file

@ -37,7 +37,9 @@ fn main() -> ExitCode {
#[derive(Debug, thiserror::Error)]
enum Failure {
#[error("usage: serve --config PATH | config init PATH --flag value ... | config check PATH")]
#[error(
"usage: serve --config PATH | config init PATH --flag value ... | config check PATH (check reads the credential from the environment it runs in: run it in the unit's environment)"
)]
Usage,
#[error(transparent)]
Config(#[from] ConfigError),

134
crates/claude-connector/src/process.rs
View file

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

151
crates/claude-connector/src/quota.rs Normal file
View file

@ -0,0 +1,151 @@
//! The plan's own usage windows, as the CLI reports them. The subscription
//! decides how much a plan may burn; the CLI's `rate_limit_event` is the only
//! place that number appears, so the latest observation of each window is kept
//! here and read out beside the connector's own spend ledger.
use std::collections::BTreeMap;
use std::sync::Mutex;
use std::time::{SystemTime, UNIX_EPOCH};
use model_connector::{PlanQuota, PlanQuotaSnapshot};
/// 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.
#[derive(Default)]
pub struct QuotaCell(Mutex<BTreeMap<String, PlanQuota>>);
impl QuotaCell {
pub fn new() -> Self {
Self::default()
}
/// Records the windows one event reported, stamping each with now. Each
/// replaces that window's earlier observation, and windows the event did
/// not mention stay as they were, stamp included.
pub fn observe(&self, windows: impl IntoIterator<Item = PlanQuota>) {
self.observe_at(unix_ms(), windows);
}
fn observe_at(&self, now: u64, windows: impl IntoIterator<Item = PlanQuota>) {
let mut held = self.0.lock().unwrap_or_else(|e| e.into_inner());
for mut quota in windows {
quota.observed_at_unix_ms = now;
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.0.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.
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 = 1_793_307_000_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);
}
#[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");
}
}

129
crates/claude-connector/src/stream.rs
View file

@ -3,10 +3,16 @@
use std::collections::BTreeMap;
use model_connector::{FailureKind, StopReason, Usage};
use model_connector::{FailureKind, PlanQuota, StopReason, Usage};
use serde::Deserialize;
use serde_json::Value;
use crate::quota::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`
/// (the kit bills its reservation) unless a result event reported it or the
/// call never reached the provider (zero).
@ -63,12 +69,18 @@ struct Line {
session_id: Option<String>,
model: Option<String>,
message: Option<Message>,
rate_limit_info: Option<RateLimitInfo>,
/// Read by hand: the shape moves with the CLI, and a field that changes
/// type must cost the quota reading, never the call.
rate_limit_info: Option<Value>,
is_error: Option<bool>,
stop_reason: Option<String>,
api_error_status: Option<Value>,
/// Milliseconds spent in API calls; zero means no call was made.
duration_api_ms: Option<u64>,
result: Option<Value>,
errors: Option<Value>,
/// The answer to a `--json-schema` request, an object.
structured_output: Option<Value>,
/// What this turn alone consumed. `modelUsage` is deliberately not read:
/// on a resumed session it carries every earlier turn too.
model_usage_since_last_result: Option<Value>,
@ -85,11 +97,7 @@ struct Block {
#[serde(rename = "type")]
kind: String,
text: Option<String>,
}
#[derive(Deserialize)]
struct RateLimitInfo {
status: Option<String>,
name: Option<String>,
}
/// One model's entry. The four counts are required: an entry that lacks one
@ -107,6 +115,13 @@ struct ModelUsage {
/// One call's stream, line by line.
pub(crate) struct StreamState<'a> {
model: &'a str,
/// Whether the request carried an output schema, so the answer is the
/// result's structured output.
structured: bool,
/// Whether that schema's top type admits an object (absent counts): the
/// structured output must be an object, never a bare scalar or array.
schema_top_is_object: bool,
quota: &'a QuotaCell,
session_id: Option<String>,
texts: Vec<String>,
message_id: Option<String>,
@ -114,9 +129,16 @@ pub(crate) struct StreamState<'a> {
}
impl<'a> StreamState<'a> {
pub(crate) fn new(model: &'a str) -> Self {
/// `schema` is the request's output schema text, when it carried one.
pub(crate) fn new(model: &'a str, schema: Option<&str>, quota: &'a QuotaCell) -> Self {
let top = schema
.and_then(|text| serde_json::from_str::<Value>(text).ok())
.and_then(|schema| schema.get("type").cloned());
Self {
model,
structured: schema.is_some(),
schema_top_is_object: top.is_none_or(|top| top.as_str() == Some("object")),
quota,
session_id: None,
texts: Vec::new(),
message_id: None,
@ -151,10 +173,11 @@ impl<'a> StreamState<'a> {
},
"assistant" => self.assistant(line.message),
"rate_limit_event" => {
self.plan_rejected |= line
.rate_limit_info
.and_then(|info| info.status)
.is_some_and(|status| status == "rejected");
if let Some(info) = &line.rate_limit_info {
self.plan_rejected |=
info.get("status").and_then(Value::as_str) == Some("rejected");
self.quota.observe(plan_windows(info));
}
Step::More
}
"result" => Step::Done(Box::new(self.result(line))),
@ -176,6 +199,9 @@ impl<'a> StreamState<'a> {
self.texts.push(text);
}
}
"tool_use"
if self.structured && block.name.as_deref() == Some(STRUCTURED_OUTPUT_TOOL) => {
}
"tool_use" => {
return fail(
FailureKind::Provider,
@ -195,7 +221,16 @@ impl<'a> StreamState<'a> {
line.api_error_status.as_ref().and_then(Value::as_u64),
Some(401 | 403 | 429)
);
let usage = read_usage(line.model_usage_since_last_result, door_refusal);
// The CLI found no credential and never called the API: an error with
// no status and no API time at all.
let no_call = !success
&& line.duration_api_ms == Some(0)
&& line.api_error_status.as_ref().is_none_or(Value::is_null);
let usage = read_usage(
line.model_usage_since_last_result,
door_refusal || no_call,
no_call,
);
if !success {
return Err(Failure {
kind: self.error_kind(&line.api_error_status, &line.result, &line.errors),
@ -210,8 +245,26 @@ impl<'a> StreamState<'a> {
usage: None,
});
};
let mut texts = std::mem::take(&mut self.texts);
if self.structured {
// The answer is the structured output alone; any prose around the
// tool call is not part of it.
let Some(output) = line
.structured_output
.filter(|value| value.is_object() && self.schema_top_is_object)
else {
return Err(Failure {
kind: FailureKind::Provider,
message: "claude-code: a schema was sent and the result carried no structured output object",
usage: Some(usage),
});
};
texts = vec![output.to_string()];
}
let stop_reason = match line.stop_reason.as_deref() {
Some("end_turn") => StopReason::EndTurn,
// The structured-output tool call is how such a turn ends.
Some("tool_use") if self.structured => StopReason::EndTurn,
Some("max_tokens") => StopReason::MaxOutputTokens,
Some("stop_sequence") => StopReason::StopSequence,
Some("refusal") => StopReason::Refusal,
@ -233,7 +286,7 @@ impl<'a> StreamState<'a> {
};
Ok(Finished {
session_id,
texts: std::mem::take(&mut self.texts),
texts,
message_id: self.message_id.take(),
stop_reason,
usage,
@ -263,20 +316,54 @@ impl<'a> StreamState<'a> {
}
}
/// The plan windows one `rate_limit_event` reports under
/// `rate_limit_info.unifiedWindows`: utilization as a fraction of the window,
/// reset as unix seconds. A window missing any of its three fields is skipped;
/// an unknown reading is never zero.
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_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 for a refusal
/// at the door (401, 403 or 429), which the provider cannot have billed; elsewhere
/// the CLI also zeroes usage on crashes after real work. `None` is billed the
/// reservation by the kit.
fn read_usage(reported: Option<Value>, door_refusal: bool) -> Option<Usage> {
let models: BTreeMap<String, ModelUsage> = serde_json::from_value(reported?).ok()?;
/// 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 billed the reservation by the kit.
fn read_usage(reported: Option<Value>, empty_is_zero: bool, absent_is_zero: bool) -> Option<Usage> {
let Some(reported) = reported else {
return absent_is_zero.then(|| sum_usage(std::iter::empty()));
};
let models: BTreeMap<String, ModelUsage> = serde_json::from_value(reported).ok()?;
if models.is_empty() {
return door_refusal.then(|| sum_usage(std::iter::empty()));
return empty_is_zero.then(|| sum_usage(std::iter::empty()));
}
Some(sum_usage(models.into_values()))
}

121
crates/claude-connector/tests/config_document.rs
View file

@ -6,6 +6,7 @@ use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::process::{Command, Output};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;
use claude_connector::{ClaudeConnectorConfig, ConfigError, load_config, write_config};
use cultcache_rs::{
@ -276,13 +277,6 @@ fn config_refuses_aliases_unpriced_models_and_second_callers() {
c.prices.insert("claude-sonnet-latest".to_string(), price());
}),
),
(
"no minor",
Box::new(|c| {
c.callers[0].allowed_models = vec!["claude-opus-5".to_string()];
c.prices.insert("claude-opus-5".to_string(), price());
}),
),
(
"uppercase family",
Box::new(|c| {
@ -343,7 +337,6 @@ fn config_refuses_aliases_unpriced_models_and_second_callers() {
"allowed_models",
"allowed_models",
"allowed_models",
"allowed_models",
"callers",
"callers",
];
@ -355,12 +348,23 @@ fn config_refuses_aliases_unpriced_models_and_second_callers() {
}
#[test]
fn a_dated_canonical_id_is_allowed() {
let dated = "claude-opus-5-1-20260101".to_string();
let mut config = valid();
config.callers[0].allowed_models = vec![dated.clone()];
config.prices.insert(dated, price());
config.validate().unwrap();
fn a_priced_id_of_any_loose_form_is_allowed_and_the_price_table_decides() {
// Real ids come with a minor, without one, and with a date either way.
for id in [
"claude-opus-5-1-20260101",
"claude-opus-5",
"claude-opus-5-20260101",
"claude-sonnet-4-5",
] {
let mut config = valid();
config.callers[0].allowed_models = vec![id.to_string()];
config.prices.insert(id.to_string(), price());
config.validate().unwrap_or_else(|e| panic!("{id}: {e}"));
// The same id with no price is not on the known list.
config.prices.remove(id);
assert_eq!(refused(&config), "allowed_models", "{id}");
}
}
#[test]
@ -543,13 +547,64 @@ fn no_refusal_echoes_a_value() {
// ---- the binary ----
/// How long any connector verb may take in these tests. `check` and `init`
/// return at once and `serve` is meant to refuse; a `serve` that starts
/// instead would run forever, so it fails a named test here, not the suite.
const VERB_DEADLINE: Duration = Duration::from_secs(20);
fn run(args: &[&str], envs: &[(&str, &Path)]) -> Output {
let mut command = Command::new(env!("CARGO_BIN_EXE_claude-connector"));
command.args(args).env_clear();
for (name, value) in envs {
command.env(name, value);
}
command.output().unwrap()
run_within(command, VERB_DEADLINE).unwrap_or_else(|| {
panic!(
"`{}` was still running after {VERB_DEADLINE:?}: it started",
args.join(" ")
)
})
}
/// The command's output, or `None` when it outlived `deadline` and was killed.
fn run_within(mut command: Command, deadline: Duration) -> Option<Output> {
use std::process::Stdio;
use std::time::Instant;
let mut child = command
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()
.unwrap();
let started = Instant::now();
while child.try_wait().unwrap().is_none() {
if started.elapsed() > deadline {
child.kill().unwrap();
child.wait().unwrap();
return None;
}
std::thread::sleep(Duration::from_millis(20));
}
Some(child.wait_with_output().unwrap())
}
#[cfg(unix)]
#[test]
fn a_process_past_its_deadline_is_killed_and_reported_as_none() {
let mut forever = Command::new("sleep");
forever.arg("60");
let started = std::time::Instant::now();
assert!(run_within(forever, Duration::from_millis(200)).is_none());
assert!(started.elapsed() < Duration::from_secs(10));
let mut quick = Command::new("true");
quick.env_clear();
assert!(
run_within(quick, Duration::from_secs(10))
.unwrap()
.status
.success()
);
}
fn shown(output: &Output) -> String {
@ -756,8 +811,42 @@ fn check_refuses_exactly_what_serve_refuses_and_names_the_field() {
assert!(init(&scratch, &flags).status.success());
let (checked, served) = check_and_serve(&scratch.path("connector.cc"), None);
assert!(!checked.status.success() && !served.status.success());
assert!(shown(&checked).contains("token_credential"));
assert!(shown(&checked).contains("credentials_directory"));
assert!(!shown(&checked).contains("token_credential"));
assert_eq!(shown(&checked), shown(&served));
// A credential directory without the credential is the other fault, and
// the token is checked last: a broken pin is named first.
let scratch = Scratch::new();
let (flags, credentials) = prepare(&scratch);
assert!(init(&scratch, &flags).status.success());
std::fs::remove_file(credentials.join("claude-token")).unwrap();
std::fs::write(scratch.path("claude"), b"other bytes").unwrap();
let (checked, _) = check_and_serve(&scratch.path("connector.cc"), Some(&credentials));
assert!(
shown(&checked).contains("claude_binary_sha256"),
"{}",
shown(&checked)
);
std::fs::write(scratch.path("claude"), b"pinned claude").unwrap();
let (checked, _) = check_and_serve(&scratch.path("connector.cc"), Some(&credentials));
assert!(
shown(&checked).contains("token_credential"),
"{}",
shown(&checked)
);
assert!(!shown(&checked).contains("credentials_directory"));
}
#[test]
fn usage_says_check_runs_in_the_units_environment() {
let output = run(&["nonsense"], &[]);
assert!(!output.status.success());
assert!(
shown(&output).contains("unit's environment"),
"{}",
shown(&output)
);
}
#[cfg(unix)]

96
crates/model-connector/src/client.rs
View file

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

260
crates/model-connector/src/envelope.rs
View file

@ -17,11 +17,14 @@ use crate::contract::{
ContractError, ENVELOPE_SCHEMA_ID, TransportInvocation, TransportKey, TransportResult,
canonical_invocation_bytes, require_id, wire_decode, wire_encode,
};
use crate::readout::SpendReadout;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
pub enum EnvelopeKind {
Invocation,
Result,
SpendReadoutRequest,
SpendReadoutResult,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
@ -49,8 +52,10 @@ pub(crate) fn binary_schema(_: &mut schemars::SchemaGenerator) -> schemars::Sche
/// The CultNet service every frame of this connector addresses.
pub const CONNECTOR_SERVICE_ID: &str = "gamecult.model.connector";
/// The one operation of that service: run one sealed invocation.
/// The operation that runs one sealed invocation.
pub const INVOKE_OPERATION: &str = "invoke";
/// The operation that reads the spend ledger and the plan's quota.
pub const SPEND_READOUT_OPERATION: &str = "spend_readout";
const PAYLOAD_ENCODING: &str = "messagepack-base64";
const OK_STATUS: &str = "ok";
@ -69,16 +74,39 @@ fn seal_payload(envelope: &TransportEnvelope) -> Result<String, ServiceError> {
Ok(STANDARD.encode(wire_encode(envelope)?))
}
/// The CultNet frame an invocation travels in: a `cultnet.operation_request.v0`
/// for [`CONNECTOR_SERVICE_ID`] and [`INVOKE_OPERATION`], its `message_id` the
/// The two operations of the connector service.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum Operation {
Invoke,
SpendReadout,
}
impl Operation {
fn name(self) -> &'static str {
match self {
Operation::Invoke => INVOKE_OPERATION,
Operation::SpendReadout => SPEND_READOUT_OPERATION,
}
}
fn from_name(name: &str) -> Option<Self> {
[Operation::Invoke, Operation::SpendReadout]
.into_iter()
.find(|operation| operation.name() == name)
}
}
/// The CultNet frame a request travels in: a `cultnet.operation_request.v0`
/// for [`CONNECTOR_SERVICE_ID`] and the named operation, its `message_id` the
/// request id and its payload the base64 of the named sealed envelope.
pub(crate) fn encode_invocation_frame(
fn request_frame(
operation: Operation,
envelope: &TransportEnvelope,
) -> Result<Vec<u8>, ServiceError> {
carriage_bytes(&CultNetMessage::OperationRequest {
message_id: envelope.request_id.clone(),
service_id: CONNECTOR_SERVICE_ID.to_string(),
operation: INVOKE_OPERATION.to_string(),
operation: operation.name().to_string(),
payload_schema: ENVELOPE_SCHEMA_ID.to_string(),
payload_encoding: PAYLOAD_ENCODING.to_string(),
payload: seal_payload(envelope)?,
@ -87,14 +115,17 @@ pub(crate) fn encode_invocation_frame(
})
}
/// The CultNet frame a result travels in: a `cultnet.operation_response.v0`
/// with status `ok`, otherwise shaped like the invocation's.
/// The CultNet frame a response travels in: a `cultnet.operation_response.v0`
/// with status `ok`, otherwise shaped like the request's.
#[cfg(any(feature = "daemon", test))]
pub(crate) fn encode_result_frame(envelope: &TransportEnvelope) -> Result<Vec<u8>, ServiceError> {
fn response_frame(
operation: Operation,
envelope: &TransportEnvelope,
) -> Result<Vec<u8>, ServiceError> {
carriage_bytes(&CultNetMessage::OperationResponse {
message_id: envelope.request_id.clone(),
service_id: CONNECTOR_SERVICE_ID.to_string(),
operation: INVOKE_OPERATION.to_string(),
operation: operation.name().to_string(),
status: OK_STATUS.to_string(),
payload_schema: ENVELOPE_SCHEMA_ID.to_string(),
payload_encoding: PAYLOAD_ENCODING.to_string(),
@ -104,20 +135,43 @@ pub(crate) fn encode_result_frame(envelope: &TransportEnvelope) -> Result<Vec<u8
})
}
pub(crate) fn encode_invocation_frame(
envelope: &TransportEnvelope,
) -> Result<Vec<u8>, ServiceError> {
request_frame(Operation::Invoke, envelope)
}
pub(crate) fn encode_spend_readout_request_frame(
envelope: &TransportEnvelope,
) -> Result<Vec<u8>, ServiceError> {
request_frame(Operation::SpendReadout, envelope)
}
#[cfg(any(feature = "daemon", test))]
pub(crate) fn encode_result_frame(envelope: &TransportEnvelope) -> Result<Vec<u8>, ServiceError> {
response_frame(Operation::Invoke, envelope)
}
#[cfg(feature = "daemon")]
pub(crate) fn encode_spend_readout_response_frame(
envelope: &TransportEnvelope,
) -> Result<Vec<u8>, ServiceError> {
response_frame(Operation::SpendReadout, envelope)
}
/// Opens the carriage of one frame. Everything the carriage says is checked
/// here, before any decryption: the service, the operation, the payload schema,
/// and that the message id is the sealed envelope's request id.
/// here, before any decryption: the service, that the operation is one this
/// service has, the payload schema, and that the message id is the sealed
/// envelope's request id.
fn open_carriage(
message_id: &str,
service_id: &str,
operation: &str,
payload_schema: &str,
payload: &str,
) -> Result<TransportEnvelope, ServiceError> {
if service_id != CONNECTOR_SERVICE_ID
|| operation != INVOKE_OPERATION
|| payload_schema != ENVELOPE_SCHEMA_ID
{
) -> Result<(Operation, TransportEnvelope), ServiceError> {
let operation = Operation::from_name(operation).ok_or(ServiceError::Carriage)?;
if service_id != CONNECTOR_SERVICE_ID || payload_schema != ENVELOPE_SCHEMA_ID {
return Err(ServiceError::Carriage);
}
let bytes = STANDARD
@ -127,7 +181,7 @@ fn open_carriage(
if envelope.request_id != message_id {
return Err(ServiceError::OuterIdentity);
}
Ok(envelope)
Ok((operation, envelope))
}
fn read_carriage(bytes: &[u8]) -> Result<CultNetMessage, ServiceError> {
@ -135,8 +189,15 @@ fn read_carriage(bytes: &[u8]) -> Result<CultNetMessage, ServiceError> {
.map_err(|_| ServiceError::Carriage)
}
/// A request frame, by the operation it names.
#[cfg(any(feature = "daemon", test))]
pub(crate) fn decode_invocation_frame(bytes: &[u8]) -> Result<TransportEnvelope, ServiceError> {
pub(crate) enum RequestFrame {
Invoke(TransportEnvelope),
SpendReadout(TransportEnvelope),
}
#[cfg(any(feature = "daemon", test))]
pub(crate) fn decode_request_frame(bytes: &[u8]) -> Result<RequestFrame, ServiceError> {
match read_carriage(bytes)? {
CultNetMessage::OperationRequest {
message_id,
@ -145,18 +206,34 @@ pub(crate) fn decode_invocation_frame(bytes: &[u8]) -> Result<TransportEnvelope,
payload_schema,
payload,
..
} => open_carriage(
&message_id,
&service_id,
&operation,
&payload_schema,
&payload,
),
} => {
match open_carriage(
&message_id,
&service_id,
&operation,
&payload_schema,
&payload,
)? {
(Operation::Invoke, envelope) => Ok(RequestFrame::Invoke(envelope)),
(Operation::SpendReadout, envelope) => Ok(RequestFrame::SpendReadout(envelope)),
}
}
_ => Err(ServiceError::Carriage),
}
}
pub(crate) fn decode_result_frame(bytes: &[u8]) -> Result<TransportEnvelope, ServiceError> {
#[cfg(test)]
pub(crate) fn decode_invocation_frame(bytes: &[u8]) -> Result<TransportEnvelope, ServiceError> {
match decode_request_frame(bytes)? {
RequestFrame::Invoke(envelope) => Ok(envelope),
RequestFrame::SpendReadout(_) => Err(ServiceError::Carriage),
}
}
fn decode_response_frame(
bytes: &[u8],
expected: Operation,
) -> Result<TransportEnvelope, ServiceError> {
match read_carriage(bytes)? {
CultNetMessage::OperationResponse {
message_id,
@ -166,17 +243,32 @@ pub(crate) fn decode_result_frame(bytes: &[u8]) -> Result<TransportEnvelope, Ser
payload_schema,
payload,
..
} if status == OK_STATUS => open_carriage(
&message_id,
&service_id,
&operation,
&payload_schema,
&payload,
),
} if status == OK_STATUS => {
match open_carriage(
&message_id,
&service_id,
&operation,
&payload_schema,
&payload,
)? {
(found, envelope) if found == expected => Ok(envelope),
_ => Err(ServiceError::Carriage),
}
}
_ => Err(ServiceError::Carriage),
}
}
pub(crate) fn decode_result_frame(bytes: &[u8]) -> Result<TransportEnvelope, ServiceError> {
decode_response_frame(bytes, Operation::Invoke)
}
pub(crate) fn decode_spend_readout_response_frame(
bytes: &[u8],
) -> Result<TransportEnvelope, ServiceError> {
decode_response_frame(bytes, Operation::SpendReadout)
}
/// The most bytes one frame holding a sealed envelope of `ciphertext_len`
/// ciphertext bytes can occupy, with ids at their longest: the frame a caller
/// payload bound (`CallerAdmission`'s `max_payload_bytes`, a bound on
@ -324,6 +416,59 @@ pub fn decrypt_result(
Ok(result)
}
/// A spend readout request carries no content: the sealed envelope's
/// authenticated identity and kind are the whole request.
pub(crate) fn encrypt_spend_readout_request(
caller_runtime_id: &str,
request_id: &str,
security: &TransportKey,
) -> Result<TransportEnvelope, ServiceError> {
encrypt_message(
caller_runtime_id,
request_id,
EnvelopeKind::SpendReadoutRequest,
&[],
security,
)
}
#[cfg(any(feature = "daemon", test))]
pub(crate) fn decrypt_spend_readout_request(
envelope: &TransportEnvelope,
security: &TransportKey,
) -> Result<(), ServiceError> {
validate_envelope(envelope, EnvelopeKind::SpendReadoutRequest)?;
decrypt_message(envelope, security).map(drop)
}
#[cfg(feature = "daemon")]
pub(crate) fn encrypt_spend_readout(
caller_runtime_id: &str,
request_id: &str,
readout: &SpendReadout,
security: &TransportKey,
) -> Result<TransportEnvelope, ServiceError> {
let plaintext = wire_encode(readout).map_err(|_| ServiceError::Encoding)?;
encrypt_message(
caller_runtime_id,
request_id,
EnvelopeKind::SpendReadoutResult,
&plaintext,
security,
)
}
pub(crate) fn decrypt_spend_readout(
envelope: &TransportEnvelope,
security: &TransportKey,
) -> Result<SpendReadout, ServiceError> {
validate_envelope(envelope, EnvelopeKind::SpendReadoutResult)?;
let plaintext = decrypt_message(envelope, security)?;
let readout: SpendReadout = wire_decode(&plaintext).map_err(|_| ServiceError::Encoding)?;
readout.validate()?;
Ok(readout)
}
#[derive(Debug)]
pub(crate) enum TransportFrameError {
Connection(std::io::Error),
@ -1138,4 +1283,53 @@ mod tests {
);
}
}
#[test]
fn a_readout_request_is_not_an_invocation_and_names_its_own_operation() {
let key = TransportKey::from_connection_secret("readout-test-key").unwrap();
let sealed = encrypt_spend_readout_request("caller-1", "readout-1", &key).unwrap();
let frame = encode_spend_readout_request_frame(&sealed).unwrap();
assert!(matches!(
decode_request_frame(&frame),
Ok(RequestFrame::SpendReadout(found)) if found == sealed
));
assert_eq!(decode_invocation_frame(&frame), Err(ServiceError::Carriage));
// Sealed as one kind, the authenticated data refuses it as another.
assert!(decrypt_spend_readout_request(&sealed, &key).is_ok());
assert_eq!(
decrypt_spend_readout(&sealed, &key).unwrap_err(),
ServiceError::Envelope
);
let mut relabelled = sealed.clone();
relabelled.message_kind = EnvelopeKind::SpendReadoutResult;
assert_eq!(
decrypt_spend_readout(&relabelled, &key).unwrap_err(),
ServiceError::Encryption
);
}
#[cfg(feature = "daemon")]
#[test]
fn a_response_is_read_as_its_own_operation_only() {
use crate::readout::PlanQuotaSnapshot;
let key = TransportKey::from_connection_secret("readout-test-key").unwrap();
let readout = SpendReadout::new(1, Vec::new(), PlanQuotaSnapshot::default());
let sealed = encrypt_spend_readout("caller-1", "readout-1", &readout, &key).unwrap();
let readout_frame = encode_spend_readout_response_frame(&sealed).unwrap();
let invoke_frame = encode_result_frame(&sealed).unwrap();
assert_eq!(
decode_spend_readout_response_frame(&readout_frame).unwrap(),
sealed
);
assert_eq!(decode_result_frame(&invoke_frame).unwrap(), sealed);
assert_eq!(
decode_result_frame(&readout_frame),
Err(ServiceError::Carriage)
);
assert_eq!(
decode_spend_readout_response_frame(&invoke_frame),
Err(ServiceError::Carriage)
);
assert_eq!(decrypt_spend_readout(&sealed, &key).unwrap(), readout);
}
}

2
crates/model-connector/src/lib.rs
View file

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

128
crates/model-connector/src/readout.rs Normal file
View file

@ -0,0 +1,128 @@
//! The spend readout: what each caller has spent and holds in its current
//! window, and the plan's own usage as the backend last saw it. The ledger is
//! the replay store's; a readout only reads it.
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use crate::contract::ContractError;
pub const SPEND_READOUT_SCHEMA_ID: &str = "gamecult.model.spend_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>,
}
/// One caller's current spend window in micro-USD. `spent` is settled charges
/// plus what a crash left counted at its full reservation; `reserved` is what
/// calls in flight hold, and both count against the ceiling.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct CallerSpend {
pub caller_id: String,
pub window_ms: u64,
pub ceiling_micro_usd: u64,
pub spent_micro_usd: u64,
pub reserved_micro_usd: u64,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct SpendReadout {
pub schema_id: String,
pub at_unix_ms: u64,
pub callers: Vec<CallerSpend>,
pub plan_quota: Vec<PlanQuota>,
}
impl SpendReadout {
pub fn new(at_unix_ms: u64, callers: Vec<CallerSpend>, plan: PlanQuotaSnapshot) -> Self {
Self {
schema_id: SPEND_READOUT_SCHEMA_ID.to_string(),
at_unix_ms,
callers,
plan_quota: plan.windows,
}
}
pub(crate) fn validate(&self) -> Result<(), ContractError> {
if self.schema_id != SPEND_READOUT_SCHEMA_ID {
return Err(ContractError::Schema);
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::contract::{wire_decode, wire_encode};
fn readout() -> SpendReadout {
SpendReadout::new(
1_000,
vec![CallerSpend {
caller_id: "eureka".into(),
window_ms: 3_600_000,
ceiling_micro_usd: 50_000_000,
spent_micro_usd: 7,
reserved_micro_usd: 3,
}],
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", "callers", "plan_quota"]);
assert_eq!(wire_decode::<SpendReadout>(&bytes).unwrap(), readout());
}
#[test]
fn an_unobserved_plan_is_none_not_zero() {
let bare = SpendReadout::new(1, Vec::new(), PlanQuotaSnapshot::default());
let decoded: SpendReadout = 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.spend_readout.v2".into();
assert!(other.validate().is_err());
assert!(readout().validate().is_ok());
}
}

109
crates/model-connector/src/serve.rs
View file

@ -12,9 +12,11 @@ use zeroize::Zeroizing;
use crate::contract::TransportResult;
use crate::envelope::{
ServiceError, TransportFrameError, decode_invocation_frame, encode_result_frame,
frame_bytes_for_payload, read_transport_frame, write_transport_frame,
RequestFrame, ServiceError, TransportFrameError, decode_request_frame, encode_result_frame,
encode_spend_readout_response_frame, frame_bytes_for_payload, read_transport_frame,
write_transport_frame,
};
use crate::readout::PlanQuotaSnapshot;
use crate::service::{
CallerAdmission, PriceTable, ProviderCall, SpendCeiling, SpendWindowMs, TransportAdmission,
TransportService,
@ -32,6 +34,11 @@ pub trait ProviderBackend: Send + Sync {
/// `serve` on: it names the provider on every receipt and decides which
/// reasoning state is foreign.
fn provider(&self) -> &str;
/// What the provider's own plan allows and has used, for the spend
/// readout. A provider with no plan to report keeps the default.
fn plan_quota(&self) -> PlanQuotaSnapshot {
PlanQuotaSnapshot::default()
}
/// Executes one admitted call. The service has already bound it to its
/// caller and recorded its replay claim. The call cannot show a foreign
/// reasoning item and builds the result's receipt itself.
@ -232,17 +239,28 @@ fn serve_connection(
.and_then(|()| stream.set_write_timeout(Some(timeout)))
.map_err(ServeError::Connection)?;
let request = read_transport_frame(&mut stream, limits.max_frame_bytes).map_err(frame_error)?;
let response = execute_inference_frame(&request, service, backend, limits)?;
let response = answer_frame(&request, service, backend, limits)?;
write_transport_frame(&mut stream, &response, limits.max_frame_bytes).map_err(frame_error)
}
fn execute_inference_frame(
fn answer_frame(
request: &[u8],
service: &Mutex<TransportService>,
backend: &dyn ProviderBackend,
limits: &ServeLimits,
) -> Result<Vec<u8>, ServeError> {
let envelope = decode_invocation_frame(request).map_err(|_| ServeError::FrameEncoding)?;
let envelope = match decode_request_frame(request).map_err(|_| ServeError::FrameEncoding)? {
RequestFrame::Invoke(envelope) => envelope,
RequestFrame::SpendReadout(envelope) => {
let plan = backend.plan_quota();
let response = service
.lock()
.map_err(|_| ServeError::ServicePoisoned)?
.spend_readout(&envelope, unix_ms()?, plan, limits.max_frame_bytes)?;
return encode_spend_readout_response_frame(&response)
.map_err(|_| ServeError::FrameEncoding);
}
};
let admission = service
.lock()
.map_err(|_| ServeError::ServicePoisoned)?
@ -495,6 +513,87 @@ mod tests {
.collect()
}
/// A counting backend that reports a plan, as a provider with one would.
struct PlanBackend {
inner: Arc<CountingBackend>,
plan: PlanQuotaSnapshot,
}
impl ProviderBackend for PlanBackend {
fn ready(&self) -> Result<(), &'static str> {
self.inner.ready()
}
fn provider(&self) -> &str {
self.inner.provider()
}
fn plan_quota(&self) -> PlanQuotaSnapshot {
self.plan.clone()
}
fn execute(&self, call: &ProviderCall) -> TransportResult {
self.inner.execute(call)
}
}
#[test]
fn serve_reads_out_spend_and_plan_quota_without_executing_anything() {
let store = TestStore::new();
let inner = CountingBackend::new(true);
let plan = PlanQuotaSnapshot {
windows: vec![crate::readout::PlanQuota {
window: "five_hour".into(),
utilization_permille: Some(80),
resets_at_unix_ms: 1_793_307_000_000,
status: "allowed".into(),
observed_at_unix_ms: 1_793_300_000_000,
}],
};
let backend = Arc::new(PlanBackend {
inner: inner.clone(),
plan: plan.clone(),
});
let (client, endpoint) = start_with(backend, &store, vec![caller(CALLER, KEY, 4)]);
client.execute(&fresh("request-1")).unwrap();
assert_eq!(inner.executions.load(Ordering::SeqCst), 1);
let before = now_ms();
let readout = client.spend_readout(CALLER).unwrap();
assert!(readout.at_unix_ms >= before && readout.at_unix_ms <= now_ms());
assert_eq!(readout.plan_quota, plan.windows);
let [row] = readout.callers.as_slice() else {
panic!("one admitted caller");
};
assert_eq!(row.caller_id, CALLER);
assert_eq!((row.ceiling_micro_usd, row.window_ms), (CEILING, WINDOW_MS));
assert!(row.spent_micro_usd > 0, "the settled call was charged");
assert_eq!(row.reserved_micro_usd, 0);
assert_eq!(
inner.executions.load(Ordering::SeqCst),
1,
"a readout runs nothing"
);
// A client without the caller's key is dropped, not answered.
let stranger = ConnectorClient::new(
endpoint,
"some-other-key",
FRAME,
Some(Duration::from_secs(5)),
)
.unwrap();
assert!(stranger.spend_readout(CALLER).is_err());
}
#[test]
fn a_backend_with_no_plan_reads_out_none() {
let store = TestStore::new();
let (client, _) = start(CountingBackend::new(true), &store);
let readout = client.spend_readout(CALLER).unwrap();
assert!(readout.plan_quota.is_empty());
}
#[test]
fn serve_refuses_a_request_over_its_max_cost_without_executing_it() {
let store = TestStore::new();

260
crates/model-connector/src/service.rs
View file

@ -14,8 +14,10 @@ use crate::contract::{
};
use crate::envelope::{
EnvelopeKind, MAX_IDENTITY_BYTES, ServiceError, TransportEnvelope, decrypt_invocation,
encode_result_frame, encrypt_result, validate_envelope,
decrypt_spend_readout_request, encode_result_frame, encode_spend_readout_response_frame,
encrypt_result, encrypt_spend_readout, validate_envelope,
};
use crate::readout::{CallerSpend, PlanQuotaSnapshot, SpendReadout};
/// What a caller may spend in one window, in micro-USD. A distinct type from
/// [`SpendWindowMs`] so the pair cannot be passed swapped:
@ -280,15 +282,15 @@ fn admission_clock(
/// it was admitted in, so nothing older is kept.
fn prune_windows(
window_spend: &mut HashMap<String, BTreeMap<u64, u64>>,
live: &HashMap<RequestIdentity, u64>,
live: &HashMap<RequestIdentity, LiveRequest>,
caller_runtime_id: &str,
floor: u64,
) {
if let Some(windows) = window_spend.get_mut(caller_runtime_id) {
windows.retain(|window, _| {
*window >= floor
|| live.iter().any(|(identity, live_window)| {
live_window == window && identity.caller_runtime_id == caller_runtime_id
|| live.iter().any(|(identity, request)| {
request.window == *window && identity.caller_runtime_id == caller_runtime_id
})
});
}
@ -298,7 +300,7 @@ fn prune_windows(
/// places that do not know the clock (open, settle).
fn prune_to_latest(
window_spend: &mut HashMap<String, BTreeMap<u64, u64>>,
live: &HashMap<RequestIdentity, u64>,
live: &HashMap<RequestIdentity, LiveRequest>,
caller_runtime_id: &str,
) {
let latest = window_spend
@ -309,6 +311,14 @@ fn prune_to_latest(
}
}
/// A request between claim and settlement: the window it was admitted in and
/// the reservation it holds there.
#[derive(Debug, Clone, Copy)]
struct LiveRequest {
window: u64,
reserved_micro_usd: u64,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct RequestIdentity {
caller_runtime_id: String,
@ -464,7 +474,7 @@ pub struct TransportService {
callers: HashMap<String, CallerAdmission>,
/// The requests between claim and settlement, each with the window it was
/// admitted in.
live: HashMap<RequestIdentity, u64>,
live: HashMap<RequestIdentity, LiveRequest>,
replay_store: CultCache,
max_expiry_skew_ms: u64,
provider: String,
@ -611,6 +621,70 @@ impl TransportService {
.unwrap_or(0)
}
/// Answers one sealed spend readout request: for each admitted caller, what
/// its current window holds against its ceiling, with the backend's plan
/// quota beside it. Read-only: nothing is reserved, recorded or pruned.
/// Any admitted caller reads every caller's row; the spend is the
/// operator's and the callers are all hers. `max_frame_bytes` bounds the
/// reply like it bounds `complete`'s.
pub fn spend_readout(
&self,
envelope: &TransportEnvelope,
now_unix_ms: u64,
plan: PlanQuotaSnapshot,
max_frame_bytes: usize,
) -> Result<TransportEnvelope, ServiceError> {
validate_envelope(envelope, EnvelopeKind::SpendReadoutRequest)?;
let asker = self
.callers
.get(&envelope.caller_runtime_id)
.ok_or(ServiceError::CallerNotAdmitted)?;
decrypt_spend_readout_request(envelope, &asker.security)?;
let mut callers: Vec<CallerSpend> = self
.callers
.values()
.map(|caller| {
// The window admission would use right now.
let window = admission_clock(
&self.window_spend,
&caller.caller_runtime_id,
now_unix_ms,
caller.spend_window_ms,
) / caller.spend_window_ms;
let held = self.window_spend(&caller.caller_runtime_id, window);
let reserved: u64 = self
.live
.iter()
.filter(|(identity, request)| {
request.window == window
&& identity.caller_runtime_id == caller.caller_runtime_id
})
.fold(0, |sum, (_, request)| {
sum.saturating_add(request.reserved_micro_usd)
});
CallerSpend {
caller_id: caller.caller_runtime_id.clone(),
window_ms: caller.spend_window_ms,
ceiling_micro_usd: caller.spend_ceiling_micro_usd,
spent_micro_usd: held.saturating_sub(reserved),
reserved_micro_usd: reserved,
}
})
.collect();
callers.sort_by(|a, b| a.caller_id.cmp(&b.caller_id));
let readout = SpendReadout::new(now_unix_ms, callers, plan);
let sealed = encrypt_spend_readout(
&envelope.caller_runtime_id,
&envelope.request_id,
&readout,
&asker.security,
)?;
if !readout_fits_frame(&sealed, max_frame_bytes)? {
return Err(ServiceError::FrameBound);
}
Ok(sealed)
}
/// Admits one request. `max_frame_bytes` is the frame every reply must
/// fit, the same bound `complete` is given: a recorded result that no
/// longer fits it (the service reopened at a smaller frame) is answered
@ -794,7 +868,13 @@ impl TransportService {
0,
reserved_micro_usd,
);
self.live.insert(identity, window);
self.live.insert(
identity,
LiveRequest {
window,
reserved_micro_usd,
},
);
Ok(TransportAdmission::Execute(Box::new(ExecutionClaim {
call: ProviderCall {
receipt: TransportReceipt {
@ -1005,6 +1085,14 @@ fn fits_frame(envelope: &TransportEnvelope, max_frame_bytes: usize) -> Result<bo
encode_result_frame(envelope).map(|bytes| bytes.len() <= max_frame_bytes)
}
/// `fits_frame` for the frame a spend readout answer travels in.
fn readout_fits_frame(
envelope: &TransportEnvelope,
max_frame_bytes: usize,
) -> Result<bool, ServiceError> {
encode_spend_readout_response_frame(envelope).map(|bytes| bytes.len() <= max_frame_bytes)
}
/// The provider named on the stand-in failure, fixed so its size is bounded.
const STAND_IN_PROVIDER: &str = "model-connector";
@ -1758,6 +1846,17 @@ mod tests {
assert_eq!(fits_frame(&envelope, size - 1), Ok(false));
}
#[test]
fn a_readout_frame_check_is_inclusive_of_the_serialized_envelope() {
let readout = SpendReadout::new(1, Vec::new(), PlanQuotaSnapshot::default());
let envelope = encrypt_spend_readout(CALLER, "sized", &readout, &security(KEY)).unwrap();
let size = encode_spend_readout_response_frame(&envelope)
.unwrap()
.len();
assert_eq!(readout_fits_frame(&envelope, size), Ok(true));
assert_eq!(readout_fits_frame(&envelope, size - 1), Ok(false));
}
fn claim_for(
service: &mut TransportService,
request_id: &str,
@ -3616,4 +3715,151 @@ mod tests {
<= sized(Refusal::ProviderDigestSubstitution)
);
}
// ---- spend readout ----
fn readout_as(
service: &TransportService,
caller: &str,
key: &str,
now: u64,
frame: usize,
) -> Result<SpendReadout, ServiceError> {
let sealed = encrypt_spend_readout_request(caller, "readout-1", &security(key)).unwrap();
let answer = service.spend_readout(&sealed, now, PlanQuotaSnapshot::default(), frame)?;
assert_eq!(answer.request_id, "readout-1");
Ok(decrypt_spend_readout(&answer, &security(key)).unwrap())
}
fn row<'a>(readout: &'a SpendReadout, caller: &str) -> &'a CallerSpend {
readout
.callers
.iter()
.find(|row| row.caller_id == caller)
.expect("the caller has a row")
}
#[test]
fn spend_readout_reports_window() {
let store = TestStore::new();
let mut service = open_service(&store, 4);
let claim = claim_at(&mut service, "request-1", 1_000);
// In flight: the reservation is held, nothing is spent.
let readout = readout_as(&service, CALLER, KEY, 1_500, FRAME).unwrap();
assert_eq!(readout.at_unix_ms, 1_500);
let mine = row(&readout, CALLER);
assert_eq!(mine.window_ms, WINDOW_MS);
assert_eq!(mine.ceiling_micro_usd, CEILING);
assert_eq!(
(mine.spent_micro_usd, mine.reserved_micro_usd),
(0, reservation())
);
// Settled: the charge replaces the reservation.
let result = text_result(claim.call(), "done");
service.complete(*claim, result, FRAME).unwrap();
let readout = readout_as(&service, CALLER, KEY, 1_500, FRAME).unwrap();
let mine = row(&readout, CALLER);
assert_eq!(
(mine.spent_micro_usd, mine.reserved_micro_usd),
(FIXTURE_USAGE_COST, 0)
);
// Every admitted caller has a row, in id order; the idle one is empty.
let ids: Vec<&str> = readout
.callers
.iter()
.map(|r| r.caller_id.as_str())
.collect();
assert_eq!(ids, ["epiphany-yggdrasil", "ghostlight-yggdrasil"]);
let idle = row(&readout, "ghostlight-yggdrasil");
assert_eq!((idle.spent_micro_usd, idle.reserved_micro_usd), (0, 0));
}
#[test]
fn a_claim_a_crash_left_reads_as_spent_not_reserved() {
let store = TestStore::new();
let mut service = open_service(&store, 4);
let _claim = claim_at(&mut service, "request-1", 1_000);
drop(service);
let service = open_service(&store, 4);
let readout = readout_as(&service, CALLER, KEY, 1_500, FRAME).unwrap();
let mine = row(&readout, CALLER);
assert_eq!(
(mine.spent_micro_usd, mine.reserved_micro_usd),
(reservation(), 0)
);
}
#[test]
fn a_readout_reads_the_ledger_and_changes_nothing_in_it() {
let store = TestStore::new();
let mut service = open_service(&store, 4);
let claim = claim_at(&mut service, "request-1", 1_000);
// Two windows later: the caller's current window is empty, and
// looking must not prune the window the live claim settles into.
let later = 2 * WINDOW_MS + 5;
let readout = readout_as(&service, CALLER, KEY, later, FRAME).unwrap();
let mine = row(&readout, CALLER);
assert_eq!((mine.spent_micro_usd, mine.reserved_micro_usd), (0, 0));
assert_eq!(service.window_spend(CALLER, 0), reservation());
let result = text_result(claim.call(), "done");
service.complete(*claim, result, FRAME).unwrap();
assert_eq!(service.window_spend(CALLER, 0), FIXTURE_USAGE_COST);
}
#[test]
fn a_readout_is_refused_to_a_stranger_a_wrong_key_and_an_invocation() {
let store = TestStore::new();
let service = open_service(&store, 4);
let sealed = |caller: &str, key: &str| {
encrypt_spend_readout_request(caller, "readout-1", &security(key)).unwrap()
};
assert_eq!(
service
.spend_readout(
&sealed("nobody", KEY),
1,
PlanQuotaSnapshot::default(),
FRAME
)
.unwrap_err(),
ServiceError::CallerNotAdmitted
);
assert_eq!(
service
.spend_readout(
&sealed(CALLER, "another-key"),
1,
PlanQuotaSnapshot::default(),
FRAME
)
.unwrap_err(),
ServiceError::Encryption
);
let invocation = invocation_at("request-1", 1_000, MAX_COST);
assert_eq!(
service
.spend_readout(
&raw_envelope(&invocation),
1,
PlanQuotaSnapshot::default(),
FRAME
)
.unwrap_err(),
ServiceError::Envelope
);
}
#[test]
fn a_readout_that_cannot_fit_the_frame_is_a_frame_bound_failure() {
let store = TestStore::new();
let service = open_service(&store, 4);
assert_eq!(
readout_as(&service, CALLER, KEY, 1_000, 64).unwrap_err(),
ServiceError::FrameBound
);
assert!(readout_as(&service, CALLER, KEY, 1_000, FRAME).is_ok());
}
}

3
crates/model-connector/tests/published_schemas.rs
View file

@ -28,6 +28,7 @@ fn published_schemas_are_the_types() {
("transport_result", render::<TransportResult>()),
("transport_receipt", render::<TransportReceipt>()),
("transport_envelope", render::<TransportEnvelope>()),
("spend_readout", render::<SpendReadout>()),
];
let dir = schemas_dir();
if std::env::var_os(REGENERATE).is_some() {
@ -53,7 +54,7 @@ fn published_schemas_are_the_types() {
expected.sort();
assert_eq!(
on_disk, expected,
"schemas/ holds exactly the five documents"
"schemas/ holds exactly the six documents"
);
}

112
schemas/spend_readout.json Normal file
View file

@ -0,0 +1,112 @@
{
"$schema": "https://json-schema.org/draft/2020-12/schema",
"title": "SpendReadout",
"type": "object",
"properties": {
"at_unix_ms": {
"type": "integer",
"format": "uint64",
"minimum": 0
},
"callers": {
"type": "array",
"items": {
"$ref": "#/$defs/CallerSpend"
}
},
"plan_quota": {
"type": "array",
"items": {
"$ref": "#/$defs/PlanQuota"
}
},
"schema_id": {
"type": "string"
}
},
"additionalProperties": false,
"required": [
"schema_id",
"at_unix_ms",
"callers",
"plan_quota"
],
"$defs": {
"CallerSpend": {
"description": "One caller's current spend window in micro-USD. `spent` is settled charges\nplus what a crash left counted at its full reservation; `reserved` is what\ncalls in flight hold, and both count against the ceiling.",
"type": "object",
"properties": {
"caller_id": {
"type": "string"
},
"ceiling_micro_usd": {
"type": "integer",
"format": "uint64",
"minimum": 0
},
"reserved_micro_usd": {
"type": "integer",
"format": "uint64",
"minimum": 0
},
"spent_micro_usd": {
"type": "integer",
"format": "uint64",
"minimum": 0
},
"window_ms": {
"type": "integer",
"format": "uint64",
"minimum": 0
}
},
"additionalProperties": false,
"required": [
"caller_id",
"window_ms",
"ceiling_micro_usd",
"spent_micro_usd",
"reserved_micro_usd"
]
},
"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"
]
}
}
}

4
schemas/transport_envelope.json
View file

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