model.rs

//! Model provider integration and label-based routing.
//!
//! This module adapts provider-specific completion APIs to the common
//! [`CompletionRequest`] and [`AgentOutput`] contract used by Anda agents.
//! Built-in providers currently include OpenAI-compatible APIs, Anthropic, and
//! Google Gemini.
//!
//! The [`Models`] registry maps model labels such as `primary`, `pro`,
//! `flash`, or `lite` to concrete [`Model`] instances. Labels let agents
//! request capability tiers without hard-coding provider model names.
//!
//! Custom providers can implement [`CompletionFeaturesDyn`] and be wrapped with
//! [`Model::with_completer`].

use anda_core::{
    AgentOutput, BoxError, BoxPinFut, CONTENT_TYPE_JSON, CompletionRequest, Json, ToolCall,
};
use arc_swap::ArcSwap;
use futures_util::StreamExt;
use serde::de::DeserializeOwned;
use serde::{Deserialize, Serialize};
use std::future::Future;
use std::time::{Duration, Instant};
use std::{
    collections::{BTreeSet, HashMap, hash_map::Entry},
    error::Error,
    fmt,
    sync::Arc,
};

pub mod anthropic;
pub(crate) mod driver;
pub mod gemini;
pub mod openai;
pub(crate) mod raw;
#[cfg(test)]
pub(crate) mod test_support;
pub mod testing;

/// Deserializes a JSON `null` as the type's default value, so providers that
/// send explicit nulls for optional counters/objects do not fail typed parses.
pub(crate) fn null_default<'de, D, T>(deserializer: D) -> Result<T, D::Error>
where
    D: serde::Deserializer<'de>,
    T: Default + Deserialize<'de>,
{
    Ok(Option::<T>::deserialize(deserializer)?.unwrap_or_default())
}

/// Generates `as_str`, `Serialize`, and `Deserialize` for an open string enum:
/// known wire values map to variants (extra input-only aliases may follow the
/// canonical value with `|`), and unknown values are preserved in the given
/// catch-all variant instead of failing the parse.
macro_rules! string_enum_serde {
    ($ty:ident, { $($wire:literal $(| $alias:literal)* => $variant:ident),+ $(,)? }, $unknown:ident) => {
        impl $ty {
            fn as_str(&self) -> &str {
                match self {
                    $(Self::$variant => $wire,)+
                    Self::$unknown(value) => value.as_str(),
                }
            }
        }

        impl serde::Serialize for $ty {
            fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
            where
                S: serde::Serializer,
            {
                serializer.serialize_str(self.as_str())
            }
        }

        impl<'de> serde::Deserialize<'de> for $ty {
            fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
            where
                D: serde::Deserializer<'de>,
            {
                let value = String::deserialize(deserializer)?;
                Ok(match value.as_str() {
                    $($wire $(| $alias)* => Self::$variant,)+
                    _ => Self::$unknown(value),
                })
            }
        }
    };
}
pub(crate) use string_enum_serde;

/// Resolves a configured endpoint, falling back to the provider default when
/// the endpoint is absent or empty.
pub(crate) fn resolve_endpoint(endpoint: Option<String>, default: &str) -> String {
    match endpoint {
        Some(endpoint) if !endpoint.is_empty() => endpoint,
        _ => default.to_string(),
    }
}

pub use reqwest;
pub use reqwest::Proxy;

use crate::APP_USER_AGENT;

pub use anda_core::ModelEffort;

const MODEL_REQUEST_MAX_RETRIES: usize = 3;
const MODEL_RETRY_BACKOFF: Duration = Duration::from_secs(1);
const MODEL_RETRY_MAX_BACKOFF: Duration = Duration::from_secs(300);
const COMPLETION_HTTP2_KEEP_ALIVE_INTERVAL: Option<Duration> = None;
const COMPLETION_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
const COMPLETION_READ_TIMEOUT: Duration = Duration::from_secs(180);
const COMPLETION_REQUEST_TIMEOUT: Duration = Duration::from_secs(600);

/// Upper bound on the number of bytes buffered from a completion response.
///
/// Guards against a runaway or malicious provider streaming an unbounded body
/// and exhausting memory. Text completion responses are far smaller than this
/// in practice.
const MAX_COMPLETION_RESPONSE_BYTES: usize = 64 * 1024 * 1024;

/// Upper bound on the response-body excerpt embedded in a completion error.
///
/// Error messages end up in logs and in `failed_reason`, so a gateway's HTML
/// error page or a malformed multi-megabyte body is cut to a diagnostic prefix.
const MAX_ERROR_BODY_BYTES: usize = 8 * 1024;

/// Returns a lossy UTF-8 excerpt of `data` bounded by [`MAX_ERROR_BODY_BYTES`],
/// for embedding a provider response body in an error message.
pub(crate) fn error_body_excerpt(data: &[u8]) -> String {
    if data.len() <= MAX_ERROR_BODY_BYTES {
        return String::from_utf8_lossy(data).into_owned();
    }

    let mut text = String::from_utf8_lossy(&data[..MAX_ERROR_BODY_BYTES]).into_owned();
    text.push_str("… [truncated]");
    text
}

/// Reads a non-success response body for diagnostics, stopping once
/// [`MAX_ERROR_BODY_BYTES`] is exceeded instead of buffering the whole body.
async fn read_error_body(response: reqwest::Response) -> Result<String, reqwest::Error> {
    let mut stream = response.bytes_stream();
    let mut body = Vec::new();
    while let Some(chunk) = stream.next().await {
        body.extend_from_slice(&chunk?);
        if body.len() > MAX_ERROR_BODY_BYTES {
            break;
        }
    }
    Ok(error_body_excerpt(&body))
}

/// Serializable configuration for constructing a model adapter.
///
/// [`Debug`] is implemented manually so the `api_key` is redacted and never
/// leaks through `{:?}` logging.
#[derive(Default, Clone, Deserialize, Serialize)]
pub struct ModelConfig {
    /// Provider family, such as `gemini`, `anthropic`, `openai-response` or `openai`.
    pub family: String,

    /// Provider-specific model name.
    pub model: String,

    /// Base URL for the provider API.
    pub api_base: String,

    /// API key used by the provider adapter.
    pub api_key: String,

    /// Optional labels for selecting this model in the engine.
    ///
    /// If omitted, the provider model name is used as the only label. Common
    /// labels include `primary`, `pro`, `flash`, `lite`, `audio`, `video`, `image`, `memory`.
    #[serde(default)]
    pub labels: Vec<String>,

    #[serde(default)]
    /// Provider context window in input tokens; `0` means unknown.
    pub context_window: usize,

    #[serde(default)]
    /// Provider maximum output tokens; `0` means unknown.
    ///
    /// When known, it caps each request's explicit output budget (see
    /// [`Model::completion`]) and replaces the default budget of adapters that
    /// always send one (Anthropic `max_tokens`, Gemini `maxOutputTokens`).
    pub max_output: usize,

    /// Optional reasoning/thinking effort for providers and models that support it.
    ///
    /// Supported config values are `minimal`, `low`, `medium`, `high`, and `max`.
    /// The effective set depends on the selected provider and model.
    #[serde(default)]
    pub effort: Option<ModelEffort>,

    /// Skips this model when loading a list of configs.
    #[serde(default)]
    pub disabled: bool,

    /// Sends Anthropic credentials with bearer authentication instead of the
    /// provider-specific API-key header.
    #[serde(default)]
    pub bearer_auth: bool,

    #[serde(default)]
    /// Whether to request streaming completions from this model.
    /// The OpenAI Responses adapter always streams and aggregates the result,
    /// regardless of this setting.
    pub stream: bool,
}

impl std::fmt::Debug for ModelConfig {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ModelConfig")
            .field("family", &self.family)
            .field("model", &self.model)
            .field("api_base", &self.api_base)
            .field(
                "api_key",
                &if self.api_key.is_empty() {
                    ""
                } else {
                    "[REDACTED]"
                },
            )
            .field("labels", &self.labels)
            .field("context_window", &self.context_window)
            .field("max_output", &self.max_output)
            .field("effort", &self.effort)
            .field("disabled", &self.disabled)
            .field("bearer_auth", &self.bearer_auth)
            .field("stream", &self.stream)
            .finish()
    }
}

