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

worker: structure runtime worker identities

Browse Source
This commit is contained in:
Hare
2026-08-05 21:26:55 +09:00
parent fd391ef705
commit a2781f57e9
9 changed files with 775 additions and 816 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker-runtime/src/catalog.rs
+3 -4
View File
@@ -1,4 +1,4 @@
use crate::identity::{WorkerId, WorkerRef};
use crate::identity::{RuntimeWorkerRef, WorkerId, WorkerRef};
use crate::interaction::WorkerInput;
use crate::profile_archive::{ProfileSourceArchive, ProfileSourceArchiveRef};
use serde::{Deserialize, Serialize};
@@ -132,9 +132,8 @@ pub struct WorkingDirectoryCleanupTarget {
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct WorkingDirectoryOccupancy {
pub runtime_id: String,
pub runtime_worker_id: u64,
pub worker_id: String,
#[serde(flatten)]
pub worker: RuntimeWorkerRef,
pub display_name: String,
pub linked_at: String,
}
crates/worker-runtime/src/identity.rs
+52
View File
@@ -30,6 +30,32 @@ impl fmt::Display for WorkerId {
}
}
/// Backend-visible Worker identity, namespaced by the Runtime that owns the Worker record.
///
/// This is intentionally distinct from [`WorkerRef`], which is meaningful only inside one
/// Runtime. Do not flatten this reference into a concatenated string for authority decisions.
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub struct RuntimeWorkerRef {
pub runtime_id: String,
pub worker_id: String,
}
impl RuntimeWorkerRef {
pub fn new(runtime_id: impl Into<String>, worker_id: impl Into<String>) -> Self {
Self {
runtime_id: runtime_id.into(),
worker_id: worker_id.into(),
}
}
pub fn local_worker_ref(&self) -> Result<WorkerRef, std::num::ParseIntError> {
self.worker_id
.parse::<u64>()
.map(WorkerId::new)
.map(WorkerRef::new)
}
}
/// Runtime-local authority reference for Worker operations.
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub struct WorkerRef {
@@ -41,3 +67,29 @@ impl WorkerRef {
Self { worker_id }
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn runtime_worker_ref_preserves_structured_identity_and_json_fields() {
let worker = RuntimeWorkerRef::new("arcadia", "30");
assert_eq!(worker.runtime_id, "arcadia");
assert_eq!(worker.worker_id, "30");
assert_eq!(
worker.local_worker_ref().unwrap(),
WorkerRef::new(WorkerId::new(30))
);
assert_eq!(
serde_json::to_value(&worker).unwrap(),
serde_json::json!({"runtime_id": "arcadia", "worker_id": "30"})
);
}
#[test]
fn runtime_worker_ref_does_not_treat_composite_text_as_local_worker_id() {
let worker = RuntimeWorkerRef::new("arcadia", "embedded-worker-runtime-5");
assert!(worker.local_worker_ref().is_err());
}
}
crates/workspace-server/src/hosts.rs
+160 -154
View File
@@ -1,5 +1,5 @@
use crate::Error;
use crate::resource_broker::BackendResourceBroker;
use crate::resource_broker::{BackendResourceBroker, BackendResourceTarget};
use chrono::Utc;
use reqwest::blocking::{Client as BlockingHttpClient, RequestBuilder};
use reqwest::header::{AUTHORIZATION, CONTENT_TYPE};
@@ -45,7 +45,9 @@ use worker_runtime::http_server::{
RuntimeHttpWorkerWorkspaceApiRequest, RuntimeHttpWorkersResponse,
RuntimeHttpWorkingDirectoriesResponse, RuntimeHttpWorkingDirectoryResponse,
};
use worker_runtime::identity::{WorkerId as EmbeddedWorkerId, WorkerRef as EmbeddedWorkerRef};
use worker_runtime::identity::{
RuntimeWorkerRef, WorkerId as EmbeddedWorkerId, WorkerRef as EmbeddedWorkerRef,
};
use worker_runtime::interaction::{
WorkerInput as EmbeddedWorkerInput, WorkerInputKind as EmbeddedWorkerInputKind,
};
@@ -235,8 +237,8 @@ pub struct WorkerCapabilitySummary {
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct WorkerSummary {
pub runtime_id: String,
pub worker_id: String,
#[serde(flatten)]
pub worker: RuntimeWorkerRef,
pub host_id: String,
/// Human-readable display name. This is not identity and may be duplicated.
pub display_name: String,
@@ -487,8 +489,8 @@ pub struct WorkerLifecycleRequest {
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct WorkerLifecycleResult {
pub state: WorkerOperationState,
pub runtime_id: String,
pub worker_id: String,
#[serde(flatten)]
pub worker: RuntimeWorkerRef,
pub diagnostics: Vec<RuntimeDiagnostic>,
}
@@ -505,8 +507,8 @@ pub enum WorkerInputKind {
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct WorkerDeleteResult {
pub state: WorkerOperationState,
pub runtime_id: String,
pub worker_id: String,
#[serde(flatten)]
pub worker: RuntimeWorkerRef,
pub deleted: bool,
pub diagnostics: Vec<RuntimeDiagnostic>,
}
@@ -529,8 +531,8 @@ pub struct WorkerCompletionsRequest {
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct WorkerCompletionsResult {
pub runtime_id: String,
pub worker_id: String,
#[serde(flatten)]
pub worker: RuntimeWorkerRef,
pub kind: protocol::CompletionKind,
pub prefix: String,
pub entries: Vec<protocol::CompletionEntry>,
@@ -540,8 +542,8 @@ pub struct WorkerCompletionsResult {
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct WorkerInputResult {
pub state: WorkerOperationState,
pub runtime_id: String,
pub worker_id: String,
#[serde(flatten)]
pub worker: RuntimeWorkerRef,
pub diagnostics: Vec<RuntimeDiagnostic>,
}
@@ -561,8 +563,7 @@ pub enum RuntimeRegistryError {
UnknownRuntime(String),
UnknownHost(String),
UnknownWorker {
runtime_id: String,
worker_id: String,
worker: RuntimeWorkerRef,
},
RuntimeOperationFailed {
runtime_id: String,
@@ -579,10 +580,10 @@ impl RuntimeRegistryError {
}
Self::UnknownRuntime(runtime_id) => format!("unknown runtime `{runtime_id}`"),
Self::UnknownHost(host_id) => format!("unknown host `{host_id}`"),
Self::UnknownWorker {
runtime_id,
worker_id,
} => format!("unknown worker `{worker_id}` in runtime `{runtime_id}`"),
Self::UnknownWorker { worker } => format!(
"unknown worker `{}` in runtime `{}`",
worker.worker_id, worker.runtime_id
),
Self::RuntimeOperationFailed { message, .. } => message.clone(),
}
}
@@ -595,13 +596,7 @@ impl RuntimeRegistryError {
},
Self::UnknownRuntime(runtime_id) => Error::UnknownRuntime(runtime_id),
Self::UnknownHost(host_id) => Error::UnknownHost(host_id),
Self::UnknownWorker {
runtime_id,
worker_id,
} => Error::UnknownWorker {
runtime_id,
worker_id,
},
Self::UnknownWorker { worker } => Error::UnknownWorker { worker },
Self::RuntimeOperationFailed {
runtime_id,
code,
@@ -798,8 +793,7 @@ pub trait WorkspaceWorkerRuntime: Send + Sync {
) -> WorkerLifecycleResult {
WorkerLifecycleResult {
state: WorkerOperationState::Unsupported,
runtime_id: self.runtime_id().to_string(),
worker_id: worker_id.to_string(),
worker: RuntimeWorkerRef::new(self.runtime_id().to_string(), worker_id.to_string()),
diagnostics: vec![diagnostic(
"worker_stop_pending",
DiagnosticSeverity::Info,
@@ -817,8 +811,7 @@ pub trait WorkspaceWorkerRuntime: Send + Sync {
) -> WorkerLifecycleResult {
WorkerLifecycleResult {
state: WorkerOperationState::Unsupported,
runtime_id: self.runtime_id().to_string(),
worker_id: worker_id.to_string(),
worker: RuntimeWorkerRef::new(self.runtime_id().to_string(), worker_id.to_string()),
diagnostics: vec![diagnostic(
"worker_cancel_pending",
DiagnosticSeverity::Info,
@@ -832,8 +825,7 @@ pub trait WorkspaceWorkerRuntime: Send + Sync {
fn delete_worker(&self, worker_id: &str) -> WorkerDeleteResult {
WorkerDeleteResult {
state: WorkerOperationState::Unsupported,
runtime_id: self.runtime_id().to_string(),
worker_id: worker_id.to_string(),
worker: RuntimeWorkerRef::new(self.runtime_id().to_string(), worker_id.to_string()),
deleted: false,
diagnostics: vec![diagnostic(
"worker_delete_unsupported",
@@ -853,8 +845,7 @@ pub trait WorkspaceWorkerRuntime: Send + Sync {
fn send_input(&self, worker_id: &str, _request: WorkerInputRequest) -> WorkerInputResult {
WorkerInputResult {
state: WorkerOperationState::Unsupported,
runtime_id: self.runtime_id().to_string(),
worker_id: worker_id.to_string(),
worker: RuntimeWorkerRef::new(self.runtime_id().to_string(), worker_id.to_string()),
diagnostics: vec![diagnostic(
"worker_input_pending",
DiagnosticSeverity::Info,
@@ -871,8 +862,7 @@ pub trait WorkspaceWorkerRuntime: Send + Sync {
request: WorkerCompletionsRequest,
) -> WorkerCompletionsResult {
WorkerCompletionsResult {
runtime_id: self.runtime_id().to_string(),
worker_id: worker_id.to_string(),
worker: RuntimeWorkerRef::new(self.runtime_id().to_string(), worker_id.to_string()),
kind: request.kind,
prefix: request.prefix,
entries: Vec::new(),
@@ -1099,11 +1089,9 @@ impl RuntimeRegistry {
}
}
pub fn worker(
&self,
runtime_id: &str,
worker_id: &str,
) -> Result<WorkerSummary, RuntimeRegistryError> {
pub fn worker(&self, worker: &RuntimeWorkerRef) -> Result<WorkerSummary, RuntimeRegistryError> {
let runtime_id = worker.runtime_id.as_str();
let worker_id = worker.worker_id.as_str();
validate_backend_identifier("runtime_id", runtime_id)?;
validate_backend_identifier("worker_id", worker_id)?;
let runtime = self.runtime(runtime_id)?;
@@ -1116,9 +1104,10 @@ impl RuntimeRegistry {
pub fn restore_worker(
&self,
runtime_id: &str,
worker_id: &str,
worker: &RuntimeWorkerRef,
) -> Result<WorkerRestoreResult, RuntimeRegistryError> {
let runtime_id = worker.runtime_id.as_str();
let worker_id = worker.worker_id.as_str();
validate_backend_identifier("runtime_id", runtime_id)?;
validate_backend_identifier("worker_id", worker_id)?;
let runtime = self.runtime(runtime_id)?;
@@ -1127,10 +1116,11 @@ impl RuntimeRegistry {
pub fn replace_worker_workspace_api(
&self,
runtime_id: &str,
worker_id: &str,
worker: &RuntimeWorkerRef,
workspace_api: WorkspaceApiRef,
) -> Result<WorkerWorkspaceApiResult, RuntimeRegistryError> {
let runtime_id = worker.runtime_id.as_str();
let worker_id = worker.worker_id.as_str();
validate_backend_identifier("runtime_id", runtime_id)?;
validate_backend_identifier("worker_id", worker_id)?;
let runtime = self.runtime(runtime_id)?;
@@ -1241,10 +1231,11 @@ impl RuntimeRegistry {
pub fn send_protocol_method(
&self,
runtime_id: &str,
worker_id: &str,
worker: &RuntimeWorkerRef,
method: protocol::Method,
) -> Result<Vec<protocol::Event>, RuntimeRegistryError> {
let runtime_id = worker.runtime_id.as_str();
let worker_id = worker.worker_id.as_str();
validate_backend_identifier("runtime_id", runtime_id)?;
validate_backend_identifier("worker_id", worker_id)?;
let runtime = self.runtime(runtime_id)?;
@@ -1261,10 +1252,11 @@ impl RuntimeRegistry {
pub fn send_input(
&self,
runtime_id: &str,
worker_id: &str,
worker: &RuntimeWorkerRef,
request: WorkerInputRequest,
) -> Result<WorkerInputResult, RuntimeRegistryError> {
let runtime_id = worker.runtime_id.as_str();
let worker_id = worker.worker_id.as_str();
validate_backend_identifier("runtime_id", runtime_id)?;
validate_backend_identifier("worker_id", worker_id)?;
let runtime = self.runtime(runtime_id)?;
@@ -1281,10 +1273,11 @@ impl RuntimeRegistry {
pub fn worker_completions(
&self,
runtime_id: &str,
worker_id: &str,
worker: &RuntimeWorkerRef,
request: WorkerCompletionsRequest,
) -> Result<WorkerCompletionsResult, RuntimeRegistryError> {
let runtime_id = worker.runtime_id.as_str();
let worker_id = worker.worker_id.as_str();
validate_backend_identifier("runtime_id", runtime_id)?;
validate_backend_identifier("worker_id", worker_id)?;
let runtime = self.runtime(runtime_id)?;
@@ -1301,10 +1294,11 @@ impl RuntimeRegistry {
pub fn stop_worker(
&self,
runtime_id: &str,
worker_id: &str,
worker: &RuntimeWorkerRef,
request: WorkerLifecycleRequest,
) -> Result<WorkerLifecycleResult, RuntimeRegistryError> {
let runtime_id = worker.runtime_id.as_str();
let worker_id = worker.worker_id.as_str();
validate_backend_identifier("runtime_id", runtime_id)?;
validate_backend_identifier("worker_id", worker_id)?;
let runtime = self.runtime(runtime_id)?;
@@ -1321,10 +1315,11 @@ impl RuntimeRegistry {
pub fn cancel_worker(
&self,
runtime_id: &str,
worker_id: &str,
worker: &RuntimeWorkerRef,
request: WorkerLifecycleRequest,
) -> Result<WorkerLifecycleResult, RuntimeRegistryError> {
let runtime_id = worker.runtime_id.as_str();
let worker_id = worker.worker_id.as_str();
validate_backend_identifier("runtime_id", runtime_id)?;
validate_backend_identifier("worker_id", worker_id)?;
let runtime = self.runtime(runtime_id)?;
@@ -1341,9 +1336,10 @@ impl RuntimeRegistry {
pub fn delete_worker(
&self,
runtime_id: &str,
worker_id: &str,
worker: &RuntimeWorkerRef,
) -> Result<WorkerDeleteResult, RuntimeRegistryError> {
let runtime_id = worker.runtime_id.as_str();
let worker_id = worker.worker_id.as_str();
validate_backend_identifier("runtime_id", runtime_id)?;
validate_backend_identifier("worker_id", worker_id)?;
let runtime = self.runtime(runtime_id)?;
@@ -1360,17 +1356,17 @@ impl RuntimeRegistry {
pub fn observation_source(
&self,
runtime_id: &str,
worker_id: &str,
worker: &RuntimeWorkerRef,
) -> Result<crate::observation::RuntimeObservationSource, RuntimeRegistryError> {
let runtime_id = worker.runtime_id.as_str();
let worker_id = worker.worker_id.as_str();
validate_backend_identifier("runtime_id", runtime_id)?;
validate_backend_identifier("worker_id", worker_id)?;
let runtime = self.runtime(runtime_id)?;
runtime
.observation_source(worker_id)
.ok_or_else(|| RuntimeRegistryError::UnknownWorker {
runtime_id: runtime_id.to_string(),
worker_id: worker_id.to_string(),
worker: worker.clone(),
})
}
@@ -1483,8 +1479,7 @@ impl EmbeddedWorkerRuntime {
true,
);
WorkerSummary {
runtime_id: self.runtime_id.clone(),
worker_id,
worker: RuntimeWorkerRef::new(&self.runtime_id, worker_id.clone()),
host_id: self.host_id.clone(),
display_name: display.display_name.clone(),
label: display.display_name,
@@ -1522,8 +1517,7 @@ impl EmbeddedWorkerRuntime {
true,
);
WorkerSummary {
runtime_id: self.runtime_id.clone(),
worker_id,
worker: RuntimeWorkerRef::new(&self.runtime_id, worker_id.clone()),
host_id: self.host_id.clone(),
display_name: display.display_name.clone(),
label: display.display_name,
@@ -1974,8 +1968,7 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime {
match self.runtime.stop_worker(&worker_ref, request.reason) {
Ok(_) => WorkerLifecycleResult {
state: WorkerOperationState::Accepted,
runtime_id: self.runtime_id.clone(),
worker_id: worker_id.to_string(),
worker: RuntimeWorkerRef::new(self.runtime_id.clone(), worker_id.to_string()),
diagnostics: Vec::new(),
},
Err(error) => embedded_lifecycle_rejected(
@@ -2018,8 +2011,7 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime {
match self.runtime.cancel_worker(&worker_ref, request.reason) {
Ok(_) => WorkerLifecycleResult {
state: WorkerOperationState::Accepted,
runtime_id: self.runtime_id.clone(),
worker_id: worker_id.to_string(),
worker: RuntimeWorkerRef::new(self.runtime_id.clone(), worker_id.to_string()),
diagnostics: Vec::new(),
},
Err(error) => embedded_lifecycle_rejected(
@@ -2034,8 +2026,7 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime {
let Some(worker_ref) = self.worker_ref(worker_id) else {
return WorkerDeleteResult {
state: WorkerOperationState::Rejected,
runtime_id: self.runtime_id.clone(),
worker_id: worker_id.to_string(),
worker: RuntimeWorkerRef::new(self.runtime_id.clone(), worker_id.to_string()),
deleted: false,
diagnostics: vec![diagnostic(
"embedded_worker_id_invalid",
@@ -2047,15 +2038,16 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime {
match self.runtime.delete_worker(&worker_ref) {
Ok(result) => WorkerDeleteResult {
state: WorkerOperationState::Accepted,
runtime_id: self.runtime_id.clone(),
worker_id: result.worker_id.to_string(),
worker: RuntimeWorkerRef::new(
self.runtime_id.clone(),
result.worker_id.to_string(),
),
deleted: result.deleted,
diagnostics: Vec::new(),
},
Err(error) => WorkerDeleteResult {
state: WorkerOperationState::Rejected,
runtime_id: self.runtime_id.clone(),
worker_id: worker_id.to_string(),
worker: RuntimeWorkerRef::new(self.runtime_id.clone(), worker_id.to_string()),
deleted: false,
diagnostics: vec![embedded_runtime_diagnostic(&error)],
},
@@ -2072,8 +2064,7 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime {
}
Some(crate::observation::RuntimeObservationSource::embedded(
crate::observation::EmbeddedRuntimeObservationSource {
runtime_id: self.runtime_id.clone(),
worker_id: worker_id.to_string(),
worker: RuntimeWorkerRef::new(&self.runtime_id, worker_id),
runtime: self.runtime.clone(),
worker_ref,
},
@@ -2096,8 +2087,7 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime {
}
let Some(worker_ref) = self.worker_ref(worker_id) else {
return Err(RuntimeRegistryError::UnknownWorker {
runtime_id: self.runtime_id.clone(),
worker_id: worker_id.to_string(),
worker: RuntimeWorkerRef::new(&self.runtime_id, worker_id),
});
};
self.runtime
@@ -2148,8 +2138,7 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime {
match self.runtime.send_input(&worker_ref, input) {
Ok(_) => WorkerInputResult {
state: WorkerOperationState::Accepted,
runtime_id: self.runtime_id.clone(),
worker_id: worker_id.to_string(),
worker: RuntimeWorkerRef::new(self.runtime_id.clone(), worker_id.to_string()),
diagnostics: Vec::new(),
},
Err(error) => embedded_input_rejected(
@@ -2167,8 +2156,7 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime {
) -> WorkerCompletionsResult {
if !self.execution_enabled {
return WorkerCompletionsResult {
runtime_id: self.runtime_id.clone(),
worker_id: worker_id.to_string(),
worker: RuntimeWorkerRef::new(self.runtime_id.clone(), worker_id.to_string()),
kind: request.kind,
prefix: request.prefix,
entries: Vec::new(),
@@ -2183,8 +2171,7 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime {
}
let Some(worker_ref) = self.worker_ref(worker_id) else {
return WorkerCompletionsResult {
runtime_id: self.runtime_id.clone(),
worker_id: worker_id.to_string(),
worker: RuntimeWorkerRef::new(self.runtime_id.clone(), worker_id.to_string()),
kind: request.kind,
prefix: request.prefix,
entries: Vec::new(),
@@ -2200,16 +2187,14 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime {
.worker_completions(&worker_ref, request.kind, &request.prefix)
{
Ok(entries) => WorkerCompletionsResult {
runtime_id: self.runtime_id.clone(),
worker_id: worker_id.to_string(),
worker: RuntimeWorkerRef::new(self.runtime_id.clone(), worker_id.to_string()),
kind: request.kind,
prefix: request.prefix,
entries,
diagnostics: Vec::new(),
},
Err(error) => WorkerCompletionsResult {
runtime_id: self.runtime_id.clone(),
worker_id: worker_id.to_string(),
worker: RuntimeWorkerRef::new(self.runtime_id.clone(), worker_id.to_string()),
kind: request.kind,
prefix: request.prefix,
entries: Vec::new(),
@@ -2572,8 +2557,7 @@ impl RemoteWorkerRuntime {
false,
);
WorkerSummary {
runtime_id: self.runtime_id.clone(),
worker_id,
worker: RuntimeWorkerRef::new(&self.runtime_id, worker_id.clone()),
host_id: self.host_id.clone(),
display_name: display.display_name.clone(),
label: display.display_name,
@@ -2615,8 +2599,7 @@ impl RemoteWorkerRuntime {
false,
);
WorkerSummary {
runtime_id: self.runtime_id.clone(),
worker_id,
worker: RuntimeWorkerRef::new(&self.runtime_id, worker_id.clone()),
host_id: self.host_id.clone(),
display_name: display.display_name.clone(),
label: display.display_name,
@@ -2655,8 +2638,7 @@ impl RemoteWorkerRuntime {
) -> WorkerLifecycleResult {
WorkerLifecycleResult {
state: WorkerOperationState::Accepted,
runtime_id: self.runtime_id.clone(),
worker_id: worker_id.to_string(),
worker: RuntimeWorkerRef::new(self.runtime_id.clone(), worker_id.to_string()),
diagnostics: vec![diagnostic(
"remote_runtime_lifecycle_accepted",
DiagnosticSeverity::Info,
@@ -3073,15 +3055,16 @@ impl WorkspaceWorkerRuntime for RemoteWorkerRuntime {
{
Ok(response) => WorkerDeleteResult {
state: WorkerOperationState::Accepted,
runtime_id: self.runtime_id.clone(),
worker_id: response.worker.worker_id.to_string(),
worker: RuntimeWorkerRef::new(
self.runtime_id.clone(),
response.worker.worker_id.to_string(),
),
deleted: response.worker.deleted,
diagnostics: Vec::new(),
},
Err(diagnostic) => WorkerDeleteResult {
state: WorkerOperationState::Rejected,
runtime_id: self.runtime_id.clone(),
worker_id: worker_id.to_string(),
worker: RuntimeWorkerRef::new(self.runtime_id.clone(), worker_id.to_string()),
deleted: false,
diagnostics: vec![diagnostic],
},
@@ -3094,8 +3077,7 @@ impl WorkspaceWorkerRuntime for RemoteWorkerRuntime {
) -> Option<crate::observation::RuntimeObservationSource> {
Some(crate::observation::RuntimeObservationSource::remote_ws(
crate::observation::RuntimeObservationSourceConfig {
runtime_id: self.runtime_id.clone(),
worker_id: worker_id.to_string(),
worker: RuntimeWorkerRef::new(&self.runtime_id, worker_id),
endpoint: self.ws_endpoint(worker_id),
bearer_token: self
.runtime_capability_token(&format!("/v1/workers/{worker_id}/protocol"))
@@ -3122,8 +3104,7 @@ impl WorkspaceWorkerRuntime for RemoteWorkerRuntime {
) {
Ok(_) => WorkerInputResult {
state: WorkerOperationState::Accepted,
runtime_id: self.runtime_id.clone(),
worker_id: worker_id.to_string(),
worker: RuntimeWorkerRef::new(self.runtime_id.clone(), worker_id.to_string()),
diagnostics: Vec::new(),
},
Err(diagnostic) => remote_input_rejected(&self.runtime_id, worker_id, diagnostic),
@@ -3144,16 +3125,14 @@ impl WorkspaceWorkerRuntime for RemoteWorkerRuntime {
&request,
) {
Ok(response) => WorkerCompletionsResult {
runtime_id: self.runtime_id.clone(),
worker_id: worker_id.to_string(),
worker: RuntimeWorkerRef::new(self.runtime_id.clone(), worker_id.to_string()),
kind: response.kind,
prefix: response.prefix,
entries: response.entries,
diagnostics: Vec::new(),
},
Err(diagnostic) => WorkerCompletionsResult {
runtime_id: self.runtime_id.clone(),
worker_id: worker_id.to_string(),
worker: RuntimeWorkerRef::new(self.runtime_id.clone(), worker_id.to_string()),
kind: request.kind,
prefix: request.prefix,
entries: Vec::new(),
@@ -3247,10 +3226,12 @@ fn profile_source_archive_http_source(
backend_base_url: &str,
) -> Result<ProfileSourceArchiveSource, String> {
let archive = profile_source_archive_for_request(request, profile)?;
let target = runtime_id
.map(BackendResourceTarget::Runtime)
.unwrap_or(BackendResourceTarget::Workspace);
let _handle = resource_broker.issue_profile_source_archive_handle(
workspace_id.to_string(),
runtime_id,
None,
target,
archive.clone(),
);
let etag = format!("\"profile-source:{}\"", archive.reference.digest);
@@ -3294,10 +3275,12 @@ fn builtin_profile_config_bundle(
let (profile_source_archive, profile_source_archive_handle) = match archive_transport {
ProfileSourceArchiveTransport::Inline => (Some(archive), None),
ProfileSourceArchiveTransport::BackendResourceHandle => {
let target = runtime_id
.map(BackendResourceTarget::Runtime)
.unwrap_or(BackendResourceTarget::Workspace);
let handle = resource_broker.issue_profile_source_archive_handle(
workspace_id.to_string(),
runtime_id,
None,
target,
archive,
);
(None, Some(handle))
@@ -3515,8 +3498,7 @@ fn embedded_input_rejected(
) -> WorkerInputResult {
WorkerInputResult {
state: WorkerOperationState::Rejected,
runtime_id: runtime_id.to_string(),
worker_id: worker_id.to_string(),
worker: RuntimeWorkerRef::new(runtime_id.to_string(), worker_id.to_string()),
diagnostics: vec![diagnostic],
}
}
@@ -3528,8 +3510,7 @@ fn remote_input_rejected(
) -> WorkerInputResult {
WorkerInputResult {
state: WorkerOperationState::Rejected,
runtime_id: runtime_id.to_string(),
worker_id: worker_id.to_string(),
worker: RuntimeWorkerRef::new(runtime_id.to_string(), worker_id.to_string()),
diagnostics: vec![diagnostic],
}
}
@@ -3541,8 +3522,7 @@ fn embedded_lifecycle_rejected(
) -> WorkerLifecycleResult {
WorkerLifecycleResult {
state: WorkerOperationState::Rejected,
runtime_id: runtime_id.to_string(),
worker_id: worker_id.to_string(),
worker: RuntimeWorkerRef::new(runtime_id.to_string(), worker_id.to_string()),
diagnostics: vec![diagnostic],
}
}
@@ -3554,8 +3534,7 @@ fn remote_lifecycle_rejected(
) -> WorkerLifecycleResult {
WorkerLifecycleResult {
state: WorkerOperationState::Rejected,
runtime_id: runtime_id.to_string(),
worker_id: worker_id.to_string(),
worker: RuntimeWorkerRef::new(runtime_id.to_string(), worker_id.to_string()),
diagnostics: vec![diagnostic],
}
}
@@ -3779,8 +3758,7 @@ fn operation_failed_or_unknown_worker(
message: diagnostic.message,
})
.unwrap_or_else(|| RuntimeRegistryError::UnknownWorker {
runtime_id: runtime_id.to_string(),
worker_id: worker_id.to_string(),
worker: RuntimeWorkerRef::new(runtime_id, worker_id),
})
}
@@ -3887,8 +3865,7 @@ fn worker_spawn_intent_label(intent: &WorkerSpawnIntent) -> &'static str {
pub fn placeholder_worker(host_id: impl Into<String>) -> WorkerSummary {
let host_id = host_id.into();
WorkerSummary {
runtime_id: "placeholder".to_string(),
worker_id: "worker-placeholder".to_string(),
worker: RuntimeWorkerRef::new("placeholder", "worker-placeholder"),
host_id,
display_name: "Worker runtime actions are not implemented".to_string(),
label: "Worker runtime actions are not implemented".to_string(),
@@ -3951,6 +3928,29 @@ mod tests {
}
}
#[test]
fn worker_summary_keeps_flat_wire_identity_while_using_structured_internal_identity() {
let summary = placeholder_worker("placeholder");
assert_eq!(
summary.worker,
RuntimeWorkerRef::new("placeholder", "worker-placeholder")
);
let value = serde_json::to_value(summary).unwrap();
assert_eq!(value["runtime_id"], "placeholder");
assert_eq!(value["worker_id"], "worker-placeholder");
assert!(value.get("worker").is_none());
let lifecycle = WorkerLifecycleResult {
state: WorkerOperationState::Accepted,
worker: RuntimeWorkerRef::new("arcadia", "30"),
diagnostics: Vec::new(),
};
let value = serde_json::to_value(lifecycle).unwrap();
assert_eq!(value["runtime_id"], "arcadia");
assert_eq!(value["worker_id"], "30");
assert!(value.get("worker").is_none());
}
#[test]
fn embedded_orchestrator_profile_enables_workdir_and_worker_authority() {
let root = tempfile::tempdir().unwrap();
@@ -4289,8 +4289,7 @@ mod tests {
runtime_id: runtime_id.to_string(),
host_id: host_id.to_string(),
workers: vec![WorkerSummary {
runtime_id: runtime_id.to_string(),
worker_id: worker_id.to_string(),
worker: RuntimeWorkerRef::new(runtime_id, worker_id),
host_id: host_id.to_string(),
display_name: label.to_string(),
label: label.to_string(),
@@ -4382,7 +4381,7 @@ mod tests {
worker: self
.workers
.iter()
.find(|worker| worker.worker_id == worker_id)
.find(|worker| worker.worker.worker_id == worker_id)
.cloned(),
diagnostics: Vec::new(),
}
@@ -4406,13 +4405,17 @@ mod tests {
)),
]);
let from_runtime_b = registry.worker("runtime-b", "shared-worker").unwrap();
assert_eq!(from_runtime_b.runtime_id, "runtime-b");
let from_runtime_b = registry
.worker(&RuntimeWorkerRef::new("runtime-b", "shared-worker"))
.unwrap();
assert_eq!(from_runtime_b.worker.runtime_id, "runtime-b");
assert_eq!(from_runtime_b.host_id, "host-b");
assert_eq!(from_runtime_b.label, "worker from runtime b");
let from_runtime_a = registry.worker("runtime-a", "shared-worker").unwrap();
assert_eq!(from_runtime_a.runtime_id, "runtime-a");
let from_runtime_a = registry
.worker(&RuntimeWorkerRef::new("runtime-a", "shared-worker"))
.unwrap();
assert_eq!(from_runtime_a.worker.runtime_id, "runtime-a");
assert_eq!(from_runtime_a.host_id, "host-a");
assert_eq!(from_runtime_a.label, "worker from runtime a");
}
@@ -4436,7 +4439,7 @@ mod tests {
let listed = registry.list_workers_for_runtime("runtime-b", 10).unwrap();
assert_eq!(listed.items.len(), 1);
assert_eq!(listed.items[0].runtime_id, "runtime-b");
assert_eq!(listed.items[0].worker.runtime_id, "runtime-b");
assert_eq!(listed.items[0].host_id, "host-b");
assert_eq!(listed.items[0].label, "worker from runtime b");
}
@@ -4455,7 +4458,9 @@ mod tests {
Some("builtin:companion")
);
let worker = registry.worker("runtime-a", "worker-a").unwrap();
let worker = registry
.worker(&RuntimeWorkerRef::new("runtime-a", "worker-a"))
.unwrap();
assert_eq!(worker.profile.as_deref(), Some("builtin:companion"));
}
@@ -4468,7 +4473,9 @@ mod tests {
"worker from runtime a",
))]);
let unknown_runtime = registry.worker("runtime-missing", "worker-a").unwrap_err();
let unknown_runtime = registry
.worker(&RuntimeWorkerRef::new("runtime-missing", "worker-a"))
.unwrap_err();
assert_eq!(
unknown_runtime,
RuntimeRegistryError::UnknownRuntime("runtime-missing".to_string())
@@ -4478,18 +4485,19 @@ mod tests {
Error::UnknownRuntime(runtime_id) if runtime_id == "runtime-missing"
));
let unknown_worker = registry.worker("runtime-a", "999").unwrap_err();
let unknown_worker = registry
.worker(&RuntimeWorkerRef::new("runtime-a", "999"))
.unwrap_err();
assert_eq!(
unknown_worker,
RuntimeRegistryError::UnknownWorker {
runtime_id: "runtime-a".to_string(),
worker_id: "999".to_string(),
worker: RuntimeWorkerRef::new("runtime-a", "999"),
}
);
assert!(matches!(
unknown_worker.into_error(),
Error::UnknownWorker { runtime_id, worker_id }
if runtime_id == "runtime-a" && worker_id == "999"
Error::UnknownWorker { worker }
if worker == RuntimeWorkerRef::new("runtime-a", "999")
));
}
@@ -4591,7 +4599,7 @@ mod tests {
assert!(worker.capabilities.can_stop);
let input = runtime.send_input(
&worker.worker_id,
&worker.worker.worker_id,
WorkerInputRequest {
kind: WorkerInputKind::User,
content: "hello".to_string(),
@@ -4603,7 +4611,7 @@ mod tests {
let deadline = std::time::Instant::now() + std::time::Duration::from_secs(2);
loop {
let detail = runtime
.worker(&worker.worker_id)
.worker(&worker.worker.worker_id)
.worker
.expect("worker detail");
if detail.state == "idle" {
@@ -4671,15 +4679,14 @@ mod tests {
.any(|evidence| evidence.kind == "embedded_runtime_backend_internal_projection")
);
let worker = spawned.worker.expect("created embedded worker");
assert_eq!(worker.runtime_id, EMBEDDED_RUNTIME_ID);
assert_eq!(worker.worker.runtime_id, EMBEDDED_RUNTIME_ID);
assert_eq!(worker.workspace.visibility, "backend_internal");
assert_eq!(worker.workspace.identity, "runtime_registry_worker");
assert_eq!(worker.implementation.kind, "embedded_worker_runtime");
assert_eq!(worker.profile.as_deref(), Some("builtin:coder"));
let input = registry
.send_input(
EMBEDDED_RUNTIME_ID,
&worker.worker_id,
&worker.worker,
WorkerInputRequest {
kind: WorkerInputKind::User,
content: "hello embedded runtime".to_string(),
@@ -4688,12 +4695,10 @@ mod tests {
)
.unwrap();
assert_eq!(input.state, WorkerOperationState::Accepted);
assert_eq!(input.runtime_id, EMBEDDED_RUNTIME_ID);
assert_eq!(input.worker_id, worker.worker_id);
assert_eq!(input.worker.runtime_id, EMBEDDED_RUNTIME_ID);
assert_eq!(input.worker.worker_id, worker.worker.worker_id);
let detail = registry
.worker(EMBEDDED_RUNTIME_ID, &worker.worker_id)
.unwrap();
let detail = registry.worker(&worker.worker).unwrap();
let json = serde_json::to_string(&(embedded_summary, worker, input, detail)).unwrap();
for forbidden in [
@@ -4871,7 +4876,7 @@ mod tests {
);
let observation = registry
.observation_source("remote:primary", "1")
.observation_source(&RuntimeWorkerRef::new("remote:primary", "1"))
.expect("remote runtime exposes backend-owned WS observation source");
let crate::observation::RuntimeObservationSource::RemoteWs(observation) = observation
else {
@@ -4883,8 +4888,8 @@ mod tests {
let workers = registry.list_workers(10);
assert_eq!(workers.items.len(), 1);
assert_eq!(workers.items[0].runtime_id, "remote:primary");
assert_eq!(workers.items[0].worker_id, "1");
assert_eq!(workers.items[0].worker.runtime_id, "remote:primary");
assert_eq!(workers.items[0].worker.worker_id, "1");
assert_eq!(
workers.items[0].implementation.kind,
"remote_worker_runtime"
@@ -4897,8 +4902,7 @@ mod tests {
let input = registry
.send_input(
"remote:primary",
"1",
&RuntimeWorkerRef::new("remote:primary", "1"),
WorkerInputRequest {
kind: WorkerInputKind::User,
content: "hello remote".to_string(),
@@ -4975,7 +4979,9 @@ mod tests {
assert_eq!(workers.items[2].state, "paused");
assert_eq!(workers.items[3].state, "idle");
let stopped_detail = registry.worker("remote:primary", "1").unwrap();
let stopped_detail = registry
.worker(&RuntimeWorkerRef::new("remote:primary", "1"))
.unwrap();
assert!(!stopped_detail.capabilities.can_stop);
assert_eq!(stopped_detail.state, "stopped");
@@ -5154,7 +5160,7 @@ mod tests {
);
let error = registry
.worker("remote:primary", "999")
.worker(&RuntimeWorkerRef::new("remote:primary", "999"))
.expect_err("auth failure is a backend operation error");
assert!(matches!(
error,
crates/workspace-server/src/lib.rs
+6 -5
View File
@@ -43,6 +43,8 @@ pub use repositories::{
pub use server::{AuthConfig, ServerConfig, WorkspaceApi, build_router, serve};
pub use store::{ControlPlaneStore, SqliteWorkspaceStore, WorkspaceRecord};
use worker_runtime::identity::RuntimeWorkerRef;
pub type Result<T> = std::result::Result<T, Error>;
#[derive(Debug, thiserror::Error)]
@@ -65,11 +67,8 @@ pub enum Error {
UnknownHost(String),
#[error("unknown runtime `{0}`")]
UnknownRuntime(String),
#[error("unknown worker `{worker_id}` in runtime `{runtime_id}`")]
UnknownWorker {
runtime_id: String,
worker_id: String,
},
#[error("unknown worker `{}` in runtime `{}`", worker.worker_id, worker.runtime_id)]
UnknownWorker { worker: RuntimeWorkerRef },
#[error("invalid runtime {kind} `{value}`")]
InvalidRuntimeIdentifier { kind: String, value: String },
#[error("worker name is reserved for a dedicated Workspace service: {0}")]
@@ -93,6 +92,8 @@ pub enum Error {
TicketAssignmentConflict(String),
#[error("Workdir attachment conflict: {0}")]
WorkdirAttachmentConflict(String),
#[error("Registry inconsistency: {0}")]
RegistryInconsistency(String),
#[error("Worker source identity is invalid: {0}")]
WorkerSourceIdentity(String),
#[error("workspace identity error: {0}")]
crates/workspace-server/src/observation.rs
+36 -66
View File
@@ -1,7 +1,7 @@
use std::collections::{BTreeMap, VecDeque};
use std::sync::Arc;
use worker_runtime::identity::WorkerRef;
use worker_runtime::identity::{RuntimeWorkerRef, WorkerRef};
use worker_runtime::observation::{WorkerObservationCursor, WorkerObservationEvent};
use axum::http::StatusCode;
@@ -15,8 +15,7 @@ use tokio_tungstenite::tungstenite::{Error as TungsteniteError, Message as Tungs
/// Backend-private source for a runtime worker observation stream.
#[derive(Clone, PartialEq, Eq)]
pub struct RuntimeObservationSourceConfig {
pub runtime_id: String,
pub worker_id: String,
pub worker: RuntimeWorkerRef,
pub endpoint: String,
pub bearer_token: Option<String>,
}
@@ -24,8 +23,8 @@ pub struct RuntimeObservationSourceConfig {
impl std::fmt::Debug for RuntimeObservationSourceConfig {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RuntimeObservationSourceConfig")
.field("runtime_id", &self.runtime_id)
.field("worker_id", &self.worker_id)
.field("runtime_id", &self.worker.runtime_id)
.field("worker_id", &self.worker.worker_id)
.field("endpoint", &"<backend-private>")
.field(
"bearer_token",
@@ -37,8 +36,7 @@ impl std::fmt::Debug for RuntimeObservationSourceConfig {
#[derive(Clone)]
pub struct EmbeddedRuntimeObservationSource {
pub runtime_id: String,
pub worker_id: String,
pub worker: RuntimeWorkerRef,
pub runtime: worker_runtime::Runtime,
pub worker_ref: WorkerRef,
}
@@ -60,15 +58,15 @@ impl RuntimeObservationSource {
pub fn runtime_id(&self) -> &str {
match self {
Self::RemoteWs(config) => &config.runtime_id,
Self::Embedded(source) => &source.runtime_id,
Self::RemoteWs(config) => &config.worker.runtime_id,
Self::Embedded(source) => &source.worker.runtime_id,
}
}
pub fn worker_id(&self) -> &str {
match self {
Self::RemoteWs(config) => &config.worker_id,
Self::Embedded(source) => &source.worker_id,
Self::RemoteWs(config) => &config.worker.worker_id,
Self::Embedded(source) => &source.worker.worker_id,
}
}
}
@@ -78,8 +76,8 @@ impl std::fmt::Debug for RuntimeObservationSource {
match self {
Self::RemoteWs(config) => formatter
.debug_struct("RemoteRuntimeObservationSource")
.field("runtime_id", &config.runtime_id)
.field("worker_id", &config.worker_id)
.field("runtime_id", &config.worker.runtime_id)
.field("worker_id", &config.worker.worker_id)
.field("endpoint", &"<backend-private>")
.field(
"bearer_token",
@@ -88,8 +86,8 @@ impl std::fmt::Debug for RuntimeObservationSource {
.finish(),
Self::Embedded(source) => formatter
.debug_struct("EmbeddedRuntimeObservationSource")
.field("runtime_id", &source.runtime_id)
.field("worker_id", &source.worker_id)
.field("runtime_id", &source.worker.runtime_id)
.field("worker_id", &source.worker.worker_id)
.finish(),
}
}
@@ -98,8 +96,8 @@ impl std::fmt::Debug for RuntimeObservationSource {
/// Event consumed from a Runtime-owned worker observation WebSocket.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct RuntimeObservationUpstreamEvent {
pub runtime_id: String,
pub worker_id: String,
#[serde(flatten)]
pub worker: RuntimeWorkerRef,
pub runtime_event_id: String,
pub payload: protocol::Event,
}
@@ -132,11 +130,7 @@ impl ObservationProxyError {
}
}
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
struct ObservationKey {
runtime_id: String,
worker_id: String,
}
type ObservationKey = RuntimeWorkerRef;
/// Backend-owned in-memory v0 observation proxy state.
#[derive(Clone)]
@@ -156,15 +150,7 @@ impl BackendObservationProxy {
pub fn new(sources: Vec<RuntimeObservationSourceConfig>) -> Self {
let sources = sources
.into_iter()
.map(|source| {
(
ObservationKey {
runtime_id: source.runtime_id.clone(),
worker_id: source.worker_id.clone(),
},
source,
)
})
.map(|source| (source.worker.clone(), source))
.collect();
Self {
sources: Arc::new(sources),
@@ -173,19 +159,16 @@ impl BackendObservationProxy {
pub fn source(
&self,
runtime_id: &str,
worker_id: &str,
worker: &RuntimeWorkerRef,
) -> Result<RuntimeObservationSource, ObservationProxyError> {
self.sources
.get(&ObservationKey {
runtime_id: runtime_id.to_string(),
worker_id: worker_id.to_string(),
})
.get(worker)
.cloned()
.map(RuntimeObservationSource::remote_ws)
.ok_or_else(|| {
ObservationProxyError::WorkerNotFound(format!(
"worker {worker_id} is not registered for runtime {runtime_id}"
"worker {} is not registered for runtime {}",
worker.worker_id, worker.runtime_id
))
})
}
@@ -209,8 +192,7 @@ fn map_runtime_connect_error(error: TungsteniteError) -> ObservationProxyError {
}
pub struct RuntimeWsObservationClient {
runtime_id: String,
worker_id: String,
worker: RuntimeWorkerRef,
stream: tokio_tungstenite::WebSocketStream<
tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
>,
@@ -240,8 +222,7 @@ impl RuntimeWsObservationClient {
.await
.map_err(map_runtime_connect_error)?;
Ok(Self {
runtime_id: source.runtime_id.clone(),
worker_id: source.worker_id.clone(),
worker: source.worker.clone(),
stream,
})
}
@@ -291,8 +272,7 @@ impl RuntimeWsObservationClient {
))
})?;
return Ok(RuntimeObservationUpstreamEvent {
runtime_id: self.runtime_id.clone(),
worker_id: self.worker_id.clone(),
worker: self.worker.clone(),
runtime_event_id: "protocol".to_string(),
payload,
});
@@ -330,8 +310,7 @@ impl RuntimeObservationClient {
}
pub struct EmbeddedObservationClient {
runtime_id: String,
worker_id: String,
worker: RuntimeWorkerRef,
worker_ref: WorkerRef,
cursor: WorkerObservationCursor,
receiver: tokio::sync::broadcast::Receiver<WorkerObservationEvent>,
@@ -346,7 +325,7 @@ impl EmbeddedObservationClient {
.map_err(|err| {
ObservationProxyError::WorkerNotFound(format!(
"embedded Worker '{}' is not observable: {err}",
source.worker_id
source.worker.worker_id
))
})?;
let receiver = source
@@ -355,7 +334,7 @@ impl EmbeddedObservationClient {
.map_err(|err| {
ObservationProxyError::WorkerNotFound(format!(
"embedded Worker '{}' observation subscription is unavailable: {err}",
source.worker_id
source.worker.worker_id
))
})?;
let mut queued = VecDeque::new();
@@ -365,12 +344,11 @@ impl EmbeddedObservationClient {
.map_err(|err| {
ObservationProxyError::WorkerNotFound(format!(
"embedded Worker '{}' snapshot is unavailable: {err}",
source.worker_id
source.worker.worker_id
))
})?;
queued.push_back(RuntimeObservationUpstreamEvent {
runtime_id: source.runtime_id.clone(),
worker_id: source.worker_id.clone(),
worker: source.worker.clone(),
runtime_event_id: "snapshot".to_string(),
payload: snapshot,
});
@@ -380,19 +358,14 @@ impl EmbeddedObservationClient {
.map_err(|err| {
ObservationProxyError::RuntimeUnavailable(format!(
"embedded Worker '{}' observation cursor is unavailable: {err}",
source.worker_id
source.worker.worker_id
))
})?
{
queued.push_back(Self::map_event(
&source.runtime_id,
&source.worker_id,
event,
));
queued.push_back(Self::map_event(&source.worker, event));
}
Ok(Self {
runtime_id: source.runtime_id.clone(),
worker_id: source.worker_id.clone(),
worker: source.worker.clone(),
worker_ref: source.worker_ref.clone(),
cursor,
receiver,
@@ -418,7 +391,7 @@ impl EmbeddedObservationClient {
"embedded runtime emitted a malformed cursor".into(),
)
})?;
return Ok(Self::map_event(&self.runtime_id, &self.worker_id, event));
return Ok(Self::map_event(&self.worker, event));
}
Ok(_) => continue,
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {
@@ -436,13 +409,11 @@ impl EmbeddedObservationClient {
}
fn map_event(
runtime_id: &str,
worker_id: &str,
worker: &RuntimeWorkerRef,
event: WorkerObservationEvent,
) -> RuntimeObservationUpstreamEvent {
RuntimeObservationUpstreamEvent {
runtime_id: runtime_id.to_string(),
worker_id: worker_id.to_string(),
worker: worker.clone(),
runtime_event_id: event.cursor.clone(),
payload: event.payload,
}
@@ -455,8 +426,7 @@ mod tests {
fn sensitive_source() -> RuntimeObservationSourceConfig {
RuntimeObservationSourceConfig {
runtime_id: "remote-runtime".to_string(),
worker_id: "worker-1".to_string(),
worker: RuntimeWorkerRef::new("remote-runtime", "worker-1"),
endpoint: "wss://remote.example.invalid/private/workers/worker-1/protocol/ws"
.to_string(),
bearer_token: Some("top-secret-bearer-token".to_string()),
crates/workspace-server/src/resource_broker.rs
+36 -28
View File
@@ -3,7 +3,7 @@ use chrono::{Duration, Utc};
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use uuid::Uuid;
use worker_runtime::identity::WorkerId;
use worker_runtime::identity::RuntimeWorkerRef;
use worker_runtime::profile_archive::ProfileSourceArchive;
use worker_runtime::resource::{
BackendResourceClient, BackendResourceError, BackendResourceFetchRequest,
@@ -17,10 +17,17 @@ pub struct BackendResourceBroker {
resources: Arc<Mutex<HashMap<String, StoredResource>>>,
}
#[derive(Clone, Copy, Debug)]
pub enum BackendResourceTarget<'a> {
Workspace,
Runtime(&'a str),
Worker(&'a RuntimeWorkerRef),
}
#[derive(Clone)]
struct StoredResource {
runtime_id: Option<String>,
worker_id: Option<String>,
worker: Option<RuntimeWorkerRef>,
handle: BackendResourceHandle,
archive: ProfileSourceArchive,
}
@@ -29,11 +36,17 @@ impl BackendResourceBroker {
pub fn issue_profile_source_archive_handle(
&self,
workspace_id: impl Into<String>,
runtime_id: Option<&str>,
worker_id: Option<&WorkerId>,
target: BackendResourceTarget<'_>,
archive: ProfileSourceArchive,
) -> BackendResourceHandle {
let workspace_id = workspace_id.into();
let (runtime_id, worker) = match target {
BackendResourceTarget::Workspace => (None, None),
BackendResourceTarget::Runtime(runtime_id) => (Some(runtime_id.to_string()), None),
BackendResourceTarget::Worker(worker) => {
(Some(worker.runtime_id.clone()), Some(worker.clone()))
}
};
let nonce = Uuid::now_v7().to_string();
let audit_correlation_id = format!("resource-fetch-{nonce}");
let expires_at = Utc::now() + Duration::minutes(15);
@@ -41,8 +54,8 @@ impl BackendResourceBroker {
kind: BackendResourceKind::ProfileSourceArchive,
workspace_id: workspace_id.clone(),
scope_id: Some("workspace-profile-source".to_string()),
runtime_id: runtime_id.map(|id| id.to_string()),
worker_id: worker_id.map(|id| id.to_string()),
runtime_id: runtime_id.clone(),
worker_id: worker.as_ref().map(|worker| worker.worker_id.clone()),
resource_id: archive.reference.id.clone(),
digest: archive.reference.digest.clone(),
operation: BackendResourceOperation::FetchArchive,
@@ -57,8 +70,8 @@ impl BackendResourceBroker {
profile_source_graph: Some(archive.reference.source_graph.clone()),
};
let stored = StoredResource {
runtime_id: runtime_id.map(|id| id.to_string()),
worker_id: worker_id.map(|id| id.to_string()),
runtime_id,
worker,
handle: handle.clone(),
archive,
};
@@ -117,8 +130,10 @@ impl BackendResourceBroker {
});
}
}
if let Some(expected_worker_id) = stored.worker_id.as_deref() {
if Some(expected_worker_id) != request.worker_id.as_deref() {
if let Some(expected_worker) = stored.worker.as_ref() {
if expected_worker.runtime_id != request.runtime_id
|| Some(expected_worker.worker_id.as_str()) != request.worker_id.as_deref()
{
return Err(BackendResourceError::Unauthorized {
message: "worker id does not match resource handle".to_string(),
});
@@ -167,7 +182,6 @@ fn verify_handle_shape(handle: &BackendResourceHandle) -> Result<(), BackendReso
mod tests {
use super::*;
use std::collections::BTreeMap;
use worker_runtime::identity::WorkerId;
use worker_runtime::profile_archive::{
ProfileSourceArchive, ProfileSourceArchiveRef, ProfileSourceGraphSummary, sha256_hex,
};
@@ -215,13 +229,13 @@ mod tests {
fn request(
handle: BackendResourceHandle,
runtime_id: &str,
worker_id: Option<&WorkerId>,
worker_id: Option<&str>,
) -> BackendResourceFetchRequest {
BackendResourceFetchRequest {
audit_correlation_id: handle.audit_correlation_id.clone(),
handle,
runtime_id: runtime_id.to_string(),
worker_id: worker_id.map(|id| id.to_string()),
worker_id: worker_id.map(str::to_string),
}
}
@@ -231,8 +245,7 @@ mod tests {
let runtime_id = "runtime-test";
let handle = broker.issue_profile_source_archive_handle(
"workspace-test",
Some(runtime_id),
None,
BackendResourceTarget::Runtime(runtime_id),
archive(),
);
let response = broker
@@ -253,8 +266,7 @@ mod tests {
let runtime_a = "runtime-a";
let handle = broker.issue_profile_source_archive_handle(
"workspace-test",
Some(runtime_a),
None,
BackendResourceTarget::Runtime(runtime_a),
archive(),
);
let err = broker
@@ -267,16 +279,15 @@ mod tests {
fn broker_rejects_worker_mismatch() {
let broker = BackendResourceBroker::default();
let runtime_id = "runtime-test";
let worker_a = WorkerId::new(1);
let worker_b = WorkerId::new(2);
let worker_a = RuntimeWorkerRef::new(runtime_id, "1");
let worker_b = RuntimeWorkerRef::new(runtime_id, "2");
let handle = broker.issue_profile_source_archive_handle(
"workspace-test",
Some(runtime_id),
Some(&worker_a),
BackendResourceTarget::Worker(&worker_a),
archive(),
);
let err = broker
.fetch_profile_source_archive(request(handle, &runtime_id, Some(&worker_b)))
.fetch_profile_source_archive(request(handle, runtime_id, Some(&worker_b.worker_id)))
.unwrap_err();
assert!(matches!(err, BackendResourceError::Unauthorized { .. }));
}
@@ -287,8 +298,7 @@ mod tests {
let runtime_id = "runtime-test";
let handle = broker.issue_profile_source_archive_handle(
"workspace-test",
Some(runtime_id),
None,
BackendResourceTarget::Runtime(runtime_id),
archive(),
);
broker
@@ -313,8 +323,7 @@ mod tests {
let runtime_id = "runtime-test";
let mut handle = broker.issue_profile_source_archive_handle(
"workspace-test",
Some(runtime_id),
None,
BackendResourceTarget::Runtime(runtime_id),
archive(),
);
handle.scope_id = Some("tampered-scope".to_string());
@@ -331,8 +340,7 @@ mod tests {
let archive = archive_with_len((DEFAULT_PROFILE_SOURCE_ARCHIVE_MAX_BYTES + 1) as usize);
let mut handle = broker.issue_profile_source_archive_handle(
"workspace-test",
Some(runtime_id),
None,
BackendResourceTarget::Runtime(runtime_id),
archive,
);
handle.max_bytes = DEFAULT_PROFILE_SOURCE_ARCHIVE_MAX_BYTES + 1024;
crates/workspace-server/src/server.rs
+376 -431
View File
File diff suppressed because it is too large Load Diff
crates/workspace-server/src/store.rs
+83 -98
View File
@@ -6,6 +6,8 @@ use async_trait::async_trait;
use rusqlite::{Connection, OptionalExtension, params};
use serde::{Deserialize, Serialize};
use worker_runtime::identity::RuntimeWorkerRef;
use crate::{Error, Result};
const WORKSPACES_V0_COLUMNS: &[&str] = &[
@@ -268,8 +270,7 @@ pub struct DeviceLoginFlowRecord {
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct WorkerRegistryRecord {
pub workspace_id: String,
pub runtime_id: String,
pub runtime_worker_id: u64,
pub worker: RuntimeWorkerRef,
pub display_name: String,
pub profile: Option<String>,
/// Retention state is explicit so `pinned` can be represented before prune exists.
@@ -287,8 +288,7 @@ pub struct TicketWorkerAssignmentRecord {
pub workspace_id: String,
pub ticket_id: String,
pub assignment_id: String,
pub runtime_id: String,
pub worker_id: String,
pub worker: RuntimeWorkerRef,
pub assigned_by: String,
pub assigned_at: String,
}
@@ -330,8 +330,7 @@ pub struct WorkdirRegistryRecord {
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct WorkerWorkdirLinkRecord {
pub workspace_id: String,
pub runtime_id: String,
pub runtime_worker_id: u64,
pub worker: RuntimeWorkerRef,
pub workdir_id: String,
pub role: String,
pub linked_at: String,
@@ -541,8 +540,7 @@ pub trait ControlPlaneStore: Send + Sync {
fn get_worker_registry(
&self,
workspace_id: &str,
runtime_id: &str,
runtime_worker_id: u64,
worker: &RuntimeWorkerRef,
) -> Result<Option<WorkerRegistryRecord>>;
fn list_worker_registry(
&self,
@@ -552,17 +550,12 @@ pub trait ControlPlaneStore: Send + Sync {
fn update_worker_retention(
&self,
workspace_id: &str,
runtime_id: &str,
runtime_worker_id: u64,
worker: &RuntimeWorkerRef,
retention_state: &str,
updated_at: &str,
) -> Result<bool>;
fn delete_worker_registry(
&self,
workspace_id: &str,
runtime_id: &str,
runtime_worker_id: u64,
) -> Result<bool>;
fn delete_worker_registry(&self, workspace_id: &str, worker: &RuntimeWorkerRef)
-> Result<bool>;
fn get_ticket_assignment_operation(
&self,
@@ -653,22 +646,19 @@ pub trait ControlPlaneStore: Send + Sync {
fn detach_worker_workdir(
&self,
workspace_id: &str,
runtime_id: &str,
runtime_worker_id: u64,
worker: &RuntimeWorkerRef,
expected_workdir_id: Option<&str>,
unlinked_at: &str,
) -> Result<Option<WorkerWorkdirLinkRecord>>;
fn worker_workdir_link_history_exists(
&self,
workspace_id: &str,
runtime_id: &str,
runtime_worker_id: u64,
worker: &RuntimeWorkerRef,
) -> Result<bool>;
fn list_worker_workdir_links(
&self,
workspace_id: &str,
runtime_id: &str,
runtime_worker_id: u64,
worker: &RuntimeWorkerRef,
) -> Result<Vec<WorkerWorkdirLinkRecord>>;
fn list_workdir_worker_links(
&self,
@@ -1637,8 +1627,8 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
updated_at = excluded.updated_at"#,
params![
record.workspace_id,
record.runtime_id,
record.runtime_worker_id,
record.worker.runtime_id,
record.worker.worker_id,
record.display_name,
record.profile,
record.retention_state,
@@ -1657,8 +1647,7 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
fn get_worker_registry(
&self,
workspace_id: &str,
runtime_id: &str,
runtime_worker_id: u64,
worker: &RuntimeWorkerRef,
) -> Result<Option<WorkerRegistryRecord>> {
self.with_conn(|conn| {
conn.query_row(
@@ -1666,7 +1655,7 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
"WHERE workspace_id = ?1 AND runtime_id = ?2 AND runtime_worker_id = ?3",
)
.as_str(),
params![workspace_id, runtime_id, runtime_worker_id],
params![workspace_id, worker.runtime_id, worker.worker_id],
read_worker_registry_record,
)
.optional()
@@ -1696,8 +1685,7 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
fn update_worker_retention(
&self,
workspace_id: &str,
runtime_id: &str,
runtime_worker_id: u64,
worker: &RuntimeWorkerRef,
retention_state: &str,
updated_at: &str,
) -> Result<bool> {
@@ -1708,8 +1696,8 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
WHERE workspace_id = ?1 AND runtime_id = ?2 AND runtime_worker_id = ?3"#,
params![
workspace_id,
runtime_id,
runtime_worker_id,
worker.runtime_id,
worker.worker_id,
retention_state,
updated_at
],
@@ -1721,8 +1709,7 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
fn delete_worker_registry(
&self,
workspace_id: &str,
runtime_id: &str,
runtime_worker_id: u64,
worker: &RuntimeWorkerRef,
) -> Result<bool> {
self.with_conn(|conn| {
let tx = conn.unchecked_transaction()?;
@@ -1730,11 +1717,11 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
r#"UPDATE worker_workdir_links
SET unlinked_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now')
WHERE workspace_id = ?1 AND runtime_id = ?2 AND runtime_worker_id = ?3 AND unlinked_at IS NULL"#,
params![workspace_id, runtime_id, runtime_worker_id],
params![workspace_id, worker.runtime_id, worker.worker_id],
)?;
let changed = tx.execute(
"DELETE FROM worker_registry WHERE workspace_id = ?1 AND runtime_id = ?2 AND runtime_worker_id = ?3",
params![workspace_id, runtime_id, runtime_worker_id],
params![workspace_id, worker.runtime_id, worker.worker_id],
)?;
tx.commit()?;
Ok(changed > 0)
@@ -1787,7 +1774,12 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
if existing.action == "assign"
&& existing.ticket_id == ticket_id
&& existing.runtime_id.as_deref() == Some(runtime_id)
&& (worker_id.is_none() || existing.worker_id.as_deref() == worker_id)
&& (worker_id.is_none()
|| existing
.worker
.as_ref()
.map(|worker| worker.worker_id.as_str())
== worker_id)
&& existing.expected_assignment_id.is_none()
&& existing.request_fingerprint.as_deref() == Some(request_fingerprint)
{
@@ -1855,8 +1847,7 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
{
if existing.action != if allow_reassign { "reassign" } else { "assign" }
|| existing.ticket_id != record.ticket_id
|| existing.runtime_id.as_deref() != Some(record.runtime_id.as_str())
|| existing.worker_id.as_deref() != Some(record.worker_id.as_str())
|| existing.worker.as_ref() != Some(&record.worker)
|| existing.expected_assignment_id.as_deref() != expected_assignment_id
{
return Err(Error::TicketAssignmentConflict(format!(
@@ -1937,8 +1928,8 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
record.workspace_id,
record.ticket_id,
record.assignment_id,
record.runtime_id,
record.worker_id,
record.worker.runtime_id,
record.worker.worker_id,
record.assigned_by,
record.assigned_at,
],
@@ -1952,8 +1943,8 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
record.workspace_id,
record.ticket_id,
record.assignment_id,
record.runtime_id,
record.worker_id,
record.worker.runtime_id,
record.worker.worker_id,
record.assigned_at,
],
)
@@ -1966,8 +1957,8 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
record.workspace_id,
record.ticket_id,
record.assignment_id,
record.runtime_id,
record.worker_id,
record.worker.runtime_id,
record.worker.worker_id,
record.assigned_at,
],
)
@@ -1976,7 +1967,7 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
return Err(map_assignment_constraint(
error,
&record.ticket_id,
&record.worker_id,
&record.worker.worker_id,
));
}
tx.execute(
@@ -2024,8 +2015,8 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
operation_id,
if allow_reassign { "reassign" } else { "assign" },
record.ticket_id,
record.runtime_id,
record.worker_id,
record.worker.runtime_id,
record.worker.worker_id,
record.assignment_id,
expected_assignment_id,
record.assigned_at,
@@ -2116,8 +2107,8 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
workspace_id,
operation_id,
ticket_id,
previous.runtime_id,
previous.worker_id,
previous.worker.runtime_id,
previous.worker.worker_id,
previous.assignment_id,
expected_assignment_id,
created_at,
@@ -2369,8 +2360,8 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
unlinked_at = NULL"#,
params![
record.workspace_id,
record.runtime_id,
record.runtime_worker_id,
record.worker.runtime_id,
record.worker.worker_id,
record.workdir_id,
record.role,
record.linked_at,
@@ -2451,8 +2442,8 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
WHERE workspace_id = ?1 AND runtime_id = ?2 AND runtime_worker_id = ?3 AND unlinked_at IS NULL"#,
params![
record.workspace_id,
record.runtime_id,
record.runtime_worker_id,
record.worker.runtime_id,
record.worker.worker_id,
],
read_worker_workdir_link_record,
)
@@ -2464,7 +2455,7 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
}
return Err(Error::WorkdirAttachmentConflict(format!(
"Worker {}:{} is already attached to Workdir {}",
record.runtime_id, record.runtime_worker_id, active.workdir_id
record.worker.runtime_id, record.worker.worker_id, active.workdir_id
)));
}
let active_for_workdir = tx
@@ -2479,7 +2470,7 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
if let Some(active) = active_for_workdir {
return Err(Error::WorkdirAttachmentConflict(format!(
"Workdir {} is already attached to Worker {}:{}",
record.workdir_id, active.runtime_id, active.runtime_worker_id
record.workdir_id, active.worker.runtime_id, active.worker.worker_id
)));
}
let write = tx.execute(
@@ -2491,8 +2482,8 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
unlinked_at = NULL"#,
params![
record.workspace_id,
record.runtime_id,
record.runtime_worker_id,
record.worker.runtime_id,
record.worker.worker_id,
record.workdir_id,
record.role,
record.linked_at,
@@ -2515,8 +2506,7 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
fn detach_worker_workdir(
&self,
workspace_id: &str,
runtime_id: &str,
runtime_worker_id: u64,
worker: &RuntimeWorkerRef,
expected_workdir_id: Option<&str>,
unlinked_at: &str,
) -> Result<Option<WorkerWorkdirLinkRecord>> {
@@ -2530,7 +2520,7 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
r#"SELECT workspace_id, runtime_id, runtime_worker_id, workdir_id, role, linked_at, unlinked_at
FROM worker_workdir_links
WHERE workspace_id = ?1 AND runtime_id = ?2 AND runtime_worker_id = ?3 AND unlinked_at IS NULL"#,
params![workspace_id, runtime_id, runtime_worker_id],
params![workspace_id, worker.runtime_id, worker.worker_id],
read_worker_workdir_link_record,
)
.optional()?;
@@ -2541,8 +2531,8 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
if let Some(expected_workdir_id) = expected_workdir_id {
if active.workdir_id != expected_workdir_id {
return Err(Error::WorkdirAttachmentConflict(format!(
"Worker {runtime_id}:{runtime_worker_id} is attached to Workdir {}, not {expected_workdir_id}",
active.workdir_id
"Worker {}:{} is attached to Workdir {}, not {expected_workdir_id}",
worker.runtime_id, worker.worker_id, active.workdir_id
)));
}
}
@@ -2550,11 +2540,12 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
r#"UPDATE worker_workdir_links
SET unlinked_at = ?4
WHERE workspace_id = ?1 AND runtime_id = ?2 AND runtime_worker_id = ?3 AND unlinked_at IS NULL"#,
params![workspace_id, runtime_id, runtime_worker_id, unlinked_at],
params![workspace_id, worker.runtime_id, worker.worker_id, unlinked_at],
)?;
if changed != 1 {
return Err(Error::WorkdirAttachmentConflict(format!(
"Worker {runtime_id}:{runtime_worker_id} attachment changed during detach"
"Worker {}:{} attachment changed during detach",
worker.runtime_id, worker.worker_id
)));
}
tx.commit()?;
@@ -2568,8 +2559,7 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
fn worker_workdir_link_history_exists(
&self,
workspace_id: &str,
runtime_id: &str,
runtime_worker_id: u64,
worker: &RuntimeWorkerRef,
) -> Result<bool> {
self.with_conn(|conn| {
let exists = conn.query_row(
@@ -2577,7 +2567,7 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
SELECT 1 FROM worker_workdir_links
WHERE workspace_id = ?1 AND runtime_id = ?2 AND runtime_worker_id = ?3
)"#,
params![workspace_id, runtime_id, runtime_worker_id],
params![workspace_id, worker.runtime_id, worker.worker_id],
|row| row.get(0),
)?;
Ok(exists)
@@ -2587,8 +2577,7 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
fn list_worker_workdir_links(
&self,
workspace_id: &str,
runtime_id: &str,
runtime_worker_id: u64,
worker: &RuntimeWorkerRef,
) -> Result<Vec<WorkerWorkdirLinkRecord>> {
self.with_conn(|conn| {
let mut stmt = conn.prepare(
@@ -2598,7 +2587,7 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
ORDER BY linked_at DESC"#,
)?;
let rows = stmt.query_map(
params![workspace_id, runtime_id, runtime_worker_id],
params![workspace_id, worker.runtime_id, worker.worker_id],
read_worker_workdir_link_record,
)?;
rows.collect::<std::result::Result<Vec<_>, _>>()
@@ -2807,8 +2796,7 @@ fn read_worker_workdir_link_record(
) -> rusqlite::Result<WorkerWorkdirLinkRecord> {
Ok(WorkerWorkdirLinkRecord {
workspace_id: row.get(0)?,
runtime_id: row.get(1)?,
runtime_worker_id: row.get(2)?,
worker: RuntimeWorkerRef::new(row.get::<_, String>(1)?, row.get::<_, u64>(2)?.to_string()),
workdir_id: row.get(3)?,
role: row.get(4)?,
linked_at: row.get(5)?,
@@ -2900,8 +2888,7 @@ fn worker_registry_select_sql(where_clause: &str) -> String {
fn read_worker_registry_record(row: &rusqlite::Row<'_>) -> rusqlite::Result<WorkerRegistryRecord> {
Ok(WorkerRegistryRecord {
workspace_id: row.get(0)?,
runtime_id: row.get(1)?,
runtime_worker_id: row.get(2)?,
worker: RuntimeWorkerRef::new(row.get::<_, String>(1)?, row.get::<_, u64>(2)?.to_string()),
display_name: row.get(3)?,
profile: row.get(4)?,
retention_state: row.get(5)?,
@@ -2933,8 +2920,7 @@ fn read_ticket_worker_assignment_record(
workspace_id: row.get(0)?,
ticket_id: row.get(1)?,
assignment_id: row.get(2)?,
runtime_id: row.get(3)?,
worker_id: row.get(4)?,
worker: RuntimeWorkerRef::new(row.get::<_, String>(3)?, row.get::<_, String>(4)?),
assigned_by: row.get(5)?,
assigned_at: row.get(6)?,
})
@@ -2960,7 +2946,7 @@ pub struct TicketAssignmentOperationRecord {
pub action: String,
pub ticket_id: String,
pub runtime_id: Option<String>,
pub worker_id: Option<String>,
pub worker: Option<RuntimeWorkerRef>,
pub assignment_id: Option<String>,
pub expected_assignment_id: Option<String>,
pub request_fingerprint: Option<String>,
@@ -2978,11 +2964,15 @@ fn read_assignment_operation(
WHERE workspace_id = ?1 AND operation_id = ?2"#,
params![workspace_id, operation_id],
|row| {
let runtime_id: Option<String> = row.get(2)?;
let worker_id: Option<String> = row.get(3)?;
Ok(TicketAssignmentOperationRecord {
action: row.get(0)?,
ticket_id: row.get(1)?,
runtime_id: row.get(2)?,
worker_id: row.get(3)?,
runtime_id: runtime_id.clone(),
worker: runtime_id
.zip(worker_id)
.map(|(runtime_id, worker_id)| RuntimeWorkerRef::new(runtime_id, worker_id)),
assignment_id: row.get(4)?,
expected_assignment_id: row.get(5)?,
request_fingerprint: row.get(6)?,
@@ -4204,8 +4194,7 @@ CREATE TABLE ticket_worker_links (ticket_id TEXT, worker_ref_key TEXT);
workspace_id: "workspace-a".to_string(),
ticket_id: "ticket-1".to_string(),
assignment_id: "assignment-1".to_string(),
runtime_id: "runtime-1".to_string(),
worker_id: "worker-1".to_string(),
worker: RuntimeWorkerRef::new("runtime-1", "worker-1"),
assigned_by: "user-1".to_string(),
assigned_at: "2026-07-31T00:00:01Z".to_string(),
};
@@ -4239,7 +4228,7 @@ CREATE TABLE ticket_worker_links (ticket_id TEXT, worker_ref_key TEXT);
.set_current_ticket_worker_assignment(
&TicketWorkerAssignmentRecord {
assignment_id: "implicit-reassign".to_string(),
worker_id: "worker-other".to_string(),
worker: RuntimeWorkerRef::new("runtime-1", "worker-other"),
..first.clone()
},
None,
@@ -4272,8 +4261,7 @@ CREATE TABLE ticket_worker_links (ticket_id TEXT, worker_ref_key TEXT);
let second = TicketWorkerAssignmentRecord {
assignment_id: "assignment-2".to_string(),
runtime_id: "runtime-2".to_string(),
worker_id: "worker-2".to_string(),
worker: RuntimeWorkerRef::new("runtime-2", "worker-2"),
assigned_by: "user-2".to_string(),
assigned_at: "2026-07-31T00:00:02Z".to_string(),
..first.clone()
@@ -4360,7 +4348,7 @@ CREATE TABLE ticket_worker_links (ticket_id TEXT, worker_ref_key TEXT);
.get_ticket_assignment_operation("workspace-a", "reserved-operation")
.unwrap()
.unwrap();
assert_eq!(pending.worker_id, None);
assert_eq!(pending.worker, None);
assert_eq!(
pending.request_fingerprint.as_deref(),
Some("sha256:reserved")
@@ -4376,8 +4364,7 @@ CREATE TABLE ticket_worker_links (ticket_id TEXT, worker_ref_key TEXT);
workspace_id: "workspace-a".to_string(),
ticket_id: "ticket-3".to_string(),
assignment_id: "assignment-3".to_string(),
runtime_id: "runtime-3".to_string(),
worker_id: "worker-3".to_string(),
worker: RuntimeWorkerRef::new("runtime-3", "worker-3"),
assigned_by: "runtime".to_string(),
assigned_at: "2026-07-31T00:00:06Z".to_string(),
};
@@ -4943,8 +4930,7 @@ CREATE TABLE ticket_assignment_operations (
let worker = WorkerRegistryRecord {
workspace_id: "local-dev".to_string(),
runtime_id: "embedded".to_string(),
runtime_worker_id: 1,
worker: RuntimeWorkerRef::new("embedded", "1"),
display_name: "Browser 1".to_string(),
profile: Some("builtin:companion".to_string()),
retention_state: "pinned".to_string(),
@@ -4996,8 +4982,7 @@ CREATE TABLE ticket_assignment_operations (
let link = WorkerWorkdirLinkRecord {
workspace_id: "local-dev".to_string(),
runtime_id: worker.runtime_id.clone(),
runtime_worker_id: worker.runtime_worker_id,
worker: worker.worker.clone(),
workdir_id: workdir.workdir_id.clone(),
role: "attachment".to_string(),
linked_at: "4".to_string(),
@@ -5033,7 +5018,7 @@ CREATE TABLE ticket_assignment_operations (
assert_eq!(
store
.get_worker_registry("local-dev", "embedded", 1)
.get_worker_registry("local-dev", &worker.worker)
.unwrap(),
Some(expected_worker.clone())
);
@@ -5049,7 +5034,7 @@ CREATE TABLE ticket_assignment_operations (
);
assert_eq!(
store
.list_worker_workdir_links("local-dev", "embedded", 1)
.list_worker_workdir_links("local-dev", &worker.worker)
.unwrap(),
vec![link.clone()]
);
@@ -5065,7 +5050,7 @@ CREATE TABLE ticket_assignment_operations (
));
let second_worker = WorkerRegistryRecord {
runtime_worker_id: 2,
worker: RuntimeWorkerRef::new("embedded", "2"),
display_name: "Browser 2".to_string(),
created_at: "5".to_string(),
updated_at: "5".to_string(),
@@ -5073,7 +5058,7 @@ CREATE TABLE ticket_assignment_operations (
};
store.upsert_worker_registry(&second_worker).unwrap();
let workdir_conflict = WorkerWorkdirLinkRecord {
runtime_worker_id: second_worker.runtime_worker_id,
worker: second_worker.worker.clone(),
linked_at: "5".to_string(),
..link.clone()
};
@@ -5082,17 +5067,17 @@ CREATE TABLE ticket_assignment_operations (
Err(Error::WorkdirAttachmentConflict(_))
));
assert!(matches!(
store.detach_worker_workdir("local-dev", "embedded", 1, Some("wrong-workdir"), "6"),
store.detach_worker_workdir("local-dev", &worker.worker, Some("wrong-workdir"), "6",),
Err(Error::WorkdirAttachmentConflict(_))
));
let detached = store
.detach_worker_workdir("local-dev", "embedded", 1, Some(&workdir.workdir_id), "6")
.detach_worker_workdir("local-dev", &worker.worker, Some(&workdir.workdir_id), "6")
.unwrap()
.unwrap();
assert_eq!(detached.unlinked_at.as_deref(), Some("6"));
assert!(
store
.worker_workdir_link_history_exists("local-dev", "embedded", 1)
.worker_workdir_link_history_exists("local-dev", &worker.worker)
.unwrap()
);
assert_eq!(
crates/workspace-server/src/workspace_subscription.rs
+23 -30
View File
@@ -8,6 +8,7 @@ use protocol::subscription::{
SubscriptionResponse, SubscriptionSnapshot, SubscriptionTerminationCode, SubscriptionWorker,
};
use tokio::sync::mpsc;
use worker_runtime::identity::RuntimeWorkerRef;
use crate::runtime_subscription::{BrokerSubscriptionEvent, RuntimeSubscriptionBroker};
use crate::server::{WorkspaceApi, connect_workspace_worker_protocol};
@@ -81,13 +82,8 @@ pub(crate) async fn serve_workspace_subscription(api: WorkspaceApi, socket: WebS
worker_id,
runtime_id: Some(runtime_id),
} => {
match connect_workspace_worker_protocol(
&api,
&runtime_id,
worker_id.as_str(),
)
.await
{
let worker = RuntimeWorkerRef::new(&runtime_id, worker_id.as_str());
match connect_workspace_worker_protocol(&api, &worker).await {
Ok(connection) => {
let methods = connection.methods.clone();
let task = tokio::spawn(run_worker_protocol(
@@ -323,16 +319,15 @@ async fn run_workspace_workers(
}
}
let mut revisions = HashMap::<String, u64>::new();
let mut initial_workers = workers
.values_mut()
.flat_map(|runtime| runtime.values_mut())
.map(|worker| {
let key = worker_key(worker.runtime_id.as_deref(), worker.worker_id.as_str());
worker.subject_revision = next_revision(&mut revisions, &key);
worker.clone()
})
.collect::<Vec<_>>();
let mut revisions = HashMap::<RuntimeWorkerRef, u64>::new();
let mut initial_workers = Vec::new();
for (runtime_id, runtime) in &mut workers {
for worker in runtime.values_mut() {
let worker_ref = RuntimeWorkerRef::new(runtime_id, worker.worker_id.as_str());
worker.subject_revision = next_revision(&mut revisions, &worker_ref);
initial_workers.push(worker.clone());
}
}
sort_workers(&mut initial_workers);
if send_frame(
&outbound,
@@ -359,8 +354,8 @@ async fn run_workspace_workers(
BrokerSubscriptionEvent::Snapshot { snapshot, .. } => {
let removed = workers.remove(&runtime_id).unwrap_or_default();
for worker in removed.values() {
let key = worker_key(Some(&runtime_id), worker.worker_id.as_str());
let revision = next_revision(&mut revisions, &key);
let worker_ref = RuntimeWorkerRef::new(&runtime_id, worker.worker_id.as_str());
let revision = next_revision(&mut revisions, &worker_ref);
if send_event(
&outbound,
&subscription_id,
@@ -379,8 +374,9 @@ async fn run_workspace_workers(
install_snapshot(&mut workers, &runtime_id, snapshot);
if let Some(current) = workers.get_mut(&runtime_id) {
for worker in current.values_mut() {
let key = worker_key(Some(&runtime_id), worker.worker_id.as_str());
let revision = next_revision(&mut revisions, &key);
let worker_ref =
RuntimeWorkerRef::new(&runtime_id, worker.worker_id.as_str());
let revision = next_revision(&mut revisions, &worker_ref);
worker.subject_revision = revision;
if send_event(
&outbound,
@@ -401,8 +397,8 @@ async fn run_workspace_workers(
BrokerSubscriptionEvent::Event { payload, .. } => match payload {
SubscriptionEventPayload::WorkerUpserted { mut worker } => {
worker.runtime_id = Some(runtime_id.clone());
let key = worker_key(Some(&runtime_id), worker.worker_id.as_str());
let revision = next_revision(&mut revisions, &key);
let worker_ref = RuntimeWorkerRef::new(&runtime_id, worker.worker_id.as_str());
let revision = next_revision(&mut revisions, &worker_ref);
worker.subject_revision = revision;
workers
.entry(runtime_id)
@@ -425,8 +421,8 @@ async fn run_workspace_workers(
.entry(runtime_id.clone())
.or_default()
.remove(worker_id.as_str());
let key = worker_key(Some(&runtime_id), worker_id.as_str());
let revision = next_revision(&mut revisions, &key);
let worker_ref = RuntimeWorkerRef::new(&runtime_id, worker_id.as_str());
let revision = next_revision(&mut revisions, &worker_ref);
if send_event(
&outbound,
&subscription_id,
@@ -534,14 +530,11 @@ async fn send_frame(
.map_err(|_| ())
}
fn next_revision(revisions: &mut HashMap<String, u64>, key: &str) -> u64 {
let revision = revisions.entry(key.to_string()).or_insert(0);
fn next_revision(revisions: &mut HashMap<RuntimeWorkerRef, u64>, worker: &RuntimeWorkerRef) -> u64 {
let revision = revisions.entry(worker.clone()).or_insert(0);
*revision = revision.saturating_add(1);
*revision
}
fn worker_key(runtime_id: Option<&str>, worker_id: &str) -> String {
format!("{}:{worker_id}", runtime_id.unwrap_or_default())
}
fn sort_workers(workers: &mut [SubscriptionWorker]) {
workers.sort_by(|left, right| {
left.runtime_id