Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

fix: fence attachment grants and lifecycle cleanup

Browse Source
This commit is contained in:
Hare
2026-09-03 07:22:07 +09:00
parent 3f6bb65eb1
commit 9dc8d9a77a
19 changed files with 1190 additions and 167 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/client/src/backend_runtime.rs
+65 -3
View File
@@ -58,18 +58,52 @@ impl BackendRuntimeTarget {
file_name: &str,
media_type: &str,
content: Vec<u8>,
) -> Result<protocol::UploadedFileRef, BackendRuntimeClientError> {
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<u8>,
) -> Result<protocol::UploadedFileRef, BackendRuntimeClientError> {
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::<AttachmentUploadGrantResponse>()
.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,
crates/session-store/src/fs_store.rs
+182 -9
View File
@@ -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<bool, StoreError> {
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<UploadedFileRef, StoreError> {
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<u64, StoreError> {
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,
crates/session-store/src/lib.rs
+1
View File
@@ -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,
crates/session-store/src/segment.rs
+2
View File
@@ -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))
}
crates/session-store/src/store.rs
+23 -1
View File
@@ -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<UploadedFileRef, StoreError> {
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<u64, StoreError> {
Ok(0)
}
/// Append a trace entry to the debug event trace file.
fn append_trace(
&self,
crates/session-store/src/uploaded_file.rs
+113 -1
View File
@@ -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<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
upload_context: Option<UploadedFileUploadContext>,
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<UploadedFileRef> {
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<Vec<UploadedFileRef>> {
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<u64> {
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<u64> {
fs::create_dir_all(dir)?;
let aggregate_lock = fs::OpenOptions::new()
crates/session-store/tests/session_test.rs
+1
View File
@@ -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(),
crates/tui/src/app.rs
+4
View File
@@ -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<String>) {
self.blocks.push(Block::Alert {
level: AlertLevel::Error,
crates/tui/src/console/mod.rs
+158 -48
View File
@@ -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<UploadedFileRef, String>;
struct ConsoleConnection<T> {
client: Client<T>,
standalone_host: Option<StandaloneHost>,
backend_target: Option<BackendRuntimeTarget>,
pending_attachments: Vec<UploadedFileRef>,
upload_tasks: Vec<tokio::task::JoinHandle<()>>,
upload_ids: HashMap<PathBuf, String>,
}
async fn upload_client_path(
target: &BackendRuntimeTarget,
path: &Path,
upload_id: &str,
) -> Result<UploadedFileRef, Box<dyn std::error::Error>> {
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<T: Socket> ConsoleConnection<T> {
standalone_host: Some(host),
backend_target: None,
pending_attachments: Vec::new(),
upload_tasks: Vec::new(),
upload_ids: HashMap::new(),
}
}
@@ -162,6 +209,8 @@ impl<T: Socket> ConsoleConnection<T> {
standalone_host: None,
backend_target: Some(target),
pending_attachments: Vec::new(),
upload_tasks: Vec::new(),
upload_ids: HashMap::new(),
}
}
@@ -188,59 +237,49 @@ impl<T: Socket> ConsoleConnection<T> {
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<UploadedFileRef, Box<dyn std::error::Error>> {
let target = self.backend_target.as_ref().ok_or_else(|| {
path: PathBuf,
result_tx: mpsc::UnboundedSender<AttachmentUploadResult>,
) -> Result<(), Box<dyn std::error::Error>> {
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::<Vec<_>>();
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<P> {
Terminal(TerminalEventResult),
Worker(P),
Upload(AttachmentUploadResult),
Tick,
}
async fn next_loop_input<P, F, T>(
term_rx: &mut mpsc::UnboundedReceiver<TerminalEventResult>,
upload_rx: &mut mpsc::UnboundedReceiver<AttachmentUploadResult>,
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<T: Socket>(
app: &mut App,
client: &mut ConsoleConnection<T>,
term_rx: &mut mpsc::UnboundedReceiver<TerminalEventResult>,
upload_tx: &mpsc::UnboundedSender<AttachmentUploadResult>,
) -> Result<bool, Box<dyn std::error::Error>> {
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<T: Socket>(
client: &mut ConsoleConnection<T>,
) -> Result<(), Box<dyn std::error::Error>> {
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<T: Socket>(
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<T: Socket>(
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<T: Socket>(
.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<T: Socket>(
app: &mut App,
client: &mut ConsoleConnection<T>,
upload_tx: &mpsc::UnboundedSender<AttachmentUploadResult>,
event: TermEvent,
) -> Result<(), Box<dyn std::error::Error>> {
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<T: Socket>(
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::<TerminalEventResult>();
let (_upload_tx, mut upload_rx) = mpsc::unbounded_channel();
assert!(matches!(
next_loop_input(
&mut rx,
&mut upload_rx,
true,
std::future::pending::<Option<u8>>(),
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::<Option<u8>>(),
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,
crates/worker-runtime/src/execution.rs
+3 -1
View File
@@ -393,6 +393,7 @@ pub trait WorkerExecutionBackend: Send + Sync + 'static {
_file_name: &str,
_media_type: &str,
_content: &[u8],
_context: Option<&session_store::UploadedFileUploadContext>,
) -> Result<UploadedFileRef, WorkerExecutionResult> {
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<UploadedFileRef, WorkerExecutionResult> {
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(
crates/worker-runtime/src/http_server.rs
+71 -1
View File
@@ -390,6 +390,16 @@ pub struct RuntimeHttpWorkerInputResponse {
pub struct RuntimeHttpUploadFileQuery {
pub file_name: String,
pub media_type: String,
#[serde(default)]
pub upload_id: Option<String>,
#[serde(default)]
pub principal_id: Option<String>,
#[serde(default)]
pub workspace_id: Option<String>,
#[serde(default)]
pub runtime_id: Option<String>,
#[serde(default)]
pub owner_worker_id: Option<String>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
@@ -1455,14 +1465,74 @@ async fn upload_worker_file(
body: Bytes,
) -> RestResult<RuntimeHttpUploadedFileResponse> {
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(
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,
crates/worker-runtime/src/lib.rs
+1
View File
@@ -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;
crates/worker-runtime/src/runtime.rs
+36 -1
View File
@@ -1224,6 +1224,41 @@ impl Runtime {
file_name: &str,
media_type: &str,
content: &[u8],
) -> Result<protocol::UploadedFileRef, RuntimeError> {
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<protocol::UploadedFileRef, RuntimeError> {
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<protocol::UploadedFileRef, RuntimeError> {
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<protocol::UploadedFileRef, RuntimeError> {
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,
crates/worker-runtime/src/worker_backend.rs
+16 -23
View File
@@ -2031,14 +2031,19 @@ where
file_name: &str,
media_type: &str,
content: &[u8],
context: Option<&session_store::UploadedFileUploadContext>,
) -> Result<protocol::UploadedFileRef, WorkerExecutionResult> {
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| {
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}"),
@@ -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),
}
}
crates/worker/src/controller.rs
+17
View File
@@ -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<UploadedFileRef, session_store::StoreError> {
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,
crates/workspace-server/src/hosts.rs
+29 -5
View File
@@ -1048,6 +1048,7 @@ pub trait WorkspaceWorkerRuntime: Send + Sync {
_file_name: &str,
_media_type: &str,
_content: &[u8],
_context: Option<&worker_runtime::UploadedFileUploadContext>,
) -> Result<protocol::UploadedFileRef, RuntimeRegistryError> {
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<protocol::UploadedFileRef, RuntimeRegistryError> {
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,15 +2590,26 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime {
file_name: &str,
media_type: &str,
content: &[u8],
context: Option<&worker_runtime::UploadedFileUploadContext>,
) -> Result<protocol::UploadedFileRef, RuntimeRegistryError> {
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 {
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(),
@@ -3665,13 +3678,24 @@ impl WorkspaceWorkerRuntime for RemoteWorkerRuntime {
file_name: &str,
media_type: &str,
content: &[u8],
context: Option<&worker_runtime::UploadedFileUploadContext>,
) -> Result<protocol::UploadedFileRef, RuntimeRegistryError> {
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::<RuntimeHttpUploadedFileResponse>(&path, content)
.map(|response| response.file)
.map_err(|diagnostic| RuntimeRegistryError::RuntimeOperationFailed {
crates/workspace-server/src/server.rs
+362 -31
View File
@@ -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<Option<protocol::UploadedFileRef>> {
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<Mutex<HashMap<String, Arc<std::sync::Mutex<()>>>>>,
workdir_remove_attempt_owner: WorkdirRemovalAttemptOwner,
worker_control_locks: Arc<Mutex<HashMap<String, Arc<tokio::sync::Mutex<()>>>>>,
attachment_upload_grants: Arc<Mutex<HashMap<String, AttachmentUploadGrant>>>,
}
#[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<String>,
}
#[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> {
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<WorkspaceApi>,
Extension(actor): Extension<RequestActor>,
AxumPath(path): AxumPath<ScopedRuntimeWorkerPath>,
Query(query): Query<WorkerFileUploadQuery>,
) -> ApiResult<Json<AttachmentUploadGrantResponse>> {
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<WorkspaceApi>,
Extension(actor): Extension<RequestActor>,
AxumPath(path): AxumPath<ScopedAttachmentUploadPath>,
body: Bytes,
) -> ApiResult<Json<WorkerFileUploadResponse>> {
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,
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,
)
.await
.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<WorkspaceApi>,
Extension(actor): Extension<RequestActor>,
AxumPath(path): AxumPath<ScopedAttachmentUploadPath>,
) -> ApiResult<Json<AttachmentUploadCancelResponse>> {
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<WorkspaceApi>,
AxumPath((runtime_id, worker_id)): AxumPath<(String, String)>,
Query(query): Query<WorkerFileUploadQuery>,
body: Bytes,
) -> ApiResult<Json<WorkerFileUploadResponse>> {
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<WorkspaceApi>,
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();
web/workspace/src/lib/workspace/console/composer-attachments.ts
+73 -15
View File
@@ -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,29 +46,72 @@ 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) => {
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(grantRequest.responseText);
if (!isUploadGrantResponse(parsed)) {
callbacks.failed("Upload grant returned an invalid response.");
return;
}
uploadId = parsed.upload_id;
} catch {
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));
}
});
request.addEventListener("load", () => {
if (request.status < 200 || request.status >= 300) {
callbacks.failed(`Upload failed (${request.status}).`);
uploadRequest.addEventListener("load", () => {
if (uploadRequest.status < 200 || uploadRequest.status >= 300) {
callbacks.failed(`Upload failed (${uploadRequest.status}).`);
return;
}
try {
const parsed: unknown = JSON.parse(request.responseText);
const parsed: unknown = JSON.parse(uploadRequest.responseText);
if (!isUploadedFileResponse(parsed)) {
callbacks.failed("Upload returned an invalid attachment reference.");
return;
@@ -75,10 +121,22 @@ export function uploadAttachment(
callbacks.failed("Upload returned an invalid response.");
}
});
request.addEventListener("error", () => callbacks.failed("Upload failed."));
request.addEventListener("abort", () => callbacks.failed("Upload cancelled."));
request.send(file);
return request;
uploadRequest.addEventListener("error", () => callbacks.failed("Upload failed."));
uploadRequest.addEventListener("abort", () => callbacks.failed("Upload cancelled."));
uploadRequest.send(file);
});
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(
web/workspace/src/routes/w/[workspaceId]/runtimes/[runtimeId]/workers/[workerId]/console/+page.svelte
+9 -4
View File
@@ -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<ComposerAttachment>): 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);
}