fix: make subject session reads resumable

2026-10-02 17:36:10 +09:00
parent ce841cfee0
commit c491cf8a48
10 changed files with 1410 additions and 233 deletions
@@ -1,5 +1,6 @@
 use std::path::Path;

+use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
 use runtime_api::{
    SessionPublicEntryKind, SessionPublicLineage, SessionPublicLineageKind,
    SessionPublicReadAvailability, SessionPublicReadMode, SessionPublicReadPage,
@@ -7,14 +8,23 @@ use runtime_api::{
    SessionPublicSearchPage, SessionPublicSearchRequest, SessionPublicToolPart,
    SessionPublicUnavailableReason, WorkerSessionArchiveManifest,
 };
+use serde::{Deserialize, Serialize};
 use session_store::{
    SessionPublicIndex, SessionPublicIndexEntry, SessionPublicIndexEntryKind,
    SessionPublicIndexLimits, SessionPublicIndexLineage, SessionPublicIndexOriginKind,
-    SessionPublicIndexReadError, SessionPublicIndexToolPart, read_session_public_index,
+    SessionPublicIndexReadError, SessionPublicIndexScanPosition, SessionPublicIndexToolPart,
+    read_session_public_index, read_session_public_index_page,
 };

 use crate::retention::WorkerSessionArchiveManifest as DomainArchiveManifest;

+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(deny_unknown_fields)]
+struct SessionPublicRuntimeCursor {
+    generation: String,
+    position: SessionPublicIndexScanPosition,
+}
+
 pub(crate) fn search(
    session_root: &Path,
    request: &SessionPublicSearchRequest,
@@ -26,29 +36,23 @@ pub(crate) fn search(
            "Session public search limit must be between 1 and 100",
        );
    }