impl ModelConfig {
    /// Builds a [`Model`] from this configuration.
    pub fn model(&self, http_client: reqwest::Client) -> Result<Model, BoxError> {
        if self.disabled {
            return Err("model is disabled".into());
        }
        if self.model.is_empty() {
            return Err(format!("{}: model name is required", self.model).into());
        }
        if self.family.is_empty() {
            return Err(format!("{}: model family is required", self.model).into());
        }
        if self.api_base.is_empty() {
            return Err(format!("{}: api_base is required", self.model).into());
        }
        if self.api_key.is_empty() {
            return Err(format!("{}: api_key is required", self.model).into());
        }

        let mut model = match self.family.as_str() {
            "gemini" => Model::with_completer(Arc::new(
                gemini::Client::new_with_client(
                    &self.api_key,
                    Some(self.api_base.clone()),
                    http_client,
                )
                .completion_model(&self.model)
                .with_stream(self.stream)
                .with_effort(self.effort)
                .with_max_output(self.max_output),
            )),
            "anthropic" => {
                let mut cli = anthropic::Client::new_with_client(
                    &self.api_key,
                    Some(self.api_base.clone()),
                    http_client,
                );
                if self.bearer_auth {
                    cli = cli.with_bearer_auth(true);
                }
                Model::with_completer(Arc::new(
                    cli.completion_model(&self.model)
                        .with_stream(self.stream)
                        .with_effort(self.effort)
                        .with_max_output(self.max_output),
                ))
            }
            "openai-response" => {
                let cli = openai::Client::new_with_client(
                    &self.api_key,
                    Some(self.api_base.clone()),
                    http_client,
                );
                Model::with_completer(Arc::new(
                    cli.completion_model_v2(&self.model)
                        .with_stream(self.stream)
                        .with_effort(self.effort),
                ))
            }
            "openai" => {
                let cli = openai::Client::new_with_client(
                    &self.api_key,
                    Some(self.api_base.clone()),
                    http_client,
                );
                if self.model.starts_with("gpt") {
                    Model::with_completer(Arc::new(
                        cli.completion_model_v2(&self.model)
                            .with_stream(self.stream)
                            .with_effort(self.effort),
                    ))
                } else {
                    Model::with_completer(Arc::new(
                        cli.completion_model(&self.model)
                            .with_stream(self.stream)
                            .with_effort(self.effort),
                    ))
                }
            }
            _ => return Err(format!("unsupported model family: {}", self.family).into()),
        };

        let labels = if self.labels.is_empty() {
            vec![self.model.to_ascii_lowercase()]
        } else {
            self.labels.clone()
        };
        model.context_window = self.context_window;
        model.max_output = self.max_output;
        Ok(model.with_labels(labels))
    }
}

/// Thread-safe model registry used by the engine.
///
/// It maintains two layers:
/// - `model`: the primary default model for general requests
/// - `models`: a label-based map for selecting specific models
///
/// The dedicated primary slot can be set explicitly via [`Models::set_model`]
/// or derived from the special label `primary` in the label map. This keeps
/// direct lookup (`get`) separate from default-routing (`get_model`).
pub struct Models {
    model: ArcSwap<Option<Model>>,
    models: ArcSwap<HashMap<String, Vec<Model>>>,
}

impl Default for Models {
    fn default() -> Self {
        Self {
            model: ArcSwap::new(Arc::new(None)),
            models: ArcSwap::new(Arc::new(HashMap::new())),
        }
    }
}

impl Models {
    /// Creates a new Models instance by cloning the internal state of another Models instance.
    pub fn from_clone(other: &Models) -> Self {
        Self {
            model: ArcSwap::new(other.model.load_full()),
            models: ArcSwap::new(Arc::new(other.models.load().as_ref().clone())),
        }
    }

    /// Replaces this registry with a clone of another [`Models`] instance.
    pub fn replace(&self, other: &Models) {
        self.model.store(other.model.load_full());
        self.models
            .store(Arc::new(other.models.load().as_ref().clone()));
    }

    /// Builds a registry from model configs by registering every resolved label.
    ///
    /// Configs explicitly marked `disabled` are skipped quietly (info level);
    /// configs that fail to build because of a misconfiguration are skipped with
    /// a warning so the problem is not lost silently.
    pub fn from_configs(configs: &[ModelConfig], http_client: reqwest::Client) -> Self {
        let models = Self::default();
        for config in configs {
            if config.disabled {
                log::info!(
                    "skipping disabled model: family={}, model={}",
                    config.family,
                    config.model
                );
                continue;
            }
            match config.model(http_client.clone()) {
                Ok(model) => models.inner_set(model.labels.clone(), model),
                Err(err) => {
                    log::warn!(
                        "skipping misconfigured model: family={}, model={}, error={}",
                        config.family,
                        config.model,
                        err
                    );
                }
            }
        }
        models
    }

    /// Returns whether a label exists in the direct lookup table.
    pub fn contains(&self, label: &str) -> bool {
        self.models.load().contains_key(&label.to_ascii_lowercase())
    }

    /// Returns the set of all registered model names across all labels.
    pub fn model_names(&self) -> BTreeSet<String> {
        self.models
            .load()
            .values()
            .flatten()
            .map(|m| m.model_name())
            .collect()
    }

    /// Sets the primary default model without mutating the label map.
    pub fn set_model(&self, model: Model) {
        self.inner_set(model.labels.clone(), model.clone());
        self.model.store(Arc::new(Some(model)));
    }

    /// Inserts or updates a single labeled model.
    ///
    /// The special label `primary` also updates the dedicated routing slot.
    /// If no primary exists yet, any inserted model is promoted
    /// to become the primary default.
    pub fn set(&self, label: String, model: Model) {
        self.inner_set(vec![label], model);
    }

    fn inner_set(&self, mut labels: Vec<String>, model: Model) {
        if self.model.load().is_none() {
            self.model.store(Arc::new(Some(model.clone())));
        }

        let model_name = model.model_name();
        labels.push(model_name.to_ascii_lowercase());
        for label in labels.iter_mut() {
            label.make_ascii_lowercase();
            if label == "primary" {
                self.model.store(Arc::new(Some(model.clone())));
            }
        }

        // rcu keeps concurrent inserts from losing each other's labels.
        self.models.rcu(|models| {
            let mut models = models.as_ref().clone();
            for label in &labels {
                match models.entry(label.clone()) {
                    Entry::Vacant(e) => {
                        e.insert(vec![model.clone()]);
                    }
                    Entry::Occupied(mut e) => {
                        e.get_mut().retain(|m| m.model_name() != model_name);
                        e.get_mut().push(model.clone());
                    }
                }
            }
            models
        });
    }

    /// Returns a model by lowercase label if it exists.
    ///
    /// This is a direct lookup only and never falls back to default routing.
    pub fn get(&self, label: &str) -> Option<Model> {
        self.models
            .load()
            .get(&label.to_ascii_lowercase())
            .and_then(|v| v.last().cloned())
    }

    /// Returns the primary model if available; otherwise returns any remaining
    /// labeled model.
    pub fn get_model(&self) -> Option<Model> {
        if let Some(m) = self.model.load().as_ref() {
            return Some(m.clone());
        }
        self.models
            .load()
            .values()
            .next()
            .and_then(|v| v.last().cloned())
    }

    /// Resolves a model for lowercase-label-aware routing.
    ///
    /// Resolution order is:
    /// - the exact label match when `label` is non-empty
    /// - the default routing result from [`Models::get_model`]
    pub fn resolve(&self, label: &str) -> Option<Model> {
        if label.is_empty() {
            return self.get_model();
        }
        self.get(label).or_else(|| self.get_model())
    }
}

/// Object-safe completion provider interface.
pub trait CompletionFeaturesDyn: Send + Sync + 'static {
    /// Performs a completion request and returns the agent-facing output.
    ///
    /// Built-in adapters wrap exhausted transient provider failures in
    /// [`ModelError`]. Use [`is_retryable_box_error`] to decide whether an upper
    /// layer should schedule a delayed retry.
    fn completion(&self, req: CompletionRequest) -> BoxPinFut<Result<AgentOutput, BoxError>>;

    /// Returns the provider model name used for diagnostics and usage reports.
    fn model_name(&self) -> String;

    /// Removes unanswered tool-call requests from `raw_history[start..]`.
    ///
    /// `raw_history` holds this provider's own message JSON, so only the
    /// provider can classify its items reliably. The runner calls this after an
    /// interrupt (steering, discard, stop) so the next request does not inherit
    /// a tool-call requirement the provider would reject as unanswered. Visible
    /// text and reasoning must stay untouched, except reasoning that is valid
    /// only together with the pruned call.
    ///
    /// The default implementation is a conservative union over the built-in
    /// provider wire shapes; override it when your provider's raw items are not
    /// covered by it.
    fn prune_unanswered_tool_calls(&self, raw_history: &mut Vec<Json>, start: usize) {
        raw::prune_unanswered_tool_calls(raw_history, start);
    }

    /// Removes completed tool interactions (calls and results) from `raw_history`.
    ///
    /// Long-lived callers invoke this at an idle boundary to reclaim
    /// context-window budget from tool payloads the model has already consumed.
    /// Visible text and reasoning must stay untouched, and provider ordering
    /// constraints must be preserved: a reasoning item that must precede its
    /// sibling, or Anthropic thinking that is valid only in its unchanged
    /// message, leaves together with the pruned call.
    ///
    /// The default implementation is a conservative union over the built-in
    /// provider wire shapes; override it when your provider's raw items are not
    /// covered by it.
    fn prune_tool_interactions(&self, raw_history: &mut Vec<Json>) {
        raw::prune_tool_interactions(raw_history);
    }

    /// Replaces inline media payloads (base64 blocks, `data:` URLs, inline audio and file data)
    /// in `raw_history` with a short text note in the provider's text shape.
    ///
    /// Callers that send attachment bytes for one task use this at an idle boundary, so the
    /// bytes are not replayed on every later request. Media referenced by remote URL or
    /// provider file id must stay, messages the model wrote keep the media it generated, and
    /// no message may be left without content. Where the provider accepts only media (Gemini
    /// function-response parts), inline media is removed and the note joins the response text.
    ///
    /// The default implementation is a conservative union over the built-in provider wire
    /// shapes; override it when your provider's raw items are not covered by it.
    fn prune_inline_media(&self, raw_history: &mut Vec<Json>) {
        raw::prune_inline_media(raw_history);
    }
}

