diff --git a/crates/session-store/src/lib.rs b/crates/session-store/src/lib.rs index 81f071fe2..50bfc4697 100644 --- a/crates/session-store/src/lib.rs +++ b/crates/session-store/src/lib.rs @@ -114,7 +114,8 @@ pub use public_snapshot::{ RetainedHistoryReadLimits, RetainedSessionAttachment, RetainedSessionIdentity, RetainedSessionSnapshot, RetainedSnapshotReadError, project_session_snapshot, read_retained_session_attachment, read_retained_session_history_page, - read_retained_session_snapshot, session_tool_attachment_id, + read_retained_session_snapshot, read_session_attachment_from_entries, + session_tool_attachment_id, }; pub use segment::{ SegmentStartState, append_entry, append_system_item, classify_logged_history_entry, diff --git a/crates/session-store/src/public_snapshot.rs b/crates/session-store/src/public_snapshot.rs index 56905054b..bc7659c12 100644 --- a/crates/session-store/src/public_snapshot.rs +++ b/crates/session-store/src/public_snapshot.rs @@ -251,30 +251,82 @@ pub fn read_retained_session_attachment( .ok_or(RetainedAttachmentReadError::ActivePointerMissing)?; let store = WorkerSessionStore::open_read_only(aggregate_root.join("session")) .map_err(map_attachment_store_error)?; - if store - .segment_log_len(segment_id) - .map_err(map_attachment_store_error)? - > max_scan_bytes - { - return Err(RetainedAttachmentReadError::ResourceLimit); - } - let entries = store - .read_all_read_only(active.session_id, segment_id) - .map_err(map_attachment_store_error)?; - for record in &entries { - let mut found = None; - visit_record_history(record, &mut |entry| { - if found.is_none() { - found = attachment_from_entry(entry, attachment_id, max_attachment_bytes); + let limits = RetainedHistoryReadLimits { + max_scan_bytes, + ..RetainedHistoryReadLimits::default() + }; + let (lineage, _, mut scanned_bytes, mut scanned_entries) = + load_adopted_lineage(&store, active.session_id, segment_id, limits) + .map_err(map_history_attachment_error)?; + for segment in &lineage { + let adopted_end = locate_adopted_segment_end( + &store, + active.session_id, + segment, + limits, + &mut scanned_bytes, + &mut scanned_entries, + ) + .map_err(map_history_attachment_error)?; + let mut reader = store + .open_retained_segment_reader(active.session_id, segment.segment_id, adopted_end) + .map_err(map_attachment_store_error)?; + loop { + let remaining = limits.max_scan_bytes.saturating_sub(scanned_bytes); + let (record, bytes) = reader + .previous_record(remaining) + .map_err(map_attachment_store_error)?; + scanned_bytes = scanned_bytes + .checked_add(bytes) + .ok_or(RetainedAttachmentReadError::ResourceLimit)?; + let Some(record) = record else { + break; + }; + scanned_entries = scanned_entries + .checked_add(persisted_entry_units(std::slice::from_ref(&record.entry))) + .ok_or(RetainedAttachmentReadError::ResourceLimit)?; + if scanned_entries > limits.max_entries { + return Err(RetainedAttachmentReadError::ResourceLimit); } - }); - if let Some(result) = found { + if let Some(result) = + session_attachment_from_record(&record.entry, attachment_id, max_attachment_bytes) + { + return result; + } + } + } + Err(RetainedAttachmentReadError::NotFound) +} + +pub fn read_session_attachment_from_entries( + entries: &[LogEntry], + attachment_id: &str, + max_attachment_bytes: u64, +) -> Result { + for record in entries { + if let Some(result) = + session_attachment_from_record(record, attachment_id, max_attachment_bytes) + { return result; } } Err(RetainedAttachmentReadError::NotFound) } +fn session_attachment_from_record( + record: &LogEntry, + attachment_id: &str, + max_attachment_bytes: u64, +) -> Option> { + let mut found = None; + visit_record_history(record, &mut |entry| { + if found.is_none() { + found = attachment_from_entry(entry, attachment_id, max_attachment_bytes); + } + }); + found +} + fn visit_record_history(record: &LogEntry, visit: &mut impl FnMut(&LoggedHistoryEntry)) { match record { LogEntry::AnnotatedSegmentStart { history, .. } @@ -1055,6 +1107,25 @@ fn map_history_store_error(error: StoreError) -> RetainedHistoryReadError { } } +fn map_history_attachment_error(error: RetainedHistoryReadError) -> RetainedAttachmentReadError { + match error { + RetainedHistoryReadError::RetentionMissing => RetainedAttachmentReadError::RetentionMissing, + RetainedHistoryReadError::ActivePointerMissing => { + RetainedAttachmentReadError::ActivePointerMissing + } + RetainedHistoryReadError::MigrationRequired => { + RetainedAttachmentReadError::MigrationRequired + } + RetainedHistoryReadError::CorruptLog | RetainedHistoryReadError::InvalidCursor => { + RetainedAttachmentReadError::CorruptLog + } + RetainedHistoryReadError::StorageUnavailable => { + RetainedAttachmentReadError::StorageUnavailable + } + RetainedHistoryReadError::ResourceLimit => RetainedAttachmentReadError::ResourceLimit, + } +} + fn map_attachment_worker_store_error(error: WorkerStoreError) -> RetainedAttachmentReadError { match error { WorkerStoreError::Io(error) if error.kind() == std::io::ErrorKind::NotFound => { @@ -3481,6 +3552,92 @@ mod tests { )); } + #[test] + fn retained_attachment_resolves_from_an_adopted_compaction_ancestor() { + let worker_name = "worker-ancestor-image"; + let session_id = crate::new_session_id(); + let source_segment = crate::new_segment_id(); + let active_segment = crate::new_segment_id(); + let image_entry = LoggedHistoryEntry { + item: LoggedItem::ToolResult { + call_id: "call-ancestor-image".into(), + summary: "ancestor screenshot".into(), + content: None, + attachments: vec![crate::logged_item::LoggedAttachment::Image { + mime_type: "image/png".into(), + data: vec![7, 8, 9], + }], + disposition: agen::tool::ToolResultDisposition::Success, + is_error: false, + }, + metadata: LoggedSessionHistoryMetadata { + entry_id: LoggedSessionHistoryEntryId("ancestor-image".into()), + origin: LoggedSessionHistoryOrigin::LegacyUnknown, + derivation: None, + }, + }; + let mut source_log = vec![segment_start(session_id, Vec::new(), None, None)]; + source_log.push(LogEntry::AnnotatedUserInput { + ts: 2, + segments: vec![Segment::Text { + content: "show image".into(), + }], + history: vec![user_message(1)], + extensions: Vec::new(), + }); + source_log.push(LogEntry::AnnotatedToolResult { + ts: 3, + entry: image_entry, + }); + source_log.push(LogEntry::TurnEnd { + ts: 4, + turn_count: 1, + }); + let active_log = vec![segment_start( + session_id, + Vec::new(), + None, + Some(SegmentOrigin { + segment_id: source_segment, + at_turn_index: 1, + }), + )]; + let (_root, aggregate_root) = persist_history_fixture( + worker_name, + session_id, + active_segment, + vec![(source_segment, source_log), (active_segment, active_log)], + ); + + let page = read_retained_session_history_page( + &aggregate_root, + worker_name, + None, + None, + RetainedHistoryReadLimits::default(), + ) + .unwrap(); + let attachment = page.turns[0] + .entries + .iter() + .find_map(|entry| match &entry.data { + SessionSnapshotEntryData::ToolResult { attachments, .. } => attachments.first(), + _ => None, + }) + .expect("ancestor attachment reference"); + let body = read_retained_session_attachment( + &aggregate_root, + worker_name, + &session_id.to_string(), + &attachment.attachment_id, + DEFAULT_RETAINED_HISTORY_MAX_SCAN_BYTES, + 10 * 1024 * 1024, + ) + .unwrap(); + assert_eq!(body.media_type, "image/png"); + assert_eq!(body.data, vec![7, 8, 9]); + } + #[test] fn large_legacy_inline_images_project_as_bounded_refs_and_fetch_from_retained_storage() { let worker_name = "worker-images"; diff --git a/crates/worker-runtime/src/worker_backend.rs b/crates/worker-runtime/src/worker_backend.rs index 36476d139..8f9918468 100644 --- a/crates/worker-runtime/src/worker_backend.rs +++ b/crates/worker-runtime/src/worker_backend.rs @@ -2448,6 +2448,24 @@ where request: crate::execution::WorkerSessionAttachmentRequest, ) -> Result { + let live = self + .workers + .lock() + .map_err(|_| session_store::RetainedAttachmentReadError::StorageUnavailable)? + .get(&request.worker_ref) + .map(|execution| execution.handle.clone()); + if let Some(handle) = live { + match handle.session_attachment( + &request.session_id, + &request.attachment_id, + 10 * 1024 * 1024, + ) { + Ok(attachment) => return Ok(attachment), + Err(session_store::RetainedAttachmentReadError::SessionMismatch) + | Err(session_store::RetainedAttachmentReadError::NotFound) => {} + Err(error) => return Err(error), + } + } self.factory.retained_session_attachment( &request.worker_ref, &request.session_id, diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index cb8cba96d..fa886c8ed 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -129,6 +129,27 @@ impl WorkerHandle { self.sink.subscribe_with_snapshot().0 } + pub fn session_attachment( + &self, + session_id: &str, + attachment_id: &str, + max_attachment_bytes: u64, + ) -> Result + { + if session_id == self.session_id.to_string() { + return session_store::read_session_attachment_from_entries( + &self.committed_entries(), + attachment_id, + max_attachment_bytes, + ); + } + self.spawned_registry + .session_attachment(session_id, attachment_id, max_attachment_bytes) + .unwrap_or(Err( + session_store::RetainedAttachmentReadError::SessionMismatch, + )) + } + pub fn snapshot_event(&self) -> Event { self.snapshot_event_with_entry_subscription().0 } diff --git a/crates/worker/src/internal_worker.rs b/crates/worker/src/internal_worker.rs index 47642a2e1..2bb7c7e69 100644 --- a/crates/worker/src/internal_worker.rs +++ b/crates/worker/src/internal_worker.rs @@ -521,6 +521,19 @@ impl InternalWorkerSessionHandle { .unwrap_or_default() } + pub(crate) fn session_attachment( + &self, + attachment_id: &str, + max_attachment_bytes: u64, + ) -> Result + { + session_store::read_session_attachment_from_entries( + &self.entries(), + attachment_id, + max_attachment_bytes, + ) + } + pub(crate) async fn send( &self, input: impl Into, @@ -1447,6 +1460,50 @@ permission = "write" assert_eq!(result.identity.kind, "test"); } + #[tokio::test] + async fn visible_internal_session_resolves_its_projected_image_attachment() { + let (handle, _events) = + test_internal_worker_session(InternalWorkerVisibility::ParentClient); + let image = vec![1, 3, 5, 7]; + let entry_id = "internal-image-entry"; + let attachment_id = + session_store::session_tool_attachment_id(entry_id, 0, "image/png", &image); + handle.publish_test_entry(LogEntry::AnnotatedToolResult { + ts: 1, + entry: session_store::LoggedHistoryEntry { + item: session_store::LoggedItem::ToolResult { + call_id: "internal-image-call".into(), + summary: "internal screenshot".into(), + content: None, + attachments: vec![session_store::logged_item::LoggedAttachment::Image { + mime_type: "image/png".into(), + data: image.clone(), + }], + disposition: agen::tool::ToolResultDisposition::Success, + is_error: false, + }, + metadata: session_store::LoggedSessionHistoryMetadata { + entry_id: session_store::LoggedSessionHistoryEntryId(entry_id.into()), + origin: session_store::LoggedSessionHistoryOrigin::LegacyUnknown, + derivation: None, + }, + }, + }); + + let projected = handle.protocol_snapshot(); + let protocol::SessionSnapshotEntryData::ToolResult { attachments, .. } = + &projected.session.entries[0].data + else { + panic!("expected projected tool result"); + }; + assert_eq!(attachments[0].attachment_id, attachment_id); + let attachment = handle + .session_attachment(&attachment_id, 1024) + .expect("internal image body"); + assert_eq!(attachment.media_type, "image/png"); + assert_eq!(attachment.data, image); + } + #[test] fn internal_turn_result_mapping_is_exhaustive() { let cases = [ diff --git a/crates/worker/src/spawn/registry.rs b/crates/worker/src/spawn/registry.rs index 5b706f792..037cba6ee 100644 --- a/crates/worker/src/spawn/registry.rs +++ b/crates/worker/src/spawn/registry.rs @@ -774,6 +774,41 @@ impl SpawnedWorkerRegistry { }); } + pub(crate) fn session_attachment( + &self, + session_id: &str, + attachment_id: &str, + max_attachment_bytes: u64, + ) -> Option< + Result< + session_store::RetainedSessionAttachment, + session_store::RetainedAttachmentReadError, + >, + > { + let internal_session = self + .internal_records + .lock() + .ok()? + .iter() + .find(|record| { + record.session.visibility() == InternalWorkerVisibility::ParentClient + && record.session.session_id_string() == session_id + }) + .map(|record| record.session.clone()); + let session = internal_session.or_else(|| { + self.service_records + .lock() + .ok()? + .iter() + .find(|record| { + record.session.visibility() == InternalWorkerVisibility::ParentClient + && record.session.session_id_string() == session_id + }) + .map(|record| record.session.clone()) + })?; + Some(session.session_attachment(attachment_id, max_attachment_bytes)) + } + pub(crate) fn internal_worker_snapshots(&self) -> Vec { let parent_session_id = self .parent_protocol @@ -1258,6 +1293,59 @@ mod tests { ) } + #[tokio::test] + async fn parent_visible_internal_attachment_resolves_by_session_identity() { + let registry = registry(); + let (visible, _visible_tx) = + record("visible", InternalWorkerVisibility::ParentClient).await; + let session_id = visible.session.session_id_string(); + let image = vec![2, 4, 6, 8]; + let entry_id = "visible-image-entry"; + let attachment_id = + session_store::session_tool_attachment_id(entry_id, 0, "image/png", &image); + visible + .session + .publish_test_entry(LogEntry::AnnotatedToolResult { + ts: 1, + entry: session_store::LoggedHistoryEntry { + item: session_store::LoggedItem::ToolResult { + call_id: "visible-image-call".into(), + summary: "visible image".into(), + content: None, + attachments: vec![session_store::logged_item::LoggedAttachment::Image { + mime_type: "image/png".into(), + data: image.clone(), + }], + disposition: agen::tool::ToolResultDisposition::Success, + is_error: false, + }, + metadata: session_store::LoggedSessionHistoryMetadata { + entry_id: session_store::LoggedSessionHistoryEntryId(entry_id.into()), + origin: session_store::LoggedSessionHistoryOrigin::LegacyUnknown, + derivation: None, + }, + }, + }); + registry.internal_records.lock().unwrap().push(visible); + + let attachment = registry + .session_attachment(&session_id, &attachment_id, 1024) + .expect("visible session") + .expect("visible attachment"); + assert_eq!(attachment.data, image); + + let (private, _private_tx) = + record("private", InternalWorkerVisibility::ServicePrivate).await; + let private_session_id = private.session.session_id_string(); + registry.internal_records.lock().unwrap().push(private); + assert!( + registry + .session_attachment(&private_session_id, &attachment_id, 1024) + .is_none(), + "service-private sessions must not become parent-fetchable" + ); + } + #[tokio::test] async fn visible_internal_worker_starts_with_own_live_snapshot_metadata() { let registry = registry();