From ec7a6f59c827d92b0039c759fe1033429b22a526 Mon Sep 17 00:00:00 2001 From: Hare Date: Fri, 2 Oct 2026 18:00:30 +0900 Subject: [PATCH] fix: complete subject session paging --- Cargo.lock | 1 + crates/session-store/src/lib.rs | 9 +- crates/session-store/src/public_index.rs | 123 ++++ crates/worker-runtime/src/session_public.rs | 132 ++--- crates/workspace-server/Cargo.toml | 1 + crates/workspace-server/src/server.rs | 589 +++++++++++++++++--- 6 files changed, 700 insertions(+), 155 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 41fb46abc..1afc18527 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6403,6 +6403,7 @@ dependencies = [ "serde_json", "serde_yaml", "server-api", + "session-store", "sha2 0.11.0", "ssh-key", "tempfile", diff --git a/crates/session-store/src/lib.rs b/crates/session-store/src/lib.rs index 0a16e0451..ac44c0be4 100644 --- a/crates/session-store/src/lib.rs +++ b/crates/session-store/src/lib.rs @@ -100,10 +100,11 @@ pub use public_index::{ DEFAULT_SESSION_PUBLIC_INDEX_MAX_BYTES, DEFAULT_SESSION_PUBLIC_INDEX_MAX_ENTRIES, DEFAULT_SESSION_PUBLIC_INDEX_MAX_SEGMENTS, SESSION_PUBLIC_INDEX_COMPACT_TEXT_MAX_BYTES, SessionPublicIndex, SessionPublicIndexEntry, SessionPublicIndexEntryKind, - SessionPublicIndexLimits, SessionPublicIndexLineage, SessionPublicIndexOriginKind, - SessionPublicIndexPage, SessionPublicIndexPageEntry, SessionPublicIndexReadError, - SessionPublicIndexScanPosition, SessionPublicIndexSegment, SessionPublicIndexToolPart, - read_session_public_index, read_session_public_index_page, + SessionPublicIndexEntryRead, SessionPublicIndexLimits, SessionPublicIndexLineage, + SessionPublicIndexOriginKind, SessionPublicIndexPage, SessionPublicIndexPageEntry, + SessionPublicIndexReadError, SessionPublicIndexScanPosition, SessionPublicIndexSegment, + SessionPublicIndexToolPart, read_session_public_index, read_session_public_index_entry, + 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 b5176e1df..520065b25 100644 --- a/crates/session-store/src/public_index.rs +++ b/crates/session-store/src/public_index.rs @@ -95,6 +95,16 @@ pub struct SessionPublicIndexPage { pub scanned_entries: usize, } +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SessionPublicIndexEntryRead { + pub session_id: String, + pub generation: String, + pub item: SessionPublicIndexPageEntry, + pub scanned_bytes: u64, + pub scanned_segments: usize, + pub scanned_entries: usize, +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct SessionPublicIndex { pub session_id: String, @@ -558,6 +568,76 @@ where }) } +/// Read one exact public entry from its selected Segment without charging +/// unrelated Segments against the caller's aggregate scan budget. The Session +/// generation still covers every persisted Segment so content-page cursors are +/// invalidated by any concurrent Session mutation. +pub fn read_session_public_index_entry( + session_root: &Path, + limits: SessionPublicIndexLimits, + expected_generation: Option<&str>, + segment_id: &str, + entry_ref: &str, +) -> Result, SessionPublicIndexReadError> { + if limits.max_bytes == 0 || limits.max_segments == 0 || limits.max_entries == 0 { + return Err(SessionPublicIndexReadError::ResourceLimit); + } + let selected_segment = match uuid::Uuid::parse_str(segment_id) { + Ok(segment_id) => segment_id, + Err(_) => return Ok(None), + }; + 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 Some(segment_index) = segment_ids + .iter() + .position(|candidate| *candidate == selected_segment) + else { + return Ok(None); + }; + let page = read_session_public_index_page( + session_root, + SessionPublicIndexLimits { + max_segments: 1, + ..limits + }, + expected_generation, + Some(SessionPublicIndexScanPosition { + segment_index, + byte_offset: 0, + entry_offset: 0, + lineage: None, + }), + 2, + |entry| entry.entry_ref == entry_ref, + )?; + let mut matches = page.entries.into_iter(); + let item = matches.next(); + if matches.next().is_some() { + return Err(SessionPublicIndexReadError::Corrupt); + } + if page + .next_position + .as_ref() + .is_some_and(|position| position.segment_index == segment_index) + { + return Err(SessionPublicIndexReadError::ResourceLimit); + } + Ok(item.map(|item| SessionPublicIndexEntryRead { + session_id: page.session_id, + generation: page.generation, + item, + scanned_bytes: page.scanned_bytes, + scanned_segments: page.scanned_segments, + scanned_entries: page.scanned_entries, + })) +} + #[derive(Debug)] struct PublicLogUnit { entries: Vec, @@ -1687,6 +1767,49 @@ mod tests { } } + #[test] + fn incremental_scan_exact_aggregate_budgets_finish_without_cursor() { + let (_temp, root, store) = create_store(); + let session_id = crate::new_session_id(); + let first_segment = segment_id(1_600); + let second_segment = segment_id(1_601); + for (segment_id, id) in [(first_segment, "entry-a"), (second_segment, "entry-b")] { + store + .create_segment( + session_id, + segment_id, + &[start( + session_id, + vec![message(id, LoggedRole::User, "ordinary")], + None, + None, + )], + ) + .unwrap(); + } + let total_bytes = store.segment_log_len(first_segment).unwrap() + + store.segment_log_len(second_segment).unwrap(); + for limits in [ + SessionPublicIndexLimits { + max_bytes: total_bytes, + ..SessionPublicIndexLimits::default() + }, + SessionPublicIndexLimits { + max_segments: 2, + ..SessionPublicIndexLimits::default() + }, + SessionPublicIndexLimits { + max_entries: 4, + ..SessionPublicIndexLimits::default() + }, + ] { + let page = + read_session_public_index_page(&root, limits, None, None, 10, |_| false).unwrap(); + assert!(!page.has_more, "exact final-range budget must be complete"); + assert!(page.next_position.is_none()); + } + } + #[test] fn incremental_scan_rejects_one_indivisible_record_over_the_byte_budget() { let (_temp, root, store) = create_store(); diff --git a/crates/worker-runtime/src/session_public.rs b/crates/worker-runtime/src/session_public.rs index c8a9c486d..48906b000 100644 --- a/crates/worker-runtime/src/session_public.rs +++ b/crates/worker-runtime/src/session_public.rs @@ -10,10 +10,10 @@ use runtime_api::{ }; use serde::{Deserialize, Serialize}; use session_store::{ - SessionPublicIndex, SessionPublicIndexEntry, SessionPublicIndexEntryKind, - SessionPublicIndexLimits, SessionPublicIndexLineage, SessionPublicIndexOriginKind, - SessionPublicIndexReadError, SessionPublicIndexScanPosition, SessionPublicIndexToolPart, - read_session_public_index, read_session_public_index_page, + SessionPublicIndexEntry, SessionPublicIndexEntryKind, SessionPublicIndexLimits, + SessionPublicIndexLineage, SessionPublicIndexOriginKind, SessionPublicIndexReadError, + SessionPublicIndexScanPosition, SessionPublicIndexToolPart, read_session_public_index_entry, + read_session_public_index_page, }; use crate::retention::WorkerSessionArchiveManifest as DomainArchiveManifest; @@ -171,57 +171,37 @@ pub(crate) fn read( "Session public read content limit must be positive", ); } - let index = match load_index( + let selected = match read_session_public_index_entry( session_root, - request.max_scan_bytes, - request.max_segments, - request.max_entries, + SessionPublicIndexLimits { + max_bytes: request.max_scan_bytes, + max_segments: request.max_segments, + max_entries: request.max_entries, + }, + request.expected_generation.as_deref(), + &request.segment_id, + &request.entry_ref, ) { - Ok(index) => index, - Err((reason, message)) => return read_unavailable(reason, message), + Ok(Some(selected)) => selected, + Ok(None) => { + return read_unavailable( + SessionPublicUnavailableReason::NotFound, + "Session entry was not found in the authorized source", + ); + } + Err(error) => { + let (reason, message) = map_index_error(error); + return read_unavailable(reason, message); + } }; - if index.session_id != request.expected_session_id { + if selected.session_id != request.expected_session_id { return read_unavailable( SessionPublicUnavailableReason::NotFound, "Session identity does not match the authorized source", ); } - if request - .expected_generation - .as_deref() - .is_some_and(|expected| expected != index.generation) - { - return read_unavailable( - SessionPublicUnavailableReason::InvalidCursor, - "Session public read generation changed", - ); - } - let Some(segment) = index - .segments - .iter() - .find(|segment| segment.segment_id == request.segment_id) - else { - return read_unavailable( - SessionPublicUnavailableReason::NotFound, - "Session entry was not found in the authorized source", - ); - }; - let mut entries = segment - .entries - .iter() - .filter(|entry| entry.entry_ref == request.entry_ref); - let Some(entry) = entries.next() else { - return read_unavailable( - SessionPublicUnavailableReason::NotFound, - "Session entry was not found in the authorized source", - ); - }; - if entries.next().is_some() { - return read_unavailable( - SessionPublicUnavailableReason::CorruptLog, - "Session entry identity is ambiguous in its Segment", - ); - } + let item = selected.item; + let entry = item.entry; let content = match request.mode { SessionPublicReadMode::Compact => &entry.compact_text, SessionPublicReadMode::Full => &entry.full_text, @@ -248,42 +228,25 @@ pub(crate) fn read( let has_more = end < content.len(); SessionPublicReadAvailability::Page { page: SessionPublicReadPage { - session_id: index.session_id, - generation: index.generation, - segment_id: segment.segment_id.clone(), + session_id: selected.session_id, + generation: selected.generation, + segment_id: item.segment_id, entry_ref: entry.entry_ref.clone(), kind: entry_kind(entry.kind), origin: entry.provenance.clone(), - lineage: lineage(&segment.lineage), + lineage: lineage(&item.lineage), mode: request.mode, content: content[request.byte_offset..end].to_string(), next_byte_offset: has_more.then_some(end), has_more, - scanned_bytes: index.scanned_bytes, - scanned_segments: index.segments.len(), - scanned_entries: index.scanned_entries, + scanned_bytes: selected.scanned_bytes, + scanned_segments: selected.scanned_segments, + scanned_entries: selected.scanned_entries, archive_manifest: archive_manifest.map(api_archive_manifest), }, } } -fn load_index( - session_root: &Path, - max_bytes: u64, - max_segments: usize, - max_entries: usize, -) -> Result { - read_session_public_index( - session_root, - SessionPublicIndexLimits { - max_bytes, - max_segments, - max_entries, - }, - ) - .map_err(map_index_error) -} - fn matches_entry( entry: &SessionPublicIndexEntry, request: &SessionPublicSearchRequest, @@ -522,6 +485,33 @@ mod tests { assert_eq!(second.items[0].compact, "needle"); assert!(!second.has_more); + let read_page = match read( + &root, + &SessionPublicReadRequest { + workspace_id: "workspace-1".into(), + source: runtime_api::SessionPublicSource::Retained { + worker_id: crate::identity::WorkerId::now_v7(), + }, + expected_session_id: session_id.to_string(), + expected_generation: Some(second.generation.clone()), + segment_id: last_segment.unwrap().to_string(), + entry_ref: "Eentry-64".into(), + mode: SessionPublicReadMode::Full, + byte_offset: 0, + max_content_bytes: 16 * 1024, + max_scan_bytes: 64 * 1024 * 1024, + max_segments: 64, + max_entries: 100_000, + }, + None, + ) { + SessionPublicReadAvailability::Page { page } => page, + other => panic!("unexpected direct read result: {other:?}"), + }; + assert_eq!(read_page.content, "needle"); + assert_eq!(read_page.segment_id, last_segment.unwrap().to_string()); + assert!(!read_page.has_more); + store .append( session_id, diff --git a/crates/workspace-server/Cargo.toml b/crates/workspace-server/Cargo.toml index 5c78de89c..94d5c0bb7 100644 --- a/crates/workspace-server/Cargo.toml +++ b/crates/workspace-server/Cargo.toml @@ -51,6 +51,7 @@ webauthn-rs = { workspace = true } [dev-dependencies] agen.workspace = true +session-store.workspace = true tempfile.workspace = true tools.workspace = true yoi = { path = "../yoi" } diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 4347072a0..e8e3c16ee 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -17034,6 +17034,7 @@ struct SubjektivSessionSearchPosition { archive_id: Option, generation: String, runtime_scan_cursor: String, + attribution_has_more: bool, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -17630,92 +17631,98 @@ fn subjektiv_session_search( while items.len() < limit && scanned_sessions < SUBJEKTIV_SESSION_MAX_SCANNED_SESSIONS { if remaining_scan_bytes == 0 || remaining_segments == 0 || remaining_entries == 0 { - has_more = true; - state.partial = true; + if has_more || state.current.is_some() { + state.partial = true; + } break; } - 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, - ) { - 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 { + let ( + attribution, + expected_generation, + runtime_scan_cursor, + cursor_source, + attribution_has_more, + ) = 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::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", + server_api::SubjektivSessionDiagnosticCode::SubjectUnavailable, + error.to_string(), ); } - ( - 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 { + 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), + Some(current.runtime_scan_cursor), + Some((current.storage, current.archive_id)), + current.attribution_has_more, + ) + } 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; + }; + let attribution_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, attribution_has_more) + }; + has_more = attribution_has_more; scanned_sessions += 1; let source = match authorized_subjektiv_session( api, @@ -17870,6 +17877,7 @@ fn subjektiv_session_search( archive_id: source.archive_id.clone(), generation: generation.clone(), runtime_scan_cursor: item_scan_cursor, + attribution_has_more, }); break; } @@ -17896,6 +17904,7 @@ fn subjektiv_session_search( archive_id: source.archive_id, generation, runtime_scan_cursor: next_scan_cursor, + attribution_has_more, }); break; } @@ -37526,6 +37535,426 @@ mod tests { )); } + #[tokio::test] + async fn subjektiv_session_search_preserves_later_session_after_runtime_resume() { + use session_store::{ + LogEntry, LoggedContentPart, LoggedHistoryEntry, LoggedItem, LoggedRole, + LoggedSessionHistoryEntryId, LoggedSessionHistoryMetadata, LoggedSessionHistoryOrigin, + Store as _, WorkerSessionStore, + }; + + fn public_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, + }, + } + } + + let workspace = tempfile::tempdir().unwrap(); + init_clean_git_workspace(workspace.path()); + let api = test_api(workspace.path()).await; + let create = |display_name: &str| CreateWorkspaceWorkerRequest { + runtime_id: EMBEDDED_WORKER_RUNTIME_ID.to_string(), + display_name: display_name.to_string(), + singleton_key: None, + profile: None, + ticket_assignment: None, + initial_submit: Vec::new(), + workdir_attachments: Vec::new(), + control_operation_id: None, + }; + let Json(controller) = create_workspace_worker( + State(api.clone()), + HeaderMap::new(), + Json(create("Session search controller")), + ) + .await + .unwrap(); + let Json(first_worker) = create_workspace_worker( + State(api.clone()), + HeaderMap::new(), + Json(create("Large Session source")), + ) + .await + .unwrap(); + let Json(second_worker) = create_workspace_worker( + State(api.clone()), + HeaderMap::new(), + Json(create("Later Session source")), + ) + .await + .unwrap(); + let controller = RuntimeWorkerRef::new(&controller.runtime_id, &controller.worker_id); + let first_worker = RuntimeWorkerRef::new(&first_worker.runtime_id, &first_worker.worker_id); + let second_worker = + RuntimeWorkerRef::new(&second_worker.runtime_id, &second_worker.worker_id); + + let first_session = session_store::new_session_id(); + let first_root = workspace + .path() + .join(".test-embedded-runtime-store/workers") + .join(&first_worker.worker_id) + .join("session"); + let first_store = WorkerSessionStore::new(&first_root).unwrap(); + for index in 0..65_u128 { + first_store + .create_segment( + first_session, + uuid::Uuid::from_u128(20_000 + index), + &[LogEntry::AnnotatedSegmentStart { + ts: 1, + session_id: first_session, + system_prompt: None, + config: agen::llm_client::RequestConfig::default(), + history: vec![public_message( + &format!("first-{index}"), + if index == 64 { "needle" } else { "ordinary" }, + )], + forked_from: None, + compacted_from: None, + }], + ) + .unwrap(); + } + let second_session = session_store::new_session_id(); + let second_root = workspace + .path() + .join(".test-embedded-runtime-store/workers") + .join(&second_worker.worker_id) + .join("session"); + WorkerSessionStore::new(&second_root) + .unwrap() + .create_segment( + second_session, + uuid::Uuid::from_u128(30_000), + &[LogEntry::AnnotatedSegmentStart { + ts: 1, + session_id: second_session, + system_prompt: None, + config: agen::llm_client::RequestConfig::default(), + history: vec![public_message("second", "needle in later Session")], + forked_from: None, + compacted_from: None, + }], + ) + .unwrap(); + + let store = open_subjektiv_store(&api).unwrap(); + let subject = store + .create_subject(crate::subjektiv::SubjectRole::new("companion").unwrap()) + .unwrap(); + for (worker, session_id, attributed_at) in [ + (&first_worker, first_session, "2026-10-02T00:00:02Z"), + (&second_worker, second_session, "2026-10-02T00:00:01Z"), + ] { + let mut attribution = crate::subjektiv::SubjectSessionAttribution::new( + &subject.id, + &worker.runtime_id, + &worker.worker_id, + session_id.to_string(), + ) + .unwrap(); + attribution.attributed_at = attributed_at.to_string(); + store.record_session_attribution(attribution).unwrap(); + api.store + .create_worker_control_grant(&WorkerControlGrantRecord { + workspace_id: TEST_WORKSPACE_ID.to_string(), + grant_id: format!("search-observe-{}", worker.worker_id), + controller: controller.clone(), + subject: worker.clone(), + relation: "subject_session".to_string(), + origin: "test".to_string(), + permissions: vec!["observe".to_string()], + operation_id: format!("search-observe-{}", worker.worker_id), + created_at: Utc::now().to_rfc3339(), + revoked_at: None, + }) + .unwrap(); + } + + let search = |cursor| { + subjektiv_session_search( + &api, + &store, + TEST_WORKSPACE_ID, + &subject.id, + &controller, + server_api::SubjektivSessionSearchRequest { + query: Some("needle".to_string()), + session_id: None, + kind: None, + tool_name: None, + tool_part: server_api::SubjektivSessionToolPart::Both, + limit: Some(1), + cursor, + }, + ) + }; + let page = |response| match response { + server_api::SubjektivSessionBackendResponse::Ok { + result: server_api::SubjektivSessionBackendResult::Search(page), + } => page, + other => panic!("unexpected Session search response: {other:?}"), + }; + + let first = page(search(None)); + assert!(first.items.is_empty()); + assert!(first.has_more); + let second = page(search(first.next_cursor)); + assert_eq!(second.items.len(), 1); + assert_eq!(second.items[0].session_id, first_session.to_string()); + assert!( + second.has_more, + "the outer attribution continuation must survive a resumed Runtime page" + ); + let third = page(search(second.next_cursor)); + assert_eq!(third.items.len(), 1); + assert_eq!(third.items[0].session_id, second_session.to_string()); + assert!(!third.has_more); + assert!(third.next_cursor.is_none()); + } + + #[tokio::test] + async fn archived_session_search_and_read_reauthorize_after_grant_revocation() { + use session_store::{ + LogEntry, LoggedContentPart, LoggedHistoryEntry, LoggedItem, LoggedRole, + LoggedSessionHistoryEntryId, LoggedSessionHistoryMetadata, LoggedSessionHistoryOrigin, + Store as _, WorkerSessionStore, + }; + + fn public_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, + }, + } + } + + let workspace = tempfile::tempdir().unwrap(); + init_clean_git_workspace(workspace.path()); + let api = test_api(workspace.path()).await; + let Json(orchestrator) = scoped_start_workspace_orchestrator( + State(api.clone()), + AxumPath(ScopedWorkspacePath { + workspace_id: TEST_WORKSPACE_ID.to_string(), + }), + ) + .await + .unwrap(); + let controller = orchestrator.worker.unwrap(); + let controller = RuntimeWorkerRef::new(&controller.runtime_id, &controller.worker_id); + let spawned = api + .spawn_workspace_worker( + EMBEDDED_WORKER_RUNTIME_ID, + WorkerSpawnRequest { + intent: WorkerSpawnIntent::WorkspaceCompanion, + requested_worker_name: Some("archived Session source".to_string()), + singleton_key: None, + acceptance: WorkerSpawnAcceptanceRequirement::RunAccepted { + expected_segments: 0, + }, + profile: worker_runtime::catalog::ProfileSelector::Builtin( + "builtin:companion".to_string(), + ), + ticket_assignment: None, + initial_submit: Vec::new(), + workdir_attachment_requests: Vec::new(), + resolved_workdir_attachment_requests: Vec::new(), + resolved_workdir_attachments: Vec::new(), + resolved_config_bundle: Some(runtime_test_bundle()), + resolved_worker_observation_enabled: false, + resolved_worker_observation_grants: Vec::new(), + resolved_workspace_api: None, + resolved_memory_settings: None, + resolved_control_operation: None, + }, + ) + .unwrap(); + let target = spawned.worker.unwrap().worker; + let stopped = api + .runtime + .stop_worker( + &target, + WorkerLifecycleRequest { + reason: Some("prepare archive revocation regression".to_string()), + ticket_assignment: None, + }, + ) + .unwrap(); + assert_eq!(stopped.state, InternalWorkerOperationState::Accepted); + + let session_id = session_store::new_session_id(); + let segment_id = uuid::Uuid::from_u128(40_000); + let session_root = workspace + .path() + .join(".test-embedded-runtime-store/workers") + .join(&target.worker_id) + .join("session"); + WorkerSessionStore::new(&session_root) + .unwrap() + .create_segment( + session_id, + segment_id, + &[LogEntry::AnnotatedSegmentStart { + ts: 1, + session_id, + system_prompt: None, + config: agen::llm_client::RequestConfig::default(), + history: vec![ + public_message("archive-a", "archived first entry"), + public_message("archive-b", "archived second entry"), + ], + forked_from: None, + compacted_from: None, + }], + ) + .unwrap(); + let subjektiv_store = open_subjektiv_store(&api).unwrap(); + let subject = subjektiv_store + .create_subject(crate::subjektiv::SubjectRole::new("companion").unwrap()) + .unwrap(); + subjektiv_store + .record_session_attribution( + crate::subjektiv::SubjectSessionAttribution::new( + &subject.id, + &target.runtime_id, + &target.worker_id, + session_id.to_string(), + ) + .unwrap(), + ) + .unwrap(); + let grant_id = "archive-revocable-observe"; + api.store + .create_worker_control_grant(&WorkerControlGrantRecord { + workspace_id: TEST_WORKSPACE_ID.to_string(), + grant_id: grant_id.to_string(), + controller: controller.clone(), + subject: target.clone(), + relation: "subject_session".to_string(), + origin: "test".to_string(), + permissions: vec!["observe".to_string(), "remove".to_string()], + operation_id: "archive-revocation-proof".to_string(), + created_at: Utc::now().to_rfc3339(), + revoked_at: None, + }) + .unwrap(); + let summary = api.runtime.worker(&target).unwrap(); + sync_worker_observation(&api, &summary).unwrap(); + let removal = WorkerRemovalService::new(&api) + .execute_async( + crate::worker_source::VerifiedWorkerMutationSource { + runtime_id: controller.runtime_id.clone(), + worker_id: controller.worker_id.clone(), + actor_kind: worker_runtime::auth::WorkerMutationActorKind::Worker, + permission: worker_runtime::auth::WORKER_REMOVE_PERMISSION.to_string(), + jti: "archive-revocation-proof".to_string(), + }, + &target.runtime_id, + &target.worker_id, + "archive Session for revocation regression", + ) + .await + .unwrap(); + assert_eq!(removal.status, StatusCode::OK.as_u16(), "{}", removal.body); + assert!( + api.config_store + .worker_session_archive(TEST_WORKSPACE_ID, &target, &session_id.to_string()) + .unwrap() + .is_some(), + "removal must commit the archive catalog before backend reads" + ); + + let search_request = |cursor| server_api::SubjektivSessionSearchRequest { + query: None, + session_id: Some(session_id.to_string()), + kind: None, + tool_name: None, + tool_part: server_api::SubjektivSessionToolPart::Both, + limit: Some(1), + cursor, + }; + let first = subjektiv_session_search( + &api, + &subjektiv_store, + TEST_WORKSPACE_ID, + &subject.id, + &controller, + search_request(None), + ); + let cursor = match first { + server_api::SubjektivSessionBackendResponse::Ok { + result: server_api::SubjektivSessionBackendResult::Search(page), + } => { + assert_eq!(page.items.len(), 1); + assert!(page.has_more); + page.next_cursor.unwrap() + } + other => panic!("unexpected archived Session search response: {other:?}"), + }; + + assert!( + api.store + .revoke_worker_control_grant(TEST_WORKSPACE_ID, grant_id, &Utc::now().to_rfc3339()) + .unwrap() + ); + assert!(matches!( + subjektiv_session_search( + &api, + &subjektiv_store, + TEST_WORKSPACE_ID, + &subject.id, + &controller, + search_request(Some(cursor)), + ), + server_api::SubjektivSessionBackendResponse::Error { + error: server_api::SubjektivSessionErrorResponse { + code: server_api::SubjektivSessionDiagnosticCode::NotFoundOrNotAuthorized, + .. + } + } + )); + assert!(matches!( + subjektiv_session_read( + &api, + &subjektiv_store, + TEST_WORKSPACE_ID, + &subject.id, + &controller, + server_api::SubjektivSessionReadRequest { + session_id: session_id.to_string(), + segment_id: segment_id.to_string(), + entry_ref: "Earchive-a".to_string(), + mode: server_api::SubjektivSessionReadMode::Full, + cursor: None, + }, + ), + server_api::SubjektivSessionBackendResponse::Error { + error: server_api::SubjektivSessionErrorResponse { + code: server_api::SubjektivSessionDiagnosticCode::NotFoundOrNotAuthorized, + .. + } + } + )); + } + #[test] fn subjektiv_session_read_budget_uses_serialized_json_and_preserves_pages() { fn response(