diff --git a/Cargo.lock b/Cargo.lock index 6edb2f68..dda192e6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4413,6 +4413,8 @@ dependencies = [ "tokio", "tracing", "unicode-normalization", + "unicode-properties", + "unicode-security", "uuid", ] @@ -5327,6 +5329,7 @@ name = "tui" version = "0.1.0" dependencies = [ "agen", + "async-trait", "base64 0.22.1", "client", "crossterm 0.28.1", @@ -5437,6 +5440,22 @@ version = "0.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7df058c713841ad818f1dc5d3fd88063241cc61f49f5fbea4b951e8cf5a8d71d" +[[package]] +name = "unicode-script" +version = "0.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "383ad40bb927465ec0ce7720e033cb4ca06912855fc35db31b5755d0de75b1ee" + +[[package]] +name = "unicode-security" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e4ddba1535dd35ed8b61c52166b7155d7f4e4b8847cec6f48e71dc66d8b5e50" +dependencies = [ + "unicode-normalization", + "unicode-script", +] + [[package]] name = "unicode-segmentation" version = "1.13.2" diff --git a/crates/session-store/Cargo.toml b/crates/session-store/Cargo.toml index 3a75f5a9..5dd8b688 100644 --- a/crates/session-store/Cargo.toml +++ b/crates/session-store/Cargo.toml @@ -17,6 +17,8 @@ thiserror = { workspace = true } protocol = { workspace = true } tracing.workspace = true unicode-normalization = "0.1.25" +unicode-properties = { version = "0.1.4", features = ["general-category"] } +unicode-security = "0.1.2" [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 d7e7274c..d51cefbc 100644 --- a/crates/session-store/src/fs_store.rs +++ b/crates/session-store/src/fs_store.rs @@ -918,6 +918,36 @@ mod tests { )); } + #[test] + fn uploaded_file_names_reject_format_mixed_script_and_confusable_forms() { + let tmp = tempfile::TempDir::new().unwrap(); + let store = FsStore::new(tmp.path()).unwrap(); + let session_id = new_session_id(); + let limits = UploadedFileLimits::default(); + + for file_name in [ + "safe\u{00ad}name.txt", + "safe\u{061c}name.txt", + "safe\u{180e}name.txt", + "safe\u{e0001}name.txt", + "p\u{0430}ypal.txt", + "\u{0440}\u{0430}\u{0443}\u{0440}\u{0430}\u{04cf}.txt", + "\u{ff26}\u{ff49}\u{ff4c}\u{ff45}.txt", + "re\u{0301}sume\u{0301}.txt", + ] { + assert!(matches!( + store.write_uploaded_file(session_id, file_name, "text/plain", b"safe", limits), + Err(StoreError::InvalidUploadedFileName) + )); + } + + for file_name in ["notes.txt", "résumé.txt", "日本語.txt", "📎.txt"] { + store + .write_uploaded_file(session_id, file_name, "text/plain", b"safe", limits) + .unwrap(); + } + } + #[test] fn paste_artifact_limits_and_corruption_fail_closed() { let tmp = tempfile::TempDir::new().unwrap(); diff --git a/crates/session-store/src/segment.rs b/crates/session-store/src/segment.rs index 0ea4fb04..e041ba7c 100644 --- a/crates/session-store/src/segment.rs +++ b/crates/session-store/src/segment.rs @@ -462,11 +462,18 @@ pub fn fork_at( ) -> Result { let entries = store.read_all(source_session_id, source_id)?; let cut = if at_turn_index == 0 { - // Branch directly after the SegmentStart (or whatever opens the - // segment), before any turn completes. + // Branch from the seeded state before any new turn completes. A typed + // input checkpoint immediately following SegmentStart is part of that + // seed and must stay atomic with its annotated history. entries .iter() - .position(|e| !matches!(e, LogEntry::AnnotatedSegmentStart { .. })) + .position(|entry| { + !matches!( + entry, + LogEntry::AnnotatedSegmentStart { .. } + | LogEntry::InputSegmentsCheckpoint { .. } + ) + }) .unwrap_or(entries.len()) } else { entries @@ -478,8 +485,9 @@ pub fn fork_at( let state = segment_log::collect_state(&entries[..cut]); let fork_id = crate::new_segment_id(); + let ts = segment_log::now_millis(); let entry = LogEntry::AnnotatedSegmentStart { - ts: segment_log::now_millis(), + ts, session_id: source_session_id, system_prompt: state.system_prompt, config: state.config, @@ -490,7 +498,14 @@ pub fn fork_at( }), compacted_from: None, }; - store.create_segment(source_session_id, fork_id, &[entry])?; + let mut fork_entries = vec![entry]; + if !state.user_segments.is_empty() { + fork_entries.push(LogEntry::InputSegmentsCheckpoint { + ts, + user_segments: state.user_segments, + }); + } + store.create_segment(source_session_id, fork_id, &fork_entries)?; Ok(fork_id) } diff --git a/crates/session-store/src/uploaded_file.rs b/crates/session-store/src/uploaded_file.rs index e8b5207f..291966bd 100644 --- a/crates/session-store/src/uploaded_file.rs +++ b/crates/session-store/src/uploaded_file.rs @@ -10,6 +10,8 @@ use protocol::{UploadedFileAvailability, UploadedFileRef}; use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; use unicode_normalization::UnicodeNormalization; +use unicode_properties::general_category::{GeneralCategory, UnicodeGeneralCategory}; +use unicode_security::{confusable_detection::skeleton, mixed_script::MixedScript}; use uuid::Uuid; use crate::StoreError; @@ -62,18 +64,26 @@ struct StoredUploadedFile { } pub(crate) fn validate_file_name(file_name: &str) -> Result<()> { + let normalized: String = file_name.nfkc().collect(); + let stem = file_name + .rsplit_once('.') + .map_or(file_name, |(stem, _)| stem); + let confusable_skeleton: String = skeleton(stem).collect(); + let ascii_confusable = stem.chars().any(|ch| !ch.is_ascii()) + && confusable_skeleton.is_ascii() + && !confusable_skeleton.eq_ignore_ascii_case(stem); + if file_name.is_empty() || file_name.chars().count() > MAX_FILE_NAME_CHARS || file_name == "." || file_name == ".." + || normalized != file_name + || !stem.is_single_script() + || ascii_confusable || 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) + || ch.general_category() == GeneralCategory::Format + || matches!(ch, '/' | '\\') }) { return Err(StoreError::InvalidUploadedFileName); diff --git a/crates/session-store/tests/session_test.rs b/crates/session-store/tests/session_test.rs index 09b80166..a8bddf1c 100644 --- a/crates/session-store/tests/session_test.rs +++ b/crates/session-store/tests/session_test.rs @@ -10,6 +10,7 @@ use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; use agen::{Engine, History}; use async_trait::async_trait; use common::MockLlmClient; +use protocol::{Segment, SessionSnapshotEntryData, UploadedFileAvailability, UploadedFileRef}; use session_store::{FsStore, LogEntry, SegmentStartState, Store, collect_state}; // ============================================================================= @@ -455,7 +456,11 @@ async fn session_fork_at_truncates_within_session() { let fork_segid = session_store::fork_at(&store, sid, segid, worker.turn_count()).unwrap(); let fork_entries = store.read_all(sid, fork_segid).unwrap(); - assert_eq!(fork_entries.len(), 1); // Just the new SegmentStart + assert_eq!(fork_entries.len(), 2); + assert!(matches!( + &fork_entries[1], + LogEntry::InputSegmentsCheckpoint { .. } + )); let fork_state = collect_state(&fork_entries); assert_eq!(fork_state.session_id, Some(sid), "fork_at inherits Session"); @@ -467,6 +472,7 @@ async fn session_fork_at_truncates_within_session() { .position(|e| matches!(e, LogEntry::TurnEnd { turn_count, .. } if *turn_count == worker.turn_count())) .expect("source segment has the matching TurnEnd"); let source_state_at_fork = collect_state(&all_entries[..=turn_end_pos]); + assert_eq!(fork_state.user_segments, source_state_at_fork.user_segments); assert_eq!(fork_state.history.len(), source_state_at_fork.history.len()); assert_eq!( fork_state.annotated_history, source_state_at_fork.annotated_history, @@ -492,6 +498,58 @@ async fn session_fork_at_truncates_within_session() { assert!(segs.contains(&fork_segid)); } +#[test] +fn rewound_fork_preserves_uploaded_file_segments_in_snapshot() { + let (_dir, store) = make_store(); + let config = RequestConfig::default(); + let (sid, segid) = session_store::create_segment( + &store, + SegmentStartState { + system_prompt: Some("System prompt"), + config: &config, + history: Vec::new(), + }, + ) + .unwrap(); + let uploaded = UploadedFileRef { + artifact_id: "uploaded-file-1".into(), + file_name: "notes.txt".into(), + media_type: "text/plain".into(), + created_at_ms: 123, + availability: UploadedFileAvailability::Available, + byte_len: 5, + sha256: "a".repeat(64), + source_entry_id: Some("entry-1".into()), + }; + let segments = vec![Segment::UploadedFile { + file: uploaded.clone(), + }]; + session_store::save_user_input( + &store, + sid, + segid, + segments.clone(), + annotated(&[Item::user_message(Segment::flatten_to_text(&segments))]), + ) + .unwrap(); + session_store::save_turn_end(&store, sid, segid, 1).unwrap(); + + let fork_segid = session_store::fork_at(&store, sid, segid, 1).unwrap(); + let fork_entries = store.read_all(sid, fork_segid).unwrap(); + let snapshot = session_store::public_snapshot::project_session_snapshot(sid, &fork_entries); + + assert!(fork_entries.iter().any(|entry| matches!( + entry, + LogEntry::InputSegmentsCheckpoint { user_segments, .. } + if user_segments == &vec![segments.clone()] + ))); + assert!(snapshot.entries.iter().any(|entry| matches!( + &entry.data, + SessionSnapshotEntryData::UserInput { segments: restored } + if restored == &segments + ))); +} + #[tokio::test] async fn session_config_changed_logged() { let (_dir, store) = make_store(); @@ -654,12 +712,19 @@ async fn nested_past_fork_leaves_ancestors_immutable() { let fork1_entries = store.read_all(sid, fork1).unwrap(); assert_eq!( fork1_entries.len(), - 1, - "fork1 is just its SegmentStart seed" + 2, + "fork1 stores its SegmentStart and typed input checkpoint" ); - // fork2's lineage points at fork1, not the root. - match &store.read_all(sid, fork2).unwrap()[0] { + // fork2's lineage points at fork1, not the root, and the typed seed remains + // intact across the nested turn-zero fork. + let fork2_entries = store.read_all(sid, fork2).unwrap(); + assert_eq!(fork2_entries.len(), 2); + assert_eq!( + collect_state(&fork2_entries).user_segments, + collect_state(&fork1_entries).user_segments + ); + match &fork2_entries[0] { LogEntry::AnnotatedSegmentStart { forked_from: Some(origin), .. diff --git a/crates/tui/Cargo.toml b/crates/tui/Cargo.toml index 67199b7e..18147a9b 100644 --- a/crates/tui/Cargo.toml +++ b/crates/tui/Cargo.toml @@ -30,4 +30,5 @@ pulldown-cmark = { version = "0.13.3", default-features = false } agen.workspace = true [dev-dependencies] +async-trait.workspace = true tempfile = { workspace = true } diff --git a/crates/tui/src/console/mod.rs b/crates/tui/src/console/mod.rs index f1402a29..b4f67f1f 100644 --- a/crates/tui/src/console/mod.rs +++ b/crates/tui/src/console/mod.rs @@ -123,11 +123,45 @@ fn copy_selection_to_terminal(app: &mut App) -> bool { type AttachmentUploadResult = Result; +fn mark_attachments_after_transport_acceptance( + pending: &mut Vec, + awaiting_acceptance: &mut Vec, +) { + awaiting_acceptance.append(pending); +} + +fn reconcile_attachment_submission( + pending: &mut Vec, + awaiting_acceptance: &mut Vec, + event: &Event, +) -> bool { + match event { + Event::UserMessage { segments, .. } => { + let had_awaiting = !awaiting_acceptance.is_empty(); + let accepted_ids = segments + .iter() + .filter_map(|segment| match segment { + Segment::UploadedFile { file } => Some(file.artifact_id.as_str()), + _ => None, + }) + .collect::>(); + awaiting_acceptance.retain(|file| !accepted_ids.contains(&file.artifact_id.as_str())); + had_awaiting && awaiting_acceptance.is_empty() + } + Event::Error { .. } if !awaiting_acceptance.is_empty() => { + pending.append(awaiting_acceptance); + false + } + _ => false, + } +} + struct ConsoleConnection { client: Client, standalone_host: Option, backend_target: Option, pending_attachments: Vec, + awaiting_attachment_acceptance: Vec, upload_tasks: Vec>, upload_ids: HashMap, } @@ -198,6 +232,7 @@ impl ConsoleConnection { standalone_host: Some(host), backend_target: None, pending_attachments: Vec::new(), + awaiting_attachment_acceptance: Vec::new(), upload_tasks: Vec::new(), upload_ids: HashMap::new(), } @@ -209,6 +244,7 @@ impl ConsoleConnection { standalone_host: None, backend_target: Some(target), pending_attachments: Vec::new(), + awaiting_attachment_acceptance: Vec::new(), upload_tasks: Vec::new(), upload_ids: HashMap::new(), } @@ -236,8 +272,7 @@ impl ConsoleConnection { } self.client.send(&prepared).await?; if carries_attachments { - self.pending_attachments.clear(); - self.upload_ids.clear(); + self.mark_attachments_awaiting_acceptance(); } Ok(()) } @@ -268,6 +303,23 @@ impl ConsoleConnection { Ok(()) } + fn observe_worker_event(&mut self, event: &Event) { + if reconcile_attachment_submission( + &mut self.pending_attachments, + &mut self.awaiting_attachment_acceptance, + event, + ) { + self.upload_ids.clear(); + } + } + + fn mark_attachments_awaiting_acceptance(&mut self) { + mark_attachments_after_transport_acceptance( + &mut self.pending_attachments, + &mut self.awaiting_attachment_acceptance, + ); + } + async fn clear_pending_attachments(&mut self) { for task in self.upload_tasks.drain(..) { task.abort(); @@ -287,6 +339,8 @@ impl ConsoleConnection { } async fn shutdown(&mut self) -> Result<(), Box> { + self.pending_attachments + .append(&mut self.awaiting_attachment_acceptance); self.clear_pending_attachments().await; if let Some(host) = self.standalone_host.take() { host.shutdown().await?; @@ -737,6 +791,7 @@ async fn drain_worker_events( match client.try_next_event()? { Some(ev) => { handled = true; + client.observe_worker_event(&ev); if let Some(method) = app.handle_worker_event(ev) { client.send(&method).await?; } @@ -811,6 +866,7 @@ async fn run_loop( }, LoopInput::Worker(event) => match event? { Some(ev) => { + client.observe_worker_event(&ev); if let Some(method) = app.handle_worker_event(ev) { client.send(&method).await?; } @@ -1349,7 +1405,10 @@ fn handle_pause_or_quit(app: &mut App) -> Option { mod tests { use super::*; use crate::text_selection::{HistoryViewport, SelectionRow}; - use protocol::{Event, RewindTarget, RewindTargetId, Segment}; + use async_trait::async_trait; + use protocol::{ + Event, RewindTarget, RewindTargetId, Segment, UploadedFileAvailability, UploadedFileRef, + }; #[test] fn standalone_console_starts_with_in_process_connection_ready() { @@ -1382,6 +1441,88 @@ mod tests { assert_eq!(attachment_media_type(Path::new("program.exe")), None); } + struct FailOnceSocket { + fail_next_send: bool, + } + + #[async_trait] + impl Socket for FailOnceSocket { + type Error = io::Error; + + async fn send(&mut self, _message: String) -> Result<(), Self::Error> { + if std::mem::take(&mut self.fail_next_send) { + Err(io::Error::new( + io::ErrorKind::ConnectionReset, + "disconnected", + )) + } else { + Ok(()) + } + } + + async fn next(&mut self) -> Result, Self::Error> { + Ok(None) + } + + fn try_next(&mut self) -> Result, Self::Error> { + Ok(None) + } + } + + #[tokio::test] + async fn attachment_submission_waits_for_authoritative_acceptance_and_can_retry() { + let file = UploadedFileRef { + artifact_id: "artifact-1".into(), + file_name: "notes.txt".into(), + media_type: "text/plain".into(), + created_at_ms: 1, + availability: UploadedFileAvailability::Available, + byte_len: 5, + sha256: "a".repeat(64), + source_entry_id: None, + }; + let mut connection = ConsoleConnection { + client: Client::new(FailOnceSocket { + fail_next_send: true, + }), + standalone_host: None, + backend_target: None, + pending_attachments: vec![file.clone()], + awaiting_attachment_acceptance: Vec::new(), + upload_tasks: Vec::new(), + upload_ids: HashMap::new(), + }; + + assert!( + connection + .send(&Method::Run { + input: vec![Segment::text("inspect")], + }) + .await + .is_err() + ); + assert_eq!(connection.pending_attachments, vec![file.clone()]); + assert!(connection.awaiting_attachment_acceptance.is_empty()); + + connection + .send(&Method::Run { + input: vec![Segment::text("inspect")], + }) + .await + .unwrap(); + assert!(connection.pending_attachments.is_empty()); + assert_eq!( + connection.awaiting_attachment_acceptance, + vec![file.clone()] + ); + + connection.observe_worker_event(&Event::UserMessage { + segments: vec![Segment::text("inspect"), Segment::UploadedFile { file }], + }); + assert!(connection.pending_attachments.is_empty()); + assert!(connection.awaiting_attachment_acceptance.is_empty()); + } + #[test] fn single_worker_mouse_capture_avoids_drag_and_all_motion_modes() { let mut ansi = String::new();