procyon 0.1.0

//! Provider dispatch.
//!
//! An enum rather than a trait object: the two methods borrow a channel and return a future, and
//! the set of providers is closed and small, so `async_trait` boxing buys nothing here.

use std::sync::LazyLock;
use std::time::Duration;

use color_eyre::{eyre::bail, Result};
use reqwest::{Client, RequestBuilder, StatusCode};
use tokio::sync::mpsc;

use crate::agent::{ContentPart, Message, StreamOutcome, ToolDefinition};
use crate::anthropic::ClaudeClient;
use crate::config::AppConfig;
use crate::openai::OpenAiClient;

// One connection pool for the process. A `reqwest::Client` owns its pool and resolver, so building
// one per client — and a client is built per sub-agent — threw away every warm connection and
// opened a fresh TLS handshake for each persona in a party round.
static HTTP: LazyLock<Client> = LazyLock::new(Client::new);

/// Returns a handle on the shared HTTP client. Cloning shares the underlying pool.
pub fn http_client() -> Client {
    HTTP.clone()
}

/// Attempts one request makes before its failure is reported.
const MAX_ATTEMPTS: u32 = 3;

/// Ceiling on a server-supplied `Retry-After`, so a provider asking for a ten minute wait does not
/// hold the whole session.
const MAX_RETRY_AFTER: Duration = Duration::from_secs(20);

/// Whether a status is worth trying again.
///
/// Rate limits and 5xx are the provider saying "not now"; every 4xx below 429 is the request
/// itself being wrong, and repeating it only wastes the quota. 408 is a gateway giving up on a
/// request the model may never have seen.
fn is_transient(status: StatusCode) -> bool {
    status == StatusCode::TOO_MANY_REQUESTS
        || status == StatusCode::REQUEST_TIMEOUT
        || status.is_server_error()
}

/// Whether a transport failure is worth trying again. A refused connection or a dropped socket may
/// be the network; a body that could not be built is us.
fn is_transient_transport(error: &reqwest::Error) -> bool {
    error.is_timeout() || error.is_connect() || error.is_request()
}

fn backoff(attempt: u32) -> Duration {
    // 500ms, 1s. Short: a turn is interactive, and the point is to survive a blip rather than to
    // wait out a sustained outage.
    Duration::from_millis(500 * 2u64.pow(attempt.saturating_sub(1)))
}

/// Sends a request, retrying the failures that are the provider's rather than ours.
///
/// Every non-2xx used to end the turn outright, which meant a single `overloaded_error` threw away
/// the whole turn — including every tool result already produced and paid for. `build` is a
/// factory rather than a `RequestBuilder` because a retry needs a fresh body.
pub async fn send_with_retry<F>(build: F) -> Result<reqwest::Response>
where
    F: Fn() -> RequestBuilder,
{
    let mut attempt = 0;

    loop {
        attempt += 1;
        let last = attempt >= MAX_ATTEMPTS;

        match build().send().await {
            Ok(response) if response.status().is_success() => return Ok(response),
            Ok(response) => {
                let status = response.status();
                if !is_transient(status) || last {
                    // Read after the retry decision: consuming the body moves the response, and a
                    // retryable status does not need its text.
                    let body = response.text().await.unwrap_or_default();
                    bail!("API error {}: {}", status, body);
                }
                let wait = retry_after(&response).unwrap_or_else(|| backoff(attempt));
                crate::diag::warn(format!(
                    "provider returned {}, retrying in {:?} (attempt {} of {})",
                    status, wait, attempt, MAX_ATTEMPTS
                ));
                tokio::time::sleep(wait).await;
            }
            Err(e) => {
                if !is_transient_transport(&e) || last {
                    return Err(e.into());
                }
                let wait = backoff(attempt);
                crate::diag::warn(format!(
                    "request failed ({}), retrying in {:?} (attempt {} of {})",
                    e, wait, attempt, MAX_ATTEMPTS
                ));
                tokio::time::sleep(wait).await;
            }
        }
    }
}

/// Honours the provider's own pacing when it states one. Only the delta-seconds form is read; the
/// HTTP-date form is rare here and a wrong parse would be worse than a backoff.
fn retry_after(response: &reqwest::Response) -> Option<Duration> {
    let seconds: u64 = response
        .headers()
        .get(reqwest::header::RETRY_AFTER)?
        .to_str()
        .ok()?
        .trim()
        .parse()
        .ok()?;

    Some(Duration::from_secs(seconds).min(MAX_RETRY_AFTER))
}

pub enum LlmClient {
    Anthropic(ClaudeClient),
    OpenAiCompatible(OpenAiClient),
}

impl LlmClient {
    /// Builds the client the config selects, resolving the credential for that provider.
    pub fn from_config(config: &AppConfig) -> Result<Self> {
        let api_key = config.get_api_key()?;

        // Dispatch is on the protocol: every provider but Anthropic speaks the OpenAI dialect,
        // differing only in the endpoint the config already resolved.
        Ok(if config.provider.is_anthropic() {
            Self::Anthropic(
                ClaudeClient::new(api_key)
                    .with_model(config.default_model.clone())
                    .with_max_tokens(config.max_tokens),
            )
        } else {
            Self::OpenAiCompatible(OpenAiClient::new(
                api_key,
                config.resolve_base_url()?,
                config.default_model.clone(),
                config.max_tokens,
            ))
        })
    }

    pub async fn send_message(
        &self,
        history: &[Message],
        tools: Option<&[ToolDefinition]>,
        system: Option<&str>,
    ) -> Result<Vec<ContentPart>> {
        match self {
            Self::Anthropic(client) => client.send_message(history, tools, system).await,
            Self::OpenAiCompatible(client) => client.send_message(history, tools, system).await,
        }
    }

    pub async fn send_message_streaming(
        &self,
        history: &[Message],
        tools: Option<&[ToolDefinition]>,
        system: Option<&str>,
        update_tx: &mpsc::UnboundedSender<String>,
    ) -> Result<StreamOutcome> {
        match self {
            Self::Anthropic(client) => {
                client
                    .send_message_streaming(history, tools, system, update_tx)
                    .await
            }
            Self::OpenAiCompatible(client) => {
                client
                    .send_message_streaming(history, tools, system, update_tx)
                    .await
            }
        }
    }
}