From 21df210a7fd580cd2d14a2dd92b2479579fa2288 Mon Sep 17 00:00:00 2001 From: Hare Date: Mon, 12 Oct 2026 02:56:42 +0900 Subject: [PATCH] fix: bound Worker protocol backpressure and connection ownership --- crates/worker-runtime/src/worker_backend.rs | 68 +- crates/workspace-server/src/server.rs | 161 ++- .../tests/protocol_backpressure_tests.rs | 920 ++++++++++++++++++ .../src/server/tests/worker_protocol_tests.rs | 210 ++++ .../workspace_protocol_backpressure_tests.rs | 140 +++ .../src/server/worker_operations.rs | 48 +- .../src/server/worker_protocol.rs | 247 +++++ .../src/workspace_subscription.rs | 392 +++++--- .../worker-protocol-backpressure.md | 112 +++ 9 files changed, 2036 insertions(+), 262 deletions(-) create mode 100644 crates/workspace-server/src/server/tests/protocol_backpressure_tests.rs create mode 100644 crates/workspace-server/src/server/tests/worker_protocol_tests.rs create mode 100644 crates/workspace-server/src/server/tests/workspace_protocol_backpressure_tests.rs create mode 100644 crates/workspace-server/src/server/worker_protocol.rs create mode 100644 docs/development/worker-protocol-backpressure.md diff --git a/crates/worker-runtime/src/worker_backend.rs b/crates/worker-runtime/src/worker_backend.rs index 63a4e8163..436131647 100644 --- a/crates/worker-runtime/src/worker_backend.rs +++ b/crates/worker-runtime/src/worker_backend.rs @@ -2537,16 +2537,24 @@ where } } method => { - let result = self.send_method( - WorkerExecutionOperation::ProtocolMethod, - execution.handle.clone(), - method, - ); - if result.is_accepted() { + // The connected protocol owns this bounded admission wait. The old + // generic adapter waiter could time out while leaving a detached send + // task alive; it might deliver later after the connection had ended. + // Join cancellation before releasing endpoint exclusion. A send which + // raced cancellation remains delivery-unknown, never automatically retried. + let handle = execution.handle.clone(); + self.run_cancellable_on_adapter_runtime(USER_INPUT_TASK_TIMEOUT, async move { + handle.send(method).await.map_err(|error| { + format!("connected Worker method delivery failed or is unknown: {error}") + })?; Ok(Vec::new()) - } else { - Err(result) - } + }) + .map_err(|message| { + WorkerExecutionResult::errored( + WorkerExecutionOperation::ProtocolMethod, + message, + ) + }) } } } @@ -7898,6 +7906,48 @@ mod tests { assert!(backend.worker_snapshot(&worker_ref).is_none()); } + #[test] + fn connected_protocol_admission_timeout_joins_pending_send_before_capacity_reopens() { + // This is the adapter's cancellation/join policy boundary used by the + // captured protocol endpoint, not an execution-side-effect rollback. + let backend = WorkerRuntimeExecutionBackend::new(FailingFactory).unwrap(); + let (sender, mut receiver) = tokio::sync::mpsc::channel(1); + sender.try_send(Method::ListPendingSubmissions).unwrap(); + let dropped = Arc::new(AtomicBool::new(false)); + struct PendingSendDrop(Arc); + impl Drop for PendingSendDrop { + fn drop(&mut self) { + self.0.store(true, Ordering::SeqCst); + } + } + let owned = PendingSendDrop(dropped.clone()); + let error = backend + .run_cancellable_on_adapter_runtime(Duration::from_millis(5), async move { + let _owned = owned; + sender + .send(Method::ListRewindTargets) + .await + .map_err(|error| error.to_string()) + }) + .unwrap_err(); + assert!(error.contains("was cancelled"), "{error}"); + assert!( + dropped.load(Ordering::SeqCst), + "must join task cancellation before returning" + ); + assert!(matches!( + receiver.try_recv(), + Ok(Method::ListPendingSubmissions) + )); + assert!( + matches!( + receiver.try_recv(), + Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) + ), + "timeout must not leave an admission task capable of a late send" + ); + } + #[test] fn absent_execution_has_typed_input_rejection_and_safe_stop() { let backend = WorkerRuntimeExecutionBackend::new(FailingFactory).unwrap(); diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index d391d9f33..cb8094c83 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -2,9 +2,12 @@ use subjektiv::api::*; mod backend_job_profile; mod subjektiv_job; mod worker_operations; +mod worker_protocol; pub(crate) use worker_operations::{ WorkerOperationContext, WorkspaceWorker, WorkspaceWorkerMethodSender, }; +pub(crate) use worker_protocol::WorkerProtocolEnd; +use worker_protocol::*; #[path = "server_drive.rs"] mod drive; @@ -29083,51 +29086,19 @@ async fn worker_protocol_ws( match worker.connect_protocol(&context).await { Ok(connection) => workspace_worker_protocol_session(connection, socket, context).await, Err(error) => { - let _ = send_protocol_event(&mut socket, &protocol_error_event(error.to_string())) - .await; - let _ = socket.send(WsMessage::Close(None)).await; + let _ = tokio::time::timeout(WORKER_PROTOCOL_TRANSPORT_TIMEOUT, async { + let _ = + send_protocol_event(&mut socket, &protocol_error_event(error.to_string())) + .await; + let _ = socket.send(WsMessage::Close(None)).await; + }) + .await; } } }) .into_response() } -async fn workspace_worker_protocol_session( - mut connection: worker_operations::WorkspaceWorkerConnection, - mut socket: WebSocket, - context: WorkerOperationContext, -) { - loop { - tokio::select! { - inbound = socket.next() => match inbound { - Some(Ok(WsMessage::Text(text))) => { - let result = match decode_method(&text) { - Ok(method) => connection.sender.send(&context, method).await.map_err(|error| error.to_string()), - Err(error) => Err(error.to_string()), - }; - if let Err(error) = result { - if !send_protocol_event(&mut socket, &protocol_error_event(error)).await { break; } - } - } - Some(Ok(WsMessage::Ping(value))) => { if socket.send(WsMessage::Pong(value)).await.is_err() { break; } } - Some(Ok(WsMessage::Pong(_))) => {} - _ => break, - }, - event = connection.events.recv() => { - let Some(event) = event else { break; }; - let shutdown = matches!(event, protocol::Event::Shutdown); - if !send_protocol_event(&mut socket, &event).await || shutdown { break; } - } - } - } - let _ = socket.send(WsMessage::Close(None)).await; -} - -pub(crate) struct WorkspaceWorkerProtocolConnection { - pub(crate) methods: tokio::sync::mpsc::Sender, - pub(crate) events: tokio::sync::mpsc::Receiver, -} - pub(crate) async fn connect_workspace_worker_protocol( api: &WorkspaceApi, worker: &RuntimeWorkerRef, @@ -29199,39 +29170,21 @@ async fn connect_remote_worker_protocol( } insert_authenticated_input_source_header(request.headers_mut(), input_source)?; let (socket, _) = - connect_async(request) + tokio::time::timeout(WORKER_PROTOCOL_TRANSPORT_TIMEOUT, connect_async(request)) .await + .map_err(|_| Error::RuntimeOperationFailed { + runtime_id: config.worker.runtime_id.clone(), + code: "worker_protocol_connect_timeout".into(), + message: "Worker protocol connection handshake timed out; no method was admitted" + .into(), + })? .map_err(|error| Error::RuntimeOperationFailed { runtime_id: config.worker.runtime_id.clone(), - code: "worker_protocol_connect_failed".to_string(), + code: "worker_protocol_connect_failed".into(), 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 }) + let (sink, stream) = socket.split(); + Ok(remote_worker_protocol_bridge(sink, stream)) } async fn connect_embedded_worker_protocol( @@ -29249,47 +29202,48 @@ async fn connect_embedded_worker_protocol( code: "embedded_worker_protocol_connect_failed".to_string(), message: error.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 { - if event_sender.send(transport.snapshot.clone()).await.is_err() { - return; - } - loop { - tokio::select! { - _ = event_sender.closed() => break, - method = method_receiver.recv() => { - let Some(method) = method else { break }; - let shutdown = matches!(&method, protocol::Method::Shutdown { .. }); - match source.runtime.send_connected_protocol_method_scoped( - &scope, &source.worker_ref, &transport, method, - ) { - Ok(direct_events) => { - for event in direct_events { - let terminal = matches!(event, protocol::Event::Shutdown); - if event_sender.send(event).await.is_err() || terminal { return; } + Ok(spawn_worker_protocol_pump( + move |mut methods, events| async move { + if let Err(reason) = enqueue_worker_protocol_event(&events, transport.snapshot.clone()) + { + return reason; + } + loop { + tokio::select! { + method = methods.recv() => { + let Some(method) = method else { return WorkerProtocolEnd::Disconnected; }; + let shutdown = matches!(&method, protocol::Method::Shutdown { .. }); + match source.runtime.send_connected_protocol_method_scoped( + &scope, &source.worker_ref, &transport, method, + ) { + Ok(direct_events) => { + for event in direct_events { + let terminal = matches!(event, protocol::Event::Shutdown); + if let Err(reason) = enqueue_worker_protocol_event(&events, event) { return reason; } + if terminal { return WorkerProtocolEnd::Disconnected; } + } + } + Err(error) => { + let terminal = embedded_protocol_dispatch_is_terminal(&error, &transport); + if let Err(reason) = enqueue_worker_protocol_event(&events, protocol_error_event(error.to_string())) { return reason; } + if terminal { return WorkerProtocolEnd::Disconnected; } } } - Err(error) => { - let terminal = embedded_protocol_dispatch_is_terminal(&error, &transport); - if event_sender.send(protocol_error_event(error.to_string())).await.is_err() || terminal { return; } + if shutdown { return WorkerProtocolEnd::Disconnected; } + } + event = transport.events.recv() => match event { + Some(event) => { + let terminal = matches!(event, protocol::Event::Shutdown); + if let Err(reason) = enqueue_worker_protocol_event(&events, event) { return reason; } + if terminal { return WorkerProtocolEnd::Disconnected; } } + // Stop, disconnect and upstream lag never cause an implicit reattach. + None => return WorkerProtocolEnd::Disconnected, } - if shutdown { break; } - } - event = transport.events.recv() => match event { - Some(event) => { - let terminal = matches!(event, protocol::Event::Shutdown); - if event_sender.send(event).await.is_err() || terminal { break; } - } - // The execution relay closes on Stop, disconnect or lag. - // Uncertain streams are terminal, never grounds to reattach. - None => break, } } - } - }); - Ok(WorkspaceWorkerProtocolConnection { methods, events }) + }, + )) } fn embedded_protocol_dispatch_is_terminal( @@ -32502,6 +32456,7 @@ mod tests { mod drive_tests; mod drive_web_roundtrip_tests; mod external_checkout_tests; + mod protocol_backpressure_tests; mod subject_spawn_tests; mod subjektiv_jobs_tests; mod ticket_evidence_tests; @@ -32511,7 +32466,9 @@ mod tests { mod workdir_delegation_tests; mod worker_drive_tests; mod worker_operations_tests; + mod worker_protocol_tests; mod workspace_admission_tests; + mod workspace_protocol_backpressure_tests; include!("server_workspace_config_tests.rs"); use super::*; use axum::body::{Body, to_bytes}; diff --git a/crates/workspace-server/src/server/tests/protocol_backpressure_tests.rs b/crates/workspace-server/src/server/tests/protocol_backpressure_tests.rs new file mode 100644 index 000000000..7a4066b39 --- /dev/null +++ b/crates/workspace-server/src/server/tests/protocol_backpressure_tests.rs @@ -0,0 +1,920 @@ +use super::*; + +const PROTOCOL_TEST_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(2); + +// Completion replies are produced synchronously by the actual attached execution. +// Unlike flooding the fixture's 32-entry broadcast bus, this reaches the Server's +// event admission boundary without first causing unrelated Runtime relay lag. +#[tokio::test] +async fn embedded_event_overflow_terminates_without_losing_the_admitted_prefix() { + let (runtime, worker_ref, execution) = runtime_with_worker_backend(); + let mut connection = connect_embedded_worker_protocol( + embedded_protocol_test_source(&runtime, &worker_ref), + embedded_protocol_test_scope(), + ) + .await + .unwrap(); + + // Leave the snapshot queued. Exactly capacity - 1 completion replies fit; + // the next reply must end this connection rather than block its bridge. + tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, async { + for index in 0..WORKER_PROTOCOL_EVENT_CAPACITY { + connection + .methods + .send(protocol_test_completion(&format!("overflow-{index}"))) + .await + .unwrap(); + } + assert_eq!( + connection.termination.ended().await, + WorkerProtocolEnd::EventOverflow + ); + connection.methods.closed().await; + }) + .await + .expect("full embedded event queue did not terminate the connection"); + + assert_embedded_protocol_snapshot( + next_embedded_protocol_event(&mut connection).await, + &worker_ref, + ); + for index in 0..WORKER_PROTOCOL_EVENT_CAPACITY - 1 { + assert!(matches!( + next_embedded_protocol_event(&mut connection).await, + protocol::Event::Completions { request_id: Some(id), .. } + if id == format!("overflow-{index}") + )); + } + assert!( + tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, connection.events.recv()) + .await + .expect("overflowed event receiver did not close") + .is_none() + ); + assert!( + connection + .methods + .try_send(protocol_test_completion("after-overflow")) + .is_err() + ); + assert_eq!( + execution.protocol_methods().len(), + WORKER_PROTOCOL_EVENT_CAPACITY + ); +} + +#[tokio::test] +async fn embedded_protocol_delivers_snapshot_then_ordered_replies_and_live_events() { + let (runtime, worker_ref, execution) = runtime_with_worker_backend(); + let mut connection = connect_embedded_worker_protocol( + embedded_protocol_test_source(&runtime, &worker_ref), + embedded_protocol_test_scope(), + ) + .await + .unwrap(); + for id in ["first", "second", "third"] { + connection + .methods + .try_send(protocol_test_completion(id)) + .unwrap(); + } + assert_embedded_protocol_snapshot( + next_embedded_protocol_event(&mut connection).await, + &worker_ref, + ); + for expected in ["first", "second", "third"] { + assert!(matches!( + next_embedded_protocol_event(&mut connection).await, + protocol::Event::Completions { request_id: Some(id), .. } if id == expected + )); + } + execution.emit_worker_event( + &worker_ref, + protocol::Event::TextDelta { + text: "live-after-replies".into(), + }, + ); + assert!(matches!( + next_embedded_protocol_event(&mut connection).await, + protocol::Event::TextDelta { text } if text == "live-after-replies" + )); + assert_eq!(connection.termination.reason(), None); +} + +#[tokio::test] +async fn embedded_overflow_requires_explicit_reconnect_and_stop_ends_the_new_connection() { + let (runtime, worker_ref, execution) = runtime_with_worker_backend(); + let source = embedded_protocol_test_source(&runtime, &worker_ref); + let mut old = connect_embedded_worker_protocol(source.clone(), embedded_protocol_test_scope()) + .await + .unwrap(); + assert_embedded_protocol_snapshot(next_embedded_protocol_event(&mut old).await, &worker_ref); + tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, async { + for index in 0..=WORKER_PROTOCOL_EVENT_CAPACITY { + old.methods + .send(protocol_test_completion(&format!("old-{index}"))) + .await + .unwrap(); + } + assert_eq!( + old.termination.ended().await, + WorkerProtocolEnd::EventOverflow + ); + old.methods.closed().await; + }) + .await + .expect("overflowed connection remained live"); + + // Overflow ends only the attached connection, not its Worker. Reconnection + // is explicit and starts from a fresh execution snapshot, not old backlog. + let mut fresh = connect_embedded_worker_protocol(source, embedded_protocol_test_scope()) + .await + .unwrap(); + assert_embedded_protocol_snapshot(next_embedded_protocol_event(&mut fresh).await, &worker_ref); + fresh + .methods + .try_send(protocol_test_completion("explicit-reconnect")) + .unwrap(); + assert!(matches!( + next_embedded_protocol_event(&mut fresh).await, + protocol::Event::Completions { request_id: Some(id), .. } if id == "explicit-reconnect" + )); + assert!( + old.methods + .try_send(protocol_test_completion("must-not-reattach")) + .is_err() + ); + assert_eq!( + old.termination.reason(), + Some(WorkerProtocolEnd::EventOverflow) + ); + assert_eq!( + execution.protocol_methods().len(), + WORKER_PROTOCOL_EVENT_CAPACITY + 2 + ); + + runtime.stop_worker(&worker_ref, None).unwrap(); + assert_eq!( + tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, fresh.termination.ended()) + .await + .expect("Stop did not end the attached protocol connection"), + WorkerProtocolEnd::Disconnected + ); + assert_embedded_protocol_closed(&mut fresh).await; +} + +async fn backpressure_browser_context(api: &WorkspaceApi) -> (WorkerOperationContext, String) { + seed_test_api_token(api.store.as_ref(), "backpressure-browser"); + let token = "backpressure-browser-session"; + let token_hash = crate::auth::token_hash(token); + api.store + .create_browser_session(&BrowserSessionRecord { + token_hash: token_hash.clone(), + session_id: "backpressure-session".into(), + user_id: "user-backpressure-browser".into(), + created_at: "2026-01-01T00:00:00Z".into(), + expires_at: "2099-01-01T00:00:00Z".into(), + revoked_at: None, + }) + .unwrap(); + let AuthConfig::Passkey { + cookie_name, + origin, + .. + } = &api.config.auth; + let mut headers = HeaderMap::new(); + headers.insert(ORIGIN, origin.parse().unwrap()); + headers.insert( + axum::http::header::COOKIE, + format!("{cookie_name}={token}").parse().unwrap(), + ); + let actor = resolve_actor(&ServerAuthApi::from(api), &headers) + .await + .unwrap() + .expect("fixture browser session must authenticate"); + ( + WorkerOperationContext::Browser { + headers, + source: authenticated_browser_input_source(&actor), + }, + token_hash, + ) +} + +async fn undrained_protocol_connection() -> ( + WorkspaceWorkerProtocolConnection, + tokio::sync::mpsc::Receiver, +) { + let (receiver_sender, receiver) = tokio::sync::oneshot::channel(); + let connection = spawn_worker_protocol_pump(move |methods, events| async move { + receiver_sender.send(methods).unwrap(); + // Keep the event endpoint live until common termination cancels this + // fixture pump. Method consumption is controlled solely by the test. + let _events = events; + std::future::pending::().await + }); + let receiver = tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, receiver) + .await + .expect("fixture pump did not expose its undrained method endpoint") + .unwrap(); + (connection, receiver) +} + +async fn full_protocol_method_queue() -> ( + WorkspaceWorkerProtocolConnection, + tokio::sync::mpsc::Receiver, +) { + let (connection, receiver) = undrained_protocol_connection().await; + for index in 0..WORKER_PROTOCOL_METHOD_CAPACITY { + connection + .methods + .try_send(protocol_test_completion(&format!("queued-{index}"))) + .unwrap(); + } + (connection, receiver) +} + +#[tokio::test] +async fn browser_method_queue_overflow_fails_fast_and_does_not_pin_session_cleanup() { + let workspace = tempfile::tempdir().unwrap(); + let (api, execution) = test_api_with_recording_backend(workspace.path()).await; + let identity = spawn_ticket_check_source(&api, "method-overflow"); + let worker = WorkspaceWorker::resolve(&api, &identity.runtime_id, &identity.worker_id).unwrap(); + let (context, _) = backpressure_browser_context(&api).await; + let (connection, mut receiver) = full_protocol_method_queue().await; + let termination = connection.termination.clone(); + let sender = + WorkspaceWorkerMethodSender::new(worker.clone(), connection.methods, termination.clone()); + let _events = connection.events; + + let result = tokio::time::timeout( + PROTOCOL_TEST_TIMEOUT, + sender.send( + &context, + protocol::Method::Cancel { + command: protocol::WorkerCommandEnvelope::new(801), + }, + ), + ) + .await + .expect("method admission blocked while holding the session guard"); + assert!(result.is_err(), "full method queue must reject Cancel too"); + assert_eq!( + termination.reason(), + Some(WorkerProtocolEnd::MethodOverflow) + ); + assert_eq!(receiver.len(), WORKER_PROTOCOL_METHOD_CAPACITY); + assert!(matches!( + receiver.try_recv().unwrap(), + protocol::Method::ListCompletions { request_id: Some(id), .. } if id == "queued-0" + )); + assert!( + sender + .send( + &context, + protocol_test_completion("must-not-retry-after-drain") + ) + .await + .is_err() + ); + assert_eq!(receiver.len(), WORKER_PROTOCOL_METHOD_CAPACITY - 1); + + let stopped = tokio::time::timeout( + PROTOCOL_TEST_TIMEOUT, + worker.stop( + &WorkerOperationContext::Backend, + WorkerLifecycleRequest { + reason: Some("cleanup-after-method-overflow".into()), + ticket_assignment: None, + }, + ), + ) + .await + .expect("full method queue pinned Stop's session cleanup") + .unwrap(); + assert_eq!(stopped.state, InternalWorkerOperationState::Accepted); + assert!(execution.protocol_methods().is_empty()); + for index in 1..WORKER_PROTOCOL_METHOD_CAPACITY { + assert!(matches!( + receiver.try_recv().unwrap(), + protocol::Method::ListCompletions { request_id: Some(id), .. } + if id == format!("queued-{index}") + )); + } + assert!(receiver.try_recv().is_err()); +} + +#[tokio::test] +async fn revoked_browser_credentials_reject_before_method_overflow_admission() { + let workspace = tempfile::tempdir().unwrap(); + let (api, execution) = test_api_with_recording_backend(workspace.path()).await; + let identity = spawn_ticket_check_source(&api, "revoked-full-queue"); + let worker = WorkspaceWorker::resolve(&api, &identity.runtime_id, &identity.worker_id).unwrap(); + let (context, token_hash) = backpressure_browser_context(&api).await; + let (connection, receiver) = full_protocol_method_queue().await; + let termination = connection.termination.clone(); + let sender = WorkspaceWorkerMethodSender::new(worker, connection.methods, termination.clone()); + let _events = connection.events; + assert!( + api.store + .revoke_browser_session(&token_hash, &now_registry_timestamp()) + .unwrap() + ); + + let result = tokio::time::timeout( + PROTOCOL_TEST_TIMEOUT, + sender.send(&context, protocol_test_completion("unauthorized")), + ) + .await + .expect("revoked credentials waited for method capacity"); + assert!(matches!(result, Err(Error::WorkspacePermissionDenied(_)))); + assert_eq!(termination.reason(), None); + assert_eq!(receiver.len(), WORKER_PROTOCOL_METHOD_CAPACITY); + assert!(execution.protocol_methods().is_empty()); +} + +#[tokio::test] +async fn cancelling_pending_browser_admission_does_not_enqueue_or_hold_the_session_guard() { + let workspace = tempfile::tempdir().unwrap(); + let (api, execution) = test_api_with_recording_backend(workspace.path()).await; + let identity = spawn_ticket_check_source(&api, "cancel-pending-admission"); + let worker = WorkspaceWorker::resolve(&api, &identity.runtime_id, &identity.worker_id).unwrap(); + let (context, _) = backpressure_browser_context(&api).await; + let (connection, mut receiver) = undrained_protocol_connection().await; + let termination = connection.termination.clone(); + let sender = WorkspaceWorkerMethodSender::new(worker, connection.methods, termination.clone()); + let _events = connection.events; + let session_lock = current_worker_session_lock(&api, &identity); + let guard = session_lock.lock().await; + { + let admission = sender.send(&context, protocol_test_completion("cancelled-admission")); + tokio::pin!(admission); + assert!(futures::poll!(admission.as_mut()).is_pending()); + // Dropping a blocked admission future must not schedule a detached retry. + } + drop(guard); + assert!(receiver.try_recv().is_err()); + let guard = tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, session_lock.lock()) + .await + .expect("cancelled admission retained the session guard"); + close_current_worker_session_locked(&api, &identity) + .await + .unwrap(); + drop(guard); + + let command = protocol::WorkerCommandEnvelope::new(802); + sender + .send(&context, protocol::Method::Cancel { command }) + .await + .unwrap(); + assert!(matches!( + receiver.try_recv().unwrap(), + protocol::Method::Cancel { command: admitted } if admitted == command + )); + assert!(receiver.try_recv().is_err()); + assert_eq!(termination.reason(), None); + assert!(execution.protocol_methods().is_empty()); +} + +// In-process browser transport: one real JSON method frame, then an idle +// connection. Its writes are immediately ready, isolating authorization waits +// from transport congestion and without crossing a complete product flow. +struct AdmissionCancellationSocket { + inbound: Option, + written: Vec, +} + +impl futures::Stream for AdmissionCancellationSocket { + type Item = std::result::Result; + + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + match self.inbound.take() { + Some(message) => std::task::Poll::Ready(Some(Ok(message))), + None => std::task::Poll::Pending, + } + } +} + +impl futures::Sink for AdmissionCancellationSocket { + type Error = std::convert::Infallible; + + fn poll_ready( + self: std::pin::Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + std::task::Poll::Ready(Ok(())) + } + + fn start_send( + mut self: std::pin::Pin<&mut Self>, + message: WsMessage, + ) -> std::result::Result<(), Self::Error> { + self.written.push(message); + Ok(()) + } + + fn poll_flush( + self: std::pin::Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + std::task::Poll::Ready(Ok(())) + } + + fn poll_close( + self: std::pin::Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + std::task::Poll::Ready(Ok(())) + } +} + +#[tokio::test] +async fn browser_protocol_termination_cancels_admission_waiting_for_session_authority() { + let workspace = tempfile::tempdir().unwrap(); + let (api, execution) = test_api_with_recording_backend(workspace.path()).await; + let identity = spawn_ticket_check_source(&api, "browser-authorization-cancellation"); + let worker = WorkspaceWorker::resolve(&api, &identity.runtime_id, &identity.worker_id).unwrap(); + let (context, _) = backpressure_browser_context(&api).await; + let (connection, mut receiver) = undrained_protocol_connection().await; + let termination = connection.termination.clone(); + let sender = WorkspaceWorkerMethodSender::new(worker, connection.methods, termination.clone()); + let connection = worker_operations::WorkspaceWorkerConnection { + sender, + events: connection.events, + }; + let session_lock = current_worker_session_lock(&api, &identity); + let guard = session_lock.lock().await; + let method = protocol::Method::Cancel { + command: protocol::WorkerCommandEnvelope::new(803), + }; + let mut socket = AdmissionCancellationSocket { + inbound: Some(WsMessage::Text( + protocol::stream::encode_method(&method).unwrap().into(), + )), + written: Vec::new(), + }; + { + let session = workspace_worker_protocol_session(connection, &mut socket, context); + tokio::pin!(session); + // SQLite browser authorization is synchronous; this poll consumes the + // frame and reaches the deliberately held session lock. + assert!(futures::poll!(session.as_mut()).is_pending()); + termination.terminate(WorkerProtocolEnd::Disconnected); + tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, session) + .await + .expect("terminal connection did not cancel pending browser admission"); + } + assert!(matches!(socket.written.last(), Some(WsMessage::Close(_)))); + assert!(receiver.try_recv().is_err()); + drop(guard); + let _guard = tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, session_lock.lock()) + .await + .expect("cancelled browser authorization pinned session cleanup"); + assert!(execution.protocol_methods().is_empty()); +} + +struct BackpressureTestServer(tokio::task::JoinHandle<()>); + +impl Drop for BackpressureTestServer { + fn drop(&mut self) { + self.0.abort(); + } +} + +#[tokio::test] +async fn browser_protocol_session_ends_on_full_method_queue_without_blocking_cleanup() { + let workspace = tempfile::tempdir().unwrap(); + let (api, execution) = test_api_with_recording_backend(workspace.path()).await; + let identity = spawn_ticket_check_source(&api, "browser-session-overflow"); + let worker = WorkspaceWorker::resolve(&api, &identity.runtime_id, &identity.worker_id).unwrap(); + let (context, _) = backpressure_browser_context(&api).await; + let (connection, receiver) = full_protocol_method_queue().await; + let termination = connection.termination.clone(); + let sender = WorkspaceWorkerMethodSender::new(worker, connection.methods, termination.clone()); + let events = connection.events; + let (session_done, finished) = tokio::sync::oneshot::channel(); + let pending = Arc::new(std::sync::Mutex::new(Some(( + worker_operations::WorkspaceWorkerConnection { sender, events }, + session_done, + )))); + let app = Router::new().route( + "/protocol/ws", + get(move |ws: WebSocketUpgrade| { + let pending = pending.clone(); + let context = context.clone(); + async move { + ws.on_upgrade(move |socket| async move { + let (connection, session_done) = pending.lock().unwrap().take().unwrap(); + workspace_worker_protocol_session(connection, socket, context).await; + let _ = session_done.send(()); + }) + } + }), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let _server = BackpressureTestServer(tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + })); + let (mut socket, _) = connect_async(format!("ws://{address}/protocol/ws")) + .await + .unwrap(); + tokio::time::timeout( + PROTOCOL_TEST_TIMEOUT, + send_protocol_test_ws_method(&mut socket, &protocol_test_completion("browser-overflow")), + ) + .await + .expect("browser method frame did not reach the loopback transport"); + assert_eq!( + tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, termination.ended()) + .await + .expect("browser session did not report method overflow"), + WorkerProtocolEnd::MethodOverflow + ); + tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, finished) + .await + .expect("browser protocol session did not finish after overflow") + .expect("browser protocol session dropped its completion signal"); + let session_lock = current_worker_session_lock(&api, &identity); + let _guard = tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, session_lock.lock()) + .await + .expect("browser method overflow pinned session cleanup"); + close_current_worker_session_locked(&api, &identity) + .await + .unwrap(); + assert_eq!(receiver.len(), WORKER_PROTOCOL_METHOD_CAPACITY); + assert!(execution.protocol_methods().is_empty()); +} + +// Accepting an output frame is not proof of delivery: this browser transport +// never finishes flushing, including the final best-effort Close frame. +struct StalledBrowserOutput { + entered: Option>, + dropped: Option>, + written: Arc>>, +} + +impl futures::Stream for StalledBrowserOutput { + type Item = std::result::Result; + + fn poll_next( + self: std::pin::Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + std::task::Poll::Pending + } +} + +impl futures::Sink for StalledBrowserOutput { + type Error = std::convert::Infallible; + + fn poll_ready( + self: std::pin::Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + std::task::Poll::Ready(Ok(())) + } + + fn start_send( + mut self: std::pin::Pin<&mut Self>, + message: WsMessage, + ) -> std::result::Result<(), Self::Error> { + self.written.lock().unwrap().push(message); + if let Some(entered) = self.entered.take() { + let _ = entered.send(()); + } + Ok(()) + } + + fn poll_flush( + self: std::pin::Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + std::task::Poll::Pending + } + + fn poll_close( + self: std::pin::Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + std::task::Poll::Pending + } +} + +impl Drop for StalledBrowserOutput { + fn drop(&mut self) { + if let Some(dropped) = self.dropped.take() { + let _ = dropped.send(()); + } + } +} + +fn assert_bounded_unknown_delivery_close(messages: &[WsMessage]) { + let Some(WsMessage::Close(Some(frame))) = messages.last() else { + panic!("terminal browser session must attempt an explicit Close frame: {messages:?}"); + }; + assert_eq!(frame.code, 1011); + assert!(frame.reason.contains("delivery unknown"), "{frame:?}"); + assert!(frame.reason.contains("no retry"), "{frame:?}"); + assert!( + frame.reason.len() < 123, + "Close reason exceeds its wire bound: {frame:?}" + ); +} + +async fn assert_same_worker_reconnect_and_stop( + worker: &WorkspaceWorker, + context: &WorkerOperationContext, +) { + let mut fresh = tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, worker.connect_protocol(context)) + .await + .expect("terminated transport retained same Worker's connection authority") + .unwrap(); + assert!(matches!( + tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, fresh.events.recv()) + .await + .expect("same Worker reconnect did not deliver a fresh snapshot"), + Some(protocol::Event::Snapshot { .. }) + )); + let stopped = tokio::time::timeout( + PROTOCOL_TEST_TIMEOUT, + worker.stop( + &WorkerOperationContext::Backend, + WorkerLifecycleRequest { + reason: Some("cleanup-after-stalled-protocol".into()), + ticket_assignment: None, + }, + ), + ) + .await + .expect("terminated transport pinned same Worker's Stop/session cleanup") + .unwrap(); + assert_eq!(stopped.state, InternalWorkerOperationState::Accepted); + assert_eq!( + tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, fresh.sender.termination.ended()) + .await + .expect("Stop did not terminate the fresh same Worker connection"), + WorkerProtocolEnd::Disconnected + ); +} + +async fn assert_stalled_browser_output_reclaimed(explicit_termination: bool) { + let workspace = tempfile::tempdir().unwrap(); + let (api, execution) = test_api_with_recording_backend(workspace.path()).await; + let identity = spawn_ticket_check_source(&api, "stalled-browser-output"); + let worker = WorkspaceWorker::resolve(&api, &identity.runtime_id, &identity.worker_id).unwrap(); + let (context, _) = backpressure_browser_context(&api).await; + // Use the real attached embedded execution snapshot as browser output. + let raw = connect_workspace_worker_protocol(&api, &identity, None) + .await + .unwrap(); + let methods = raw.methods.clone(); + let termination = raw.termination; + let connection = worker_operations::WorkspaceWorkerConnection { + sender: WorkspaceWorkerMethodSender::new(worker.clone(), raw.methods, termination.clone()), + events: raw.events, + }; + let (entered_sender, entered) = tokio::sync::oneshot::channel(); + let (dropped_sender, dropped) = tokio::sync::oneshot::channel(); + let written = Arc::new(std::sync::Mutex::new(Vec::new())); + let socket = StalledBrowserOutput { + entered: Some(entered_sender), + dropped: Some(dropped_sender), + written: written.clone(), + }; + let started = tokio::time::Instant::now(); + { + let session = workspace_worker_protocol_session(connection, socket, context.clone()); + tokio::pin!(session); + assert!(futures::poll!(session.as_mut()).is_pending()); + // Poll the real session alongside its output-acceptance oneshot rather + // than assuming how many scheduler turns the attached pump needs. + let output_accepted = async { + tokio::select! { + result = entered => result.expect("browser sink dropped before receiving output"), + _ = session.as_mut() => panic!("browser session ended before accepting output"), + } + }; + tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, output_accepted) + .await + .expect("attached snapshot never reached the stalled browser output"); + { + let messages = written.lock().unwrap(); + assert_eq!(messages.len(), 1); + let WsMessage::Text(text) = &messages[0] else { + panic!("first browser output must be the attached snapshot"); + }; + assert!(matches!( + protocol::stream::decode_event(text.as_str()).unwrap(), + protocol::Event::Snapshot { .. } + )); + } + + if explicit_termination { + termination.terminate(WorkerProtocolEnd::Disconnected); + } else { + tokio::time::advance(std::time::Duration::from_secs(9)).await; + assert!(futures::poll!(session.as_mut()).is_pending()); + assert_eq!(termination.reason(), None); + tokio::time::advance(std::time::Duration::from_secs(1)).await; + } + assert!(futures::poll!(session.as_mut()).is_pending()); + assert_eq!( + termination.reason(), + Some(if explicit_termination { + WorkerProtocolEnd::Disconnected + } else { + WorkerProtocolEnd::TransportTimeout + }), + "browser end must publish the concrete common termination cause" + ); + { + let messages = written.lock().unwrap(); + assert_eq!(messages.len(), 2, "must not retry uncertain browser output"); + assert_bounded_unknown_delivery_close(&messages); + } + tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, methods.closed()) + .await + .expect("browser termination did not cancel its actual attached pump"); + let session_lock = current_worker_session_lock(&api, &identity); + let guard = tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, session_lock.lock()) + .await + .expect("blocked browser Close pinned session ownership"); + close_current_worker_session_locked(&api, &identity) + .await + .unwrap(); + drop(guard); + + // The terminal Close itself is stalled. It must get a separate bounded + // attempt, not become a second unbounded output wait. + tokio::time::advance(std::time::Duration::from_secs(9)).await; + assert!(futures::poll!(session.as_mut()).is_pending()); + tokio::time::advance(std::time::Duration::from_secs(1)).await; + tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, session.as_mut()) + .await + .expect("stalled terminal Close did not reach its deadline"); + } + tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, dropped) + .await + .expect("completed browser session retained its stalled sink") + .expect("browser sink dropped without its ownership signal"); + assert_eq!(written.lock().unwrap().len(), 2); + assert_eq!( + tokio::time::Instant::now() - started, + std::time::Duration::from_secs(if explicit_termination { 10 } else { 20 }) + ); + assert!(execution.protocol_methods().is_empty()); + assert_same_worker_reconnect_and_stop(&worker, &context).await; +} + +#[tokio::test(start_paused = true)] +async fn browser_blocked_output_deadline_ends_pump_and_bounds_the_terminal_close() { + assert_stalled_browser_output_reclaimed(false).await; +} + +#[tokio::test(start_paused = true)] +async fn browser_common_termination_preempts_blocked_output_and_bounds_the_terminal_close() { + assert_stalled_browser_output_reclaimed(true).await; +} + +#[tokio::test] +async fn embedded_disconnect_and_last_owner_drop_release_the_execution_event_subscription() { + for explicit_termination in [false, true] { + let (runtime, worker_ref, execution) = runtime_with_worker_backend(); + let endpoint = execution + .contexts + .lock() + .unwrap() + .get(&worker_ref) + .cloned() + .unwrap(); + let mut connection = connect_embedded_worker_protocol( + embedded_protocol_test_source(&runtime, &worker_ref), + embedded_protocol_test_scope(), + ) + .await + .unwrap(); + assert_embedded_protocol_snapshot( + next_embedded_protocol_event(&mut connection).await, + &worker_ref, + ); + assert_eq!(endpoint.events.receiver_count(), 1); + let WorkspaceWorkerProtocolConnection { + methods, + mut events, + termination, + } = connection; + let last_owner = termination.clone(); + drop(termination); + assert!(!methods.is_closed()); + if explicit_termination { + last_owner.terminate(WorkerProtocolEnd::Disconnected); + // Keep the owner alive until watch cancellation is observed so + // abort-on-drop cannot hide a broken explicit-disconnect path. + tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, methods.closed()) + .await + .expect("embedded explicit disconnect did not reclaim its pump"); + } + // Both channel endpoints remain held. Without explicit termination, + // only dropping the last client owner can cancel this attached pump. + drop(last_owner); + tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, methods.closed()) + .await + .expect("embedded last-owner drop did not reclaim its method pump"); + assert!( + tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, events.recv()) + .await + .expect("reclaimed embedded event endpoint stayed open") + .is_none() + ); + // The fixture's execution relay notices its dropped receiver on its + // next event. Wake that relay once, then observe subscription closure. + let _ = endpoint.events.send(protocol::Event::TextDelta { + text: "execution-relay-close-probe".into(), + }); + tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, endpoint.events.closed()) + .await + .expect("dropped embedded transport retained its execution event subscription"); + assert_eq!(endpoint.events.receiver_count(), 0); + assert!(execution.protocol_methods().is_empty()); + runtime.stop_worker(&worker_ref, None).unwrap(); + } +} + +#[tokio::test(start_paused = true)] +async fn stalled_remote_bridge_browser_cancel_overflow_reclaims_wire_and_allows_same_worker_progress() + { + let workspace = tempfile::tempdir().unwrap(); + let (api, execution) = test_api_with_recording_backend(workspace.path()).await; + let identity = spawn_ticket_check_source(&api, "remote-browser-overflow"); + let worker = WorkspaceWorker::resolve(&api, &identity.runtime_id, &identity.worker_id).unwrap(); + let (context, _) = backpressure_browser_context(&api).await; + let (raw, entered, dropped, sent) = super::worker_protocol_tests::stalled_remote_bridge(); + raw.methods + .try_send(protocol_test_completion("wire-delivery-unknown")) + .unwrap(); + tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, entered) + .await + .expect("actual remote pump did not enter its stalled wire send") + .unwrap(); + for index in 0..WORKER_PROTOCOL_METHOD_CAPACITY { + raw.methods + .try_send(protocol_test_completion(&format!("remote-queued-{index}"))) + .unwrap(); + } + let methods = raw.methods.clone(); + let termination = raw.termination; + let connection = worker_operations::WorkspaceWorkerConnection { + sender: WorkspaceWorkerMethodSender::new(worker.clone(), raw.methods, termination.clone()), + events: raw.events, + }; + let cancel = protocol::Method::Cancel { + command: protocol::WorkerCommandEnvelope::new(804), + }; + let mut socket = AdmissionCancellationSocket { + inbound: Some(WsMessage::Text( + protocol::stream::encode_method(&cancel).unwrap().into(), + )), + written: Vec::new(), + }; + // This real browser session decodes Cancel and passes the authenticated + // context through the actual method sender into the saturated remote pump. + tokio::time::timeout( + PROTOCOL_TEST_TIMEOUT, + workspace_worker_protocol_session(connection, &mut socket, context.clone()), + ) + .await + .expect("browser Cancel waited for the stalled remote wire instead of failing admission"); + assert_eq!( + termination.reason(), + Some(WorkerProtocolEnd::MethodOverflow) + ); + assert!(socket.inbound.is_none()); + assert_bounded_unknown_delivery_close(&socket.written); + tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, dropped) + .await + .expect("method overflow did not drop the stalled remote sink") + .unwrap(); + assert!(methods.is_closed()); + { + let messages = sent.lock().unwrap(); + assert_eq!( + messages.len(), + 1, + "must not retry or drain unknown remote delivery" + ); + let TungsteniteMessage::Text(text) = &messages[0] else { + panic!("remote pump must send an actual encoded method frame"); + }; + assert!(matches!( + decode_method(text.as_str()).unwrap(), + protocol::Method::ListCompletions { request_id: Some(id), .. } + if id == "wire-delivery-unknown" + )); + } + assert!(execution.protocol_methods().is_empty()); + assert_same_worker_reconnect_and_stop(&worker, &context).await; +} diff --git a/crates/workspace-server/src/server/tests/worker_protocol_tests.rs b/crates/workspace-server/src/server/tests/worker_protocol_tests.rs new file mode 100644 index 000000000..a66c68322 --- /dev/null +++ b/crates/workspace-server/src/server/tests/worker_protocol_tests.rs @@ -0,0 +1,210 @@ +use super::*; +use std::pin::Pin; +use std::task::{Context, Poll}; +use tokio::sync::oneshot; + +// The real bridge owns this sink exactly as it owns a TCP WebSocket sink. It +// accepts bytes but never finishes flushing: delivery is deliberately unknown. +struct StalledWire { + entered: Option>, + dropped: Option>, + sent: std::sync::Arc>>, +} +impl futures::Sink for StalledWire { + type Error = tokio_tungstenite::tungstenite::Error; + fn poll_ready( + self: Pin<&mut Self>, + _: &mut Context<'_>, + ) -> Poll> { + Poll::Ready(Ok(())) + } + fn start_send( + mut self: Pin<&mut Self>, + message: TungsteniteMessage, + ) -> std::result::Result<(), Self::Error> { + self.sent.lock().unwrap().push(message); + if let Some(entered) = self.entered.take() { + let _ = entered.send(()); + } + Ok(()) + } + fn poll_flush( + self: Pin<&mut Self>, + _: &mut Context<'_>, + ) -> Poll> { + Poll::Pending + } + fn poll_close( + self: Pin<&mut Self>, + _: &mut Context<'_>, + ) -> Poll> { + Poll::Pending + } +} +impl Drop for StalledWire { + fn drop(&mut self) { + if let Some(dropped) = self.dropped.take() { + let _ = dropped.send(()); + } + } +} +pub(super) fn stalled_remote_bridge() -> ( + WorkspaceWorkerProtocolConnection, + oneshot::Receiver<()>, + oneshot::Receiver<()>, + std::sync::Arc>>, +) { + let (entered_tx, entered) = oneshot::channel(); + let (dropped_tx, dropped) = oneshot::channel(); + let sent = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); + let sink = StalledWire { + entered: Some(entered_tx), + dropped: Some(dropped_tx), + sent: sent.clone(), + }; + let connection = remote_worker_protocol_bridge(sink, futures::stream::pending()); + (connection, entered, dropped, sent) +} + +#[tokio::test(start_paused = true)] +async fn remote_pending_method_timeout_closes_both_queues_and_never_retries_unknown_delivery() { + let (mut connection, entered, dropped, sent) = stalled_remote_bridge(); + connection + .methods + .try_send(protocol_test_completion("delivery-unknown")) + .unwrap(); + entered.await.unwrap(); + for _ in 0..WORKER_PROTOCOL_METHOD_CAPACITY { + connection + .methods + .try_send(protocol::Method::ListPendingSubmissions) + .unwrap(); + } + // Queues are saturated while the actual bridge is stalled in a wire send. + assert!( + connection + .methods + .try_send(protocol::Method::ListRewindTargets) + .is_err() + ); + tokio::time::advance(std::time::Duration::from_secs(10)).await; + assert_eq!( + connection.termination.ended().await, + WorkerProtocolEnd::TransportTimeout + ); + dropped.await.unwrap(); + assert!(connection.methods.is_closed()); + assert!(connection.events.recv().await.is_none()); + assert_eq!( + sent.lock().unwrap().len(), + 1, + "must not retry or drain uncertain pending methods" + ); +} + +#[tokio::test] +async fn remote_disconnect_and_last_capability_drop_reclaim_a_pending_wire_send() { + for explicit in [false, true] { + let (connection, entered, dropped, sent) = stalled_remote_bridge(); + connection + .methods + .try_send(protocol_test_completion("cancelled-wire")) + .unwrap(); + entered.await.unwrap(); + if explicit { + connection + .termination + .terminate(WorkerProtocolEnd::Disconnected); + } + drop(connection); + tokio::time::timeout(std::time::Duration::from_secs(2), dropped) + .await + .unwrap() + .unwrap(); + assert_eq!(sent.lock().unwrap().len(), 1); + } +} + +#[tokio::test] +async fn remote_actual_connector_overflow_ends_connection_and_preserves_delivered_event_order() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (release_tx, release) = oneshot::channel(); + let (wire_closed_tx, wire_closed) = oneshot::channel(); + let peer = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut socket = tokio_tungstenite::accept_hdr_async( + stream, + |request: &tokio_tungstenite::tungstenite::handshake::server::Request, response| { + assert_eq!( + request.headers()["authorization"], + "Bearer test-worker-protocol" + ); + assert_eq!( + request.headers()[protocol::AUTHENTICATED_ACCOUNT_ID_HEADER], + "test-account" + ); + Ok(response) + }, + ) + .await + .unwrap(); + release.await.unwrap(); + for i in 0..=WORKER_PROTOCOL_EVENT_CAPACITY { + let event = protocol::Event::TextDelta { + text: i.to_string(), + }; + let text = protocol::stream::encode_event(&event).unwrap(); + // TCP may observe closure before all writes flush; either way the + // consumer must reach the same explicit bounded-overflow end. + if socket + .send(TungsteniteMessage::Text(text.into())) + .await + .is_err() + { + break; + } + } + while let Some(Ok(message)) = socket.next().await { + if matches!(message, TungsteniteMessage::Close(_)) { + break; + } + } + let _ = wire_closed_tx.send(()); + }); + let config = RuntimeObservationSourceConfig { + worker: RuntimeWorkerRef::new("test-runtime", "test-worker"), + endpoint: format!("ws://{addr}/protocol/ws"), + bearer_token: Some("test-worker-protocol".into()), + }; + let mut connection = connect_remote_worker_protocol( + config, + Some(&protocol::AuthenticatedInputSource::Account { + account_id: "test-account".into(), + }), + ) + .await + .unwrap(); + release_tx.send(()).unwrap(); + assert_eq!( + tokio::time::timeout( + std::time::Duration::from_secs(2), + connection.termination.ended() + ) + .await + .unwrap(), + WorkerProtocolEnd::EventOverflow + ); + for i in 0..WORKER_PROTOCOL_EVENT_CAPACITY { + assert!( + matches!(connection.events.recv().await, Some(protocol::Event::TextDelta { text }) if text == i.to_string()) + ); + } + assert!(connection.events.recv().await.is_none()); + assert!(connection.methods.is_closed()); + tokio::time::timeout(std::time::Duration::from_secs(2), wire_closed) + .await + .unwrap() + .unwrap(); + peer.await.unwrap(); +} diff --git a/crates/workspace-server/src/server/tests/workspace_protocol_backpressure_tests.rs b/crates/workspace-server/src/server/tests/workspace_protocol_backpressure_tests.rs new file mode 100644 index 000000000..add77016e --- /dev/null +++ b/crates/workspace-server/src/server/tests/workspace_protocol_backpressure_tests.rs @@ -0,0 +1,140 @@ +use super::*; +use protocol::subscription::{ + EventSubscriptionSelector, SubscriptionFrame, SubscriptionFramePayload, SubscriptionRequest, + SubscriptionRequestId, +}; +use std::pin::Pin; +use std::task::{Context, Poll}; +use tokio::sync::{mpsc, oneshot}; + +// A controlled browser wire. Receiving remains independent of a blocked writer, +// reproducing the real multiplexer split without filling an OS socket by timing. +struct CongestedBrowser { + inbound: mpsc::Receiver, + entered: Option>, + dropped: Option>, +} +impl futures::Stream for CongestedBrowser { + type Item = std::result::Result; + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.inbound.poll_recv(cx).map(|message| message.map(Ok)) + } +} +impl futures::Sink for CongestedBrowser { + type Error = std::convert::Infallible; + fn poll_ready( + self: Pin<&mut Self>, + _: &mut Context<'_>, + ) -> Poll> { + Poll::Ready(Ok(())) + } + fn start_send(mut self: Pin<&mut Self>, _: WsMessage) -> std::result::Result<(), Self::Error> { + if let Some(entered) = self.entered.take() { + let _ = entered.send(()); + } + Ok(()) + } + fn poll_flush( + self: Pin<&mut Self>, + _: &mut Context<'_>, + ) -> Poll> { + Poll::Pending + } + fn poll_close( + self: Pin<&mut Self>, + _: &mut Context<'_>, + ) -> Poll> { + Poll::Pending + } +} +impl Drop for CongestedBrowser { + fn drop(&mut self) { + if let Some(dropped) = self.dropped.take() { + let _ = dropped.send(()); + } + } +} + +#[tokio::test(start_paused = true)] +async fn actual_workspace_multiplexer_stalled_writer_timeout_or_disconnect_reclaims_connection_tasks() + { + for exit in ["timeout", "disconnect", "caller-cancel"] { + let workspace = tempfile::tempdir().unwrap(); + let (api, _) = test_api_with_recording_backend(workspace.path()).await; + let identity = spawn_ticket_check_source(&api, "multiplexer-stalled-writer"); + let (inbound, receiver) = mpsc::channel(300); + let (entered_tx, entered) = oneshot::channel(); + let (dropped_tx, dropped) = oneshot::channel(); + let socket = CongestedBrowser { + inbound: receiver, + entered: Some(entered_tx), + dropped: Some(dropped_tx), + }; + let frame = SubscriptionFrame::new(SubscriptionFramePayload::Request( + SubscriptionRequest::SubscribeEvents { + request_id: SubscriptionRequestId::new("saturated-browser").unwrap(), + selector: EventSubscriptionSelector::WorkspaceWorkers, + }, + )); + inbound + .try_send(WsMessage::Text( + serde_json::to_string(&frame).unwrap().into(), + )) + .unwrap(); + let task = tokio::spawn(crate::workspace_subscription::serve_workspace_subscription( + api.clone(), + socket, + WorkerOperationContext::Backend, + )); + entered.await.unwrap(); // Actual subscription snapshot reached the blocked wire. + match exit { + "timeout" => { + // Fill the control queue as well; the receiving side may now be + // waiting for Pong capacity. Writer failure must preempt that wait. + for _ in 0..258 { + inbound + .try_send(WsMessage::Ping(Vec::new().into())) + .unwrap(); + } + tokio::time::advance(std::time::Duration::from_secs(10)).await; + task.await.unwrap(); + } + "disconnect" => { + drop(inbound); + task.await.unwrap(); + } + "caller-cancel" => { + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + } + _ => unreachable!(), + } + dropped.await.unwrap(); + let worker = + WorkspaceWorker::resolve(&api, &identity.runtime_id, &identity.worker_id).unwrap(); + let mut fresh = worker + .connect_protocol(&WorkerOperationContext::Backend) + .await + .unwrap(); + assert!(matches!( + fresh.events.recv().await, + Some(protocol::Event::Snapshot { .. }) + )); + let stopped = worker + .stop( + &WorkerOperationContext::Backend, + WorkerLifecycleRequest { + reason: Some("cleanup-after-multiplexer-exit".into()), + ticket_assignment: None, + }, + ) + .await + .unwrap(); + assert_eq!(stopped.state, InternalWorkerOperationState::Accepted); + assert!( + current_worker_session_lock(&api, &identity) + .try_lock() + .is_ok() + ); + } +} diff --git a/crates/workspace-server/src/server/worker_operations.rs b/crates/workspace-server/src/server/worker_operations.rs index fe84fb7f2..35b60148f 100644 --- a/crates/workspace-server/src/server/worker_operations.rs +++ b/crates/workspace-server/src/server/worker_operations.rs @@ -470,10 +470,11 @@ impl WorkspaceWorker { let connection = connect_workspace_worker_protocol(&self.api, &self.identity, source).await?; Ok(WorkspaceWorkerConnection { - sender: WorkspaceWorkerMethodSender { - worker: self.clone(), - methods: connection.methods, - }, + sender: WorkspaceWorkerMethodSender::new( + self.clone(), + connection.methods, + connection.termination, + ), events: connection.events, }) } @@ -485,12 +486,27 @@ impl WorkspaceWorker { pub(crate) struct WorkspaceWorkerMethodSender { worker: WorkspaceWorker, methods: tokio::sync::mpsc::Sender, + pub(super) termination: WorkerProtocolTermination, } pub(crate) struct WorkspaceWorkerConnection { pub(crate) sender: WorkspaceWorkerMethodSender, pub(crate) events: tokio::sync::mpsc::Receiver, } impl WorkspaceWorkerMethodSender { + pub(crate) fn end_status(&self) -> tokio::sync::watch::Receiver> { + self.termination.subscribe() + } + pub(super) fn new( + worker: WorkspaceWorker, + methods: tokio::sync::mpsc::Sender, + termination: WorkerProtocolTermination, + ) -> Self { + Self { + worker, + methods, + termination, + } + } pub(crate) async fn send( &self, context: &WorkerOperationContext, @@ -604,15 +620,25 @@ impl WorkspaceWorkerMethodSender { }; let method = authorize_browser_worker_method(method, source) .map_err(|message| Error::InvalidInput(message.into()))?; - self.methods.send(method).await.map_err(|_| { + // Session exclusion fences attach/stop only through admission, never delivery. + let result = if self.termination.reason().is_some() { + Err(tokio::sync::mpsc::error::TrySendError::Closed(method)) + } else { + self.methods.try_send(method) + }; + result.map_err(|error| { + let (code, message) = match error { + tokio::sync::mpsc::error::TrySendError::Full(_) => { + self.termination.terminate(WorkerProtocolEnd::MethodOverflow); + ("worker_protocol_overflow", "Worker method queue overflow; this method was not admitted. Earlier pending methods may have unknown delivery; reconnect with a fresh snapshot, no automatic retry") + } + tokio::sync::mpsc::error::TrySendError::Closed(_) => + ("worker_protocol_closed", "Worker connection closed; pending delivery may be unknown; no implicit Restore or retry was performed"), + }; Error::RuntimeOperationFailed { runtime_id: self.worker.identity.runtime_id.clone(), - code: "worker_protocol_closed".into(), - message: - "Worker execution connection closed; no implicit Restore or retry was performed" - .into(), - } - .into() + code: code.into(), message: message.into(), + }.into() }) } } diff --git a/crates/workspace-server/src/server/worker_protocol.rs b/crates/workspace-server/src/server/worker_protocol.rs new file mode 100644 index 000000000..c1a5ed9e0 --- /dev/null +++ b/crates/workspace-server/src/server/worker_protocol.rs @@ -0,0 +1,247 @@ +//! Bounded operational protocol connections. Queue admission never waits for delivery. +//! Overflow ends the connection: reconnect/resubscribe with a fresh snapshot, never replay +//! a method whose transport delivery or execution outcome may be unknown. +use super::*; +use futures::{Sink, Stream}; +use std::future::Future; +use tokio::sync::{mpsc, watch}; + +pub(super) const WORKER_PROTOCOL_METHOD_CAPACITY: usize = 256; +pub(super) const WORKER_PROTOCOL_EVENT_CAPACITY: usize = 512; +pub(super) const WORKER_PROTOCOL_TRANSPORT_TIMEOUT: std::time::Duration = + std::time::Duration::from_secs(10); + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum WorkerProtocolEnd { + MethodOverflow, + EventOverflow, + TransportTimeout, + Disconnected, +} +impl WorkerProtocolEnd { + fn message(self) -> &'static str { + match self { + Self::MethodOverflow => "Worker protocol method queue overflow", + Self::EventOverflow => "Worker protocol event queue overflow/lag", + Self::TransportTimeout => "Worker protocol transport timed out", + Self::Disconnected => "Worker protocol disconnected", + } + } +} + +struct WorkerProtocolPump(tokio::task::AbortHandle); +impl Drop for WorkerProtocolPump { + fn drop(&mut self) { + self.0.abort(); + } +} + +/// The last client capability owns the pump. The pump holds only the watch sender, +/// not this owner, so dropping a cancelled connection cannot leave an orphan task. +#[derive(Clone)] +pub(crate) struct WorkerProtocolTermination { + state: watch::Sender>, + _pump: std::sync::Arc, +} +impl WorkerProtocolTermination { + pub(crate) fn subscribe(&self) -> watch::Receiver> { + self.state.subscribe() + } + pub(crate) fn reason(&self) -> Option { + *self.state.borrow() + } + pub(crate) fn terminate(&self, reason: WorkerProtocolEnd) { + finish(&self.state, reason); + } + pub(crate) async fn ended(&self) -> WorkerProtocolEnd { + ended(&mut self.state.subscribe()).await + } +} +fn finish(state: &watch::Sender>, reason: WorkerProtocolEnd) { + state.send_if_modified(|current| { + if current.is_some() { + return false; + } + *current = Some(reason); + true + }); +} +async fn ended(state: &mut watch::Receiver>) -> WorkerProtocolEnd { + loop { + if let Some(reason) = *state.borrow_and_update() { + return reason; + } + if state.changed().await.is_err() { + return WorkerProtocolEnd::Disconnected; + } + } +} + +pub(crate) struct WorkspaceWorkerProtocolConnection { + pub(crate) methods: mpsc::Sender, + pub(crate) events: mpsc::Receiver, + pub(crate) termination: WorkerProtocolTermination, +} + +pub(super) fn spawn_worker_protocol_pump(pump: F) -> WorkspaceWorkerProtocolConnection +where + F: FnOnce(mpsc::Receiver, mpsc::Sender) -> Fut, + Fut: Future + Send + 'static, +{ + let (methods, receiver) = mpsc::channel(WORKER_PROTOCOL_METHOD_CAPACITY); + let (sender, events) = mpsc::channel(WORKER_PROTOCOL_EVENT_CAPACITY); + let (state, mut shutdown) = watch::channel(None); + let task_state = state.clone(); + let closed = sender.clone(); + let future = pump(receiver, sender); + let task = tokio::spawn(async move { + let reason = tokio::select! { + biased; + reason = ended(&mut shutdown) => reason, + _ = closed.closed() => WorkerProtocolEnd::Disconnected, + reason = future => reason, + }; + finish(&task_state, reason); + // Dropping the future releases both queues and the captured execution/WS endpoint. + }); + WorkspaceWorkerProtocolConnection { + methods, + events, + termination: WorkerProtocolTermination { + state, + _pump: std::sync::Arc::new(WorkerProtocolPump(task.abort_handle())), + }, + } +} + +pub(super) fn enqueue_worker_protocol_event( + sender: &mpsc::Sender, + event: protocol::Event, +) -> std::result::Result<(), WorkerProtocolEnd> { + sender.try_send(event).map_err(|error| match error { + mpsc::error::TrySendError::Full(_) => WorkerProtocolEnd::EventOverflow, + mpsc::error::TrySendError::Closed(_) => WorkerProtocolEnd::Disconnected, + }) +} + +/// Actual remote bridge pump, also usable with a bounded in-process wire fixture. +/// A transport send may have partially succeeded when interrupted: no retry or success +/// acknowledgement is synthesized. All transport waits have a deadline; cancellation +/// of the common owner drops the pending send and both halves of this one connection. +pub(super) fn remote_worker_protocol_bridge( + mut sink: S, + mut stream: R, +) -> WorkspaceWorkerProtocolConnection +where + S: Sink + Unpin + Send + 'static, + R: Stream< + Item = std::result::Result, + > + Unpin + + Send + + 'static, + S::Error: Send, +{ + spawn_worker_protocol_pump(move |mut methods, events| async move { + loop { + tokio::select! { + method = methods.recv() => { + let Some(method) = method else { return WorkerProtocolEnd::Disconnected; }; + let Ok(text) = protocol::stream::encode_method(&method) else { return WorkerProtocolEnd::Disconnected; }; + match tokio::time::timeout(WORKER_PROTOCOL_TRANSPORT_TIMEOUT, sink.send(TungsteniteMessage::Text(text.into()))).await { + Ok(Ok(())) => {}, + Ok(Err(_)) => return WorkerProtocolEnd::Disconnected, + Err(_) => return WorkerProtocolEnd::TransportTimeout, + } + }, + message = stream.next() => match message { + Some(Ok(TungsteniteMessage::Text(text))) => { + let Ok(event) = protocol::stream::decode_event(text.as_ref()) else { return WorkerProtocolEnd::Disconnected; }; + if let Err(reason) = enqueue_worker_protocol_event(&events, event) { return reason; } + } + Some(Ok(TungsteniteMessage::Ping(value))) => { + match tokio::time::timeout(WORKER_PROTOCOL_TRANSPORT_TIMEOUT, sink.send(TungsteniteMessage::Pong(value))).await { + Ok(Ok(())) => {}, + Ok(Err(_)) => return WorkerProtocolEnd::Disconnected, + Err(_) => return WorkerProtocolEnd::TransportTimeout, + } + } + Some(Ok(TungsteniteMessage::Pong(_))) => {}, + _ => return WorkerProtocolEnd::Disconnected, + } + } + } + }) +} + +/// Browser admission can await authorization/session ownership, but never method queue +/// capacity. Common termination preempts that wait. Browser writes are bounded too. +pub(super) async fn workspace_worker_protocol_session( + mut connection: worker_operations::WorkspaceWorkerConnection, + mut socket: S, + context: WorkerOperationContext, +) where + S: Stream> + Sink + Unpin, +{ + let termination = connection.sender.termination.clone(); + let session = async { + loop { + tokio::select! { + inbound = socket.next() => match inbound { + Some(Ok(WsMessage::Text(text))) => { + let result = match decode_method(&text) { + Ok(method) => connection.sender.send(&context, method).await.map_err(|error| error.to_string()), + Err(error) => Err(error.to_string()), + }; + if let Err(error) = result { + let event = protocol_error_event(error); + let Ok(text) = protocol::stream::encode_event(&event) else { return WorkerProtocolEnd::Disconnected; }; + if let Err(reason) = browser_send(&mut socket, WsMessage::Text(text.into())).await { return reason; } + } + } + Some(Ok(WsMessage::Ping(value))) => { + if let Err(reason) = browser_send(&mut socket, WsMessage::Pong(value)).await { return reason; } + } + Some(Ok(WsMessage::Pong(_))) => {}, + _ => return WorkerProtocolEnd::Disconnected, + }, + event = connection.events.recv() => { + let Some(event) = event else { return WorkerProtocolEnd::Disconnected; }; + let shutdown = matches!(event, protocol::Event::Shutdown); + let Ok(text) = protocol::stream::encode_event(&event) else { return WorkerProtocolEnd::Disconnected; }; + if let Err(reason) = browser_send(&mut socket, WsMessage::Text(text.into())).await { return reason; } + if shutdown { return WorkerProtocolEnd::Disconnected; } + } + } + } + }; + let reason = tokio::select! { + biased; + reason = termination.ended() => reason, + reason = session => reason, + }; + termination.terminate(reason); + // Best-effort explicit close even if the downstream transport is also blocked. + // The close reason deliberately makes no claim that admitted methods did not run. + let _ = tokio::time::timeout( + WORKER_PROTOCOL_TRANSPORT_TIMEOUT, + socket.send(WsMessage::Close(Some(axum::extract::ws::CloseFrame { + code: 1011, + reason: format!( + "{}; pending delivery unknown; resync snapshot; no retry", + reason.message() + ) + .into(), + }))), + ) + .await; +} +async fn browser_send + Unpin>( + socket: &mut S, + message: WsMessage, +) -> std::result::Result<(), WorkerProtocolEnd> { + match tokio::time::timeout(WORKER_PROTOCOL_TRANSPORT_TIMEOUT, socket.send(message)).await { + Ok(Ok(())) => Ok(()), + Ok(Err(_)) => Err(WorkerProtocolEnd::Disconnected), + Err(_) => Err(WorkerProtocolEnd::TransportTimeout), + } +} diff --git a/crates/workspace-server/src/workspace_subscription.rs b/crates/workspace-server/src/workspace_subscription.rs index 9e7698a30..8b0e0b818 100644 --- a/crates/workspace-server/src/workspace_subscription.rs +++ b/crates/workspace-server/src/workspace_subscription.rs @@ -1,6 +1,6 @@ use std::collections::HashMap; -use axum::extract::ws::{Message as WsMessage, WebSocket}; +use axum::extract::ws::Message as WsMessage; use futures::{SinkExt, StreamExt}; use protocol::subscription::{ EventSubscriptionSelector, SubscriptionEvent, SubscriptionEventPayload, SubscriptionFrame, @@ -12,27 +12,43 @@ use worker_runtime::identity::RuntimeWorkerRef; use crate::runtime_subscription::RuntimeSubscriptionBroker; use crate::server::{ - WorkerOperationContext, WorkspaceApi, WorkspaceWorker, WorkspaceWorkerMethodSender, + WorkerOperationContext, WorkerProtocolEnd, WorkspaceApi, WorkspaceWorker, + WorkspaceWorkerMethodSender, }; const OUTBOUND_CAPACITY: usize = 256; const WORKER_PROTOCOL_SNAPSHOT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); +// Every spawned writer/subscription is connection-owned, including caller cancellation. +struct SubscriptionTask(tokio::task::JoinHandle<()>); +impl Drop for SubscriptionTask { + fn drop(&mut self) { + self.0.abort(); + } +} struct ActiveSubscription { - task: tokio::task::JoinHandle<()>, + task: SubscriptionTask, methods: Option, } -pub(crate) async fn serve_workspace_subscription( +pub(crate) async fn serve_workspace_subscription( api: WorkspaceApi, - socket: WebSocket, + socket: S, context: WorkerOperationContext, -) { +) where + S: futures::Stream> + + futures::Sink + + Unpin + + Send + + 'static, + S::Error: Send, +{ let broker = api.runtime_subscription_broker().clone(); let (mut socket_sender, mut socket_receiver) = socket.split(); let (control_outbound, mut control_receiver) = mpsc::channel::(OUTBOUND_CAPACITY); let (protocol_outbound, mut protocol_receiver) = mpsc::channel::(OUTBOUND_CAPACITY); - let writer = tokio::spawn(async move { + let (writer_done, finished) = tokio::sync::oneshot::channel(); + let mut writer = SubscriptionTask(tokio::spawn(async move { loop { let message = tokio::select! { biased; @@ -40,160 +56,179 @@ pub(crate) async fn serve_workspace_subscription( message = protocol_receiver.recv() => message, }; let Some(message) = message else { break }; - if socket_sender.send(message).await.is_err() { + if !matches!( + tokio::time::timeout( + WORKER_PROTOCOL_SNAPSHOT_TIMEOUT, + socket_sender.send(message) + ) + .await, + Ok(Ok(())) + ) { break; } } - }); + let _ = writer_done.send(()); + })); let mut next_subscription_id = 1_u64; let mut subscriptions = HashMap::::new(); - while let Some(message) = socket_receiver.next().await { - let Ok(message) = message else { break }; - match message { - WsMessage::Text(text) => { - let Ok(frame) = serde_json::from_str::(text.as_str()) else { - break; - }; - if frame.validate().is_err() { - break; - } - subscriptions.retain(|_, subscription| !subscription.task.is_finished()); - match frame.payload { - SubscriptionFramePayload::Request(SubscriptionRequest::SubscribeEvents { - request_id, - selector, - }) => { - 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); - match selector { - EventSubscriptionSelector::WorkspaceWorkers => { - let task = tokio::spawn(run_workspace_workers( - api.clone(), - 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), - } => { - let connection = match WorkspaceWorker::resolve( - &api, - &runtime_id, - worker_id.as_str(), - ) { - Ok(worker) => worker - .connect_protocol(&context) - .await - .map_err(|error| error.to_string()), - Err(error) => Err(error.to_string()), - }; - match connection { - Ok(connection) => { - let methods = connection.sender.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 receiving = async { + while let Some(message) = socket_receiver.next().await { + let Ok(message) = message else { break }; + match message { + WsMessage::Text(text) => { + let Ok(frame) = serde_json::from_str::(text.as_str()) else { + break; + }; + if frame.validate().is_err() { + break; + } + subscriptions.retain(|_, subscription| !subscription.task.0.is_finished()); + match frame.payload { + SubscriptionFramePayload::Request( + SubscriptionRequest::SubscribeEvents { + request_id, + selector, + }, + ) => { + 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); + match selector { + EventSubscriptionSelector::WorkspaceWorkers => { + let task = tokio::spawn(run_workspace_workers( + api.clone(), + broker.clone(), + request_id, + subscription_id.clone(), + control_outbound.clone(), + )); + subscriptions.insert( + subscription_id, + ActiveSubscription { + task: SubscriptionTask(task), + methods: None, + }, + ); + } + EventSubscriptionSelector::WorkerProtocol { + worker_id, + runtime_id: Some(runtime_id), + } => { + let connection = match WorkspaceWorker::resolve( + &api, + &runtime_id, + worker_id.as_str(), + ) { + Ok(worker) => worker + .connect_protocol(&context) + .await + .map_err(|error| error.to_string()), + Err(error) => Err(error.to_string()), + }; + match connection { + Ok(connection) => { + let methods = connection.sender.clone(); + let task = tokio::spawn(run_worker_protocol( + request_id, + subscription_id.clone(), + runtime_id, + worker_id, + connection.events, + methods.end_status(), + control_outbound.clone(), + protocol_outbound.clone(), + )); + subscriptions.insert( + subscription_id, + ActiveSubscription { + task: SubscriptionTask(task), + methods: Some(methods), + }, + ); + } + Err(error) => { + let _ = send_rejected( + &control_outbound, + request_id, + SubscriptionRejectionCode::ResourceNotFound, + error.to_string(), + ) + .await; + } } } - } - _ => { - let _ = send_rejected( + _ => { + 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; + } } } - } - SubscriptionFramePayload::Request(SubscriptionRequest::UnsubscribeEvents { - request_id, - subscription_id, - }) => { - if let Some(subscription) = subscriptions.remove(&subscription_id) { - subscription.task.abort(); + SubscriptionFramePayload::Request( + SubscriptionRequest::UnsubscribeEvents { + request_id, + subscription_id, + }, + ) => { + if let Some(mut subscription) = subscriptions.remove(&subscription_id) { + subscription.task.0.abort(); + let _ = (&mut subscription.task.0).await; + } + if send_frame( + &control_outbound, + SubscriptionFrame::new(SubscriptionFramePayload::Response( + SubscriptionResponse::Unsubscribed { + request_id, + subscription_id, + }, + )), + ) + .await + .is_err() + { + break; + } } - if send_frame( - &control_outbound, - SubscriptionFrame::new(SubscriptionFramePayload::Response( - SubscriptionResponse::Unsubscribed { - request_id, - subscription_id, - }, - )), - ) - .await - .is_err() - { + SubscriptionFramePayload::WorkerProtocol(message) => { + let Some(methods) = subscriptions + .get(&message.subscription_id) + .and_then(|value| value.methods.clone()) + else { + break; + }; + if methods.send(&context, message.method).await.is_err() { + break; + } + } + SubscriptionFramePayload::Response(_) + | SubscriptionFramePayload::Event(_) => { break; } } - SubscriptionFramePayload::WorkerProtocol(message) => { - let Some(methods) = subscriptions - .get(&message.subscription_id) - .and_then(|value| value.methods.clone()) - else { - break; - }; - if methods.send(&context, message.method).await.is_err() { - break; - } - } - SubscriptionFramePayload::Response(_) | SubscriptionFramePayload::Event(_) => { + } + WsMessage::Ping(value) => { + if control_outbound.send(WsMessage::Pong(value)).await.is_err() { break; } } + WsMessage::Pong(_) => {} + WsMessage::Close(_) | WsMessage::Binary(_) => break, } - WsMessage::Ping(value) => { - if control_outbound.send(WsMessage::Pong(value)).await.is_err() { - break; - } - } - WsMessage::Pong(_) => {} - WsMessage::Close(_) | WsMessage::Binary(_) => break, } + }; + // Writer termination also preempts method authorization/control queue waits. + tokio::select! { _ = receiving => {}, _ = finished => {} } + for (_, mut subscription) in subscriptions { + subscription.task.0.abort(); + let _ = (&mut subscription.task.0).await; } - - for (_, subscription) in subscriptions { - subscription.task.abort(); - } - drop(control_outbound); - drop(protocol_outbound); - let _ = writer.await; + // Never wait for a congested writer to drain after disconnect or admission failure. + writer.0.abort(); + let _ = (&mut writer.0).await; } async fn send_rejected( @@ -222,6 +257,7 @@ async fn run_worker_protocol( runtime_id: String, worker_id: protocol::subscription::SubscriptionWorkerId, mut events: mpsc::Receiver, + end_status: tokio::sync::watch::Receiver>, control_outbound: mpsc::Sender, protocol_outbound: mpsc::Sender, ) { @@ -313,8 +349,8 @@ async fn run_worker_protocol( SubscriptionFrame::new(SubscriptionFramePayload::Event( SubscriptionEvent::SubscriptionClosed { subscription_id, - code: SubscriptionTerminationCode::ResourceGone, - message: "Worker protocol stream closed".to_string(), + code: if matches!(*end_status.borrow(), Some(WorkerProtocolEnd::EventOverflow | WorkerProtocolEnd::MethodOverflow)) { SubscriptionTerminationCode::Lagged } else { SubscriptionTerminationCode::ResourceGone }, + message: "Worker protocol stream ended; pending method delivery may be unknown; resubscribe with a fresh snapshot, no automatic retry".to_string(), }, )), ) @@ -646,6 +682,7 @@ mod tests { "runtime-1".to_string(), worker_id(), events, + tokio::sync::watch::channel(None).1, control_outbound, protocol_outbound, ) @@ -660,6 +697,81 @@ mod tests { response } + #[tokio::test] + async fn worker_bridge_overflow_is_serialized_as_lagged_not_resource_disappearance() { + for (reason, expected) in [ + ( + WorkerProtocolEnd::EventOverflow, + SubscriptionTerminationCode::Lagged, + ), + ( + WorkerProtocolEnd::MethodOverflow, + SubscriptionTerminationCode::Lagged, + ), + ( + WorkerProtocolEnd::Disconnected, + SubscriptionTerminationCode::ResourceGone, + ), + ] { + let (event_sender, events) = mpsc::channel(1); + event_sender + .try_send(protocol::Event::Snapshot { + session: protocol::SessionSnapshot { + entries: vec![], + pending_submissions: Default::default(), + }, + greeting: protocol::Greeting { + worker_name: "worker-1".into(), + cwd: String::new(), + provider: "fixture".into(), + model: "fixture".into(), + reasoning: None, + scope_summary: String::new(), + tools: vec![], + context_window: 0, + context_tokens: 0, + context_usage: None, + }, + state: protocol::WorkerStateSnapshot::initial(), + in_flight: Default::default(), + internal_workers: vec![], + }) + .unwrap(); + drop(event_sender); + let (control, mut receiver) = mpsc::channel(2); + let (protocol, _events) = mpsc::channel(1); + run_worker_protocol( + request_id(), + subscription_id(), + "runtime-1".into(), + worker_id(), + events, + tokio::sync::watch::channel(Some(reason)).1, + control, + protocol, + ) + .await; + let WsMessage::Text(snapshot) = receiver.recv().await.unwrap() else { + panic!("expected snapshot response"); + }; + let snapshot: SubscriptionFrame = serde_json::from_str(snapshot.as_str()).unwrap(); + assert!(matches!( + snapshot.payload, + SubscriptionFramePayload::Response(SubscriptionResponse::Subscribed { .. }) + )); + let WsMessage::Text(closed) = receiver.recv().await.unwrap() else { + panic!("expected terminal control frame"); + }; + let frame: SubscriptionFrame = serde_json::from_str(closed.as_str()).unwrap(); + frame.validate().unwrap(); + assert!( + matches!(frame.payload, SubscriptionFramePayload::Event(SubscriptionEvent::SubscriptionClosed { code, message, .. }) + if code == expected && message.contains("unknown") && message.contains("fresh snapshot")), + "{closed}" + ); + } + } + #[tokio::test] async fn worker_protocol_rejects_a_non_snapshot_initial_event() { let (event_sender, events) = mpsc::channel(1); diff --git a/docs/development/worker-protocol-backpressure.md b/docs/development/worker-protocol-backpressure.md new file mode 100644 index 000000000..c7a22264c --- /dev/null +++ b/docs/development/worker-protocol-backpressure.md @@ -0,0 +1,112 @@ +# Worker operational protocol backpressure + +The Workspace Worker operational connection is separate from Workspace request +admission (T-732). Both the dedicated Worker WebSocket and the Workspace +multiplexer use the identity-bound `WorkspaceWorkerMethodSender` and the same +remote/embedded bridge. This does not change Ticket assignment, general input +routing, or database execution placement. + +## Ownership and admission + +- Each bridge owns one captured remote WebSocket or embedded execution endpoint. + It never follows a restored execution under the same Worker identity. +- Methods have 256 queue slots; events have 512. Browser authorization, input + source attribution, interactive-Worker checks, and session exclusion remain at + method admission. Queue admission is `try_send`, not a transport-capacity wait + under the session guard. Success here means **queued**, not delivered/executed. +- Event enqueue never waits for its consumer. Snapshot is first; direct replies + and live events retain the existing bridge ordering. No additional event pump + can reorder them. Overflow ends the connection rather than skipping an event + and continuing an uncertain stream. +- The first terminal cause wins: method overflow, event overflow/lag, transport + timeout, or disconnect. Terminal state cannot reopen when capacity becomes + available. It preempts pending browser admission and remote transport waits. +- The last client capability owns an abort-on-drop pump handle. The pump does not + retain that capability, avoiding an ownership cycle. Event-consumer closure + also cancels the pump even when method capabilities remain. Cancellation drops + the bounded channels and captured transport; it does not stop the Worker. + +A rejected full-queue method is known **not admitted to that queue**. Previously +queued methods, interrupted wire sends, and execution calls which have begun may +have unknown delivery or side effects. They are never converted to success or a +claim of non-execution, and are never automatically retried. Cancel is subject to +this same rule: it either enters the bounded queue or the connection explicitly +ends; there is no silent special-case drop or unbounded control lane. + +## Transport deadlines and termination + +Remote handshake and sends, browser writes (including Pong/error), and Workspace +multiplexer writes have 10-second deadlines. The dedicated browser session makes +one separately bounded Close attempt (at most another 10 seconds), with code 1011 +and a short reason requiring snapshot resynchronization and warning of unknown +pending delivery. A completely blocked wire cannot guarantee delivery of the +Close/error frame; releasing the socket is the terminal signal in that case. + +The multiplexer retains separate bounded control and high-volume protocol lanes +and control priority. Writer termination cancels the receive/admission loop, even +if that loop is awaiting control capacity or session ownership. Receive exit +aborts/joins the writer and subscription tasks instead of waiting indefinitely +for a slow writer to drain. Caller cancellation has abort-on-drop ownership. +Bridge overflow is exposed as `SubscriptionClosed { code: Lagged }`, not normal +resource disappearance, when the terminal control frame can be delivered. + +Embedded dispatch is an existing synchronous Runtime API, not a new database or +executor redesign. Its connected-method adapter wait uses the existing +10-second cancellation-and-join waiter rather than leaving a detached send after +timeout. Existing completion and lifecycle ownership/fencing are retained. A +synchronous embedded dispatch already begun is not immediately preemptible by +browser disconnect; its bounded adapter wait settles before endpoint exclusion +is released. Cancellation/join is resource reclamation, **not rollback of an +already delivered method**. + +## Recovery + +Reconnect explicitly, resubscribe where using the multiplexer, and replace local +state with the fresh snapshot before live delivery (T-415). Do not replay old +method queues or transient protocol deltas. An ended connection never implicitly +Restores, reattaches, or retargets the Worker. Slow connections do not hold Server +session exclusion during transport waits, so other connections and Stop/session +cleanup can proceed. + +## Evidence and limits + +The initial static audit used `develop@a9a67fd6`; the same wait edges were present +at `develop@98023b98`. This is not evidence of a production incident. Controlled +fixtures exercise the actual remote connector, embedded bridge, browser method +admission/session, and Workspace multiplexer. Reinstating only the old bounded +queue waits makes the remote/embedded overflow and browser session-exclusion +fixtures time out; the corrected paths terminate without guessed sleeps. + +Fixtures use channel/barrier ownership, controllable sinks, loopback WebSockets, +and paused time for transport deadlines. They cover ordered admitted prefixes, +auth rejection before admission, Cancel overflow, pending admission cancellation, +unknown wire delivery without retry, same-Worker fresh connection/Stop, and +transport/channel reclamation. They are not full browser/Runtime process E2E or a +production load benchmark. Existing attach/stop/restore and serialization tests +remain required alongside these regressions. + +### Implementation validation (T-747) + +- `cargo check --offline` at Workspace root: passed (normal TUI/Server closure). +- `cargo test -p yoi-workspace-server --offline --lib protocol_`: 29 passed, + including the real bridges, browser admission/ownership, and stalled multiplexer. +- `cargo test -p yoi-workspace-server --offline`: 867 library tests passed, one + existing ignored test, 11 binary tests passed, including the lag-code wire + classification matrix (`--lib workspace_subscription::tests`: 4 passed). +- `cargo test -p worker-runtime --offline`: 416 library and 19 tests in each of + the two binaries passed, including connected admission cancellation/join and + existing endpoint fencing/lifecycle regressions. +- `cargo test -p protocol --offline`: 80 passed (wire contracts unchanged). +- `cargo fmt --all -- --check`, `git diff --check HEAD`, and + `env -u LD_LIBRARY_PATH deno run --allow-read --allow-run=git tools/check-state-contract-names.ts`: + passed; naming checker reported 0 unapproved / 0 stale exceptions. + +Additional Runtime compile probes with `--no-default-features`, with only +`http-server,fs-store`, and with only `ws-server` failed in pre-existing feature +boundaries (respectively missing gated request dependencies, missing futures +imports/dependency, and missing filesystem-backed retention methods). Replacing +only the changed Runtime file with its exact `98023b98` content produced the same +error groups in all three probes. These configurations are not claimed validated; +this patch does not repair their unrelated feature wiring. The normal +`ws-server,fs-store` configuration is covered by root check and whole-crate tests. +