refactor(server): unify registry-bound workspace worker operations
This commit is contained in:
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,864 @@
|
||||
use super::*;
|
||||
|
||||
fn operation_input(kind: WorkerInputKind) -> WorkerInputRequest {
|
||||
WorkerInputRequest {
|
||||
kind,
|
||||
content: "operation-boundary-test".into(),
|
||||
submission_request_id: None,
|
||||
segments: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn operation_lifecycle() -> WorkerLifecycleRequest {
|
||||
WorkerLifecycleRequest {
|
||||
reason: Some("operation-boundary-test".into()),
|
||||
ticket_assignment: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn operation_cleanup_candidate(identity: &RuntimeWorkerRef) -> CleanupWorkerCandidate {
|
||||
CleanupWorkerCandidate {
|
||||
target_id: format!("worker:{}:{}", identity.runtime_id, identity.worker_id),
|
||||
action: server_api::CleanupTargetKind::WorkerDelete,
|
||||
worker_id: identity.worker_id.clone(),
|
||||
runtime_worker_id: identity.worker_id.clone(),
|
||||
runtime_id: identity.runtime_id.clone(),
|
||||
reason: "operation-boundary-test".into(),
|
||||
blocking_reason: None,
|
||||
pinned: false,
|
||||
retention_state: "normal".into(),
|
||||
linked_workdir_ids: Vec::new(),
|
||||
running_linked: false,
|
||||
estimated_reclaim_bytes: None,
|
||||
}
|
||||
}
|
||||
|
||||
async fn browser_operation_context(api: &WorkspaceApi) -> (WorkerOperationContext, String) {
|
||||
seed_test_api_token(api.store.as_ref(), "worker-operation-browser");
|
||||
let token = "worker-operation-browser-session";
|
||||
api.store
|
||||
.create_browser_session(&BrowserSessionRecord {
|
||||
token_hash: crate::auth::token_hash(token),
|
||||
session_id: "worker-operation-session".into(),
|
||||
user_id: "user-worker-operation-browser".into(),
|
||||
created_at: "2026-01-01T00:00:00Z".into(),
|
||||
expires_at: "2099-01-01T00:00:00Z".into(),
|
||||
revoked_at: None,
|
||||
})
|
||||
.unwrap();
|
||||
let AuthConfig::Passkey {
|
||||
cookie_name,
|
||||
origin,
|
||||
..
|
||||
} = &api.config.auth;
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert(ORIGIN, origin.parse().unwrap());
|
||||
headers.insert(
|
||||
axum::http::header::COOKIE,
|
||||
format!("{cookie_name}={token}").parse().unwrap(),
|
||||
);
|
||||
let actor = resolve_actor(&ServerAuthApi::from(api), &headers)
|
||||
.await
|
||||
.unwrap()
|
||||
.expect("valid browser session");
|
||||
(
|
||||
WorkerOperationContext::Browser {
|
||||
headers,
|
||||
source: authenticated_browser_input_source(&actor),
|
||||
},
|
||||
crate::auth::token_hash(token),
|
||||
)
|
||||
}
|
||||
|
||||
fn browser_operation_methods() -> Vec<protocol::Method> {
|
||||
vec![
|
||||
protocol::Method::Pause {
|
||||
command: protocol::WorkerCommandEnvelope::new(701),
|
||||
},
|
||||
protocol::Method::Resume {
|
||||
command: protocol::WorkerCommandEnvelope::new(702),
|
||||
},
|
||||
protocol::Method::Cancel {
|
||||
command: protocol::WorkerCommandEnvelope::new(703),
|
||||
},
|
||||
protocol::Method::Compact {
|
||||
command: protocol::WorkerCommandEnvelope::new(704),
|
||||
},
|
||||
protocol::Method::Submit {
|
||||
submission_request_id: "worker-operation-submit".into(),
|
||||
input: vec![protocol::Segment::Text {
|
||||
content: "must not reach execution".into(),
|
||||
}],
|
||||
},
|
||||
protocol::Method::Notify {
|
||||
notification_request_id: "worker-operation-notify".into(),
|
||||
message: "must not reach execution".into(),
|
||||
},
|
||||
protocol_test_completion("worker-operation-completion"),
|
||||
]
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn handle_resource_keys_preserve_identity_for_equal_local_ids_across_runtimes() {
|
||||
let workspace = tempfile::tempdir().unwrap();
|
||||
let api = test_api(workspace.path()).await;
|
||||
let a = RuntimeWorkerRef::new("offline-runtime-a", "shared-local-id");
|
||||
let b = RuntimeWorkerRef::new("offline-runtime-b", "shared-local-id");
|
||||
for identity in [&a, &b] {
|
||||
seed_worker_source_member(&api, &identity.runtime_id, &identity.worker_id);
|
||||
}
|
||||
let key = |identity: &RuntimeWorkerRef| {
|
||||
api.store
|
||||
.worker_resource_key(TEST_WORKSPACE_ID, identity)
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
};
|
||||
assert_ne!(key(&a), key(&b));
|
||||
|
||||
for identity in [&a, &b] {
|
||||
let direct =
|
||||
WorkspaceWorker::resolve(&api, &identity.runtime_id, &identity.worker_id).unwrap();
|
||||
let canonical =
|
||||
WorkspaceWorker::resolve(&api, &identity.runtime_id, &key(identity)).unwrap();
|
||||
let pinned = canonical
|
||||
.set_pinned(&WorkerOperationContext::Backend, true)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(pinned.runtime_id, identity.runtime_id);
|
||||
assert_eq!(pinned.worker_id, identity.worker_id);
|
||||
let unpinned = direct
|
||||
.set_pinned(&WorkerOperationContext::Backend, false)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(unpinned.runtime_id, pinned.runtime_id);
|
||||
assert_eq!(unpinned.worker_id, pinned.worker_id);
|
||||
}
|
||||
assert!(matches!(
|
||||
WorkspaceWorker::resolve(&api, &b.runtime_id, &key(&a)),
|
||||
Err(Error::UnknownWorker { .. })
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn handle_resolution_rejects_unknown_foreign_and_internal_only_workers() {
|
||||
let workspace = tempfile::tempdir().unwrap();
|
||||
let api = test_api(workspace.path()).await;
|
||||
let known = spawn_ticket_check_source(&api, "operation-resolution-parent");
|
||||
let mut foreign = api
|
||||
.store
|
||||
.get_worker_registry(TEST_WORKSPACE_ID, &known)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let mut other_workspace = api
|
||||
.store
|
||||
.get_workspace(TEST_WORKSPACE_ID)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
other_workspace.workspace_id = "other-workspace".into();
|
||||
api.store.upsert_workspace(&other_workspace).await.unwrap();
|
||||
foreign.workspace_id = other_workspace.workspace_id;
|
||||
foreign.worker.worker_id = "foreign-only-worker".into();
|
||||
api.store.upsert_worker_registry(&foreign).unwrap();
|
||||
assert!(
|
||||
api.store
|
||||
.get_worker_registry(TEST_WORKSPACE_ID, &foreign.worker)
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
|
||||
let RuntimeObservationSource::Embedded(source) =
|
||||
api.runtime.observation_source(&known).unwrap()
|
||||
else {
|
||||
panic!("expected embedded fixture");
|
||||
};
|
||||
let internal = protocol::InternalWorkerRef {
|
||||
session_id: "internal-only-session".into(),
|
||||
parent_session_id: None,
|
||||
name: "internal-only-child".into(),
|
||||
kind: protocol::InternalWorkerKind::SubWorker,
|
||||
};
|
||||
source
|
||||
.runtime
|
||||
.observe_worker_event(
|
||||
&source.worker_ref,
|
||||
protocol::Event::InternalWorker {
|
||||
worker: internal.clone(),
|
||||
revision: 1,
|
||||
event: Box::new(protocol::Event::WorkerState {
|
||||
snapshot: protocol::WorkerStateSnapshot::initial(),
|
||||
}),
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
for reference in [
|
||||
"unknown-worker",
|
||||
&foreign.worker.worker_id,
|
||||
&internal.session_id,
|
||||
&internal.name,
|
||||
] {
|
||||
assert!(
|
||||
matches!(
|
||||
WorkspaceWorker::resolve(&api, &known.runtime_id, reference),
|
||||
Err(Error::UnknownWorker { .. })
|
||||
),
|
||||
"unexpected registry handle for {reference}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn offline_handle_can_pin_but_deleted_handle_cannot_mutate_or_connect() {
|
||||
let workspace = tempfile::tempdir().unwrap();
|
||||
let api = test_api(workspace.path()).await;
|
||||
let identity = RuntimeWorkerRef::new("unregistered-offline-runtime", "42");
|
||||
seed_worker_source_member(&api, &identity.runtime_id, &identity.worker_id);
|
||||
assert!(
|
||||
api.runtime.worker(&identity).is_err(),
|
||||
"fixture must have no live Runtime"
|
||||
);
|
||||
let handle = WorkspaceWorker::resolve(&api, &identity.runtime_id, &identity.worker_id).unwrap();
|
||||
let retained = handle
|
||||
.set_pinned(&WorkerOperationContext::Backend, true)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(retained.pinned);
|
||||
assert_eq!(
|
||||
api.store
|
||||
.get_worker_registry(TEST_WORKSPACE_ID, &identity)
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.retention_state,
|
||||
"pinned"
|
||||
);
|
||||
|
||||
let mut wrong_candidate = operation_cleanup_candidate(&identity);
|
||||
wrong_candidate.runtime_worker_id = "different-worker".into();
|
||||
assert!(matches!(
|
||||
handle.delete_from_plan(&WorkerOperationContext::Backend, &wrong_candidate)
|
||||
.await.unwrap_err().error,
|
||||
Error::InvalidInput(message) if message.contains("resolved Worker")
|
||||
));
|
||||
|
||||
api.store
|
||||
.delete_worker_registry(TEST_WORKSPACE_ID, &identity)
|
||||
.unwrap();
|
||||
assert!(
|
||||
matches!(handle.set_pinned(&WorkerOperationContext::Backend, false)
|
||||
.await.unwrap_err().error, Error::UnknownWorker { worker } if worker == identity)
|
||||
);
|
||||
assert!(matches!(
|
||||
handle
|
||||
.restore(&WorkerOperationContext::Backend)
|
||||
.await
|
||||
.unwrap_err()
|
||||
.error,
|
||||
Error::UnknownWorker { .. }
|
||||
));
|
||||
assert!(matches!(
|
||||
handle
|
||||
.stop(&WorkerOperationContext::Backend, operation_lifecycle())
|
||||
.await
|
||||
.unwrap_err()
|
||||
.error,
|
||||
Error::UnknownWorker { .. }
|
||||
));
|
||||
assert!(matches!(
|
||||
handle
|
||||
.cancel(&WorkerOperationContext::Backend, operation_lifecycle())
|
||||
.await
|
||||
.unwrap_err()
|
||||
.error,
|
||||
Error::UnknownWorker { .. }
|
||||
));
|
||||
for kind in [WorkerInputKind::User, WorkerInputKind::Notify] {
|
||||
assert!(matches!(
|
||||
handle
|
||||
.input(&WorkerOperationContext::Backend, operation_input(kind))
|
||||
.await
|
||||
.unwrap_err()
|
||||
.error,
|
||||
Error::UnknownWorker { .. }
|
||||
));
|
||||
}
|
||||
assert!(matches!(
|
||||
handle
|
||||
.delete_from_plan(
|
||||
&WorkerOperationContext::Backend,
|
||||
&operation_cleanup_candidate(&identity)
|
||||
)
|
||||
.await
|
||||
.unwrap_err()
|
||||
.error,
|
||||
Error::UnknownWorker { .. }
|
||||
));
|
||||
assert!(matches!(
|
||||
handle
|
||||
.connect_protocol(&WorkerOperationContext::Backend)
|
||||
.await,
|
||||
Err(Error::UnknownWorker { .. })
|
||||
));
|
||||
assert!(
|
||||
api.store
|
||||
.get_worker_registry(TEST_WORKSPACE_ID, &identity)
|
||||
.unwrap()
|
||||
.is_none(),
|
||||
"a stale handle must not recreate its registry row"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn worker_control_handle_rechecks_revoked_grants_for_each_operation() {
|
||||
let workspace = tempfile::tempdir().unwrap();
|
||||
let api = test_api(workspace.path()).await;
|
||||
let controller = RuntimeWorkerRef::new("offline-runtime", "controller");
|
||||
let subject = RuntimeWorkerRef::new("offline-runtime", "subject");
|
||||
for identity in [&controller, &subject] {
|
||||
seed_worker_source_member(&api, &identity.runtime_id, &identity.worker_id);
|
||||
}
|
||||
api.store
|
||||
.create_worker_control_grant(&WorkerControlGrantRecord {
|
||||
workspace_id: TEST_WORKSPACE_ID.into(),
|
||||
grant_id: "operation-grant".into(),
|
||||
controller: controller.clone(),
|
||||
subject: subject.clone(),
|
||||
relation: "spawned".into(),
|
||||
origin: "test".into(),
|
||||
permissions: [
|
||||
"restore",
|
||||
"stop",
|
||||
"cancel",
|
||||
"send_input",
|
||||
"notify",
|
||||
"pin",
|
||||
"remove",
|
||||
]
|
||||
.into_iter()
|
||||
.map(str::to_owned)
|
||||
.collect(),
|
||||
operation_id: "operation-grant-test".into(),
|
||||
created_at: now_registry_timestamp(),
|
||||
revoked_at: None,
|
||||
})
|
||||
.unwrap();
|
||||
let handle = WorkspaceWorker::resolve(&api, &subject.runtime_id, &subject.worker_id).unwrap();
|
||||
let context = WorkerOperationContext::WorkerControl { controller };
|
||||
handle.set_pinned(&context, true).await.unwrap();
|
||||
api.store
|
||||
.revoke_worker_control_grant(
|
||||
TEST_WORKSPACE_ID,
|
||||
"operation-grant",
|
||||
&now_registry_timestamp(),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
for permission in [
|
||||
"restore",
|
||||
"stop",
|
||||
"cancel",
|
||||
"send_input",
|
||||
"notify",
|
||||
"pin",
|
||||
"remove",
|
||||
] {
|
||||
assert!(
|
||||
matches!(
|
||||
handle.authorize_operation(&context, permission).await,
|
||||
Err(ApiError {
|
||||
error: Error::UnknownWorker { .. },
|
||||
..
|
||||
})
|
||||
),
|
||||
"revoked grant authorized {permission}"
|
||||
);
|
||||
}
|
||||
assert!(matches!(
|
||||
handle.restore(&context).await.unwrap_err().error,
|
||||
Error::UnknownWorker { .. }
|
||||
));
|
||||
assert!(matches!(
|
||||
handle
|
||||
.stop(&context, operation_lifecycle())
|
||||
.await
|
||||
.unwrap_err()
|
||||
.error,
|
||||
Error::UnknownWorker { .. }
|
||||
));
|
||||
assert!(matches!(
|
||||
handle
|
||||
.cancel(&context, operation_lifecycle())
|
||||
.await
|
||||
.unwrap_err()
|
||||
.error,
|
||||
Error::UnknownWorker { .. }
|
||||
));
|
||||
for kind in [WorkerInputKind::User, WorkerInputKind::Notify] {
|
||||
assert!(matches!(
|
||||
handle
|
||||
.input(&context, operation_input(kind))
|
||||
.await
|
||||
.unwrap_err()
|
||||
.error,
|
||||
Error::UnknownWorker { .. }
|
||||
));
|
||||
}
|
||||
assert!(matches!(
|
||||
handle.set_pinned(&context, false).await.unwrap_err().error,
|
||||
Error::UnknownWorker { .. }
|
||||
));
|
||||
assert!(matches!(
|
||||
handle
|
||||
.delete_from_plan(&context, &operation_cleanup_candidate(&subject))
|
||||
.await
|
||||
.unwrap_err()
|
||||
.error,
|
||||
Error::UnknownWorker { .. }
|
||||
));
|
||||
assert_eq!(
|
||||
api.store
|
||||
.get_worker_registry(TEST_WORKSPACE_ID, &subject)
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.retention_state,
|
||||
"pinned"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn generic_and_scoped_input_aliases_match_handle_idle_only_and_notify_semantics() {
|
||||
let fixture = guarded_spawn_fixture().await;
|
||||
let response = post_guarded_spawn(
|
||||
&fixture,
|
||||
guarded_spawn_payload(&fixture, "operation-alias-parity"),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap();
|
||||
let created: BrowserCreateWorkerResponse = serde_json::from_slice(&bytes).unwrap();
|
||||
let api = fixture.api.clone();
|
||||
let runtime = &fixture.runtime;
|
||||
let subject = RuntimeWorkerRef::new(created.runtime_id, created.worker_id);
|
||||
let controller = fixture.controller.clone();
|
||||
let key = api
|
||||
.store
|
||||
.worker_resource_key(TEST_WORKSPACE_ID, &subject)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let handle = WorkspaceWorker::resolve(&api, &subject.runtime_id, &key).unwrap();
|
||||
let context = WorkerOperationContext::WorkerControl { controller };
|
||||
let idle_only = handle
|
||||
.input(&context, operation_input(WorkerInputKind::User))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(idle_only.state, InternalWorkerOperationState::Rejected);
|
||||
let notify = handle
|
||||
.input(&context, operation_input(WorkerInputKind::Notify))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(notify.state, InternalWorkerOperationState::Accepted);
|
||||
|
||||
let request = || server_api::RuntimeWorkerInputRequest {
|
||||
kind: Some("user".into()),
|
||||
content: "human input".into(),
|
||||
segments: None,
|
||||
};
|
||||
let Json(generic) = send_runtime_worker_input(
|
||||
State(api.clone()),
|
||||
AxumPath((subject.runtime_id.clone(), key.clone())),
|
||||
Json(request()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let Json(scoped) = scoped_send_runtime_worker_input(
|
||||
State(api.clone()),
|
||||
AxumPath(ScopedRuntimeWorkerPath {
|
||||
workspace_id: TEST_WORKSPACE_ID.into(),
|
||||
worker: RuntimeWorkerRef::new(&subject.runtime_id, &key),
|
||||
}),
|
||||
Json(request()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(generic.state, scoped.state);
|
||||
assert_eq!(generic.state, server_api::WorkerOperationState::Accepted);
|
||||
assert_eq!(
|
||||
runtime
|
||||
.input_requests
|
||||
.lock()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.map(|request| request.kind.clone())
|
||||
.collect::<Vec<_>>(),
|
||||
vec![
|
||||
WorkerInputKind::UserIfIdle,
|
||||
WorkerInputKind::Notify,
|
||||
WorkerInputKind::User,
|
||||
WorkerInputKind::User
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn browser_protocol_sender_rejects_job_input_and_all_live_methods() {
|
||||
let workspace = tempfile::tempdir().unwrap();
|
||||
let (api, execution) = test_api_with_recording_backend(workspace.path()).await;
|
||||
let request = BackendJobRequest {
|
||||
job_id: "worker-operation-read-only-job".into(),
|
||||
purpose: "operation_boundary_test".into(),
|
||||
input_revision: "revision-1".into(),
|
||||
input_ref: "test://worker-operation/input".into(),
|
||||
input: serde_json::json!({"immutable": true}),
|
||||
instruction: "Return a structured result.".into(),
|
||||
profile: "builtin:backend-job".into(),
|
||||
source_worker: None,
|
||||
notification_target: None,
|
||||
limits: crate::backend_job::BackendJobLimits::default(),
|
||||
};
|
||||
let dispatched = api.dispatch_backend_job(&request).unwrap();
|
||||
let identity = dispatched.attempt.worker.unwrap();
|
||||
assert_eq!(execution.take_inputs().len(), 1);
|
||||
let handle = WorkspaceWorker::resolve(&api, &identity.runtime_id, &identity.worker_id).unwrap();
|
||||
let (context, _) = browser_operation_context(&api).await;
|
||||
let mut connection = handle.connect_protocol(&context).await.unwrap();
|
||||
assert!(matches!(
|
||||
tokio::time::timeout(std::time::Duration::from_secs(2), connection.events.recv())
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap(),
|
||||
protocol::Event::Snapshot { .. }
|
||||
));
|
||||
for method in browser_operation_methods() {
|
||||
assert!(matches!(connection.sender.send(&context, method).await,
|
||||
Err(Error::InvalidInput(message)) if message.contains("read-only")));
|
||||
}
|
||||
for kind in [WorkerInputKind::User, WorkerInputKind::Notify] {
|
||||
assert!(
|
||||
matches!(handle.input(&context, operation_input(kind)).await.unwrap_err().error,
|
||||
Error::InvalidInput(message) if message.contains("read-only"))
|
||||
);
|
||||
}
|
||||
assert!(execution.protocol_methods().is_empty());
|
||||
assert!(execution.take_inputs().is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn connected_browser_sender_rechecks_revoked_credentials_before_each_method() {
|
||||
let workspace = tempfile::tempdir().unwrap();
|
||||
let (api, execution) = test_api_with_recording_backend(workspace.path()).await;
|
||||
let identity = spawn_ticket_check_source(&api, "operation-browser-revocation");
|
||||
let handle = WorkspaceWorker::resolve(&api, &identity.runtime_id, &identity.worker_id).unwrap();
|
||||
let (context, token_hash) = browser_operation_context(&api).await;
|
||||
let mut connection = handle.connect_protocol(&context).await.unwrap();
|
||||
assert!(matches!(
|
||||
tokio::time::timeout(std::time::Duration::from_secs(2), connection.events.recv())
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap(),
|
||||
protocol::Event::Snapshot { .. }
|
||||
));
|
||||
// Prove these credentials were usable before revocation, then keep the same sender.
|
||||
handle.set_pinned(&context, true).await.unwrap();
|
||||
assert!(
|
||||
api.store
|
||||
.revoke_browser_session(&token_hash, &now_registry_timestamp())
|
||||
.unwrap()
|
||||
);
|
||||
assert!(matches!(
|
||||
handle.connect_protocol(&context).await,
|
||||
Err(Error::WorkspacePermissionDenied(_))
|
||||
));
|
||||
for method in browser_operation_methods() {
|
||||
assert!(matches!(
|
||||
connection.sender.send(&context, method).await,
|
||||
Err(Error::WorkspacePermissionDenied(_))
|
||||
));
|
||||
}
|
||||
assert!(matches!(
|
||||
handle.set_pinned(&context, false).await.unwrap_err().error,
|
||||
Error::WorkspacePermissionDenied(_)
|
||||
));
|
||||
assert!(
|
||||
execution.protocol_methods().is_empty(),
|
||||
"revoked credentials must not reach the execution channel"
|
||||
);
|
||||
assert_eq!(
|
||||
api.store
|
||||
.get_worker_registry(TEST_WORKSPACE_ID, &identity)
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.retention_state,
|
||||
"pinned"
|
||||
);
|
||||
}
|
||||
|
||||
fn operation_request_context(
|
||||
source: Option<server_api::ServerRuntimeSource>,
|
||||
) -> server_api::ServerRequestContext {
|
||||
server_api::ServerRequestContext {
|
||||
actor: None,
|
||||
worker_source: None,
|
||||
runtime_source: source,
|
||||
origin: None,
|
||||
transport_headers: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn request_context_uses_verified_subject_not_headers_or_account_fallback() {
|
||||
let workspace = tempfile::tempdir().unwrap();
|
||||
let api = test_api(workspace.path()).await;
|
||||
let mut headers_only = operation_request_context(None);
|
||||
headers_only.transport_headers = vec![
|
||||
("x-yoi-runtime-id".into(), b"runtime-a".to_vec()),
|
||||
("x-yoi-worker-id".into(), b"controller-a".to_vec()),
|
||||
];
|
||||
assert!(WorkerOperationContext::from_request(&api, &headers_only).is_err());
|
||||
let mut runtime_only = operation_request_context(Some(server_api::ServerRuntimeSource {
|
||||
runtime_id: "runtime-a".into(),
|
||||
worker_id: None,
|
||||
}));
|
||||
runtime_only.actor = Some(test_browser_request_actor());
|
||||
assert!(WorkerOperationContext::from_request(&api, &runtime_only).is_err());
|
||||
let mut mismatched = operation_request_context(Some(server_api::ServerRuntimeSource {
|
||||
runtime_id: "runtime-a".into(),
|
||||
worker_id: Some("controller-a".into()),
|
||||
}));
|
||||
mismatched.actor = Some(test_browser_request_actor());
|
||||
mismatched.transport_headers = vec![("x-yoi-worker-id".into(), b"controller-b".to_vec())];
|
||||
assert!(WorkerOperationContext::from_request(&api, &mismatched).is_err());
|
||||
seed_worker_source_member(&api, "runtime-a", "controller-a");
|
||||
let context = operation_request_context(Some(server_api::ServerRuntimeSource {
|
||||
runtime_id: "runtime-a".into(),
|
||||
worker_id: Some("controller-a".into()),
|
||||
}));
|
||||
assert!(
|
||||
matches!(WorkerOperationContext::from_request(&api, &context).unwrap(),
|
||||
WorkerOperationContext::WorkerControl { controller } if controller == RuntimeWorkerRef::new("runtime-a", "controller-a"))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn generic_management_aliases_cannot_bypass_worker_grants_or_target_bound_remove() {
|
||||
use server_api::ServerApi;
|
||||
let workspace = tempfile::tempdir().unwrap();
|
||||
let api = test_api(workspace.path()).await;
|
||||
seed_worker_source_member(&api, "offline-runtime", "controller");
|
||||
seed_worker_source_member(&api, "offline-runtime", "subject");
|
||||
let service = ServerApiContractService::Workspace(api.clone());
|
||||
let context = || {
|
||||
operation_request_context(Some(server_api::ServerRuntimeSource {
|
||||
runtime_id: "offline-runtime".into(),
|
||||
worker_id: Some("controller".into()),
|
||||
}))
|
||||
};
|
||||
let input = || server_api::RuntimeWorkerInputRequest {
|
||||
kind: Some("user".into()),
|
||||
content: "must not execute".into(),
|
||||
segments: None,
|
||||
};
|
||||
assert!(
|
||||
service
|
||||
.runtime_worker_input_alias(
|
||||
context(),
|
||||
"offline-runtime".into(),
|
||||
"subject".into(),
|
||||
input()
|
||||
)
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
assert!(
|
||||
service
|
||||
.runtime_worker_input(
|
||||
context(),
|
||||
TEST_WORKSPACE_ID.into(),
|
||||
"offline-runtime".into(),
|
||||
"subject".into(),
|
||||
input()
|
||||
)
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
assert!(
|
||||
service
|
||||
.runtime_worker_restore_alias(context(), "offline-runtime".into(), "subject".into())
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
let lifecycle = || server_api::RuntimeWorkerLifecycleRequest {
|
||||
reason: None,
|
||||
ticket_assignment: None,
|
||||
};
|
||||
assert!(
|
||||
service
|
||||
.runtime_worker_stop_alias(
|
||||
context(),
|
||||
"offline-runtime".into(),
|
||||
"subject".into(),
|
||||
lifecycle()
|
||||
)
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
assert!(
|
||||
service
|
||||
.runtime_worker_cancel_alias(
|
||||
context(),
|
||||
"offline-runtime".into(),
|
||||
"subject".into(),
|
||||
lifecycle()
|
||||
)
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
assert!(
|
||||
service
|
||||
.runtime_worker_pin(
|
||||
context(),
|
||||
TEST_WORKSPACE_ID.into(),
|
||||
"offline-runtime".into(),
|
||||
"subject".into()
|
||||
)
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
let cleanup = service
|
||||
.runtime_cleanup_execute(
|
||||
context(),
|
||||
TEST_WORKSPACE_ID.into(),
|
||||
"offline-runtime".into(),
|
||||
ExecuteRuntimeCleanupRequest {
|
||||
expected_plan_revision: "irrelevant".into(),
|
||||
expected_plan_digest: "irrelevant".into(),
|
||||
worker_target_ids: Vec::new(),
|
||||
workdir_target_ids: Vec::new(),
|
||||
confirm_dirty_discard_target_ids: Vec::new(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(cleanup.message.contains("target-bound WorkerRemove"));
|
||||
assert_eq!(
|
||||
api.store
|
||||
.get_worker_registry(
|
||||
TEST_WORKSPACE_ID,
|
||||
&RuntimeWorkerRef::new("offline-runtime", "subject")
|
||||
)
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.retention_state,
|
||||
"normal"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn uncertain_unary_input_keeps_a_distinct_diagnostic_without_claiming_acceptance() {
|
||||
let result = worker_input_result_to_api(WorkerInputResult {
|
||||
state: InternalWorkerOperationState::Rejected,
|
||||
disposition: WorkerInputDisposition::Unknown,
|
||||
worker: RuntimeWorkerRef::new("runtime-a", "worker-a"),
|
||||
runtime_run_id: None,
|
||||
diagnostics: Vec::new(),
|
||||
});
|
||||
assert_eq!(result.state, server_api::WorkerOperationState::Rejected);
|
||||
assert!(
|
||||
result
|
||||
.diagnostics
|
||||
.iter()
|
||||
.any(|diagnostic| diagnostic.code == "worker_input_outcome_unknown")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn grant_replacement_while_admission_waits_cannot_use_another_grants_lock() {
|
||||
let workspace = tempfile::tempdir().unwrap();
|
||||
let api = test_api(workspace.path()).await;
|
||||
let controller = RuntimeWorkerRef::new("offline-runtime", "controller");
|
||||
let subject = RuntimeWorkerRef::new("offline-runtime", "subject");
|
||||
for identity in [&controller, &subject] {
|
||||
seed_worker_source_member(&api, &identity.runtime_id, &identity.worker_id);
|
||||
}
|
||||
let grant = |id: &str| WorkerControlGrantRecord {
|
||||
workspace_id: TEST_WORKSPACE_ID.into(),
|
||||
grant_id: id.into(),
|
||||
controller: controller.clone(),
|
||||
subject: subject.clone(),
|
||||
relation: "spawned".into(),
|
||||
origin: "test".into(),
|
||||
permissions: vec!["pin".into()],
|
||||
operation_id: id.into(),
|
||||
created_at: now_registry_timestamp(),
|
||||
revoked_at: None,
|
||||
};
|
||||
api.store
|
||||
.create_worker_control_grant(&grant("original-grant"))
|
||||
.unwrap();
|
||||
let guard = worker_control_lock(&api, "original-grant")
|
||||
.lock_owned()
|
||||
.await;
|
||||
let handle = WorkspaceWorker::resolve(&api, &subject.runtime_id, &subject.worker_id).unwrap();
|
||||
let context = WorkerOperationContext::WorkerControl {
|
||||
controller: controller.clone(),
|
||||
};
|
||||
let admission = handle.set_pinned(&context, true);
|
||||
tokio::pin!(admission);
|
||||
// Poll deterministically to the held lock before replacing grant authority.
|
||||
assert!(futures::poll!(admission.as_mut()).is_pending());
|
||||
api.store
|
||||
.revoke_worker_control_grant(
|
||||
TEST_WORKSPACE_ID,
|
||||
"original-grant",
|
||||
&now_registry_timestamp(),
|
||||
)
|
||||
.unwrap();
|
||||
api.store
|
||||
.create_worker_control_grant(&grant("replacement-grant"))
|
||||
.unwrap();
|
||||
drop(guard);
|
||||
assert!(matches!(
|
||||
admission.await.unwrap_err().error,
|
||||
Error::RepositoryConflict(_)
|
||||
));
|
||||
assert_eq!(
|
||||
api.store
|
||||
.get_worker_registry(TEST_WORKSPACE_ID, &subject)
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.retention_state,
|
||||
"normal"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn browser_protocol_connection_and_cleanup_reject_cookie_origin_mismatch() {
|
||||
let workspace = tempfile::tempdir().unwrap();
|
||||
let (api, execution) = test_api_with_recording_backend(workspace.path()).await;
|
||||
let identity = spawn_ticket_check_source(&api, "operation-origin");
|
||||
let handle = WorkspaceWorker::resolve(&api, &identity.runtime_id, &identity.worker_id).unwrap();
|
||||
let (mut context, _) = browser_operation_context(&api).await;
|
||||
let WorkerOperationContext::Browser { headers, .. } = &mut context else {
|
||||
unreachable!()
|
||||
};
|
||||
headers.insert(ORIGIN, "https://untrusted.example".parse().unwrap());
|
||||
assert!(matches!(
|
||||
handle.connect_protocol(&context).await,
|
||||
Err(Error::WorkspacePermissionDenied(_))
|
||||
));
|
||||
let result = execute_runtime_cleanup_with_context(
|
||||
&api,
|
||||
EMBEDDED_WORKER_RUNTIME_ID,
|
||||
ExecuteRuntimeCleanupRequest {
|
||||
expected_plan_revision: "irrelevant".into(),
|
||||
expected_plan_digest: "irrelevant".into(),
|
||||
worker_target_ids: Vec::new(),
|
||||
workdir_target_ids: Vec::new(),
|
||||
confirm_dirty_discard_target_ids: Vec::new(),
|
||||
},
|
||||
&context,
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(matches!(result.error, Error::WorkspacePermissionDenied(_)));
|
||||
assert!(execution.protocol_methods().is_empty());
|
||||
}
|
||||
@@ -0,0 +1,504 @@
|
||||
//! Workspace identity-bound operation reception. Runtime/controller remain execution authority.
|
||||
use super::*;
|
||||
|
||||
/// Caller authority supplied by authenticated Server adapters, never public JSON/identity hints.
|
||||
/// The unchecked Backend variant exists only for lower-level regression fixtures.
|
||||
#[derive(Clone)]
|
||||
pub(crate) enum WorkerOperationContext {
|
||||
#[cfg(test)]
|
||||
Backend,
|
||||
WorkerControl {
|
||||
controller: RuntimeWorkerRef,
|
||||
},
|
||||
Browser {
|
||||
headers: HeaderMap,
|
||||
source: protocol::AuthenticatedInputSource,
|
||||
},
|
||||
}
|
||||
|
||||
impl WorkerOperationContext {
|
||||
pub(super) async fn authorize_workspace(&self, api: &WorkspaceApi) -> ApiResult<()> {
|
||||
if let Self::Browser { headers, source } = self {
|
||||
let actor = resolve_actor(&ServerAuthApi::from(api), headers)
|
||||
.await?
|
||||
.ok_or_else(|| {
|
||||
Error::WorkspacePermissionDenied(
|
||||
"Worker connection credentials were revoked".into(),
|
||||
)
|
||||
})?;
|
||||
if matches!(actor.auth_method, ActorAuthMethod::BrowserSession) {
|
||||
let AuthConfig::Passkey { origin, .. } = &api.config.auth;
|
||||
if headers.get(ORIGIN).and_then(|value| value.to_str().ok())
|
||||
!= Some(origin.as_str())
|
||||
{
|
||||
return Err(Error::WorkspacePermissionDenied(
|
||||
"cross-origin Worker mutation denied".into(),
|
||||
)
|
||||
.into());
|
||||
}
|
||||
}
|
||||
if authenticated_browser_input_source(&actor) != *source {
|
||||
return Err(Error::WorkspacePermissionDenied(
|
||||
"Worker connection Account changed".into(),
|
||||
)
|
||||
.into());
|
||||
}
|
||||
if !api
|
||||
.store
|
||||
.get_workspace(api.workspace_id())
|
||||
.await?
|
||||
.is_some_and(|workspace| workspace.state == "active")
|
||||
{
|
||||
return Err(Error::WorkspacePermissionDenied(
|
||||
"Workspace is not accepting Worker mutations".into(),
|
||||
)
|
||||
.into());
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
pub(super) fn from_request(
|
||||
api: &WorkspaceApi,
|
||||
context: &server_api::ServerRequestContext,
|
||||
) -> std::result::Result<Self, server_api::RepositoryApiError> {
|
||||
if let Some(source) = context.runtime_source.as_ref() {
|
||||
// Check hints against middleware's verified subject without replaying its proof
|
||||
// or requiring the target Runtime to be online just to represent identity.
|
||||
current_worker_contract_headers(context)?;
|
||||
let controller = RuntimeWorkerRef::new(
|
||||
&source.runtime_id,
|
||||
source
|
||||
.worker_id
|
||||
.as_deref()
|
||||
.expect("source binding required a Worker identity"),
|
||||
);
|
||||
if api
|
||||
.store
|
||||
.get_worker_registry(api.workspace_id(), &controller)
|
||||
.map_err(|error| ApiError::from(error).into_repository_api_error())?
|
||||
.is_none()
|
||||
{
|
||||
return Err(ApiError::from(Error::UnknownWorker { worker: controller })
|
||||
.into_repository_api_error());
|
||||
}
|
||||
return Ok(Self::WorkerControl { controller });
|
||||
}
|
||||
let actor = context.actor.as_ref().ok_or_else(|| {
|
||||
ApiError::from(Error::WorkspacePermissionDenied(
|
||||
"authenticated Worker operation source required".into(),
|
||||
))
|
||||
.into_repository_api_error()
|
||||
})?;
|
||||
let headers = contract_request_headers(context)?;
|
||||
Ok(Self::Browser {
|
||||
headers,
|
||||
source: authenticated_browser_input_source(actor),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// A registry identity, not an authorization lease or a snapshot of live execution.
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct WorkspaceWorker {
|
||||
api: WorkspaceApi,
|
||||
identity: RuntimeWorkerRef,
|
||||
}
|
||||
|
||||
impl WorkspaceWorker {
|
||||
pub(crate) fn resolve(api: &WorkspaceApi, runtime_id: &str, reference: &str) -> Result<Self> {
|
||||
let identity = resolve_workspace_worker_reference(api, runtime_id, reference)
|
||||
.map_err(|error| error.error)?;
|
||||
Ok(Self {
|
||||
api: api.clone(),
|
||||
identity,
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn identity(&self) -> &RuntimeWorkerRef {
|
||||
&self.identity
|
||||
}
|
||||
|
||||
fn current_record(&self) -> ApiResult<WorkerRegistryRecord> {
|
||||
self.api
|
||||
.store
|
||||
.get_worker_registry(self.api.workspace_id(), &self.identity)?
|
||||
.ok_or_else(|| {
|
||||
Error::UnknownWorker {
|
||||
worker: self.identity.clone(),
|
||||
}
|
||||
.into()
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) async fn authorize_operation(
|
||||
&self,
|
||||
context: &WorkerOperationContext,
|
||||
permission: &str,
|
||||
) -> ApiResult<Option<tokio::sync::OwnedMutexGuard<()>>> {
|
||||
self.current_record()?;
|
||||
match context {
|
||||
#[cfg(test)]
|
||||
WorkerOperationContext::Backend => Ok(None),
|
||||
WorkerOperationContext::Browser { .. } => {
|
||||
context.authorize_workspace(&self.api).await?;
|
||||
Ok(None)
|
||||
}
|
||||
WorkerOperationContext::WorkerControl { controller } => {
|
||||
if self
|
||||
.api
|
||||
.store
|
||||
.get_worker_registry(self.api.workspace_id(), controller)?
|
||||
.is_none()
|
||||
{
|
||||
return Err(Error::UnknownWorker {
|
||||
worker: controller.clone(),
|
||||
}
|
||||
.into());
|
||||
}
|
||||
let grant = authorize_known_worker_permission(
|
||||
&self.api,
|
||||
self.api.workspace_id(),
|
||||
controller,
|
||||
&self.identity,
|
||||
permission,
|
||||
)?;
|
||||
let guard = worker_control_lock(&self.api, &grant.grant_id)
|
||||
.lock_owned()
|
||||
.await;
|
||||
self.current_record()?;
|
||||
let current_grant = authorize_known_worker_permission(
|
||||
&self.api,
|
||||
self.api.workspace_id(),
|
||||
controller,
|
||||
&self.identity,
|
||||
permission,
|
||||
)?;
|
||||
if current_grant.grant_id != grant.grant_id {
|
||||
return Err(Error::RepositoryConflict(
|
||||
"Worker control grant changed while waiting for admission".into(),
|
||||
)
|
||||
.into());
|
||||
}
|
||||
Ok(Some(guard))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn restore(
|
||||
&self,
|
||||
context: &WorkerOperationContext,
|
||||
) -> ApiResult<InternalWorkerRestoreResult> {
|
||||
let _guard = self.authorize_operation(context, "restore").await?;
|
||||
// Preserve singleton, Workdir, credentials, bindings and pending operation reconciliation.
|
||||
self.api.restore_workspace_worker(&self.identity)
|
||||
}
|
||||
|
||||
pub(super) async fn stop(
|
||||
&self,
|
||||
context: &WorkerOperationContext,
|
||||
request: WorkerLifecycleRequest,
|
||||
) -> ApiResult<WorkerLifecycleResult> {
|
||||
let _guard = self.authorize_operation(context, "stop").await?;
|
||||
let result = self
|
||||
.api
|
||||
.runtime
|
||||
.stop_worker(&self.identity, request)
|
||||
.map_err(|e| e.into_error())?;
|
||||
if result.state != InternalWorkerOperationState::Accepted {
|
||||
return Ok(result);
|
||||
}
|
||||
parse_runtime_worker_id_for_registry(&self.identity.worker_id)?;
|
||||
let session_lock = current_worker_session_lock(&self.api, &self.identity);
|
||||
let _session_guard = session_lock.lock().await;
|
||||
close_current_worker_session_locked(&self.api, &self.identity).await?;
|
||||
if let Some(record) = self
|
||||
.api
|
||||
.store
|
||||
.get_worker_registry(self.api.workspace_id(), &self.identity)?
|
||||
{
|
||||
sync_linked_workdir_after_worker_stop(&self.api, &self.identity.runtime_id, &record)?;
|
||||
}
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
pub(super) async fn cancel(
|
||||
&self,
|
||||
context: &WorkerOperationContext,
|
||||
request: WorkerLifecycleRequest,
|
||||
) -> ApiResult<WorkerLifecycleResult> {
|
||||
let _guard = self.authorize_operation(context, "cancel").await?;
|
||||
Ok(self
|
||||
.api
|
||||
.runtime
|
||||
.cancel_worker(&self.identity, request)
|
||||
.map_err(|e| e.into_error())?)
|
||||
}
|
||||
|
||||
pub(super) async fn input(
|
||||
&self,
|
||||
context: &WorkerOperationContext,
|
||||
mut request: WorkerInputRequest,
|
||||
) -> ApiResult<WorkerInputResult> {
|
||||
let permission = if request.kind == WorkerInputKind::Notify {
|
||||
"notify"
|
||||
} else {
|
||||
"send_input"
|
||||
};
|
||||
let _guard = self.authorize_operation(context, permission).await?;
|
||||
self.ensure_interactive()?;
|
||||
if matches!(context, WorkerOperationContext::WorkerControl { .. })
|
||||
&& request.kind == WorkerInputKind::User
|
||||
{
|
||||
request.kind = WorkerInputKind::UserIfIdle;
|
||||
}
|
||||
Ok(self
|
||||
.api
|
||||
.runtime
|
||||
.send_input(&self.identity, request)
|
||||
.map_err(|e| e.into_error())?)
|
||||
}
|
||||
|
||||
fn ensure_interactive(&self) -> ApiResult<()> {
|
||||
if self
|
||||
.api
|
||||
.store
|
||||
.worker_registry_projection(self.api.workspace_id(), &self.identity)?
|
||||
.and_then(|projection| projection.job)
|
||||
.is_some()
|
||||
{
|
||||
return Err(Error::InvalidInput("Backend Job Workers accept only their Backend-owned immutable input; Console is read-only".into()).into());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(super) async fn set_pinned(
|
||||
&self,
|
||||
context: &WorkerOperationContext,
|
||||
pinned: bool,
|
||||
) -> ApiResult<WorkerRetentionResponse> {
|
||||
let _guard = self.authorize_operation(context, "pin").await?;
|
||||
let retention_state = if pinned { "pinned" } else { "normal" };
|
||||
let changed = self.api.worker_projection.publish_ordered(|| {
|
||||
self.api
|
||||
.store
|
||||
.update_worker_retention(
|
||||
self.api.workspace_id(),
|
||||
&self.identity,
|
||||
retention_state,
|
||||
now_registry_timestamp().as_str(),
|
||||
)
|
||||
.map(|commit| (!commit.changes.is_empty(), commit))
|
||||
})?;
|
||||
if !changed {
|
||||
return Err(cleanup_api_error(
|
||||
&self.identity.runtime_id,
|
||||
"workspace_worker_retention_unknown_worker",
|
||||
"Worker is not known to the Backend registry",
|
||||
));
|
||||
}
|
||||
Ok(WorkerRetentionResponse {
|
||||
workspace_id: self.api.workspace_id().to_owned(),
|
||||
runtime_id: self.identity.runtime_id.clone(),
|
||||
worker_id: self.identity.worker_id.clone(),
|
||||
pinned,
|
||||
retention_state: retention_state.into(),
|
||||
})
|
||||
}
|
||||
|
||||
/// The caller owns the full cleanup-plan revision/digest check. The removal service
|
||||
/// rechecks live assignment/retention/use at its existing commit boundary.
|
||||
pub(super) async fn delete_from_plan(
|
||||
&self,
|
||||
context: &WorkerOperationContext,
|
||||
candidate: &CleanupWorkerCandidate,
|
||||
) -> ApiResult<()> {
|
||||
let _guard = self.authorize_operation(context, "remove").await?;
|
||||
if candidate.runtime_id != self.identity.runtime_id
|
||||
|| candidate.runtime_worker_id != self.identity.worker_id
|
||||
{
|
||||
return Err(Error::InvalidInput(
|
||||
"cleanup candidate does not address the resolved Worker".into(),
|
||||
)
|
||||
.into());
|
||||
}
|
||||
WorkerRemovalService::new(&self.api)
|
||||
.execute_cleanup_removal(candidate)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(super) async fn remove(
|
||||
&self,
|
||||
source: crate::worker_source::VerifiedWorkerMutationSource,
|
||||
reason: &str,
|
||||
) -> std::result::Result<worker::WorkspaceResponse, String> {
|
||||
self.current_record()
|
||||
.map_err(|error| error.error.to_string())?;
|
||||
WorkerRemovalService::new(&self.api)
|
||||
.execute_async(
|
||||
source,
|
||||
&self.identity.runtime_id,
|
||||
&self.identity.worker_id,
|
||||
reason,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn connect_protocol(
|
||||
&self,
|
||||
context: &WorkerOperationContext,
|
||||
) -> Result<WorkspaceWorkerConnection> {
|
||||
let _guard = self
|
||||
.authorize_operation(context, "observe")
|
||||
.await
|
||||
.map_err(|error| error.error)?;
|
||||
let source = match context {
|
||||
WorkerOperationContext::Browser { source, .. } => Some(source),
|
||||
_ => None,
|
||||
};
|
||||
let connection =
|
||||
connect_workspace_worker_protocol(&self.api, &self.identity, source).await?;
|
||||
Ok(WorkspaceWorkerConnection {
|
||||
sender: WorkspaceWorkerMethodSender {
|
||||
worker: self.clone(),
|
||||
methods: connection.methods,
|
||||
},
|
||||
events: connection.events,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// An execution connection is bound to exactly the identity which created it. No caller
|
||||
/// can pair a different target handle with a subscription's raw sending capability.
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct WorkspaceWorkerMethodSender {
|
||||
worker: WorkspaceWorker,
|
||||
methods: tokio::sync::mpsc::Sender<protocol::Method>,
|
||||
}
|
||||
pub(crate) struct WorkspaceWorkerConnection {
|
||||
pub(crate) sender: WorkspaceWorkerMethodSender,
|
||||
pub(crate) events: tokio::sync::mpsc::Receiver<protocol::Event>,
|
||||
}
|
||||
impl WorkspaceWorkerMethodSender {
|
||||
pub(crate) async fn send(
|
||||
&self,
|
||||
context: &WorkerOperationContext,
|
||||
method: protocol::Method,
|
||||
) -> Result<()> {
|
||||
match method {
|
||||
protocol::Method::Pause { command } => self.pause(context, command).await,
|
||||
protocol::Method::Resume { command } => self.resume(context, command).await,
|
||||
protocol::Method::Cancel { command } => self.cancel(context, command).await,
|
||||
protocol::Method::Compact { command } => self.compact(context, command).await,
|
||||
protocol::Method::Submit {
|
||||
submission_request_id,
|
||||
input,
|
||||
} => self.submit(context, submission_request_id, input).await,
|
||||
protocol::Method::Notify {
|
||||
notification_request_id,
|
||||
message,
|
||||
} => self.notify(context, notification_request_id, message).await,
|
||||
other => self.dispatch(context, other).await,
|
||||
}
|
||||
}
|
||||
async fn dispatch(
|
||||
&self,
|
||||
context: &WorkerOperationContext,
|
||||
method: protocol::Method,
|
||||
) -> Result<()> {
|
||||
self.send_checked(context, method)
|
||||
.await
|
||||
.map_err(|error| error.error)
|
||||
}
|
||||
pub(crate) async fn pause(
|
||||
&self,
|
||||
context: &WorkerOperationContext,
|
||||
command: protocol::WorkerCommandEnvelope,
|
||||
) -> Result<()> {
|
||||
self.dispatch(context, protocol::Method::Pause { command })
|
||||
.await
|
||||
}
|
||||
pub(crate) async fn resume(
|
||||
&self,
|
||||
context: &WorkerOperationContext,
|
||||
command: protocol::WorkerCommandEnvelope,
|
||||
) -> Result<()> {
|
||||
self.dispatch(context, protocol::Method::Resume { command })
|
||||
.await
|
||||
}
|
||||
pub(crate) async fn cancel(
|
||||
&self,
|
||||
context: &WorkerOperationContext,
|
||||
command: protocol::WorkerCommandEnvelope,
|
||||
) -> Result<()> {
|
||||
self.dispatch(context, protocol::Method::Cancel { command })
|
||||
.await
|
||||
}
|
||||
pub(crate) async fn compact(
|
||||
&self,
|
||||
context: &WorkerOperationContext,
|
||||
command: protocol::WorkerCommandEnvelope,
|
||||
) -> Result<()> {
|
||||
self.dispatch(context, protocol::Method::Compact { command })
|
||||
.await
|
||||
}
|
||||
pub(crate) async fn submit(
|
||||
&self,
|
||||
context: &WorkerOperationContext,
|
||||
submission_request_id: String,
|
||||
input: Vec<protocol::Segment>,
|
||||
) -> Result<()> {
|
||||
self.dispatch(
|
||||
context,
|
||||
protocol::Method::Submit {
|
||||
submission_request_id,
|
||||
input,
|
||||
},
|
||||
)
|
||||
.await
|
||||
}
|
||||
pub(crate) async fn notify(
|
||||
&self,
|
||||
context: &WorkerOperationContext,
|
||||
notification_request_id: String,
|
||||
message: String,
|
||||
) -> Result<()> {
|
||||
self.dispatch(
|
||||
context,
|
||||
protocol::Method::Notify {
|
||||
notification_request_id,
|
||||
message,
|
||||
},
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn send_checked(
|
||||
&self,
|
||||
context: &WorkerOperationContext,
|
||||
method: protocol::Method,
|
||||
) -> ApiResult<()> {
|
||||
let _guard = self
|
||||
.worker
|
||||
.authorize_operation(context, "send_input")
|
||||
.await?;
|
||||
self.worker.ensure_interactive()?;
|
||||
let WorkerOperationContext::Browser { source, .. } = context else {
|
||||
return Err(Error::WorkspacePermissionDenied(
|
||||
"Browser Worker protocol context required".into(),
|
||||
)
|
||||
.into());
|
||||
};
|
||||
let method = authorize_browser_worker_method(method, source)
|
||||
.map_err(|message| Error::InvalidInput(message.into()))?;
|
||||
self.methods.send(method).await.map_err(|_| {
|
||||
Error::RuntimeOperationFailed {
|
||||
runtime_id: self.worker.identity.runtime_id.clone(),
|
||||
code: "worker_protocol_closed".into(),
|
||||
message:
|
||||
"Worker execution connection closed; no implicit Restore or retry was performed"
|
||||
.into(),
|
||||
}
|
||||
.into()
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -12,7 +12,7 @@ use worker_runtime::identity::RuntimeWorkerRef;
|
||||
|
||||
use crate::runtime_subscription::RuntimeSubscriptionBroker;
|
||||
use crate::server::{
|
||||
WorkspaceApi, authorize_browser_worker_method, connect_workspace_worker_protocol,
|
||||
WorkerOperationContext, WorkspaceApi, WorkspaceWorker, WorkspaceWorkerMethodSender,
|
||||
};
|
||||
|
||||
const OUTBOUND_CAPACITY: usize = 256;
|
||||
@@ -20,13 +20,13 @@ const WORKER_PROTOCOL_SNAPSHOT_TIMEOUT: std::time::Duration = std::time::Duratio
|
||||
|
||||
struct ActiveSubscription {
|
||||
task: tokio::task::JoinHandle<()>,
|
||||
methods: Option<mpsc::Sender<protocol::Method>>,
|
||||
methods: Option<WorkspaceWorkerMethodSender>,
|
||||
}
|
||||
|
||||
pub(crate) async fn serve_workspace_subscription(
|
||||
api: WorkspaceApi,
|
||||
socket: WebSocket,
|
||||
input_source: protocol::AuthenticatedInputSource,
|
||||
context: WorkerOperationContext,
|
||||
) {
|
||||
let broker = api.runtime_subscription_broker().clone();
|
||||
let (mut socket_sender, mut socket_receiver) = socket.split();
|
||||
@@ -90,16 +90,20 @@ pub(crate) async fn serve_workspace_subscription(
|
||||
worker_id,
|
||||
runtime_id: Some(runtime_id),
|
||||
} => {
|
||||
let worker = RuntimeWorkerRef::new(&runtime_id, worker_id.as_str());
|
||||
match connect_workspace_worker_protocol(
|
||||
let connection = match WorkspaceWorker::resolve(
|
||||
&api,
|
||||
&worker,
|
||||
Some(&input_source),
|
||||
)
|
||||
.await
|
||||
{
|
||||
&runtime_id,
|
||||
worker_id.as_str(),
|
||||
) {
|
||||
Ok(worker) => worker
|
||||
.connect_protocol(&context)
|
||||
.await
|
||||
.map_err(|error| error.to_string()),
|
||||
Err(error) => Err(error.to_string()),
|
||||
};
|
||||
match connection {
|
||||
Ok(connection) => {
|
||||
let methods = connection.methods.clone();
|
||||
let methods = connection.sender.clone();
|
||||
let task = tokio::spawn(run_worker_protocol(
|
||||
request_id,
|
||||
subscription_id.clone(),
|
||||
@@ -165,12 +169,7 @@ pub(crate) async fn serve_workspace_subscription(
|
||||
else {
|
||||
break;
|
||||
};
|
||||
let Ok(method) =
|
||||
authorize_browser_worker_method(message.method, &input_source)
|
||||
else {
|
||||
break;
|
||||
};
|
||||
if methods.send(method).await.is_err() {
|
||||
if methods.send(&context, message.method).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user