server: multiplex worker protocol subscriptions
This commit is contained in:
@@ -248,6 +248,8 @@ pub enum EventSubscriptionSelector {
|
||||
},
|
||||
WorkerProtocol {
|
||||
worker_id: SubscriptionWorkerId,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
runtime_id: Option<String>,
|
||||
},
|
||||
/// Server-derived Workspace projection. Workspace identity comes from the
|
||||
/// authenticated connection and is deliberately absent from this selector.
|
||||
@@ -260,7 +262,16 @@ impl EventSubscriptionSelector {
|
||||
pub fn validate(&self) -> Result<(), SubscriptionValidationError> {
|
||||
match self {
|
||||
Self::WorkerLifecycle { worker_ids } => worker_ids.validate(),
|
||||
Self::WorkerProtocol { worker_id } => worker_id.validate(),
|
||||
Self::WorkerProtocol {
|
||||
worker_id,
|
||||
runtime_id,
|
||||
} => {
|
||||
worker_id.validate()?;
|
||||
if let Some(runtime_id) = runtime_id {
|
||||
validate_identifier("runtime_id", runtime_id, MAX_RESOURCE_ID_BYTES)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
Self::RuntimeWorkers | Self::WorkspaceWorkers | Self::WorkspaceWorkdirs => Ok(()),
|
||||
}
|
||||
}
|
||||
@@ -271,6 +282,7 @@ impl EventSubscriptionSelector {
|
||||
Self::WorkerLifecycle { worker_ids } => worker_ids.contains(worker_id),
|
||||
Self::WorkerProtocol {
|
||||
worker_id: selected,
|
||||
..
|
||||
} => selected == worker_id,
|
||||
Self::WorkspaceWorkdirs => false,
|
||||
}
|
||||
@@ -312,6 +324,7 @@ pub enum SubscriptionFramePayload {
|
||||
Request(SubscriptionRequest),
|
||||
Response(SubscriptionResponse),
|
||||
Event(SubscriptionEvent),
|
||||
WorkerProtocol(SubscriptionWorkerProtocolMethod),
|
||||
}
|
||||
|
||||
impl SubscriptionFramePayload {
|
||||
@@ -320,10 +333,24 @@ impl SubscriptionFramePayload {
|
||||
Self::Request(request) => request.validate(),
|
||||
Self::Response(response) => response.validate(),
|
||||
Self::Event(event) => event.validate(),
|
||||
Self::WorkerProtocol(message) => message.validate(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
|
||||
pub struct SubscriptionWorkerProtocolMethod {
|
||||
pub subscription_id: SubscriptionId,
|
||||
pub method: crate::Method,
|
||||
}
|
||||
|
||||
impl SubscriptionWorkerProtocolMethod {
|
||||
pub fn validate(&self) -> Result<(), SubscriptionValidationError> {
|
||||
self.subscription_id.validate()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
|
||||
#[serde(tag = "method", content = "params", rename_all = "snake_case")]
|
||||
@@ -617,6 +644,7 @@ impl SubscriptionSnapshot {
|
||||
(
|
||||
EventSubscriptionSelector::WorkerProtocol {
|
||||
worker_id: selected,
|
||||
..
|
||||
},
|
||||
Self::WorkerProtocol { worker_id, .. },
|
||||
) if selected == worker_id => worker_id.validate(),
|
||||
@@ -703,6 +731,7 @@ impl SubscriptionEventPayload {
|
||||
(
|
||||
EventSubscriptionSelector::WorkerProtocol {
|
||||
worker_id: selected,
|
||||
..
|
||||
},
|
||||
Self::WorkerProtocol { worker_id, .. },
|
||||
) if selected == worker_id => Ok(()),
|
||||
@@ -891,6 +920,37 @@ mod tests {
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn worker_protocol_method_uses_subscription_lane() {
|
||||
let frame = SubscriptionFrame::new(SubscriptionFramePayload::WorkerProtocol(
|
||||
SubscriptionWorkerProtocolMethod {
|
||||
subscription_id: subscription_id(),
|
||||
method: crate::Method::ListCompletions {
|
||||
kind: crate::CompletionKind::File,
|
||||
prefix: "src/".to_string(),
|
||||
},
|
||||
},
|
||||
));
|
||||
frame.validate().unwrap();
|
||||
assert_eq!(
|
||||
serde_json::to_value(frame).unwrap(),
|
||||
serde_json::json!({
|
||||
"protocol_version": 1,
|
||||
"frame": "worker_protocol",
|
||||
"message": {
|
||||
"subscription_id": "subscription-1",
|
||||
"method": {
|
||||
"method": "list_completions",
|
||||
"params": {
|
||||
"kind": "file",
|
||||
"prefix": "src/"
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn workspace_snapshot_allows_equal_local_worker_ids_from_distinct_runtimes() {
|
||||
let mut first = worker("1");
|
||||
|
||||
@@ -12,7 +12,8 @@ use crate::{
|
||||
SubscriptionFramePayload, SubscriptionId, SubscriptionRejectionCode, SubscriptionRequest,
|
||||
SubscriptionRequestId, SubscriptionResponse, SubscriptionSnapshot,
|
||||
SubscriptionTerminationCode, SubscriptionWorkdir, SubscriptionWorkdirId,
|
||||
SubscriptionWorker, SubscriptionWorkerId, SubscriptionWorkerIds, SubscriptionWorkerState,
|
||||
SubscriptionWorker, SubscriptionWorkerId, SubscriptionWorkerIds,
|
||||
SubscriptionWorkerProtocolMethod, SubscriptionWorkerState,
|
||||
},
|
||||
};
|
||||
|
||||
@@ -71,6 +72,7 @@ pub fn generated_protocol_types() -> String {
|
||||
push_decl::<SubscriptionRejectionCode>(&cfg, &mut output);
|
||||
push_decl::<SubscriptionTerminationCode>(&cfg, &mut output);
|
||||
push_decl::<SubscriptionRequest>(&cfg, &mut output);
|
||||
push_decl::<SubscriptionWorkerProtocolMethod>(&cfg, &mut output);
|
||||
push_decl::<SubscriptionResponse>(&cfg, &mut output);
|
||||
push_decl::<SubscriptionEvent>(&cfg, &mut output);
|
||||
push_decl::<SubscriptionFramePayload>(&cfg, &mut output);
|
||||
|
||||
@@ -772,7 +772,9 @@ async fn handle_frame(
|
||||
send_subscribe(socket, state, selector).await?;
|
||||
}
|
||||
}
|
||||
SubscriptionFramePayload::Request(_) => return Err(()),
|
||||
SubscriptionFramePayload::Request(_) | SubscriptionFramePayload::WorkerProtocol(_) => {
|
||||
return Err(());
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -3807,10 +3807,9 @@ async fn scoped_workspace_protocol_ws(
|
||||
{
|
||||
return Err(StatusCode::FORBIDDEN.into_response());
|
||||
}
|
||||
let broker = api.runtime_subscription_broker().clone();
|
||||
Ok(ws
|
||||
.on_upgrade(move |socket| {
|
||||
crate::workspace_subscription::serve_workspace_subscription(broker, socket)
|
||||
crate::workspace_subscription::serve_workspace_subscription(api, socket)
|
||||
})
|
||||
.into_response())
|
||||
}
|
||||
@@ -6685,6 +6684,131 @@ async fn worker_protocol_ws(
|
||||
ws.on_upgrade(move |socket| worker_protocol_ws_session(source, socket))
|
||||
}
|
||||
|
||||
pub(crate) struct WorkspaceWorkerProtocolConnection {
|
||||
pub(crate) methods: tokio::sync::mpsc::Sender<protocol::Method>,
|
||||
pub(crate) events: tokio::sync::mpsc::Receiver<protocol::Event>,
|
||||
}
|
||||
|
||||
pub(crate) async fn connect_workspace_worker_protocol(
|
||||
api: &WorkspaceApi,
|
||||
runtime_id: &str,
|
||||
worker_id: &str,
|
||||
) -> Result<WorkspaceWorkerProtocolConnection> {
|
||||
let source = match api.observation_proxy.source(runtime_id, worker_id) {
|
||||
Ok(source) => source,
|
||||
Err(ObservationProxyError::WorkerNotFound(_)) => api
|
||||
.runtime
|
||||
.observation_source(runtime_id, worker_id)
|
||||
.map_err(|error| error.into_error())?,
|
||||
Err(error) => {
|
||||
return Err(Error::RuntimeOperationFailed {
|
||||
runtime_id: runtime_id.to_string(),
|
||||
code: error.code().to_string(),
|
||||
message: error.message().to_string(),
|
||||
});
|
||||
}
|
||||
};
|
||||
match source {
|
||||
RuntimeObservationSource::RemoteWs(config) => connect_remote_worker_protocol(config).await,
|
||||
RuntimeObservationSource::Embedded(source) => {
|
||||
connect_embedded_worker_protocol(source).await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn connect_remote_worker_protocol(
|
||||
config: RuntimeObservationSourceConfig,
|
||||
) -> Result<WorkspaceWorkerProtocolConnection> {
|
||||
let mut request = config
|
||||
.endpoint
|
||||
.clone()
|
||||
.into_client_request()
|
||||
.map_err(|error| Error::Config(format!("invalid Runtime protocol endpoint: {error}")))?;
|
||||
if let Some(token) = &config.bearer_token {
|
||||
request.headers_mut().insert(
|
||||
"authorization",
|
||||
format!("Bearer {token}").parse().map_err(|error| {
|
||||
Error::Config(format!("invalid Runtime authorization: {error}"))
|
||||
})?,
|
||||
);
|
||||
}
|
||||
let (socket, _) =
|
||||
connect_async(request)
|
||||
.await
|
||||
.map_err(|error| Error::RuntimeOperationFailed {
|
||||
runtime_id: config.runtime_id.clone(),
|
||||
code: "worker_protocol_connect_failed".to_string(),
|
||||
message: error.to_string(),
|
||||
})?;
|
||||
let (mut sink, mut stream) = socket.split();
|
||||
let (methods, mut method_receiver) = tokio::sync::mpsc::channel(256);
|
||||
let (event_sender, events) = tokio::sync::mpsc::channel(512);
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
tokio::select! {
|
||||
method = method_receiver.recv() => {
|
||||
let Some(method) = method else { break };
|
||||
let Ok(text) = protocol::stream::encode_method(&method) else { break };
|
||||
if sink.send(TungsteniteMessage::Text(text.into())).await.is_err() { break; }
|
||||
}
|
||||
message = stream.next() => match message {
|
||||
Some(Ok(TungsteniteMessage::Text(text))) => {
|
||||
let Ok(event) = protocol::stream::decode_event(text.as_ref()) else { break; };
|
||||
if event_sender.send(event).await.is_err() { break; }
|
||||
}
|
||||
Some(Ok(TungsteniteMessage::Ping(value))) => {
|
||||
if sink.send(TungsteniteMessage::Pong(value)).await.is_err() { break; }
|
||||
}
|
||||
Some(Ok(TungsteniteMessage::Pong(_))) => {}
|
||||
_ => break,
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
Ok(WorkspaceWorkerProtocolConnection { methods, events })
|
||||
}
|
||||
|
||||
async fn connect_embedded_worker_protocol(
|
||||
source: crate::observation::EmbeddedRuntimeObservationSource,
|
||||
) -> Result<WorkspaceWorkerProtocolConnection> {
|
||||
let mut upstream =
|
||||
RuntimeObservationClient::connect(&RuntimeObservationSource::Embedded(source.clone()))
|
||||
.await
|
||||
.map_err(|error| Error::RuntimeOperationFailed {
|
||||
runtime_id: source.runtime_id.clone(),
|
||||
code: error.code().to_string(),
|
||||
message: error.message().to_string(),
|
||||
})?;
|
||||
let (methods, mut method_receiver) = tokio::sync::mpsc::channel(256);
|
||||
let (event_sender, events) = tokio::sync::mpsc::channel(512);
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
tokio::select! {
|
||||
method = method_receiver.recv() => {
|
||||
let Some(method) = method else { break };
|
||||
match source.runtime.send_protocol_method(&source.worker_ref, method) {
|
||||
Ok(direct_events) => {
|
||||
for event in direct_events {
|
||||
if event_sender.send(event).await.is_err() { return; }
|
||||
}
|
||||
}
|
||||
Err(error) => {
|
||||
if event_sender.send(protocol_error_event(error.to_string())).await.is_err() { return; }
|
||||
}
|
||||
}
|
||||
}
|
||||
event = upstream.next_event() => match event {
|
||||
Ok(event) => {
|
||||
if event_sender.send(event.payload).await.is_err() { break; }
|
||||
}
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
Ok(WorkspaceWorkerProtocolConnection { methods, events })
|
||||
}
|
||||
|
||||
async fn worker_protocol_ws_session(source: RuntimeObservationSource, socket: WebSocket) {
|
||||
match source {
|
||||
RuntimeObservationSource::RemoteWs(config) => {
|
||||
@@ -12716,6 +12840,29 @@ mod tests {
|
||||
.unwrap()
|
||||
.to_string();
|
||||
|
||||
let spawn_request = WorkerSpawnRequest {
|
||||
intent: WorkerSpawnIntent::WorkspaceCompanion,
|
||||
requested_worker_name: Some("multiplexed-console".to_string()),
|
||||
acceptance: WorkerSpawnAcceptanceRequirement::RunAccepted {
|
||||
expected_segments: 0,
|
||||
},
|
||||
profile: worker_runtime::catalog::ProfileSelector::Builtin(
|
||||
"builtin:companion".to_string(),
|
||||
),
|
||||
ticket_assignment: None,
|
||||
initial_input: None,
|
||||
working_directory_request: None,
|
||||
resolved_working_directory_request: None,
|
||||
resolved_working_directory: None,
|
||||
resolved_config_bundle: Some(runtime_test_bundle()),
|
||||
resolved_workspace_api: None,
|
||||
};
|
||||
let spawned = api
|
||||
.spawn_workspace_worker(EMBEDDED_WORKER_RUNTIME_ID, spawn_request)
|
||||
.unwrap();
|
||||
assert_eq!(spawned.state, WorkerOperationState::Accepted);
|
||||
let worker_id = spawned.worker.unwrap().worker_id;
|
||||
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let app = build_router(api);
|
||||
@@ -12757,6 +12904,141 @@ mod tests {
|
||||
}
|
||||
)
|
||||
));
|
||||
|
||||
let subscribe_protocol = protocol::subscription::SubscriptionFrame::new(
|
||||
protocol::subscription::SubscriptionFramePayload::Request(
|
||||
protocol::subscription::SubscriptionRequest::SubscribeEvents {
|
||||
request_id: protocol::subscription::SubscriptionRequestId::new("request-2")
|
||||
.unwrap(),
|
||||
selector: protocol::subscription::EventSubscriptionSelector::WorkerProtocol {
|
||||
worker_id: protocol::subscription::SubscriptionWorkerId::new(
|
||||
worker_id.clone(),
|
||||
)
|
||||
.unwrap(),
|
||||
runtime_id: Some(EMBEDDED_WORKER_RUNTIME_ID.to_string()),
|
||||
},
|
||||
},
|
||||
),
|
||||
);
|
||||
socket
|
||||
.send(Message::Text(
|
||||
serde_json::to_string(&subscribe_protocol).unwrap().into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
let protocol_subscription_id = loop {
|
||||
let Message::Text(text) = socket.next().await.unwrap().unwrap() else {
|
||||
continue;
|
||||
};
|
||||
let frame: protocol::subscription::SubscriptionFrame =
|
||||
serde_json::from_str(text.as_str()).unwrap();
|
||||
if let protocol::subscription::SubscriptionFramePayload::Response(
|
||||
protocol::subscription::SubscriptionResponse::Subscribed {
|
||||
subscription_id,
|
||||
selector:
|
||||
protocol::subscription::EventSubscriptionSelector::WorkerProtocol { .. },
|
||||
..
|
||||
},
|
||||
) = frame.payload
|
||||
{
|
||||
break subscription_id;
|
||||
}
|
||||
};
|
||||
let second_subscribe = protocol::subscription::SubscriptionFrame::new(
|
||||
protocol::subscription::SubscriptionFramePayload::Request(
|
||||
protocol::subscription::SubscriptionRequest::SubscribeEvents {
|
||||
request_id: protocol::subscription::SubscriptionRequestId::new("request-3")
|
||||
.unwrap(),
|
||||
selector: protocol::subscription::EventSubscriptionSelector::WorkerProtocol {
|
||||
worker_id: protocol::subscription::SubscriptionWorkerId::new(
|
||||
worker_id.clone(),
|
||||
)
|
||||
.unwrap(),
|
||||
runtime_id: Some(EMBEDDED_WORKER_RUNTIME_ID.to_string()),
|
||||
},
|
||||
},
|
||||
),
|
||||
);
|
||||
socket
|
||||
.send(Message::Text(
|
||||
serde_json::to_string(&second_subscribe).unwrap().into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
let second_protocol_subscription_id = loop {
|
||||
let Message::Text(text) = socket.next().await.unwrap().unwrap() else {
|
||||
continue;
|
||||
};
|
||||
let frame: protocol::subscription::SubscriptionFrame =
|
||||
serde_json::from_str(text.as_str()).unwrap();
|
||||
if let protocol::subscription::SubscriptionFramePayload::Response(
|
||||
protocol::subscription::SubscriptionResponse::Subscribed {
|
||||
subscription_id,
|
||||
selector:
|
||||
protocol::subscription::EventSubscriptionSelector::WorkerProtocol { .. },
|
||||
..
|
||||
},
|
||||
) = frame.payload
|
||||
{
|
||||
break subscription_id;
|
||||
}
|
||||
};
|
||||
assert_ne!(protocol_subscription_id, second_protocol_subscription_id);
|
||||
let unsubscribe = protocol::subscription::SubscriptionFrame::new(
|
||||
protocol::subscription::SubscriptionFramePayload::Request(
|
||||
protocol::subscription::SubscriptionRequest::UnsubscribeEvents {
|
||||
request_id: protocol::subscription::SubscriptionRequestId::new("request-4")
|
||||
.unwrap(),
|
||||
subscription_id: protocol_subscription_id,
|
||||
},
|
||||
),
|
||||
);
|
||||
socket
|
||||
.send(Message::Text(
|
||||
serde_json::to_string(&unsubscribe).unwrap().into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let method = protocol::subscription::SubscriptionFrame::new(
|
||||
protocol::subscription::SubscriptionFramePayload::WorkerProtocol(
|
||||
protocol::subscription::SubscriptionWorkerProtocolMethod {
|
||||
subscription_id: second_protocol_subscription_id.clone(),
|
||||
method: protocol::Method::ListCompletions {
|
||||
kind: protocol::CompletionKind::File,
|
||||
prefix: String::new(),
|
||||
},
|
||||
},
|
||||
),
|
||||
);
|
||||
socket
|
||||
.send(Message::Text(
|
||||
serde_json::to_string(&method).unwrap().into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
loop {
|
||||
let Message::Text(text) = socket.next().await.unwrap().unwrap() else {
|
||||
continue;
|
||||
};
|
||||
let frame: protocol::subscription::SubscriptionFrame =
|
||||
serde_json::from_str(text.as_str()).unwrap();
|
||||
if matches!(
|
||||
frame.payload,
|
||||
protocol::subscription::SubscriptionFramePayload::Event(
|
||||
protocol::subscription::SubscriptionEvent::Event {
|
||||
subscription_id,
|
||||
payload: protocol::subscription::SubscriptionEventPayload::WorkerProtocol {
|
||||
event: protocol::Event::Completions { .. },
|
||||
..
|
||||
},
|
||||
..
|
||||
}
|
||||
) if subscription_id == second_protocol_subscription_id
|
||||
) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
server.abort();
|
||||
}
|
||||
|
||||
|
||||
@@ -10,24 +10,35 @@ use protocol::subscription::{
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use crate::runtime_subscription::{BrokerSubscriptionEvent, RuntimeSubscriptionBroker};
|
||||
use crate::server::{WorkspaceApi, connect_workspace_worker_protocol};
|
||||
|
||||
const OUTBOUND_CAPACITY: usize = 256;
|
||||
|
||||
pub(crate) async fn serve_workspace_subscription(
|
||||
broker: RuntimeSubscriptionBroker,
|
||||
socket: WebSocket,
|
||||
) {
|
||||
struct ActiveSubscription {
|
||||
task: tokio::task::JoinHandle<()>,
|
||||
methods: Option<mpsc::Sender<protocol::Method>>,
|
||||
}
|
||||
|
||||
pub(crate) async fn serve_workspace_subscription(api: WorkspaceApi, socket: WebSocket) {
|
||||
let broker = api.runtime_subscription_broker().clone();
|
||||
let (mut socket_sender, mut socket_receiver) = socket.split();
|
||||
let (outbound, mut outbound_receiver) = mpsc::channel::<WsMessage>(OUTBOUND_CAPACITY);
|
||||
let (control_outbound, mut control_receiver) = mpsc::channel::<WsMessage>(OUTBOUND_CAPACITY);
|
||||
let (protocol_outbound, mut protocol_receiver) = mpsc::channel::<WsMessage>(OUTBOUND_CAPACITY);
|
||||
let writer = tokio::spawn(async move {
|
||||
while let Some(message) = outbound_receiver.recv().await {
|
||||
loop {
|
||||
let message = tokio::select! {
|
||||
biased;
|
||||
message = control_receiver.recv() => message,
|
||||
message = protocol_receiver.recv() => message,
|
||||
};
|
||||
let Some(message) = message else { break };
|
||||
if socket_sender.send(message).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
});
|
||||
let mut next_subscription_id = 1_u64;
|
||||
let mut subscriptions = HashMap::<SubscriptionId, tokio::task::JoinHandle<()>>::new();
|
||||
let mut subscriptions = HashMap::<SubscriptionId, ActiveSubscription>::new();
|
||||
|
||||
while let Some(message) = socket_receiver.next().await {
|
||||
let Ok(message) = message else { break };
|
||||
@@ -39,50 +50,91 @@ pub(crate) async fn serve_workspace_subscription(
|
||||
if frame.validate().is_err() {
|
||||
break;
|
||||
}
|
||||
let SubscriptionFramePayload::Request(request) = frame.payload else {
|
||||
break;
|
||||
};
|
||||
subscriptions.retain(|_, task| !task.is_finished());
|
||||
match request {
|
||||
SubscriptionRequest::SubscribeEvents {
|
||||
subscriptions.retain(|_, subscription| !subscription.task.is_finished());
|
||||
match frame.payload {
|
||||
SubscriptionFramePayload::Request(SubscriptionRequest::SubscribeEvents {
|
||||
request_id,
|
||||
selector,
|
||||
} => {
|
||||
if selector != EventSubscriptionSelector::WorkspaceWorkers {
|
||||
let _ = send_frame(&outbound, SubscriptionFrame::new(
|
||||
SubscriptionFramePayload::Response(
|
||||
SubscriptionResponse::SubscriptionRejected {
|
||||
request_id,
|
||||
subscription_id: None,
|
||||
code: SubscriptionRejectionCode::UnsupportedSelector,
|
||||
message: "Workspace clients may subscribe only to workspace_workers on this endpoint".to_string(),
|
||||
},
|
||||
),
|
||||
)).await;
|
||||
continue;
|
||||
}
|
||||
}) => {
|
||||
let subscription_id = SubscriptionId::new(format!(
|
||||
"workspace-subscription-{next_subscription_id}"
|
||||
))
|
||||
.expect("generated Workspace subscription id is valid");
|
||||
next_subscription_id = next_subscription_id.saturating_add(1);
|
||||
let task = tokio::spawn(run_workspace_workers(
|
||||
broker.clone(),
|
||||
request_id,
|
||||
subscription_id.clone(),
|
||||
outbound.clone(),
|
||||
));
|
||||
subscriptions.insert(subscription_id, task);
|
||||
match selector {
|
||||
EventSubscriptionSelector::WorkspaceWorkers => {
|
||||
let task = tokio::spawn(run_workspace_workers(
|
||||
broker.clone(),
|
||||
request_id,
|
||||
subscription_id.clone(),
|
||||
control_outbound.clone(),
|
||||
));
|
||||
subscriptions.insert(
|
||||
subscription_id,
|
||||
ActiveSubscription {
|
||||
task,
|
||||
methods: None,
|
||||
},
|
||||
);
|
||||
}
|
||||
EventSubscriptionSelector::WorkerProtocol {
|
||||
worker_id,
|
||||
runtime_id: Some(runtime_id),
|
||||
} => {
|
||||
match connect_workspace_worker_protocol(
|
||||
&api,
|
||||
&runtime_id,
|
||||
worker_id.as_str(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(connection) => {
|
||||
let methods = connection.methods.clone();
|
||||
let task = tokio::spawn(run_worker_protocol(
|
||||
request_id,
|
||||
subscription_id.clone(),
|
||||
runtime_id,
|
||||
worker_id,
|
||||
connection.events,
|
||||
control_outbound.clone(),
|
||||
protocol_outbound.clone(),
|
||||
));
|
||||
subscriptions.insert(
|
||||
subscription_id,
|
||||
ActiveSubscription {
|
||||
task,
|
||||
methods: Some(methods),
|
||||
},
|
||||
);
|
||||
}
|
||||
Err(error) => {
|
||||
let _ = send_rejected(
|
||||
&control_outbound,
|
||||
request_id,
|
||||
SubscriptionRejectionCode::ResourceNotFound,
|
||||
error.to_string(),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
let _ = send_rejected(
|
||||
&control_outbound, request_id, SubscriptionRejectionCode::UnsupportedSelector,
|
||||
"Workspace clients may subscribe only to workspace_workers or a runtime-scoped worker_protocol selector".to_string(),
|
||||
).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
SubscriptionRequest::UnsubscribeEvents {
|
||||
SubscriptionFramePayload::Request(SubscriptionRequest::UnsubscribeEvents {
|
||||
request_id,
|
||||
subscription_id,
|
||||
} => {
|
||||
if let Some(task) = subscriptions.remove(&subscription_id) {
|
||||
task.abort();
|
||||
}) => {
|
||||
if let Some(subscription) = subscriptions.remove(&subscription_id) {
|
||||
subscription.task.abort();
|
||||
}
|
||||
if send_frame(
|
||||
&outbound,
|
||||
&control_outbound,
|
||||
SubscriptionFrame::new(SubscriptionFramePayload::Response(
|
||||
SubscriptionResponse::Unsubscribed {
|
||||
request_id,
|
||||
@@ -96,10 +148,24 @@ pub(crate) async fn serve_workspace_subscription(
|
||||
break;
|
||||
}
|
||||
}
|
||||
SubscriptionFramePayload::WorkerProtocol(message) => {
|
||||
let Some(methods) = subscriptions
|
||||
.get(&message.subscription_id)
|
||||
.and_then(|value| value.methods.clone())
|
||||
else {
|
||||
break;
|
||||
};
|
||||
if methods.send(message.method).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
SubscriptionFramePayload::Response(_) | SubscriptionFramePayload::Event(_) => {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
WsMessage::Ping(value) => {
|
||||
if outbound.send(WsMessage::Pong(value)).await.is_err() {
|
||||
if control_outbound.send(WsMessage::Pong(value)).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
@@ -108,13 +174,108 @@ pub(crate) async fn serve_workspace_subscription(
|
||||
}
|
||||
}
|
||||
|
||||
for (_, task) in subscriptions {
|
||||
task.abort();
|
||||
for (_, subscription) in subscriptions {
|
||||
subscription.task.abort();
|
||||
}
|
||||
drop(outbound);
|
||||
drop(control_outbound);
|
||||
drop(protocol_outbound);
|
||||
let _ = writer.await;
|
||||
}
|
||||
|
||||
async fn send_rejected(
|
||||
outbound: &mpsc::Sender<WsMessage>,
|
||||
request_id: protocol::subscription::SubscriptionRequestId,
|
||||
code: SubscriptionRejectionCode,
|
||||
message: String,
|
||||
) -> Result<(), ()> {
|
||||
send_frame(
|
||||
outbound,
|
||||
SubscriptionFrame::new(SubscriptionFramePayload::Response(
|
||||
SubscriptionResponse::SubscriptionRejected {
|
||||
request_id,
|
||||
subscription_id: None,
|
||||
code,
|
||||
message,
|
||||
},
|
||||
)),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn run_worker_protocol(
|
||||
request_id: protocol::subscription::SubscriptionRequestId,
|
||||
subscription_id: SubscriptionId,
|
||||
runtime_id: String,
|
||||
worker_id: protocol::subscription::SubscriptionWorkerId,
|
||||
mut events: mpsc::Receiver<protocol::Event>,
|
||||
control_outbound: mpsc::Sender<WsMessage>,
|
||||
protocol_outbound: mpsc::Sender<WsMessage>,
|
||||
) {
|
||||
if send_frame(
|
||||
&control_outbound,
|
||||
SubscriptionFrame::new(SubscriptionFramePayload::Response(
|
||||
SubscriptionResponse::Subscribed {
|
||||
request_id,
|
||||
subscription_id: subscription_id.clone(),
|
||||
selector: EventSubscriptionSelector::WorkerProtocol {
|
||||
worker_id: worker_id.clone(),
|
||||
runtime_id: Some(runtime_id),
|
||||
},
|
||||
snapshot_revision: 0,
|
||||
snapshot: SubscriptionSnapshot::WorkerProtocol {
|
||||
worker_id: worker_id.clone(),
|
||||
events: Vec::new(),
|
||||
},
|
||||
},
|
||||
)),
|
||||
)
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
return;
|
||||
}
|
||||
let mut subject_revision = 0_u64;
|
||||
while let Some(event) = events.recv().await {
|
||||
subject_revision = subject_revision.saturating_add(1);
|
||||
let frame =
|
||||
SubscriptionFrame::new(SubscriptionFramePayload::Event(SubscriptionEvent::Event {
|
||||
subscription_id: subscription_id.clone(),
|
||||
subject_revision,
|
||||
payload: SubscriptionEventPayload::WorkerProtocol {
|
||||
worker_id: worker_id.clone(),
|
||||
event,
|
||||
},
|
||||
}));
|
||||
if try_send_frame(&protocol_outbound, frame).is_err() {
|
||||
let _ = send_frame(
|
||||
&control_outbound,
|
||||
SubscriptionFrame::new(SubscriptionFramePayload::Event(
|
||||
SubscriptionEvent::SubscriptionClosed {
|
||||
subscription_id: subscription_id.clone(),
|
||||
code: SubscriptionTerminationCode::Lagged,
|
||||
message:
|
||||
"Worker protocol subscriber lagged; resubscribe for a fresh snapshot"
|
||||
.to_string(),
|
||||
},
|
||||
)),
|
||||
)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
let _ = send_frame(
|
||||
&control_outbound,
|
||||
SubscriptionFrame::new(SubscriptionFramePayload::Event(
|
||||
SubscriptionEvent::SubscriptionClosed {
|
||||
subscription_id,
|
||||
code: SubscriptionTerminationCode::ResourceGone,
|
||||
message: "Worker protocol stream closed".to_string(),
|
||||
},
|
||||
)),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
async fn run_workspace_workers(
|
||||
broker: RuntimeSubscriptionBroker,
|
||||
request_id: protocol::subscription::SubscriptionRequestId,
|
||||
@@ -352,6 +513,14 @@ async fn send_event(
|
||||
.await
|
||||
}
|
||||
|
||||
fn try_send_frame(outbound: &mpsc::Sender<WsMessage>, frame: SubscriptionFrame) -> Result<(), ()> {
|
||||
frame.validate().map_err(|_| ())?;
|
||||
let text = serde_json::to_string(&frame).map_err(|_| ())?;
|
||||
outbound
|
||||
.try_send(WsMessage::Text(text.into()))
|
||||
.map_err(|_| ())
|
||||
}
|
||||
|
||||
async fn send_frame(
|
||||
outbound: &mpsc::Sender<WsMessage>,
|
||||
frame: SubscriptionFrame,
|
||||
|
||||
@@ -111,7 +111,7 @@ export type SubscriptionWorkerIds = Array<SubscriptionWorkerId>;
|
||||
|
||||
export type SubscriptionWorkerState = "idle" | "running" | "paused" | "stopped" | "cancelled";
|
||||
|
||||
export type EventSubscriptionSelector = { "topic": "runtime_workers" } | { "topic": "worker_lifecycle", worker_ids: SubscriptionWorkerIds, } | { "topic": "worker_protocol", worker_id: SubscriptionWorkerId, } | { "topic": "workspace_workers" } | { "topic": "workspace_workdirs" };
|
||||
export type EventSubscriptionSelector = { "topic": "runtime_workers" } | { "topic": "worker_lifecycle", worker_ids: SubscriptionWorkerIds, } | { "topic": "worker_protocol", worker_id: SubscriptionWorkerId, runtime_id?: string | null, } | { "topic": "workspace_workers" } | { "topic": "workspace_workdirs" };
|
||||
|
||||
export type SubscriptionWorker = { worker_id: SubscriptionWorkerId,
|
||||
/**
|
||||
@@ -136,13 +136,15 @@ export type SubscriptionTerminationCode = "lagged" | "resource_gone" | "unauthor
|
||||
|
||||
export type SubscriptionRequest = { "method": "subscribe_events", "params": { request_id: SubscriptionRequestId, selector: EventSubscriptionSelector, } } | { "method": "unsubscribe_events", "params": { request_id: SubscriptionRequestId, subscription_id: SubscriptionId, } };
|
||||
|
||||
export type SubscriptionWorkerProtocolMethod = { subscription_id: SubscriptionId, method: Method, };
|
||||
|
||||
export type SubscriptionResponse = { "result": "subscribed", "payload": { request_id: SubscriptionRequestId, subscription_id: SubscriptionId, selector: EventSubscriptionSelector, snapshot_revision: number, snapshot: SubscriptionSnapshot, } } | { "result": "unsubscribed", "payload": { request_id: SubscriptionRequestId, subscription_id: SubscriptionId, } } | { "result": "subscription_rejected", "payload": { request_id: SubscriptionRequestId, subscription_id?: SubscriptionId | null, code: SubscriptionRejectionCode, message: string, } };
|
||||
|
||||
export type SubscriptionEvent = { "event": "event", "data": { subscription_id: SubscriptionId, subject_revision: number, payload: SubscriptionEventPayload, } } | { "event": "subscription_closed", "data": { subscription_id: SubscriptionId, code: SubscriptionTerminationCode, message: string, } };
|
||||
|
||||
export type SubscriptionFramePayload = { "frame": "request", "message": SubscriptionRequest } | { "frame": "response", "message": SubscriptionResponse } | { "frame": "event", "message": SubscriptionEvent };
|
||||
export type SubscriptionFramePayload = { "frame": "request", "message": SubscriptionRequest } | { "frame": "response", "message": SubscriptionResponse } | { "frame": "event", "message": SubscriptionEvent } | { "frame": "worker_protocol", "message": SubscriptionWorkerProtocolMethod };
|
||||
|
||||
export type SubscriptionFrame = { protocol_version: number, } & ({ "frame": "request", "message": SubscriptionRequest } | { "frame": "response", "message": SubscriptionResponse } | { "frame": "event", "message": SubscriptionEvent });
|
||||
export type SubscriptionFrame = { protocol_version: number, } & ({ "frame": "request", "message": SubscriptionRequest } | { "frame": "response", "message": SubscriptionResponse } | { "frame": "event", "message": SubscriptionEvent } | { "frame": "worker_protocol", "message": SubscriptionWorkerProtocolMethod });
|
||||
|
||||
export type Method = { "method": "run", "params": { input: Array<Segment>, } } | { "method": "notify", "params": { message: string, auto_run?: boolean, } } | { "method": "worker_event", "params": WorkerEvent } | { "method": "resume" } | { "method": "cancel" } | { "method": "pause" } | { "method": "compact" } | { "method": "list_rewind_targets" } | { "method": "rewind_to", "params": { target: RewindTargetId, expected_head_entries: number, } } | { "method": "shutdown" } | { "method": "list_completions", "params": { kind: CompletionKind, prefix: string, } } | { "method": "list_workers" } | { "method": "restore_worker", "params": { name: string, } } | { "method": "register_peer", "params": { name: string, } };
|
||||
|
||||
|
||||
Reference in New Issue
Block a user