1use std::collections::HashMap;
9use std::time::Duration;
10
11use async_trait::async_trait;
12use futures::StreamExt;
13use serde::{Deserialize, Serialize};
14
15use supercode_interchange::{ChatMessage, FunctionCall, Role, ToolCall};
16
17use crate::{CachePlan, ChatRequest, Result, RuntimeError as Error, ToolSchema, Usage};
18
19const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
22
23const READ_IDLE_TIMEOUT: Duration = Duration::from_secs(120);
29
30const MAX_RETRIES: u32 = 2;
34
35const RETRY_BACKOFF_BASE: Duration = Duration::from_millis(500);
38
39#[derive(Debug, Clone, Copy)]
45#[doc(hidden)]
46pub struct HttpOptions {
47 pub(crate) connect_timeout: Duration,
48 pub(crate) read_idle_timeout: Duration,
49 pub(crate) max_retries: u32,
50 pub(crate) retry_backoff_base: Duration,
51}
52
53impl Default for HttpOptions {
54 fn default() -> Self {
55 HttpOptions {
56 connect_timeout: CONNECT_TIMEOUT,
57 read_idle_timeout: READ_IDLE_TIMEOUT,
58 max_retries: MAX_RETRIES,
59 retry_backoff_base: RETRY_BACKOFF_BASE,
60 }
61 }
62}
63
64impl HttpOptions {
65 #[doc(hidden)]
78 pub fn from_retry_config(
79 enabled: bool,
80 max_retries: Option<u32>,
81 base_delay_ms: Option<u64>,
82 ) -> HttpOptions {
83 let base = HttpOptions::default();
84 HttpOptions {
85 max_retries: if enabled {
86 max_retries.unwrap_or(base.max_retries)
87 } else {
88 0
89 },
90 retry_backoff_base: base_delay_ms
91 .map(Duration::from_millis)
92 .unwrap_or(base.retry_backoff_base),
93 ..base
94 }
95 }
96}
97
98#[derive(Debug, Clone, PartialEq, Eq)]
107pub struct RetryNotice {
108 pub attempt: u32,
110 pub delay_ms: u64,
112 pub reason: String,
114}
115
116#[derive(Debug, Default)]
125pub struct RetryLog {
126 notices: std::sync::Mutex<Vec<RetryNotice>>,
127}
128
129impl RetryLog {
130 pub fn record(&self, notice: RetryNotice) {
132 self.notices
133 .lock()
134 .unwrap_or_else(std::sync::PoisonError::into_inner)
135 .push(notice);
136 }
137
138 pub fn drain(&self) -> Vec<RetryNotice> {
140 std::mem::take(
141 &mut *self
142 .notices
143 .lock()
144 .unwrap_or_else(std::sync::PoisonError::into_inner),
145 )
146 }
147}
148
149pub(crate) fn build_request_body(req: &ChatRequest, stream: bool) -> serde_json::Value {
153 use serde_json::json;
154 let mut body = json!({
155 "model": req.model,
156 "messages": req.messages,
157 "stream": stream,
158 });
159 let obj = body.as_object_mut().unwrap();
160 if !req.tools.is_empty() {
161 obj.insert(
162 "tools".into(),
163 serde_json::to_value(req.tools.iter().map(WireTool::from).collect::<Vec<_>>()).unwrap(),
164 );
165 }
166 if let Some(t) = req.temperature {
167 obj.insert("temperature".into(), json!(t));
168 }
169 if let Some(m) = req.max_tokens {
170 obj.insert("max_tokens".into(), json!(m));
171 }
172 if let Some(e) = &req.effort {
173 obj.insert("reasoning_effort".into(), json!(e));
174 }
175 if let Some(rf) = &req.response_format {
176 obj.insert("response_format".into(), rf.clone());
177 }
178 if let Some(tier) = &req.service_tier {
181 obj.insert("service_tier".into(), json!(tier));
182 }
183 if let Some(budget) = req.thinking_budget {
190 let entry = obj
191 .entry("reasoning".to_string())
192 .or_insert_with(|| json!({}));
193 if let Some(o) = entry.as_object_mut() {
194 o.insert("max_tokens".into(), json!(budget));
195 }
196 }
197 if stream {
198 obj.insert("stream_options".into(), json!({"include_usage": true}));
199 }
200 for (k, v) in &req.extra_body {
202 obj.insert(k.clone(), v.clone());
203 }
204 body
205}
206
207#[doc(hidden)]
233pub fn apply_cache_plan(
234 messages: &[ChatMessage],
235 plan: CachePlan,
236 imported_prefix_len: Option<usize>,
237) -> Vec<ChatMessage> {
238 let mut out = messages.to_vec();
239 if !matches!(plan, CachePlan::ImportedPrefix) {
240 return out;
241 }
242 let Some(len) = imported_prefix_len.filter(|&n| n > 0) else {
243 return out;
244 };
245 let last = len - 1;
246 let mut targets = vec![0usize];
247 if last != 0 {
248 targets.push(last);
249 }
250 for idx in targets {
251 if let Some(msg) = out.get_mut(idx) {
252 annotate_cache_breakpoint(msg);
253 }
254 }
255 out
256}
257
258fn annotate_cache_breakpoint(msg: &mut ChatMessage) {
261 let cache_control = serde_json::json!({"type": "ephemeral"});
262 if let Some(parts) = msg.content_parts.as_mut() {
263 if let Some(text_part) = parts
265 .iter_mut()
266 .rev()
267 .find(|p| p.get("type").and_then(serde_json::Value::as_str) == Some("text"))
268 {
269 if let Some(obj) = text_part.as_object_mut() {
270 obj.insert("cache_control".to_string(), cache_control);
271 }
272 }
273 return;
274 }
275 let text = msg.content.take().unwrap_or_default();
276 msg.content_parts = Some(vec![serde_json::json!({
277 "type": "text",
278 "text": text,
279 "cache_control": cache_control,
280 })]);
281}
282
283#[doc(hidden)]
298pub fn tier_change_is_cache_bust(previous: Option<u64>, current: u64) -> bool {
299 previous.is_some_and(|p| p != current)
300}
301
302pub(crate) const CACHE_TTL_SECS: i64 = 300;
309
310pub(crate) const CACHE_MISS_RATIO_THRESHOLD: f64 = 0.10;
316
317pub(crate) const CACHE_STALE_DISPROVE_RATIO_THRESHOLD: f64 = 1.0 - CACHE_MISS_RATIO_THRESHOLD;
335
336#[doc(hidden)]
354pub fn is_anthropic_family_model(model: &str) -> bool {
355 model.starts_with("anthropic/") || model.starts_with("claude-") || model.starts_with("claude/")
356}
357
358#[derive(Debug, Clone, Copy, PartialEq)]
361#[doc(hidden)]
362pub enum CacheColdReason {
363 Stale {
368 idle_secs: i64,
370 },
371 Miss {
376 cached_tokens: u64,
378 prompt_tokens: u64,
380 },
381}
382
383impl CacheColdReason {
384 #[doc(hidden)]
386 pub fn message(&self) -> String {
387 match self {
388 CacheColdReason::Stale { idle_secs } => format!(
389 "cache likely cold — this turn was sent {}m{:02}s after the cache was last \
390 refreshed (Anthropic's ephemeral prompt cache expires after 5m idle) — this \
391 turn likely paid full input cost for the cached prefix",
392 idle_secs / 60,
393 idle_secs % 60,
394 ),
395 CacheColdReason::Miss {
396 cached_tokens,
397 prompt_tokens,
398 } => format!(
399 "unexpected cache miss — only {cached_tokens}/{prompt_tokens} prompt tokens \
400 were served from cache this turn even though reuse was expected — this turn \
401 likely paid full input cost for the cached prefix",
402 ),
403 }
404 }
405}
406
407#[doc(hidden)]
455pub fn cache_cold_reason(
456 will_annotate: bool,
457 cache_established: bool,
458 idle_secs: Option<i64>,
459 usage: &Usage,
460) -> Option<CacheColdReason> {
461 if !will_annotate {
462 return None;
463 }
464 if let Some(idle_secs) = idle_secs {
465 if idle_secs >= CACHE_TTL_SECS {
466 let disproven_by_usage = usage
467 .prompt_tokens_details
468 .filter(|_| usage.prompt_tokens > 0)
469 .is_some_and(|details| {
470 details.cached_tokens as f64 / usage.prompt_tokens as f64
471 >= CACHE_STALE_DISPROVE_RATIO_THRESHOLD
472 });
473 if !disproven_by_usage {
474 return Some(CacheColdReason::Stale { idle_secs });
475 }
476 }
477 }
478 if !cache_established {
479 return None;
482 }
483 let details = usage.prompt_tokens_details?;
484 if usage.prompt_tokens == 0 {
485 return None;
488 }
489 let ratio = details.cached_tokens as f64 / usage.prompt_tokens as f64;
490 if ratio < CACHE_MISS_RATIO_THRESHOLD {
491 return Some(CacheColdReason::Miss {
492 cached_tokens: details.cached_tokens,
493 prompt_tokens: usage.prompt_tokens,
494 });
495 }
496 None
497}
498
499#[async_trait]
502pub trait Provider: Send + Sync {
503 async fn complete(
506 &self,
507 req: &ChatRequest,
508 on_delta: &(dyn for<'a> Fn(&'a str) + Send + Sync),
509 ) -> Result<(ChatMessage, Usage)>;
510}
511
512pub struct OpenAiProvider {
515 client: reqwest::Client,
516 base_url: String,
517 api_key: String,
518 extra_headers: HashMap<String, String>,
519 http_options: HttpOptions,
520 retry_log: Option<std::sync::Arc<RetryLog>>,
524}
525
526impl OpenAiProvider {
527 pub fn new(
529 base_url: impl Into<String>,
530 api_key: impl Into<String>,
531 extra_headers: HashMap<String, String>,
532 ) -> Self {
533 Self::new_with_options(base_url, api_key, extra_headers, HttpOptions::default())
534 }
535
536 #[doc(hidden)]
541 pub fn new_with_options(
542 base_url: impl Into<String>,
543 api_key: impl Into<String>,
544 extra_headers: HashMap<String, String>,
545 http_options: HttpOptions,
546 ) -> Self {
547 OpenAiProvider {
548 client: reqwest::Client::builder()
549 .connect_timeout(http_options.connect_timeout)
550 .read_timeout(http_options.read_idle_timeout)
551 .build()
552 .expect("static reqwest client config cannot fail"),
553 base_url: base_url.into(),
554 api_key: api_key.into(),
555 extra_headers,
556 http_options,
557 retry_log: None,
558 }
559 }
560
561 pub fn with_retry_log(mut self, log: std::sync::Arc<RetryLog>) -> Self {
564 self.retry_log = Some(log);
565 self
566 }
567
568 fn endpoint(&self) -> String {
569 format!("{}/chat/completions", self.base_url.trim_end_matches('/'))
570 }
571
572 async fn send_with_retry(&self, wire: &serde_json::Value) -> Result<reqwest::Response> {
579 let mut attempt = 0u32;
580 loop {
581 let mut builder = self
582 .client
583 .post(self.endpoint())
584 .bearer_auth(&self.api_key)
585 .header("Content-Type", "application/json");
586 for (k, v) in &self.extra_headers {
587 builder = builder.header(k, v);
588 }
589
590 let sent = builder.json(wire).send().await;
591 let (retryable, result): (bool, Result<reqwest::Response>) = match sent {
592 Err(e) => (true, Err(Error::from(e))),
593 Ok(resp) => {
594 let status = resp.status();
595 if status.is_success() {
596 (false, Ok(resp))
597 } else if status.is_server_error() {
598 let body = resp.text().await.unwrap_or_default();
599 (
600 true,
601 Err(Error::Provider {
602 status: status.as_u16(),
603 body: truncate(&body, 2000),
604 }),
605 )
606 } else {
607 let body = resp.text().await.unwrap_or_default();
608 (
609 false,
610 Err(Error::Provider {
611 status: status.as_u16(),
612 body: truncate(&body, 2000),
613 }),
614 )
615 }
616 }
617 };
618
619 if !retryable || attempt >= self.http_options.max_retries {
620 return result;
621 }
622 let backoff = self.http_options.retry_backoff_base * 2u32.pow(attempt);
623 if let Some(log) = &self.retry_log {
628 log.record(RetryNotice {
629 attempt,
630 delay_ms: backoff.as_millis() as u64,
631 reason: match &result {
632 Err(e) => e.to_string(),
633 Ok(_) => String::new(),
634 },
635 });
636 }
637 tokio::time::sleep(backoff).await;
638 attempt += 1;
639 }
640 }
641}
642
643#[async_trait]
644impl Provider for OpenAiProvider {
645 async fn complete(
646 &self,
647 req: &ChatRequest,
648 on_delta: &(dyn for<'a> Fn(&'a str) + Send + Sync),
649 ) -> Result<(ChatMessage, Usage)> {
650 let wire = build_request_body(req, true);
651
652 let resp = self.send_with_retry(&wire).await?;
653
654 let mut acc = Accumulator::default();
655 let mut buf: Vec<u8> = Vec::new();
661 let mut deltas: Vec<String> = Vec::new();
662 let mut stream = resp.bytes_stream();
663 while let Some(chunk) = stream.next().await {
664 let bytes = chunk?;
665 buf.extend_from_slice(&bytes);
666 drain_sse_lines(&mut buf, &mut acc, &mut deltas)?;
667 for d in deltas.drain(..) {
668 on_delta(&d);
669 }
670 }
671 let tail = String::from_utf8_lossy(&buf);
673 if !tail.trim().is_empty() {
674 handle_sse_line(tail.trim(), &mut acc, &mut deltas)?;
675 for d in deltas.drain(..) {
676 on_delta(&d);
677 }
678 }
679
680 Ok((acc.to_message(), acc_usage(&acc)))
681 }
682}
683
684#[derive(Default)]
687struct Accumulator {
688 content: String,
689 tool_calls: Vec<ToolCallAccum>,
690 usage: Usage,
691 served_model: Option<String>,
696}
697
698#[derive(Default)]
699struct ToolCallAccum {
700 id: String,
701 name: String,
702 arguments: String,
703}
704
705impl Accumulator {
706 fn ensure(&mut self, index: usize) -> &mut ToolCallAccum {
707 while self.tool_calls.len() <= index {
708 self.tool_calls.push(ToolCallAccum::default());
709 }
710 &mut self.tool_calls[index]
711 }
712
713 fn to_message(&self) -> ChatMessage {
714 let calls: Vec<ToolCall> = self
715 .tool_calls
716 .iter()
717 .filter(|c| !c.id.is_empty() || !c.name.is_empty())
718 .map(|c| ToolCall {
719 id: c.id.clone(),
720 kind: "function".to_string(),
721 function: FunctionCall {
722 name: c.name.clone(),
723 arguments: c.arguments.clone(),
724 },
725 })
726 .collect();
727 ChatMessage {
728 role: Role::Assistant,
729 content: (!self.content.is_empty()).then(|| self.content.clone()),
730 content_parts: None,
731 tool_calls: (!calls.is_empty()).then_some(calls),
732 tool_call_id: None,
733 name: None,
734 metadata: match &self.served_model {
735 Some(model) => {
736 let mut m = std::collections::BTreeMap::new();
737 m.insert(SERVED_MODEL_KEY.to_string(), model.clone());
738 m
739 }
740 None => Default::default(),
741 },
742 }
743 }
744}
745
746pub const SERVED_MODEL_KEY: &str = "served_model";
750
751fn acc_usage(acc: &Accumulator) -> Usage {
752 acc.usage.clone()
753}
754
755fn drain_sse_lines(
756 buf: &mut Vec<u8>,
757 acc: &mut Accumulator,
758 deltas: &mut Vec<String>,
759) -> Result<()> {
760 while let Some(pos) = buf.iter().position(|&b| b == b'\n') {
761 let line: Vec<u8> = buf.drain(..=pos).collect();
762 let line = String::from_utf8_lossy(&line);
763 handle_sse_line(line.trim(), acc, deltas)?;
764 }
765 Ok(())
766}
767
768fn handle_sse_line(line: &str, acc: &mut Accumulator, deltas: &mut Vec<String>) -> Result<()> {
769 let Some(data) = line.strip_prefix("data:") else {
770 return Ok(());
771 };
772 let data = data.trim();
773 if data.is_empty() || data == "[DONE]" {
774 return Ok(());
775 }
776 let chunk: StreamChunk = match serde_json::from_str(data) {
777 Ok(c) => c,
778 Err(_) => return Ok(()), };
780 if let Some(u) = chunk.usage {
781 acc.usage = u;
782 }
783 if let Some(model) = chunk.model {
784 if !model.is_empty() {
785 acc.served_model = Some(model);
786 }
787 }
788 for choice in chunk.choices {
789 if let Some(text) = choice.delta.content {
790 if !text.is_empty() {
791 acc.content.push_str(&text);
792 deltas.push(text);
793 }
794 }
795 for tc in choice.delta.tool_calls.unwrap_or_default() {
796 let slot = acc.ensure(tc.index);
797 if let Some(id) = tc.id {
798 slot.id = id;
799 }
800 if let Some(f) = tc.function {
801 if let Some(name) = f.name {
802 slot.name.push_str(&name);
803 }
804 if let Some(args) = f.arguments {
805 slot.arguments.push_str(&args);
806 }
807 }
808 }
809 }
810 Ok(())
811}
812
813fn truncate(s: &str, max: usize) -> String {
814 if s.len() <= max {
815 s.to_string()
816 } else {
817 let mut end = max;
820 while end > 0 && !s.is_char_boundary(end) {
821 end -= 1;
822 }
823 format!("{}…", &s[..end])
824 }
825}
826
827#[derive(Serialize)]
830struct WireTool<'a> {
831 #[serde(rename = "type")]
832 kind: &'static str,
833 function: WireFunction<'a>,
834}
835
836#[derive(Serialize)]
837struct WireFunction<'a> {
838 name: &'a str,
839 description: &'a str,
840 parameters: &'a serde_json::Value,
841}
842
843impl<'a> From<&'a ToolSchema> for WireTool<'a> {
844 fn from(t: &'a ToolSchema) -> Self {
845 WireTool {
846 kind: "function",
847 function: WireFunction {
848 name: &t.name,
849 description: &t.description,
850 parameters: &t.parameters,
851 },
852 }
853 }
854}
855
856#[derive(Deserialize)]
857struct StreamChunk {
858 #[serde(default)]
859 choices: Vec<StreamChoice>,
860 #[serde(default)]
861 usage: Option<Usage>,
862 #[serde(default)]
865 model: Option<String>,
866}
867
868#[derive(Deserialize)]
869struct StreamChoice {
870 delta: Delta,
871}
872
873#[derive(Deserialize)]
874struct Delta {
875 #[serde(default)]
876 content: Option<String>,
877 #[serde(default)]
878 tool_calls: Option<Vec<ToolCallDelta>>,
879}
880
881#[derive(Deserialize)]
882struct ToolCallDelta {
883 #[serde(default)]
884 index: usize,
885 #[serde(default)]
886 id: Option<String>,
887 #[serde(default)]
888 function: Option<FnDelta>,
889}
890
891#[derive(Deserialize)]
892struct FnDelta {
893 #[serde(default)]
894 name: Option<String>,
895 #[serde(default)]
896 arguments: Option<String>,
897}
898
899#[cfg(test)]
900mod tests {
901 use super::*;
902 use crate::PromptTokensDetails;
903 use supercode_interchange::ChatMessage;
904
905 #[test]
906 fn request_body_includes_effort_format_and_passthrough() {
907 let mut req = ChatRequest::new("m", vec![ChatMessage::user("hi")]);
908 req.effort = Some("high".into());
909 req.response_format =
910 Some(serde_json::json!({"type": "json_schema", "json_schema": {"name": "x"}}));
911 req.extra_body.insert(
912 "cache_control".into(),
913 serde_json::json!({"type": "ephemeral"}),
914 );
915 req.extra_body.insert(
916 "provider".into(),
917 serde_json::json!({"order": ["anthropic"]}),
918 );
919
920 let body = build_request_body(&req, false);
921 assert_eq!(body["model"], "m");
922 assert_eq!(body["reasoning_effort"], "high");
923 assert_eq!(body["response_format"]["type"], "json_schema");
924 assert_eq!(body["cache_control"]["type"], "ephemeral");
925 assert_eq!(body["provider"]["order"][0], "anthropic");
926 assert!(body.get("stream_options").is_none());
928 }
929
930 #[test]
931 fn extra_body_overrides_modeled_fields() {
932 let mut req = ChatRequest::new("m", vec![ChatMessage::user("hi")]);
933 req.max_tokens = Some(100);
934 req.extra_body
935 .insert("max_tokens".into(), serde_json::json!(999));
936 let body = build_request_body(&req, true);
937 assert_eq!(body["max_tokens"], 999, "extra_body wins");
938 assert_eq!(body["stream_options"]["include_usage"], true);
939 }
940
941 #[test]
950 fn cache_plan_annotates_system_and_last_imported_message_only() {
951 let messages = vec![
952 ChatMessage::system("sys"),
953 ChatMessage::user("u1"),
954 ChatMessage::assistant("a1"),
955 ChatMessage::user("u2"),
956 ];
957 let mut req = ChatRequest::new("m", messages);
958 req.messages = apply_cache_plan(&req.messages, crate::CachePlan::ImportedPrefix, Some(3));
959
960 let body = build_request_body(&req, false);
961 let msgs = body["messages"].as_array().unwrap();
962 assert_eq!(msgs.len(), 4, "annotation must not change message count");
963
964 assert_eq!(
965 msgs[0]["content"][0]["cache_control"]["type"], "ephemeral",
966 "breakpoint 1: system message"
967 );
968 assert_eq!(
969 msgs[2]["content"][0]["cache_control"]["type"], "ephemeral",
970 "breakpoint 2: last message of the imported prefix (a1)"
971 );
972 assert_eq!(
973 msgs[2]["content"][0]["text"], "a1",
974 "annotated text must be byte-identical to the original content"
975 );
976
977 for (i, m) in msgs.iter().enumerate() {
979 if i == 0 || i == 2 {
980 continue;
981 }
982 let has_cc = match &m["content"] {
983 serde_json::Value::Array(parts) => {
984 parts.iter().any(|p| p.get("cache_control").is_some())
985 }
986 serde_json::Value::String(_) => false,
987 _ => false,
988 };
989 assert!(!has_cc, "message {i} must not carry cache_control: {m:?}");
990 }
991 }
992
993 #[test]
994 fn cache_plan_off_never_annotates() {
995 let messages = vec![ChatMessage::system("sys"), ChatMessage::user("u1")];
996 let out = apply_cache_plan(&messages, crate::CachePlan::Off, Some(2));
997 assert_eq!(out[0].content_parts, None);
998 assert_eq!(out[1].content_parts, None);
999 }
1000
1001 #[test]
1002 fn tier_change_is_cache_bust_truth_table() {
1003 assert!(!tier_change_is_cache_bust(None, 42));
1005 assert!(!tier_change_is_cache_bust(Some(42), 42));
1007 assert!(tier_change_is_cache_bust(Some(42), 7));
1009 }
1010
1011 #[test]
1012 fn cache_plan_dedupes_when_prefix_is_only_the_system_message() {
1013 let messages = vec![ChatMessage::system("sys"), ChatMessage::user("u1")];
1016 let out = apply_cache_plan(&messages, crate::CachePlan::ImportedPrefix, Some(1));
1017 assert!(out[0].content_parts.is_some());
1018 assert_eq!(out[1].content_parts, None);
1019 }
1020
1021 #[test]
1022 fn cache_plan_annotates_last_text_part_of_already_multimodal_message() {
1023 let imported_last = ChatMessage::user_with_images("caption", &["https://x/y.png".into()]);
1024 assert_eq!(
1026 imported_last.content_parts.as_ref().unwrap()[0]["type"],
1027 "text"
1028 );
1029 let messages = vec![ChatMessage::system("sys"), imported_last];
1030 let out = apply_cache_plan(&messages, crate::CachePlan::ImportedPrefix, Some(2));
1031 let parts = out[1].content_parts.as_ref().unwrap();
1032 assert_eq!(parts[0]["cache_control"]["type"], "ephemeral");
1033 assert_eq!(parts[0]["text"], "caption");
1034 assert!(
1035 parts[1].get("cache_control").is_none(),
1036 "the image_url part must not be annotated"
1037 );
1038 }
1039
1040 #[test]
1044 fn usage_parses_prompt_tokens_details_cached_tokens() {
1045 let acc = drain(&[
1046 r#"data: {"choices":[{"delta":{"content":"hi"}}]}"#,
1047 r#"data: {"usage":{"prompt_tokens":100,"completion_tokens":5,"prompt_tokens_details":{"cached_tokens":90}}}"#,
1048 "data: [DONE]",
1049 ]);
1050 assert_eq!(acc.usage.prompt_tokens, 100);
1051 let details = acc.usage.prompt_tokens_details.expect("details present");
1052 assert_eq!(details.cached_tokens, 90);
1053 }
1054
1055 fn warm_usage() -> Usage {
1058 Usage {
1061 prompt_tokens: 1000,
1062 completion_tokens: 20,
1063 total_tokens: 1020,
1064 prompt_tokens_details: Some(PromptTokensDetails { cached_tokens: 950 }),
1065 }
1066 }
1067
1068 fn cold_usage() -> Usage {
1069 Usage {
1071 prompt_tokens: 1000,
1072 completion_tokens: 20,
1073 total_tokens: 1020,
1074 prompt_tokens_details: Some(PromptTokensDetails { cached_tokens: 3 }),
1075 }
1076 }
1077
1078 fn moderate_usage() -> Usage {
1085 Usage {
1086 prompt_tokens: 1000,
1087 completion_tokens: 20,
1088 total_tokens: 1020,
1089 prompt_tokens_details: Some(PromptTokensDetails { cached_tokens: 500 }),
1090 }
1091 }
1092
1093 #[test]
1097 fn cache_cold_reason_never_fires_when_not_annotated() {
1098 assert_eq!(
1099 cache_cold_reason(false, true, Some(10_000), &cold_usage()),
1100 None
1101 );
1102 assert_eq!(cache_cold_reason(false, false, None, &cold_usage()), None);
1103 }
1104
1105 #[test]
1111 fn cache_cold_reason_first_annotated_request_never_reports_miss() {
1112 assert_eq!(
1113 cache_cold_reason(true, false, Some(1), &cold_usage()),
1114 None,
1115 "first write: a near-zero cache-read ratio is expected, not a miss"
1116 );
1117 }
1118
1119 #[test]
1123 fn cache_cold_reason_silent_on_warm_back_to_back_turn() {
1124 assert_eq!(cache_cold_reason(true, true, Some(5), &warm_usage()), None);
1125 assert_eq!(cache_cold_reason(true, true, None, &warm_usage()), None);
1128 }
1129
1130 #[test]
1139 fn cache_cold_reason_fires_stale_on_first_turn_of_a_resumed_idle_session() {
1140 assert_eq!(
1141 cache_cold_reason(true, false, Some(20 * 60), &cold_usage()),
1142 Some(CacheColdReason::Stale { idle_secs: 20 * 60 })
1143 );
1144 }
1145
1146 #[test]
1160 fn cache_cold_reason_stale_suppressed_when_usage_disproves_it() {
1161 assert_eq!(
1162 cache_cold_reason(true, false, Some(20 * 60), &warm_usage()),
1163 None,
1164 "cache_established == false, but usage still disproves staleness"
1165 );
1166 assert_eq!(
1167 cache_cold_reason(true, true, Some(CACHE_TTL_SECS), &warm_usage()),
1168 None,
1169 "cache_established == true, at the TTL boundary, usage disproves staleness"
1170 );
1171 }
1172
1173 #[test]
1180 fn cache_cold_reason_stale_disprove_threshold_boundary() {
1181 let at_bar = Usage {
1182 prompt_tokens: 1000,
1183 completion_tokens: 1,
1184 total_tokens: 1001,
1185 prompt_tokens_details: Some(PromptTokensDetails { cached_tokens: 900 }), };
1187 assert_eq!(
1188 cache_cold_reason(true, false, Some(CACHE_TTL_SECS), &at_bar),
1189 None,
1190 "exactly at the disprove bar suppresses Stale"
1191 );
1192
1193 let just_under = Usage {
1194 prompt_tokens_details: Some(PromptTokensDetails { cached_tokens: 899 }),
1195 ..at_bar
1196 };
1197 assert_eq!(
1198 cache_cold_reason(true, false, Some(CACHE_TTL_SECS), &just_under),
1199 Some(CacheColdReason::Stale {
1200 idle_secs: CACHE_TTL_SECS
1201 }),
1202 "one token under the disprove bar must not suppress Stale"
1203 );
1204 }
1205
1206 #[test]
1212 fn cache_cold_reason_stale_not_suppressed_by_ambiguous_usage() {
1213 assert_eq!(
1214 cache_cold_reason(true, false, Some(20 * 60), &moderate_usage()),
1215 Some(CacheColdReason::Stale { idle_secs: 20 * 60 })
1216 );
1217 }
1218
1219 #[test]
1223 fn cache_cold_reason_stale_not_suppressed_by_missing_usage_details() {
1224 let no_details = Usage {
1225 prompt_tokens: 1000,
1226 completion_tokens: 20,
1227 total_tokens: 1020,
1228 prompt_tokens_details: None,
1229 };
1230 assert_eq!(
1231 cache_cold_reason(true, false, Some(20 * 60), &no_details),
1232 Some(CacheColdReason::Stale { idle_secs: 20 * 60 })
1233 );
1234 }
1235
1236 #[test]
1242 fn cache_cold_reason_fires_stale_at_ttl_boundary() {
1243 assert_eq!(
1244 cache_cold_reason(true, true, Some(CACHE_TTL_SECS), &moderate_usage()),
1245 Some(CacheColdReason::Stale {
1246 idle_secs: CACHE_TTL_SECS
1247 })
1248 );
1249 assert_eq!(
1250 cache_cold_reason(true, true, Some(CACHE_TTL_SECS - 1), &moderate_usage()),
1251 None,
1252 "one second under the TTL must not fire"
1253 );
1254 }
1255
1256 #[test]
1259 fn cache_cold_reason_fires_miss_on_low_ratio_inside_ttl() {
1260 assert_eq!(
1261 cache_cold_reason(true, true, Some(1), &cold_usage()),
1262 Some(CacheColdReason::Miss {
1263 cached_tokens: 3,
1264 prompt_tokens: 1000,
1265 })
1266 );
1267 }
1268
1269 #[test]
1272 fn cache_cold_reason_ratio_threshold_is_exclusive() {
1273 let at_threshold = Usage {
1274 prompt_tokens: 1000,
1275 completion_tokens: 1,
1276 total_tokens: 1001,
1277 prompt_tokens_details: Some(PromptTokensDetails {
1278 cached_tokens: 100, }),
1280 };
1281 assert_eq!(cache_cold_reason(true, true, Some(1), &at_threshold), None);
1282
1283 let just_under = Usage {
1284 prompt_tokens_details: Some(PromptTokensDetails { cached_tokens: 99 }),
1285 ..at_threshold
1286 };
1287 assert!(cache_cold_reason(true, true, Some(1), &just_under).is_some());
1288 }
1289
1290 #[test]
1294 fn cache_cold_reason_no_verdict_without_usage_details() {
1295 let usage = Usage {
1296 prompt_tokens: 1000,
1297 completion_tokens: 5,
1298 total_tokens: 1005,
1299 prompt_tokens_details: None,
1300 };
1301 assert_eq!(cache_cold_reason(true, true, Some(1), &usage), None);
1302 }
1303
1304 #[test]
1307 fn cache_cold_reason_no_verdict_on_zero_prompt_tokens() {
1308 let usage = Usage {
1309 prompt_tokens: 0,
1310 completion_tokens: 5,
1311 total_tokens: 5,
1312 prompt_tokens_details: Some(PromptTokensDetails { cached_tokens: 0 }),
1313 };
1314 assert_eq!(cache_cold_reason(true, true, Some(1), &usage), None);
1315 }
1316
1317 #[test]
1323 fn is_anthropic_family_model_recognizes_vendor_prefixed_slugs() {
1324 assert!(is_anthropic_family_model("anthropic/claude-opus-4-8"));
1325 assert!(is_anthropic_family_model("anthropic/claude-sonnet-4-6"));
1326 assert!(is_anthropic_family_model("anthropic/claude-haiku-4-5"));
1327 }
1328
1329 #[test]
1333 fn is_anthropic_family_model_recognizes_bare_claude_slugs() {
1334 assert!(is_anthropic_family_model("claude-opus-4-8"));
1335 assert!(is_anthropic_family_model("claude-3-5-sonnet-20241022"));
1336 }
1337
1338 #[test]
1342 fn is_anthropic_family_model_rejects_other_known_vendors() {
1343 assert!(!is_anthropic_family_model("openai/gpt-5"));
1344 assert!(!is_anthropic_family_model("openai/gpt-5.5"));
1345 assert!(!is_anthropic_family_model("google/gemini-2.5-pro"));
1346 assert!(!is_anthropic_family_model("deepseek/deepseek-v4-pro"));
1347 assert!(!is_anthropic_family_model("meta-llama/llama-4-maverick"));
1348 }
1349
1350 #[test]
1353 fn is_anthropic_family_model_does_not_assume_unknown_slugs() {
1354 assert!(!is_anthropic_family_model("my-custom-local-model"));
1355 assert!(!is_anthropic_family_model(""));
1356 }
1357
1358 #[test]
1359 fn truncate_never_splits_a_codepoint() {
1360 let s = "é".repeat(2000); let out = truncate(&s, 2001); assert!(out.ends_with('…'));
1364 assert!(out.len() <= 2001 + '…'.len_utf8());
1365 }
1366
1367 fn drain(lines: &[&str]) -> Accumulator {
1369 let mut acc = Accumulator::default();
1370 let mut deltas = Vec::new();
1371 let mut buf: Vec<u8> = Vec::new();
1372 for l in lines {
1373 buf.extend_from_slice(l.as_bytes());
1374 buf.push(b'\n');
1375 }
1376 drain_sse_lines(&mut buf, &mut acc, &mut deltas).unwrap();
1377 acc
1378 }
1379
1380 #[test]
1381 fn streaming_assembles_tool_calls_and_usage_across_deltas() {
1382 let acc = drain(&[
1383 r#"data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"read_"}}]}}]}"#,
1384 r#"data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"name":"file","arguments":"{\"path\":"}}]}}]}"#,
1385 r#"data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"a\"}"}}]}}]}"#,
1386 r#"data: {"choices":[{"delta":{"content":"done"}}]}"#,
1387 r#"data: {"usage":{"prompt_tokens":3,"completion_tokens":5}}"#,
1388 "data: [DONE]",
1389 ]);
1390 let msg = acc.to_message();
1391 let calls = msg.tool_calls.expect("tool calls");
1392 assert_eq!(calls.len(), 1);
1393 assert_eq!(calls[0].id, "call_1");
1394 assert_eq!(
1395 calls[0].function.name, "read_file",
1396 "name spread over deltas"
1397 );
1398 assert_eq!(calls[0].function.arguments, r#"{"path":"a"}"#);
1399 assert_eq!(msg.content.as_deref(), Some("done"));
1400 assert_eq!(acc.usage.completion_tokens, 5);
1401 }
1402
1403 #[test]
1404 fn streaming_tolerates_done_keepalive_and_blank_lines() {
1405 let acc = drain(&[
1407 "",
1408 ": keep-alive",
1409 r#"data: {"choices":[{"delta":{"content":"hi"}}]}"#,
1410 "data: not-json",
1411 "data: [DONE]",
1412 ]);
1413 assert_eq!(acc.to_message().content.as_deref(), Some("hi"));
1414 }
1415
1416 #[tokio::test]
1417 async fn non_success_status_becomes_provider_error() {
1418 use crate::RuntimeError as Error;
1419 use tokio::io::{AsyncReadExt, AsyncWriteExt};
1420
1421 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
1422 let addr = listener.local_addr().unwrap();
1423 let server = tokio::spawn(async move {
1424 let (mut sock, _) = listener.accept().await.unwrap();
1425 let mut buf = [0u8; 2048];
1426 let _ = sock.read(&mut buf).await;
1427 let body = r#"{"error":{"message":"bad key"}}"#;
1428 let resp = format!(
1429 "HTTP/1.1 401 Unauthorized\r\nContent-Length: {}\r\nContent-Type: application/json\r\nConnection: close\r\n\r\n{}",
1430 body.len(),
1431 body
1432 );
1433 sock.write_all(resp.as_bytes()).await.unwrap();
1434 sock.flush().await.unwrap();
1435 });
1436
1437 let provider = OpenAiProvider::new(format!("http://{addr}"), "k", HashMap::new());
1438 let req = ChatRequest::new("m", vec![ChatMessage::user("hi")]);
1439 let err = provider.complete(&req, &|_: &str| {}).await.unwrap_err();
1440 match err {
1441 Error::Provider { status, body } => {
1442 assert_eq!(status, 401);
1443 assert!(body.contains("bad key"), "body: {body}");
1444 }
1445 other => panic!("expected Provider error, got: {other:?}"),
1446 }
1447 server.await.unwrap();
1448 }
1449
1450 #[tokio::test]
1451 async fn streams_a_200_response_into_a_message() {
1452 use std::sync::{Arc, Mutex};
1453 use tokio::io::{AsyncReadExt, AsyncWriteExt};
1454
1455 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
1456 let addr = listener.local_addr().unwrap();
1457 let server = tokio::spawn(async move {
1458 let (mut sock, _) = listener.accept().await.unwrap();
1459 let mut buf = [0u8; 2048];
1460 let _ = sock.read(&mut buf).await;
1461 let sse = "data: {\"choices\":[{\"delta\":{\"content\":\"hel\"}}]}\n\n\
1462 data: {\"choices\":[{\"delta\":{\"content\":\"lo\"}}]}\n\n\
1463 data: [DONE]\n\n";
1464 let resp = format!(
1465 "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
1466 sse.len(),
1467 sse
1468 );
1469 sock.write_all(resp.as_bytes()).await.unwrap();
1470 sock.flush().await.unwrap();
1471 });
1472
1473 let provider = OpenAiProvider::new(format!("http://{addr}"), "k", HashMap::new());
1474 let req = ChatRequest::new("m", vec![ChatMessage::user("hi")]);
1475 let seen = Arc::new(Mutex::new(String::new()));
1476 let seen2 = seen.clone();
1477 let on_delta = move |s: &str| seen2.lock().unwrap().push_str(s);
1478 let (msg, _usage) = provider.complete(&req, &on_delta).await.unwrap();
1479 assert_eq!(msg.content.as_deref(), Some("hello"));
1480 assert_eq!(*seen.lock().unwrap(), "hello", "deltas streamed live");
1481 server.await.unwrap();
1482 }
1483
1484 #[test]
1487 fn from_retry_config_unset_is_byte_identical_to_default() {
1488 let opts = HttpOptions::from_retry_config(true, None, None);
1489 let default = HttpOptions::default();
1490 assert_eq!(opts.max_retries, default.max_retries);
1491 assert_eq!(opts.retry_backoff_base, default.retry_backoff_base);
1492 assert_eq!(opts.connect_timeout, default.connect_timeout);
1493 assert_eq!(opts.read_idle_timeout, default.read_idle_timeout);
1494 }
1495
1496 #[test]
1497 fn from_retry_config_disabled_forces_zero_retries() {
1498 let opts = HttpOptions::from_retry_config(false, None, None);
1499 assert_eq!(opts.max_retries, 0);
1500 assert_eq!(
1503 opts.retry_backoff_base,
1504 HttpOptions::default().retry_backoff_base
1505 );
1506 }
1507
1508 #[test]
1509 fn from_retry_config_disabled_with_explicit_max_retries_still_forces_zero() {
1510 let opts = HttpOptions::from_retry_config(false, Some(5), None);
1513 assert_eq!(opts.max_retries, 0);
1514 }
1515
1516 #[test]
1517 fn from_retry_config_overrides_apply_when_enabled() {
1518 let opts = HttpOptions::from_retry_config(true, Some(7), Some(1234));
1519 assert_eq!(opts.max_retries, 7);
1520 assert_eq!(opts.retry_backoff_base, Duration::from_millis(1234));
1521 }
1522
1523 #[test]
1524 fn from_retry_config_partial_override_leaves_the_other_at_default() {
1525 let opts = HttpOptions::from_retry_config(true, Some(9), None);
1526 assert_eq!(opts.max_retries, 9);
1527 assert_eq!(
1528 opts.retry_backoff_base,
1529 HttpOptions::default().retry_backoff_base
1530 );
1531 }
1532
1533 fn test_http_options() -> HttpOptions {
1536 HttpOptions {
1537 connect_timeout: Duration::from_millis(250),
1538 read_idle_timeout: Duration::from_millis(250),
1539 max_retries: 2,
1540 retry_backoff_base: Duration::from_millis(10),
1541 }
1542 }
1543
1544 #[tokio::test]
1545 async fn hung_connection_errors_via_read_timeout_within_bounded_time() {
1546 use tokio::io::AsyncReadExt;
1547
1548 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
1549 let addr = listener.local_addr().unwrap();
1550 let server = tokio::spawn(async move {
1555 loop {
1556 let Ok((mut sock, _)) = listener.accept().await else {
1557 break;
1558 };
1559 tokio::spawn(async move {
1560 let mut buf = [0u8; 2048];
1561 let _ = sock.read(&mut buf).await;
1562 tokio::time::sleep(Duration::from_secs(2)).await;
1565 });
1566 }
1567 });
1568
1569 let provider = OpenAiProvider::new_with_options(
1570 format!("http://{addr}"),
1571 "k",
1572 HashMap::new(),
1573 test_http_options(),
1574 );
1575 let req = ChatRequest::new("m", vec![ChatMessage::user("hi")]);
1576
1577 let outcome = tokio::time::timeout(Duration::from_secs(5), async {
1584 provider.complete(&req, &|_: &str| {}).await
1585 })
1586 .await
1587 .expect("complete() must return within the outer bound, not hang forever");
1588
1589 match outcome {
1590 Err(Error::Http(_)) => {}
1591 other => panic!("expected Err(Error::Http(_)) from the read timeout, got: {other:?}"),
1592 }
1593
1594 server.abort();
1595 }
1596
1597 #[tokio::test]
1598 async fn retries_503_then_succeeds_with_exactly_two_requests() {
1599 use std::sync::atomic::{AtomicUsize, Ordering};
1600 use std::sync::Arc;
1601 use tokio::io::{AsyncReadExt, AsyncWriteExt};
1602
1603 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
1604 let addr = listener.local_addr().unwrap();
1605 let connections = Arc::new(AtomicUsize::new(0));
1606 let connections2 = connections.clone();
1607 let server = tokio::spawn(async move {
1608 for _ in 0..2 {
1609 let (mut sock, _) = listener.accept().await.unwrap();
1610 let n = connections2.fetch_add(1, Ordering::SeqCst) + 1;
1611 let mut buf = [0u8; 2048];
1612 let _ = sock.read(&mut buf).await;
1613 if n == 1 {
1614 let resp =
1615 "HTTP/1.1 503 Service Unavailable\r\nContent-Length: 0\r\nConnection: close\r\n\r\n";
1616 sock.write_all(resp.as_bytes()).await.unwrap();
1617 } else {
1618 let sse = "data: {\"choices\":[{\"delta\":{\"content\":\"hel\"}}]}\n\n\
1619 data: {\"choices\":[{\"delta\":{\"content\":\"lo\"}}]}\n\n\
1620 data: [DONE]\n\n";
1621 let resp = format!(
1622 "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
1623 sse.len(),
1624 sse
1625 );
1626 sock.write_all(resp.as_bytes()).await.unwrap();
1627 }
1628 sock.flush().await.unwrap();
1629 }
1630 });
1631
1632 let provider = OpenAiProvider::new_with_options(
1633 format!("http://{addr}"),
1634 "k",
1635 HashMap::new(),
1636 test_http_options(),
1637 );
1638 let req = ChatRequest::new("m", vec![ChatMessage::user("hi")]);
1639 let (msg, _usage) = provider.complete(&req, &|_: &str| {}).await.unwrap();
1640 assert_eq!(msg.content.as_deref(), Some("hello"));
1641 server.await.unwrap();
1642 assert_eq!(
1643 connections.load(Ordering::SeqCst),
1644 2,
1645 "exactly 2 requests made: one 503, one successful retry"
1646 );
1647 }
1648
1649 #[test]
1650 fn streaming_decodes_multibyte_across_chunk_boundaries() {
1651 let line = "data: {\"choices\":[{\"delta\":{\"content\":\"héllo🌍\"}}]}\n";
1654 let bytes = line.as_bytes();
1655 let mut deltas = Vec::new();
1656 for split in 1..bytes.len() {
1658 let mut acc = Accumulator::default();
1659 let mut buf: Vec<u8> = Vec::new();
1660 buf.extend_from_slice(&bytes[..split]);
1661 drain_sse_lines(&mut buf, &mut acc, &mut deltas).unwrap();
1662 buf.extend_from_slice(&bytes[split..]);
1663 drain_sse_lines(&mut buf, &mut acc, &mut deltas).unwrap();
1664 assert_eq!(acc.content, "héllo🌍", "split at byte {split}");
1665 assert!(!acc.content.contains('\u{FFFD}'));
1666 }
1667 }
1668}