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

feat: observe retained worker sessions without restore

Browse Source
This commit is contained in:
Hare
2026-09-18 04:27:38 +09:00
parent ae7f0f0446
commit 7a61a26670
26 changed files with 834 additions and 124 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker-runtime/src/execution.rs
+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,
crates/worker-runtime/src/http_server.rs
+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);
crates/worker-runtime/src/http_server/runtime_management_api.rs
+29
View File
@@ -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,
crates/worker-runtime/src/runtime.rs
+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();
crates/worker-runtime/src/worker_backend.rs
+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,