fix: resume exact reads within large segments

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(
+    let mut position = SessionPublicIndexScanPosition {
-        session_root,
+        segment_index,
-        SessionPublicIndexLimits {
+        byte_offset: 0,
-            max_segments: 1,
+        entry_offset: 0,
-            ..limits
+        lineage: None,
-        },
+    };
-        expected_generation,
+    let mut generation = expected_generation.map(str::to_string);
-        Some(SessionPublicIndexScanPosition {
+    let mut selected = None;
-            segment_index,
+    let mut scanned_bytes = 0_u64;
-            byte_offset: 0,
+    let mut scanned_segments = 0_usize;
-            entry_offset: 0,
+    let mut scanned_entries = 0_usize;
-            lineage: None,
+    loop {
-        }),
+        let page = read_session_public_index_page(
-        2,
+            session_root,
-        |entry| entry.entry_ref == entry_ref,
+            SessionPublicIndexLimits {
-    )?;
+                max_segments: 1,
-    let mut matches = page.entries.into_iter();
+                ..limits
-    let item = matches.next();
+            },
-    if matches.next().is_some() {
+            generation.as_deref(),
-        return Err(SessionPublicIndexReadError::Corrupt);
+            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
+    Ok(selected.map(|item| SessionPublicIndexEntryRead {
-        .next_position
+        session_id: session_id.to_string(),
-        .as_ref()
+        generation: generation.expect("a successful index page always has a generation"),
-        .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,
         item,
-        scanned_bytes: page.scanned_bytes,
+        scanned_bytes,
-        scanned_segments: page.scanned_segments,
+        scanned_segments,
-        scanned_entries: page.scanned_entries,
+        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();

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);
+    }
 }