Skip to main content

kunobi_decision/client/
builder.rs

1//! Client construction.
2
3use std::fmt;
4use std::future::Future;
5use std::net::IpAddr;
6use std::sync::Arc;
7use std::sync::atomic::AtomicU64;
8use std::time::Duration;
9
10use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
11use secrecy::SecretString;
12use tokio::sync::Semaphore;
13
14use super::env::read_env;
15use super::{Client, Inner, Provider};
16use crate::credentials::{BoxError, CredentialProvider, Credentials, FnProvider, bearer_header};
17use crate::error::{Error, Result};
18use crate::retry::{DEFAULT_TIMEOUT, DEFAULT_TOTAL_TIMEOUT, RetryPolicy};
19
20/// Builder for [`Client`].
21///
22/// Values set here take precedence over environment variables, then defaults.
23/// Blank environment values are ignored.
24#[derive(Default)]
25pub struct ClientBuilder {
26    provider: Provider,
27    credentials: Option<Credentials>,
28    base_url: Option<String>,
29    allow_insecure_http: bool,
30    log_bodies: bool,
31    default_model: Option<String>,
32    retry: Option<RetryPolicy>,
33    timeout: Option<Duration>,
34    /// Outer `None` means "not set, use the default"; inner `None` means the
35    /// caller removed the bound.
36    total_timeout: Option<Option<Duration>>,
37    max_concurrent_requests: Option<usize>,
38    default_headers: HeaderMap,
39    http_client: Option<reqwest::Client>,
40}
41
42impl ClientBuilder {
43    /// Select the access provider. TypeSafe is the default.
44    pub fn provider(mut self, provider: Provider) -> Self {
45        self.provider = provider;
46        self
47    }
48
49    /// Authenticate with an API key; falls back to the selected provider's key environment variable.
50    ///
51    /// The key is held as a [`SecretString`] and wiped from memory when the client
52    /// is dropped. Don't embed an API key in software that runs on user machines;
53    /// use [`ClientBuilder::credential_provider`] with your own backend instead.
54    pub fn api_key(mut self, api_key: impl Into<SecretString>) -> Self {
55        self.credentials = Some(Credentials::ApiKey(api_key.into()));
56        self
57    }
58
59    /// Authenticate with a bearer token from a provider, asked before every attempt.
60    ///
61    /// Replaces any API key set on the builder, and the environment key is not read.
62    pub fn credential_provider(mut self, provider: impl CredentialProvider) -> Self {
63        self.credentials = Some(Credentials::Provider(Box::new(provider)));
64        self
65    }
66
67    /// Authenticate with a bearer token from an async closure, asked before every attempt.
68    ///
69    /// ```no_run
70    /// # fn example() -> kunobi_decision::Result<()> {
71    /// use kunobi_decision::Client;
72    ///
73    /// let client = Client::builder()
74    ///     .base_url("https://jev.example.com")
75    ///     .credentials_fn(|| async { Ok::<_, std::io::Error>(String::from("short-lived-token")) })
76    ///     .build()?;
77    /// # Ok(()) }
78    /// ```
79    pub fn credentials_fn<F, Fut, T, E>(self, provider: F) -> Self
80    where
81        F: Fn() -> Fut + Send + Sync + 'static,
82        Fut: Future<Output = std::result::Result<T, E>> + Send + 'static,
83        T: Into<SecretString> + 'static,
84        E: Into<BoxError> + 'static,
85    {
86        self.credential_provider(FnProvider(provider))
87    }
88
89    /// API root; defaults to the selected provider's root.
90    /// TypeSafe reads `TYPESAFE_BASE_URL`; OpenRouter reads `OPENROUTER_BASE_URL`.
91    ///
92    /// Must be https. Plain http is accepted only for loopback hosts, unless
93    /// [`ClientBuilder::allow_insecure_http`] is set.
94    pub fn base_url(mut self, base_url: impl Into<String>) -> Self {
95        self.base_url = Some(base_url.into());
96        self
97    }
98
99    /// Accept a plain-http `base_url` on any host. Credentials then travel unencrypted.
100    ///
101    /// Meant for trusted private networks, such as a sidecar on the same pod.
102    pub fn allow_insecure_http(mut self, allow: bool) -> Self {
103        self.allow_insecure_http = allow;
104        self
105    }
106
107    /// Include request and response bodies in `debug` logs. Default: off.
108    ///
109    /// Bodies carry the state you send and the answers you get, which may include
110    /// personal data. Headers are always logged with credentials redacted.
111    pub fn log_bodies(mut self, log_bodies: bool) -> Self {
112        self.log_bodies = log_bodies;
113        self
114    }
115
116    /// Default model; defaults to the selected provider's model.
117    /// TypeSafe reads `TYPESAFE_DEFAULT_MODEL`; OpenRouter reads `OPENROUTER_DEFAULT_MODEL`.
118    pub fn default_model(mut self, model: impl Into<String>) -> Self {
119        self.default_model = Some(model.into());
120        self
121    }
122
123    /// Retry policy; defaults to [`RetryPolicy::default`].
124    pub fn retry(mut self, retry: RetryPolicy) -> Self {
125        self.retry = Some(retry);
126        self
127    }
128
129    /// Timeout per attempt, including the response body.
130    /// Default: [`DEFAULT_TIMEOUT`](crate::DEFAULT_TIMEOUT), 5 s.
131    ///
132    /// A credential provider gets the same timeout. To bound a whole call, retries
133    /// included, use [`ClientBuilder::total_timeout`].
134    pub fn timeout(mut self, timeout: Duration) -> Self {
135        self.timeout = Some(timeout);
136        self
137    }
138
139    /// Upper bound for a whole call, including credentials, retries and backoff.
140    /// Default: [`DEFAULT_TOTAL_TIMEOUT`](crate::DEFAULT_TOTAL_TIMEOUT), 10 s.
141    ///
142    /// Pass `None` to remove the bound and let the retry policy run to its end.
143    /// Per-call [`Call::total_timeout`](crate::Call::total_timeout) takes precedence.
144    pub fn total_timeout(mut self, total_timeout: impl Into<Option<Duration>>) -> Self {
145        self.total_timeout = Some(total_timeout.into());
146        self
147    }
148
149    /// Limit how many attempts this client and its clones send at once. Default: no limit.
150    ///
151    /// Calls over the limit wait for a slot; the wait counts against
152    /// [`total_timeout`](Self::total_timeout) but not the per-attempt timeout. A slot is
153    /// held for one attempt and released during retry backoff.
154    pub fn max_concurrent_requests(mut self, max: usize) -> Self {
155        self.max_concurrent_requests = Some(max);
156        self
157    }
158
159    /// Add a header sent with every request. Per-call headers take precedence.
160    pub fn default_header(mut self, name: HeaderName, value: HeaderValue) -> Self {
161        self.default_headers.insert(name, value);
162        self
163    }
164
165    /// Add headers sent with every request, replacing earlier values for the same names.
166    pub fn default_headers(mut self, headers: HeaderMap) -> Self {
167        self.default_headers.extend(headers);
168        self
169    }
170
171    /// Use a preconfigured HTTP client, for proxies, custom TLS or tests.
172    ///
173    /// Leave its own timeout unset: the client applies [`ClientBuilder::timeout`] per attempt.
174    pub fn http_client(mut self, http: reqwest::Client) -> Self {
175        self.http_client = Some(http);
176        self
177    }
178
179    /// Build the client.
180    ///
181    /// Fails when credentials are missing or configuration is invalid.
182    pub fn build(self) -> Result<Client> {
183        self.build_with_env(read_env)
184    }
185
186    fn build_with_env(self, env: impl Fn(&str) -> Option<String>) -> Result<Client> {
187        self.provider.validate()?;
188        let api_key_env = self.provider.api_key_env();
189        let credentials = match self.credentials {
190            Some(credentials) => credentials,
191            None => Credentials::ApiKey(SecretString::from(env(api_key_env).ok_or_else(|| {
192                Error::Config(format!(
193                    "No credentials were provided. Call `ClientBuilder::api_key` or \
194                     `ClientBuilder::credential_provider`, or set the {api_key_env} environment variable."
195                ))
196            })?)),
197        };
198        if let Credentials::ApiKey(key) = &credentials {
199            bearer_header(key)
200                .map_err(|message| Error::Config(format!("The API key {message}.")))?;
201        }
202
203        let base_url = self
204            .base_url
205            .or_else(|| self.provider.base_url_env().and_then(&env))
206            .unwrap_or_else(|| self.provider.base_url().to_owned());
207        let base_url = base_url.trim().trim_end_matches('/').to_owned();
208        validate_base_url(&base_url, self.allow_insecure_http)?;
209
210        let default_model = self
211            .default_model
212            .or_else(|| self.provider.default_model_env().and_then(&env))
213            .unwrap_or_else(|| self.provider.default_model().to_owned());
214
215        let retry = self.retry.unwrap_or_default();
216        retry.validate()?;
217        let timeout = validate_timeout(self.timeout.unwrap_or(DEFAULT_TIMEOUT))?;
218        let total_timeout = self.total_timeout.unwrap_or(Some(DEFAULT_TOTAL_TIMEOUT));
219        if total_timeout.is_some_and(|total| total.is_zero()) {
220            return Err(Error::Config(
221                "`total_timeout` must be a positive duration, got 0.".into(),
222            ));
223        }
224
225        let limiter = match self.max_concurrent_requests {
226            Some(max) if max == 0 || max > Semaphore::MAX_PERMITS => {
227                return Err(Error::Config(format!(
228                    "`max_concurrent_requests` must be between 1 and {}, got {max}.",
229                    Semaphore::MAX_PERMITS
230                )));
231            }
232            Some(max) => Some((Semaphore::new(max), max)),
233            None => None,
234        };
235
236        let http = match self.http_client {
237            Some(http) => http,
238            None if !TLS_BACKEND && base_url.starts_with("https:") => {
239                return Err(Error::Config(
240                    "kunobi-decision was built without a TLS backend, so it cannot reach an https \
241                     `base_url`. Enable the `rustls` or `native-tls` feature, or pass a client \
242                     through `ClientBuilder::http_client`."
243                        .into(),
244                ));
245            }
246            None => reqwest::Client::builder()
247                .redirect(reqwest::redirect::Policy::none())
248                .build()
249                .map_err(|err| Error::Config(format!("Could not create the HTTP client: {err}")))?,
250        };
251
252        Ok(Client {
253            inner: Arc::new(Inner {
254                provider: self.provider,
255                credentials,
256                base_url,
257                log_bodies: self.log_bodies,
258                default_model,
259                retry,
260                timeout,
261                total_timeout,
262                limiter,
263                default_headers: self.default_headers,
264                http,
265                request_count: AtomicU64::new(0),
266            }),
267        })
268    }
269}
270
271/// Whether the default HTTP client can speak TLS.
272const TLS_BACKEND: bool = cfg!(any(feature = "rustls", feature = "native-tls"));
273
274/// Require https, allowing plain http only for loopback hosts or when explicitly allowed.
275fn validate_base_url(base_url: &str, allow_insecure_http: bool) -> Result<()> {
276    let url = reqwest::Url::parse(base_url).map_err(|_| {
277        Error::Config(format!(
278            "`base_url` must be an https URL, got \"{base_url}\"."
279        ))
280    })?;
281    if !url.username().is_empty()
282        || url.password().is_some()
283        || url.query().is_some()
284        || url.fragment().is_some()
285    {
286        return Err(Error::Config(
287            "`base_url` must not contain credentials, a query or a fragment.".into(),
288        ));
289    }
290    match url.scheme() {
291        "https" => Ok(()),
292        "http" if allow_insecure_http || is_loopback(url.host_str().unwrap_or_default()) => Ok(()),
293        "http" => Err(Error::Config(format!(
294            "`base_url` \"{base_url}\" uses plain http, which would send credentials unencrypted. \
295             Use https, a loopback host, or `ClientBuilder::allow_insecure_http(true)`."
296        ))),
297        _ => Err(Error::Config(format!(
298            "`base_url` must be an https URL, got \"{base_url}\"."
299        ))),
300    }
301}
302
303fn is_loopback(host: &str) -> bool {
304    host.eq_ignore_ascii_case("localhost")
305        || host
306            .trim_start_matches('[')
307            .trim_end_matches(']')
308            .parse::<IpAddr>()
309            .is_ok_and(|ip| ip.is_loopback())
310}
311
312pub(crate) fn validate_timeout(timeout: Duration) -> Result<Duration> {
313    if timeout.is_zero() {
314        return Err(Error::Config(
315            "`timeout` must be a positive duration, got 0.".into(),
316        ));
317    }
318    Ok(timeout)
319}
320
321impl fmt::Debug for ClientBuilder {
322    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
323        f.debug_struct("ClientBuilder")
324            .field("provider", &self.provider)
325            .field("credentials", &self.credentials)
326            .field("base_url", &self.base_url)
327            .field("allow_insecure_http", &self.allow_insecure_http)
328            .field("log_bodies", &self.log_bodies)
329            .field("default_model", &self.default_model)
330            .field("retry", &self.retry)
331            .field("timeout", &self.timeout)
332            .field("total_timeout", &self.total_timeout)
333            .field("max_concurrent_requests", &self.max_concurrent_requests)
334            .finish_non_exhaustive()
335    }
336}
337
338#[cfg(test)]
339mod tests {
340    use super::*;
341    use crate::client::env::non_blank;
342    use crate::{DEFAULT_BASE_URL, DEFAULT_MODEL, ENV_API_KEY, ENV_BASE_URL, ENV_DEFAULT_MODEL};
343    use std::collections::HashMap;
344
345    fn env(pairs: &[(&str, &str)]) -> impl Fn(&str) -> Option<String> {
346        let map: HashMap<String, String> = pairs
347            .iter()
348            .map(|(k, v)| (k.to_string(), v.to_string()))
349            .collect();
350        move |name| non_blank(map.get(name).cloned())
351    }
352
353    /// A builder that works in every feature set: with no TLS backend compiled in, the
354    /// default client refuses https, which is not what these tests are about.
355    fn builder() -> ClientBuilder {
356        ClientBuilder::default().http_client(reqwest::Client::new())
357    }
358
359    async fn authorization(client: &Client) -> String {
360        client
361            .inner
362            .credentials
363            .authorization(Duration::from_secs(1))
364            .await
365            .unwrap()
366            .to_str()
367            .unwrap()
368            .to_owned()
369    }
370
371    #[tokio::test]
372    async fn provider_credentials_and_defaults_are_isolated_from_typesafe_environment() {
373        for provider in [
374            Provider::Liquid,
375            Provider::Vercel,
376            Provider::OpenRouter,
377            Provider::Cloudflare {
378                account_id: "account".into(),
379            },
380        ] {
381            let client = builder()
382                .provider(provider.clone())
383                .build_with_env(env(&[
384                    (provider.api_key_env(), "provider-key"),
385                    (ENV_API_KEY, "wrong-key"),
386                    (ENV_BASE_URL, "https://wrong.example"),
387                    (ENV_DEFAULT_MODEL, "wrong-model"),
388                ]))
389                .unwrap();
390            assert_eq!(client.base_url(), provider.base_url());
391            assert_eq!(client.default_model(), provider.default_model());
392            assert_eq!(authorization(&client).await, "Bearer provider-key");
393            assert!(
394                builder()
395                    .provider(provider)
396                    .build_with_env(env(&[(ENV_API_KEY, "wrong-key")]))
397                    .is_err()
398            );
399        }
400    }
401
402    #[tokio::test]
403    async fn openrouter_environment_is_scoped_and_explicit_settings_take_precedence() {
404        let vars = env(&[
405            ("OPENROUTER_API_KEY", "router-key"),
406            ("OPENROUTER_BASE_URL", "http://localhost:9000/api"),
407            ("OPENROUTER_DEFAULT_MODEL", "typesafe/jev-1.13"),
408            (ENV_API_KEY, "typesafe-key"),
409        ]);
410        let c = builder()
411            .provider(Provider::OpenRouter)
412            .build_with_env(&vars)
413            .unwrap();
414        assert_eq!(c.base_url(), "http://localhost:9000/api");
415        assert_eq!(c.default_model(), "typesafe/jev-1.13");
416        assert_eq!(authorization(&c).await, "Bearer router-key");
417        let c = builder()
418            .provider(Provider::OpenRouter)
419            .api_key("explicit-key")
420            .base_url("http://localhost:9001/api")
421            .default_model("liquid/d1")
422            .build_with_env(&vars)
423            .unwrap();
424        assert_eq!(c.base_url(), "http://localhost:9001/api");
425        assert_eq!(c.default_model(), "liquid/d1");
426        assert_eq!(authorization(&c).await, "Bearer explicit-key");
427        let c = builder().build_with_env(vars).unwrap();
428        assert_eq!(c.provider(), &Provider::TypeSafe);
429        assert_eq!(c.base_url(), DEFAULT_BASE_URL);
430        assert_eq!(c.default_model(), DEFAULT_MODEL);
431        assert_eq!(authorization(&c).await, "Bearer typesafe-key");
432    }
433
434    #[test]
435    fn requires_credentials() {
436        let err = builder().build_with_env(env(&[])).unwrap_err();
437        assert!(err.to_string().contains(ENV_API_KEY), "{err}");
438        let blank_env = builder().build_with_env(env(&[(ENV_API_KEY, "  ")]));
439        assert!(blank_env.is_err());
440        let blank_code = builder()
441            .api_key(" ")
442            .build_with_env(env(&[(ENV_API_KEY, "k")]));
443        assert_eq!(blank_code.unwrap_err().to_string(), "The API key is blank.");
444    }
445
446    #[tokio::test]
447    async fn falls_back_to_environment_then_defaults() {
448        let client = builder()
449            .build_with_env(env(&[(ENV_API_KEY, " k ")]))
450            .unwrap();
451        assert_eq!(client.base_url(), DEFAULT_BASE_URL);
452        assert_eq!(client.default_model(), DEFAULT_MODEL);
453        assert_eq!(client.retry(), &RetryPolicy::default());
454        assert_eq!(client.timeout(), DEFAULT_TIMEOUT);
455        assert_eq!(authorization(&client).await, "Bearer k");
456        assert!(!client.inner.log_bodies);
457
458        let from_env = builder()
459            .build_with_env(env(&[
460                (ENV_API_KEY, "k"),
461                (ENV_BASE_URL, "http://localhost:8080//"),
462                (ENV_DEFAULT_MODEL, "jev-test"),
463            ]))
464            .unwrap();
465        assert_eq!(from_env.base_url(), "http://localhost:8080");
466        assert_eq!(from_env.default_model(), "jev-test");
467    }
468
469    #[tokio::test]
470    async fn code_takes_precedence_over_environment() {
471        let client = builder()
472            .api_key("from-code")
473            .base_url("https://example.com/")
474            .default_model("jev-code")
475            .build_with_env(env(&[
476                (ENV_API_KEY, "from-env"),
477                (ENV_BASE_URL, "https://env.example.com"),
478                (ENV_DEFAULT_MODEL, "jev-env"),
479            ]))
480            .unwrap();
481        assert_eq!(authorization(&client).await, "Bearer from-code");
482        assert_eq!(client.base_url(), "https://example.com");
483        assert_eq!(client.default_model(), "jev-code");
484    }
485
486    #[tokio::test]
487    async fn a_provider_replaces_api_keys_from_code_and_environment() {
488        let client = builder()
489            .api_key("from-code")
490            .credentials_fn(|| async { Ok::<_, BoxError>("from-provider") })
491            .build_with_env(env(&[(ENV_API_KEY, "from-env")]))
492            .unwrap();
493        assert_eq!(authorization(&client).await, "Bearer from-provider");
494    }
495
496    #[test]
497    fn https_needs_a_tls_backend_or_a_custom_client() {
498        let default_client = ClientBuilder::default()
499            .api_key("k")
500            .build_with_env(env(&[]));
501        assert_eq!(default_client.is_ok(), TLS_BACKEND);
502        if !TLS_BACKEND {
503            let err = default_client.unwrap_err();
504            assert!(err.to_string().contains("without a TLS backend"), "{err}");
505        }
506        let custom = ClientBuilder::default()
507            .api_key("k")
508            .http_client(reqwest::Client::new())
509            .build_with_env(env(&[]));
510        assert!(custom.is_ok());
511    }
512
513    #[test]
514    fn plain_http_is_limited_to_loopback_hosts() {
515        let build = |url: &str, allow: bool| {
516            builder()
517                .api_key("k")
518                .base_url(url)
519                .allow_insecure_http(allow)
520                .build_with_env(env(&[]))
521        };
522        for ok in [
523            "https://api.typesafe.ai",
524            "http://localhost:1234",
525            "http://LOCALHOST",
526            "http://127.0.0.1:9",
527            "http://127.8.0.1",
528            "http://[::1]:80",
529        ] {
530            assert!(build(ok, false).is_ok(), "{ok}");
531        }
532        for insecure in [
533            "http://api.typesafe.ai",
534            "http://10.0.0.5",
535            "http://localhost.evil.com",
536            "http://[::2]",
537        ] {
538            let err = build(insecure, false).unwrap_err();
539            assert!(err.to_string().contains("plain http"), "{insecure}: {err}");
540            assert!(build(insecure, true).is_ok(), "{insecure} with opt-out");
541        }
542        for invalid in ["ftp://x", "not a url", ""] {
543            assert!(build(invalid, true).is_err(), "{invalid}");
544        }
545    }
546
547    #[test]
548    fn plain_http_from_the_environment_is_refused_too() {
549        let err = builder()
550            .build_with_env(env(&[
551                (ENV_API_KEY, "k"),
552                (ENV_BASE_URL, "http://attacker.example"),
553            ]))
554            .unwrap_err();
555        assert!(err.to_string().contains("plain http"), "{err}");
556    }
557
558    #[test]
559    fn rejects_invalid_configuration() {
560        let build = |builder: ClientBuilder| builder.api_key("k").build_with_env(env(&[]));
561        let err = build(builder().timeout(Duration::ZERO)).unwrap_err();
562        assert!(err.to_string().contains("timeout"));
563        let err = build(builder().retry(RetryPolicy {
564            backoff_jitter: 2.0,
565            ..RetryPolicy::default()
566        }))
567        .unwrap_err();
568        assert!(err.to_string().contains("retry.backoff_jitter"));
569        let err = builder()
570            .api_key("bad\nkey")
571            .build_with_env(env(&[]))
572            .unwrap_err();
573        assert!(err.to_string().contains("HTTP header"));
574        assert!(
575            !err.to_string().contains("bad"),
576            "the key must not leak: {err}"
577        );
578    }
579
580    #[test]
581    fn debug_output_hides_credentials() {
582        let builder = builder().api_key("sk-secret");
583        assert!(!format!("{builder:?}").contains("sk-secret"));
584        let client = builder.build_with_env(env(&[])).unwrap();
585        assert!(!format!("{client:?}").contains("sk-secret"));
586    }
587}