Resolve Console images across session lineage

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,

crates/session-store/src/public_snapshot.rs

@@ -251,28 +251,80 @@ 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
+    let limits = RetainedHistoryReadLimits {
-        .segment_log_len(segment_id)
+        max_scan_bytes,
-        .map_err(map_attachment_store_error)?
+        ..RetainedHistoryReadLimits::default()
-        > max_scan_bytes
+    };
-    {
+    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);
             }
-    let entries = store
+            if let Some(result) =
-        .read_all_read_only(active.session_id, segment_id)
+                session_attachment_from_record(&record.entry, attachment_id, max_attachment_bytes)
-        .map_err(map_attachment_store_error)?;
+            {
-    for record in &entries {
+                return result;
+            }
+        }
+    }
+    Err(RetainedAttachmentReadError::NotFound)
+}
+
+pub fn read_session_attachment_from_entries(
+    entries: &[LogEntry],
+    attachment_id: &str,
+    max_attachment_bytes: u64,
+) -> Result<RetainedSessionAttachment, RetainedAttachmentReadError> {
+    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<Result<RetainedSessionAttachment, RetainedAttachmentReadError>> {
     let mut found = None;
     visit_record_history(record, &mut |entry| {
         if found.is_none() {
             found = attachment_from_entry(entry, attachment_id, max_attachment_bytes);
         }
     });
-        if let Some(result) = found {
+    found
-            return result;
-        }
-    }
-    Err(RetainedAttachmentReadError::NotFound)
 }
 
 fn visit_record_history(record: &LogEntry, visit: &mut impl FnMut(&LoggedHistoryEntry)) {
@@ -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";

crates/worker-runtime/src/worker_backend.rs

@@ -2448,6 +2448,24 @@ where
         request: crate::execution::WorkerSessionAttachmentRequest,
     ) -> Result<session_store::RetainedSessionAttachment, session_store::RetainedAttachmentReadError>
     {
+        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,

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<session_store::RetainedSessionAttachment, session_store::RetainedAttachmentReadError>
+    {
+        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
     }

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::RetainedSessionAttachment, session_store::RetainedAttachmentReadError>
+    {
+        session_store::read_session_attachment_from_entries(
+            &self.entries(),
+            attachment_id,
+            max_attachment_bytes,
+        )
+    }
+
     pub(crate) async fn send(
         &self,
         input: impl Into<String>,
@@ -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 = [

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<InternalWorkerSnapshot> {
         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();