From 34506f9cec624ddc41f7d3f576e80bfde4be2e01 Mon Sep 17 00:00:00 2001 From: asahoo Date: Tue, 7 Jul 2026 16:09:56 -0500 Subject: [PATCH 1/3] fix: delete coglet-managed output files after upload Coglet creates output files (IOBase writes, oversized JSON spills) in /tmp/coglet/predictions/{id}/outputs/ but never deletes them. This can cause stale-output bugs where models return outputs from previous predictions if an overwrite fails (issue #1434). Add a managed bool field to the FileOutput IPC protocol message so the orchestrator knows which files are safe to delete (coglet-created) vs user-authored Path outputs (returned by reference, must not be deleted). The orchestrator deletes managed files immediately after reading their bytes into memory. A backstop remove_dir_all in remove_prediction cleans up the entire prediction directory to catch files from aborted uploads or cancelled predictions. Closes #1434 --- crates/coglet-python/src/predictor.rs | 4 +- crates/coglet/src/bridge/protocol.rs | 37 +++++++++++++++++ ...__slot_file_output_managed_serializes.snap | 13 ++++++ ...slot_file_output_oversized_serializes.snap | 12 ++++++ ...slot_file_output_unmanaged_serializes.snap | 12 ++++++ crates/coglet/src/orchestrator.rs | 7 +++- crates/coglet/src/service.rs | 40 +++++++++++++++++-- crates/coglet/src/worker.rs | 13 +++++- 8 files changed, 130 insertions(+), 8 deletions(-) create mode 100644 crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_file_output_managed_serializes.snap create mode 100644 crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_file_output_oversized_serializes.snap create mode 100644 crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_file_output_unmanaged_serializes.snap diff --git a/crates/coglet-python/src/predictor.rs b/crates/coglet-python/src/predictor.rs index 6142e48e24..69de338bb3 100644 --- a/crates/coglet-python/src/predictor.rs +++ b/crates/coglet-python/src/predictor.rs @@ -195,7 +195,7 @@ fn send_output_item( .and_then(|p| p.extract()) .map_err(|e| PredictionError::Failed(format!("Failed to get fspath: {}", e)))?; slot_sender - .send_file_output(std::path::PathBuf::from(path_str), None) + .send_file_output(std::path::PathBuf::from(path_str), None, false) .map_err(|e| PredictionError::Failed(format!("Failed to send file output: {}", e)))?; return Ok(()); } @@ -954,7 +954,7 @@ impl PythonPredictor { .and_then(|p| p.extract()) .map_err(|e| PredictionError::Failed(format!("Failed to get fspath: {}", e)))?; slot_sender - .send_file_output(std::path::PathBuf::from(path_str), None) + .send_file_output(std::path::PathBuf::from(path_str), None, false) .map_err(|e| { PredictionError::Failed(format!("Failed to send file output: {}", e)) })?; diff --git a/crates/coglet/src/bridge/protocol.rs b/crates/coglet/src/bridge/protocol.rs index 0ef46a4e68..68d30e828c 100644 --- a/crates/coglet/src/bridge/protocol.rs +++ b/crates/coglet/src/bridge/protocol.rs @@ -327,6 +327,10 @@ pub enum SlotResponse { /// Explicit MIME type from the predictor. Falls back to mime_guess when None. #[serde(skip_serializing_if = "Option::is_none")] mime_type: Option, + /// True if Coglet created this file (IOBase write or oversized spill). + /// False for user-authored Path outputs — must not be deleted. + #[serde(default)] + managed: bool, }, /// Streaming output chunk for generator and iterator output. @@ -592,6 +596,39 @@ mod tests { insta::assert_json_snapshot!(resp); } + #[test] + fn slot_file_output_managed_serializes() { + let resp = SlotResponse::FileOutput { + filename: "/tmp/coglet/predictions/pred_123/outputs/0.png".to_string(), + kind: FileOutputKind::FileType, + mime_type: Some("image/png".to_string()), + managed: true, + }; + insta::assert_json_snapshot!(resp); + } + + #[test] + fn slot_file_output_unmanaged_serializes() { + let resp = SlotResponse::FileOutput { + filename: "/home/user/model/output.wav".to_string(), + kind: FileOutputKind::FileType, + mime_type: None, + managed: false, + }; + insta::assert_json_snapshot!(resp); + } + + #[test] + fn slot_file_output_oversized_serializes() { + let resp = SlotResponse::FileOutput { + filename: "/tmp/coglet/predictions/pred_123/outputs/spill_abc.json".to_string(), + kind: FileOutputKind::Oversized, + mime_type: None, + managed: true, + }; + insta::assert_json_snapshot!(resp); + } + #[test] fn slot_metric_replace_serializes() { let resp = SlotResponse::Metric { diff --git a/crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_file_output_managed_serializes.snap b/crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_file_output_managed_serializes.snap new file mode 100644 index 0000000000..9acbbb9cd5 --- /dev/null +++ b/crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_file_output_managed_serializes.snap @@ -0,0 +1,13 @@ +--- +source: coglet/src/bridge/protocol.rs +expression: resp +--- +{ + "type": "file_output", + "filename": "/tmp/coglet/predictions/pred_123/outputs/0.png", + "kind": { + "type": "file_type" + }, + "mime_type": "image/png", + "managed": true +} diff --git a/crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_file_output_oversized_serializes.snap b/crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_file_output_oversized_serializes.snap new file mode 100644 index 0000000000..447bfeefab --- /dev/null +++ b/crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_file_output_oversized_serializes.snap @@ -0,0 +1,12 @@ +--- +source: coglet/src/bridge/protocol.rs +expression: resp +--- +{ + "type": "file_output", + "filename": "/tmp/coglet/predictions/pred_123/outputs/spill_abc.json", + "kind": { + "type": "oversized" + }, + "managed": true +} diff --git a/crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_file_output_unmanaged_serializes.snap b/crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_file_output_unmanaged_serializes.snap new file mode 100644 index 0000000000..248622d623 --- /dev/null +++ b/crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_file_output_unmanaged_serializes.snap @@ -0,0 +1,12 @@ +--- +source: coglet/src/bridge/protocol.rs +expression: resp +--- +{ + "type": "file_output", + "filename": "/home/user/model/output.wav", + "kind": { + "type": "file_type" + }, + "managed": false +} diff --git a/crates/coglet/src/orchestrator.rs b/crates/coglet/src/orchestrator.rs index b49ba1fc5e..3a8994bf9b 100644 --- a/crates/coglet/src/orchestrator.rs +++ b/crates/coglet/src/orchestrator.rs @@ -1072,7 +1072,7 @@ async fn run_event_loop( predictions.remove(&slot_id); } } - Ok(SlotResponse::FileOutput { filename, kind, mime_type }) => { + Ok(SlotResponse::FileOutput { filename, kind, mime_type, managed }) => { tracing::debug!(%slot_id, %filename, ?kind, "FileOutput received"); let bytes = match std::fs::read(&filename) { Ok(b) => b, @@ -1081,6 +1081,11 @@ async fn run_event_loop( continue; } }; + if managed + && let Err(e) = std::fs::remove_file(&filename) + { + tracing::debug!(%slot_id, %filename, error = %e, "Failed to delete managed output file"); + } match kind { FileOutputKind::Oversized => { let output: serde_json::Value = match serde_json::from_slice(&bytes) { diff --git a/crates/coglet/src/service.rs b/crates/coglet/src/service.rs index 460a09e09a..0e6c4b1224 100644 --- a/crates/coglet/src/service.rs +++ b/crates/coglet/src/service.rs @@ -657,8 +657,7 @@ impl PredictionService { .await; // Create per-prediction dirs for file-based inputs/outputs - let prediction_dir = - std::path::PathBuf::from("/tmp/coglet/predictions").join(&prediction_id); + let prediction_dir = prediction_dir_for(&prediction_id); let output_dir = prediction_dir.join("outputs"); let input_dir = prediction_dir.join("inputs"); std::fs::create_dir_all(&output_dir) @@ -797,9 +796,20 @@ impl PredictionService { } } - /// Remove a prediction from the DashMap after completion. + /// Remove a prediction from the DashMap after completion and clean up + /// its on-disk prediction directory as a backstop for any output files + /// that were not deleted individually (e.g. aborted uploads, cancelled + /// predictions). pub fn remove_prediction(&self, id: &str) { self.predictions.remove(id); + let dir = prediction_dir_for(id); + match std::fs::remove_dir_all(&dir) { + Ok(()) => {} + Err(e) if e.kind() == std::io::ErrorKind::NotFound => {} + Err(e) => { + tracing::debug!(prediction_id = %id, error = %e, "Failed to remove prediction dir") + } + } } pub fn trigger_shutdown(&self) { @@ -811,6 +821,10 @@ impl PredictionService { } } +fn prediction_dir_for(id: &str) -> std::path::PathBuf { + std::path::PathBuf::from("/tmp/coglet/predictions").join(id) +} + fn spawn_orchestrator_cancel(orch: Arc, id: String) { let Ok(handle) = tokio::runtime::Handle::try_current() else { tracing::warn!(prediction_id = %id, "No tokio runtime available to cancel prediction"); @@ -1723,6 +1737,26 @@ mod tests { assert!(!svc.prediction_exists("test-remove")); } + #[test] + fn remove_prediction_deletes_prediction_dir() { + let svc = PredictionService::new_no_pool(); + let dir = prediction_dir_for("test-dir-cleanup"); + std::fs::create_dir_all(dir.join("outputs")).unwrap(); + std::fs::create_dir_all(dir.join("inputs")).unwrap(); + std::fs::write(dir.join("outputs").join("0.png"), b"fake").unwrap(); + + svc.remove_prediction("test-dir-cleanup"); + + assert!(!dir.exists(), "prediction dir should be removed"); + } + + #[test] + fn remove_prediction_missing_dir_is_ok() { + let svc = PredictionService::new_no_pool(); + // Should not panic or error when the prediction dir doesn't exist + svc.remove_prediction("nonexistent-prediction"); + } + #[test] fn build_slot_request_small_input_inline() { let dir = tempfile::tempdir().unwrap(); diff --git a/crates/coglet/src/worker.rs b/crates/coglet/src/worker.rs index c59a26cc7a..bc0645b62a 100644 --- a/crates/coglet/src/worker.rs +++ b/crates/coglet/src/worker.rs @@ -201,14 +201,21 @@ impl SlotSender { ) -> io::Result<()> { let path = self.next_output_path(extension); std::fs::write(&path, data)?; - self.send_file_output(path, mime_type) + self.send_file_output(path, mime_type, true) } /// Send a file-typed output (e.g. Path, File return types). /// /// The file is already on disk at `path` — we just send the path reference. /// `mime_type` is an explicit MIME type; when None the parent guesses from extension. - pub fn send_file_output(&self, path: PathBuf, mime_type: Option) -> io::Result<()> { + /// `managed` is true when Coglet created the file (safe to delete after consumption); + /// false for user-authored Path outputs. + pub fn send_file_output( + &self, + path: PathBuf, + mime_type: Option, + managed: bool, + ) -> io::Result<()> { let filename = path .to_str() .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "non-UTF-8 path"))? @@ -217,6 +224,7 @@ impl SlotSender { filename, kind: FileOutputKind::FileType, mime_type, + managed, }; self.tx .send(msg) @@ -265,6 +273,7 @@ fn build_output_message( filename, kind: FileOutputKind::Oversized, mime_type: None, + managed: true, }) } else { Ok(SlotResponse::OutputChunk { output, index }) From a6119b204be570b227607f57dc46f265fe3f47ba Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Mon, 13 Jul 2026 11:13:26 -0500 Subject: [PATCH 2/3] fix: safe ownership for returned paths and prediction ID hardening Addresses review feedback on PR #3096 by validating prediction IDs, rejecting duplicate active IDs, hardening remove_dir_all, and safely handling returned cog.Path files via the managed output directory. --- crates/coglet-python/src/predictor.rs | 8 +- crates/coglet/src/service.rs | 195 +++++++++++++++++++-- crates/coglet/src/transport/http/routes.rs | 195 +++++++++++++++++---- crates/coglet/src/worker.rs | 88 ++++++++++ docs/llms.txt | 3 +- docs/python.md | 3 +- 6 files changed, 445 insertions(+), 47 deletions(-) diff --git a/crates/coglet-python/src/predictor.rs b/crates/coglet-python/src/predictor.rs index 69de338bb3..6fc6a335e1 100644 --- a/crates/coglet-python/src/predictor.rs +++ b/crates/coglet-python/src/predictor.rs @@ -189,13 +189,15 @@ fn send_output_item( .map_err(|e| PredictionError::Failed(format!("Failed to get io.IOBase: {}", e)))?; if item.is_instance(&pathlike).unwrap_or(false) { - // Path output — file already on disk, send path reference + // Path output — file already on disk. send_user_file_output copies + // external paths into the managed output directory and deletes them + // after upload, so Coglet never deletes arbitrary user-owned files. let path_str: String = item .call_method0("__fspath__") .and_then(|p| p.extract()) .map_err(|e| PredictionError::Failed(format!("Failed to get fspath: {}", e)))?; slot_sender - .send_file_output(std::path::PathBuf::from(path_str), None, false) + .send_user_file_output(std::path::PathBuf::from(path_str), None) .map_err(|e| PredictionError::Failed(format!("Failed to send file output: {}", e)))?; return Ok(()); } @@ -954,7 +956,7 @@ impl PythonPredictor { .and_then(|p| p.extract()) .map_err(|e| PredictionError::Failed(format!("Failed to get fspath: {}", e)))?; slot_sender - .send_file_output(std::path::PathBuf::from(path_str), None, false) + .send_user_file_output(std::path::PathBuf::from(path_str), None) .map_err(|e| { PredictionError::Failed(format!("Failed to send file output: {}", e)) })?; diff --git a/crates/coglet/src/service.rs b/crates/coglet/src/service.rs index 0e6c4b1224..0b6a7c723f 100644 --- a/crates/coglet/src/service.rs +++ b/crates/coglet/src/service.rs @@ -51,6 +51,10 @@ pub enum CreatePredictionError { NotReady, #[error("At capacity (no slots available)")] AtCapacity, + #[error("A prediction with this ID already exists")] + DuplicateId, + #[error("Invalid prediction ID: {0}")] + InvalidId(&'static str), } const MAX_STREAM_SUBSCRIBERS: usize = STREAM_CHANNEL_CAPACITY; @@ -518,6 +522,8 @@ impl PredictionService { webhook: Option, cancel_on_stream_drop: bool, ) -> Result<(PredictionHandle, UnregisteredPredictionSlot), CreatePredictionError> { + validate_prediction_id(&id).map_err(CreatePredictionError::InvalidId)?; + let health = *self.health.read().await; if health != Health::Ready { return Err(CreatePredictionError::NotReady); @@ -537,16 +543,21 @@ impl PredictionService { let slot = PredictionSlot::new(prediction, permit, idle_rx); let prediction_arc = slot.prediction(); - // Register in DashMap — this is the single source of truth - self.predictions.insert( - id.clone(), - PredictionEntry { - prediction: prediction_arc, - cancel_token: cancel_token.clone(), - input, - cancel_on_stream_drop, - }, - ); + // Register atomically in DashMap — duplicate active IDs are rejected to + // prevent one prediction from clobbering or deleting another's state. + match self.predictions.entry(id.clone()) { + dashmap::Entry::Vacant(entry) => { + entry.insert(PredictionEntry { + prediction: prediction_arc, + cancel_token: cancel_token.clone(), + input, + cancel_on_stream_drop, + }); + } + dashmap::Entry::Occupied(_) => { + return Err(CreatePredictionError::DuplicateId); + } + } let handle = PredictionHandle { id, cancel_token }; @@ -802,7 +813,18 @@ impl PredictionService { /// predictions). pub fn remove_prediction(&self, id: &str) { self.predictions.remove(id); + + // Defense in depth: validate the ID and ensure we only delete a direct + // child of the prediction root, even though the HTTP layer also validates. + if let Err(e) = validate_prediction_id(id) { + tracing::warn!(prediction_id = %id, error = e, "Refusing to remove prediction dir: invalid ID"); + return; + } let dir = prediction_dir_for(id); + if !is_direct_prediction_child(&dir) { + tracing::warn!(prediction_id = %id, dir = %dir.display(), "Refusing to remove prediction dir: not a direct child of prediction root"); + return; + } match std::fs::remove_dir_all(&dir) { Ok(()) => {} Err(e) if e.kind() == std::io::ErrorKind::NotFound => {} @@ -821,8 +843,48 @@ impl PredictionService { } } +const PREDICTION_ROOT: &str = "/tmp/coglet/predictions"; + +/// Validate a client-supplied prediction ID. +/// +/// IDs must be non-empty, must not contain path separators or traversal +/// sequences, and must not start with a dot. This prevents client-controlled +/// IDs from escaping the prediction root directory. +pub fn validate_prediction_id(id: &str) -> Result<(), &'static str> { + if id.is_empty() { + return Err("prediction ID must not be empty"); + } + if id.starts_with('.') { + return Err("prediction ID must not start with a dot"); + } + if id.contains('/') || id.contains('\\') { + return Err("prediction ID must not contain path separators"); + } + if id.contains("..") { + return Err("prediction ID must not contain traversal sequences"); + } + Ok(()) +} + fn prediction_dir_for(id: &str) -> std::path::PathBuf { - std::path::PathBuf::from("/tmp/coglet/predictions").join(id) + std::path::PathBuf::from(PREDICTION_ROOT).join(id) +} + +/// Verify that `dir` is a direct child of the prediction root. +/// +/// Used as a defense-in-depth check before recursive deletion. Returns true if +/// the canonicalized parent equals the canonicalized prediction root. +fn is_direct_prediction_child(dir: &std::path::Path) -> bool { + let Ok(canonical_dir) = std::fs::canonicalize(dir) else { + return false; + }; + let Some(parent) = canonical_dir.parent() else { + return false; + }; + let Ok(root) = std::fs::canonicalize(PREDICTION_ROOT) else { + return false; + }; + parent == root } fn spawn_orchestrator_cancel(orch: Arc, id: String) { @@ -1757,6 +1819,117 @@ mod tests { svc.remove_prediction("nonexistent-prediction"); } + #[test] + fn remove_prediction_does_not_delete_outside_prediction_root() { + let svc = PredictionService::new_no_pool(); + // Create a directory that looks like it could be a prediction dir but + // is outside the prediction root. + let outside_dir = std::path::PathBuf::from("/tmp/coglet_outside_test"); + std::fs::create_dir_all(&outside_dir).unwrap(); + std::fs::write(outside_dir.join("file.txt"), b"keep").unwrap(); + + // Absolute ID that would resolve to /tmp/coglet_outside_test if not defended. + svc.remove_prediction("/tmp/coglet_outside_test"); + + assert!( + outside_dir.exists(), + "directory outside prediction root should not be removed" + ); + + // Cleanup + std::fs::remove_dir_all(&outside_dir).unwrap(); + } + + #[test] + fn remove_prediction_does_not_delete_traversal_targets() { + let svc = PredictionService::new_no_pool(); + let victim = std::path::PathBuf::from("/tmp/coglet_victim_test"); + std::fs::create_dir_all(&victim).unwrap(); + + svc.remove_prediction("../coglet_victim_test"); + + assert!(victim.exists(), "traversal target should not be removed"); + + std::fs::remove_dir_all(&victim).unwrap(); + } + + #[tokio::test] + async fn submit_prediction_rejects_duplicate_id() { + let svc = PredictionService::new_no_pool(); + let pool = create_test_pool(2).await; + let orchestrator = Arc::new(MockOrchestrator::new()); + + svc.set_orchestrator(pool, orchestrator).await; + svc.set_health(Health::Ready).await; + + let (_handle, _slot) = svc + .submit_prediction( + "duplicate-id".to_string(), + serde_json::json!({}), + None, + false, + ) + .await + .unwrap(); + + let result = svc + .submit_prediction( + "duplicate-id".to_string(), + serde_json::json!({}), + None, + false, + ) + .await; + + assert!(matches!(result, Err(CreatePredictionError::DuplicateId))); + } + + #[tokio::test] + async fn submit_prediction_rejects_invalid_id() { + let svc = PredictionService::new_no_pool(); + let pool = create_test_pool(1).await; + let orchestrator = Arc::new(MockOrchestrator::new()); + + svc.set_orchestrator(pool, orchestrator).await; + svc.set_health(Health::Ready).await; + + let invalid_ids = vec![ + "", + "foo/bar", + "foo\\bar", + "../escape", + ".hidden", + "foo..bar", + ]; + + for id in invalid_ids { + let result = svc + .submit_prediction(id.to_string(), serde_json::json!({}), None, false) + .await; + assert!( + matches!(result, Err(CreatePredictionError::InvalidId(_))), + "expected invalid ID error for {:?}", + id + ); + } + } + + #[test] + fn validate_prediction_id_accepts_safe_ids() { + assert!(validate_prediction_id("pred_123").is_ok()); + assert!(validate_prediction_id("uuid-style-id").is_ok()); + assert!(validate_prediction_id("with.dots").is_ok()); + } + + #[test] + fn validate_prediction_id_rejects_dangerous_ids() { + assert!(validate_prediction_id("").is_err()); + assert!(validate_prediction_id("foo/bar").is_err()); + assert!(validate_prediction_id("foo\\bar").is_err()); + assert!(validate_prediction_id("../up").is_err()); + assert!(validate_prediction_id(".hidden").is_err()); + } + #[test] fn build_slot_request_small_input_inline() { let dir = tempfile::tempdir().unwrap(); diff --git a/crates/coglet/src/transport/http/routes.rs b/crates/coglet/src/transport/http/routes.rs index 5904737391..6b42d3e678 100644 --- a/crates/coglet/src/transport/http/routes.rs +++ b/crates/coglet/src/transport/http/routes.rs @@ -23,7 +23,7 @@ use crate::prediction::SharedPredictionStreamEvent; use crate::predictor::PredictionError; use crate::service::{ CreatePredictionError, HealthSnapshot, PredictionService, PredictionStreamSubscription, - SubscribePredictionStreamError, + SubscribePredictionStreamError, validate_prediction_id, }; use crate::version::VersionInfo; use crate::webhook::{TraceContext, WebhookConfig, WebhookEventType, WebhookSender}; @@ -288,6 +288,20 @@ fn extract_trace_context(headers: &HeaderMap) -> TraceContext { } } +fn invalid_prediction_id_response(msg: &str) -> Response { + ( + StatusCode::UNPROCESSABLE_ENTITY, + Json(serde_json::json!({ + "detail": [{ + "loc": ["body", "id"], + "msg": msg, + "type": "value_error" + }] + })), + ) + .into_response() +} + async fn create_prediction( State(service): State>, headers: HeaderMap, @@ -300,7 +314,15 @@ async fn create_prediction( webhook: None, webhook_events_filter: default_webhook_events_filter(), }); - let prediction_id = request.id.unwrap_or_else(generate_prediction_id); + let prediction_id = match request.id { + Some(id) => { + if let Err(msg) = validate_prediction_id(&id) { + return invalid_prediction_id_response(msg); + } + id + } + None => generate_prediction_id(), + }; let response_mode = prediction_response_mode(&headers); let trace_context = extract_trace_context(&headers); create_prediction_with_id( @@ -323,6 +345,10 @@ async fn create_prediction_idempotent( headers: HeaderMap, body: Option>, ) -> Response { + if let Err(msg) = validate_prediction_id(&prediction_id) { + return invalid_prediction_id_response(msg); + } + let request = body.map(|Json(r)| r).unwrap_or_else(|| PredictionRequest { id: None, input: serde_json::json!({}), @@ -331,20 +357,23 @@ async fn create_prediction_idempotent( webhook_events_filter: default_webhook_events_filter(), }); - if let Some(ref req_id) = request.id - && req_id != &prediction_id - { - return ( - StatusCode::UNPROCESSABLE_ENTITY, - Json(serde_json::json!({ - "detail": [{ - "loc": ["body", "id"], - "msg": "prediction ID must match the ID supplied in the URL", - "type": "value_error" - }] - })), - ) - .into_response(); + if let Some(ref req_id) = request.id { + if let Err(msg) = validate_prediction_id(req_id) { + return invalid_prediction_id_response(msg); + } + if req_id != &prediction_id { + return ( + StatusCode::UNPROCESSABLE_ENTITY, + Json(serde_json::json!({ + "detail": [{ + "loc": ["body", "id"], + "msg": "prediction ID must match the ID supplied in the URL", + "type": "value_error" + }] + })), + ) + .into_response(); + } } let response_mode = prediction_response_mode(&headers); @@ -488,6 +517,19 @@ async fn create_prediction_with_id( ) .into_response(); } + Err(CreatePredictionError::DuplicateId) => { + return ( + StatusCode::CONFLICT, + Json(serde_json::json!({ + "error": "A prediction with this ID already exists", + "status": "failed" + })), + ) + .into_response(); + } + Err(CreatePredictionError::InvalidId(msg)) => { + return invalid_prediction_id_response(msg); + } }; let prediction = unregistered_slot.prediction(); @@ -823,7 +865,15 @@ async fn create_training( webhook: None, webhook_events_filter: default_webhook_events_filter(), }); - let prediction_id = request.id.unwrap_or_else(generate_prediction_id); + let prediction_id = match request.id { + Some(id) => { + if let Err(msg) = validate_prediction_id(&id) { + return invalid_prediction_id_response(msg); + } + id + } + None => generate_prediction_id(), + }; let response_mode = json_response_mode(&headers); let trace_context = extract_trace_context(&headers); create_prediction_with_id( @@ -850,6 +900,10 @@ async fn create_training_idempotent( return training_streaming_not_supported_response(); } + if let Err(msg) = validate_prediction_id(&training_id) { + return invalid_prediction_id_response(msg); + } + let request = body.map(|Json(r)| r).unwrap_or_else(|| PredictionRequest { id: None, input: serde_json::json!({}), @@ -858,20 +912,23 @@ async fn create_training_idempotent( webhook_events_filter: default_webhook_events_filter(), }); - if let Some(ref req_id) = request.id - && req_id != &training_id - { - return ( - StatusCode::UNPROCESSABLE_ENTITY, - Json(serde_json::json!({ - "detail": [{ - "loc": ["body", "id"], - "msg": "training ID must match the ID supplied in the URL", - "type": "value_error" - }] - })), - ) - .into_response(); + if let Some(ref req_id) = request.id { + if let Err(msg) = validate_prediction_id(req_id) { + return invalid_prediction_id_response(msg); + } + if req_id != &training_id { + return ( + StatusCode::UNPROCESSABLE_ENTITY, + Json(serde_json::json!({ + "detail": [{ + "loc": ["body", "id"], + "msg": "training ID must match the ID supplied in the URL", + "type": "value_error" + }] + })), + ) + .into_response(); + } } // Idempotent: return existing state if already submitted @@ -1578,6 +1635,82 @@ mod tests { ); } + #[tokio::test] + async fn prediction_rejects_invalid_id_in_body() { + let service = create_ready_service().await; + let app = routes(service); + + let response = app + .oneshot( + Request::post("/predictions") + .header("content-type", "application/json") + .body(Body::from(r#"{"id":"foo/bar","input":{}}"#)) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::UNPROCESSABLE_ENTITY); + } + + #[tokio::test] + async fn prediction_rejects_invalid_id_in_url() { + let service = create_ready_service().await; + let app = routes(service); + + let response = app + .oneshot( + Request::put("/predictions/.hidden") + .header("content-type", "application/json") + .body(Body::from(r#"{"input":{}}"#)) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::UNPROCESSABLE_ENTITY); + } + + #[tokio::test] + async fn prediction_rejects_duplicate_id() { + // Use an orchestrator that never completes so the first prediction stays + // registered long enough for the duplicate request to be rejected. + let service = Arc::new(PredictionService::new_no_pool()); + let pool = create_test_pool(2).await; + let orchestrator = Arc::new(MockOrchestrator::never_complete()); + service.set_orchestrator(pool, orchestrator).await; + service.set_health(Health::Ready).await; + + let app = routes(Arc::clone(&service)); + let response = app + .oneshot( + Request::post("/predictions") + .header("content-type", "application/json") + .header("prefer", "respond-async") + .body(Body::from(r#"{"id":"dup-1","input":{}}"#)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::ACCEPTED); + + // Small delay to let the async task register the prediction + tokio::time::sleep(tokio::time::Duration::from_millis(10)).await; + + let app2 = routes(service); + let response = app2 + .oneshot( + Request::post("/predictions") + .header("content-type", "application/json") + .body(Body::from(r#"{"id":"dup-1","input":{}}"#)) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::CONFLICT); + } + #[tokio::test] async fn prediction_at_capacity() { let service = Arc::new(PredictionService::new_no_pool()); diff --git a/crates/coglet/src/worker.rs b/crates/coglet/src/worker.rs index bc0645b62a..1242e3a183 100644 --- a/crates/coglet/src/worker.rs +++ b/crates/coglet/src/worker.rs @@ -164,6 +164,11 @@ impl SlotSender { } } + /// Directory where Coglet-managed output files for this prediction live. + pub fn output_dir(&self) -> &std::path::Path { + &self.output_dir + } + fn next_output_index(&self) -> u64 { self.output_counter.fetch_add(1, Ordering::Relaxed) } @@ -231,6 +236,33 @@ impl SlotSender { .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "slot channel closed")) } + /// Send a user-returned file path, ensuring Coglet only deletes files it owns. + /// + /// If the path is already inside the Cog-managed output directory, send it as + /// `managed: true` so the orchestrator deletes it after upload. If it lives + /// outside the output directory, copy it to a unique path inside the output + /// directory and send the copy as managed. This preserves the documented + /// behavior that returned files are cleaned up while never deleting arbitrary + /// user-owned paths. + pub fn send_user_file_output( + &self, + path: PathBuf, + mime_type: Option, + ) -> io::Result<()> { + if path.starts_with(&self.output_dir) { + return self.send_file_output(path, mime_type, true); + } + + let ext = path + .extension() + .and_then(|e| e.to_str()) + .unwrap_or("bin") + .to_string(); + let dest = self.next_output_path(&ext); + std::fs::copy(&path, &dest)?; + self.send_file_output(dest, mime_type, true) + } + /// Send a user metric to the parent process. pub fn send_metric( &self, @@ -982,4 +1014,60 @@ mod tests { let config = WorkerConfig::default(); assert_eq!(config.num_slots, 1); } + + #[test] + fn send_user_file_output_copies_external_path() { + let output_dir = tempfile::tempdir().unwrap(); + let (tx, mut rx) = mpsc::unbounded_channel::(); + let sender = SlotSender::new(tx, output_dir.path().to_path_buf()); + + // Create a file outside the managed output dir + let external = tempfile::NamedTempFile::with_suffix(".txt").unwrap(); + std::fs::write(external.path(), b"hello").unwrap(); + + sender + .send_user_file_output(external.path().to_path_buf(), None) + .unwrap(); + + let msg = rx.try_recv().unwrap(); + match msg { + SlotResponse::FileOutput { + filename, + managed: true, + .. + } => { + assert!( + filename.starts_with(output_dir.path().to_str().unwrap()), + "external path should be copied into output dir" + ); + assert!(std::path::Path::new(&filename).exists()); + assert_eq!(std::fs::read(&filename).unwrap(), b"hello"); + } + _ => panic!("expected managed FileOutput"), + } + } + + #[test] + fn send_user_file_output_passes_through_output_dir_path() { + let output_dir = tempfile::tempdir().unwrap(); + let (tx, mut rx) = mpsc::unbounded_channel::(); + let sender = SlotSender::new(tx, output_dir.path().to_path_buf()); + + let inside = output_dir.path().join("inside.txt"); + std::fs::write(&inside, b"world").unwrap(); + + sender.send_user_file_output(inside.clone(), None).unwrap(); + + let msg = rx.try_recv().unwrap(); + match msg { + SlotResponse::FileOutput { + filename, + managed: true, + .. + } => { + assert_eq!(filename, inside.to_str().unwrap()); + } + _ => panic!("expected managed FileOutput for path inside output dir"), + } + } } diff --git a/docs/llms.txt b/docs/llms.txt index e62311b0d6..737e1dcaea 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -2771,7 +2771,8 @@ class Runner(BaseRunner): upscaled_image = do_some_processing(image) # To output cog.Path objects the file needs to exist, so create a temporary file first. - # This file will automatically be deleted by Cog after it has been returned. + # Cog copies returned files into a managed output directory, uploads the copy, + # and deletes it afterward. The original file is left untouched. output_path = Path(tempfile.mkdtemp()) / "upscaled.png" upscaled_image.save(output_path) return Path(output_path) diff --git a/docs/python.md b/docs/python.md index d717cf56f2..eb7d9e37a3 100644 --- a/docs/python.md +++ b/docs/python.md @@ -537,7 +537,8 @@ class Runner(BaseRunner): upscaled_image = do_some_processing(image) # To output cog.Path objects the file needs to exist, so create a temporary file first. - # This file will automatically be deleted by Cog after it has been returned. + # Cog copies returned files into a managed output directory, uploads the copy, + # and deletes it afterward. The original file is left untouched. output_path = Path(tempfile.mkdtemp()) / "upscaled.png" upscaled_image.save(output_path) return Path(output_path) From 8d5a4d52456fe160082cef0e8fa8d5074075ad9c Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Thu, 16 Jul 2026 16:21:15 -0500 Subject: [PATCH 3/3] fix: complete managed output cleanup --- crates/coglet-python/src/predictor.rs | 70 +++++++--- crates/coglet/src/bridge/protocol.rs | 3 +- crates/coglet/src/orchestrator.rs | 39 +++++- crates/coglet/src/service.rs | 152 +++++++++++++++------ crates/coglet/src/transport/http/routes.rs | 2 +- crates/coglet/src/worker.rs | 50 +++---- docs/llms.txt | 3 +- docs/python.md | 3 +- 8 files changed, 221 insertions(+), 101 deletions(-) diff --git a/crates/coglet-python/src/predictor.rs b/crates/coglet-python/src/predictor.rs index 6fc6a335e1..edb137ed23 100644 --- a/crates/coglet-python/src/predictor.rs +++ b/crates/coglet-python/src/predictor.rs @@ -166,9 +166,22 @@ fn submit_async_coroutine( /// Send a single output item over IPC, routing file outputs to disk. /// -/// For Path outputs (os.PathLike): sends the existing file path via send_file_output. +/// For Path outputs (os.PathLike): transfers the file to Coglet-managed storage. /// For IOBase outputs: reads bytes, writes to output_dir via write_file_output. /// For everything else: processes through make_encodeable + upload_files, then send_output. +fn send_path_output( + item: &Bound<'_, PyAny>, + slot_sender: &SlotSender, +) -> Result<(), PredictionError> { + let path_str: String = item + .call_method0("__fspath__") + .and_then(|path| path.extract()) + .map_err(|e| PredictionError::Failed(format!("Failed to get fspath: {}", e)))?; + slot_sender + .send_user_file_output(std::path::PathBuf::from(path_str), None) + .map_err(|e| PredictionError::Failed(format!("Failed to send file output: {}", e))) +} + fn send_output_item( py: Python<'_>, item: &Bound<'_, PyAny>, @@ -189,16 +202,9 @@ fn send_output_item( .map_err(|e| PredictionError::Failed(format!("Failed to get io.IOBase: {}", e)))?; if item.is_instance(&pathlike).unwrap_or(false) { - // Path output — file already on disk. send_user_file_output copies - // external paths into the managed output directory and deletes them - // after upload, so Coglet never deletes arbitrary user-owned files. - let path_str: String = item - .call_method0("__fspath__") - .and_then(|p| p.extract()) - .map_err(|e| PredictionError::Failed(format!("Failed to get fspath: {}", e)))?; - slot_sender - .send_user_file_output(std::path::PathBuf::from(path_str), None) - .map_err(|e| PredictionError::Failed(format!("Failed to send file output: {}", e)))?; + // Returning a Path transfers ownership to Coglet, which moves the data + // into managed storage and deletes it after consumption. + send_path_output(item, slot_sender)?; return Ok(()); } @@ -951,15 +957,7 @@ impl PythonPredictor { .map_err(|e| PredictionError::Failed(format!("Failed to get io.IOBase: {}", e)))?; if result.is_instance(&pathlike).unwrap_or(false) { - let path_str: String = result - .call_method0("__fspath__") - .and_then(|p| p.extract()) - .map_err(|e| PredictionError::Failed(format!("Failed to get fspath: {}", e)))?; - slot_sender - .send_user_file_output(std::path::PathBuf::from(path_str), None) - .map_err(|e| { - PredictionError::Failed(format!("Failed to send file output: {}", e)) - })?; + send_path_output(result, slot_sender)?; return Ok(PredictionOutput::Single(serde_json::Value::Null)); } @@ -1380,6 +1378,7 @@ mod tests { use std::path::PathBuf; + use coglet_core::bridge::protocol::SlotResponse; use pyo3::types::PyList; fn add_python_sdk_path(py: Python<'_>) { @@ -1437,6 +1436,37 @@ sys.modules.setdefault('requests', requests) }) } + #[test] + fn returned_path_transfers_ownership_to_slot_sender() { + pyo3::Python::initialize(); + let source_dir = tempfile::tempdir().unwrap(); + let source = source_dir.path().join("output.txt"); + std::fs::write(&source, b"output").unwrap(); + let output_dir = tempfile::tempdir().unwrap(); + let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); + let sender = SlotSender::new(tx, output_dir.path().to_path_buf()); + + Python::attach(|py| { + let pathlib = py.import("pathlib").unwrap(); + let path = pathlib + .getattr("Path") + .unwrap() + .call1((source.to_str().unwrap(),)) + .unwrap(); + send_path_output(&path, &sender).unwrap(); + }); + + assert!(!source.exists()); + match rx.try_recv().unwrap() { + SlotResponse::FileOutput { + filename, + managed: true, + .. + } => assert_eq!(std::fs::read(filename).unwrap(), b"output"), + response => panic!("expected managed FileOutput, got {response:?}"), + } + } + #[test] fn class_with_run_loads() { let predictor = load_predictor_source( diff --git a/crates/coglet/src/bridge/protocol.rs b/crates/coglet/src/bridge/protocol.rs index 68d30e828c..0e0ba4aba1 100644 --- a/crates/coglet/src/bridge/protocol.rs +++ b/crates/coglet/src/bridge/protocol.rs @@ -327,8 +327,7 @@ pub enum SlotResponse { /// Explicit MIME type from the predictor. Falls back to mime_guess when None. #[serde(skip_serializing_if = "Option::is_none")] mime_type: Option, - /// True if Coglet created this file (IOBase write or oversized spill). - /// False for user-authored Path outputs — must not be deleted. + /// True if Coglet owns this file and may delete it after consumption. #[serde(default)] managed: bool, }, diff --git a/crates/coglet/src/orchestrator.rs b/crates/coglet/src/orchestrator.rs index 3a8994bf9b..96921923e6 100644 --- a/crates/coglet/src/orchestrator.rs +++ b/crates/coglet/src/orchestrator.rs @@ -84,6 +84,14 @@ fn ensure_trailing_slash(s: &str) -> String { } } +fn read_file_output(filename: &str, managed: bool) -> std::io::Result> { + let bytes = std::fs::read(filename)?; + if managed && let Err(e) = std::fs::remove_file(filename) { + tracing::debug!(%filename, error = %e, "Failed to delete managed output file"); + } + Ok(bytes) +} + /// Try to lock a prediction mutex. /// On poison: logs error, recovers to fail the prediction, returns None. /// Caller should remove the prediction from tracking if None is returned. @@ -1074,18 +1082,13 @@ async fn run_event_loop( } Ok(SlotResponse::FileOutput { filename, kind, mime_type, managed }) => { tracing::debug!(%slot_id, %filename, ?kind, "FileOutput received"); - let bytes = match std::fs::read(&filename) { + let bytes = match read_file_output(&filename, managed) { Ok(b) => b, Err(e) => { tracing::error!(%slot_id, %filename, error = %e, "Failed to read FileOutput"); continue; } }; - if managed - && let Err(e) = std::fs::remove_file(&filename) - { - tracing::debug!(%slot_id, %filename, error = %e, "Failed to delete managed output file"); - } match kind { FileOutputKind::Oversized => { let output: serde_json::Value = match serde_json::from_slice(&bytes) { @@ -1314,6 +1317,30 @@ mod tests { assert!(!pending.contains("overflow")); } + #[test] + fn read_file_output_deletes_managed_file() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("managed.txt"); + std::fs::write(&path, b"managed").unwrap(); + + let bytes = read_file_output(path.to_str().unwrap(), true).unwrap(); + + assert_eq!(bytes, b"managed"); + assert!(!path.exists()); + } + + #[test] + fn read_file_output_preserves_unmanaged_file() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("unmanaged.txt"); + std::fs::write(&path, b"unmanaged").unwrap(); + + let bytes = read_file_output(path.to_str().unwrap(), false).unwrap(); + + assert_eq!(bytes, b"unmanaged"); + assert!(path.exists()); + } + #[test] fn wrap_outputs_schema_array_single_item() { // List[Path] with num_outputs=1 → ["url"] not "url" diff --git a/crates/coglet/src/service.rs b/crates/coglet/src/service.rs index 0b6a7c723f..5f1a510929 100644 --- a/crates/coglet/src/service.rs +++ b/crates/coglet/src/service.rs @@ -533,6 +533,15 @@ impl PredictionService { let pool = self.pool().await; let pool = pool.as_ref().ok_or(CreatePredictionError::NotReady)?; + // Reserve the ID before acquiring a permit so a rejected duplicate + // cannot consume slot capacity. + let prediction_entry = match self.predictions.entry(id.clone()) { + dashmap::Entry::Vacant(entry) => entry, + dashmap::Entry::Occupied(_) => { + return Err(CreatePredictionError::DuplicateId); + } + }; + let permit = pool .try_acquire() .ok_or(CreatePredictionError::AtCapacity)?; @@ -543,21 +552,12 @@ impl PredictionService { let slot = PredictionSlot::new(prediction, permit, idle_rx); let prediction_arc = slot.prediction(); - // Register atomically in DashMap — duplicate active IDs are rejected to - // prevent one prediction from clobbering or deleting another's state. - match self.predictions.entry(id.clone()) { - dashmap::Entry::Vacant(entry) => { - entry.insert(PredictionEntry { - prediction: prediction_arc, - cancel_token: cancel_token.clone(), - input, - cancel_on_stream_drop, - }); - } - dashmap::Entry::Occupied(_) => { - return Err(CreatePredictionError::DuplicateId); - } - } + prediction_entry.insert(PredictionEntry { + prediction: prediction_arc, + cancel_token: cancel_token.clone(), + input, + cancel_on_stream_drop, + }); let handle = PredictionHandle { id, cancel_token }; @@ -812,25 +812,47 @@ impl PredictionService { /// that were not deleted individually (e.g. aborted uploads, cancelled /// predictions). pub fn remove_prediction(&self, id: &str) { - self.predictions.remove(id); - // Defense in depth: validate the ID and ensure we only delete a direct // child of the prediction root, even though the HTTP layer also validates. if let Err(e) = validate_prediction_id(id) { tracing::warn!(prediction_id = %id, error = e, "Refusing to remove prediction dir: invalid ID"); return; } + + // Keep the ID reserved until cleanup completes so a new prediction with + // the same ID cannot create files that this call then removes. + let prediction_entry = self.predictions.entry(id.to_string()); let dir = prediction_dir_for(id); - if !is_direct_prediction_child(&dir) { - tracing::warn!(prediction_id = %id, dir = %dir.display(), "Refusing to remove prediction dir: not a direct child of prediction root"); - return; - } - match std::fs::remove_dir_all(&dir) { - Ok(()) => {} - Err(e) if e.kind() == std::io::ErrorKind::NotFound => {} + let cleanup_complete = match std::fs::symlink_metadata(&dir) { + Err(e) if e.kind() == std::io::ErrorKind::NotFound => true, Err(e) => { - tracing::debug!(prediction_id = %id, error = %e, "Failed to remove prediction dir") + tracing::debug!(prediction_id = %id, error = %e, "Failed to inspect prediction dir"); + false } + Ok(metadata) if metadata.file_type().is_symlink() => match std::fs::remove_file(&dir) { + Ok(()) => true, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => true, + Err(e) => { + tracing::debug!(prediction_id = %id, error = %e, "Failed to remove prediction dir symlink"); + false + } + }, + Ok(_) if !is_direct_prediction_child(&dir) => { + tracing::warn!(prediction_id = %id, dir = %dir.display(), "Refusing to remove prediction dir: not a direct child of prediction root"); + false + } + Ok(_) => match std::fs::remove_dir_all(&dir) { + Ok(()) => true, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => true, + Err(e) => { + tracing::debug!(prediction_id = %id, error = %e, "Failed to remove prediction dir"); + false + } + }, + }; + + if cleanup_complete && let dashmap::Entry::Occupied(entry) = prediction_entry { + entry.remove(); } } @@ -847,21 +869,18 @@ const PREDICTION_ROOT: &str = "/tmp/coglet/predictions"; /// Validate a client-supplied prediction ID. /// -/// IDs must be non-empty, must not contain path separators or traversal -/// sequences, and must not start with a dot. This prevents client-controlled -/// IDs from escaping the prediction root directory. +/// IDs must be non-empty direct path components. This prevents client-controlled +/// IDs from escaping the prediction root directory without rejecting otherwise +/// safe IDs containing dots. pub fn validate_prediction_id(id: &str) -> Result<(), &'static str> { if id.is_empty() { return Err("prediction ID must not be empty"); } - if id.starts_with('.') { - return Err("prediction ID must not start with a dot"); - } if id.contains('/') || id.contains('\\') { return Err("prediction ID must not contain path separators"); } - if id.contains("..") { - return Err("prediction ID must not contain traversal sequences"); + if id == "." || id == ".." { + return Err("prediction ID must not be . or .."); } Ok(()) } @@ -1853,13 +1872,54 @@ mod tests { std::fs::remove_dir_all(&victim).unwrap(); } + #[cfg(unix)] + #[test] + fn remove_prediction_unlinks_directory_symlink_without_following_it() { + let svc = PredictionService::new_no_pool(); + let id = format!("symlink-test-{}", uuid::Uuid::new_v4()); + let dir = prediction_dir_for(&id); + std::fs::create_dir_all(PREDICTION_ROOT).unwrap(); + let victim = tempfile::tempdir().unwrap(); + std::fs::write(victim.path().join("keep.txt"), b"keep").unwrap(); + std::os::unix::fs::symlink(victim.path(), &dir).unwrap(); + + svc.remove_prediction(&id); + + assert!(!dir.exists(), "prediction symlink should be removed"); + assert!(victim.path().join("keep.txt").exists()); + } + + #[tokio::test] + async fn remove_prediction_retains_id_when_cleanup_fails() { + let svc = PredictionService::new_no_pool(); + let pool = create_test_pool(1).await; + let orchestrator = Arc::new(MockOrchestrator::new()); + svc.set_orchestrator(pool, orchestrator).await; + svc.set_health(Health::Ready).await; + let id = format!("cleanup-failure-{}", uuid::Uuid::new_v4()); + let (_handle, _slot) = svc + .submit_prediction(id.clone(), serde_json::json!({}), None, false) + .await + .unwrap(); + let dir = prediction_dir_for(&id); + std::fs::create_dir_all(PREDICTION_ROOT).unwrap(); + std::fs::write(&dir, b"not a directory").unwrap(); + + svc.remove_prediction(&id); + + assert!(svc.prediction_exists(&id)); + std::fs::remove_file(dir).unwrap(); + svc.remove_prediction(&id); + assert!(!svc.prediction_exists(&id)); + } + #[tokio::test] async fn submit_prediction_rejects_duplicate_id() { let svc = PredictionService::new_no_pool(); let pool = create_test_pool(2).await; let orchestrator = Arc::new(MockOrchestrator::new()); - svc.set_orchestrator(pool, orchestrator).await; + svc.set_orchestrator(Arc::clone(&pool), orchestrator).await; svc.set_health(Health::Ready).await; let (_handle, _slot) = svc @@ -1882,6 +1942,17 @@ mod tests { .await; assert!(matches!(result, Err(CreatePredictionError::DuplicateId))); + assert_eq!(pool.available(), 1, "duplicate must not consume a permit"); + + let unrelated = svc + .submit_prediction( + "unrelated-id".to_string(), + serde_json::json!({}), + None, + false, + ) + .await; + assert!(unrelated.is_ok(), "remaining permit should still be usable"); } #[tokio::test] @@ -1893,14 +1964,7 @@ mod tests { svc.set_orchestrator(pool, orchestrator).await; svc.set_health(Health::Ready).await; - let invalid_ids = vec![ - "", - "foo/bar", - "foo\\bar", - "../escape", - ".hidden", - "foo..bar", - ]; + let invalid_ids = vec!["", ".", "..", "foo/bar", "foo\\bar", "../escape"]; for id in invalid_ids { let result = svc @@ -1919,6 +1983,7 @@ mod tests { assert!(validate_prediction_id("pred_123").is_ok()); assert!(validate_prediction_id("uuid-style-id").is_ok()); assert!(validate_prediction_id("with.dots").is_ok()); + assert!(validate_prediction_id(".hidden").is_ok()); } #[test] @@ -1927,7 +1992,8 @@ mod tests { assert!(validate_prediction_id("foo/bar").is_err()); assert!(validate_prediction_id("foo\\bar").is_err()); assert!(validate_prediction_id("../up").is_err()); - assert!(validate_prediction_id(".hidden").is_err()); + assert!(validate_prediction_id(".").is_err()); + assert!(validate_prediction_id("..").is_err()); } #[test] diff --git a/crates/coglet/src/transport/http/routes.rs b/crates/coglet/src/transport/http/routes.rs index 6b42d3e678..f64feb2659 100644 --- a/crates/coglet/src/transport/http/routes.rs +++ b/crates/coglet/src/transport/http/routes.rs @@ -1660,7 +1660,7 @@ mod tests { let response = app .oneshot( - Request::put("/predictions/.hidden") + Request::put("/predictions/foo%5Cbar") .header("content-type", "application/json") .body(Body::from(r#"{"input":{}}"#)) .unwrap(), diff --git a/crates/coglet/src/worker.rs b/crates/coglet/src/worker.rs index 1242e3a183..9464e93e7e 100644 --- a/crates/coglet/src/worker.rs +++ b/crates/coglet/src/worker.rs @@ -164,11 +164,6 @@ impl SlotSender { } } - /// Directory where Coglet-managed output files for this prediction live. - pub fn output_dir(&self) -> &std::path::Path { - &self.output_dir - } - fn next_output_index(&self) -> u64 { self.output_counter.fetch_add(1, Ordering::Relaxed) } @@ -213,8 +208,7 @@ impl SlotSender { /// /// The file is already on disk at `path` — we just send the path reference. /// `mime_type` is an explicit MIME type; when None the parent guesses from extension. - /// `managed` is true when Coglet created the file (safe to delete after consumption); - /// false for user-authored Path outputs. + /// `managed` is true when Coglet owns the file and may delete it after consumption. pub fn send_file_output( &self, path: PathBuf, @@ -236,30 +230,30 @@ impl SlotSender { .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "slot channel closed")) } - /// Send a user-returned file path, ensuring Coglet only deletes files it owns. + /// Transfer a user-returned file into Coglet's managed output directory. /// - /// If the path is already inside the Cog-managed output directory, send it as - /// `managed: true` so the orchestrator deletes it after upload. If it lives - /// outside the output directory, copy it to a unique path inside the output - /// directory and send the copy as managed. This preserves the documented - /// behavior that returned files are cleaned up while never deleting arbitrary - /// user-owned paths. + /// Returning a path hands ownership to Coglet. Copying before unlinking works + /// across filesystems while ensuring the source is only removed after the + /// managed copy has been written successfully. pub fn send_user_file_output( &self, path: PathBuf, mime_type: Option, ) -> io::Result<()> { - if path.starts_with(&self.output_dir) { - return self.send_file_output(path, mime_type, true); - } - let ext = path .extension() .and_then(|e| e.to_str()) .unwrap_or("bin") .to_string(); - let dest = self.next_output_path(&ext); + let mut dest = self.next_output_path(&ext); + while dest == path || dest.exists() { + dest = self.next_output_path(&ext); + } std::fs::copy(&path, &dest)?; + if let Err(e) = std::fs::remove_file(&path) { + let _ = std::fs::remove_file(&dest); + return Err(e); + } self.send_file_output(dest, mime_type, true) } @@ -1016,19 +1010,20 @@ mod tests { } #[test] - fn send_user_file_output_copies_external_path() { + fn send_user_file_output_transfers_external_path() { let output_dir = tempfile::tempdir().unwrap(); let (tx, mut rx) = mpsc::unbounded_channel::(); let sender = SlotSender::new(tx, output_dir.path().to_path_buf()); - // Create a file outside the managed output dir let external = tempfile::NamedTempFile::with_suffix(".txt").unwrap(); std::fs::write(external.path(), b"hello").unwrap(); + let external_path = external.path().to_path_buf(); sender - .send_user_file_output(external.path().to_path_buf(), None) + .send_user_file_output(external_path.clone(), None) .unwrap(); + assert!(!external_path.exists(), "returned path should be removed"); let msg = rx.try_recv().unwrap(); match msg { SlotResponse::FileOutput { @@ -1045,10 +1040,13 @@ mod tests { } _ => panic!("expected managed FileOutput"), } + + assert!(sender.send_user_file_output(external_path, None).is_err()); + assert!(rx.try_recv().is_err(), "stale path must not be sent again"); } #[test] - fn send_user_file_output_passes_through_output_dir_path() { + fn send_user_file_output_transfers_output_dir_path() { let output_dir = tempfile::tempdir().unwrap(); let (tx, mut rx) = mpsc::unbounded_channel::(); let sender = SlotSender::new(tx, output_dir.path().to_path_buf()); @@ -1058,6 +1056,7 @@ mod tests { sender.send_user_file_output(inside.clone(), None).unwrap(); + assert!(!inside.exists(), "returned path should be removed"); let msg = rx.try_recv().unwrap(); match msg { SlotResponse::FileOutput { @@ -1065,9 +1064,10 @@ mod tests { managed: true, .. } => { - assert_eq!(filename, inside.to_str().unwrap()); + assert_ne!(filename, inside.to_str().unwrap()); + assert_eq!(std::fs::read(filename).unwrap(), b"world"); } - _ => panic!("expected managed FileOutput for path inside output dir"), + _ => panic!("expected managed FileOutput"), } } } diff --git a/docs/llms.txt b/docs/llms.txt index 8de057a814..e8419e39a8 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -2784,8 +2784,7 @@ class Runner(BaseRunner): upscaled_image = do_some_processing(image) # To output cog.Path objects the file needs to exist, so create a temporary file first. - # Cog copies returned files into a managed output directory, uploads the copy, - # and deletes it afterward. The original file is left untouched. + # This file will automatically be deleted by Cog after it has been returned. output_path = Path(tempfile.mkdtemp()) / "upscaled.png" upscaled_image.save(output_path) return Path(output_path) diff --git a/docs/python.md b/docs/python.md index eb7d9e37a3..d717cf56f2 100644 --- a/docs/python.md +++ b/docs/python.md @@ -537,8 +537,7 @@ class Runner(BaseRunner): upscaled_image = do_some_processing(image) # To output cog.Path objects the file needs to exist, so create a temporary file first. - # Cog copies returned files into a managed output directory, uploads the copy, - # and deletes it afterward. The original file is left untouched. + # This file will automatically be deleted by Cog after it has been returned. output_path = Path(tempfile.mkdtemp()) / "upscaled.png" upscaled_image.save(output_path) return Path(output_path)