From 0762370219795ab2fbcce3f1509ee99c51c776f5 Mon Sep 17 00:00:00 2001 From: Meta Vens Date: Thu, 29 Oct 2026 18:58:37 +0100 Subject: [PATCH] claude-connector: output_schema_json lowers to --json-schema and the structured output is the one AssistantText; token on a pipe fd, never in the child env; QuotaCell keeps the plan windows from rate_limit_event; check polish (loose model-id form with the price table as the known list, token checked last and unset CREDENTIALS_DIRECTORY named, usage names the unit's environment, deadline on verbs in tests) Co-Authored-By: Claude Sonnet 5.5 --- README.md | 15 +- crates/claude-connector/src/backend.rs | 54 ++- crates/claude-connector/src/backend/tests.rs | 366 ++++++++++++++++-- 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 | 46 ++- crates/claude-connector/src/quota.rs | 107 +++++ crates/claude-connector/src/stream.rs | 85 +++- .../claude-connector/tests/config_document.rs | 121 +++++- 12 files changed, 764 insertions(+), 115 deletions(-) create mode 100644 crates/claude-connector/src/quota.rs 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..e10d6be 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.is_some(), + &self.quota, + ); child.kill(); child.wait(); let tail = child.stderr_tail(); @@ -231,14 +261,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, + structured: bool, + quota: &QuotaCell, ) -> Result { - let mut state = StreamState::new(model); + let mut state = StreamState::new(model, structured, quota); let finished = loop { match child.read_line(deadline) { StdoutRead::Line(line) => { @@ -295,12 +331,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..da4fea4 100644 --- a/crates/claude-connector/src/backend/tests.rs +++ b/crates/claude-connector/src/backend/tests.rs @@ -13,9 +13,12 @@ use model_connector::{ 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":1793307000,"rateLimitType":"five_hour","overageStatus":"rejected","isUsingOverage":false,"unifiedWindows":{"five_hour":{"utilization":0.08,"resetsAt":1793307000,"status":"allowed"},"seven_day":{"utilization":0.3,"resetsAt":1793656800,"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,179 @@ 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 rig = rig(vec![structured_script()]); + let (kind, message, _) = failure_of(&rig.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 rig = rig(vec![Script::lines(vec![ + init_line(MODEL, SESSION), + structured_call_line("Bash"), + structured_result_line(), + ])]); + let (kind, message, _) = failure_of(&rig.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 rig = rig(vec![Script::lines(vec![ + init_line(MODEL, SESSION), + text_line("msg_1", "hi"), + result_line("tool_use", USAGE), + ])]); + let (kind, message, _) = failure_of(&rig.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() +} + +fn window(name: &str, permille: u32, resets_at_unix_ms: u64) -> model_connector::PlanQuota { + model_connector::PlanQuota { + window: name.to_string(), + utilization_permille: permille, + resets_at_unix_ms, + status: "allowed".to_string(), + } +} + +#[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!( + snapshot.windows, + vec![ + window("five_hour", 80, 1_793_307_000_000), + window("seven_day", 300, 1_793_656_800_000), + ] + ); + assert!(snapshot.observed_at_unix_ms.unwrap() >= before); + // 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":1793307000,"status":"rejected"},"seven_day":{"resetsAt":1793656800,"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, 1_793_307_000_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. + assert_eq!(quota_of(&rig).windows, vec![rejected]); +} + +#[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":1793310000,"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!( + quota_of(&rig).windows, + vec![ + window("five_hour", 500, 1_793_310_000_000), + window("seven_day", 300, 1_793_656_800_000), + ] + ); +} + #[cfg(unix)] mod real { use super::*; @@ -1892,11 +2144,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 +2154,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 +2214,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 +2229,55 @@ 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_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 huge = "t".repeat(1 << 20); + let refused = port.spawn(&argv, &env, fake.dir.path(), Vec::new(), &huge); + assert!(refused.is_err()); + } + #[test] fn the_output_caps_hold_at_their_exact_sizes() { let lines = good_lines(); 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..1b35fc4 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,31 @@ 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; + + /// Puts the pipe's read end at [`TOKEN_FD`] in the child. Runs between + /// fork and exec, so it calls nothing but `dup2` and `fcntl`. + fn place_token_fd(source: RawFd) -> io::Result<()> { + // SAFETY: dup2 and fcntl take integers only and are async-signal-safe. + unsafe { + if source == TOKEN_FD { + // dup2 onto itself changes nothing, so the close-on-exec flag + // the pipe was created with must be cleared by hand. + let flags = libc::fcntl(TOKEN_FD, libc::F_GETFD); + if flags < 0 || libc::fcntl(TOKEN_FD, libc::F_SETFD, flags & !libc::FD_CLOEXEC) < 0 + { + return Err(io::Error::last_os_error()); + } + } else if libc::dup2(source, TOKEN_FD) < 0 { + return Err(io::Error::last_os_error()); + } + } + Ok(()) + } /// Spawns each child in its own process group, so that killing the call /// kills what the child started too. @@ -78,7 +106,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 +124,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_token_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_token_fd(source)) }; let mut child = command.spawn()?; + drop(reader); let (stdin_pipe, stdout, stderr) = (child.stdin.take(), child.stdout.take(), child.stderr.take()); diff --git a/crates/claude-connector/src/quota.rs b/crates/claude-connector/src/quota.rs new file mode 100644 index 0000000..3c02d2c --- /dev/null +++ b/crates/claude-connector/src/quota.rs @@ -0,0 +1,107 @@ +//! 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 one seen 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. +#[derive(Default)] +pub struct QuotaCell(Mutex); + +#[derive(Default)] +struct Observed { + windows: BTreeMap, + observed_at_unix_ms: Option, +} + +impl QuotaCell { + pub fn new() -> Self { + Self::default() + } + + /// Records the windows one event reported: each replaces that window's + /// earlier observation, and windows the event did not mention stay. An + /// event that reported none changes nothing, not even the time. + pub fn observe(&self, windows: impl IntoIterator) { + let mut held = self.0.lock().unwrap_or_else(|e| e.into_inner()); + let mut any = false; + for quota in windows { + held.windows.insert(quota.window.clone(), quota); + any = true; + } + if any { + held.observed_at_unix_ms = Some(unix_ms()); + } + } + + pub fn snapshot(&self) -> PlanQuotaSnapshot { + let held = self.0.lock().unwrap_or_else(|e| e.into_inner()); + PlanQuotaSnapshot { + windows: held.windows.values().cloned().collect(), + observed_at_unix_ms: held.observed_at_unix_ms, + } + } +} + +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::*; + + fn quota(window: &str, permille: u32) -> PlanQuota { + PlanQuota { + window: window.to_string(), + utilization_permille: permille, + resets_at_unix_ms: 1_793_307_000_000, + status: "allowed".to_string(), + } + } + + #[test] + fn an_empty_cell_has_observed_nothing() { + let snapshot = QuotaCell::new().snapshot(); + assert!(snapshot.windows.is_empty()); + assert_eq!(snapshot.observed_at_unix_ms, None); + } + + #[test] + fn a_window_keeps_only_its_latest_observation_and_others_stay() { + let cell = QuotaCell::new(); + let before = unix_ms(); + cell.observe([quota("five_hour", 80), quota("seven_day", 300)]); + cell.observe([quota("five_hour", 120)]); + let after = unix_ms(); + let snapshot = cell.snapshot(); + assert_eq!( + snapshot.windows, + vec![quota("five_hour", 120), quota("seven_day", 300)] + ); + let at = snapshot.observed_at_unix_ms.expect("stamped"); + assert!((before..=after).contains(&at)); + } + + #[test] + fn an_event_with_no_windows_does_not_move_the_stamp() { + let cell = QuotaCell::new(); + cell.observe(Vec::new()); + assert_eq!(cell.snapshot().observed_at_unix_ms, None); + cell.observe([quota("five_hour", 80)]); + let first = cell.snapshot().observed_at_unix_ms; + std::thread::sleep(std::time::Duration::from_millis(3)); + cell.observe(Vec::new()); + assert_eq!(cell.snapshot().observed_at_unix_ms, first); + } +} diff --git a/crates/claude-connector/src/stream.rs b/crates/claude-connector/src/stream.rs index e99935f..a75a37f 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,16 @@ 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, 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 +95,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 +113,10 @@ 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, + quota: &'a QuotaCell, session_id: Option, texts: Vec, message_id: Option, @@ -114,9 +124,11 @@ pub(crate) struct StreamState<'a> { } impl<'a> StreamState<'a> { - pub(crate) fn new(model: &'a str) -> Self { + pub(crate) fn new(model: &'a str, structured: bool, quota: &'a QuotaCell) -> Self { Self { model, + structured, + quota, session_id: None, texts: Vec::new(), message_id: None, @@ -151,10 +163,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 +189,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, @@ -210,8 +226,23 @@ 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_null()) else { + return Err(Failure { + kind: FailureKind::Provider, + message: "claude-code: a schema was sent and the result carried no structured output", + 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 +264,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,6 +294,34 @@ 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: (utilization * 1000.0).round() as u32, + resets_at_unix_ms: resets_at.saturating_mul(1000), + status: status.to_string(), + }) + }) + .collect() +} + fn fail(kind: FailureKind, message: &'static str) -> Step { Step::Done(Box::new(Err(Failure::cost_unknown(kind, message)))) } 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)]