fix: fence attachment grants and lifecycle cleanup
This commit is contained in:
@@ -58,18 +58,52 @@ impl BackendRuntimeTarget {
|
|||||||
file_name: &str,
|
file_name: &str,
|
||||||
media_type: &str,
|
media_type: &str,
|
||||||
content: Vec<u8>,
|
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> {
|
) -> Result<protocol::UploadedFileRef, BackendRuntimeClientError> {
|
||||||
let api = BackendApiClient::from_stored_token(&self.base_url)?;
|
let api = BackendApiClient::from_stored_token(&self.base_url)?;
|
||||||
let path = format!(
|
let worker_path = format!(
|
||||||
"/api/w/{}/runtimes/{}/workers/{}/attachments?file_name={}&media_type={}",
|
"/api/w/{}/runtimes/{}/workers/{}",
|
||||||
path_segment_encode(&self.workspace_id),
|
path_segment_encode(&self.workspace_id),
|
||||||
path_segment_encode(&self.runtime_id),
|
path_segment_encode(&self.runtime_id),
|
||||||
path_segment_encode(&self.worker_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(file_name),
|
||||||
path_segment_encode(media_type),
|
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
|
let response = api
|
||||||
.request(HttpMethod::POST, &path)?
|
.request(HttpMethod::PUT, &upload_path)?
|
||||||
.body(content)
|
.body(content)
|
||||||
.send()
|
.send()
|
||||||
.await
|
.await
|
||||||
@@ -82,6 +116,27 @@ impl BackendRuntimeTarget {
|
|||||||
.map_err(BackendRuntimeClientError::Http)
|
.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(
|
pub async fn delete_uploaded_file(
|
||||||
&self,
|
&self,
|
||||||
artifact_id: &str,
|
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)]
|
#[derive(Deserialize)]
|
||||||
struct UploadedFileResponse {
|
struct UploadedFileResponse {
|
||||||
file: protocol::UploadedFileRef,
|
file: protocol::UploadedFileRef,
|
||||||
|
|||||||
@@ -20,10 +20,13 @@ use crate::paste_artifact::{read_from_dir, write_to_dir};
|
|||||||
use crate::segment_log::LogEntry;
|
use crate::segment_log::LogEntry;
|
||||||
use crate::store::{Store, StoreError};
|
use crate::store::{Store, StoreError};
|
||||||
use crate::uploaded_file::{
|
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,
|
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 protocol::{PasteArtifactRef, UploadedFileRef};
|
||||||
use std::fs;
|
use std::fs;
|
||||||
use std::io::{Read, Seek, SeekFrom, Write};
|
use std::io::{Read, Seek, SeekFrom, Write};
|
||||||
@@ -119,6 +122,40 @@ impl FsStore {
|
|||||||
self.session_dir(session_id).join("artifacts").join("paste")
|
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)]
|
#[cfg(test)]
|
||||||
fn paste_artifact_path(&self, session_id: SessionId, artifact_id: &str) -> PathBuf {
|
fn paste_artifact_path(&self, session_id: SessionId, artifact_id: &str) -> PathBuf {
|
||||||
self.paste_artifact_dir(session_id)
|
self.paste_artifact_dir(session_id)
|
||||||
@@ -410,6 +447,30 @@ impl Store for FsStore {
|
|||||||
file_name,
|
file_name,
|
||||||
media_type,
|
media_type,
|
||||||
content,
|
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,
|
limits,
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
@@ -440,11 +501,21 @@ impl Store for FsStore {
|
|||||||
.append_lock
|
.append_lock
|
||||||
.lock()
|
.lock()
|
||||||
.map_err(|_| std::io::Error::other("session store append lock was poisoned"))?;
|
.map_err(|_| std::io::Error::other("session store append lock was poisoned"))?;
|
||||||
bind_uploaded_file(
|
let dir = self.paste_artifact_dir(session_id);
|
||||||
&self.paste_artifact_dir(session_id),
|
match bind_uploaded_file(&dir, reference, source_entry_id) {
|
||||||
reference,
|
Err(StoreError::ArtifactAlreadyCommitted) => {
|
||||||
source_entry_id,
|
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(
|
fn delete_uploaded_file(
|
||||||
@@ -464,7 +535,37 @@ impl Store for FsStore {
|
|||||||
.append_lock
|
.append_lock
|
||||||
.lock()
|
.lock()
|
||||||
.map_err(|_| std::io::Error::other("session store append lock was poisoned"))?;
|
.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(
|
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]
|
#[test]
|
||||||
fn uploaded_files_are_session_scoped_integrity_checked_and_removable() {
|
fn uploaded_files_are_session_scoped_integrity_checked_and_removable() {
|
||||||
let tmp = tempfile::TempDir::new().unwrap();
|
let tmp = tempfile::TempDir::new().unwrap();
|
||||||
@@ -715,18 +855,51 @@ mod tests {
|
|||||||
store.write_uploaded_file(session_id, "README.txt", "text/plain", b"y", limits),
|
store.write_uploaded_file(session_id, "README.txt", "text/plain", b"y", limits),
|
||||||
Err(StoreError::InvalidUploadedFileName)
|
Err(StoreError::InvalidUploadedFileName)
|
||||||
));
|
));
|
||||||
|
store
|
||||||
|
.bind_uploaded_file(session_id, &pending, "entry-from-failed-submit")
|
||||||
|
.unwrap();
|
||||||
let bound = store
|
let bound = store
|
||||||
.bind_uploaded_file(session_id, &pending, "entry-upload")
|
.bind_uploaded_file(session_id, &pending, "entry-upload")
|
||||||
.unwrap();
|
.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
|
let other = store
|
||||||
.write_uploaded_file(session_id, "other.txt", "text/plain", b"z", limits)
|
.write_uploaded_file(session_id, "other.txt", "text/plain", b"z", limits)
|
||||||
.unwrap();
|
.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!(
|
assert_eq!(
|
||||||
store.delete_uncommitted_uploaded_files(session_id).unwrap(),
|
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, &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");
|
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
|
store
|
||||||
.write_paste_artifact(
|
.write_paste_artifact(
|
||||||
session_id,
|
session_id,
|
||||||
|
|||||||
@@ -70,6 +70,7 @@ pub use system_item::{
|
|||||||
pub use uploaded_file::{
|
pub use uploaded_file::{
|
||||||
DEFAULT_MAX_FILES_PER_SUBMISSION, DEFAULT_MAX_SESSION_ARTIFACT_BYTES,
|
DEFAULT_MAX_FILES_PER_SUBMISSION, DEFAULT_MAX_SESSION_ARTIFACT_BYTES,
|
||||||
DEFAULT_MAX_SESSION_UPLOADED_FILES, DEFAULT_MAX_UPLOADED_FILE_BYTES, UploadedFileLimits,
|
DEFAULT_MAX_SESSION_UPLOADED_FILES, DEFAULT_MAX_UPLOADED_FILE_BYTES, UploadedFileLimits,
|
||||||
|
UploadedFileUploadContext,
|
||||||
};
|
};
|
||||||
pub use worker_metadata::{
|
pub use worker_metadata::{
|
||||||
CombinedStore, FsWorkerStore, WorkerActiveSegmentRef, WorkerAggregateStore, WorkerMetadata,
|
CombinedStore, FsWorkerStore, WorkerActiveSegmentRef, WorkerAggregateStore, WorkerMetadata,
|
||||||
|
|||||||
@@ -420,6 +420,7 @@ pub fn save_config_changed(
|
|||||||
/// [`fork_at`] or [`ensure_head_or_fork`] instead.
|
/// [`fork_at`] or [`ensure_head_or_fork`] instead.
|
||||||
pub fn fork(
|
pub fn fork(
|
||||||
store: &impl Store,
|
store: &impl Store,
|
||||||
|
source_session_id: SessionId,
|
||||||
state: SegmentStartState<'_>,
|
state: SegmentStartState<'_>,
|
||||||
) -> Result<(SessionId, SegmentId), StoreError> {
|
) -> Result<(SessionId, SegmentId), StoreError> {
|
||||||
let session_id = crate::new_session_id();
|
let session_id = crate::new_session_id();
|
||||||
@@ -434,6 +435,7 @@ pub fn fork(
|
|||||||
compacted_from: None,
|
compacted_from: None,
|
||||||
};
|
};
|
||||||
store.create_segment(session_id, fork_id, &[entry])?;
|
store.create_segment(session_id, fork_id, &[entry])?;
|
||||||
|
store.copy_committed_uploaded_files(source_session_id, session_id)?;
|
||||||
Ok((session_id, fork_id))
|
Ok((session_id, fork_id))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -13,7 +13,9 @@
|
|||||||
|
|
||||||
use crate::event_trace::TraceEntry;
|
use crate::event_trace::TraceEntry;
|
||||||
use crate::segment_log::LogEntry;
|
use crate::segment_log::LogEntry;
|
||||||
use crate::{PasteArtifactLimits, SegmentId, SessionId, UploadedFileLimits};
|
use crate::{
|
||||||
|
PasteArtifactLimits, SegmentId, SessionId, UploadedFileLimits, UploadedFileUploadContext,
|
||||||
|
};
|
||||||
use protocol::{PasteArtifactRef, UploadedFileRef};
|
use protocol::{PasteArtifactRef, UploadedFileRef};
|
||||||
|
|
||||||
/// Errors from the persistence store.
|
/// Errors from the persistence store.
|
||||||
@@ -186,6 +188,18 @@ pub trait Store: Send + Sync {
|
|||||||
Err(StoreError::PasteArtifactUnsupported)
|
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`.
|
/// Read and integrity-check an uploaded file owned by `session_id`.
|
||||||
fn read_uploaded_file(
|
fn read_uploaded_file(
|
||||||
&self,
|
&self,
|
||||||
@@ -225,6 +239,14 @@ pub trait Store: Send + Sync {
|
|||||||
Ok(0)
|
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.
|
/// Append a trace entry to the debug event trace file.
|
||||||
fn append_trace(
|
fn append_trace(
|
||||||
&self,
|
&self,
|
||||||
|
|||||||
@@ -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)]
|
#[derive(Debug, Serialize, Deserialize)]
|
||||||
struct StoredUploadedFile {
|
struct StoredUploadedFile {
|
||||||
file_name: String,
|
file_name: String,
|
||||||
@@ -47,6 +56,8 @@ struct StoredUploadedFile {
|
|||||||
sha256: String,
|
sha256: String,
|
||||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||||
source_entry_id: Option<String>,
|
source_entry_id: Option<String>,
|
||||||
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||||
|
upload_context: Option<UploadedFileUploadContext>,
|
||||||
content_base64: String,
|
content_base64: String,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -186,6 +197,7 @@ pub(crate) fn write_uploaded_file(
|
|||||||
file_name: &str,
|
file_name: &str,
|
||||||
media_type: &str,
|
media_type: &str,
|
||||||
content: &[u8],
|
content: &[u8],
|
||||||
|
context: Option<&UploadedFileUploadContext>,
|
||||||
limits: UploadedFileLimits,
|
limits: UploadedFileLimits,
|
||||||
) -> Result<UploadedFileRef> {
|
) -> Result<UploadedFileRef> {
|
||||||
validate_file_name(file_name)?;
|
validate_file_name(file_name)?;
|
||||||
@@ -226,6 +238,7 @@ pub(crate) fn write_uploaded_file(
|
|||||||
if stored.media_type == media_type
|
if stored.media_type == media_type
|
||||||
&& stored.byte_len == byte_len
|
&& stored.byte_len == byte_len
|
||||||
&& stored.sha256 == sha256
|
&& stored.sha256 == sha256
|
||||||
|
&& stored.upload_context.as_ref() == context
|
||||||
{
|
{
|
||||||
let artifact_id = path
|
let artifact_id = path
|
||||||
.file_name()
|
.file_name()
|
||||||
@@ -264,6 +277,7 @@ pub(crate) fn write_uploaded_file(
|
|||||||
byte_len,
|
byte_len,
|
||||||
sha256: sha256.clone(),
|
sha256: sha256.clone(),
|
||||||
source_entry_id: None,
|
source_entry_id: None,
|
||||||
|
upload_context: context.cloned(),
|
||||||
content_base64: BASE64.encode(content),
|
content_base64: BASE64.encode(content),
|
||||||
};
|
};
|
||||||
let path = record_path(dir, &artifact_id)?;
|
let path = record_path(dir, &artifact_id)?;
|
||||||
@@ -324,6 +338,30 @@ pub(crate) fn read_uploaded_file(dir: &Path, reference: &UploadedFileRef) -> Res
|
|||||||
Ok(content)
|
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(
|
pub(crate) fn bind_uploaded_file(
|
||||||
dir: &Path,
|
dir: &Path,
|
||||||
reference: &UploadedFileRef,
|
reference: &UploadedFileRef,
|
||||||
@@ -338,7 +376,15 @@ pub(crate) fn bind_uploaded_file(
|
|||||||
.write(true)
|
.write(true)
|
||||||
.open(dir.join(".aggregate.lock"))?;
|
.open(dir.join(".aggregate.lock"))?;
|
||||||
FileExt::lock_exclusive(&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 path = record_path(dir, &reference.artifact_id)?;
|
||||||
let mut stored: StoredUploadedFile = serde_json::from_slice(&fs::read(&path)?)?;
|
let mut stored: StoredUploadedFile = serde_json::from_slice(&fs::read(&path)?)?;
|
||||||
if stored.source_entry_id.is_some() {
|
if stored.source_entry_id.is_some() {
|
||||||
@@ -353,6 +399,72 @@ pub(crate) fn bind_uploaded_file(
|
|||||||
Ok(bound)
|
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> {
|
pub(crate) fn delete_uncommitted_uploaded_files(dir: &Path) -> Result<u64> {
|
||||||
fs::create_dir_all(dir)?;
|
fs::create_dir_all(dir)?;
|
||||||
let aggregate_lock = fs::OpenOptions::new()
|
let aggregate_lock = fs::OpenOptions::new()
|
||||||
|
|||||||
@@ -406,6 +406,7 @@ async fn session_fork_creates_new_session() {
|
|||||||
let original_history_len = worker.history().len();
|
let original_history_len = worker.history().len();
|
||||||
let (fork_sid, fork_segid) = session_store::fork(
|
let (fork_sid, fork_segid) = session_store::fork(
|
||||||
&store,
|
&store,
|
||||||
|
sid,
|
||||||
SegmentStartState {
|
SegmentStartState {
|
||||||
system_prompt: worker.get_system_prompt(),
|
system_prompt: worker.get_system_prompt(),
|
||||||
config: worker.request_config(),
|
config: worker.request_config(),
|
||||||
|
|||||||
@@ -928,6 +928,10 @@ impl App {
|
|||||||
Some(self.method_for_run(queued.segments))
|
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>) {
|
pub fn push_error(&mut self, message: impl Into<String>) {
|
||||||
self.blocks.push(Block::Alert {
|
self.blocks.push(Block::Alert {
|
||||||
level: AlertLevel::Error,
|
level: AlertLevel::Error,
|
||||||
|
|||||||
+158
-48
@@ -1,3 +1,4 @@
|
|||||||
|
use std::collections::HashMap;
|
||||||
use std::fmt;
|
use std::fmt;
|
||||||
use std::future::Future;
|
use std::future::Future;
|
||||||
use std::io;
|
use std::io;
|
||||||
@@ -120,11 +121,55 @@ fn copy_selection_to_terminal(app: &mut App) -> bool {
|
|||||||
copy_selection_to_writer(app, &mut stdout)
|
copy_selection_to_writer(app, &mut stdout)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
type AttachmentUploadResult = Result<UploadedFileRef, String>;
|
||||||
|
|
||||||
struct ConsoleConnection<T> {
|
struct ConsoleConnection<T> {
|
||||||
client: Client<T>,
|
client: Client<T>,
|
||||||
standalone_host: Option<StandaloneHost>,
|
standalone_host: Option<StandaloneHost>,
|
||||||
backend_target: Option<BackendRuntimeTarget>,
|
backend_target: Option<BackendRuntimeTarget>,
|
||||||
pending_attachments: Vec<UploadedFileRef>,
|
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> {
|
fn attachment_media_type(path: &Path) -> Option<&'static str> {
|
||||||
@@ -153,6 +198,8 @@ impl<T: Socket> ConsoleConnection<T> {
|
|||||||
standalone_host: Some(host),
|
standalone_host: Some(host),
|
||||||
backend_target: None,
|
backend_target: None,
|
||||||
pending_attachments: Vec::new(),
|
pending_attachments: Vec::new(),
|
||||||
|
upload_tasks: Vec::new(),
|
||||||
|
upload_ids: HashMap::new(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -162,6 +209,8 @@ impl<T: Socket> ConsoleConnection<T> {
|
|||||||
standalone_host: None,
|
standalone_host: None,
|
||||||
backend_target: Some(target),
|
backend_target: Some(target),
|
||||||
pending_attachments: Vec::new(),
|
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?;
|
self.client.send(&prepared).await?;
|
||||||
if carries_attachments {
|
if carries_attachments {
|
||||||
self.pending_attachments.clear();
|
self.pending_attachments.clear();
|
||||||
|
self.upload_ids.clear();
|
||||||
}
|
}
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn upload_path(
|
fn start_upload(
|
||||||
&mut self,
|
&mut self,
|
||||||
path: &Path,
|
path: PathBuf,
|
||||||
) -> Result<UploadedFileRef, Box<dyn std::error::Error>> {
|
result_tx: mpsc::UnboundedSender<AttachmentUploadResult>,
|
||||||
let target = self.backend_target.as_ref().ok_or_else(|| {
|
) -> Result<(), Box<dyn std::error::Error>> {
|
||||||
|
let target = self.backend_target.clone().ok_or_else(|| {
|
||||||
io::Error::new(
|
io::Error::new(
|
||||||
io::ErrorKind::Unsupported,
|
io::ErrorKind::Unsupported,
|
||||||
"client-local file upload is available only for Backend Workers",
|
"client-local file upload is available only for Backend Workers",
|
||||||
)
|
)
|
||||||
})?;
|
})?;
|
||||||
let metadata = tokio::fs::metadata(path).await?;
|
self.upload_tasks.retain(|task| !task.is_finished());
|
||||||
if !metadata.is_file() {
|
let upload_id = self
|
||||||
return Err(io::Error::new(
|
.upload_ids
|
||||||
io::ErrorKind::InvalidInput,
|
.entry(path.clone())
|
||||||
"attachment path is not a file",
|
.or_insert_with(|| uuid::Uuid::now_v7().to_string())
|
||||||
)
|
.clone();
|
||||||
.into());
|
self.upload_tasks.push(tokio::spawn(async move {
|
||||||
}
|
let result = upload_client_path(&target, &path, &upload_id)
|
||||||
if metadata.len() > 10 * 1024 * 1024 {
|
.await
|
||||||
return Err(io::Error::new(
|
.map_err(|error| error.to_string());
|
||||||
io::ErrorKind::InvalidInput,
|
let _ = result_tx.send(result);
|
||||||
"attachment exceeds the 10 MiB limit",
|
}));
|
||||||
)
|
Ok(())
|
||||||
.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)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn clear_pending_attachments(&mut self) {
|
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);
|
let references = std::mem::take(&mut self.pending_attachments);
|
||||||
if let Some(target) = &self.backend_target {
|
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 {
|
for reference in references {
|
||||||
let _ = target.delete_uploaded_file(&reference.artifact_id).await;
|
let _ = target.delete_uploaded_file(&reference.artifact_id).await;
|
||||||
}
|
}
|
||||||
@@ -626,11 +665,13 @@ enum E2eRewindInput {
|
|||||||
enum LoopInput<P> {
|
enum LoopInput<P> {
|
||||||
Terminal(TerminalEventResult),
|
Terminal(TerminalEventResult),
|
||||||
Worker(P),
|
Worker(P),
|
||||||
|
Upload(AttachmentUploadResult),
|
||||||
Tick,
|
Tick,
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn next_loop_input<P, F, T>(
|
async fn next_loop_input<P, F, T>(
|
||||||
term_rx: &mut mpsc::UnboundedReceiver<TerminalEventResult>,
|
term_rx: &mut mpsc::UnboundedReceiver<TerminalEventResult>,
|
||||||
|
upload_rx: &mut mpsc::UnboundedReceiver<AttachmentUploadResult>,
|
||||||
connected: bool,
|
connected: bool,
|
||||||
pod_next: F,
|
pod_next: F,
|
||||||
animate: bool,
|
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),
|
event = pod_next, if connected => LoopInput::Worker(event),
|
||||||
_ = animation_tick, if animate => LoopInput::Tick,
|
_ = animation_tick, if animate => LoopInput::Tick,
|
||||||
}
|
}
|
||||||
@@ -660,13 +704,14 @@ async fn drain_terminal_events<T: Socket>(
|
|||||||
app: &mut App,
|
app: &mut App,
|
||||||
client: &mut ConsoleConnection<T>,
|
client: &mut ConsoleConnection<T>,
|
||||||
term_rx: &mut mpsc::UnboundedReceiver<TerminalEventResult>,
|
term_rx: &mut mpsc::UnboundedReceiver<TerminalEventResult>,
|
||||||
|
upload_tx: &mpsc::UnboundedSender<AttachmentUploadResult>,
|
||||||
) -> Result<bool, Box<dyn std::error::Error>> {
|
) -> Result<bool, Box<dyn std::error::Error>> {
|
||||||
let mut handled = false;
|
let mut handled = false;
|
||||||
for _ in 0..TERMINAL_EVENT_DRAIN_LIMIT {
|
for _ in 0..TERMINAL_EVENT_DRAIN_LIMIT {
|
||||||
match term_rx.try_recv() {
|
match term_rx.try_recv() {
|
||||||
Ok(event) => {
|
Ok(event) => {
|
||||||
handled = true;
|
handled = true;
|
||||||
handle_terminal_event(app, client, event?).await?;
|
handle_terminal_event(app, client, upload_tx, event?).await?;
|
||||||
if app.quit {
|
if app.quit {
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
@@ -708,6 +753,7 @@ async fn run_loop<T: Socket>(
|
|||||||
client: &mut ConsoleConnection<T>,
|
client: &mut ConsoleConnection<T>,
|
||||||
) -> Result<(), Box<dyn std::error::Error>> {
|
) -> Result<(), Box<dyn std::error::Error>> {
|
||||||
let (_terminal_reader, mut term_rx) = TerminalEventReader::spawn()?;
|
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));
|
let mut animation_tick = tokio::time::interval(Duration::from_millis(80));
|
||||||
animation_tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
animation_tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
||||||
|
|
||||||
@@ -718,7 +764,8 @@ async fn run_loop<T: Socket>(
|
|||||||
break;
|
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 {
|
if app.quit {
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
@@ -730,6 +777,7 @@ async fn run_loop<T: Socket>(
|
|||||||
|
|
||||||
match next_loop_input(
|
match next_loop_input(
|
||||||
&mut term_rx,
|
&mut term_rx,
|
||||||
|
&mut upload_rx,
|
||||||
app.connected,
|
app.connected,
|
||||||
client.next_event(),
|
client.next_event(),
|
||||||
app.running,
|
app.running,
|
||||||
@@ -738,8 +786,29 @@ async fn run_loop<T: Socket>(
|
|||||||
.await
|
.await
|
||||||
{
|
{
|
||||||
LoopInput::Terminal(term_event) => {
|
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? {
|
LoopInput::Worker(event) => match event? {
|
||||||
Some(ev) => {
|
Some(ev) => {
|
||||||
if let Some(method) = app.handle_worker_event(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>(
|
async fn handle_terminal_event<T: Socket>(
|
||||||
app: &mut App,
|
app: &mut App,
|
||||||
client: &mut ConsoleConnection<T>,
|
client: &mut ConsoleConnection<T>,
|
||||||
|
upload_tx: &mpsc::UnboundedSender<AttachmentUploadResult>,
|
||||||
event: TermEvent,
|
event: TermEvent,
|
||||||
) -> Result<(), Box<dyn std::error::Error>> {
|
) -> Result<(), Box<dyn std::error::Error>> {
|
||||||
match event {
|
match event {
|
||||||
TermEvent::Key(key) => {
|
TermEvent::Key(key) => {
|
||||||
if let Some(method) = handle_key(app, key) {
|
if let Some(method) = handle_key(app, key) {
|
||||||
if let Some(path) = attachment_command_path(&method) {
|
if let Some(path) = attachment_command_path(&method) {
|
||||||
match client.upload_path(&path).await {
|
match client.start_upload(path, upload_tx.clone()) {
|
||||||
Ok(reference) => app.flash_actionbar_notice(
|
Ok(()) => app.flash_actionbar_notice(
|
||||||
format!(
|
"Uploading attachment… Use /clear-attachments to cancel.",
|
||||||
"Attached {} ({} bytes); it will be sent with the next message.",
|
|
||||||
reference.file_name, reference.byte_len
|
|
||||||
),
|
|
||||||
ActionbarNoticeLevel::Info,
|
ActionbarNoticeLevel::Info,
|
||||||
ActionbarNoticeSource::Tui,
|
ActionbarNoticeSource::Tui,
|
||||||
Duration::from_secs(6),
|
Duration::from_secs(30),
|
||||||
),
|
),
|
||||||
Err(error) => {
|
Err(error) => {
|
||||||
app.push_error(format!("Attachment upload failed: {error}"));
|
app.push_error(format!("Attachment upload failed: {error}"));
|
||||||
@@ -814,7 +881,12 @@ async fn handle_terminal_event<T: Socket>(
|
|||||||
Duration::from_secs(4),
|
Duration::from_secs(4),
|
||||||
);
|
);
|
||||||
} else {
|
} else {
|
||||||
|
let sends_attachments = matches!(method, Method::Run { .. })
|
||||||
|
&& !client.pending_attachments.is_empty();
|
||||||
client.send(&method).await?;
|
client.send(&method).await?;
|
||||||
|
if sends_attachments {
|
||||||
|
app.clear_actionbar_notice();
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -1420,10 +1492,12 @@ mod tests {
|
|||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn animation_tick_wakes_loop_while_running() {
|
async fn animation_tick_wakes_loop_while_running() {
|
||||||
let (_tx, mut rx) = mpsc::unbounded_channel::<TerminalEventResult>();
|
let (_tx, mut rx) = mpsc::unbounded_channel::<TerminalEventResult>();
|
||||||
|
let (_upload_tx, mut upload_rx) = mpsc::unbounded_channel();
|
||||||
|
|
||||||
assert!(matches!(
|
assert!(matches!(
|
||||||
next_loop_input(
|
next_loop_input(
|
||||||
&mut rx,
|
&mut rx,
|
||||||
|
&mut upload_rx,
|
||||||
true,
|
true,
|
||||||
std::future::pending::<Option<u8>>(),
|
std::future::pending::<Option<u8>>(),
|
||||||
true,
|
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]
|
#[tokio::test]
|
||||||
async fn terminal_event_is_selected_before_ready_worker_event() {
|
async fn terminal_event_is_selected_before_ready_worker_event() {
|
||||||
let (tx, mut rx) = mpsc::unbounded_channel();
|
let (tx, mut rx) = mpsc::unbounded_channel();
|
||||||
|
let (_upload_tx, mut upload_rx) = mpsc::unbounded_channel();
|
||||||
tx.send(Ok(TermEvent::Key(KeyEvent::new(
|
tx.send(Ok(TermEvent::Key(KeyEvent::new(
|
||||||
KeyCode::Char('x'),
|
KeyCode::Char('x'),
|
||||||
KeyModifiers::NONE,
|
KeyModifiers::NONE,
|
||||||
@@ -1445,6 +1551,7 @@ mod tests {
|
|||||||
|
|
||||||
match next_loop_input(
|
match next_loop_input(
|
||||||
&mut rx,
|
&mut rx,
|
||||||
|
&mut upload_rx,
|
||||||
true,
|
true,
|
||||||
std::future::ready(Some(())),
|
std::future::ready(Some(())),
|
||||||
false,
|
false,
|
||||||
@@ -1462,9 +1569,11 @@ mod tests {
|
|||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn terminal_event_is_preserved_after_worker_event_wins() {
|
async fn terminal_event_is_preserved_after_worker_event_wins() {
|
||||||
let (tx, mut rx) = mpsc::unbounded_channel();
|
let (tx, mut rx) = mpsc::unbounded_channel();
|
||||||
|
let (_upload_tx, mut upload_rx) = mpsc::unbounded_channel();
|
||||||
|
|
||||||
match next_loop_input(
|
match next_loop_input(
|
||||||
&mut rx,
|
&mut rx,
|
||||||
|
&mut upload_rx,
|
||||||
true,
|
true,
|
||||||
std::future::ready(Some(1_u8)),
|
std::future::ready(Some(1_u8)),
|
||||||
false,
|
false,
|
||||||
@@ -1484,6 +1593,7 @@ mod tests {
|
|||||||
|
|
||||||
match next_loop_input(
|
match next_loop_input(
|
||||||
&mut rx,
|
&mut rx,
|
||||||
|
&mut upload_rx,
|
||||||
true,
|
true,
|
||||||
std::future::ready(Some(2_u8)),
|
std::future::ready(Some(2_u8)),
|
||||||
false,
|
false,
|
||||||
|
|||||||
@@ -393,6 +393,7 @@ pub trait WorkerExecutionBackend: Send + Sync + 'static {
|
|||||||
_file_name: &str,
|
_file_name: &str,
|
||||||
_media_type: &str,
|
_media_type: &str,
|
||||||
_content: &[u8],
|
_content: &[u8],
|
||||||
|
_context: Option<&session_store::UploadedFileUploadContext>,
|
||||||
) -> Result<UploadedFileRef, WorkerExecutionResult> {
|
) -> Result<UploadedFileRef, WorkerExecutionResult> {
|
||||||
Err(WorkerExecutionResult::unsupported(
|
Err(WorkerExecutionResult::unsupported(
|
||||||
WorkerExecutionOperation::UploadFile,
|
WorkerExecutionOperation::UploadFile,
|
||||||
@@ -546,9 +547,10 @@ impl WorkerExecutionBackendRef {
|
|||||||
file_name: &str,
|
file_name: &str,
|
||||||
media_type: &str,
|
media_type: &str,
|
||||||
content: &[u8],
|
content: &[u8],
|
||||||
|
context: Option<&session_store::UploadedFileUploadContext>,
|
||||||
) -> Result<UploadedFileRef, WorkerExecutionResult> {
|
) -> Result<UploadedFileRef, WorkerExecutionResult> {
|
||||||
self.backend
|
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(
|
pub(crate) fn delete_uploaded_file(
|
||||||
|
|||||||
@@ -390,6 +390,16 @@ pub struct RuntimeHttpWorkerInputResponse {
|
|||||||
pub struct RuntimeHttpUploadFileQuery {
|
pub struct RuntimeHttpUploadFileQuery {
|
||||||
pub file_name: String,
|
pub file_name: String,
|
||||||
pub media_type: 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)]
|
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||||
@@ -1455,14 +1465,74 @@ async fn upload_worker_file(
|
|||||||
body: Bytes,
|
body: Bytes,
|
||||||
) -> RestResult<RuntimeHttpUploadedFileResponse> {
|
) -> RestResult<RuntimeHttpUploadedFileResponse> {
|
||||||
let worker_ref = worker_ref_for(&state.runtime, worker_id)?;
|
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())? {
|
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,
|
&scope,
|
||||||
&worker_ref,
|
&worker_ref,
|
||||||
&query.file_name,
|
&query.file_name,
|
||||||
&query.media_type,
|
&query.media_type,
|
||||||
&body,
|
&body,
|
||||||
),
|
),
|
||||||
|
}
|
||||||
|
}
|
||||||
None => state.runtime.upload_worker_file(
|
None => state.runtime.upload_worker_file(
|
||||||
&worker_ref,
|
&worker_ref,
|
||||||
&query.file_name,
|
&query.file_name,
|
||||||
|
|||||||
@@ -33,3 +33,4 @@ pub mod working_directory;
|
|||||||
pub use fs_store::{FsRuntimeStore, FsRuntimeStoreOptions};
|
pub use fs_store::{FsRuntimeStore, FsRuntimeStoreOptions};
|
||||||
pub use management::RuntimeOptions;
|
pub use management::RuntimeOptions;
|
||||||
pub use runtime::{Runtime, RuntimeWorkspaceScope};
|
pub use runtime::{Runtime, RuntimeWorkspaceScope};
|
||||||
|
pub use session_store::UploadedFileUploadContext;
|
||||||
|
|||||||
@@ -1224,6 +1224,41 @@ impl Runtime {
|
|||||||
file_name: &str,
|
file_name: &str,
|
||||||
media_type: &str,
|
media_type: &str,
|
||||||
content: &[u8],
|
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> {
|
) -> Result<protocol::UploadedFileRef, RuntimeError> {
|
||||||
let (backend, handle) = {
|
let (backend, handle) = {
|
||||||
let state = self.lock()?;
|
let state = self.lock()?;
|
||||||
@@ -1244,7 +1279,7 @@ impl Runtime {
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
backend
|
backend
|
||||||
.upload_file(&handle, file_name, media_type, content)
|
.upload_file(&handle, file_name, media_type, content, context)
|
||||||
.map_err(|result| RuntimeError::WorkerExecutionRejected {
|
.map_err(|result| RuntimeError::WorkerExecutionRejected {
|
||||||
worker_id: worker_ref.worker_id.clone(),
|
worker_id: worker_ref.worker_id.clone(),
|
||||||
operation: result.operation,
|
operation: result.operation,
|
||||||
|
|||||||
@@ -2031,14 +2031,19 @@ where
|
|||||||
file_name: &str,
|
file_name: &str,
|
||||||
media_type: &str,
|
media_type: &str,
|
||||||
content: &[u8],
|
content: &[u8],
|
||||||
|
context: Option<&session_store::UploadedFileUploadContext>,
|
||||||
) -> Result<protocol::UploadedFileRef, WorkerExecutionResult> {
|
) -> Result<protocol::UploadedFileRef, WorkerExecutionResult> {
|
||||||
let (worker, _, _) = self.get_execution(handle).map_err(|mut result| {
|
let (worker, _, _) = self.get_execution(handle).map_err(|mut result| {
|
||||||
result.operation = WorkerExecutionOperation::UploadFile;
|
result.operation = WorkerExecutionOperation::UploadFile;
|
||||||
result
|
result
|
||||||
})?;
|
})?;
|
||||||
worker
|
let uploaded = match context {
|
||||||
.upload_file(file_name, media_type, content)
|
Some(context) => {
|
||||||
.map_err(|error| {
|
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(
|
WorkerExecutionResult::rejected(
|
||||||
WorkerExecutionOperation::UploadFile,
|
WorkerExecutionOperation::UploadFile,
|
||||||
format!("uploaded_file_rejected: {error}"),
|
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() {
|
let execution = match self.workers.lock() {
|
||||||
Ok(mut workers) => workers.remove(handle.worker_ref()),
|
Ok(mut workers) => workers.remove(handle.worker_ref()),
|
||||||
Err(_) => {
|
Err(_) => {
|
||||||
@@ -2179,6 +2165,7 @@ where
|
|||||||
"execution handle does not reference a live Worker",
|
"execution handle does not reference a live Worker",
|
||||||
);
|
);
|
||||||
};
|
};
|
||||||
|
let artifact_cleanup = execution.handle.clone();
|
||||||
let shutdown = execution.shutdown.clone();
|
let shutdown = execution.shutdown.clone();
|
||||||
let result = self.send_method(
|
let result = self.send_method(
|
||||||
WorkerExecutionOperation::Stop,
|
WorkerExecutionOperation::Stop,
|
||||||
@@ -2198,7 +2185,13 @@ where
|
|||||||
}
|
}
|
||||||
Ok(())
|
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),
|
Err(message) => WorkerExecutionResult::errored(WorkerExecutionOperation::Stop, message),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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(
|
pub fn delete_uploaded_file(
|
||||||
&self,
|
&self,
|
||||||
artifact_id: &str,
|
artifact_id: &str,
|
||||||
|
|||||||
@@ -1048,6 +1048,7 @@ pub trait WorkspaceWorkerRuntime: Send + Sync {
|
|||||||
_file_name: &str,
|
_file_name: &str,
|
||||||
_media_type: &str,
|
_media_type: &str,
|
||||||
_content: &[u8],
|
_content: &[u8],
|
||||||
|
_context: Option<&worker_runtime::UploadedFileUploadContext>,
|
||||||
) -> Result<protocol::UploadedFileRef, RuntimeRegistryError> {
|
) -> Result<protocol::UploadedFileRef, RuntimeRegistryError> {
|
||||||
Err(RuntimeRegistryError::RuntimeOperationFailed {
|
Err(RuntimeRegistryError::RuntimeOperationFailed {
|
||||||
runtime_id: self.runtime_id().to_string(),
|
runtime_id: self.runtime_id().to_string(),
|
||||||
@@ -1578,6 +1579,7 @@ impl RuntimeRegistry {
|
|||||||
file_name: &str,
|
file_name: &str,
|
||||||
media_type: &str,
|
media_type: &str,
|
||||||
content: &[u8],
|
content: &[u8],
|
||||||
|
context: Option<&worker_runtime::UploadedFileUploadContext>,
|
||||||
) -> Result<protocol::UploadedFileRef, RuntimeRegistryError> {
|
) -> Result<protocol::UploadedFileRef, RuntimeRegistryError> {
|
||||||
let runtime_id = worker.runtime_id.as_str();
|
let runtime_id = worker.runtime_id.as_str();
|
||||||
let worker_id = worker.worker_id.as_str();
|
let worker_id = worker.worker_id.as_str();
|
||||||
@@ -1592,7 +1594,7 @@ impl RuntimeRegistry {
|
|||||||
lookup.diagnostics,
|
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(
|
pub fn delete_worker_uploaded_file(
|
||||||
@@ -2588,15 +2590,26 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime {
|
|||||||
file_name: &str,
|
file_name: &str,
|
||||||
media_type: &str,
|
media_type: &str,
|
||||||
content: &[u8],
|
content: &[u8],
|
||||||
|
context: Option<&worker_runtime::UploadedFileUploadContext>,
|
||||||
) -> Result<protocol::UploadedFileRef, RuntimeRegistryError> {
|
) -> Result<protocol::UploadedFileRef, RuntimeRegistryError> {
|
||||||
let worker_ref =
|
let worker_ref =
|
||||||
self.worker_ref(worker_id)
|
self.worker_ref(worker_id)
|
||||||
.ok_or_else(|| RuntimeRegistryError::UnknownWorker {
|
.ok_or_else(|| RuntimeRegistryError::UnknownWorker {
|
||||||
worker: RuntimeWorkerRef::new(&self.runtime_id, worker_id),
|
worker: RuntimeWorkerRef::new(&self.runtime_id, worker_id),
|
||||||
})?;
|
})?;
|
||||||
self.runtime
|
let uploaded = match context {
|
||||||
.upload_worker_file(&worker_ref, file_name, media_type, content)
|
Some(context) => self.runtime.upload_worker_file_with_context(
|
||||||
.map_err(|error| RuntimeRegistryError::RuntimeOperationFailed {
|
&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(),
|
runtime_id: self.runtime_id.clone(),
|
||||||
code: "embedded_worker_file_upload_failed".to_string(),
|
code: "embedded_worker_file_upload_failed".to_string(),
|
||||||
message: error.to_string(),
|
message: error.to_string(),
|
||||||
@@ -3665,13 +3678,24 @@ impl WorkspaceWorkerRuntime for RemoteWorkerRuntime {
|
|||||||
file_name: &str,
|
file_name: &str,
|
||||||
media_type: &str,
|
media_type: &str,
|
||||||
content: &[u8],
|
content: &[u8],
|
||||||
|
context: Option<&worker_runtime::UploadedFileUploadContext>,
|
||||||
) -> Result<protocol::UploadedFileRef, RuntimeRegistryError> {
|
) -> Result<protocol::UploadedFileRef, RuntimeRegistryError> {
|
||||||
let path = format!(
|
let mut path = format!(
|
||||||
"/v1/workers/{}/attachments?file_name={}&media_type={}",
|
"/v1/workers/{}/attachments?file_name={}&media_type={}",
|
||||||
url_path_segment_encode(worker_id),
|
url_path_segment_encode(worker_id),
|
||||||
url_query_value_encode(file_name),
|
url_query_value_encode(file_name),
|
||||||
url_query_value_encode(media_type),
|
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)
|
self.post_bytes::<RuntimeHttpUploadedFileResponse>(&path, content)
|
||||||
.map(|response| response.file)
|
.map(|response| response.file)
|
||||||
.map_err(|diagnostic| RuntimeRegistryError::RuntimeOperationFailed {
|
.map_err(|diagnostic| RuntimeRegistryError::RuntimeOperationFailed {
|
||||||
|
|||||||
@@ -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)]
|
#[derive(Clone)]
|
||||||
pub struct WorkspaceApi {
|
pub struct WorkspaceApi {
|
||||||
pub(crate) config: ServerConfig,
|
pub(crate) config: ServerConfig,
|
||||||
@@ -524,6 +591,7 @@ pub struct WorkspaceApi {
|
|||||||
workdir_remove_locks: Arc<Mutex<HashMap<String, Arc<std::sync::Mutex<()>>>>>,
|
workdir_remove_locks: Arc<Mutex<HashMap<String, Arc<std::sync::Mutex<()>>>>>,
|
||||||
workdir_remove_attempt_owner: WorkdirRemovalAttemptOwner,
|
workdir_remove_attempt_owner: WorkdirRemovalAttemptOwner,
|
||||||
worker_control_locks: Arc<Mutex<HashMap<String, Arc<tokio::sync::Mutex<()>>>>>,
|
worker_control_locks: Arc<Mutex<HashMap<String, Arc<tokio::sync::Mutex<()>>>>>,
|
||||||
|
attachment_upload_grants: Arc<Mutex<HashMap<String, AttachmentUploadGrant>>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Clone)]
|
#[derive(Clone)]
|
||||||
@@ -1645,6 +1713,7 @@ impl WorkspaceApi {
|
|||||||
workdir_remove_locks: Arc::new(Mutex::new(HashMap::new())),
|
workdir_remove_locks: Arc::new(Mutex::new(HashMap::new())),
|
||||||
workdir_remove_attempt_owner: current_workdir_removal_attempt_owner()?,
|
workdir_remove_attempt_owner: current_workdir_removal_attempt_owner()?,
|
||||||
worker_control_locks: Arc::new(Mutex::new(HashMap::new())),
|
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 {
|
if let Some(dispatcher) = worker_remove_dispatcher {
|
||||||
dispatcher
|
dispatcher
|
||||||
@@ -2681,21 +2750,19 @@ fn build_inner_router(api: WorkspaceApi) -> Router {
|
|||||||
"/api/runtimes/{runtime_id}/workers/{worker_id}/input",
|
"/api/runtimes/{runtime_id}/workers/{worker_id}/input",
|
||||||
post(send_runtime_worker_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(
|
.route(
|
||||||
"/api/w/{workspace_id}/runtimes/{runtime_id}/workers/{worker_id}/input",
|
"/api/w/{workspace_id}/runtimes/{runtime_id}/workers/{worker_id}/input",
|
||||||
post(scoped_send_runtime_worker_input),
|
post(scoped_send_runtime_worker_input),
|
||||||
)
|
)
|
||||||
.route(
|
.route(
|
||||||
"/api/w/{workspace_id}/runtimes/{runtime_id}/workers/{worker_id}/attachments",
|
"/api/w/{workspace_id}/runtimes/{runtime_id}/workers/{worker_id}/attachment-upload-grants",
|
||||||
post(scoped_upload_runtime_worker_file).layer(DefaultBodyLimit::max(MAX_WORKER_FILE_UPLOAD_BYTES)),
|
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(
|
.route(
|
||||||
"/api/w/{workspace_id}/runtimes/{runtime_id}/workers/{worker_id}/attachments/{artifact_id}",
|
"/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 {
|
struct WorkerFileUploadQuery {
|
||||||
file_name: String,
|
file_name: String,
|
||||||
media_type: String,
|
media_type: String,
|
||||||
|
#[serde(default)]
|
||||||
|
upload_id: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Serialize)]
|
#[derive(Debug, Serialize)]
|
||||||
@@ -10768,20 +10837,225 @@ struct WorkerFileDeleteResponse {
|
|||||||
deleted: bool,
|
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>,
|
State(api): State<WorkspaceApi>,
|
||||||
|
Extension(actor): Extension<RequestActor>,
|
||||||
AxumPath(path): AxumPath<ScopedRuntimeWorkerPath>,
|
AxumPath(path): AxumPath<ScopedRuntimeWorkerPath>,
|
||||||
Query(query): Query<WorkerFileUploadQuery>,
|
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,
|
body: Bytes,
|
||||||
) -> ApiResult<Json<WorkerFileUploadResponse>> {
|
) -> ApiResult<Json<WorkerFileUploadResponse>> {
|
||||||
validate_workspace_scope(&api, &path.workspace_id)?;
|
validate_workspace_scope(&api, &path.workspace_id)?;
|
||||||
upload_runtime_worker_file(
|
let body_sha256 = attachment_body_sha256(&body);
|
||||||
State(api),
|
let grant = {
|
||||||
AxumPath((path.worker.runtime_id, path.worker.worker_id)),
|
let mut grants = api
|
||||||
Query(query),
|
.attachment_upload_grants
|
||||||
body,
|
.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(
|
async fn scoped_delete_runtime_worker_uploaded_file(
|
||||||
@@ -13323,20 +13597,6 @@ async fn send_runtime_worker_input(
|
|||||||
Ok(Json(result))
|
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(
|
async fn delete_runtime_worker_uploaded_file(
|
||||||
State(api): State<WorkspaceApi>,
|
State(api): State<WorkspaceApi>,
|
||||||
AxumPath((runtime_id, worker_id, artifact_id)): AxumPath<(String, String, String)>,
|
AxumPath((runtime_id, worker_id, artifact_id)): AxumPath<(String, String, String)>,
|
||||||
@@ -16084,6 +16344,77 @@ mod tests {
|
|||||||
SqliteWorkspaceStore, TrustedRuntimeRecord, UserRecord, WorkspaceRecord,
|
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]
|
#[tokio::test]
|
||||||
async fn command_session_survives_attachment_refresh_until_worker_revocation() {
|
async fn command_session_survives_attachment_refresh_until_worker_revocation() {
|
||||||
let directory = tempfile::tempdir().unwrap();
|
let directory = tempfile::tempdir().unwrap();
|
||||||
|
|||||||
@@ -5,15 +5,18 @@ export const MAX_FILES_PER_SUBMISSION = 8;
|
|||||||
|
|
||||||
export type AttachmentUploadState = "uploading" | "uploaded" | "failed";
|
export type AttachmentUploadState = "uploading" | "uploaded" | "failed";
|
||||||
|
|
||||||
|
export type AttachmentUploadHandle = { abort(): void };
|
||||||
|
|
||||||
export type ComposerAttachment = {
|
export type ComposerAttachment = {
|
||||||
id: number;
|
id: number;
|
||||||
file: File;
|
file: File;
|
||||||
uploadPath: string;
|
uploadPath: string;
|
||||||
|
uploadId: string;
|
||||||
state: AttachmentUploadState;
|
state: AttachmentUploadState;
|
||||||
progress: number;
|
progress: number;
|
||||||
reference: UploadedFileRef | null;
|
reference: UploadedFileRef | null;
|
||||||
error: string | null;
|
error: string | null;
|
||||||
request: XMLHttpRequest | null;
|
request: AttachmentUploadHandle | null;
|
||||||
};
|
};
|
||||||
|
|
||||||
export function acceptedAttachmentMediaType(mediaType: string): boolean {
|
export function acceptedAttachmentMediaType(mediaType: string): boolean {
|
||||||
@@ -43,29 +46,72 @@ export type AttachmentUploadCallbacks = {
|
|||||||
};
|
};
|
||||||
|
|
||||||
export function uploadAttachment(
|
export function uploadAttachment(
|
||||||
path: string,
|
workerPath: string,
|
||||||
file: File,
|
file: File,
|
||||||
|
uploadId: string,
|
||||||
callbacks: AttachmentUploadCallbacks,
|
callbacks: AttachmentUploadCallbacks,
|
||||||
): XMLHttpRequest {
|
): AttachmentUploadHandle {
|
||||||
const request = new XMLHttpRequest();
|
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({
|
const query = new URLSearchParams({
|
||||||
file_name: file.name,
|
file_name: file.name,
|
||||||
media_type: file.type,
|
media_type: file.type,
|
||||||
|
upload_id: uploadId,
|
||||||
});
|
});
|
||||||
request.open("POST", `${path}?${query.toString()}`);
|
const grantRequest = new XMLHttpRequest();
|
||||||
request.setRequestHeader("content-type", "application/octet-stream");
|
activeRequest = grantRequest;
|
||||||
request.upload.addEventListener("progress", (event) => {
|
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) {
|
if (event.lengthComputable && event.total > 0) {
|
||||||
callbacks.progress(Math.min(1, event.loaded / event.total));
|
callbacks.progress(Math.min(1, event.loaded / event.total));
|
||||||
}
|
}
|
||||||
});
|
});
|
||||||
request.addEventListener("load", () => {
|
uploadRequest.addEventListener("load", () => {
|
||||||
if (request.status < 200 || request.status >= 300) {
|
if (uploadRequest.status < 200 || uploadRequest.status >= 300) {
|
||||||
callbacks.failed(`Upload failed (${request.status}).`);
|
callbacks.failed(`Upload failed (${uploadRequest.status}).`);
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
try {
|
try {
|
||||||
const parsed: unknown = JSON.parse(request.responseText);
|
const parsed: unknown = JSON.parse(uploadRequest.responseText);
|
||||||
if (!isUploadedFileResponse(parsed)) {
|
if (!isUploadedFileResponse(parsed)) {
|
||||||
callbacks.failed("Upload returned an invalid attachment reference.");
|
callbacks.failed("Upload returned an invalid attachment reference.");
|
||||||
return;
|
return;
|
||||||
@@ -75,10 +121,22 @@ export function uploadAttachment(
|
|||||||
callbacks.failed("Upload returned an invalid response.");
|
callbacks.failed("Upload returned an invalid response.");
|
||||||
}
|
}
|
||||||
});
|
});
|
||||||
request.addEventListener("error", () => callbacks.failed("Upload failed."));
|
uploadRequest.addEventListener("error", () => callbacks.failed("Upload failed."));
|
||||||
request.addEventListener("abort", () => callbacks.failed("Upload cancelled."));
|
uploadRequest.addEventListener("abort", () => callbacks.failed("Upload cancelled."));
|
||||||
request.send(file);
|
uploadRequest.send(file);
|
||||||
return request;
|
});
|
||||||
|
grantRequest.addEventListener("error", () => callbacks.failed("Upload grant failed."));
|
||||||
|
grantRequest.addEventListener("abort", () => callbacks.failed("Upload cancelled."));
|
||||||
|
grantRequest.send();
|
||||||
|
return handle;
|
||||||
|
}
|
||||||
|
|
||||||
|
function isUploadGrantResponse(
|
||||||
|
value: unknown,
|
||||||
|
): value is { upload_id: string; expires_at_ms: number } {
|
||||||
|
return !!value && typeof value === "object" &&
|
||||||
|
"upload_id" in value && typeof value.upload_id === "string" &&
|
||||||
|
"expires_at_ms" in value && typeof value.expires_at_ms === "number";
|
||||||
}
|
}
|
||||||
|
|
||||||
function isUploadedFileResponse(
|
function isUploadedFileResponse(
|
||||||
|
|||||||
+9
-4
@@ -599,7 +599,7 @@
|
|||||||
attachment.request?.abort();
|
attachment.request?.abort();
|
||||||
if (attachment.reference) {
|
if (attachment.reference) {
|
||||||
void fetch(
|
void fetch(
|
||||||
`${attachment.uploadPath}/${encodeURIComponent(attachment.reference.artifact_id)}`,
|
`${attachment.uploadPath}/attachments/${encodeURIComponent(attachment.reference.artifact_id)}`,
|
||||||
{ method: "DELETE" },
|
{ method: "DELETE" },
|
||||||
).catch(() => undefined);
|
).catch(() => undefined);
|
||||||
}
|
}
|
||||||
@@ -639,7 +639,7 @@
|
|||||||
}
|
}
|
||||||
|
|
||||||
function attachmentPath(): string {
|
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 {
|
function updateAttachment(id: number, update: Partial<ComposerAttachment>): void {
|
||||||
@@ -661,7 +661,11 @@
|
|||||||
error: null,
|
error: null,
|
||||||
request: 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 }),
|
progress: (progress) => updateAttachment(attachment.id, { progress }),
|
||||||
complete: (reference) =>
|
complete: (reference) =>
|
||||||
updateAttachment(attachment.id, {
|
updateAttachment(attachment.id, {
|
||||||
@@ -688,6 +692,7 @@
|
|||||||
id: nextAttachmentId++,
|
id: nextAttachmentId++,
|
||||||
file,
|
file,
|
||||||
uploadPath: attachmentPath(),
|
uploadPath: attachmentPath(),
|
||||||
|
uploadId: crypto.randomUUID(),
|
||||||
state: "uploading",
|
state: "uploading",
|
||||||
progress: 0,
|
progress: 0,
|
||||||
reference: null,
|
reference: null,
|
||||||
@@ -703,7 +708,7 @@
|
|||||||
attachment.request?.abort();
|
attachment.request?.abort();
|
||||||
attachments = attachments.filter((candidate) => candidate.id !== attachment.id);
|
attachments = attachments.filter((candidate) => candidate.id !== attachment.id);
|
||||||
if (attachment.reference) {
|
if (attachment.reference) {
|
||||||
await fetch(`${attachment.uploadPath}/${encodeURIComponent(attachment.reference.artifact_id)}`, {
|
await fetch(`${attachment.uploadPath}/attachments/${encodeURIComponent(attachment.reference.artifact_id)}`, {
|
||||||
method: "DELETE",
|
method: "DELETE",
|
||||||
}).catch(() => undefined);
|
}).catch(() => undefined);
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user