jev_sdk/

client.rs

//! The TypeSafe API client.

use reqwest::header::{self, HeaderMap, HeaderValue};
use serde::Serialize;
use serde_json::{Map, Value};
use std::fmt;
use std::net::IpAddr;
use std::time::Duration;

use crate::error::{ApiError, Error};
use crate::models::ListModelsResponse;
use crate::question::Question;
use crate::response::SystemOneResponse;
use crate::retry::RetryPolicy;

/// Default API root.
pub const DEFAULT_BASE_URL: &str = "https://api.typesafe.ai";
/// Default model used when none is configured.
pub const DEFAULT_MODEL: &str = "jev-latest";
/// Default request timeout.
pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(60);

/// Environment variable for the API key.
pub const ENV_API_KEY: &str = "TYPESAFE_API_KEY";
/// Environment variable for the default model.
pub const ENV_DEFAULT_MODEL: &str = "TYPESAFE_DEFAULT_MODEL";
/// Environment variable for the API root.
pub const ENV_BASE_URL: &str = "TYPESAFE_BASE_URL";

const USER_AGENT: &str = concat!("jev-sdk/", env!("CARGO_PKG_VERSION"));

/// Upper bound for a successful response body. Real System One responses are
/// a few KiB; this is generous headroom. Guards against memory exhaustion
/// from a hostile or buggy endpoint (only reachable via a custom base URL).
const MAX_SUCCESS_BODY: u64 = 16 * 1024 * 1024; // 16 MiB

/// Upper bound for how much of an *error* response body is read. Anything
/// past this is discarded; the status code alone drives error handling.
const MAX_ERROR_BODY: usize = 64 * 1024; // 64 KiB

/// An asynchronous client for the TypeSafe AI System One API.
///
/// ```
/// use jev_sdk::{Noul, TypeSafeClient};
///
/// # async fn example() -> Result<(), jev_sdk::Error> {
/// let client = TypeSafeClient::new("apikey_...")?;
/// let response = client
///     .system_one(
///         "I was charged twice. Please help.",
///         [("billing", Noul::new("Is this about billing?"))],
///     )
///     .await?;
/// println!("{:?}", response.noul("billing").unwrap().noul);
/// # Ok(())
/// # }
/// ```
#[derive(Clone)]
pub struct TypeSafeClient {
    http: reqwest::Client,
    /// `Bearer <key>`. Marked sensitive *and* excluded from `Debug` below;
    /// never log or print this value.
    auth: HeaderValue,
    base_url: String,
    model: String,
    retry: RetryPolicy,
}

// Manual impl so the API key can never leak through formatting, even if the
// `HeaderValue` sensitive flag regresses.
impl fmt::Debug for TypeSafeClient {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("TypeSafeClient")
            .field("auth", &"[redacted]")
            .field("base_url", &self.base_url)
            .field("model", &self.model)
            .field("retry", &self.retry)
            .finish_non_exhaustive()
    }
}

impl TypeSafeClient {
    /// Create a client with the given API key and default settings.
    pub fn new(api_key: impl Into<String>) -> Result<Self, Error> {
        Self::builder().api_key(api_key).build()
    }

    /// Create a client from environment variables.
    ///
    /// * [`ENV_API_KEY`] (`TYPESAFE_API_KEY`) — required.
    /// * [`ENV_DEFAULT_MODEL`] (`TYPESAFE_DEFAULT_MODEL`) — optional.
    /// * [`ENV_BASE_URL`] (`TYPESAFE_BASE_URL`) — optional.
    pub fn from_env() -> Result<Self, Error> {
        Self::builder().build()
    }

    /// Start building a client with custom settings.
    pub fn builder() -> ClientBuilder {
        ClientBuilder::default()
    }

    /// The model this client uses by default.
    pub fn model(&self) -> &str {
        &self.model
    }

