feat: observe retained worker sessions without restore

This commit is contained in:
2026-09-18 04:27:38 +09:00
parent ae7f0f0446
commit 7a61a26670
26 changed files with 834 additions and 124 deletions
Generated
+16
View File
@@ -650,8 +650,10 @@ dependencies = [
"futures",
"protocol",
"reqwest",
"runtime-api",
"serde",
"serde_json",
"server-api",
"tempfile",
"thiserror 2.0.18",
"ticket",
@@ -3714,6 +3716,7 @@ dependencies = [
"manifest",
"protocol",
"serde",
"serde_json",
"tokio",
"workdir",
"workspace-api",
@@ -4128,6 +4131,18 @@ dependencies = [
"syn 2.0.117",
]
[[package]]
name = "server-api"
version = "0.1.0"
dependencies = [
"api-macros",
"runtime-api",
"serde",
"serde_json",
"tokio",
"workspace-api",
]
[[package]]
name = "session-analytics"
version = "0.1.0"
@@ -6291,6 +6306,7 @@ dependencies = [
"serde",
"serde_json",
"serde_yaml",
"server-api",
"sha2 0.11.0",
"ssh-key",
"tempfile",
+3
View File
@@ -14,6 +14,7 @@ members = [
"crates/worker",
"crates/worker-runtime",
"crates/runtime-api",
"crates/server-api",
"crates/plugin-pdk",
"crates/yoi",
"crates/protocol",
@@ -49,6 +50,7 @@ default-members = [
"crates/worker",
"crates/worker-runtime",
"crates/runtime-api",
"crates/server-api",
"crates/plugin-pdk",
"crates/yoi",
"crates/protocol",
@@ -104,6 +106,7 @@ session-analytics = { path = "crates/session-analytics" }
session-store = { path = "crates/session-store" }
standalone = { path = "crates/standalone" }
runtime-api = { path = "crates/runtime-api" }
server-api = { path = "crates/server-api" }
secrets = { path = "crates/secrets" }
tools = { path = "crates/tools" }
config-source = { path = "crates/config-source" }
+2
View File
@@ -8,11 +8,13 @@ license.workspace = true
async-trait.workspace = true
chrono = { version = "0.4", default-features = false, features = ["clock"] }
protocol = { workspace = true }
runtime-api.workspace = true
ticket = { workspace = true }
futures = { workspace = true }
reqwest = { version = "0.13", default-features = false, features = ["blocking", "json", "native-tls"] }
serde = { workspace = true }
serde_json = { workspace = true }
server-api.workspace = true
thiserror = { workspace = true }
tokio = { workspace = true, features = ["rt", "macros", "net", "io-util", "sync", "time"] }
tokio-tungstenite = { workspace = true }
+3 -1
View File
@@ -28,7 +28,7 @@ pub use workspace_api::{
WorkingDirectorySummary as BackendWorkingDirectorySummary,
};
#[derive(Debug, Clone, PartialEq, Eq)]
#[derive(Debug, Clone, PartialEq)]
pub struct BackendRuntimeTarget {
/// Workspace Backend API root URL, for example `http://127.0.0.1:8787`.
/// This is intentionally the Backend endpoint, not a Runtime endpoint.
@@ -39,6 +39,7 @@ pub struct BackendRuntimeTarget {
pub runtime_id: String,
/// Backend-owned Worker identity used as path authority.
pub worker_id: String,
pub initial_snapshot: Option<protocol::SessionSnapshot>,
}
impl BackendRuntimeTarget {
@@ -53,6 +54,7 @@ impl BackendRuntimeTarget {
workspace_id: workspace_id.into(),
runtime_id: runtime_id.into(),
worker_id: worker_id.into(),
initial_snapshot: None,
}
}
+2
View File
@@ -40,6 +40,8 @@ pub use backend_workspace::{
list_backend_workspaces_blocking,
};
pub use client::{Client, ClientError};
pub use runtime_api::{WorkerSessionAvailability, WorkerSessionUnavailableReason};
pub use server_api::{ServerApi, ServerApiClient, WorkspaceWorkerSessionResponse};
pub use target::{
BackendTarget, BackendWorkerLaunch, Dashboard, ResolvedTarget, StandaloneTarget,
StandaloneWorkerListIntent, StandaloneWorkerResumeIntent, Target, TargetError, TargetKind,
+1 -1
View File
@@ -134,7 +134,7 @@ pub struct WorkerList {
pub include_stopped: bool,
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[derive(Debug, Clone, PartialEq)]
pub struct WorkerConnection {
pub target: BackendRuntimeTarget,
}
+1
View File
@@ -14,4 +14,5 @@ workspace-api.workspace = true
workdir.workspace = true
[dev-dependencies]
serde_json.workspace = true
tokio = { workspace = true, features = ["macros", "net", "rt-multi-thread"] }
+66 -1
View File
@@ -532,6 +532,44 @@ pub struct WorkerRetentionExecutionResult {
pub diagnostics_retained: bool,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct WorkerSessionRequest {
pub workspace_id: String,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct RetainedSessionIdentity {
pub session_id: String,
pub segment_id: String,
pub entry_count: u64,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum WorkerSessionUnavailableReason {
RetentionMissing,
ActivePointerMissing,
CorruptLog,
MigrationRequired,
RetentionExpired,
SnapshotTooLarge,
StorageUnavailable,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "availability", rename_all = "snake_case")]
pub enum WorkerSessionAvailability {
LiveProtocol,
RetainedSnapshot {
identity: RetainedSessionIdentity,
snapshot: protocol::SessionSnapshot,
},
Unavailable {
reason: WorkerSessionUnavailableReason,
message: String,
},
}
#[api(reqwest, axum)]
pub trait RuntimeApi {
#[get("/v1/ping", status = 200, error_status = 400)]
@@ -555,6 +593,13 @@ pub trait RuntimeApi {
#[path] worker_id: String,
) -> Result<WorkerResponse, RuntimeApiError>;
#[get("/v1/workers/{worker_id}/session", status = 200, error_status = 400)]
async fn worker_session(
&self,
#[path] worker_id: String,
#[query] request: WorkerSessionRequest,
) -> Result<WorkerSessionAvailability, RuntimeApiError>;
#[post("/v1/workers", status = 200, error_status = 400)]
async fn create_worker(
&self,
@@ -827,6 +872,14 @@ mod tests {
Err(test_error(404))
}
async fn worker_session(
&self,
_worker_id: String,
_request: WorkerSessionRequest,
) -> Result<WorkerSessionAvailability, RuntimeApiError> {
Ok(WorkerSessionAvailability::LiveProtocol)
}
async fn create_worker(
&self,
_request: CreateWorkerRequest,
@@ -929,10 +982,22 @@ mod tests {
server.abort();
}
#[test]
fn worker_session_availability_has_closed_tagged_wire_shape() {
let value = serde_json::to_value(WorkerSessionAvailability::Unavailable {
reason: WorkerSessionUnavailableReason::SnapshotTooLarge,
message: "snapshot exceeds limit".to_string(),
})
.unwrap();
assert_eq!(value["availability"], "unavailable");
assert_eq!(value["reason"], "snapshot_too_large");
assert_eq!(value["message"], "snapshot exceeds limit");
}
#[test]
fn contract_inventory_is_complete_and_unique() {
let operations = RuntimeApiMetadata::OPERATIONS;
assert_eq!(operations.len(), 14);
assert_eq!(operations.len(), 15);
let mut routes = operations
.iter()
.map(|operation| (format!("{:?}", operation.method), operation.path))
+14
View File
@@ -0,0 +1,14 @@
[package]
name = "server-api"
version = "0.1.0"
edition = "2024"
[dependencies]
api-macros = { workspace = true, features = ["reqwest", "axum"] }
runtime-api.workspace = true
serde = { workspace = true, features = ["derive"] }
workspace-api.workspace = true
[dev-dependencies]
serde_json.workspace = true
tokio = { workspace = true }
+59
View File
@@ -0,0 +1,59 @@
//! Generated Workspace Server REST contract for public clients.
use api_macros::api;
pub use api_macros::axum as server_support;
pub use api_macros::reqwest as client_support;
pub use api_macros::{ApiContract, HttpMethod};
use serde::{Deserialize, Serialize};
use workspace_api::WorkspaceWorkerSubject;
pub type ServerApiError = runtime_api::RuntimeApiError;
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct WorkspaceWorkerSessionResponse {
pub subject: WorkspaceWorkerSubject,
#[serde(flatten)]
pub observation: runtime_api::WorkerSessionAvailability,
}
#[api(reqwest, axum)]
pub trait ServerApi {
#[get(
"/api/w/{workspace_id}/runtimes/{runtime_id}/workers/{worker_id}/session",
status = 200,
error_status = 400
)]
async fn worker_session(
&self,
#[path] workspace_id: String,
#[path] runtime_id: String,
#[path] worker_id: String,
) -> Result<WorkspaceWorkerSessionResponse, ServerApiError>;
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn worker_session_contract_and_flattened_availability_are_stable() {
let operations = ServerApiMetadata::OPERATIONS;
assert_eq!(operations.len(), 1);
assert_eq!(operations[0].method, HttpMethod::Get);
assert_eq!(
operations[0].path,
"/api/w/{workspace_id}/runtimes/{runtime_id}/workers/{worker_id}/session"
);
let response = WorkspaceWorkerSessionResponse {
subject: WorkspaceWorkerSubject::RuntimeWorker {
runtime_id: "runtime-a".to_string(),
worker_id: "worker-a".to_string(),
},
observation: runtime_api::WorkerSessionAvailability::LiveProtocol,
};
let value = serde_json::to_value(response).unwrap();
assert_eq!(value["availability"], "live_protocol");
assert_eq!(value["subject"]["kind"], "runtime_worker");
}
}
+4
View File
@@ -57,6 +57,10 @@ pub use history::{
};
pub use logged_item::{LoggedContentPart, LoggedItem, LoggedRole, from_logged, to_logged};
pub use paste_artifact::PasteArtifactLimits;
pub use public_snapshot::{
DEFAULT_RETAINED_SNAPSHOT_MAX_BYTES, RetainedSessionIdentity, RetainedSessionSnapshot,
RetainedSnapshotReadError, project_session_snapshot, read_retained_session_snapshot,
};
pub use segment::{
SegmentStartState, append_entry, append_system_item, classify_logged_history_entry,
create_compacted_segment, create_segment, create_segment_with_ids, ensure_head_or_fork, fork,
+170 -1
View File
@@ -7,11 +7,119 @@ use protocol::{
SessionSnapshotEntry, SessionSnapshotEntryData, SessionToolAttachment,
};
use std::path::Path;
use thiserror::Error;
use crate::{
LogEntry, LoggedContentPart, LoggedHistoryEntry, LoggedItem, LoggedRole,
LoggedSessionHistoryOrigin, SessionId, SystemItem,
LoggedSessionHistoryOrigin, SessionId, Store, StoreError, SystemItem, WorkerAggregateStore,
WorkerMetadataStore, WorkerSessionStore, WorkerStoreError,
};
/// Projection limit leaves one MiB for the typed API envelope so a public
/// retained-session response remains below the 16 MiB response ceiling.
pub const DEFAULT_RETAINED_SNAPSHOT_MAX_BYTES: u64 = 15 * 1024 * 1024;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RetainedSessionIdentity {
pub session_id: String,
pub segment_id: String,
pub entry_count: u64,
}
#[derive(Debug, Clone, PartialEq)]
pub struct RetainedSessionSnapshot {
pub identity: RetainedSessionIdentity,
pub snapshot: SessionSnapshot,
}
#[derive(Debug, Error)]
pub enum RetainedSnapshotReadError {
#[error("retained worker state is unavailable")]
RetentionMissing,
#[error("retained worker state has no active session pointer")]
ActivePointerMissing,
#[error("retained session requires migration")]
MigrationRequired,
#[error("retained session log is corrupt")]
CorruptLog,
#[error("retained session snapshot exceeds the observation limit")]
SnapshotTooLarge,
}
/// Read the active retained Session without creating, migrating, restoring, or
/// rewriting any persisted state.
pub fn read_retained_session_snapshot(
aggregate_root: &Path,
worker_name: &str,
max_bytes: u64,
) -> Result<RetainedSessionSnapshot, RetainedSnapshotReadError> {
let aggregate = WorkerAggregateStore::open_read_only(aggregate_root, worker_name)
.map_err(map_worker_store_error)?;
let metadata = aggregate
.read_by_name(worker_name)
.map_err(map_worker_store_error)?
.ok_or(RetainedSnapshotReadError::RetentionMissing)?;
let active = metadata
.active
.ok_or(RetainedSnapshotReadError::ActivePointerMissing)?;
let segment_id = active
.segment_id
.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 {
return Err(RetainedSnapshotReadError::SnapshotTooLarge);
}
let entries = store
.read_all(active.session_id, segment_id)
.map_err(map_store_error)?;
let snapshot = project_session_snapshot(active.session_id, &entries);
if serde_json::to_vec(&snapshot)
.map_err(|_| RetainedSnapshotReadError::CorruptLog)?
.len() as u64
> max_bytes
{
return Err(RetainedSnapshotReadError::SnapshotTooLarge);
}
Ok(RetainedSessionSnapshot {
identity: RetainedSessionIdentity {
session_id: active.session_id.to_string(),
segment_id: segment_id.to_string(),
entry_count: entries.len() as u64,
},
snapshot,
})
}
fn map_worker_store_error(error: WorkerStoreError) -> RetainedSnapshotReadError {
match error {
WorkerStoreError::Io(error) if error.kind() == std::io::ErrorKind::NotFound => {
RetainedSnapshotReadError::RetentionMissing
}
WorkerStoreError::Serde(_) | WorkerStoreError::InvalidWorkerName(_) => {
RetainedSnapshotReadError::CorruptLog
}
WorkerStoreError::Io(_) => RetainedSnapshotReadError::CorruptLog,
}
}
fn map_store_error(error: StoreError) -> RetainedSnapshotReadError {
match error {
StoreError::Io(error) if error.kind() == std::io::ErrorKind::NotFound => {
RetainedSnapshotReadError::RetentionMissing
}
StoreError::Corrupt { message, .. } if message.contains("requires migration") => {
RetainedSnapshotReadError::MigrationRequired
}
StoreError::Io(_) | StoreError::Serde(_) | StoreError::Corrupt { .. } => {
RetainedSnapshotReadError::CorruptLog
}
_ => RetainedSnapshotReadError::CorruptLog,
}
}
/// Project a complete current-segment log. A valid segment starts with one
/// canonical annotated SegmentStart record; malformed partial input uses the
/// nil session only to keep the public failure projection deterministic.
@@ -559,4 +667,65 @@ mod tests {
SessionEntryProvenance::ModelOutput
);
}
#[test]
fn retained_snapshot_read_projects_active_segment_without_rewriting_files() {
let root = tempfile::tempdir().unwrap();
let aggregate_root = root.path().join("worker-a");
let aggregate = WorkerAggregateStore::new(&aggregate_root, "worker-a").unwrap();
let session_id = crate::new_session_id();
let segment_id = crate::new_segment_id();
aggregate
.write(&crate::WorkerMetadata::new(
"worker-a",
Some(crate::WorkerActiveSegmentRef::active_segment(
session_id, segment_id,
)),
))
.unwrap();
let session = WorkerSessionStore::new(aggregate_root.join("session")).unwrap();
session.create_segment(session_id, segment_id, &[]).unwrap();
let metadata_before = std::fs::read(aggregate_root.join("metadata.json")).unwrap();
let manifest_before = std::fs::read(aggregate_root.join("session/session.json")).unwrap();
let retained = read_retained_session_snapshot(
&aggregate_root,
"worker-a",
DEFAULT_RETAINED_SNAPSHOT_MAX_BYTES,
)
.unwrap();
assert_eq!(retained.identity.session_id, session_id.to_string());
assert_eq!(retained.identity.segment_id, segment_id.to_string());
assert_eq!(retained.identity.entry_count, 0);
assert!(retained.snapshot.entries.is_empty());
assert_eq!(
std::fs::read(aggregate_root.join("metadata.json")).unwrap(),
metadata_before
);
assert_eq!(
std::fs::read(aggregate_root.join("session/session.json")).unwrap(),
manifest_before
);
assert!(matches!(
read_retained_session_snapshot(&aggregate_root, "worker-a", 0),
Err(RetainedSnapshotReadError::SnapshotTooLarge)
));
}
#[test]
fn retained_snapshot_read_does_not_create_missing_aggregate() {
let root = tempfile::tempdir().unwrap();
let aggregate = root.path().join("missing-worker");
let error = read_retained_session_snapshot(
&aggregate,
"missing-worker",
DEFAULT_RETAINED_SNAPSHOT_MAX_BYTES,
)
.unwrap_err();
assert!(matches!(error, RetainedSnapshotReadError::RetentionMissing));
assert!(!aggregate.exists());
}
}
@@ -383,6 +383,27 @@ impl WorkerAggregateStore {
})
}
/// Open an existing aggregate without creating or rewriting retained files.
pub fn open_read_only(
root: impl Into<PathBuf>,
worker_name: impl Into<String>,
) -> Result<Self, WorkerStoreError> {
let root = root.into();
let worker_name = worker_name.into();
validate_worker_name(&worker_name)?;
if !root.is_dir() {
return Err(WorkerStoreError::Io(std::io::Error::new(
std::io::ErrorKind::NotFound,
"worker aggregate is not retained",
)));
}
Ok(Self {
update_lock: metadata_lock(&root),
root,
worker_name,
})
}
fn validate_name(&self, worker_name: &str) -> Result<(), WorkerStoreError> {
validate_worker_name(worker_name)?;
if worker_name == self.worker_name {
@@ -84,6 +84,32 @@ impl WorkerSessionStore {
})
}
/// Open an already-retained Session without creating directories or migrating
/// persisted data. Observation paths must never mutate retained state.
pub fn open_read_only(root: impl Into<PathBuf>) -> Result<Self, StoreError> {
let root = root.into();
let bytes = fs::read(root.join(SESSION_FILE))?;
let manifest: SessionManifest = serde_json::from_slice(&bytes)?;
if manifest.schema_version != SESSION_SCHEMA_VERSION {
return Err(StoreError::Corrupt {
line: 0,
message: format!(
"Worker Session schema version {} requires migration; expected {}",
manifest.schema_version, SESSION_SCHEMA_VERSION
),
});
}
Ok(Self {
root,
session_id: Arc::new(Mutex::new(Some(manifest.session_id))),
append_lock: Arc::new(Mutex::new(())),
})
}
pub fn segment_log_len(&self, segment_id: SegmentId) -> Result<u64, StoreError> {
Ok(fs::metadata(self.log_path(segment_id))?.len())
}
pub fn root_dir(&self) -> &Path {
&self.root
}
+1
View File
@@ -27,6 +27,7 @@ session-store = { workspace = true }
ticket = { workspace = true }
serde = { workspace = true, features = ["derive"] }
pulldown-cmark = { version = "0.13.3", default-features = false }
async-trait.workspace = true
agen.workspace = true
[dev-dependencies]
+25 -111
View File
@@ -3,9 +3,8 @@ use std::io;
use std::time::Duration;
use client::{
BackendRuntimeListTarget, BackendWorkerRestoreResponse, BackendWorkerRestoreState,
BackendWorkerSummary, list_backend_stopped_workers, list_backend_workers,
restore_backend_worker,
BackendRuntimeListTarget, BackendWorkerSummary, ServerApiClient, WorkerSessionAvailability,
list_backend_stopped_workers, list_backend_workers,
};
use crossterm::event::{self, Event as TermEvent, KeyCode, KeyEventKind, KeyModifiers};
use ratatui::Frame;
@@ -81,62 +80,39 @@ pub(crate) async fn run(
}
WorkerPickerResult::Selected(selected) => selected,
};
let worker = if selected.state == "stopped" {
let restore_target = target
.runtime_target(selected.runtime_id.clone(), selected.worker_id.clone())
let mut attach_target = target
.runtime_target(selected.runtime_id.clone(), selected.worker_id.clone())
.map_err(|error| io::Error::other(error.to_string()))?;
if selected.state == "stopped" {
let api = ServerApiClient::builder(&target.base_url)
.map_err(|error| io::Error::other(error.to_string()))?
.build()
.map_err(|error| io::Error::other(error.to_string()))?;
let restore = restore_backend_worker(&restore_target)
let response = api
.worker_session(
attach_target.workspace_id.clone(),
selected.runtime_id.clone(),
selected.worker_id.clone(),
)
.await
.map_err(|error| {
io::Error::other(format!(
"failed to restore Backend worker {}/{}: {error}",
"failed to observe retained Worker Session {}/{}: {error}",
selected.runtime_id, selected.worker_id
))
})?;
restored_worker(restore).map_err(|error| {
io::Error::other(format!(
"failed to restore Backend worker {}/{}: {error}",
selected.runtime_id, selected.worker_id
))
})?
} else {
selected
};
let attach_target = target
.runtime_target(worker.runtime_id, worker.worker_id)
.map_err(|error| io::Error::other(error.to_string()))?;
attach_target.initial_snapshot = match response.observation {
WorkerSessionAvailability::RetainedSnapshot { snapshot, .. } => Some(snapshot),
WorkerSessionAvailability::LiveProtocol => None,
WorkerSessionAvailability::Unavailable { message, .. } => {
return Err(io::Error::other(message).into());
}
};
}
return console::run_backend_runtime(attach_target).await;
}
}
fn restored_worker(response: BackendWorkerRestoreResponse) -> Result<BackendWorkerSummary, String> {
if response.result.state != BackendWorkerRestoreState::Accepted {
let diagnostics = response
.result
.diagnostics
.iter()
.map(|diagnostic| format!("{}: {}", diagnostic.code, diagnostic.message))
.collect::<Vec<_>>()
.join("; ");
let state = match response.result.state {
BackendWorkerRestoreState::Accepted => unreachable!(),
BackendWorkerRestoreState::Rejected => "rejected",
BackendWorkerRestoreState::RolledBack => "rolled back",
BackendWorkerRestoreState::ReconciliationRequired => "reconciliation required",
};
return Err(if diagnostics.is_empty() {
format!("restore was {state} without a diagnostic")
} else {
format!("restore was {state}: {diagnostics}")
});
}
response
.result
.worker
.ok_or_else(|| "restore was accepted without a Worker snapshot".to_string())
}
fn dedup_workers(workers: &mut Vec<BackendWorkerSummary>) {
let mut seen = std::collections::HashSet::new();
workers.retain(|worker| seen.insert((worker.runtime_id.clone(), worker.worker_id.clone())));
@@ -437,8 +413,7 @@ fn working_directory_text(worker: &BackendWorkerSummary) -> String {
mod tests {
use super::*;
use client::{
BackendDiagnostic, BackendDiagnosticSeverity, BackendWorkerCapabilitySummary,
BackendWorkerImplementationSummary, BackendWorkerRestoreResult,
BackendWorkerCapabilitySummary, BackendWorkerImplementationSummary,
BackendWorkerWorkspaceSummary,
};
@@ -494,67 +469,6 @@ mod tests {
text_width(&text[..byte_offset])
}
fn restore_response(
state: BackendWorkerRestoreState,
worker: Option<BackendWorkerSummary>,
diagnostics: Vec<BackendDiagnostic>,
) -> BackendWorkerRestoreResponse {
BackendWorkerRestoreResponse {
workspace_id: "workspace-a".to_string(),
runtime_id: "runtime-a".to_string(),
worker_id: "worker-a".to_string(),
result: BackendWorkerRestoreResult {
state,
worker,
diagnostics,
},
}
}
#[test]
fn rejected_restore_surfaces_diagnostic_instead_of_attaching_selected_worker() {
let error = restored_worker(restore_response(
BackendWorkerRestoreState::Rejected,
None,
vec![BackendDiagnostic {
code: "working_directory_not_found".to_string(),
severity: BackendDiagnosticSeverity::Error,
message: "working directory was not found".to_string(),
}],
))
.expect_err("rejected restore must not produce a Worker to attach");
assert_eq!(
error,
"restore was rejected: working_directory_not_found: working directory was not found"
);
}
#[test]
fn accepted_restore_requires_returned_worker_snapshot() {
let error = restored_worker(restore_response(
BackendWorkerRestoreState::Accepted,
None,
Vec::new(),
))
.expect_err("accepted restore without a Worker must not attach the stale selection");
assert_eq!(error, "restore was accepted without a Worker snapshot");
}
#[test]
fn accepted_restore_returns_authoritative_worker_snapshot() {
let worker = worker("runtime-a", "worker-a", Some("builtin:companion"));
let restored = restored_worker(restore_response(
BackendWorkerRestoreState::Accepted,
Some(worker.clone()),
Vec::new(),
))
.expect("accepted restore should return its Worker snapshot");
assert_eq!(restored, worker);
}
#[test]
fn worker_row_orders_and_simplifies_columns() {
let mut worker = worker("runtime-a", "worker-b", Some("builtin:coder"));
+65 -4
View File
@@ -457,20 +457,81 @@ async fn run_standalone_host(
pub(crate) async fn run_backend_runtime(
target: BackendRuntimeTarget,
) -> Result<(), Box<dyn std::error::Error>> {
if let Some(snapshot) = target.initial_snapshot.clone() {
let client = Client::new(RetainedSnapshotSocket::new(snapshot)?);
return run_backend_console(target, client, false).await;
}
let client = connect_backend_runtime(target.clone()).await?;
run_backend_console(target, client, true).await
}
async fn run_backend_console<T: Socket>(
target: BackendRuntimeTarget,
client: Client<T>,
connected: bool,
) -> Result<(), Box<dyn std::error::Error>> {
let worker_label = target.display_label();
let attachment_target = target.clone();
let client = connect_backend_runtime(target).await?;
let mut terminal = enter_fullscreen()?;
let workspace_root = std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from("."));
let mut app = App::new_with_persistent_input_history(worker_label, &workspace_root);
app.connected = true;
let mut connection = ConsoleConnection::with_backend_target(client, attachment_target);
app.connected = connected;
let mut connection = ConsoleConnection::with_backend_target(client, target);
let result = run_loop(&mut terminal, &mut app, &mut connection).await;
let _ = leave_fullscreen(&mut terminal);
result
}
struct RetainedSnapshotSocket {
event: Option<String>,
}
impl RetainedSnapshotSocket {
fn new(snapshot: protocol::SessionSnapshot) -> Result<Self, serde_json::Error> {
Ok(Self {
event: Some(serde_json::to_string(&Event::Snapshot {
session: snapshot,
greeting: protocol::Greeting {
worker_name: "retained worker".to_string(),
cwd: String::new(),
provider: "retained session".to_string(),
model: String::new(),
scope_summary: "read-only retained session".to_string(),
tools: Vec::new(),
context_window: 0,
context_tokens: 0,
},
state: protocol::WorkerStateSnapshot {
last_command_id: 0,
state: protocol::WorkerState::Idle,
},
in_flight: protocol::InFlightSnapshot::default(),
internal_workers: Vec::new(),
})?),
})
}
}
#[async_trait::async_trait]
impl Socket for RetainedSnapshotSocket {
type Error = std::convert::Infallible;
async fn send(&mut self, _message: String) -> Result<(), Self::Error> {
Ok(())
}
async fn next(&mut self) -> Result<Option<String>, Self::Error> {
if self.event.is_some() {
return Ok(self.event.take());
}
std::future::pending::<Result<Option<String>, Self::Error>>().await
}
fn try_next(&mut self) -> Result<Option<String>, Self::Error> {
Ok(self.event.take())
}
}
fn enter_fullscreen() -> Result<ConsoleTerminal, Box<dyn std::error::Error>> {
let mut stdout = io::stdout();
// Enable button-event tracking so the transcript can own drag selection;
+22
View File
@@ -318,6 +318,11 @@ impl WorkerExecutionSpawnResult {
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct WorkerSessionObservationRequest {
pub worker_ref: WorkerRef,
}
pub trait WorkerExecutionBackend: Send + Sync + 'static {
fn backend_id(&self) -> &str;
@@ -328,6 +333,16 @@ pub trait WorkerExecutionBackend: Send + Sync + 'static {
Err("execution backend does not support Workspace Config fetching".to_string())
}
fn worker_session(
&self,
_request: WorkerSessionObservationRequest,
) -> runtime_api::WorkerSessionAvailability {
runtime_api::WorkerSessionAvailability::Unavailable {
reason: runtime_api::WorkerSessionUnavailableReason::StorageUnavailable,
message: "retained session storage is unavailable".to_string(),
}
}
fn spawn_worker(&self, request: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult;
/// Read-only restore validation performed before any live controller,
@@ -525,6 +540,13 @@ impl WorkerExecutionBackendRef {
self.backend.fetch_workspace_config(request)
}
pub(crate) fn worker_session(
&self,
request: WorkerSessionObservationRequest,
) -> runtime_api::WorkerSessionAvailability {
self.backend.worker_session(request)
}
pub(crate) fn spawn_worker(
&self,
request: WorkerExecutionSpawnRequest,
+2 -1
View File
@@ -2566,7 +2566,7 @@ pub enum RuntimeHttpServerError {
mod tests {
use super::*;
use crate::auth::RuntimeIdentityMaterial;
use crate::catalog::{ConfigBundleRef, ProfileSelector, WorkerStatus, WorkspaceApiRef};
use crate::catalog::{ConfigBundleRef, ProfileSelector, WorkspaceApiRef};
use crate::config_bundle::{
ConfigBundle, ConfigBundleMetadata, ConfigBundleProvenance, ConfigProfileDescriptor,
};
@@ -2588,6 +2588,7 @@ mod tests {
use manifest::{Scope, SharedScope};
use sha2::Digest as _;
use tower::ServiceExt;
use workspace_api::WorkerRestoreState;
#[derive(Clone)]
struct TestRuntimeAuthorizer(String);
@@ -231,6 +231,35 @@ impl runtime_api::RuntimeApi for RuntimeManagementApi {
Ok(runtime_api::WorkerResponse { worker })
}
async fn worker_session(
&self,
worker_id: String,
request: runtime_api::WorkerSessionRequest,
) -> Result<runtime_api::WorkerSessionAvailability, runtime_api::RuntimeApiError> {
let worker_ref = worker_ref_for(&self.state.runtime, worker_id).map_err(api_error)?;
let scope = auth_workspace_scope(&self.state, auth_extension().as_ref())
.map_err(api_error)?
.ok_or_else(|| {
runtime_api::RuntimeApiError::new(
StatusCode::FORBIDDEN.as_u16(),
"runtime_worker_session_scope_required",
"Worker Session observation requires a Workspace-scoped capability",
)
})?;
if scope.workspace_id != request.workspace_id {
return Err(runtime_api::RuntimeApiError::new(
StatusCode::FORBIDDEN.as_u16(),
"runtime_worker_session_workspace_scope_mismatch",
"Worker Session Workspace scope does not match the authenticated capability",
));
}
self.state
.runtime
.worker_session_scoped(&scope, &worker_ref)
.map_err(RuntimeHttpRestError::runtime)
.map_err(api_error)
}
async fn create_worker(
&self,
value: runtime_api::CreateWorkerRequest,
+67 -1
View File
@@ -15,7 +15,8 @@ use crate::execution::WorkerExecutionRestoreRequest;
use crate::execution::{
WorkerExecutionBackend, WorkerExecutionBackendRef, WorkerExecutionHandle,
WorkerExecutionOperation, WorkerExecutionResult, WorkerExecutionSpawnRequest,
WorkerExecutionSpawnResult, WorkspaceConfigFetchRequest, WorkspaceConfigFetchResult,
WorkerExecutionSpawnResult, WorkerSessionObservationRequest, WorkspaceConfigFetchRequest,
WorkspaceConfigFetchResult,
};
#[cfg(feature = "fs-store")]
use crate::fs_store::{
@@ -1424,6 +1425,37 @@ impl Runtime {
Ok(state.worker(worker_ref)?.detail())
}
/// Observe a Worker's live protocol availability or its retained Session.
/// The same lifecycle lock used by stop/restore fences the decision and read.
pub fn worker_session_scoped(
&self,
scope: &RuntimeWorkspaceScope,
worker_ref: &WorkerRef,
) -> Result<runtime_api::WorkerSessionAvailability, RuntimeError> {
let operation_lock = self.worker_operation_lock(worker_ref.worker_id)?;
let _operation_guard = operation_lock
.lock()
.map_err(|_| RuntimeError::StatePoisoned)?;
self.ensure_worker_in_workspace(scope, worker_ref)?;
let backend = {
let state = self.lock()?;
let worker = state.worker(worker_ref)?;
if worker.execution_handle.is_some() && worker.status.is_active() {
return Ok(runtime_api::WorkerSessionAvailability::LiveProtocol);
}
state.execution_backend.clone()
};
Ok(match backend {
Some(backend) => backend.worker_session(WorkerSessionObservationRequest {
worker_ref: worker_ref.clone(),
}),
None => runtime_api::WorkerSessionAvailability::Unavailable {
reason: runtime_api::WorkerSessionUnavailableReason::StorageUnavailable,
message: "retained session storage is unavailable".to_string(),
},
})
}
/// Attach a live execution through a workspace-scoped Runtime authorization context.
pub fn restore_worker_scoped(
&self,
@@ -6743,6 +6775,40 @@ mod tests {
);
}
#[test]
fn session_observation_distinguishes_live_and_retained_without_restore() {
let (runtime, backend) = runtime_and_backend();
let scope = scope("local", "server-a");
let detail = runtime
.create_worker_scoped(&scope, task_request("observe without restore"))
.unwrap();
assert!(matches!(
runtime
.worker_session_scoped(&scope, &detail.worker_ref)
.unwrap(),
runtime_api::WorkerSessionAvailability::LiveProtocol
));
runtime
.send_protocol_method(
&detail.worker_ref,
Method::Shutdown {
command: test_command(),
},
)
.unwrap();
let observation = runtime
.worker_session_scoped(&scope, &detail.worker_ref)
.unwrap();
assert!(matches!(
observation,
runtime_api::WorkerSessionAvailability::Unavailable { .. }
));
assert_eq!(*backend.restore_count.lock().unwrap(), 0);
}
#[test]
fn stopped_worker_rejects_input_until_explicitly_restored() {
let (runtime, backend) = runtime_and_backend();
+65 -1
View File
@@ -24,7 +24,8 @@ use crate::config_bundle::{ConfigBundle, workspace_config_etag};
use crate::execution::{
WorkerExecutionBackend, WorkerExecutionHandle, WorkerExecutionOperation,
WorkerExecutionRestoreRequest, WorkerExecutionResult, WorkerExecutionSpawnRequest,
WorkerExecutionSpawnResult, WorkspaceConfigFetchRequest, WorkspaceConfigFetchResult,
WorkerExecutionSpawnResult, WorkerSessionObservationRequest, WorkspaceConfigFetchRequest,
WorkspaceConfigFetchResult,
};
use crate::identity::WorkerRef;
use crate::interaction::{WorkerInput, WorkerInputKind};
@@ -119,6 +120,14 @@ pub trait RuntimeWorkerFactory: Send + Sync + 'static {
Err("Runtime Worker factory does not support Workspace Config fetching".to_string())
}
fn retained_session_snapshot(
&self,
_worker_ref: &WorkerRef,
) -> Result<session_store::RetainedSessionSnapshot, session_store::RetainedSnapshotReadError>
{
Err(session_store::RetainedSnapshotReadError::RetentionMissing)
}
async fn spawn_controller(
&self,
request: WorkerExecutionSpawnRequest,
@@ -807,6 +816,22 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
fetch_workspace_config_http(&request, client).await
}
fn retained_session_snapshot(
&self,
worker_ref: &WorkerRef,
) -> Result<session_store::RetainedSessionSnapshot, session_store::RetainedSnapshotReadError>
{
let Some(root) = self.worker_aggregate_root.as_ref() else {
return Err(session_store::RetainedSnapshotReadError::RetentionMissing);
};
let worker_name = Self::runtime_worker_name_for_ref(worker_ref);
session_store::read_retained_session_snapshot(
&root.join(&worker_name),
&worker_name,
session_store::DEFAULT_RETAINED_SNAPSHOT_MAX_BYTES,
)
}
fn observe_workspace_prompt_projection(
&self,
projection: worker::WorkspacePromptProjection,
@@ -1923,6 +1948,45 @@ where
self.run_on_adapter_runtime(async move { factory.fetch_workspace_config(request).await })
}
fn worker_session(
&self,
request: WorkerSessionObservationRequest,
) -> runtime_api::WorkerSessionAvailability {
match self.factory.retained_session_snapshot(&request.worker_ref) {
Ok(retained) => runtime_api::WorkerSessionAvailability::RetainedSnapshot {
identity: runtime_api::RetainedSessionIdentity {
session_id: retained.identity.session_id,
segment_id: retained.identity.segment_id,
entry_count: retained.identity.entry_count,
},
snapshot: retained.snapshot,
},
Err(error) => {
let reason = match error {
session_store::RetainedSnapshotReadError::RetentionMissing => {
runtime_api::WorkerSessionUnavailableReason::RetentionMissing
}
session_store::RetainedSnapshotReadError::ActivePointerMissing => {
runtime_api::WorkerSessionUnavailableReason::ActivePointerMissing
}
session_store::RetainedSnapshotReadError::MigrationRequired => {
runtime_api::WorkerSessionUnavailableReason::MigrationRequired
}
session_store::RetainedSnapshotReadError::CorruptLog => {
runtime_api::WorkerSessionUnavailableReason::CorruptLog
}
session_store::RetainedSnapshotReadError::SnapshotTooLarge => {
runtime_api::WorkerSessionUnavailableReason::SnapshotTooLarge
}
};
runtime_api::WorkerSessionAvailability::Unavailable {
reason,
message: error.to_string(),
}
}
}
}
fn observe_workspace_prompt_projection(
&self,
projection: worker::WorkspacePromptProjection,
+1
View File
@@ -24,6 +24,7 @@ flow = { path = "../flow" }
manifest.workspace = true
protocol = { workspace = true }
runtime-api.workspace = true
server-api.workspace = true
project-record.workspace = true
reqwest = { version = "0.13", default-features = false, features = ["blocking", "json", "native-tls"] }
ring.workspace = true
+62
View File
@@ -1132,6 +1132,17 @@ pub trait WorkspaceWorkerRuntime: Send + Sync {
))
}
fn worker_session(
&self,
_workspace_id: &str,
_worker_ref: &EmbeddedWorkerRef,
) -> Result<runtime_api::WorkerSessionAvailability, String> {
Ok(runtime_api::WorkerSessionAvailability::Unavailable {
reason: runtime_api::WorkerSessionUnavailableReason::StorageUnavailable,
message: "retained session storage is unavailable".to_string(),
})
}
fn observation_source(
&self,
_worker_id: &str,
@@ -1922,6 +1933,29 @@ impl RuntimeRegistry {
})
}
pub fn worker_session(
&self,
workspace_id: &str,
worker: &RuntimeWorkerRef,
) -> Result<runtime_api::WorkerSessionAvailability, RuntimeRegistryError> {
validate_backend_identifier("runtime_id", &worker.runtime_id)?;
validate_backend_identifier("worker_id", &worker.worker_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(workspace_id, &worker_ref)
.map_err(|message| RuntimeRegistryError::RuntimeOperationFailed {
runtime_id: worker.runtime_id.clone(),
code: "worker_session_observation_failed".to_string(),
message,
})
}
pub fn observation_source(
&self,
worker: &RuntimeWorkerRef,
@@ -2773,6 +2807,17 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime {
.map_err(|error| error.to_string())
}
fn worker_session(
&self,
workspace_id: &str,
worker_ref: &EmbeddedWorkerRef,
) -> Result<runtime_api::WorkerSessionAvailability, String> {
let scope = RuntimeWorkspaceScope::new(workspace_id, "embedded-backend");
self.runtime
.worker_session_scoped(&scope, worker_ref)
.map_err(|error| error.to_string())
}
fn observation_source(
&self,
worker_id: &str,
@@ -4712,6 +4757,23 @@ impl WorkspaceWorkerRuntime for RemoteWorkerRuntime {
.map_err(|diagnostic| diagnostic.message)
}
fn worker_session(
&self,
workspace_id: &str,
worker_ref: &EmbeddedWorkerRef,
) -> Result<runtime_api::WorkerSessionAvailability, String> {
let worker_id = worker_ref.worker_id.to_string();
let request = runtime_api::WorkerSessionRequest {
workspace_id: workspace_id.to_string(),
};
self.run_runtime_api(
self.request_timeout,
MAX_REMOTE_RUNTIME_RESPONSE_BYTES,
move |client| async move { client.worker_session(worker_id, request).await },
)
.map_err(|diagnostic| diagnostic.message)
}
fn observation_source(
&self,
worker_id: &str,
+39
View File
@@ -3040,6 +3040,42 @@ fn build_server_auth_router(api: ServerAuthApi) -> Router {
.with_state(api)
}
impl server_api::ServerApi for WorkspaceApi {
async fn worker_session(
&self,
workspace_id: String,
runtime_id: String,
worker_id: String,
) -> std::result::Result<server_api::WorkspaceWorkerSessionResponse, server_api::ServerApiError>
{
if workspace_id != self.config.workspace_id {
return Err(runtime_api::RuntimeApiError::new(
StatusCode::NOT_FOUND.as_u16(),
"workspace_not_found",
"Workspace was not found",
));
}
let worker = RuntimeWorkerRef::new(runtime_id.clone(), worker_id.clone());
let observation = self
.runtime
.worker_session(&workspace_id, &worker)
.map_err(|_error| {
runtime_api::RuntimeApiError::new(
StatusCode::BAD_GATEWAY.as_u16(),
"worker_session_observation_failed",
"Worker Session observation failed",
)
})?;
Ok(server_api::WorkspaceWorkerSessionResponse {
subject: workspace_api::WorkspaceWorkerSubject::RuntimeWorker {
runtime_id,
worker_id,
},
observation,
})
}
}
fn build_inner_router(api: WorkspaceApi) -> Router {
let auth = build_server_auth_router(ServerAuthApi::from(&api));
let scoped_ticket_relations_query_path =
@@ -3632,6 +3668,9 @@ fn build_inner_router(api: WorkspaceApi) -> Router {
"/api/w/{workspace_id}/hosts/{host_id}/workers",
get(scoped_list_host_workers),
)
.merge(
server_api::ServerApiAxum::router(api.clone()).with_state::<WorkspaceApi>(()),
)
.fallback(get(static_or_spa_fallback))
.with_state(api);
auth.merge(workspace)
@@ -42,6 +42,7 @@
PendingSubmissionsSnapshot,
RewindTarget,
Segment,
SessionSnapshot,
} from "$lib/generated/protocol";
import {
MAX_FILES_PER_SUBMISSION,
@@ -933,8 +934,73 @@
target: ConsoleTarget,
) {
if (!targetWorker || targetWorker.state === "stopped") {
protocolState = "closed";
return;
protocolState = "connecting";
const controller = new AbortController();
void fetch(
workerApiPath(
`/runtimes/${encodeURIComponent(target.runtimeId)}/workers/${encodeURIComponent(target.workerId)}/session`,
),
{ signal: controller.signal, headers: { accept: "application/json" } },
)
.then(async (response) => {
if (!response.ok) {
throw new Error(`failed to load retained Worker Session (${response.status})`);
}
return (await response.json()) as
| { availability: "live_protocol" }
| { availability: "retained_snapshot"; snapshot: SessionSnapshot }
| { availability: "unavailable"; message: string };
})
.then((observation) => {
if (token !== reloadToken || controller.signal.aborted) return;
if (observation.availability === "unavailable") {
throw new Error(observation.message);
}
if (observation.availability === "live_protocol") {
throw new Error("Worker became live; reload to connect to its protocol stream.");
}
handleIncomingProtocolEvent({
event: "snapshot",
data: {
session: observation.snapshot,
greeting: {
worker_name: "retained worker",
cwd: "",
provider: "retained session",
model: "",
scope_summary: "read-only retained session",
tools: [],
context_window: 0,
context_tokens: 0,
},
state: { last_command_id: 0, state: { kind: "idle" } },
in_flight: { responses: [], commands: [] },
internal_workers: [],
},
} as ProtocolEvent);
protocolState = "closed";
streamDiagnostics = [
...streamDiagnostics,
{
code: "retained_session_read_only",
severity: "info",
message: "Showing a read-only retained Session snapshot.",
},
];
})
.catch((error) => {
if (token !== reloadToken || controller.signal.aborted) return;
protocolState = "error";
streamDiagnostics = [
...streamDiagnostics,
{
code: "retained_session_unavailable",
severity: "warning",
message: error instanceof Error ? error.message : String(error),
},
];
});
return () => controller.abort();
}
protocolState = "connecting";
const subscription = workspaceMultiplexer(target.workspaceId).subscribe(