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

feat: add session-owned uploaded file attachments

Browse Source
This commit is contained in:
Hare
2026-09-03 04:58:11 +09:00
parent d87441448e
commit 09a33e7283
17 changed files with 1416 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/protocol/src/lib.rs
+83
View File
@@ -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 {
crates/protocol/src/typescript.rs
+4 -1
View File
@@ -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);
crates/session-store/src/fs_store.rs
+150 -2
View File
@@ -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();
crates/session-store/src/lib.rs
+5
View File
@@ -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,
crates/session-store/src/paste_artifact.rs
+37 -18
View File
@@ -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())
})?;
crates/session-store/src/store.rs
+73 -2
View File
@@ -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,
crates/session-store/src/uploaded_file.rs
+299
View File
@@ -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()),
}
}
crates/worker-runtime/src/execution.rs
+46 -1
View File
@@ -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,
crates/worker-runtime/src/http_server.rs
+93 -2
View File
@@ -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 {
crates/worker-runtime/src/runtime.rs
+99
View File
@@ -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() {
crates/worker-runtime/src/worker_backend.rs
+45
View File
@@ -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,
crates/worker/src/controller.rs
+31 -1
View File
@@ -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 {
crates/worker/src/paste_artifact_tool.rs
+106 -10
View File
@@ -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();
crates/worker/src/worker.rs
+85
View File
@@ -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;
crates/workspace-server/src/hosts.rs
+164
View File
@@ -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,
crates/workspace-server/src/server.rs
+91 -3
View File
@@ -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)>,