Merge pull request #6 from gpu-mode/feat/streaming-response

Feat: streaming response

Author: S1ro1

Cargo.lock

@@ -20,5 +20,7 @@ serde_yaml = "0.9"
 webbrowser = "0.8"
 base64-url = "3.0.0"
 urlencoding = "2.1.3"
+bytes = "1.10.1"
+futures-util = "0.3.31"
 
 

Cargo.toml

@@ -6,6 +6,7 @@ use serde_json::Value;
 use std::env;
 use std::path::Path;
 use std::time::Duration;
+use tokio::io::AsyncWriteExt;
 
 use crate::models::{GpuItem, LeaderboardItem};
 
@@ -125,22 +126,105 @@ pub async fn submit_solution<P: AsRef<Path>>(
     let resp = client
         .post(&url)
         .multipart(form)
-        .timeout(Duration::from_secs(180))
+        .timeout(Duration::from_secs(300))
         .send()
         .await?;
 
     let status = resp.status();
     if !status.is_success() {
         let error_text = resp.text().await?;
-        return Err(anyhow!("Server returned status {}: {}", status, error_text));
+        let detail = serde_json::from_str::<Value>(&error_text)
+            .ok()
+            .and_then(|v| v.get("detail").and_then(|d| d.as_str()).map(str::to_string));
+
+        return Err(anyhow!(
+            "Server returned status {}: {}",
+            status,
+            detail.unwrap_or(error_text)
+        ));
     }
 
-    let result: Value = resp.json().await?;
-
-    let pretty_result = match result.get("results") {
-        Some(result_obj) => serde_json::to_string_pretty(result_obj)?,
-        None => return Err(anyhow!("Invalid response structure")),
-    };
-
-    Ok(pretty_result)
+    if resp
+        .headers()
+        .get(reqwest::header::CONTENT_TYPE)
+        .and_then(|v| v.to_str().ok())
+        .map_or(false, |s| s.starts_with("text/event-stream"))
+    {
+        let mut resp = resp;
+        let mut buffer = String::new();
+        let mut stderr = tokio::io::stderr();
+
+        while let Some(chunk) = resp.chunk().await? {
+            buffer.push_str(&String::from_utf8_lossy(&chunk));
+
+            while let Some(pos) = buffer.find("\n\n") {
+                let message_str = buffer.drain(..pos + 2).collect::<String>();
+                let mut event_type = None;
+                let mut data_json = None;
+
+                for line in message_str.lines() {
+                    if line.starts_with("event:") {
+                        event_type = Some(line["event:".len()..].trim());
+                    } else if line.starts_with("data:") {
+                        data_json = Some(line["data:".len()..].trim());
+                    }
+                }
+
+                if let (Some(event), Some(data)) = (event_type, data_json) {
+                    match event {
+                        "status" => (),
+                        "result" => {
+                            let result_val: Value = serde_json::from_str(data)?;
+                            let pretty_result = match result_val.get("results") {
+                                Some(result_obj) => serde_json::to_string_pretty(result_obj)?,
+                                None => {
+                                    return Err(anyhow!(
+                                        "Invalid 'result' event structure: missing 'results' field"
+                                    ))
+                                }
+                            };
+                            return Ok(pretty_result);
+                        }
+                        "error" => {
+                            let error_val: Value = serde_json::from_str(data)?;
+                            let detail = error_val
+                                .get("detail")
+                                .and_then(|d| d.as_str())
+                                .unwrap_or("Unknown server error");
+                            let status_code = error_val.get("status_code").and_then(|s| s.as_i64());
+                            let raw_error = error_val.get("raw_error").and_then(|e| e.as_str());
+
+                            let mut error_msg = format!("Server processing error: {}", detail);
+                            if let Some(sc) = status_code {
+                                error_msg.push_str(&format!(" (Status Code: {})", sc));
+                            }
+                            if let Some(re) = raw_error {
+                                error_msg.push_str(&format!(" | Raw Error: {}", re));
+                            }
+
+                            return Err(anyhow!(error_msg));
+                        }
+                        _ => {
+                            stderr
+                                .write_all(
+                                    format!("Ignoring unknown SSE event: {}\n", event).as_bytes(),
+                                )
+                                .await?;
+                            stderr.flush().await?;
+                        }
+                    }
+                }
+            }
+        }
+        Err(anyhow!(
+            "Stream ended unexpectedly without a final result or error event."
+        ))
+    } else {
+        let result: Value = resp.json().await?;
+        let pretty_result = match result.get("results") {
+            Some(result_obj) => serde_json::to_string_pretty(result_obj)?,
+            None => return Err(anyhow!("Invalid non-streaming response structure")),
+        };
+        Ok(pretty_result)
+    }
 }