zesdex/src/service/provider.rs

//! Blocking HTTP client for OpenAI/Anthropic-compatible chat completion APIs,
//! supporting both non-streaming and SSE-streaming requests with automatic retry.

use std::time::Duration;
use anyhow::Result;

use crate::app::runtime::stream::{SseParser, StreamEvent};
use crate::app::runtime::stream::turn::StreamedTurn;
use crate::dto::chat::message::ChatMessage;
use crate::dto::provider::request::{ChatRequest, StreamOptions, ToolDef};

pub(crate) const DEFAULT_BASE_URL: &str = "https://opencode.ai/zen/v1";
const DEFAULT_MODEL: &str = "deepseek-v4-flash-free";
pub const DEFAULT_API_KEY: &str = "";
const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
const REQUEST_TIMEOUT: Duration = Duration::from_mins(1);

/// Blocking HTTP client for a single LLM provider endpoint.
///
/// Holds the reqwest client, credentials, and model/base URL selection used
/// by both the non-streaming and streaming chat completion calls.
pub struct LlmClient {
    pub client: reqwest::blocking::Client,
    pub api_key: String,
    pub base_url: String,
    pub model: String,
}

impl LlmClient {
    /// Construct a client, falling back to built-in defaults for empty inputs.
    ///
    /// Flow: empty `api_key/model` → substitute defaults → build reqwest client
    /// with connect/request timeouts → if TLS config fails, retry with just
    /// request timeout (no connect timeout) → normalize `base_url`.
    ///
    /// Why: empty strings are treated as "unset" rather than errors so callers
    /// can pass through unconfigured settings without special-casing them.
    /// Timeouts are always enforced — the pure-default-client fallback is only
    /// used as a last resort when even the no-connect-timeout build fails.
    pub fn new(mut api_key: String, model: String, base_url: Option<String>) -> Self {
        if api_key.is_empty() {
            api_key = DEFAULT_API_KEY.to_string();
        }
        let model = if model.is_empty() {
            DEFAULT_MODEL.to_string()
        } else {
            model
        };
        let client = match reqwest::blocking::Client::builder()
            .timeout(REQUEST_TIMEOUT)
            .connect_timeout(CONNECT_TIMEOUT)
            .build()
        {
            Ok(c) => c,
            Err(e) => {
                tracing::warn!(
                    "failed to build reqwest client with connect timeout: {}. \
                     retrying without connect timeout",
                    e,
                );
                match reqwest::blocking::Client::builder()
                    .timeout(REQUEST_TIMEOUT)
                    .build()
                {
                    Ok(c) => c,
                    Err(e2) => {
                        tracing::warn!(
                            "also failed: {}. using default client (no configured timeouts)",
                            e2,
                        );
                        reqwest::blocking::Client::new()
                    }
                }
            }
        };
        LlmClient {
            client,
            api_key,
            base_url: base_url.filter(|s| !s.is_empty()).unwrap_or_else(|| DEFAULT_BASE_URL.to_string()),
            model,
        }
    }

