1use std::collections::HashMap;
9use std::time::Duration;
10
11use async_trait::async_trait;
12use futures::StreamExt;
13use serde::{Deserialize, Serialize};
14
15use crate::config::CachePlan;
16use crate::error::{Error, Result};
17use crate::message::{ChatMessage, FunctionCall, Role, ToolCall};
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)]
45pub(crate) struct HttpOptions {
46 pub(crate) connect_timeout: Duration,
47 pub(crate) read_idle_timeout: Duration,
48 pub(crate) max_retries: u32,
49 pub(crate) retry_backoff_base: Duration,
50}
51
52impl Default for HttpOptions {
53 fn default() -> Self {
54 HttpOptions {
55 connect_timeout: CONNECT_TIMEOUT,
56 read_idle_timeout: READ_IDLE_TIMEOUT,
57 max_retries: MAX_RETRIES,
58 retry_backoff_base: RETRY_BACKOFF_BASE,
59 }
60 }
61}
62
63impl HttpOptions {
64 pub(crate) fn from_retry_config(
77 enabled: bool,
78 max_retries: Option<u32>,
79 base_delay_ms: Option<u64>,
80 ) -> HttpOptions {
81 let base = HttpOptions::default();
82 HttpOptions {
83 max_retries: if enabled {
84 max_retries.unwrap_or(base.max_retries)
85 } else {
86 0
87 },
88 retry_backoff_base: base_delay_ms
89 .map(Duration::from_millis)
90 .unwrap_or(base.retry_backoff_base),
91 ..base
92 }
93 }
94}
95
96#[derive(Debug, Clone, PartialEq, Serialize)]
98#[non_exhaustive]
99pub struct ToolSchema {
100 pub name: String,
102 pub description: String,
104 pub parameters: serde_json::Value,
106}
107
108impl ToolSchema {
109 pub fn new(
115 name: impl Into<String>,
116 description: impl Into<String>,
117 parameters: serde_json::Value,
118 ) -> Self {
119 ToolSchema {
120 name: name.into(),
121 description: description.into(),
122 parameters,
123 }
124 }
125}
126
127#[derive(Debug, Clone, PartialEq)]
129#[non_exhaustive]
130pub struct ChatRequest {
131 pub model: String,
133 pub messages: Vec<ChatMessage>,
135 pub tools: Vec<ToolSchema>,
137 pub temperature: Option<f32>,
139 pub max_tokens: Option<u32>,
141 pub effort: Option<String>,
143 pub response_format: Option<serde_json::Value>,
146 pub extra_body: serde_json::Map<String, serde_json::Value>,
150}
151
152impl ChatRequest {
153 pub fn new(model: impl Into<String>, messages: Vec<ChatMessage>) -> Self {
155 ChatRequest {
156 model: model.into(),
157 messages,
158 tools: Vec::new(),
159 temperature: None,
160 max_tokens: None,
161 effort: None,
162 response_format: None,
163 extra_body: serde_json::Map::new(),
164 }
165 }
166}
167
168pub(crate) fn build_request_body(req: &ChatRequest, stream: bool) -> serde_json::Value {
172 use serde_json::json;
173 let mut body = json!({
174 "model": req.model,
175 "messages": req.messages,
176 "stream": stream,
177 });
178 let obj = body.as_object_mut().unwrap();
179 if !req.tools.is_empty() {
180 obj.insert(
181 "tools".into(),
182 serde_json::to_value(req.tools.iter().map(WireTool::from).collect::<Vec<_>>()).unwrap(),
183 );
184 }
185 if let Some(t) = req.temperature {
186 obj.insert("temperature".into(), json!(t));
187 }
188 if let Some(m) = req.max_tokens {
189 obj.insert("max_tokens".into(), json!(m));
190 }
191 if let Some(e) = &req.effort {
192 obj.insert("reasoning_effort".into(), json!(e));
193 }
194 if let Some(rf) = &req.response_format {
195 obj.insert("response_format".into(), rf.clone());
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
207pub(crate) fn apply_cache_plan(
233 messages: &[ChatMessage],
234 plan: CachePlan,
235 imported_prefix_len: Option<usize>,
236) -> Vec<ChatMessage> {
237 let mut out = messages.to_vec();
238 if !matches!(plan, CachePlan::ImportedPrefix) {
239 return out;
240 }
241 let Some(len) = imported_prefix_len.filter(|&n| n > 0) else {
242 return out;
243 };
244 let last = len - 1;
245 let mut targets = vec![0usize];
246 if last != 0 {
247 targets.push(last);
248 }
249 for idx in targets {
250 if let Some(msg) = out.get_mut(idx) {
251 annotate_cache_breakpoint(msg);
252 }
253 }
254 out
255}
256
257fn annotate_cache_breakpoint(msg: &mut ChatMessage) {
260 let cache_control = serde_json::json!({"type": "ephemeral"});
261 if let Some(parts) = msg.content_parts.as_mut() {
262 if let Some(text_part) = parts
264 .iter_mut()
265 .rev()
266 .find(|p| p.get("type").and_then(serde_json::Value::as_str) == Some("text"))
267 {
268 if let Some(obj) = text_part.as_object_mut() {
269 obj.insert("cache_control".to_string(), cache_control);
270 }
271 }
272 return;
273 }
274 let text = msg.content.take().unwrap_or_default();
275 msg.content_parts = Some(vec![serde_json::json!({
276 "type": "text",
277 "text": text,
278 "cache_control": cache_control,
279 })]);
280}
281
282pub(crate) fn tier_change_is_cache_bust(previous: Option<u64>, current: u64) -> bool {
297 previous.is_some_and(|p| p != current)
298}
299
300#[derive(Debug, Clone, Default, Deserialize)]
302pub struct Usage {
303 #[serde(default)]
305 pub prompt_tokens: u64,
306 #[serde(default)]
308 pub completion_tokens: u64,
309 #[serde(default)]
311 pub total_tokens: u64,
312 #[serde(default)]
314 pub prompt_tokens_details: Option<PromptTokensDetails>,
315}
316
317#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Deserialize)]
321pub struct PromptTokensDetails {
322 #[serde(default)]
324 pub cached_tokens: u64,
325}
326
327pub(crate) const CACHE_TTL_SECS: i64 = 300;
334
335pub(crate) const CACHE_MISS_RATIO_THRESHOLD: f64 = 0.10;
341
342pub(crate) const CACHE_STALE_DISPROVE_RATIO_THRESHOLD: f64 = 1.0 - CACHE_MISS_RATIO_THRESHOLD;
360
361pub(crate) fn is_anthropic_family_model(model: &str) -> bool {
379 model.starts_with("anthropic/") || model.starts_with("claude-") || model.starts_with("claude/")
380}
381
382#[derive(Debug, Clone, Copy, PartialEq)]
385pub(crate) enum CacheColdReason {
386 Stale {
391 idle_secs: i64,
393 },
394 Miss {
399 cached_tokens: u64,
401 prompt_tokens: u64,
403 },
404}
405
406impl CacheColdReason {
407 pub(crate) fn message(&self) -> String {
409 match self {
410 CacheColdReason::Stale { idle_secs } => format!(
411 "cache likely cold — this turn was sent {}m{:02}s after the cache was last \
412 refreshed (Anthropic's ephemeral prompt cache expires after 5m idle) — this \
413 turn likely paid full input cost for the cached prefix",
414 idle_secs / 60,
415 idle_secs % 60,
416 ),
417 CacheColdReason::Miss {
418 cached_tokens,
419 prompt_tokens,
420 } => format!(
421 "unexpected cache miss — only {cached_tokens}/{prompt_tokens} prompt tokens \
422 were served from cache this turn even though reuse was expected — this turn \
423 likely paid full input cost for the cached prefix",
424 ),
425 }
426 }
427}
428
429pub(crate) fn cache_cold_reason(
477 will_annotate: bool,
478 cache_established: bool,
479 idle_secs: Option<i64>,
480 usage: &Usage,
481) -> Option<CacheColdReason> {
482 if !will_annotate {
483 return None;
484 }
485 if let Some(idle_secs) = idle_secs {
486 if idle_secs >= CACHE_TTL_SECS {
487 let disproven_by_usage = usage
488 .prompt_tokens_details
489 .filter(|_| usage.prompt_tokens > 0)
490 .is_some_and(|details| {
491 details.cached_tokens as f64 / usage.prompt_tokens as f64
492 >= CACHE_STALE_DISPROVE_RATIO_THRESHOLD
493 });
494 if !disproven_by_usage {
495 return Some(CacheColdReason::Stale { idle_secs });
496 }
497 }
498 }
499 if !cache_established {
500 return None;
503 }
504 let details = usage.prompt_tokens_details?;
505 if usage.prompt_tokens == 0 {
506 return None;
509 }
510 let ratio = details.cached_tokens as f64 / usage.prompt_tokens as f64;
511 if ratio < CACHE_MISS_RATIO_THRESHOLD {
512 return Some(CacheColdReason::Miss {
513 cached_tokens: details.cached_tokens,
514 prompt_tokens: usage.prompt_tokens,
515 });
516 }
517 None
518}
519
520const KNOWN_MODEL_CONTEXT_LIMITS: &[(&str, u64)] = &[
531 ("z-ai/glm-5.2", 1_048_576),
532 ("deepseek/deepseek-v4-flash", 1_048_576),
533 ("deepseek/deepseek-v4-pro", 1_048_576),
534 ("google/gemini-2.5-pro", 1_048_576),
535 ("meta-llama/llama-4-maverick", 1_048_576),
536 ("openai/gpt-5.5", 400_000),
537 ("openai/gpt-5", 400_000),
538 ("anthropic/claude-opus-4-8", 500_000),
539 ("anthropic/claude-sonnet-4-6", 500_000),
540 ("anthropic/claude-haiku-4-5", 200_000),
541];
542
543pub const UNKNOWN_MODEL_CONTEXT_FLOOR: u64 = 200_000;
551
552pub fn model_context_limit(model: &str) -> Option<u64> {
556 KNOWN_MODEL_CONTEXT_LIMITS
557 .iter()
558 .find(|(slug, _)| *slug == model)
559 .map(|(_, limit)| *limit)
560}
561
562#[async_trait]
565pub trait Provider: Send + Sync {
566 async fn complete(
569 &self,
570 req: &ChatRequest,
571 on_delta: &(dyn for<'a> Fn(&'a str) + Send + Sync),
572 ) -> Result<(ChatMessage, Usage)>;
573}
574
575pub struct OpenAiProvider {
578 client: reqwest::Client,
579 base_url: String,
580 api_key: String,
581 extra_headers: HashMap<String, String>,
582 http_options: HttpOptions,
583}
584
585impl OpenAiProvider {
586 pub fn new(
588 base_url: impl Into<String>,
589 api_key: impl Into<String>,
590 extra_headers: HashMap<String, String>,
591 ) -> Self {
592 Self::new_with_options(base_url, api_key, extra_headers, HttpOptions::default())
593 }
594
595 pub(crate) fn new_with_options(
600 base_url: impl Into<String>,
601 api_key: impl Into<String>,
602 extra_headers: HashMap<String, String>,
603 http_options: HttpOptions,
604 ) -> Self {
605 OpenAiProvider {
606 client: reqwest::Client::builder()
607 .connect_timeout(http_options.connect_timeout)
608 .read_timeout(http_options.read_idle_timeout)
609 .build()
610 .expect("static reqwest client config cannot fail"),
611 base_url: base_url.into(),
612 api_key: api_key.into(),
613 extra_headers,
614 http_options,
615 }
616 }
617
618 fn endpoint(&self) -> String {
619 format!("{}/chat/completions", self.base_url.trim_end_matches('/'))
620 }
621
622 async fn send_with_retry(&self, wire: &serde_json::Value) -> Result<reqwest::Response> {
629 let mut attempt = 0u32;
630 loop {
631 let mut builder = self
632 .client
633 .post(self.endpoint())
634 .bearer_auth(&self.api_key)
635 .header("Content-Type", "application/json");
636 for (k, v) in &self.extra_headers {
637 builder = builder.header(k, v);
638 }
639
640 let sent = builder.json(wire).send().await;
641 let (retryable, result): (bool, Result<reqwest::Response>) = match sent {
642 Err(e) => (true, Err(Error::from(e))),
643 Ok(resp) => {
644 let status = resp.status();
645 if status.is_success() {
646 (false, Ok(resp))
647 } else if status.is_server_error() {
648 let body = resp.text().await.unwrap_or_default();
649 (
650 true,
651 Err(Error::Provider {
652 status: status.as_u16(),
653 body: truncate(&body, 2000),
654 }),
655 )
656 } else {
657 let body = resp.text().await.unwrap_or_default();
658 (
659 false,
660 Err(Error::Provider {
661 status: status.as_u16(),
662 body: truncate(&body, 2000),
663 }),
664 )
665 }
666 }
667 };
668
669 if !retryable || attempt >= self.http_options.max_retries {
670 return result;
671 }
672 let backoff = self.http_options.retry_backoff_base * 2u32.pow(attempt);
673 tokio::time::sleep(backoff).await;
674 attempt += 1;
675 }
676 }
677}
678
679#[async_trait]
680impl Provider for OpenAiProvider {
681 async fn complete(
682 &self,
683 req: &ChatRequest,
684 on_delta: &(dyn for<'a> Fn(&'a str) + Send + Sync),
685 ) -> Result<(ChatMessage, Usage)> {
686 let wire = build_request_body(req, true);
687
688 let resp = self.send_with_retry(&wire).await?;
689
690 let mut acc = Accumulator::default();
691 let mut buf: Vec<u8> = Vec::new();
697 let mut deltas: Vec<String> = Vec::new();
698 let mut stream = resp.bytes_stream();
699 while let Some(chunk) = stream.next().await {
700 let bytes = chunk?;
701 buf.extend_from_slice(&bytes);
702 drain_sse_lines(&mut buf, &mut acc, &mut deltas)?;
703 for d in deltas.drain(..) {
704 on_delta(&d);
705 }
706 }
707 let tail = String::from_utf8_lossy(&buf);
709 if !tail.trim().is_empty() {
710 handle_sse_line(tail.trim(), &mut acc, &mut deltas)?;
711 for d in deltas.drain(..) {
712 on_delta(&d);
713 }
714 }
715
716 Ok((acc.to_message(), acc_usage(&acc)))
717 }
718}
719
720#[derive(Default)]
723struct Accumulator {
724 content: String,
725 tool_calls: Vec<ToolCallAccum>,
726 usage: Usage,
727}
728
729#[derive(Default)]
730struct ToolCallAccum {
731 id: String,
732 name: String,
733 arguments: String,
734}
735
736impl Accumulator {
737 fn ensure(&mut self, index: usize) -> &mut ToolCallAccum {
738 while self.tool_calls.len() <= index {
739 self.tool_calls.push(ToolCallAccum::default());
740 }
741 &mut self.tool_calls[index]
742 }
743
744 fn to_message(&self) -> ChatMessage {
745 let calls: Vec<ToolCall> = self
746 .tool_calls
747 .iter()
748 .filter(|c| !c.id.is_empty() || !c.name.is_empty())
749 .map(|c| ToolCall {
750 id: c.id.clone(),
751 kind: "function".to_string(),
752 function: FunctionCall {
753 name: c.name.clone(),
754 arguments: c.arguments.clone(),
755 },
756 })
757 .collect();
758 ChatMessage {
759 role: Role::Assistant,
760 content: (!self.content.is_empty()).then(|| self.content.clone()),
761 content_parts: None,
762 tool_calls: (!calls.is_empty()).then_some(calls),
763 tool_call_id: None,
764 name: None,
765 metadata: Default::default(),
766 }
767 }
768}
769
770fn acc_usage(acc: &Accumulator) -> Usage {
771 acc.usage.clone()
772}
773
774fn drain_sse_lines(
775 buf: &mut Vec<u8>,
776 acc: &mut Accumulator,
777 deltas: &mut Vec<String>,
778) -> Result<()> {
779 while let Some(pos) = buf.iter().position(|&b| b == b'\n') {
780 let line: Vec<u8> = buf.drain(..=pos).collect();
781 let line = String::from_utf8_lossy(&line);
782 handle_sse_line(line.trim(), acc, deltas)?;
783 }
784 Ok(())
785}
786
787fn handle_sse_line(line: &str, acc: &mut Accumulator, deltas: &mut Vec<String>) -> Result<()> {
788 let Some(data) = line.strip_prefix("data:") else {
789 return Ok(());
790 };
791 let data = data.trim();
792 if data.is_empty() || data == "[DONE]" {
793 return Ok(());
794 }
795 let chunk: StreamChunk = match serde_json::from_str(data) {
796 Ok(c) => c,
797 Err(_) => return Ok(()), };
799 if let Some(u) = chunk.usage {
800 acc.usage = u;
801 }
802 for choice in chunk.choices {
803 if let Some(text) = choice.delta.content {
804 if !text.is_empty() {
805 acc.content.push_str(&text);
806 deltas.push(text);
807 }
808 }
809 for tc in choice.delta.tool_calls.unwrap_or_default() {
810 let slot = acc.ensure(tc.index);
811 if let Some(id) = tc.id {
812 slot.id = id;
813 }
814 if let Some(f) = tc.function {
815 if let Some(name) = f.name {
816 slot.name.push_str(&name);
817 }
818 if let Some(args) = f.arguments {
819 slot.arguments.push_str(&args);
820 }
821 }
822 }
823 }
824 Ok(())
825}
826
827fn truncate(s: &str, max: usize) -> String {
828 if s.len() <= max {
829 s.to_string()
830 } else {
831 let mut end = max;
834 while end > 0 && !s.is_char_boundary(end) {
835 end -= 1;
836 }
837 format!("{}…", &s[..end])
838 }
839}
840
841#[derive(Serialize)]
844struct WireTool<'a> {
845 #[serde(rename = "type")]
846 kind: &'static str,
847 function: WireFunction<'a>,
848}
849
850#[derive(Serialize)]
851struct WireFunction<'a> {
852 name: &'a str,
853 description: &'a str,
854 parameters: &'a serde_json::Value,
855}
856
857impl<'a> From<&'a ToolSchema> for WireTool<'a> {
858 fn from(t: &'a ToolSchema) -> Self {
859 WireTool {
860 kind: "function",
861 function: WireFunction {
862 name: &t.name,
863 description: &t.description,
864 parameters: &t.parameters,
865 },
866 }
867 }
868}
869
870#[derive(Deserialize)]
871struct StreamChunk {
872 #[serde(default)]
873 choices: Vec<StreamChoice>,
874 #[serde(default)]
875 usage: Option<Usage>,
876}
877
878#[derive(Deserialize)]
879struct StreamChoice {
880 delta: Delta,
881}
882
883#[derive(Deserialize)]
884struct Delta {
885 #[serde(default)]
886 content: Option<String>,
887 #[serde(default)]
888 tool_calls: Option<Vec<ToolCallDelta>>,
889}
890
891#[derive(Deserialize)]
892struct ToolCallDelta {
893 #[serde(default)]
894 index: usize,
895 #[serde(default)]
896 id: Option<String>,
897 #[serde(default)]
898 function: Option<FnDelta>,
899}
900
901#[derive(Deserialize)]
902struct FnDelta {
903 #[serde(default)]
904 name: Option<String>,
905 #[serde(default)]
906 arguments: Option<String>,
907}
908
909#[cfg(test)]
910mod tests {
911 use super::*;
912 use crate::message::ChatMessage;
913
914 #[test]
915 fn request_body_includes_effort_format_and_passthrough() {
916 let mut req = ChatRequest::new("m", vec![ChatMessage::user("hi")]);
917 req.effort = Some("high".into());
918 req.response_format =
919 Some(serde_json::json!({"type": "json_schema", "json_schema": {"name": "x"}}));
920 req.extra_body.insert(
921 "cache_control".into(),
922 serde_json::json!({"type": "ephemeral"}),
923 );
924 req.extra_body.insert(
925 "provider".into(),
926 serde_json::json!({"order": ["anthropic"]}),
927 );
928
929 let body = build_request_body(&req, false);
930 assert_eq!(body["model"], "m");
931 assert_eq!(body["reasoning_effort"], "high");
932 assert_eq!(body["response_format"]["type"], "json_schema");
933 assert_eq!(body["cache_control"]["type"], "ephemeral");
934 assert_eq!(body["provider"]["order"][0], "anthropic");
935 assert!(body.get("stream_options").is_none());
937 }
938
939 #[test]
940 fn extra_body_overrides_modeled_fields() {
941 let mut req = ChatRequest::new("m", vec![ChatMessage::user("hi")]);
942 req.max_tokens = Some(100);
943 req.extra_body
944 .insert("max_tokens".into(), serde_json::json!(999));
945 let body = build_request_body(&req, true);
946 assert_eq!(body["max_tokens"], 999, "extra_body wins");
947 assert_eq!(body["stream_options"]["include_usage"], true);
948 }
949
950 #[test]
959 fn cache_plan_annotates_system_and_last_imported_message_only() {
960 let messages = vec![
961 ChatMessage::system("sys"),
962 ChatMessage::user("u1"),
963 ChatMessage::assistant("a1"),
964 ChatMessage::user("u2"),
965 ];
966 let mut req = ChatRequest::new("m", messages);
967 req.messages = apply_cache_plan(&req.messages, crate::CachePlan::ImportedPrefix, Some(3));
968
969 let body = build_request_body(&req, false);
970 let msgs = body["messages"].as_array().unwrap();
971 assert_eq!(msgs.len(), 4, "annotation must not change message count");
972
973 assert_eq!(
974 msgs[0]["content"][0]["cache_control"]["type"], "ephemeral",
975 "breakpoint 1: system message"
976 );
977 assert_eq!(
978 msgs[2]["content"][0]["cache_control"]["type"], "ephemeral",
979 "breakpoint 2: last message of the imported prefix (a1)"
980 );
981 assert_eq!(
982 msgs[2]["content"][0]["text"], "a1",
983 "annotated text must be byte-identical to the original content"
984 );
985
986 for (i, m) in msgs.iter().enumerate() {
988 if i == 0 || i == 2 {
989 continue;
990 }
991 let has_cc = match &m["content"] {
992 serde_json::Value::Array(parts) => {
993 parts.iter().any(|p| p.get("cache_control").is_some())
994 }
995 serde_json::Value::String(_) => false,
996 _ => false,
997 };
998 assert!(!has_cc, "message {i} must not carry cache_control: {m:?}");
999 }
1000 }
1001
1002 #[test]
1003 fn cache_plan_off_never_annotates() {
1004 let messages = vec![ChatMessage::system("sys"), ChatMessage::user("u1")];
1005 let out = apply_cache_plan(&messages, crate::CachePlan::Off, Some(2));
1006 assert_eq!(out[0].content_parts, None);
1007 assert_eq!(out[1].content_parts, None);
1008 }
1009
1010 #[test]
1011 fn tier_change_is_cache_bust_truth_table() {
1012 assert!(!tier_change_is_cache_bust(None, 42));
1014 assert!(!tier_change_is_cache_bust(Some(42), 42));
1016 assert!(tier_change_is_cache_bust(Some(42), 7));
1018 }
1019
1020 #[test]
1021 fn cache_plan_dedupes_when_prefix_is_only_the_system_message() {
1022 let messages = vec![ChatMessage::system("sys"), ChatMessage::user("u1")];
1025 let out = apply_cache_plan(&messages, crate::CachePlan::ImportedPrefix, Some(1));
1026 assert!(out[0].content_parts.is_some());
1027 assert_eq!(out[1].content_parts, None);
1028 }
1029
1030 #[test]
1031 fn cache_plan_annotates_last_text_part_of_already_multimodal_message() {
1032 let imported_last = ChatMessage::user_with_images("caption", &["https://x/y.png".into()]);
1033 assert_eq!(
1035 imported_last.content_parts.as_ref().unwrap()[0]["type"],
1036 "text"
1037 );
1038 let messages = vec![ChatMessage::system("sys"), imported_last];
1039 let out = apply_cache_plan(&messages, crate::CachePlan::ImportedPrefix, Some(2));
1040 let parts = out[1].content_parts.as_ref().unwrap();
1041 assert_eq!(parts[0]["cache_control"]["type"], "ephemeral");
1042 assert_eq!(parts[0]["text"], "caption");
1043 assert!(
1044 parts[1].get("cache_control").is_none(),
1045 "the image_url part must not be annotated"
1046 );
1047 }
1048
1049 #[test]
1053 fn usage_parses_prompt_tokens_details_cached_tokens() {
1054 let acc = drain(&[
1055 r#"data: {"choices":[{"delta":{"content":"hi"}}]}"#,
1056 r#"data: {"usage":{"prompt_tokens":100,"completion_tokens":5,"prompt_tokens_details":{"cached_tokens":90}}}"#,
1057 "data: [DONE]",
1058 ]);
1059 assert_eq!(acc.usage.prompt_tokens, 100);
1060 let details = acc.usage.prompt_tokens_details.expect("details present");
1061 assert_eq!(details.cached_tokens, 90);
1062 }
1063
1064 fn warm_usage() -> Usage {
1067 Usage {
1070 prompt_tokens: 1000,
1071 completion_tokens: 20,
1072 total_tokens: 1020,
1073 prompt_tokens_details: Some(PromptTokensDetails { cached_tokens: 950 }),
1074 }
1075 }
1076
1077 fn cold_usage() -> Usage {
1078 Usage {
1080 prompt_tokens: 1000,
1081 completion_tokens: 20,
1082 total_tokens: 1020,
1083 prompt_tokens_details: Some(PromptTokensDetails { cached_tokens: 3 }),
1084 }
1085 }
1086
1087 fn moderate_usage() -> Usage {
1094 Usage {
1095 prompt_tokens: 1000,
1096 completion_tokens: 20,
1097 total_tokens: 1020,
1098 prompt_tokens_details: Some(PromptTokensDetails { cached_tokens: 500 }),
1099 }
1100 }
1101
1102 #[test]
1106 fn cache_cold_reason_never_fires_when_not_annotated() {
1107 assert_eq!(
1108 cache_cold_reason(false, true, Some(10_000), &cold_usage()),
1109 None
1110 );
1111 assert_eq!(cache_cold_reason(false, false, None, &cold_usage()), None);
1112 }
1113
1114 #[test]
1120 fn cache_cold_reason_first_annotated_request_never_reports_miss() {
1121 assert_eq!(
1122 cache_cold_reason(true, false, Some(1), &cold_usage()),
1123 None,
1124 "first write: a near-zero cache-read ratio is expected, not a miss"
1125 );
1126 }
1127
1128 #[test]
1132 fn cache_cold_reason_silent_on_warm_back_to_back_turn() {
1133 assert_eq!(cache_cold_reason(true, true, Some(5), &warm_usage()), None);
1134 assert_eq!(cache_cold_reason(true, true, None, &warm_usage()), None);
1137 }
1138
1139 #[test]
1148 fn cache_cold_reason_fires_stale_on_first_turn_of_a_resumed_idle_session() {
1149 assert_eq!(
1150 cache_cold_reason(true, false, Some(20 * 60), &cold_usage()),
1151 Some(CacheColdReason::Stale { idle_secs: 20 * 60 })
1152 );
1153 }
1154
1155 #[test]
1169 fn cache_cold_reason_stale_suppressed_when_usage_disproves_it() {
1170 assert_eq!(
1171 cache_cold_reason(true, false, Some(20 * 60), &warm_usage()),
1172 None,
1173 "cache_established == false, but usage still disproves staleness"
1174 );
1175 assert_eq!(
1176 cache_cold_reason(true, true, Some(CACHE_TTL_SECS), &warm_usage()),
1177 None,
1178 "cache_established == true, at the TTL boundary, usage disproves staleness"
1179 );
1180 }
1181
1182 #[test]
1189 fn cache_cold_reason_stale_disprove_threshold_boundary() {
1190 let at_bar = Usage {
1191 prompt_tokens: 1000,
1192 completion_tokens: 1,
1193 total_tokens: 1001,
1194 prompt_tokens_details: Some(PromptTokensDetails { cached_tokens: 900 }), };
1196 assert_eq!(
1197 cache_cold_reason(true, false, Some(CACHE_TTL_SECS), &at_bar),
1198 None,
1199 "exactly at the disprove bar suppresses Stale"
1200 );
1201
1202 let just_under = Usage {
1203 prompt_tokens_details: Some(PromptTokensDetails { cached_tokens: 899 }),
1204 ..at_bar
1205 };
1206 assert_eq!(
1207 cache_cold_reason(true, false, Some(CACHE_TTL_SECS), &just_under),
1208 Some(CacheColdReason::Stale {
1209 idle_secs: CACHE_TTL_SECS
1210 }),
1211 "one token under the disprove bar must not suppress Stale"
1212 );
1213 }
1214
1215 #[test]
1221 fn cache_cold_reason_stale_not_suppressed_by_ambiguous_usage() {
1222 assert_eq!(
1223 cache_cold_reason(true, false, Some(20 * 60), &moderate_usage()),
1224 Some(CacheColdReason::Stale { idle_secs: 20 * 60 })
1225 );
1226 }
1227
1228 #[test]
1232 fn cache_cold_reason_stale_not_suppressed_by_missing_usage_details() {
1233 let no_details = Usage {
1234 prompt_tokens: 1000,
1235 completion_tokens: 20,
1236 total_tokens: 1020,
1237 prompt_tokens_details: None,
1238 };
1239 assert_eq!(
1240 cache_cold_reason(true, false, Some(20 * 60), &no_details),
1241 Some(CacheColdReason::Stale { idle_secs: 20 * 60 })
1242 );
1243 }
1244
1245 #[test]
1251 fn cache_cold_reason_fires_stale_at_ttl_boundary() {
1252 assert_eq!(
1253 cache_cold_reason(true, true, Some(CACHE_TTL_SECS), &moderate_usage()),
1254 Some(CacheColdReason::Stale {
1255 idle_secs: CACHE_TTL_SECS
1256 })
1257 );
1258 assert_eq!(
1259 cache_cold_reason(true, true, Some(CACHE_TTL_SECS - 1), &moderate_usage()),
1260 None,
1261 "one second under the TTL must not fire"
1262 );
1263 }
1264
1265 #[test]
1268 fn cache_cold_reason_fires_miss_on_low_ratio_inside_ttl() {
1269 assert_eq!(
1270 cache_cold_reason(true, true, Some(1), &cold_usage()),
1271 Some(CacheColdReason::Miss {
1272 cached_tokens: 3,
1273 prompt_tokens: 1000,
1274 })
1275 );
1276 }
1277
1278 #[test]
1281 fn cache_cold_reason_ratio_threshold_is_exclusive() {
1282 let at_threshold = Usage {
1283 prompt_tokens: 1000,
1284 completion_tokens: 1,
1285 total_tokens: 1001,
1286 prompt_tokens_details: Some(PromptTokensDetails {
1287 cached_tokens: 100, }),
1289 };
1290 assert_eq!(cache_cold_reason(true, true, Some(1), &at_threshold), None);
1291
1292 let just_under = Usage {
1293 prompt_tokens_details: Some(PromptTokensDetails { cached_tokens: 99 }),
1294 ..at_threshold
1295 };
1296 assert!(cache_cold_reason(true, true, Some(1), &just_under).is_some());
1297 }
1298
1299 #[test]
1303 fn cache_cold_reason_no_verdict_without_usage_details() {
1304 let usage = Usage {
1305 prompt_tokens: 1000,
1306 completion_tokens: 5,
1307 total_tokens: 1005,
1308 prompt_tokens_details: None,
1309 };
1310 assert_eq!(cache_cold_reason(true, true, Some(1), &usage), None);
1311 }
1312
1313 #[test]
1316 fn cache_cold_reason_no_verdict_on_zero_prompt_tokens() {
1317 let usage = Usage {
1318 prompt_tokens: 0,
1319 completion_tokens: 5,
1320 total_tokens: 5,
1321 prompt_tokens_details: Some(PromptTokensDetails { cached_tokens: 0 }),
1322 };
1323 assert_eq!(cache_cold_reason(true, true, Some(1), &usage), None);
1324 }
1325
1326 #[test]
1332 fn is_anthropic_family_model_recognizes_vendor_prefixed_slugs() {
1333 assert!(is_anthropic_family_model("anthropic/claude-opus-4-8"));
1334 assert!(is_anthropic_family_model("anthropic/claude-sonnet-4-6"));
1335 assert!(is_anthropic_family_model("anthropic/claude-haiku-4-5"));
1336 }
1337
1338 #[test]
1342 fn is_anthropic_family_model_recognizes_bare_claude_slugs() {
1343 assert!(is_anthropic_family_model("claude-opus-4-8"));
1344 assert!(is_anthropic_family_model("claude-3-5-sonnet-20241022"));
1345 }
1346
1347 #[test]
1351 fn is_anthropic_family_model_rejects_other_known_vendors() {
1352 assert!(!is_anthropic_family_model("openai/gpt-5"));
1353 assert!(!is_anthropic_family_model("openai/gpt-5.5"));
1354 assert!(!is_anthropic_family_model("google/gemini-2.5-pro"));
1355 assert!(!is_anthropic_family_model("deepseek/deepseek-v4-pro"));
1356 assert!(!is_anthropic_family_model("meta-llama/llama-4-maverick"));
1357 }
1358
1359 #[test]
1362 fn is_anthropic_family_model_does_not_assume_unknown_slugs() {
1363 assert!(!is_anthropic_family_model("my-custom-local-model"));
1364 assert!(!is_anthropic_family_model(""));
1365 }
1366
1367 #[test]
1368 fn truncate_never_splits_a_codepoint() {
1369 let s = "é".repeat(2000); let out = truncate(&s, 2001); assert!(out.ends_with('…'));
1373 assert!(out.len() <= 2001 + '…'.len_utf8());
1374 }
1375
1376 fn drain(lines: &[&str]) -> Accumulator {
1378 let mut acc = Accumulator::default();
1379 let mut deltas = Vec::new();
1380 let mut buf: Vec<u8> = Vec::new();
1381 for l in lines {
1382 buf.extend_from_slice(l.as_bytes());
1383 buf.push(b'\n');
1384 }
1385 drain_sse_lines(&mut buf, &mut acc, &mut deltas).unwrap();
1386 acc
1387 }
1388
1389 #[test]
1390 fn streaming_assembles_tool_calls_and_usage_across_deltas() {
1391 let acc = drain(&[
1392 r#"data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"read_"}}]}}]}"#,
1393 r#"data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"name":"file","arguments":"{\"path\":"}}]}}]}"#,
1394 r#"data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"a\"}"}}]}}]}"#,
1395 r#"data: {"choices":[{"delta":{"content":"done"}}]}"#,
1396 r#"data: {"usage":{"prompt_tokens":3,"completion_tokens":5}}"#,
1397 "data: [DONE]",
1398 ]);
1399 let msg = acc.to_message();
1400 let calls = msg.tool_calls.expect("tool calls");
1401 assert_eq!(calls.len(), 1);
1402 assert_eq!(calls[0].id, "call_1");
1403 assert_eq!(
1404 calls[0].function.name, "read_file",
1405 "name spread over deltas"
1406 );
1407 assert_eq!(calls[0].function.arguments, r#"{"path":"a"}"#);
1408 assert_eq!(msg.content.as_deref(), Some("done"));
1409 assert_eq!(acc.usage.completion_tokens, 5);
1410 }
1411
1412 #[test]
1413 fn streaming_tolerates_done_keepalive_and_blank_lines() {
1414 let acc = drain(&[
1416 "",
1417 ": keep-alive",
1418 r#"data: {"choices":[{"delta":{"content":"hi"}}]}"#,
1419 "data: not-json",
1420 "data: [DONE]",
1421 ]);
1422 assert_eq!(acc.to_message().content.as_deref(), Some("hi"));
1423 }
1424
1425 #[tokio::test]
1426 async fn non_success_status_becomes_provider_error() {
1427 use crate::error::Error;
1428 use tokio::io::{AsyncReadExt, AsyncWriteExt};
1429
1430 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
1431 let addr = listener.local_addr().unwrap();
1432 let server = tokio::spawn(async move {
1433 let (mut sock, _) = listener.accept().await.unwrap();
1434 let mut buf = [0u8; 2048];
1435 let _ = sock.read(&mut buf).await;
1436 let body = r#"{"error":{"message":"bad key"}}"#;
1437 let resp = format!(
1438 "HTTP/1.1 401 Unauthorized\r\nContent-Length: {}\r\nContent-Type: application/json\r\nConnection: close\r\n\r\n{}",
1439 body.len(),
1440 body
1441 );
1442 sock.write_all(resp.as_bytes()).await.unwrap();
1443 sock.flush().await.unwrap();
1444 });
1445
1446 let provider = OpenAiProvider::new(format!("http://{addr}"), "k", HashMap::new());
1447 let req = ChatRequest::new("m", vec![ChatMessage::user("hi")]);
1448 let err = provider.complete(&req, &|_: &str| {}).await.unwrap_err();
1449 match err {
1450 Error::Provider { status, body } => {
1451 assert_eq!(status, 401);
1452 assert!(body.contains("bad key"), "body: {body}");
1453 }
1454 other => panic!("expected Provider error, got: {other:?}"),
1455 }
1456 server.await.unwrap();
1457 }
1458
1459 #[tokio::test]
1460 async fn streams_a_200_response_into_a_message() {
1461 use std::sync::{Arc, Mutex};
1462 use tokio::io::{AsyncReadExt, AsyncWriteExt};
1463
1464 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
1465 let addr = listener.local_addr().unwrap();
1466 let server = tokio::spawn(async move {
1467 let (mut sock, _) = listener.accept().await.unwrap();
1468 let mut buf = [0u8; 2048];
1469 let _ = sock.read(&mut buf).await;
1470 let sse = "data: {\"choices\":[{\"delta\":{\"content\":\"hel\"}}]}\n\n\
1471 data: {\"choices\":[{\"delta\":{\"content\":\"lo\"}}]}\n\n\
1472 data: [DONE]\n\n";
1473 let resp = format!(
1474 "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
1475 sse.len(),
1476 sse
1477 );
1478 sock.write_all(resp.as_bytes()).await.unwrap();
1479 sock.flush().await.unwrap();
1480 });
1481
1482 let provider = OpenAiProvider::new(format!("http://{addr}"), "k", HashMap::new());
1483 let req = ChatRequest::new("m", vec![ChatMessage::user("hi")]);
1484 let seen = Arc::new(Mutex::new(String::new()));
1485 let seen2 = seen.clone();
1486 let on_delta = move |s: &str| seen2.lock().unwrap().push_str(s);
1487 let (msg, _usage) = provider.complete(&req, &on_delta).await.unwrap();
1488 assert_eq!(msg.content.as_deref(), Some("hello"));
1489 assert_eq!(*seen.lock().unwrap(), "hello", "deltas streamed live");
1490 server.await.unwrap();
1491 }
1492
1493 #[test]
1496 fn from_retry_config_unset_is_byte_identical_to_default() {
1497 let opts = HttpOptions::from_retry_config(true, None, None);
1498 let default = HttpOptions::default();
1499 assert_eq!(opts.max_retries, default.max_retries);
1500 assert_eq!(opts.retry_backoff_base, default.retry_backoff_base);
1501 assert_eq!(opts.connect_timeout, default.connect_timeout);
1502 assert_eq!(opts.read_idle_timeout, default.read_idle_timeout);
1503 }
1504
1505 #[test]
1506 fn from_retry_config_disabled_forces_zero_retries() {
1507 let opts = HttpOptions::from_retry_config(false, None, None);
1508 assert_eq!(opts.max_retries, 0);
1509 assert_eq!(
1512 opts.retry_backoff_base,
1513 HttpOptions::default().retry_backoff_base
1514 );
1515 }
1516
1517 #[test]
1518 fn from_retry_config_disabled_with_explicit_max_retries_still_forces_zero() {
1519 let opts = HttpOptions::from_retry_config(false, Some(5), None);
1522 assert_eq!(opts.max_retries, 0);
1523 }
1524
1525 #[test]
1526 fn from_retry_config_overrides_apply_when_enabled() {
1527 let opts = HttpOptions::from_retry_config(true, Some(7), Some(1234));
1528 assert_eq!(opts.max_retries, 7);
1529 assert_eq!(opts.retry_backoff_base, Duration::from_millis(1234));
1530 }
1531
1532 #[test]
1533 fn from_retry_config_partial_override_leaves_the_other_at_default() {
1534 let opts = HttpOptions::from_retry_config(true, Some(9), None);
1535 assert_eq!(opts.max_retries, 9);
1536 assert_eq!(
1537 opts.retry_backoff_base,
1538 HttpOptions::default().retry_backoff_base
1539 );
1540 }
1541
1542 fn test_http_options() -> HttpOptions {
1545 HttpOptions {
1546 connect_timeout: Duration::from_millis(250),
1547 read_idle_timeout: Duration::from_millis(250),
1548 max_retries: 2,
1549 retry_backoff_base: Duration::from_millis(10),
1550 }
1551 }
1552
1553 #[tokio::test]
1554 async fn hung_connection_errors_via_read_timeout_within_bounded_time() {
1555 use tokio::io::AsyncReadExt;
1556
1557 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
1558 let addr = listener.local_addr().unwrap();
1559 let server = tokio::spawn(async move {
1564 loop {
1565 let Ok((mut sock, _)) = listener.accept().await else {
1566 break;
1567 };
1568 tokio::spawn(async move {
1569 let mut buf = [0u8; 2048];
1570 let _ = sock.read(&mut buf).await;
1571 tokio::time::sleep(Duration::from_secs(2)).await;
1574 });
1575 }
1576 });
1577
1578 let provider = OpenAiProvider::new_with_options(
1579 format!("http://{addr}"),
1580 "k",
1581 HashMap::new(),
1582 test_http_options(),
1583 );
1584 let req = ChatRequest::new("m", vec![ChatMessage::user("hi")]);
1585
1586 let outcome = tokio::time::timeout(Duration::from_secs(5), async {
1593 provider.complete(&req, &|_: &str| {}).await
1594 })
1595 .await
1596 .expect("complete() must return within the outer bound, not hang forever");
1597
1598 match outcome {
1599 Err(Error::Http(_)) => {}
1600 other => panic!("expected Err(Error::Http(_)) from the read timeout, got: {other:?}"),
1601 }
1602
1603 server.abort();
1604 }
1605
1606 #[tokio::test]
1607 async fn retries_503_then_succeeds_with_exactly_two_requests() {
1608 use std::sync::atomic::{AtomicUsize, Ordering};
1609 use std::sync::Arc;
1610 use tokio::io::{AsyncReadExt, AsyncWriteExt};
1611
1612 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
1613 let addr = listener.local_addr().unwrap();
1614 let connections = Arc::new(AtomicUsize::new(0));
1615 let connections2 = connections.clone();
1616 let server = tokio::spawn(async move {
1617 for _ in 0..2 {
1618 let (mut sock, _) = listener.accept().await.unwrap();
1619 let n = connections2.fetch_add(1, Ordering::SeqCst) + 1;
1620 let mut buf = [0u8; 2048];
1621 let _ = sock.read(&mut buf).await;
1622 if n == 1 {
1623 let resp =
1624 "HTTP/1.1 503 Service Unavailable\r\nContent-Length: 0\r\nConnection: close\r\n\r\n";
1625 sock.write_all(resp.as_bytes()).await.unwrap();
1626 } else {
1627 let sse = "data: {\"choices\":[{\"delta\":{\"content\":\"hel\"}}]}\n\n\
1628 data: {\"choices\":[{\"delta\":{\"content\":\"lo\"}}]}\n\n\
1629 data: [DONE]\n\n";
1630 let resp = format!(
1631 "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
1632 sse.len(),
1633 sse
1634 );
1635 sock.write_all(resp.as_bytes()).await.unwrap();
1636 }
1637 sock.flush().await.unwrap();
1638 }
1639 });
1640
1641 let provider = OpenAiProvider::new_with_options(
1642 format!("http://{addr}"),
1643 "k",
1644 HashMap::new(),
1645 test_http_options(),
1646 );
1647 let req = ChatRequest::new("m", vec![ChatMessage::user("hi")]);
1648 let (msg, _usage) = provider.complete(&req, &|_: &str| {}).await.unwrap();
1649 assert_eq!(msg.content.as_deref(), Some("hello"));
1650 server.await.unwrap();
1651 assert_eq!(
1652 connections.load(Ordering::SeqCst),
1653 2,
1654 "exactly 2 requests made: one 503, one successful retry"
1655 );
1656 }
1657
1658 #[test]
1659 fn streaming_decodes_multibyte_across_chunk_boundaries() {
1660 let line = "data: {\"choices\":[{\"delta\":{\"content\":\"héllo🌍\"}}]}\n";
1663 let bytes = line.as_bytes();
1664 let mut deltas = Vec::new();
1665 for split in 1..bytes.len() {
1667 let mut acc = Accumulator::default();
1668 let mut buf: Vec<u8> = Vec::new();
1669 buf.extend_from_slice(&bytes[..split]);
1670 drain_sse_lines(&mut buf, &mut acc, &mut deltas).unwrap();
1671 buf.extend_from_slice(&bytes[split..]);
1672 drain_sse_lines(&mut buf, &mut acc, &mut deltas).unwrap();
1673 assert_eq!(acc.content, "héllo🌍", "split at byte {split}");
1674 assert!(!acc.content.contains('\u{FFFD}'));
1675 }
1676 }
1677}