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

feat: integrate worker submit queue

Browse Source
This commit is contained in:
Hare
2026-09-06 04:24:07 +09:00
parent e7803d1aba 9da20d15da
commit 3ed1545c3c
51 changed files with 5022 additions and 929 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/workspace-server/src/hosts.rs
+9 -7
View File
@@ -538,7 +538,7 @@ fn initial_worker_input(segments: &[Segment]) -> Option<EmbeddedWorkerInput> {
Some(EmbeddedWorkerInput {
kind: EmbeddedWorkerInputKind::User,
content: Segment::flatten_to_text(segments),
submission_id: None,
submission_request_id: None,
segments: Some(segments.to_vec()),
})
}
@@ -2706,7 +2706,7 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime {
WorkerInputKind::RegisterPeer => EmbeddedWorkerInputKind::RegisterPeer,
},
content: request.content,
submission_id: None,
submission_request_id: None,
segments: request.segments,
};
match self.runtime.send_input(&worker_ref, input) {
@@ -3934,7 +3934,7 @@ impl WorkspaceWorkerRuntime for RemoteWorkerRuntime {
WorkerInputKind::RegisterPeer => EmbeddedWorkerInputKind::RegisterPeer,
},
content: request.content,
submission_id: None,
submission_request_id: None,
segments: request.segments,
};
match self.post_json::<_, RuntimeHttpWorkerInputResponse>(
@@ -5195,7 +5195,7 @@ mod tests {
"missing test context",
);
};
let submission_id = input.submission_id.clone();
let submission_request_id = input.submission_request_id.clone();
let content = input.content;
std::thread::spawn(move || {
std::thread::sleep(std::time::Duration::from_millis(10));
@@ -5212,11 +5212,13 @@ mod tests {
status: protocol::WorkerStatus::Idle,
});
});
if let Some(submission_id) = submission_id {
worker_runtime::execution::WorkerExecutionResult::accepted_input_committed(
if let Some(submission_request_id) = submission_request_id {
worker_runtime::execution::WorkerExecutionResult::accepted_submission(
worker_runtime::execution::WorkerExecutionOperation::Input,
WorkerExecutionRunState::Busy,
submission_id,
submission_request_id,
uuid::Uuid::now_v7().to_string(),
protocol::SubmissionDisposition::Started,
)
} else {
worker_runtime::execution::WorkerExecutionResult::accepted(
crates/workspace-server/src/runtime_subscription_tests.rs
+5 -3
View File
@@ -32,11 +32,13 @@ impl WorkerExecutionBackend for TestExecutionBackend {
_handle: &WorkerExecutionHandle,
input: worker_runtime::interaction::WorkerInput,
) -> WorkerExecutionResult {
if let Some(submission_id) = input.submission_id {
WorkerExecutionResult::accepted_input_committed(
if let Some(submission_request_id) = input.submission_request_id {
WorkerExecutionResult::accepted_submission(
WorkerExecutionOperation::Input,
WorkerExecutionRunState::Busy,
submission_id,
submission_request_id,
uuid::Uuid::now_v7().to_string(),
protocol::SubmissionDisposition::Started,
)
} else {
WorkerExecutionResult::accepted(
crates/workspace-server/src/server.rs
+196 -33
View File
@@ -8574,13 +8574,15 @@ async fn scoped_list_runtimes(
async fn scoped_workspace_protocol_ws(
State(api): State<WorkspaceApi>,
Extension(actor): Extension<RequestActor>,
AxumPath(workspace_id): AxumPath<String>,
ws: axum::extract::ws::WebSocketUpgrade,
) -> std::result::Result<Response, Response> {
validate_workspace_scope(&api, &workspace_id).map_err(|error| error.into_response())?;
let input_source = authenticated_browser_input_source(&actor);
Ok(ws
.on_upgrade(move |socket| {
crate::workspace_subscription::serve_workspace_subscription(api, socket)
crate::workspace_subscription::serve_workspace_subscription(api, socket, input_source)
})
.into_response())
}
@@ -9286,7 +9288,7 @@ async fn scoped_capture_worker_observation_session(
return Err(ApiError::from(Error::UnknownWorker { worker: target }));
}
let mut connection = connect_workspace_worker_protocol(&api, &target).await?;
let mut connection = connect_workspace_worker_protocol(&api, &target, None).await?;
let event = tokio::time::timeout(std::time::Duration::from_secs(10), connection.events.recv())
.await
.map_err(|_| {
@@ -11495,6 +11497,7 @@ async fn scoped_cancel_runtime_worker(
async fn scoped_worker_protocol_ws(
ws: WebSocketUpgrade,
State(api): State<WorkspaceApi>,
Extension(actor): Extension<RequestActor>,
AxumPath(path): AxumPath<ScopedRuntimeWorkerPath>,
) -> Response {
if let Err(err) = validate_workspace_scope(&api, &path.workspace_id) {
@@ -11502,6 +11505,7 @@ async fn scoped_worker_protocol_ws(
}
worker_protocol_ws(
State(api),
Extension(actor),
AxumPath((path.worker.runtime_id, path.worker.worker_id)),
ws,
)
@@ -13896,8 +13900,45 @@ async fn cancel_runtime_worker(
Ok(Json(result))
}
fn authenticated_browser_input_source(actor: &RequestActor) -> protocol::AuthenticatedInputSource {
protocol::AuthenticatedInputSource::Account {
account_id: actor.account_id.clone(),
}
}
pub(crate) fn authorize_browser_worker_method(
method: protocol::Method,
source: &protocol::AuthenticatedInputSource,
) -> std::result::Result<protocol::Method, &'static str> {
match method {
protocol::Method::Submit {
submission_request_id,
input,
} => Ok(protocol::Method::SubmitTracked {
submission_request_id,
input,
source: source.clone(),
}),
protocol::Method::Notify {
notification_request_id,
message,
auto_run,
} => Ok(protocol::Method::NotifyTracked {
notification_request_id,
message,
auto_run,
source: source.clone(),
}),
protocol::Method::SubmitTracked { .. } | protocol::Method::NotifyTracked { .. } => {
Err("authenticated Worker input source is server-owned")
}
other => Ok(other),
}
}
async fn worker_protocol_ws(
State(api): State<WorkspaceApi>,
Extension(actor): Extension<RequestActor>,
AxumPath((runtime_id, worker_id)): AxumPath<(String, String)>,
ws: WebSocketUpgrade,
) -> impl IntoResponse {
@@ -13921,7 +13962,8 @@ async fn worker_protocol_ws(
.into_response();
}
};
ws.on_upgrade(move |socket| worker_protocol_ws_session(source, socket))
let input_source = authenticated_browser_input_source(&actor);
ws.on_upgrade(move |socket| worker_protocol_ws_session(source, socket, input_source))
}
pub(crate) struct WorkspaceWorkerProtocolConnection {
@@ -13932,6 +13974,7 @@ pub(crate) struct WorkspaceWorkerProtocolConnection {
pub(crate) async fn connect_workspace_worker_protocol(
api: &WorkspaceApi,
worker: &RuntimeWorkerRef,
input_source: Option<&protocol::AuthenticatedInputSource>,
) -> Result<WorkspaceWorkerProtocolConnection> {
let source = match api.observation_proxy.source(worker) {
Ok(source) => source,
@@ -13948,15 +13991,39 @@ pub(crate) async fn connect_workspace_worker_protocol(
}
};
match source {
RuntimeObservationSource::RemoteWs(config) => connect_remote_worker_protocol(config).await,
RuntimeObservationSource::RemoteWs(config) => {
connect_remote_worker_protocol(config, input_source).await
}
RuntimeObservationSource::Embedded(source) => {
connect_embedded_worker_protocol(source).await
}
}
}
fn insert_authenticated_input_source_header(
headers: &mut HeaderMap,
input_source: Option<&protocol::AuthenticatedInputSource>,
) -> Result<()> {
let Some(input_source) = input_source else {
return Ok(());
};
let protocol::AuthenticatedInputSource::Account { account_id } = input_source else {
return Err(Error::Config(
"remote Worker protocol transport supports only Account input source".into(),
));
};
headers.insert(
protocol::AUTHENTICATED_ACCOUNT_ID_HEADER,
account_id.parse().map_err(|error| {
Error::Config(format!("invalid authenticated Account identity: {error}"))
})?,
);
Ok(())
}
async fn connect_remote_worker_protocol(
config: RuntimeObservationSourceConfig,
input_source: Option<&protocol::AuthenticatedInputSource>,
) -> Result<WorkspaceWorkerProtocolConnection> {
let mut request = config
.endpoint
@@ -13971,6 +14038,7 @@ async fn connect_remote_worker_protocol(
})?,
);
}
insert_authenticated_input_source_header(request.headers_mut(), input_source)?;
let (socket, _) =
connect_async(request)
.await
@@ -14048,13 +14116,17 @@ async fn connect_embedded_worker_protocol(
Ok(WorkspaceWorkerProtocolConnection { methods, events })
}
async fn worker_protocol_ws_session(source: RuntimeObservationSource, socket: WebSocket) {
async fn worker_protocol_ws_session(
source: RuntimeObservationSource,
socket: WebSocket,
input_source: protocol::AuthenticatedInputSource,
) {
match source {
RuntimeObservationSource::RemoteWs(config) => {
remote_worker_protocol_ws_session(config, socket).await;
remote_worker_protocol_ws_session(config, socket, input_source).await;
}
RuntimeObservationSource::Embedded(source) => {
embedded_worker_protocol_ws_session(source, socket).await;
embedded_worker_protocol_ws_session(source, socket, input_source).await;
}
}
}
@@ -14062,6 +14134,7 @@ async fn worker_protocol_ws_session(source: RuntimeObservationSource, socket: We
async fn remote_worker_protocol_ws_session(
config: RuntimeObservationSourceConfig,
socket: WebSocket,
input_source: protocol::AuthenticatedInputSource,
) {
let mut request = match config.endpoint.clone().into_client_request() {
Ok(request) => request,
@@ -14089,6 +14162,16 @@ async fn remote_worker_protocol_ws_session(
}
}
}
if let Err(error) =
insert_authenticated_input_source_header(request.headers_mut(), Some(&input_source))
{
let mut socket = socket;
let event = protocol_error_event(format!(
"failed to build authenticated Account identity header: {error}"
));
let _ = send_protocol_event(&mut socket, &event).await;
return;
}
let (upstream, _) = match connect_async(request).await {
Ok(connection) => connection,
@@ -14110,14 +14193,33 @@ async fn remote_worker_protocol_ws_session(
inbound = client_stream.next() => {
match inbound {
Some(Ok(WsMessage::Text(text))) => {
if upstream_sink.send(TungsteniteMessage::Text(text.to_string().into())).await.is_err() {
let method = match protocol::stream::decode_method(text.as_ref()) {
Ok(method) => match authorize_browser_worker_method(method, &input_source) {
Ok(method) => method,
Err(message) => {
if let Ok(event) = protocol::stream::encode_event(&protocol_error_event(message)) {
let _ = client_sink.send(WsMessage::Text(event.into())).await;
}
break;
}
},
Err(error) => {
if let Ok(event) = protocol::stream::encode_event(&protocol_error_event(error.to_string())) {
let _ = client_sink.send(WsMessage::Text(event.into())).await;
}
break;
}
};
let Ok(method) = protocol::stream::encode_method(&method) else { break };
if upstream_sink.send(TungsteniteMessage::Text(method.into())).await.is_err() {
break;
}
}
Some(Ok(WsMessage::Binary(binary))) => {
if upstream_sink.send(TungsteniteMessage::Binary(binary.to_vec().into())).await.is_err() {
break;
Some(Ok(WsMessage::Binary(_))) => {
if let Ok(event) = protocol::stream::encode_event(&protocol_error_event("binary Worker methods are not accepted")) {
let _ = client_sink.send(WsMessage::Text(event.into())).await;
}
break;
}
Some(Ok(WsMessage::Close(_))) | None => {
let _ = upstream_sink.send(TungsteniteMessage::Close(None)).await;
@@ -14173,6 +14275,7 @@ async fn remote_worker_protocol_ws_session(
async fn embedded_worker_protocol_ws_session(
source: crate::observation::EmbeddedRuntimeObservationSource,
mut socket: WebSocket,
input_source: protocol::AuthenticatedInputSource,
) {
let mut upstream = match RuntimeObservationClient::connect(&RuntimeObservationSource::Embedded(
source.clone(),
@@ -14192,24 +14295,32 @@ async fn embedded_worker_protocol_ws_session(
inbound = socket.next() => {
match inbound {
Some(Ok(WsMessage::Text(text))) => match decode_method(&text) {
Ok(method) => match source.runtime.send_protocol_method(&source.worker_ref, method) {
Ok(events) => {
for event in events {
Ok(method) => match authorize_browser_worker_method(method, &input_source) {
Ok(method) => match source.runtime.send_protocol_method(&source.worker_ref, method) {
Ok(events) => {
for event in events {
if !send_protocol_event(&mut socket, &event).await {
return;
}
}
}
Err(error) => {
let event = protocol_error_event(error.to_string());
if !send_protocol_event(&mut socket, &event).await {
return;
}
}
}
Err(error) => {
let event = protocol_error_event(error.to_string());
if !send_protocol_event(&mut socket, &event).await {
return;
}
},
Err(message) => {
let event = protocol_error_event(message);
let _ = send_protocol_event(&mut socket, &event).await;
return;
}
},
Err(error) => {
let event =
protocol_error_event(format!("malformed protocol method frame: {error}"));
let event = protocol_error_event(format!(
"malformed protocol method frame: {error}"
));
if !send_protocol_event(&mut socket, &event).await {
return;
}
@@ -16219,6 +16330,48 @@ mod tests {
&tail[..end]
}
#[test]
fn browser_worker_methods_receive_server_owned_account_source() {
let source = protocol::AuthenticatedInputSource::Account {
account_id: "account-1".into(),
};
let method = authorize_browser_worker_method(
protocol::Method::Submit {
submission_request_id: "request-1".into(),
input: vec![protocol::Segment::text("hello")],
},
&source,
)
.unwrap();
assert!(matches!(
method,
protocol::Method::SubmitTracked {
source: protocol::AuthenticatedInputSource::Account { ref account_id },
..
} if account_id == "account-1"
));
assert!(authorize_browser_worker_method(method, &source).is_err());
}
#[test]
fn remote_worker_protocol_header_preserves_authenticated_account_source() {
let mut headers = HeaderMap::new();
insert_authenticated_input_source_header(
&mut headers,
Some(&protocol::AuthenticatedInputSource::Account {
account_id: "account-1".into(),
}),
)
.unwrap();
assert_eq!(
headers
.get(protocol::AUTHENTICATED_ACCOUNT_ID_HEADER)
.unwrap(),
"account-1"
);
}
#[test]
fn merge_request_http_paths_observe_refs_through_runtime_provider_authority() {
let source = include_str!("server.rs");
@@ -18468,7 +18621,7 @@ mod tests {
.get(handle.worker_ref())
.cloned()
.expect("execution context");
let submission_id = input.submission_id.clone();
let submission_request_id = input.submission_request_id.clone();
let content = input.content.clone();
std::thread::spawn(move || {
std::thread::sleep(std::time::Duration::from_millis(25));
@@ -18476,11 +18629,13 @@ mod tests {
text: format!("server companion echoed: {content}"),
});
});
if let Some(submission_id) = submission_id {
worker_runtime::execution::WorkerExecutionResult::accepted_input_committed(
if let Some(submission_request_id) = submission_request_id {
worker_runtime::execution::WorkerExecutionResult::accepted_submission(
worker_runtime::execution::WorkerExecutionOperation::Input,
worker_runtime::execution::WorkerExecutionRunState::Idle,
submission_id,
submission_request_id,
uuid::Uuid::now_v7().to_string(),
protocol::SubmissionDisposition::Started,
)
} else {
worker_runtime::execution::WorkerExecutionResult::accepted(
@@ -27069,6 +27224,16 @@ mod tests {
(runtime, worker_ref, endpoint)
}
fn test_browser_request_actor() -> RequestActor {
RequestActor {
user_id: "test-user".into(),
account_id: format!("account-{TEST_WORKSPACE_ID}"),
handle: "test".into(),
display_name: "Test".into(),
auth_method: ActorAuthMethod::BrowserSession,
}
}
async fn spawn_workspace_proxy(
source: RuntimeObservationSourceConfig,
) -> (String, tempfile::TempDir) {
@@ -27087,11 +27252,8 @@ mod tests {
.unwrap();
let app_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let app_addr = app_listener.local_addr().unwrap();
tokio::spawn(async move {
axum::serve(app_listener, build_inner_router(api))
.await
.unwrap()
});
let app = build_inner_router(api).layer(Extension(test_browser_request_actor()));
tokio::spawn(async move { axum::serve(app_listener, app).await.unwrap() });
(
format!("ws://{app_addr}/api/runtimes/{runtime_id}/workers/{worker_id}/protocol/ws"),
dir,
@@ -27103,7 +27265,8 @@ mod tests {
let dir = tempfile::tempdir().unwrap();
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let app = build_inner_router(test_api(dir.path()).await);
let app = build_inner_router(test_api(dir.path()).await)
.layer(Extension(test_browser_request_actor()));
let server = tokio::spawn(async move {
let _ = axum::serve(listener, app).await;
});
@@ -27151,7 +27314,7 @@ mod tests {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let app = build_inner_router(api);
let app = build_inner_router(api).layer(Extension(test_browser_request_actor()));
let server = tokio::spawn(async move {
let _ = axum::serve(listener, app).await;
});
crates/workspace-server/src/workspace_subscription.rs
+21 -4
View File
@@ -11,7 +11,9 @@ use tokio::sync::mpsc;
use worker_runtime::identity::RuntimeWorkerRef;
use crate::runtime_subscription::{BrokerSubscriptionEvent, RuntimeSubscriptionBroker};
use crate::server::{WorkspaceApi, connect_workspace_worker_protocol};
use crate::server::{
WorkspaceApi, authorize_browser_worker_method, connect_workspace_worker_protocol,
};
use crate::store::WorkspaceResourceKind;
const OUTBOUND_CAPACITY: usize = 256;
@@ -21,7 +23,11 @@ struct ActiveSubscription {
methods: Option<mpsc::Sender<protocol::Method>>,
}
pub(crate) async fn serve_workspace_subscription(api: WorkspaceApi, socket: WebSocket) {
pub(crate) async fn serve_workspace_subscription(
api: WorkspaceApi,
socket: WebSocket,
input_source: protocol::AuthenticatedInputSource,
) {
let broker = api.runtime_subscription_broker().clone();
let (mut socket_sender, mut socket_receiver) = socket.split();
let (control_outbound, mut control_receiver) = mpsc::channel::<WsMessage>(OUTBOUND_CAPACITY);
@@ -85,7 +91,13 @@ pub(crate) async fn serve_workspace_subscription(api: WorkspaceApi, socket: WebS
runtime_id: Some(runtime_id),
} => {
let worker = RuntimeWorkerRef::new(&runtime_id, worker_id.as_str());
match connect_workspace_worker_protocol(&api, &worker).await {
match connect_workspace_worker_protocol(
&api,
&worker,
Some(&input_source),
)
.await
{
Ok(connection) => {
let methods = connection.methods.clone();
let task = tokio::spawn(run_worker_protocol(
@@ -153,7 +165,12 @@ pub(crate) async fn serve_workspace_subscription(api: WorkspaceApi, socket: WebS
else {
break;
};
if methods.send(message.method).await.is_err() {
let Ok(method) =
authorize_browser_worker_method(message.method, &input_source)
else {
break;
};
if methods.send(method).await.is_err() {
break;
}
}