From 04e296a4ef8047f558b9834c8b743b7e7abaf6bd Mon Sep 17 00:00:00 2001 From: Hare Date: Tue, 1 Sep 2026 17:26:58 +0900 Subject: [PATCH] feat: store large paste inputs as artifacts --- Cargo.lock | 1 + crates/protocol/src/lib.rs | 56 +++ crates/protocol/src/typescript.rs | 10 +- crates/session-store/Cargo.toml | 1 + crates/session-store/src/fs_store.rs | 124 ++++++- crates/session-store/src/lib.rs | 2 + crates/session-store/src/paste_artifact.rs | 156 ++++++++ crates/session-store/src/store.rs | 35 +- crates/session-store/src/worker_metadata.rs | 18 + .../session-store/src/worker_session_store.rs | 73 +++- crates/tui/src/input.rs | 54 ++- crates/tui/src/ui.rs | 11 + crates/worker/src/controller.rs | 14 + crates/worker/src/lib.rs | 1 + crates/worker/src/paste_artifact_tool.rs | 347 ++++++++++++++++++ crates/worker/src/session_history.rs | 16 + crates/worker/src/worker.rs | 230 +++++++++++- crates/worker/tests/controller_test.rs | 14 +- web/workspace/src/lib/generated/protocol.ts | 4 +- .../src/lib/workspace/console/model.test.ts | 18 + .../src/lib/workspace/console/model.ts | 2 + 21 files changed, 1164 insertions(+), 23 deletions(-) create mode 100644 crates/session-store/src/paste_artifact.rs create mode 100644 crates/worker/src/paste_artifact_tool.rs diff --git a/Cargo.lock b/Cargo.lock index c597c3be..b64bd69c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4406,6 +4406,7 @@ dependencies = [ "protocol", "serde", "serde_json", + "sha2 0.11.0", "tempfile", "thiserror 2.0.18", "tokio", diff --git a/crates/protocol/src/lib.rs b/crates/protocol/src/lib.rs index 9d494b5d..2d3ef4e7 100644 --- a/crates/protocol/src/lib.rs +++ b/crates/protocol/src/lib.rs @@ -193,6 +193,23 @@ impl WorkerEvent { /// variants — emits an alert and inserts a `[unknown input segment]` /// placeholder into the LLM context so neither user nor LLM is blind to /// the dropped intent. +/// Session-owned reference to a large pasted-input artifact. +/// +/// The reference contains only bounded integrity and provenance metadata. The +/// artifact body remains in session storage and is available to the model only +/// through the scoped paste-artifact tools installed by Worker. +#[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 PasteArtifactRef { + pub artifact_id: String, + pub byte_len: u64, + pub char_count: u64, + pub line_count: u64, + pub sha256: String, + pub source_entry_id: String, +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[cfg_attr(feature = "typescript", derive(ts_rs::TS))] #[cfg_attr(feature = "json-schema", derive(schemars::JsonSchema))] @@ -210,6 +227,10 @@ pub enum Segment { lines: u32, content: String, }, + /// Internal reference produced when Worker stores a large `Paste` before + /// committing input. Clients may receive this in history/event projections; + /// the body is intentionally absent. + PasteArtifact { artifact: PasteArtifactRef }, /// `@` file-system reference. Worker resolves readable files to /// `[File: ]` attachments and readable normal directories to shallow /// `[Dir: ]` listings; the flattened user text keeps the literal @@ -250,6 +271,18 @@ impl Segment { match seg { Segment::Text { content } => out.push_str(content), Segment::Paste { content, .. } => out.push_str(content), + Segment::PasteArtifact { artifact } => { + use std::fmt::Write as _; + let _ = write!( + out, + "[Large paste stored as artifact {}: {} bytes, {} chars, {} lines, sha256 {}; use SearchInputArtifact and ReadInputArtifact to inspect it]", + artifact.artifact_id, + artifact.byte_len, + artifact.char_count, + artifact.line_count, + artifact.sha256 + ); + } Segment::FileRef { path } => { out.push('@'); out.push_str(path); @@ -1202,6 +1235,29 @@ mod tests { } } + #[test] + fn paste_artifact_segment_roundtrips_without_body() { + let artifact = PasteArtifactRef { + artifact_id: "019ca7c8-57b6-7f05-8edf-524147aba7b2".to_string(), + byte_len: 65_536, + char_count: 65_530, + line_count: 200, + sha256: "a".repeat(64), + source_entry_id: "entry-1".to_string(), + }; + let segment = Segment::PasteArtifact { + artifact: artifact.clone(), + }; + let json = serde_json::to_string(&segment).unwrap(); + assert!(!json.contains("pasted body")); + assert_eq!(serde_json::from_str::(&json).unwrap(), segment); + let projected = Segment::flatten_to_text(&[segment]); + assert!(projected.contains(&artifact.artifact_id)); + assert!(projected.contains("SearchInputArtifact")); + assert!(projected.contains("ReadInputArtifact")); + assert!(!projected.contains("pasted body")); + } + #[test] fn method_run_flow_segment_roundtrip() { let method = Method::Run { diff --git a/crates/protocol/src/typescript.rs b/crates/protocol/src/typescript.rs index bf75e324..ef7a6cf3 100644 --- a/crates/protocol/src/typescript.rs +++ b/crates/protocol/src/typescript.rs @@ -7,10 +7,11 @@ use crate::{ CommandStreamSlice, CompactionLifecycle, CompactionLifecycleState, CompletionEntry, CompletionKind, ErrorCode, Event, Greeting, InFlightBlock, InFlightSnapshot, InFlightToolCallState, InternalWorkerKind, InternalWorkerRef, InternalWorkerSnapshot, - InvokeKind, MemoryWorkerEvent, Method, Permission, RewindSummary, RewindTarget, RewindTargetId, - RunResult, ScopeRule, Segment, SessionContentPart, SessionEntryProvenance, SessionMessageRole, - SessionSnapshot, SessionSnapshotEntry, SessionSnapshotEntryData, SessionToolAttachment, - ToolResultDisposition, TurnResult, WorkerEvent, WorkerStatus, + InvokeKind, MemoryWorkerEvent, Method, PasteArtifactRef, Permission, RewindSummary, + RewindTarget, RewindTargetId, RunResult, ScopeRule, Segment, SessionContentPart, + SessionEntryProvenance, SessionMessageRole, SessionSnapshot, SessionSnapshotEntry, + SessionSnapshotEntryData, SessionToolAttachment, ToolResultDisposition, TurnResult, + WorkerEvent, WorkerStatus, subscription::{ EventSubscriptionSelector, SubscriptionEvent, SubscriptionEventPayload, SubscriptionFrame, SubscriptionFramePayload, SubscriptionId, SubscriptionRejectionCode, SubscriptionRequest, @@ -78,6 +79,7 @@ pub fn generated_protocol_types() -> String { push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); + push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); diff --git a/crates/session-store/Cargo.toml b/crates/session-store/Cargo.toml index 02f16fdb..a4f1528f 100644 --- a/crates/session-store/Cargo.toml +++ b/crates/session-store/Cargo.toml @@ -10,6 +10,7 @@ base64.workspace = true agen = { workspace = true } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } +sha2.workspace = true uuid = { workspace = true, features = ["v7", "serde"] } thiserror = { workspace = true } protocol = { workspace = true } diff --git a/crates/session-store/src/fs_store.rs b/crates/session-store/src/fs_store.rs index cf0edab0..673bc35b 100644 --- a/crates/session-store/src/fs_store.rs +++ b/crates/session-store/src/fs_store.rs @@ -16,9 +16,11 @@ //! enumerable by the picker. 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::{SegmentId, SessionId}; +use crate::{PasteArtifactLimits, SegmentId, SessionId}; +use protocol::PasteArtifactRef; use std::fs; use std::io::{Read, Seek, SeekFrom, Write}; use std::path::{Path, PathBuf}; @@ -109,6 +111,16 @@ impl FsStore { .join(format!("{segment_id}.trace.jsonl")) } + fn paste_artifact_dir(&self, session_id: SessionId) -> PathBuf { + self.session_dir(session_id).join("artifacts").join("paste") + } + + #[cfg(test)] + fn paste_artifact_path(&self, session_id: SessionId, artifact_id: &str) -> PathBuf { + self.paste_artifact_dir(session_id) + .join(format!("{artifact_id}.json")) + } + fn append_line(&self, path: &Path, line: &str) -> Result<(), StoreError> { let _guard = self .append_lock @@ -350,6 +362,33 @@ impl Store for FsStore { Ok(complete.lines().filter(|l| !l.trim().is_empty()).count()) } + fn write_paste_artifact( + &self, + session_id: SessionId, + source_entry_id: &str, + content: &str, + limits: PasteArtifactLimits, + ) -> Result { + let _guard = self + .append_lock + .lock() + .map_err(|_| std::io::Error::other("session store append lock was poisoned"))?; + write_to_dir( + &self.paste_artifact_dir(session_id), + source_entry_id, + content, + limits, + ) + } + + fn read_paste_artifact( + &self, + session_id: SessionId, + artifact_id: &str, + ) -> Result<(PasteArtifactRef, String), StoreError> { + read_from_dir(&self.paste_artifact_dir(session_id), artifact_id) + } + fn append_trace( &self, session_id: SessionId, @@ -398,4 +437,87 @@ mod tests { store.create_segment(session_id, segment_id, &[]).unwrap(); assert!(store.session_modified_at(session_id).unwrap().is_some()); } + + #[test] + fn paste_artifacts_are_atomic_integrity_checked_and_session_scoped() { + let tmp = tempfile::TempDir::new().unwrap(); + let store = FsStore::new(tmp.path()).unwrap(); + let owner = new_session_id(); + let other = new_session_id(); + let content = "αβγ\nsecond line\n"; + let reference = store + .write_paste_artifact(owner, "entry-1", content, PasteArtifactLimits::default()) + .unwrap(); + + assert_eq!(reference.byte_len, content.len() as u64); + assert_eq!(reference.char_count, content.chars().count() as u64); + assert_eq!(reference.source_entry_id, "entry-1"); + assert_eq!( + store + .read_paste_artifact(owner, &reference.artifact_id) + .unwrap() + .1, + content + ); + assert!(matches!( + store.read_paste_artifact(other, &reference.artifact_id), + Err(StoreError::PasteArtifactNotFound(_)) + )); + assert!( + self::fs::read_dir(store.paste_artifact_dir(owner)) + .unwrap() + .all(|entry| !entry + .unwrap() + .file_name() + .to_string_lossy() + .ends_with(".tmp")) + ); + let very_large = "z".repeat(1024 * 1024); + let very_large_ref = store + .write_paste_artifact( + owner, + "entry-2", + &very_large, + PasteArtifactLimits::default(), + ) + .unwrap(); + assert_eq!( + store + .read_paste_artifact(owner, &very_large_ref.artifact_id) + .unwrap() + .1, + very_large + ); + } + + #[test] + fn paste_artifact_limits_and_corruption_fail_closed() { + let tmp = tempfile::TempDir::new().unwrap(); + let store = FsStore::new(tmp.path()).unwrap(); + let session_id = new_session_id(); + let limits = PasteArtifactLimits { + max_artifact_bytes: 5, + max_session_bytes: 8, + }; + let first = store + .write_paste_artifact(session_id, "entry-1", "1234", limits) + .unwrap(); + assert!(matches!( + store.write_paste_artifact(session_id, "entry-2", "56789", limits), + Err(StoreError::PasteArtifactLimit(_)) + )); + assert!(matches!( + store.write_paste_artifact(session_id, "entry-2", "5678", limits), + Ok(_) + )); + std::fs::write( + store.paste_artifact_path(session_id, &first.artifact_id), + b"{}", + ) + .unwrap(); + assert!(matches!( + store.read_paste_artifact(session_id, &first.artifact_id), + Err(StoreError::Serde(_)) | Err(StoreError::PasteArtifactIntegrity(_)) + )); + } } diff --git a/crates/session-store/src/lib.rs b/crates/session-store/src/lib.rs index 411a9aa0..87b59ae6 100644 --- a/crates/session-store/src/lib.rs +++ b/crates/session-store/src/lib.rs @@ -35,6 +35,7 @@ pub mod fs_store; pub mod history; mod legacy_session_log; pub mod logged_item; +mod paste_artifact; pub mod public_snapshot; pub mod segment; pub mod segment_log; @@ -53,6 +54,7 @@ pub use history::{ LoggedWorkerSubject, }; pub use logged_item::{LoggedContentPart, LoggedItem, LoggedRole, from_logged, to_logged}; +pub use paste_artifact::PasteArtifactLimits; pub use segment::{ SegmentStartState, append_entry, append_system_item, classify_logged_history_entry, create_compacted_segment, create_segment, create_segment_with_ids, ensure_head_or_fork, fork, diff --git a/crates/session-store/src/paste_artifact.rs b/crates/session-store/src/paste_artifact.rs new file mode 100644 index 00000000..1144787c --- /dev/null +++ b/crates/session-store/src/paste_artifact.rs @@ -0,0 +1,156 @@ +//! Session-owned storage for large pasted-input artifacts. + +use std::fs; +use std::io::Write as _; +use std::path::Path; + +use protocol::PasteArtifactRef; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; + +use crate::StoreError; + +/// Bounded storage policy applied before a large paste becomes durable input. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct PasteArtifactLimits { + pub max_artifact_bytes: u64, + pub max_session_bytes: u64, +} + +impl Default for PasteArtifactLimits { + fn default() -> Self { + Self { + max_artifact_bytes: 8 * 1024 * 1024, + max_session_bytes: 64 * 1024 * 1024, + } + } +} + +/// Integrity-bearing on-disk record. The body and metadata are committed in one +/// atomic file replacement so readers never observe a half-written artifact. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct StoredPasteArtifact { + pub reference: PasteArtifactRef, + pub content: String, +} + +pub(crate) fn write_to_dir( + artifact_dir: &Path, + source_entry_id: &str, + content: &str, + limits: PasteArtifactLimits, +) -> Result { + let byte_len = content.len() as u64; + if byte_len > limits.max_artifact_bytes { + return Err(StoreError::PasteArtifactLimit(format!( + "artifact has {byte_len} bytes; maximum is {}", + limits.max_artifact_bytes + ))); + } + fs::create_dir_all(artifact_dir)?; + let mut aggregate = 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)?; + aggregate = aggregate + .checked_add(stored.reference.byte_len) + .ok_or_else(|| { + StoreError::PasteArtifactLimit("session aggregate size overflow".to_string()) + })?; + } + let projected = aggregate.checked_add(byte_len).ok_or_else(|| { + StoreError::PasteArtifactLimit("session aggregate size overflow".to_string()) + })?; + if projected > limits.max_session_bytes { + return Err(StoreError::PasteArtifactLimit(format!( + "session artifacts would use {projected} bytes; maximum is {}", + limits.max_session_bytes + ))); + } + + let artifact_id = uuid::Uuid::now_v7().to_string(); + let reference = PasteArtifactRef { + artifact_id: artifact_id.clone(), + byte_len, + char_count: content.chars().count() as u64, + line_count: line_count(content), + sha256: sha256_hex(content), + source_entry_id: source_entry_id.to_string(), + }; + let bytes = serde_json::to_vec(&StoredPasteArtifact { + reference: reference.clone(), + content: content.to_string(), + })?; + let target = artifact_dir.join(format!("{artifact_id}.json")); + let temporary = artifact_dir.join(format!(".{artifact_id}.tmp")); + let mut file = fs::OpenOptions::new() + .create_new(true) + .write(true) + .open(&temporary)?; + if let Err(error) = file.write_all(&bytes).and_then(|_| file.sync_all()) { + let _ = fs::remove_file(&temporary); + return Err(error.into()); + } + if let Err(error) = fs::rename(&temporary, &target) { + let _ = fs::remove_file(&temporary); + return Err(error.into()); + } + if let Ok(directory) = fs::File::open(artifact_dir) { + directory.sync_all()?; + } + Ok(reference) +} + +pub(crate) fn read_from_dir( + artifact_dir: &Path, + artifact_id: &str, +) -> Result<(PasteArtifactRef, String), StoreError> { + let parsed = uuid::Uuid::parse_str(artifact_id) + .map_err(|_| StoreError::PasteArtifactNotFound(artifact_id.to_string()))?; + if parsed.to_string() != artifact_id { + return Err(StoreError::PasteArtifactNotFound(artifact_id.to_string())); + } + let path = artifact_dir.join(format!("{artifact_id}.json")); + let bytes = match fs::read(path) { + Ok(bytes) => bytes, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => { + return Err(StoreError::PasteArtifactNotFound(artifact_id.to_string())); + } + Err(error) => return Err(error.into()), + }; + let stored: StoredPasteArtifact = serde_json::from_slice(&bytes)?; + verify(&stored, artifact_id)?; + Ok((stored.reference, stored.content)) +} + +fn verify(stored: &StoredPasteArtifact, artifact_id: &str) -> Result<(), StoreError> { + let actual_digest = sha256_hex(&stored.content); + if stored.reference.artifact_id != artifact_id + || stored.reference.byte_len != stored.content.len() as u64 + || stored.reference.char_count != stored.content.chars().count() as u64 + || stored.reference.line_count != line_count(&stored.content) + || stored.reference.sha256 != actual_digest + { + return Err(StoreError::PasteArtifactIntegrity(artifact_id.to_string())); + } + Ok(()) +} + +fn sha256_hex(content: &str) -> String { + Sha256::digest(content.as_bytes()) + .iter() + .map(|byte| format!("{byte:02x}")) + .collect() +} + +fn line_count(content: &str) -> u64 { + if content.is_empty() { + 0 + } else { + content.lines().count().max(1) as u64 + } +} diff --git a/crates/session-store/src/store.rs b/crates/session-store/src/store.rs index 7510d225..142e580d 100644 --- a/crates/session-store/src/store.rs +++ b/crates/session-store/src/store.rs @@ -13,7 +13,8 @@ use crate::event_trace::TraceEntry; use crate::segment_log::LogEntry; -use crate::{SegmentId, SessionId}; +use crate::{PasteArtifactLimits, SegmentId, SessionId}; +use protocol::PasteArtifactRef; /// Errors from the persistence store. #[derive(Debug, thiserror::Error)] @@ -29,6 +30,18 @@ pub enum StoreError { #[error("log corrupted at line {line}: {message}")] Corrupt { line: usize, message: String }, + + #[error("paste artifact storage is unavailable")] + PasteArtifactUnsupported, + + #[error("paste artifact not found: {0}")] + PasteArtifactNotFound(String), + + #[error("paste artifact integrity check failed: {0}")] + PasteArtifactIntegrity(String), + + #[error("paste artifact size limit exceeded: {0}")] + PasteArtifactLimit(String), } /// Sync persistence backend for segment logs. @@ -117,6 +130,26 @@ pub trait Store: Send + Sync { segment_id: SegmentId, ) -> Result; + /// Store a large paste before its reference is committed to history. + fn write_paste_artifact( + &self, + _session_id: SessionId, + _source_entry_id: &str, + _content: &str, + _limits: PasteArtifactLimits, + ) -> Result { + Err(StoreError::PasteArtifactUnsupported) + } + + /// Read and verify one artifact owned by `session_id`. + fn read_paste_artifact( + &self, + _session_id: SessionId, + _artifact_id: &str, + ) -> Result<(PasteArtifactRef, String), StoreError> { + Err(StoreError::PasteArtifactUnsupported) + } + /// Append a trace entry to the debug event trace file. fn append_trace( &self, diff --git a/crates/session-store/src/worker_metadata.rs b/crates/session-store/src/worker_metadata.rs index 77b82f7e..ec63f563 100644 --- a/crates/session-store/src/worker_metadata.rs +++ b/crates/session-store/src/worker_metadata.rs @@ -608,6 +608,24 @@ where ) -> Result { self.session_store.read_entry_count(session_id, segment_id) } + fn write_paste_artifact( + &self, + session_id: SessionId, + source_entry_id: &str, + content: &str, + limits: crate::PasteArtifactLimits, + ) -> Result { + self.session_store + .write_paste_artifact(session_id, source_entry_id, content, limits) + } + fn read_paste_artifact( + &self, + session_id: SessionId, + artifact_id: &str, + ) -> Result<(protocol::PasteArtifactRef, String), crate::StoreError> { + self.session_store + .read_paste_artifact(session_id, artifact_id) + } fn append_trace( &self, session_id: SessionId, diff --git a/crates/session-store/src/worker_session_store.rs b/crates/session-store/src/worker_session_store.rs index dcc12820..ee00d508 100644 --- a/crates/session-store/src/worker_session_store.rs +++ b/crates/session-store/src/worker_session_store.rs @@ -10,9 +10,11 @@ //! every later operation must use that same ID. 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::{SegmentId, SessionId}; +use crate::{PasteArtifactLimits, SegmentId, SessionId}; +use protocol::PasteArtifactRef; use serde::{Deserialize, Serialize}; use std::fs::{self, File, OpenOptions}; use std::io::{Read, Seek, SeekFrom, Write}; @@ -25,6 +27,7 @@ const PREVIOUS_SESSION_SCHEMA_VERSION: u32 = 2; const LEGACY_SESSION_SCHEMA_VERSION: u32 = 1; const SESSION_FILE: &str = "session.json"; const SEGMENTS_DIR: &str = "segments"; +const PASTE_ARTIFACTS_DIR: &str = "artifacts/paste"; #[derive(Clone)] pub struct WorkerSessionStore { @@ -317,6 +320,35 @@ impl Store for WorkerSessionStore { .count()) } + fn write_paste_artifact( + &self, + session_id: SessionId, + source_entry_id: &str, + content: &str, + limits: PasteArtifactLimits, + ) -> Result { + self.ensure_session(session_id, true)?; + let _guard = self + .append_lock + .lock() + .map_err(|_| std::io::Error::other("Worker Session append lock was poisoned"))?; + write_to_dir( + &self.root.join(PASTE_ARTIFACTS_DIR), + source_entry_id, + content, + limits, + ) + } + + fn read_paste_artifact( + &self, + session_id: SessionId, + artifact_id: &str, + ) -> Result<(PasteArtifactRef, String), StoreError> { + self.ensure_session(session_id, false)?; + read_from_dir(&self.root.join(PASTE_ARTIFACTS_DIR), artifact_id) + } + fn append_trace( &self, session_id: SessionId, @@ -601,6 +633,45 @@ mod tests { assert_eq!(store.list_sessions().unwrap(), vec![session_id]); } + #[test] + fn worker_session_store_keeps_paste_artifacts_inside_retention_root() { + let root = tempfile::tempdir().unwrap(); + let store = WorkerSessionStore::new(root.path().join("session")).unwrap(); + let session_id = new_session_id(); + store + .create_segment(session_id, new_segment_id(), &[]) + .unwrap(); + let content = "large paste body\n終端\n"; + let reference = store + .write_paste_artifact( + session_id, + "entry-1", + content, + PasteArtifactLimits::default(), + ) + .unwrap(); + + assert!( + root.path() + .join(format!( + "session/{PASTE_ARTIFACTS_DIR}/{}.json", + reference.artifact_id + )) + .is_file() + ); + assert_eq!( + store + .read_paste_artifact(session_id, &reference.artifact_id) + .unwrap() + .1, + content + ); + assert!(matches!( + store.read_paste_artifact(new_session_id(), &reference.artifact_id), + Err(StoreError::Corrupt { .. }) + )); + } + #[test] fn schema_v1_logs_are_rewritten_and_promoted_to_v3() { let root = tempfile::tempdir().unwrap(); diff --git a/crates/tui/src/input.rs b/crates/tui/src/input.rs index c8d0f8a2..5ed9d894 100644 --- a/crates/tui/src/input.rs +++ b/crates/tui/src/input.rs @@ -61,6 +61,7 @@ impl FlowRefAtom { pub enum Atom { Char(char), Paste(PasteRef), + PasteArtifact(protocol::PasteArtifactRef), FileRef(FileRefAtom), FlowRef(FlowRefAtom), } @@ -72,6 +73,13 @@ impl Atom { match self { Atom::Char(_) => None, Atom::Paste(p) => Some((Style::default().fg(Color::Magenta), p.label())), + Atom::PasteArtifact(artifact) => Some(( + Style::default().fg(Color::Magenta), + format!( + "[Paste artifact {} | {} chars, {} lines]", + artifact.artifact_id, artifact.char_count, artifact.line_count + ), + )), Atom::FileRef(r) => Some((Style::default().fg(Color::Cyan), r.label())), Atom::FlowRef(r) => Some((Style::default().fg(Color::Yellow), r.label())), } @@ -102,7 +110,9 @@ enum WordKind { fn atom_class(atom: &Atom) -> AtomClass { match atom { Atom::Char(c) => char_class(*c), - Atom::Paste(_) | Atom::FileRef(_) | Atom::FlowRef(_) => AtomClass::Chip, + Atom::Paste(_) | Atom::PasteArtifact(_) | Atom::FileRef(_) | Atom::FlowRef(_) => { + AtomClass::Chip + } } } @@ -190,6 +200,9 @@ impl InputBuffer { content: content.clone(), })); } + protocol::Segment::PasteArtifact { artifact } => { + self.atoms.push(Atom::PasteArtifact(artifact.clone())); + } protocol::Segment::FileRef { path } => { self.atoms .push(Atom::FileRef(FileRefAtom { path: path.clone() })); @@ -225,6 +238,13 @@ impl InputBuffer { match atom { Atom::Char(c) => text.push(*c), Atom::Paste(paste) => text.push_str(&paste.content), + Atom::PasteArtifact(artifact) => { + text.push_str(&protocol::Segment::flatten_to_text(&[ + protocol::Segment::PasteArtifact { + artifact: artifact.clone(), + }, + ])) + } Atom::FileRef(file) => text.push_str(&file.path), Atom::FlowRef(flow) => text.push_str(&flow.selector), } @@ -497,6 +517,12 @@ impl InputBuffer { content: p.content.clone(), }); } + Atom::PasteArtifact(artifact) => { + flush_text(&mut buf, &mut out); + out.push(protocol::Segment::PasteArtifact { + artifact: artifact.clone(), + }); + } Atom::FileRef(r) => { flush_text(&mut buf, &mut out); out.push(protocol::Segment::FileRef { @@ -902,6 +928,28 @@ mod submit_segments_tests { } } + #[test] + fn restored_paste_artifact_remains_a_typed_segment() { + let artifact = protocol::PasteArtifactRef { + artifact_id: "019ca7c8-57b6-7f05-8edf-524147aba7b2".to_string(), + byte_len: 65_536, + char_count: 65_530, + line_count: 200, + sha256: "a".repeat(64), + source_entry_id: "entry-1".to_string(), + }; + let original = Segment::PasteArtifact { + artifact: artifact.clone(), + }; + let mut buf = InputBuffer::new(); + buf.replace_with_segments(std::slice::from_ref(&original)); + + assert_eq!( + buf.submit_segments(), + vec![Segment::PasteArtifact { artifact }] + ); + } + #[test] fn empty_buffer_yields_empty_segments() { let buf = InputBuffer::new(); @@ -1219,7 +1267,9 @@ mod word_motion_tests { for a in &buf.atoms { match a { Atom::Char(c) => out.push(*c), - Atom::Paste(_) | Atom::FileRef(_) | Atom::FlowRef(_) => out.push_str("

"), + Atom::Paste(_) | Atom::PasteArtifact(_) | Atom::FileRef(_) | Atom::FlowRef(_) => { + out.push_str("

") + } } } out diff --git a/crates/tui/src/ui.rs b/crates/tui/src/ui.rs index bd5ae01d..df9c5f6b 100644 --- a/crates/tui/src/ui.rs +++ b/crates/tui/src/ui.rs @@ -1296,6 +1296,13 @@ fn chip_span_for(seg: &Segment, fallback: Style) -> (Style, String) { Style::default().fg(Color::Magenta), format!("[Clipboard #{id} | {chars} chars, {line_count} lines]"), ), + Segment::PasteArtifact { artifact } => ( + Style::default().fg(Color::Magenta), + format!( + "[Paste artifact {} | {} chars, {} lines]", + artifact.artifact_id, artifact.char_count, artifact.line_count + ), + ), Segment::FileRef { path } => (Style::default().fg(Color::Cyan), format!("@{path}")), Segment::Flow { selector } => ( Style::default().fg(Color::Yellow), @@ -1314,6 +1321,10 @@ fn segment_display_text(seg: &Segment) -> String { Segment::Paste { id, chars, lines, .. } => format!("[Clipboard #{id} | {chars} chars, {lines} lines]"), + Segment::PasteArtifact { artifact } => format!( + "[Paste artifact {} | {} chars, {} lines]", + artifact.artifact_id, artifact.char_count, artifact.line_count + ), Segment::FileRef { path } => format!("@{path}"), Segment::Flow { selector } => format!("[Flow: {selector}]"), Segment::Unknown => "[unknown segment]".to_owned(), diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index 7d35739e..aaca6aa9 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -898,6 +898,20 @@ where crate::spawn::tool::ParentNotificationTarget::Buffer(worker.notify_buffer_handle()) }); let prompts = worker.prompts().clone(); + let paste_store = worker.store().clone(); + let paste_session_id = worker.session_id(); + worker + .engine_mut() + .register_tool(crate::paste_artifact_tool::search_input_artifact_tool( + paste_store.clone(), + paste_session_id, + )); + worker + .engine_mut() + .register_tool(crate::paste_artifact_tool::read_input_artifact_tool( + paste_store, + paste_session_id, + )); // Resolve the existing Worker–Workdir binding into the domain provider. // Tools only consume the provider handle; they do not own its root, cwd, // scope, or lifecycle. No-workdir Workers expose no local tools. diff --git a/crates/worker/src/lib.rs b/crates/worker/src/lib.rs index 166bdeb3..7db1f37d 100644 --- a/crates/worker/src/lib.rs +++ b/crates/worker/src/lib.rs @@ -9,6 +9,7 @@ pub mod hook; pub(crate) mod in_flight; pub mod ipc; pub mod model_client; +mod paste_artifact_tool; pub mod prompt; pub mod runtime; pub mod runtime_command; diff --git a/crates/worker/src/paste_artifact_tool.rs b/crates/worker/src/paste_artifact_tool.rs new file mode 100644 index 00000000..b61af5d7 --- /dev/null +++ b/crates/worker/src/paste_artifact_tool.rs @@ -0,0 +1,347 @@ +//! Bounded model-facing access to session-owned large paste artifacts. + +use std::sync::Arc; + +use agen::tool::{Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolMeta, ToolOutput}; +use async_trait::async_trait; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; +use session_store::{SessionId, Store, StoreError}; + +const MAX_QUERY_BYTES: usize = 256; +const DEFAULT_SEARCH_RESULTS: usize = 20; +const MAX_SEARCH_RESULTS: usize = 100; +const MAX_SNIPPET_CHARS: usize = 300; +const DEFAULT_READ_BYTES: usize = 8 * 1024; +const MAX_READ_BYTES: usize = 16 * 1024; + +const SEARCH_DESCRIPTION: &str = "Search one large pasted-input artifact owned by the current Worker. Returns bounded matching line snippets; never returns the whole artifact."; +const READ_DESCRIPTION: &str = "Read a bounded UTF-8 byte range from one large pasted-input artifact owned by the current Worker. Use next_offset for repeated calls instead of requesting the whole artifact."; + +#[derive(Clone)] +struct ArtifactAccess { + store: St, + session_id: SessionId, +} + +#[derive(Debug, Deserialize, JsonSchema)] +struct SearchInputArtifactInput { + /// Opaque artifact id from a large-paste history reference. + artifact_id: String, + /// Literal case-sensitive text to find. + query: String, + /// Maximum matching lines to return (1..=100). + max_results: Option, +} + +#[derive(Debug, Serialize)] +struct SearchInputArtifactOutput { + artifact_id: String, + matches: Vec, + truncated: bool, +} + +#[derive(Debug, Serialize)] +struct SearchMatch { + line: u64, + byte_offset: u64, + snippet: String, +} + +struct SearchInputArtifactTool { + access: ArtifactAccess, +} + +#[async_trait] +impl Tool for SearchInputArtifactTool +where + St: Store + Clone + Send + Sync + 'static, +{ + async fn execute( + &self, + input_json: &str, + _context: ToolExecutionContext, + ) -> Result { + let input: SearchInputArtifactInput = + serde_json::from_str(input_json).map_err(|error| { + ToolError::InvalidArgument(format!("invalid SearchInputArtifact input: {error}")) + })?; + if input.query.is_empty() || input.query.len() > MAX_QUERY_BYTES { + return Err(ToolError::InvalidArgument( + "query must contain 1..=256 UTF-8 bytes".to_string(), + )); + } + let max_results = input + .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 mut matches = Vec::new(); + let mut truncated = false; + let mut byte_offset = 0_u64; + for (index, raw_line) in content.split_inclusive('\n').enumerate() { + let line = raw_line.strip_suffix('\n').unwrap_or(raw_line); + if line.contains(&input.query) { + if matches.len() == max_results { + truncated = true; + break; + } + matches.push(SearchMatch { + line: index as u64 + 1, + byte_offset, + snippet: truncate_chars(line, MAX_SNIPPET_CHARS), + }); + } + byte_offset += raw_line.len() as u64; + } + json_output( + format!("Found {} matching pasted-input line(s).", matches.len()), + &SearchInputArtifactOutput { + artifact_id: input.artifact_id, + matches, + truncated, + }, + ) + } +} + +#[derive(Debug, Deserialize, JsonSchema)] +struct ReadInputArtifactInput { + /// Opaque artifact id from a large-paste history reference. + artifact_id: String, + /// UTF-8 byte offset to start reading. Defaults to 0 and must be a character boundary. + offset: Option, + /// Maximum UTF-8 bytes to return (4..=16384). Defaults to 8192. + max_bytes: Option, +} + +#[derive(Debug, Serialize)] +struct ReadInputArtifactOutput { + artifact_id: String, + offset: u64, + content: String, + next_offset: Option, + truncated: bool, +} + +struct ReadInputArtifactTool { + access: ArtifactAccess, +} + +#[async_trait] +impl Tool for ReadInputArtifactTool +where + St: Store + Clone + Send + Sync + 'static, +{ + async fn execute( + &self, + input_json: &str, + _context: ToolExecutionContext, + ) -> Result { + let input: ReadInputArtifactInput = serde_json::from_str(input_json).map_err(|error| { + ToolError::InvalidArgument(format!("invalid ReadInputArtifact input: {error}")) + })?; + let offset = input.offset.unwrap_or(0); + let max_bytes = input + .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 offset = usize::try_from(offset).map_err(|_| { + ToolError::InvalidArgument("offset exceeds the artifact size".to_string()) + })?; + if offset > content.len() || !content.is_char_boundary(offset) { + return Err(ToolError::InvalidArgument( + "offset must be a UTF-8 character boundary within the artifact".to_string(), + )); + } + let mut end = offset.saturating_add(max_bytes).min(content.len()); + while end > offset && !content.is_char_boundary(end) { + end -= 1; + } + let output = content[offset..end].to_string(); + let next_offset = (end < content.len()).then_some(end as u64); + let truncated = next_offset.is_some(); + + json_output( + format!("Read {} pasted-input byte(s).", output.len()), + &ReadInputArtifactOutput { + artifact_id: input.artifact_id, + offset: offset as u64, + content: output, + next_offset, + truncated, + }, + ) + } +} + +pub(crate) fn search_input_artifact_tool(store: St, session_id: SessionId) -> ToolDefinition +where + St: Store + Clone + Send + Sync + 'static, +{ + Arc::new(move || { + let schema = serde_json::to_value(schemars::schema_for!(SearchInputArtifactInput)) + .unwrap_or_else(|_| serde_json::json!({})); + let meta = ToolMeta::new("SearchInputArtifact") + .description(SEARCH_DESCRIPTION) + .input_schema(schema); + let tool: Arc = Arc::new(SearchInputArtifactTool { + access: ArtifactAccess { + store: store.clone(), + session_id, + }, + }); + (meta, tool) + }) +} + +pub(crate) fn read_input_artifact_tool(store: St, session_id: SessionId) -> ToolDefinition +where + St: Store + Clone + Send + Sync + 'static, +{ + Arc::new(move || { + let schema = serde_json::to_value(schemars::schema_for!(ReadInputArtifactInput)) + .unwrap_or_else(|_| serde_json::json!({})); + let meta = ToolMeta::new("ReadInputArtifact") + .description(READ_DESCRIPTION) + .input_schema(schema); + let tool: Arc = Arc::new(ReadInputArtifactTool { + access: ArtifactAccess { + store: store.clone(), + session_id, + }, + }); + (meta, tool) + }) +} + +fn json_output(summary: String, value: &impl Serialize) -> Result { + let content = serde_json::to_string(value) + .map_err(|error| ToolError::ExecutionFailed(error.to_string()))?; + Ok(ToolOutput { + summary, + content: Some(content), + attachments: Vec::new(), + }) +} + +fn tool_store_error(error: StoreError) -> ToolError { + let message = match error { + StoreError::PasteArtifactNotFound(_) => "paste artifact not found", + StoreError::PasteArtifactIntegrity(_) | StoreError::Corrupt { .. } => { + "paste artifact failed its integrity check" + } + StoreError::PasteArtifactUnsupported => "paste artifact storage is unavailable", + _ => "paste artifact is unavailable", + }; + ToolError::ExecutionFailed(message.to_string()) +} + +fn truncate_chars(value: &str, limit: usize) -> String { + let mut chars = value.chars(); + let truncated = chars.by_ref().take(limit).collect::(); + if chars.next().is_some() { + format!("{truncated}…") + } else { + truncated + } +} + +#[cfg(test)] +mod tests { + use agen::tool::ToolExecutionContext; + use session_store::{FsStore, PasteArtifactLimits, Store, new_session_id}; + + use super::*; + + #[tokio::test] + async fn search_and_read_are_bounded_and_owner_scoped() { + let temp = tempfile::TempDir::new().unwrap(); + let store = FsStore::new(temp.path()).unwrap(); + let owner = new_session_id(); + let other = new_session_id(); + let content = (0..700) + .map(|index| format!("line {index}: needle {}", "x".repeat(80))) + .collect::>() + .join("\n"); + let artifact = store + .write_paste_artifact(owner, "entry-1", &content, PasteArtifactLimits::default()) + .unwrap(); + + let search = SearchInputArtifactTool { + access: ArtifactAccess { + store: store.clone(), + session_id: owner, + }, + }; + let searched = search + .execute( + &serde_json::json!({ + "artifact_id": artifact.artifact_id, + "query": "needle", + "max_results": 3 + }) + .to_string(), + ToolExecutionContext::default(), + ) + .await + .unwrap(); + let searched: serde_json::Value = + serde_json::from_str(searched.content.as_deref().unwrap()).unwrap(); + assert_eq!(searched["matches"].as_array().unwrap().len(), 3); + assert_eq!(searched["truncated"], true); + + let read = ReadInputArtifactTool { + access: ArtifactAccess { + store: store.clone(), + session_id: owner, + }, + }; + let read_output = read + .execute( + &serde_json::json!({ + "artifact_id": artifact.artifact_id, + "offset": 2, + "max_bytes": 999999 + }) + .to_string(), + ToolExecutionContext::default(), + ) + .await + .unwrap(); + let read_output: serde_json::Value = + serde_json::from_str(read_output.content.as_deref().unwrap()).unwrap(); + assert!(read_output["content"].as_str().unwrap().len() <= MAX_READ_BYTES); + assert_eq!(read_output["truncated"], true); + assert!(read_output["next_offset"].as_u64().is_some()); + + let foreign = ReadInputArtifactTool { + access: ArtifactAccess { + store, + session_id: other, + }, + }; + let error = foreign + .execute( + &serde_json::json!({ + "artifact_id": artifact.artifact_id, + "offset": 0, + "max_bytes": 1 + }) + .to_string(), + ToolExecutionContext::default(), + ) + .await + .unwrap_err(); + assert!(error.to_string().contains("paste artifact not found")); + } +} diff --git a/crates/worker/src/session_history.rs b/crates/worker/src/session_history.rs index 33ac6eed..42042918 100644 --- a/crates/worker/src/session_history.rs +++ b/crates/worker/src/session_history.rs @@ -62,6 +62,7 @@ pub(crate) fn metadata( } } +#[cfg(test)] pub(crate) fn history_entry( item: Item, origin: WorkerHistoryProvenance, @@ -69,6 +70,21 @@ pub(crate) fn history_entry( HistoryEntry::new(item, metadata(origin, None)) } +pub(crate) fn history_entry_with_id( + item: Item, + entry_id: SessionHistoryEntryId, + origin: WorkerHistoryProvenance, +) -> HistoryEntry { + HistoryEntry::new( + item, + SessionHistoryMetadata { + entry_id, + origin, + derivation: None, + }, + ) +} + pub(crate) fn to_logged_history_entry( entry: &HistoryEntry, ) -> LoggedHistoryEntry { diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index cf2e579e..11548ad8 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -15,8 +15,8 @@ use agen::{ }; use arc_swap::ArcSwap; use session_store::{ - LogEntry, PromptRenderProvenance, SegmentId, SessionExtension, SessionId, Store, StoreError, - SystemItem, segment_log, + LogEntry, PasteArtifactLimits, PromptRenderProvenance, SegmentId, SessionExtension, SessionId, + Store, StoreError, SystemItem, segment_log, }; use session_store::{ WorkerActiveSegmentRef, WorkerMetadata, WorkerMetadataStore, WorkerReclaimedChild, @@ -25,10 +25,12 @@ use session_store::{ use tracing::{info, warn}; use crate::segment_log_sink::SegmentLogSink; +#[cfg(test)] +use crate::session_history::history_entry; use crate::session_history::{ - SessionHistoryDerivation, SessionHistoryMetadata, WorkerHistoryProvenance, history_entry, - metadata as new_history_metadata, restore_history_entries, to_logged_history_entry, - worker_subject, + SessionHistoryDerivation, SessionHistoryEntryId, SessionHistoryMetadata, + WorkerHistoryProvenance, history_entry_with_id, metadata as new_history_metadata, + restore_history_entries, to_logged_history_entry, worker_subject, }; use manifest::{ @@ -58,6 +60,7 @@ use crate::internal_worker::{ }; const COMPACTION_EXTENSION_DOMAIN: &str = "yoi.compaction"; +const LARGE_PASTE_INLINE_MAX_BYTES: usize = 32 * 1024; const WORKER_ORCHESTRATION_INSTRUCTION_ID: &str = "worker.orchestration"; const WORKER_ORCHESTRATION_PROMPT_REF: &str = "common.worker_orchestration"; @@ -2797,7 +2800,15 @@ impl Worker { St: Clone + 'static, F: FnOnce(), { - let (input, pending_flow_state, flow_projection) = self.prepare_flow_input(input)?; + let (mut input, pending_flow_state, flow_projection) = self.prepare_flow_input(input)?; + let projected_entry_ids = if flow_projection.is_some() { + (0..input.len()) + .map(|_| SessionHistoryEntryId::new()) + .collect::>() + } else { + vec![SessionHistoryEntryId::new()] + }; + self.materialize_large_pastes(&mut input, &projected_entry_ids, flow_projection.is_some())?; if let Some(state) = pending_flow_state.as_ref() { let payload = serde_json::to_value(state).map_err(|error| { WorkerError::FlowInput(format!("serialize Flow runtime state: {error}")) @@ -2830,7 +2841,8 @@ impl Worker { trigger: protocol::InvokeKind::UserSend, })?; - let projected_input = self.projected_input_history(&input, flow_projection.as_ref()); + let projected_input = + self.projected_input_history(&input, flow_projection.as_ref(), &projected_entry_ids); // Persist original typed segments together with the exact ordered // model-visible item+origin projection before any entry becomes live. @@ -3064,17 +3076,61 @@ impl Worker { Ok(()) } + fn materialize_large_pastes( + &self, + input: &mut [Segment], + projected_entry_ids: &[SessionHistoryEntryId], + one_entry_per_segment: bool, + ) -> Result<(), WorkerError> { + for (index, segment) in input.iter_mut().enumerate() { + if let Segment::PasteArtifact { artifact } = segment { + let (stored, _) = self + .store + .read_paste_artifact(self.session_id(), &artifact.artifact_id)?; + if &stored != artifact { + return Err(WorkerError::Store(StoreError::PasteArtifactIntegrity( + artifact.artifact_id.clone(), + ))); + } + continue; + } + let Segment::Paste { content, .. } = segment else { + continue; + }; + if content.len() <= LARGE_PASTE_INLINE_MAX_BYTES { + continue; + } + 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 paste") + .0 + .as_str(); + let artifact = self.store.write_paste_artifact( + self.session_id(), + source_entry_id, + content, + PasteArtifactLimits::default(), + )?; + *segment = Segment::PasteArtifact { artifact }; + } + Ok(()) + } + fn projected_input_history( &self, input: &[Segment], flow_projection: Option<&PreparedFlowProjection>, + entry_ids: &[SessionHistoryEntryId], ) -> Vec> { if let Some(flow) = flow_projection { return input .iter() - .map(|segment| match segment { - Segment::Flow { .. } => history_entry( + .zip(entry_ids) + .map(|(segment, entry_id)| match segment { + Segment::Flow { .. } => history_entry_with_id( Item::user_message(flow.instructions.clone()), + entry_id.clone(), WorkerHistoryProvenance::FlowInstruction { selector: flow.selector.clone(), definition_id: flow.definition_id.clone(), @@ -3083,8 +3139,9 @@ impl Worker { state_id: flow.state_id.clone(), }, ), - other => history_entry( + other => history_entry_with_id( Item::user_message(Segment::flatten_to_text(std::slice::from_ref(other))), + entry_id.clone(), // Current public submit transport does not carry a // trusted account/Worker subject envelope. Fail closed // instead of promoting role=user to HumanInput. @@ -3094,8 +3151,12 @@ impl Worker { .collect(); } - vec![history_entry( + vec![history_entry_with_id( Item::user_message(Segment::flatten_to_text(input)), + entry_ids + .first() + .expect("projected Worker input always has one entry id") + .clone(), WorkerHistoryProvenance::LegacyUnknown, )] } @@ -6405,6 +6466,11 @@ fn preview_segments(segments: &[Segment]) -> String { match segment { Segment::Text { content } => preview.push_str(content.trim()), Segment::Paste { content, .. } => preview.push_str(content.trim()), + Segment::PasteArtifact { artifact } => { + preview.push_str("[Large paste artifact: "); + preview.push_str(&artifact.artifact_id); + preview.push(']'); + } Segment::FileRef { path } => { preview.push('@'); preview.push_str(path); @@ -7909,7 +7975,9 @@ mod build_summary_prompt_tests { FLOW_RUNTIME_EXTENSION_DOMAIN, serde_json::to_value(&state).unwrap(), ); - let projected = worker.projected_input_history(&segments, projection.as_ref()); + let projected_ids = vec![SessionHistoryEntryId::new(), SessionHistoryEntryId::new()]; + let projected = + worker.projected_input_history(&segments, projection.as_ref(), &projected_ids); worker .commit_entry(LogEntry::AnnotatedUserInput { ts: segment_log::now_millis(), @@ -7981,6 +8049,144 @@ mod build_summary_prompt_tests { ); } + #[tokio::test] + async fn large_paste_is_stored_before_compact_history_is_committed() { + let (_dir, worker) = rewind_test_worker().await; + let exact = "x".repeat(LARGE_PASTE_INLINE_MAX_BYTES); + let exact_ids = vec![SessionHistoryEntryId::new()]; + let mut exact_input = vec![Segment::Paste { + id: 1, + chars: exact.len() as u32, + lines: 1, + content: exact.clone(), + }]; + worker + .materialize_large_pastes(&mut exact_input, &exact_ids, false) + .unwrap(); + assert!(matches!(&exact_input[0], Segment::Paste { content, .. } if content == &exact)); + let mut empty_input = vec![Segment::Paste { + id: 0, + chars: 0, + lines: 0, + content: String::new(), + }]; + worker + .materialize_large_pastes(&mut empty_input, &exact_ids, false) + .unwrap(); + assert!(matches!( + &empty_input[0], + Segment::Paste { content, .. } if content.is_empty() + )); + + let body = format!("{}\n終端\n", "多".repeat(12_000)); + let entry_id = SessionHistoryEntryId::new(); + let mut input = vec![Segment::Paste { + id: 2, + chars: body.chars().count() as u32, + lines: 3, + content: body.clone(), + }]; + worker + .materialize_large_pastes(&mut input, std::slice::from_ref(&entry_id), false) + .unwrap(); + let artifact = match &input[0] { + Segment::PasteArtifact { artifact } => artifact.clone(), + other => panic!("expected stored paste reference, got {other:?}"), + }; + assert_eq!(artifact.source_entry_id, entry_id.0); + assert_eq!(artifact.byte_len, body.len() as u64); + assert_eq!(artifact.char_count, body.chars().count() as u64); + assert_eq!( + worker + .store + .read_paste_artifact(worker.session_id(), &artifact.artifact_id) + .unwrap() + .1, + body + ); + worker + .materialize_large_pastes(&mut input, &[SessionHistoryEntryId::new()], false) + .unwrap(); + assert!(matches!( + &input[0], + Segment::PasteArtifact { artifact: retained } + if retained.source_entry_id == artifact.source_entry_id + )); + + let history = worker.projected_input_history(&input, None, &[entry_id]); + assert!(!history[0].item.as_text().unwrap().contains("終端")); + append_test_entry( + &worker, + LogEntry::Invoke { + ts: segment_log::now_millis(), + trigger: protocol::InvokeKind::UserSend, + }, + ); + worker + .commit_entry(LogEntry::AnnotatedUserInput { + ts: segment_log::now_millis(), + segments: input.clone(), + extensions: Vec::new(), + history: history.iter().map(to_logged_history_entry).collect(), + }) + .unwrap(); + let location = worker.segment_state.location(); + let entries = worker + .store + .read_all(location.session_id, location.segment_id) + .unwrap(); + let persisted = serde_json::to_string(&entries).unwrap(); + assert!(!persisted.contains("終端")); + let state = session_store::collect_state(&entries); + assert!(matches!( + &state.user_segments[0][0], + Segment::PasteArtifact { artifact: restored } + if restored.artifact_id == artifact.artifact_id + )); + } + + #[tokio::test] + async fn large_paste_storage_failure_commits_no_input() { + let temp = tempfile::TempDir::new().unwrap(); + let store = session_store::FsStore::new(temp.path()).unwrap(); + let mut worker = Worker::new( + minimal_manifest(), + Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(NoopClient), + store.clone(), + WorkerWorkspaceContext::unavailable(None, "test unavailable"), + WorkerFilesystemAuthority::None, + Scope::empty(), + ) + .await + .unwrap(); + worker.ensure_segment_head().unwrap(); + std::fs::write( + temp.path() + .join(worker.session_id().to_string()) + .join("artifacts"), + "block artifact directory creation", + ) + .unwrap(); + let result = worker + .run(vec![Segment::Paste { + id: 1, + chars: (LARGE_PASTE_INLINE_MAX_BYTES + 1) as u32, + lines: 1, + content: "x".repeat(LARGE_PASTE_INLINE_MAX_BYTES + 1), + }]) + .await; + assert!(matches!(result, Err(WorkerError::Store(StoreError::Io(_))))); + let location = worker.segment_state.location(); + let entries = store + .read_all(location.session_id, location.segment_id) + .unwrap(); + assert!( + !entries + .iter() + .any(|entry| matches!(entry, LogEntry::AnnotatedUserInput { .. })) + ); + } + async fn rewind_test_worker() -> ( tempfile::TempDir, Worker, diff --git a/crates/worker/tests/controller_test.rs b/crates/worker/tests/controller_test.rs index 7b13e752..0a1b8e31 100644 --- a/crates/worker/tests/controller_test.rs +++ b/crates/worker/tests/controller_test.rs @@ -613,7 +613,19 @@ async fn feature_flags_default_to_core_tool_surface_only() { let request = wait_for_captured_request(&client_for_assert).await; let names = request_tool_names(&request); - assert_eq!(names, vec!["Bash", "Edit", "Glob", "Grep", "Read", "Write"]); + assert_eq!( + names, + vec![ + "Bash", + "Edit", + "Glob", + "Grep", + "Read", + "ReadInputArtifact", + "SearchInputArtifact", + "Write", + ] + ); assert!(!names.iter().any(|name| name == "TaskCreate")); assert!(!names.iter().any(|name| name == "WebSearch")); assert!(!names.iter().any(|name| name == "SubWorkerSpawn")); diff --git a/web/workspace/src/lib/generated/protocol.ts b/web/workspace/src/lib/generated/protocol.ts index d25d2eb3..7c7530e0 100644 --- a/web/workspace/src/lib/generated/protocol.ts +++ b/web/workspace/src/lib/generated/protocol.ts @@ -133,7 +133,9 @@ message: string, */ timestamp_ms: number, }; -export type Segment = { "kind": "text", content: string, } | { "kind": "paste", id: number, chars: number, lines: number, content: string, } | { "kind": "file_ref", path: string, } | { "kind": "flow", selector: string, } | { "kind": "unknown" }; +export type PasteArtifactRef = { artifact_id: string, 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 WorkerEvent = { "kind": "turn_ended", worker_name: string, } | { "kind": "errored", worker_name: string, message: string, } | { "kind": "shut_down", worker_name: string, } | { "kind": "scope_sub_delegated", /** diff --git a/web/workspace/src/lib/workspace/console/model.test.ts b/web/workspace/src/lib/workspace/console/model.test.ts index 94c2c018..e9579ca0 100644 --- a/web/workspace/src/lib/workspace/console/model.test.ts +++ b/web/workspace/src/lib/workspace/console/model.test.ts @@ -137,6 +137,24 @@ function snapshotEvent(cwd: string, entries: unknown[] = []): Event { }; } +Deno.test("large paste segments project compact artifact metadata", () => { + const body = "secret pasted body"; + const text = segmentsToText([{ + kind: "paste_artifact", + artifact: { + artifact_id: "019ca7c8-57b6-7f05-8edf-524147aba7b2", + byte_len: 65536, + char_count: 65530, + line_count: 200, + sha256: "a".repeat(64), + source_entry_id: "entry-1", + }, + }]); + assert(text.includes("019ca7c8-57b6-7f05-8edf-524147aba7b2"), "artifact id is visible"); + assert(text.includes("65536 bytes"), "bounded size metadata is visible"); + assert(!text.includes(body), "artifact body is not projected"); +}); + Deno.test("console routing projects live errors but not completion replies", () => { const errorEvent = { event: "error", diff --git a/web/workspace/src/lib/workspace/console/model.ts b/web/workspace/src/lib/workspace/console/model.ts index 7a965cda..597f9133 100644 --- a/web/workspace/src/lib/workspace/console/model.ts +++ b/web/workspace/src/lib/workspace/console/model.ts @@ -1069,6 +1069,8 @@ export function segmentsToText(segments: Segment[]): string { case "paste": return segment.content || `[paste ${segment.id}: ${segment.chars} chars / ${segment.lines} lines]`; + case "paste_artifact": + return `[Large paste artifact ${segment.artifact.artifact_id}: ${segment.artifact.byte_len} bytes, sha256 ${segment.artifact.sha256}]`; case "file_ref": return `@file ${segment.path}`; case "unknown":