    /// Answer named questions about text or structured `state`.
    ///
    /// `questions` maps question ids (you choose them; answers come back under
    /// the same keys) to [`Question`]s. `Noul`, `Choice`, and `Score` all
    /// convert into `Question`, and the different types can be mixed freely in
    /// one call — every question is evaluated in parallel against the same
    /// state.
    pub async fn system_one<S, I, K, Q>(
        &self,
        state: S,
        questions: I,
    ) -> Result<SystemOneResponse, Error>
    where
        S: Serialize,
        I: IntoIterator<Item = (K, Q)>,
        K: Into<String>,
        Q: Into<Question>,
    {
        let model = self.model.clone();
        self.system_one_with_model(&model, state, questions).await
    }

    /// Like [`TypeSafeClient::system_one`], with a per-call model override.
    pub async fn system_one_with_model<S, I, K, Q>(
        &self,
        model: &str,
        state: S,
        questions: I,
    ) -> Result<SystemOneResponse, Error>
    where
        S: Serialize,
        I: IntoIterator<Item = (K, Q)>,
        K: Into<String>,
        Q: Into<Question>,
    {
        let questions = build_questions_map(questions)?;
        let state =
            serde_json::to_value(state).map_err(|e| Error::InvalidRequest(e.to_string()))?;

        let mut body = Map::new();
        body.insert("state".to_owned(), state);
        body.insert("model".to_owned(), Value::String(model.to_owned()));
        body.insert("questions".to_owned(), Value::Object(questions));

        let url = self.endpoint("/v1/systemone");
        let body = Value::Object(body);
        self.send_with_retry(|| self.http.post(&url).json(&body))
            .await
    }

    /// List the models available to the account.
    pub async fn list_models(&self) -> Result<ListModelsResponse, Error> {
        let url = self.endpoint("/v1/models");
        self.send_with_retry(|| self.http.get(&url)).await
    }

    fn endpoint(&self, path: &str) -> String {
        format!("{}{path}", self.base_url.trim_end_matches('/'))
    }

    async fn send_with_retry<T, F>(&self, make_request: F) -> Result<T, Error>
    where
        T: serde::de::DeserializeOwned,
        F: Fn() -> reqwest::RequestBuilder,
    {
        let mut attempt = 0u32;
        loop {
            let request = make_request()
                .header(header::AUTHORIZATION, self.auth.clone())
                .header(header::ACCEPT, "application/json");
            match request.send().await {
                Ok(resp) => {
                    let status = resp.status();
                    if status.is_success() {
                        let body = read_body_limited(resp, MAX_SUCCESS_BODY).await?;
                        return serde_json::from_slice::<T>(&body)
                            .map_err(|e| Error::Decode(e.to_string()));
                    }
                    let headers = resp.headers().clone();
                    let body = read_body_prefix(resp, MAX_ERROR_BODY).await;
                    let raw = String::from_utf8_lossy(&body);
                    let err = ApiError::from_response(status, &raw);
                    if attempt < self.retry.max_retries && self.retry.should_retry_status(status) {
                        let wait = self.retry.delay_for(attempt, &headers);
                        tokio::time::sleep(wait).await;
                        attempt += 1;
                        continue;
                    }
                    return Err(Error::Api(err));
                }
                Err(e) => {
                    let transient = e.is_connect() || e.is_timeout();
                    if attempt < self.retry.max_retries && transient {
                        let wait = self.retry.backoff_for(attempt);
                        tokio::time::sleep(wait).await;
                        attempt += 1;
                        continue;
                    }
                    return Err(Error::Connection(e.to_string()));
                }
            }
        }
    }
}

/// Builds [`TypeSafeClient`]s.
///
/// Explicit options take precedence over environment variables, which take
/// precedence over defaults. Empty or whitespace-only environment values are
/// ignored.
#[derive(Default)]
pub struct ClientBuilder {
    api_key: Option<String>,
    model: Option<String>,
    base_url: Option<String>,
    timeout: Option<Duration>,
    retry: Option<RetryPolicy>,
    default_headers: Option<HeaderMap>,
    http_client: Option<reqwest::Client>,
}

