timvw_rust_runtime/
provider.rs

1//! Send a generated client to another provider of the same API: [`ProviderLayer`] moves each request to the provider's
2//! base URL and adds the headers, query parameters and path changes the provider needs.
3//!
4//! The layer is provider-neutral: it knows no provider by name. A crate for one API (OpenAI-compatible providers, for
5//! example) builds its presets on it.
6
7use std::fmt;
8use std::sync::Arc;
9use std::task::{Context, Poll};
10
11use futures_util::future::{self, Either, Ready};
12use futures_util::TryFutureExt;
13use http::header::{HeaderName, HeaderValue};
14use http::{HeaderMap, Uri};
15use serde::Serialize;
16use tower::{Layer, Service};
17use url::Url;
18
19use crate::request_options::Applied;
20use crate::{BoxError, Error};
21
22/// Rewrites a path relative to the base URL, see [`ProviderLayer::map_path`].
23type MapPath = Arc<dyn Fn(&str) -> String + Send + Sync>;
24
25/// A tower layer that sends requests to a provider: another host (and path prefix) serving the same API, with its own
26/// headers, query parameters and paths.
27///
28/// A request is the provider's when its URL starts with the provider's base URL, or with the base declared with
29/// [`ProviderLayer::rewrite_from`], which is then replaced with the provider's: same scheme, host and port, and a path
30/// that starts with the base's path at a segment boundary (`/v1` holds `/v1` and `/v1/models`, not `/v10`). To such a
31/// request the layer, in this order:
32///
33/// 1. replaces the base (scheme, host, port and path prefix) with the provider's, keeping the rest of the path and the
34///    query;
35/// 2. rewrites the rest of the path with [`ProviderLayer::map_path`], if set;
36/// 3. sets the [headers](ProviderLayer::header), each replacing every value the request has under its name;
37/// 4. sets the [query parameters](ProviderLayer::query), each replacing the parameters of its name the request has.
38///
39/// Any other request (a download from another host, say) is passed on unchanged, so the provider's credentials are
40/// never sent elsewhere.
41///
42/// Per-call [`RequestOptions`](crate::RequestOptions) win over the layer:
43///
44/// - a header the options set is not replaced by the layer's header of that name;
45/// - a query parameter the options added is not replaced by the layer's parameter of that name (with an `api-version`
46///   in both, the options' is sent, once);
47/// - a call the options sent to a [base URL](crate::RequestOptions::base_url) of their own is not moved: the layer
48///   ignores [`ProviderLayer::rewrite_from`] for it. When that base is the provider's (or under it), the layer still
49///   rewrites the path and adds its headers and query parameters; any other base gets nothing from the layer, so its
50///   credentials do not go there.
51///
52/// There are two ways to point a generated client at a provider:
53///
54/// - Set `Configuration.base_path` to the provider's base URL and add `ProviderLayer::new(provider_base)` for the
55///   headers, query parameters and paths. The requests already start with the provider's base.
56/// - Keep the generated `base_path` and declare it with `.rewrite_from(generated_base)`: the layer moves every request
57///   under it to the provider. This suits a `Client` shared by several `Configuration`s, or a crate that hands out a
58///   ready-made client for a provider.
59///
60/// ```
61/// use timvw_rust_runtime::provider::ProviderLayer;
62///
63/// // An OpenAI-compatible provider under another base URL, with headers of its own.
64/// let openrouter_like = ProviderLayer::new("https://openrouter.example/api/v1")?
65///     .rewrite_from("https://api.openai.com/v1")?
66///     .sensitive_header("authorization", "Bearer sk-or-123")?
67///     .header("x-title", "my app")?;
68///
69/// assert_eq!(
70///     openrouter_like.rewrite_uri(&"https://api.openai.com/v1/models?limit=2".parse()?)?.to_string(),
71///     "https://openrouter.example/api/v1/models?limit=2"
72/// );
73///
74/// // A provider that puts a deployment in the path and wants an `api-version` on every request.
75/// let azure_like = ProviderLayer::new("https://my-resource.example.com/openai")?
76///     .rewrite_from("https://api.openai.com/v1")?
77///     .sensitive_header("api-key", "secret")?
78///     .query(&[("api-version", "2024-10-21")])?
79///     .map_path(|path| format!("/deployments/my-deployment{path}"));
80///
81/// assert_eq!(
82///     azure_like.rewrite_uri(&"https://api.openai.com/v1/chat/completions".parse()?)?.to_string(),
83///     "https://my-resource.example.com/openai/deployments/my-deployment/chat/completions?api-version=2024-10-21"
84/// );
85/// # Ok::<(), Box<dyn std::error::Error>>(())
86/// ```
87///
88/// Add it with [`Client::layer`](crate::Client::layer), to `Client::default()` or to a generated `Configuration.client`.
89/// Added last, it sees each request first and rewrites it once, before the retries.
90#[derive(Clone)]
91pub struct ProviderLayer {
92    base: Url,
93    rewrite_from: Option<Url>,
94    headers: HeaderMap,
95    query: Vec<(String, String)>,
96    map_path: Option<MapPath>,
97}
98
99impl ProviderLayer {
100    /// A provider whose API is at `base_url`, such as `https://openrouter.ai/api/v1`. Fails, with an error whose
101    /// [`is_builder()`](Error::is_builder) is true, on a URL that is not absolute or that has a query or fragment.
102    pub fn new(base_url: &str) -> Result<Self, Error> {
103        Ok(Self {
104            base: parse_base(base_url)?,
105            rewrite_from: None,
106            headers: HeaderMap::new(),
107            query: Vec::new(),
108            map_path: None,
109        })
110    }
111
112    /// Moves the requests under `base_url`, typically a generated client's default `base_path` such as
113    /// `https://api.openai.com/v1`, to the provider: `base_url` is replaced with the provider's base URL, and the rest
114    /// of the path and the query are kept. Fails like [`ProviderLayer::new`]. Calling it again replaces the base.
115    pub fn rewrite_from(mut self, base_url: &str) -> Result<Self, Error> {
116        self.rewrite_from = Some(parse_base(base_url)?);
117        Ok(self)
118    }
119
120    /// Sets a header on every request to the provider, replacing every value the request has under that name (the
121    /// generated client's included). Calling it again for the same name replaces the value. Fails, with an error
122    /// whose [`is_builder()`](Error::is_builder) is true, on an invalid name or value.
123    pub fn header<K, V>(mut self, name: K, value: V) -> Result<Self, Error>
124    where
125        HeaderName: TryFrom<K>,
126        <HeaderName as TryFrom<K>>::Error: Into<http::Error>,
127        HeaderValue: TryFrom<V>,
128        <HeaderValue as TryFrom<V>>::Error: Into<http::Error>,
129    {
130        let (name, value) = header(name, value)?;
131        self.headers.insert(name, value);
132        Ok(self)
133    }
134
135    /// Like [`ProviderLayer::header`], and marks the value sensitive (see [`HeaderValue::set_sensitive`]): for an API
136    /// key, so that it is left out of `Debug` output.
137    pub fn sensitive_header<K, V>(mut self, name: K, value: V) -> Result<Self, Error>
138    where
139        HeaderName: TryFrom<K>,
140        <HeaderName as TryFrom<K>>::Error: Into<http::Error>,
141        HeaderValue: TryFrom<V>,
142        <HeaderValue as TryFrom<V>>::Error: Into<http::Error>,
143    {
144        let (name, mut value) = header(name, value)?;
145        value.set_sensitive(true);
146        self.headers.insert(name, value);
147        Ok(self)
148    }
149
150    /// Sets every header of a map, replacing the values set before under its names. A name with several values keeps
151    /// them all. Values marked sensitive stay sensitive.
152    pub fn headers(mut self, headers: HeaderMap) -> Self {
153        self.headers.extend(headers);
154        self
155    }
156
157    /// Sets query parameters on every request to the provider, such as `api-version`, serialized as
158    /// `application/x-www-form-urlencoded`: a sequence of pairs, a map or a struct. Each replaces the parameters of its
159    /// name the request has; a name given twice here is sent twice. Fails, with an error whose
160    /// [`is_builder()`](Error::is_builder) is true, on a value that is not a flat list of pairs.
161    pub fn query<T: Serialize + ?Sized>(mut self, query: &T) -> Result<Self, Error> {
162        let encoded = serde_urlencoded::to_string(query).map_err(Error::builder)?;
163        let pairs: Vec<(String, String)> = url::form_urlencoded::parse(encoded.as_bytes())
164            .into_owned()
165            .collect();
166        self.query
167            .retain(|(name, _)| !pairs.iter().any(|(new, _)| new == name));
168        self.query.extend(pairs);
169        Ok(self)
170    }
171
172    /// Rewrites the path of each request to the provider, relative to the base URL: `f` gets the path after the base
173    /// (`/chat/completions` for `https://api.openai.com/v1/chat/completions` with base `https://api.openai.com/v1`;
174    /// empty for the base itself), percent-encoded, and returns the new one, which is put after the provider's base.
175    pub fn map_path<F>(mut self, f: F) -> Self
176    where
177        F: Fn(&str) -> String + Send + Sync + 'static,
178    {
179        self.map_path = Some(Arc::new(f));
180        self
181    }
182
183    /// The provider's base URL.
184    pub fn base_url(&self) -> &str {
185        self.base.as_str()
186    }
187
188    /// The headers set on each request to the provider.
189    pub fn header_map(&self) -> &HeaderMap {
190        &self.headers
191    }
192
193    /// The query parameters set on each request to the provider, decoded, in order.
194    pub fn query_pairs(&self) -> &[(String, String)] {
195        &self.query
196    }
197
198    /// The URL a request to `uri` is sent to: moved to the provider, its path rewritten and its query parameters set,
199    /// or `uri` unchanged when it is not the provider's. Fails, with an error whose
200    /// [`is_builder()`](Error::is_builder) is true, when the rewritten URL is invalid.
201    pub fn rewrite_uri(&self, uri: &Uri) -> Result<Uri, Error> {
202        Ok(self.rewrite(uri, None)?.unwrap_or_else(|| uri.clone()))
203    }
204
205    /// The rewritten URL, or `None` when the request is not the provider's. `applied` is what per-call options set.
206    fn rewrite(&self, uri: &Uri, applied: Option<&Applied>) -> Result<Option<Uri>, Error> {
207        let Ok(mut url) = Url::parse(&uri.to_string()) else {
208            return Ok(None);
209        };
210        // A call the options sent to a base of their own is not moved.
211        let rewrite_from = self
212            .rewrite_from
213            .as_ref()
214            .filter(|_| !applied.is_some_and(|applied| applied.base));
215        let rest = rewrite_from
216            .into_iter()
217            .chain([&self.base])
218            .find_map(|base| rest_of_path(&url, base));
219        let Some(rest) = rest else {
220            return Ok(None);
221        };
222        let rest = match &self.map_path {
223            Some(map_path) => map_path(&rest),
224            None => rest,
225        };
226        let mut path = trimmed_path(&self.base).to_owned();
227        path.push_str(&rest);
228        // Scheme, host and port of the provider; the scheme is http(s) on both sides, so these do not fail.
229        let _ = url.set_scheme(self.base.scheme());
230        url.set_host(self.base.host_str()).map_err(Error::builder)?;
231        let _ = url.set_port(self.base.port());
232        url.set_path(if path.is_empty() { "/" } else { &path });
233        self.set_query(&mut url, applied);
234        url.as_str()
235            .parse::<Uri>()
236            .map(Some)
237            .map_err(Error::builder)
238    }
239
240    fn set_query(&self, url: &mut Url, applied: Option<&Applied>) {
241        // The options' parameters win: the layer leaves their names alone.
242        let query: Vec<&(String, String)> = self
243            .query
244            .iter()
245            .filter(|(name, _)| !applied.is_some_and(|applied| applied.has_query(name)))
246            .collect();
247        if query.is_empty() {
248            return;
249        }
250        let replaced = url
251            .query_pairs()
252            .any(|(name, _)| query.iter().any(|(new, _)| *new == name));
253        if replaced {
254            // Re-encode the query without the replaced names.
255            let kept: Vec<(String, String)> = url
256                .query_pairs()
257                .filter(|(name, _)| !query.iter().any(|(new, _)| new == name))
258                .map(|(name, value)| (name.into_owned(), value.into_owned()))
259                .collect();
260            url.set_query(None);
261            url.query_pairs_mut().extend_pairs(kept);
262        }
263        url.query_pairs_mut()
264            .extend_pairs(query.iter().map(|(name, value)| (name, value)));
265    }
266}
267
268impl fmt::Debug for ProviderLayer {
269    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
270        f.debug_struct("ProviderLayer")
271            .field("base", &self.base.as_str())
272            .field("rewrite_from", &self.rewrite_from.as_ref().map(Url::as_str))
273            .field("headers", &self.headers)
274            .field("query", &self.query)
275            .field("map_path", &self.map_path.is_some())
276            .finish()
277    }
278}
279
280impl<S> Layer<S> for ProviderLayer {
281    type Service = Provider<S>;
282
283    fn layer(&self, inner: S) -> Self::Service {
284        Provider {
285            inner,
286            provider: Arc::new(self.clone()),
287        }
288    }
289}
290
291/// The service of a [`ProviderLayer`].
292#[derive(Clone)]
293pub struct Provider<S> {
294    inner: S,
295    provider: Arc<ProviderLayer>,
296}
297
298impl<S: fmt::Debug> fmt::Debug for Provider<S> {
299    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
300        f.debug_struct("Provider")
301            .field("inner", &self.inner)
302            .field("provider", &self.provider)
303            .finish()
304    }
305}
306
307impl<S, B> Service<http::Request<B>> for Provider<S>
308where
309    S: Service<http::Request<B>>,
310    S::Error: Into<BoxError>,
311{
312    type Response = S::Response;
313    type Error = BoxError;
314    #[allow(clippy::type_complexity)]
315    type Future = Either<
316        Ready<Result<S::Response, BoxError>>,
317        futures_util::future::MapErr<S::Future, fn(S::Error) -> BoxError>,
318    >;
319
320    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
321        self.inner.poll_ready(cx).map_err(Into::into)
322    }
323
324    fn call(&mut self, mut request: http::Request<B>) -> Self::Future {
325        let applied = request.extensions().get::<Applied>().cloned();
326        match self.provider.rewrite(request.uri(), applied.as_ref()) {
327            Err(error) => Either::Left(future::ready(Err(error.into()))),
328            Ok(None) => Either::Right(
329                self.inner
330                    .call(request)
331                    .map_err(into_box_error::<S::Error> as fn(_) -> _),
332            ),
333            Ok(Some(uri)) => {
334                *request.uri_mut() = uri;
335                match &applied {
336                    None => request.headers_mut().extend(self.provider.headers.clone()),
337                    Some(applied) => {
338                        // The options' headers win: the layer leaves their names alone.
339                        let mut kept = HeaderMap::new();
340                        for (name, value) in &self.provider.headers {
341                            if !applied.has_header(name) {
342                                kept.append(name.clone(), value.clone());
343                            }
344                        }
345                        request.headers_mut().extend(kept);
346                    }
347                }
348                Either::Right(
349                    self.inner
350                        .call(request)
351                        .map_err(into_box_error::<S::Error> as fn(_) -> _),
352                )
353            }
354        }
355    }
356}
357
358fn into_box_error<E: Into<BoxError>>(error: E) -> BoxError {
359    error.into()
360}
361
362fn header<K, V>(name: K, value: V) -> Result<(HeaderName, HeaderValue), Error>
363where
364    HeaderName: TryFrom<K>,
365    <HeaderName as TryFrom<K>>::Error: Into<http::Error>,
366    HeaderValue: TryFrom<V>,
367    <HeaderValue as TryFrom<V>>::Error: Into<http::Error>,
368{
369    let name = HeaderName::try_from(name).map_err(|e| Error::builder(e.into()))?;
370    let value = HeaderValue::try_from(value).map_err(|e| Error::builder(e.into()))?;
371    Ok((name, value))
372}
373
374/// A base URL: absolute, with a host, without a query or fragment.
375pub(crate) fn parse_base(base_url: &str) -> Result<Url, Error> {
376    let url = Url::parse(base_url).map_err(Error::builder)?;
377    if url.cannot_be_a_base() || url.host_str().is_none() {
378        return Err(Error::builder(format!(
379            "not a base URL with a host: {base_url}"
380        )));
381    }
382    if url.query().is_some() || url.fragment().is_some() {
383        return Err(Error::builder(format!(
384            "a base URL has no query or fragment (set query parameters with `query`): {base_url}"
385        )));
386    }
387    Ok(url)
388}
389
390/// The path of a base URL without its trailing `/`: empty for the root.
391pub(crate) fn trimmed_path(base: &Url) -> &str {
392    base.path().trim_end_matches('/')
393}
394
395/// The path of `url` after `base`'s, when `url` is under `base`: the same scheme, host and port, and a path that starts
396/// with `base`'s at a segment boundary.
397pub(crate) fn rest_of_path(url: &Url, base: &Url) -> Option<String> {
398    if url.scheme() != base.scheme()
399        || url.host_str() != base.host_str()
400        || url.port_or_known_default() != base.port_or_known_default()
401    {
402        return None;
403    }
404    let rest = url.path().strip_prefix(trimmed_path(base))?;
405    (rest.is_empty() || rest.starts_with('/')).then(|| rest.to_owned())
406}
407
408#[cfg(test)]
409mod tests {
410    use super::*;
411
412    fn rewrite(layer: &ProviderLayer, uri: &str) -> String {
413        layer
414            .rewrite_uri(&uri.parse().unwrap())
415            .unwrap()
416            .to_string()
417    }
418
419    #[test]
420    fn moves_requests_under_the_declared_base() {
421        let layer = ProviderLayer::new("https://openrouter.ai/api/v1/")
422            .unwrap()
423            .rewrite_from("https://api.openai.com/v1")
424            .unwrap();
425        assert_eq!(
426            rewrite(&layer, "https://api.openai.com/v1/chat/completions"),
427            "https://openrouter.ai/api/v1/chat/completions"
428        );
429        assert_eq!(
430            rewrite(
431                &layer,
432                "https://api.openai.com:443/v1/models?limit=2&after=a%20b"
433            ),
434            "https://openrouter.ai/api/v1/models?limit=2&after=a%20b"
435        );
436        assert_eq!(
437            rewrite(&layer, "https://api.openai.com/v1"),
438            "https://openrouter.ai/api/v1"
439        );
440        // Not under the base: unchanged.
441        for uri in [
442            "https://api.openai.com/v10/models",
443            "https://api.openai.com/models",
444            "http://api.openai.com/v1/models",
445            "https://api.openai.com:8443/v1/models",
446            "https://files.example.com/v1/download",
447        ] {
448            assert_eq!(rewrite(&layer, uri), uri);
449        }
450    }
451
452    #[test]
453    fn a_request_already_at_the_provider_gets_its_path_mapped_once() {
454        let layer = ProviderLayer::new("http://localhost:11434/v1")
455            .unwrap()
456            .map_path(|path| format!("/x{path}"));
457        assert_eq!(
458            rewrite(&layer, "http://localhost:11434/v1/models"),
459            "http://localhost:11434/v1/x/models"
460        );
461        assert_eq!(
462            rewrite(&layer, "http://localhost:11435/v1/models"),
463            "http://localhost:11435/v1/models"
464        );
465    }
466
467    #[test]
468    fn root_bases() {
469        let layer = ProviderLayer::new("https://proxy.example")
470            .unwrap()
471            .rewrite_from("https://api.example.com/")
472            .unwrap();
473        assert_eq!(
474            rewrite(&layer, "https://api.example.com/items?a=1"),
475            "https://proxy.example/items?a=1"
476        );
477        assert_eq!(
478            rewrite(&layer, "https://api.example.com"),
479            "https://proxy.example/"
480        );
481        let layer = ProviderLayer::new("https://proxy.example/prefix")
482            .unwrap()
483            .rewrite_from("https://api.example.com")
484            .unwrap();
485        assert_eq!(
486            rewrite(&layer, "https://api.example.com/items"),
487            "https://proxy.example/prefix/items"
488        );
489    }
490
491    #[test]
492    fn query_parameters_replace_those_of_their_name() {
493        let layer = ProviderLayer::new("https://p.example/openai")
494            .unwrap()
495            .query(&[("api-version", "old")])
496            .unwrap()
497            .query(&[("api-version", "2024-10-21"), ("tag", "a"), ("tag", "b")])
498            .unwrap();
499        assert_eq!(
500            layer.query_pairs(),
501            [
502                ("api-version".to_owned(), "2024-10-21".to_owned()),
503                ("tag".to_owned(), "a".to_owned()),
504                ("tag".to_owned(), "b".to_owned()),
505            ]
506        );
507        assert_eq!(
508            rewrite(&layer, "https://p.example/openai/models"),
509            "https://p.example/openai/models?api-version=2024-10-21&tag=a&tag=b"
510        );
511        // Existing ones kept as sent, unless replaced.
512        assert_eq!(
513            rewrite(&layer, "https://p.example/openai/models?limit=2&q=a%20b"),
514            "https://p.example/openai/models?limit=2&q=a%20b&api-version=2024-10-21&tag=a&tag=b"
515        );
516        assert_eq!(
517            rewrite(
518                &layer,
519                "https://p.example/openai/models?api-version=x&limit=2&tag=z"
520            ),
521            "https://p.example/openai/models?limit=2&api-version=2024-10-21&tag=a&tag=b"
522        );
523    }
524
525    #[test]
526    fn invalid_input_is_a_builder_error() {
527        for base in [
528            "not a url",
529            "mailto:someone@example.com",
530            "https://p.example/v1?a=1",
531            "https://p.example/v1#top",
532        ] {
533            assert!(ProviderLayer::new(base).unwrap_err().is_builder(), "{base}");
534        }
535        let layer = ProviderLayer::new("https://p.example").unwrap();
536        assert!(layer.clone().header("x bad", "v").unwrap_err().is_builder());
537        assert!(layer
538            .clone()
539            .sensitive_header("x-key", "line\nbreak")
540            .unwrap_err()
541            .is_builder());
542        assert!(layer.query(&[("a", [1, 2])]).unwrap_err().is_builder());
543    }
544
545    #[test]
546    fn debug_hides_sensitive_headers() {
547        let layer = ProviderLayer::new("https://p.example")
548            .unwrap()
549            .sensitive_header("api-key", "secret-123")
550            .unwrap()
551            .header("x-title", "app")
552            .unwrap();
553        let debug = format!("{layer:?}");
554        assert!(!debug.contains("secret-123"), "{debug}");
555        assert!(debug.contains("app"), "{debug}");
556        assert!(layer.header_map()["api-key"].is_sensitive());
557        assert_eq!(layer.base_url(), "https://p.example/");
558    }
559}
timvw_rust_runtime/provider.rs

timvw_rust_runtime/
provider.rs
