Skip to main content

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}