    /// Send a non-streaming chat completion request and return the assistant's reply.
    ///
    /// Flow: build request → POST with retry loop (up to 10 attempts, 2s backoff)
    /// → parse JSON response → extract first choice's message and token usage.
    ///
    /// Why: retries transient failures but aborts immediately on 401/403, since
    /// those indicate a bad API key that retrying won't fix.
    ///
    /// Return: `Err` if all retries are exhausted, an auth error occurs, or the
    /// response has no choices.
    pub fn chat_with_tools_non_streaming(
        &self,
        messages: &[ChatMessage],
        tools: Option<Vec<ToolDef>>,
    ) -> Result<(ChatMessage, Option<(u64, u64)>)> {
        let req = ChatRequest {
            model: self.model.clone(),
            messages: messages.to_vec(),
            max_tokens: Some(4096),
            temperature: Some(0.7),
            tools,
            stream: Some(false),
            top_p: None,
            stop: None,
            stream_options: None,
        };

        let url = format!("{}/chat/completions", self.base_url);
        let max_retries = 10;
        let mut attempt = 0;

        loop {
            attempt += 1;

            let mut http_req = self.client
                .post(&url)
                .header("Content-Type", "application/json");

            if !self.api_key.is_empty() {
                http_req = http_req.header("Authorization", format!("Bearer {}", self.api_key));
            }

            let result = (|| -> Result<(ChatMessage, Option<(u64, u64)>)> {
                let resp = http_req.json(&req).send().map_err(|e| {
                    if e.is_timeout() {
                        anyhow::anyhow!("API request timed out after {REQUEST_TIMEOUT:?}. Check your network or try again.")
                    } else if e.is_connect() {
                        anyhow::anyhow!("Could not connect to {}. Is the URL correct and is the service reachable?", self.base_url)
                    } else {
                        anyhow::anyhow!("API request failed: {e}")
                    }
                })?;

                if !resp.status().is_success() {
                    let status = resp.status();
                    let body = resp.text().unwrap_or_default();
                    anyhow::bail!("API error {} from {}: {}", status, self.base_url, body);
                }

                let data: crate::dto::provider::response::ChatResponse = resp.json()?;
                let usage = data.usage.map(|u| {
                    (u64::from(u.prompt_tokens.unwrap_or(0)), u64::from(u.completion_tokens.unwrap_or(0)))
                });
                let message = data
                    .choices
                    .into_iter()
                    .next()
                    .map(|c| c.message)
                    .ok_or_else(|| anyhow::anyhow!("API response had no choices"))?;
                Ok((message, usage))
            })();

            match result {
                Ok((msg, usage)) => return Ok((msg, usage)),
                Err(e) => {
                    let err_str = e.to_string();
                    let err_lower = err_str.to_lowercase();
                    let is_auth_error = err_str.contains("401") || err_str.contains("403")
                        || err_lower.contains("unauthorized")
                        || err_lower.contains("forbidden")
                        || err_lower.contains("authentication failed");
                    if attempt >= max_retries || is_auth_error {
                        return Err(e);
                    }
                    tracing::warn!("Warning: {}. Retrying {}/{}...", e, attempt, max_retries);
                    std::thread::sleep(Duration::from_secs(2));
                }
            }
        }
    }

    /// Streaming variant of `chat_with_tools`. Feeds SSE chunks into an `SseParser` /
    /// `StreamedTurn` and invokes `on_event` for every parsed `StreamEvent` as it arrives,
    /// so the caller can push incremental UI updates in real time. Returns the fully
    /// assembled assistant message plus token usage (prompt, completion) if the server
    /// reported it. Retries the whole request only if no event has been observed yet
    /// (once tokens start arriving, a partial turn cannot be safely replayed).
    pub fn chat_with_tools_streaming(
        &self,
        messages: &[ChatMessage],
        tools: Option<Vec<ToolDef>>,
        temperature: Option<f32>,
        max_tokens: Option<u32>,
        mut on_event: impl FnMut(&StreamEvent) -> bool,
    ) -> Result<(ChatMessage, Option<(u64, u64)>)> {
        let req = ChatRequest {
            model: self.model.clone(),
            messages: messages.to_vec(),
            max_tokens: Some(max_tokens.unwrap_or(4096)),
            temperature: Some(temperature.unwrap_or(0.7)),
            tools,
            stream: Some(true),
            top_p: None,
            stop: None,
            stream_options: Some(StreamOptions { include_usage: true }),
        };

        let url = format!("{}/chat/completions", self.base_url);
        // Fewer retries on streaming because `run_agent_turn` has a
        // non-streaming fallback that also retries. Combined total is
        // capped implicitly by the per-turn timeout and step limits.
        let max_retries = 3;
        let mut attempt = 0;
        let mut started = false;

        loop {
            attempt += 1;
            let mut wrapped = |event: &StreamEvent| -> bool {
                started = true;
                on_event(event)
            };
            match self.try_stream_once(&req, &url, &mut wrapped) {
                Ok(result) => return Ok(result),
                Err(e) => {
                    let err_str = e.to_string();
                    let err_lower = err_str.to_lowercase();
                    let is_auth_error = err_str.contains("401") || err_str.contains("403")
                        || err_lower.contains("unauthorized")
                        || err_lower.contains("forbidden")
                        || err_lower.contains("authentication failed");
                    if started || attempt >= max_retries || is_auth_error {
                        return Err(e);
                    }
                    tracing::warn!("Warning: {}. Retrying {}/{}...", e, attempt, max_retries);
                    std::thread::sleep(Duration::from_secs(2));
                }
            }
        }
    }

