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

feat: guarded Worker Restore from shared Web actions (T-701)

Browse Source
This commit is contained in:
Hare
2026-10-06 17:45:54 +09:00
parent 1a6d391f83
commit fdce12a47b
54 changed files with 7214 additions and 527 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker-runtime/src/catalog.rs
+4
View File
@@ -334,6 +334,8 @@ pub(crate) enum WorkerRestoreIntent {
/// Lightweight persisted Worker identity projection.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct WorkerSummary {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub restore_observation_token: Option<String>,
pub worker_ref: WorkerRef,
pub worker_id: WorkerId,
pub status: WorkerStatus,
@@ -359,6 +361,8 @@ pub struct WorkerSummary {
/// Full persisted Worker identity and lifecycle detail.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct WorkerDetail {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub restore_observation_token: Option<String>,
pub worker_ref: WorkerRef,
pub worker_id: WorkerId,
pub status: WorkerStatus,
crates/worker-runtime/src/error.rs
+3
View File
@@ -12,6 +12,9 @@ pub enum RuntimeError {
#[error("initial worker input must be user input, got {kind}")]
InvalidInitialInputKind { kind: String },
#[error("worker {worker_id} Restore observation conflicts with the current operation")]
RestoreObservationConflict { worker_id: WorkerId },
#[error("worker {worker_id} was not found")]
WorkerNotFound { worker_id: WorkerId },
crates/worker-runtime/src/fs_store.rs
+7
View File
@@ -520,6 +520,7 @@ pub(crate) enum PersistedWorkerExecutionState {
#[derive(Clone, Debug)]
pub(crate) struct PersistedWorkerRecord {
pub(crate) restore_guard: crate::runtime::RestoreGuard,
pub(crate) worker_ref: WorkerRef,
pub(crate) worker_id: WorkerId,
pub(crate) profile: ProfileSelector,
@@ -1819,6 +1820,8 @@ impl RuntimeSnapshot {
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct WorkerAggregateRecord {
#[serde(default)]
restore_guard: crate::runtime::RestoreGuard,
schema_version: u32,
worker_ref: WorkerRef,
worker_id: WorkerId,
@@ -1874,6 +1877,7 @@ impl WorkerAggregateRecord {
};
Self {
schema_version: SCHEMA_VERSION,
restore_guard: worker.restore_guard.clone(),
worker_ref: worker.worker_ref.clone(),
worker_id: worker.worker_id,
profile: worker.profile.clone(),
@@ -1902,6 +1906,7 @@ impl WorkerAggregateRecord {
}
fn validate(self, path: &Path) -> Result<PersistedWorkerRecord, RuntimeError> {
let restore_guard = self.restore_guard;
let last_finished_submission_request_id = self.last_finished_submission_request_id;
let identity = WorkerIdentityRecord {
schema_version: self.schema_version,
@@ -2007,6 +2012,7 @@ impl WorkerAggregateRecord {
WorkerExecutionStateRecord::Unavailable => PersistedWorkerExecutionState::Unavailable,
};
let mut persisted = identity.into_persisted(execution_state);
persisted.restore_guard = restore_guard;
persisted.last_finished_submission_request_id = last_finished_submission_request_id;
Ok(persisted)
}
@@ -2180,6 +2186,7 @@ impl WorkerIdentityRecord {
self.workdir_attachments
};
PersistedWorkerRecord {
restore_guard: crate::runtime::RestoreGuard::default(),
worker_ref: self.worker_ref,
worker_id: self.worker_id,
profile: self.profile,
crates/worker-runtime/src/http_server.rs
+136 -7
View File
@@ -2302,6 +2302,7 @@ fn required_runtime_permission(method: &Method, path: &str) -> Option<&'static s
}
if path.ends_with("/input")
|| path.ends_with("/restore")
|| path.ends_with("/restore/coordinate")
|| (path.contains("/attachments") && *method != Method::GET)
|| path.contains("/workdir-attachments")
{
@@ -2514,6 +2515,7 @@ fn status_for_runtime_error(error: &RuntimeError) -> StatusCode {
StatusCode::CONFLICT
}
RuntimeError::RuntimeStopped
| RuntimeError::RestoreObservationConflict { .. }
| RuntimeError::RuntimeStoreAlreadyOpen { .. }
| RuntimeError::WorkerExecutionUnavailable { .. }
| RuntimeError::ExecutionBackendUnavailable { .. }
@@ -2538,6 +2540,9 @@ fn status_for_runtime_error(error: &RuntimeError) -> StatusCode {
fn code_for_runtime_error(error: &RuntimeError) -> String {
match error {
RuntimeError::RuntimeStopped => "runtime_stopped".to_string(),
RuntimeError::RestoreObservationConflict { .. } => {
"restore_observation_conflict".to_string()
}
RuntimeError::RuntimeStoreAlreadyOpen { .. } => "runtime_store_already_open".to_string(),
RuntimeError::WorkerNotFound { .. } => "worker_not_found".to_string(),
RuntimeError::WorkerExecutionUnavailable { .. } => {
@@ -3642,7 +3647,17 @@ mod tests {
client
.restore_worker(
generated_worker_id.clone(),
runtime_api::EmptyObjectRequest::default(),
runtime_api::WorkerRestoreRequest {
expected_observation_token: client
.get_worker(generated_worker_id.clone())
.await
.unwrap()
.worker
.restore_observation_token
.unwrap(),
request_id: "generated-restore".into(),
preparation: Some(runtime_api::WorkerRestorePreparation::default()),
},
)
.await
.unwrap();
@@ -3786,18 +3801,86 @@ mod tests {
let stop: RuntimeHttpWorkerLifecycleResponse = read_json(response).await;
assert_eq!(stop.ack.worker_ref, created.worker.worker_ref);
let response = authed_empty_request(
let restore_uri = format!("/v1/workers/{}/restore", created.worker.worker_id);
let missing = authed_empty_request(app.clone(), Method::POST, &restore_uri, token).await;
assert_eq!(missing.status(), StatusCode::BAD_REQUEST);
let stale = runtime_api::WorkerRestoreRequest {
expected_observation_token: created.worker.restore_observation_token.clone().unwrap(),
request_id: "stale-http-restore".into(),
preparation: None,
};
let response =
authed_json_request(app.clone(), Method::POST, &restore_uri, token, &stale).await;
assert_eq!(response.status(), StatusCode::CONFLICT);
let conflict: RuntimeHttpErrorResponse = read_json(response).await;
assert_eq!(conflict.error.code, "restore_observation_conflict");
let mut restore_request = runtime.test_restore_request(&created.worker.worker_ref);
let coordinate_uri = format!("{restore_uri}/coordinate");
let mut coordination = runtime_api::WorkerRestoreCoordinationRequest {
expected_observation_token: restore_request.expected_observation_token.clone(),
request_id: restore_request.request_id.clone(),
preparation: None,
};
let before_lookup = runtime.worker_detail(&created.worker.worker_ref).unwrap();
let response = authed_json_request(
app.clone(),
Method::POST,
&format!("/v1/workers/{}/restore", created.worker.worker_id),
&coordinate_uri,
token,
&coordination,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let missing_owner: runtime_api::WorkerRestoreCoordinationResponse =
read_json(response).await;
assert!(missing_owner.result.is_none() && missing_owner.preparation.is_none());
assert_eq!(
runtime.worker_detail(&created.worker.worker_ref).unwrap(),
before_lookup
);
let response = authed_json_request(
app.clone(),
Method::POST,
&coordinate_uri,
"wrong-token",
&coordination,
)
.await;
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
coordination.preparation = Some(runtime_api::WorkerRestorePreparation::default());
let response = authed_json_request(
app.clone(),
Method::POST,
&coordinate_uri,
token,
&coordination,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let admitted: runtime_api::WorkerRestoreCoordinationResponse = read_json(response).await;
assert!(admitted.result.is_none() && admitted.preparation.is_some());
assert_eq!(
runtime
.worker_detail(&created.worker.worker_ref)
.unwrap()
.status,
crate::catalog::WorkerStatus::Stopped
);
restore_request.preparation = admitted.preparation;
let response = authed_json_request(
app.clone(),
Method::POST,
&restore_uri,
token,
&restore_request,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let restored: runtime_api::WorkerRestoreResponse = read_json(response).await;
assert_eq!(restored.state, WorkerRestoreState::Accepted);
assert_eq!(
restored.worker.unwrap().status,
restored.worker.as_ref().unwrap().status,
runtime_api::WorkerStatus::Idle
);
@@ -3810,6 +3893,47 @@ mod tests {
.await;
assert_eq!(response.status(), StatusCode::OK);
coordination.preparation = None;
let response = authed_json_request(
app.clone(),
Method::POST,
&coordinate_uri,
token,
&coordination,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let recovered: runtime_api::WorkerRestoreCoordinationResponse = read_json(response).await;
assert!(recovered.preparation.is_none());
assert_eq!(recovered.result, Some(restored.clone()));
restore_request.preparation = Some(serde_json::from_value(serde_json::json!({
"workspace_api": { "workspace_id": "must-not-bind", "base_url": "https://fresh.invalid.example" },
"workdir_attachments": [],
"repository_access": []
})).unwrap());
let response = authed_json_request(
app.clone(),
Method::POST,
&restore_uri,
token,
&restore_request,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let replayed: runtime_api::WorkerRestoreResponse = read_json(response).await;
assert_eq!(
serde_json::to_value(replayed).unwrap(),
serde_json::to_value(restored).unwrap()
);
assert_eq!(
runtime
.worker_detail(&created.worker.worker_ref)
.unwrap()
.status,
crate::catalog::WorkerStatus::Stopped
);
let response = authed_empty_request(
app.clone(),
Method::POST,
@@ -3892,11 +4016,12 @@ mod tests {
assert_eq!(listed.workers[0].status, crate::catalog::WorkerStatus::Idle);
assert!(listed.workers[0].worker_state.is_none());
let response = authed_empty_request(
let response = authed_json_request(
app.clone(),
Method::POST,
&format!("{worker_uri}/restore"),
token,
&runtime.test_restore_request(&created.worker_ref),
)
.await;
assert_eq!(response.status(), StatusCode::OK);
@@ -4864,7 +4989,9 @@ mod ws_tests {
let (mut old_stream, _) = connect_async(authed_ws_request(&url)).await.unwrap();
let _ = next_frame(&mut old_stream).await;
runtime.stop_worker(&worker_ref, None).unwrap();
runtime.restore_worker(&worker_ref).unwrap();
runtime
.restore_worker(&worker_ref, runtime.test_restore_request(&worker_ref))
.unwrap();
backend.publish(
&worker_ref,
protocol::Event::TextDelta {
@@ -4946,7 +5073,9 @@ mod ws_tests {
);
backend.publish(&worker_ref, protocol::Event::Shutdown);
runtime.stop_worker(&worker_ref, None).unwrap();
runtime.restore_worker(&worker_ref).unwrap();
runtime
.restore_worker(&worker_ref, runtime.test_restore_request(&worker_ref))
.unwrap();
let resume_url = format!("{url}?cursor={}", old.cursor);
let (mut stream, _) = connect_async(authed_ws_request(&resume_url)).await.unwrap();
assert!(matches!(
crates/worker-runtime/src/http_server/runtime_management_api.rs
+63 -10
View File
@@ -462,21 +462,74 @@ impl runtime_api::RuntimeApi for RuntimeManagementApi {
Ok(runtime_api::WorkerLifecycleResponse { ack })
}
async fn coordinate_worker_restore(
&self,
worker_id: String,
request: runtime_api::WorkerRestoreCoordinationRequest,
) -> Result<runtime_api::WorkerRestoreCoordinationResponse, 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)?;
let runtime = self.state.runtime.clone();
let (result, preparation) = tokio::task::spawn_blocking(move || {
runtime.coordinate_worker_restore_operation(&worker_ref, request, scope.as_ref())
})
.await
.map_err(|error| {
runtime_api::RuntimeApiError::new(500, "runtime_restore_task_failed", error.to_string())
})?
.map_err(RuntimeHttpRestError::runtime)
.map_err(api_error)?;
let result = result
.map(|result| -> Result<_, runtime_api::RuntimeApiError> {
Ok(runtime_api::WorkerRestoreResponse {
state: match result.state {
server_api::WorkerRestoreState::Accepted => {
runtime_api::WorkerRestoreState::Accepted
}
server_api::WorkerRestoreState::Rejected => {
runtime_api::WorkerRestoreState::Rejected
}
server_api::WorkerRestoreState::RolledBack => {
runtime_api::WorkerRestoreState::RolledBack
}
server_api::WorkerRestoreState::ReconciliationRequired => {
runtime_api::WorkerRestoreState::ReconciliationRequired
}
},
worker: result.worker.map(response).transpose()?,
reason_code: result.reason_code,
message: result.message,
})
})
.transpose()?;
Ok(runtime_api::WorkerRestoreCoordinationResponse {
result,
preparation,
})
}
async fn restore_worker(
&self,
worker_id: String,
_request: runtime_api::EmptyObjectRequest,
request: runtime_api::WorkerRestoreRequest,
) -> Result<runtime_api::WorkerRestoreResponse, runtime_api::RuntimeApiError> {
let worker_ref = worker_ref_for(&self.state.runtime, worker_id).map_err(api_error)?;
let result = match auth_workspace_scope(&self.state, auth_extension().as_ref())
.map_err(api_error)?
{
Some(scope) => self
.state
.runtime
.restore_worker_operation_scoped(&scope, &worker_ref),
None => self.state.runtime.restore_worker_operation(&worker_ref),
}
let scope =
auth_workspace_scope(&self.state, auth_extension().as_ref()).map_err(api_error)?;
let runtime = self.state.runtime.clone();
let result = tokio::task::spawn_blocking(move || match scope {
Some(scope) => runtime.restore_worker_operation_scoped(&scope, &worker_ref, request),
None => runtime.restore_worker_operation(&worker_ref, request),
})
.await
.map_err(|error| {
runtime_api::RuntimeApiError::new(
StatusCode::INTERNAL_SERVER_ERROR.as_u16(),
"runtime_restore_task_failed",
error.to_string(),
)
})?
.map_err(RuntimeHttpRestError::runtime)
.map_err(api_error)?;
Ok(runtime_api::WorkerRestoreResponse {
crates/worker-runtime/src/runtime.rs
+2152 -99
View File
File diff suppressed because it is too large Load Diff
crates/worker-runtime/src/worker_backend.rs
+20 -3
View File
@@ -7394,7 +7394,14 @@ mod tests {
.create_worker(create_request("restart cleanup"))
.unwrap();
runtime.stop_worker(&worker.worker_ref, None).unwrap();
assert!(runtime.restore_worker(&worker.worker_ref).is_err());
assert!(
runtime
.restore_worker(
&worker.worker_ref,
runtime.test_restore_request(&worker.worker_ref)
)
.is_err()
);
assert!(runtime.stop_worker(&worker.worker_ref, None).is_err());
let prior_restore_id = restore_id.lock().unwrap().unwrap();
let aggregate = runtime_store
@@ -8904,7 +8911,12 @@ mod tests {
.unwrap();
assert!(detached.workdir_attachments.is_empty());
runtime.stop_worker(&detail.worker_ref, None).unwrap();
let restored_without_workdir = runtime.restore_worker(&detail.worker_ref).unwrap();
let restored_without_workdir = runtime
.restore_worker(
&detail.worker_ref,
runtime.test_restore_request(&detail.worker_ref),
)
.unwrap();
assert!(restored_without_workdir.workdir_attachments.is_empty());
assert!(
!repo.path().join(".yoi").exists(),
@@ -9191,7 +9203,12 @@ mod tests {
.contains_key(&detail.worker_ref)
);
runtime.restore_worker(&detail.worker_ref).unwrap();
runtime
.restore_worker(
&detail.worker_ref,
runtime.test_restore_request(&detail.worker_ref),
)
.unwrap();
assert_eq!(
runtime.worker_detail(&detail.worker_ref).unwrap().status,
crate::catalog::WorkerStatus::Idle
crates/worker-runtime/src/worker_source.rs
+7 -1
View File
@@ -555,9 +555,15 @@ impl WorkspaceClient for RuntimeOwnedWorkspaceClient {
WorkspaceServerOperation::WorkerControlRestore {
runtime_id,
worker_id,
request,
} => generated_workspace_response(
client
.worker_control_restore(workspace_id, runtime_id, worker_id)
.worker_control_restore(
workspace_id,
runtime_id,
worker_id,
request,
)
.await,
),
WorkspaceServerOperation::WorkerObservationSessions => {