mirror of
https://github.com/codewhale-hq/Codewhale.git
synced 2026-11-08 08:04:25 -10:00
wip: runtime audit slice (uncompiled; to be split)
This commit is contained in:
1 parent
8b6e87bc35
commit
6633405673
25 files changed
+1561
-395
No files matched your search
@@ -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(_))
|
||||
));
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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
@@ -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:?}"
|
||||
|
||||
@@ -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"
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
@@ -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) => {}
|
||||
|
||||
@@ -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);
|
||||
};
|
||||
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
@@ -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"));
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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.
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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, ¤t).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();
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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 =
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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]
|
||||
|
||||
Reference in new issue
Block a user