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/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;