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