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

fix: preserve attachment lifecycle boundaries

Browse Source
This commit is contained in:
Hare
2026-09-03 08:10:35 +09:00
parent 9dc8d9a77a
commit a2e1a3d939
8 changed files with 302 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/session-store/Cargo.toml
+2
View File
@@ -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 }
crates/session-store/src/fs_store.rs
+30
View File
@@ -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();
crates/session-store/src/segment.rs
+20 -5
View File
@@ -462,11 +462,18 @@ pub fn fork_at(
) -> Result<SegmentId, StoreError> {
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)
}
crates/session-store/src/uploaded_file.rs
+16 -6
View File
@@ -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);
crates/session-store/tests/session_test.rs
+70 -5
View File
@@ -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),
..
crates/tui/Cargo.toml
+1
View File
@@ -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 }
crates/tui/src/console/mod.rs
+144 -3
View File
@@ -123,11 +123,45 @@ fn copy_selection_to_terminal(app: &mut App) -> bool {
type AttachmentUploadResult = Result<UploadedFileRef, String>;
fn mark_attachments_after_transport_acceptance(
pending: &mut Vec<UploadedFileRef>,
awaiting_acceptance: &mut Vec<UploadedFileRef>,
) {
awaiting_acceptance.append(pending);
}
fn reconcile_attachment_submission(
pending: &mut Vec<UploadedFileRef>,
awaiting_acceptance: &mut Vec<UploadedFileRef>,
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::<Vec<_>>();
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<T> {
client: Client<T>,
standalone_host: Option<StandaloneHost>,
backend_target: Option<BackendRuntimeTarget>,
pending_attachments: Vec<UploadedFileRef>,
awaiting_attachment_acceptance: Vec<UploadedFileRef>,
upload_tasks: Vec<tokio::task::JoinHandle<()>>,
upload_ids: HashMap<PathBuf, String>,
}
@@ -198,6 +232,7 @@ impl<T: Socket> ConsoleConnection<T> {
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<T: Socket> ConsoleConnection<T> {
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<T: Socket> ConsoleConnection<T> {
}
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<T: Socket> ConsoleConnection<T> {
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<T: Socket> ConsoleConnection<T> {
}
async fn shutdown(&mut self) -> Result<(), Box<dyn std::error::Error>> {
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<T: Socket>(
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<T: Socket>(
},
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<Method> {
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<Option<String>, Self::Error> {
Ok(None)
}
fn try_next(&mut self) -> Result<Option<String>, 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();