feat: add uv version to user agent (#2136)

## Summary Closes #1977 This allows us to send uv's version in the `uv-client` User Agent header. Here's how request headers look like to a server now: ``` ... Accept: application/vnd.pypi.simple.v1+json, application/vnd.pypi.simple.v1+html;q=0.2, text/html;q=0.01 User-Agent: uv/0.1.13 ... ``` ~~I went for a mix of Option 1 and 2 from #1977.~~ Open to alternative naming as well, not tied too strongly here to the names picked. ~~Another possibility for this new crate is that we can use it to consolidate metadata that exists across crates to ultimately be able to create linehaul information described in #1958, but I haven't looked into what those changes might look like.~~  ## Test Plan  Added initial tests in the new crate to exercise its public API and added a new test to uv-client to validate the headers using a 1-time disposable server.
astral-sh · Mar 4, 2024 · 93f5609 · 93f5609
1 parent fda6914
commit 93f5609
Show file tree

Hide file tree

Showing 8 changed files with 110 additions and 2 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/uv-client/Cargo.toml b/crates/uv-client/Cargo.toml
@@ -15,6 +15,7 @@ uv-auth = { path = "../uv-auth" }
 uv-cache = { path = "../uv-cache" }
 uv-fs = { path = "../uv-fs", features = ["tokio"] }
 uv-normalize = { path = "../uv-normalize" }
+uv-version = { path = "../uv-version" }
 uv-warnings = { path = "../uv-warnings" }
 pypi-types = { path = "../pypi-types" }
 
@@ -48,5 +49,6 @@ urlencoding = { workspace = true }
 
 [dev-dependencies]
 anyhow = { workspace = true }
+hyper = { version = "0.14.28", features = ["server", "http1"] }
 insta = { version = "1.35.1" }
 tokio = { workspace = true, features = ["fs", "macros"] }
diff --git a/crates/uv-client/src/registry_client.rs b/crates/uv-client/src/registry_client.rs
@@ -6,7 +6,6 @@ use std::str::FromStr;
 
 use async_http_range_reader::AsyncHttpRangeReader;
 use futures::{FutureExt, TryStreamExt};
-
 use http::HeaderMap;
 use reqwest::{Client, ClientBuilder, Response, StatusCode};
 use reqwest_retry::policies::ExponentialBackoff;
@@ -25,6 +24,7 @@ use pypi_types::{Metadata21, SimpleJson};
 use uv_auth::safe_copy_url_auth;
 use uv_cache::{Cache, CacheBucket, WheelCache};
 use uv_normalize::PackageName;
+use uv_version::version;
 use uv_warnings::warn_user_once;
 
 use crate::cached_client::CacheControl;
@@ -88,6 +88,9 @@ impl RegistryClientBuilder {
     }
 
     pub fn build(self) -> RegistryClient {
+        // Create user agent.
+        let user_agent_string = format!("uv/{}", version());
+
         // Timeout options, matching https://doc.rust-lang.org/nightly/cargo/reference/config.html#httptimeout
         // `UV_REQUEST_TIMEOUT` is provided for backwards compatibility with v0.1.6
         let default_timeout = 5 * 60;
@@ -108,7 +111,7 @@ impl RegistryClientBuilder {
         let client_raw = self.client.unwrap_or_else(|| {
             // Disallow any connections.
             let client_core = ClientBuilder::new()
-                .user_agent("uv")
+                .user_agent(user_agent_string)
                 .pool_max_idle_per_host(20)
                 .timeout(std::time::Duration::from_secs(timeout));
 

diff --git a/crates/uv-client/tests/user_agent_version.rs b/crates/uv-client/tests/user_agent_version.rs
@@ -0,0 +1,63 @@
+use anyhow::Result;
+use futures::future;
+use hyper::header::USER_AGENT;
+use hyper::server::conn::Http;
+use hyper::service::service_fn;
+use hyper::{Body, Request, Response};
+use tokio::net::TcpListener;
+
+use uv_cache::Cache;
+use uv_client::RegistryClientBuilder;
+use uv_version::version;
+
+#[tokio::test]
+async fn test_user_agent_has_version() -> Result<()> {
+    // Set up the TCP listener on a random available port
+    let listener = TcpListener::bind("127.0.0.1:0").await?;
+    let addr = listener.local_addr()?;
+
+    // Spawn the server loop in a background task
+    tokio::spawn(async move {
+        let svc = service_fn(move |req: Request<Body>| {
+            // Get User Agent Header and send it back in the response
+            let user_agent = req
+                .headers()
+                .get(USER_AGENT)
+                .and_then(|v| v.to_str().ok())
+                .map(|s| s.to_string())
+                .unwrap_or_default(); // Empty Default
+            future::ok::<_, hyper::Error>(Response::new(Body::from(user_agent)))
+        });
+        // Start Hyper Server
+        let (socket, _) = listener.accept().await.unwrap();
+        Http::new()
+            .http1_keep_alive(false)
+            .serve_connection(socket, svc)
+            .with_upgrades()
+            .await
+            .expect("Server Started");
+    });
+
+    // Initialize uv-client
+    let cache = Cache::temp()?;
+    let client = RegistryClientBuilder::new(cache).build();
+
+    // Send request to our dummy server
+    let res = client
+        .cached_client()
+        .uncached()
+        .get(format!("http://{addr}"))
+        .send()
+        .await?;
+
+    // Check the HTTP status
+    assert!(res.status().is_success());
+
+    // Check User Agent
+    let body = res.text().await?;
+
+    // Verify body matches regex
+    assert_eq!(body, format!("uv/{}", version()));
+
+    Ok(())
+}
diff --git a/crates/uv-version/Cargo.toml b/crates/uv-version/Cargo.toml
@@ -0,0 +1,15 @@
+[package]
+name = "uv-version"
+version = "0.1.14"
+edition = { workspace = true }
+rust-version = { workspace = true }
+homepage = { workspace = true }
+documentation = { workspace = true }
+repository = { workspace = true }
+authors = { workspace = true }
+license = { workspace = true }
+
+[lints]
+workspace = true
+
+[dependencies]
diff --git a/crates/uv-version/src/lib.rs b/crates/uv-version/src/lib.rs
@@ -0,0 +1,16 @@
+/// Return the application version.
+///
+/// This should be in sync with uv's version based on the Crate version.
+pub fn version() -> &'static str {
+    env!("CARGO_PKG_VERSION")
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_get_version() {
+        assert_eq!(version().to_string(), env!("CARGO_PKG_VERSION").to_string());
+    }
+}
diff --git a/crates/uv/Cargo.toml b/crates/uv/Cargo.toml
@@ -34,6 +34,7 @@ uv-interpreter = { path = "../uv-interpreter" }
 uv-normalize = { path = "../uv-normalize" }
 uv-resolver = { path = "../uv-resolver", features = ["clap"] }
 uv-traits = { path = "../uv-traits" }
+uv-version = { path = "../uv-version" }
 uv-virtualenv = { path = "../uv-virtualenv" }
 uv-warnings = { path = "../uv-warnings" }
 

diff --git a/pyproject.toml b/pyproject.toml
@@ -64,4 +64,5 @@ changelog_contributors = false
 version_files = [
   "README.md",
   "crates/uv/Cargo.toml",
+  "crates/uv-version/Cargo.toml",
 ]