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

timvw_rust_runtime/
provider.rs