// Manual impl so the API key can never leak through formatting.
impl fmt::Debug for ClientBuilder {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ClientBuilder")
            .field("api_key", &self.api_key.as_ref().map(|_| "[redacted]"))
            .field("model", &self.model)
            .field("base_url", &self.base_url)
            .field("timeout", &self.timeout)
            .field("retry", &self.retry)
            .field("default_headers", &self.default_headers)
            .finish_non_exhaustive()
    }
}

impl ClientBuilder {
    /// Set the API key. Required unless [`ENV_API_KEY`] is set.
    pub fn api_key(mut self, api_key: impl Into<String>) -> Self {
        self.api_key = Some(api_key.into());
        self
    }

    /// Set the default model (default: `jev-latest`).
    pub fn model(mut self, model: impl Into<String>) -> Self {
        self.model = Some(model.into());
        self
    }

    /// Set the API root (default: `https://api.typesafe.ai`).
    pub fn base_url(mut self, base_url: impl Into<String>) -> Self {
        self.base_url = Some(base_url.into());
        self
    }

    /// Set the per-request timeout (default: 60s).
    pub fn timeout(mut self, timeout: Duration) -> Self {
        self.timeout = Some(timeout);
        self
    }

    /// Set the retry policy (default: [`RetryPolicy::default`]).
    pub fn retry(mut self, retry: RetryPolicy) -> Self {
        self.retry = Some(retry);
        self
    }

    /// Additional headers to send with every request. `Authorization`,
    /// `Accept`, and `Content-Type` remain managed by the SDK.
    pub fn default_headers(mut self, headers: HeaderMap) -> Self {
        self.default_headers = Some(headers);
        self
    }

    /// Supply a preconfigured `reqwest` client.
    ///
    /// The `timeout` and `default_headers` options are ignored when a client
    /// is supplied, and the supplied client's configuration wins. The caller
    /// is then responsible for its hardening (TLS verification, redirect
    /// policy, timeouts) — the SDK's defaults (redirects disabled, 60s
    /// timeout, rustls with certificate verification) no longer apply.
    pub fn http_client(mut self, client: reqwest::Client) -> Self {
        self.http_client = Some(client);
        self
    }

    /// Build the client.
    pub fn build(self) -> Result<TypeSafeClient, Error> {
        let api_key = resolve(self.api_key, ENV_API_KEY).ok_or_else(|| {
            Error::Config(format!(
                "missing API key: set {ENV_API_KEY} or use ClientBuilder::api_key"
            ))
        })?;
        if api_key.trim().is_empty() {
            return Err(Error::Config("API key is empty".into()));
        }

        let base_url =
            resolve(self.base_url, ENV_BASE_URL).unwrap_or_else(|| DEFAULT_BASE_URL.to_owned());
        validate_base_url(&base_url)?;

        let model =
            resolve(self.model, ENV_DEFAULT_MODEL).unwrap_or_else(|| DEFAULT_MODEL.to_owned());

        let mut auth = HeaderValue::from_str(&format!("Bearer {api_key}"))
            .map_err(|_| Error::Config("API key contains invalid header characters".into()))?;
        auth.set_sensitive(true);

        let http = match self.http_client {
            Some(client) => client,
            None => {
                let mut builder = reqwest::Client::builder()
                    .user_agent(USER_AGENT)
                    .timeout(self.timeout.unwrap_or(DEFAULT_TIMEOUT))
                    // API endpoints never legitimately redirect. Refusing to
                    // follow them removes any risk of the Authorization
                    // header being forwarded to an unexpected host.
                    .redirect(reqwest::redirect::Policy::none());
                if let Some(headers) = self.default_headers {
                    builder = builder.default_headers(headers);
                }
                builder
                    .build()
                    .map_err(|e| Error::Config(format!("failed to build HTTP client: {e}")))?
            }
        };

        Ok(TypeSafeClient {
            http,
            auth,
            base_url,
            model,
            retry: self.retry.unwrap_or_default(),
        })
    }
}

