From 6cb62bc78142bf35544a9fe07620aa1ae75012bd Mon Sep 17 00:00:00 2001
From: nicodes <nicozessoules@gmail.com>
Date: Sat, 31 Oct 2026 11:21:05 +0100
Subject: [PATCH] fix: identify timeouts while reading job responses

---
 src/client.rs     | 17 ++++++++++++++++-
 tests/cli_test.py | 15 +++++++++++++--
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index 2039a96..fc6ae04 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -75,6 +75,7 @@ pub fn execute(
     if let Some(remaining) = remaining {
         request = request.timeout(remaining);
     }
+    let started = Instant::now();
     let mut response = match request.send() {
         Ok(response) => response,
         Err(error) if error.is_timeout() => return Err("request timeout".into()),
@@ -88,7 +89,21 @@ pub fn execute(
         .unwrap_or("")
         .to_owned();
     let mut bytes = Vec::new();
-    std::io::Read::take(&mut response, (512 << 20) + 1).read_to_end(&mut bytes)?;
+    if let Err(error) = std::io::Read::take(&mut response, (512 << 20) + 1).read_to_end(&mut bytes)
+    {
+        // The blocking response wraps body timeouts in io::Error; unlike send(),
+        // its Display text does not reliably include the underlying timeout.
+        if error.kind() == std::io::ErrorKind::TimedOut
+            || error
+                .get_ref()
+                .and_then(|source| source.downcast_ref::<reqwest::Error>())
+                .is_some_and(reqwest::Error::is_timeout)
+            || remaining.is_some_and(|budget| started.elapsed() >= budget)
+        {
+            return Err("request timeout".into());
+        }
+        return Err(error.into());
+    }
     if bytes.len() > 512 << 20 {
         return Err("API response exceeds 512 MiB".into());
     }
diff --git a/tests/cli_test.py b/tests/cli_test.py
index b46e7d6..205b809 100644
--- a/tests/cli_test.py
+++ b/tests/cli_test.py
@@ -10,6 +10,7 @@
 import subprocess
 import tempfile
 import threading
+import time
 import unittest
 from urllib.parse import parse_qs, urlsplit
 
@@ -18,7 +19,7 @@
 
 
 @contextmanager
-def api(responder=None):
+def api(responder=None, body_delay=0):
     calls=[]
     class Handler(BaseHTTPRequestHandler):
         def log_message(self,*_args):
@@ -33,7 +34,11 @@ def respond(self):
             self.send_header('Content-Type',media)
             self.send_header('Content-Length',str(len(data)))
             self.end_headers()
-            self.wfile.write(data)
+            time.sleep(body_delay)
+            try:
+                self.wfile.write(data)
+            except (BrokenPipeError, ConnectionResetError, ConnectionAbortedError):
+                pass  # A timed-out client is expected to close this connection.
         do_GET=respond
         do_POST=respond
         do_PATCH=respond
@@ -150,6 +155,12 @@ def test_timeout_and_api_failures_return_nonzero(self):
             result=self.request(url,'balance')
             self.assertEqual(result.returncode,1)
             self.assertIn('API error 500: fixture failed',result.stderr)
+    def test_timeout_after_response_headers_is_identified(self):
+        with api(lambda calls:(200,dict(job_id='job-1') if len(calls)==1 else dict(status='processing'),'application/json'), body_delay=.5) as (url,calls):
+            result=self.request(url,'generate-image-v2','--body-json','{}','--wait','--timeout','100ms')
+            self.assertNotEqual(result.returncode,0)
+            self.assertIn('timeout',result.stderr)
+            self.assertEqual(len(calls),2)
     def test_token_precedence_flag_environment_then_dotenv(self):
         (self.root/'.env').write_text('export PIXELLAB_API_KEY="file-token"\n')
         with api() as (url,calls):