-    let index = match load_index(
-        session_root,
-        request.max_scan_bytes,
-        request.max_segments,
-        request.max_entries,
-    ) {
-        Ok(index) => index,
-        Err((reason, message)) => return unavailable(reason, message),
+    let cursor = match request.scan_cursor.as_deref().map(decode_scan_cursor) {
+        Some(Ok(cursor)) => Some(cursor),
+        Some(Err(())) => {
+            return unavailable(
+                SessionPublicUnavailableReason::InvalidCursor,
+                "Session public search cursor is invalid",
+            );
+        }
+        None => None,
    };
-    if index.session_id != request.expected_session_id {
-        return unavailable(
-            SessionPublicUnavailableReason::NotFound,
-            "Session identity does not match the authorized source",
-        );
-    }
-    if request
-        .expected_generation
-        .as_deref()
-        .is_some_and(|expected| expected != index.generation)
+    if let (Some(expected), Some(cursor)) =
+        (request.expected_generation.as_deref(), cursor.as_ref())
+        && expected != cursor.generation
    {
        return unavailable(
            SessionPublicUnavailableReason::InvalidCursor,
-            "Session public search generation changed",
+            "Session public search cursor generation changed",
        );
    }

@@ -58,51 +62,104 @@ pub(crate) fn search(
        .map(str::trim)
        .filter(|value| !value.is_empty())
        .map(str::to_lowercase);
-    let mut matches = Vec::new();
-    for segment in &index.segments {
-        for entry in &segment.entries {
-            if !matches_entry(entry, request, query.as_deref()) {
-                continue;
-            }
-            matches.push(SessionPublicSearchItem {
-                segment_id: segment.segment_id.clone(),
-                entry_ref: entry.entry_ref.clone(),
-                kind: entry_kind(entry.kind),
-                origin: entry.provenance.clone(),
-                tool_name: entry.tool_name.clone(),
-                tool_part: entry.tool_part.map(tool_part),
-                compact: entry.compact_text.clone(),
-                lineage: lineage(&segment.lineage),
-            });
+    let expected_generation = request
+        .expected_generation
+        .as_deref()
+        .or_else(|| cursor.as_ref().map(|cursor| cursor.generation.as_str()));
+    let index = match read_session_public_index_page(
+        session_root,
+        SessionPublicIndexLimits {
+            max_bytes: request.max_scan_bytes,
+            max_segments: request.max_segments,
+            max_entries: request.max_entries,
+        },
+        expected_generation,
+        cursor.as_ref().map(|cursor| cursor.position.clone()),
+        request.limit,
+        |entry| matches_entry(entry, request, query.as_deref()),
+    ) {
+        Ok(index) => index,
+        Err(error) => {
+            let (reason, message) = map_index_error(error);
+            return unavailable(reason, message);
        }
-    }
-    if request.offset > matches.len() {
+    };
+    if index.session_id != request.expected_session_id {
        return unavailable(
-            SessionPublicUnavailableReason::InvalidCursor,
-            "Session public search position is stale",
+            SessionPublicUnavailableReason::NotFound,
+            "Session identity does not match the authorized source",
        );
    }
-    let end = request
-        .offset
-        .saturating_add(request.limit)
-        .min(matches.len());
-    let items = matches[request.offset..end].to_vec();
-    let has_more = end < matches.len();
+
+    let mut items = Vec::with_capacity(index.entries.len());
+    for item in index.entries {
+        let scan_cursor = match encode_scan_cursor(&SessionPublicRuntimeCursor {
+            generation: index.generation.clone(),
+            position: item.scan_from,
+        }) {
+            Ok(cursor) => cursor,
+            Err(()) => {
+                return unavailable(
+                    SessionPublicUnavailableReason::StorageUnavailable,
+                    "Session public search cursor could not be encoded",
+                );
+            }
+        };
+        let entry = item.entry;
+        items.push(SessionPublicSearchItem {
+            segment_id: item.segment_id,
+            entry_ref: entry.entry_ref,
+            kind: entry_kind(entry.kind),
+            origin: entry.provenance,
+            tool_name: entry.tool_name,
+            tool_part: entry.tool_part.map(tool_part),
+            compact: entry.compact_text,
+            compact_truncated: entry.compact_text_truncated,
+            lineage: lineage(&item.lineage),
+            scan_cursor,
+        });
+    }
+    let next_scan_cursor = match index.next_position {
+        Some(position) => match encode_scan_cursor(&SessionPublicRuntimeCursor {
+            generation: index.generation.clone(),
+            position,
+        }) {
+            Ok(cursor) => Some(cursor),
+            Err(()) => {
+                return unavailable(
+                    SessionPublicUnavailableReason::StorageUnavailable,
+                    "Session public search cursor could not be encoded",
+                );
+            }
+        },
+        None => None,
+    };
    SessionPublicSearchAvailability::Page {
        page: SessionPublicSearchPage {
            session_id: index.session_id,
            generation: index.generation,
            items,
-            next_offset: has_more.then_some(end),
-            has_more,
+            next_scan_cursor,
+            has_more: index.has_more,
            scanned_bytes: index.scanned_bytes,
-            scanned_segments: index.segments.len(),
+            scanned_segments: index.scanned_segments,
            scanned_entries: index.scanned_entries,
            archive_manifest: archive_manifest.map(api_archive_manifest),
        },
    }
 }

+fn encode_scan_cursor(cursor: &SessionPublicRuntimeCursor) -> Result<String, ()> {
+    serde_json::to_vec(cursor)
+        .map(|bytes| URL_SAFE_NO_PAD.encode(bytes))
+        .map_err(|_| ())
+}
+
+fn decode_scan_cursor(cursor: &str) -> Result<SessionPublicRuntimeCursor, ()> {
+    let bytes = URL_SAFE_NO_PAD.decode(cursor).map_err(|_| ())?;
+    serde_json::from_slice(&bytes).map_err(|_| ())
+}
+
 pub(crate) fn read(
    session_root: &Path,
    request: &SessionPublicReadRequest,
@@ -309,6 +366,10 @@ fn map_index_error(
            SessionPublicUnavailableReason::ResourceLimit,
            "Session scan exceeds a resource limit",
        ),
+        SessionPublicIndexReadError::StaleCursor => (
+            SessionPublicUnavailableReason::InvalidCursor,
+            "Session public search cursor is stale",
+        ),
    }
 }

@@ -355,3 +416,132 @@ fn api_archive_manifest(manifest: &DomainArchiveManifest) -> WorkerSessionArchiv
        input_fingerprint: manifest.input_fingerprint.clone(),
    }
 }
