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/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),