From 07fcc249fef034e98485c9a56cd42bed1b15c8e7 Mon Sep 17 00:00:00 2001 From: Hare Date: Fri, 2 Oct 2026 18:11:51 +0900 Subject: [PATCH] fix: resume exact reads within large segments --- crates/session-store/src/public_index.rs | 154 +++++++++++++++----- crates/worker-runtime/src/session_public.rs | 85 +++++++++++ 2 files changed, 206 insertions(+), 33 deletions(-) diff --git a/crates/session-store/src/public_index.rs b/crates/session-store/src/public_index.rs index 520065b25..d569c09b7 100644 --- a/crates/session-store/src/public_index.rs +++ b/crates/session-store/src/public_index.rs @@ -600,41 +600,56 @@ pub fn read_session_public_index_entry( else { return Ok(None); }; - let page = read_session_public_index_page( - session_root, - SessionPublicIndexLimits { - max_segments: 1, - ..limits - }, - expected_generation, - Some(SessionPublicIndexScanPosition { - segment_index, - byte_offset: 0, - entry_offset: 0, - lineage: None, - }), - 2, - |entry| entry.entry_ref == entry_ref, - )?; - let mut matches = page.entries.into_iter(); - let item = matches.next(); - if matches.next().is_some() { - return Err(SessionPublicIndexReadError::Corrupt); + let mut position = SessionPublicIndexScanPosition { + segment_index, + byte_offset: 0, + entry_offset: 0, + lineage: None, + }; + let mut generation = expected_generation.map(str::to_string); + let mut selected = None; + let mut scanned_bytes = 0_u64; + let mut scanned_segments = 0_usize; + let mut scanned_entries = 0_usize; + loop { + let page = read_session_public_index_page( + session_root, + SessionPublicIndexLimits { + max_segments: 1, + ..limits + }, + generation.as_deref(), + Some(position.clone()), + 2, + |entry| entry.entry_ref == entry_ref, + )?; + generation.get_or_insert_with(|| page.generation.clone()); + scanned_bytes = scanned_bytes.saturating_add(page.scanned_bytes); + scanned_segments = scanned_segments.saturating_add(page.scanned_segments); + scanned_entries = scanned_entries.saturating_add(page.scanned_entries); + for item in page.entries { + if selected.replace(item).is_some() { + return Err(SessionPublicIndexReadError::Corrupt); + } + } + let Some(next_position) = page.next_position else { + break; + }; + if next_position.segment_index != segment_index { + break; + } + if next_position == position { + return Err(SessionPublicIndexReadError::ResourceLimit); + } + position = next_position; } - if page - .next_position - .as_ref() - .is_some_and(|position| position.segment_index == segment_index) - { - return Err(SessionPublicIndexReadError::ResourceLimit); - } - Ok(item.map(|item| SessionPublicIndexEntryRead { - session_id: page.session_id, - generation: page.generation, + Ok(selected.map(|item| SessionPublicIndexEntryRead { + session_id: session_id.to_string(), + generation: generation.expect("a successful index page always has a generation"), item, - scanned_bytes: page.scanned_bytes, - scanned_segments: page.scanned_segments, - scanned_entries: page.scanned_entries, + scanned_bytes, + scanned_segments, + scanned_entries, })) } @@ -1810,6 +1825,79 @@ mod tests { } } + #[test] + fn exact_entry_read_resumes_within_selected_segment_across_byte_and_entry_budgets() { + let (_temp, root, store) = create_store(); + let session_id = crate::new_session_id(); + let segment_id = segment_id(1_700); + store + .create_segment( + session_id, + segment_id, + &[start( + session_id, + vec![message("seed", LoggedRole::User, "ordinary seed")], + None, + None, + )], + ) + .unwrap(); + for index in 0..8 { + let id = if index == 7 { + "target".to_string() + } else { + format!("ordinary-{index}") + }; + store + .append( + session_id, + segment_id, + &LogEntry::AnnotatedAssistantItem { + ts: 2 + index, + entry: message( + &id, + LoggedRole::Assistant, + if index == 7 { "needle" } else { "ordinary" }, + ), + }, + ) + .unwrap(); + } + let (_, first_record_bytes, _) = store + .read_first_log_record_read_only_bounded(session_id, segment_id, u64::MAX) + .unwrap(); + + let by_bytes = read_session_public_index_entry( + &root, + SessionPublicIndexLimits { + max_bytes: first_record_bytes, + ..SessionPublicIndexLimits::default() + }, + None, + &segment_id.to_string(), + "Etarget", + ) + .unwrap() + .unwrap(); + assert_eq!(by_bytes.item.entry.full_text, "needle"); + assert!(by_bytes.scanned_bytes > first_record_bytes); + + let by_entries = read_session_public_index_entry( + &root, + SessionPublicIndexLimits { + max_entries: 2, + ..SessionPublicIndexLimits::default() + }, + None, + &segment_id.to_string(), + "Etarget", + ) + .unwrap() + .unwrap(); + assert_eq!(by_entries.item.entry.full_text, "needle"); + assert!(by_entries.scanned_entries > 2); + } + #[test] fn incremental_scan_rejects_one_indivisible_record_over_the_byte_budget() { let (_temp, root, store) = create_store(); diff --git a/crates/worker-runtime/src/session_public.rs b/crates/worker-runtime/src/session_public.rs index 48906b000..60b154ac9 100644 --- a/crates/worker-runtime/src/session_public.rs +++ b/crates/worker-runtime/src/session_public.rs @@ -534,4 +534,89 @@ mod tests { } )); } + + #[test] + fn search_then_read_crosses_entry_budget_within_one_segment() { + let temp = tempfile::TempDir::new().unwrap(); + let root = temp.path().join("session"); + let store = WorkerSessionStore::new(&root).unwrap(); + let session_id = session_store::new_session_id(); + let segment_id = uuid::Uuid::from_u128(20_000); + store + .create_segment( + session_id, + segment_id, + &[LogEntry::AnnotatedSegmentStart { + ts: 1, + session_id, + system_prompt: None, + config: RequestConfig::default(), + history: vec![message("seed", "ordinary")], + forked_from: None, + compacted_from: None, + }], + ) + .unwrap(); + for index in 0..4 { + let id = if index == 3 { + "target".to_string() + } else { + format!("ordinary-{index}") + }; + store + .append( + session_id, + segment_id, + &LogEntry::AnnotatedAssistantItem { + ts: 2 + index, + entry: message(&id, if index == 3 { "needle" } else { "ordinary" }), + }, + ) + .unwrap(); + } + + let mut generation = None; + let mut scan_cursor = None; + let found = loop { + let mut search_request = request(session_id, generation.clone(), scan_cursor.clone()); + search_request.max_entries = 2; + let page = match search(&root, &search_request, None) { + SessionPublicSearchAvailability::Page { page } => page, + other => panic!("unexpected search result: {other:?}"), + }; + if let Some(item) = page.items.into_iter().next() { + break (page.generation, item); + } + assert!(page.has_more); + generation = Some(page.generation); + scan_cursor = page.next_scan_cursor; + }; + assert_eq!(found.1.entry_ref, "Etarget"); + + let read_page = match read( + &root, + &SessionPublicReadRequest { + workspace_id: "workspace-1".into(), + source: runtime_api::SessionPublicSource::Retained { + worker_id: crate::identity::WorkerId::now_v7(), + }, + expected_session_id: session_id.to_string(), + expected_generation: Some(found.0), + segment_id: segment_id.to_string(), + entry_ref: found.1.entry_ref, + mode: SessionPublicReadMode::Full, + byte_offset: 0, + max_content_bytes: 16 * 1024, + max_scan_bytes: 64 * 1024 * 1024, + max_segments: 64, + max_entries: 2, + }, + None, + ) { + SessionPublicReadAvailability::Page { page } => page, + other => panic!("unexpected direct read result: {other:?}"), + }; + assert_eq!(read_page.content, "needle"); + assert!(read_page.scanned_entries > 2); + } }