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

Separate Console image bodies from session snapshots

Browse Source
This commit is contained in:
Hare
2026-10-04 12:56:40 +09:00
parent f7b722246d
commit 2eb9418a4e
20 changed files with 1021 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/protocol/src/lib.rs
+4 -3
View File
@@ -1037,10 +1037,11 @@ pub enum SessionContentPart {
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[cfg_attr(feature = "json-schema", derive(schemars::JsonSchema))]
pub struct SessionToolAttachment {
/// Stable, bounded identity derived from the owning Session entry, attachment
/// position, and immutable image bytes. It is a locator, not a bearer secret.
pub attachment_id: String,
pub media_type: String,
/// Base64-encoded durable attachment body. Public snapshots preserve the
/// committed multimodal value instead of replacing it with placeholder text.
pub data_base64: String,
pub byte_len: u64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
crates/runtime-api/src/lib.rs
+8 -1
View File
@@ -1085,6 +1085,8 @@ pub const RUNTIME_ROUTE_WORKER_PROTOCOL_WS: &str = "/v1/workers/{worker_id}/prot
pub const RUNTIME_ROUTE_WORKER_ATTACHMENTS: &str = "/v1/workers/{worker_id}/attachments";
pub const RUNTIME_ROUTE_WORKER_ATTACHMENT: &str =
"/v1/workers/{worker_id}/attachments/{artifact_id}";
pub const RUNTIME_ROUTE_WORKER_SESSION_ATTACHMENT: &str =
"/v1/workers/{worker_id}/sessions/{session_id}/attachments/{attachment_id}";
/// Intentionally out-of-contract Runtime routes. This inventory keeps manual paths visible rather
/// than letting them be mistaken for management-contract omissions.
@@ -1189,6 +1191,11 @@ pub const REMAINING_RUNTIME_ROUTES: &[RemainingRuntimeRoute] = &[
path: RUNTIME_ROUTE_WORKER_ATTACHMENT,
reason: "Binary attachment transport",
},
RemainingRuntimeRoute {
method: "GET",
path: RUNTIME_ROUTE_WORKER_SESSION_ATTACHMENT,
reason: "Session image attachment transport",
},
];
#[cfg(test)]
@@ -1478,7 +1485,7 @@ mod tests {
#[test]
fn remaining_route_inventory_does_not_overlap_contract() {
assert_eq!(REMAINING_RUNTIME_ROUTES.len(), 20);
assert_eq!(REMAINING_RUNTIME_ROUTES.len(), 21);
for remaining in REMAINING_RUNTIME_ROUTES {
assert!(!RuntimeApiMetadata::OPERATIONS.iter().any(|operation| {
format!("{:?}", operation.method).eq_ignore_ascii_case(remaining.method)
crates/session-store/src/lib.rs
+5 -3
View File
@@ -110,9 +110,11 @@ pub use public_snapshot::{
DEFAULT_RETAINED_HISTORY_MAX_ENTRIES, DEFAULT_RETAINED_HISTORY_MAX_RESPONSE_BYTES,
DEFAULT_RETAINED_HISTORY_MAX_SCAN_BYTES, DEFAULT_RETAINED_HISTORY_MAX_SEGMENTS,
DEFAULT_RETAINED_HISTORY_PAGE_TURNS, DEFAULT_RETAINED_SNAPSHOT_MAX_BYTES,
MAX_RETAINED_HISTORY_PAGE_TURNS, RetainedHistoryReadError, RetainedHistoryReadLimits,
RetainedSessionIdentity, RetainedSessionSnapshot, RetainedSnapshotReadError,
project_session_snapshot, read_retained_session_history_page, read_retained_session_snapshot,
MAX_RETAINED_HISTORY_PAGE_TURNS, RetainedAttachmentReadError, RetainedHistoryReadError,
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,
};
pub use segment::{
SegmentStartState, append_entry, append_system_item, classify_logged_history_entry,
crates/session-store/src/public_snapshot.rs
+298 -10
View File
@@ -1,7 +1,4 @@
use base64::{
Engine as _,
engine::general_purpose::{STANDARD as BASE64, URL_SAFE_NO_PAD},
};
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use protocol::{
Segment, SessionContentPart, SessionConversationTurn, SessionEntryProvenance,
SessionHistoryLineageBoundary, SessionHistoryPage, SessionMessageRole, SessionSnapshot,
@@ -54,6 +51,32 @@ pub enum RetainedSnapshotReadError {
SnapshotTooLarge,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RetainedSessionAttachment {
pub media_type: String,
pub data: Vec<u8>,
}
#[derive(Debug, Error, PartialEq, Eq)]
pub enum RetainedAttachmentReadError {
#[error("retained worker state is unavailable")]
RetentionMissing,
#[error("retained worker state has no active session pointer")]
ActivePointerMissing,
#[error("retained session identity does not match")]
SessionMismatch,
#[error("session attachment was not found")]
NotFound,
#[error("retained session requires migration")]
MigrationRequired,
#[error("retained session log is corrupt")]
CorruptLog,
#[error("retained session storage is unavailable")]
StorageUnavailable,
#[error("retained session attachment exceeds a resource limit")]
ResourceLimit,
}
pub const DEFAULT_RETAINED_HISTORY_PAGE_TURNS: usize = 5;
pub const MAX_RETAINED_HISTORY_PAGE_TURNS: usize = 5;
pub const DEFAULT_RETAINED_HISTORY_MAX_SCAN_BYTES: u64 = 64 * 1024 * 1024;
@@ -175,7 +198,9 @@ pub fn read_retained_session_snapshot(
.ok_or(RetainedSnapshotReadError::ActivePointerMissing)?;
let store = WorkerSessionStore::open_read_only(aggregate_root.join("session"))
.map_err(map_store_error)?;
if store.segment_log_len(segment_id).map_err(map_store_error)? > max_bytes {
if store.segment_log_len(segment_id).map_err(map_store_error)?
> DEFAULT_RETAINED_HISTORY_MAX_SCAN_BYTES
{
return Err(RetainedSnapshotReadError::SnapshotTooLarge);
}
let entries = store
@@ -199,6 +224,101 @@ pub fn read_retained_session_snapshot(
})
}
/// Resolve one immutable image body from the active retained Session. The
/// caller supplies only public identities; storage paths never cross this API.
pub fn read_retained_session_attachment(
aggregate_root: &Path,
worker_name: &str,
expected_session_id: &str,
attachment_id: &str,
max_scan_bytes: u64,
max_attachment_bytes: u64,
) -> Result<RetainedSessionAttachment, RetainedAttachmentReadError> {
let aggregate = WorkerAggregateStore::open_read_only(aggregate_root, worker_name)
.map_err(map_attachment_worker_store_error)?;
let metadata = aggregate
.read_read_only()
.map_err(map_attachment_worker_store_error)?
.ok_or(RetainedAttachmentReadError::RetentionMissing)?;
let active = metadata
.active
.ok_or(RetainedAttachmentReadError::ActivePointerMissing)?;
if active.session_id.to_string() != expected_session_id {
return Err(RetainedAttachmentReadError::SessionMismatch);
}
let segment_id = active
.segment_id
.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);
}
});
if let Some(result) = found {
return result;
}
}
Err(RetainedAttachmentReadError::NotFound)
}
fn visit_record_history(record: &LogEntry, visit: &mut impl FnMut(&LoggedHistoryEntry)) {
match record {
LogEntry::AnnotatedSegmentStart { history, .. }
| LogEntry::AnnotatedUserInput { history, .. } => {
for entry in history {
visit(entry);
}
}
LogEntry::AnnotatedAssistantItem { entry, .. }
| LogEntry::AnnotatedToolResult { entry, .. } => visit(entry),
_ => {}
}
}
fn attachment_from_entry(
entry: &LoggedHistoryEntry,
requested_id: &str,
max_attachment_bytes: u64,
) -> Option<Result<RetainedSessionAttachment, RetainedAttachmentReadError>> {
let LoggedItem::ToolResult { attachments, .. } = &entry.item else {
return None;
};
for (attachment_index, attachment) in attachments.iter().enumerate() {
let crate::logged_item::LoggedAttachment::Image { mime_type, data } = attachment;
if session_tool_attachment_id(
&entry.metadata.entry_id.0,
attachment_index,
mime_type,
data,
) != requested_id
{
continue;
}
if data.len() as u64 > max_attachment_bytes {
return Some(Err(RetainedAttachmentReadError::ResourceLimit));
}
return Some(Ok(RetainedSessionAttachment {
media_type: mime_type.clone(),
data: data.clone(),
}));
}
None
}
/// Read one bounded backward page from the retained active Segment's adopted
/// lineage without creating, migrating, restoring, or rewriting persisted state.
///
@@ -935,6 +1055,35 @@ fn map_history_store_error(error: StoreError) -> RetainedHistoryReadError {
}
}
fn map_attachment_worker_store_error(error: WorkerStoreError) -> RetainedAttachmentReadError {
match error {
WorkerStoreError::Io(error) if error.kind() == std::io::ErrorKind::NotFound => {
RetainedAttachmentReadError::RetentionMissing
}
WorkerStoreError::Serde(_) | WorkerStoreError::InvalidWorkerName(_) => {
RetainedAttachmentReadError::CorruptLog
}
WorkerStoreError::Io(_) => RetainedAttachmentReadError::StorageUnavailable,
}
}
fn map_attachment_store_error(error: StoreError) -> RetainedAttachmentReadError {
match error {
StoreError::Io(error) if error.kind() == std::io::ErrorKind::NotFound => {
RetainedAttachmentReadError::RetentionMissing
}
StoreError::Corrupt { message, .. } if message.contains("requires migration") => {
RetainedAttachmentReadError::MigrationRequired
}
StoreError::ReadLimitExceeded => RetainedAttachmentReadError::ResourceLimit,
StoreError::Io(_) => RetainedAttachmentReadError::StorageUnavailable,
StoreError::Serde(_) | StoreError::Corrupt { .. } | StoreError::NotFound(_) => {
RetainedAttachmentReadError::CorruptLog
}
_ => RetainedAttachmentReadError::CorruptLog,
}
}
fn map_worker_store_error(error: WorkerStoreError) -> RetainedSnapshotReadError {
match error {
WorkerStoreError::Io(error) if error.kind() == std::io::ErrorKind::NotFound => {
@@ -1028,7 +1177,7 @@ pub(crate) fn project_session_snapshot_for_segment(
} => extend_history(&mut entries, history, Some(segments), *ts),
LogEntry::AnnotatedAssistantItem { ts, entry }
| LogEntry::AnnotatedToolResult { ts, entry } => {
if let Some(data) = project_item(&entry.item) {
if let Some(data) = project_item(&entry.metadata.entry_id.0, &entry.item) {
entries.push(history_entry(entry, *ts, data));
}
}
@@ -1188,7 +1337,7 @@ fn extend_history(
segments: input_segments.cloned().unwrap_or_default(),
}
} else {
let Some(data) = project_item(&entry.item) else {
let Some(data) = project_item(&entry.metadata.entry_id.0, &entry.item) else {
continue;
};
data
@@ -1306,7 +1455,25 @@ fn provenance(origin: &LoggedSessionHistoryOrigin) -> SessionEntryProvenance {
}
}
fn project_item(item: &LoggedItem) -> Option<SessionSnapshotEntryData> {
pub fn session_tool_attachment_id(
entry_id: &str,
attachment_index: usize,
media_type: &str,
data: &[u8],
) -> String {
let mut hash = Sha256::new();
hash.update(b"yoi-session-tool-attachment-v1\0");
hash.update((entry_id.len() as u64).to_be_bytes());
hash.update(entry_id.as_bytes());
hash.update((attachment_index as u64).to_be_bytes());
hash.update((media_type.len() as u64).to_be_bytes());
hash.update(media_type.as_bytes());
hash.update((data.len() as u64).to_be_bytes());
hash.update(data);
URL_SAFE_NO_PAD.encode(hash.finalize())
}
fn project_item(entry_id: &str, item: &LoggedItem) -> Option<SessionSnapshotEntryData> {
match item {
LoggedItem::Message { role, content } => {
let role = match role {
@@ -1355,11 +1522,18 @@ fn project_item(item: &LoggedItem) -> Option<SessionSnapshotEntryData> {
is_error: *is_error,
attachments: attachments
.iter()
.map(|attachment| match attachment {
.enumerate()
.map(|(attachment_index, attachment)| match attachment {
crate::logged_item::LoggedAttachment::Image { mime_type, data } => {
SessionToolAttachment {
attachment_id: session_tool_attachment_id(
entry_id,
attachment_index,
mime_type,
data,
),
media_type: mime_type.clone(),
data_base64: BASE64.encode(data),
byte_len: data.len() as u64,
}
}
})
@@ -3307,6 +3481,120 @@ mod tests {
));
}
#[test]
fn large_legacy_inline_images_project_as_bounded_refs_and_fetch_from_retained_storage() {
let worker_name = "worker-images";
let session_id = crate::new_session_id();
let segment_id = crate::new_segment_id();
// Approximate the decoded byte sizes of the four W-296 screenshots.
let sizes = [3_475_194_usize, 3_459_051, 5_735_340, 5_703_924];
let attachments = sizes
.iter()
.enumerate()
.map(
|(index, size)| crate::logged_item::LoggedAttachment::Image {
mime_type: "image/png".into(),
data: vec![index as u8 + 1; *size],
},
)
.collect();
let tool_result = LoggedHistoryEntry {
item: LoggedItem::ToolResult {
call_id: "call-images".into(),
summary: "four screenshots".into(),
content: None,
attachments,
disposition: agen::tool::ToolResultDisposition::Success,
is_error: false,
},
metadata: LoggedSessionHistoryMetadata {
entry_id: LoggedSessionHistoryEntryId("tool-images".into()),
origin: LoggedSessionHistoryOrigin::LegacyUnknown,
derivation: None,
},
};
let live = project_current_session_snapshot(&[LogEntry::AnnotatedToolResult {
ts: 2,
entry: tool_result.clone(),
}]);
let log = vec![LogEntry::AnnotatedSegmentStart {
ts: 1,
session_id,
system_prompt: None,
config: RequestConfig::default(),
history: vec![tool_result],
forked_from: None,
compacted_from: None,
}];
let (_root, aggregate_root) =
persist_history_fixture(worker_name, session_id, segment_id, vec![(segment_id, log)]);
let retained = read_retained_session_snapshot(
&aggregate_root,
worker_name,
DEFAULT_RETAINED_SNAPSHOT_MAX_BYTES,
)
.unwrap();
let SessionSnapshotEntryData::ToolResult { attachments, .. } =
&retained.snapshot.entries[0].data
else {
panic!("expected tool result");
};
assert_eq!(attachments.len(), 4);
let SessionSnapshotEntryData::ToolResult {
attachments: live_attachments,
..
} = &live.entries[0].data
else {
panic!("expected live tool result");
};
assert_eq!(live_attachments, attachments);
let snapshot_json = serde_json::to_vec(&retained.snapshot).unwrap();
assert!(snapshot_json.len() < 16 * 1024);
assert!(!String::from_utf8_lossy(&snapshot_json).contains("data_base64"));
for (index, attachment) in attachments.iter().enumerate() {
assert_eq!(attachment.byte_len, sizes[index] as u64);
assert_eq!(attachment.attachment_id.len(), 43);
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.len(), sizes[index]);
assert!(body.data.iter().all(|byte| *byte == index as u8 + 1));
}
assert_eq!(
read_retained_session_attachment(
&aggregate_root,
worker_name,
&crate::new_session_id().to_string(),
&attachments[0].attachment_id,
DEFAULT_RETAINED_HISTORY_MAX_SCAN_BYTES,
10 * 1024 * 1024,
)
.unwrap_err(),
RetainedAttachmentReadError::SessionMismatch
);
assert_eq!(
read_retained_session_attachment(
&aggregate_root,
worker_name,
&session_id.to_string(),
"../../metadata.json",
DEFAULT_RETAINED_HISTORY_MAX_SCAN_BYTES,
10 * 1024 * 1024,
)
.unwrap_err(),
RetainedAttachmentReadError::NotFound
);
}
#[test]
fn retained_reads_do_not_create_missing_aggregate() {
let root = tempfile::tempdir().unwrap();
crates/worker-runtime/src/execution.rs
+23
View File
@@ -514,6 +514,13 @@ pub struct WorkerSessionHistoryRequest {
pub limit: Option<usize>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct WorkerSessionAttachmentRequest {
pub worker_ref: WorkerRef,
pub session_id: String,
pub attachment_id: String,
}
pub trait WorkerExecutionBackend: Send + Sync + 'static {
fn backend_id(&self) -> &str;
@@ -544,6 +551,14 @@ pub trait WorkerExecutionBackend: Send + Sync + 'static {
}
}
fn worker_session_attachment(
&self,
_request: WorkerSessionAttachmentRequest,
) -> Result<session_store::RetainedSessionAttachment, session_store::RetainedAttachmentReadError>
{
Err(session_store::RetainedAttachmentReadError::RetentionMissing)
}
fn spawn_worker(&self, request: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult;
/// Read-only restore validation performed before any live controller,
@@ -780,6 +795,14 @@ impl WorkerExecutionBackendRef {
self.backend.worker_session_history(request)
}
pub(crate) fn worker_session_attachment(
&self,
request: WorkerSessionAttachmentRequest,
) -> Result<session_store::RetainedSessionAttachment, session_store::RetainedAttachmentReadError>
{
self.backend.worker_session_attachment(request)
}
pub(crate) fn spawn_worker(
&self,
request: WorkerExecutionSpawnRequest,
crates/worker-runtime/src/http_server.rs
+69 -1
View File
@@ -336,6 +336,13 @@ fn runtime_http_router_with_auth_and_ssh_keyscan_program(
.route(
remaining_route(runtime_api::RUNTIME_ROUTE_WORKER_ATTACHMENT, &["DELETE"]),
delete(delete_worker_uploaded_file),
)
.route(
remaining_route(
runtime_api::RUNTIME_ROUTE_WORKER_SESSION_ATTACHMENT,
&["GET"],
),
get(get_worker_session_attachment),
);
#[cfg(feature = "ws-server")]
@@ -1980,6 +1987,60 @@ async fn delete_worker_uploaded_file(
}))
}
async fn get_worker_session_attachment(
State(state): State<RuntimeHttpState>,
auth: Option<Extension<RuntimeAuthContext>>,
Path((worker_id, session_id, attachment_id)): Path<(String, String, String)>,
Query(request): Query<runtime_api::WorkerSessionRequest>,
) -> Result<Response, RuntimeHttpRestError> {
let worker_ref = worker_ref_for(&state.runtime, worker_id)?;
let scope = auth_workspace_scope(&state, auth.as_ref())?.ok_or_else(|| {
RuntimeHttpRestError::new(
StatusCode::FORBIDDEN,
"runtime_worker_session_scope_required",
"Worker Session attachment reads require a Workspace-scoped capability",
)
})?;
if scope.workspace_id != request.workspace_id {
return Err(RuntimeHttpRestError::new(
StatusCode::FORBIDDEN,
"runtime_worker_session_workspace_scope_mismatch",
"Worker Session Workspace scope does not match the authenticated capability",
));
}
let attachment = state
.runtime
.worker_session_attachment_scoped(&scope, &worker_ref, session_id, attachment_id)
.map_err(|error| {
use session_store::RetainedAttachmentReadError as Error;
let (status, code) = match error {
Error::RetentionMissing | Error::ActivePointerMissing => {
(StatusCode::GONE, "session_attachment_expired")
}
Error::SessionMismatch | Error::NotFound => {
(StatusCode::NOT_FOUND, "session_attachment_not_found")
}
Error::ResourceLimit => (
StatusCode::PAYLOAD_TOO_LARGE,
"session_attachment_too_large",
),
Error::MigrationRequired | Error::CorruptLog | Error::StorageUnavailable => (
StatusCode::SERVICE_UNAVAILABLE,
"session_attachment_unavailable",
),
};
RuntimeHttpRestError::new(status, code, error.to_string())
})?;
let content_type = axum::http::HeaderValue::from_str(&attachment.media_type).map_err(|_| {
RuntimeHttpRestError::new(
StatusCode::SERVICE_UNAVAILABLE,
"session_attachment_media_type_invalid",
"Session attachment media type is invalid",
)
})?;
Ok(([(header::CONTENT_TYPE, content_type)], attachment.data).into_response())
}
#[allow(dead_code)]
async fn stop_worker(
State(state): State<RuntimeHttpState>,
@@ -2289,7 +2350,7 @@ fn required_runtime_permission(method: &Method, path: &str) -> Option<&'static s
}
if path.ends_with("/input")
|| path.ends_with("/restore")
|| path.contains("/attachments")
|| (path.contains("/attachments") && *method != Method::GET)
|| path.contains("/workdir-attachments")
{
return Some("workers:input");
@@ -3031,6 +3092,13 @@ mod tests {
),
Some("workers:input")
);
assert_eq!(
required_runtime_permission(
&Method::GET,
"/v1/workers/7/sessions/session-1/attachments/attachment-1"
),
Some("workers:read")
);
assert_eq!(
required_runtime_permission(&Method::POST, "/v1/workers/7/workdir-attachments"),
Some("workers:input")
crates/worker-runtime/src/runtime.rs
+33
View File
@@ -1673,6 +1673,39 @@ impl Runtime {
})
}
pub fn worker_session_attachment_scoped(
&self,
scope: &RuntimeWorkspaceScope,
worker_ref: &WorkerRef,
session_id: String,
attachment_id: String,
) -> Result<session_store::RetainedSessionAttachment, session_store::RetainedAttachmentReadError>
{
let operation_lock = self
.worker_operation_lock(worker_ref.worker_id)
.map_err(|_| session_store::RetainedAttachmentReadError::RetentionMissing)?;
let _operation_guard = operation_lock
.lock()
.map_err(|_| session_store::RetainedAttachmentReadError::StorageUnavailable)?;
self.ensure_worker_in_workspace(scope, worker_ref)
.map_err(|_| session_store::RetainedAttachmentReadError::NotFound)?;
let backend = {
let state = self
.lock()
.map_err(|_| session_store::RetainedAttachmentReadError::StorageUnavailable)?;
state
.worker(worker_ref)
.map_err(|_| session_store::RetainedAttachmentReadError::NotFound)?;
state.execution_backend.clone()
}
.ok_or(session_store::RetainedAttachmentReadError::StorageUnavailable)?;
backend.worker_session_attachment(crate::execution::WorkerSessionAttachmentRequest {
worker_ref: worker_ref.clone(),
session_id,
attachment_id,
})
}
/// Attach a live execution through a workspace-scoped Runtime authorization context.
pub fn restore_worker_scoped(
&self,
crates/worker-runtime/src/worker_backend.rs
+43
View File
@@ -186,6 +186,16 @@ pub trait RuntimeWorkerFactory: Send + Sync + 'static {
Err(session_store::RetainedHistoryReadError::RetentionMissing)
}
fn retained_session_attachment(
&self,
_worker_ref: &WorkerRef,
_session_id: &str,
_attachment_id: &str,
) -> Result<session_store::RetainedSessionAttachment, session_store::RetainedAttachmentReadError>
{
Err(session_store::RetainedAttachmentReadError::RetentionMissing)
}
async fn spawn_controller(
&self,
request: WorkerExecutionSpawnRequest,
@@ -1116,6 +1126,27 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
)
}
fn retained_session_attachment(
&self,
worker_ref: &WorkerRef,
session_id: &str,
attachment_id: &str,
) -> Result<session_store::RetainedSessionAttachment, session_store::RetainedAttachmentReadError>
{
let aggregate_dir = self
.worker_aggregate_dir(worker_ref)
.map_err(|_| session_store::RetainedAttachmentReadError::RetentionMissing)?;
let worker_name = Self::runtime_worker_name_for_ref(worker_ref);
session_store::read_retained_session_attachment(
&aggregate_dir,
&worker_name,
session_id,
attachment_id,
session_store::DEFAULT_RETAINED_HISTORY_MAX_SCAN_BYTES,
10 * 1024 * 1024,
)
}
fn observe_workspace_prompt_projection(
&self,
projection: worker::WorkspacePromptProjection,
@@ -2412,6 +2443,18 @@ where
}
}
fn worker_session_attachment(
&self,
request: crate::execution::WorkerSessionAttachmentRequest,
) -> Result<session_store::RetainedSessionAttachment, session_store::RetainedAttachmentReadError>
{
self.factory.retained_session_attachment(
&request.worker_ref,
&request.session_id,
&request.attachment_id,
)
}
fn observe_workspace_prompt_projection(
&self,
projection: worker::WorkspacePromptProjection,
crates/workspace-server/Cargo.toml
+1 -1
View File
@@ -40,6 +40,7 @@ tokio-tungstenite = { workspace = true, features = ["rustls-tls-webpki-roots"] }
worker.workspace = true
workdir = { workspace = true, features = ["http-client"] }
worker-runtime.workspace = true
session-store.workspace = true
toml.workspace = true
tracing.workspace = true
tracing-subscriber.workspace = true
@@ -51,7 +52,6 @@ webauthn-rs = { workspace = true }
[dev-dependencies]
agen.workspace = true
session-store.workspace = true
tempfile.workspace = true
tools.workspace = true
yoi = { path = "../yoi" }
crates/workspace-server/src/hosts.rs
+173 -1
View File
@@ -770,6 +770,12 @@ pub struct WorkerProxyConnectPoint {
pub diagnostics: Vec<RuntimeDiagnostic>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RuntimeSessionAttachment {
pub media_type: String,
pub data: Vec<u8>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum RuntimeRegistryError {
InvalidIdentifier {
@@ -1220,6 +1226,23 @@ pub trait WorkspaceWorkerRuntime: Send + Sync {
})
}
fn worker_session_attachment(
&self,
_workspace_id: &str,
worker_ref: &EmbeddedWorkerRef,
_session_id: &str,
_attachment_id: &str,
) -> Result<RuntimeSessionAttachment, RuntimeRegistryError> {
Err(RuntimeRegistryError::RuntimeOperationFailed {
runtime_id: self.runtime_id().to_string(),
code: "session_attachment_unsupported".to_string(),
message: format!(
"Session attachment reads are not supported for Worker `{}`",
worker_ref.worker_id
),
})
}
fn session_public_search(
&self,
_request: runtime_api::SessionPublicSearchRequest,
@@ -2096,6 +2119,27 @@ impl RuntimeRegistry {
})
}
pub fn worker_session_attachment(
&self,
workspace_id: &str,
worker: &RuntimeWorkerRef,
session_id: &str,
attachment_id: &str,
) -> Result<RuntimeSessionAttachment, RuntimeRegistryError> {
validate_backend_identifier("runtime_id", &worker.runtime_id)?;
validate_backend_identifier("worker_id", &worker.worker_id)?;
validate_backend_identifier("session_id", session_id)?;
validate_backend_identifier("attachment_id", attachment_id)?;
let runtime = self.runtime(&worker.runtime_id)?;
let worker_ref =
EmbeddedWorkerRef::new(EmbeddedWorkerId::parse(&worker.worker_id).ok_or_else(
|| RuntimeRegistryError::UnknownWorker {
worker: worker.clone(),
},
)?);
runtime.worker_session_attachment(workspace_id, &worker_ref, session_id, attachment_id)
}
pub fn session_public_search(
&self,
runtime_id: &str,
@@ -3017,6 +3061,46 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime {
.map_err(|error| error.to_string())
}
fn worker_session_attachment(
&self,
workspace_id: &str,
worker_ref: &EmbeddedWorkerRef,
session_id: &str,
attachment_id: &str,
) -> Result<RuntimeSessionAttachment, RuntimeRegistryError> {
let scope = RuntimeWorkspaceScope::new(workspace_id, "embedded-backend");
self.runtime
.worker_session_attachment_scoped(
&scope,
worker_ref,
session_id.to_string(),
attachment_id.to_string(),
)
.map(|attachment| RuntimeSessionAttachment {
media_type: attachment.media_type,
data: attachment.data,
})
.map_err(|error| RuntimeRegistryError::RuntimeOperationFailed {
runtime_id: self.runtime_id.clone(),
code: match error {
session_store::RetainedAttachmentReadError::RetentionMissing
| session_store::RetainedAttachmentReadError::ActivePointerMissing => {
"session_attachment_expired"
}
session_store::RetainedAttachmentReadError::SessionMismatch
| session_store::RetainedAttachmentReadError::NotFound => {
"session_attachment_not_found"
}
session_store::RetainedAttachmentReadError::ResourceLimit => {
"session_attachment_too_large"
}
_ => "session_attachment_unavailable",
}
.to_string(),
message: error.to_string(),
})
}
fn session_public_search(
&self,
request: runtime_api::SessionPublicSearchRequest,
@@ -3732,7 +3816,7 @@ fn workspace_runtime_operation(method: &str, path_and_query: &str) -> &'static s
}
if path.ends_with("/input")
|| path.ends_with("/restore")
|| path.contains("/attachments")
|| (path.contains("/attachments") && method != "GET")
|| path.contains("/workdir-attachments")
{
return "workers:input";
@@ -4059,6 +4143,64 @@ impl RemoteWorkerRuntime {
self.send_json(path, "GET", &[], self.http.get(self.endpoint(path)))
}
fn get_session_attachment(
&self,
path: &str,
) -> Result<RuntimeSessionAttachment, RuntimeDiagnostic> {
let runtime_id = self.runtime_id.clone();
let workspace_id = self.workspace_id.clone();
let bearer_token = self.bearer_token.clone();
let capability_token = match &self.workspace_authorization {
Some(authorization) => Some(authorization.issue(
"GET",
path,
workspace_runtime_operation("GET", path),
worker_id_from_remote_path(path).as_deref(),
&[],
)?),
None => None,
};
let request = self.http.get(self.endpoint(path));
run_blocking_http(move || {
let mut request = request.header(RUNTIME_WORKSPACE_SCOPE_HEADER, &workspace_id);
if let Some(token) = capability_token.as_deref().or(bearer_token.as_deref()) {
request = request.header(AUTHORIZATION, format!("Bearer {token}"));
}
let response = request
.send()
.map_err(|error| remote_reqwest_diagnostic(&runtime_id, error))?;
let status = response.status();
if !status.is_success() {
return Err(remote_http_status_diagnostic(&runtime_id, status, response));
}
let media_type = response
.headers()
.get(CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.unwrap_or("application/octet-stream")
.to_string();
let mut data = Vec::new();
response
.take((MAX_REMOTE_RUNTIME_RESPONSE_BYTES + 1) as u64)
.read_to_end(&mut data)
.map_err(|_| {
diagnostic(
"remote_runtime_response_read_failed",
HostDiagnosticSeverity::Error,
"Remote Runtime attachment response could not be read".to_string(),
)
})?;
if data.len() > MAX_REMOTE_RUNTIME_RESPONSE_BYTES {
return Err(diagnostic(
"remote_runtime_response_too_large",
HostDiagnosticSeverity::Error,
"Remote Runtime attachment response exceeded the allowed size".to_string(),
));
}
Ok(RuntimeSessionAttachment { media_type, data })
})
}
fn post_json<B, T>(&self, path: &str, body: &B) -> Result<T, RuntimeDiagnostic>
where
B: Serialize + ?Sized,
@@ -5093,6 +5235,29 @@ impl WorkspaceWorkerRuntime for RemoteWorkerRuntime {
.map_err(|diagnostic| diagnostic.message)
}
fn worker_session_attachment(
&self,
workspace_id: &str,
worker_ref: &EmbeddedWorkerRef,
session_id: &str,
attachment_id: &str,
) -> Result<RuntimeSessionAttachment, RuntimeRegistryError> {
let path = format!(
"/v1/workers/{}/sessions/{}/attachments/{}?workspace_id={}",
url_path_segment_encode(&worker_ref.worker_id.to_string()),
url_path_segment_encode(session_id),
url_path_segment_encode(attachment_id),
url_query_value_encode(workspace_id),
);
self.get_session_attachment(&path).map_err(|diagnostic| {
RuntimeRegistryError::RuntimeOperationFailed {
runtime_id: self.runtime_id.clone(),
code: diagnostic.code,
message: diagnostic.message,
}
})
}
fn session_public_search(
&self,
request: runtime_api::SessionPublicSearchRequest,
@@ -7648,6 +7813,13 @@ mod tests {
),
"workers:input"
);
assert_eq!(
workspace_runtime_operation(
"GET",
&format!("/v1/workers/{worker_id}/sessions/session-1/attachments/attachment-1")
),
"workers:read"
);
assert_eq!(
workspace_runtime_operation("POST", SSH_HOST_KEY_PROBE_PATH),
SSH_HOST_KEY_PROBE_OPERATION
crates/workspace-server/src/server.rs
+72 -2
View File
@@ -7,9 +7,9 @@ use std::sync::{Arc, Mutex, RwLock, Weak};
use async_trait::async_trait;
use axum::extract::ws::{CloseFrame, Message as WsMessage, WebSocket, WebSocketUpgrade};
use axum::extract::{DefaultBodyLimit, Extension, Path as AxumPath, Query, Request, State};
use axum::http::header::{CACHE_CONTROL, CONTENT_TYPE, LOCATION, ORIGIN};
#[cfg(test)]
use axum::http::header::{CACHE_CONTROL, ETAG, IF_NONE_MATCH};
use axum::http::header::{CONTENT_TYPE, LOCATION, ORIGIN};
use axum::http::header::{ETAG, IF_NONE_MATCH};
use axum::http::{HeaderMap, HeaderName, HeaderValue, Method, StatusCode, Uri};
use axum::middleware::{self, Next};
use axum::response::{IntoResponse, Response};
@@ -10268,6 +10268,72 @@ impl server_api::ServerApi for ServerApiContractService {
}
}
async fn scoped_worker_session_attachment(
State(api): State<WorkspaceApi>,
AxumPath((workspace_id, runtime_id, worker_id, session_id, attachment_id)): AxumPath<(
String,
String,
String,
String,
String,
)>,
Extension(_actor): Extension<RequestActor>,
) -> Response {
if workspace_id != api.config.workspace_id {
return (
StatusCode::NOT_FOUND,
Json(serde_json::json!({"code": "workspace_not_found"})),
)
.into_response();
}
let worker = RuntimeWorkerRef::new(runtime_id, worker_id);
match api
.runtime
.worker_session_attachment(&workspace_id, &worker, &session_id, &attachment_id)
{
Ok(attachment) => {
let Ok(content_type) = HeaderValue::from_str(&attachment.media_type) else {
return (
StatusCode::SERVICE_UNAVAILABLE,
Json(serde_json::json!({"code": "session_attachment_media_type_invalid"})),
)
.into_response();
};
(
[
(CONTENT_TYPE, content_type),
(
CACHE_CONTROL,
HeaderValue::from_static("private, max-age=3600"),
),
],
attachment.data,
)
.into_response()
}
Err(crate::hosts::RuntimeRegistryError::RuntimeOperationFailed {
code, message, ..
}) => {
let status = match code.as_str() {
"session_attachment_expired" => StatusCode::GONE,
"session_attachment_not_found" => StatusCode::NOT_FOUND,
"session_attachment_too_large" => StatusCode::PAYLOAD_TOO_LARGE,
_ => StatusCode::BAD_GATEWAY,
};
(
status,
Json(serde_json::json!({"code": code, "message": message})),
)
.into_response()
}
Err(_) => (
StatusCode::NOT_FOUND,
Json(serde_json::json!({"code": "session_attachment_not_found"})),
)
.into_response(),
}
}
fn build_inner_router(api: WorkspaceApi) -> Router {
let contract_service = ServerApiContractService::Workspace(api.clone());
let auth = generated_auth_contract_router(contract_service.clone());
@@ -10296,6 +10362,10 @@ fn build_inner_router(api: WorkspaceApi) -> Router {
_,
_,
>(workspace, scoped_worker_protocol_ws)
.route(
"/api/w/{workspace_id}/runtimes/{runtime_id}/workers/{worker_id}/sessions/{session_id}/attachments/{attachment_id}",
get(scoped_worker_session_attachment),
)
.route(
"/api/w/{workspace_id}/workers/self/backend-job-result",
post(scoped_submit_backend_job_result),