wip: runtime audit slice (uncompiled; to be split)

This commit is contained in:
CodeWhale Bot committed 2026-09-30 00:49:45 -07:00
1 parent 8b6e87bc35
commit 6633405673
25 files changed
+1561 -395

No files matched your search

+65 -12
View File
@@ -1,10 +1,22 @@
//! Compiled-regex cache for patterns a user or model supplies (tool search,
//! purge). Those patterns are untrusted input, so compilation is bounded:
//! the pattern length is capped, the compiled program and lazy DFA have
//! explicit size limits, compilation runs outside the cache lock, and a
//! rejected pattern is remembered so repeating it costs a lookup.
use std::num::NonZeroUsize;
use std::sync::{Mutex, OnceLock};
use lru::LruCache;
use regex::Regex;
use regex::{Regex, RegexBuilder};
const DEFAULT_USER_REGEX_CACHE_CAPACITY: usize = 64;
/// Longest pattern accepted, in bytes. Search and purge patterns are short.
pub const MAX_USER_REGEX_PATTERN_BYTES: usize = 4 * 1024;
/// Compiled program limit (the `regex` default is 10 MiB).
const USER_REGEX_SIZE_LIMIT: usize = 1024 * 1024;
/// Lazy DFA cache limit per matcher.
const USER_REGEX_DFA_SIZE_LIMIT: usize = 1024 * 1024;
static USER_REGEX_CACHE: OnceLock<UserRegexCache> = OnceLock::new();
@@ -16,8 +28,23 @@ fn user_regex_cache() -> &'static UserRegexCache {
USER_REGEX_CACHE.get_or_init(UserRegexCache::new)
}
fn compile_bounded(pattern: &str) -> Result<Regex, regex::Error> {
if pattern.len() > MAX_USER_REGEX_PATTERN_BYTES {
return Err(regex::Error::Syntax(format!(
"pattern is {} bytes; the limit is {MAX_USER_REGEX_PATTERN_BYTES}",
pattern.len()
)));
}
RegexBuilder::new(pattern)
.size_limit(USER_REGEX_SIZE_LIMIT)
.dfa_size_limit(USER_REGEX_DFA_SIZE_LIMIT)
.build()
}
struct UserRegexCache {
inner: Mutex<LruCache<String, Regex>>,
/// Compiled patterns and rejected ones alike, so a repeated bad pattern
/// is not recompiled.
inner: Mutex<LruCache<String, Result<Regex, regex::Error>>>,
}
impl UserRegexCache {
@@ -34,16 +61,22 @@ impl UserRegexCache {
}
fn compile(&self, pattern: &str) -> Result<Regex, regex::Error> {
let Ok(mut cache) = self.inner.lock() else {
return Regex::new(pattern);
};
if let Some(regex) = cache.get(pattern) {
return Ok(regex.clone());
// An over-long pattern is refused before it is hashed or stored.
if pattern.len() > MAX_USER_REGEX_PATTERN_BYTES {
return compile_bounded(pattern);
}
let regex = Regex::new(pattern)?;
cache.put(pattern.to_string(), regex.clone());
Ok(regex)
if let Ok(mut cache) = self.inner.lock()
&& let Some(cached) = cache.get(pattern)
{
return cached.clone();
}
// Compile without the lock: one slow pattern must not stall every
// other caller. Two racing misses both compile; the result is equal.
let compiled = compile_bounded(pattern);
if let Ok(mut cache) = self.inner.lock() {
cache.put(pattern.to_string(), compiled.clone());
}
compiled
}
#[cfg(test)]
@@ -88,11 +121,31 @@ mod tests {
}
#[test]
fn invalid_pattern_is_not_cached() {
fn invalid_pattern_is_remembered_as_invalid() {
let cache = UserRegexCache::with_capacity(NonZeroUsize::new(2).unwrap());
assert!(cache.compile("[").is_err());
assert!(cache.contains("["));
assert!(cache.compile("[").is_err());
assert_eq!(cache.len(), 1);
}
#[test]
fn oversized_patterns_are_refused_without_being_cached() {
let cache = UserRegexCache::with_capacity(NonZeroUsize::new(2).unwrap());
let long = "a".repeat(MAX_USER_REGEX_PATTERN_BYTES + 1);
assert!(cache.compile(&long).is_err());
assert_eq!(cache.len(), 0);
// Within the length cap, a pattern whose compiled form exceeds the
// program limit is refused rather than built.
let explosive = r"(?:\w{100}){100}";
assert!(explosive.len() <= MAX_USER_REGEX_PATTERN_BYTES);
assert!(matches!(
cache.compile(explosive),
Err(regex::Error::CompiledTooBig(_))
));
}
}
+56 -25
View File
@@ -23,10 +23,20 @@
//! supposed to be invisible. If a future feature ever needs per-engine
//! retry surfaces, swap this for an `Arc<RwLock<...>>` carried on the
//! `EngineHandle`; the public API stays the same.
//!
//! The `Retry-After` pause is process-wide too, but keyed by provider scope:
//! every request to the rate-limited route waits it out, and requests to any
//! other route (another provider, a local runtime) do not.
use std::collections::HashMap;
use std::sync::{Mutex, OnceLock};
use std::time::{Duration, Instant};
/// Rate-limit pause deadlines, one per provider scope (see
/// [`note_rate_limit`]). A 429 from one provider must not stall requests to
/// another provider, a local runtime, or a sub-agent on a different route.
type RateLimitPauses = HashMap<String, Instant>;
/// One in-flight retry attempt. `deadline` is the wall-clock time the
/// next request will fire — the UI subtracts `Instant::now()` from it
/// to render a live countdown.
@@ -97,10 +107,10 @@ fn with_state<R>(f: impl FnOnce(&mut RetryState) -> R) -> R {
}
#[cfg(not(any(test, all(feature = "test-thread-scoped-state", debug_assertions))))]
fn with_rate_limit<R>(f: impl FnOnce(&mut Option<Instant>) -> R) -> R {
static STATE: OnceLock<Mutex<Option<Instant>>> = OnceLock::new();
fn with_rate_limit<R>(f: impl FnOnce(&mut RateLimitPauses) -> R) -> R {
static STATE: OnceLock<Mutex<RateLimitPauses>> = OnceLock::new();
let mut state = STATE
.get_or_init(|| Mutex::new(None))
.get_or_init(|| Mutex::new(HashMap::new()))
.lock()
.unwrap_or_else(|error| error.into_inner());
f(&mut state)
@@ -140,13 +150,11 @@ fn with_state<R>(f: impl FnOnce(&mut RetryState) -> R) -> R {
}
#[cfg(any(test, all(feature = "test-thread-scoped-state", debug_assertions)))]
fn with_rate_limit<R>(f: impl FnOnce(&mut Option<Instant>) -> R) -> R {
#[allow(clippy::type_complexity)]
static STATE: OnceLock<
Mutex<std::collections::HashMap<std::thread::ThreadId, Option<Instant>>>,
> = OnceLock::new();
fn with_rate_limit<R>(f: impl FnOnce(&mut RateLimitPauses) -> R) -> R {
static STATE: OnceLock<Mutex<HashMap<std::thread::ThreadId, RateLimitPauses>>> =
OnceLock::new();
let mut by_thread = STATE
.get_or_init(|| Mutex::new(std::collections::HashMap::new()))
.get_or_init(|| Mutex::new(HashMap::new()))
.lock()
.unwrap_or_else(|error| error.into_inner());
f(by_thread.entry(std::thread::current().id()).or_default())
@@ -160,26 +168,32 @@ pub fn snapshot() -> RetryState {
with_state(|state| state.clone())
}
/// Extend the provider-wide rate-limit pause window. This is separate from
/// the footer banner so one successful concurrent request cannot clear another
/// request's active `Retry-After` window.
pub fn note_rate_limit(delay: Duration) {
let deadline = Instant::now() + delay;
with_rate_limit(|current| {
if current.is_none_or(|existing| existing < deadline) {
*current = Some(deadline);
/// Extend the rate-limit pause window for one provider `scope` (the caller's
/// route identity and host). This is separate from the footer banner so one
/// successful concurrent request cannot clear another request's active
/// `Retry-After` window, and it is keyed so a 429 from one provider never
/// pauses requests to another. Expired scopes are dropped here, so the map
/// holds at most the providers currently rate limited.
pub fn note_rate_limit(scope: &str, delay: Duration) {
let now = Instant::now();
let deadline = now + delay;
with_rate_limit(|pauses| {
pauses.retain(|_, existing| *existing > now);
let current = pauses.entry(scope.to_string()).or_insert(deadline);
if *current < deadline {
*current = deadline;
}
});
}
/// Remaining provider-wide rate-limit pause, if any.
/// Remaining rate-limit pause for `scope`, if any.
#[must_use]
pub fn rate_limit_remaining() -> Option<Duration> {
pub fn rate_limit_remaining(scope: &str) -> Option<Duration> {
let now = Instant::now();
with_rate_limit(|current| match *current {
with_rate_limit(|pauses| match pauses.get(scope).copied() {
Some(deadline) if deadline > now => Some(deadline.duration_since(now)),
Some(_) => {
*current = None;
pauses.remove(scope);
None
}
None => None,
@@ -220,9 +234,10 @@ pub fn clear() {
with_state(|state| *state = RetryState::Idle);
}
/// Drop every scope's rate-limit pause.
#[cfg(any(test, feature = "test-support"))]
pub fn clear_rate_limit() {
with_rate_limit(|current| *current = None);
with_rate_limit(HashMap::clear);
}
/// Test helper: serialize tests that touch the global state so cargo's
@@ -300,13 +315,29 @@ mod tests {
#[test]
fn rate_limit_deadline_survives_banner_clear() {
let _g = setup();
note_rate_limit(Duration::from_secs(5));
note_rate_limit("provider-a", Duration::from_secs(5));
start(1, Duration::from_secs(5), "rate limited");
succeeded();
assert!(
rate_limit_remaining().is_some(),
"provider-wide rate limit pause must not be cleared by an unrelated success"
rate_limit_remaining("provider-a").is_some(),
"provider rate limit pause must not be cleared by an unrelated success"
);
clear_rate_limit();
}
#[test]
fn one_providers_rate_limit_does_not_pause_another() {
let _g = setup();
note_rate_limit("[email protected]", Duration::from_secs(30));
assert!(rate_limit_remaining("[email protected]").is_some());
assert_eq!(rate_limit_remaining("[email protected]:11434"), None);
// A shorter Retry-After never shortens an existing window.
note_rate_limit("[email protected]", Duration::from_secs(1));
assert!(
rate_limit_remaining("[email protected]").unwrap() > Duration::from_secs(20)
);
clear_rate_limit();
assert_eq!(rate_limit_remaining("[email protected]"), None);
}
}
+20 -4
View File
@@ -77,6 +77,11 @@ const MAX_ACP_SESSIONS: usize = 64;
/// broken client cannot strand the stdio server forever after cancellation.
const ACP_PERMISSION_CANCEL_GRACE: Duration = Duration::from_secs(2);
/// How long a cancelled tool gets to observe its token and wind down (kill a
/// child, flush a write) before the turn stops waiting and drops it. A tool
/// that ignores its token must not hold the stdio server forever.
const ACP_TOOL_CANCEL_GRACE: Duration = Duration::from_secs(5);
/// Agent-originated JSON-RPC request ids have their own namespace. Strings
/// avoid the client-specific numeric response-id compatibility shim used for
/// replies to client-originated requests.
@@ -1226,11 +1231,22 @@ where
write_jsonrpc_result(writer, msg_id, json!(null)).await?;
}
cancel_token.cancel();
// Give the tool a chance to observe the token and
// wind down (e.g. kill a running child process)
// before we drop it.
// Give the tool a bounded chance to observe the
// token and wind down (e.g. kill a running child
// process) before we drop it.
cancelled = true;
break (&mut exec_fut).await;
break match tokio::time::timeout(
ACP_TOOL_CANCEL_GRACE,
&mut exec_fut,
)
.await
{
Ok(result) => result,
Err(_) => Err(ToolError::cancelled(format!(
"the tool did not stop within {}s of cancellation and was abandoned",
ACP_TOOL_CANCEL_GRACE.as_secs()
))),
};
}
if let Some(msg_id) = msg_id {
let msg_id = response_id_policy.response_id(msg_id);
+33 -12
View File
@@ -3765,17 +3765,22 @@ impl CodewhaleClient {
return self.send_with_isolated_retry(build, disclosure).await;
}
let retry_cfg: LlmRetryConfig = self.retry.clone().into();
let pause_scope = self.rate_limit_scope();
let callback_scope = pause_scope.clone();
let request_result = with_retry(
&retry_cfg,
|| {
let request = build();
let pause_scope = pause_scope.as_str();
async move {
// Sleep in bounded slices rather than the full remaining
// window: the pause is process-global, so a concurrent
// `clear_rate_limit()` (or a shortened deadline) must
// release requests that are already waiting instead of
// stranding them for the whole original window.
while let Some(delay) = crate::retry_status::rate_limit_remaining() {
// window: the pause is shared by every request to this
// route, so a concurrent `clear_rate_limit()` (or a
// shortened deadline) must release requests that are
// already waiting instead of stranding them for the whole
// original window.
while let Some(delay) = crate::retry_status::rate_limit_remaining(pause_scope)
{
tokio::time::sleep(delay.min(RATE_LIMIT_PAUSE_RECHECK_INTERVAL)).await;
}
self.wait_for_rate_limit().await;
@@ -3793,7 +3798,7 @@ impl CodewhaleClient {
Err(self.http_error_with_route_context(status.as_u16(), &body, retry_after))
}
},
Some(Box::new(|err, attempt, delay| {
Some(Box::new(move |err, attempt, delay| {
let (reason_label, human_reason) = retry_reason_label_and_human(err);
logging::warn(format!(
"HTTP retry reason={} attempt={} delay={:.2}s",
@@ -3802,7 +3807,7 @@ impl CodewhaleClient {
delay.as_secs_f64(),
));
if matches!(err, LlmError::RateLimited { .. }) {
crate::retry_status::note_rate_limit(delay);
crate::retry_status::note_rate_limit(&callback_scope, delay);
}
crate::retry_status::start(attempt + 1, delay, human_reason);
})),
@@ -3818,6 +3823,7 @@ impl CodewhaleClient {
Err(err) => {
if let LlmError::RateLimited { retry_after, .. } = &err.last_error {
crate::retry_status::note_rate_limit(
&pause_scope,
retry_after
.unwrap_or_else(|| retry_cfg.delay_for_attempt(retry_cfg.max_retries)),
);
@@ -3838,6 +3844,21 @@ impl CodewhaleClient {
}
}
/// Key for this route's shared `Retry-After` pause: the configured route
/// identity plus the host it reaches. A 429 from one provider pauses only
/// requests that would hit the same limit, never another provider, a
/// local runtime, or a sub-agent on a different route.
fn rate_limit_scope(&self) -> String {
let route = if self.provider_identity.is_empty() {
self.api_provider.as_str()
} else {
self.provider_identity.as_str()
};
let host = crate::llm_client::base_url_authority(&self.base_url)
.unwrap_or_else(|| redact_url_for_display(&self.base_url));
format!("{route}@{host}")
}
/// The same bounded transport retry policy without process-global retry
/// banners, provider-wide pause cells, or shared connection-health writes.
/// Used only by the Auto classifier during read-only request inspection.
@@ -6596,7 +6617,6 @@ mod tests {
crate::retry_status::clear();
crate::retry_status::clear_rate_limit();
crate::retry_status::start(7, Duration::from_secs(60), "foreground sentinel");
crate::retry_status::note_rate_limit(Duration::from_secs(60));
let server = MockServer::start().await;
Mock::given(method("POST"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
@@ -6615,6 +6635,7 @@ mod tests {
.await;
let client = deepseek_request_boundary_client("https://api.deepseek.com/v1", server.uri());
crate::retry_status::note_rate_limit(&client.rate_limit_scope(), Duration::from_secs(60));
let request = MessageRequest {
model: "deepseek-v4-pro".to_string(),
messages: vec![Message {
@@ -6681,7 +6702,7 @@ mod tests {
state => panic!("isolated success mutated retry state: {state:?}"),
}
assert!(
crate::retry_status::rate_limit_remaining().is_some(),
crate::retry_status::rate_limit_remaining(&client.rate_limit_scope()).is_some(),
"isolated success must not clear the foreground provider pause"
);
crate::retry_status::clear();
@@ -6697,7 +6718,6 @@ mod tests {
crate::retry_status::clear();
crate::retry_status::clear_rate_limit();
crate::retry_status::start(9, Duration::from_secs(60), "foreground sentinel 429");
crate::retry_status::note_rate_limit(Duration::from_secs(60));
let server = MockServer::start().await;
Mock::given(method("POST"))
@@ -6713,6 +6733,7 @@ mod tests {
deepseek_request_boundary_client("https://api.deepseek.com/v1", server.uri());
client.retry.enabled = false;
client.retry.max_retries = 0;
crate::retry_status::note_rate_limit(&client.rate_limit_scope(), Duration::from_secs(60));
let request = MessageRequest {
model: "deepseek-v4-pro".to_string(),
messages: vec![Message {
@@ -6751,8 +6772,8 @@ mod tests {
}
state => panic!("isolated 429 mutated retry state: {state:?}"),
}
let remaining =
crate::retry_status::rate_limit_remaining().expect("foreground provider pause remains");
let remaining = crate::retry_status::rate_limit_remaining(&client.rate_limit_scope())
.expect("foreground provider pause remains");
assert!(
remaining < Duration::from_secs(70),
"classifier Retry-After must not extend the global pause: {remaining:?}"
+191 -83
View File
@@ -36,8 +36,8 @@ use codewhale_models::{ContentBlock, MessageRequest, MessageResponse, StreamEven
use super::prepared::WireDialect;
use super::role_placement::{RolePlacement, role_placement};
use super::wire::{extract_sse_data_value, next_sse_line};
use super::{CodewhaleClient, ERROR_BODY_MAX_BYTES, bounded_error_text};
use super::wire::{extract_sse_data_value, next_sse_line, push_sse_event_data};
use super::CodewhaleClient;
/// Maximum `cache_control` breakpoints Anthropic accepts per request.
const MAX_CACHE_BREAKPOINTS: usize = 4;
@@ -213,78 +213,64 @@ impl CodewhaleClient {
body
}
/// Non-streaming send through the shared typed retry path
/// (`send_with_retry`): 429/5xx/transport failures retry with backoff and
/// honor `Retry-After`, and a final failure stays a downcastable
/// `LlmError` so the engine can classify auth, rate-limit, context and
/// invalid-request failures like every other wire.
async fn send_anthropic_request(&self, url: &str, body: &Value) -> Result<reqwest::Response> {
let url = self.messages_transport_url(url);
self.wait_for_rate_limit().await;
let response = self
.http_client
.post(&url)
.header("Accept", "text/event-stream")
.json(body)
.send()
.await
.context("Anthropic Messages API request failed")?;
self.check_anthropic_response(response).await
}
/// Shared status/error-envelope handling for streaming and
/// non-streaming Messages responses.
async fn check_anthropic_response(
&self,
response: reqwest::Response,
) -> Result<reqwest::Response> {
let status = response.status();
if !status.is_success() {
let raw = bounded_error_text(response, ERROR_BODY_MAX_BYTES).await;
let (error_type, message) = parse_anthropic_error_envelope(&raw);
self.mark_request_failure(&format!("anthropic status={status}"))
.await;
anyhow::bail!("Anthropic API error (HTTP {status} {error_type}): {message}");
}
self.mark_request_success().await;
Ok(response)
let request_body =
serde_json::to_vec(body).context("Failed to serialize Anthropic Messages request")?;
self.send_with_retry(|| {
self.http_client
.post(&url)
.header(reqwest::header::CONTENT_TYPE, "application/json")
.header("Accept", "text/event-stream")
.body(request_body.clone())
})
.await
.context("Anthropic Messages API request failed")
}
/// Open the streaming Messages request through the shared stream-entry
/// transport policy: bounded header wait, dual-client selection, and at
/// most one HTTP/1.1 fallback retry on a classified H2 header stall.
/// Wire-specific request construction (headers, endpoint, body) stays
/// here at the adapter edge.
/// Inside each open attempt the provider retry loop (`send_with_retry`)
/// handles rate limits and transient upstream failures before any stream
/// body exists, as the Chat and Responses adapters do. Wire-specific
/// request construction (headers, endpoint, body) stays here at the
/// adapter edge.
async fn open_anthropic_stream_response(
&self,
url: &str,
body: &Value,
) -> Result<reqwest::Response> {
let url = self.messages_transport_url(url);
let request_body =
serde_json::to_vec(body).context("Failed to serialize Anthropic Messages request")?;
let open_req = self.stream_open_request();
let opened = super::stream_entry::open_sse_response(&open_req, |policy| {
super::stream_entry::open_sse_response(&open_req, |policy| {
let url = url.clone();
let request_body = request_body.clone();
async move {
self.wait_for_rate_limit().await;
let client = super::stream_entry::client_for_policy(
&self.http_client,
self.http1_fallback_client(),
policy,
);
client
.post(&url)
.header("Accept", "text/event-stream")
.json(body)
.send()
.await
.context("Anthropic Messages API request failed")
self.send_with_retry(|| {
client
.post(&url)
.header(reqwest::header::CONTENT_TYPE, "application/json")
.header("Accept", "text/event-stream")
.body(request_body.clone())
})
.await
.context("Anthropic Messages API request failed")
}
})
.await;
let response = match opened {
Ok(response) => response,
Err(err) => {
self.mark_request_failure(&format!("anthropic stream open: {err}"))
.await;
return Err(err);
}
};
self.check_anthropic_response(response).await
.await
}
/// Handle a streaming Messages API request.
@@ -312,6 +298,9 @@ impl CodewhaleClient {
// multi-byte UTF-8 char (CJK/emoji) split across HTTP/2 DATA is
// never corrupted to U+FFFD. Genuine invalid bytes fail closed.
let mut buffer: Vec<u8> = Vec::new();
// `data:` fields of the event being assembled, dispatched at the
// blank line that ends it (or at stream end).
let mut event_data = String::new();
let stream_start = std::time::Instant::now();
let mut last_chunk_at = std::time::Instant::now();
let mut bytes_received: usize = 0;
@@ -350,22 +339,36 @@ impl CodewhaleClient {
}
loop {
let line = match next_sse_line(&mut buffer, ended) {
Ok(Some(line)) => line,
let data = match next_sse_line(&mut buffer, ended) {
// A blank line ends the event.
Ok(Some(line)) if line.is_empty() => std::mem::take(&mut event_data),
Ok(Some(line)) => {
// `event:` lines are redundant (the data payload
// carries `type`) and comment/heartbeat lines are
// ignorable.
if let Some(value) = extract_sse_data_value(&line)
&& let Err(err) = push_sse_event_data(&mut event_data, value)
{
yield Err(anyhow::anyhow!("{err}"));
return;
}
continue;
}
// The final event may arrive without its blank line.
Ok(None) if ended && !event_data.is_empty() => {
std::mem::take(&mut event_data)
}
Ok(None) => break,
Err(err) => {
yield Err(anyhow::anyhow!("{err}"));
return;
}
};
// `event:` lines are redundant (the data payload carries
// `type`) and comment/heartbeat lines are ignorable.
let Some(data) = extract_sse_data_value(&line) else {
if data.is_empty() {
continue;
};
}
match convert_anthropic_sse_data(data) {
match convert_anthropic_sse_data(&data) {
Some(Ok(StreamEvent::Error { error })) => {
let (error_type, message) = anthropic_error_fields(&error);
yield Err(anyhow::anyhow!(
@@ -391,6 +394,12 @@ impl CodewhaleClient {
break;
}
}
// Only `message_stop` (returned above) or a provider error proves
// the response is whole. A bare HTTP EOF is truncation, and ending
// the stream quietly would hand the turn a partial answer as done.
yield Err(anyhow::anyhow!(
"Anthropic Messages stream closed before message_stop"
));
};
Ok(Box::pin(stream))
@@ -994,14 +1003,6 @@ fn parse_anthropic_usage(usage: &Value) -> Usage {
/// Extract `error.type` / `error.message` from an Anthropic error envelope
/// (`{"type":"error","error":{"type":...,"message":...}}`), falling back to
/// the raw body so nothing is swallowed.
fn parse_anthropic_error_envelope(raw: &str) -> (String, String) {
let Ok(value) = serde_json::from_str::<Value>(raw) else {
return ("unknown".to_string(), raw.to_string());
};
let error = value.get("error").unwrap_or(&value);
anthropic_error_fields(error)
}
fn anthropic_error_fields(error: &Value) -> (String, String) {
let error_type = error
.get("type")
@@ -2027,19 +2028,6 @@ mod tests {
assert_eq!(usage.prompt_cache_write_tokens, Some(2045));
}
#[test]
fn error_envelope_parses_type_and_message() {
let (error_type, message) = parse_anthropic_error_envelope(
r#"{"type":"error","error":{"type":"rate_limit_error","message":"Too many requests"},"request_id":"req_1"}"#,
);
assert_eq!(error_type, "rate_limit_error");
assert_eq!(message, "Too many requests");
let (error_type, message) = parse_anthropic_error_envelope("upstream blew up");
assert_eq!(error_type, "unknown");
assert_eq!(message, "upstream blew up");
}
#[test]
fn data_url_image_becomes_a_base64_source_not_a_url_source() {
// Anthropic rejects a `data:` URL under `{"type":"url"}`. This is the
@@ -2248,6 +2236,73 @@ mod tests {
assert!(saw_stop, "message_stop should arrive through the seam");
}
async fn collect_anthropic_stream(body: &'static str) -> Vec<Result<StreamEvent>> {
use futures_util::StreamExt;
use wiremock::matchers::{method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/v1/messages"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("Content-Type", "text/event-stream")
.set_body_string(body),
)
.mount(&server)
.await;
let client = deepseek_test_client(&server.uri());
let stream = client
.handle_anthropic_stream(
&client
.prepare_outbound_request(request_with("deepseek-v4", None, None, None), true)
.expect("anthropic request prepares"),
)
.await
.expect("stream opens");
tokio::time::timeout(std::time::Duration::from_secs(5), stream.collect())
.await
.expect("stream ends")
}
#[tokio::test]
async fn anthropic_stream_joins_multiline_data_fields_into_one_event() {
use codewhale_models::Delta;
// One JSON payload split across two `data:` fields of a single event.
let events = collect_anthropic_stream(concat!(
"event: content_block_delta\n",
"data: {\"type\":\"content_block_delta\",\"index\":0,\n",
"data: \"delta\":{\"type\":\"text_delta\",\"text\":\"joined\"}}\n\n",
"data: {\"type\":\"message_stop\"}\n\n",
))
.await;
assert!(
events.iter().any(|event| matches!(
event,
Ok(StreamEvent::ContentBlockDelta {
delta: Delta::TextDelta { text },
..
}) if text == "joined"
)),
"the split event was lost: {events:?}"
);
assert!(matches!(events.last(), Some(Ok(StreamEvent::MessageStop))));
}
#[tokio::test]
async fn anthropic_stream_eof_without_message_stop_is_an_error() {
let events = collect_anthropic_stream(
"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"partial\"}}\n\n",
)
.await;
let last = events.last().expect("the delta and a terminal item");
assert!(
last.as_ref()
.is_err_and(|error| error.to_string().contains("closed before message_stop")),
"a truncated stream must not end quietly: {events:?}"
);
}
/// Fault injection (#6184): a provider that answers the headers and then
/// sends nothing fails the stream at the first-byte bound with a
/// distinct error and a `crashes/` stall record, instead of holding the
@@ -2344,10 +2399,63 @@ mod tests {
Ok(_) => panic!("auth errors must fail fast"),
Err(err) => err,
};
let text = err.to_string();
// Typed through the shared classifier, provider message kept.
assert!(
text.contains("HTTP 401") && text.contains("authentication_error"),
"error envelope should be preserved: {text}"
matches!(
err.chain()
.find_map(|cause| cause.downcast_ref::<crate::llm_client::LlmError>()),
Some(crate::llm_client::LlmError::AuthenticationError(_))
),
"a 401 must stay a typed authentication error: {err:#}"
);
assert!(format!("{err:#}").contains("bad key"), "{err:#}");
}
#[tokio::test]
async fn anthropic_rate_limit_is_retried_honoring_retry_after() {
use wiremock::matchers::{method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/v1/messages"))
.respond_with(
ResponseTemplate::new(429)
.insert_header("retry-after", "0")
.set_body_string(
"{\"type\":\"error\",\"error\":{\"type\":\"rate_limit_error\",\"message\":\"slow down\"}}",
),
)
.up_to_n_times(1)
.expect(1)
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path("/v1/messages"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("Content-Type", "text/event-stream")
.set_body_string("data: {\"type\":\"message_stop\"}\n\n"),
)
.expect(1)
.mount(&server)
.await;
let mut client = deepseek_test_client(&server.uri());
client.retry.enabled = true;
client.retry.max_retries = 2;
client.retry.initial_delay = 0.0;
client.retry.max_delay = 0.0;
let stream = client
.handle_anthropic_stream(
&client
.prepare_outbound_request(request_with("deepseek-v4", None, None, None), true)
.expect("anthropic request prepares"),
)
.await;
assert!(
stream.is_ok(),
"a 429 before the stream body must be retried, not surfaced"
);
}
+14 -15
View File
@@ -33,7 +33,7 @@ use codewhale_models::{
use super::prepared::WireDialect;
use super::role_placement::{RolePlacement, role_placement};
use super::wire::{extract_sse_data_value, flush_sse_line, take_sse_line};
use super::wire::{extract_sse_data_value, flush_sse_line, push_sse_event_data, take_sse_line};
use super::{
CodewhaleClient, ERROR_BODY_MAX_BYTES, SSE_BACKPRESSURE_HIGH_WATERMARK,
SSE_BACKPRESSURE_SLEEP_MS, SSE_MAX_LINES_PER_CHUNK, acquire_stream_buffer,
@@ -1603,15 +1603,13 @@ impl CodewhaleClient {
continue;
}
if let Some(data) = extract_sse_data_value(&line) {
// The SSE spec joins multiple `data:` fields within one
// event with '\n'; concatenating with no separator would
// yield `{…}{…}` and fail JSON parsing, silently dropping
// the frame.
if !line_buf.is_empty() {
line_buf.push('\n');
}
line_buf.push_str(data);
if let Some(data) = extract_sse_data_value(&line)
&& let Err(err) = push_sse_event_data(&mut line_buf, data)
{
decode_failed = true;
stream_failed = true;
yield Err(anyhow::anyhow!("{err}"));
break 'stream;
}
// Ignore other SSE fields (event:, id:, retry:)
@@ -1639,11 +1637,12 @@ impl CodewhaleClient {
if !saw_done && !decode_failed {
match flush_sse_line(&mut byte_buf) {
Ok(Some(line)) => {
if let Some(data) = extract_sse_data_value(&line) {
if !line_buf.is_empty() {
line_buf.push('\n');
}
line_buf.push_str(data);
if let Some(data) = extract_sse_data_value(&line)
&& let Err(err) = push_sse_event_data(&mut line_buf, data)
{
decode_failed = true;
stream_failed = true;
yield Err(anyhow::anyhow!("{err}"));
}
}
Ok(None) => {}
+35 -12
View File
@@ -21,7 +21,7 @@ use codewhale_models::{
use super::prepared::WireDialect;
use super::role_placement::{RolePlacement, role_placement};
use super::wire::{extract_sse_data_value, next_sse_line};
use super::wire::{extract_sse_data_value, next_sse_line, push_sse_event_data};
use super::{
CodewhaleClient, ERROR_BODY_MAX_BYTES, bounded_error_text, from_api_tool_name,
system_to_instructions, to_api_tool_name,
@@ -255,6 +255,9 @@ impl CodewhaleClient {
// multi-byte UTF-8 char split across HTTP/2 DATA is never
// corrupted to U+FFFD. Genuine invalid bytes fail closed.
let mut buffer: Vec<u8> = Vec::new();
// `data:` fields of the event being assembled, dispatched at the
// blank line that ends it (or at stream end).
let mut event_data = String::new();
let mut done = false;
let mut ended = false;
let mut content_block_counter: u32 = 0;
@@ -297,25 +300,37 @@ impl CodewhaleClient {
// Process complete SSE lines, and the unterminated tail at stream end.
loop {
let line = match next_sse_line(&mut buffer, ended) {
Ok(Some(line)) => line,
let data = match next_sse_line(&mut buffer, ended) {
// A blank line ends the event.
Ok(Some(line)) if line.is_empty() => std::mem::take(&mut event_data),
Ok(Some(line)) => {
if line.starts_with(':') {
// SSE comment keep-alive: the provider is alive (#6184).
yield Ok(StreamEvent::Ping);
} else if let Some(value) = extract_sse_data_value(&line)
&& let Err(err) = push_sse_event_data(&mut event_data, value)
{
yield Err(anyhow::anyhow!("{err}"));
return;
}
continue;
}
// The final event may arrive without its blank line.
Ok(None) if ended && !event_data.is_empty() => {
std::mem::take(&mut event_data)
}
Ok(None) => break,
Err(err) => {
yield Err(anyhow::anyhow!("{err}"));
return;
}
};
if line.is_empty() {
continue;
}
if line.starts_with(':') {
// SSE comment keep-alive: the provider is alive (#6184).
yield Ok(StreamEvent::Ping);
if data.is_empty() {
continue;
}
if let Some(data) = extract_sse_data_value(&line) {
{
let data = data.as_str();
if data == "[DONE]" {
done = true;
break;
@@ -561,7 +576,15 @@ impl CodewhaleClient {
}
}
// Emit MessageStop.
// Only `[DONE]` or `response.completed`/`response.incomplete`
// proves the response is whole. A bare HTTP EOF is truncation, and
// a MessageStop here would hand the turn a partial answer as done.
if !done {
yield Err(anyhow::anyhow!(
"Responses stream closed before response.completed or [DONE]"
));
return;
}
yield Ok(StreamEvent::MessageStop);
};
+74
View File
@@ -234,6 +234,80 @@ async fn responses_stream_retries_upstream_499_before_streaming() {
assert_eq!(attempts.load(Ordering::SeqCst), 2);
}
async fn collect_responses_stream(sse_body: &'static str) -> Vec<Result<StreamEvent>> {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path(CODEX_RESPONSES_PATH))
.respond_with(
ResponseTemplate::new(200)
.insert_header("Content-Type", "text/event-stream")
.set_body_string(sse_body),
)
.mount(&server)
.await;
let client = {
let _env_lock = crate::test_support::lock_test_env();
let _codex_token =
crate::test_support::EnvVarGuard::set("OPENAI_CODEX_ACCESS_TOKEN", "test-token");
let _legacy_codex_token = crate::test_support::EnvVarGuard::remove("CODEX_ACCESS_TOKEN");
CodewhaleClient::new(&test_codex_config(&server)).unwrap()
};
let stream = client
.handle_responses_stream(
&client
.prepare_outbound_request(minimal_responses_request(), true)
.expect("responses request prepares"),
)
.await
.expect("Responses stream opens");
tokio::time::timeout(std::time::Duration::from_secs(5), stream.collect())
.await
.expect("stream ends")
}
#[tokio::test]
async fn responses_stream_eof_without_a_terminal_event_is_an_error() {
let events = collect_responses_stream(concat!(
"data: {\"type\":\"response.output_item.added\",\"item\":{\"type\":\"message\",\"id\":\"m\"}}\n\n",
"data: {\"type\":\"response.output_text.delta\",\"delta\":\"partial\"}\n\n",
))
.await;
assert!(
!events
.iter()
.any(|event| matches!(event, Ok(StreamEvent::MessageStop))),
"a truncated stream must not report MessageStop: {events:?}"
);
assert!(
events.last().is_some_and(|event| event
.as_ref()
.is_err_and(|error| error.to_string().contains("closed before response.completed"))),
"{events:?}"
);
}
#[tokio::test]
async fn responses_stream_joins_multiline_data_fields_into_one_event() {
let events = collect_responses_stream(concat!(
"data: {\"type\":\"response.output_item.added\",\"item\":{\"type\":\"message\",\"id\":\"m\"}}\n\n",
"data: {\"type\":\"response.output_text.delta\",\n",
"data: \"delta\":\"joined\"}\n\n",
"data: [DONE]\n\n",
))
.await;
assert!(
events.iter().any(|event| matches!(
event,
Ok(StreamEvent::ContentBlockDelta {
delta: Delta::TextDelta { text },
..
}) if text == "joined"
)),
"the split event was lost: {events:?}"
);
assert!(matches!(events.last(), Some(Ok(StreamEvent::MessageStop))));
}
#[tokio::test]
async fn responses_stream_finishes_on_semantic_terminal_event_without_done_marker() {
let server = MockServer::start().await;
+30
View File
@@ -98,6 +98,10 @@ pub(super) fn extract_sse_data_value(line: &str) -> Option<&str> {
/// Hard ceiling for one pending SSE line, matching the MCP frame limit.
/// This does not bound the size of a multi-line event or the whole stream.
const MAX_SSE_LINE_BYTES: usize = 8 * 1024 * 1024;
/// Hard ceiling for one assembled SSE event (its `data:` fields joined), so a
/// peer that never sends the blank line ending an event cannot grow memory
/// without bound. Same size as the single-line ceiling.
const MAX_SSE_EVENT_BYTES: usize = MAX_SSE_LINE_BYTES;
/// Invalid or oversized SSE line (or unterminated flush).
///
@@ -109,6 +113,7 @@ const MAX_SSE_LINE_BYTES: usize = 8 * 1024 * 1024;
pub(super) enum SseLineError {
InvalidUtf8 { valid_up_to: usize },
TooLong,
EventTooLong,
}
impl std::fmt::Display for SseLineError {
@@ -121,6 +126,10 @@ impl std::fmt::Display for SseLineError {
f,
"SSE line exceeded {MAX_SSE_LINE_BYTES} bytes (8 MiB) — aborting stream"
),
Self::EventTooLong => write!(
f,
"SSE event exceeded {MAX_SSE_EVENT_BYTES} bytes (8 MiB) — aborting stream"
),
}
}
}
@@ -184,6 +193,27 @@ pub(super) fn flush_sse_line(buffer: &mut Vec<u8>) -> Result<Option<String>, Sse
decoded.map(|line| (!line.is_empty()).then_some(line))
}
/// Append one `data:` field to the event being assembled. The SSE spec joins
/// the `data:` fields of one event with '\n' and dispatches the event at the
/// blank line that ends it; a provider may split one JSON payload across
/// several fields, so parsing each field alone would drop the event. Fails
/// closed (clearing the event) past [`MAX_SSE_EVENT_BYTES`].
pub(super) fn push_sse_event_data(event: &mut String, data: &str) -> Result<(), SseLineError> {
let needed = event
.len()
.saturating_add(usize::from(!event.is_empty()))
.saturating_add(data.len());
if needed > MAX_SSE_EVENT_BYTES {
event.clear();
return Err(SseLineError::EventTooLong);
}
if !event.is_empty() {
event.push('\n');
}
event.push_str(data);
Ok(())
}
/// Next decoded SSE line. When `at_end` is false, wait for `\n`. When `at_end`
/// is true, also flush an unterminated tail (stream closed).
pub(super) fn next_sse_line(
+114 -47
View File
@@ -21,7 +21,6 @@
//! provider-accepted active observation.
use std::fs;
use std::io::Write;
use std::path::{Path, PathBuf};
use std::process::Command;
use std::time::{SystemTime, UNIX_EPOCH};
@@ -288,25 +287,13 @@ impl CloudJobStore {
}
/// Persist a job atomically. Never writes credentials.
///
/// The record is written through a unique temporary file (never a fixed
/// `.json.tmp` two writers would share) and renamed into place while the
/// job's lock is held, so it cannot interleave with a
/// [`Self::save_unless_canceled`] or a cancel.
pub fn save(&self, job: &CloudJob) -> Result<()> {
fs::create_dir_all(&self.root).context("failed to create cloud-jobs directory")?;
let path = self.job_path(&job.id)?;
let tmp = path.with_extension("json.tmp");
let body = serde_json::to_vec_pretty(job).context("failed to encode cloud job")?;
{
let mut file =
fs::File::create(&tmp).context("failed to start a private cloud job write")?;
file.write_all(&body)
.context("failed to write the cloud job record")?;
file.sync_all().ok();
}
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let _ = fs::set_permissions(&tmp, fs::Permissions::from_mode(0o600));
}
fs::rename(&tmp, &path).context("failed to commit the cloud job record")?;
Ok(())
self.with_job_lock(&job.id, || self.write_record(job))
}
/// Cancel-authoritative save: refuse to overwrite a `canceled` record.
@@ -319,19 +306,49 @@ impl CloudJobStore {
/// persisted record is already `canceled`, `Ok(true)` after a normal
/// save.
///
/// This is load-check-save: the store is file-backed with no
/// cross-process lock, so the check cannot remove the load→save window
/// entirely — it narrows the clobber window from a whole phase (seconds
/// to minutes) to the span of one save, which is the single-writer
/// discipline this store assumes elsewhere.
/// The load, check and save run under the job's advisory file lock, and
/// [`cancel_job`] flips the record under the same lock, so the check and
/// the write are one step with respect to a cancel in another process.
pub fn save_unless_canceled(&self, job: &CloudJob) -> Result<bool> {
if let Ok(current) = self.load(&job.id)
&& current.status == CloudJobStatus::Canceled
{
return Ok(false);
}
self.save(job)?;
Ok(true)
self.with_job_lock(&job.id, || {
if let Ok(current) = self.load(&job.id)
&& current.status == CloudJobStatus::Canceled
{
return Ok(false);
}
self.write_record(job)?;
Ok(true)
})
}
/// Run `f` holding the job's exclusive advisory lock
/// (`<root>/<id>.json.lock`). The lock is released when `f` returns.
///
/// Limitation: the lock is advisory, so it serializes Codewhale writers
/// only; it is not reentrant, so `f` must not call a locking method.
fn with_job_lock<T>(&self, id: &str, f: impl FnOnce() -> Result<T>) -> Result<T> {
let path = self.job_path(id)?;
fs::create_dir_all(&self.root).context("failed to create cloud-jobs directory")?;
let lock_file = fs::OpenOptions::new()
.create(true)
.truncate(false)
.read(true)
.write(true)
.open(path.with_extension("json.lock"))
.context("failed to open the cloud job lock")?;
let mut lock = fd_lock::RwLock::new(lock_file);
let _guard = lock
.write()
.context("failed to lock the cloud job record")?;
f()
}
/// Write the record through a unique private temporary file. Callers
/// hold the job lock.
fn write_record(&self, job: &CloudJob) -> Result<()> {
let path = self.job_path(&job.id)?;
let body = serde_json::to_vec_pretty(job).context("failed to encode cloud job")?;
crate::utils::write_atomic(&path, &body).context("failed to commit the cloud job record")
}
/// Load one job by id.
@@ -789,23 +806,32 @@ pub fn cancel_job(
id: &str,
launcher: &dyn DaytonaLauncher,
) -> Result<CloudJob> {
let mut job = store.load(id)?;
if matches!(
job.status,
CloudJobStatus::Canceled | CloudJobStatus::Failed | CloudJobStatus::Refused
) {
return Ok(job);
}
let had_sandbox = job.sandbox_id.is_some();
let sandbox_may_exist = had_sandbox || job.sandbox_pending;
job.status = CloudJobStatus::Canceled;
job.finished_unix = Some(unix_now());
job.note = if sandbox_may_exist {
"Canceled locally; the cloud agent sandbox is being torn down.".to_string()
} else {
"Canceled locally before a sandbox was created.".to_string()
// The flip is one locked read-modify-write, so a runner phase save cannot
// land between this load and this write (and lose either the cancel or
// the sandbox id the runner just recorded).
let flipped = store.with_job_lock(id, || {
let mut job = store.load(id)?;
if matches!(
job.status,
CloudJobStatus::Canceled | CloudJobStatus::Failed | CloudJobStatus::Refused
) {
return Ok((job, false));
}
let sandbox_may_exist = job.sandbox_id.is_some() || job.sandbox_pending;
job.status = CloudJobStatus::Canceled;
job.finished_unix = Some(unix_now());
job.note = if sandbox_may_exist {
"Canceled locally; the cloud agent sandbox is being torn down.".to_string()
} else {
"Canceled locally before a sandbox was created.".to_string()
};
store.write_record(&job)?;
Ok((job, true))
})?;
let mut job = match flipped {
(job, false) => return Ok(job),
(job, true) => job,
};
store.save(&job)?;
if let Some(sandbox_id) = job.sandbox_id.clone() {
let receipt = SandboxReceipt {
sandbox_id,
@@ -3181,6 +3207,47 @@ mod tests {
assert_eq!(persisted.finished_unix, Some(10_000_060));
}
#[test]
fn a_phase_save_waits_for_a_cancel_holding_the_job_lock() {
let temp = tempfile::tempdir().unwrap();
let store = CloudJobStore::from_path(temp.path().join("jobs"));
let job = stored_job(CloudJobStatus::Running, 10_000_000);
store.save(&job).unwrap();
let mut runner_copy = job.clone();
runner_copy.status = CloudJobStatus::OpeningPr;
// A cancel in progress holds the lock between its load and its write;
// the runner's check-and-save must not slip into that window.
let outcome = store
.with_job_lock(&job.id, || {
let runner_store = store.clone();
let runner = std::thread::spawn(move || {
runner_store.save_unless_canceled(&runner_copy).unwrap()
});
std::thread::sleep(std::time::Duration::from_millis(200));
assert!(
!runner.is_finished(),
"the phase save ran inside the cancel's read-modify-write"
);
let mut canceled = store.load(&job.id)?;
canceled.status = CloudJobStatus::Canceled;
store.write_record(&canceled)?;
Ok(runner)
})
.unwrap()
.join()
.unwrap();
assert!(!outcome, "the phase save must see the cancel and stand down");
assert_eq!(
store.load(&job.id).unwrap().status,
CloudJobStatus::Canceled
);
assert!(
!temp.path().join("jobs").join(format!("{}.json.tmp", job.id)).exists(),
"records are written through unique temporaries"
);
}
#[test]
fn sweep_fails_stale_active_jobs_and_tears_down_their_sandboxes() {
let temp = tempfile::tempdir().unwrap();
+18 -14
View File
@@ -350,8 +350,9 @@ pub enum ComputerMeterError {
/// Interval ends before it starts.
#[error("A Computer meter interval cannot end before it starts.")]
IntervalReversed,
/// Admission has expired before the interval started.
#[error("The Computer admission expired before this interval.")]
/// Admission expired before the interval ended. An interval must lie
/// wholly inside the admission; the tail after expiry was never authorized.
#[error("The Computer admission expired before this interval ended.")]
AdmissionExpired,
/// Meter or catalog revision is not the pinned v3 revision.
#[error("{message}")]
@@ -479,7 +480,10 @@ pub fn issue_computer_meter_receipt(
return Err(ComputerMeterError::IntervalReversed);
}
let expires = parse_timestamp(&admission.expires_at, "expiresAt")?;
if started >= expires {
// `ended_at` is exclusive, so an interval ending exactly at expiry is
// wholly authorized; one that runs past it is refused rather than clipped,
// because a clipped receipt would no longer restate the provider interval.
if started >= expires || ended > expires {
return Err(ComputerMeterError::AdmissionExpired);
}
let accepted_seconds = elapsed_whole_seconds(started, ended);
@@ -539,21 +543,21 @@ pub fn issue_computer_meter_receipt(
Ok(receipt)
}
/// Exact replay of an existing receipt. Identity and digest must match.
/// Exact replay of an existing receipt. Every field must match (only the
/// `replay_of` back-reference may differ), and the incoming digest and id must
/// be the ones its own fields produce: a caller-supplied digest is never
/// trusted to stand for fields it does not cover.
pub fn assert_computer_meter_receipt_replay(
existing: &ComputerMeterReceipt,
incoming: &ComputerMeterReceipt,
) -> Result<(), ComputerMeterError> {
if existing.receipt_id == incoming.receipt_id
&& existing.binding_digest == incoming.binding_digest
&& existing.admission_id == incoming.admission_id
&& existing.account_id == incoming.account_id
&& existing.provider_event_ref == incoming.provider_event_ref
&& existing.accepted_seconds == incoming.accepted_seconds
&& existing.standard_equivalent_seconds == incoming.standard_equivalent_seconds
&& existing.profile_id == incoming.profile_id
&& existing.multiplier == incoming.multiplier
{
let self_consistent = receipt_binding_digest(incoming) == incoming.binding_digest
&& receipt_id_for(incoming) == incoming.receipt_id;
let same_terms = ComputerMeterReceipt {
replay_of: existing.replay_of.clone(),
..incoming.clone()
} == *existing;
if self_consistent && same_terms {
return Ok(());
}
Err(ComputerMeterError::ReplayConflict)
+55
View File
@@ -190,6 +190,61 @@ fn exact_replay_is_idempotent_and_a_mutated_binding_conflicts() {
);
}
#[test]
fn replay_rejects_fields_the_supplied_digest_does_not_vouch_for() {
let case = fixtures()
.into_iter()
.find(|case| case.id == "accepted-standard-8")
.unwrap();
let first = issue_case(&case).unwrap();
// Same id and digest, different interval: the digest no longer matches
// the fields it claims to bind.
let mut shifted = first.clone();
shifted.started_at = "2026-08-31T13:00:00.000Z".to_string();
shifted.ended_at = "2026-08-31T13:10:00.000Z".to_string();
assert_eq!(
assert_computer_meter_receipt_replay(&first, &shifted)
.unwrap_err()
.code(),
"computer_meter_receipt_replay_conflict"
);
// A field outside the digest must still match exactly.
let mut restated = first.clone();
restated.provider_allocation.accepted = false;
assert_eq!(
assert_computer_meter_receipt_replay(&first, &restated)
.unwrap_err()
.code(),
"computer_meter_receipt_replay_conflict"
);
}
#[test]
fn an_interval_that_runs_past_admission_expiry_is_refused() {
let case = fixtures()
.into_iter()
.find(|case| case.id == "accepted-standard-8")
.unwrap();
let mut admission_request = case.admission.clone().unwrap();
let observation = case.observation.clone().unwrap();
// Starts before expiry, ends after it.
admission_request.expires_at = "2026-08-31T12:05:00.000Z".to_string();
let admission = bind_computer_admission(admission_request.clone()).unwrap();
assert_eq!(
issue_computer_meter_receipt(&admission, observation.clone())
.unwrap_err()
.code(),
"computer_admission_expired"
);
// Ending exactly at expiry is inside the admission (end is exclusive).
admission_request.expires_at = observation.ended_at.clone();
let admission = bind_computer_admission(admission_request).unwrap();
assert!(issue_computer_meter_receipt(&admission, observation).is_ok());
}
#[test]
fn admission_binds_profile_multiplier_account_and_expiry_before_dispatch() {
let admission = bind_computer_admission(ComputerAdmissionRequest {
+201 -80
View File
@@ -1311,7 +1311,7 @@ struct ScorecardArgs {
#[arg(long, value_name = "FILE")]
baseline: Option<PathBuf>,
/// Regression threshold, in percent increase over the baseline.
#[arg(long, default_value_t = 5.0)]
#[arg(long, default_value_t = 5.0, value_parser = parse_regression_threshold)]
threshold: f64,
/// Emit machine-readable JSON instead of the human summary.
#[arg(long, default_value_t = false)]
@@ -1995,7 +1995,7 @@ fn run_with_args(args: Vec<String>) -> Result<()> {
plugin_registry = Some(discovery.registry_for_workspace(&workspace));
plugin_discovery = Some(discovery);
},
warn_on_workspace_dotenv_result,
|| warn_on_workspace_dotenv_result(&workspace),
);
let plugin_discovery = plugin_discovery
.expect("plugin discovery initialization must precede workspace dotenv loading");
@@ -2961,8 +2961,8 @@ struct WorkspaceDotenvReport {
/// MCP servers, plugin trust, executable lookup, sandbox/approval posture, or
/// network destinations. Shell-exported values and config/CLI arguments remain
/// the explicit surfaces for those controls.
fn warn_on_workspace_dotenv_result() {
match load_workspace_dotenv_credentials() {
fn warn_on_workspace_dotenv_result(workspace: &Path) {
match load_workspace_dotenv_credentials(workspace) {
Ok(Some(report)) if !report.ignored.is_empty() => {
eprintln!(
"Codewhale ignored non-credential settings in {}: {}. Use config.toml, CLI flags, or the launching shell for control settings.",
@@ -3002,21 +3002,31 @@ fn display_env_key_set(keys: &BTreeSet<String>) -> String {
labels.join(", ")
}
fn load_workspace_dotenv_credentials() -> Result<Option<WorkspaceDotenvReport>> {
let Some(path) = find_workspace_dotenv()? else {
fn load_workspace_dotenv_credentials(workspace: &Path) -> Result<Option<WorkspaceDotenvReport>> {
let Some(path) = find_workspace_dotenv(workspace)? else {
return Ok(None);
};
load_workspace_dotenv_credentials_from_path(&path).map(Some)
}
fn find_workspace_dotenv() -> Result<Option<PathBuf>> {
let cwd = std::env::current_dir().context("could not resolve the current workspace")?;
let boundary = cwd
/// The nearest `.env` from the resolved launch workspace (`--workspace`, else
/// the current directory) up to its repository root. Searching from the
/// process directory instead would load another tree's credentials when the
/// two differ.
fn find_workspace_dotenv(workspace: &Path) -> Result<Option<PathBuf>> {
let start = if workspace.is_absolute() {
workspace.to_path_buf()
} else {
std::env::current_dir()
.context("could not resolve the current workspace")?
.join(workspace)
};
let boundary = start
.ancestors()
.find(|ancestor| std::fs::symlink_metadata(ancestor.join(".git")).is_ok())
.unwrap_or(cwd.as_path());
.unwrap_or(start.as_path());
for ancestor in cwd.ancestors() {
for ancestor in start.ancestors() {
let candidate = ancestor.join(".env");
match std::fs::symlink_metadata(&candidate) {
Ok(_) => return Ok(Some(candidate)),
@@ -3345,6 +3355,20 @@ fn run_eval(args: EvalArgs) -> Result<()> {
}
}
/// A regression gate threshold must be a finite percentage: `NaN` compares
/// false against every change (a gate that always passes) and infinity
/// disables the gate outright.
fn parse_regression_threshold(raw: &str) -> Result<f64, String> {
let value: f64 = raw
.trim()
.parse()
.map_err(|_| format!("`{raw}` is not a number"))?;
if !value.is_finite() {
return Err(format!("`{raw}` is not a finite percentage"));
}
Ok(value)
}
/// Score a run's token/cache/cost from recorded turns and (optionally) flag
/// regressions against a committed baseline. Offline: reads recorded usage from
/// a JSON file, reuses the pricing layer, never calls a model. Exits non-zero
@@ -11191,10 +11215,26 @@ fn run_sandbox_command(args: SandboxArgs) -> Result<()> {
.stdout(Stdio::piped())
.stderr(Stdio::piped());
child_env::apply_to_command(&mut cmd, child_env::string_map_env(&exec_env.env));
// Lead a process group so the timeout ends everything the command started.
#[cfg(unix)]
{
use std::os::unix::process::CommandExt as _;
cmd.process_group(0);
}
let mut child = cmd
.spawn()
.map_err(|e| anyhow::anyhow!("Failed to run command: {e}"))?;
// The sandbox run is the tree's lifetime: dropping `tree` on return ends
// anything the command left running, as `contained_output` does.
let tree = match crate::process_tree::ProcessTree::attach(&child) {
Ok(tree) => tree,
Err(error) => {
let _ = child.kill();
let _ = child.wait();
bail!("Failed to contain the sandboxed command: {error}");
}
};
let stdout_handle = child
.stdout
.take()
@@ -11204,48 +11244,75 @@ fn run_sandbox_command(args: SandboxArgs) -> Result<()> {
.take()
.ok_or_else(|| anyhow::anyhow!("stderr unavailable"))?;
let timeout = exec_env.timeout;
let stdout_thread = std::thread::spawn(move || {
// Output streams straight through instead of being buffered whole: a
// command's output size is unbounded, so only a bounded stderr tail is
// kept for sandbox-denial detection.
const STDERR_TAIL_BYTES: usize = 64 * 1024;
let (done_tx, done_rx) = std::sync::mpsc::channel::<()>();
let stdout_done = done_tx.clone();
std::thread::spawn(move || {
let mut reader = stdout_handle;
let mut buf = Vec::new();
let _ = reader.read_to_end(&mut buf);
buf
let _ = io::copy(&mut reader, &mut io::stdout());
let _ = stdout_done.send(());
});
let stderr_thread = std::thread::spawn(move || {
let stderr_tail = Arc::new(std::sync::Mutex::new(Vec::<u8>::new()));
let tail = Arc::clone(&stderr_tail);
std::thread::spawn(move || {
let mut reader = stderr_handle;
let mut buf = Vec::new();
let _ = reader.read_to_end(&mut buf);
buf
let mut chunk = [0_u8; 8192];
loop {
let read = match reader.read(&mut chunk) {
Ok(0) => break,
Ok(read) => read,
Err(error) if error.kind() == io::ErrorKind::Interrupted => continue,
Err(_) => break,
};
let _ = io::stderr().write_all(&chunk[..read]);
if let Ok(mut tail) = tail.lock() {
tail.extend_from_slice(&chunk[..read]);
let excess = tail.len().saturating_sub(STDERR_TAIL_BYTES);
tail.drain(..excess);
}
}
let _ = done_tx.send(());
});
if let Some(status) = child.wait_timeout(timeout)? {
let stdout = stdout_thread.join().unwrap_or_default();
let stderr = stderr_thread.join().unwrap_or_default();
let stderr_str = String::from_utf8_lossy(&stderr);
let exit_code = status.code().unwrap_or(-1);
let sandbox_type = exec_env.sandbox_type;
let sandbox_denied = SandboxManager::was_denied(sandbox_type, exit_code, &stderr_str);
if !stdout.is_empty() {
print!("{}", String::from_utf8_lossy(&stdout));
}
if !stderr.is_empty() {
eprint!("{stderr_str}");
}
if sandbox_denied {
eprintln!(
"{}",
SandboxManager::denial_message(sandbox_type, &stderr_str)
);
}
if !status.success() {
bail!("Command failed with exit code {exit_code}");
}
} else {
let timeout = exec_env.timeout;
let deadline = Instant::now() + timeout;
let Some(status) = child.wait_timeout(timeout)? else {
let _ = tree.kill();
let _ = child.kill();
let _ = child.wait();
bail!("Command timed out after {}ms", timeout.as_millis());
};
// A descendant may still hold the output pipes. Let it finish inside the
// same budget, then end the tree so the drain always reaches EOF.
let mut drained = 0;
while drained < 2 {
match done_rx.recv_timeout(deadline.saturating_duration_since(Instant::now())) {
Ok(()) => drained += 1,
Err(_) => break,
}
}
if drained < 2 {
let _ = tree.kill();
while drained < 2 && done_rx.recv_timeout(Duration::from_secs(1)).is_ok() {
drained += 1;
}
}
let stderr = stderr_tail.lock().map(|tail| tail.clone()).unwrap_or_default();
let stderr_str = String::from_utf8_lossy(&stderr);
let exit_code = status.code().unwrap_or(-1);
let sandbox_type = exec_env.sandbox_type;
if SandboxManager::was_denied(sandbox_type, exit_code, &stderr_str) {
eprintln!(
"{}",
SandboxManager::denial_message(sandbox_type, &stderr_str)
);
}
if !status.success() {
bail!("Command failed with exit code {exit_code}");
}
Ok(())
}
@@ -11600,7 +11667,19 @@ fn preserve_interrupted_checkpoint_for_explicit_resume(launch_workspace: &Path)
/// else falls back to the global value.
#[cfg(test)]
fn merge_project_config(config: &mut Config, workspace: &Path) {
merge_project_config_with_approval_baseline(config, workspace, None);
merge_project_config_with_approval_baseline(config, workspace, None)
.expect("project config applies");
}
/// A project config that exists but cannot be applied is an error, not an
/// absent file: it may be the thing tightening approval, sandbox or shell for
/// this workspace, and launching on the looser user baseline without it would
/// fail open. The reason never quotes file contents.
fn project_config_unusable(path: &Path, reason: &str) -> anyhow::Error {
anyhow!(
"Project config {} could not be applied ({reason}), so its approval, sandbox and shell restrictions are not in effect. Fix the file, or launch with --no-project-config to ignore it.",
path.display()
)
}
/// Apply project config while evaluating approval tightening against the
@@ -11611,7 +11690,7 @@ fn merge_project_config_with_approval_baseline(
config: &mut Config,
workspace: &Path,
saved_permission_posture: Option<&str>,
) {
) -> Result<()> {
// When the workspace is the user's home directory, the project-scope
// config file is also the global config file. Skip the merge to avoid
// redundant processing and a misleading "project-scope config key
@@ -11623,46 +11702,41 @@ fn merge_project_config_with_approval_baseline(
)
&& w == h
{
return;
return Ok(());
}
// v0.8.44: prefer .codewhale/config.toml, fall back to .deepseek/
let path = workspace
let primary = workspace
.join(codewhale_config::CODEWHALE_APP_DIR)
.join("config.toml");
let raw = match read_project_config_file(&path) {
Ok(Some(r)) => r,
let (path, raw) = match read_project_config_file(&primary) {
Ok(Some(raw)) => (primary, raw),
Ok(None) => {
let legacy = workspace
.join(codewhale_config::LEGACY_APP_DIR)
.join("config.toml");
match read_project_config_file(&legacy) {
Ok(Some(r)) => r,
Ok(None) => return,
Err(err) => {
eprintln!(
"warning: failed to read project-scope config {}: {err}",
legacy.display()
);
return;
}
Ok(Some(raw)) => (legacy, raw),
Ok(None) => return Ok(()),
Err(err) => return Err(project_config_unusable(&legacy, &err.to_string())),
}
}
Err(err) => {
eprintln!(
"warning: failed to read project-scope config {}: {err}",
path.display()
);
return;
}
Err(err) => return Err(project_config_unusable(&primary, &err.to_string())),
};
let project: toml::Value = match toml::from_str(&raw) {
Ok(v) => v,
Err(_) => return,
};
let table = match project.as_table() {
Some(t) => t,
None => return,
let project: toml::Value = toml::from_str(&raw).map_err(|err| {
// Position only: the parser's message can quote the offending value.
let reason = err.span().map_or_else(
|| "invalid TOML".to_string(),
|span| {
let prefix = &raw.as_bytes()[..span.start.min(raw.len())];
let line = prefix.iter().filter(|byte| **byte == b'\n').count() + 1;
format!("invalid TOML at line {line}")
},
);
project_config_unusable(&path, &reason)
})?;
let Some(table) = project.as_table() else {
return Err(project_config_unusable(&path, "not a TOML table"));
};
// #417: dangerous keys are denied at project scope. A malicious
@@ -11778,6 +11852,7 @@ fn merge_project_config_with_approval_baseline(
(See #417.)"
);
}
Ok(())
}
/// Maximum bytes read from a project config file. Configs are kilobytes.
@@ -12011,7 +12086,7 @@ async fn run_interactive_with_notice(
&mut merged_config,
&workspace,
saved_permission_posture.as_deref(),
);
)?;
}
if resume_session_id.is_none() {
let explicit_route_override = crate::config::explicit_launch_provider_override().is_some()
@@ -18756,6 +18831,42 @@ api_key = "test-only-key"
}
}
#[test]
fn scorecard_rejects_a_threshold_that_cannot_gate() {
for bad in ["NaN", "inf", "-inf"] {
assert!(
Cli::try_parse_from(["codewhale", "scorecard", "--input", "t.json", "--threshold", bad])
.is_err(),
"--threshold {bad} must be refused"
);
}
assert!(
Cli::try_parse_from(["codewhale", "scorecard", "--input", "t.json", "--threshold", "2.5"])
.is_ok()
);
}
#[test]
fn workspace_dotenv_is_found_from_the_launch_workspace_not_the_process_directory() {
let workspace = tempfile::tempdir().expect("workspace tempdir");
let nested = workspace.path().join("crates/app");
std::fs::create_dir_all(&nested).expect("mkdir nested");
std::fs::create_dir_all(workspace.path().join(".git")).expect("mkdir .git");
std::fs::write(workspace.path().join(".env"), "OPENAI_API_KEY=workspace\n")
.expect("write .env");
// `--workspace <dir>` from a process directory in another tree.
assert_eq!(
find_workspace_dotenv(workspace.path()).expect("search"),
Some(workspace.path().join(".env"))
);
// Nested launch directory walks up to the repository root.
assert_eq!(
find_workspace_dotenv(&nested).expect("search"),
Some(workspace.path().join(".env"))
);
}
#[test]
fn workspace_dotenv_loads_only_provider_credentials_and_preserves_shell_values() {
let _lock = crate::test_support::lock_test_env();
@@ -19994,7 +20105,9 @@ mod project_config_tests {
..Config::default()
};
merge_project_config(&mut config, workspace.path());
let error = merge_project_config_with_approval_baseline(&mut config, workspace.path(), None)
.expect_err("a symlinked primary project config must stop the launch");
assert!(error.to_string().contains("--no-project-config"), "{error}");
assert_eq!(
config.default_text_model.as_deref(),
@@ -20601,14 +20714,22 @@ max_subagents = -3
}
#[test]
fn project_overlay_skips_malformed_toml() {
let tmp = workspace_with_project_config("this is not valid TOML !!");
fn project_overlay_refuses_malformed_toml_instead_of_dropping_its_restrictions() {
let tmp = workspace_with_project_config(
"approval_policy = \"on-request\"\nallow_shell = false\nthis is not valid TOML !!",
);
let mut config = Config {
provider: Some("codewhale".to_string()),
..Config::default()
};
merge_project_config(&mut config, tmp.path());
// Untouched on parse error — better to fall back to global than crash.
// A broken file may be the one tightening this workspace; launching on
// the looser user baseline without it would fail open.
let error = merge_project_config_with_approval_baseline(&mut config, tmp.path(), None)
.expect_err("a malformed project config must stop the launch");
let message = error.to_string();
assert!(message.contains("invalid TOML at line 3"), "{message}");
assert!(message.contains("--no-project-config"), "{message}");
assert!(!message.contains("this is not valid"), "{message}");
assert_eq!(config.provider.as_deref(), Some("codewhale"));
}
+11
View File
@@ -1222,6 +1222,17 @@ pub type RetryCallback = Box<dyn Fn(&LlmError, u32, Duration) + Send + Sync>;
/// * `Ok(T)` - The successful result from the operation
/// * `Err(RetryError)` - All retries exhausted or non-retryable error encountered
///
/// # Known limitation: ambiguous failures are replayed
///
/// A timeout or connection loss after the request was written is retried
/// like a connect failure, so a provider that already accepted the first
/// attempt may bill a second completion. Every caller sends model inference
/// (messages, FIM, translation, speech, provider web search): a replay costs
/// compute but has no external side effect, and the provider APIs used here
/// expose no idempotency key for these requests that could dedupe it. An
/// operation with an external side effect must not be retried through this
/// helper without an idempotency key the server honors.
///
/// # Example
///
/// ```ignore
+45 -8
View File
@@ -826,13 +826,17 @@ impl McpOAuthRuntime {
/// the rotated token while keeping our dead grant would make the next
/// `invalid_grant` compare an "unchanged" store and delete the newer
/// valid credential.
///
/// The comparison and the delete are one step under the secret store's
/// entry lock (the same lock every save takes), so a login that lands
/// between them is never the credential that gets deleted.
async fn clear_stored_tokens(&self, reason: &str) -> Result<()> {
let held = { self.inner.last_tokens.lock().await.take() };
let Some(held) = held else {
return Ok(());
};
match load_oauth_tokens(&self.inner.server_name, &self.inner.url)? {
Some(stored) if stored != held => {
match delete_oauth_tokens_if_held(&self.inner.server_name, &self.inner.url, &held)? {
Some(stored) => {
tracing::debug!(
target: "mcp",
server = %self.inner.server_name,
@@ -845,8 +849,7 @@ impl McpOAuthRuntime {
*self.inner.last_tokens.lock().await = Some(stored);
*self.inner.rejection.lock().await = None;
}
_ => {
delete_oauth_tokens(&self.inner.server_name, &self.inner.url)?;
None => {
*self.inner.rejection.lock().await = Some(reason.to_string());
}
}
@@ -863,8 +866,9 @@ impl McpOAuthRuntime {
};
let Some(credentials) = credentials else {
let mut last = self.inner.last_tokens.lock().await;
if last.take().is_some() {
delete_oauth_tokens(&self.inner.server_name, &self.inner.url)?;
if let Some(previous) = last.take() {
// Only our own credential goes; a newer login stays.
delete_oauth_tokens_if_held(&self.inner.server_name, &self.inner.url, &previous)?;
}
return Ok(());
};
@@ -1444,9 +1448,42 @@ pub(crate) fn load_oauth_tokens(
else {
return Ok(None);
};
let mut tokens = parse_stored_oauth_tokens(&serialized, server_name)?;
decode_stored_oauth_tokens(&serialized, server_name).map(Some)
}
fn decode_stored_oauth_tokens(serialized: &str, server_name: &str) -> Result<StoredMcpOAuthTokens> {
let mut tokens = parse_stored_oauth_tokens(serialized, server_name)?;
refresh_expires_in_from_timestamp(&mut tokens);
Ok(Some(tokens))
Ok(tokens)
}
/// Delete the stored credential only while it is still exactly `held`,
/// comparing and deleting under the store's entry lock. Returns the different
/// credential found in its place (left untouched), or `None` when the entry
/// was `held` and is now gone, or was already absent. An unreadable entry is
/// an error and is left in place, as a plain load would report it.
fn delete_oauth_tokens_if_held(
server_name: &str,
url: &str,
held: &StoredMcpOAuthTokens,
) -> Result<Option<StoredMcpOAuthTokens>> {
let secrets = codewhale_secrets::Secrets::auto_detect();
let key = store_key(server_name, url);
secrets
.with_entry_transaction(&key, |current| {
let Some(serialized) = current.as_deref() else {
return Ok(Ok(None));
};
Ok(match decode_stored_oauth_tokens(serialized, server_name) {
Ok(stored) if stored == *held => {
*current = None;
Ok(None)
}
Ok(stored) => Ok(Some(stored)),
Err(error) => Err(error),
})
})
.with_context(|| format!("clearing the MCP OAuth token for '{server_name}'"))?
}
fn parse_stored_oauth_tokens(serialized: &str, server_name: &str) -> Result<StoredMcpOAuthTokens> {
+58 -18
View File
@@ -28,6 +28,13 @@ pub(super) struct StdioTransport {
pub(super) authority_cancel_watch: Option<tokio::task::JoinHandle<()>>,
/// Holds reviewed executable/script handles for the process lifetime.
pub(super) _reviewed_launch: Option<super::ReviewedStdioLaunch>,
/// Everything the server started: its own process group on Unix, a Job
/// Object on Windows. Termination signals the whole tree, so a server's
/// grandchild (an `npx` wrapper's node, a shell's background job) dies
/// with it on shutdown and on plugin revocation. The last clone's drop
/// kills whatever is left. `None` only for a transport built around an
/// uncontained child (tests).
pub(super) process_tree: Option<Arc<crate::process_tree::ProcessTree>>,
}
/// How long `StdioTransport::shutdown` waits for the child to exit on SIGTERM
@@ -172,6 +179,9 @@ impl StdioTransport {
child_env::string_map_env(&expanded_env),
);
}
// Lead a process group of its own so teardown can reach descendants.
#[cfg(unix)]
cmd.process_group(0);
let mut child = cmd.spawn().map_err(|error| {
let message = if error.kind() == std::io::ErrorKind::NotFound
@@ -195,6 +205,16 @@ impl StdioTransport {
anyhow::Error::new(error).context(message)
})?;
let process_tree = match crate::process_tree::ProcessTree::attach_tokio(&child) {
Ok(tree) => Arc::new(tree),
Err(error) => {
let _ = child.start_kill();
return Err(anyhow::Error::new(error).context(format!(
"MCP server {server_name} could not be contained with its child processes"
)));
}
};
let stdin = child.stdin.take().context("Failed to get MCP stdin")?;
let stdout = child.stdout.take().context("Failed to get MCP stdout")?;
let stderr = child.stderr.take().context("Failed to get MCP stderr")?;
@@ -218,9 +238,10 @@ impl StdioTransport {
let child = Arc::new(TokioMutex::new(child));
let authority_cancel_watch = config.reviewed_plugin.as_ref().map(|_| {
let watched_child = Arc::clone(&child);
let watched_tree = Arc::clone(&process_tree);
tokio::spawn(async move {
cancel_token.cancelled().await;
terminate_child_for_authority_change(&watched_child).await;
terminate_child_for_authority_change(&watched_child, &watched_tree).await;
})
});
@@ -232,6 +253,7 @@ impl StdioTransport {
stderr_tail,
authority_cancel_watch,
_reviewed_launch: reviewed_launch,
process_tree: Some(process_tree),
})
}
}
@@ -315,19 +337,22 @@ async fn format_stderr_context(tail: &StderrTail) -> Option<String> {
))
}
/// Best-effort SIGTERM. On Unix uses `libc::kill`; on Windows there's no
/// equivalent so we let `kill_on_drop` (TerminateProcess) handle it via the
/// subsequent Drop. Returns whether a signal was actually sent.
fn send_sigterm(child: &Child) -> bool {
/// Best-effort SIGTERM. On Unix uses `libc::kill`, addressed to the child's
/// whole process group when it leads one (`contained`); on Windows there's no
/// equivalent so we let `kill_on_drop` (TerminateProcess) and the Job Object
/// handle it. Returns whether a signal was actually sent.
fn send_sigterm(child: &Child, contained: bool) -> bool {
#[cfg(unix)]
{
if let Some(pid) = child.id() {
// SAFETY: pid was just obtained from `child.id()`. `libc::kill`
// with `SIGTERM` is async-signal-safe and never observes invalid
// memory. Worst case (pid wrap / process already gone) returns
// ESRCH, which we deliberately ignore.
let pid = pid as i32;
let target = if contained { -pid } else { pid };
// SAFETY: pid was just obtained from `child.id()` of an unreaped
// child, so neither it nor the group it leads can have been
// recycled. `libc::kill` with `SIGTERM` is async-signal-safe and
// never observes invalid memory. ESRCH is deliberately ignored.
unsafe {
let _ = libc::kill(pid as i32, libc::SIGTERM);
let _ = libc::kill(target, libc::SIGTERM);
}
return true;
}
@@ -335,27 +360,34 @@ fn send_sigterm(child: &Child) -> bool {
}
#[cfg(not(unix))]
{
let _ = child;
let _ = (child, contained);
false
}
}
async fn terminate_child_for_authority_change(child: &Arc<TokioMutex<Child>>) {
async fn terminate_child_for_authority_change(
child: &Arc<TokioMutex<Child>>,
tree: &crate::process_tree::ProcessTree,
) {
let mut child = child.lock().await;
terminate_child(&mut child).await;
terminate_child(&mut child, Some(tree)).await;
}
async fn terminate_child(child: &mut Child) {
async fn terminate_child(child: &mut Child, tree: Option<&crate::process_tree::ProcessTree>) {
// Reap an already-exited child before resolving its PID. Until it is
// reaped, the OS cannot recycle that identity; after it is reaped there is
// nothing left to signal. This avoids a PID-only watcher ever targeting an
// unrelated process after rapid PID reuse.
if child.try_wait().is_ok_and(|status| status.is_some()) {
// The server is gone, but what it started may not be.
if let Some(tree) = tree {
let _ = tree.kill();
}
return;
}
#[cfg(unix)]
send_sigterm(child);
send_sigterm(child, tree.is_some());
#[cfg(not(unix))]
let _ = child.start_kill();
@@ -369,6 +401,10 @@ async fn terminate_child(child: &mut Child) {
let _ = child.wait().await;
}
}
// Descendants that outlived the grace (or ignored SIGTERM) go now.
if let Some(tree) = tree {
let _ = tree.kill();
}
}
#[async_trait::async_trait]
@@ -450,7 +486,7 @@ impl McpTransport for StdioTransport {
/// then force termination and reap the child as the backstop.
async fn shutdown(&mut self) {
let mut child = self.child.lock().await;
terminate_child(&mut child).await;
terminate_child(&mut child, self.process_tree.as_deref()).await;
}
}
@@ -466,18 +502,22 @@ impl Drop for StdioTransport {
if let Ok(runtime) = tokio::runtime::Handle::try_current() {
let child = Arc::clone(&self.child);
let reviewed_launch = self._reviewed_launch.take();
// The task owns the tree so it is not killed before the grace.
let tree = self.process_tree.take();
runtime.spawn(async move {
let _reviewed_launch = reviewed_launch;
let mut child = child.lock().await;
terminate_child(&mut child).await;
terminate_child(&mut child, tree.as_deref()).await;
});
return;
}
if let Ok(mut child) = self.child.try_lock()
&& !child.try_wait().is_ok_and(|status| status.is_some())
{
send_sigterm(&child);
send_sigterm(&child, self.process_tree.is_some());
}
// No runtime: dropping the tree below SIGKILLs the group / closes the
// job, the same backstop `kill_on_drop` gives the direct child.
}
}
+47
View File
@@ -1946,6 +1946,51 @@ async fn plugin_stdio_does_not_surface_reviewed_child_stderr() {
assert!(!error.contains("ARBITRARY_PLUGIN_CREDENTIAL"));
}
/// Shutdown reaches what the server started, not only the server: a
/// background grandchild (an `npx` wrapper's node, a shell job) must not
/// survive the transport.
#[cfg(unix)]
#[tokio::test]
async fn stdio_shutdown_also_terminates_the_servers_grandchildren() {
let mut config = test_server_config();
config.command = Some("sh".to_string());
config.args = vec![
"-c".to_string(),
"sleep 300 & echo $!; wait".to_string(),
];
let mut transport = StdioTransport::spawn(
"grandparent",
"sh",
&config,
tokio_util::sync::CancellationToken::new(),
)
.unwrap();
let line = tokio::time::timeout(Duration::from_secs(5), transport.recv())
.await
.expect("grandchild pid line")
.unwrap();
let grandchild: i32 = String::from_utf8(line).unwrap().trim().parse().unwrap();
transport.shutdown().await;
let deadline = tokio::time::Instant::now() + Duration::from_secs(3);
loop {
// SAFETY: signal 0 only probes whether the pid still exists.
if unsafe { libc::kill(grandchild, 0) } != 0 {
break;
}
if tokio::time::Instant::now() >= deadline {
// Do not leak the sleeper past a failing run.
// SAFETY: plain kill(2) of the pid this test started.
unsafe {
libc::kill(grandchild, libc::SIGKILL);
}
panic!("the MCP server's grandchild outlived the transport");
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
}
/// #6187: a crashed stdio child must stop reading as "ready" before any
/// call is in flight — `is_ready` probes the child, so the pool rebuilds
/// the connection on the next use instead of handing the dead transport
@@ -5120,6 +5165,7 @@ async fn stdio_transport_shutdown_terminates_child() {
stderr_tail: StderrTail::new(),
authority_cancel_watch: None,
_reviewed_launch: None,
process_tree: None,
};
// shutdown() should send SIGTERM and complete within the grace window.
@@ -5249,6 +5295,7 @@ async fn stdio_transport_recv_error_includes_stderr_tail() {
stderr_tail,
authority_cancel_watch: None,
_reviewed_launch: None,
process_tree: None,
};
// Give the subprocess time to write its stderr line and exit.
+179 -12
View File
@@ -46,6 +46,14 @@ pub const ENV_MODELS_DEV_PATH: &str = "CODEWHALE_MODELS_DEV_PATH";
pub const ENV_DISABLE_FETCH: &str = "CODEWHALE_DISABLE_MODELS_DEV_FETCH";
const CACHE_SCHEMA_VERSION: u32 = 1;
/// Largest catalog body accepted from the network, an override file, or the
/// disk cache. The public catalog is a few MiB; anything past this is refused
/// instead of being read whole into memory.
const MAX_CATALOG_BYTES: usize = 32 * 1024 * 1024;
/// Clock skew tolerated on a cache timestamp. A `fetched_at` further in the
/// future than this is not "age zero": it is untrustworthy, so the cache is
/// published as stale and the next refresh fetches.
const MAX_CACHE_CLOCK_SKEW_SECS: u64 = 5 * 60;
/// Provenance / freshness of the Models.dev live layer for UI chips (#4187).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
@@ -227,11 +235,16 @@ pub fn maybe_load_persisted_cache() {
let Some(cache) = load_cache_file(&path) else {
return;
};
let age = now_unix().saturating_sub(cache.fetched_at);
let freshness = if age > DEFAULT_MODELS_DEV_TTL_SECS {
ModelsDevFreshness::Stale
} else {
// Live only when the cache is recent by a trustworthy clock *and* came
// from the source this process would fetch now. A cache written for
// another mirror or override file still publishes (offline startups keep
// rows) but as stale, so the next refresh replaces it.
let freshness = if within_ttl(cache.fetched_at, now_unix())
&& cache.source_fingerprint == current_source_fingerprint()
{
ModelsDevFreshness::Live
} else {
ModelsDevFreshness::Stale
};
if let Err(err) = publish_from_body(
&cache.body,
@@ -269,7 +282,7 @@ pub async fn refresh(force_network: bool) -> Result<usize, ModelsDevRefreshError
if current.freshness == ModelsDevFreshness::Live
&& current
.fetched_at
.is_some_and(|ts| now_unix().saturating_sub(ts) < DEFAULT_MODELS_DEV_TTL_SECS)
.is_some_and(|ts| within_ttl(ts, now_unix()))
{
return Ok(current.offering_count);
}
@@ -302,6 +315,29 @@ pub async fn refresh(force_network: bool) -> Result<usize, ModelsDevRefreshError
Ok(count)
}
/// Whether a payload fetched at `fetched_at` is still inside the TTL at
/// `now`. A timestamp beyond the tolerated clock skew is never fresh: a
/// saturating age would read it as zero and suppress refresh until the clock
/// caught up.
fn within_ttl(fetched_at: u64, now: u64) -> bool {
if fetched_at > now.saturating_add(MAX_CACHE_CLOCK_SKEW_SECS) {
return false;
}
now.saturating_sub(fetched_at) <= DEFAULT_MODELS_DEV_TTL_SECS
}
/// Fingerprint of the source a refresh would read right now: the override
/// file when `CODEWHALE_MODELS_DEV_PATH` is set, else the catalog URL. Same
/// derivation `refresh` stores in the cache.
fn current_source_fingerprint() -> String {
match std::env::var(ENV_MODELS_DEV_PATH) {
Ok(path) if !path.trim().is_empty() => {
base_url_fingerprint(&format!("file:{}", Path::new(path.trim()).display()))
}
_ => base_url_fingerprint(&resolve_catalog_url()),
}
}
/// Best-effort background refresh: never panics, never blocks callers.
pub fn spawn_background_refresh() {
if env_truthy(ENV_DISABLE_FETCH) && std::env::var(ENV_MODELS_DEV_PATH).is_err() {
@@ -328,10 +364,9 @@ pub fn spawn_background_refresh() {
}
async fn refresh_from_path(path: &Path) -> Result<usize, ModelsDevRefreshError> {
let body = match tokio::fs::read_to_string(path).await {
let body = match read_catalog_file(path).await {
Ok(body) => body,
Err(err) => {
let mapped = ModelsDevRefreshError::Io(err.to_string());
Err(mapped) => {
mark_failed(mapped.clone());
return Err(mapped);
}
@@ -375,10 +410,56 @@ async fn fetch_catalog_body(url: &str) -> Result<String, ModelsDevRefreshError>
return Err(ModelsDevRefreshError::HttpStatus(status.as_u16()));
}
response
.text()
read_catalog_response(response, MAX_CATALOG_BYTES).await
}
fn oversized_catalog(limit: usize) -> ModelsDevRefreshError {
ModelsDevRefreshError::InvalidResponse(format!("catalog exceeds {limit} bytes"))
}
/// Read a catalog response body without ever holding more than `limit` bytes.
async fn read_catalog_response(
mut response: reqwest::Response,
limit: usize,
) -> Result<String, ModelsDevRefreshError> {
if response
.content_length()
.is_some_and(|length| length > limit as u64)
{
return Err(oversized_catalog(limit));
}
let mut bytes = Vec::new();
while let Some(chunk) = response
.chunk()
.await
.map_err(|err| ModelsDevRefreshError::Network(err.to_string()))
.map_err(|err| ModelsDevRefreshError::Network(err.to_string()))?
{
if bytes.len().saturating_add(chunk.len()) > limit {
return Err(oversized_catalog(limit));
}
bytes.extend_from_slice(&chunk);
}
String::from_utf8(bytes)
.map_err(|_| ModelsDevRefreshError::InvalidResponse("catalog is not UTF-8".into()))
}
/// Read an override catalog file, refusing one larger than the catalog limit.
async fn read_catalog_file(path: &Path) -> Result<String, ModelsDevRefreshError> {
use tokio::io::AsyncReadExt as _;
let file = tokio::fs::File::open(path)
.await
.map_err(|err| ModelsDevRefreshError::Io(err.to_string()))?;
let mut bytes = Vec::new();
file.take(MAX_CATALOG_BYTES as u64 + 1)
.read_to_end(&mut bytes)
.await
.map_err(|err| ModelsDevRefreshError::Io(err.to_string()))?;
if bytes.len() > MAX_CATALOG_BYTES {
return Err(oversized_catalog(MAX_CATALOG_BYTES));
}
String::from_utf8(bytes)
.map_err(|_| ModelsDevRefreshError::InvalidResponse("catalog is not UTF-8".into()))
}
fn publish_from_body(
@@ -430,7 +511,20 @@ fn mark_failed(err: ModelsDevRefreshError) {
/// with an escaped body) formats. Returns metadata and the *unescaped* body
/// without copying it in the v2 path.
fn load_cache_file(path: &Path) -> Option<PersistedModelsDevCache> {
let bytes = std::fs::read(path).ok()?;
use std::io::Read as _;
// Header line plus body; a larger file is not ours to trust or to read
// whole, so the bundled rows stand instead.
const MAX_CACHE_FILE_BYTES: u64 = MAX_CATALOG_BYTES as u64 + 64 * 1024;
let mut bytes = Vec::new();
std::fs::File::open(path)
.ok()?
.take(MAX_CACHE_FILE_BYTES + 1)
.read_to_end(&mut bytes)
.ok()?;
if bytes.len() as u64 > MAX_CACHE_FILE_BYTES {
return None;
}
// v2: single-line JSON header terminated by a newline, then the verbatim
// catalog body. One small parse, zero body copies.
if bytes.first() == Some(&b'{') && bytes.contains(&b'\n') {
@@ -735,6 +829,79 @@ mod tests {
clear_live_snapshot();
}
#[test]
fn a_future_or_foreign_cache_is_never_live() {
let now = 2_000_000_000;
assert!(within_ttl(now - 60, now));
assert!(within_ttl(now + 30, now), "small skew is tolerated");
assert!(
!within_ttl(now + 3 * 24 * 60 * 60, now),
"a future timestamp is not age zero"
);
assert!(!within_ttl(now - DEFAULT_MODELS_DEV_TTL_SECS - 1, now));
let _lock = lock_test_env();
let _live = lock_live_snapshot();
clear_live_snapshot();
let dir = tempfile::tempdir().expect("tempdir");
let home = dir.path().join("home");
let _home = EnvVarGuard::set("CODEWHALE_HOME", &home);
let _path = EnvVarGuard::remove(ENV_MODELS_DEV_PATH);
let _url = EnvVarGuard::remove(ENV_MODELS_DEV_URL);
let cache_dir = home.join("catalog");
std::fs::create_dir_all(&cache_dir).expect("mkdir");
let cache = cache_dir.join(CACHE_FILE);
// Fresh by the clock, but fetched from another source.
let foreign = PersistedModelsDevCache {
schema_version: CACHE_SCHEMA_VERSION,
fetched_at: now_unix(),
source_fingerprint: base_url_fingerprint("https://mirror.example/catalog.json"),
source_label: "https://mirror.example/catalog.json".into(),
body: FIXTURE.into(),
};
save_cache_file(&cache, &foreign).expect("save");
maybe_load_persisted_cache();
assert_eq!(status().freshness, ModelsDevFreshness::Stale);
// From the current source, but stamped far in the future.
let future = PersistedModelsDevCache {
fetched_at: now_unix() + 30 * 24 * 60 * 60,
source_fingerprint: base_url_fingerprint(MODELS_DEV_CATALOG_URL),
source_label: MODELS_DEV_CATALOG_URL.into(),
..foreign.clone()
};
save_cache_file(&cache, &future).expect("save");
maybe_load_persisted_cache();
assert_eq!(status().freshness, ModelsDevFreshness::Stale);
// The same cache stamped now is live.
let current = PersistedModelsDevCache {
fetched_at: now_unix(),
..future
};
save_cache_file(&cache, &current).expect("save");
maybe_load_persisted_cache();
assert_eq!(status().freshness, ModelsDevFreshness::Live);
clear_live_snapshot();
}
#[test]
fn an_oversized_override_catalog_is_refused() {
let dir = tempfile::tempdir().expect("tempdir");
let path = dir.path().join("huge.json");
let file = std::fs::File::create(&path).expect("create");
file.set_len(MAX_CATALOG_BYTES as u64 + 1).expect("sparse size");
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("runtime");
let error = rt
.block_on(read_catalog_file(&path))
.expect_err("oversized catalog must be refused");
assert!(error.to_string().contains("exceeds"), "{error}");
}
#[test]
fn network_failure_keeps_prior_rows_and_marks_failed() {
let _lock = lock_test_env();
+86 -29
View File
@@ -165,6 +165,10 @@ static MERGED_CACHE: RwLock<Option<MergedCacheEntry>> = RwLock::new(None);
/// 600+ row OpenRouter catalog for every route candidate.
static RUNTIME_RESOLVER_CACHE: RwLock<BTreeMap<String, RuntimeResolverCacheEntry>> =
RwLock::new(BTreeMap::new());
/// Most distinct endpoints the resolver cache holds at once. Entries from an
/// older catalog generation can never hit again and are dropped on insert;
/// this caps the endpoints of the current generation.
const MAX_RUNTIME_RESOLVER_CACHE_ENTRIES: usize = 64;
#[derive(Clone)]
struct RuntimeResolverCacheEntry {
@@ -408,27 +412,61 @@ pub fn live_catalog_origin(provider: ApiProvider, wire_model_id: &str) -> Option
let Ok(guard) = LIVE_SNAPSHOT.read() else {
return None;
};
let matches = |row: &CatalogOffering| {
row.provider.eq_ignore_ascii_case(catalog_id)
&& row.wire_model_id.eq_ignore_ascii_case(needle)
};
if guard
.per_provider
.get(&owner)
.is_some_and(|snap| snap.offerings.iter().any(matches))
{
return Some(LiveSource::PerProvider);
}
if guard
.models_dev
.as_ref()
.is_some_and(|snap| snap.offerings.iter().any(|row| matches(row)))
{
return Some(LiveSource::ModelsDev);
// Wire ids are opaque and case-sensitive: an exact row in either
// partition decides before a case-folded one does.
for exact in [true, false] {
let matches = |row: &CatalogOffering| {
row.provider.eq_ignore_ascii_case(catalog_id)
&& if exact {
row.wire_model_id == needle
} else {
row.wire_model_id.eq_ignore_ascii_case(needle)
}
};
if guard
.per_provider
.get(&owner)
.is_some_and(|snap| snap.offerings.iter().any(matches))
{
return Some(LiveSource::PerProvider);
}
if guard
.models_dev
.as_ref()
.is_some_and(|snap| snap.offerings.iter().any(|row| matches(row)))
{
return Some(LiveSource::ModelsDev);
}
}
None
}
/// The row whose wire id is exactly `needle`, else the one wire id equal to
/// it ignoring ASCII case. Wire ids are opaque and case-sensitive, so the
/// fallback only forgives a case slip when it cannot choose between two
/// distinct ids; otherwise it answers nothing rather than another model's
/// metadata.
fn find_wire_model<'a>(
rows: impl IntoIterator<Item = &'a CatalogOffering>,
needle: &str,
) -> Option<&'a CatalogOffering> {
let mut folded: Option<&'a CatalogOffering> = None;
let mut ambiguous = false;
for row in rows {
if row.wire_model_id == needle {
return Some(row);
}
if row.wire_model_id.eq_ignore_ascii_case(needle) {
match folded {
None => folded = Some(row),
Some(first) if first.wire_model_id != row.wire_model_id => ambiguous = true,
Some(_) => {}
}
}
}
folded.filter(|_| !ambiguous)
}
/// Serialize tests that mutate the process-wide live snapshot.
///
/// Lock ordering: this takes the test env barrier FIRST (skipped when the
@@ -938,6 +976,12 @@ pub(crate) fn runtime_catalog_resolver_for_identity(
let resolver = RouteResolver::from_offerings(route_offerings.into_values().collect());
if let Ok(mut cache) = RUNTIME_RESOLVER_CACHE.write() {
cache.retain(|_, entry| {
entry.generation == generation && entry.cloud_generation == cloud_generation
});
if cache.len() >= MAX_RUNTIME_RESOLVER_CACHE_ENTRIES && !cache.contains_key(&cache_key) {
cache.pop_first();
}
cache.insert(
cache_key,
RuntimeResolverCacheEntry {
@@ -1110,14 +1154,14 @@ pub fn catalog_offering_for_model_identity(
return None;
}
if provider == ApiProvider::Custom {
return exact_custom_offerings(catalog_id.as_ref())
.into_iter()
.find(|row| row.wire_model_id.eq_ignore_ascii_case(needle));
let rows = exact_custom_offerings(catalog_id.as_ref());
return find_wire_model(&rows, needle).cloned();
}
offerings_for_provider_identity(&merged_snapshot(), catalog_id.as_ref())
.into_iter()
.find(|row| row.wire_model_id.eq_ignore_ascii_case(needle))
.cloned()
find_wire_model(
offerings_for_provider_identity(&merged_snapshot(), catalog_id.as_ref()),
needle,
)
.cloned()
}
/// Metadata from the exact route, without borrowing another endpoint's live facts.
@@ -1252,11 +1296,7 @@ pub fn bundled_catalog_offering_for_model(
if needle.is_empty() {
return None;
}
bundled_snapshot()
.offerings_for_provider(catalog_id)
.into_iter()
.find(|row| row.wire_model_id.eq_ignore_ascii_case(needle))
.cloned()
find_wire_model(bundled_snapshot().offerings_for_provider(catalog_id), needle).cloned()
}
/// Count of merged-catalog models for one provider (catalog view / dashboard).
@@ -1860,6 +1900,23 @@ mod tests {
use crate::config::{DEFAULT_TOGETHER_FLASH_MODEL, DEFAULT_TOGETHER_MODEL};
use codewhale_config::catalog::CatalogSource;
#[test]
fn wire_model_lookup_is_exact_first_and_never_guesses_between_case_variants() {
let row = |id: &str| CatalogOffering {
provider: "p".to_string(),
wire_model_id: id.to_string(),
..Default::default()
};
let rows = vec![row("org/Model-A"), row("org/model-a"), row("org/Model-B")];
let found = |needle: &str| {
find_wire_model(&rows, needle).map(|row| row.wire_model_id.clone())
};
assert_eq!(found("org/model-a").as_deref(), Some("org/model-a"));
assert_eq!(found("org/Model-A").as_deref(), Some("org/Model-A"));
assert_eq!(found("ORG/MODEL-A"), None, "two distinct ids fold together");
assert_eq!(found("org/model-b").as_deref(), Some("org/Model-B"));
}
fn catalog_test_config(first_url: &str, second_url: &str) -> Config {
use crate::config::{ProviderConfig, ProvidersConfig};
Config {
+60 -5
View File
@@ -295,6 +295,7 @@ fn lighthouse_plan(inputs: &DeployInputs) -> Vec<ProvisionStep> {
fn azure_plan(inputs: &DeployInputs) -> Vec<ProvisionStep> {
let rg = format!("{}-rg", inputs.instance_name);
let vault = format!("{}-kv", inputs.instance_name);
let nsg = format!("{}-nsg", inputs.instance_name);
let provider_secret = format!("codewhale-{}-key", inputs.provider_slug);
vec![
ProvisionStep::new(
@@ -355,20 +356,40 @@ fn azure_plan(inputs: &DeployInputs) -> Vec<ProvisionStep> {
"--custom-data",
"cloud-init.yaml",
"--assign-identity",
// A Linux `vm create` otherwise adds an NSG rule admitting SSH
// from any address; the next step adds the caller-only rule.
"--nsg",
&nsg,
"--nsg-rule",
"NONE",
],
),
ProvisionStep::new(
"Scope the NSG to SSH (22) from the caller IP only; 7878 stays on 127.0.0.1",
"Allow SSH (22) from the caller IP only (replace <your-public-ip>); 7878 stays on 127.0.0.1",
"az",
&[
"vm",
"open-port",
"network",
"nsg",
"rule",
"create",
"--resource-group",
&rg,
"--nsg-name",
&nsg,
"--name",
&inputs.instance_name,
"--port",
"allow-ssh-from-caller",
"--priority",
"1000",
"--direction",
"Inbound",
"--access",
"Allow",
"--protocol",
"Tcp",
"--destination-port-ranges",
"22",
"--source-address-prefixes",
"<your-public-ip>/32",
],
),
]
@@ -540,6 +561,40 @@ mod tests {
);
}
#[test]
fn azure_plan_admits_ssh_from_the_caller_only() {
let steps = (AZURE.plan)(&DeployInputs::default());
let arg_after = |step: &ProvisionStep, flag: &str| {
let at = step.args.iter().position(|arg| arg == flag)?;
step.args.get(at + 1).cloned()
};
// `open-port` has no source restriction: it admits the whole internet.
assert!(
steps.iter().all(|s| !s.args.iter().any(|a| a == "open-port")),
"Azure plan must not open a port to every source"
);
let create = steps
.iter()
.find(|s| s.args.starts_with(&["vm".to_string(), "create".to_string()]))
.expect("vm create step");
assert_eq!(arg_after(create, "--nsg-rule").as_deref(), Some("NONE"));
let ssh_rules: Vec<_> = steps
.iter()
.filter(|s| arg_after(s, "--destination-port-ranges").as_deref() == Some("22"))
.collect();
assert_eq!(ssh_rules.len(), 1, "exactly one SSH rule");
let source = arg_after(ssh_rules[0], "--source-address-prefixes").unwrap();
assert!(
source.ends_with("/32") && !source.starts_with('*') && !source.starts_with("0.0.0.0"),
"SSH source must be a single caller address, got {source}"
);
assert_eq!(
arg_after(ssh_rules[0], "--nsg-name"),
arg_after(create, "--nsg"),
"the rule must land on the NSG the VM was created with"
);
}
#[test]
fn display_command_redacts_secret_args() {
let mut step =
+42 -1
View File
@@ -355,7 +355,14 @@ fn classify(
) -> BillingPresentation {
match provider {
ApiProvider::Ollama | ApiProvider::Sglang | ApiProvider::Vllm => BillingPresentation::Local,
ApiProvider::OpenaiCodex => BillingPresentation::Subscription("Codex OAuth quota"),
// The Codex OAuth quota is a fact of the ChatGPT backend, not of the
// provider name: a custom `[providers.openai_codex] base_url` (a
// gateway, a proxy) sells its own terms, so it is Unknown rather than
// a subscription that would hide metered spend.
ApiProvider::OpenaiCodex if is_chatgpt_codex_backend(base_url) => {
BillingPresentation::Subscription("Codex OAuth quota")
}
ApiProvider::OpenaiCodex => BillingPresentation::Unknown,
ApiProvider::OpencodeGo => BillingPresentation::Subscription("OpenCode Go quota"),
// StepFun already reduces an endpoint to a non-secret billing surface
// and fails closed on anything it does not recognize.
@@ -608,6 +615,20 @@ fn stepfun_billing_for_endpoint(base_url: Option<&str>) -> BillingPresentation {
}
}
/// The ChatGPT backend the Codex OAuth route ships with
/// (`https://chatgpt.com/backend-api`, or a path under it).
fn is_chatgpt_codex_backend(base_url: &str) -> bool {
let Ok(url) = reqwest::Url::parse(base_url.trim()) else {
return false;
};
let path = url.path().trim_end_matches('/');
url.scheme() == "https"
&& url.host_str() == Some("chatgpt.com")
&& url.port().is_none()
&& url.username().is_empty()
&& (path == "/backend-api" || path.starts_with("/backend-api/"))
}
fn is_zai_coding_plan_endpoint(base_url: &str) -> bool {
base_url
.trim()
@@ -2401,6 +2422,26 @@ mod tests {
),
BillingPresentation::Subscription("Codex OAuth quota")
);
// The OAuth token pointed anywhere else proves no Codex quota.
for elsewhere in [
"https://codex-gateway.example.com/backend-api",
"https://chatgpt.com.example.net/backend-api",
"http://chatgpt.com/backend-api",
"https://chatgpt.com/v1",
"",
] {
assert_eq!(
for_dispatched_route(
&config,
DispatchedRoute {
provider: ApiProvider::OpenaiCodex,
base_url: elsewhere,
},
),
BillingPresentation::Unknown,
"{elsewhere:?}"
);
}
}
#[test]
+9 -1
View File
@@ -105,10 +105,18 @@ pub(super) struct WorkspaceGitMetadata {
pub(super) dirty: bool,
}
/// Status is several blocking `git` spawns, so it runs off the async workers
/// (#6149) like every other git read on this surface. It deliberately uses
/// normal Git (filters, untracked settings) rather than the review command,
/// so its counts match `git status` and the operator writes (see `git.rs`).
pub(super) async fn workspace_status(
State(state): State<RuntimeApiState>,
) -> Result<Json<WorkspaceStatusResponse>, ApiError> {
Ok(Json(collect_workspace_status(&state.workspace)))
let workspace = state.workspace.clone();
tokio::task::spawn_blocking(move || collect_workspace_status(&workspace))
.await
.map(Json)
.map_err(|_| ApiError::internal("workspace status failed"))
}
pub(super) fn collect_workspace_status(workspace: &FsPath) -> WorkspaceStatusResponse {
+59 -2
View File
@@ -648,12 +648,22 @@ impl ScorecardMetrics {
/// Flag metrics that grew more than `threshold_pct` over `baseline`. Cost
/// and token counts are "lower is better", so only *increases* are
/// regressions. (Cache-hit ratio is the opposite, reported separately.)
///
/// Fails closed on values that cannot be compared: a non-finite threshold
/// admits no growth at all (it is treated as 0%), and a metric whose
/// change is `NaN` counts as regressed. Every `NaN` comparison is false,
/// so without this a `NaN` threshold would pass any run.
#[must_use]
pub fn regressions_against(
&self,
baseline: &ScorecardMetrics,
threshold_pct: f64,
) -> Vec<Regression> {
let threshold_pct = if threshold_pct.is_finite() {
threshold_pct
} else {
0.0
};
let mut out = Vec::new();
// A partial/unknown subtotal is not comparable to a complete baseline,
// but losing completeness is itself a regression. Otherwise removing
@@ -711,7 +721,7 @@ impl ScorecardMetrics {
let drop_pct = (baseline.cache_hit_ratio - self.cache_hit_ratio)
/ baseline.cache_hit_ratio
* 100.0;
if drop_pct > threshold_pct {
if drop_pct.is_nan() || drop_pct > threshold_pct {
out.push(Regression {
metric: "cache_hit_ratio_drop".to_string(),
baseline: baseline.cache_hit_ratio,
@@ -733,7 +743,7 @@ fn push_regression(
) {
if base > 0.0 {
let pct = (cur - base) / base * 100.0;
if pct > threshold_pct {
if pct.is_nan() || pct > threshold_pct {
out.push(Regression {
metric: metric.to_string(),
baseline: base,
@@ -1908,6 +1918,53 @@ mod tests {
assert!(!names.contains(&"total_input_tokens")); // under threshold
}
#[test]
fn a_non_finite_threshold_cannot_pass_a_regressed_run() {
let baseline = ScorecardMetrics {
turns: 1,
money_metered_turns: 1,
unpriced_turns: 0,
cny_unpriced_turns: 0,
cost_complete: true,
cny_cost_complete: true,
unpriced_classes: Vec::new(),
total_input_tokens: 1000,
total_output_tokens: 1000,
total_cache_read_tokens: 0,
total_cache_write_tokens: 0,
total_reasoning_tokens: 0,
total_cost_usd: 0.10,
total_cost_cny: 0.7,
cache_hit_ratio: 0.5,
};
let current = ScorecardMetrics {
total_output_tokens: 2000,
..baseline.clone()
};
for threshold in [f64::NAN, f64::INFINITY] {
let names: Vec<String> = current
.regressions_against(&baseline, threshold)
.into_iter()
.map(|r| r.metric)
.collect();
assert!(
names.iter().any(|name| name == "total_output_tokens"),
"threshold {threshold} passed a +100% regression: {names:?}"
);
}
// A metric whose change cannot be computed is not a pass either.
let unknown = ScorecardMetrics {
cache_hit_ratio: f64::NAN,
..baseline.clone()
};
assert!(
unknown
.regressions_against(&baseline, 5.0)
.iter()
.any(|r| r.metric == "cache_hit_ratio_drop")
);
}
#[test]
fn regression_flags_loss_of_cost_completeness_without_comparing_subtotals() {
let baseline = ScorecardMetrics {
+16 -7
View File
@@ -32,14 +32,15 @@ pub fn tool_output_status(messages: &[Message], artifacts: &[ArtifactRecord]) ->
for message in messages {
for block in &message.content {
if let ContentBlock::ToolResult { content, .. } = block {
if looks_like_receipt(content) {
// The prefix is text any tool output can carry, so it never
// exempts a result from the size check: a "receipt" above the
// threshold is raw pressure like any other large result.
let chars = content.chars().count();
if chars > RAW_TOOL_OUTPUT_RECEIPT_THRESHOLD_CHARS {
status.raw_large_count += 1;
status.raw_large_chars = status.raw_large_chars.saturating_add(chars);
} else if looks_like_receipt(content) {
status.receipt_count += 1;
} else {
let chars = content.chars().count();
if chars > RAW_TOOL_OUTPUT_RECEIPT_THRESHOLD_CHARS {
status.raw_large_count += 1;
status.raw_large_chars = status.raw_large_chars.saturating_add(chars);
}
}
}
}
@@ -151,4 +152,12 @@ mod tests {
assert!(rendered.contains("compact receipt"));
assert!(rendered.contains("artifact"));
}
#[test]
fn a_receipt_prefix_does_not_hide_a_large_raw_result() {
let forged = format!("[TOOL_OUTPUT_RECEIPT]\n{}", "RAW\n".repeat(4_000));
let status = tool_output_status(&[tool_result_message("call-forged", &forged)], &[]);
assert_eq!(status.raw_large_count, 1);
assert_eq!(status.receipt_count, 0);
}
}
+43 -8
View File
@@ -112,15 +112,16 @@ fn reasoning_effort_for_route_selection(
}
}
/// Settle the classifier usage of a turn whose route then failed, into the
/// cost scope that was current when the classifier call was made. A fresh
/// token taken here would bill it to whatever session is current by now:
/// after a `/new` or session load during the call, a different one.
fn settle_failed_parent_route(
cost_scope: crate::cost_status::CostScopeToken,
error: String,
initial_routed_usage: &crate::cost_status::RuntimeUsageBatch,
) -> String {
crate::cost_status::report_runtime_usage_batch(
crate::cost_status::scope_token(),
None,
initial_routed_usage,
);
crate::cost_status::report_runtime_usage_batch(cost_scope, None, initial_routed_usage);
error
}
@@ -141,6 +142,9 @@ fn settle_failed_parent_route(
pub(crate) async fn plan_turn_route(
request: TurnRoutePlanRequest<'_>,
) -> Result<PlannedTurnRoute, String> {
// Taken before the classifier call so its usage settles into the scope
// it was spent in, not the one current when a failure is noticed.
let cost_scope = crate::cost_status::scope_token();
let mut auto_selection = if request.should_auto_resolve {
Some(
crate::model_routing::resolve_auto_route_with_inventory_for_session_and_cache_policy(
@@ -210,6 +214,7 @@ pub(crate) async fn plan_turn_route(
Ok(route) => route,
Err(err) => {
return Err(settle_failed_parent_route(
cost_scope,
err.to_string(),
&initial_routed_usage,
));
@@ -219,7 +224,11 @@ pub(crate) async fn plan_turn_route(
match turn_route.preflight() {
Ok(route) => route,
Err(err) => {
return Err(settle_failed_parent_route(err, &initial_routed_usage));
return Err(settle_failed_parent_route(
cost_scope,
err,
&initial_routed_usage,
));
}
}
} else {
@@ -357,17 +366,43 @@ mod tests {
dropped_records: 1,
};
let scope = crate::cost_status::scope_token();
assert_eq!(
settle_failed_parent_route("route failed".to_string(), &batch),
settle_failed_parent_route(scope, "route failed".to_string(), &batch),
"route failed"
);
settle_failed_parent_route("route failed".to_string(), &batch);
settle_failed_parent_route(scope, "route failed".to_string(), &batch);
let pending = crate::cost_status::drain();
assert_eq!(
pending.usage_source_fingerprints.len(),
2,
"both exact classifier outcomes persist, and replay is idempotent"
);
// The classifier ran in one session; `/new` closed it before the
// route failed. Its cost belongs to the closed session, never the new.
let spent_in = crate::cost_status::scope_token();
let _closed = crate::cost_status::close_current_scope();
let batch = crate::cost_status::RuntimeUsageBatch {
records: batch
.records
.iter()
.cloned()
.map(|mut record| {
record.source_id = "auto-router:plan-usage-after-new".to_string();
record
})
.collect(),
drop_records: Vec::new(),
dropped_records: 0,
};
settle_failed_parent_route(spent_in, "route failed".to_string(), &batch);
assert!(
crate::cost_status::drain()
.usage_source_fingerprints
.is_empty(),
"classifier cost leaked into the session opened after it was spent"
);
}
#[test]