Draft mid-streaming error retry

crates/uv/tests/it/network.rs

@@ -1,8 +1,13 @@
-use std::{env, io};
-
 use assert_fs::fixture::{ChildPath, FileWriteStr, PathChild};
 use http::StatusCode;
+use hyper::body::Bytes;
+use hyper::server::conn::http1;
+use hyper::service::service_fn;
+use hyper::{Method, Request, Response};
+use hyper_util::rt::TokioIo;
 use serde_json::json;
+use std::{convert::Infallible, env, io};
+use tokio::net::TcpListener;
 use wiremock::matchers::method;
 use wiremock::{Mock, MockServer, ResponseTemplate};
 
@@ -36,6 +41,41 @@ async fn io_error_server() -> (MockServer, String) {
     (server, mock_server_uri)
 }
 
+/// Creates a server that sends partial HTTP response data, then drops the connection.
+async fn mid_stream_io_error_server() -> (tokio::task::JoinHandle<()>, String) {
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = listener.local_addr().unwrap();
+
+    let server_task = tokio::spawn(async move {
+        // Handle multiple connections (for retries)
+        loop {
+            if let Ok((mut socket, _)) = listener.accept().await {
+                // Spawn a task for each connection to handle them concurrently
+                tokio::spawn(async move {
+                    // Read the incoming HTTP request (we don't parse it, just consume it)
+                    let mut buffer = [0; 1024];
+                    let _ = socket.read(&mut buffer).await;
+
+                    // Send a partial HTTP response - start with valid headers but incomplete body
+                    let partial_response = b"HTTP/1.1 200 OK\r\nContent-Length: 1000000\r\nContent-Type: application/octet-stream\r\n\r\n";
+                    let _ = socket.write_all(partial_response).await;
+
+                    // Send some initial data
+                    let partial_data = b"PK\x03\x04"; // Start of a ZIP file (wheel file)
+                    let _ = socket.write_all(partial_data).await;
+
+                    // Wait a bit then drop the connection abruptly
+                    tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
+                    drop(socket);
+                });
+            }
+        }
+    });
+
+    let server_uri = format!("http://{addr}");
+    (server_task, server_uri)
+}
+
 /// Check the simple index error message when the server returns HTTP status 500, a retryable error.
 #[tokio::test]
 async fn simple_http_500() {
@@ -195,6 +235,45 @@ async fn direct_url_io_error() {
     ");
 }
 
+/// Check the direct package URL error message when the server sends partial data then errors.
+#[tokio::test]
+async fn direct_url_mid_stream_io_error() {
+    let context = TestContext::new("3.12");
+
+    let (server_task, mock_server_uri) = mid_stream_io_error_server().await;
+
+    let tqdm_url = format!(
+        "{mock_server_uri}/packages/d0/30/dc54f88dd4a2b5dc8a0279bdd7270e735851848b762aeb1c1184ed1f6b14/tqdm-4.67.1-py3-none-any.whl"
+    );
+
+    tokio::task::spawn_blocking(move || {
+        let filters = vec![(mock_server_uri.as_str(), "[SERVER]")];
+        uv_snapshot!(filters, context
+            .pip_install()
+            .arg(format!("tqdm @ {tqdm_url}")), @r"
+    success: false
+    exit_code: 1
+    ----- stdout -----
+
+    ----- stderr -----
+      × Failed to download `tqdm @ [SERVER]/packages/d0/30/dc54f88dd4a2b5dc8a0279bdd7270e735851848b762aeb1c1184ed1f6b14/tqdm-4.67.1-py3-none-any.whl`
+      ├─▶ Request failed after 3 retries
+      ├─▶ Failed to read metadata: `[SERVER]/packages/d0/30/dc54f88dd4a2b5dc8a0279bdd7270e735851848b762aeb1c1184ed1f6b14/tqdm-4.67.1-py3-none-any.whl`
+      ├─▶ Failed to read from zip file
+      ├─▶ an upstream reader returned an error: error decoding response body
+      ├─▶ error decoding response body
+      ├─▶ request or response body error
+      ├─▶ error reading a body from connection
+      ╰─▶ end of file before message length reached
+    ");
+    })
+    .await
+    .unwrap();
+
+    // Clean up the server task
+    server_task.abort();
+}
+
 fn write_python_downloads_json(context: &TestContext, mock_server_uri: &String) -> ChildPath {
     let python_downloads_json = context.temp_dir.child("python_downloads.json");
     let interpreter = json!({