From 480cf941c8b56ff34cd641fdcccbb72bcdce80c6 Mon Sep 17 00:00:00 2001 From: Eureka Date: Thu, 29 Oct 2026 20:18:35 +0100 Subject: [PATCH] claude-turn-schema-quota: schema-typed Persona replies via --json-schema, token on fd 3 never in the child env, per-window plan-quota capture, communal spend readout, not-logged-in bills zero (squash of eureka-body/claude-turn-schema-quota e867fad..429b344; Soul s1, s2) Co-Authored-By: Claude Opus 5.5 --- README.md | 15 +- crates/claude-connector/src/backend.rs | 76 ++- crates/claude-connector/src/backend/tests.rs | 523 ++++++++++++++++-- crates/claude-connector/src/daemon.rs | 37 +- crates/claude-connector/src/handle.rs | 5 - crates/claude-connector/src/launch.rs | 37 +- crates/claude-connector/src/lib.rs | 2 + crates/claude-connector/src/main.rs | 4 +- crates/claude-connector/src/process.rs | 134 ++++- crates/claude-connector/src/quota.rs | 151 +++++ crates/claude-connector/src/stream.rs | 129 ++++- .../claude-connector/tests/config_document.rs | 121 +++- crates/model-connector/src/client.rs | 96 +++- crates/model-connector/src/envelope.rs | 260 +++++++-- crates/model-connector/src/lib.rs | 2 + crates/model-connector/src/readout.rs | 128 +++++ crates/model-connector/src/serve.rs | 109 +++- crates/model-connector/src/service.rs | 260 ++++++++- .../tests/published_schemas.rs | 3 +- schemas/spend_readout.json | 112 ++++ schemas/transport_envelope.json | 4 +- 21 files changed, 2021 insertions(+), 187 deletions(-) create mode 100644 crates/claude-connector/src/quota.rs create mode 100644 crates/model-connector/src/readout.rs create mode 100644 schemas/spend_readout.json diff --git a/README.md b/README.md index 339e793..3a12abd 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/crates/claude-connector/src/backend.rs b/crates/claude-connector/src/backend.rs index c5a8f81..1963e1d 100644 --- a/crates/claude-connector/src/backend.rs +++ b/crates/claude-connector/src/backend.rs @@ -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>; + /// 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, sink: Box, tokens: Box, + quota: Arc, } 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 { + 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 { - 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(()) } diff --git a/crates/claude-connector/src/backend/tests.rs b/crates/claude-connector/src/backend/tests.rs index 21a8c72..2e3b993 100644 --- a/crates/claude-connector/src/backend/tests.rs +++ b/crates/claude-connector/src/backend/tests.rs @@ -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, + /// 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, + token: &str, ) -> io::Result> { 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, + _: &str, ) -> io::Result> { 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 { + 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 = 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")); +} diff --git a/crates/claude-connector/src/daemon.rs b/crates/claude-connector/src/daemon.rs index 87450a8..674975d 100644 --- a/crates/claude-connector/src/daemon.rs +++ b/crates/claude-connector/src/daemon.rs @@ -68,7 +68,8 @@ pub struct ClaudeConnectorConfig { /// Exactly one caller until claude-handle-caller-binding lands. #[cultcache(key = 17)] pub callers: Vec, - /// 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, } @@ -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---` -/// 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--`, then an optional `-` 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 diff --git a/crates/claude-connector/src/handle.rs b/crates/claude-connector/src/handle.rs index c20fda0..222b74b 100644 --- a/crates/claude-connector/src/handle.rs +++ b/crates/claude-connector/src/handle.rs @@ -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() diff --git a/crates/claude-connector/src/launch.rs b/crates/claude-connector/src/launch.rs index c94134c..b233c09 100644 --- a/crates/claude-connector/src/launch.rs +++ b/crates/claude-connector/src/launch.rs @@ -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)>); +/// `Debug` lists names only. +pub struct ChildEnv(Vec<(&'static str, String)>); impl ChildEnv { pub fn iter(&self) -> impl Iterator { @@ -37,21 +40,17 @@ impl fmt::Debug for ChildEnv { } } -pub(crate) fn child_env( - config: &ClaudeCodeConfig, - request: &ProviderRequest, - token: Zeroizing, -) -> 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(), diff --git a/crates/claude-connector/src/lib.rs b/crates/claude-connector/src/lib.rs index ef3e404..327a078 100644 --- a/crates/claude-connector/src/lib.rs +++ b/crates/claude-connector/src/lib.rs @@ -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; diff --git a/crates/claude-connector/src/main.rs b/crates/claude-connector/src/main.rs index ece39cc..c58e7f5 100644 --- a/crates/claude-connector/src/main.rs +++ b/crates/claude-connector/src/main.rs @@ -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), diff --git a/crates/claude-connector/src/process.rs b/crates/claude-connector/src/process.rs index 450f388..2421e0e 100644 --- a/crates/claude-connector/src/process.rs +++ b/crates/claude-connector/src/process.rs @@ -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, + token: &str, ) -> io::Result>; } @@ -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 { + 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, + token: &str, ) -> io::Result> { + 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()); + } + } } diff --git a/crates/claude-connector/src/quota.rs b/crates/claude-connector/src/quota.rs new file mode 100644 index 0000000..76f0ce4 --- /dev/null +++ b/crates/claude-connector/src/quota.rs @@ -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>); + +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) { + self.observe_at(unix_ms(), windows); + } + + fn observe_at(&self, now: u64, windows: impl IntoIterator) { + 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"); + } +} diff --git a/crates/claude-connector/src/stream.rs b/crates/claude-connector/src/stream.rs index e99935f..f712892 100644 --- a/crates/claude-connector/src/stream.rs +++ b/crates/claude-connector/src/stream.rs @@ -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, model: Option, message: Option, - rate_limit_info: Option, + /// 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, is_error: Option, stop_reason: Option, api_error_status: Option, + /// Milliseconds spent in API calls; zero means no call was made. + duration_api_ms: Option, result: Option, errors: Option, + /// The answer to a `--json-schema` request, an object. + structured_output: Option, /// 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, @@ -85,11 +97,7 @@ struct Block { #[serde(rename = "type")] kind: String, text: Option, -} - -#[derive(Deserialize)] -struct RateLimitInfo { - status: Option, + name: Option, } /// 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, texts: Vec, message_id: Option, @@ -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::(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 { + 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, door_refusal: bool) -> Option { - let models: BTreeMap = 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, empty_is_zero: bool, absent_is_zero: bool) -> Option { + let Some(reported) = reported else { + return absent_is_zero.then(|| sum_usage(std::iter::empty())); + }; + let models: BTreeMap = 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())) } diff --git a/crates/claude-connector/tests/config_document.rs b/crates/claude-connector/tests/config_document.rs index 4bc35d4..141f471 100644 --- a/crates/claude-connector/tests/config_document.rs +++ b/crates/claude-connector/tests/config_document.rs @@ -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 { + 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)] diff --git a/crates/model-connector/src/client.rs b/crates/model-connector/src/client.rs index 1421a5b..47b6c94 100644 --- a/crates/model-connector/src/client.rs +++ b/crates/model-connector/src/client.rs @@ -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 { + 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::() + ); + 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, 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(); + } + } } diff --git a/crates/model-connector/src/envelope.rs b/crates/model-connector/src/envelope.rs index 40cfd09..70d8ed2 100644 --- a/crates/model-connector/src/envelope.rs +++ b/crates/model-connector/src/envelope.rs @@ -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 { 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 { + [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, 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, ServiceError> { +fn response_frame( + operation: Operation, + envelope: &TransportEnvelope, +) -> Result, 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 Result, ServiceError> { + request_frame(Operation::Invoke, envelope) +} + +pub(crate) fn encode_spend_readout_request_frame( + envelope: &TransportEnvelope, +) -> Result, ServiceError> { + request_frame(Operation::SpendReadout, envelope) +} + +#[cfg(any(feature = "daemon", test))] +pub(crate) fn encode_result_frame(envelope: &TransportEnvelope) -> Result, ServiceError> { + response_frame(Operation::Invoke, envelope) +} + +#[cfg(feature = "daemon")] +pub(crate) fn encode_spend_readout_response_frame( + envelope: &TransportEnvelope, +) -> Result, 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 { - 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 { @@ -135,8 +189,15 @@ fn read_carriage(bytes: &[u8]) -> Result { .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 { +pub(crate) enum RequestFrame { + Invoke(TransportEnvelope), + SpendReadout(TransportEnvelope), +} + +#[cfg(any(feature = "daemon", test))] +pub(crate) fn decode_request_frame(bytes: &[u8]) -> Result { match read_carriage(bytes)? { CultNetMessage::OperationRequest { message_id, @@ -145,18 +206,34 @@ pub(crate) fn decode_invocation_frame(bytes: &[u8]) -> Result 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 { +#[cfg(test)] +pub(crate) fn decode_invocation_frame(bytes: &[u8]) -> Result { + match decode_request_frame(bytes)? { + RequestFrame::Invoke(envelope) => Ok(envelope), + RequestFrame::SpendReadout(_) => Err(ServiceError::Carriage), + } +} + +fn decode_response_frame( + bytes: &[u8], + expected: Operation, +) -> Result { match read_carriage(bytes)? { CultNetMessage::OperationResponse { message_id, @@ -166,17 +243,32 @@ pub(crate) fn decode_result_frame(bytes: &[u8]) -> Result 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 { + decode_response_frame(bytes, Operation::Invoke) +} + +pub(crate) fn decode_spend_readout_response_frame( + bytes: &[u8], +) -> Result { + 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 { + 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 { + 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 { + 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); + } } diff --git a/crates/model-connector/src/lib.rs b/crates/model-connector/src/lib.rs index ba1642b..c55202a 100644 --- a/crates/model-connector/src/lib.rs +++ b/crates/model-connector/src/lib.rs @@ -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::*; diff --git a/crates/model-connector/src/readout.rs b/crates/model-connector/src/readout.rs new file mode 100644 index 0000000..36d4cb9 --- /dev/null +++ b/crates/model-connector/src/readout.rs @@ -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, + 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, +} + +/// 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, + pub plan_quota: Vec, +} + +impl SpendReadout { + pub fn new(at_unix_ms: u64, callers: Vec, 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::(&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()); + } +} diff --git a/crates/model-connector/src/serve.rs b/crates/model-connector/src/serve.rs index d6723a4..26f5e2d 100644 --- a/crates/model-connector/src/serve.rs +++ b/crates/model-connector/src/serve.rs @@ -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, backend: &dyn ProviderBackend, limits: &ServeLimits, ) -> Result, 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, + 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(); diff --git a/crates/model-connector/src/service.rs b/crates/model-connector/src/service.rs index 6d10cb9..c269916 100644 --- a/crates/model-connector/src/service.rs +++ b/crates/model-connector/src/service.rs @@ -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>, - live: &HashMap, + live: &HashMap, 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>, - live: &HashMap, + live: &HashMap, 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, /// The requests between claim and settlement, each with the window it was /// admitted in. - live: HashMap, + live: HashMap, 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 { + 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 = 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 Result { + 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 { + 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()); + } } diff --git a/crates/model-connector/tests/published_schemas.rs b/crates/model-connector/tests/published_schemas.rs index 15ce197..cb21b01 100644 --- a/crates/model-connector/tests/published_schemas.rs +++ b/crates/model-connector/tests/published_schemas.rs @@ -28,6 +28,7 @@ fn published_schemas_are_the_types() { ("transport_result", render::()), ("transport_receipt", render::()), ("transport_envelope", render::()), + ("spend_readout", render::()), ]; 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" ); } diff --git a/schemas/spend_readout.json b/schemas/spend_readout.json new file mode 100644 index 0000000..5a5d23d --- /dev/null +++ b/schemas/spend_readout.json @@ -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" + ] + } + } +} diff --git a/schemas/transport_envelope.json b/schemas/transport_envelope.json index 69a97d3..8b062e3 100644 --- a/schemas/transport_envelope.json +++ b/schemas/transport_envelope.json @@ -46,7 +46,9 @@ "type": "string", "enum": [ "Invocation", - "Result" + "Result", + "SpendReadoutRequest", + "SpendReadoutResult" ] } }