diff --git a/Cargo.lock b/Cargo.lock index 4c6a67e7..6edb2f68 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4412,6 +4412,7 @@ dependencies = [ "thiserror 2.0.18", "tokio", "tracing", + "unicode-normalization", "uuid", ] diff --git a/crates/session-store/Cargo.toml b/crates/session-store/Cargo.toml index 1811c8e3..3a75f5a9 100644 --- a/crates/session-store/Cargo.toml +++ b/crates/session-store/Cargo.toml @@ -16,6 +16,7 @@ uuid = { workspace = true, features = ["v7", "serde"] } thiserror = { workspace = true } protocol = { workspace = true } tracing.workspace = true +unicode-normalization = "0.1.25" [dev-dependencies] async-trait = { workspace = true } diff --git a/crates/session-store/src/fs_store.rs b/crates/session-store/src/fs_store.rs index 4bae850f..9efccd0a 100644 --- a/crates/session-store/src/fs_store.rs +++ b/crates/session-store/src/fs_store.rs @@ -20,8 +20,8 @@ use crate::paste_artifact::{read_from_dir, write_to_dir}; use crate::segment_log::LogEntry; use crate::store::{Store, StoreError}; use crate::uploaded_file::{ - bind_uploaded_file, delete_uploaded_file, read_uploaded_file, read_uploaded_file_by_id, - write_uploaded_file, + bind_uploaded_file, delete_uncommitted_uploaded_files, delete_uploaded_file, + read_uploaded_file, read_uploaded_file_by_id, write_uploaded_file, }; use crate::{PasteArtifactLimits, SegmentId, SessionId, UploadedFileLimits}; use protocol::{PasteArtifactRef, UploadedFileRef}; @@ -459,6 +459,14 @@ impl Store for FsStore { delete_uploaded_file(&self.paste_artifact_dir(session_id), artifact_id) } + fn delete_uncommitted_uploaded_files(&self, session_id: SessionId) -> Result { + let _guard = self + .append_lock + .lock() + .map_err(|_| std::io::Error::other("session store append lock was poisoned"))?; + delete_uncommitted_uploaded_files(&self.paste_artifact_dir(session_id)) + } + fn append_trace( &self, session_id: SessionId, @@ -678,6 +686,47 @@ mod tests { store.write_uploaded_file(session_id, "notes.txt", "not a type", b"x", limits), Err(StoreError::InvalidUploadedFileMediaType) )); + assert!(matches!( + store.write_uploaded_file( + session_id, + "safe\u{202e}txt.exe", + "text/plain", + b"x", + limits + ), + Err(StoreError::InvalidUploadedFileName) + )); + assert!(matches!( + store.write_uploaded_file(session_id, "image.png", "image/png", b"not a png", limits), + Err(StoreError::ArtifactIntegrityMismatch) + )); + let pending = store + .write_uploaded_file(session_id, "Readme.txt", "text/plain", b"x", limits) + .unwrap(); + let replay = store + .write_uploaded_file(session_id, "Readme.txt", "text/plain", b"x", limits) + .unwrap(); + assert_eq!(replay.artifact_id, pending.artifact_id); + assert!(matches!( + store.write_uploaded_file(session_id, "README.txt", "text/plain", b"changed", limits), + Err(StoreError::InvalidUploadedFileName) + )); + assert!(matches!( + store.write_uploaded_file(session_id, "ＲＥＡＤＭＥ.txt", "text/plain", b"y", limits), + Err(StoreError::InvalidUploadedFileName) + )); + let bound = store + .bind_uploaded_file(session_id, &pending, "entry-upload") + .unwrap(); + let other = store + .write_uploaded_file(session_id, "other.txt", "text/plain", b"z", limits) + .unwrap(); + assert_eq!( + store.delete_uncommitted_uploaded_files(session_id).unwrap(), + 1 + ); + assert!(store.read_uploaded_file(session_id, &other).is_err()); + assert_eq!(store.read_uploaded_file(session_id, &bound).unwrap(), b"x"); store .write_paste_artifact( session_id, diff --git a/crates/session-store/src/lib.rs b/crates/session-store/src/lib.rs index d2ea196f..1d4a5f31 100644 --- a/crates/session-store/src/lib.rs +++ b/crates/session-store/src/lib.rs @@ -69,7 +69,7 @@ pub use system_item::{ }; pub use uploaded_file::{ DEFAULT_MAX_FILES_PER_SUBMISSION, DEFAULT_MAX_SESSION_ARTIFACT_BYTES, - DEFAULT_MAX_UPLOADED_FILE_BYTES, UploadedFileLimits, + DEFAULT_MAX_SESSION_UPLOADED_FILES, DEFAULT_MAX_UPLOADED_FILE_BYTES, UploadedFileLimits, }; pub use worker_metadata::{ CombinedStore, FsWorkerStore, WorkerActiveSegmentRef, WorkerAggregateStore, WorkerMetadata, diff --git a/crates/session-store/src/public_snapshot.rs b/crates/session-store/src/public_snapshot.rs index 3fdb32e1..f3b61754 100644 --- a/crates/session-store/src/public_snapshot.rs +++ b/crates/session-store/src/public_snapshot.rs @@ -41,6 +41,24 @@ pub fn project_session_snapshot(session_id: SessionId, log: &[LogEntry]) -> Sess entries.clear(); extend_history(&mut entries, history, None, *ts); } + LogEntry::InputSegmentsCheckpoint { user_segments, .. } => { + let mut segments = user_segments.iter(); + for entry in &mut entries { + let is_user = matches!( + &entry.data, + SessionSnapshotEntryData::UserInput { .. } + | SessionSnapshotEntryData::Message { + role: SessionMessageRole::User, + .. + } + ); + if is_user && let Some(checkpoint) = segments.next() { + entry.data = SessionSnapshotEntryData::UserInput { + segments: checkpoint.clone(), + }; + } + } + } LogEntry::AnnotatedUserInput { ts, segments, @@ -357,6 +375,63 @@ mod tests { assert!(json.contains("visible")); } + #[test] + fn compacted_checkpoint_restores_uploaded_file_segments() { + let session_id = crate::new_session_id(); + let user_entry_id = LoggedSessionHistoryEntryId::new(); + let file = protocol::UploadedFileRef { + artifact_id: "019ca7c8-57b6-7f05-8edf-524147aba7b3".into(), + file_name: "notes.md".into(), + media_type: "text/markdown".into(), + created_at_ms: 7, + availability: protocol::UploadedFileAvailability::Available, + byte_len: 12, + sha256: "a".repeat(64), + source_entry_id: Some(user_entry_id.0.clone()), + }; + let segment = Segment::UploadedFile { file }; + let log = vec![ + LogEntry::AnnotatedSegmentStart { + ts: 10, + session_id, + system_prompt: None, + config: RequestConfig::default(), + history: vec![LoggedHistoryEntry { + item: LoggedItem::Message { + role: LoggedRole::User, + content: vec![LoggedContentPart::Text { + text: "[Attached file: notes.md]".into(), + }], + }, + metadata: LoggedSessionHistoryMetadata { + entry_id: user_entry_id, + origin: LoggedSessionHistoryOrigin::HumanInput { + account_id: "account-1".into(), + }, + derivation: None, + }, + }], + forked_from: None, + compacted_from: Some(crate::SegmentOrigin { + segment_id: crate::new_segment_id(), + at_turn_index: 1, + }), + }, + LogEntry::InputSegmentsCheckpoint { + ts: 10, + user_segments: vec![vec![segment.clone()]], + }, + ]; + + let snapshot = project_current_session_snapshot(&log); + assert_eq!( + snapshot.entries[0].data, + SessionSnapshotEntryData::UserInput { + segments: vec![segment] + } + ); + } + #[test] fn annotated_user_input_attaches_segments_to_first_user_role_entry_for_any_origin() { let session_id = crate::new_session_id(); diff --git a/crates/session-store/src/segment_log.rs b/crates/session-store/src/segment_log.rs index 73e250ea..9daf85c2 100644 --- a/crates/session-store/src/segment_log.rs +++ b/crates/session-store/src/segment_log.rs @@ -63,6 +63,14 @@ pub enum LogEntry { compacted_from: Option, }, + /// Typed user-segment projection accompanying a compacted or forked + /// SegmentStart history snapshot. This keeps attachment identity and + /// metadata aligned with retained user entries without embedding bodies. + InputSegmentsCheckpoint { + ts: u64, + user_segments: Vec>, + }, + /// IDLE → active marker. Records the start of a new self-driving /// cycle (Invoke range). The range extends implicitly until the /// next `Invoke` entry; this entry carries the trigger only — the @@ -273,6 +281,9 @@ pub fn collect_state(entries: &[LogEntry]) -> RestoredState { .map(|entry| Item::from(entry.item)) .collect(); } + LogEntry::InputSegmentsCheckpoint { user_segments, .. } => { + state.user_segments = user_segments.clone(); + } LogEntry::Invoke { .. } => { // A terminal run record below clears or refines this. If the // log ends first, restore must treat the turn as interrupted. diff --git a/crates/session-store/src/store.rs b/crates/session-store/src/store.rs index f498ac71..757a6b15 100644 --- a/crates/session-store/src/store.rs +++ b/crates/session-store/src/store.rs @@ -221,6 +221,10 @@ pub trait Store: Send + Sync { Err(StoreError::PasteArtifactUnsupported) } + fn delete_uncommitted_uploaded_files(&self, _session_id: SessionId) -> Result { + Ok(0) + } + /// Append a trace entry to the debug event trace file. fn append_trace( &self, diff --git a/crates/session-store/src/uploaded_file.rs b/crates/session-store/src/uploaded_file.rs index 456a6973..05af45a9 100644 --- a/crates/session-store/src/uploaded_file.rs +++ b/crates/session-store/src/uploaded_file.rs @@ -9,6 +9,7 @@ use fs4::fs_std::FileExt; use protocol::{UploadedFileAvailability, UploadedFileRef}; use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; +use unicode_normalization::UnicodeNormalization; use uuid::Uuid; use crate::StoreError; @@ -18,6 +19,7 @@ type Result = std::result::Result; 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; +pub const DEFAULT_MAX_SESSION_UPLOADED_FILES: u64 = 256; const MAX_FILE_NAME_CHARS: usize = 255; const MAX_MEDIA_TYPE_BYTES: usize = 127; @@ -53,9 +55,15 @@ pub(crate) fn validate_file_name(file_name: &str) -> Result<()> { || file_name.chars().count() > MAX_FILE_NAME_CHARS || file_name == "." || file_name == ".." - || file_name - .chars() - .any(|ch| ch.is_control() || matches!(ch, '/' | '\\')) + || file_name.chars().any(|ch| { + ch.is_control() + || matches!( + ch, + '/' | '\\' | '\u{200b}' | '\u{200c}' | '\u{200d}' | '\u{2060}' | '\u{feff}' + ) + || ('\u{202a}'..='\u{202e}').contains(&ch) + || ('\u{2066}'..='\u{2069}').contains(&ch) + }) { return Err(StoreError::InvalidUploadedFileName); } @@ -96,6 +104,35 @@ pub(crate) fn validate_media_type(media_type: &str) -> Result<()> { Ok(()) } +fn normalized_file_name(file_name: &str) -> String { + file_name.nfkc().flat_map(char::to_lowercase).collect() +} + +fn validate_content(media_type: &str, content: &[u8]) -> Result<()> { + if content.is_empty() { + return Err(StoreError::InvalidUploadedFileMediaType); + } + let matches_declared_type = if media_type.starts_with("text/") { + std::str::from_utf8(content).is_ok() + } else { + match media_type { + "application/json" => serde_json::from_slice::(content).is_ok(), + "application/pdf" => content.starts_with(b"%PDF-"), + "image/png" => content.starts_with(b"\x89PNG\r\n\x1a\n"), + "image/jpeg" => content.starts_with(&[0xff, 0xd8, 0xff]), + "image/gif" => content.starts_with(b"GIF87a") || content.starts_with(b"GIF89a"), + "image/webp" => { + content.len() >= 12 && content.starts_with(b"RIFF") && &content[8..12] == b"WEBP" + } + _ => false, + } + }; + if !matches_declared_type { + return Err(StoreError::ArtifactIntegrityMismatch); + } + Ok(()) +} + fn record_path(dir: &Path, artifact_id: &str) -> Result { let id = Uuid::parse_str(artifact_id).map_err(|_| StoreError::InvalidArtifactId)?; Ok(dir.join(format!("{id}.file.json"))) @@ -133,7 +170,7 @@ pub(crate) fn stored_uploaded_file_usage(dir: &Path) -> Result<(u64, u64)> { { continue; } - let stored: StoredUploadedFile = serde_json::from_slice(&fs::read(path)?)?; + let stored: StoredUploadedFile = serde_json::from_slice(&fs::read(&path)?)?; bytes = bytes .checked_add(stored.byte_len) .ok_or(StoreError::ArtifactQuotaExceeded)?; @@ -153,7 +190,9 @@ pub(crate) fn write_uploaded_file( ) -> Result { validate_file_name(file_name)?; validate_media_type(media_type)?; + validate_content(media_type, content)?; let byte_len = u64::try_from(content.len()).map_err(|_| StoreError::ArtifactTooLarge)?; + let sha256 = digest(content); if byte_len > limits.max_file_bytes { return Err(StoreError::ArtifactTooLarge); } @@ -166,7 +205,48 @@ pub(crate) fn write_uploaded_file( .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)?; + let (file_bytes, file_count) = stored_uploaded_file_usage(dir)?; + if file_count >= DEFAULT_MAX_SESSION_UPLOADED_FILES { + return Err(StoreError::ArtifactQuotaExceeded); + } + let normalized_name = normalized_file_name(file_name); + for entry in fs::read_dir(dir)? { + let path = entry?.path(); + if !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)?)?; + if stored.source_entry_id.is_none() + && normalized_file_name(&stored.file_name) == normalized_name + { + if stored.media_type == media_type + && stored.byte_len == byte_len + && stored.sha256 == sha256 + { + let artifact_id = path + .file_name() + .and_then(|name| name.to_str()) + .and_then(|name| name.strip_suffix(".file.json")) + .ok_or(StoreError::InvalidArtifactId)? + .to_string(); + return Ok(UploadedFileRef { + artifact_id, + 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: None, + }); + } + return Err(StoreError::InvalidUploadedFileName); + } + } if paste_bytes .checked_add(file_bytes) .and_then(|total| total.checked_add(byte_len)) @@ -177,7 +257,6 @@ pub(crate) fn write_uploaded_file( 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(), @@ -274,6 +353,36 @@ pub(crate) fn bind_uploaded_file( Ok(bound) } +pub(crate) fn delete_uncommitted_uploaded_files(dir: &Path) -> Result { + fs::create_dir_all(dir)?; + let aggregate_lock = fs::OpenOptions::new() + .create(true) + .read(true) + .write(true) + .open(dir.join(".aggregate.lock"))?; + FileExt::lock_exclusive(&aggregate_lock)?; + let mut removed = 0_u64; + for entry in fs::read_dir(dir)? { + let entry = entry?; + let path = entry.path(); + if !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)?)?; + if stored.source_entry_id.is_none() { + fs::remove_file(path)?; + removed = removed + .checked_add(1) + .ok_or(StoreError::ArtifactQuotaExceeded)?; + } + } + Ok(removed) +} + pub(crate) fn delete_uploaded_file(dir: &Path, artifact_id: &str) -> Result { fs::create_dir_all(dir)?; let aggregate_lock = fs::OpenOptions::new() diff --git a/crates/tui/src/app.rs b/crates/tui/src/app.rs index fd69315f..d3b2396d 100644 --- a/crates/tui/src/app.rs +++ b/crates/tui/src/app.rs @@ -928,14 +928,6 @@ impl App { Some(self.method_for_run(queued.segments)) } - pub fn push_notice(&mut self, message: impl Into) { - self.blocks.push(Block::Alert { - level: AlertLevel::Warn, - source: AlertSource::Worker, - message: message.into(), - }); - } - pub fn push_error(&mut self, message: impl Into) { self.blocks.push(Block::Alert { level: AlertLevel::Error, diff --git a/crates/tui/src/console/mod.rs b/crates/tui/src/console/mod.rs index 7f4e393b..79355685 100644 --- a/crates/tui/src/console/mod.rs +++ b/crates/tui/src/console/mod.rs @@ -792,17 +792,27 @@ async fn handle_terminal_event( if let Some(method) = handle_key(app, key) { if let Some(path) = attachment_command_path(&method) { match client.upload_path(&path).await { - Ok(reference) => app.push_notice(format!( - "Attached {} ({} bytes); it will be sent with the next message.", - reference.file_name, reference.byte_len - )), + Ok(reference) => app.flash_actionbar_notice( + format!( + "Attached {} ({} bytes); it will be sent with the next message.", + reference.file_name, reference.byte_len + ), + ActionbarNoticeLevel::Info, + ActionbarNoticeSource::Tui, + Duration::from_secs(6), + ), Err(error) => { app.push_error(format!("Attachment upload failed: {error}")); } } } else if is_clear_attachments_command(&method) { client.clear_pending_attachments().await; - app.push_notice("Removed pending attachments."); + app.flash_actionbar_notice( + "Removed pending attachments.", + ActionbarNoticeLevel::Info, + ActionbarNoticeSource::Tui, + Duration::from_secs(4), + ); } else { client.send(&method).await?; } diff --git a/crates/worker-runtime/src/worker_backend.rs b/crates/worker-runtime/src/worker_backend.rs index 3fd8a56b..23275a29 100644 --- a/crates/worker-runtime/src/worker_backend.rs +++ b/crates/worker-runtime/src/worker_backend.rs @@ -2145,6 +2145,25 @@ where ), ); } + let cleanup_handle = match self.workers.lock() { + Ok(workers) => workers + .get(handle.worker_ref()) + .map(|execution| execution.handle.clone()), + Err(_) => { + return WorkerExecutionResult::errored( + WorkerExecutionOperation::Stop, + "worker adapter registry lock is poisoned", + ); + } + }; + if let Some(worker) = cleanup_handle + && let Err(error) = worker.delete_uncommitted_uploaded_files() + { + return WorkerExecutionResult::errored( + WorkerExecutionOperation::Stop, + format!("uploaded_file_cleanup_failed: {error}"), + ); + } let execution = match self.workers.lock() { Ok(mut workers) => workers.remove(handle.worker_ref()), Err(_) => { diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index 5d0f98c9..cb2de8fd 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -88,6 +88,11 @@ impl WorkerHandle { .delete_uploaded_file(self.session_id, artifact_id) } + pub fn delete_uncommitted_uploaded_files(&self) -> Result { + self.artifact_store + .delete_uncommitted_uploaded_files(self.session_id) + } + pub fn subscribe(&self) -> broadcast::Receiver { self.working_event_tx.subscribe() } diff --git a/crates/worker/src/paste_artifact_tool.rs b/crates/worker/src/paste_artifact_tool.rs index 9747c074..9c1598e2 100644 --- a/crates/worker/src/paste_artifact_tool.rs +++ b/crates/worker/src/paste_artifact_tool.rs @@ -340,7 +340,7 @@ mod tests { owner, "image.png", "image/png", - &[0xff, 0xd8, 0x00], + b"\x89PNG\r\n\x1a\nbody", session_store::UploadedFileLimits::default(), ) .unwrap(); diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 302a26b2..62568ac0 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -3331,7 +3331,13 @@ impl Worker { }), compacted_from: None, }; - let mut initial_entries = vec![entry.clone()]; + let mut initial_entries = vec![ + entry.clone(), + LogEntry::InputSegmentsCheckpoint { + ts: segment_log::now_millis(), + user_segments: self.user_segments.clone(), + }, + ]; if let Some(checkpoint) = active_run_checkpoint_entry(w.active_run_turn_count(), w.turn_count()) { @@ -4295,6 +4301,13 @@ impl Worker { }) .collect::>(); + let retained_user_segments = self + .user_segments + .iter() + .skip(self.user_segments.len().saturating_sub(retained_user_msgs)) + .cloned() + .collect::>(); + // Build the SegmentStart entry for the new compacted segment. // Inherits the source Segment's session_id so the compacted // lineage stays grouped under the same Session. Atomically @@ -4320,7 +4333,13 @@ impl Worker { at_turn_index: source_turn_count, }), }; - let mut initial_entries = vec![entry.clone()]; + let mut initial_entries = vec![ + entry.clone(), + LogEntry::InputSegmentsCheckpoint { + ts: segment_log::now_millis(), + user_segments: retained_user_segments.clone(), + }, + ]; if let Some(checkpoint) = active_run_checkpoint_entry(w.active_run_turn_count(), source_turn_count) { @@ -4372,10 +4391,7 @@ impl Worker { // segments; only the user_messages surviving in retained_items // keep them. They are always the trailing K entries of // `self.user_segments` because submissions are appended in order. - let drop_n = self.user_segments.len().saturating_sub(retained_user_msgs); - if drop_n > 0 { - self.user_segments.drain(..drop_n); - } + self.user_segments = retained_user_segments; self.session.replace_history(compacted_history_entries); // Compaction-introduced system messages are part of the new