/// Placeholder implementation that returns errors for completion requests.
#[derive(Clone, Debug)]
pub struct NotImplemented;

impl CompletionFeaturesDyn for NotImplemented {
    fn model_name(&self) -> String {
        "not_implemented".to_string()
    }

    fn completion(&self, _req: CompletionRequest) -> BoxPinFut<Result<AgentOutput, BoxError>> {
        Box::pin(futures::future::ready(Err("not implemented".into())))
    }
}

/// Mock implementation for tests and examples.
#[derive(Clone, Debug)]
pub struct MockImplemented;

impl CompletionFeaturesDyn for MockImplemented {
    fn model_name(&self) -> String {
        "mock_implemented".to_string()
    }

    fn completion(&self, req: CompletionRequest) -> BoxPinFut<Result<AgentOutput, BoxError>> {
        Box::pin(futures::future::ready(Ok(AgentOutput {
            content: req.prompt.clone(),
            tool_calls: req
                .tools
                .iter()
                .filter_map(|tool| {
                    if req.prompt.is_empty() {
                        return None;
                    }
                    Some(ToolCall {
                        name: tool.name.clone(),
                        args: serde_json::from_str(&req.prompt).unwrap_or_default(),
                        call_id: None,
                        result: None,
                        remote_id: None,
                    })
                })
                .collect(),
            ..Default::default()
        })))
    }
}

/// Concrete model entry registered with the engine.
#[derive(Clone)]
pub struct Model {
    /// Completion provider implementation.
    pub completer: Arc<dyn CompletionFeaturesDyn>,

    /// Labels that can route requests to this model.
    pub labels: Vec<String>,

    /// Context window in input tokens; `0` means unknown.
    pub context_window: usize,

    /// Maximum output tokens; `0` means unknown. A known value caps each
    /// request's explicit `max_output_tokens`.
    pub max_output: usize,
}

impl Model {
    /// Creates a model from a completion provider.
    pub fn new(completer: Arc<dyn CompletionFeaturesDyn>) -> Self {
        Self {
            completer,
            labels: Vec::new(),
            context_window: 0,
            max_output: 0,
        }
    }

    /// Creates a model from a completion provider.
    pub fn with_completer(completer: Arc<dyn CompletionFeaturesDyn>) -> Self {
        Self::new(completer)
    }

    /// Assigns labels used by [`Models`] for routing.
    pub fn with_labels(mut self, labels: Vec<String>) -> Self {
        self.labels = labels;
        self
    }

    /// Creates a model whose completion calls return `not implemented` errors.
    pub fn not_implemented() -> Self {
        Self::new(Arc::new(NotImplemented))
    }

    /// Creates a model with deterministic mock completion behavior for tests.
    pub fn mock_implemented() -> Self {
        Self::new(Arc::new(MockImplemented))
    }

    /// Returns the provider model name for this model.
    pub fn model_name(&self) -> String {
        self.completer.model_name()
    }

    /// Executes a completion request with the underlying provider.
    ///
    /// A known [`Model::max_output`] caps an explicit `max_output_tokens`, so a request never
    /// asks the provider for more than the configured model accepts. An unset budget stays
    /// unset: defaulting it is the provider adapter's decision (see
    /// [`ModelConfig::model`]), since some APIs reject an output limit they do not need.
    pub async fn completion(&self, mut req: CompletionRequest) -> Result<AgentOutput, BoxError> {
        if self.max_output > 0
            && let Some(tokens) = req.max_output_tokens.as_mut()
        {
            *tokens = (*tokens).min(self.max_output);
        }
        self.completer.completion(req).await
    }

    /// Removes unanswered tool-call requests from `raw_history[start..]` using
    /// the provider's own wire-format knowledge.
    ///
    /// See [`CompletionFeaturesDyn::prune_unanswered_tool_calls`].
    pub fn prune_unanswered_tool_calls(&self, raw_history: &mut Vec<Json>, start: usize) {
        self.completer
            .prune_unanswered_tool_calls(raw_history, start);
    }

    /// Removes completed tool interactions from `raw_history` using the
    /// provider's own wire-format knowledge.
    ///
    /// See [`CompletionFeaturesDyn::prune_tool_interactions`].
    pub fn prune_tool_interactions(&self, raw_history: &mut Vec<Json>) {
        self.completer.prune_tool_interactions(raw_history);
    }

    /// Replaces inline media payloads in `raw_history` using the provider's own wire-format
    /// knowledge.
    ///
    /// See [`CompletionFeaturesDyn::prune_inline_media`].
    pub fn prune_inline_media(&self, raw_history: &mut Vec<Json>) {
        self.completer.prune_inline_media(raw_history);
    }
}

/// Error returned by built-in model adapters when the caller can inspect retry
/// semantics after the SDK-level retry has already been attempted.
#[derive(Debug)]
pub struct ModelError {
    message: String,
    retryable: bool,
    status: Option<http::StatusCode>,
    retry_after: Option<Duration>,
    source: Option<BoxError>,
}

impl ModelError {
    /// Creates a non-retryable model error.
    pub fn new(message: impl Into<String>) -> Self {
        Self {
            message: message.into(),
            retryable: false,
            status: None,
            retry_after: None,
            source: None,
        }
    }

    /// Marks whether this error should be considered retryable by the caller.
    pub fn with_retryable(mut self, retryable: bool) -> Self {
        self.retryable = retryable;
        self
    }

    /// Attaches the upstream HTTP status, if the provider returned one.
    pub fn with_status(mut self, status: http::StatusCode) -> Self {
        self.status = Some(status);
        self
    }

    /// Attaches a suggested retry delay from the upstream response.
    pub fn with_retry_after(mut self, retry_after: Option<Duration>) -> Self {
        self.retry_after = retry_after;
        self
    }

    /// Attaches the lower-level transport/read error.
    pub fn with_source(mut self, source: BoxError) -> Self {
        self.source = Some(source);
        self
    }

    /// Returns true when the upper layer may choose a delayed retry.
    pub fn is_retryable(&self) -> bool {
        self.retryable
    }

    /// Returns the upstream HTTP status, when present.
    pub fn status(&self) -> Option<http::StatusCode> {
        self.status
    }

    /// Returns the upstream retry delay, when present.
    pub fn retry_after(&self) -> Option<Duration> {
        self.retry_after
    }
}

impl fmt::Display for ModelError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(&self.message)
    }
}

