a2a_llm/lib.rs
1//! Provider-neutral vocabulary for chat completions, plus the providers that
2//! speak it.
3//!
4//! [`LlmProvider`] is the port: [`chat_completion`](LlmProvider::chat_completion)
5//! and [`chat_completion_stream`](LlmProvider::chat_completion_stream) over
6//! [`LlmRequest`] / [`LlmResponse`]. [`openai`] covers OpenAI and every
7//! OpenAI-compatible endpoint (OpenRouter, vLLM, llama.cpp); [`gemini`] covers
8//! Google's API. [`provider_from_env`] picks one from the environment.
9//!
10//! The types are deliberately not tied to A2A. [`ToolCall`] and
11//! [`ToolDefinition`] are the tool-calling vocabulary shared with the MCP
12//! bridge, which is why they live in their own crate rather than inside an
13//! agent framework.
14
15use async_trait::async_trait;
16use futures::stream::BoxStream;
17use serde::{Deserialize, Serialize};
18use serde_json::Value;
19
20pub mod gemini;
21pub mod openai;
22pub mod provider;
23pub mod tool_call;
24
25pub use provider::{
26 LlmConfigError, LlmSettings, PROVIDER_ENV_VARS, ReasoningPlan, SUPPORTED_PROVIDERS,
27 SelectedLlm, provider_from_env, provider_from_settings,
28};
29pub use tool_call::{PartialToolCall, ToolCallAccumulator};
30
31/// The environment, as this crate reads it when building a provider.
32///
33/// Passed in rather than read directly so the selection rules can be tested
34/// without mutating the process environment, which would race other tests.
35#[derive(Clone, Copy)]
36pub(crate) struct Env<'a>(&'a dyn Fn(&str) -> Option<String>);
37
38impl<'a> Env<'a> {
39 /// A stand-in environment. Only tests need one; production reads
40 /// [`Env::os`].
41 #[cfg(test)]
42 pub(crate) fn new(lookup: &'a dyn Fn(&str) -> Option<String>) -> Self {
43 Self(lookup)
44 }
45
46 /// The process environment.
47 pub(crate) fn os() -> Env<'static> {
48 const LOOKUP: &dyn Fn(&str) -> Option<String> = &os_lookup;
49 Env(LOOKUP)
50 }
51
52 /// A variable set to whitespace reads as unset. `.env` files leave those
53 /// behind, and an empty `OPENROUTER_API_KEY` would otherwise select a
54 /// provider that cannot authenticate.
55 pub(crate) fn get(&self, key: &str) -> Option<String> {
56 (self.0)(key)
57 .map(|value| value.trim().to_string())
58 .filter(|value| !value.is_empty())
59 }
60}
61
62fn os_lookup(key: &str) -> Option<String> {
63 std::env::var(key).ok()
64}
65
66/// Represents an error returned by an LLM provider.
67#[derive(Debug, thiserror::Error)]
68pub enum LlmError {
69 #[error("API error: {0}")]
70 ApiError(String),
71 /// The request was larger than the model's context window.
72 ///
73 /// Separate from [`LlmError::ApiError`] because it is the one API failure a
74 /// caller can act on: drop history and try again. Folded into `ApiError` it
75 /// reached the handler as `A2AError::Internal("LLM error: API error (400)
76 /// …")` and simply failed the task.
77 ///
78 /// The numbers are carried where the provider's body names them, because
79 /// they are the difference between an answer and a bound: a caller advising
80 /// "set your ceiling below N" from its own token *estimate* hands out a
81 /// number with the estimator's error on it, while `context_window` is the
82 /// limit the provider actually enforced.
83 #[error("context length exceeded: {detail}")]
84 ContextLengthExceeded {
85 /// The provider's message as received (label, status, body).
86 detail: String,
87 /// Prompt tokens as the provider counted them, where the body names
88 /// them. The provider's count, never an estimate; `None` must not read
89 /// as zero (the same rule as [`TokenUsage`]).
90 prompt_tokens: Option<u32>,
91 /// The context window the provider enforced, where the body names it.
92 context_window: Option<u32>,
93 },
94 #[error("Network error: {0}")]
95 NetworkError(String),
96 #[error("Serialization error: {0}")]
97 SerializationError(String),
98 #[error("Provider error: {0}")]
99 ProviderError(String),
100}
101
102/// An error and everything under it, as one line.
103///
104/// `reqwest::Error`'s `Display` omits its source chain, so a DNS failure, a
105/// refused connection and an untrusted certificate all read as `error sending
106/// request for url (…)` — which is what made a TLS-intercepting proxy
107/// indistinguishable from the network being down, and cost a full investigation
108/// (see `NOTES.md`). The certificate error was one `source()` away the whole
109/// time. Takes `dyn Error` so the SSE stream's wrapper is covered by the same
110/// rule.
111pub(crate) fn describe_transport_error(error: &dyn std::error::Error) -> String {
112 let mut message = error.to_string();
113 let mut source = error.source();
114 while let Some(cause) = source {
115 message.push_str(": ");
116 message.push_str(&cause.to_string());
117 source = cause.source();
118 }
119 message
120}
121
122/// Substrings that identify an over-long request in a provider's error body.
123///
124/// Providers disagree on both the status code and the shape, and several return
125/// a plain 400 with prose, so matching on text is the only thing that works
126/// across all of them. Checked lowercase.
127const CONTEXT_LENGTH_MARKERS: [&str; 8] = [
128 // OpenAI (`"code": "context_length_exceeded"`), and OpenRouter passes it through.
129 "context_length_exceeded",
130 // OpenAI / OpenRouter prose, and most OpenAI-compatible servers.
131 "maximum context length",
132 "context length",
133 // The same concept spelled "size", which is llama.cpp's word for it and
134 // matched none of the above. Kept as a bare marker, symmetric with
135 // "context length", so a build whose error `type` differs from the one
136 // below is still recognized by its prose.
137 "context size",
138 // llama.cpp, vLLM.
139 "too many tokens",
140 "exceeds the maximum",
141 // llama.cpp b10524's `type`, for the same overflow:
142 // `{"type":"exceed_context_size_error","message":"request (40089 tokens)
143 // exceeds the available context size (32768 tokens)"}`. Matched on the
144 // `type` as well as the prose above, that being the half of the body least
145 // likely to be reworded.
146 "exceed_context_size_error",
147 // Gemini: INVALID_ARGUMENT naming the input token count.
148 "input token count",
149];
150
151/// Classify a provider's failure body, so an over-long request becomes
152/// [`LlmError::ContextLengthExceeded`] rather than an opaque API error.
153///
154/// One function for both providers and both code paths (streaming and not) so
155/// they classify identically — it owns the `{label} ({status}): {body}`
156/// format too, which used to be written at every call site. Takes the **raw**
157/// body because the numbers a context refusal carries are read from it before
158/// anything is flattened into prose.
159pub(crate) fn classify_api_error(label: &str, status: reqwest::StatusCode, body: &str) -> LlmError {
160 let message = format!("{label} ({status}): {body}");
161 let haystack = message.to_lowercase();
162 if CONTEXT_LENGTH_MARKERS
163 .iter()
164 .any(|marker| haystack.contains(marker))
165 {
166 let (prompt_tokens, context_window) = context_refusal_numbers(body);
167 return LlmError::ContextLengthExceeded {
168 detail: message,
169 prompt_tokens,
170 context_window,
171 };
172 }
173 LlmError::ApiError(message)
174}
175
176/// What the provider's body says about a refused request: (prompt tokens as it
177/// counted them, the window it enforced).
178///
179/// JSON fields first — llama.cpp returns `n_prompt_tokens` and `n_ctx` beside
180/// the message, inside the `error` object — then the two prose shapes the
181/// fixtures pin: OpenAI's "maximum context length is <N>" names the window,
182/// Gemini's "input token count (<N>)" names the count. Nothing speculative: a
183/// shape with no fixture is not parsed, because a wrong number is worse than
184/// none — these values are what a caller repairs its config with.
185fn context_refusal_numbers(body: &str) -> (Option<u32>, Option<u32>) {
186 let json: Option<Value> = serde_json::from_str(body).ok();
187 let field = |name: &str| {
188 let json = json.as_ref()?;
189 json.get(name)
190 .or_else(|| json.get("error").and_then(|error| error.get(name)))
191 .and_then(Value::as_u64)
192 .and_then(|n| u32::try_from(n).ok())
193 };
194 let prose = body.to_lowercase();
195 let prompt_tokens =
196 field("n_prompt_tokens").or_else(|| number_after(&prose, "input token count ("));
197 let context_window =
198 field("n_ctx").or_else(|| number_after(&prose, "maximum context length is "));
199 (prompt_tokens, context_window)
200}
201
202/// The unsigned number immediately following `pattern` in `haystack`, if any.
203fn number_after(haystack: &str, pattern: &str) -> Option<u32> {
204 let start = haystack.find(pattern)? + pattern.len();
205 let digits: &str = haystack[start..]
206 .split(|c: char| !c.is_ascii_digit())
207 .next()?;
208 digits.parse().ok()
209}
210
211/// Whether the model behind an endpoint accepts the reasoning parameter its
212/// adapter sends.
213///
214/// Support turns on the *model*, not the provider: `reasoning_effort` is a 400
215/// on `gpt-4o-mini` and mandatory on `gpt-5-pro`, and Gemini's `thinkingLevel`
216/// is accepted by some 2.5-generation models and refused by others. A table of
217/// model names answers that for the models it lists and goes stale with every
218/// release, so the parameter is sent instead, the refusal is read back off the
219/// 400, and the request is retried once without it.
220///
221/// The answer is remembered here, so only the first refused call of a process
222/// pays the extra round trip — and only once the retry has confirmed it, since a
223/// 400 that names the field can still be about something else. Shared across
224/// clones of a provider, since a provider is cloned per handler and they all
225/// call the same model.
226#[derive(Clone, Default)]
227pub(crate) struct ReasoningSupport(std::sync::Arc<std::sync::atomic::AtomicBool>);
228
229impl ReasoningSupport {
230 /// Whether a refusal has already been seen. Nothing is sent after one.
231 pub(crate) fn refused(&self) -> bool {
232 self.0.load(std::sync::atomic::Ordering::Relaxed)
233 }
234
235 /// Remember that this endpoint refused the parameter. `Relaxed` because a
236 /// racing caller re-sending once is the cost of being wrong.
237 pub(crate) fn record_refusal(&self) {
238 self.0.store(true, std::sync::atomic::Ordering::Relaxed);
239 }
240}
241
242/// Whether a failed request is the endpoint refusing the reasoning parameter,
243/// rather than the request failing on its own merits.
244///
245/// Both providers answer 400 and name the field: OpenAI with
246/// `Unsupported parameter: 'reasoning_effort' is not supported with this model`,
247/// Gemini with `Unknown name "thinkingLevel": Cannot find field`. A refused
248/// *value* reads the same way (`Invalid value: 'none'. Supported values are …`)
249/// and wants the same recovery, so the test is the field name rather than the
250/// prose. `field` is the one word every such message contains — `reasoning`,
251/// `thinking` — which is specific enough because this is only asked when that
252/// field was actually sent.
253pub(crate) fn refuses_reasoning(status: u16, body: &str, field: &str) -> bool {
254 status == 400 && body.to_lowercase().contains(field)
255}
256
257/// Tokens a provider reported for one request.
258///
259/// Reported rather than estimated: a caller's own token estimate decides what to
260/// send, and this says what it actually cost. Every field is optional because
261/// providers disagree on which they return, and a missing count must not read as
262/// zero.
263#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
264pub struct TokenUsage {
265 /// Tokens in the request, including the system prompt and tool definitions.
266 pub prompt_tokens: Option<u32>,
267 /// Tokens the model generated, excluding reasoning where a provider splits
268 /// them out.
269 pub completion_tokens: Option<u32>,
270 /// Reasoning tokens, where the provider reports them separately. Billed, and
271 /// invisible in `completion_tokens` on most providers.
272 pub reasoning_tokens: Option<u32>,
273 /// The provider's own total. Not derived from the fields above — a provider
274 /// that reports only this one is common, and a total that disagrees with the
275 /// parts is the provider's answer, not ours to correct.
276 pub total_tokens: Option<u32>,
277}
278
279impl TokenUsage {
280 /// Whether the provider reported anything at all. A response carrying no
281 /// counts is `Some(TokenUsage::default())` nowhere — it is `None`.
282 pub fn is_empty(&self) -> bool {
283 *self == Self::default()
284 }
285}
286
287impl std::fmt::Display for TokenUsage {
288 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
289 let field = |value: Option<u32>| match value {
290 Some(count) => count.to_string(),
291 None => "?".to_string(),
292 };
293 write!(
294 f,
295 "prompt={} completion={} reasoning={} total={}",
296 field(self.prompt_tokens),
297 field(self.completion_tokens),
298 field(self.reasoning_tokens),
299 field(self.total_tokens)
300 )
301 }
302}
303
304/// The role of the message sender.
305#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
306#[serde(rename_all = "lowercase")]
307pub enum MessageRole {
308 System,
309 User,
310 Assistant,
311 Tool,
312}
313
314/// Defines a tool (function) available for the LLM to call.
315#[derive(Debug, Clone, Serialize, Deserialize)]
316pub struct ToolDefinition {
317 pub name: String,
318 pub description: String,
319 pub parameters: Value, // JSON Schema representation of arguments
320}
321
322/// Represents a specific tool invocation requested by the LLM.
323#[derive(Debug, Clone, Serialize, Deserialize)]
324pub struct ToolCall {
325 pub id: String, // ID of the tool call
326 pub name: String,
327 pub arguments: String, // Stringified JSON arguments
328}
329
330/// A single message in a chat conversation.
331#[derive(Debug, Clone, Serialize, Deserialize)]
332pub struct ChatMessage {
333 pub role: MessageRole,
334 #[serde(skip_serializing_if = "Option::is_none")]
335 pub content: Option<String>,
336 #[serde(skip_serializing_if = "Option::is_none")]
337 pub tool_calls: Option<Vec<ToolCall>>,
338 #[serde(skip_serializing_if = "Option::is_none")]
339 pub tool_call_id: Option<String>,
340 #[serde(skip_serializing_if = "Option::is_none")]
341 pub name: Option<String>,
342}
343
344impl ChatMessage {
345 pub fn system(content: impl Into<String>) -> Self {
346 Self {
347 role: MessageRole::System,
348 content: Some(content.into()),
349 tool_calls: None,
350 tool_call_id: None,
351 name: None,
352 }
353 }
354
355 pub fn user(content: impl Into<String>) -> Self {
356 Self {
357 role: MessageRole::User,
358 content: Some(content.into()),
359 tool_calls: None,
360 tool_call_id: None,
361 name: None,
362 }
363 }
364
365 pub fn assistant(content: impl Into<String>) -> Self {
366 Self {
367 role: MessageRole::Assistant,
368 content: Some(content.into()),
369 tool_calls: None,
370 tool_call_id: None,
371 name: None,
372 }
373 }
374
375 pub fn tool_result(
376 tool_call_id: impl Into<String>,
377 name: impl Into<String>,
378 content: impl Into<String>,
379 ) -> Self {
380 Self {
381 role: MessageRole::Tool,
382 content: Some(content.into()),
383 tool_calls: None,
384 tool_call_id: Some(tool_call_id.into()),
385 name: Some(name.into()),
386 }
387 }
388}
389
390/// How hard a reasoning model should think, when reasoning is requested.
391#[derive(Debug, Clone, Copy, PartialEq, Eq)]
392pub enum ReasoningEffort {
393 Low,
394 Medium,
395 High,
396}
397
398impl ReasoningEffort {
399 /// The wire token used by OpenRouter's `reasoning.effort`.
400 pub fn as_str(self) -> &'static str {
401 match self {
402 ReasoningEffort::Low => "low",
403 ReasoningEffort::Medium => "medium",
404 ReasoningEffort::High => "high",
405 }
406 }
407}
408
409/// What to ask a reasoning-capable model to do with its thinking.
410///
411/// This is a *request*, not a capability: `Some(_)` says what the caller wants
412/// and `None` says nothing at all, leaving the model's own default alone.
413/// Whether the endpoint can carry it is the provider's business — a provider
414/// that cannot say this on the wire drops it rather than making every caller
415/// ask first.
416///
417/// Where the model does honour it, its thinking comes back on a separate channel
418/// ([`LlmResponse::reasoning`] / [`LlmStreamEvent::Reasoning`]).
419#[derive(Debug, Clone, Copy, PartialEq, Eq)]
420pub enum Reasoning {
421 /// Don't think — for models that let reasoning be turned off. The one
422 /// setting a small, fast model usually wants, and the one an effort-only
423 /// knob cannot express.
424 Off,
425 /// Think at one of the provider's named effort levels.
426 Effort(ReasoningEffort),
427 /// Think within a hard budget of reasoning tokens.
428 Budget(u32),
429}
430
431/// What a host's config or environment may spell, listed once so the parser,
432/// the error message, and the docs cannot drift apart.
433const REASONING_EXPECTED: &str =
434 r#""off", "low", "medium", "high", or a number of reasoning tokens"#;
435
436impl std::str::FromStr for Reasoning {
437 type Err = String;
438
439 /// Parses the tokens a host config or `OPENROUTER_REASONING` accepts:
440 /// `off`, `low`, `medium`, `high`, or a plain token budget (`2000`).
441 fn from_str(s: &str) -> Result<Self, Self::Err> {
442 match s.trim() {
443 "off" => Ok(Reasoning::Off),
444 "low" => Ok(Reasoning::Effort(ReasoningEffort::Low)),
445 "medium" => Ok(Reasoning::Effort(ReasoningEffort::Medium)),
446 "high" => Ok(Reasoning::Effort(ReasoningEffort::High)),
447 budget => budget
448 .parse()
449 .map(Reasoning::Budget)
450 .map_err(|_| format!("expected {REASONING_EXPECTED}; got {s:?}")),
451 }
452 }
453}
454
455impl std::fmt::Display for Reasoning {
456 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
457 match self {
458 Reasoning::Off => f.write_str("off"),
459 Reasoning::Effort(effort) => f.write_str(effort.as_str()),
460 Reasoning::Budget(tokens) => write!(f, "{tokens}"),
461 }
462 }
463}
464
465impl Serialize for Reasoning {
466 fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
467 match self {
468 // A budget round-trips as the number it was written as; the levels
469 // as their token. Both are what a host config spells.
470 Reasoning::Budget(tokens) => serializer.serialize_u32(*tokens),
471 level => serializer.serialize_str(&level.to_string()),
472 }
473 }
474}
475
476impl<'de> Deserialize<'de> for Reasoning {
477 /// Accepts a level (`"high"`) or a token budget (`2000`) — one parser for
478 /// every host, so a bad value is refused the same way with the same message
479 /// wherever it was written.
480 fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
481 use serde::de::{Error, Unexpected, Visitor};
482
483 struct ReasoningVisitor;
484
485 impl Visitor<'_> for ReasoningVisitor {
486 type Value = Reasoning;
487
488 fn expecting(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
489 f.write_str(REASONING_EXPECTED)
490 }
491
492 fn visit_str<E: Error>(self, value: &str) -> Result<Reasoning, E> {
493 value
494 .parse()
495 .map_err(|_| E::invalid_value(Unexpected::Str(value), &self))
496 }
497
498 fn visit_u64<E: Error>(self, value: u64) -> Result<Reasoning, E> {
499 u32::try_from(value)
500 .map(Reasoning::Budget)
501 .map_err(|_| E::invalid_value(Unexpected::Unsigned(value), &self))
502 }
503
504 fn visit_i64<E: Error>(self, value: i64) -> Result<Reasoning, E> {
505 u32::try_from(value)
506 .map(Reasoning::Budget)
507 .map_err(|_| E::invalid_value(Unexpected::Signed(value), &self))
508 }
509 }
510
511 deserializer.deserialize_any(ReasoningVisitor)
512 }
513}
514
515/// A request to an LLM provider for chat completion.
516#[derive(Debug, Clone)]
517pub struct LlmRequest {
518 pub messages: Vec<ChatMessage>,
519 pub tools: Option<Vec<ToolDefinition>>,
520 pub temperature: Option<f32>,
521 pub max_tokens: Option<u32>,
522 pub force_json: bool,
523 /// What this request asks of a reasoning model; `None` defers to whatever
524 /// default the provider was configured with, and then to the model's own.
525 pub reasoning: Option<Reasoning>,
526}
527
528impl LlmRequest {
529 pub fn new(messages: Vec<ChatMessage>) -> Self {
530 Self {
531 messages,
532 tools: None,
533 temperature: None,
534 max_tokens: None,
535 force_json: false,
536 reasoning: None,
537 }
538 }
539
540 pub fn reasoning(mut self, reasoning: Reasoning) -> Self {
541 self.reasoning = Some(reasoning);
542 self
543 }
544
545 pub fn temperature(mut self, temp: f32) -> Self {
546 self.temperature = Some(temp);
547 self
548 }
549
550 pub fn max_tokens(mut self, tokens: u32) -> Self {
551 self.max_tokens = Some(tokens);
552 self
553 }
554
555 pub fn tools(mut self, tools: Vec<ToolDefinition>) -> Self {
556 self.tools = Some(tools);
557 self
558 }
559
560 pub fn force_json(mut self, force: bool) -> Self {
561 self.force_json = force;
562 self
563 }
564}
565
566/// A response from an LLM provider.
567#[derive(Debug, Clone)]
568pub struct LlmResponse {
569 pub content: Option<String>,
570 pub tool_calls: Option<Vec<ToolCall>>,
571 /// Reasoning-model "thinking" text, when the provider exposes it separately
572 /// from the answer (e.g. OpenRouter's `reasoning`, Zhipu/GLM's
573 /// `reasoning_content`). `None` for providers that don't surface it.
574 pub reasoning: Option<String>,
575 /// What the provider says the request cost. `None` when it reported nothing.
576 pub usage: Option<TokenUsage>,
577}
578
579/// An event emitted during a streaming LLM response.
580#[derive(Debug, Clone)]
581pub enum LlmStreamEvent {
582 ContentChunk(String),
583 /// A chunk of reasoning-model "thinking" text, distinct from the answer
584 /// content (e.g. OpenRouter's `reasoning` / Zhipu's `reasoning_content`).
585 Reasoning(String),
586 ToolCallChunk {
587 id: String,
588 name: Option<String>,
589 arguments: String,
590 },
591 ToolCall(ToolCall),
592 /// What the request cost, as reported by the provider. Terminal: it arrives
593 /// in the final chunk, after the content. Absent on endpoints that do not
594 /// report usage while streaming — see `OpenAiConfig::stream_usage`.
595 Usage(TokenUsage),
596}
597
598/// Trait defining a generic LLM provider for standardizing AI integration across agents.
599#[async_trait]
600pub trait LlmProvider: Send + Sync {
601 /// Generates a chat completion based on the provided request.
602 async fn chat_completion(&self, request: LlmRequest) -> Result<LlmResponse, LlmError>;
603
604 /// Generates a streaming chat completion.
605 async fn chat_completion_stream(
606 &self,
607 request: LlmRequest,
608 ) -> Result<BoxStream<'static, Result<LlmStreamEvent, LlmError>>, LlmError>;
609}
610
611#[cfg(test)]
612mod error_tests {
613 use super::*;
614
615 fn status(code: u16) -> reqwest::StatusCode {
616 reqwest::StatusCode::from_u16(code).unwrap()
617 }
618
619 /// The one API failure a caller can act on has to be tellable from the rest,
620 /// across the shapes the providers actually return — and where a body names
621 /// the numbers, they come through, because they are what a caller repairs
622 /// its config with (an estimate-derived bound has the estimator's error on
623 /// it; `n_ctx` is the limit the provider actually enforced).
624 #[test]
625 fn an_over_long_request_is_classified_as_a_context_length_failure() {
626 // (label, body, provider-counted prompt tokens, enforced window)
627 let cases = [
628 (
629 "OpenAI API error",
630 r#"{"error":{"message":"This model's maximum context length is 128000 tokens","code":"context_length_exceeded"}}"#,
631 None,
632 Some(128000),
633 ),
634 (
635 "OpenAI stream error",
636 "Requested 200000 tokens, exceeds the maximum for this model",
637 // No fixture-pinned shape names these numbers, and a shape
638 // without a fixture is not parsed.
639 None,
640 None,
641 ),
642 (
643 "Gemini API error",
644 r#"{"error":{"status":"INVALID_ARGUMENT","message":"The input token count (1200000) exceeds the maximum"}}"#,
645 Some(1200000),
646 None,
647 ),
648 ("OpenAI API error", "too many tokens in prompt", None, None),
649 // llama.cpp b10524, verbatim. None of the markers above match it
650 // (it says "context size", not "context length", and "exceeds the
651 // available", not "exceeds the maximum") — and it is the one body
652 // carrying both numbers as JSON fields.
653 (
654 "OpenAI stream error",
655 r#"{"error":{"code":400,"message":"request (40089 tokens) exceeds the available context size (32768 tokens), try increasing it","type":"exceed_context_size_error","n_prompt_tokens":40089,"n_ctx":32768}}"#,
656 Some(40089),
657 Some(32768),
658 ),
659 ];
660 for (label, body, prompt, window) in cases {
661 match classify_api_error(label, status(400), body) {
662 LlmError::ContextLengthExceeded {
663 detail,
664 prompt_tokens,
665 context_window,
666 } => {
667 assert!(detail.contains(label) && detail.contains(body), "{detail}");
668 assert_eq!(prompt_tokens, prompt, "prompt tokens for: {body}");
669 assert_eq!(context_window, window, "window for: {body}");
670 }
671 other => panic!("should classify as context length: {body} → {other:?}"),
672 }
673 }
674 }
675
676 /// Everything else stays an ordinary API error. Classifying a bad key as
677 /// "too long" would send the handler into a compaction loop it can never win.
678 #[test]
679 fn other_failures_stay_api_errors() {
680 let cases = [
681 (
682 "OpenAI API error",
683 401,
684 r#"{"error":{"message":"Incorrect API key provided"}}"#,
685 ),
686 ("OpenAI API error", 429, "Rate limit reached for requests"),
687 ("Gemini API error", 503, "The model is overloaded"),
688 ];
689 for (label, code, body) in cases {
690 assert!(
691 matches!(
692 classify_api_error(label, status(code), body),
693 LlmError::ApiError(_)
694 ),
695 "should stay an API error: {body}"
696 );
697 }
698 }
699
700 /// The recovery turns on recognizing this one failure, so it is checked
701 /// against the bodies the two APIs actually answer with — an unsupported
702 /// parameter, an unknown field, and a value the model does not offer, which
703 /// wants the same retry as the other two.
704 #[test]
705 fn a_refused_reasoning_parameter_is_recognized_from_the_body() {
706 let openai = [
707 r#"{"error":{"message":"Unsupported parameter: 'reasoning_effort' is not supported with this model.","type":"invalid_request_error","param":"reasoning_effort","code":"unsupported_parameter"}}"#,
708 r#"{"error":{"message":"Invalid value: 'none'. Supported values are: 'low', 'medium' and 'high'.","type":"invalid_request_error","param":"reasoning_effort","code":"invalid_value"}}"#,
709 r#"{"error":{"message":"Unrecognized request argument supplied: reasoning_effort"}}"#,
710 ];
711 for body in openai {
712 assert!(refuses_reasoning(400, body, "reasoning"), "{body}");
713 }
714
715 let gemini = [
716 r#"{"error":{"code":400,"message":"Invalid JSON payload received. Unknown name \"thinkingLevel\" at 'generation_config': Cannot find field.","status":"INVALID_ARGUMENT"}}"#,
717 r#"{"error":{"code":400,"message":"Budget 128 is invalid. thinkingBudget must be 0 or in the range [128, 32768]","status":"INVALID_ARGUMENT"}}"#,
718 ];
719 for body in gemini {
720 assert!(refuses_reasoning(400, body, "thinking"), "{body}");
721 }
722 }
723
724 /// Everything else is the request failing on its own merits, and retrying it
725 /// without reasoning would waste a round trip and hide the real cause. A 5xx
726 /// that happens to name the field is the same: dropping the setting on an
727 /// outage would leave the model thinking at its default long after.
728 #[test]
729 fn other_failures_are_not_read_as_a_refusal() {
730 assert!(!refuses_reasoning(
731 400,
732 r#"{"error":{"message":"Incorrect API key provided"}}"#,
733 "reasoning"
734 ));
735 assert!(!refuses_reasoning(
736 429,
737 r#"{"error":{"message":"Rate limit reached"}}"#,
738 "reasoning"
739 ));
740 assert!(!refuses_reasoning(
741 503,
742 r#"{"error":{"message":"reasoning_effort backend unavailable"}}"#,
743 "reasoning"
744 ));
745 }
746
747 /// A refusal is remembered once, and every clone of the provider sees it:
748 /// they all call the same model.
749 #[test]
750 fn a_recorded_refusal_is_shared() {
751 let support = ReasoningSupport::default();
752 let clone = support.clone();
753 assert!(!support.refused());
754 clone.record_refusal();
755 assert!(support.refused());
756 }
757
758 #[test]
759 fn usage_with_nothing_reported_reads_as_empty() {
760 assert!(TokenUsage::default().is_empty());
761 assert!(
762 !TokenUsage {
763 prompt_tokens: Some(0),
764 ..Default::default()
765 }
766 .is_empty(),
767 "a reported zero is a report, not an absence"
768 );
769 }
770}