diff --git a/crates/coglet-python/src/predictor.rs b/crates/coglet-python/src/predictor.rs index 6142e48e24..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,14 +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 path reference - 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) - .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(()); } @@ -949,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_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)); } @@ -1378,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<'_>) { @@ -1435,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 0ef46a4e68..0e0ba4aba1 100644 --- a/crates/coglet/src/bridge/protocol.rs +++ b/crates/coglet/src/bridge/protocol.rs @@ -327,6 +327,9 @@ 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 owns this file and may delete it after consumption. + #[serde(default)] + managed: bool, }, /// Streaming output chunk for generator and iterator output. @@ -592,6 +595,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..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. @@ -1072,9 +1080,9 @@ 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) { + let bytes = match read_file_output(&filename, managed) { Ok(b) => b, Err(e) => { tracing::error!(%slot_id, %filename, error = %e, "Failed to read FileOutput"); @@ -1309,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 460a09e09a..5f1a510929 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); @@ -527,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)?; @@ -537,16 +552,12 @@ 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, - }, - ); + prediction_entry.insert(PredictionEntry { + prediction: prediction_arc, + cancel_token: cancel_token.clone(), + input, + cancel_on_stream_drop, + }); let handle = PredictionHandle { id, cancel_token }; @@ -657,8 +668,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 +807,53 @@ 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); + // 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); + 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 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(); + } } pub fn trigger_shutdown(&self) { @@ -811,6 +865,47 @@ impl PredictionService { } } +const PREDICTION_ROOT: &str = "/tmp/coglet/predictions"; + +/// Validate a client-supplied prediction ID. +/// +/// 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.contains('/') || id.contains('\\') { + return Err("prediction ID must not contain path separators"); + } + if id == "." || id == ".." { + return Err("prediction ID must not be . or .."); + } + Ok(()) +} + +fn prediction_dir_for(id: &str) -> std::path::PathBuf { + 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) { let Ok(handle) = tokio::runtime::Handle::try_current() else { tracing::warn!(prediction_id = %id, "No tokio runtime available to cancel prediction"); @@ -1723,6 +1818,184 @@ 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 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(); + } + + #[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(Arc::clone(&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))); + 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] + 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"]; + + 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()); + assert!(validate_prediction_id(".hidden").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(".").is_err()); + assert!(validate_prediction_id("..").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..f64feb2659 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/foo%5Cbar") + .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 c59a26cc7a..9464e93e7e 100644 --- a/crates/coglet/src/worker.rs +++ b/crates/coglet/src/worker.rs @@ -201,14 +201,20 @@ 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 owns the file and may delete it after consumption. + 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,12 +223,40 @@ impl SlotSender { filename, kind: FileOutputKind::FileType, mime_type, + managed, }; self.tx .send(msg) .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "slot channel closed")) } + /// Transfer a user-returned file into Coglet's managed output directory. + /// + /// 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<()> { + let ext = path + .extension() + .and_then(|e| e.to_str()) + .unwrap_or("bin") + .to_string(); + 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) + } + /// Send a user metric to the parent process. pub fn send_metric( &self, @@ -265,6 +299,7 @@ fn build_output_message( filename, kind: FileOutputKind::Oversized, mime_type: None, + managed: true, }) } else { Ok(SlotResponse::OutputChunk { output, index }) @@ -973,4 +1008,66 @@ mod tests { let config = WorkerConfig::default(); assert_eq!(config.num_slots, 1); } + + #[test] + 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()); + + 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.clone(), None) + .unwrap(); + + assert!(!external_path.exists(), "returned path should be removed"); + 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"), + } + + 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_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()); + + let inside = output_dir.path().join("inside.txt"); + std::fs::write(&inside, b"world").unwrap(); + + 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 { + filename, + managed: true, + .. + } => { + assert_ne!(filename, inside.to_str().unwrap()); + assert_eq!(std::fs::read(filename).unwrap(), b"world"); + } + _ => panic!("expected managed FileOutput"), + } + } }