/// Resolve an explicit option against an environment variable; blank values
/// lose.
fn resolve(explicit: Option<String>, env: &str) -> Option<String> {
    explicit.or_else(|| std::env::var(env).ok()).and_then(|v| {
        let trimmed = v.trim();
        (!trimmed.is_empty()).then_some(v)
    })
}

/// Require the base URL to be valid and safe to send credentials to: HTTPS
/// everywhere, with plain HTTP allowed only for loopback hosts (local test
/// servers).
fn validate_base_url(base_url: &str) -> Result<(), Error> {
    let url = reqwest::Url::parse(base_url)
        .map_err(|_| Error::Config(format!("invalid base URL: {base_url:?}")))?;
    if url.scheme() == "https" {
        return Ok(());
    }
    if url.scheme() == "http" && url.host_str().is_some_and(is_loopback_host) {
        return Ok(());
    }
    Err(Error::Config(format!(
        "base URL {base_url:?} must use https (http is only allowed for loopback hosts); \
         refusing to send the API key over plaintext"
    )))
}

/// Whether a host is loopback: `localhost` (including subdomains) or a
/// loopback IP literal. `url::Url::host_str` keeps the brackets on IPv6
/// literals, so strip them before parsing.
fn is_loopback_host(host: &str) -> bool {
    if host == "localhost" || host.ends_with(".localhost") {
        return true;
    }
    let bare = host.strip_prefix('[').and_then(|h| h.strip_suffix(']'));
    let candidate = bare.unwrap_or(host);
    candidate.parse::<IpAddr>().is_ok_and(|ip| ip.is_loopback())
}

/// Read a successful response body, failing if it exceeds `limit` bytes.
/// Works even when the server omits or lies about `Content-Length`.
async fn read_body_limited(mut resp: reqwest::Response, limit: u64) -> Result<Vec<u8>, Error> {
    if resp.content_length().is_some_and(|len| len > limit) {
        return Err(Error::Decode(format!(
            "response body too large (Content-Length exceeds {limit} bytes)"
        )));
    }
    let mut buf = Vec::new();
    while let Some(chunk) = resp
        .chunk()
        .await
        .map_err(|e| Error::Connection(e.to_string()))?
    {
        if buf.len() as u64 + chunk.len() as u64 > limit {
            return Err(Error::Decode(format!(
                "response body exceeded the {limit}-byte limit"
            )));
        }
        buf.extend_from_slice(&chunk);
    }
    Ok(buf)
}

/// Read at most `limit` bytes of an error response body; the rest is
/// discarded. Error handling must not fail just because the body is huge.
async fn read_body_prefix(mut resp: reqwest::Response, limit: usize) -> Vec<u8> {
    let mut buf = Vec::new();
    loop {
        let remaining = limit.saturating_sub(buf.len());
        if remaining == 0 {
            break;
        }
        match resp.chunk().await {
            Ok(Some(chunk)) => buf.extend_from_slice(&chunk[..chunk.len().min(remaining)]),
            Ok(None) | Err(_) => break,
        }
    }
    buf
}