+
+#[cfg(test)]
+mod tests {
+    use agen::llm_client::RequestConfig;
+    use session_store::{
+        LogEntry, LoggedContentPart, LoggedHistoryEntry, LoggedItem, LoggedRole,
+        LoggedSessionHistoryEntryId, LoggedSessionHistoryMetadata, LoggedSessionHistoryOrigin,
+        Store, WorkerSessionStore,
+    };
+
+    use super::*;
+
+    fn message(id: &str, text: &str) -> LoggedHistoryEntry {
+        LoggedHistoryEntry {
+            item: LoggedItem::Message {
+                role: LoggedRole::User,
+                content: vec![LoggedContentPart::Text { text: text.into() }],
+            },
+            metadata: LoggedSessionHistoryMetadata {
+                entry_id: LoggedSessionHistoryEntryId(id.to_string()),
+                origin: LoggedSessionHistoryOrigin::HumanInput {
+                    account_id: "account-1".to_string(),
+                },
+                derivation: None,
+            },
+        }
+    }
+
+    fn request(
+        session_id: session_store::SessionId,
+        generation: Option<String>,
+        scan_cursor: Option<String>,
+    ) -> SessionPublicSearchRequest {
+        SessionPublicSearchRequest {
+            workspace_id: "workspace-1".to_string(),
+            source: runtime_api::SessionPublicSource::Retained {
+                worker_id: crate::identity::WorkerId::now_v7(),
+            },
+            expected_session_id: session_id.to_string(),
+            expected_generation: generation,
+            query: Some("needle".to_string()),
+            kind: None,
+            tool_name: None,
+            tool_part: SessionPublicToolPart::Both,
+            scan_cursor,
+            limit: 20,
+            max_scan_bytes: 64 * 1024 * 1024,
+            max_segments: 64,
+            max_entries: 100_000,
+        }
+    }
+
+    #[test]
+    fn search_continues_beyond_segment_budget_and_stales_after_mutation() {
+        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 mut last_segment = None;
+        for index in 0..65_u128 {
+            let segment_id = uuid::Uuid::from_u128(10_000 + index);
+            last_segment = Some(segment_id);
+            store
+                .create_segment(
+                    session_id,
+                    segment_id,
+                    &[LogEntry::AnnotatedSegmentStart {
+                        ts: 1,
+                        session_id,
+                        system_prompt: None,
+                        config: RequestConfig::default(),
+                        history: vec![message(
+                            &format!("entry-{index}"),
+                            if index == 64 { "needle" } else { "ordinary" },
+                        )],
+                        forked_from: None,
+                        compacted_from: None,
+                    }],
+                )
+                .unwrap();
+        }
+
+        let first = match search(&root, &request(session_id, None, None), None) {
+            SessionPublicSearchAvailability::Page { page } => page,
+            other => panic!("unexpected first search result: {other:?}"),
+        };
+        assert!(first.items.is_empty());
+        assert!(first.has_more);
+        let cursor = first.next_scan_cursor.clone().unwrap();
+
+        let second = match search(
+            &root,
+            &request(
+                session_id,
+                Some(first.generation.clone()),
+                Some(cursor.clone()),
+            ),
+            None,
+        ) {
+            SessionPublicSearchAvailability::Page { page } => page,
+            other => panic!("unexpected continued search result: {other:?}"),
+        };
+        assert_eq!(second.items.len(), 1);
+        assert_eq!(second.items[0].compact, "needle");
+        assert!(!second.has_more);
+
+        store
+            .append(
+                session_id,
+                last_segment.unwrap(),
+                &LogEntry::AnnotatedAssistantItem {
+                    ts: 2,
+                    entry: message("mutation", "changed"),
+                },
+            )
+            .unwrap();
+        assert!(matches!(
+            search(
+                &root,
+                &request(session_id, Some(first.generation), Some(cursor)),
+                None,
+            ),
+            SessionPublicSearchAvailability::Unavailable {
+                reason: SessionPublicUnavailableReason::InvalidCursor,
+                ..
+            }
+        ));
+    }
+}