Resolve Console images across session lineage

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