feat: feat(cluster-protocol): add the Rust WebSocket binding and cloud data-plane contract

crates/openengine-cluster-server/src/websocket.rs

@@ -4,7 +4,7 @@
 //! ([`crate::stdio::ConnectionState`], `classify_ndjson_line`, `dispatch_classified_line`) so
 //! results, events, and errors stay byte-equivalent between the stdio and WebSocket bindings.
 //! Framing rules unique to WebSocket -- binary rejection, the 1,048,576 UTF-8 byte bound, and
-//! best-effort `$/cancelRequest` -- live only here; `stdio::serve_ndjson` itself is untouched.
+//! race-free `$/cancelRequest` -- live only here; `stdio::serve_ndjson` itself is untouched.
 
 use std::collections::HashMap;
 use std::io;
@@ -36,9 +36,12 @@ pub const MAX_FRAME_BYTES: usize = 1_048_576;
 /// Bounded per-connection outbound queue, matching `stdio::serve_ndjson`'s bound.
 const OUTBOUND_QUEUE_CAPACITY: usize = 256;
 
-/// Per-connection registry of best-effort abort handles for in-flight passthrough (unary) request
-/// tasks, keyed by their [`RequestId`]. `$/cancelRequest` looks an id up here and aborts it; an
-/// absent id (unknown, or already completed and self-removed) is silently a no-op. Establishing
+/// Per-connection registry of abort handles for in-flight passthrough (unary) request tasks, keyed
+/// by their [`RequestId`]. Registration happens-before the spawned task can perform any work --
+/// including completing -- via `spawn_passthrough`'s readiness gate, and release is unconditional
+/// on every exit path (normal completion, cancellation, or panic) via [`PassthroughGuard`]'s
+/// `Drop` impl, so no stale or leaked entry is possible. `$/cancelRequest` looks an id up here and
+/// aborts it; an absent id (unknown, or already completed) is silently a no-op. Establishing
 /// requests (`watch`/`logs`/`agent/attach`) are not registered here -- their in-flight lifetime is
 /// already covered by `subscription/cancel` once established, and is intentionally short.
 type CancelRegistry = Arc<Mutex<HashMap<RequestId, AbortHandle>>>;