impl Error for ModelError {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        self.source
            .as_deref()
            .map(|source| source as &(dyn Error + 'static))
    }
}

/// Walks an error chain and returns the first value produced by `f`.
fn find_in_error_chain<T>(
    error: &(dyn Error + 'static),
    f: impl Fn(&(dyn Error + 'static)) -> Option<T>,
) -> Option<T> {
    let mut current = Some(error);
    while let Some(error) = current {
        if let Some(found) = f(error) {
            return Some(found);
        }
        current = error.source();
    }
    None
}

/// Returns true if the error chain carries a retryable model error signal.
///
/// This function and its siblings ([`is_retryable_box_error`],
/// [`model_error_status`], [`model_error_retry_after`]) are the public
/// inspection surface of [`ModelError`] for embedding applications: built-in
/// adapters have already exhausted their own transport retries when they
/// return one, and the engine deliberately does not retry again on its own —
/// whether to schedule a delayed retry is an application-level policy.
pub fn is_retryable_model_error(error: &(dyn Error + 'static)) -> bool {
    find_in_error_chain(error, |error| {
        let retryable = error
            .downcast_ref::<ModelError>()
            .is_some_and(ModelError::is_retryable)
            || error
                .downcast_ref::<reqwest::Error>()
                .is_some_and(is_retryable_reqwest_error);
        retryable.then_some(())
    })
    .is_some()
}

/// Convenience wrapper for callers that keep completion errors as [`BoxError`].
pub fn is_retryable_box_error(error: &BoxError) -> bool {
    is_retryable_model_error(error.as_ref() as &(dyn Error + 'static))
}

/// Returns the first [`ModelError`] HTTP status found in the error chain.
pub fn model_error_status(error: &(dyn Error + 'static)) -> Option<http::StatusCode> {
    find_in_error_chain(error, |error| error.downcast_ref::<ModelError>()?.status())
}

/// Returns the first upstream retry delay found in the error chain.
pub fn model_error_retry_after(error: &(dyn Error + 'static)) -> Option<Duration> {
    find_in_error_chain(error, |error| {
        error.downcast_ref::<ModelError>()?.retry_after()
    })
}

/// Statuses that are transient enough for one immediate SDK retry and for an
/// upper-layer delayed retry after the SDK retry has been exhausted.
pub fn is_retryable_status(status: http::StatusCode) -> bool {
    matches!(
        status,
        http::StatusCode::REQUEST_TIMEOUT
            | http::StatusCode::TOO_MANY_REQUESTS
            | http::StatusCode::INTERNAL_SERVER_ERROR
            | http::StatusCode::BAD_GATEWAY
            | http::StatusCode::SERVICE_UNAVAILABLE
            | http::StatusCode::GATEWAY_TIMEOUT
    ) || status.as_u16() == 529
}

pub(crate) fn is_retryable_reqwest_error(err: &reqwest::Error) -> bool {
    err.is_timeout()
        || err.is_connect()
        || err.is_body()
        || err.is_decode()
        || err.status().is_some_and(is_retryable_status)
}

/// Formats an error together with its source chain, e.g.
/// "error decoding response body: request or response body error: operation
/// timed out". reqwest 0.12.2+ stopped including sources in `Display`, so the
/// top-level message alone hides the root cause of transport failures behind
/// a generic phrase.
pub(crate) fn format_error_chain(err: &(dyn Error + 'static)) -> String {
    let mut message = err.to_string();
    let mut source = err.source();
    while let Some(err) = source {
        let text = err.to_string();
        // Some errors repeat their source in `Display`; skip duplicates.
        if !message.contains(&text) {
            message.push_str(": ");
            message.push_str(&text);
        }
        source = err.source();
    }
    message
}

/// Extracts the upstream request id from response headers for error context.
/// Covers the header names used by OpenAI-compatible APIs, Anthropic, and
/// common gateway/CDN fronts.
fn upstream_request_id(headers: &http::HeaderMap) -> Option<String> {
    ["x-request-id", "request-id", "x-amzn-requestid", "cf-ray"]
        .into_iter()
        .find_map(|name| headers.get(name)?.to_str().ok())
        .map(str::to_string)
}

pub(crate) fn completion_transport_error(
    model: &str,
    action: &str,
    err: reqwest::Error,
) -> BoxError {
    let retryable = is_retryable_reqwest_error(&err);
    let message = format!(
        "{action}, model: {model}, error: {}",
        format_error_chain(&err)
    );
    Box::new(
        ModelError::new(message)
            .with_retryable(retryable)
            .with_source(Box::new(err)),
    )
}

pub(crate) async fn read_completion_response_bytes(
    response: reqwest::Response,
    model: &str,
) -> Result<bytes::Bytes, BoxError> {
    let request_id = upstream_request_id(response.headers());
    // Reject up front when the declared length already exceeds the cap.
    if let Some(len) = response.content_length()
        && len > MAX_COMPLETION_RESPONSE_BYTES as u64
    {
        return Err(completion_response_too_large(
            model,
            request_id.as_deref(),
            len as usize,
        ));
    }

    // Stream the body so the size cap is enforced even when no `Content-Length`
    // is provided.
    let mut stream = response.bytes_stream();
    let mut body: Vec<u8> = Vec::new();
    while let Some(chunk) = stream.next().await {
        let chunk = chunk.map_err(|err| {
            let action = format!(
                "Failed to read completion response (request_id: {})",
                request_id.as_deref().unwrap_or("-")
            );
            completion_transport_error(model, &action, err)
        })?;
        if body.len() + chunk.len() > MAX_COMPLETION_RESPONSE_BYTES {
            return Err(completion_response_too_large(
                model,
                request_id.as_deref(),
                body.len() + chunk.len(),
            ));
        }
        body.extend_from_slice(&chunk);
    }
    Ok(body.into())
}

/// Builds the error returned when a completion response exceeds the size cap.
fn completion_response_too_large(
    model: &str,
    request_id: Option<&str>,
    received: usize,
) -> BoxError {
    Box::new(ModelError::new(format!(
        "Completion response too large (model: {}, request_id: {}, received: {} bytes, limit: {} bytes)",
        model,
        request_id.unwrap_or("-"),
        received,
        MAX_COMPLETION_RESPONSE_BYTES,
    )))
}

pub(crate) async fn execute_completion_request_with_retry<T, BuildRequest, HandleResponse, Fut>(
    model: &str,
    build_request: BuildRequest,
    handle_response: HandleResponse,
) -> Result<T, BoxError>
where
    BuildRequest: Fn() -> reqwest::RequestBuilder,
    HandleResponse: Fn(reqwest::Response) -> Fut,
    Fut: Future<Output = Result<T, BoxError>>,
{
    for attempt in 0..=MODEL_REQUEST_MAX_RETRIES {
        let response = match build_request().send().await {
            Ok(response) => response,
            Err(err) => {
                let retryable = is_retryable_reqwest_error(&err);
                let message = format!(
                    "Failed to send completion request, model: {}, error: {}",
                    model,
                    format_error_chain(&err)
                );
                if retryable && attempt < MODEL_REQUEST_MAX_RETRIES {
                    log_completion_retry(model, attempt + 1);
                    backoff_before_retry(None).await;
                    continue;
                }

                return Err(Box::new(
                    ModelError::new(message)
                        .with_retryable(retryable)
                        .with_source(Box::new(err)),
                ));
            }
        };

        let status = response.status();
        if status.is_success() {
            match handle_response(response).await {
                Ok(output) => return Ok(output),
                Err(err) if is_retryable_box_error(&err) && attempt < MODEL_REQUEST_MAX_RETRIES => {
                    log_completion_retry(model, attempt + 1);
                    backoff_before_retry(None).await;
                    continue;
                }
                Err(err) => return Err(err),
            }
        }

        let retryable = is_retryable_status(status);
        let retry_after = retry_after_duration(response.headers());
        let body = match read_error_body(response).await {
            Ok(body) => body,
            Err(err) => {
                let retryable = retryable || is_retryable_reqwest_error(&err);
                let message = format!(
                    "Completion failed, model: {}, status: {}; failed to read error body: {}",
                    model,
                    status,
                    format_error_chain(&err)
                );
                if retryable && attempt < MODEL_REQUEST_MAX_RETRIES {
                    log_completion_retry(model, attempt + 1);
                    backoff_before_retry(retry_after).await;
                    continue;
                }

                return Err(Box::new(
                    ModelError::new(message)
                        .with_retryable(retryable)
                        .with_status(status)
                        .with_retry_after(retry_after)
                        .with_source(Box::new(err)),
                ));
            }
        };
        let message = format!(
            "Completion failed, model: {}, status: {}, body: {}",
            model, status, body
        );

        if retryable && attempt < MODEL_REQUEST_MAX_RETRIES {
            log_completion_retry(model, attempt + 1);
            backoff_before_retry(retry_after).await;
            continue;
        }

        return Err(Box::new(
            ModelError::new(message)
                .with_retryable(retryable)
                .with_status(status)
                .with_retry_after(retry_after),
        ));
    }

    unreachable!("completion retry loop always returns before exhausting attempts")
}

/// Sleeps before an in-SDK retry so transient overload (429/5xx/connection
/// flaps) has a chance to clear. The upstream `Retry-After` hint is honored up
/// to the configured retry cap; longer waits are the responsibility of upper
/// layers, which receive the hint via [`ModelError::retry_after`].
async fn backoff_before_retry(retry_after: Option<Duration>) {
    let delay = retry_after
        .unwrap_or(MODEL_RETRY_BACKOFF)
        .min(MODEL_RETRY_MAX_BACKOFF);
    tokio::time::sleep(delay).await;
}

fn retry_after_duration(headers: &http::HeaderMap) -> Option<Duration> {
    let value = headers
        .get(http::header::RETRY_AFTER)?
        .to_str()
        .ok()?
        .trim();
    if let Ok(seconds) = value.parse::<u64>() {
        return Some(Duration::from_secs(seconds));
    }

    // HTTP-date form, e.g. "Wed, 21 Oct 2026 07:28:00 GMT", common from
    // gateways and CDNs in front of model providers.
    let when = chrono::DateTime::parse_from_rfc2822(value).ok()?;
    (when.with_timezone(&chrono::Utc) - chrono::Utc::now())
        .to_std()
        .ok()
}

fn log_completion_retry(model: &str, retry: usize) {
    // The returned error retains diagnostics; routine retry logs omit upstream bodies.
    log::warn!(
        "Retrying completion request, model: {}, retry: {}/{}",
        model,
        retry,
        MODEL_REQUEST_MAX_RETRIES
    );
}

/// Host matcher that accepts every provider host.
#[derive(Clone, Copy, Debug)]
pub struct AnyHost;

impl PartialEq<&str> for AnyHost {
    fn eq(&self, _other: &&str) -> bool {
        true
    }
}

/// Creates a reqwest client builder with Anda Engine defaults.
pub fn request_client_builder() -> reqwest::ClientBuilder {
    reqwest::Client::builder()
        .use_rustls_tls()
        .https_only(true)
        .retry(
            reqwest::retry::for_host(AnyHost)
                .max_retries_per_request(1)
                .classify_fn(|req_rep| {
                    let is_idempotent = matches!(
                        req_rep.method(),
                        &http::Method::GET
                            | &http::Method::HEAD
                            | &http::Method::OPTIONS
                            | &http::Method::TRACE
                            | &http::Method::PUT
                            | &http::Method::DELETE
                    );

                    if !is_idempotent {
                        return req_rep.success();
                    }

                    if req_rep.error().is_some() {
                        return req_rep.retryable();
                    }

                    match req_rep.status() {
                        Some(status) if is_retryable_status(status) => req_rep.retryable(),
                        _ => req_rep.success(),
                    }
                }),
        )
        // Do not use HTTP/2 PINGs as the liveness detector for model SSE
        // streams. Some provider/CDN edges can keep a long reasoning stream
        // alive while delaying PING ACKs; hyper then closes the connection and
        // reqwest reports "error decoding response body: ... operation timed
        // out" even though the body was still progressing. The per-read body
        // timeout below is the stall detector for completions.
        .http2_keep_alive_interval(COMPLETION_HTTP2_KEEP_ALIVE_INTERVAL)
        .connect_timeout(COMPLETION_CONNECT_TIMEOUT)
        // Read (idle) timeout is the authoritative stall detector for streamed
        // completions: it resets on every chunk, so a long-but-progressing
        // reasoning stream of unknown size is never killed, while a connection
        // that goes silent is failed promptly with clear attribution.
        .read_timeout(COMPLETION_READ_TIMEOUT)
        // Total request timeout, including the streamed body. Heavy reasoning
        // completions can run for many minutes; provider SDKs default to 10
        // minutes.
        .timeout(COMPLETION_REQUEST_TIMEOUT)
        .user_agent(APP_USER_AGENT)
        .default_headers({
            let mut headers = reqwest::header::HeaderMap::new();
            let ct: http::HeaderValue = http::HeaderValue::from_static(CONTENT_TYPE_JSON);
            headers.insert(http::header::CONTENT_TYPE, ct.clone());
            headers.insert(http::header::ACCEPT, ct);
            headers
        })
}

const SSE_DONE_MARKER: &[u8] = b"data: [DONE]";

pub(crate) async fn read_completion_stream<T>(
    response: reqwest::Response,
    model: &str,
) -> Result<(Vec<T>, bool), BoxError>
where
    T: DeserializeOwned,
{
    let request_id = upstream_request_id(response.headers());
    let started = Instant::now();
    let mut body = Vec::new();
    let mut scanned: usize = 0;
    let mut stream = response.bytes_stream();

    while let Some(chunk) = stream.next().await {
        let chunk = chunk.map_err(|err| {
            // Mid-stream failures are reported with enough context to tell a
            // client-side timeout (elapsed near the request deadline) apart
            // from an upstream/gateway abort, and to follow up with the
            // provider via the request id.
            let action = format!(
                "Failed to read streaming completion response (request_id: {}, received: {} bytes, elapsed: {:.1?})",
                request_id.as_deref().unwrap_or("-"),
                body.len(),
                started.elapsed(),
            );
            completion_transport_error(model, &action, err)
        })?;
        if body.len() + chunk.len() > MAX_COMPLETION_RESPONSE_BYTES {
            return Err(completion_response_too_large(
                model,
                request_id.as_deref(),
                body.len() + chunk.len(),
            ));
        }
        body.extend_from_slice(&chunk);
        // Only scan the unscanned tail (with marker-sized overlap), so long
        // streams are not rescanned from the start on every chunk.
        let start = scanned.saturating_sub(SSE_DONE_MARKER.len());
        if body_contains_sse_done(&body, start) {
            return Ok((parse_streaming_json_events(&body, model)?, true));
        }
        scanned = body.len();
    }

    Ok((parse_streaming_json_events(&body, model)?, false))
}

#[cfg(test)]
async fn read_sse_json_events<T: DeserializeOwned>(
    response: reqwest::Response,
    model: &str,
) -> Result<Vec<T>, BoxError> {
    read_completion_stream(response, model)
        .await
        .map(|(events, _)| events)
}

/// Returns true when a `data: [DONE]` line exists at or after `from`.
///
/// The marker must be anchored at the start of an SSE line (buffer start or a
/// preceding `\n`). Generated content inside a JSON string can legitimately
/// contain the marker text, and must not terminate the stream early.
fn body_contains_sse_done(body: &[u8], from: usize) -> bool {
    if from == 0 && body.starts_with(SSE_DONE_MARKER) {
        return true;
    }
    body[from..]
        .windows(SSE_DONE_MARKER.len() + 1)
        .any(|window| window[0] == b'\n' && &window[1..] == SSE_DONE_MARKER)
}

fn parse_streaming_json_events<T>(body: &[u8], model: &str) -> Result<Vec<T>, BoxError>
where
    T: DeserializeOwned,
{
    let body = std::str::from_utf8(body).map_err(|err| {
        format!(
            "Invalid UTF-8 in streaming completion response, model: {}, error: {}",
            model, err
        )
    })?;
    let body = body.strip_prefix('\u{feff}').unwrap_or(body);

    if !looks_like_sse(body) {
        return parse_json_event_payload(body, model);
    }

    let mut data = String::new();
    let mut events = Vec::new();

    for line in body.lines() {
        let line = line.strip_suffix('\r').unwrap_or(line);
        handle_sse_text_line(line, &mut data, &mut events, model)?;
    }
    flush_sse_data(&mut data, &mut events, model)?;

    Ok(events)
}

fn looks_like_sse(body: &str) -> bool {
    body.lines().any(|line| {
        let line = line.strip_prefix('\u{feff}').unwrap_or(line);
        line.starts_with("data:")
            || line.starts_with("event:")
            || line.starts_with("id:")
            || line.starts_with("retry:")
            || line.starts_with(':')
    })
}

fn handle_sse_text_line<T>(
    line: &str,
    data: &mut String,
    events: &mut Vec<T>,
    model: &str,
) -> Result<(), BoxError>
where
    T: DeserializeOwned,
{
    if line.is_empty() {
        return flush_sse_data(data, events, model);
    }
    if line.starts_with(':') {
        return Ok(());
    }

    let Some(value) = line.strip_prefix("data:") else {
        return Ok(());
    };
    let value = value.strip_prefix(' ').unwrap_or(value);
    if !data.is_empty() {
        data.push('\n');
    }
    data.push_str(value);
    Ok(())
}

fn flush_sse_data<T>(data: &mut String, events: &mut Vec<T>, model: &str) -> Result<(), BoxError>
where
    T: DeserializeOwned,
{
    let value = data.trim_end();
    if value.is_empty() || value == "[DONE]" {
        data.clear();
        return Ok(());
    }

    let event = serde_json::from_str::<T>(value).map_err(|err| {
        format!(
            "Invalid streaming completion event, model: {}, error: {}, body: {}",
            model, err, value
        )
    })?;
    events.push(event);
    data.clear();
    Ok(())
}

fn parse_json_event_payload<T>(body: &str, model: &str) -> Result<Vec<T>, BoxError>
where
    T: DeserializeOwned,
{
    let value = body.trim().strip_prefix('\u{feff}').unwrap_or(body.trim());
    if value.is_empty() || value == "[DONE]" {
        return Ok(Vec::new());
    }

    if value.starts_with('[')
        && let Ok(events) = serde_json::from_str::<Vec<T>>(value)
    {
        return Ok(events);
    }

    match serde_json::from_str::<T>(value) {
        Ok(event) => Ok(vec![event]),
        Err(single_err) => match serde_json::from_str::<Vec<T>>(value) {
            Ok(events) => Ok(events),
            Err(array_err) => {
                let mut events = Vec::new();
                let mut saw_line = false;
                for line in value.lines() {
                    let line = line.trim();
                    if line.is_empty() || line == "[DONE]" {
                        continue;
                    }
                    saw_line = true;
                    let event = serde_json::from_str::<T>(line).map_err(|line_err| {
                        format!(
                            "Invalid streaming completion event, model: {}, error: {}, body: {}",
                            model, line_err, line
                        )
                    })?;
                    events.push(event);
                }

                if saw_line {
                    return Ok(events);
                }

                Err(format!(
                    "Invalid streaming completion event, model: {}, error: {}; array error: {}, body: {}",
                    model, single_err, array_err, value
                )
                .into())
            }
        },
    }
}

pub(crate) fn streaming_completion_request(
    request: reqwest::RequestBuilder,
) -> reqwest::RequestBuilder {
    request
        .header(reqwest::header::ACCEPT, "text/event-stream")
        .header(reqwest::header::ACCEPT_ENCODING, "identity")
        // Streaming model completions can run much longer than generic HTTP
        // calls made with the same client. Set the request-level total timeout
        // explicitly so downstream clients with shorter shared timeouts do not
        // abort a progressing SSE body before the model timeout budget.
        .timeout(COMPLETION_REQUEST_TIMEOUT)
}

#[cfg(test)]
mod tests {
    use super::*;
    use anda_core::FunctionDefinition;
    use http::{HeaderMap, HeaderValue, StatusCode};
    use tokio::io::{AsyncReadExt, AsyncWriteExt};

    #[tokio::test]
    async fn prematurely_closed_provider_streams_are_retried() {
        use serde_json::json;
        let chat = json!({"id":"r","model":"test","choices":[{"index":0,"delta":{"role":"assistant","content":"ok"}}]});
        let gemini = json!({"candidates":[{"content":{"role":"model","parts":[{"text":"ok"}]}}]});
        let anthropic = json!({"type":"message_start","message":{"id":"r","type":"message","role":"assistant","model":"test","content":[{"type":"text","text":"ok"}],"usage":{}}});
        let cases = [
            (
                "openai",
                vec![chat],
                vec![json!({"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]})],
            ),
            (
                "gemini",
                vec![gemini],
                vec![json!({"candidates":[{"finishReason":"STOP"}]})],
            ),
            (
                "anthropic",
                vec![anthropic],
                vec![
                    json!({"type":"message_delta","delta":{"stop_reason":"end_turn"}}),
                    json!({"type":"message_stop"}),
                ],
            ),
        ];
        for (family, initial, terminal) in cases {
            let encode = |events: Vec<Json>| {
                events
                    .into_iter()
                    .map(|event| format!("data: {event}\n\n"))
                    .collect::<String>()
                    .into_bytes()
            };
            let incomplete = test_support::MockResponse {
                status: StatusCode::OK,
                headers: test_support::sse_headers(),
                body: encode(initial.clone()),
            };
            let complete = test_support::MockResponse {
                body: encode(initial.into_iter().chain(terminal).collect()),
                ..incomplete.clone()
            };
            let (endpoint, state) =
                test_support::spawn_retry_mock_server(vec![incomplete, complete]).await;
            let config = ModelConfig {
                family: family.into(),
                model: "test".into(),
                api_base: endpoint,
                api_key: "fake".into(),
                stream: true,
                ..Default::default()
            };
            let model = config.model(test_support::no_proxy_client()).unwrap();
            let output = model
                .completion(CompletionRequest {
                    prompt: "hello".into(),
                    ..Default::default()
                })
                .await
                .unwrap();
            assert_eq!(output.content, "ok", "{family}");
            assert!(output.failed_reason.is_none(), "{family}");
            assert_eq!(state.lock().unwrap().1, 2, "{family}");
        }
    }

    #[derive(Clone)]
    struct TestCompleter {
        name: &'static str,
    }

    impl CompletionFeaturesDyn for TestCompleter {
        fn completion(&self, _req: CompletionRequest) -> BoxPinFut<Result<AgentOutput, BoxError>> {
            Box::pin(futures::future::ready(Ok(AgentOutput::default())))
        }

        fn model_name(&self) -> String {
            self.name.to_string()
        }
    }

    fn test_model(name: &'static str) -> Model {
        Model::new(Arc::new(TestCompleter { name }))
    }

    fn http_client() -> reqwest::Client {
        test_support::no_proxy_client()
    }

    async fn spawn_truncated_sse_after_done_server() -> String {
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        tokio::spawn(async move {
            let (mut socket, _) = listener.accept().await.unwrap();
            let mut request = [0; 1024];
            let _ = socket.read(&mut request).await;
            socket
                .write_all(
                    b"HTTP/1.1 200 OK\r\n\
                      Content-Type: text/event-stream\r\n\
                      Content-Length: 4096\r\n\
                      Connection: close\r\n\
                      \r\n\
                      data: {\"a\":1}\n\n\
                      data: [DONE]\n\n",
                )
                .await
                .unwrap();
            let _ = socket.shutdown().await;
        });
        format!("http://{addr}")
    }

    /// Sends an event whose JSON content embeds the literal `data: [DONE]`
    /// text in an early chunk, then a second event and the real terminator
    /// in a later chunk.
    async fn spawn_sse_with_done_marker_in_content_server() -> String {
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        tokio::spawn(async move {
            let (mut socket, _) = listener.accept().await.unwrap();
            let mut request = [0; 1024];
            let _ = socket.read(&mut request).await;
            socket
                .write_all(
                    b"HTTP/1.1 200 OK\r\n\
                      Content-Type: text/event-stream\r\n\
                      Connection: close\r\n\
                      \r\n\
                      data: {\"text\":\"sse ends with data: [DONE]\"}\n\n",
                )
                .await
                .unwrap();
            socket.flush().await.unwrap();
            tokio::time::sleep(Duration::from_millis(50)).await;
            socket
                .write_all(b"data: {\"b\":2}\n\ndata: [DONE]\n\n")
                .await
                .unwrap();
            let _ = socket.shutdown().await;
        });
        format!("http://{addr}")
    }

    async fn spawn_stalling_sse_body_server() -> String {
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        tokio::spawn(async move {
            let (mut socket, _) = listener.accept().await.unwrap();
            let mut request = [0; 1024];
            let _ = socket.read(&mut request).await;
            socket
                .write_all(
                    b"HTTP/1.1 200 OK\r\n\
                      Content-Type: text/event-stream\r\n\
                      Connection: close\r\n\
                      \r\n\
                      data: {\"a\":1}\n\n",
                )
                .await
                .unwrap();
            socket.flush().await.unwrap();
            tokio::time::sleep(Duration::from_millis(500)).await;
            let _ = socket.shutdown().await;
        });
        format!("http://{addr}")
    }

    fn retry_count(state: &test_support::RetryState) -> usize {
        state.lock().unwrap().1
    }

    fn model_config(family: &str, model: &str) -> ModelConfig {
        ModelConfig {
            family: family.to_string(),
            model: model.to_string(),
            api_base: "https://example.com".to_string(),
            api_key: "test-key".to_string(),
            ..Default::default()
        }
    }

    #[test]
    fn model_effort_serializes_config_values() {
        let config: ModelConfig = serde_json::from_value(serde_json::json!({
            "family": "openai",
            "model": "gpt-5",
            "api_base": "http://localhost",
            "api_key": "test-key",
            "effort": "max"
        }))
        .unwrap();

        assert_eq!(config.effort, Some(ModelEffort::Max));
        assert_eq!(
            serde_json::to_value(ModelEffort::Minimal).unwrap(),
            "minimal"
        );
    }

    #[test]
    fn models_default_is_empty() {
        let models = Models::default();

        assert!(models.get_model().is_none());
        assert!(models.get("missing").is_none());
        assert!(models.resolve("missing").is_none());
    }

    #[test]
    fn set_model_sets_primary_without_registering_a_label() {
        let models = Models::default();
        models.set_model(test_model("primary"));

        assert_eq!(
            models
                .get_model()
                .expect("primary model should exist")
                .model_name(),
            "primary"
        );
        assert!(models.get("primary").is_some());
    }

    #[test]
    fn set_promotes_first_inserted_model_to_primary() {
        let models = Models::default();
        models.set("x".to_string(), test_model("X"));

        assert_eq!(
            models.get("x").expect("label x should exist").model_name(),
            "X"
        );
        assert_eq!(
            models
                .get_model()
                .expect("primary model should be initialized")
                .model_name(),
            "X"
        );
    }

    #[test]
    fn fallback_label_has_no_default_routing_semantics() {
        let models = Models::default();
        models.set_model(test_model("primary"));
        models.set("fallback".to_string(), test_model("fallback"));

        assert_eq!(
            models
                .get("fallback")
                .expect("fallback is still a normal label")
                .model_name(),
            "fallback"
        );
        assert_eq!(
            models
                .get_model()
                .expect("primary model should stay the default")
                .model_name(),
            "primary"
        );
        assert_eq!(
            models
                .resolve("unknown")
                .expect("missing label should use default routing")
                .model_name(),
            "primary"
        );
    }

    #[test]
    fn resolve_prefers_exact_label_then_default() {
        let models = Models::default();
        models.set_model(test_model("primary"));
        models.set("flash".to_string(), test_model("flash"));

        assert_eq!(
            models
                .resolve("flash")
                .expect("exact label should win")
                .model_name(),
            "flash"
        );
        assert_eq!(
            models
                .resolve("missing")
                .expect("missing label should use default routing")
                .model_name(),
            "primary"
        );
        assert_eq!(
            models
                .resolve("")
                .expect("empty label should use default routing")
                .model_name(),
            "primary"
        );
    }

    #[test]
    fn model_config_validates_required_fields_and_builds_supported_families() {
        let client = http_client();

        let mut config = model_config("openai", "gpt-5");
        config.disabled = true;
        let Err(err) = config.model(client.clone()) else {
            panic!("disabled model should fail");
        };
        assert!(err.to_string().contains("disabled"));

        for (field, config) in [
            (
                "model name",
                ModelConfig {
                    model: String::new(),
                    ..model_config("openai", "gpt-5")
                },
            ),
            (
                "model family",
                ModelConfig {
                    family: String::new(),
                    ..model_config("openai", "gpt-5")
                },
            ),
            (
                "api_base",
                ModelConfig {
                    api_base: String::new(),
                    ..model_config("openai", "gpt-5")
                },
            ),
            (
                "api_key",
                ModelConfig {
                    api_key: String::new(),
                    ..model_config("openai", "gpt-5")
                },
            ),
        ] {
            let Err(err) = config.model(client.clone()) else {
                panic!("{field} should fail");
            };
            let err = err.to_string();
            assert!(err.contains(field), "{field}: {err}");
        }

        let Err(err) = model_config("unknown", "m").model(client.clone()) else {
            panic!("unsupported family should fail");
        };
        assert!(err.to_string().contains("unsupported model family"));

        let mut gemini = model_config("gemini", "gemini-2.5-pro");
        gemini.context_window = 123;
        gemini.max_output = 45;
        let model = gemini.model(client.clone()).unwrap();
        assert_eq!(model.model_name(), "gemini-2.5-pro");
        assert_eq!(model.labels, vec!["gemini-2.5-pro"]);
        assert_eq!(model.context_window, 123);
        assert_eq!(model.max_output, 45);

        let mut anthropic = model_config("anthropic", "claude-sonnet-4-5");
        anthropic.labels = vec!["pro".to_string(), "primary".to_string()];
        anthropic.bearer_auth = true;
        anthropic.stream = true;
        anthropic.effort = Some(ModelEffort::High);
        let model = anthropic.model(client.clone()).unwrap();
        assert_eq!(model.model_name(), "claude-sonnet-4-5");
        assert_eq!(model.labels, vec!["pro", "primary"]);

        let model = model_config("openai", "gpt-5")
            .model(client.clone())
            .unwrap();
        assert_eq!(model.model_name(), "gpt-5");
        let model = model_config("openai", "deepseek-chat")
            .model(client)
            .unwrap();
        assert_eq!(model.model_name(), "deepseek-chat");
    }

    #[test]
    fn models_registry_clones_names_replaces_labels_and_loads_configs() {
        let models = Models::default();
        models.set_model(test_model("flash-v1").with_labels(vec!["FAST".into()]));
        assert!(models.contains("fast"));
        assert_eq!(
            models.model_names(),
            BTreeSet::from(["flash-v1".to_string()])
        );

        models.set("flash".to_string(), test_model("flash-v2"));
        assert!(models.contains("flash"));
        assert_eq!(models.get("FLASH").unwrap().model_name(), "flash-v2");
        assert_eq!(
            models.model_names(),
            BTreeSet::from(["flash-v1".to_string(), "flash-v2".to_string()])
        );

        models.set("primary".to_string(), test_model("primary-v2"));
        assert_eq!(models.get_model().unwrap().model_name(), "primary-v2");

        let cloned = Models::from_clone(&models);
        assert_eq!(cloned.get("primary").unwrap().model_name(), "primary-v2");
        assert_eq!(
            cloned.resolve("missing").unwrap().model_name(),
            "primary-v2"
        );

        let replacement = Models::default();
        replacement.set_model(test_model("replacement-primary").with_labels(vec!["next".into()]));
        let replaced = Models::default();
        replaced.set("old".to_string(), test_model("old"));
        replaced.replace(&replacement);
        assert!(!replaced.contains("old"));
        assert!(replaced.contains("next"));
        assert_eq!(
            replaced.get_model().unwrap().model_name(),
            "replacement-primary"
        );

        replacement.set("later".to_string(), test_model("later"));
        assert!(!replaced.contains("later"));

        let configs = vec![
            ModelConfig {
                labels: vec!["primary".to_string()],
                ..model_config("openai", "gpt-5")
            },
            ModelConfig {
                disabled: true,
                ..model_config("openai", "disabled")
            },
        ];
        let loaded = Models::from_configs(&configs, http_client());
        assert!(loaded.contains("primary"));
        assert!(!loaded.contains("disabled"));
        assert_eq!(loaded.get_model().unwrap().model_name(), "gpt-5");
    }

    #[tokio::test]
    async fn model_max_output_caps_explicit_request_output_budget() {
        let completer = testing::ScriptedCompleter::new("capped").into_arc();
        let mut model = Model::with_completer(completer.clone());

        // Unknown model limit: the request budget passes through untouched.
        model
            .completion(CompletionRequest::default())
            .await
            .unwrap();

        model.max_output = 32_000;
        model
            .completion(CompletionRequest::default())
            .await
            .unwrap();
        for requested in [8_000, 64_000] {
            model
                .completion(CompletionRequest {
                    max_output_tokens: Some(requested),
                    ..Default::default()
                })
                .await
                .unwrap();
        }

        let sent = completer
            .requests()
            .into_iter()
            .map(|req| req.max_output_tokens)
            .collect::<Vec<_>>();
        // An unset budget is left to the adapter; explicit budgets are capped.
        assert_eq!(sent, vec![None, None, Some(8_000), Some(32_000)]);
    }

    #[tokio::test]
    async fn model_completion_placeholders_and_mock_tool_calls_are_stable() {
        let not_implemented = Model::not_implemented();
        assert_eq!(not_implemented.model_name(), "not_implemented");
        let err = not_implemented
            .completion(CompletionRequest::default())
            .await
            .unwrap_err();
        assert!(err.to_string().contains("not implemented"));

        let mock = Model::mock_implemented().with_labels(vec!["mock".into()]);
        assert_eq!(mock.model_name(), "mock_implemented");
        let output = mock
            .completion(CompletionRequest {
                prompt: "{\"q\":\"anda\"}".to_string(),
                tools: vec![FunctionDefinition {
                    name: "lookup".to_string(),
                    ..Default::default()
                }],
                ..Default::default()
            })
            .await
            .unwrap();
        assert_eq!(output.content, "{\"q\":\"anda\"}");
        assert_eq!(output.tool_calls.len(), 1);
        assert_eq!(output.tool_calls[0].name, "lookup");
        assert_eq!(output.tool_calls[0].args["q"], "anda");

        let output = mock
            .completion(CompletionRequest {
                prompt: String::new(),
                tools: vec![FunctionDefinition {
                    name: "lookup".to_string(),
                    ..Default::default()
                }],
                ..Default::default()
            })
            .await
            .unwrap();
        assert!(output.tool_calls.is_empty());
    }

    #[test]
    fn streaming_json_event_parser_accepts_bom_sse_ndjson_and_arrays() {
        let events = parse_streaming_json_events::<serde_json::Value>(
            b"\xef\xbb\xbfdata: {\"a\":1}\n\ndata: [DONE]\n\n",
            "test-model",
        )
        .unwrap();
        assert_eq!(events, vec![serde_json::json!({"a": 1})]);

        let events = parse_streaming_json_events::<serde_json::Value>(
            b"{\"a\":1}\n{\"b\":2}\n[DONE]\n",
            "test-model",
        )
        .unwrap();
        assert_eq!(
            events,
            vec![serde_json::json!({"a": 1}), serde_json::json!({"b": 2})]
        );

        let events =
            parse_streaming_json_events::<serde_json::Value>(br#"[{"a":1},{"b":2}]"#, "test-model")
                .unwrap();
        assert_eq!(
            events,
            vec![serde_json::json!({"a": 1}), serde_json::json!({"b": 2})]
        );
    }

    #[tokio::test]
    async fn streaming_reader_ignores_mislabelled_content_encoding() {
        let mut headers = HeaderMap::new();
        headers.insert(
            http::header::CONTENT_TYPE,
            HeaderValue::from_static("text/event-stream"),
        );
        let (endpoint, _) =
            test_support::spawn_retry_mock_server(vec![test_support::MockResponse {
                status: StatusCode::OK,
                headers,
                body: b"data: {\"a\":1}\n\ndata: [DONE]\n\n".to_vec(),
            }])
            .await;
        let client = request_client_builder()
            .https_only(false)
            .no_proxy()
            .build()
            .unwrap();
        let response = client.get(endpoint).send().await.unwrap();

        let events = read_sse_json_events::<serde_json::Value>(response, "test-model")
            .await
            .unwrap();

        assert_eq!(events, vec![serde_json::json!({"a": 1})]);
    }

    #[tokio::test]
    async fn streaming_reader_returns_after_done_before_late_body_error() {
        let endpoint = spawn_truncated_sse_after_done_server().await;
        let client = request_client_builder()
            .https_only(false)
            .no_proxy()
            .build()
            .unwrap();
        let response = client.get(endpoint).send().await.unwrap();

        let events = read_sse_json_events::<serde_json::Value>(response, "test-model")
            .await
            .unwrap();

        assert_eq!(events, vec![serde_json::json!({"a": 1})]);
    }

    #[test]
    fn sse_done_detection_is_line_anchored() {
        assert!(body_contains_sse_done(b"data: [DONE]\n\n", 0));
        assert!(body_contains_sse_done(
            b"data: {\"a\":1}\n\ndata: [DONE]\n\n",
            0
        ));
        // The marker text inside generated JSON content must not terminate
        // the stream.
        assert!(!body_contains_sse_done(
            b"data: {\"text\":\"sse ends with data: [DONE]\"}\n\n",
            0
        ));
    }

    #[tokio::test]
    async fn streaming_reader_is_not_truncated_by_done_marker_in_content() {
        let endpoint = spawn_sse_with_done_marker_in_content_server().await;
        let client = request_client_builder()
            .https_only(false)
            .no_proxy()
            .build()
            .unwrap();
        let response = client.get(endpoint).send().await.unwrap();

        let events = read_sse_json_events::<serde_json::Value>(response, "test-model")
            .await
            .unwrap();

        assert_eq!(
            events,
            vec![
                serde_json::json!({"text": "sse ends with data: [DONE]"}),
                serde_json::json!({"b": 2})
            ]
        );
    }

    #[test]
    fn completion_transport_timeouts_are_streaming_safe() {
        // The observed failure hit at ~118s with body bytes already received;
        // HTTP/2 PING ACK timeouts must not be able to abort such a stream
        // before the explicit idle body timeout can make that decision.
        assert_eq!(COMPLETION_HTTP2_KEEP_ALIVE_INTERVAL, None);
        assert!(COMPLETION_READ_TIMEOUT > Duration::from_secs(118));
        assert!(COMPLETION_READ_TIMEOUT < COMPLETION_REQUEST_TIMEOUT);
        assert_eq!(COMPLETION_REQUEST_TIMEOUT, Duration::from_secs(600));
    }

    #[test]
    fn streaming_completion_request_overrides_short_client_total_timeout() {
        let client = reqwest::Client::builder()
            .no_proxy()
            .timeout(Duration::from_millis(100))
            .build()
            .unwrap();
        let request = streaming_completion_request(client.get("https://example.com"))
            .build()
            .unwrap();

        assert_eq!(request.timeout(), Some(&COMPLETION_REQUEST_TIMEOUT));
    }

    #[tokio::test]
    async fn streaming_reader_body_idle_timeout_is_retryable() {
        let endpoint = spawn_stalling_sse_body_server().await;
        let client = request_client_builder()
            .https_only(false)
            .no_proxy()
            .read_timeout(Duration::from_millis(100))
            .timeout(Duration::from_secs(5))
            .build()
            .unwrap();
        let response = client.get(endpoint).send().await.unwrap();

        let err = tokio::time::timeout(
            Duration::from_secs(2),
            read_sse_json_events::<serde_json::Value>(response, "test-model"),
        )
        .await
        .expect("body read timeout should fire")
        .unwrap_err();

        let message = err.to_string();
        assert!(
            message.contains("Failed to read streaming completion response"),
            "{message}"
        );
        assert!(message.contains("received:"), "{message}");
        assert!(message.contains("operation timed out"), "{message}");
        assert!(is_retryable_box_error(&err));
    }

    #[test]
    fn retry_after_parses_seconds_and_http_date() {
        let mut headers = HeaderMap::new();
        headers.insert(http::header::RETRY_AFTER, HeaderValue::from_static("42"));
        assert_eq!(
            retry_after_duration(&headers),
            Some(Duration::from_secs(42))
        );

        let when = chrono::Utc::now() + chrono::Duration::seconds(90);
        headers.insert(
            http::header::RETRY_AFTER,
            HeaderValue::from_str(&when.to_rfc2822()).unwrap(),
        );
        let parsed = retry_after_duration(&headers).expect("http-date should parse");
        assert!(parsed <= Duration::from_secs(90));
        assert!(parsed >= Duration::from_secs(80));

        // A date in the past yields no delay hint.
        let when = chrono::Utc::now() - chrono::Duration::seconds(90);
        headers.insert(
            http::header::RETRY_AFTER,
            HeaderValue::from_str(&when.to_rfc2822()).unwrap(),
        );
        assert_eq!(retry_after_duration(&headers), None);

        headers.insert(
            http::header::RETRY_AFTER,
            HeaderValue::from_static("not-a-date"),
        );
        assert_eq!(retry_after_duration(&headers), None);
    }

    #[tokio::test]
    async fn custom_client_streaming_decode_errors_are_retryable() {
        let mut headers = HeaderMap::new();
        headers.insert(
            http::header::CONTENT_TYPE,
            HeaderValue::from_static("text/event-stream"),
        );
        headers.insert(
            http::header::CONTENT_ENCODING,
            HeaderValue::from_static("gzip"),
        );
        let (endpoint, _) =
            test_support::spawn_retry_mock_server(vec![test_support::MockResponse {
                status: StatusCode::OK,
                headers,
                body: b"data: {\"a\":1}\n\ndata: [DONE]\n\n".to_vec(),
            }])
            .await;
        let client = reqwest::Client::builder().no_proxy().build().unwrap();
        let response = streaming_completion_request(client.get(endpoint))
            .send()
            .await
            .unwrap();

        let err = read_sse_json_events::<serde_json::Value>(response, "test-model")
            .await
            .unwrap_err();

        let message = err.to_string();
        assert!(message.contains("error decoding response body"));
        // The reported error carries stream context and the decode source
        // chain, which reqwest's `Display` alone no longer exposes.
        assert!(message.contains("received: 0 bytes"), "{message}");
        assert!(message.contains("request_id: -"), "{message}");
        assert!(
            message.contains("error decoding response body: "),
            "{message}"
        );
        assert!(is_retryable_box_error(&err));
    }

    #[test]
    fn error_chain_formatting_appends_unique_sources() {
        let root = std::io::Error::new(std::io::ErrorKind::TimedOut, "operation timed out");
        let outer =
            ModelError::new("error decoding response body".to_string()).with_source(Box::new(root));
        assert_eq!(
            format_error_chain(&outer),
            "error decoding response body: operation timed out"
        );

        // A source already repeated in the message is not appended twice.
        let root = std::io::Error::new(std::io::ErrorKind::TimedOut, "operation timed out");
        let outer = ModelError::new("request failed: operation timed out".to_string())
            .with_source(Box::new(root));
        assert_eq!(
            format_error_chain(&outer),
            "request failed: operation timed out"
        );
    }

    #[tokio::test]
    async fn completion_error_bodies_are_truncated_for_diagnostics() {
        assert_eq!(error_body_excerpt(b"short body"), "short body");
        let excerpt = error_body_excerpt(&vec![b'x'; MAX_ERROR_BODY_BYTES * 4]);
        assert!(excerpt.ends_with("… [truncated]"));
        assert!(excerpt.len() < MAX_ERROR_BODY_BYTES + 32);

        let (endpoint, _) =
            test_support::spawn_retry_mock_server(vec![test_support::MockResponse {
                status: StatusCode::BAD_REQUEST,
                headers: HeaderMap::new(),
                body: vec![b'e'; 1024 * 1024],
            }])
            .await;
        let client = http_client();
        let err = execute_completion_request_with_retry(
            "error-body-test",
            || client.post(&endpoint),
            |response| async { read_completion_response_bytes(response, "error-body-test").await },
        )
        .await
        .unwrap_err();
        let message = err.to_string();
        assert!(message.contains("status: 400"));
        assert!(message.ends_with("… [truncated]"));
        assert!(message.len() < MAX_ERROR_BODY_BYTES + 256);
    }

    #[test]
    fn upstream_request_id_checks_known_headers() {
        let mut headers = HeaderMap::new();
        assert_eq!(upstream_request_id(&headers), None);

        headers.insert("cf-ray", HeaderValue::from_static("ray-123"));
        assert_eq!(upstream_request_id(&headers), Some("ray-123".to_string()));

        headers.insert("x-request-id", HeaderValue::from_static("req-456"));
        assert_eq!(upstream_request_id(&headers), Some("req-456".to_string()));
    }

    #[tokio::test]
    async fn completion_request_retries_transient_errors_and_exposes_retry_signal() {
        let mut headers = HeaderMap::new();
        headers.insert(http::header::RETRY_AFTER, HeaderValue::from_static("0"));
        let (endpoint, state) = test_support::spawn_retry_mock_server(vec![
            test_support::MockResponse {
                status: StatusCode::TOO_MANY_REQUESTS,
                headers,
                body: b"rate limited".to_vec(),
            },
            test_support::MockResponse {
                status: StatusCode::OK,
                headers: HeaderMap::new(),
                body: b"ok".to_vec(),
            },
        ])
        .await;
        let client = http_client();

        let body = execute_completion_request_with_retry(
            "retry-test",
            || client.post(&endpoint),
            |response| async { read_completion_response_bytes(response, "retry-test").await },
        )
        .await
        .unwrap();

        assert_eq!(&body[..], b"ok");
        assert_eq!(retry_count(&state), 2);

        let mut retry_now = HeaderMap::new();
        retry_now.insert(http::header::RETRY_AFTER, HeaderValue::from_static("0"));
        let mut final_headers = HeaderMap::new();
        final_headers.insert(http::header::RETRY_AFTER, HeaderValue::from_static("45"));
        let (endpoint, state) = test_support::spawn_retry_mock_server(vec![
            test_support::MockResponse {
                status: StatusCode::TOO_MANY_REQUESTS,
                headers: retry_now.clone(),
                body: b"first limit".to_vec(),
            },
            test_support::MockResponse {
                status: StatusCode::TOO_MANY_REQUESTS,
                headers: retry_now.clone(),
                body: b"still limited".to_vec(),
            },
            test_support::MockResponse {
                status: StatusCode::TOO_MANY_REQUESTS,
                headers: retry_now,
                body: b"limited again".to_vec(),
            },
            test_support::MockResponse {
                status: StatusCode::TOO_MANY_REQUESTS,
                headers: final_headers,
                body: b"finally limited".to_vec(),
            },
        ])
        .await;
        let err = execute_completion_request_with_retry(
            "retry-test",
            || client.post(&endpoint),
            |response| async { read_completion_response_bytes(response, "retry-test").await },
        )
        .await
        .unwrap_err();
        let err_ref = err.as_ref() as &(dyn Error + 'static);

        assert_eq!(retry_count(&state), MODEL_REQUEST_MAX_RETRIES + 1);
        assert!(is_retryable_box_error(&err));
        assert_eq!(
            model_error_status(err_ref),
            Some(StatusCode::TOO_MANY_REQUESTS)
        );
        assert_eq!(
            model_error_retry_after(err_ref),
            Some(Duration::from_secs(45))
        );

        let (endpoint, state) =
            test_support::spawn_retry_mock_server(vec![test_support::MockResponse {
                status: StatusCode::BAD_REQUEST,
                headers: HeaderMap::new(),
                body: b"bad request".to_vec(),
            }])
            .await;
        let err = execute_completion_request_with_retry(
            "retry-test",
            || client.post(&endpoint),
            |response| async { read_completion_response_bytes(response, "retry-test").await },
        )
        .await
        .unwrap_err();

        assert_eq!(retry_count(&state), 1);
        assert!(!is_retryable_box_error(&err));
    }
}