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

fix: make subject session reads resumable

Browse Source
This commit is contained in:
Hare
2026-10-02 17:36:10 +09:00
parent ce841cfee0
commit c491cf8a48
10 changed files with 1410 additions and 233 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/runtime-api/src/lib.rs
+7 -2
View File
@@ -755,7 +755,8 @@ pub struct SessionPublicSearchRequest {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub tool_name: Option<String>,
pub tool_part: SessionPublicToolPart,
pub offset: usize,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub scan_cursor: Option<String>,
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<SessionPublicToolPart>,
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<SessionPublicSearchItem>,
pub next_offset: Option<usize>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub next_scan_cursor: Option<String>,
pub has_more: bool,
pub scanned_bytes: u64,
pub scanned_segments: usize,
crates/session-store/src/lib.rs
+3 -2
View File
@@ -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,
crates/session-store/src/public_index.rs
+645 -4
View File
@@ -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<SessionPublicIndexLineage>,
}
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<SessionPublicIndexPageEntry>,
pub next_position: Option<SessionPublicIndexScanPosition>,
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<String>,
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<F>(
session_root: &Path,
limits: SessionPublicIndexLimits,
expected_generation: Option<&str>,
position: Option<SessionPublicIndexScanPosition>,
max_results: usize,
mut matches: F,
) -> Result<SessionPublicIndexPage, SessionPublicIndexReadError>
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<LogEntry>,
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<RetainedLogRecord>,
) -> Result<PublicLogUnit, SessionPublicIndexReadError> {
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<String, SessionPublicIndexReadError> {
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();
crates/session-store/src/worker_session_store.rs
+16
View File
@@ -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<u64, StoreError> {
Ok(fs::metadata(self.log_path(segment_id))?.len())
}
crates/worker-runtime/src/session_public.rs
+239 -49
View File
@@ -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,
..
}
));
}
}
crates/worker/src/feature/builtin/subjektiv_session.rs
+38
View File
@@ -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]
crates/workspace-server/src/latest_schema.sql
+1 -1
View File
@@ -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);
crates/workspace-server/src/retention.rs
+52
View File
@@ -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(
crates/workspace-server/src/server.rs
+397 -172
View File
@@ -17033,7 +17033,7 @@ struct SubjektivSessionSearchPosition {
storage: server_api::SubjektivSessionStorage,
archive_id: Option<String>,
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(&current.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(&current.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<usize>,
) -> std::result::Result<usize, server_api::SubjektivSessionBackendResponse> {
@@ -17983,6 +18033,13 @@ fn subjektiv_session_error(
}
}
fn subjektiv_session_response_fits_model_budget(
response: &server_api::SubjektivSessionBackendResponse,
) -> std::result::Result<bool, serde_json::Error> {
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::<String>();
assert_eq!(chunks, content);
}
#[tokio::test]
async fn subjektiv_subject_worker_and_staging_derive_scope_from_singleton_authority() {
let workspace = tempfile::tempdir().unwrap();
crates/workspace-server/src/store.rs
+12 -3
View File
@@ -7451,14 +7451,22 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
grant_id: &str,
revoked_at: &str,
) -> Result<bool> {
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,