@@ -53,7 +56,7 @@ pub fn websocket_config() -> WebSocketConfig {
 
 /// Serves one already-handshaken WebSocket connection: demultiplexes unary requests and `watch`/
 /// `logs`/`agent/attach` subscriptions sharing this connection exactly like `stdio::serve_ndjson`,
-/// plus per-connection best-effort `$/cancelRequest`. Binary frames close with code 1003;
+/// plus per-connection race-free `$/cancelRequest`. Binary frames close with code 1003;
 /// oversized or capacity-rejected text frames close with code 1009. Never returns an `Err`:
 /// transport failures close the connection and this simply returns once torn down.
 pub async fn serve_websocket<B, S>(
@@ -144,7 +147,7 @@ enum FrameOutcome {
 }
 
 /// Per-frame handling context: the shared [`DispatchCtx`] plus the two handles unique to this
-/// WebSocket binding (the best-effort `$/cancelRequest` registry and the deterministic-close
+/// WebSocket binding (the race-free `$/cancelRequest` registry and the deterministic-close
 /// signal), bundled so [`handle_message`] and its helpers take one argument instead of an
 /// ever-growing list.
 struct WsCtx<'a, B> {
@@ -202,8 +205,22 @@ where
     FrameOutcome::Continue
 }
 
-/// Spawns the passthrough (non-subscription) request task for an admission-approved line, best-
-/// effort registering its abort handle under `id` in `cancel_registry` for `$/cancelRequest`.
+/// Spawns the passthrough (non-subscription) request task for an admission-approved line,
+/// registering its abort handle under `id` in `cancel_registry` for `$/cancelRequest`. The spawned
+/// task cannot perform any work -- including completing -- until the caller has finished inserting
+/// that handle: a request that carries an `id` gates its own start on a one-shot readiness signal
+/// the caller only fires after the insert, making registration happen-before any task work
+/// deterministically regardless of which worker thread the scheduler picks it up on. A bare
+/// notification (`id.is_none()`) has nothing to register and keeps starting immediately.
+///
+/// [`PassthroughGuard`] is constructed here -- synchronously, on the connection task -- and moved
+/// into the spawned future's captured environment rather than built inside the task's own body.
+/// This matters because a task can be aborted before the runtime ever polls it even once (e.g. a
+/// `$/cancelRequest` that lands while still parked behind the readiness gate above): such a task
+/// never executes a single statement of its own body, so a guard created *inside* that body would
+/// never come into existence and cleanup would silently never run. A value captured by an `async
+/// move` block, by contrast, is part of the future's state from the moment the block is
+/// constructed, so Rust drops it normally when the never-polled future itself is dropped.
 fn spawn_passthrough<B>(
     ctx: &mut WsCtx<'_, B>,
     id: Option<RequestId>,
@@ -215,56 +232,90 @@ fn spawn_passthrough<B>(
     let task_dispatcher = ctx.dispatch.dispatcher.clone();
     let task_state = ctx.dispatch.state.clone();
     let task_cancel_registry = Arc::clone(ctx.cancel_registry);
-    let spawn_id = id.clone();
+    let guard = id.clone().map(|id| PassthroughGuard {
+        id,
+        state: task_state.clone(),
+        cancel_registry: Arc::clone(&task_cancel_registry),
+    });
+    let (ready_tx, ready_rx) = if id.is_some() {
+        let (tx, rx) = oneshot::channel::<()>();
+        (Some(tx), Some(rx))
+    } else {
+        (None, None)
+    };
     let abort_handle = ctx.dispatch.tasks.spawn(async move {
+        if let Some(ready_rx) = ready_rx {
+            let _ = ready_rx.await;
+        }
         let _permit = permit;
         run_passthrough_request(PassthroughRequest {
             dispatcher: task_dispatcher,
-            id,
             line,
             state: task_state,
-            cancel_registry: task_cancel_registry,
+            guard,
         })
         .await;
     });
-    // Best-effort: a pathologically fast dispatch can self-remove (see `run_passthrough_request`)
-    // before this insert runs, leaving a harmless stale entry -- `$/cancelRequest` never claims a
-    // rollback guarantee, and the entry is bounded by this connection's lifetime and task-slot cap.
-    if let Some(spawn_id) = spawn_id {
-        ctx.cancel_registry.lock().insert(spawn_id, abort_handle);
+    // Test-only: widens the window between the spawned task being handed to the scheduler and
+    // this registration completing, so `tests::fast_completion_race_cannot_leave_stale_registry_entry`
+    // can deterministically manufacture -- rather than hope to get lucky on -- the exact adverse
+    // interleaving the readiness gate above exists to close. `cfg(test)` only activates while
+    // compiling this crate's own lib unit tests, so this never runs in a real binary or in
+    // `tests/websocket.rs`'s integration coverage (a separate crate that links this one normally).
+    #[cfg(test)]
+    std::thread::sleep(std::time::Duration::from_millis(2));
+    if let Some(id) = id {
+        ctx.cancel_registry.lock().insert(id, abort_handle);
+    }
+    if let Some(ready_tx) = ready_tx {
+        let _ = ready_tx.send(());
     }
 }
 
 /// Grouped arguments for [`run_passthrough_request`], keeping that function's argument count
 /// reasonable.
 struct PassthroughRequest<B> {
     dispatcher: Dispatcher<B>,
-    id: Option<RequestId>,
     line: String,
     state: ConnectionState,
+    guard: Option<PassthroughGuard>,
+}
+
+/// Releases a passthrough request's `in_flight_ids` and `cancel_registry` entries exactly once, on
+/// every exit path -- including `$/cancelRequest`'s task abortion, which drops this future in
+/// place (whether at its current await point, or -- if aborted before ever being polled -- as part
+/// of dropping the future's captured-but-never-executed initial state) instead of resuming past
+/// it, so cleanup cannot depend on falling off the end of `run_passthrough_request`.
+struct PassthroughGuard {
+    id: RequestId,
+    state: ConnectionState,
     cancel_registry: CancelRegistry,
 }
 
+impl Drop for PassthroughGuard {
+    fn drop(&mut self) {
+        self.state.in_flight_ids.lock().remove(&self.id);
+        self.cancel_registry.lock().remove(&self.id);
+    }
+}
+
 /// Dispatches a non-subscription request or notification frame, releasing its in-flight id and
-/// best-effort cancel registration (if any) once the backend call returns and before the response
-/// is enqueued. Mirrors `stdio::run_passthrough_request`, plus the `cancel_registry` cleanup that
-/// binding has no notion of.
+/// cancel registration (if any, via the already-constructed [`PassthroughGuard`]) once the backend
+/// call returns and before the response is enqueued -- unconditionally, even if the task is
+/// aborted mid-flight. Mirrors `stdio::run_passthrough_request`, plus the `cancel_registry` cleanup
+/// that binding has no notion of.
 async fn run_passthrough_request<B>(request: PassthroughRequest<B>)
 where
     B: ClusterBackend,
 {
     let PassthroughRequest {
         dispatcher,
-        id,
         line,
         state,
-        cancel_registry,
+        guard,
     } = request;
     let response = dispatcher.dispatch(&line).await;
-    if let Some(id) = &id {
-        state.in_flight_ids.lock().remove(id);
-        cancel_registry.lock().remove(id);
-    }
+    drop(guard);
     let _ = state.outbound_tx.send(response).await;
 }
 
@@ -331,3 +382,6 @@ async fn run_writer<Si>(
     }
     let _ = sink.close().await;
 }
+
+#[cfg(test)]
+mod tests;

crates/openengine-cluster-server/src/websocket/tests.rs

@@ -0,0 +1,105 @@
+use std::time::Duration;
+
+use openengine_cluster_protocol::RunId;
+use serde_json::{json, Value};
+
+use super::*;
+use crate::stdio::{
+    classify_ndjson_line, dispatch_classified_line, new_connection_setup, ConnectionSetup,
+    DispatchCtx, LineDispatch,
+};
+use crate::watch::fixtures::{FixtureBackend, FixtureStore};
+use crate::ConnectionContext;
+
+/// Regression test for the spawn-before-register race in `spawn_passthrough`: pre-fix, the newly
+/// spawned passthrough task starts running immediately once handed to the scheduler and, once its
+/// (near-instant, in-memory) dispatch resolves, removes its own `cancel_registry` entry as part of
+/// normal cleanup -- a no-op if that happens before the caller's `insert` for the same id has run.
+/// Nothing in the pre-fix code establishes a happens-before relationship between "the task is
+/// schedulable" and "the caller's insert completes", so if the removal is observed first, the
+/// caller's later insert leaves a stale handle for an already-finished request that nothing will
+/// ever clean up.
+///
+/// Real scheduling alone essentially never reproduces this: the caller attempts its insert within
+/// nanoseconds of spawning (no other work in between), while the spawned task needs a genuine
+/// cross-thread scheduling round-trip before it can even begin dispatch, so the safe
+/// insert-then-remove ordering dominates overwhelmingly by default (independently confirmed by
+/// running an all-real-timing version of this test hundreds of times against the pre-fix code
+/// without once observing the race). To make the test actually exercise the interleaving the issue
+/// requires proving fixed, `spawn_passthrough` widens that window under `#[cfg(test)]` only (a
+/// no-op in every real binary and in `tests/websocket.rs`'s separate-crate integration coverage),
+/// giving the spawned task's cleanup a deterministic chance to run first.
+///
+/// Post-fix, `spawn_passthrough`'s readiness gate means the spawned task cannot attempt *any* work
+/// -- including reaching its own cleanup -- until a readiness signal fires, which is sent only
+/// after the caller's insert has already completed. Widening the same window therefore cannot move
+/// the outcome post-fix: the spawned task is still parked behind the gate regardless of how long
+/// the caller takes to reach its insert, so registration remains happens-before any request work by
+/// construction. That asymmetry -- reliably broken pre-fix, structurally unbreakable post-fix -- is
+/// exactly what this test proves.
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+async fn fast_completion_race_cannot_leave_stale_registry_entry() {
+    let store = Arc::new(FixtureStore::new(RunId::new("run-1"), Vec::new(), 8));
+    let dispatcher = Dispatcher::new(FixtureBackend::new(store), ConnectionContext::default());
+    let (outbound_tx, mut outbound_rx) = mpsc::channel::<String>(256);
+    let ConnectionSetup {
+        task_slots,
+        mut tasks,
+        state,
+        ..
+    } = new_connection_setup(&outbound_tx);
+    let cancel_registry: CancelRegistry = Arc::new(Mutex::new(HashMap::new()));
+
+    const TRIALS: i64 = 20;
+    for trial in 0..TRIALS {
+        let id = RequestId::Integer(trial);
+        let line =
+            json!({"jsonrpc": "2.0", "id": trial, "method": "get", "params": {}}).to_string();
+
+        let kind = classify_ndjson_line(&line);
+        let mut dispatch_ctx = DispatchCtx {
+            dispatcher: &dispatcher,
+            state: &state,
+            task_slots: &task_slots,
+            tasks: &mut tasks,
+        };
+        let LineDispatch::Passthrough { id: got_id, permit } =
+            dispatch_classified_line(kind, &mut dispatch_ctx).await
+        else {
+            panic!("expected a passthrough dispatch for a `get` request");
+        };
+        assert_eq!(got_id.as_ref(), Some(&id));
+
+        let mut close_tx: Option<oneshot::Sender<CloseFrame>> = None;
+        let mut ctx = WsCtx {
+            dispatch: DispatchCtx {
+                dispatcher: &dispatcher,
+                state: &state,
+                task_slots: &task_slots,
+                tasks: &mut tasks,
+            },
+            cancel_registry: &cancel_registry,
+            close_tx: &mut close_tx,
+        };
+        spawn_passthrough(&mut ctx, got_id, permit, line);
+
+        // Drain this trial's task to completion (its post-dispatch cleanup, whichever way the
+        // race landed, has therefore already run) before inspecting registry state.
+        while tasks.join_next().await.is_some() {}
+
+        let response = tokio::time::timeout(Duration::from_secs(1), outbound_rx.recv())
+            .await
+            .expect("the request must receive its response promptly")
+            .expect("outbound channel must not close mid-test");
+        let response: Value = serde_json::from_str(&response).unwrap();
+        assert_eq!(response["id"], trial);
+        assert!(response.get("result").is_some(), "{response}");
+
+        assert!(
+            cancel_registry.lock().is_empty(),
+            "trial {trial}: a stale cancel_registry entry survived a fast completion racing the \
+             caller's registration -- insert is not guaranteed to happen-before the spawned task \
+             can observe/complete its own cleanup"
+        );
+    }
+}

crates/openengine-cluster-server/tests/websocket.rs

@@ -216,6 +216,107 @@ impl GatedHarnessSpawn for Harness {
     }
 }
 
+#[tokio::test]
+async fn cancel_request_releases_in_flight_id_and_permits_id_reuse() {
+    let (mut harness, gate) = spawn_gated_harness::<Harness>().await;
+
+    // Blocks on the gate: this request never completes, so its `in_flight_ids`/`cancel_registry`
+    // entries can only be released via the cancellation path, never via normal fall-through.
+    harness.send_get(1).await;
+
+    let cancel = json!({
+        "jsonrpc": "2.0",
+        "method": "$/cancelRequest",
+        "params": {"id": 1},
+    })
+    .to_string();
+    send_text(&mut harness.client, cancel).await;
+
+    // `AbortHandle::abort()` only schedules the target task's cancellation; the guard that
+    // releases `in_flight_ids` runs whenever the runtime actually gets around to polling (and
+    // dropping) that task, not synchronously when `abort()` returns. Poll for the id becoming
+    // reusable via bounded cooperative retries rather than a fixed sleep: a duplicate rejection is
+    // a harmless, synchronous no-op (it never touches `in_flight_ids`, so resending is safe), and
+    // it always arrives promptly since it never waits on the still-closed gate -- whereas an
+    // *accepted* retry's task immediately blocks on that same still-closed gate exactly like the
+    // original did, so it produces no prompt response at all. That absence, not its content,
+    // is what distinguishes "admitted" from "still rejected" here, deterministically regardless
+    // of scheduler load. The attempt cap still fails the test if the id is never released -- the
+    // exact pre-fix bug (a cancelled request permanently reserving its id).
+    let mut accepted = false;
+    for attempt in 0..100_000 {
+        harness.send_get(1).await;
+        match tokio::time::timeout(Duration::from_millis(50), harness.recv_value()).await {
+            Ok(response) => {
+                assert_eq!(response["id"], 1);
+                assert_eq!(
+                    response["error"]["data"]["code"], "DUPLICATE_REQUEST_ID",
+                    "a prompt response before the gate is released can only be a duplicate \
+                     rejection (a real result would block on the still-closed gate); got \
+                     {response}"
+                );
+                assert!(
+                    attempt < 99_999,
+                    "id 1 was never released after cancellation settled -- it is still \
+                     permanently reserved in in_flight_ids"
+                );
+                tokio::task::yield_now().await;
+            }
+            Err(_) => {
+                // No prompt response: this retry was admitted and its task is now parked on the
+                // gate, exactly like the original was before cancellation.
+                accepted = true;
+                break;
+            }
+        }
+    }
+    assert!(accepted, "id 1 was never accepted for reuse");
+
+    // Releasing the gate now only unblocks the accepted retry's still-pending backend call; no
+    // response frame was ever emitted for the original cancelled request, so this is the only
+    // frame left to receive.
+    gate.notify_one();
+    let second = tokio::time::timeout(Duration::from_secs(1), harness.recv_value())
+        .await
+        .expect("the accepted retry must complete once the gate is released");
+    assert_eq!(second["id"], 1);
+    assert!(second.get("result").is_some(), "{second}");
+
+    shut_down(harness).await;
+}
+
+#[tokio::test]
+async fn fast_completions_do_not_leak_or_corrupt_cancel_registry() {
+    let store = Arc::new(FixtureStore::new(RunId::new("run-1"), Vec::new(), 8));
+    let mut harness = spawn_server(FixtureBackend::new(store)).await;
+
+    // Ungated backend: each `get` resolves as fast as the runtime allows. Repeatedly reusing the
+    // same id and firing a `$/cancelRequest` immediately after each completion -- when the id is
+    // already fully released and the notification can only ever hit an unknown-or-completed id --
+    // is a high-volume regression/sanity net for that exact "silent no-op" contract and for the
+    // general fast-completion-plus-cancellation interaction: every iteration must still receive
+    // its own correct, uncorrupted response and the connection must never misbehave (unexpected
+    // close, mismatched response, panic, deadlock) across 200 rapid cycles.
+    for _ in 0..200 {
+        send_text(&mut harness.client, request_text(1, "get", json!({}))).await;
+        let response = tokio::time::timeout(Duration::from_secs(1), recv_json(&mut harness.client))
+            .await
+            .expect("each fast completion must receive its own uncorrupted response");
+        assert_eq!(response["id"], 1);
+        assert!(response.get("result").is_some(), "{response}");
+
+        let cancel = json!({
+            "jsonrpc": "2.0",
+            "method": "$/cancelRequest",
+            "params": {"id": 1},
+        })
+        .to_string();
+        send_text(&mut harness.client, cancel).await;
+    }
+
+    shut_down(harness).await;
+}
+
 #[tokio::test]
 async fn duplicate_in_flight_request_ids_are_rejected() {
     let (mut harness, gate) = spawn_gated_harness::<Harness>().await;