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

Resolve Console images across session lineage

Browse Source
This commit is contained in:
Hare
2026-10-04 13:14:36 +09:00
parent 2eb9418a4e
commit 773907462b
6 changed files with 360 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/session-store/src/lib.rs
+2 -1
View File
@@ -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
+171 -14
View File
@@ -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
.segment_log_len(segment_id)
.map_err(map_attachment_store_error)?
> max_scan_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);
}
let entries = store
.read_all_read_only(active.session_id, segment_id)
.map_err(map_attachment_store_error)?;
for record in &entries {
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<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 {
return result;
}
}
Err(RetainedAttachmentReadError::NotFound)
found
}
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
+18
View File
@@ -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
+21
View File
@@ -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
+57
View File
@@ -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
+88
View File
@@ -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();