fn build_questions_map<I, K, Q>(questions: I) -> Result<Map<String, Value>, Error>
where
    I: IntoIterator<Item = (K, Q)>,
    K: Into<String>,
    Q: Into<Question>,
{
    let mut map = Map::new();
    for (key, question) in questions {
        let question = question.into();
        if let Question::Score(score) = &question {
            if score.criteria.is_empty() {
                return Err(Error::InvalidRequest(format!(
                    "score question {:?} must have at least one criteria level",
                    key.into()
                )));
            }
        }
        let key = key.into();
        let value = serde_json::to_value(&question)
            .map_err(|e| Error::InvalidRequest(format!("question {key:?}: {e}")))?;
        map.insert(key, value);
    }
    if map.is_empty() {
        return Err(Error::InvalidRequest(
            "questions must not be empty".to_owned(),
        ));
    }
    Ok(map)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::question::{Choice, Noul, Score};

    #[test]
    fn builder_requires_api_key() {
        // Ensure the env var cannot leak in from the test environment.
        let key_present = std::env::var(ENV_API_KEY).is_ok();
        let result = TypeSafeClient::builder().build();
        if key_present {
            assert!(result.is_ok());
        } else {
            assert!(matches!(result, Err(Error::Config(_))));
        }
    }

    #[test]
    fn builder_rejects_bad_base_url() {
        let result = TypeSafeClient::builder()
            .api_key("k")
            .base_url("not a url")
            .build();
        assert!(matches!(result, Err(Error::Config(_))));
    }

    #[test]
    fn builder_rejects_plaintext_http_base_url() {
        // Plain HTTP to a remote host would leak the API key.
        let result = TypeSafeClient::builder()
            .api_key("k")
            .base_url("http://api.typesafe.ai")
            .build();
        let err = result.unwrap_err();
        assert!(matches!(err, Error::Config(_)));
        assert!(err.to_string().contains("https"), "{err}");

        // Other schemes are rejected too.
        let result = TypeSafeClient::builder()
            .api_key("k")
            .base_url("ftp://example.com")
            .build();
        assert!(matches!(result.unwrap_err(), Error::Config(_)));
    }

    #[test]
    fn builder_allows_loopback_http_for_testing() {
        for url in [
            "http://localhost:8080",
            "http://127.0.0.1:8080",
            "http://[::1]",
            "http://test.localhost",
            "https://api.example.com",
        ] {
            TypeSafeClient::builder()
                .api_key("k")
                .base_url(url)
                .build()
                .unwrap_or_else(|e| panic!("{url} should be accepted: {e}"));
        }
    }

    #[test]
    fn builder_rejects_key_with_control_characters() {
        // A newline in the key must never reach a header (CRLF injection).
        let result = TypeSafeClient::builder()
            .api_key("good\r\nX-Injected: bad")
            .build();
        assert!(matches!(result.unwrap_err(), Error::Config(_)));
    }

    #[test]
    fn debug_formatting_never_exposes_the_api_key() {
        let secret = "apikey_supersecretvalue";
        let builder = TypeSafeClient::builder().api_key(secret);
        let client = TypeSafeClient::new(secret).unwrap();

        for rendered in [format!("{builder:?}"), format!("{client:?}")] {
            assert!(
                !rendered.contains(secret),
                "secret leaked into Debug output: {rendered}"
            );
            assert!(rendered.contains("[redacted]"), "{rendered}");
        }
    }

    #[test]
    fn rejects_empty_questions() {
        let questions: Vec<(String, Question)> = vec![];
        let err = build_questions_map(questions).unwrap_err();
        assert!(matches!(err, Error::InvalidRequest(_)));
    }

    #[test]
    fn rejects_empty_score_criteria() {
        let questions = vec![(
            "s",
            Question::from(Score::new("rate", Vec::<String>::new())),
        )];
        let err = build_questions_map(questions).unwrap_err();
        assert!(matches!(err, Error::InvalidRequest(_)));
    }

    #[test]
    fn builds_mixed_question_map() {
        let questions = vec![
            (
                "tone",
                Question::from(Choice::options("Tone?", ["calm", "angry"])),
            ),
            ("billing", Question::from(Noul::new("Billing?"))),
        ];
        let map = build_questions_map(questions).unwrap();
        assert_eq!(map.len(), 2);
        assert_eq!(map["tone"]["type"], "choice");
        assert_eq!(map["billing"]["type"], "noul");
        // Insertion order preserved.
        let keys: Vec<_> = map.keys().collect();
        assert_eq!(keys, ["tone", "billing"]);
    }
}