[API server][Rust frontend][Bug fix] Delete a refused request's uploads on /generate and /v1/images/edits

mstar/api_server/entrypoint.py

@@ -1519,17 +1519,21 @@ def _ws_modalities(value, key: str) -> list[str] | None:
     raise ValueError(f"{key} must be a comma-separated string or a list of strings")
 
 
+def _remove_uploads(file_paths: dict[str, list[str]]) -> None:
+    for ps in file_paths.values():
+        for p in ps:
+            try:
+                Path(p).unlink(missing_ok=True)
+            except OSError:
+                pass
+
+
 def _schedule_upload_cleanup(file_paths: dict[str, list[str]], delay_s: float = 60.0) -> None:
     """Delete a request's uploaded files once the preprocess worker has had
     time to read them. Shared by the HTTP and WebSocket routes."""
     def _cleanup(paths: dict[str, list[str]]) -> None:
         time.sleep(delay_s)
-        for ps in paths.values():
-            for p in ps:
-                try:
-                    Path(p).unlink(missing_ok=True)
-                except OSError:
-                    pass
+        _remove_uploads(paths)
     threading.Thread(target=_cleanup, args=(file_paths,), daemon=True).start()
 
 
@@ -1747,6 +1751,9 @@ async def generate(
         for f in files:
             modality = _detect_modality(f.filename or "")
             if modality == "unknown":
+                # a refused request is never submitted, so the deferred cleanup
+                # below never sees what it saved
+                _remove_uploads(file_paths)
                 raise HTTPException(
                     status_code=400,
                     detail=f"Cannot determine modality for file: {f.filename}",
@@ -1783,6 +1790,7 @@ async def generate(
         in_mods = [p.modality for p in parts]
 
     if "video_frame" in in_mods:
+        _remove_uploads(file_paths)
         raise HTTPException(
             status_code=400,
             detail="'video_frame' is an output-only modality",
@@ -1791,12 +1799,14 @@ async def generate(
     try:
         parsed_kwargs = json.loads(model_kwargs) if model_kwargs else None
     except json.JSONDecodeError as e:
+        _remove_uploads(file_paths)
         raise HTTPException(
             status_code=400,
             detail="model_kwargs must be valid JSON",
         ) from e
 
     if parsed_kwargs is not None and not isinstance(parsed_kwargs, dict):
+        _remove_uploads(file_paths)
         raise HTTPException(
             status_code=400,
             detail="model_kwargs must be a JSON object",

mstar/api_server/openai/serving_images.py

@@ -87,7 +87,14 @@ async def create_image_edit(
     image_path = upload_dir / f"{uuid4().hex}{ext}"
     image_path.write_bytes(image_bytes)
 
-    args = adapter.image_edit_to_request(prompt, str(image_path), model_kwargs)
+    try:
+        args = adapter.image_edit_to_request(prompt, str(image_path), model_kwargs)
+    except Exception as exc:
+        image_path.unlink(missing_ok=True)
+        if isinstance(exc, NotImplementedError):
+            # a model with images but no editing: the client's error, a 400
+            raise ValueError(str(exc)) from exc
+        raise
     request_id = rid("img")
 
     api.submit_request(

rust/server/src/main.rs

@@ -891,7 +891,10 @@ async fn images_edits(State(st): State<AppState>, mut mp: Multipart) -> Response
 
     let args = match adapter.image_edit_to_request(&prompt, &image_path.to_string_lossy(), extra) {
         Ok(a) => a,
-        Err(e) => return error(400, &e, "invalid_request_error"),
+        Err(e) => {
+            let _ = std::fs::remove_file(&image_path);
+            return error(400, &e, "invalid_request_error");
+        }
     };
     schedule_upload_cleanup(&st.upload_dir, &args);
     let request_id = rid("img");
@@ -946,7 +949,7 @@ async fn generate(State(st): State<AppState>, req: axum::extract::Request) -> Re
     let mut tokenize = false;
     let mut mk_raw: Option<String> = None;
     let mut request_id: Option<String> = None;
-    let mut file_paths: BTreeMap<String, Vec<String>> = BTreeMap::new();
+    let mut uploads = SavedUploads::new(&st.upload_dir);
 
     // The FastAPI endpoint this replaces uses `Form(...)` fields, which
     // Starlette parses from EITHER `multipart/form-data` (with file uploads)
@@ -987,7 +990,7 @@ async fn generate(State(st): State<AppState>, req: axum::extract::Request) -> Re
         }
         return generate_finish(
             st, text, in_mods_raw, out_mods_raw, streaming, tokenize, mk_raw,
-            request_id, file_paths,
+            request_id, uploads,
         )
         .await;
     }
@@ -1045,7 +1048,8 @@ async fn generate(State(st): State<AppState>, req: axum::extract::Request) -> Re
             if let Err(e) = std::fs::write(&save_path, &bytes) {
                 return detail_error(500, &e.to_string());
             }
-            file_paths
+            uploads
+                .paths
                 .entry(modality.to_string())
                 .or_default()
                 .push(save_path.to_string_lossy().into_owned());
@@ -1081,11 +1085,36 @@ async fn generate(State(st): State<AppState>, req: axum::extract::Request) -> Re
 
     generate_finish(
         st, text, in_mods_raw, out_mods_raw, streaming, tokenize, mk_raw,
-        request_id, file_paths,
+        request_id, uploads,
     )
     .await
 }
 
+/// The files a `/generate` request saved under the upload directory, until
+/// the request takes them with [`SavedUploads::keep`]. A refusal drops them
+/// unkept, which deletes them: the deferred cleanup only covers submitted
+/// requests.
+struct SavedUploads {
+    upload_dir: PathBuf,
+    paths: BTreeMap<String, Vec<String>>,
+}
+
+impl SavedUploads {
+    fn new(upload_dir: &std::path::Path) -> Self {
+        Self { upload_dir: upload_dir.to_path_buf(), paths: BTreeMap::new() }
+    }
+
+    fn keep(mut self) -> BTreeMap<String, Vec<String>> {
+        std::mem::take(&mut self.paths)
+    }
+}
+
+impl Drop for SavedUploads {
+    fn drop(&mut self) {
+        remove_uploads(&self.upload_dir, &self.paths);
+    }
+}
+
 /// Shared tail of `/generate` once the form fields are parsed (from either
 /// multipart or urlencoded): build the request and stream NDJSON.
 #[allow(clippy::too_many_arguments)]
@@ -1098,7 +1127,7 @@ async fn generate_finish(
     tokenize: bool,
     mk_raw: Option<String>,
     request_id: Option<String>,
-    file_paths: BTreeMap<String, Vec<String>>,
+    uploads: SavedUploads,
 ) -> Response {
     let output_modalities: Vec<String> = out_mods_raw
         .split(',')
@@ -1114,15 +1143,14 @@ async fn generate_finish(
             .map(str::to_string)
             .collect(),
         None => {
-            let mut m: Vec<String> = file_paths.keys().cloned().collect();
+            let mut m: Vec<String> = uploads.paths.keys().cloned().collect();
             if text.as_ref().map(|t| !t.is_empty()).unwrap_or(false) {
                 m.push("text".to_string());
             }
             m
         }
     };
     if let Some(message) = video_frame_refusal(&input_modalities, &output_modalities, streaming) {
-        remove_uploads(&st.upload_dir, &file_paths);
         return detail_error(400, message);
     }
     let model_kwargs: Map<String, Value> = match mk_raw {
@@ -1150,7 +1178,7 @@ async fn generate_finish(
     let args = SubmitArgs {
         text,
         tokens,
-        file_paths,
+        file_paths: uploads.keep(),
         input_modalities,
         output_modalities,
         model_kwargs,
@@ -1410,9 +1438,11 @@ mod tests {
         st: AppState, input: Option<&str>, output: &str, streaming: bool,
         file_paths: BTreeMap<String, Vec<String>>,
     ) -> (StatusCode, Vec<u8>) {
+        let mut uploads = SavedUploads::new(&st.upload_dir);
+        uploads.paths = file_paths;
         let resp = generate_finish(
             st, Some("hi".into()), input.map(str::to_string), output.into(), streaming,
-            false, None, None, file_paths,
+            false, None, None, uploads,
         )
         .await;
         let (status, _, data) = body(resp).await;
@@ -1453,6 +1483,91 @@ mod tests {
         let _ = std::fs::remove_dir_all(&st.upload_dir);
     }
 
+    fn multipart_request(uri: &str, parts: &[(&str, Option<&str>, &[u8])]) -> axum::extract::Request {
+        let mut payload = Vec::new();
+        for (name, filename, data) in parts {
+            payload.extend_from_slice(b"--XBOUNDARY\r\n");
+            let disposition = match filename {
+                Some(f) => format!("Content-Disposition: form-data; name=\"{name}\"; filename=\"{f}\"\r\n\r\n"),
+                None => format!("Content-Disposition: form-data; name=\"{name}\"\r\n\r\n"),
+            };
+            payload.extend_from_slice(disposition.as_bytes());
+            payload.extend_from_slice(data);
+            payload.extend_from_slice(b"\r\n");
+        }
+        payload.extend_from_slice(b"--XBOUNDARY--\r\n");
+        axum::http::Request::builder()
+            .method("POST")
+            .uri(uri)
+            .header(CONTENT_TYPE, "multipart/form-data; boundary=XBOUNDARY")
+            .body(Body::from(payload))
+            .unwrap()
+    }
+
+    #[tokio::test]
+    async fn every_refusal_after_the_form_leaves_no_upload_behind() {
+        let mut st = state(None);
+        st.upload_dir = socket_dir();
+        for (field, value, detail) in [
+            ("model_kwargs", "[1]", "model_kwargs must be a JSON object"),
+            ("model_kwargs", "{", "model_kwargs JSON: EOF while parsing an object at line 1 column 1"),
+            ("tokenize", "true", "tokenize=true requires the server to be started with MSTAR_TOKENIZER"),
+            ("input_modalities", "video_frame", "'video_frame' is an output-only modality"),
+        ] {
+            let req = multipart_request("/generate", &[
+                ("files", Some("clip.wav"), b"RIFF"),
+                (field, None, value.as_bytes()),
+            ]);
+
+            let (status, _, data) = body(generate(State(st.clone()), req).await).await;
+
+            let body: Value = serde_json::from_slice(&data).unwrap();
+            assert_eq!((status, body), (StatusCode::BAD_REQUEST, json!({"detail": detail})), "{field}={value}");
+            assert_eq!(std::fs::read_dir(&st.upload_dir).unwrap().count(), 0, "{field}={value}");
+        }
+        let _ = std::fs::remove_dir_all(&st.upload_dir);
+    }
+
+    #[tokio::test]
+    async fn a_file_refused_mid_form_takes_the_saved_ones_with_it() {
+        let mut st = state(None);
+        st.upload_dir = socket_dir();
+        let req = multipart_request("/generate", &[
+            ("files", Some("clip.wav"), b"RIFF"),
+            ("files", Some("notes.xyz"), b"?"),
+        ]);
+
+        let (status, _, data) = body(generate(State(st.clone()), req).await).await;
+
+        let body: Value = serde_json::from_slice(&data).unwrap();
+        assert_eq!(
+            (status, body),
+            (StatusCode::BAD_REQUEST, json!({"detail": "Cannot determine modality for file: notes.xyz"}))
+        );
+        assert_eq!(std::fs::read_dir(&st.upload_dir).unwrap().count(), 0);
+        let _ = std::fs::remove_dir_all(&st.upload_dir);
+    }
+
+    #[tokio::test]
+    async fn an_edit_the_model_cannot_do_leaves_no_upload_behind() {
+        let mut st = state(None);
+        st.upload_dir = socket_dir();
+        st.python_adapter = "Cosmos3EdgeAdapter".into();
+        let req = multipart_request("/v1/images/edits", &[
+            ("image", Some("cat.png"), b"\x89PNG"),
+            ("prompt", None, b"a hat"),
+        ]);
+        let mp = Multipart::from_request(req, &st).await.unwrap();
+
+        let (status, _, data) = body(images_edits(State(st.clone()), mp).await).await;
+
+        let body: Value = serde_json::from_slice(&data).unwrap();
+        assert_eq!(status, StatusCode::BAD_REQUEST);
+        assert_eq!(body["error"]["message"], "image editing is not supported by this model");
+        assert_eq!(std::fs::read_dir(&st.upload_dir).unwrap().count(), 0);
+        let _ = std::fs::remove_dir_all(&st.upload_dir);
+    }
+
     #[test]
     fn max_images_parses_as_python_does() {
         for (raw, cap) in [

test/modular/test_generate_route.py

@@ -112,3 +112,32 @@ def submit_request(self, **kwargs):
 
     assert response.status_code == 500
     assert response.json()["detail"].startswith("downstream failure")
+
+
+class _UploadServer(_RecordingServer):
+    def __init__(self, upload_dir):
+        super().__init__()
+        self.upload_dir = upload_dir
+
+
+@pytest.mark.parametrize("files,fields,detail", [
+    ([("clip.wav", b"RIFF"), ("notes.xyz", b"?")], {}, "Cannot determine modality for file: notes.xyz"),
+    ([("clip.wav", b"RIFF")], {"input_modalities": "video_frame"}, "'video_frame' is an output-only modality"),
+    ([("clip.wav", b"RIFF")], {"model_kwargs": "{"}, "model_kwargs must be valid JSON"),
+    ([("clip.wav", b"RIFF")], {"model_kwargs": "[1]"}, "model_kwargs must be a JSON object"),
+], ids=["unknown-file", "video-frame-input", "bad-json", "not-an-object"])
+def test_a_refused_generate_leaves_no_upload_behind(monkeypatch, tmp_path, files, fields, detail):
+    """A request refused after its files are saved is never submitted, so the
+    deferred cleanup never sees them; the refusal deletes them."""
+    server = _UploadServer(tmp_path)
+    monkeypatch.setattr(entrypoint, "api_server", server)
+
+    response = TestClient(entrypoint.app).post(
+        "/generate",
+        data={"text": "hi", **fields},
+        files=[("files", (name, data, "application/octet-stream")) for name, data in files],
+    )
+
+    assert (response.status_code, response.json()) == (400, {"detail": detail})
+    assert server.streaming is None
+    assert list(tmp_path.iterdir()) == []

test/modular/test_openai_router.py

@@ -645,6 +645,22 @@ def reject(**kw):
     assert resp.status_code == 400
 
 
+@pytest.mark.parametrize("model", ["cosmos3", "cosmos3_edge"])
+def test_an_edit_the_model_cannot_do_is_a_400_and_leaves_no_upload(client_and_stub, model):
+    """A model with image generation but no editing refuses the edit as the
+    Rust frontend does, and the saved image goes with the refusal."""
+    client, stub = client_and_stub
+    stub.model_name = model
+    resp = client.post(
+        "/v1/images/edits",
+        files={"image": ("in.png", b"\x89PNGinput", "image/png")},
+        data={"prompt": "make it neon"},
+    )
+    assert resp.status_code == 400
+    assert resp.json()["error"]["message"] == "image editing is not supported by this model"
+    assert stub.last_submit is None and not any(stub.upload_dir.iterdir())
+
+
 def _refuse(**kw):
     raise ValueError("size must be 'WxH'")