use std::collections::HashMap;
use std::time::Duration;
use async_trait::async_trait;
use futures::StreamExt;
use serde::{Deserialize, Serialize};
use supercode_interchange::{ChatMessage, FunctionCall, Role, ToolCall};
use crate::{CachePlan, ChatRequest, Result, RuntimeError as Error, ToolSchema, Usage};
const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
const READ_IDLE_TIMEOUT: Duration = Duration::from_secs(120);
const MAX_RETRIES: u32 = 2;
const RETRY_BACKOFF_BASE: Duration = Duration::from_millis(500);
#[derive(Debug, Clone, Copy)]
#[doc(hidden)]
pub struct HttpOptions {
pub(crate) connect_timeout: Duration,
pub(crate) read_idle_timeout: Duration,
pub(crate) max_retries: u32,
pub(crate) retry_backoff_base: Duration,
}
impl Default for HttpOptions {
fn default() -> Self {
HttpOptions {
connect_timeout: CONNECT_TIMEOUT,
read_idle_timeout: READ_IDLE_TIMEOUT,
max_retries: MAX_RETRIES,
retry_backoff_base: RETRY_BACKOFF_BASE,
}
}
}
impl HttpOptions {
#[doc(hidden)]
pub fn from_retry_config(
enabled: bool,
max_retries: Option<u32>,
base_delay_ms: Option<u64>,
) -> HttpOptions {
let base = HttpOptions::default();
HttpOptions {
max_retries: if enabled {
max_retries.unwrap_or(base.max_retries)
} else {
0
},
retry_backoff_base: base_delay_ms
.map(Duration::from_millis)
.unwrap_or(base.retry_backoff_base),
..base
}
}
}
pub(crate) fn build_request_body(req: &ChatRequest, stream: bool) -> serde_json::Value {
use serde_json::json;
let mut body = json!({
"model": req.model,
"messages": req.messages,
"stream": stream,
});
let obj = body.as_object_mut().unwrap();
if !req.tools.is_empty() {
obj.insert(
"tools".into(),
serde_json::to_value(req.tools.iter().map(WireTool::from).collect::<Vec<_>>()).unwrap(),
);
}
if let Some(t) = req.temperature {
obj.insert("temperature".into(), json!(t));
}
if let Some(m) = req.max_tokens {
obj.insert("max_tokens".into(), json!(m));
}
if let Some(e) = &req.effort {
obj.insert("reasoning_effort".into(), json!(e));
}
if let Some(rf) = &req.response_format {
obj.insert("response_format".into(), rf.clone());
}
if stream {
obj.insert("stream_options".into(), json!({"include_usage": true}));
}
for (k, v) in &req.extra_body {
obj.insert(k.clone(), v.clone());
}
body
}
#[doc(hidden)]
pub fn apply_cache_plan(
messages: &[ChatMessage],
plan: CachePlan,
imported_prefix_len: Option<usize>,
) -> Vec<ChatMessage> {
let mut out = messages.to_vec();
if !matches!(plan, CachePlan::ImportedPrefix) {
return out;
}
let Some(len) = imported_prefix_len.filter(|&n| n > 0) else {
return out;
};
let last = len - 1;
let mut targets = vec![0usize];
if last != 0 {
targets.push(last);
}
for idx in targets {
if let Some(msg) = out.get_mut(idx) {
annotate_cache_breakpoint(msg);
}
}
out
}
fn annotate_cache_breakpoint(msg: &mut ChatMessage) {
let cache_control = serde_json::json!({"type": "ephemeral"});
if let Some(parts) = msg.content_parts.as_mut() {
if let Some(text_part) = parts
.iter_mut()
.rev()
.find(|p| p.get("type").and_then(serde_json::Value::as_str) == Some("text"))
{
if let Some(obj) = text_part.as_object_mut() {
obj.insert("cache_control".to_string(), cache_control);
}
}
return;
}
let text = msg.content.take().unwrap_or_default();
msg.content_parts = Some(vec![serde_json::json!({
"type": "text",
"text": text,
"cache_control": cache_control,
})]);
}
#[doc(hidden)]
pub fn tier_change_is_cache_bust(previous: Option<u64>, current: u64) -> bool {
previous.is_some_and(|p| p != current)
}
pub(crate) const CACHE_TTL_SECS: i64 = 300;
pub(crate) const CACHE_MISS_RATIO_THRESHOLD: f64 = 0.10;
pub(crate) const CACHE_STALE_DISPROVE_RATIO_THRESHOLD: f64 = 1.0 - CACHE_MISS_RATIO_THRESHOLD;
#[doc(hidden)]
pub fn is_anthropic_family_model(model: &str) -> bool {
model.starts_with("anthropic/") || model.starts_with("claude-") || model.starts_with("claude/")
}
#[derive(Debug, Clone, Copy, PartialEq)]
#[doc(hidden)]
pub enum CacheColdReason {
Stale {
idle_secs: i64,
},
Miss {
cached_tokens: u64,
prompt_tokens: u64,
},
}
impl CacheColdReason {
#[doc(hidden)]
pub fn message(&self) -> String {
match self {
CacheColdReason::Stale { idle_secs } => format!(
"cache likely cold — this turn was sent {}m{:02}s after the cache was last \
refreshed (Anthropic's ephemeral prompt cache expires after 5m idle) — this \
turn likely paid full input cost for the cached prefix",
idle_secs / 60,
idle_secs % 60,
),
CacheColdReason::Miss {
cached_tokens,
prompt_tokens,
} => format!(
"unexpected cache miss — only {cached_tokens}/{prompt_tokens} prompt tokens \
were served from cache this turn even though reuse was expected — this turn \
likely paid full input cost for the cached prefix",
),
}
}
}
#[doc(hidden)]
pub fn cache_cold_reason(
will_annotate: bool,
cache_established: bool,
idle_secs: Option<i64>,
usage: &Usage,
) -> Option<CacheColdReason> {
if !will_annotate {
return None;
}
if let Some(idle_secs) = idle_secs {
if idle_secs >= CACHE_TTL_SECS {
let disproven_by_usage = usage
.prompt_tokens_details
.filter(|_| usage.prompt_tokens > 0)
.is_some_and(|details| {
details.cached_tokens as f64 / usage.prompt_tokens as f64
>= CACHE_STALE_DISPROVE_RATIO_THRESHOLD
});
if !disproven_by_usage {
return Some(CacheColdReason::Stale { idle_secs });
}
}
}
if !cache_established {
return None;
}
let details = usage.prompt_tokens_details?;
if usage.prompt_tokens == 0 {
return None;
}
let ratio = details.cached_tokens as f64 / usage.prompt_tokens as f64;
if ratio < CACHE_MISS_RATIO_THRESHOLD {
return Some(CacheColdReason::Miss {
cached_tokens: details.cached_tokens,
prompt_tokens: usage.prompt_tokens,
});
}
None
}
#[async_trait]
pub trait Provider: Send + Sync {
async fn complete(
&self,
req: &ChatRequest,
on_delta: &(dyn for<'a> Fn(&'a str) + Send + Sync),
) -> Result<(ChatMessage, Usage)>;
}
pub struct OpenAiProvider {
client: reqwest::Client,
base_url: String,
api_key: String,
extra_headers: HashMap<String, String>,
http_options: HttpOptions,
}
impl OpenAiProvider {
pub fn new(
base_url: impl Into<String>,
api_key: impl Into<String>,
extra_headers: HashMap<String, String>,
) -> Self {
Self::new_with_options(base_url, api_key, extra_headers, HttpOptions::default())
}
#[doc(hidden)]
pub fn new_with_options(
base_url: impl Into<String>,
api_key: impl Into<String>,
extra_headers: HashMap<String, String>,
http_options: HttpOptions,
) -> Self {
OpenAiProvider {
client: reqwest::Client::builder()
.connect_timeout(http_options.connect_timeout)
.read_timeout(http_options.read_idle_timeout)
.build()
.expect("static reqwest client config cannot fail"),
base_url: base_url.into(),
api_key: api_key.into(),
extra_headers,
http_options,
}
}
fn endpoint(&self) -> String {
format!("{}/chat/completions", self.base_url.trim_end_matches('/'))
}
async fn send_with_retry(&self, wire: &serde_json::Value) -> Result<reqwest::Response> {
let mut attempt = 0u32;
loop {
let mut builder = self
.client
.post(self.endpoint())
.bearer_auth(&self.api_key)
.header("Content-Type", "application/json");
for (k, v) in &self.extra_headers {
builder = builder.header(k, v);
}
let sent = builder.json(wire).send().await;
let (retryable, result): (bool, Result<reqwest::Response>) = match sent {
Err(e) => (true, Err(Error::from(e))),
Ok(resp) => {
let status = resp.status();
if status.is_success() {
(false, Ok(resp))
} else if status.is_server_error() {
let body = resp.text().await.unwrap_or_default();
(
true,
Err(Error::Provider {
status: status.as_u16(),
body: truncate(&body, 2000),
}),
)
} else {
let body = resp.text().await.unwrap_or_default();
(
false,
Err(Error::Provider {
status: status.as_u16(),
body: truncate(&body, 2000),
}),
)
}
}
};
if !retryable || attempt >= self.http_options.max_retries {
return result;
}
let backoff = self.http_options.retry_backoff_base * 2u32.pow(attempt);
tokio::time::sleep(backoff).await;
attempt += 1;
}
}
}
#[async_trait]
impl Provider for OpenAiProvider {
async fn complete(
&self,
req: &ChatRequest,
on_delta: &(dyn for<'a> Fn(&'a str) + Send + Sync),
) -> Result<(ChatMessage, Usage)> {
let wire = build_request_body(req, true);
let resp = self.send_with_retry(&wire).await?;
let mut acc = Accumulator::default();
let mut buf: Vec<u8> = Vec::new();
let mut deltas: Vec<String> = Vec::new();
let mut stream = resp.bytes_stream();
while let Some(chunk) = stream.next().await {
let bytes = chunk?;
buf.extend_from_slice(&bytes);
drain_sse_lines(&mut buf, &mut acc, &mut deltas)?;
for d in deltas.drain(..) {
on_delta(&d);
}
}
let tail = String::from_utf8_lossy(&buf);
if !tail.trim().is_empty() {
handle_sse_line(tail.trim(), &mut acc, &mut deltas)?;
for d in deltas.drain(..) {
on_delta(&d);
}
}
Ok((acc.to_message(), acc_usage(&acc)))
}
}
#[derive(Default)]
struct Accumulator {
content: String,
tool_calls: Vec<ToolCallAccum>,
usage: Usage,
}
#[derive(Default)]
struct ToolCallAccum {
id: String,
name: String,
arguments: String,
}
impl Accumulator {
fn ensure(&mut self, index: usize) -> &mut ToolCallAccum {
while self.tool_calls.len() <= index {
self.tool_calls.push(ToolCallAccum::default());
}
&mut self.tool_calls[index]
}
fn to_message(&self) -> ChatMessage {
let calls: Vec<ToolCall> = self
.tool_calls
.iter()
.filter(|c| !c.id.is_empty() || !c.name.is_empty())
.map(|c| ToolCall {
id: c.id.clone(),
kind: "function".to_string(),
function: FunctionCall {
name: c.name.clone(),
arguments: c.arguments.clone(),
},
})
.collect();
ChatMessage {
role: Role::Assistant,
content: (!self.content.is_empty()).then(|| self.content.clone()),
content_parts: None,
tool_calls: (!calls.is_empty()).then_some(calls),
tool_call_id: None,
name: None,
metadata: Default::default(),
}
}
}
fn acc_usage(acc: &Accumulator) -> Usage {
acc.usage.clone()
}
fn drain_sse_lines(
buf: &mut Vec<u8>,
acc: &mut Accumulator,
deltas: &mut Vec<String>,
) -> Result<()> {
while let Some(pos) = buf.iter().position(|&b| b == b'\n') {
let line: Vec<u8> = buf.drain(..=pos).collect();
let line = String::from_utf8_lossy(&line);
handle_sse_line(line.trim(), acc, deltas)?;
}
Ok(())
}
fn handle_sse_line(line: &str, acc: &mut Accumulator, deltas: &mut Vec<String>) -> Result<()> {
let Some(data) = line.strip_prefix("data:") else {
return Ok(());
};
let data = data.trim();
if data.is_empty() || data == "[DONE]" {
return Ok(());
}
let chunk: StreamChunk = match serde_json::from_str(data) {
Ok(c) => c,
Err(_) => return Ok(()), };
if let Some(u) = chunk.usage {
acc.usage = u;
}
for choice in chunk.choices {
if let Some(text) = choice.delta.content {
if !text.is_empty() {
acc.content.push_str(&text);
deltas.push(text);
}
}
for tc in choice.delta.tool_calls.unwrap_or_default() {
let slot = acc.ensure(tc.index);
if let Some(id) = tc.id {
slot.id = id;
}
if let Some(f) = tc.function {
if let Some(name) = f.name {
slot.name.push_str(&name);
}
if let Some(args) = f.arguments {
slot.arguments.push_str(&args);
}
}
}
}
Ok(())
}
fn truncate(s: &str, max: usize) -> String {
if s.len() <= max {
s.to_string()
} else {
let mut end = max;
while end > 0 && !s.is_char_boundary(end) {
end -= 1;
}
format!("{}…", &s[..end])
}
}
#[derive(Serialize)]
struct WireTool<'a> {
#[serde(rename = "type")]
kind: &'static str,
function: WireFunction<'a>,
}
#[derive(Serialize)]
struct WireFunction<'a> {
name: &'a str,
description: &'a str,
parameters: &'a serde_json::Value,
}
impl<'a> From<&'a ToolSchema> for WireTool<'a> {
fn from(t: &'a ToolSchema) -> Self {
WireTool {
kind: "function",
function: WireFunction {
name: &t.name,
description: &t.description,
parameters: &t.parameters,
},
}
}
}
#[derive(Deserialize)]
struct StreamChunk {
#[serde(default)]
choices: Vec<StreamChoice>,
#[serde(default)]
usage: Option<Usage>,
}
#[derive(Deserialize)]
struct StreamChoice {
delta: Delta,
}
#[derive(Deserialize)]
struct Delta {
#[serde(default)]
content: Option<String>,
#[serde(default)]
tool_calls: Option<Vec<ToolCallDelta>>,
}
#[derive(Deserialize)]
struct ToolCallDelta {
#[serde(default)]
index: usize,
#[serde(default)]
id: Option<String>,
#[serde(default)]
function: Option<FnDelta>,
}
#[derive(Deserialize)]
struct FnDelta {
#[serde(default)]
name: Option<String>,
#[serde(default)]
arguments: Option<String>,
}
#[cfg(test)]
mod tests {
use super::*;
use crate::PromptTokensDetails;
use supercode_interchange::ChatMessage;
#[test]
fn request_body_includes_effort_format_and_passthrough() {
let mut req = ChatRequest::new("m", vec![ChatMessage::user("hi")]);
req.effort = Some("high".into());
req.response_format =
Some(serde_json::json!({"type": "json_schema", "json_schema": {"name": "x"}}));
req.extra_body.insert(
"cache_control".into(),
serde_json::json!({"type": "ephemeral"}),
);
req.extra_body.insert(
"provider".into(),
serde_json::json!({"order": ["anthropic"]}),
);
let body = build_request_body(&req, false);
assert_eq!(body["model"], "m");
assert_eq!(body["reasoning_effort"], "high");
assert_eq!(body["response_format"]["type"], "json_schema");
assert_eq!(body["cache_control"]["type"], "ephemeral");
assert_eq!(body["provider"]["order"][0], "anthropic");
assert!(body.get("stream_options").is_none());
}
#[test]
fn extra_body_overrides_modeled_fields() {
let mut req = ChatRequest::new("m", vec![ChatMessage::user("hi")]);
req.max_tokens = Some(100);
req.extra_body
.insert("max_tokens".into(), serde_json::json!(999));
let body = build_request_body(&req, true);
assert_eq!(body["max_tokens"], 999, "extra_body wins");
assert_eq!(body["stream_options"]["include_usage"], true);
}
#[test]
fn cache_plan_annotates_system_and_last_imported_message_only() {
let messages = vec![
ChatMessage::system("sys"),
ChatMessage::user("u1"),
ChatMessage::assistant("a1"),
ChatMessage::user("u2"),
];
let mut req = ChatRequest::new("m", messages);
req.messages = apply_cache_plan(&req.messages, crate::CachePlan::ImportedPrefix, Some(3));
let body = build_request_body(&req, false);
let msgs = body["messages"].as_array().unwrap();
assert_eq!(msgs.len(), 4, "annotation must not change message count");
assert_eq!(
msgs[0]["content"][0]["cache_control"]["type"], "ephemeral",
"breakpoint 1: system message"
);
assert_eq!(
msgs[2]["content"][0]["cache_control"]["type"], "ephemeral",
"breakpoint 2: last message of the imported prefix (a1)"
);
assert_eq!(
msgs[2]["content"][0]["text"], "a1",
"annotated text must be byte-identical to the original content"
);
for (i, m) in msgs.iter().enumerate() {
if i == 0 || i == 2 {
continue;
}
let has_cc = match &m["content"] {
serde_json::Value::Array(parts) => {
parts.iter().any(|p| p.get("cache_control").is_some())
}
serde_json::Value::String(_) => false,
_ => false,
};
assert!(!has_cc, "message {i} must not carry cache_control: {m:?}");
}
}
#[test]
fn cache_plan_off_never_annotates() {
let messages = vec![ChatMessage::system("sys"), ChatMessage::user("u1")];
let out = apply_cache_plan(&messages, crate::CachePlan::Off, Some(2));
assert_eq!(out[0].content_parts, None);
assert_eq!(out[1].content_parts, None);
}
#[test]
fn tier_change_is_cache_bust_truth_table() {
assert!(!tier_change_is_cache_bust(None, 42));
assert!(!tier_change_is_cache_bust(Some(42), 42));
assert!(tier_change_is_cache_bust(Some(42), 7));
}
#[test]
fn cache_plan_dedupes_when_prefix_is_only_the_system_message() {
let messages = vec![ChatMessage::system("sys"), ChatMessage::user("u1")];
let out = apply_cache_plan(&messages, crate::CachePlan::ImportedPrefix, Some(1));
assert!(out[0].content_parts.is_some());
assert_eq!(out[1].content_parts, None);
}
#[test]
fn cache_plan_annotates_last_text_part_of_already_multimodal_message() {
let imported_last = ChatMessage::user_with_images("caption", &["https://x/y.png".into()]);
assert_eq!(
imported_last.content_parts.as_ref().unwrap()[0]["type"],
"text"
);
let messages = vec![ChatMessage::system("sys"), imported_last];
let out = apply_cache_plan(&messages, crate::CachePlan::ImportedPrefix, Some(2));
let parts = out[1].content_parts.as_ref().unwrap();
assert_eq!(parts[0]["cache_control"]["type"], "ephemeral");
assert_eq!(parts[0]["text"], "caption");
assert!(
parts[1].get("cache_control").is_none(),
"the image_url part must not be annotated"
);
}
#[test]
fn usage_parses_prompt_tokens_details_cached_tokens() {
let acc = drain(&[
r#"data: {"choices":[{"delta":{"content":"hi"}}]}"#,
r#"data: {"usage":{"prompt_tokens":100,"completion_tokens":5,"prompt_tokens_details":{"cached_tokens":90}}}"#,
"data: [DONE]",
]);
assert_eq!(acc.usage.prompt_tokens, 100);
let details = acc.usage.prompt_tokens_details.expect("details present");
assert_eq!(details.cached_tokens, 90);
}
fn warm_usage() -> Usage {
Usage {
prompt_tokens: 1000,
completion_tokens: 20,
total_tokens: 1020,
prompt_tokens_details: Some(PromptTokensDetails { cached_tokens: 950 }),
}
}
fn cold_usage() -> Usage {
Usage {
prompt_tokens: 1000,
completion_tokens: 20,
total_tokens: 1020,
prompt_tokens_details: Some(PromptTokensDetails { cached_tokens: 3 }),
}
}
fn moderate_usage() -> Usage {
Usage {
prompt_tokens: 1000,
completion_tokens: 20,
total_tokens: 1020,
prompt_tokens_details: Some(PromptTokensDetails { cached_tokens: 500 }),
}
}
#[test]
fn cache_cold_reason_never_fires_when_not_annotated() {
assert_eq!(
cache_cold_reason(false, true, Some(10_000), &cold_usage()),
None
);
assert_eq!(cache_cold_reason(false, false, None, &cold_usage()), None);
}
#[test]
fn cache_cold_reason_first_annotated_request_never_reports_miss() {
assert_eq!(
cache_cold_reason(true, false, Some(1), &cold_usage()),
None,
"first write: a near-zero cache-read ratio is expected, not a miss"
);
}
#[test]
fn cache_cold_reason_silent_on_warm_back_to_back_turn() {
assert_eq!(cache_cold_reason(true, true, Some(5), &warm_usage()), None);
assert_eq!(cache_cold_reason(true, true, None, &warm_usage()), None);
}
#[test]
fn cache_cold_reason_fires_stale_on_first_turn_of_a_resumed_idle_session() {
assert_eq!(
cache_cold_reason(true, false, Some(20 * 60), &cold_usage()),
Some(CacheColdReason::Stale { idle_secs: 20 * 60 })
);
}
#[test]
fn cache_cold_reason_stale_suppressed_when_usage_disproves_it() {
assert_eq!(
cache_cold_reason(true, false, Some(20 * 60), &warm_usage()),
None,
"cache_established == false, but usage still disproves staleness"
);
assert_eq!(
cache_cold_reason(true, true, Some(CACHE_TTL_SECS), &warm_usage()),
None,
"cache_established == true, at the TTL boundary, usage disproves staleness"
);
}
#[test]
fn cache_cold_reason_stale_disprove_threshold_boundary() {
let at_bar = Usage {
prompt_tokens: 1000,
completion_tokens: 1,
total_tokens: 1001,
prompt_tokens_details: Some(PromptTokensDetails { cached_tokens: 900 }), };
assert_eq!(
cache_cold_reason(true, false, Some(CACHE_TTL_SECS), &at_bar),
None,
"exactly at the disprove bar suppresses Stale"
);
let just_under = Usage {
prompt_tokens_details: Some(PromptTokensDetails { cached_tokens: 899 }),
..at_bar
};
assert_eq!(
cache_cold_reason(true, false, Some(CACHE_TTL_SECS), &just_under),
Some(CacheColdReason::Stale {
idle_secs: CACHE_TTL_SECS
}),
"one token under the disprove bar must not suppress Stale"
);
}
#[test]
fn cache_cold_reason_stale_not_suppressed_by_ambiguous_usage() {
assert_eq!(
cache_cold_reason(true, false, Some(20 * 60), &moderate_usage()),
Some(CacheColdReason::Stale { idle_secs: 20 * 60 })
);
}
#[test]
fn cache_cold_reason_stale_not_suppressed_by_missing_usage_details() {
let no_details = Usage {
prompt_tokens: 1000,
completion_tokens: 20,
total_tokens: 1020,
prompt_tokens_details: None,
};
assert_eq!(
cache_cold_reason(true, false, Some(20 * 60), &no_details),
Some(CacheColdReason::Stale { idle_secs: 20 * 60 })
);
}
#[test]
fn cache_cold_reason_fires_stale_at_ttl_boundary() {
assert_eq!(
cache_cold_reason(true, true, Some(CACHE_TTL_SECS), &moderate_usage()),
Some(CacheColdReason::Stale {
idle_secs: CACHE_TTL_SECS
})
);
assert_eq!(
cache_cold_reason(true, true, Some(CACHE_TTL_SECS - 1), &moderate_usage()),
None,
"one second under the TTL must not fire"
);
}
#[test]
fn cache_cold_reason_fires_miss_on_low_ratio_inside_ttl() {
assert_eq!(
cache_cold_reason(true, true, Some(1), &cold_usage()),
Some(CacheColdReason::Miss {
cached_tokens: 3,
prompt_tokens: 1000,
})
);
}
#[test]
fn cache_cold_reason_ratio_threshold_is_exclusive() {
let at_threshold = Usage {
prompt_tokens: 1000,
completion_tokens: 1,
total_tokens: 1001,
prompt_tokens_details: Some(PromptTokensDetails {
cached_tokens: 100, }),
};
assert_eq!(cache_cold_reason(true, true, Some(1), &at_threshold), None);
let just_under = Usage {
prompt_tokens_details: Some(PromptTokensDetails { cached_tokens: 99 }),
..at_threshold
};
assert!(cache_cold_reason(true, true, Some(1), &just_under).is_some());
}
#[test]
fn cache_cold_reason_no_verdict_without_usage_details() {
let usage = Usage {
prompt_tokens: 1000,
completion_tokens: 5,
total_tokens: 1005,
prompt_tokens_details: None,
};
assert_eq!(cache_cold_reason(true, true, Some(1), &usage), None);
}
#[test]
fn cache_cold_reason_no_verdict_on_zero_prompt_tokens() {
let usage = Usage {
prompt_tokens: 0,
completion_tokens: 5,
total_tokens: 5,
prompt_tokens_details: Some(PromptTokensDetails { cached_tokens: 0 }),
};
assert_eq!(cache_cold_reason(true, true, Some(1), &usage), None);
}
#[test]
fn is_anthropic_family_model_recognizes_vendor_prefixed_slugs() {
assert!(is_anthropic_family_model("anthropic/claude-opus-4-8"));
assert!(is_anthropic_family_model("anthropic/claude-sonnet-4-6"));
assert!(is_anthropic_family_model("anthropic/claude-haiku-4-5"));
}
#[test]
fn is_anthropic_family_model_recognizes_bare_claude_slugs() {
assert!(is_anthropic_family_model("claude-opus-4-8"));
assert!(is_anthropic_family_model("claude-3-5-sonnet-20241022"));
}
#[test]
fn is_anthropic_family_model_rejects_other_known_vendors() {
assert!(!is_anthropic_family_model("openai/gpt-5"));
assert!(!is_anthropic_family_model("openai/gpt-5.5"));
assert!(!is_anthropic_family_model("google/gemini-2.5-pro"));
assert!(!is_anthropic_family_model("deepseek/deepseek-v4-pro"));
assert!(!is_anthropic_family_model("meta-llama/llama-4-maverick"));
}
#[test]
fn is_anthropic_family_model_does_not_assume_unknown_slugs() {
assert!(!is_anthropic_family_model("my-custom-local-model"));
assert!(!is_anthropic_family_model(""));
}
#[test]
fn truncate_never_splits_a_codepoint() {
let s = "é".repeat(2000); let out = truncate(&s, 2001); assert!(out.ends_with('…'));
assert!(out.len() <= 2001 + '…'.len_utf8());
}
fn drain(lines: &[&str]) -> Accumulator {
let mut acc = Accumulator::default();
let mut deltas = Vec::new();
let mut buf: Vec<u8> = Vec::new();
for l in lines {
buf.extend_from_slice(l.as_bytes());
buf.push(b'\n');
}
drain_sse_lines(&mut buf, &mut acc, &mut deltas).unwrap();
acc
}
#[test]
fn streaming_assembles_tool_calls_and_usage_across_deltas() {
let acc = drain(&[
r#"data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"read_"}}]}}]}"#,
r#"data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"name":"file","arguments":"{\"path\":"}}]}}]}"#,
r#"data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"a\"}"}}]}}]}"#,
r#"data: {"choices":[{"delta":{"content":"done"}}]}"#,
r#"data: {"usage":{"prompt_tokens":3,"completion_tokens":5}}"#,
"data: [DONE]",
]);
let msg = acc.to_message();
let calls = msg.tool_calls.expect("tool calls");
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].id, "call_1");
assert_eq!(
calls[0].function.name, "read_file",
"name spread over deltas"
);
assert_eq!(calls[0].function.arguments, r#"{"path":"a"}"#);
assert_eq!(msg.content.as_deref(), Some("done"));
assert_eq!(acc.usage.completion_tokens, 5);
}
#[test]
fn streaming_tolerates_done_keepalive_and_blank_lines() {
let acc = drain(&[
"",
": keep-alive",
r#"data: {"choices":[{"delta":{"content":"hi"}}]}"#,
"data: not-json",
"data: [DONE]",
]);
assert_eq!(acc.to_message().content.as_deref(), Some("hi"));
}
#[tokio::test]
async fn non_success_status_becomes_provider_error() {
use crate::RuntimeError as Error;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (mut sock, _) = listener.accept().await.unwrap();
let mut buf = [0u8; 2048];
let _ = sock.read(&mut buf).await;
let body = r#"{"error":{"message":"bad key"}}"#;
let resp = format!(
"HTTP/1.1 401 Unauthorized\r\nContent-Length: {}\r\nContent-Type: application/json\r\nConnection: close\r\n\r\n{}",
body.len(),
body
);
sock.write_all(resp.as_bytes()).await.unwrap();
sock.flush().await.unwrap();
});
let provider = OpenAiProvider::new(format!("http://{addr}"), "k", HashMap::new());
let req = ChatRequest::new("m", vec![ChatMessage::user("hi")]);
let err = provider.complete(&req, &|_: &str| {}).await.unwrap_err();
match err {
Error::Provider { status, body } => {
assert_eq!(status, 401);
assert!(body.contains("bad key"), "body: {body}");
}
other => panic!("expected Provider error, got: {other:?}"),
}
server.await.unwrap();
}
#[tokio::test]
async fn streams_a_200_response_into_a_message() {
use std::sync::{Arc, Mutex};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (mut sock, _) = listener.accept().await.unwrap();
let mut buf = [0u8; 2048];
let _ = sock.read(&mut buf).await;
let sse = "data: {\"choices\":[{\"delta\":{\"content\":\"hel\"}}]}\n\n\
data: {\"choices\":[{\"delta\":{\"content\":\"lo\"}}]}\n\n\
data: [DONE]\n\n";
let resp = format!(
"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
sse.len(),
sse
);
sock.write_all(resp.as_bytes()).await.unwrap();
sock.flush().await.unwrap();
});
let provider = OpenAiProvider::new(format!("http://{addr}"), "k", HashMap::new());
let req = ChatRequest::new("m", vec![ChatMessage::user("hi")]);
let seen = Arc::new(Mutex::new(String::new()));
let seen2 = seen.clone();
let on_delta = move |s: &str| seen2.lock().unwrap().push_str(s);
let (msg, _usage) = provider.complete(&req, &on_delta).await.unwrap();
assert_eq!(msg.content.as_deref(), Some("hello"));
assert_eq!(*seen.lock().unwrap(), "hello", "deltas streamed live");
server.await.unwrap();
}
#[test]
fn from_retry_config_unset_is_byte_identical_to_default() {
let opts = HttpOptions::from_retry_config(true, None, None);
let default = HttpOptions::default();
assert_eq!(opts.max_retries, default.max_retries);
assert_eq!(opts.retry_backoff_base, default.retry_backoff_base);
assert_eq!(opts.connect_timeout, default.connect_timeout);
assert_eq!(opts.read_idle_timeout, default.read_idle_timeout);
}
#[test]
fn from_retry_config_disabled_forces_zero_retries() {
let opts = HttpOptions::from_retry_config(false, None, None);
assert_eq!(opts.max_retries, 0);
assert_eq!(
opts.retry_backoff_base,
HttpOptions::default().retry_backoff_base
);
}
#[test]
fn from_retry_config_disabled_with_explicit_max_retries_still_forces_zero() {
let opts = HttpOptions::from_retry_config(false, Some(5), None);
assert_eq!(opts.max_retries, 0);
}
#[test]
fn from_retry_config_overrides_apply_when_enabled() {
let opts = HttpOptions::from_retry_config(true, Some(7), Some(1234));
assert_eq!(opts.max_retries, 7);
assert_eq!(opts.retry_backoff_base, Duration::from_millis(1234));
}
#[test]
fn from_retry_config_partial_override_leaves_the_other_at_default() {
let opts = HttpOptions::from_retry_config(true, Some(9), None);
assert_eq!(opts.max_retries, 9);
assert_eq!(
opts.retry_backoff_base,
HttpOptions::default().retry_backoff_base
);
}
fn test_http_options() -> HttpOptions {
HttpOptions {
connect_timeout: Duration::from_millis(250),
read_idle_timeout: Duration::from_millis(250),
max_retries: 2,
retry_backoff_base: Duration::from_millis(10),
}
}
#[tokio::test]
async fn hung_connection_errors_via_read_timeout_within_bounded_time() {
use tokio::io::AsyncReadExt;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
loop {
let Ok((mut sock, _)) = listener.accept().await else {
break;
};
tokio::spawn(async move {
let mut buf = [0u8; 2048];
let _ = sock.read(&mut buf).await;
tokio::time::sleep(Duration::from_secs(2)).await;
});
}
});
let provider = OpenAiProvider::new_with_options(
format!("http://{addr}"),
"k",
HashMap::new(),
test_http_options(),
);
let req = ChatRequest::new("m", vec![ChatMessage::user("hi")]);
let outcome = tokio::time::timeout(Duration::from_secs(5), async {
provider.complete(&req, &|_: &str| {}).await
})
.await
.expect("complete() must return within the outer bound, not hang forever");
match outcome {
Err(Error::Http(_)) => {}
other => panic!("expected Err(Error::Http(_)) from the read timeout, got: {other:?}"),
}
server.abort();
}
#[tokio::test]
async fn retries_503_then_succeeds_with_exactly_two_requests() {
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let connections = Arc::new(AtomicUsize::new(0));
let connections2 = connections.clone();
let server = tokio::spawn(async move {
for _ in 0..2 {
let (mut sock, _) = listener.accept().await.unwrap();
let n = connections2.fetch_add(1, Ordering::SeqCst) + 1;
let mut buf = [0u8; 2048];
let _ = sock.read(&mut buf).await;
if n == 1 {
let resp =
"HTTP/1.1 503 Service Unavailable\r\nContent-Length: 0\r\nConnection: close\r\n\r\n";
sock.write_all(resp.as_bytes()).await.unwrap();
} else {
let sse = "data: {\"choices\":[{\"delta\":{\"content\":\"hel\"}}]}\n\n\
data: {\"choices\":[{\"delta\":{\"content\":\"lo\"}}]}\n\n\
data: [DONE]\n\n";
let resp = format!(
"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
sse.len(),
sse
);
sock.write_all(resp.as_bytes()).await.unwrap();
}
sock.flush().await.unwrap();
}
});
let provider = OpenAiProvider::new_with_options(
format!("http://{addr}"),
"k",
HashMap::new(),
test_http_options(),
);
let req = ChatRequest::new("m", vec![ChatMessage::user("hi")]);
let (msg, _usage) = provider.complete(&req, &|_: &str| {}).await.unwrap();
assert_eq!(msg.content.as_deref(), Some("hello"));
server.await.unwrap();
assert_eq!(
connections.load(Ordering::SeqCst),
2,
"exactly 2 requests made: one 503, one successful retry"
);
}
#[test]
fn streaming_decodes_multibyte_across_chunk_boundaries() {
let line = "data: {\"choices\":[{\"delta\":{\"content\":\"héllo🌍\"}}]}\n";
let bytes = line.as_bytes();
let mut deltas = Vec::new();
for split in 1..bytes.len() {
let mut acc = Accumulator::default();
let mut buf: Vec<u8> = Vec::new();
buf.extend_from_slice(&bytes[..split]);
drain_sse_lines(&mut buf, &mut acc, &mut deltas).unwrap();
buf.extend_from_slice(&bytes[split..]);
drain_sse_lines(&mut buf, &mut acc, &mut deltas).unwrap();
assert_eq!(acc.content, "héllo🌍", "split at byte {split}");
assert!(!acc.content.contains('\u{FFFD}'));
}
}
}