sie_sdk/client/
mod.rs

1//! The client, its builder, and the one request path every endpoint shares.
2
3mod batches;
4mod connections;
5mod encode;
6mod files;
7pub(crate) mod generate;
8pub mod jobs;
9pub(crate) mod meta;
10mod pools;
11pub mod stream;
12#[cfg(feature = "watch")]
13mod watch;
14
15pub use batches::{BatchCreate, BatchListRequest, Batches};
16pub use connections::{ConnectionAdd, Connections};
17pub use encode::{EncodeRequest, ExtractRequest, ScoreRequest};
18pub use files::{FileListRequest, FileUpload, Files, SortOrder};
19pub use generate::{ChatRequest, GenerateRequest, ResponseInput, ResponsesRequest};
20pub use jobs::{
21    JobItem, JobSink, JobSource, JobSubmit, Jobs, chunk_is_retrievable, connection_name,
22    decode_chunk_payload, require_connection_name, require_connection_schema_policy,
23    require_connector_idempotency_key,
24};
25pub use pools::{PoolCreate, Pools};
26pub use stream::ChunkStream;
27#[cfg(feature = "watch")]
28pub use watch::WatchMode;
29
30use std::collections::HashMap;
31use std::sync::{Arc, Mutex};
32use std::time::Duration;
33
34use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
35use reqwest::{Method, Url};
36use serde_json::Value;
37use tokio::sync::Semaphore;
38
39use crate::error::{Error, Result, TransportErrorKind};
40use crate::http::headers::Origin;
41use crate::http::{HttpResponse, PreparedRequest, headers, version};
42use crate::retry::{Decision, RequestOptions, RetryPolicy, RetryState};
43
44/// Machine profile and pool a request should be routed to.
45#[derive(Debug, Clone, Default)]
46pub(crate) struct Routing {
47    pub pool: Option<String>,
48    pub profile: Option<String>,
49}
50
51pub(crate) struct Inner {
52    http: reqwest::Client,
53    /// Always ends in `/` so relative paths join predictably.
54    base_url: Url,
55    base_origin: Option<Origin>,
56    control_plane_url: Option<Url>,
57    org: Option<String>,
58    timeout: Duration,
59    defaults: RequestOptions,
60    default_options: Option<Value>,
61    edge_headers: Vec<(HeaderName, HeaderValue)>,
62    /// Kept alongside the reqwest defaults for the WebSocket handshake, which builds its
63    /// own request rather than going through reqwest.
64    #[cfg(feature = "watch")]
65    http_authorization: Option<HeaderValue>,
66    concurrency: Option<Semaphore>,
67    /// Background lease renewals, one per pool this client created.
68    leases: Mutex<HashMap<String, tokio::task::JoinHandle<()>>>,
69}
70
71impl Drop for Inner {
72    fn drop(&mut self) {
73        if let Ok(leases) = self.leases.lock() {
74            for handle in leases.values() {
75                handle.abort();
76            }
77        }
78    }
79}
80
81/// A handle to one SIE server.
82///
83/// Cloning is cheap and shares the underlying connection pool, so a `Client` can be stored
84/// once and used from every task.
85#[derive(Clone)]
86pub struct Client {
87    pub(crate) inner: Arc<Inner>,
88}
89
90impl std::fmt::Debug for Client {
91    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
92        f.debug_struct("Client")
93            .field("base_url", &self.inner.base_url.as_str())
94            .field("timeout", &self.inner.timeout)
95            .finish_non_exhaustive()
96    }
97}
98
99impl Client {
100    /// A client with default settings.
101    pub fn new(base_url: impl AsRef<str>) -> Result<Self> {
102        ClientBuilder::new(base_url).build()
103    }
104
105    /// Start configuring a client.
106    pub fn builder(base_url: impl AsRef<str>) -> ClientBuilder {
107        ClientBuilder::new(base_url)
108    }
109
110    /// The server root, always with a trailing slash.
111    pub fn base_url(&self) -> &str {
112        self.inner.base_url.as_str()
113    }
114
115    /// Per-call defaults this client applies when a request builder does not override them.
116    pub fn default_options(&self) -> &RequestOptions {
117        &self.inner.defaults
118    }
119
120    pub(crate) fn timeout(&self) -> Duration {
121        self.inner.timeout
122    }
123
124    /// Resolve a path against the base URL.
125    pub(crate) fn url(&self, path: &str) -> Result<Url> {
126        self.inner
127            .base_url
128            .join(path.trim_start_matches('/'))
129            .map_err(|err| Error::invalid(format!("invalid request path {path:?}: {err}")))
130    }
131
132    /// The control-plane base for the connections namespace.
133    pub(crate) fn control_plane(&self) -> Result<(&Url, &str)> {
134        let url = self.inner.control_plane_url.as_ref().ok_or_else(|| {
135            Error::invalid("connections require a control_plane_url on the client builder")
136        })?;
137        let org =
138            self.inner.org.as_deref().ok_or_else(|| {
139                Error::invalid("connections require an org on the client builder")
140            })?;
141        Ok((url, org))
142    }
143
144    /// Split a `gpu` parameter into routing headers, falling back to the client default.
145    pub(crate) fn routing(&self, gpu: Option<&str>) -> Routing {
146        let resolved = gpu.or(self.inner.defaults.gpu.as_deref());
147        match resolved {
148            Some(value) => {
149                let (pool, profile) = version::parse_gpu_param(value);
150                Routing {
151                    pool: pool.map(str::to_string),
152                    profile: Some(profile.to_string()).filter(|p| !p.is_empty()),
153                }
154            }
155            None => Routing::default(),
156        }
157    }
158
159    /// Merge per-call runtime options over the client defaults.
160    pub(crate) fn merge_options(&self, options: Option<&Value>) -> Option<Value> {
161        match (self.inner.default_options.as_ref(), options) {
162            (None, other) => other.cloned(),
163            (Some(defaults), None) => Some(defaults.clone()),
164            (Some(defaults), Some(overrides)) => {
165                let mut merged = defaults.as_object().cloned().unwrap_or_default();
166                if let Some(overrides) = overrides.as_object() {
167                    for (key, value) in overrides {
168                        merged.insert(key.clone(), value.clone());
169                    }
170                }
171                Some(Value::Object(merged))
172            }
173        }
174    }
175
176    /// Build the per-call options for a request builder's defaults.
177    pub(crate) fn request_options(&self) -> RequestOptions {
178        self.inner.defaults.clone()
179    }
180
181    /// Options for the metadata endpoints, which never wait for capacity.
182    pub(crate) fn metadata_options(&self) -> RequestOptions {
183        RequestOptions {
184            gpu: None,
185            wait_for_capacity: false,
186            provision_timeout: self.inner.timeout.max(Duration::from_secs(1)),
187            max_oom_retries: 0,
188        }
189    }
190
191    /// The `Authorization` value this client sends, when it has one.
192    #[cfg(feature = "watch")]
193    pub(crate) fn authorization_header(&self) -> Option<HeaderValue> {
194        self.inner.http_authorization.clone()
195    }
196
197    /// The validated edge headers, for transports that cannot reuse the reqwest client.
198    #[cfg(feature = "watch")]
199    pub(crate) fn edge_headers(&self) -> &[(HeaderName, HeaderValue)] {
200        &self.inner.edge_headers
201    }
202
203    /// Whether a `ws`/`wss` URL is the WebSocket counterpart of the base origin.
204    #[cfg(feature = "watch")]
205    pub(crate) fn websocket_matches_base_origin(&self, url: &Url) -> bool {
206        !self.inner.edge_headers.is_empty()
207            && self
208                .inner
209                .base_origin
210                .as_ref()
211                .is_some_and(|origin| origin.matches_websocket(url))
212    }
213
214    pub(crate) fn edge_headers_apply_to(&self, url: &Url) -> bool {
215        !self.inner.edge_headers.is_empty()
216            && self
217                .inner
218                .base_origin
219                .as_ref()
220                .is_some_and(|origin| Origin::parse(url).is_some_and(|target| target == *origin))
221    }
222
223    fn build_attempt(
224        &self,
225        request: &PreparedRequest,
226        timeout: Duration,
227        allow_edge: bool,
228    ) -> reqwest::RequestBuilder {
229        let mut builder = self
230            .inner
231            .http
232            .request(request.method.clone(), request.url.clone())
233            .timeout(timeout)
234            .headers(request.headers.clone());
235
236        if allow_edge && self.edge_headers_apply_to(&request.url) {
237            for (name, value) in &self.inner.edge_headers {
238                builder = builder.header(name.clone(), value.clone());
239            }
240        }
241        match &request.body {
242            Some(body) => builder.body(body.clone()),
243            None => builder,
244        }
245    }
246
247    async fn permit(&self) -> Option<tokio::sync::SemaphorePermit<'_>> {
248        match &self.inner.concurrency {
249            // The semaphore lives as long as the client, so acquire cannot fail.
250            Some(semaphore) => semaphore.acquire().await.ok(),
251            None => None,
252        }
253    }
254
255    /// Send a request, retrying per the state's policy, and buffer the response.
256    pub(crate) async fn send(
257        &self,
258        request: PreparedRequest,
259        state: &mut RetryState,
260    ) -> Result<HttpResponse> {
261        loop {
262            let timeout = state.attempt_timeout(self.inner.timeout)?;
263            let outcome = {
264                let _permit = self.permit().await;
265                match self.build_attempt(&request, timeout, true).send().await {
266                    Ok(response) => HttpResponse::read(response).await,
267                    Err(error) => Err(error),
268                }
269            };
270
271            let response = match outcome {
272                Ok(response) => response,
273                Err(error) => {
274                    let delay = state.on_transport_error(&error, self.base_url())?;
275                    tokio::time::sleep(delay).await;
276                    continue;
277                }
278            };
279
280            match state.on_response(&response)? {
281                Decision::Accept => {
282                    Self::check_version(&response);
283                    return Ok(response);
284                }
285                Decision::Retry(delay) => tokio::time::sleep(delay).await,
286            }
287        }
288    }
289
290    /// Send a request whose successful response is consumed as a stream.
291    ///
292    /// Failures are buffered first so the shared error and retry handling sees a complete
293    /// body, exactly as the buffered path does.
294    pub(crate) async fn send_streaming(
295        &self,
296        request: &PreparedRequest,
297        state: &mut RetryState,
298    ) -> Result<StreamAttempt> {
299        loop {
300            let timeout = state.attempt_timeout(self.inner.timeout)?;
301            // A stream is read after `send` returns, so the permit cannot be held for the
302            // lifetime of the body without deadlocking a caller that reads them serially.
303            let attempt = {
304                let _permit = self.permit().await;
305                self.build_attempt(request, timeout, true).send().await
306            };
307
308            let response = match attempt {
309                Ok(response) => response,
310                Err(error) => {
311                    let delay = state.on_transport_error(&error, self.base_url())?;
312                    tokio::time::sleep(delay).await;
313                    continue;
314                }
315            };
316
317            if response.status().is_success() {
318                let buffered = HttpResponse {
319                    status: response.status().as_u16(),
320                    headers: response.headers().clone(),
321                    body: bytes::Bytes::new(),
322                };
323                Self::check_version(&buffered);
324                return Ok(StreamAttempt {
325                    headers: buffered.headers,
326                    response,
327                });
328            }
329
330            let buffered = HttpResponse::read(response).await.map_err(|error| {
331                Error::connection(TransportErrorKind::MidFlight, error.to_string(), error)
332            })?;
333            match state.on_response(&buffered)? {
334                // A non-success status that the policy accepts is still an error here: the
335                // caller asked for a stream and the server did not open one.
336                Decision::Accept => {
337                    return Err(crate::wire::handle_error(&buffered, None, state.retries()));
338                }
339                Decision::Retry(delay) => tokio::time::sleep(delay).await,
340            }
341        }
342    }
343
344    /// Warn once per process when the server and SDK versions have drifted.
345    fn check_version(response: &HttpResponse) {
346        if let Some(server_version) = response.header(headers::SERVER_VERSION) {
347            version::warn_once(server_version);
348        }
349    }
350
351    /// A single request with no retry policy, used by endpoints that own their own loop.
352    pub(crate) async fn send_once(
353        &self,
354        request: PreparedRequest,
355        policy: RetryPolicy,
356    ) -> Result<HttpResponse> {
357        let mut state = RetryState::new(policy, &self.metadata_options(), None);
358        self.send(request, &mut state).await
359    }
360
361    /// A single request whose timeout is raised to at least `floor`.
362    ///
363    /// File and batch transfers move whole payloads, so the per-request timeout that suits
364    /// an inference call is too tight for them.
365    pub(crate) async fn send_with_timeout(
366        &self,
367        request: PreparedRequest,
368        policy: RetryPolicy,
369        floor: Duration,
370    ) -> Result<HttpResponse> {
371        let options = RequestOptions {
372            provision_timeout: self.inner.timeout.max(floor),
373            ..self.metadata_options()
374        };
375        let mut state = RetryState::new(policy, &options, None);
376        self.send(request, &mut state).await
377    }
378
379    /// Build a request with the SDK's standard headers already applied.
380    pub(crate) fn request(&self, method: Method, path: &str) -> Result<PreparedRequest> {
381        Ok(PreparedRequest::new(method, self.url(path)?))
382    }
383}
384
385/// A stream response that passed the retry gauntlet.
386pub(crate) struct StreamAttempt {
387    pub headers: HeaderMap,
388    pub response: reqwest::Response,
389}
390
391/// Configures a [`Client`].
392pub struct ClientBuilder {
393    base_url: String,
394    api_key: Option<String>,
395    timeout: Duration,
396    gpu: Option<String>,
397    options: Option<Value>,
398    max_connections: Option<usize>,
399    max_concurrency: Option<usize>,
400    control_plane_url: Option<String>,
401    org: Option<String>,
402    base_url_headers: HashMap<String, String>,
403    wait_for_capacity: bool,
404    provision_timeout: Duration,
405    max_oom_retries: u32,
406}
407
408impl ClientBuilder {
409    fn new(base_url: impl AsRef<str>) -> Self {
410        let defaults = RequestOptions::default();
411        Self {
412            base_url: base_url.as_ref().to_string(),
413            api_key: None,
414            timeout: Duration::from_secs(30),
415            gpu: None,
416            options: None,
417            max_connections: None,
418            max_concurrency: None,
419            control_plane_url: None,
420            org: None,
421            base_url_headers: HashMap::new(),
422            wait_for_capacity: defaults.wait_for_capacity,
423            provision_timeout: defaults.provision_timeout,
424            max_oom_retries: defaults.max_oom_retries,
425        }
426    }
427
428    /// Bearer token sent as `Authorization` on every request to the server.
429    pub fn api_key(mut self, api_key: impl Into<String>) -> Self {
430        self.api_key = Some(api_key.into());
431        self
432    }
433
434    /// Per-attempt HTTP timeout. Defaults to 30 seconds.
435    pub fn timeout(mut self, timeout: Duration) -> Self {
436        self.timeout = timeout;
437        self
438    }
439
440    /// Default machine profile, optionally pool-qualified as `"pool/profile"`.
441    pub fn gpu(mut self, gpu: impl Into<String>) -> Self {
442        self.gpu = Some(gpu.into());
443        self
444    }
445
446    /// Default runtime options, shallow-merged under any per-call options.
447    pub fn options(mut self, options: Value) -> Self {
448        self.options = Some(options);
449        self
450    }
451
452    /// Cap on pooled connections.
453    pub fn max_connections(mut self, max: usize) -> Self {
454        self.max_connections = Some(max);
455        self
456    }
457
458    /// Cap on requests in flight at once.
459    pub fn max_concurrency(mut self, max: usize) -> Self {
460        self.max_concurrency = Some(max);
461        self
462    }
463
464    /// Control-plane root, required by the connections namespace.
465    pub fn control_plane_url(mut self, url: impl Into<String>) -> Self {
466        self.control_plane_url = Some(url.into());
467        self
468    }
469
470    /// Organisation slug, required by the connections namespace.
471    pub fn org(mut self, org: impl Into<String>) -> Self {
472        self.org = Some(org.into());
473        self
474    }
475
476    /// Extra headers for an HTTP edge in front of the gateway.
477    ///
478    /// These are credentials: they are sent only to the exact origin of `base_url`, which
479    /// must therefore be `https` without embedded userinfo.
480    pub fn base_url_headers(mut self, headers: HashMap<String, String>) -> Self {
481        self.base_url_headers = headers;
482        self
483    }
484
485    /// Whether calls wait out provisioning by default.
486    pub fn wait_for_capacity(mut self, wait: bool) -> Self {
487        self.wait_for_capacity = wait;
488        self
489    }
490
491    /// Default wall-clock budget for a call including its retries.
492    pub fn provision_timeout(mut self, timeout: Duration) -> Self {
493        self.provision_timeout = timeout;
494        self
495    }
496
497    /// Default cap on `RESOURCE_EXHAUSTED` retries.
498    pub fn max_oom_retries(mut self, retries: u32) -> Self {
499        self.max_oom_retries = retries;
500        self
501    }
502
503    /// Validate the configuration and construct the client.
504    pub fn build(self) -> Result<Client> {
505        let trimmed = self.base_url.trim_end_matches('/');
506        let base_url = Url::parse(&format!("{trimmed}/")).map_err(|err| {
507            Error::invalid(format!("invalid base_url {:?}: {err}", self.base_url))
508        })?;
509        let base_origin = Origin::parse(&base_url);
510
511        let edge_pairs = headers::validate_base_url_headers(&self.base_url_headers)?;
512        if !edge_pairs.is_empty()
513            && !base_origin
514                .as_ref()
515                .is_some_and(Origin::accepts_credentials)
516        {
517            return Err(Error::invalid(
518                "base_url_headers require an absolute https base_url without embedded credentials",
519            ));
520        }
521        let mut edge_headers = Vec::with_capacity(edge_pairs.len());
522        for (name, value) in edge_pairs {
523            let name = HeaderName::from_bytes(name.as_bytes())
524                .map_err(|err| Error::invalid(format!("invalid base_url_headers name: {err}")))?;
525            let value = HeaderValue::from_str(&value)
526                .map_err(|err| Error::invalid(format!("invalid base_url_headers value: {err}")))?;
527            edge_headers.push((name, value));
528        }
529
530        let control_plane_url = self
531            .control_plane_url
532            .as_deref()
533            .map(|url| {
534                Url::parse(&format!("{}/", url.trim_end_matches('/'))).map_err(|err| {
535                    Error::invalid(format!("invalid control_plane_url {url:?}: {err}"))
536                })
537            })
538            .transpose()?;
539
540        let mut default_headers = HeaderMap::new();
541        default_headers.insert(
542            HeaderName::from_static("x-sie-sdk-version"),
543            HeaderValue::from_static(version::SDK_VERSION),
544        );
545        #[cfg(feature = "watch")]
546        let mut http_authorization = None;
547        if let Some(api_key) = &self.api_key {
548            let mut value = HeaderValue::from_str(&format!("Bearer {api_key}")).map_err(|_| {
549                Error::invalid("api_key contains characters that cannot be sent in a header")
550            })?;
551            value.set_sensitive(true);
552            #[cfg(feature = "watch")]
553            {
554                http_authorization = Some(value.clone());
555            }
556            default_headers.insert(reqwest::header::AUTHORIZATION, value);
557        }
558
559        let mut http = reqwest::Client::builder()
560            // Never follow redirects: a redirect off-origin would leak the Authorization
561            // and edge headers to whatever host the response names.
562            .redirect(reqwest::redirect::Policy::none())
563            .default_headers(default_headers);
564        if let Some(max) = self.max_connections {
565            http = http.pool_max_idle_per_host(max);
566        }
567        let http = http
568            .build()
569            .map_err(|err| Error::invalid(format!("could not build the HTTP client: {err}")))?;
570
571        Ok(Client {
572            inner: Arc::new(Inner {
573                http,
574                base_url,
575                base_origin,
576                control_plane_url,
577                org: self.org,
578                timeout: self.timeout,
579                defaults: RequestOptions {
580                    gpu: self.gpu,
581                    wait_for_capacity: self.wait_for_capacity,
582                    provision_timeout: self.provision_timeout,
583                    max_oom_retries: self.max_oom_retries,
584                },
585                default_options: self.options,
586                edge_headers,
587                #[cfg(feature = "watch")]
588                http_authorization,
589                concurrency: self.max_concurrency.map(Semaphore::new),
590                leases: Mutex::new(HashMap::new()),
591            }),
592        })
593    }
594}
595
596#[cfg(test)]
597mod tests {
598    use super::*;
599    use serde_json::json;
600
601    fn client() -> Client {
602        Client::new("https://sie.example.com").unwrap()
603    }
604
605    #[test]
606    fn base_url_always_ends_in_a_slash() {
607        assert_eq!(client().base_url(), "https://sie.example.com/");
608        assert_eq!(
609            Client::new("https://sie.example.com///")
610                .unwrap()
611                .base_url(),
612            "https://sie.example.com/"
613        );
614    }
615
616    #[test]
617    fn paths_join_onto_the_base_url() {
618        let client = Client::new("https://sie.example.com/prefix").unwrap();
619        assert_eq!(
620            client.url("/v1/models").unwrap().as_str(),
621            "https://sie.example.com/prefix/v1/models"
622        );
623        assert_eq!(
624            client.url("v1/models").unwrap().as_str(),
625            "https://sie.example.com/prefix/v1/models"
626        );
627    }
628
629    #[test]
630    fn edge_headers_require_an_https_origin() {
631        let headers = HashMap::from([("Modal-Key".to_string(), "k".to_string())]);
632        assert!(
633            Client::builder("https://sie.example.com")
634                .base_url_headers(headers.clone())
635                .build()
636                .is_ok()
637        );
638        let err = Client::builder("http://localhost:8080")
639            .base_url_headers(headers)
640            .build()
641            .unwrap_err();
642        assert!(err.to_string().contains("https base_url"), "{err}");
643    }
644
645    #[test]
646    fn edge_headers_are_scoped_to_the_base_origin() {
647        let client = Client::builder("https://sie.example.com")
648            .base_url_headers(HashMap::from([("Modal-Key".to_string(), "k".to_string())]))
649            .build()
650            .unwrap();
651        assert!(
652            client.edge_headers_apply_to(&Url::parse("https://sie.example.com/v1/models").unwrap())
653        );
654        assert!(
655            !client
656                .edge_headers_apply_to(&Url::parse("https://other.example.com/v1/models").unwrap())
657        );
658        assert!(
659            !client.edge_headers_apply_to(&Url::parse("https://sie.example.com:8443/v1").unwrap())
660        );
661        assert!(!client.edge_headers_apply_to(&Url::parse("http://sie.example.com/v1").unwrap()));
662    }
663
664    #[test]
665    fn routing_splits_pool_from_profile() {
666        let client = Client::builder("https://sie.example.com")
667            .gpu("prod/l4")
668            .build()
669            .unwrap();
670        let routing = client.routing(None);
671        assert_eq!(routing.pool.as_deref(), Some("prod"));
672        assert_eq!(routing.profile.as_deref(), Some("l4"));
673
674        let overridden = client.routing(Some("a100-80gb"));
675        assert!(overridden.pool.is_none());
676        assert_eq!(overridden.profile.as_deref(), Some("a100-80gb"));
677
678        let bare = Client::new("https://sie.example.com")
679            .unwrap()
680            .routing(None);
681        assert!(bare.pool.is_none() && bare.profile.is_none());
682    }
683
684    #[test]
685    fn per_call_options_override_client_defaults() {
686        let client = Client::builder("https://sie.example.com")
687            .options(json!({"is_query": true, "keep": 1}))
688            .build()
689            .unwrap();
690        let merged = client
691            .merge_options(Some(&json!({"is_query": false})))
692            .unwrap();
693        assert_eq!(merged, json!({"is_query": false, "keep": 1}));
694        assert_eq!(
695            client.merge_options(None).unwrap(),
696            json!({"is_query": true, "keep": 1})
697        );
698        assert!(
699            Client::new("https://x.example.com")
700                .unwrap()
701                .merge_options(None)
702                .is_none()
703        );
704    }
705
706    #[test]
707    fn connections_need_a_control_plane_and_org() {
708        assert!(client().control_plane().is_err());
709        let configured = Client::builder("https://sie.example.com")
710            .control_plane_url("https://cp.example.com/")
711            .org("acme")
712            .build()
713            .unwrap();
714        let (url, org) = configured.control_plane().unwrap();
715        assert_eq!(url.as_str(), "https://cp.example.com/");
716        assert_eq!(org, "acme");
717    }
718
719    #[test]
720    fn rejects_a_malformed_base_url() {
721        assert!(Client::new("not a url").is_err());
722    }
723}
sie_sdk/client/mod.rs

sie_sdk/client/
mod.rs
