diff --git a/crates/runtime-api/src/lib.rs b/crates/runtime-api/src/lib.rs index 7eac0c840..a10b33361 100644 --- a/crates/runtime-api/src/lib.rs +++ b/crates/runtime-api/src/lib.rs @@ -755,7 +755,8 @@ pub struct SessionPublicSearchRequest { #[serde(default, skip_serializing_if = "Option::is_none")] pub tool_name: Option, pub tool_part: SessionPublicToolPart, - pub offset: usize, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub scan_cursor: Option, pub limit: usize, pub max_scan_bytes: u64, pub max_segments: usize, @@ -789,7 +790,10 @@ pub struct SessionPublicSearchItem { #[serde(default, skip_serializing_if = "Option::is_none")] pub tool_part: Option, pub compact: String, + pub compact_truncated: bool, pub lineage: SessionPublicLineage, + /// Opaque Runtime continuation immediately before this result. + pub scan_cursor: String, } #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] @@ -797,7 +801,8 @@ pub struct SessionPublicSearchPage { pub session_id: String, pub generation: String, pub items: Vec, - pub next_offset: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub next_scan_cursor: Option, pub has_more: bool, pub scanned_bytes: u64, pub scanned_segments: usize, diff --git a/crates/session-store/src/lib.rs b/crates/session-store/src/lib.rs index a3f9a4207..0a16e0451 100644 --- a/crates/session-store/src/lib.rs +++ b/crates/session-store/src/lib.rs @@ -101,8 +101,9 @@ pub use public_index::{ DEFAULT_SESSION_PUBLIC_INDEX_MAX_SEGMENTS, SESSION_PUBLIC_INDEX_COMPACT_TEXT_MAX_BYTES, SessionPublicIndex, SessionPublicIndexEntry, SessionPublicIndexEntryKind, SessionPublicIndexLimits, SessionPublicIndexLineage, SessionPublicIndexOriginKind, - SessionPublicIndexReadError, SessionPublicIndexSegment, SessionPublicIndexToolPart, - read_session_public_index, + SessionPublicIndexPage, SessionPublicIndexPageEntry, SessionPublicIndexReadError, + SessionPublicIndexScanPosition, SessionPublicIndexSegment, SessionPublicIndexToolPart, + read_session_public_index, read_session_public_index_page, }; pub use public_snapshot::{ DEFAULT_RETAINED_HISTORY_MAX_ENTRIES, DEFAULT_RETAINED_HISTORY_MAX_RESPONSE_BYTES, diff --git a/crates/session-store/src/public_index.rs b/crates/session-store/src/public_index.rs index 5ed7693f5..b5176e1df 100644 --- a/crates/session-store/src/public_index.rs +++ b/crates/session-store/src/public_index.rs @@ -10,6 +10,7 @@ use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; use thiserror::Error; +use crate::worker_session_store::RetainedLogRecord; use crate::{LogEntry, SegmentId, SegmentOrigin, SessionId, StoreError, WorkerSessionStore}; pub const DEFAULT_SESSION_PUBLIC_INDEX_MAX_BYTES: u64 = 64 * 1024 * 1024; @@ -50,6 +51,48 @@ pub enum SessionPublicIndexReadError { Storage, #[error("Worker Session public index exceeds a resource limit")] ResourceLimit, + #[error("Worker Session public index cursor is stale")] + StaleCursor, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct SessionPublicIndexScanPosition { + pub segment_index: usize, + pub byte_offset: u64, + pub entry_offset: usize, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub lineage: Option, +} + +impl Default for SessionPublicIndexScanPosition { + fn default() -> Self { + Self { + segment_index: 0, + byte_offset: 0, + entry_offset: 0, + lineage: None, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SessionPublicIndexPageEntry { + pub segment_id: String, + pub lineage: SessionPublicIndexLineage, + pub entry: SessionPublicIndexEntry, + pub scan_from: SessionPublicIndexScanPosition, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SessionPublicIndexPage { + pub session_id: String, + pub generation: String, + pub entries: Vec, + pub next_position: Option, + pub has_more: bool, + pub scanned_bytes: u64, + pub scanned_segments: usize, + pub scanned_entries: usize, } #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] @@ -121,6 +164,7 @@ pub struct SessionPublicIndexEntry { #[serde(default, skip_serializing_if = "Option::is_none")] pub tool_name: Option, pub compact_text: String, + pub compact_text_truncated: bool, pub full_text: String, } @@ -246,6 +290,421 @@ pub fn read_session_public_index( }) } +/// Incrementally scan the canonical public Session index. +/// +/// Continuations stop only at a retained log-unit boundary. Ordinary records are +/// independent units; a run from `Invoke`/`RunResumed` through its terminal +/// record remains one unit so uncommitted output never crosses the public +/// projection. A single unit that cannot fit the caller's byte or entry budget +/// is therefore the only scan shape that returns [`SessionPublicIndexReadError::ResourceLimit`]. +pub fn read_session_public_index_page( + session_root: &Path, + limits: SessionPublicIndexLimits, + expected_generation: Option<&str>, + position: Option, + max_results: usize, + mut matches: F, +) -> Result +where + F: FnMut(&SessionPublicIndexEntry) -> bool, +{ + if limits.max_bytes == 0 + || limits.max_segments == 0 + || limits.max_entries == 0 + || max_results == 0 + { + return Err(SessionPublicIndexReadError::ResourceLimit); + } + + let store = WorkerSessionStore::open_read_only(session_root).map_err(map_store_open_error)?; + let session_id = store + .session_id() + .map_err(map_store_error)? + .ok_or(SessionPublicIndexReadError::Missing)?; + let segment_ids = store + .list_segments_read_only(session_id) + .map_err(map_store_error)?; + let generation = scan_generation(&store, session_id, &segment_ids)?; + if expected_generation.is_some_and(|expected| expected != generation) { + return Err(SessionPublicIndexReadError::StaleCursor); + } + + let mut position = position.unwrap_or_default(); + if position.segment_index > segment_ids.len() + || (position.segment_index == segment_ids.len() + && (position.byte_offset != 0 + || position.entry_offset != 0 + || position.lineage.is_some())) + { + return Err(SessionPublicIndexReadError::StaleCursor); + } + + let mut scanned_bytes = 0_u64; + let mut scanned_segments = 0_usize; + let mut scanned_entries = 0_usize; + let mut output = Vec::new(); + let mut made_progress = false; + let mut stopped = false; + + while position.segment_index < segment_ids.len() && output.len() < max_results { + if scanned_segments == limits.max_segments { + stopped = true; + break; + } + let segment_id = segment_ids[position.segment_index]; + scanned_segments += 1; + + let (file_len, lineage, mut first_for_unit) = if position.byte_offset == 0 { + let remaining = limits.max_bytes.saturating_sub(scanned_bytes); + let (first, first_bytes, file_len) = match store + .read_first_log_record_read_only_bounded(session_id, segment_id, remaining) + { + Ok(value) => value, + Err(StoreError::ReadLimitExceeded) if made_progress => { + stopped = true; + break; + } + Err(error) => return Err(map_store_error(error)), + }; + let Some(first) = first else { + return Err(SessionPublicIndexReadError::Corrupt); + }; + scanned_bytes = scanned_bytes + .checked_add(first_bytes) + .ok_or(SessionPublicIndexReadError::ResourceLimit)?; + let (lineage, origin) = + validate_segment_start(session_id, segment_id, std::slice::from_ref(&first.entry))?; + if position + .lineage + .as_ref() + .is_some_and(|expected| expected != &lineage) + { + return Err(SessionPublicIndexReadError::StaleCursor); + } + validate_incremental_origin(segment_id, origin.as_ref(), &segment_ids)?; + (file_len, lineage, Some(first)) + } else { + let lineage = position + .lineage + .clone() + .ok_or(SessionPublicIndexReadError::StaleCursor)?; + validate_incremental_lineage(segment_id, &lineage, &segment_ids)?; + let (file_len, _) = store + .segment_log_observation(segment_id) + .map_err(map_store_error)?; + (file_len, lineage, None) + }; + + if position.byte_offset > file_len + || (position.byte_offset == file_len && position.entry_offset != 0) + { + return Err(SessionPublicIndexReadError::StaleCursor); + } + if position.byte_offset == file_len { + position = next_segment_position(position.segment_index); + made_progress = true; + continue; + } + loop { + if position.byte_offset >= file_len || output.len() == max_results { + break; + } + let unit_start = position.byte_offset; + let budget_before_unit = (scanned_bytes, scanned_entries); + let unit = match read_public_unit( + &store, + session_id, + segment_id, + unit_start, + limits.max_bytes.saturating_sub(scanned_bytes), + first_for_unit.take(), + ) { + Ok(unit) => unit, + Err(SessionPublicIndexReadError::ResourceLimit) if made_progress => { + scanned_bytes = budget_before_unit.0; + scanned_entries = budget_before_unit.1; + stopped = true; + break; + } + Err(error) => return Err(error), + }; + if unit_start > 0 + && unit + .entries + .iter() + .any(|entry| matches!(entry, LogEntry::AnnotatedSegmentStart { .. })) + { + return Err(SessionPublicIndexReadError::Corrupt); + } + let unit_entries = persisted_entry_units(&unit.entries); + if scanned_entries.saturating_add(unit_entries) > limits.max_entries { + if made_progress { + stopped = true; + break; + } + return Err(SessionPublicIndexReadError::ResourceLimit); + } + scanned_bytes = scanned_bytes + .checked_add(unit.bytes) + .ok_or(SessionPublicIndexReadError::ResourceLimit)?; + scanned_entries = scanned_entries + .checked_add(unit_entries) + .ok_or(SessionPublicIndexReadError::ResourceLimit)?; + + let committed = committed_conversation_records(&unit.entries); + let snapshot = crate::public_snapshot::project_session_snapshot_for_segment( + session_id, + Some(segment_id), + &committed, + ); + let projected = project_index_entries(segment_id, snapshot.entries)?; + if position.entry_offset > projected.len() + || (position.entry_offset != 0 && position.entry_offset == projected.len()) + { + return Err(SessionPublicIndexReadError::StaleCursor); + } + for (entry_index, entry) in projected + .into_iter() + .enumerate() + .skip(position.entry_offset) + { + let resume_at = if entry_index + 1 < unit.public_entry_count { + SessionPublicIndexScanPosition { + segment_index: position.segment_index, + byte_offset: unit_start, + entry_offset: entry_index + 1, + lineage: Some(lineage.clone()), + } + } else if unit.end_offset == file_len { + next_segment_position(position.segment_index) + } else { + SessionPublicIndexScanPosition { + segment_index: position.segment_index, + byte_offset: unit.end_offset, + entry_offset: 0, + lineage: Some(lineage.clone()), + } + }; + if matches(&entry) { + output.push(SessionPublicIndexPageEntry { + segment_id: segment_id.to_string(), + lineage: lineage.clone(), + entry, + scan_from: SessionPublicIndexScanPosition { + segment_index: position.segment_index, + byte_offset: unit_start, + entry_offset: entry_index, + lineage: Some(lineage.clone()), + }, + }); + if output.len() == max_results { + position = resume_at; + break; + } + } + } + made_progress = true; + if output.len() == max_results { + break; + } + position = if unit.end_offset == file_len { + next_segment_position(position.segment_index) + } else { + SessionPublicIndexScanPosition { + segment_index: position.segment_index, + byte_offset: unit.end_offset, + entry_offset: 0, + lineage: Some(lineage.clone()), + } + }; + if position.segment_index >= segment_ids.len() + || position.segment_index + != segment_ids + .iter() + .position(|candidate| *candidate == segment_id) + .unwrap_or(position.segment_index) + { + break; + } + } + if stopped || output.len() == max_results { + break; + } + if position.segment_index < segment_ids.len() + && segment_ids[position.segment_index] == segment_id + { + position = next_segment_position(position.segment_index); + } + } + + let final_segment_ids = store + .list_segments_read_only(session_id) + .map_err(map_store_error)?; + let final_generation = scan_generation(&store, session_id, &final_segment_ids)?; + if final_generation != generation { + return Err(SessionPublicIndexReadError::StaleCursor); + } + + let has_more = stopped || position.segment_index < segment_ids.len(); + Ok(SessionPublicIndexPage { + session_id: session_id.to_string(), + generation, + entries: output, + next_position: has_more.then_some(position), + has_more, + scanned_bytes, + scanned_segments, + scanned_entries, + }) +} + +#[derive(Debug)] +struct PublicLogUnit { + entries: Vec, + end_offset: u64, + bytes: u64, + public_entry_count: usize, +} + +fn read_public_unit( + store: &WorkerSessionStore, + session_id: SessionId, + segment_id: SegmentId, + start_offset: u64, + max_bytes: u64, + first: Option, +) -> Result { + let (first, first_bytes) = match first { + Some(record) => (Some(record), 0), + None => store + .read_next_log_record_read_only_bounded(session_id, segment_id, start_offset, max_bytes) + .map_err(map_store_error)?, + }; + let Some(first) = first else { + return Err(SessionPublicIndexReadError::StaleCursor); + }; + if first.start_offset != start_offset { + return Err(SessionPublicIndexReadError::StaleCursor); + } + let mut bytes = first_bytes; + let mut end_offset = first.end_offset; + let run = matches!( + first.entry, + LogEntry::Invoke { .. } | LogEntry::RunResumed { .. } + ); + let mut entries = vec![first.entry]; + if run { + loop { + let (next, read_bytes) = store + .read_next_log_record_read_only_bounded( + session_id, + segment_id, + end_offset, + max_bytes.saturating_sub(bytes), + ) + .map_err(map_store_error)?; + let Some(next) = next else { + break; + }; + bytes = bytes + .checked_add(read_bytes) + .ok_or(SessionPublicIndexReadError::ResourceLimit)?; + end_offset = next.end_offset; + let terminal = matches!( + next.entry, + LogEntry::RunCompleted { .. } + | LogEntry::RunYielded { .. } + | LogEntry::RunCancelled { .. } + | LogEntry::RunErrored { .. } + | LogEntry::PausedTurnAbandoned { .. } + ); + entries.push(next.entry); + if terminal { + break; + } + } + } + let committed = committed_conversation_records(&entries); + let snapshot = crate::public_snapshot::project_session_snapshot_for_segment( + session_id, + Some(segment_id), + &committed, + ); + let public_entry_count = project_index_entries(segment_id, snapshot.entries)?.len(); + Ok(PublicLogUnit { + entries, + end_offset, + bytes, + public_entry_count, + }) +} + +fn next_segment_position(segment_index: usize) -> SessionPublicIndexScanPosition { + SessionPublicIndexScanPosition { + segment_index: segment_index.saturating_add(1), + byte_offset: 0, + entry_offset: 0, + lineage: None, + } +} + +fn validate_incremental_lineage( + segment_id: SegmentId, + lineage: &SessionPublicIndexLineage, + segment_ids: &[SegmentId], +) -> Result<(), SessionPublicIndexReadError> { + let parent = match lineage.origin_kind { + SessionPublicIndexOriginKind::Root => { + if lineage.parent_segment_id.is_some() || lineage.parent_turn_index.is_some() { + return Err(SessionPublicIndexReadError::StaleCursor); + } + return Ok(()); + } + SessionPublicIndexOriginKind::Fork | SessionPublicIndexOriginKind::Compact => lineage + .parent_segment_id + .as_deref() + .and_then(|value| uuid::Uuid::parse_str(value).ok()) + .ok_or(SessionPublicIndexReadError::StaleCursor)?, + }; + if parent == segment_id || !segment_ids.contains(&parent) || lineage.parent_turn_index.is_none() + { + return Err(SessionPublicIndexReadError::StaleCursor); + } + Ok(()) +} + +fn validate_incremental_origin( + segment_id: SegmentId, + origin: Option<&SegmentOrigin>, + segment_ids: &[SegmentId], +) -> Result<(), SessionPublicIndexReadError> { + if origin.is_some_and(|origin| { + origin.segment_id == segment_id || !segment_ids.contains(&origin.segment_id) + }) { + return Err(SessionPublicIndexReadError::Corrupt); + } + Ok(()) +} + +fn scan_generation( + store: &WorkerSessionStore, + session_id: SessionId, + segment_ids: &[SegmentId], +) -> Result { + let mut generation = Sha256::new(); + generation.update(b"yoi-session-public-index-scan-v1\0"); + generation.update(session_id.as_bytes()); + for segment_id in segment_ids { + let (len, modified_nanos) = store + .segment_log_observation(*segment_id) + .map_err(map_store_error)?; + generation.update(segment_id.as_bytes()); + generation.update(len.to_be_bytes()); + generation.update(modified_nanos.to_be_bytes()); + } + Ok(URL_SAFE_NO_PAD.encode(generation.finalize())) +} + fn validate_segment_start( expected_session_id: SessionId, segment_id: SegmentId, @@ -457,6 +916,7 @@ fn project_index_entries( _ => continue, }; + let (compact_text, compact_text_truncated) = bounded_compact_text(&compact_source); projected.push(SessionPublicIndexEntry { segment_id: segment_id.to_string(), entry_ref, @@ -465,10 +925,8 @@ fn project_index_entries( timestamp: entry.timestamp, tool_part, tool_name, - compact_text: truncate_utf8( - &compact_source, - SESSION_PUBLIC_INDEX_COMPACT_TEXT_MAX_BYTES, - ), + compact_text, + compact_text_truncated, full_text, }); } @@ -476,6 +934,13 @@ fn project_index_entries( Ok(projected) } +fn bounded_compact_text(value: &str) -> (String, bool) { + ( + truncate_utf8(value, SESSION_PUBLIC_INDEX_COMPACT_TEXT_MAX_BYTES), + value.len() > SESSION_PUBLIC_INDEX_COMPACT_TEXT_MAX_BYTES, + ) +} + fn content_parts_text(parts: &[SessionContentPart]) -> String { let mut output = String::new(); for part in parts { @@ -855,6 +1320,19 @@ mod tests { } } + #[test] + fn compact_text_reports_utf8_safe_truncation() { + let exact = "x".repeat(SESSION_PUBLIC_INDEX_COMPACT_TEXT_MAX_BYTES); + assert_eq!(bounded_compact_text(&exact), (exact, false)); + + let oversized = "界".repeat(SESSION_PUBLIC_INDEX_COMPACT_TEXT_MAX_BYTES); + let (compact, truncated) = bounded_compact_text(&oversized); + assert!(truncated); + assert!(compact.len() <= SESSION_PUBLIC_INDEX_COMPACT_TEXT_MAX_BYTES); + assert!(compact.ends_with('…')); + assert!(compact.is_char_boundary(compact.len())); + } + #[test] fn excludes_uncommitted_cancelled_and_failed_outputs_but_keeps_committed_outcomes() { let (_temp, root, store) = create_store(); @@ -1079,6 +1557,169 @@ mod tests { assert_ne!(replaced.generation, added.generation); } + #[test] + fn incremental_scan_finds_hit_beyond_segment_budget_and_stales_after_append() { + let (_temp, root, store) = create_store(); + let session_id = crate::new_session_id(); + for index in 0..65_u128 { + let text = if index == 64 { + "needle beyond first page" + } else { + "ordinary entry" + }; + store + .create_segment( + session_id, + segment_id(1_000 + index), + &[start( + session_id, + vec![message(&format!("entry-{index}"), LoggedRole::User, text)], + None, + None, + )], + ) + .unwrap(); + } + let limits = SessionPublicIndexLimits { + max_segments: 64, + ..SessionPublicIndexLimits::default() + }; + let first = read_session_public_index_page(&root, limits, None, None, 10, |entry| { + entry.full_text.contains("needle") + }) + .unwrap(); + assert!(first.entries.is_empty()); + assert!(first.has_more); + assert_eq!(first.scanned_segments, 64); + let next_position = first.next_position.clone().unwrap(); + + let second = read_session_public_index_page( + &root, + limits, + Some(&first.generation), + Some(next_position.clone()), + 10, + |entry| entry.full_text.contains("needle"), + ) + .unwrap(); + assert_eq!(second.entries.len(), 1); + assert_eq!( + second.entries[0].entry.full_text, + "needle beyond first page" + ); + assert!(!second.has_more); + assert!(second.next_position.is_none()); + + store + .append( + session_id, + segment_id(1_064), + &LogEntry::AnnotatedAssistantItem { + ts: 2, + entry: message("mutated", LoggedRole::Assistant, "changed"), + }, + ) + .unwrap(); + assert_eq!( + read_session_public_index_page( + &root, + limits, + Some(&first.generation), + Some(next_position), + 10, + |_| true, + ), + Err(SessionPublicIndexReadError::StaleCursor) + ); + } + + #[test] + fn incremental_scan_resumes_across_byte_and_entry_budgets() { + let (_temp, root, store) = create_store(); + let session_id = crate::new_session_id(); + let first_segment = segment_id(1_500); + let second_segment = segment_id(1_501); + for (segment_id, id, text) in [ + (first_segment, "entry-a", "aaaaaa"), + (second_segment, "entry-b", "needle"), + ] { + store + .create_segment( + session_id, + segment_id, + &[start( + session_id, + vec![message(id, LoggedRole::User, text)], + None, + None, + )], + ) + .unwrap(); + } + let first_len = store.segment_log_len(first_segment).unwrap(); + for limits in [ + SessionPublicIndexLimits { + max_bytes: first_len, + ..SessionPublicIndexLimits::default() + }, + SessionPublicIndexLimits { + max_entries: 2, + ..SessionPublicIndexLimits::default() + }, + ] { + let first = read_session_public_index_page(&root, limits, None, None, 10, |entry| { + entry.full_text == "needle" + }) + .unwrap(); + assert!(first.entries.is_empty()); + assert!(first.has_more); + let second = read_session_public_index_page( + &root, + limits, + Some(&first.generation), + first.next_position, + 10, + |entry| entry.full_text == "needle", + ) + .unwrap(); + assert_eq!(second.entries.len(), 1); + assert!(!second.has_more); + } + } + + #[test] + fn incremental_scan_rejects_one_indivisible_record_over_the_byte_budget() { + let (_temp, root, store) = create_store(); + let session_id = crate::new_session_id(); + store + .create_segment( + session_id, + segment_id(2_000), + &[start( + session_id, + vec![message("large-entry", LoggedRole::User, &"x".repeat(8_192))], + None, + None, + )], + ) + .unwrap(); + + assert_eq!( + read_session_public_index_page( + &root, + SessionPublicIndexLimits { + max_bytes: 512, + ..SessionPublicIndexLimits::default() + }, + None, + None, + 10, + |_| true, + ), + Err(SessionPublicIndexReadError::ResourceLimit) + ); + } + #[test] fn enforces_byte_segment_and_entry_bounds_with_typed_errors() { let (_temp, root, store) = create_store(); diff --git a/crates/session-store/src/worker_session_store.rs b/crates/session-store/src/worker_session_store.rs index f6820cff1..3711524e5 100644 --- a/crates/session-store/src/worker_session_store.rs +++ b/crates/session-store/src/worker_session_store.rs @@ -485,6 +485,22 @@ impl WorkerSessionStore { }) } + pub(crate) fn segment_log_observation( + &self, + segment_id: SegmentId, + ) -> Result<(u64, u128), StoreError> { + let metadata = fs::metadata(self.log_path(segment_id))?; + let modified = metadata + .modified()? + .duration_since(SystemTime::UNIX_EPOCH) + .map_err(|_| StoreError::Corrupt { + line: 0, + message: "retained Segment modification time predates the Unix epoch".to_string(), + })? + .as_nanos(); + Ok((metadata.len(), modified)) + } + pub fn segment_log_len(&self, segment_id: SegmentId) -> Result { Ok(fs::metadata(self.log_path(segment_id))?.len()) } diff --git a/crates/worker-runtime/src/session_public.rs b/crates/worker-runtime/src/session_public.rs index 1ff39224c..c8a9c486d 100644 --- a/crates/worker-runtime/src/session_public.rs +++ b/crates/worker-runtime/src/session_public.rs @@ -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 { + serde_json::to_vec(cursor) + .map(|bytes| URL_SAFE_NO_PAD.encode(bytes)) + .map_err(|_| ()) +} + +fn decode_scan_cursor(cursor: &str) -> Result { + 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, + scan_cursor: Option, + ) -> 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, + .. + } + )); + } +} diff --git a/crates/worker/src/feature/builtin/subjektiv_session.rs b/crates/worker/src/feature/builtin/subjektiv_session.rs index d5f101477..3f9dfab14 100644 --- a/crates/worker/src/feature/builtin/subjektiv_session.rs +++ b/crates/worker/src/feature/builtin/subjektiv_session.rs @@ -730,6 +730,44 @@ mod tests { } } )); + + let escaped_content = "\0".repeat(8 * 1024); + let paged_read = SubjektivSessionBackendResponse::Ok { + result: SubjektivSessionBackendResult::Read(server_api::SubjektivSessionReadResponse { + session_id: "session-1".into(), + segment_id: "segment-1".into(), + entry_ref: "Eentry-1".into(), + kind: server_api::SubjektivSessionEntryKind::User, + origin: protocol::SessionEntryProvenance::HumanInput, + lineage: server_api::SubjektivSessionLineage { + kind: server_api::SubjektivSessionLineageKind::Root, + parent_segment_id: None, + at_turn_index: None, + }, + mode: server_api::SubjektivSessionReadMode::Full, + content: escaped_content.clone(), + truncated: true, + next_cursor: Some("next-page".into()), + has_more: true, + }), + }; + let output = model_visible_response(paged_read).unwrap(); + let content = output.content.unwrap(); + assert!(content.len() <= SUBJEKTIV_SESSION_MAX_TOOL_CONTENT_BYTES); + let visible: SubjektivSessionBackendResponse = serde_json::from_str(&content).unwrap(); + assert!(matches!( + visible, + SubjektivSessionBackendResponse::Ok { + result: SubjektivSessionBackendResult::Read( + server_api::SubjektivSessionReadResponse { + content, + next_cursor: Some(cursor), + has_more: true, + .. + } + ) + } if content == escaped_content && cursor == "next-page" + )); } #[tokio::test] diff --git a/crates/workspace-server/src/latest_schema.sql b/crates/workspace-server/src/latest_schema.sql index c640193f9..3a24911fc 100644 --- a/crates/workspace-server/src/latest_schema.sql +++ b/crates/workspace-server/src/latest_schema.sql @@ -961,7 +961,7 @@ CREATE TABLE worker_session_archive_observe_grants ( workspace_id TEXT NOT NULL, archive_id TEXT NOT NULL, controller_runtime_id TEXT NOT NULL, controller_worker_id TEXT NOT NULL, subject_runtime_id TEXT NOT NULL, subject_worker_id TEXT NOT NULL, - source_grant_id TEXT NOT NULL, granted_at TEXT NOT NULL, + source_grant_id TEXT NOT NULL, granted_at TEXT NOT NULL, revoked_at TEXT, PRIMARY KEY(workspace_id,archive_id,controller_runtime_id,controller_worker_id,source_grant_id), FOREIGN KEY(workspace_id) REFERENCES workspaces(workspace_id) ON DELETE CASCADE, FOREIGN KEY(archive_id) REFERENCES worker_session_archives(archive_id) ON DELETE CASCADE); diff --git a/crates/workspace-server/src/retention.rs b/crates/workspace-server/src/retention.rs index cd726fa0f..bb36636f8 100644 --- a/crates/workspace-server/src/retention.rs +++ b/crates/workspace-server/src/retention.rs @@ -709,6 +709,7 @@ impl SqliteWorkspaceStore { WHERE workspace_id = ?1 AND archive_id = ?2 AND controller_runtime_id = ?3 AND controller_worker_id = ?4 AND subject_runtime_id = ?5 AND subject_worker_id = ?6 + AND revoked_at IS NULL )", params![ workspace_id, @@ -1301,6 +1302,32 @@ mod tests { #[test] fn purge_tombstone_commit_is_idempotent() { let s = setup(); + let archived_grant_id = "archive-observe-grant"; + let archived_controller = + RuntimeWorkerRef::new("r", WorkerId::from_legacy_u64(2).to_string()); + s.with_conn(|conn| { + conn.execute( + "INSERT INTO worker_registry(\ + workspace_id,worker_id,runtime_id,display_name,profile,retention_state,created_at,updated_at\ + ) VALUES('w',?1,'r','controller','builtin:companion','normal','created','controller-rev1')", + [&archived_controller.worker_id], + )?; + Ok(()) + }) + .unwrap(); + s.create_worker_control_grant(&crate::store::WorkerControlGrantRecord { + workspace_id: "w".into(), + grant_id: archived_grant_id.into(), + controller: archived_controller.clone(), + subject: req().worker, + relation: "subject_session".into(), + origin: "test".into(), + permissions: vec!["observe".into()], + operation_id: archived_grant_id.into(), + created_at: "2026-10-02T00:00:00Z".into(), + revoked_at: None, + }) + .unwrap(); s.with_conn(|conn| { conn.execute("INSERT INTO typed_tickets(workspace_id,ticket_id,slug,title,status,kind,priority,body,workflow_state,workflow_state_explicit) VALUES('w','ticket-old','ticket-old','Old Ticket','open','task','normal','','planning',1)", [])?; conn.execute("INSERT INTO worker_registry(workspace_id,worker_id,runtime_id,display_name,profile,retention_state,created_at,updated_at) VALUES('w','1','r','old worker','builtin:coder','normal','created','rev1')", [])?; @@ -1351,6 +1378,31 @@ mod tests { .state, WorkerRemovalPlanState::Succeeded ); + assert!( + s.worker_session_archive_observable_by( + "w", + p.archive_id.as_deref().unwrap(), + &archived_controller, + &p.worker, + ) + .unwrap(), + "the explicit observe grant survives Worker removal for the committed archive" + ); + assert!( + s.revoke_worker_control_grant("w", archived_grant_id, "2026-10-02T00:01:00Z") + .unwrap(), + "the existing grant revocation authority must also revoke its archive projection" + ); + assert!( + !s.worker_session_archive_observable_by( + "w", + p.archive_id.as_deref().unwrap(), + &archived_controller, + &p.worker, + ) + .unwrap(), + "every later cursor page and direct Read must observe archive grant revocation" + ); assert!(s.worker_tombstone("w", &p.worker).unwrap().is_some()); let reservation_state: String = s.with_conn(|conn| { conn.query_row( diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 0858b92a8..4347072a0 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -17033,7 +17033,7 @@ struct SubjektivSessionSearchPosition { storage: server_api::SubjektivSessionStorage, archive_id: Option, generation: String, - offset: usize, + runtime_scan_cursor: String, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -17258,9 +17258,7 @@ fn subjektiv_session_list( } } - let has_more = session_id.is_none() - && (store_has_more || scanned == SUBJEKTIV_SESSION_MAX_LIMIT) - && after.is_some(); + let has_more = session_id.is_none() && store_has_more && after.is_some(); let next_cursor = if has_more { let (after_attributed_at, after_session_id) = after.expect("checked above"); match encode_subjektiv_cursor( @@ -17404,102 +17402,129 @@ fn subjektiv_session_read( "Session storage changed while the read cursor was active", ); } - let request = runtime_api::SessionPublicReadRequest { - workspace_id: workspace_id.to_string(), - source: source.source.clone(), - expected_session_id: session_id.clone(), - expected_generation: cursor.as_ref().map(|cursor| cursor.generation.clone()), - segment_id: segment_id.clone(), - entry_ref: entry_ref.clone(), - mode: subjektiv_session_read_mode(input.mode), - byte_offset: cursor.as_ref().map_or(0, |cursor| cursor.byte_offset), - max_content_bytes: server_api::SUBJEKTIV_SESSION_MAX_READ_CONTENT_BYTES, - max_scan_bytes: SUBJEKTIV_SESSION_MAX_SCAN_BYTES, - max_segments: SUBJEKTIV_SESSION_MAX_SCAN_SEGMENTS, - max_entries: SUBJEKTIV_SESSION_MAX_SCAN_ENTRIES, - }; - let availability = match api - .runtime - .session_public_read(&attribution.runtime_id, request) - { - Ok(availability) => availability, - Err(error) => return subjektiv_session_runtime_error(error), - }; - let page = match availability { - runtime_api::SessionPublicReadAvailability::Page { page } => page, - runtime_api::SessionPublicReadAvailability::Unavailable { reason, message } => { - return subjektiv_session_error(subjektiv_session_diagnostic_code(reason), message); + let expected_mode = subjektiv_session_read_mode(input.mode); + let byte_offset = cursor.as_ref().map_or(0, |cursor| cursor.byte_offset); + let mut expected_generation = cursor.as_ref().map(|cursor| cursor.generation.clone()); + let mut max_content_bytes = server_api::SUBJEKTIV_SESSION_MAX_READ_CONTENT_BYTES; + loop { + let request = runtime_api::SessionPublicReadRequest { + workspace_id: workspace_id.to_string(), + source: source.source.clone(), + expected_session_id: session_id.clone(), + expected_generation: expected_generation.clone(), + segment_id: segment_id.clone(), + entry_ref: entry_ref.clone(), + mode: expected_mode, + byte_offset, + max_content_bytes, + max_scan_bytes: SUBJEKTIV_SESSION_MAX_SCAN_BYTES, + max_segments: SUBJEKTIV_SESSION_MAX_SCAN_SEGMENTS, + max_entries: SUBJEKTIV_SESSION_MAX_SCAN_ENTRIES, + }; + let availability = match api + .runtime + .session_public_read(&attribution.runtime_id, request) + { + Ok(availability) => availability, + Err(error) => return subjektiv_session_runtime_error(error), + }; + let page = match availability { + runtime_api::SessionPublicReadAvailability::Page { page } => page, + runtime_api::SessionPublicReadAvailability::Unavailable { reason, message } => { + return subjektiv_session_error(subjektiv_session_diagnostic_code(reason), message); + } + }; + if !subjektiv_session_archive_manifest_matches(&source, page.archive_manifest.as_ref()) { + return subjektiv_session_error( + server_api::SubjektivSessionDiagnosticCode::ArchiveIncomplete, + "Runtime Session archive does not match the committed catalog", + ); } - }; - if !subjektiv_session_archive_manifest_matches(&source, page.archive_manifest.as_ref()) { - return subjektiv_session_error( - server_api::SubjektivSessionDiagnosticCode::ArchiveIncomplete, - "Runtime Session archive does not match the committed catalog", - ); - } - if page.session_id != session_id - || page.segment_id != segment_id - || page.entry_ref != entry_ref - || page.mode != subjektiv_session_read_mode(input.mode) - { - return subjektiv_session_error( - server_api::SubjektivSessionDiagnosticCode::CorruptLog, - "Runtime Session read response did not match the selected entry", - ); - } - let next_cursor = if page.has_more { - let Some(byte_offset) = page.next_byte_offset else { + if page.session_id != session_id + || page.segment_id != segment_id + || page.entry_ref != entry_ref + || page.mode != expected_mode + { return subjektiv_session_error( server_api::SubjektivSessionDiagnosticCode::CorruptLog, - "Runtime Session read response omitted its continuation position", + "Runtime Session read response did not match the selected entry", ); + } + if expected_generation.is_none() { + expected_generation = Some(page.generation.clone()); + } + let next_cursor = if page.has_more { + let Some(next_byte_offset) = page.next_byte_offset else { + return subjektiv_session_error( + server_api::SubjektivSessionDiagnosticCode::CorruptLog, + "Runtime Session read response omitted its continuation position", + ); + }; + match encode_subjektiv_cursor( + "session-read", + &SubjektivSessionReadCursor { + workspace_id: workspace_id.to_string(), + subject_id: subject_id.to_string(), + controller_runtime_id: controller.runtime_id.clone(), + controller_worker_id: controller.worker_id.clone(), + session_id: session_id.clone(), + segment_id: segment_id.clone(), + entry_ref: entry_ref.clone(), + mode: input.mode, + runtime_id: attribution.runtime_id.clone(), + worker_id: attribution.worker_id.clone(), + storage: source.storage, + archive_id: source.archive_id.clone(), + generation: page.generation.clone(), + byte_offset: next_byte_offset, + }, + ) { + Ok(cursor) => Some(cursor), + Err(error) => { + return subjektiv_session_error( + server_api::SubjektivSessionDiagnosticCode::SubjectUnavailable, + error.error.to_string(), + ); + } + } + } else { + None }; - match encode_subjektiv_cursor( - "session-read", - &SubjektivSessionReadCursor { - workspace_id: workspace_id.to_string(), - subject_id: subject_id.to_string(), - controller_runtime_id: controller.runtime_id.clone(), - controller_worker_id: controller.worker_id.clone(), - session_id: session_id.clone(), - segment_id: segment_id.clone(), - entry_ref: entry_ref.clone(), - mode: input.mode, - runtime_id: attribution.runtime_id, - worker_id: attribution.worker_id, - storage: source.storage, - archive_id: source.archive_id, - generation: page.generation.clone(), - byte_offset, - }, - ) { - Ok(cursor) => Some(cursor), + let response = server_api::SubjektivSessionBackendResponse::Ok { + result: server_api::SubjektivSessionBackendResult::Read( + server_api::SubjektivSessionReadResponse { + session_id: session_id.clone(), + segment_id: segment_id.clone(), + entry_ref: entry_ref.clone(), + kind: subjektiv_session_entry_kind(page.kind), + origin: page.origin, + lineage: subjektiv_session_lineage(page.lineage), + mode: input.mode, + content: page.content, + truncated: page.has_more, + next_cursor, + has_more: page.has_more, + }, + ), + }; + match subjektiv_session_response_fits_model_budget(&response) { + Ok(true) => return response, Err(error) => { return subjektiv_session_error( server_api::SubjektivSessionDiagnosticCode::SubjectUnavailable, - error.error.to_string(), + format!("Session read response could not be serialized: {error}"), + ); + } + Ok(_) if max_content_bytes > 4 => { + max_content_bytes = (max_content_bytes / 2).max(4); + } + Ok(_) => { + return subjektiv_session_error( + server_api::SubjektivSessionDiagnosticCode::ResourceLimit, + "One Session read response exceeds the model-visible JSON budget", ); } } - } else { - None - }; - server_api::SubjektivSessionBackendResponse::Ok { - result: server_api::SubjektivSessionBackendResult::Read( - server_api::SubjektivSessionReadResponse { - session_id, - segment_id, - entry_ref, - kind: subjektiv_session_entry_kind(page.kind), - origin: page.origin, - lineage: subjektiv_session_lineage(page.lineage), - mode: input.mode, - content: page.content, - truncated: page.has_more, - next_cursor, - has_more: page.has_more, - }, - ), } } @@ -17609,87 +17634,88 @@ fn subjektiv_session_search( state.partial = true; break; } - let (attribution, expected_generation, offset, cursor_source) = if let Some(current) = - state.current.take() - { - let page = match store.subject_session_attribution_page( - subject_id, - Some(¤t.session_id), - &state.snapshot_at, - None, - 1, - ) { - Ok(page) => page, - Err(error) => { - return subjektiv_session_error( - server_api::SubjektivSessionDiagnosticCode::SubjectUnavailable, - error.to_string(), - ); - } - }; - let Some(attribution) = page.items.into_iter().next() else { - return subjektiv_session_error( - server_api::SubjektivSessionDiagnosticCode::StaleCursor, - "Session search attribution is no longer available", - ); - }; - if attribution.attributed_at != current.attributed_at - || attribution.runtime_id != current.runtime_id - || attribution.worker_id != current.worker_id - { - return subjektiv_session_error( - server_api::SubjektivSessionDiagnosticCode::StaleCursor, - "Session search attribution changed", - ); - } - ( - attribution, - Some(current.generation), - current.offset, - Some((current.storage, current.archive_id)), - ) - } else { - let page = match store.subject_session_attribution_page( - subject_id, - state.session_id.as_deref(), - &state.snapshot_at, - match ( - state.after_attributed_at.as_deref(), - state.after_session_id.as_deref(), + let (attribution, expected_generation, runtime_scan_cursor, cursor_source) = + if let Some(current) = state.current.take() { + let page = match store.subject_session_attribution_page( + subject_id, + Some(¤t.session_id), + &state.snapshot_at, + None, + 1, ) { - (Some(attributed_at), Some(session_id)) => Some((attributed_at, session_id)), - (None, None) => None, - _ => { + Ok(page) => page, + Err(error) => { return subjektiv_session_error( - server_api::SubjektivSessionDiagnosticCode::StaleCursor, - "Session search cursor position is incomplete", + server_api::SubjektivSessionDiagnosticCode::SubjectUnavailable, + error.to_string(), ); } - }, - 1, - ) { - Ok(page) => page, - Err(error) => { + }; + let Some(attribution) = page.items.into_iter().next() else { return subjektiv_session_error( - server_api::SubjektivSessionDiagnosticCode::SubjectUnavailable, - error.to_string(), + server_api::SubjektivSessionDiagnosticCode::StaleCursor, + "Session search attribution is no longer available", + ); + }; + if attribution.attributed_at != current.attributed_at + || attribution.runtime_id != current.runtime_id + || attribution.worker_id != current.worker_id + { + return subjektiv_session_error( + server_api::SubjektivSessionDiagnosticCode::StaleCursor, + "Session search attribution changed", ); } + ( + attribution, + Some(current.generation), + Some(current.runtime_scan_cursor), + Some((current.storage, current.archive_id)), + ) + } else { + let page = match store.subject_session_attribution_page( + subject_id, + state.session_id.as_deref(), + &state.snapshot_at, + match ( + state.after_attributed_at.as_deref(), + state.after_session_id.as_deref(), + ) { + (Some(attributed_at), Some(session_id)) => { + Some((attributed_at, session_id)) + } + (None, None) => None, + _ => { + return subjektiv_session_error( + server_api::SubjektivSessionDiagnosticCode::StaleCursor, + "Session search cursor position is incomplete", + ); + } + }, + 1, + ) { + Ok(page) => page, + Err(error) => { + return subjektiv_session_error( + server_api::SubjektivSessionDiagnosticCode::SubjectUnavailable, + error.to_string(), + ); + } + }; + let Some(attribution) = page.items.into_iter().next() else { + if state.session_id.is_some() { + return subjektiv_session_error( + server_api::SubjektivSessionDiagnosticCode::NotFoundOrNotAuthorized, + "Session was not found or is not authorized", + ); + } + break; + }; + has_more = page.has_more; + state.after_attributed_at = Some(attribution.attributed_at.clone()); + state.after_session_id = Some(attribution.session_id.clone()); + (attribution, None, None, None) }; - let Some(attribution) = page.items.into_iter().next() else { - if state.session_id.is_some() { - return subjektiv_session_error( - server_api::SubjektivSessionDiagnosticCode::NotFoundOrNotAuthorized, - "Session was not found or is not authorized", - ); - } - break; - }; - has_more = page.has_more; - state.after_attributed_at = Some(attribution.attributed_at.clone()); - state.after_session_id = Some(attribution.session_id.clone()); - (attribution, None, 0, None) - }; scanned_sessions += 1; let source = match authorized_subjektiv_session( api, @@ -17744,7 +17770,7 @@ fn subjektiv_session_search( kind: input.kind.map(runtime_session_entry_kind), tool_name: tool_name.clone(), tool_part: runtime_session_tool_part(input.tool_part), - offset, + scan_cursor: runtime_scan_cursor, limit: limit.saturating_sub(items.len()), max_scan_bytes: remaining_scan_bytes, max_segments: remaining_segments, @@ -17806,7 +17832,9 @@ fn subjektiv_session_search( let generation = page.generation.clone(); let mut consumed = 0usize; for item in page.items { - let (snippet, snippet_truncated) = truncate_subjektiv_session_snippet(&item.compact); + let item_scan_cursor = item.scan_cursor; + let (snippet, locally_truncated) = truncate_subjektiv_session_snippet(&item.compact); + let snippet_truncated = item.compact_truncated || locally_truncated; let projected = server_api::SubjektivSessionSearchItem { session_id: attribution.session_id.clone(), segment_id: item.segment_id, @@ -17841,7 +17869,7 @@ fn subjektiv_session_search( storage: source.storage, archive_id: source.archive_id.clone(), generation: generation.clone(), - offset: offset.saturating_add(consumed), + runtime_scan_cursor: item_scan_cursor, }); break; } @@ -17852,9 +17880,12 @@ fn subjektiv_session_search( break; } if page.has_more { - let next_offset = page - .next_offset - .unwrap_or_else(|| offset.saturating_add(consumed)); + let Some(next_scan_cursor) = page.next_scan_cursor else { + return subjektiv_session_error( + server_api::SubjektivSessionDiagnosticCode::CorruptLog, + "Runtime Session search response omitted its continuation position", + ); + }; has_more = true; state.current = Some(SubjektivSessionSearchPosition { attributed_at: attribution.attributed_at, @@ -17864,7 +17895,7 @@ fn subjektiv_session_search( storage: source.storage, archive_id: source.archive_id, generation, - offset: next_offset, + runtime_scan_cursor: next_scan_cursor, }); break; } @@ -17873,13 +17904,18 @@ fn subjektiv_session_search( break; } if items.len() >= limit { - has_more = true; break; } } - if scanned_sessions >= SUBJEKTIV_SESSION_MAX_SCANNED_SESSIONS && state.session_id.is_none() { - has_more = true; + let (has_more, scan_budget_partial) = subjektiv_session_search_completion( + has_more, + items.len(), + limit, + scanned_sessions, + state.session_id.is_none(), + ); + if scan_budget_partial { state.partial = true; } if !issues.is_empty() && successful_sources == 0 && !has_more { @@ -17916,6 +17952,20 @@ fn subjektiv_session_search( } } +fn subjektiv_session_search_completion( + has_unscanned: bool, + item_count: usize, + limit: usize, + scanned_sessions: usize, + cross_session: bool, +) -> (bool, bool) { + debug_assert!(item_count <= limit); + let scan_budget_partial = cross_session + && scanned_sessions >= SUBJEKTIV_SESSION_MAX_SCANNED_SESSIONS + && has_unscanned; + (has_unscanned, scan_budget_partial) +} + fn subjektiv_session_limit( limit: Option, ) -> std::result::Result { @@ -17983,6 +18033,13 @@ fn subjektiv_session_error( } } +fn subjektiv_session_response_fits_model_budget( + response: &server_api::SubjektivSessionBackendResponse, +) -> std::result::Result { + serde_json::to_vec(response) + .map(|serialized| serialized.len() <= server_api::SUBJEKTIV_SESSION_MAX_TOOL_CONTENT_BYTES) +} + fn bounded_subjektiv_session_message(message: &str) -> String { truncate_subjektiv_session_utf8(message, 512) } @@ -37353,6 +37410,174 @@ mod tests { )); } + #[test] + fn subjektiv_session_search_exact_boundaries_are_complete() { + assert_eq!( + subjektiv_session_search_completion( + false, + 20, + 20, + SUBJEKTIV_SESSION_MAX_SCANNED_SESSIONS, + true, + ), + (false, false), + "exactly limit results at the end must not invent a continuation" + ); + assert_eq!( + subjektiv_session_search_completion( + false, + 0, + 20, + SUBJEKTIV_SESSION_MAX_SCANNED_SESSIONS, + true, + ), + (false, false), + "exactly 20 exhausted Sessions must remain complete" + ); + assert_eq!( + subjektiv_session_search_completion( + true, + 0, + 20, + SUBJEKTIV_SESSION_MAX_SCANNED_SESSIONS, + true, + ), + (true, true), + "a real unscanned twenty-first Session must remain partial" + ); + } + + #[tokio::test] + async fn subjektiv_session_list_exactly_twenty_sessions_has_no_false_continuation() { + let workspace = tempfile::tempdir().unwrap(); + init_clean_git_workspace(workspace.path()); + let api = test_api(workspace.path()).await; + let Json(created) = create_workspace_worker( + State(api.clone()), + HeaderMap::new(), + Json(CreateWorkspaceWorkerRequest { + runtime_id: EMBEDDED_WORKER_RUNTIME_ID.to_string(), + display_name: "Session boundary Worker".to_string(), + singleton_key: None, + profile: None, + ticket_assignment: None, + initial_submit: Vec::new(), + workdir_attachments: Vec::new(), + control_operation_id: None, + }), + ) + .await + .unwrap(); + let controller = RuntimeWorkerRef::new(&created.runtime_id, &created.worker_id); + let store = open_subjektiv_store(&api).unwrap(); + let subject = store + .create_subject(crate::subjektiv::SubjectRole::new("companion").unwrap()) + .unwrap(); + for index in 0..20 { + let mut attribution = crate::subjektiv::SubjectSessionAttribution::new( + &subject.id, + &controller.runtime_id, + &controller.worker_id, + format!("session-{index:02}"), + ) + .unwrap(); + attribution.attributed_at = format!("2026-10-02T00:00:{index:02}Z"); + store.record_session_attribution(attribution).unwrap(); + } + api.store + .create_worker_control_grant(&WorkerControlGrantRecord { + workspace_id: TEST_WORKSPACE_ID.to_string(), + grant_id: "session-boundary-observe".to_string(), + controller: controller.clone(), + subject: controller.clone(), + relation: "subject_session".to_string(), + origin: "test".to_string(), + permissions: vec!["observe".to_string()], + operation_id: "session-boundary-observe".to_string(), + created_at: Utc::now().to_rfc3339(), + revoked_at: None, + }) + .unwrap(); + + let response = subjektiv_session_list( + &api, + &store, + TEST_WORKSPACE_ID, + &subject.id, + &controller, + server_api::SubjektivSessionListRequest { + session_id: None, + storage: server_api::SubjektivSessionStorageFilter::All, + limit: Some(20), + cursor: None, + }, + ); + assert!(matches!( + response, + server_api::SubjektivSessionBackendResponse::Ok { + result: server_api::SubjektivSessionBackendResult::List( + server_api::SubjektivSessionListResponse { + ref items, + next_cursor: None, + has_more: false, + } + ) + } if items.len() == 20 + )); + } + + #[test] + fn subjektiv_session_read_budget_uses_serialized_json_and_preserves_pages() { + fn response( + content: String, + has_more: bool, + ) -> server_api::SubjektivSessionBackendResponse { + server_api::SubjektivSessionBackendResponse::Ok { + result: server_api::SubjektivSessionBackendResult::Read( + server_api::SubjektivSessionReadResponse { + session_id: "session-a".into(), + segment_id: "segment-a".into(), + entry_ref: "Eentry-a".into(), + kind: server_api::SubjektivSessionEntryKind::User, + origin: protocol::SessionEntryProvenance::HumanInput, + lineage: server_api::SubjektivSessionLineage { + kind: server_api::SubjektivSessionLineageKind::Root, + parent_segment_id: None, + at_turn_index: None, + }, + mode: server_api::SubjektivSessionReadMode::Full, + content, + truncated: has_more, + next_cursor: has_more.then(|| "opaque-cursor".into()), + has_more, + }, + ), + } + } + + let content = "\0".repeat(server_api::SUBJEKTIV_SESSION_MAX_READ_CONTENT_BYTES); + assert!( + !subjektiv_session_response_fits_model_budget(&response(content.clone(), true)) + .unwrap(), + "the raw 16 KiB limit is not safe for JSON-escaped content" + ); + let split = content.len() / 2; + let first = response(content[..split].to_string(), true); + let second = response(content[split..].to_string(), false); + assert!(subjektiv_session_response_fits_model_budget(&first).unwrap()); + assert!(subjektiv_session_response_fits_model_budget(&second).unwrap()); + let chunks = [first, second] + .into_iter() + .map(|response| match response { + server_api::SubjektivSessionBackendResponse::Ok { + result: server_api::SubjektivSessionBackendResult::Read(read), + } => read.content, + _ => unreachable!(), + }) + .collect::(); + assert_eq!(chunks, content); + } + #[tokio::test] async fn subjektiv_subject_worker_and_staging_derive_scope_from_singleton_authority() { let workspace = tempfile::tempdir().unwrap(); diff --git a/crates/workspace-server/src/store.rs b/crates/workspace-server/src/store.rs index 5cb03fd1b..d2e4b76f9 100644 --- a/crates/workspace-server/src/store.rs +++ b/crates/workspace-server/src/store.rs @@ -7451,14 +7451,22 @@ impl ControlPlaneStore for SqliteWorkspaceStore { grant_id: &str, revoked_at: &str, ) -> Result { - self.with_conn(|conn| { - let changed = conn.execute( + self.with_conn_mut(|conn| { + let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?; + let active_changed = tx.execute( r#"UPDATE worker_control_grants SET revoked_at = ?3 WHERE workspace_id = ?1 AND grant_id = ?2 AND revoked_at IS NULL"#, params![workspace_id, grant_id, revoked_at], )?; - Ok(changed > 0) + let archived_changed = tx.execute( + r#"UPDATE worker_session_archive_observe_grants + SET revoked_at = ?3 + WHERE workspace_id = ?1 AND source_grant_id = ?2 AND revoked_at IS NULL"#, + params![workspace_id, grant_id, revoked_at], + )?; + tx.commit()?; + Ok(active_changed > 0 || archived_changed > 0) }) } @@ -14173,6 +14181,7 @@ fn migrate_archive_observe_grants_v75_to_v76(conn: &Connection) -> Result<()> { subject_worker_id TEXT NOT NULL, source_grant_id TEXT NOT NULL, granted_at TEXT NOT NULL, + revoked_at TEXT, PRIMARY KEY ( workspace_id, archive_id, controller_runtime_id, controller_worker_id,