1use crate::errors::AppError;
31use crate::retry::AttemptOutcome;
32use secrecy::{ExposeSecret, SecretBox};
33use serde::{Deserialize, Serialize};
34use std::time::Duration;
35
36use crate::constants::DEFAULT_OPENROUTER_CHAT_URL;
37const DEFAULT_TIMEOUT_SECS: u64 = 600;
41const DEFAULT_CONNECT_TIMEOUT_SECS: u64 = 10;
42
43const SCHEMA_NAME: &str = "enrich_output";
46
47#[derive(Serialize)]
48struct ChatRequest<'a> {
49 model: &'a str,
50 messages: Vec<ChatMessage<'a>>,
51 response_format: ResponseFormat,
52 provider: ProviderPrefs,
53 #[serde(skip_serializing_if = "Option::is_none")]
54 reasoning: Option<ReasoningPrefs>,
55 #[serde(skip_serializing_if = "Option::is_none")]
56 max_tokens: Option<u32>,
57}
58
59#[derive(Serialize)]
60struct ChatMessage<'a> {
61 role: &'a str,
62 content: String,
63}
64
65#[derive(Serialize)]
66struct ResponseFormat {
67 #[serde(rename = "type")]
68 format_type: &'static str,
69 json_schema: JsonSchemaSpec,
70}
71
72#[derive(Serialize)]
73struct JsonSchemaSpec {
74 name: &'static str,
75 strict: bool,
76 schema: serde_json::Value,
77}
78
79#[derive(Serialize)]
80struct ProviderPrefs {
81 require_parameters: bool,
82}
83
84#[derive(Serialize)]
85struct ReasoningPrefs {
86 enabled: bool,
87}
88
89#[derive(Deserialize)]
90struct ChatResponse {
91 #[serde(default)]
92 choices: Vec<Choice>,
93 #[serde(default)]
94 usage: Option<Usage>,
95 #[serde(default)]
100 error: Option<crate::openrouter_http::ApiError>,
101}
102
103#[derive(Deserialize)]
104struct Choice {
105 message: RespMessage,
106 #[serde(default)]
110 finish_reason: Option<String>,
111}
112
113#[derive(Deserialize)]
114struct RespMessage {
115 #[serde(default)]
116 content: Option<String>,
117}
118
119#[derive(Deserialize)]
120struct Usage {
121 #[serde(default)]
122 cost: Option<f64>,
123 #[serde(default)]
126 prompt_tokens: Option<u32>,
127 #[serde(default)]
130 completion_tokens: Option<u32>,
131}
132
133#[derive(Debug)]
140pub struct ChatCompletion {
141 pub value: serde_json::Value,
143 pub cost_usd: f64,
145 pub finish_reason: Option<String>,
147 pub prompt_tokens: Option<u32>,
149 pub completion_tokens: Option<u32>,
151}
152
153#[derive(Debug)]
168pub struct ChatError {
169 pub source: AppError,
171 pub finish_reason: Option<String>,
174 pub prompt_tokens: Option<u32>,
177 pub completion_tokens: Option<u32>,
180 pub retry_class: AttemptOutcome,
183}
184
185impl ChatError {
186 fn new(source: AppError, retry_class: AttemptOutcome) -> Self {
190 Self {
191 source,
192 finish_reason: None,
193 prompt_tokens: None,
194 completion_tokens: None,
195 retry_class,
196 }
197 }
198
199 fn with_diagnostics(
203 source: AppError,
204 finish_reason: Option<String>,
205 prompt_tokens: Option<u32>,
206 completion_tokens: Option<u32>,
207 retry_class: AttemptOutcome,
208 ) -> Self {
209 Self {
210 source,
211 finish_reason,
212 prompt_tokens,
213 completion_tokens,
214 retry_class,
215 }
216 }
217}
218
219impl std::fmt::Display for ChatError {
220 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
221 std::fmt::Display::fmt(&self.source, f)
222 }
223}
224
225impl std::error::Error for ChatError {
226 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
227 Some(&self.source)
228 }
229}
230
231pub struct OpenRouterChatClient {
234 client: reqwest::Client,
235 api_key: SecretBox<String>,
236 model: String,
237 base_url: String,
240}
241
242impl OpenRouterChatClient {
243 pub fn new(
248 api_key: SecretBox<String>,
249 model: String,
250 timeout_secs: u64,
251 ) -> Result<Self, AppError> {
252 let base_url =
253 crate::runtime_config::openrouter_chat_url(DEFAULT_OPENROUTER_CHAT_URL);
254 Self::new_with_base_url(api_key, model, timeout_secs, base_url)
255 }
256
257 pub fn new_with_base_url(
259 api_key: SecretBox<String>,
260 model: String,
261 timeout_secs: u64,
262 base_url: String,
263 ) -> Result<Self, AppError> {
264 let timeout_secs = if timeout_secs == 0 {
265 DEFAULT_TIMEOUT_SECS
266 } else {
267 timeout_secs
268 };
269 let client = reqwest::Client::builder()
270 .timeout(Duration::from_secs(timeout_secs))
271 .connect_timeout(Duration::from_secs(DEFAULT_CONNECT_TIMEOUT_SECS))
272 .user_agent(concat!("sqlite-graphrag/", env!("CARGO_PKG_VERSION")))
273 .build()
274 .map_err(|e| AppError::Validation(format!("failed to build HTTP client: {e}")))?;
275
276 Ok(Self {
277 client,
278 api_key,
279 model,
280 base_url,
281 })
282 }
283
284 #[cfg(test)]
286 pub fn new_with_url(
287 api_key: SecretBox<String>,
288 model: String,
289 base_url: String,
290 timeout_secs: u64,
291 ) -> Result<Self, AppError> {
292 Self::new_with_base_url(api_key, model, timeout_secs, base_url)
293 }
294
295 pub fn model(&self) -> &str {
297 &self.model
298 }
299
300 pub async fn complete(
323 &self,
324 system_prompt: &str,
325 input_text: &str,
326 schema_str: &str,
327 max_tokens: Option<u32>,
328 ) -> Result<ChatCompletion, ChatError> {
329 let schema: serde_json::Value = serde_json::from_str(schema_str).map_err(|e| {
333 ChatError::new(
334 AppError::Validation(format!("invalid JSON schema for OpenRouter request: {e}")),
335 AttemptOutcome::HardFailure,
336 )
337 })?;
338
339 let mut current_max_tokens = max_tokens;
340
341 for length_attempt in 0..=crate::constants::ENRICH_MAX_LENGTH_RETRIES {
342 let response = self
343 .complete_one_attempt(&schema, system_prompt, input_text, current_max_tokens)
344 .await?;
345
346 let finish_reason = response
347 .choices
348 .first()
349 .and_then(|c| c.finish_reason.clone());
350 let prompt_tokens = response.usage.as_ref().and_then(|u| u.prompt_tokens);
351 let completion_tokens = response.usage.as_ref().and_then(|u| u.completion_tokens);
352
353 let truncated = finish_reason.as_deref() == Some("length");
354 let retries_left = length_attempt < crate::constants::ENRICH_MAX_LENGTH_RETRIES;
355
356 if truncated && retries_left {
357 let next_max_tokens = grow_max_tokens(current_max_tokens);
358 tracing::warn!(
359 model = %self.model,
360 attempt = length_attempt,
361 previous_max_tokens = ?current_max_tokens,
362 next_max_tokens,
363 "OpenRouter completion truncated (finish_reason=length); \
364 retrying with a larger max_tokens budget"
365 );
366 current_max_tokens = Some(next_max_tokens);
367 continue;
368 }
369
370 if truncated {
371 tracing::warn!(
372 model = %self.model,
373 max_length_retries = crate::constants::ENRICH_MAX_LENGTH_RETRIES,
374 max_tokens = ?current_max_tokens,
375 "OpenRouter completion still truncated after exhausting \
376 max_tokens growth"
377 );
378 }
379
380 return self.finish_completion(
381 response,
382 finish_reason,
383 prompt_tokens,
384 completion_tokens,
385 );
386 }
387
388 unreachable!("loop always returns within ENRICH_MAX_LENGTH_RETRIES + 1 iterations")
389 }
390
391 async fn complete_one_attempt(
396 &self,
397 schema: &serde_json::Value,
398 system_prompt: &str,
399 input_text: &str,
400 max_tokens: Option<u32>,
401 ) -> Result<ChatResponse, ChatError> {
402 let primary = self.build_request(
409 schema.clone(),
410 system_prompt,
411 input_text,
412 max_tokens,
413 Some(ReasoningPrefs { enabled: false }),
414 );
415 match self.execute_with_retry(&primary).await {
416 Ok(r) => Ok(r),
417 Err(first_err) => {
418 if reasoning_disable_rejected(&first_err) {
419 tracing::warn!(
420 model = %self.model,
421 "model rejected reasoning.enabled=false (mandatory); \
422 retrying once with reasoning omitted"
423 );
424 let fallback = self.build_request(
425 schema.clone(),
426 system_prompt,
427 input_text,
428 max_tokens,
429 None,
430 );
431 match self.execute_with_retry(&fallback).await {
432 Ok(r) => Ok(r),
433 Err(_) => Err(first_err),
434 }
435 } else {
436 Err(first_err)
437 }
438 }
439 }
440 }
441
442 fn finish_completion(
458 &self,
459 response: ChatResponse,
460 finish_reason: Option<String>,
461 prompt_tokens: Option<u32>,
462 completion_tokens: Option<u32>,
463 ) -> Result<ChatCompletion, ChatError> {
464 let content = response
465 .choices
466 .into_iter()
467 .next()
468 .and_then(|c| c.message.content)
469 .filter(|c| !c.trim().is_empty())
470 .ok_or_else(|| {
471 AppError::Validation(format!(
472 "model '{}' returned no structured content (incompatible with \
473 structured outputs, or refused the request)",
474 self.model
475 ))
476 })
477 .map_err(|e| {
478 ChatError::with_diagnostics(
479 e,
480 finish_reason.clone(),
481 prompt_tokens,
482 completion_tokens,
483 AttemptOutcome::Transient,
484 )
485 })?;
486
487 let value = crate::json_repair::repair_to_value(&content).map_err(|e| {
493 ChatError::with_diagnostics(
494 AppError::Validation(format!(
495 "model '{}' returned content that could not be parsed even after \
496 JSON repair: {e}",
497 self.model
498 )),
499 finish_reason.clone(),
500 prompt_tokens,
501 completion_tokens,
502 AttemptOutcome::Transient,
503 )
504 })?;
505
506 if !value.is_object() {
514 return Err(ChatError::with_diagnostics(
515 AppError::Validation(format!(
516 "model '{}' returned non-object JSON after repair (got {}); \
517 likely a refusal or malformed structured output",
518 self.model,
519 json_shape_name(&value)
520 )),
521 finish_reason,
522 prompt_tokens,
523 completion_tokens,
524 AttemptOutcome::Transient,
525 ));
526 }
527
528 let cost = response.usage.and_then(|u| u.cost).unwrap_or(0.0);
529
530 Ok(ChatCompletion {
531 value,
532 cost_usd: cost,
533 finish_reason,
534 prompt_tokens,
535 completion_tokens,
536 })
537 }
538
539 fn build_request<'a>(
543 &'a self,
544 schema: serde_json::Value,
545 system_prompt: &str,
546 input_text: &str,
547 max_tokens: Option<u32>,
548 reasoning: Option<ReasoningPrefs>,
549 ) -> ChatRequest<'a> {
550 let mut messages = Vec::with_capacity(2);
551 messages.push(ChatMessage {
552 role: "system",
553 content: system_prompt.to_string(),
554 });
555 if !input_text.is_empty() {
556 messages.push(ChatMessage {
557 role: "user",
558 content: input_text.to_string(),
559 });
560 }
561 ChatRequest {
562 model: &self.model,
563 messages,
564 response_format: ResponseFormat {
565 format_type: "json_schema",
566 json_schema: JsonSchemaSpec {
567 name: SCHEMA_NAME,
568 strict: true,
569 schema,
570 },
571 },
572 provider: ProviderPrefs {
573 require_parameters: true,
574 },
575 reasoning,
576 max_tokens,
577 }
578 }
579
580 async fn execute_with_retry(
586 &self,
587 request: &ChatRequest<'_>,
588 ) -> Result<ChatResponse, ChatError> {
589 let mut last_err: Option<ChatError> = None;
590
591 for attempt in 0..crate::openrouter_http::MAX_RETRIES {
592 let result = self
593 .client
594 .post(&self.base_url)
595 .header(
596 "Authorization",
597 format!("Bearer {}", self.api_key.expose_secret()),
598 )
599 .json(request)
600 .send()
601 .await;
602
603 let resp = match result {
604 Ok(r) => r,
605 Err(e) if e.is_timeout() => {
606 return Err(ChatError::new(
607 AppError::Validation("OpenRouter chat request timed out".into()),
608 AttemptOutcome::Transient,
609 ));
610 }
611 Err(e) => {
612 last_err = Some(ChatError::new(
613 AppError::Validation(format!("HTTP request failed: {e}")),
614 AttemptOutcome::Transient,
615 ));
616 crate::openrouter_http::backoff(attempt).await;
617 continue;
618 }
619 };
620
621 let status = resp.status();
622
623 if status.is_success() {
624 let body = resp.text().await.map_err(|e| {
625 ChatError::new(
626 AppError::Validation(format!("failed to read response body: {e}")),
627 AttemptOutcome::Transient,
628 )
629 })?;
630 match serde_json::from_str::<ChatResponse>(&body) {
631 Ok(parsed) => {
632 if let Some(api_err) = parsed.error {
637 let retry_class =
638 crate::openrouter_http::provider_error_retry_class(&api_err);
639 return Err(ChatError::new(
640 AppError::ProviderError {
641 code: api_err.code_string(),
642 message: api_err.message,
643 },
644 retry_class,
645 ));
646 }
647 return Ok(parsed);
648 }
649 Err(e) => {
650 tracing::warn!(
651 attempt,
652 body_len = body.len(),
653 "HTTP 200 but parse failed (retrying): {e}"
654 );
655 last_err = Some(ChatError::new(
656 AppError::Validation(format!("failed to parse chat response: {e}")),
657 AttemptOutcome::Transient,
658 ));
659 crate::openrouter_http::backoff(attempt).await;
660 continue;
661 }
662 }
663 }
664
665 if status.as_u16() == 401 {
666 return Err(ChatError::new(
667 AppError::Validation("invalid OpenRouter API key (HTTP 401)".into()),
668 AttemptOutcome::HardFailure,
669 ));
670 }
671
672 if status.as_u16() == 400 || status.as_u16() == 404 {
673 let body = resp.text().await.unwrap_or_default();
674 return Err(ChatError::new(
675 AppError::Validation(format!(
676 "OpenRouter returned {status} for model '{}': {body}",
677 self.model
678 )),
679 AttemptOutcome::HardFailure,
680 ));
681 }
682
683 if status.as_u16() == 429 {
684 let retry_after = resp
685 .headers()
686 .get("retry-after")
687 .and_then(|v| v.to_str().ok())
688 .and_then(|v| v.parse::<u64>().ok())
689 .unwrap_or(2);
690 tracing::warn!(
691 attempt,
692 retry_after_secs = retry_after,
693 "OpenRouter rate limited, waiting"
694 );
695 last_err = Some(ChatError::new(
700 AppError::RateLimited {
701 detail: format!("OpenRouter HTTP 429 (retry-after {retry_after}s)"),
702 },
703 AttemptOutcome::Transient,
704 ));
705 tokio::time::sleep(Duration::from_secs(retry_after)).await;
706 continue;
707 }
708
709 if status.is_server_error() {
710 tracing::warn!(attempt, status = %status, "OpenRouter server error, retrying");
711 last_err = Some(ChatError::new(
712 AppError::Validation(format!("OpenRouter server error: {status}")),
713 AttemptOutcome::Transient,
714 ));
715 crate::openrouter_http::backoff(attempt).await;
716 continue;
717 }
718
719 let body = resp.text().await.unwrap_or_default();
720 return Err(ChatError::new(
721 AppError::Validation(format!("unexpected HTTP {status}: {body}")),
722 crate::openrouter_http::status_retry_class(status),
723 ));
724 }
725
726 Err(last_err.unwrap_or_else(|| {
731 ChatError::new(
732 AppError::Validation("max retries exceeded for OpenRouter chat request".into()),
733 AttemptOutcome::Transient,
734 )
735 }))
736 }
737}
738
739fn grow_max_tokens(current: Option<u32>) -> u32 {
746 let base = current.unwrap_or(crate::constants::ENRICH_INITIAL_MAX_TOKENS);
747 base.saturating_mul(crate::constants::ENRICH_MAX_TOKENS_GROWTH_FACTOR)
748 .min(crate::constants::ENRICH_MAX_TOKENS_CEILING)
749}
750
751fn reasoning_disable_rejected(err: &ChatError) -> bool {
761 let msg = err.source.to_string().to_lowercase();
762 msg.contains("400") && msg.contains("reasoning")
763}
764
765fn json_shape_name(value: &serde_json::Value) -> &'static str {
768 match value {
769 serde_json::Value::Null => "null",
770 serde_json::Value::Bool(_) => "boolean",
771 serde_json::Value::Number(_) => "number",
772 serde_json::Value::String(_) => "string",
773 serde_json::Value::Array(_) => "array",
774 serde_json::Value::Object(_) => "object",
775 }
776}
777
778#[cfg(test)]
779mod tests {
780 use super::*;
781 use serde_json::json;
782 use wiremock::matchers::{body_partial_json, method, path};
783 use wiremock::{Mock, MockServer, ResponseTemplate};
784
785 const TEST_SCHEMA: &str = r#"{"type":"object"}"#;
786
787 fn key() -> SecretBox<String> {
788 SecretBox::new(Box::new("test-key".to_string()))
789 }
790
791 fn success_body(content: &str, cost: Option<f64>) -> serde_json::Value {
796 success_body_with_finish(content, cost, "stop")
797 }
798
799 fn success_body_with_finish(
802 content: &str,
803 cost: Option<f64>,
804 finish_reason: &str,
805 ) -> serde_json::Value {
806 let mut body = json!({
807 "choices": [{ "message": { "content": content }, "finish_reason": finish_reason }]
808 });
809 if let Some(c) = cost {
810 body["usage"] = json!({ "cost": c });
811 }
812 body
813 }
814
815 async fn client_for(server: &MockServer, model: &str) -> OpenRouterChatClient {
816 OpenRouterChatClient::new_with_url(
817 key(),
818 model.to_string(),
819 format!("{}/chat/completions", server.uri()),
820 30,
821 )
822 .expect("test client builds")
823 }
824
825 #[test]
826 fn new_builds_client_and_binds_model() {
827 let client = OpenRouterChatClient::new(key(), "z-ai/glm-5.2".to_string(), 30)
828 .expect("client builds");
829 assert_eq!(client.model(), "z-ai/glm-5.2");
830 }
831
832 #[test]
833 fn new_defaults_base_url_to_public_endpoint() {
834 let client = OpenRouterChatClient::new(key(), "z-ai/glm-5.2".to_string(), 30)
835 .expect("client builds");
836 assert_eq!(client.base_url, DEFAULT_OPENROUTER_CHAT_URL);
837 }
838
839 #[test]
840 fn request_serializes_with_strict_schema_and_disabled_reasoning() {
841 let request = ChatRequest {
842 model: "deepseek/deepseek-v4-flash",
843 messages: vec![ChatMessage {
844 role: "system",
845 content: "extract".to_string(),
846 }],
847 response_format: ResponseFormat {
848 format_type: "json_schema",
849 json_schema: JsonSchemaSpec {
850 name: SCHEMA_NAME,
851 strict: true,
852 schema: serde_json::json!({"type": "object"}),
853 },
854 },
855 provider: ProviderPrefs {
856 require_parameters: true,
857 },
858 reasoning: Some(ReasoningPrefs { enabled: false }),
859 max_tokens: None,
860 };
861 let json = serde_json::to_value(&request).expect("serializes");
862 assert_eq!(json["response_format"]["type"], "json_schema");
863 assert_eq!(json["response_format"]["json_schema"]["strict"], true);
864 assert_eq!(json["provider"]["require_parameters"], true);
865 assert_eq!(json["reasoning"]["enabled"], false);
866 assert!(json.get("max_tokens").is_none());
868 }
869
870 #[test]
871 fn grow_max_tokens_uses_initial_default_when_current_is_none() {
872 assert_eq!(
873 grow_max_tokens(None),
874 crate::constants::ENRICH_INITIAL_MAX_TOKENS
875 * crate::constants::ENRICH_MAX_TOKENS_GROWTH_FACTOR
876 );
877 }
878
879 #[test]
880 fn grow_max_tokens_caps_at_ceiling() {
881 assert_eq!(
882 grow_max_tokens(Some(crate::constants::ENRICH_MAX_TOKENS_CEILING)),
883 crate::constants::ENRICH_MAX_TOKENS_CEILING
884 );
885 assert_eq!(
886 grow_max_tokens(Some(u32::MAX)),
887 crate::constants::ENRICH_MAX_TOKENS_CEILING
888 );
889 }
890
891 #[tokio::test]
892 async fn complete_sends_wellformed_request_and_parses_content() {
893 let server = MockServer::start().await;
894 Mock::given(method("POST"))
895 .and(path("/chat/completions"))
896 .and(body_partial_json(json!({
897 "model": "deepseek/deepseek-v4-flash",
898 "response_format": {
899 "type": "json_schema",
900 "json_schema": { "name": "enrich_output", "strict": true }
901 },
902 "provider": { "require_parameters": true },
903 "reasoning": { "enabled": false }
904 })))
905 .respond_with(ResponseTemplate::new(200).set_body_json(success_body(
906 r#"{"entities":[],"relationships":[]}"#,
907 Some(0.0023),
908 )))
909 .expect(1)
910 .mount(&server)
911 .await;
912
913 let client = client_for(&server, "deepseek/deepseek-v4-flash").await;
914 let completion = client
915 .complete("system", "input", TEST_SCHEMA, None)
916 .await
917 .expect("completion succeeds");
918
919 assert_eq!(
920 completion.value,
921 json!({"entities": [], "relationships": []})
922 );
923 assert!((completion.cost_usd - 0.0023).abs() < f64::EPSILON);
924 assert_eq!(completion.finish_reason.as_deref(), Some("stop"));
925 }
926
927 #[tokio::test]
928 async fn complete_defaults_cost_to_zero_when_usage_absent() {
929 let server = MockServer::start().await;
930 Mock::given(method("POST"))
931 .respond_with(
932 ResponseTemplate::new(200).set_body_json(success_body(r#"{"entities":[]}"#, None)),
933 )
934 .mount(&server)
935 .await;
936
937 let client = client_for(&server, "z-ai/glm-5.2").await;
938 let completion = client
939 .complete("system", "", TEST_SCHEMA, Some(4096))
940 .await
941 .expect("completion succeeds");
942 assert_eq!(completion.cost_usd, 0.0);
943 }
944
945 #[tokio::test]
946 async fn complete_retries_on_429_honouring_retry_after() {
947 let server = MockServer::start().await;
948 Mock::given(method("POST"))
949 .respond_with(ResponseTemplate::new(429).insert_header("retry-after", "1"))
950 .up_to_n_times(1)
951 .expect(1)
952 .mount(&server)
953 .await;
954 Mock::given(method("POST"))
955 .respond_with(
956 ResponseTemplate::new(200).set_body_json(success_body(r#"{"ok":true}"#, Some(0.0))),
957 )
958 .expect(1)
959 .mount(&server)
960 .await;
961
962 let client = client_for(&server, "minimax/minimax-m3").await;
963 let completion = client
964 .complete("system", "input", TEST_SCHEMA, None)
965 .await
966 .expect("retried completion succeeds");
967 assert_eq!(completion.value, json!({"ok": true}));
968 }
969
970 #[tokio::test]
971 async fn complete_retries_on_5xx_with_backoff() {
972 let server = MockServer::start().await;
973 Mock::given(method("POST"))
974 .respond_with(ResponseTemplate::new(503))
975 .up_to_n_times(1)
976 .expect(1)
977 .mount(&server)
978 .await;
979 Mock::given(method("POST"))
980 .respond_with(
981 ResponseTemplate::new(200).set_body_json(success_body(r#"{"ok":1}"#, Some(0.0))),
982 )
983 .expect(1)
984 .mount(&server)
985 .await;
986
987 let client = client_for(&server, "openai/gpt-oss-120b").await;
988 let completion = client
989 .complete("system", "input", TEST_SCHEMA, None)
990 .await
991 .expect("retried completion succeeds");
992 assert_eq!(completion.value, json!({"ok": 1}));
993 }
994
995 #[tokio::test]
996 async fn complete_401_is_permanent_without_retry() {
997 let server = MockServer::start().await;
998 Mock::given(method("POST"))
999 .respond_with(ResponseTemplate::new(401))
1000 .expect(1)
1001 .mount(&server)
1002 .await;
1003
1004 let client = client_for(&server, "z-ai/glm-5.2").await;
1005 let err = client
1006 .complete("system", "input", TEST_SCHEMA, None)
1007 .await
1008 .expect_err("401 is an error");
1009 assert!(err.to_string().contains("401"), "got: {err}");
1010 assert_eq!(err.retry_class, AttemptOutcome::HardFailure);
1011 }
1012
1013 #[tokio::test]
1014 async fn complete_400_returns_body_and_model_without_retry() {
1015 let server = MockServer::start().await;
1016 Mock::given(method("POST"))
1017 .respond_with(ResponseTemplate::new(400).set_body_string("schema not supported"))
1018 .expect(1)
1019 .mount(&server)
1020 .await;
1021
1022 let client = client_for(&server, "xiaomi/mimo-v2.5").await;
1023 let err = client
1024 .complete("system", "input", TEST_SCHEMA, None)
1025 .await
1026 .expect_err("400 is an error");
1027 let msg = err.to_string();
1028 assert!(msg.contains("400"), "got: {msg}");
1029 assert!(msg.contains("xiaomi/mimo-v2.5"), "got: {msg}");
1030 assert!(msg.contains("schema not supported"), "got: {msg}");
1031 assert_eq!(err.retry_class, AttemptOutcome::HardFailure);
1032 }
1033
1034 #[tokio::test]
1035 async fn complete_empty_choices_errors_citing_model() {
1036 let server = MockServer::start().await;
1037 Mock::given(method("POST"))
1038 .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "choices": [] })))
1039 .mount(&server)
1040 .await;
1041
1042 let client = client_for(&server, "minimax/minimax-m2.7").await;
1043 let err = client
1044 .complete("system", "input", TEST_SCHEMA, None)
1045 .await
1046 .expect_err("empty choices is an error");
1047 let msg = err.to_string();
1048 assert!(msg.contains("minimax/minimax-m2.7"), "got: {msg}");
1049 assert!(msg.contains("no structured content"), "got: {msg}");
1050 assert_eq!(err.finish_reason, None);
1051 assert_eq!(
1052 err.retry_class,
1053 AttemptOutcome::Transient,
1054 "no-content is a model hiccup, not a permanent rejection"
1055 );
1056 }
1057
1058 #[tokio::test]
1059 async fn complete_empty_content_errors() {
1060 let server = MockServer::start().await;
1061 Mock::given(method("POST"))
1062 .respond_with(ResponseTemplate::new(200).set_body_json(success_body(" ", Some(0.0))))
1063 .mount(&server)
1064 .await;
1065
1066 let client = client_for(&server, "z-ai/glm-5.2:nitro").await;
1067 let err = client
1068 .complete("system", "input", TEST_SCHEMA, None)
1069 .await
1070 .expect_err("blank content is an error");
1071 assert!(
1072 err.to_string().contains("no structured content"),
1073 "got: {err}"
1074 );
1075 }
1076
1077 #[tokio::test]
1078 async fn complete_non_json_content_errors_as_incompatible() {
1079 let server = MockServer::start().await;
1083 Mock::given(method("POST"))
1084 .respond_with(
1085 ResponseTemplate::new(200)
1086 .set_body_json(success_body("this is not json", Some(0.0))),
1087 )
1088 .mount(&server)
1089 .await;
1090
1091 let client = client_for(&server, "google/gemini-3.1-flash-lite").await;
1092 let err = client
1093 .complete("system", "input", TEST_SCHEMA, None)
1094 .await
1095 .expect_err("non-json content is an error");
1096 let msg = err.to_string();
1097 assert!(msg.contains("non-object JSON after repair"), "got: {msg}");
1098 assert!(msg.contains("google/gemini-3.1-flash-lite"), "got: {msg}");
1099 }
1100
1101 #[tokio::test]
1102 async fn complete_repairs_markdown_fenced_object() {
1103 let server = MockServer::start().await;
1107 Mock::given(method("POST"))
1108 .respond_with(ResponseTemplate::new(200).set_body_json(success_body(
1109 "```json\n{\"entities\":[\"rust\"],\"relationships\":[]}\n```",
1110 Some(0.0),
1111 )))
1112 .mount(&server)
1113 .await;
1114
1115 let client = client_for(&server, "deepseek/deepseek-v4-flash").await;
1116 let completion = client
1117 .complete("system", "input", TEST_SCHEMA, None)
1118 .await
1119 .expect("fenced object is repaired");
1120 assert_eq!(
1121 completion.value,
1122 json!({"entities": ["rust"], "relationships": []})
1123 );
1124 }
1125
1126 #[tokio::test]
1127 async fn complete_rejects_invalid_schema_before_network() {
1128 let client = OpenRouterChatClient::new_with_url(
1130 key(),
1131 "z-ai/glm-5.2".to_string(),
1132 "http://127.0.0.1:1/chat/completions".to_string(),
1133 30,
1134 )
1135 .expect("client builds");
1136 let err = client
1137 .complete("system", "input", "{not valid json", None)
1138 .await
1139 .expect_err("invalid schema is rejected");
1140 assert!(
1141 err.to_string().contains("invalid JSON schema"),
1142 "got: {err}"
1143 );
1144 assert_eq!(
1145 err.retry_class,
1146 AttemptOutcome::HardFailure,
1147 "a malformed schema is a permanent caller error"
1148 );
1149 }
1150
1151 #[tokio::test]
1152 async fn complete_retries_with_reasoning_omitted_when_mandatory() {
1153 let server = MockServer::start().await;
1154 Mock::given(method("POST"))
1158 .respond_with(
1159 ResponseTemplate::new(400).set_body_string(
1160 "reasoning is mandatory for this model and cannot be disabled",
1161 ),
1162 )
1163 .up_to_n_times(1)
1164 .expect(1)
1165 .mount(&server)
1166 .await;
1167 Mock::given(method("POST"))
1169 .respond_with(ResponseTemplate::new(200).set_body_json(success_body(
1170 r#"{"entities":[],"relationships":[]}"#,
1171 Some(0.0),
1172 )))
1173 .expect(1)
1174 .mount(&server)
1175 .await;
1176
1177 let client = client_for(&server, "minimax/minimax-m2.7").await;
1178 let completion = client
1179 .complete("system", "input", TEST_SCHEMA, None)
1180 .await
1181 .expect("fallback completion succeeds");
1182 assert_eq!(
1183 completion.value,
1184 json!({"entities": [], "relationships": []})
1185 );
1186
1187 let requests = server
1190 .received_requests()
1191 .await
1192 .expect("request recording is enabled");
1193 assert_eq!(requests.len(), 2, "expected primary + fallback requests");
1194 let first: serde_json::Value =
1195 serde_json::from_slice(&requests[0].body).expect("first request body is JSON");
1196 let second: serde_json::Value =
1197 serde_json::from_slice(&requests[1].body).expect("second request body is JSON");
1198 assert_eq!(
1199 first["reasoning"]["enabled"],
1200 json!(false),
1201 "primary request must send reasoning.enabled=false"
1202 );
1203 assert!(
1204 second.get("reasoning").is_none(),
1205 "fallback request must omit the reasoning field, got: {second}"
1206 );
1207 }
1208
1209 #[tokio::test]
1210 async fn complete_honours_configured_timeout() {
1211 let server = MockServer::start().await;
1215 Mock::given(method("POST"))
1216 .respond_with(
1217 ResponseTemplate::new(200)
1218 .set_delay(std::time::Duration::from_secs(2))
1219 .set_body_json(success_body(r#"{"ok":1}"#, Some(0.0))),
1220 )
1221 .mount(&server)
1222 .await;
1223
1224 let client = OpenRouterChatClient::new_with_url(
1225 key(),
1226 "z-ai/glm-5.2".to_string(),
1227 format!("{}/chat/completions", server.uri()),
1228 1,
1229 )
1230 .expect("client builds");
1231 let err = client
1232 .complete("system", "input", TEST_SCHEMA, None)
1233 .await
1234 .expect_err("request exceeds the 1s timeout");
1235 assert!(err.to_string().contains("timed out"), "got: {err}");
1236 }
1237
1238 #[tokio::test]
1239 async fn complete_surfaces_provider_error_in_200_body() {
1240 let server = MockServer::start().await;
1244 Mock::given(method("POST"))
1245 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
1246 "error": { "code": 400, "message": "context length exceeded" }
1247 })))
1248 .mount(&server)
1249 .await;
1250
1251 let client = client_for(&server, "deepseek/deepseek-v4-flash").await;
1252 let err = client
1253 .complete("system", "input", TEST_SCHEMA, None)
1254 .await
1255 .expect_err("provider error must surface");
1256 let msg = err.to_string();
1257 assert!(msg.contains("context length exceeded"), "got: {msg}");
1258 assert!(
1259 !msg.contains("no structured content"),
1260 "must not mask as empty choices: {msg}"
1261 );
1262 assert!(
1263 !msg.contains("missing field"),
1264 "must not mask as a missing field: {msg}"
1265 );
1266 assert_eq!(
1267 err.retry_class,
1268 AttemptOutcome::HardFailure,
1269 "code 400 (context length exceeded) is a permanent provider rejection"
1270 );
1271 }
1272
1273 #[tokio::test]
1274 async fn complete_classifies_provider_error_429_code_as_transient() {
1275 let server = MockServer::start().await;
1279 Mock::given(method("POST"))
1280 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
1281 "error": { "code": 429, "message": "rate limited" }
1282 })))
1283 .mount(&server)
1284 .await;
1285
1286 let client = client_for(&server, "deepseek/deepseek-v4-flash").await;
1287 let err = client
1288 .complete("system", "input", TEST_SCHEMA, None)
1289 .await
1290 .expect_err("provider rate-limit error must surface");
1291 assert_eq!(err.retry_class, AttemptOutcome::Transient);
1292 }
1293
1294 #[tokio::test]
1295 async fn complete_classifies_exhausted_5xx_retries_as_transient() {
1296 let server = MockServer::start().await;
1300 Mock::given(method("POST"))
1301 .respond_with(ResponseTemplate::new(503))
1302 .mount(&server)
1303 .await;
1304
1305 let client = client_for(&server, "openai/gpt-oss-120b").await;
1306 let err = client
1307 .complete("system", "input", TEST_SCHEMA, None)
1308 .await
1309 .expect_err("persistent 5xx exhausts retries");
1310 assert_eq!(err.retry_class, AttemptOutcome::Transient);
1311 }
1312
1313 #[tokio::test]
1314 async fn complete_regrows_max_tokens_and_retries_on_length_truncation() {
1315 let server = MockServer::start().await;
1321 Mock::given(method("POST"))
1322 .respond_with(
1323 ResponseTemplate::new(200).set_body_json(success_body_with_finish(
1324 r#"{"entities":["trunc"#,
1325 Some(0.001),
1326 "length",
1327 )),
1328 )
1329 .up_to_n_times(1)
1330 .expect(1)
1331 .mount(&server)
1332 .await;
1333 Mock::given(method("POST"))
1334 .respond_with(
1335 ResponseTemplate::new(200).set_body_json(success_body_with_finish(
1336 r#"{"entities":["rust"],"relationships":[]}"#,
1337 Some(0.002),
1338 "stop",
1339 )),
1340 )
1341 .expect(1)
1342 .mount(&server)
1343 .await;
1344
1345 let client = client_for(&server, "deepseek/deepseek-v4-flash:nitro").await;
1346 let completion = client
1347 .complete("system", "input", TEST_SCHEMA, Some(64))
1348 .await
1349 .expect("second attempt with grown max_tokens succeeds");
1350
1351 assert_eq!(
1352 completion.value,
1353 json!({"entities": ["rust"], "relationships": []})
1354 );
1355 assert_eq!(completion.finish_reason.as_deref(), Some("stop"));
1356
1357 let requests = server
1358 .received_requests()
1359 .await
1360 .expect("request recording is enabled");
1361 assert_eq!(requests.len(), 2, "expected exactly one regrowth retry");
1362 let first: serde_json::Value =
1363 serde_json::from_slice(&requests[0].body).expect("first request body is JSON");
1364 let second: serde_json::Value =
1365 serde_json::from_slice(&requests[1].body).expect("second request body is JSON");
1366 assert_eq!(first["max_tokens"], json!(64));
1367 assert_eq!(
1368 second["max_tokens"],
1369 json!(64 * crate::constants::ENRICH_MAX_TOKENS_GROWTH_FACTOR),
1370 "max_tokens must grow by ENRICH_MAX_TOKENS_GROWTH_FACTOR before the retry"
1371 );
1372 }
1373
1374 #[tokio::test]
1375 async fn complete_captures_finish_reason_and_tokens_on_success() {
1376 let server = MockServer::start().await;
1379 Mock::given(method("POST"))
1380 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
1381 "choices": [{
1382 "message": { "content": r#"{"ok":true}"# },
1383 "finish_reason": "stop"
1384 }],
1385 "usage": { "cost": 0.001, "prompt_tokens": 120, "completion_tokens": 30 }
1386 })))
1387 .mount(&server)
1388 .await;
1389
1390 let client = client_for(&server, "z-ai/glm-5.2").await;
1391 let completion = client
1392 .complete("system", "input", TEST_SCHEMA, None)
1393 .await
1394 .expect("completion succeeds");
1395
1396 assert_eq!(completion.finish_reason.as_deref(), Some("stop"));
1397 assert_eq!(completion.prompt_tokens, Some(120));
1398 assert_eq!(completion.completion_tokens, Some(30));
1399 }
1400
1401 #[tokio::test]
1402 async fn complete_gives_up_after_exhausting_length_retries() {
1403 let server = MockServer::start().await;
1408 Mock::given(method("POST"))
1409 .respond_with(
1410 ResponseTemplate::new(200).set_body_json(success_body_with_finish(
1411 r#"[1, 2, 3"#,
1412 Some(0.0),
1413 "length",
1414 )),
1415 )
1416 .mount(&server)
1417 .await;
1418
1419 let client = client_for(&server, "deepseek/deepseek-v4-flash:nitro").await;
1420 let err = client
1421 .complete("system", "input", TEST_SCHEMA, Some(64))
1422 .await
1423 .expect_err("exhausted length retries must fail");
1424 assert_eq!(err.finish_reason.as_deref(), Some("length"));
1425 assert_eq!(
1426 err.retry_class,
1427 AttemptOutcome::Transient,
1428 "a repeatedly truncated response is a bounded-retry hiccup, not permanent"
1429 );
1430
1431 let requests = server
1432 .received_requests()
1433 .await
1434 .expect("request recording is enabled");
1435 assert_eq!(
1436 requests.len() as u32,
1437 crate::constants::ENRICH_MAX_LENGTH_RETRIES + 1,
1438 "expected the primary attempt plus every regrowth retry"
1439 );
1440 }
1441}