From 9dc8d9a77a3bf3313737f67f41b2431abe9a21db Mon Sep 17 00:00:00 2001 From: Hare Date: Thu, 3 Sep 2026 07:22:07 +0900 Subject: [PATCH] fix: fence attachment grants and lifecycle cleanup --- crates/client/src/backend_runtime.rs | 68 ++- crates/session-store/src/fs_store.rs | 191 ++++++++- crates/session-store/src/lib.rs | 1 + crates/session-store/src/segment.rs | 2 + crates/session-store/src/store.rs | 24 +- crates/session-store/src/uploaded_file.rs | 114 ++++- crates/session-store/tests/session_test.rs | 1 + crates/tui/src/app.rs | 4 + crates/tui/src/console/mod.rs | 206 ++++++--- crates/worker-runtime/src/execution.rs | 4 +- crates/worker-runtime/src/http_server.rs | 84 +++- crates/worker-runtime/src/lib.rs | 1 + crates/worker-runtime/src/runtime.rs | 37 +- crates/worker-runtime/src/worker_backend.rs | 49 +-- crates/worker/src/controller.rs | 17 + crates/workspace-server/src/hosts.rs | 42 +- crates/workspace-server/src/server.rs | 395 ++++++++++++++++-- .../workspace/console/composer-attachments.ts | 104 ++++- .../workers/[workerId]/console/+page.svelte | 13 +- 19 files changed, 1190 insertions(+), 167 deletions(-) diff --git a/crates/client/src/backend_runtime.rs b/crates/client/src/backend_runtime.rs index c0461858..0d84072f 100644 --- a/crates/client/src/backend_runtime.rs +++ b/crates/client/src/backend_runtime.rs @@ -58,18 +58,52 @@ impl BackendRuntimeTarget { file_name: &str, media_type: &str, content: Vec, + ) -> Result { + self.upload_file_with_id( + &uuid::Uuid::now_v7().to_string(), + file_name, + media_type, + content, + ) + .await + } + + pub async fn upload_file_with_id( + &self, + upload_id: &str, + file_name: &str, + media_type: &str, + content: Vec, ) -> Result { let api = BackendApiClient::from_stored_token(&self.base_url)?; - let path = format!( - "/api/w/{}/runtimes/{}/workers/{}/attachments?file_name={}&media_type={}", + let worker_path = format!( + "/api/w/{}/runtimes/{}/workers/{}", path_segment_encode(&self.workspace_id), path_segment_encode(&self.runtime_id), path_segment_encode(&self.worker_id), + ); + let grant_path = format!( + "{worker_path}/attachment-upload-grants?file_name={}&media_type={}&upload_id={}", path_segment_encode(file_name), path_segment_encode(media_type), + path_segment_encode(&upload_id), + ); + let grant_response = api + .request(HttpMethod::POST, &grant_path)? + .send() + .await + .map_err(BackendRuntimeClientError::Http)?; + api.check_status(grant_response.status())?; + let grant = grant_response + .json::() + .await + .map_err(BackendRuntimeClientError::Http)?; + let upload_path = format!( + "{worker_path}/attachment-uploads/{}", + path_segment_encode(&grant.upload_id), ); let response = api - .request(HttpMethod::POST, &path)? + .request(HttpMethod::PUT, &upload_path)? .body(content) .send() .await @@ -82,6 +116,27 @@ impl BackendRuntimeTarget { .map_err(BackendRuntimeClientError::Http) } + pub async fn cancel_file_upload( + &self, + upload_id: &str, + ) -> Result<(), BackendRuntimeClientError> { + let api = BackendApiClient::from_stored_token(&self.base_url)?; + let path = format!( + "/api/w/{}/runtimes/{}/workers/{}/attachment-uploads/{}", + path_segment_encode(&self.workspace_id), + path_segment_encode(&self.runtime_id), + path_segment_encode(&self.worker_id), + path_segment_encode(upload_id), + ); + let response = api + .request(HttpMethod::DELETE, &path)? + .send() + .await + .map_err(BackendRuntimeClientError::Http)?; + api.check_status(response.status())?; + Ok(()) + } + pub async fn delete_uploaded_file( &self, artifact_id: &str, @@ -104,6 +159,13 @@ impl BackendRuntimeTarget { } } +#[derive(Deserialize)] +struct AttachmentUploadGrantResponse { + upload_id: String, + #[allow(dead_code)] + expires_at_ms: u64, +} + #[derive(Deserialize)] struct UploadedFileResponse { file: protocol::UploadedFileRef, diff --git a/crates/session-store/src/fs_store.rs b/crates/session-store/src/fs_store.rs index 9efccd0a..d7e7274c 100644 --- a/crates/session-store/src/fs_store.rs +++ b/crates/session-store/src/fs_store.rs @@ -20,10 +20,13 @@ use crate::paste_artifact::{read_from_dir, write_to_dir}; use crate::segment_log::LogEntry; use crate::store::{Store, StoreError}; use crate::uploaded_file::{ - bind_uploaded_file, delete_uncommitted_uploaded_files, delete_uploaded_file, + bind_uploaded_file, clear_uploaded_file_binding, copy_committed_uploaded_files, + delete_uncommitted_uploaded_files, delete_uploaded_file, list_uploaded_file_refs, read_uploaded_file, read_uploaded_file_by_id, write_uploaded_file, }; -use crate::{PasteArtifactLimits, SegmentId, SessionId, UploadedFileLimits}; +use crate::{ + PasteArtifactLimits, SegmentId, SessionId, UploadedFileLimits, UploadedFileUploadContext, +}; use protocol::{PasteArtifactRef, UploadedFileRef}; use std::fs; use std::io::{Read, Seek, SeekFrom, Write}; @@ -119,6 +122,40 @@ impl FsStore { self.session_dir(session_id).join("artifacts").join("paste") } + fn uploaded_file_is_referenced( + &self, + session_id: SessionId, + artifact_id: &str, + ) -> Result { + fn segments_contain(segments: &[protocol::Segment], artifact_id: &str) -> bool { + segments.iter().any(|segment| { + matches!( + segment, + protocol::Segment::UploadedFile { file } + if file.artifact_id == artifact_id + ) + }) + } + + for segment_id in self.list_segments(session_id)? { + for entry in self.read_all(session_id, segment_id)? { + let referenced = match entry { + LogEntry::AnnotatedUserInput { segments, .. } => { + segments_contain(&segments, artifact_id) + } + LogEntry::InputSegmentsCheckpoint { user_segments, .. } => user_segments + .iter() + .any(|segments| segments_contain(segments, artifact_id)), + _ => false, + }; + if referenced { + return Ok(true); + } + } + } + Ok(false) + } + #[cfg(test)] fn paste_artifact_path(&self, session_id: SessionId, artifact_id: &str) -> PathBuf { self.paste_artifact_dir(session_id) @@ -410,6 +447,30 @@ impl Store for FsStore { file_name, media_type, content, + None, + limits, + ) + } + + fn write_uploaded_file_with_context( + &self, + session_id: SessionId, + file_name: &str, + media_type: &str, + content: &[u8], + context: &UploadedFileUploadContext, + limits: UploadedFileLimits, + ) -> Result { + let _guard = self + .append_lock + .lock() + .map_err(|_| std::io::Error::other("session store append lock was poisoned"))?; + write_uploaded_file( + &self.paste_artifact_dir(session_id), + file_name, + media_type, + content, + Some(context), limits, ) } @@ -440,11 +501,21 @@ impl Store for FsStore { .append_lock .lock() .map_err(|_| std::io::Error::other("session store append lock was poisoned"))?; - bind_uploaded_file( - &self.paste_artifact_dir(session_id), - reference, - source_entry_id, - ) + let dir = self.paste_artifact_dir(session_id); + match bind_uploaded_file(&dir, reference, source_entry_id) { + Err(StoreError::ArtifactAlreadyCommitted) => { + let (stored, _) = read_uploaded_file_by_id(&dir, &reference.artifact_id)?; + let previous_source = stored + .source_entry_id + .ok_or(StoreError::ArtifactIntegrityMismatch)?; + if self.uploaded_file_is_referenced(session_id, &reference.artifact_id)? { + return Err(StoreError::ArtifactAlreadyCommitted); + } + clear_uploaded_file_binding(&dir, &reference.artifact_id, &previous_source)?; + bind_uploaded_file(&dir, reference, source_entry_id) + } + result => result, + } } fn delete_uploaded_file( @@ -464,7 +535,37 @@ impl Store for FsStore { .append_lock .lock() .map_err(|_| std::io::Error::other("session store append lock was poisoned"))?; - delete_uncommitted_uploaded_files(&self.paste_artifact_dir(session_id)) + let dir = self.paste_artifact_dir(session_id); + let mut removed = delete_uncommitted_uploaded_files(&dir)?; + for reference in list_uploaded_file_refs(&dir)? { + let Some(source_entry_id) = reference.source_entry_id.as_deref() else { + continue; + }; + if !self.uploaded_file_is_referenced(session_id, &reference.artifact_id)? { + clear_uploaded_file_binding(&dir, &reference.artifact_id, source_entry_id)?; + if delete_uploaded_file(&dir, &reference.artifact_id)? { + removed = removed + .checked_add(1) + .ok_or(StoreError::ArtifactQuotaExceeded)?; + } + } + } + Ok(removed) + } + + fn copy_committed_uploaded_files( + &self, + source_session_id: SessionId, + target_session_id: SessionId, + ) -> Result { + let _guard = self + .append_lock + .lock() + .map_err(|_| std::io::Error::other("session store append lock was poisoned"))?; + copy_committed_uploaded_files( + &self.paste_artifact_dir(source_session_id), + &self.paste_artifact_dir(target_session_id), + ) } fn append_trace( @@ -626,6 +727,45 @@ mod tests { ); } + #[test] + fn uploaded_file_persists_trusted_upload_context_without_projecting_it() { + let tmp = tempfile::TempDir::new().unwrap(); + let store = FsStore::new(tmp.path()).unwrap(); + let session_id = new_session_id(); + let context = UploadedFileUploadContext { + upload_id: "upload-1".into(), + principal_id: "account-1".into(), + workspace_id: "workspace-1".into(), + runtime_id: "runtime-1".into(), + worker_id: "worker-1".into(), + }; + let reference = store + .write_uploaded_file_with_context( + session_id, + "notes.txt", + "text/plain", + b"hello", + &context, + UploadedFileLimits::default(), + ) + .unwrap(); + let raw = fs::read_to_string( + store + .paste_artifact_dir(session_id) + .join(format!("{}.file.json", reference.artifact_id)), + ) + .unwrap(); + assert!(raw.contains("account-1")); + assert!(raw.contains("workspace-1")); + assert!(raw.contains("runtime-1")); + assert!(raw.contains("worker-1")); + assert!( + !serde_json::to_string(&reference) + .unwrap() + .contains("account-1") + ); + } + #[test] fn uploaded_files_are_session_scoped_integrity_checked_and_removable() { let tmp = tempfile::TempDir::new().unwrap(); @@ -715,18 +855,51 @@ mod tests { store.write_uploaded_file(session_id, "ＲＥＡＤＭＥ.txt", "text/plain", b"y", limits), Err(StoreError::InvalidUploadedFileName) )); + store + .bind_uploaded_file(session_id, &pending, "entry-from-failed-submit") + .unwrap(); let bound = store .bind_uploaded_file(session_id, &pending, "entry-upload") .unwrap(); + store + .create_segment( + session_id, + new_segment_id(), + &[LogEntry::InputSegmentsCheckpoint { + ts: 1, + user_segments: vec![vec![protocol::Segment::UploadedFile { + file: bound.clone(), + }]], + }], + ) + .unwrap(); let other = store .write_uploaded_file(session_id, "other.txt", "text/plain", b"z", limits) .unwrap(); + let stale = store + .write_uploaded_file(session_id, "stale.txt", "text/plain", b"s", limits) + .unwrap(); + store + .bind_uploaded_file(session_id, &stale, "entry-never-committed") + .unwrap(); assert_eq!( store.delete_uncommitted_uploaded_files(session_id).unwrap(), - 1 + 2 ); assert!(store.read_uploaded_file(session_id, &other).is_err()); + assert!(store.read_uploaded_file(session_id, &stale).is_err()); assert_eq!(store.read_uploaded_file(session_id, &bound).unwrap(), b"x"); + let fork_session_id = new_session_id(); + assert_eq!( + store + .copy_committed_uploaded_files(session_id, fork_session_id) + .unwrap(), + 1 + ); + assert_eq!( + store.read_uploaded_file(fork_session_id, &bound).unwrap(), + b"x" + ); store .write_paste_artifact( session_id, diff --git a/crates/session-store/src/lib.rs b/crates/session-store/src/lib.rs index 1d4a5f31..9e34c5aa 100644 --- a/crates/session-store/src/lib.rs +++ b/crates/session-store/src/lib.rs @@ -70,6 +70,7 @@ pub use system_item::{ pub use uploaded_file::{ DEFAULT_MAX_FILES_PER_SUBMISSION, DEFAULT_MAX_SESSION_ARTIFACT_BYTES, DEFAULT_MAX_SESSION_UPLOADED_FILES, DEFAULT_MAX_UPLOADED_FILE_BYTES, UploadedFileLimits, + UploadedFileUploadContext, }; pub use worker_metadata::{ CombinedStore, FsWorkerStore, WorkerActiveSegmentRef, WorkerAggregateStore, WorkerMetadata, diff --git a/crates/session-store/src/segment.rs b/crates/session-store/src/segment.rs index 017cb970..0ea4fb04 100644 --- a/crates/session-store/src/segment.rs +++ b/crates/session-store/src/segment.rs @@ -420,6 +420,7 @@ pub fn save_config_changed( /// [`fork_at`] or [`ensure_head_or_fork`] instead. pub fn fork( store: &impl Store, + source_session_id: SessionId, state: SegmentStartState<'_>, ) -> Result<(SessionId, SegmentId), StoreError> { let session_id = crate::new_session_id(); @@ -434,6 +435,7 @@ pub fn fork( compacted_from: None, }; store.create_segment(session_id, fork_id, &[entry])?; + store.copy_committed_uploaded_files(source_session_id, session_id)?; Ok((session_id, fork_id)) } diff --git a/crates/session-store/src/store.rs b/crates/session-store/src/store.rs index 757a6b15..a22fd6c7 100644 --- a/crates/session-store/src/store.rs +++ b/crates/session-store/src/store.rs @@ -13,7 +13,9 @@ use crate::event_trace::TraceEntry; use crate::segment_log::LogEntry; -use crate::{PasteArtifactLimits, SegmentId, SessionId, UploadedFileLimits}; +use crate::{ + PasteArtifactLimits, SegmentId, SessionId, UploadedFileLimits, UploadedFileUploadContext, +}; use protocol::{PasteArtifactRef, UploadedFileRef}; /// Errors from the persistence store. @@ -186,6 +188,18 @@ pub trait Store: Send + Sync { Err(StoreError::PasteArtifactUnsupported) } + fn write_uploaded_file_with_context( + &self, + session_id: SessionId, + file_name: &str, + media_type: &str, + content: &[u8], + _context: &UploadedFileUploadContext, + limits: UploadedFileLimits, + ) -> Result { + self.write_uploaded_file(session_id, file_name, media_type, content, limits) + } + /// Read and integrity-check an uploaded file owned by `session_id`. fn read_uploaded_file( &self, @@ -225,6 +239,14 @@ pub trait Store: Send + Sync { Ok(0) } + fn copy_committed_uploaded_files( + &self, + _source_session_id: SessionId, + _target_session_id: SessionId, + ) -> Result { + Ok(0) + } + /// Append a trace entry to the debug event trace file. fn append_trace( &self, diff --git a/crates/session-store/src/uploaded_file.rs b/crates/session-store/src/uploaded_file.rs index 05af45a9..e8b5207f 100644 --- a/crates/session-store/src/uploaded_file.rs +++ b/crates/session-store/src/uploaded_file.rs @@ -38,6 +38,15 @@ impl Default for UploadedFileLimits { } } +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct UploadedFileUploadContext { + pub upload_id: String, + pub principal_id: String, + pub workspace_id: String, + pub runtime_id: String, + pub worker_id: String, +} + #[derive(Debug, Serialize, Deserialize)] struct StoredUploadedFile { file_name: String, @@ -47,6 +56,8 @@ struct StoredUploadedFile { sha256: String, #[serde(default, skip_serializing_if = "Option::is_none")] source_entry_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + upload_context: Option, content_base64: String, } @@ -186,6 +197,7 @@ pub(crate) fn write_uploaded_file( file_name: &str, media_type: &str, content: &[u8], + context: Option<&UploadedFileUploadContext>, limits: UploadedFileLimits, ) -> Result { validate_file_name(file_name)?; @@ -226,6 +238,7 @@ pub(crate) fn write_uploaded_file( if stored.media_type == media_type && stored.byte_len == byte_len && stored.sha256 == sha256 + && stored.upload_context.as_ref() == context { let artifact_id = path .file_name() @@ -264,6 +277,7 @@ pub(crate) fn write_uploaded_file( byte_len, sha256: sha256.clone(), source_entry_id: None, + upload_context: context.cloned(), content_base64: BASE64.encode(content), }; let path = record_path(dir, &artifact_id)?; @@ -324,6 +338,30 @@ pub(crate) fn read_uploaded_file(dir: &Path, reference: &UploadedFileRef) -> Res Ok(content) } +pub(crate) fn clear_uploaded_file_binding( + dir: &Path, + artifact_id: &str, + expected_source_entry_id: &str, +) -> Result<()> { + fs::create_dir_all(dir)?; + let aggregate_lock = fs::OpenOptions::new() + .create(true) + .read(true) + .write(true) + .open(dir.join(".aggregate.lock"))?; + FileExt::lock_exclusive(&aggregate_lock)?; + let path = record_path(dir, artifact_id)?; + let mut stored: StoredUploadedFile = serde_json::from_slice(&fs::read(&path)?)?; + if stored.source_entry_id.as_deref() != Some(expected_source_entry_id) { + return Err(StoreError::ArtifactIntegrityMismatch); + } + stored.source_entry_id = None; + let temp = dir.join(format!(".{artifact_id}.file.unbind.tmp")); + fs::write(&temp, serde_json::to_vec(&stored)?)?; + fs::rename(temp, path)?; + Ok(()) +} + pub(crate) fn bind_uploaded_file( dir: &Path, reference: &UploadedFileRef, @@ -338,7 +376,15 @@ pub(crate) fn bind_uploaded_file( .write(true) .open(dir.join(".aggregate.lock"))?; FileExt::lock_exclusive(&aggregate_lock)?; - read_uploaded_file(dir, reference)?; + let (stored_reference, _) = read_uploaded_file_by_id(dir, &reference.artifact_id)?; + if stored_reference.file_name != reference.file_name + || stored_reference.media_type != reference.media_type + || stored_reference.created_at_ms != reference.created_at_ms + || stored_reference.byte_len != reference.byte_len + || stored_reference.sha256 != reference.sha256 + { + return Err(StoreError::ArtifactIntegrityMismatch); + } let path = record_path(dir, &reference.artifact_id)?; let mut stored: StoredUploadedFile = serde_json::from_slice(&fs::read(&path)?)?; if stored.source_entry_id.is_some() { @@ -353,6 +399,72 @@ pub(crate) fn bind_uploaded_file( Ok(bound) } +pub(crate) fn list_uploaded_file_refs(dir: &Path) -> Result> { + if !dir.exists() { + return Ok(Vec::new()); + } + let mut refs = Vec::new(); + for entry in fs::read_dir(dir)? { + let path = entry?.path(); + let Some(artifact_id) = path + .file_name() + .and_then(|name| name.to_str()) + .and_then(|name| name.strip_suffix(".file.json")) + else { + continue; + }; + refs.push(read_uploaded_file_by_id(dir, artifact_id)?.0); + } + Ok(refs) +} + +pub(crate) fn copy_committed_uploaded_files(source_dir: &Path, target_dir: &Path) -> Result { + if !source_dir.exists() { + return Ok(0); + } + fs::create_dir_all(target_dir)?; + let target_lock = fs::OpenOptions::new() + .create(true) + .read(true) + .write(true) + .open(target_dir.join(".aggregate.lock"))?; + FileExt::lock_exclusive(&target_lock)?; + let mut copied = 0_u64; + for entry in fs::read_dir(source_dir)? { + let entry = entry?; + let path = entry.path(); + let Some(name) = path.file_name().and_then(|name| name.to_str()) else { + continue; + }; + if !name.ends_with(".file.json") { + continue; + } + let bytes = fs::read(&path)?; + let stored: StoredUploadedFile = serde_json::from_slice(&bytes)?; + if stored.source_entry_id.is_none() { + continue; + } + let target = target_dir.join(name); + if target.exists() { + let existing: StoredUploadedFile = serde_json::from_slice(&fs::read(&target)?)?; + if existing.sha256 != stored.sha256 + || existing.file_name != stored.file_name + || existing.source_entry_id != stored.source_entry_id + { + return Err(StoreError::ArtifactIntegrityMismatch); + } + continue; + } + let temp = target_dir.join(format!(".{name}.copy.tmp")); + fs::write(&temp, &bytes)?; + fs::rename(temp, target)?; + copied = copied + .checked_add(1) + .ok_or(StoreError::ArtifactQuotaExceeded)?; + } + Ok(copied) +} + pub(crate) fn delete_uncommitted_uploaded_files(dir: &Path) -> Result { fs::create_dir_all(dir)?; let aggregate_lock = fs::OpenOptions::new() diff --git a/crates/session-store/tests/session_test.rs b/crates/session-store/tests/session_test.rs index 898bc4bd..09b80166 100644 --- a/crates/session-store/tests/session_test.rs +++ b/crates/session-store/tests/session_test.rs @@ -406,6 +406,7 @@ async fn session_fork_creates_new_session() { let original_history_len = worker.history().len(); let (fork_sid, fork_segid) = session_store::fork( &store, + sid, SegmentStartState { system_prompt: worker.get_system_prompt(), config: worker.request_config(), diff --git a/crates/tui/src/app.rs b/crates/tui/src/app.rs index d3b2396d..2b963b4d 100644 --- a/crates/tui/src/app.rs +++ b/crates/tui/src/app.rs @@ -928,6 +928,10 @@ impl App { Some(self.method_for_run(queued.segments)) } + pub fn clear_actionbar_notice(&mut self) { + self.actionbar_notice = None; + } + pub fn push_error(&mut self, message: impl Into) { self.blocks.push(Block::Alert { level: AlertLevel::Error, diff --git a/crates/tui/src/console/mod.rs b/crates/tui/src/console/mod.rs index 79355685..f1402a29 100644 --- a/crates/tui/src/console/mod.rs +++ b/crates/tui/src/console/mod.rs @@ -1,3 +1,4 @@ +use std::collections::HashMap; use std::fmt; use std::future::Future; use std::io; @@ -120,11 +121,55 @@ fn copy_selection_to_terminal(app: &mut App) -> bool { copy_selection_to_writer(app, &mut stdout) } +type AttachmentUploadResult = Result; + struct ConsoleConnection { client: Client, standalone_host: Option, backend_target: Option, pending_attachments: Vec, + upload_tasks: Vec>, + upload_ids: HashMap, +} + +async fn upload_client_path( + target: &BackendRuntimeTarget, + path: &Path, + upload_id: &str, +) -> Result> { + let metadata = tokio::fs::metadata(path).await?; + if !metadata.is_file() { + return Err( + io::Error::new(io::ErrorKind::InvalidInput, "attachment path is not a file").into(), + ); + } + if metadata.len() > 10 * 1024 * 1024 { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "attachment exceeds the 10 MiB limit", + ) + .into()); + } + let file_name = path + .file_name() + .and_then(|name| name.to_str()) + .ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidInput, + "attachment file name is not valid UTF-8", + ) + })?; + let media_type = attachment_media_type(path).ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidInput, + "attachment file type is not supported", + ) + })?; + let bytes = tokio::fs::read(path).await?; + target + .upload_file_with_id(upload_id, file_name, media_type, bytes) + .await + .map_err(Into::into) } fn attachment_media_type(path: &Path) -> Option<&'static str> { @@ -153,6 +198,8 @@ impl ConsoleConnection { standalone_host: Some(host), backend_target: None, pending_attachments: Vec::new(), + upload_tasks: Vec::new(), + upload_ids: HashMap::new(), } } @@ -162,6 +209,8 @@ impl ConsoleConnection { standalone_host: None, backend_target: Some(target), pending_attachments: Vec::new(), + upload_tasks: Vec::new(), + upload_ids: HashMap::new(), } } @@ -188,59 +237,49 @@ impl ConsoleConnection { self.client.send(&prepared).await?; if carries_attachments { self.pending_attachments.clear(); + self.upload_ids.clear(); } Ok(()) } - async fn upload_path( + fn start_upload( &mut self, - path: &Path, - ) -> Result> { - let target = self.backend_target.as_ref().ok_or_else(|| { + path: PathBuf, + result_tx: mpsc::UnboundedSender, + ) -> Result<(), Box> { + let target = self.backend_target.clone().ok_or_else(|| { io::Error::new( io::ErrorKind::Unsupported, "client-local file upload is available only for Backend Workers", ) })?; - let metadata = tokio::fs::metadata(path).await?; - if !metadata.is_file() { - return Err(io::Error::new( - io::ErrorKind::InvalidInput, - "attachment path is not a file", - ) - .into()); - } - if metadata.len() > 10 * 1024 * 1024 { - return Err(io::Error::new( - io::ErrorKind::InvalidInput, - "attachment exceeds the 10 MiB limit", - ) - .into()); - } - let file_name = path - .file_name() - .and_then(|name| name.to_str()) - .ok_or_else(|| { - io::Error::new( - io::ErrorKind::InvalidInput, - "attachment file name is not valid UTF-8", - ) - })?; - let media_type = attachment_media_type(path).ok_or_else(|| { - io::Error::new( - io::ErrorKind::InvalidInput, - "attachment file type is not supported", - ) - })?; - let bytes = tokio::fs::read(path).await?; - let reference = target.upload_file(file_name, media_type, bytes).await?; - self.pending_attachments.push(reference.clone()); - Ok(reference) + self.upload_tasks.retain(|task| !task.is_finished()); + let upload_id = self + .upload_ids + .entry(path.clone()) + .or_insert_with(|| uuid::Uuid::now_v7().to_string()) + .clone(); + self.upload_tasks.push(tokio::spawn(async move { + let result = upload_client_path(&target, &path, &upload_id) + .await + .map_err(|error| error.to_string()); + let _ = result_tx.send(result); + })); + Ok(()) } async fn clear_pending_attachments(&mut self) { + for task in self.upload_tasks.drain(..) { + task.abort(); + } + let upload_ids = std::mem::take(&mut self.upload_ids) + .into_values() + .collect::>(); let references = std::mem::take(&mut self.pending_attachments); if let Some(target) = &self.backend_target { + for upload_id in upload_ids { + let _ = target.cancel_file_upload(&upload_id).await; + } for reference in references { let _ = target.delete_uploaded_file(&reference.artifact_id).await; } @@ -626,11 +665,13 @@ enum E2eRewindInput { enum LoopInput

{ Terminal(TerminalEventResult), Worker(P), + Upload(AttachmentUploadResult), Tick, } async fn next_loop_input( term_rx: &mut mpsc::UnboundedReceiver, + upload_rx: &mut mpsc::UnboundedReceiver, connected: bool, pod_next: F, animate: bool, @@ -651,6 +692,9 @@ where )) })) } + upload = upload_rx.recv() => { + LoopInput::Upload(upload.unwrap_or_else(|| Err("attachment upload queue stopped".into()))) + } event = pod_next, if connected => LoopInput::Worker(event), _ = animation_tick, if animate => LoopInput::Tick, } @@ -660,13 +704,14 @@ async fn drain_terminal_events( app: &mut App, client: &mut ConsoleConnection, term_rx: &mut mpsc::UnboundedReceiver, + upload_tx: &mpsc::UnboundedSender, ) -> Result> { let mut handled = false; for _ in 0..TERMINAL_EVENT_DRAIN_LIMIT { match term_rx.try_recv() { Ok(event) => { handled = true; - handle_terminal_event(app, client, event?).await?; + handle_terminal_event(app, client, upload_tx, event?).await?; if app.quit { break; } @@ -708,6 +753,7 @@ async fn run_loop( client: &mut ConsoleConnection, ) -> Result<(), Box> { let (_terminal_reader, mut term_rx) = TerminalEventReader::spawn()?; + let (upload_tx, mut upload_rx) = mpsc::unbounded_channel(); let mut animation_tick = tokio::time::interval(Duration::from_millis(80)); animation_tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); @@ -718,7 +764,8 @@ async fn run_loop( break; } - let handled_term_event = drain_terminal_events(app, client, &mut term_rx).await?; + let handled_term_event = + drain_terminal_events(app, client, &mut term_rx, &upload_tx).await?; if app.quit { break; } @@ -730,6 +777,7 @@ async fn run_loop( match next_loop_input( &mut term_rx, + &mut upload_rx, app.connected, client.next_event(), app.running, @@ -738,8 +786,29 @@ async fn run_loop( .await { LoopInput::Terminal(term_event) => { - handle_terminal_event(app, client, term_event?).await?; + handle_terminal_event(app, client, &upload_tx, term_event?).await?; } + LoopInput::Upload(result) => match result { + Ok(reference) => { + app.flash_actionbar_notice( + format!( + "[{} · {} bytes · ready] Send a message or use /clear-attachments.", + reference.file_name, reference.byte_len + ), + ActionbarNoticeLevel::Info, + ActionbarNoticeSource::Tui, + Duration::from_secs(60 * 60), + ); + if !client + .pending_attachments + .iter() + .any(|pending| pending.artifact_id == reference.artifact_id) + { + client.pending_attachments.push(reference); + } + } + Err(error) => app.push_error(format!("Attachment upload failed: {error}")), + }, LoopInput::Worker(event) => match event? { Some(ev) => { if let Some(method) = app.handle_worker_event(ev) { @@ -785,21 +854,19 @@ fn is_clear_attachments_command(method: &Method) -> bool { async fn handle_terminal_event( app: &mut App, client: &mut ConsoleConnection, + upload_tx: &mpsc::UnboundedSender, event: TermEvent, ) -> Result<(), Box> { match event { TermEvent::Key(key) => { if let Some(method) = handle_key(app, key) { if let Some(path) = attachment_command_path(&method) { - match client.upload_path(&path).await { - Ok(reference) => app.flash_actionbar_notice( - format!( - "Attached {} ({} bytes); it will be sent with the next message.", - reference.file_name, reference.byte_len - ), + match client.start_upload(path, upload_tx.clone()) { + Ok(()) => app.flash_actionbar_notice( + "Uploading attachment… Use /clear-attachments to cancel.", ActionbarNoticeLevel::Info, ActionbarNoticeSource::Tui, - Duration::from_secs(6), + Duration::from_secs(30), ), Err(error) => { app.push_error(format!("Attachment upload failed: {error}")); @@ -814,7 +881,12 @@ async fn handle_terminal_event( Duration::from_secs(4), ); } else { + let sends_attachments = matches!(method, Method::Run { .. }) + && !client.pending_attachments.is_empty(); client.send(&method).await?; + if sends_attachments { + app.clear_actionbar_notice(); + } } } } @@ -1420,10 +1492,12 @@ mod tests { #[tokio::test] async fn animation_tick_wakes_loop_while_running() { let (_tx, mut rx) = mpsc::unbounded_channel::(); + let (_upload_tx, mut upload_rx) = mpsc::unbounded_channel(); assert!(matches!( next_loop_input( &mut rx, + &mut upload_rx, true, std::future::pending::>(), true, @@ -1434,9 +1508,41 @@ mod tests { )); } + #[tokio::test] + async fn attachment_upload_completion_wakes_console_loop() { + let (_terminal_tx, mut terminal_rx) = mpsc::unbounded_channel(); + let (upload_tx, mut upload_rx) = mpsc::unbounded_channel(); + let file = UploadedFileRef { + artifact_id: "019ca7c8-57b6-7f05-8edf-524147aba7b3".into(), + file_name: "notes.txt".into(), + media_type: "text/plain".into(), + created_at_ms: 1, + availability: protocol::UploadedFileAvailability::Available, + byte_len: 1, + sha256: "a".repeat(64), + source_entry_id: None, + }; + upload_tx.send(Ok(file.clone())).unwrap(); + + match next_loop_input( + &mut terminal_rx, + &mut upload_rx, + true, + std::future::pending::>(), + false, + std::future::pending::<()>(), + ) + .await + { + LoopInput::Upload(Ok(received)) => assert_eq!(received, file), + _ => panic!("expected attachment upload result"), + } + } + #[tokio::test] async fn terminal_event_is_selected_before_ready_worker_event() { let (tx, mut rx) = mpsc::unbounded_channel(); + let (_upload_tx, mut upload_rx) = mpsc::unbounded_channel(); tx.send(Ok(TermEvent::Key(KeyEvent::new( KeyCode::Char('x'), KeyModifiers::NONE, @@ -1445,6 +1551,7 @@ mod tests { match next_loop_input( &mut rx, + &mut upload_rx, true, std::future::ready(Some(())), false, @@ -1462,9 +1569,11 @@ mod tests { #[tokio::test] async fn terminal_event_is_preserved_after_worker_event_wins() { let (tx, mut rx) = mpsc::unbounded_channel(); + let (_upload_tx, mut upload_rx) = mpsc::unbounded_channel(); match next_loop_input( &mut rx, + &mut upload_rx, true, std::future::ready(Some(1_u8)), false, @@ -1484,6 +1593,7 @@ mod tests { match next_loop_input( &mut rx, + &mut upload_rx, true, std::future::ready(Some(2_u8)), false, diff --git a/crates/worker-runtime/src/execution.rs b/crates/worker-runtime/src/execution.rs index 78384c85..ff24ecbb 100644 --- a/crates/worker-runtime/src/execution.rs +++ b/crates/worker-runtime/src/execution.rs @@ -393,6 +393,7 @@ pub trait WorkerExecutionBackend: Send + Sync + 'static { _file_name: &str, _media_type: &str, _content: &[u8], + _context: Option<&session_store::UploadedFileUploadContext>, ) -> Result { Err(WorkerExecutionResult::unsupported( WorkerExecutionOperation::UploadFile, @@ -546,9 +547,10 @@ impl WorkerExecutionBackendRef { file_name: &str, media_type: &str, content: &[u8], + context: Option<&session_store::UploadedFileUploadContext>, ) -> Result { self.backend - .upload_file(handle, file_name, media_type, content) + .upload_file(handle, file_name, media_type, content, context) } pub(crate) fn delete_uploaded_file( diff --git a/crates/worker-runtime/src/http_server.rs b/crates/worker-runtime/src/http_server.rs index 4c0dd7a8..1e26482e 100644 --- a/crates/worker-runtime/src/http_server.rs +++ b/crates/worker-runtime/src/http_server.rs @@ -390,6 +390,16 @@ pub struct RuntimeHttpWorkerInputResponse { pub struct RuntimeHttpUploadFileQuery { pub file_name: String, pub media_type: String, + #[serde(default)] + pub upload_id: Option, + #[serde(default)] + pub principal_id: Option, + #[serde(default)] + pub workspace_id: Option, + #[serde(default)] + pub runtime_id: Option, + #[serde(default)] + pub owner_worker_id: Option, } #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] @@ -1455,14 +1465,74 @@ async fn upload_worker_file( body: Bytes, ) -> RestResult { let worker_ref = worker_ref_for(&state.runtime, worker_id)?; + let context = match ( + query.upload_id, + query.principal_id, + query.workspace_id, + query.runtime_id, + query.owner_worker_id, + ) { + (None, None, None, None, None) => None, + ( + Some(upload_id), + Some(principal_id), + Some(workspace_id), + Some(runtime_id), + Some(owner_worker_id), + ) => { + if owner_worker_id != worker_ref.worker_id.to_string() { + return Err(RuntimeHttpRestError::new( + StatusCode::FORBIDDEN, + "uploaded_file_owner_mismatch", + "uploaded file context does not match the target Worker", + )); + } + Some(session_store::UploadedFileUploadContext { + upload_id, + principal_id, + workspace_id, + runtime_id, + worker_id: owner_worker_id, + }) + } + _ => { + return Err(RuntimeHttpRestError::new( + StatusCode::BAD_REQUEST, + "uploaded_file_context_incomplete", + "uploaded file context fields must be provided together", + )); + } + }; let file = match auth_workspace_scope(&state, auth.as_ref())? { - Some(scope) => state.runtime.upload_worker_file_scoped( - &scope, - &worker_ref, - &query.file_name, - &query.media_type, - &body, - ), + Some(scope) => { + if context + .as_ref() + .is_some_and(|context| context.workspace_id != scope.workspace_id) + { + return Err(RuntimeHttpRestError::new( + StatusCode::FORBIDDEN, + "uploaded_file_workspace_mismatch", + "uploaded file context does not match the authenticated Workspace", + )); + } + match context.as_ref() { + Some(context) => state.runtime.upload_worker_file_with_context_scoped( + &scope, + &worker_ref, + &query.file_name, + &query.media_type, + &body, + context, + ), + None => state.runtime.upload_worker_file_scoped( + &scope, + &worker_ref, + &query.file_name, + &query.media_type, + &body, + ), + } + } None => state.runtime.upload_worker_file( &worker_ref, &query.file_name, diff --git a/crates/worker-runtime/src/lib.rs b/crates/worker-runtime/src/lib.rs index 44817026..8f481903 100644 --- a/crates/worker-runtime/src/lib.rs +++ b/crates/worker-runtime/src/lib.rs @@ -33,3 +33,4 @@ pub mod working_directory; pub use fs_store::{FsRuntimeStore, FsRuntimeStoreOptions}; pub use management::RuntimeOptions; pub use runtime::{Runtime, RuntimeWorkspaceScope}; +pub use session_store::UploadedFileUploadContext; diff --git a/crates/worker-runtime/src/runtime.rs b/crates/worker-runtime/src/runtime.rs index 31394209..dad770a2 100644 --- a/crates/worker-runtime/src/runtime.rs +++ b/crates/worker-runtime/src/runtime.rs @@ -1224,6 +1224,41 @@ impl Runtime { file_name: &str, media_type: &str, content: &[u8], + ) -> Result { + self.upload_worker_file_inner(worker_ref, file_name, media_type, content, None) + } + + pub fn upload_worker_file_with_context_scoped( + &self, + scope: &RuntimeWorkspaceScope, + worker_ref: &WorkerRef, + file_name: &str, + media_type: &str, + content: &[u8], + context: &session_store::UploadedFileUploadContext, + ) -> Result { + self.ensure_worker_in_workspace(scope, worker_ref)?; + self.upload_worker_file_inner(worker_ref, file_name, media_type, content, Some(context)) + } + + pub fn upload_worker_file_with_context( + &self, + worker_ref: &WorkerRef, + file_name: &str, + media_type: &str, + content: &[u8], + context: &session_store::UploadedFileUploadContext, + ) -> Result { + self.upload_worker_file_inner(worker_ref, file_name, media_type, content, Some(context)) + } + + fn upload_worker_file_inner( + &self, + worker_ref: &WorkerRef, + file_name: &str, + media_type: &str, + content: &[u8], + context: Option<&session_store::UploadedFileUploadContext>, ) -> Result { let (backend, handle) = { let state = self.lock()?; @@ -1244,7 +1279,7 @@ impl Runtime { } }; backend - .upload_file(&handle, file_name, media_type, content) + .upload_file(&handle, file_name, media_type, content, context) .map_err(|result| RuntimeError::WorkerExecutionRejected { worker_id: worker_ref.worker_id.clone(), operation: result.operation, diff --git a/crates/worker-runtime/src/worker_backend.rs b/crates/worker-runtime/src/worker_backend.rs index 23275a29..e9141379 100644 --- a/crates/worker-runtime/src/worker_backend.rs +++ b/crates/worker-runtime/src/worker_backend.rs @@ -2031,19 +2031,24 @@ where file_name: &str, media_type: &str, content: &[u8], + context: Option<&session_store::UploadedFileUploadContext>, ) -> Result { let (worker, _, _) = self.get_execution(handle).map_err(|mut result| { result.operation = WorkerExecutionOperation::UploadFile; result })?; - worker - .upload_file(file_name, media_type, content) - .map_err(|error| { - WorkerExecutionResult::rejected( - WorkerExecutionOperation::UploadFile, - format!("uploaded_file_rejected: {error}"), - ) - }) + let uploaded = match context { + Some(context) => { + worker.upload_file_with_context(file_name, media_type, content, context) + } + None => worker.upload_file(file_name, media_type, content), + }; + uploaded.map_err(|error| { + WorkerExecutionResult::rejected( + WorkerExecutionOperation::UploadFile, + format!("uploaded_file_rejected: {error}"), + ) + }) } fn delete_uploaded_file( @@ -2145,25 +2150,6 @@ where ), ); } - let cleanup_handle = match self.workers.lock() { - Ok(workers) => workers - .get(handle.worker_ref()) - .map(|execution| execution.handle.clone()), - Err(_) => { - return WorkerExecutionResult::errored( - WorkerExecutionOperation::Stop, - "worker adapter registry lock is poisoned", - ); - } - }; - if let Some(worker) = cleanup_handle - && let Err(error) = worker.delete_uncommitted_uploaded_files() - { - return WorkerExecutionResult::errored( - WorkerExecutionOperation::Stop, - format!("uploaded_file_cleanup_failed: {error}"), - ); - } let execution = match self.workers.lock() { Ok(mut workers) => workers.remove(handle.worker_ref()), Err(_) => { @@ -2179,6 +2165,7 @@ where "execution handle does not reference a live Worker", ); }; + let artifact_cleanup = execution.handle.clone(); let shutdown = execution.shutdown.clone(); let result = self.send_method( WorkerExecutionOperation::Stop, @@ -2198,7 +2185,13 @@ where } Ok(()) }) { - Ok(()) => result, + Ok(()) => match artifact_cleanup.delete_uncommitted_uploaded_files() { + Ok(_) => result, + Err(error) => WorkerExecutionResult::errored( + WorkerExecutionOperation::Stop, + format!("uploaded_file_cleanup_failed: {error}"), + ), + }, Err(message) => WorkerExecutionResult::errored(WorkerExecutionOperation::Stop, message), } } diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index cb2de8fd..5d7ae971 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -80,6 +80,23 @@ impl WorkerHandle { ) } + pub fn upload_file_with_context( + &self, + file_name: &str, + media_type: &str, + content: &[u8], + context: &session_store::UploadedFileUploadContext, + ) -> Result { + self.artifact_store.write_uploaded_file_with_context( + self.session_id, + file_name, + media_type, + content, + context, + session_store::UploadedFileLimits::default(), + ) + } + pub fn delete_uploaded_file( &self, artifact_id: &str, diff --git a/crates/workspace-server/src/hosts.rs b/crates/workspace-server/src/hosts.rs index c6a3b5af..cc15ba3b 100644 --- a/crates/workspace-server/src/hosts.rs +++ b/crates/workspace-server/src/hosts.rs @@ -1048,6 +1048,7 @@ pub trait WorkspaceWorkerRuntime: Send + Sync { _file_name: &str, _media_type: &str, _content: &[u8], + _context: Option<&worker_runtime::UploadedFileUploadContext>, ) -> Result { Err(RuntimeRegistryError::RuntimeOperationFailed { runtime_id: self.runtime_id().to_string(), @@ -1578,6 +1579,7 @@ impl RuntimeRegistry { file_name: &str, media_type: &str, content: &[u8], + context: Option<&worker_runtime::UploadedFileUploadContext>, ) -> Result { let runtime_id = worker.runtime_id.as_str(); let worker_id = worker.worker_id.as_str(); @@ -1592,7 +1594,7 @@ impl RuntimeRegistry { lookup.diagnostics, )); } - runtime.upload_worker_file(worker_id, file_name, media_type, content) + runtime.upload_worker_file(worker_id, file_name, media_type, content, context) } pub fn delete_worker_uploaded_file( @@ -2588,19 +2590,30 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime { file_name: &str, media_type: &str, content: &[u8], + context: Option<&worker_runtime::UploadedFileUploadContext>, ) -> Result { let worker_ref = self.worker_ref(worker_id) .ok_or_else(|| RuntimeRegistryError::UnknownWorker { worker: RuntimeWorkerRef::new(&self.runtime_id, worker_id), })?; - self.runtime - .upload_worker_file(&worker_ref, file_name, media_type, content) - .map_err(|error| RuntimeRegistryError::RuntimeOperationFailed { - runtime_id: self.runtime_id.clone(), - code: "embedded_worker_file_upload_failed".to_string(), - message: error.to_string(), - }) + let uploaded = match context { + Some(context) => self.runtime.upload_worker_file_with_context( + &worker_ref, + file_name, + media_type, + content, + context, + ), + None => self + .runtime + .upload_worker_file(&worker_ref, file_name, media_type, content), + }; + uploaded.map_err(|error| RuntimeRegistryError::RuntimeOperationFailed { + runtime_id: self.runtime_id.clone(), + code: "embedded_worker_file_upload_failed".to_string(), + message: error.to_string(), + }) } fn delete_worker_uploaded_file( @@ -3665,13 +3678,24 @@ impl WorkspaceWorkerRuntime for RemoteWorkerRuntime { file_name: &str, media_type: &str, content: &[u8], + context: Option<&worker_runtime::UploadedFileUploadContext>, ) -> Result { - let path = format!( + let mut path = format!( "/v1/workers/{}/attachments?file_name={}&media_type={}", url_path_segment_encode(worker_id), url_query_value_encode(file_name), url_query_value_encode(media_type), ); + if let Some(context) = context { + path.push_str(&format!( + "&upload_id={}&principal_id={}&workspace_id={}&runtime_id={}&owner_worker_id={}", + url_query_value_encode(&context.upload_id), + url_query_value_encode(&context.principal_id), + url_query_value_encode(&context.workspace_id), + url_query_value_encode(&context.runtime_id), + url_query_value_encode(&context.worker_id), + )); + } self.post_bytes::(&path, content) .map(|response| response.file) .map_err(|diagnostic| RuntimeRegistryError::RuntimeOperationFailed { diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index d938bb9f..e8f6d1ba 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -502,6 +502,73 @@ async fn close_worker_workdir_sessions( } } +const ATTACHMENT_UPLOAD_GRANT_TTL_MS: u64 = 5 * 60 * 1_000; + +#[derive(Clone)] +struct AttachmentUploadGrant { + workspace_id: String, + runtime_id: String, + worker_id: String, + account_id: String, + file_name: String, + media_type: String, + expires_at_ms: u64, + state: AttachmentUploadGrantState, +} + +#[derive(Clone)] +enum AttachmentUploadGrantState { + Pending, + Uploading, + Cancelled, + Completed(protocol::UploadedFileRef), +} + +impl AttachmentUploadGrant { + fn claim( + &mut self, + workspace_id: &str, + runtime_id: &str, + worker_id: &str, + account_id: &str, + now_ms: u64, + body_sha256: &str, + ) -> Result> { + if self.workspace_id != workspace_id + || self.runtime_id != runtime_id + || self.worker_id != worker_id + || self.account_id != account_id + { + return Err(Error::WorkspacePermissionDenied( + "attachment upload grant scope mismatch".into(), + )); + } + if self.expires_at_ms <= now_ms { + return Err(Error::InvalidInput( + "attachment upload grant expired".into(), + )); + } + match &self.state { + AttachmentUploadGrantState::Pending => { + self.state = AttachmentUploadGrantState::Uploading; + Ok(None) + } + AttachmentUploadGrantState::Uploading => Err(Error::RepositoryConflict( + "attachment upload is already in progress".into(), + )), + AttachmentUploadGrantState::Cancelled => Err(Error::RepositoryConflict( + "attachment upload grant was cancelled".into(), + )), + AttachmentUploadGrantState::Completed(file) if file.sha256 == body_sha256 => { + Ok(Some(file.clone())) + } + AttachmentUploadGrantState::Completed(_) => Err(Error::RepositoryConflict( + "attachment upload grant was already consumed".into(), + )), + } + } +} + #[derive(Clone)] pub struct WorkspaceApi { pub(crate) config: ServerConfig, @@ -524,6 +591,7 @@ pub struct WorkspaceApi { workdir_remove_locks: Arc>>>>, workdir_remove_attempt_owner: WorkdirRemovalAttemptOwner, worker_control_locks: Arc>>>>, + attachment_upload_grants: Arc>>, } #[derive(Clone)] @@ -1645,6 +1713,7 @@ impl WorkspaceApi { workdir_remove_locks: Arc::new(Mutex::new(HashMap::new())), workdir_remove_attempt_owner: current_workdir_removal_attempt_owner()?, worker_control_locks: Arc::new(Mutex::new(HashMap::new())), + attachment_upload_grants: Arc::new(Mutex::new(HashMap::new())), }; if let Some(dispatcher) = worker_remove_dispatcher { dispatcher @@ -2681,21 +2750,19 @@ fn build_inner_router(api: WorkspaceApi) -> Router { "/api/runtimes/{runtime_id}/workers/{worker_id}/input", post(send_runtime_worker_input), ) - .route( - "/api/runtimes/{runtime_id}/workers/{worker_id}/attachments", - post(upload_runtime_worker_file).layer(DefaultBodyLimit::max(MAX_WORKER_FILE_UPLOAD_BYTES)), - ) - .route( - "/api/runtimes/{runtime_id}/workers/{worker_id}/attachments/{artifact_id}", - delete(delete_runtime_worker_uploaded_file), - ) .route( "/api/w/{workspace_id}/runtimes/{runtime_id}/workers/{worker_id}/input", post(scoped_send_runtime_worker_input), ) .route( - "/api/w/{workspace_id}/runtimes/{runtime_id}/workers/{worker_id}/attachments", - post(scoped_upload_runtime_worker_file).layer(DefaultBodyLimit::max(MAX_WORKER_FILE_UPLOAD_BYTES)), + "/api/w/{workspace_id}/runtimes/{runtime_id}/workers/{worker_id}/attachment-upload-grants", + post(scoped_create_attachment_upload_grant), + ) + .route( + "/api/w/{workspace_id}/runtimes/{runtime_id}/workers/{worker_id}/attachment-uploads/{upload_id}", + put(scoped_upload_runtime_worker_file) + .delete(scoped_cancel_attachment_upload) + .layer(DefaultBodyLimit::max(MAX_WORKER_FILE_UPLOAD_BYTES)), ) .route( "/api/w/{workspace_id}/runtimes/{runtime_id}/workers/{worker_id}/attachments/{artifact_id}", @@ -10756,6 +10823,8 @@ async fn scoped_execute_runtime_cleanup( struct WorkerFileUploadQuery { file_name: String, media_type: String, + #[serde(default)] + upload_id: Option, } #[derive(Debug, Serialize)] @@ -10768,20 +10837,225 @@ struct WorkerFileDeleteResponse { deleted: bool, } -async fn scoped_upload_runtime_worker_file( +#[derive(Debug, Serialize, Deserialize)] +struct AttachmentUploadGrantResponse { + upload_id: String, + expires_at_ms: u64, +} + +#[derive(Debug, Deserialize)] +struct ScopedAttachmentUploadPath { + workspace_id: String, + runtime_id: String, + worker_id: String, + upload_id: String, +} + +fn attachment_body_sha256(body: &[u8]) -> String { + Sha256::digest(body) + .iter() + .map(|byte| format!("{byte:02x}")) + .collect() +} + +fn attachment_upload_now_ms() -> Result { + u64::try_from( + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map_err(|_| Error::InvalidInput("system clock predates UNIX epoch".into()))? + .as_millis(), + ) + .map_err(|_| Error::InvalidInput("attachment upload timestamp overflow".into())) +} + +async fn scoped_create_attachment_upload_grant( State(api): State, + Extension(actor): Extension, AxumPath(path): AxumPath, Query(query): Query, +) -> ApiResult> { + validate_workspace_scope(&api, &path.workspace_id)?; + resolve_workspace_worker_reference(&api, &path.worker.runtime_id, &path.worker.worker_id)?; + let now = attachment_upload_now_ms()?; + let expires_at_ms = now.saturating_add(ATTACHMENT_UPLOAD_GRANT_TTL_MS); + let upload_id = match query.upload_id { + Some(upload_id) => Uuid::parse_str(&upload_id) + .map_err(|_| Error::InvalidInput("attachment upload id must be a UUID".into()))? + .to_string(), + None => Uuid::now_v7().to_string(), + }; + let mut grants = api + .attachment_upload_grants + .lock() + .map_err(|_| Error::Store("attachment upload grant lock is poisoned".into()))?; + grants.retain(|_, grant| grant.expires_at_ms >= now); + if let Some(existing) = grants.get(&upload_id) { + if existing.workspace_id == path.workspace_id + && existing.runtime_id == path.worker.runtime_id + && existing.worker_id == path.worker.worker_id + && existing.account_id == actor.account_id + && existing.file_name == query.file_name + && existing.media_type == query.media_type + { + return Ok(Json(AttachmentUploadGrantResponse { + upload_id, + expires_at_ms: existing.expires_at_ms, + })); + } + return Err(Error::RepositoryConflict( + "attachment upload id was reused with different intent".into(), + ) + .into()); + } + grants.insert( + upload_id.clone(), + AttachmentUploadGrant { + workspace_id: path.workspace_id, + runtime_id: path.worker.runtime_id, + worker_id: path.worker.worker_id, + account_id: actor.account_id, + file_name: query.file_name, + media_type: query.media_type, + expires_at_ms, + state: AttachmentUploadGrantState::Pending, + }, + ); + Ok(Json(AttachmentUploadGrantResponse { + upload_id, + expires_at_ms, + })) +} + +async fn scoped_upload_runtime_worker_file( + State(api): State, + Extension(actor): Extension, + AxumPath(path): AxumPath, body: Bytes, ) -> ApiResult> { validate_workspace_scope(&api, &path.workspace_id)?; - upload_runtime_worker_file( - State(api), - AxumPath((path.worker.runtime_id, path.worker.worker_id)), - Query(query), - body, - ) - .await + let body_sha256 = attachment_body_sha256(&body); + let grant = { + let mut grants = api + .attachment_upload_grants + .lock() + .map_err(|_| Error::Store("attachment upload grant lock is poisoned".into()))?; + let grant = grants + .get_mut(&path.upload_id) + .ok_or_else(|| Error::InvalidInput("attachment upload grant is unknown".into()))?; + if let Some(file) = grant + .claim( + &path.workspace_id, + &path.runtime_id, + &path.worker_id, + &actor.account_id, + attachment_upload_now_ms()?, + &body_sha256, + ) + .map_err(ApiError::from)? + { + return Ok(Json(WorkerFileUploadResponse { file })); + } + grant.clone() + }; + + let worker = resolve_workspace_worker_reference(&api, &path.runtime_id, &path.worker_id)?; + let upload_context = worker_runtime::UploadedFileUploadContext { + upload_id: path.upload_id.clone(), + principal_id: grant.account_id.clone(), + workspace_id: grant.workspace_id.clone(), + runtime_id: grant.runtime_id.clone(), + worker_id: grant.worker_id.clone(), + }; + let uploaded = api.runtime.upload_worker_file( + &worker, + &grant.file_name, + &grant.media_type, + &body, + Some(&upload_context), + ); + let mut grants = api + .attachment_upload_grants + .lock() + .map_err(|_| Error::Store("attachment upload grant lock is poisoned".into()))?; + let current = grants + .get_mut(&path.upload_id) + .ok_or_else(|| Error::RepositoryConflict("attachment upload grant disappeared".into()))?; + if matches!(current.state, AttachmentUploadGrantState::Cancelled) { + let uploaded_file = uploaded.ok(); + grants.remove(&path.upload_id); + drop(grants); + if let Some(file) = uploaded_file { + api.runtime + .delete_worker_uploaded_file(&worker, &file.artifact_id) + .map_err(|error| error.into_error())?; + } + return Err(Error::RepositoryConflict("attachment upload was cancelled".into()).into()); + } + match uploaded { + Ok(file) => { + current.state = AttachmentUploadGrantState::Completed(file.clone()); + Ok(Json(WorkerFileUploadResponse { file })) + } + Err(error) => { + current.state = AttachmentUploadGrantState::Pending; + Err(error.into_error().into()) + } + } +} + +#[derive(Debug, Serialize)] +struct AttachmentUploadCancelResponse { + cancelled: bool, +} + +async fn scoped_cancel_attachment_upload( + State(api): State, + Extension(actor): Extension, + AxumPath(path): AxumPath, +) -> ApiResult> { + validate_workspace_scope(&api, &path.workspace_id)?; + let completed = { + let mut grants = api + .attachment_upload_grants + .lock() + .map_err(|_| Error::Store("attachment upload grant lock is poisoned".into()))?; + let grant = grants + .get_mut(&path.upload_id) + .ok_or_else(|| Error::InvalidInput("attachment upload grant is unknown".into()))?; + if grant.workspace_id != path.workspace_id + || grant.runtime_id != path.runtime_id + || grant.worker_id != path.worker_id + || grant.account_id != actor.account_id + { + return Err(Error::WorkspacePermissionDenied( + "attachment upload grant scope mismatch".into(), + ) + .into()); + } + match &grant.state { + AttachmentUploadGrantState::Pending => { + grants.remove(&path.upload_id); + None + } + AttachmentUploadGrantState::Uploading => { + grant.state = AttachmentUploadGrantState::Cancelled; + None + } + AttachmentUploadGrantState::Cancelled => None, + AttachmentUploadGrantState::Completed(file) => { + let file = file.clone(); + grants.remove(&path.upload_id); + Some(file) + } + } + }; + if let Some(file) = completed { + let worker = resolve_workspace_worker_reference(&api, &path.runtime_id, &path.worker_id)?; + api.runtime + .delete_worker_uploaded_file(&worker, &file.artifact_id) + .map_err(|error| error.into_error())?; + } + Ok(Json(AttachmentUploadCancelResponse { cancelled: true })) } async fn scoped_delete_runtime_worker_uploaded_file( @@ -13323,20 +13597,6 @@ async fn send_runtime_worker_input( Ok(Json(result)) } -async fn upload_runtime_worker_file( - State(api): State, - AxumPath((runtime_id, worker_id)): AxumPath<(String, String)>, - Query(query): Query, - body: Bytes, -) -> ApiResult> { - let worker = resolve_workspace_worker_reference(&api, &runtime_id, &worker_id)?; - let file = api - .runtime - .upload_worker_file(&worker, &query.file_name, &query.media_type, &body) - .map_err(|err| err.into_error())?; - Ok(Json(WorkerFileUploadResponse { file })) -} - async fn delete_runtime_worker_uploaded_file( State(api): State, AxumPath((runtime_id, worker_id, artifact_id)): AxumPath<(String, String, String)>, @@ -16084,6 +16344,77 @@ mod tests { SqliteWorkspaceStore, TrustedRuntimeRecord, UserRecord, WorkspaceRecord, }; + fn completed_upload_file(sha256: &str) -> protocol::UploadedFileRef { + protocol::UploadedFileRef { + artifact_id: "019ca7c8-57b6-7f05-8edf-524147aba7b3".into(), + file_name: "notes.txt".into(), + media_type: "text/plain".into(), + created_at_ms: 1, + availability: protocol::UploadedFileAvailability::Available, + byte_len: 1, + sha256: sha256.into(), + source_entry_id: None, + } + } + + fn pending_upload_grant() -> AttachmentUploadGrant { + AttachmentUploadGrant { + workspace_id: "workspace-1".into(), + runtime_id: "runtime-1".into(), + worker_id: "worker-1".into(), + account_id: "account-1".into(), + file_name: "notes.txt".into(), + media_type: "text/plain".into(), + expires_at_ms: 100, + state: AttachmentUploadGrantState::Pending, + } + } + + #[test] + fn attachment_upload_grant_is_scoped_expiring_and_idempotent() { + let mut grant = pending_upload_grant(); + assert!(matches!( + grant.claim("workspace-2", "runtime-1", "worker-1", "account-1", 1, "a"), + Err(Error::WorkspacePermissionDenied(_)) + )); + assert!(matches!( + pending_upload_grant().claim( + "workspace-1", + "runtime-1", + "worker-1", + "account-1", + 101, + "a" + ), + Err(Error::InvalidInput(_)) + )); + + assert!( + grant + .claim("workspace-1", "runtime-1", "worker-1", "account-1", 1, "a") + .unwrap() + .is_none() + ); + assert!(matches!( + grant.claim("workspace-1", "runtime-1", "worker-1", "account-1", 1, "a"), + Err(Error::RepositoryConflict(_)) + )); + + grant.state = AttachmentUploadGrantState::Completed(completed_upload_file("a")); + assert_eq!( + grant + .claim("workspace-1", "runtime-1", "worker-1", "account-1", 1, "a") + .unwrap() + .unwrap() + .artifact_id, + "019ca7c8-57b6-7f05-8edf-524147aba7b3" + ); + assert!(matches!( + grant.claim("workspace-1", "runtime-1", "worker-1", "account-1", 1, "b"), + Err(Error::RepositoryConflict(_)) + )); + } + #[tokio::test] async fn command_session_survives_attachment_refresh_until_worker_revocation() { let directory = tempfile::tempdir().unwrap(); diff --git a/web/workspace/src/lib/workspace/console/composer-attachments.ts b/web/workspace/src/lib/workspace/console/composer-attachments.ts index e50dacad..45d179a9 100644 --- a/web/workspace/src/lib/workspace/console/composer-attachments.ts +++ b/web/workspace/src/lib/workspace/console/composer-attachments.ts @@ -5,15 +5,18 @@ export const MAX_FILES_PER_SUBMISSION = 8; export type AttachmentUploadState = "uploading" | "uploaded" | "failed"; +export type AttachmentUploadHandle = { abort(): void }; + export type ComposerAttachment = { id: number; file: File; uploadPath: string; + uploadId: string; state: AttachmentUploadState; progress: number; reference: UploadedFileRef | null; error: string | null; - request: XMLHttpRequest | null; + request: AttachmentUploadHandle | null; }; export function acceptedAttachmentMediaType(mediaType: string): boolean { @@ -43,42 +46,97 @@ export type AttachmentUploadCallbacks = { }; export function uploadAttachment( - path: string, + workerPath: string, file: File, + uploadId: string, callbacks: AttachmentUploadCallbacks, -): XMLHttpRequest { - const request = new XMLHttpRequest(); +): AttachmentUploadHandle { + let activeRequest: XMLHttpRequest | null = null; + let aborted = false; + const handle: AttachmentUploadHandle = { + abort() { + aborted = true; + activeRequest?.abort(); + void fetch( + `${workerPath}/attachment-uploads/${encodeURIComponent(uploadId)}`, + { method: "DELETE" }, + ).catch(() => undefined); + }, + }; const query = new URLSearchParams({ file_name: file.name, media_type: file.type, + upload_id: uploadId, }); - request.open("POST", `${path}?${query.toString()}`); - request.setRequestHeader("content-type", "application/octet-stream"); - request.upload.addEventListener("progress", (event) => { - if (event.lengthComputable && event.total > 0) { - callbacks.progress(Math.min(1, event.loaded / event.total)); - } - }); - request.addEventListener("load", () => { - if (request.status < 200 || request.status >= 300) { - callbacks.failed(`Upload failed (${request.status}).`); + const grantRequest = new XMLHttpRequest(); + activeRequest = grantRequest; + grantRequest.open( + "POST", + `${workerPath}/attachment-upload-grants?${query.toString()}`, + ); + grantRequest.addEventListener("load", () => { + if (aborted) return; + if (grantRequest.status < 200 || grantRequest.status >= 300) { + callbacks.failed(`Upload grant failed (${grantRequest.status}).`); return; } + let uploadId: string; try { - const parsed: unknown = JSON.parse(request.responseText); - if (!isUploadedFileResponse(parsed)) { - callbacks.failed("Upload returned an invalid attachment reference."); + const parsed: unknown = JSON.parse(grantRequest.responseText); + if (!isUploadGrantResponse(parsed)) { + callbacks.failed("Upload grant returned an invalid response."); return; } - callbacks.complete(parsed.file); + uploadId = parsed.upload_id; } catch { - callbacks.failed("Upload returned an invalid response."); + callbacks.failed("Upload grant returned an invalid response."); + return; } + + const uploadRequest = new XMLHttpRequest(); + activeRequest = uploadRequest; + uploadRequest.open( + "PUT", + `${workerPath}/attachment-uploads/${encodeURIComponent(uploadId)}`, + ); + uploadRequest.setRequestHeader("content-type", "application/octet-stream"); + uploadRequest.upload.addEventListener("progress", (event) => { + if (event.lengthComputable && event.total > 0) { + callbacks.progress(Math.min(1, event.loaded / event.total)); + } + }); + uploadRequest.addEventListener("load", () => { + if (uploadRequest.status < 200 || uploadRequest.status >= 300) { + callbacks.failed(`Upload failed (${uploadRequest.status}).`); + return; + } + try { + const parsed: unknown = JSON.parse(uploadRequest.responseText); + if (!isUploadedFileResponse(parsed)) { + callbacks.failed("Upload returned an invalid attachment reference."); + return; + } + callbacks.complete(parsed.file); + } catch { + callbacks.failed("Upload returned an invalid response."); + } + }); + uploadRequest.addEventListener("error", () => callbacks.failed("Upload failed.")); + uploadRequest.addEventListener("abort", () => callbacks.failed("Upload cancelled.")); + uploadRequest.send(file); }); - request.addEventListener("error", () => callbacks.failed("Upload failed.")); - request.addEventListener("abort", () => callbacks.failed("Upload cancelled.")); - request.send(file); - return request; + grantRequest.addEventListener("error", () => callbacks.failed("Upload grant failed.")); + grantRequest.addEventListener("abort", () => callbacks.failed("Upload cancelled.")); + grantRequest.send(); + return handle; +} + +function isUploadGrantResponse( + value: unknown, +): value is { upload_id: string; expires_at_ms: number } { + return !!value && typeof value === "object" && + "upload_id" in value && typeof value.upload_id === "string" && + "expires_at_ms" in value && typeof value.expires_at_ms === "number"; } function isUploadedFileResponse( diff --git a/web/workspace/src/routes/w/[workspaceId]/runtimes/[runtimeId]/workers/[workerId]/console/+page.svelte b/web/workspace/src/routes/w/[workspaceId]/runtimes/[runtimeId]/workers/[workerId]/console/+page.svelte index 6b20031e..7d2ab430 100644 --- a/web/workspace/src/routes/w/[workspaceId]/runtimes/[runtimeId]/workers/[workerId]/console/+page.svelte +++ b/web/workspace/src/routes/w/[workspaceId]/runtimes/[runtimeId]/workers/[workerId]/console/+page.svelte @@ -599,7 +599,7 @@ attachment.request?.abort(); if (attachment.reference) { void fetch( - `${attachment.uploadPath}/${encodeURIComponent(attachment.reference.artifact_id)}`, + `${attachment.uploadPath}/attachments/${encodeURIComponent(attachment.reference.artifact_id)}`, { method: "DELETE" }, ).catch(() => undefined); } @@ -639,7 +639,7 @@ } function attachmentPath(): string { - return `/api/w/${encodeURIComponent(workspaceId)}/runtimes/${encodeURIComponent(runtimeId)}/workers/${encodeURIComponent(workerId)}/attachments`; + return `/api/w/${encodeURIComponent(workspaceId)}/runtimes/${encodeURIComponent(runtimeId)}/workers/${encodeURIComponent(workerId)}`; } function updateAttachment(id: number, update: Partial): void { @@ -661,7 +661,11 @@ error: null, request: null, }); - const request = uploadAttachment(attachment.uploadPath, attachment.file, { + const request = uploadAttachment( + attachment.uploadPath, + attachment.file, + attachment.uploadId, + { progress: (progress) => updateAttachment(attachment.id, { progress }), complete: (reference) => updateAttachment(attachment.id, { @@ -688,6 +692,7 @@ id: nextAttachmentId++, file, uploadPath: attachmentPath(), + uploadId: crypto.randomUUID(), state: "uploading", progress: 0, reference: null, @@ -703,7 +708,7 @@ attachment.request?.abort(); attachments = attachments.filter((candidate) => candidate.id !== attachment.id); if (attachment.reference) { - await fetch(`${attachment.uploadPath}/${encodeURIComponent(attachment.reference.artifact_id)}`, { + await fetch(`${attachment.uploadPath}/attachments/${encodeURIComponent(attachment.reference.artifact_id)}`, { method: "DELETE", }).catch(() => undefined); }