From 09a33e72834665314cf3e5d0e39ec88668f72bef Mon Sep 17 00:00:00 2001 From: Hare Date: Thu, 3 Sep 2026 04:58:11 +0900 Subject: [PATCH] feat: add session-owned uploaded file attachments --- crates/protocol/src/lib.rs | 83 ++++++ crates/protocol/src/typescript.rs | 5 +- crates/session-store/src/fs_store.rs | 152 +++++++++- crates/session-store/src/lib.rs | 5 + crates/session-store/src/paste_artifact.rs | 55 ++-- crates/session-store/src/store.rs | 75 ++++- crates/session-store/src/uploaded_file.rs | 299 ++++++++++++++++++++ crates/worker-runtime/src/execution.rs | 47 ++- crates/worker-runtime/src/http_server.rs | 95 ++++++- crates/worker-runtime/src/runtime.rs | 99 +++++++ crates/worker-runtime/src/worker_backend.rs | 45 +++ crates/worker/src/controller.rs | 32 ++- crates/worker/src/paste_artifact_tool.rs | 116 +++++++- crates/worker/src/worker.rs | 85 ++++++ crates/workspace-server/src/hosts.rs | 164 +++++++++++ crates/workspace-server/src/server.rs | 94 +++++- web/workspace/src/lib/generated/protocol.ts | 6 +- 17 files changed, 1416 insertions(+), 41 deletions(-) create mode 100644 crates/session-store/src/uploaded_file.rs diff --git a/crates/protocol/src/lib.rs b/crates/protocol/src/lib.rs index 9eb50b13..9f0ab12c 100644 --- a/crates/protocol/src/lib.rs +++ b/crates/protocol/src/lib.rs @@ -251,6 +251,48 @@ pub struct PasteArtifactRef { pub source_entry_id: String, } +/// Availability recorded for an uploaded client-local file. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +#[cfg_attr(feature = "json-schema", derive(schemars::JsonSchema))] +#[serde(rename_all = "snake_case")] +pub enum UploadedFileAvailability { + Available, + Unavailable, + IntegrityFailed, +} + +impl UploadedFileAvailability { + pub fn as_str(self) -> &'static str { + match self { + Self::Available => "available", + Self::Unavailable => "unavailable", + Self::IntegrityFailed => "integrity_failed", + } + } +} + +/// Session-owned immutable reference to a client-local uploaded file. +/// +/// Upload transports return an unbound reference. Worker fills +/// `source_entry_id` immediately before the containing user input is committed; +/// committed Session Log and public snapshot records therefore always retain +/// the durable source-entry identity without storing the file body. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +#[cfg_attr(feature = "json-schema", derive(schemars::JsonSchema))] +pub struct UploadedFileRef { + pub artifact_id: String, + pub file_name: String, + pub media_type: String, + pub created_at_ms: u64, + pub availability: UploadedFileAvailability, + pub byte_len: u64, + pub sha256: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub source_entry_id: Option, +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[cfg_attr(feature = "typescript", derive(ts_rs::TS))] #[cfg_attr(feature = "json-schema", derive(schemars::JsonSchema))] @@ -272,6 +314,10 @@ pub enum Segment { /// committing input. Clients may receive this in history/event projections; /// the body is intentionally absent. PasteArtifact { artifact: PasteArtifactRef }, + /// Client-local file uploaded into the owning Worker session before submit. + /// The Session Log stores only this immutable reference, never file bytes or + /// the client's local path. + UploadedFile { file: UploadedFileRef }, /// `@` file-system reference. Worker resolves readable files to /// `[File: ]` attachments and readable normal directories to shallow /// `[Dir: ]` listings; the flattened user text keeps the literal @@ -327,6 +373,20 @@ impl Segment { artifact.sha256 ); } + Segment::UploadedFile { file } => { + use std::fmt::Write as _; + let _ = write!( + out, + "[Attached file {} stored as input artifact {}: {} bytes, {}, {}, created at {} ms, sha256 {}; use SearchInputArtifact and ReadInputArtifact for supported text content]", + file.file_name, + file.artifact_id, + file.byte_len, + file.media_type, + file.availability.as_str(), + file.created_at_ms, + file.sha256 + ); + } Segment::FileRef { path } => { out.push('@'); out.push_str(path); @@ -1305,6 +1365,29 @@ mod tests { assert!(!projected.contains("pasted body")); } + #[test] + fn uploaded_file_segment_roundtrips_without_path_or_body() { + let file = UploadedFileRef { + artifact_id: "019ca7c8-57b6-7f05-8edf-524147aba7b3".to_string(), + file_name: "notes.md".to_string(), + media_type: "text/markdown".to_string(), + created_at_ms: 1_700_000_000_001, + availability: UploadedFileAvailability::Available, + byte_len: 128, + sha256: "b".repeat(64), + source_entry_id: Some("entry-2".to_string()), + }; + let segment = Segment::UploadedFile { file: file.clone() }; + let json = serde_json::to_string(&segment).unwrap(); + assert!(!json.contains("/home/user/private")); + assert!(!json.contains("file body")); + assert_eq!(serde_json::from_str::(&json).unwrap(), segment); + let projected = Segment::flatten_to_text(&[segment]); + assert!(projected.contains("notes.md")); + assert!(projected.contains(&file.artifact_id)); + assert!(projected.contains("ReadInputArtifact")); + } + #[test] fn method_run_flow_segment_roundtrip() { let method = Method::Run { diff --git a/crates/protocol/src/typescript.rs b/crates/protocol/src/typescript.rs index 4e3e9db7..bdac40ca 100644 --- a/crates/protocol/src/typescript.rs +++ b/crates/protocol/src/typescript.rs @@ -11,7 +11,8 @@ use crate::{ PasteArtifactRef, Permission, RewindSummary, RewindTarget, RewindTargetId, RunResult, ScopeRule, Segment, SessionContentPart, SessionEntryProvenance, SessionMessageRole, SessionSnapshot, SessionSnapshotEntry, SessionSnapshotEntryData, SessionToolAttachment, - ToolResultDisposition, TurnResult, WorkerEvent, WorkerStatus, + ToolResultDisposition, TurnResult, UploadedFileAvailability, UploadedFileRef, WorkerEvent, + WorkerStatus, subscription::{ EventSubscriptionSelector, SubscriptionEvent, SubscriptionEventPayload, SubscriptionFrame, SubscriptionFramePayload, SubscriptionId, SubscriptionRejectionCode, SubscriptionRequest, @@ -59,6 +60,8 @@ pub fn generated_protocol_types() -> String { push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); + push_decl::(&cfg, &mut output); + push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); diff --git a/crates/session-store/src/fs_store.rs b/crates/session-store/src/fs_store.rs index 60bf04a1..4bae850f 100644 --- a/crates/session-store/src/fs_store.rs +++ b/crates/session-store/src/fs_store.rs @@ -19,8 +19,12 @@ use crate::event_trace::TraceEntry; use crate::paste_artifact::{read_from_dir, write_to_dir}; use crate::segment_log::LogEntry; use crate::store::{Store, StoreError}; -use crate::{PasteArtifactLimits, SegmentId, SessionId}; -use protocol::PasteArtifactRef; +use crate::uploaded_file::{ + bind_uploaded_file, delete_uploaded_file, read_uploaded_file, read_uploaded_file_by_id, + write_uploaded_file, +}; +use crate::{PasteArtifactLimits, SegmentId, SessionId, UploadedFileLimits}; +use protocol::{PasteArtifactRef, UploadedFileRef}; use std::fs; use std::io::{Read, Seek, SeekFrom, Write}; use std::path::{Path, PathBuf}; @@ -389,6 +393,72 @@ impl Store for FsStore { read_from_dir(&self.paste_artifact_dir(session_id), artifact_id) } + fn write_uploaded_file( + &self, + session_id: SessionId, + file_name: &str, + media_type: &str, + content: &[u8], + 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, + limits, + ) + } + + fn read_uploaded_file( + &self, + session_id: SessionId, + reference: &UploadedFileRef, + ) -> Result, StoreError> { + read_uploaded_file(&self.paste_artifact_dir(session_id), reference) + } + + fn read_uploaded_file_by_id( + &self, + session_id: SessionId, + artifact_id: &str, + ) -> Result<(UploadedFileRef, Vec), StoreError> { + read_uploaded_file_by_id(&self.paste_artifact_dir(session_id), artifact_id) + } + + fn bind_uploaded_file( + &self, + session_id: SessionId, + reference: &UploadedFileRef, + source_entry_id: &str, + ) -> Result { + let _guard = self + .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, + ) + } + + fn delete_uploaded_file( + &self, + session_id: SessionId, + artifact_id: &str, + ) -> Result { + let _guard = self + .append_lock + .lock() + .map_err(|_| std::io::Error::other("session store append lock was poisoned"))?; + delete_uploaded_file(&self.paste_artifact_dir(session_id), artifact_id) + } + fn append_trace( &self, session_id: SessionId, @@ -548,6 +618,84 @@ mod tests { ); } + #[test] + fn uploaded_files_are_session_scoped_integrity_checked_and_removable() { + let tmp = tempfile::TempDir::new().unwrap(); + let store = FsStore::new(tmp.path()).unwrap(); + let owner = new_session_id(); + let other = new_session_id(); + let limits = UploadedFileLimits { + max_file_bytes: 16, + max_session_bytes: 16, + }; + let reference = store + .write_uploaded_file(owner, "notes.txt", "text/plain", b"hello", limits) + .unwrap(); + + assert_eq!(reference.file_name, "notes.txt"); + assert_eq!(reference.media_type, "text/plain"); + assert_eq!(reference.byte_len, 5); + assert_eq!(reference.source_entry_id, None); + assert_eq!( + store.read_uploaded_file(owner, &reference).unwrap(), + b"hello" + ); + assert!(store.read_uploaded_file(other, &reference).is_err()); + + let mut forged = reference.clone(); + forged.file_name = "other.txt".to_string(); + assert!(matches!( + store.read_uploaded_file(owner, &forged), + Err(StoreError::ArtifactIntegrityMismatch) + )); + assert!( + store + .delete_uploaded_file(owner, &reference.artifact_id) + .unwrap() + ); + assert!( + !store + .delete_uploaded_file(owner, &reference.artifact_id) + .unwrap() + ); + assert!(store.read_uploaded_file(owner, &reference).is_err()); + } + + #[test] + fn uploaded_file_validation_and_shared_quota_fail_closed() { + let tmp = tempfile::TempDir::new().unwrap(); + let store = FsStore::new(tmp.path()).unwrap(); + let session_id = new_session_id(); + let limits = UploadedFileLimits { + max_file_bytes: 8, + max_session_bytes: 8, + }; + assert!(matches!( + store.write_uploaded_file(session_id, "../secret", "text/plain", b"x", limits), + Err(StoreError::InvalidUploadedFileName) + )); + assert!(matches!( + store.write_uploaded_file(session_id, "notes.txt", "not a type", b"x", limits), + Err(StoreError::InvalidUploadedFileMediaType) + )); + store + .write_paste_artifact( + session_id, + "entry-1", + "1234", + PasteArtifactLimits { + max_artifact_bytes: 8, + max_session_bytes: 8, + max_session_artifacts: 4, + }, + ) + .unwrap(); + assert!(matches!( + store.write_uploaded_file(session_id, "notes.txt", "text/plain", b"56789", limits), + Err(StoreError::ArtifactQuotaExceeded) + )); + } + #[test] fn paste_artifact_limits_and_corruption_fail_closed() { let tmp = tempfile::TempDir::new().unwrap(); diff --git a/crates/session-store/src/lib.rs b/crates/session-store/src/lib.rs index 87b59ae6..d2ea196f 100644 --- a/crates/session-store/src/lib.rs +++ b/crates/session-store/src/lib.rs @@ -41,6 +41,7 @@ pub mod segment; pub mod segment_log; pub mod store; pub mod system_item; +pub mod uploaded_file; pub mod worker_metadata; pub mod worker_session_store; @@ -66,6 +67,10 @@ pub use store::{Store, StoreError}; pub use system_item::{ PromptRenderProvenance, SystemItem, SystemReminder, SystemReminderSource, render_worker_event, }; +pub use uploaded_file::{ + DEFAULT_MAX_FILES_PER_SUBMISSION, DEFAULT_MAX_SESSION_ARTIFACT_BYTES, + DEFAULT_MAX_UPLOADED_FILE_BYTES, UploadedFileLimits, +}; pub use worker_metadata::{ CombinedStore, FsWorkerStore, WorkerActiveSegmentRef, WorkerAggregateStore, WorkerMetadata, WorkerMetadataStore, WorkerPeer, WorkerReclaimedChild, WorkerSpawnedChild, diff --git a/crates/session-store/src/paste_artifact.rs b/crates/session-store/src/paste_artifact.rs index c54f83b4..783d5f4c 100644 --- a/crates/session-store/src/paste_artifact.rs +++ b/crates/session-store/src/paste_artifact.rs @@ -38,6 +38,34 @@ pub(crate) struct StoredPasteArtifact { pub content: String, } +pub(crate) fn stored_paste_usage(artifact_dir: &Path) -> Result<(u64, u64), StoreError> { + if !artifact_dir.exists() { + return Ok((0, 0)); + } + let mut aggregate = 0_u64; + let mut artifact_count = 0_u64; + for entry in fs::read_dir(artifact_dir)? { + let path = entry?.path(); + let Some(name) = path.file_name().and_then(|value| value.to_str()) else { + continue; + }; + if !name.ends_with(".json") || name.ends_with(".file.json") { + continue; + } + let stored: StoredPasteArtifact = serde_json::from_slice(&fs::read(&path)?)?; + verify(&stored, &stored.reference.artifact_id)?; + artifact_count = artifact_count.checked_add(1).ok_or_else(|| { + StoreError::PasteArtifactLimit("session artifact count overflow".to_string()) + })?; + aggregate = aggregate + .checked_add(stored.reference.byte_len) + .ok_or_else(|| { + StoreError::PasteArtifactLimit("session aggregate size overflow".to_string()) + })?; + } + Ok((aggregate, artifact_count)) +} + pub(crate) fn write_to_dir( artifact_dir: &Path, source_entry_id: &str, @@ -58,24 +86,15 @@ pub(crate) fn write_to_dir( .write(true) .open(artifact_dir.join(".aggregate.lock"))?; FileExt::lock_exclusive(&aggregate_lock)?; - let mut aggregate = 0_u64; - let mut artifact_count = 0_u64; - for entry in fs::read_dir(artifact_dir)? { - let path = entry?.path(); - if path.extension().and_then(|value| value.to_str()) != Some("json") { - continue; - } - let stored: StoredPasteArtifact = serde_json::from_slice(&fs::read(&path)?)?; - verify(&stored, &stored.reference.artifact_id)?; - artifact_count = artifact_count.checked_add(1).ok_or_else(|| { - StoreError::PasteArtifactLimit("session artifact count overflow".to_string()) - })?; - aggregate = aggregate - .checked_add(stored.reference.byte_len) - .ok_or_else(|| { - StoreError::PasteArtifactLimit("session aggregate size overflow".to_string()) - })?; - } + let (paste_bytes, artifact_count) = stored_paste_usage(artifact_dir)?; + let (uploaded_bytes, uploaded_count) = + crate::uploaded_file::stored_uploaded_file_usage(artifact_dir)?; + let aggregate = paste_bytes.checked_add(uploaded_bytes).ok_or_else(|| { + StoreError::PasteArtifactLimit("session aggregate size overflow".to_string()) + })?; + let artifact_count = artifact_count.checked_add(uploaded_count).ok_or_else(|| { + StoreError::PasteArtifactLimit("session artifact count overflow".to_string()) + })?; let projected = aggregate.checked_add(byte_len).ok_or_else(|| { StoreError::PasteArtifactLimit("session aggregate size overflow".to_string()) })?; diff --git a/crates/session-store/src/store.rs b/crates/session-store/src/store.rs index 142e580d..f498ac71 100644 --- a/crates/session-store/src/store.rs +++ b/crates/session-store/src/store.rs @@ -13,8 +13,8 @@ use crate::event_trace::TraceEntry; use crate::segment_log::LogEntry; -use crate::{PasteArtifactLimits, SegmentId, SessionId}; -use protocol::PasteArtifactRef; +use crate::{PasteArtifactLimits, SegmentId, SessionId, UploadedFileLimits}; +use protocol::{PasteArtifactRef, UploadedFileRef}; /// Errors from the persistence store. #[derive(Debug, thiserror::Error)] @@ -42,6 +42,30 @@ pub enum StoreError { #[error("paste artifact size limit exceeded: {0}")] PasteArtifactLimit(String), + + #[error("uploaded file is too large")] + ArtifactTooLarge, + + #[error("session artifact aggregate quota exceeded")] + ArtifactQuotaExceeded, + + #[error("uploaded file reference integrity check failed")] + ArtifactIntegrityMismatch, + + #[error("uploaded file name is invalid")] + InvalidUploadedFileName, + + #[error("uploaded file media type is invalid")] + InvalidUploadedFileMediaType, + + #[error("uploaded file is already committed to session history")] + ArtifactAlreadyCommitted, + + #[error("artifact id is invalid")] + InvalidArtifactId, + + #[error("artifact timestamp is invalid")] + InvalidTimestamp, } /// Sync persistence backend for segment logs. @@ -150,6 +174,53 @@ pub trait Store: Send + Sync { Err(StoreError::PasteArtifactUnsupported) } + /// Persist a client-local file before a submission references it. + fn write_uploaded_file( + &self, + _session_id: SessionId, + _file_name: &str, + _media_type: &str, + _content: &[u8], + _limits: UploadedFileLimits, + ) -> Result { + Err(StoreError::PasteArtifactUnsupported) + } + + /// Read and integrity-check an uploaded file owned by `session_id`. + fn read_uploaded_file( + &self, + _session_id: SessionId, + _reference: &UploadedFileRef, + ) -> Result, StoreError> { + Err(StoreError::PasteArtifactUnsupported) + } + + fn read_uploaded_file_by_id( + &self, + _session_id: SessionId, + _artifact_id: &str, + ) -> Result<(UploadedFileRef, Vec), StoreError> { + Err(StoreError::PasteArtifactUnsupported) + } + + fn bind_uploaded_file( + &self, + _session_id: SessionId, + _reference: &UploadedFileRef, + _source_entry_id: &str, + ) -> Result { + Err(StoreError::PasteArtifactUnsupported) + } + + /// Delete an uncommitted uploaded file owned by `session_id`. + fn delete_uploaded_file( + &self, + _session_id: SessionId, + _artifact_id: &str, + ) -> Result { + Err(StoreError::PasteArtifactUnsupported) + } + /// 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 new file mode 100644 index 00000000..456a6973 --- /dev/null +++ b/crates/session-store/src/uploaded_file.rs @@ -0,0 +1,299 @@ +use std::{ + fs, + path::Path, + time::{SystemTime, UNIX_EPOCH}, +}; + +use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64}; +use fs4::fs_std::FileExt; +use protocol::{UploadedFileAvailability, UploadedFileRef}; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; +use uuid::Uuid; + +use crate::StoreError; + +type Result = std::result::Result; + +pub const DEFAULT_MAX_UPLOADED_FILE_BYTES: u64 = 10 * 1024 * 1024; +pub const DEFAULT_MAX_SESSION_ARTIFACT_BYTES: u64 = 32 * 1024 * 1024; +pub const DEFAULT_MAX_FILES_PER_SUBMISSION: usize = 8; +const MAX_FILE_NAME_CHARS: usize = 255; +const MAX_MEDIA_TYPE_BYTES: usize = 127; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct UploadedFileLimits { + pub max_file_bytes: u64, + pub max_session_bytes: u64, +} + +impl Default for UploadedFileLimits { + fn default() -> Self { + Self { + max_file_bytes: DEFAULT_MAX_UPLOADED_FILE_BYTES, + max_session_bytes: DEFAULT_MAX_SESSION_ARTIFACT_BYTES, + } + } +} + +#[derive(Debug, Serialize, Deserialize)] +struct StoredUploadedFile { + file_name: String, + media_type: String, + created_at_ms: u64, + byte_len: u64, + sha256: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + source_entry_id: Option, + content_base64: String, +} + +pub(crate) fn validate_file_name(file_name: &str) -> Result<()> { + if file_name.is_empty() + || file_name.chars().count() > MAX_FILE_NAME_CHARS + || file_name == "." + || file_name == ".." + || file_name + .chars() + .any(|ch| ch.is_control() || matches!(ch, '/' | '\\')) + { + return Err(StoreError::InvalidUploadedFileName); + } + Ok(()) +} + +pub(crate) fn validate_media_type(media_type: &str) -> Result<()> { + let valid = !media_type.is_empty() + && media_type.len() <= MAX_MEDIA_TYPE_BYTES + && media_type.is_ascii() + && !media_type + .bytes() + .any(|byte| byte.is_ascii_control() || byte == b' ') + && media_type.split_once('/').is_some_and(|(kind, subtype)| { + !kind.is_empty() + && !subtype.is_empty() + && kind.bytes().chain(subtype.bytes()).all(|byte| { + byte.is_ascii_alphanumeric() + || matches!( + byte, + b'!' | b'#' | b'$' | b'&' | b'^' | b'_' | b'.' | b'+' | b'-' + ) + }) + }); + let allowed = media_type.starts_with("text/") + || matches!( + media_type, + "application/json" + | "application/pdf" + | "image/png" + | "image/jpeg" + | "image/gif" + | "image/webp" + ); + if !valid || !allowed { + return Err(StoreError::InvalidUploadedFileMediaType); + } + Ok(()) +} + +fn record_path(dir: &Path, artifact_id: &str) -> Result { + let id = Uuid::parse_str(artifact_id).map_err(|_| StoreError::InvalidArtifactId)?; + Ok(dir.join(format!("{id}.file.json"))) +} + +fn now_ms() -> Result { + let value = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_err(|_| StoreError::InvalidTimestamp)? + .as_millis(); + u64::try_from(value).map_err(|_| StoreError::InvalidTimestamp) +} + +fn digest(bytes: &[u8]) -> String { + Sha256::digest(bytes) + .iter() + .map(|byte| format!("{byte:02x}")) + .collect() +} + +pub(crate) fn stored_uploaded_file_usage(dir: &Path) -> Result<(u64, u64)> { + if !dir.exists() { + return Ok((0, 0)); + } + let mut bytes = 0_u64; + let mut count = 0_u64; + for entry in fs::read_dir(dir)? { + let entry = entry?; + let path = entry.path(); + if !entry.file_type()?.is_file() + || !path + .file_name() + .and_then(|name| name.to_str()) + .is_some_and(|name| name.ends_with(".file.json")) + { + continue; + } + let stored: StoredUploadedFile = serde_json::from_slice(&fs::read(path)?)?; + bytes = bytes + .checked_add(stored.byte_len) + .ok_or(StoreError::ArtifactQuotaExceeded)?; + count = count + .checked_add(1) + .ok_or(StoreError::ArtifactQuotaExceeded)?; + } + Ok((bytes, count)) +} + +pub(crate) fn write_uploaded_file( + dir: &Path, + file_name: &str, + media_type: &str, + content: &[u8], + limits: UploadedFileLimits, +) -> Result { + validate_file_name(file_name)?; + validate_media_type(media_type)?; + let byte_len = u64::try_from(content.len()).map_err(|_| StoreError::ArtifactTooLarge)?; + if byte_len > limits.max_file_bytes { + return Err(StoreError::ArtifactTooLarge); + } + + 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 (paste_bytes, _) = crate::paste_artifact::stored_paste_usage(dir)?; + let (file_bytes, _) = stored_uploaded_file_usage(dir)?; + if paste_bytes + .checked_add(file_bytes) + .and_then(|total| total.checked_add(byte_len)) + .is_none_or(|total| total > limits.max_session_bytes) + { + return Err(StoreError::ArtifactQuotaExceeded); + } + + let artifact_id = Uuid::now_v7().to_string(); + let created_at_ms = now_ms()?; + let sha256 = digest(content); + let stored = StoredUploadedFile { + file_name: file_name.to_owned(), + media_type: media_type.to_owned(), + created_at_ms, + byte_len, + sha256: sha256.clone(), + source_entry_id: None, + content_base64: BASE64.encode(content), + }; + let path = record_path(dir, &artifact_id)?; + let temp = dir.join(format!(".{artifact_id}.file.tmp")); + fs::write(&temp, serde_json::to_vec(&stored)?)?; + fs::rename(&temp, &path)?; + + Ok(UploadedFileRef { + artifact_id, + file_name: file_name.to_owned(), + media_type: media_type.to_owned(), + created_at_ms, + availability: UploadedFileAvailability::Available, + byte_len, + sha256, + source_entry_id: None, + }) +} + +pub(crate) fn read_uploaded_file_by_id( + dir: &Path, + artifact_id: &str, +) -> Result<(UploadedFileRef, Vec)> { + let stored: StoredUploadedFile = + serde_json::from_slice(&fs::read(record_path(dir, artifact_id)?)?)?; + let content = BASE64 + .decode(&stored.content_base64) + .map_err(|_| StoreError::ArtifactIntegrityMismatch)?; + if u64::try_from(content.len()).ok() != Some(stored.byte_len) + || digest(&content) != stored.sha256 + { + return Err(StoreError::ArtifactIntegrityMismatch); + } + let reference = UploadedFileRef { + artifact_id: artifact_id.to_owned(), + file_name: stored.file_name, + media_type: stored.media_type, + created_at_ms: stored.created_at_ms, + availability: UploadedFileAvailability::Available, + byte_len: stored.byte_len, + sha256: stored.sha256, + source_entry_id: stored.source_entry_id, + }; + Ok((reference, content)) +} + +pub(crate) fn read_uploaded_file(dir: &Path, reference: &UploadedFileRef) -> Result> { + let (stored_reference, content) = 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 + || stored_reference.source_entry_id != reference.source_entry_id + { + return Err(StoreError::ArtifactIntegrityMismatch); + } + Ok(content) +} + +pub(crate) fn bind_uploaded_file( + dir: &Path, + reference: &UploadedFileRef, + source_entry_id: &str, +) -> Result { + if source_entry_id.is_empty() || reference.source_entry_id.is_some() { + return Err(StoreError::ArtifactIntegrityMismatch); + } + let aggregate_lock = fs::OpenOptions::new() + .create(true) + .read(true) + .write(true) + .open(dir.join(".aggregate.lock"))?; + FileExt::lock_exclusive(&aggregate_lock)?; + read_uploaded_file(dir, reference)?; + 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() { + return Err(StoreError::ArtifactAlreadyCommitted); + } + stored.source_entry_id = Some(source_entry_id.to_owned()); + let temp = dir.join(format!(".{}.file.bind.tmp", reference.artifact_id)); + fs::write(&temp, serde_json::to_vec(&stored)?)?; + fs::rename(&temp, path)?; + let mut bound = reference.clone(); + bound.source_entry_id = Some(source_entry_id.to_owned()); + Ok(bound) +} + +pub(crate) fn delete_uploaded_file(dir: &Path, artifact_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 stored = match fs::read(&path) { + Ok(bytes) => serde_json::from_slice::(&bytes)?, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(false), + Err(error) => return Err(error.into()), + }; + if stored.source_entry_id.is_some() { + return Err(StoreError::ArtifactAlreadyCommitted); + } + match fs::remove_file(path) { + Ok(()) => Ok(true), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(false), + Err(error) => Err(error.into()), + } +} diff --git a/crates/worker-runtime/src/execution.rs b/crates/worker-runtime/src/execution.rs index c37f9792..78384c85 100644 --- a/crates/worker-runtime/src/execution.rs +++ b/crates/worker-runtime/src/execution.rs @@ -8,7 +8,7 @@ use crate::interaction::WorkerInput; #[cfg(feature = "ws-server")] use crate::observation::WorkerObservationEvent; use crate::working_directory::{WorkingDirectoryBinding, WorkingDirectoryDiagnostic}; -use protocol::Method; +use protocol::{Method, UploadedFileRef}; use serde::{Deserialize, Serialize}; use std::fmt; use std::sync::Arc; @@ -33,6 +33,8 @@ pub enum WorkerExecutionOperation { Spawn, Restore, Input, + UploadFile, + DeleteUploadedFile, ProtocolMethod, Stop, Cancel, @@ -385,6 +387,30 @@ pub trait WorkerExecutionBackend: Send + Sync + 'static { input: WorkerInput, ) -> WorkerExecutionResult; + fn upload_file( + &self, + _handle: &WorkerExecutionHandle, + _file_name: &str, + _media_type: &str, + _content: &[u8], + ) -> Result { + Err(WorkerExecutionResult::unsupported( + WorkerExecutionOperation::UploadFile, + "execution backend does not support file upload", + )) + } + + fn delete_uploaded_file( + &self, + _handle: &WorkerExecutionHandle, + _artifact_id: &str, + ) -> WorkerExecutionResult { + WorkerExecutionResult::unsupported( + WorkerExecutionOperation::DeleteUploadedFile, + "execution backend does not support uploaded-file deletion", + ) + } + fn dispatch_method( &self, _handle: &WorkerExecutionHandle, @@ -514,6 +540,25 @@ impl WorkerExecutionBackendRef { self.backend.dispatch_input(handle, input) } + pub(crate) fn upload_file( + &self, + handle: &WorkerExecutionHandle, + file_name: &str, + media_type: &str, + content: &[u8], + ) -> Result { + self.backend + .upload_file(handle, file_name, media_type, content) + } + + pub(crate) fn delete_uploaded_file( + &self, + handle: &WorkerExecutionHandle, + artifact_id: &str, + ) -> WorkerExecutionResult { + self.backend.delete_uploaded_file(handle, artifact_id) + } + pub(crate) fn dispatch_method( &self, handle: &WorkerExecutionHandle, diff --git a/crates/worker-runtime/src/http_server.rs b/crates/worker-runtime/src/http_server.rs index a8486523..4c0dd7a8 100644 --- a/crates/worker-runtime/src/http_server.rs +++ b/crates/worker-runtime/src/http_server.rs @@ -32,7 +32,7 @@ use axum::body::{Body, Bytes}; use axum::extract::rejection::{JsonRejection, QueryRejection}; #[cfg(feature = "ws-server")] use axum::extract::ws::{Message as WsMessage, WebSocket, WebSocketUpgrade}; -use axum::extract::{Extension, Path, Query, State}; +use axum::extract::{DefaultBodyLimit, Extension, Path, Query, State}; use axum::http::{Method, Request, StatusCode, header}; use axum::middleware::{self, Next}; use axum::response::{IntoResponse, Response}; @@ -238,6 +238,14 @@ fn runtime_http_router_with_optional_auth( post(execute_worker_retention), ) .route("/v1/workers/{worker_id}/input", post(send_worker_input)) + .route( + "/v1/workers/{worker_id}/attachments", + post(upload_worker_file).layer(DefaultBodyLimit::max(MAX_WORKER_FILE_UPLOAD_BYTES)), + ) + .route( + "/v1/workers/{worker_id}/attachments/{artifact_id}", + delete(delete_worker_uploaded_file), + ) .route("/v1/workers/{worker_id}/restore", post(restore_worker)) .route( "/v1/workers/{worker_id}/workspace-api", @@ -263,6 +271,9 @@ fn runtime_http_router_with_optional_auth( .layer(middleware::from_fn_with_state(state, require_runtime_auth)) } +pub const MAX_WORKER_FILE_UPLOAD_BYTES: usize = + session_store::DEFAULT_MAX_UPLOADED_FILE_BYTES as usize; + #[derive(Clone)] struct RuntimeHttpState { runtime: Runtime, @@ -375,6 +386,22 @@ pub struct RuntimeHttpWorkerInputResponse { pub ack: WorkerInteractionAck, } +#[derive(Clone, Debug, Deserialize)] +pub struct RuntimeHttpUploadFileQuery { + pub file_name: String, + pub media_type: String, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct RuntimeHttpUploadedFileResponse { + pub file: protocol::UploadedFileRef, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct RuntimeHttpUploadedFileDeleteResponse { + pub deleted: bool, +} + #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] pub struct RuntimeHttpWorkerCompletionsRequest { pub kind: protocol::CompletionKind, @@ -1420,6 +1447,55 @@ async fn worker_completions( })) } +async fn upload_worker_file( + State(state): State, + auth: Option>, + Path(worker_id): Path, + Query(query): Query, + body: Bytes, +) -> RestResult { + let worker_ref = worker_ref_for(&state.runtime, worker_id)?; + 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, + ), + None => state.runtime.upload_worker_file( + &worker_ref, + &query.file_name, + &query.media_type, + &body, + ), + } + .map_err(RuntimeHttpRestError::runtime)?; + Ok(Json(RuntimeHttpUploadedFileResponse { file })) +} + +async fn delete_worker_uploaded_file( + State(state): State, + auth: Option>, + Path((worker_id, artifact_id)): Path<(String, String)>, +) -> RestResult { + let worker_ref = worker_ref_for(&state.runtime, worker_id)?; + match auth_workspace_scope(&state, auth.as_ref())? { + Some(scope) => { + state + .runtime + .delete_worker_uploaded_file_scoped(&scope, &worker_ref, &artifact_id) + } + None => state + .runtime + .delete_worker_uploaded_file(&worker_ref, &artifact_id), + } + .map_err(RuntimeHttpRestError::runtime)?; + Ok(Json(RuntimeHttpUploadedFileDeleteResponse { + deleted: true, + })) +} + async fn stop_worker( State(state): State, auth: Option>, @@ -1621,7 +1697,7 @@ fn required_runtime_permission(method: &Method, path: &str) -> Option<&'static s if path.ends_with("/workspace-api") { return Some("workers:create"); } - if path.ends_with("/input") || path.ends_with("/restore") { + if path.ends_with("/input") || path.ends_with("/restore") || path.contains("/attachments") { return Some("workers:input"); } if path.ends_with("/stop") || path.ends_with("/cancel") { @@ -1882,6 +1958,21 @@ mod tests { WorkdirPath, WorkdirSessionCapabilities, }; + #[test] + fn attachment_routes_require_worker_input_permission() { + assert_eq!( + required_runtime_permission(&Method::POST, "/v1/workers/7/attachments"), + Some("workers:input") + ); + assert_eq!( + required_runtime_permission( + &Method::DELETE, + "/v1/workers/7/attachments/019ca7c8-57b6-7f05-8edf-524147aba7b3" + ), + Some("workers:input") + ); + } + fn test_bundle(profile: ProfileSelector) -> ConfigBundle { ConfigBundle { metadata: ConfigBundleMetadata { diff --git a/crates/worker-runtime/src/runtime.rs b/crates/worker-runtime/src/runtime.rs index 6d3d997a..31394209 100644 --- a/crates/worker-runtime/src/runtime.rs +++ b/crates/worker-runtime/src/runtime.rs @@ -1205,6 +1205,103 @@ impl Runtime { }) } + /// Store a client-local file in the owning Worker session before input submit. + pub fn upload_worker_file_scoped( + &self, + scope: &RuntimeWorkspaceScope, + worker_ref: &WorkerRef, + file_name: &str, + media_type: &str, + content: &[u8], + ) -> Result { + self.ensure_worker_in_workspace(scope, worker_ref)?; + self.upload_worker_file(worker_ref, file_name, media_type, content) + } + + pub fn upload_worker_file( + &self, + worker_ref: &WorkerRef, + file_name: &str, + media_type: &str, + content: &[u8], + ) -> Result { + let (backend, handle) = { + let state = self.lock()?; + state.ensure_running()?; + state.ensure_worker_ref(worker_ref)?; + let worker = state.worker(worker_ref)?; + match ( + state.execution_backend.clone(), + worker.execution_handle.clone(), + ) { + (Some(backend), Some(handle)) => (backend, handle), + _ => { + return Err(RuntimeError::WorkerExecutionUnavailable { + worker_id: worker_ref.worker_id.clone(), + message: "worker has no live execution handle".to_string(), + }); + } + } + }; + backend + .upload_file(&handle, file_name, media_type, content) + .map_err(|result| RuntimeError::WorkerExecutionRejected { + worker_id: worker_ref.worker_id.clone(), + operation: result.operation, + outcome: result.outcome, + message: result.message_or_default(), + result, + }) + } + + /// Delete an unsubmitted uploaded file from the owning Worker session. + pub fn delete_worker_uploaded_file_scoped( + &self, + scope: &RuntimeWorkspaceScope, + worker_ref: &WorkerRef, + artifact_id: &str, + ) -> Result<(), RuntimeError> { + self.ensure_worker_in_workspace(scope, worker_ref)?; + self.delete_worker_uploaded_file(worker_ref, artifact_id) + } + + pub fn delete_worker_uploaded_file( + &self, + worker_ref: &WorkerRef, + artifact_id: &str, + ) -> Result<(), RuntimeError> { + let (backend, handle) = { + let state = self.lock()?; + state.ensure_running()?; + state.ensure_worker_ref(worker_ref)?; + let worker = state.worker(worker_ref)?; + match ( + state.execution_backend.clone(), + worker.execution_handle.clone(), + ) { + (Some(backend), Some(handle)) => (backend, handle), + _ => { + return Err(RuntimeError::WorkerExecutionUnavailable { + worker_id: worker_ref.worker_id.clone(), + message: "worker has no live execution handle".to_string(), + }); + } + } + }; + let result = backend.delete_uploaded_file(&handle, artifact_id); + if result.is_accepted() { + Ok(()) + } else { + Err(RuntimeError::WorkerExecutionRejected { + worker_id: worker_ref.worker_id.clone(), + operation: result.operation, + outcome: result.outcome, + message: result.message_or_default(), + result, + }) + } + } + /// Return live completion entries through a workspace-scoped Runtime authorization context. pub fn worker_completions_scoped( &self, @@ -1387,6 +1484,8 @@ impl Runtime { WorkerExecutionOperation::Spawn | WorkerExecutionOperation::Restore | WorkerExecutionOperation::Input + | WorkerExecutionOperation::UploadFile + | WorkerExecutionOperation::DeleteUploadedFile | WorkerExecutionOperation::ProtocolMethod => return Ok(()), }; if result.is_accepted() { diff --git a/crates/worker-runtime/src/worker_backend.rs b/crates/worker-runtime/src/worker_backend.rs index 5f587c1b..3fd8a56b 100644 --- a/crates/worker-runtime/src/worker_backend.rs +++ b/crates/worker-runtime/src/worker_backend.rs @@ -2025,6 +2025,51 @@ where result } + fn upload_file( + &self, + handle: &WorkerExecutionHandle, + file_name: &str, + media_type: &str, + content: &[u8], + ) -> 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}"), + ) + }) + } + + fn delete_uploaded_file( + &self, + handle: &WorkerExecutionHandle, + artifact_id: &str, + ) -> WorkerExecutionResult { + let (worker, _, _) = match self.get_execution(handle) { + Ok(execution) => execution, + Err(mut result) => { + result.operation = WorkerExecutionOperation::DeleteUploadedFile; + return result; + } + }; + match worker.delete_uploaded_file(artifact_id) { + Ok(_) => WorkerExecutionResult::accepted( + WorkerExecutionOperation::DeleteUploadedFile, + WorkerExecutionRunState::Idle, + ), + Err(error) => WorkerExecutionResult::rejected( + WorkerExecutionOperation::DeleteUploadedFile, + format!("uploaded_file_delete_rejected: {error}"), + ), + } + } + fn dispatch_method( &self, handle: &WorkerExecutionHandle, diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index aaca6aa9..5d0f98c9 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -31,7 +31,8 @@ use protocol::{ AlertLevel, AlertSource, CommandEvent as ProtocolCommandEvent, CommandSnapshot as ProtocolCommandSnapshot, CommandStatus as ProtocolCommandStatus, CommandStream as ProtocolCommandStream, CommandStreamSlice as ProtocolCommandStreamSlice, - ErrorCode, Event, Method, RewindTargetId, RunResult, Segment, TurnResult, WorkerStatus, + ErrorCode, Event, Method, RewindTargetId, RunResult, Segment, TurnResult, UploadedFileRef, + WorkerStatus, }; use workdir::{ CommandEvent as WorkdirCommandEvent, CommandSnapshot as WorkdirCommandSnapshot, @@ -55,6 +56,8 @@ pub struct WorkerHandle { /// subsequent commits (Event::Entry) on the receiver. pub sink: SegmentLogSink, spawned_registry: Arc, + artifact_store: Arc, + session_id: session_store::SessionId, } impl WorkerHandle { @@ -62,6 +65,29 @@ impl WorkerHandle { self.method_tx.send(method).await } + pub fn upload_file( + &self, + file_name: &str, + media_type: &str, + content: &[u8], + ) -> Result { + self.artifact_store.write_uploaded_file( + self.session_id, + file_name, + media_type, + content, + session_store::UploadedFileLimits::default(), + ) + } + + pub fn delete_uploaded_file( + &self, + artifact_id: &str, + ) -> Result { + self.artifact_store + .delete_uploaded_file(self.session_id, artifact_id) + } + pub fn subscribe(&self) -> broadcast::Receiver { self.working_event_tx.subscribe() } @@ -503,6 +529,8 @@ impl WorkerController { runtime_dir.write_manifest(&manifest_toml).await?; runtime_dir.write_status(&shared_state).await?; + let artifact_store: Arc = Arc::new(worker.store().clone()); + let session_id = worker.session_id(); let handle = WorkerHandle { method_tx, working_event_tx: working_event_tx.clone(), @@ -512,6 +540,8 @@ impl WorkerController { in_flight: in_flight.clone(), sink: worker.sink(), spawned_registry: spawned_registry.clone(), + artifact_store, + session_id, }; let socket_server = match transport { diff --git a/crates/worker/src/paste_artifact_tool.rs b/crates/worker/src/paste_artifact_tool.rs index b61af5d7..9747c074 100644 --- a/crates/worker/src/paste_artifact_tool.rs +++ b/crates/worker/src/paste_artifact_tool.rs @@ -4,6 +4,7 @@ use std::sync::Arc; use agen::tool::{Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolMeta, ToolOutput}; use async_trait::async_trait; +use protocol::{PasteArtifactAvailability, PasteArtifactMediaType, PasteArtifactRef}; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; use session_store::{SessionId, Store, StoreError}; @@ -75,11 +76,8 @@ where .max_results .unwrap_or(DEFAULT_SEARCH_RESULTS) .clamp(1, MAX_SEARCH_RESULTS); - let (_, content) = self - .access - .store - .read_paste_artifact(self.access.session_id, &input.artifact_id) - .map_err(tool_store_error)?; + let (_, content) = + read_artifact_text(&self.access, &input.artifact_id).map_err(tool_store_error)?; let mut matches = Vec::new(); let mut truncated = false; let mut byte_offset = 0_u64; @@ -150,11 +148,8 @@ where .max_bytes .unwrap_or(DEFAULT_READ_BYTES) .clamp(4, MAX_READ_BYTES); - let (_, content) = self - .access - .store - .read_paste_artifact(self.access.session_id, &input.artifact_id) - .map_err(tool_store_error)?; + let (_, content) = + read_artifact_text(&self.access, &input.artifact_id).map_err(tool_store_error)?; let offset = usize::try_from(offset).map_err(|_| { ToolError::InvalidArgument("offset exceeds the artifact size".to_string()) })?; @@ -234,6 +229,47 @@ fn json_output(summary: String, value: &impl Serialize) -> Result( + access: &ArtifactAccess, + artifact_id: &str, +) -> Result<(PasteArtifactRef, String), StoreError> { + match access + .store + .read_paste_artifact(access.session_id, artifact_id) + { + Ok(result) => Ok(result), + Err(paste_error) => { + let (file, bytes) = match access + .store + .read_uploaded_file_by_id(access.session_id, artifact_id) + { + Ok(result) => result, + Err(_) => return Err(paste_error), + }; + let content = + String::from_utf8(bytes).map_err(|_| StoreError::ArtifactIntegrityMismatch)?; + let char_count = + u64::try_from(content.chars().count()).map_err(|_| StoreError::ArtifactTooLarge)?; + let line_count = + u64::try_from(content.lines().count()).map_err(|_| StoreError::ArtifactTooLarge)?; + Ok(( + PasteArtifactRef { + artifact_id: file.artifact_id, + created_at_ms: file.created_at_ms, + media_type: PasteArtifactMediaType::TextPlainUtf8, + availability: PasteArtifactAvailability::Available, + byte_len: file.byte_len, + char_count, + line_count, + sha256: file.sha256, + source_entry_id: file.source_entry_id.unwrap_or_default(), + }, + content, + )) + } + } +} + fn tool_store_error(error: StoreError) -> ToolError { let message = match error { StoreError::PasteArtifactNotFound(_) => "paste artifact not found", @@ -263,6 +299,66 @@ mod tests { use super::*; + #[tokio::test] + async fn read_input_artifact_reads_uploaded_text_but_rejects_binary_content() { + let temp = tempfile::TempDir::new().unwrap(); + let store = FsStore::new(temp.path()).unwrap(); + let owner = new_session_id(); + let text = store + .write_uploaded_file( + owner, + "notes.md", + "text/markdown", + b"alpha\nbeta", + session_store::UploadedFileLimits::default(), + ) + .unwrap(); + let read = ReadInputArtifactTool { + access: ArtifactAccess { + store: store.clone(), + session_id: owner, + }, + }; + let output = read + .execute( + &serde_json::json!({ + "artifact_id": text.artifact_id, + "offset": 0, + "max_bytes": 64 + }) + .to_string(), + ToolExecutionContext::default(), + ) + .await + .unwrap(); + let output: serde_json::Value = + serde_json::from_str(output.content.as_deref().unwrap()).unwrap(); + assert_eq!(output["content"], "alpha\nbeta"); + + let binary = store + .write_uploaded_file( + owner, + "image.png", + "image/png", + &[0xff, 0xd8, 0x00], + session_store::UploadedFileLimits::default(), + ) + .unwrap(); + let error = read + .execute( + &serde_json::json!({ + "artifact_id": binary.artifact_id, + "offset": 0, + "max_bytes": 64 + }) + .to_string(), + ToolExecutionContext::default(), + ) + .await + .unwrap_err(); + assert!(error.to_string().contains("unavailable")); + } + #[tokio::test] async fn search_and_read_are_bounded_and_owner_scoped() { let temp = tempfile::TempDir::new().unwrap(); diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 11548ad8..302a26b2 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -3082,7 +3082,32 @@ impl Worker { projected_entry_ids: &[SessionHistoryEntryId], one_entry_per_segment: bool, ) -> Result<(), WorkerError> { + let uploaded_file_count = input + .iter() + .filter(|segment| matches!(segment, Segment::UploadedFile { .. })) + .count(); + if uploaded_file_count > session_store::DEFAULT_MAX_FILES_PER_SUBMISSION { + return Err(WorkerError::Store(StoreError::ArtifactQuotaExceeded)); + } for (index, segment) in input.iter_mut().enumerate() { + if let Segment::UploadedFile { file } = segment { + if file.source_entry_id.is_some() + || file.availability != protocol::UploadedFileAvailability::Available + { + return Err(WorkerError::Store(StoreError::ArtifactIntegrityMismatch)); + } + let entry_index = if one_entry_per_segment { index } else { 0 }; + let source_entry_id = projected_entry_ids + .get(entry_index) + .expect("projected input id exists for every uploaded file") + .0 + .clone(); + *file = self + .store + .bind_uploaded_file(self.session_id(), file, &source_entry_id) + .map_err(WorkerError::Store)?; + continue; + } if let Segment::PasteArtifact { artifact } = segment { let (stored, _) = self .store @@ -6471,6 +6496,11 @@ fn preview_segments(segments: &[Segment]) -> String { preview.push_str(&artifact.artifact_id); preview.push(']'); } + Segment::UploadedFile { file } => { + preview.push_str("[Attached file: "); + preview.push_str(&file.file_name); + preview.push(']'); + } Segment::FileRef { path } => { preview.push('@'); preview.push_str(path); @@ -8049,6 +8079,61 @@ mod build_summary_prompt_tests { ); } + #[tokio::test] + async fn uploaded_file_is_verified_and_bound_to_projected_entry_before_commit() { + let (_dir, worker) = rewind_test_worker().await; + let reference = worker + .store + .write_uploaded_file( + worker.session_id(), + "notes.md", + "text/markdown", + b"# private body", + session_store::UploadedFileLimits::default(), + ) + .unwrap(); + let entry_id = SessionHistoryEntryId::new(); + let mut input = vec![Segment::UploadedFile { file: reference }]; + + worker + .materialize_large_pastes(&mut input, std::slice::from_ref(&entry_id), false) + .unwrap(); + let file = match &input[0] { + Segment::UploadedFile { file } => file, + other => panic!("expected uploaded file, got {other:?}"), + }; + assert_eq!(file.source_entry_id.as_deref(), Some(entry_id.0.as_str())); + assert_eq!( + worker + .store + .read_uploaded_file(worker.session_id(), file) + .unwrap(), + b"# private body" + ); + assert!(matches!( + worker + .store + .delete_uploaded_file(worker.session_id(), &file.artifact_id), + Err(StoreError::ArtifactAlreadyCommitted) + )); + let projected = worker.projected_input_history(&input, None, &[entry_id]); + let text = projected[0].item.as_text().unwrap(); + assert!(text.contains("notes.md")); + assert!(text.contains(&file.artifact_id)); + assert!(!text.contains("private body")); + + let mut forged = input.clone(); + let Segment::UploadedFile { file } = &mut forged[0] else { + unreachable!(); + }; + file.source_entry_id = None; + file.sha256 = "0".repeat(64); + assert!(matches!( + worker.materialize_large_pastes(&mut forged, &[SessionHistoryEntryId::new()], false), + Err(WorkerError::Store(StoreError::ArtifactIntegrityMismatch)) + )); + } + #[tokio::test] async fn large_paste_is_stored_before_compact_history_is_committed() { let (_dir, worker) = rewind_test_worker().await; diff --git a/crates/workspace-server/src/hosts.rs b/crates/workspace-server/src/hosts.rs index 5bdd083f..c6a3b5af 100644 --- a/crates/workspace-server/src/hosts.rs +++ b/crates/workspace-server/src/hosts.rs @@ -40,6 +40,7 @@ use worker_runtime::fs_store::FsRuntimeStoreOptions; use worker_runtime::http_server::{ RuntimeHttpConfigBundleAvailabilityResponse, RuntimeHttpConfigBundleSyncRequest, RuntimeHttpErrorResponse, RuntimeHttpRepositoryAccessResponse, RuntimeHttpSummaryResponse, + RuntimeHttpUploadedFileDeleteResponse, RuntimeHttpUploadedFileResponse, RuntimeHttpWorkerCompletionsRequest, RuntimeHttpWorkerCompletionsResponse, RuntimeHttpWorkerDeleteResponse, RuntimeHttpWorkerInputResponse, RuntimeHttpWorkerLifecycleRequest, RuntimeHttpWorkerLifecycleResponse, @@ -1041,6 +1042,34 @@ pub trait WorkspaceWorkerRuntime: Send + Sync { } } + fn upload_worker_file( + &self, + worker_id: &str, + _file_name: &str, + _media_type: &str, + _content: &[u8], + ) -> Result { + Err(RuntimeRegistryError::RuntimeOperationFailed { + runtime_id: self.runtime_id().to_string(), + code: "worker_file_upload_unsupported".to_string(), + message: format!("runtime does not support file upload for worker `{worker_id}`"), + }) + } + + fn delete_worker_uploaded_file( + &self, + worker_id: &str, + _artifact_id: &str, + ) -> Result<(), RuntimeRegistryError> { + Err(RuntimeRegistryError::RuntimeOperationFailed { + runtime_id: self.runtime_id().to_string(), + code: "worker_file_delete_unsupported".to_string(), + message: format!( + "runtime does not support uploaded-file deletion for worker `{worker_id}`" + ), + }) + } + fn worker_completions( &self, worker_id: &str, @@ -1543,6 +1572,50 @@ impl RuntimeRegistry { Ok(runtime.send_input(worker_id, request)) } + pub fn upload_worker_file( + &self, + worker: &RuntimeWorkerRef, + file_name: &str, + media_type: &str, + content: &[u8], + ) -> Result { + let runtime_id = worker.runtime_id.as_str(); + let worker_id = worker.worker_id.as_str(); + validate_backend_identifier("runtime_id", runtime_id)?; + validate_backend_identifier("worker_id", worker_id)?; + let runtime = self.runtime(runtime_id)?; + let lookup = runtime.worker(worker_id); + if lookup.worker.is_none() { + return Err(operation_failed_or_unknown_worker( + runtime_id, + worker_id, + lookup.diagnostics, + )); + } + runtime.upload_worker_file(worker_id, file_name, media_type, content) + } + + pub fn delete_worker_uploaded_file( + &self, + worker: &RuntimeWorkerRef, + artifact_id: &str, + ) -> Result<(), RuntimeRegistryError> { + let runtime_id = worker.runtime_id.as_str(); + let worker_id = worker.worker_id.as_str(); + validate_backend_identifier("runtime_id", runtime_id)?; + validate_backend_identifier("worker_id", worker_id)?; + let runtime = self.runtime(runtime_id)?; + let lookup = runtime.worker(worker_id); + if lookup.worker.is_none() { + return Err(operation_failed_or_unknown_worker( + runtime_id, + worker_id, + lookup.diagnostics, + )); + } + runtime.delete_worker_uploaded_file(worker_id, artifact_id) + } + pub fn worker_completions( &self, worker: &RuntimeWorkerRef, @@ -2509,6 +2582,46 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime { } } + fn upload_worker_file( + &self, + worker_id: &str, + file_name: &str, + media_type: &str, + content: &[u8], + ) -> 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(), + }) + } + + fn delete_worker_uploaded_file( + &self, + worker_id: &str, + artifact_id: &str, + ) -> Result<(), RuntimeRegistryError> { + let worker_ref = + self.worker_ref(worker_id) + .ok_or_else(|| RuntimeRegistryError::UnknownWorker { + worker: RuntimeWorkerRef::new(&self.runtime_id, worker_id), + })?; + self.runtime + .delete_worker_uploaded_file(&worker_ref, artifact_id) + .map_err(|error| RuntimeRegistryError::RuntimeOperationFailed { + runtime_id: self.runtime_id.clone(), + code: "embedded_worker_file_delete_failed".to_string(), + message: error.to_string(), + }) + } + fn worker_completions( &self, worker_id: &str, @@ -2848,6 +2961,16 @@ impl RemoteWorkerRuntime { self.send_json(path, self.http.post(self.endpoint(path)).json(body)) } + fn post_bytes(&self, path: &str, body: &[u8]) -> Result + where + T: DeserializeOwned + Send + 'static, + { + self.send_json( + path, + self.http.post(self.endpoint(path)).body(body.to_vec()), + ) + } + fn delete_json(&self, path: &str) -> Result where T: DeserializeOwned + Send + 'static, @@ -3536,6 +3659,47 @@ impl WorkspaceWorkerRuntime for RemoteWorkerRuntime { } } + fn upload_worker_file( + &self, + worker_id: &str, + file_name: &str, + media_type: &str, + content: &[u8], + ) -> Result { + let 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), + ); + self.post_bytes::(&path, content) + .map(|response| response.file) + .map_err(|diagnostic| RuntimeRegistryError::RuntimeOperationFailed { + runtime_id: self.runtime_id.clone(), + code: diagnostic.code, + message: diagnostic.message, + }) + } + + fn delete_worker_uploaded_file( + &self, + worker_id: &str, + artifact_id: &str, + ) -> Result<(), RuntimeRegistryError> { + let path = format!( + "/v1/workers/{}/attachments/{}", + url_path_segment_encode(worker_id), + url_path_segment_encode(artifact_id), + ); + self.delete_json::(&path) + .map(|_| ()) + .map_err(|diagnostic| RuntimeRegistryError::RuntimeOperationFailed { + runtime_id: self.runtime_id.clone(), + code: diagnostic.code, + message: diagnostic.message, + }) + } + fn worker_completions( &self, worker_id: &str, diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 2ba7ae69..d938bb9f 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -3,8 +3,9 @@ use std::path::{Component, Path, PathBuf}; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Mutex, Weak}; +use axum::body::Bytes; use axum::extract::ws::{Message as WsMessage, WebSocket, WebSocketUpgrade}; -use axum::extract::{Extension, Path as AxumPath, Query, Request, State}; +use axum::extract::{DefaultBodyLimit, Extension, Path as AxumPath, Query, Request, State}; use axum::http::header::{CONTENT_TYPE, ETAG, IF_NONE_MATCH, LOCATION, ORIGIN, SET_COOKIE}; use axum::http::{HeaderMap, Method, StatusCode, Uri}; use axum::middleware::{self, Next}; @@ -156,8 +157,9 @@ use worker_runtime::catalog::{ }; use worker_runtime::config_bundle::ConfigBundle; use worker_runtime::http_server::{ - RuntimeHttpConfigBundleAvailabilityResponse, RuntimeHttpConfigBundlesResponse, - RuntimeHttpSummaryResponse, RuntimeHttpWorkerResponse, RuntimeHttpWorkersResponse, + MAX_WORKER_FILE_UPLOAD_BYTES, RuntimeHttpConfigBundleAvailabilityResponse, + RuntimeHttpConfigBundlesResponse, RuntimeHttpSummaryResponse, RuntimeHttpWorkerResponse, + RuntimeHttpWorkersResponse, }; use worker_runtime::identity::{RuntimeWorkerRef, WorkerId}; @@ -2679,10 +2681,26 @@ 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)), + ) + .route( + "/api/w/{workspace_id}/runtimes/{runtime_id}/workers/{worker_id}/attachments/{artifact_id}", + delete(scoped_delete_runtime_worker_uploaded_file), + ) .route( "/api/runtimes/{runtime_id}/workers/{worker_id}/completions", post(runtime_worker_completions), @@ -10733,6 +10751,51 @@ async fn scoped_execute_runtime_cleanup( Ok(Json(response)) } +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +struct WorkerFileUploadQuery { + file_name: String, + media_type: String, +} + +#[derive(Debug, Serialize)] +struct WorkerFileUploadResponse { + file: protocol::UploadedFileRef, +} + +#[derive(Debug, Serialize)] +struct WorkerFileDeleteResponse { + deleted: bool, +} + +async fn scoped_upload_runtime_worker_file( + State(api): State, + AxumPath(path): AxumPath, + Query(query): Query, + 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 +} + +async fn scoped_delete_runtime_worker_uploaded_file( + State(api): State, + AxumPath((path, artifact_id)): AxumPath<(ScopedRuntimeWorkerPath, String)>, +) -> ApiResult> { + validate_workspace_scope(&api, &path.workspace_id)?; + delete_runtime_worker_uploaded_file( + State(api), + AxumPath((path.worker.runtime_id, path.worker.worker_id, artifact_id)), + ) + .await +} + async fn scoped_send_runtime_worker_input( State(api): State, AxumPath(path): AxumPath, @@ -13260,6 +13323,31 @@ 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)>, +) -> ApiResult> { + let worker = resolve_workspace_worker_reference(&api, &runtime_id, &worker_id)?; + api.runtime + .delete_worker_uploaded_file(&worker, &artifact_id) + .map_err(|err| err.into_error())?; + Ok(Json(WorkerFileDeleteResponse { deleted: true })) +} + async fn runtime_worker_completions( State(api): State, AxumPath((runtime_id, worker_id)): AxumPath<(String, String)>, diff --git a/web/workspace/src/lib/generated/protocol.ts b/web/workspace/src/lib/generated/protocol.ts index 8100c319..a619d63f 100644 --- a/web/workspace/src/lib/generated/protocol.ts +++ b/web/workspace/src/lib/generated/protocol.ts @@ -42,6 +42,10 @@ export type CompactionLifecycle = { schema_version: number, compaction_id: strin */ started_at_ms: number, ended_at_ms?: number | null, summary?: string | null, error?: string | null, new_segment_id?: string | null, }; +export type UploadedFileAvailability = "available" | "unavailable" | "integrity_failed"; + +export type UploadedFileRef = { artifact_id: string, file_name: string, media_type: string, created_at_ms: number, availability: UploadedFileAvailability, byte_len: number, sha256: string, source_entry_id?: string | null, }; + export type ScopeRule = { /** * Target path. Must be absolute by the time a `Scope` is built from @@ -144,7 +148,7 @@ export type PasteArtifactRef = { artifact_id: string, created_at_ms: number, med */ availability: PasteArtifactAvailability, byte_len: number, char_count: number, line_count: number, sha256: string, source_entry_id: string, }; -export type Segment = { "kind": "text", content: string, } | { "kind": "paste", id: number, chars: number, lines: number, content: string, } | { "kind": "paste_artifact", artifact: PasteArtifactRef, } | { "kind": "file_ref", path: string, } | { "kind": "flow", selector: string, } | { "kind": "unknown" }; +export type Segment = { "kind": "text", content: string, } | { "kind": "paste", id: number, chars: number, lines: number, content: string, } | { "kind": "paste_artifact", artifact: PasteArtifactRef, } | { "kind": "uploaded_file", file: UploadedFileRef, } | { "kind": "file_ref", path: string, } | { "kind": "flow", selector: string, } | { "kind": "unknown" }; export type WorkerEvent = { "kind": "turn_ended", worker_name: string, } | { "kind": "errored", worker_name: string, message: string, } | { "kind": "shut_down", worker_name: string, } | { "kind": "scope_sub_delegated", /**