    /// Perform one streaming chat completion request, parsing SSE events until completion.
    ///
    /// Flow: POST → read body in chunks → advance past valid UTF-8 boundary →
    /// feed into `SseParser` → dispatch each `StreamEvent` to `on_event` and
    /// accumulate in `StreamedTurn` → return assembled assistant message on `Done`.
    ///
    /// Why: chunk-by-chunk UTF-8-aware reads avoid splitting multi-byte sequences;
    /// returns `aborted` error if `on_event` returns false so the caller can cancel.
    ///
    /// Return: assembled message + optional usage on success, `Err` on read
    /// failure, non-2xx status, or callback-initiated abort.
    fn try_stream_once(
        &self,
        req: &ChatRequest,
        url: &str,
        on_event: &mut dyn FnMut(&StreamEvent) -> bool,
    ) -> Result<(ChatMessage, Option<(u64, u64)>)> {
        use std::io::Read;

        let mut http_req = self.client
            .post(url)
            .header("Content-Type", "application/json");
        if !self.api_key.is_empty() {
            http_req = http_req.header("Authorization", format!("Bearer {}", self.api_key));
        }

        let resp = http_req.json(req).send().map_err(|e| {
            if e.is_timeout() {
                anyhow::anyhow!("API request timed out after {REQUEST_TIMEOUT:?}. Check your network or try again.")
            } else if e.is_connect() {
                anyhow::anyhow!("Could not connect to {}. Is the URL correct and is the service reachable?", self.base_url)
            } else {
                anyhow::anyhow!("API request failed: {e}")
            }
        })?;

        if !resp.status().is_success() {
            let status = resp.status();
            let body = resp.text().unwrap_or_default();
            anyhow::bail!("API error {} from {}: {}", status, self.base_url, body);
        }

        let mut turn = StreamedTurn::new();
        let mut usage: Option<(u64, u64)> = None;
        let mut parser = SseParser::new();
        let mut reader = resp;
        let mut byte_buf: Vec<u8> = Vec::new();
        let mut chunk_buf = [0u8; 4096];

        loop {
            let n = reader.read(&mut chunk_buf)
                .map_err(|e| anyhow::anyhow!("stream read error: {e}"))?;
            if n == 0 {
                break;
            }
            byte_buf.extend_from_slice(&chunk_buf[..n]);
            let valid_len = match std::str::from_utf8(&byte_buf) {
                Ok(s) => s.len(),
                Err(e) => e.valid_up_to(),
            };
            if valid_len == 0 {
                continue;
            }
            let text = String::from_utf8_lossy(&byte_buf[..valid_len]).into_owned();
            byte_buf.drain(..valid_len);

            for event in parser.feed(&text) {
                if !on_event(&event) {
                    anyhow::bail!("aborted");
                }
                match &event {
                    StreamEvent::Usage { prompt_tokens, completion_tokens, .. } => {
                        usage = Some((*prompt_tokens, *completion_tokens));
                    }
                    StreamEvent::Error(msg) => {
                        anyhow::bail!("stream error: {msg}");
                    }
                    StreamEvent::Done => {
                        turn.apply_event(&event);
                        turn.done_received = true;
                        return Ok((turn.build_assistant_message(), usage));
                    }
                    _ => turn.apply_event(&event),
                }
            }
        }

        // The connection closed without an explicit `[DONE]` event. Some
        // providers legitimately omit it, so EOF alone isn't an error —
        // but if it leaves a tool call's arguments as unparsable JSON, the
        // response was truncated mid-generation, not finished. Report that
        // honestly instead of silently double-stringifying the fragment
        // into a tool call that will misbehave (e.g. a `write` call with a
        // half-written file body).
        if let Some((name, err)) = turn.incomplete_tool_call() {
            anyhow::bail!(
                "stream ended before tool call '{name}' arguments were complete: {err}"
            );
        }

        turn.is_complete = true;
        Ok((turn.build_assistant_message(), usage))
    }
}