fix(server): bound embedded protocol and Feature shutdown ownership

crates/workspace-server/src/server.rs

@@ -2651,16 +2651,7 @@ impl WorkspaceServerApi {
     async fn shutdown_feature_storage(&self) -> Result<()> {
         let apis = self.apis.lock().await.values().cloned().collect::<Vec<_>>();
         for api in apis {
-            api.worker_projection.shutdown_async().await;
-            let storage = api.feature_storage.clone();
-            api.store
-                .db_executor()
-                .run(move || {
-                    storage.shutdown().map_err(|error| {
-                        Error::Store(format!("failed to shut down Feature storage: {error}"))
-                    })
-                })
-                .await?;
+            shutdown_workspace_feature_storage(&api).await?;
         }
         Ok(())
     }
@@ -12271,15 +12262,28 @@ pub async fn serve(
     let orchestrator_hook = tokio::spawn(run_orchestrator_turn_end_hook(api.clone()));
     let result = axum::serve(listener, build_router(api.clone())).await;
     orchestrator_hook.abort();
-    let storage_result = api
-        .feature_storage
-        .shutdown()
-        .map_err(|error| Error::Store(format!("failed to shut down Feature storage: {error}")));
+    let _ = orchestrator_hook.await;
+    let storage_result = shutdown_workspace_feature_storage(&api).await;
     result?;
     storage_result?;
     Ok(())
 }
 
+// Feature shutdown fences new operations and waits for synchronous connection
+// holders. Accepted close retains ownership after its async waiter disappears.
+async fn shutdown_workspace_feature_storage(api: &WorkspaceApi) -> Result<()> {
+    api.worker_projection.shutdown_async().await;
+    let storage = api.feature_storage.clone();
+    api.store
+        .db_executor()
+        .run(move || {
+            storage.shutdown().map_err(|error| {
+                Error::Store(format!("failed to shut down Feature storage: {error}"))
+            })
+        })
+        .await
+}
+
 #[derive(Debug, Serialize, Deserialize)]
 pub struct ProjectListResponse<T> {
     pub workspace_id: String,
@@ -31749,7 +31753,7 @@ pub(crate) async fn connect_workspace_worker_protocol(
         RuntimeObservationSource::Embedded(source) => {
             let scope =
                 worker_runtime::RuntimeWorkspaceScope::new(api.workspace_id(), "embedded-backend");
-            connect_embedded_worker_protocol(source, scope).await
+            connect_embedded_worker_protocol(source, scope, api.host_executors.host()?).await
         }
     }
 }
@@ -31845,18 +31849,25 @@ async fn connect_remote_worker_protocol(
 async fn connect_embedded_worker_protocol(
     source: crate::observation::EmbeddedRuntimeObservationSource,
     scope: worker_runtime::RuntimeWorkspaceScope,
+    executor: crate::db_executor::DbExecutor,
 ) -> Result<WorkspaceWorkerProtocolConnection> {
-    // Capture one real execution endpoint, including its greeting/snapshot.
-    // The WorkerRef-wide Workspace observation bus is a separate catalogue
-    // channel and must not reconnect this operational session after Restore.
-    let mut transport = source
-        .runtime
-        .attach_worker_protocol_scoped(&scope, &source.worker_ref)
-        .map_err(|error| Error::RuntimeOperationFailed {
-            runtime_id: source.worker.runtime_id.clone(),
-            code: "embedded_worker_protocol_connect_failed".to_string(),
-            message: error.to_string(),
-        })?;
+    // Endpoint capture takes Runtime operation/state locks. Keep it, and each
+    // intact method (including fs-store commit/publication), off Tokio and the
+    // sole DB lane. A discarded reply retains the endpoint until cleanup.
+    let capture_source = source.clone();
+    let capture_scope = scope.clone();
+    let mut transport = executor
+        .run(move || {
+            capture_source
+                .runtime
+                .attach_worker_protocol_scoped(&capture_scope, &capture_source.worker_ref)
+                .map_err(|error| Error::RuntimeOperationFailed {
+                    runtime_id: capture_source.worker.runtime_id.clone(),
+                    code: "embedded_worker_protocol_connect_failed".to_string(),
+                    message: error.to_string(),
+                })
+        })
+        .await?;
     Ok(spawn_worker_protocol_pump(
         move |mut methods, events| async move {
             if let Err(reason) = enqueue_worker_protocol_event(&events, transport.snapshot.clone())
@@ -31868,9 +31879,26 @@ async fn connect_embedded_worker_protocol(
                     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,
-                        ) {
+                        let method_source = source.clone();
+                        let method_scope = scope.clone();
+                        let owned = executor.run(move || {
+                            let result = method_source.runtime.send_connected_protocol_method_scoped(
+                                &method_scope, &method_source.worker_ref, &transport, method,
+                            );
+                            // Endpoint validation is synchronous too.
+                            let terminal = result.as_ref().err().is_some_and(|error|
+                                embedded_protocol_dispatch_is_terminal(error, &transport));
+                            Ok::<_, Error>((transport, result, terminal))
+                        }).await;
+                        let (returned, result, terminal) = match owned {
+                            Ok(value) => value,
+                            Err(error) => {
+                                let _ = enqueue_worker_protocol_event(&events, protocol_error_event(error.to_string()));
+                                return WorkerProtocolEnd::Disconnected;
+                            }
+                        };
+                        transport = returned;
+                        match result {
                             Ok(direct_events) => {
                                 for event in direct_events {
                                     let terminal = matches!(event, protocol::Event::Shutdown);
@@ -31879,7 +31907,6 @@ async fn connect_embedded_worker_protocol(
                                 }
                             }
                             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; }
                             }
@@ -31892,7 +31919,6 @@ async fn connect_embedded_worker_protocol(
                             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,
                     }
                 }
@@ -35184,6 +35210,7 @@ mod tests {
     mod drive_api_tests;
     mod drive_tests;
     mod drive_web_roundtrip_tests;
+    mod embedded_owned_execution_tests;
     mod external_checkout_tests;
     mod protocol_backpressure_tests;
     mod subject_spawn_tests;
@@ -45338,6 +45365,7 @@ mod tests {
     // The receiver belongs to one mock execution, not to a Runtime WorkerRecord.
     // Bound protocol senders become terminal when Stop drops this execution.
     struct DeterministicWorkerExecution {
+        dispatch_hook: std::sync::Mutex<Option<Box<dyn FnOnce() + Send>>>,
         context: worker_runtime::execution::WorkerExecutionContext,
         state: std::sync::Mutex<protocol::WorkerStateSnapshot>,
         method_tx: tokio::sync::mpsc::Sender<protocol::Method>,
@@ -45357,6 +45385,7 @@ mod tests {
             let (method_tx, method_rx) = tokio::sync::mpsc::channel(32);
             let (events, _) = tokio::sync::broadcast::channel(32);
             Self {
+                dispatch_hook: std::sync::Mutex::new(None),
                 context,
                 state: std::sync::Mutex::new(protocol::WorkerStateSnapshot::initial()),
                 method_tx,
@@ -45380,6 +45409,9 @@ mod tests {
             Vec<protocol::Event>,
             worker_runtime::execution::WorkerExecutionResult,
         > {
+            if let Some(hook) = self.dispatch_hook.lock().unwrap().take() {
+                hook();
+            }
             // Drive the actual mock receiver synchronously; no method is routed
             // through a WorkerRef lookup after a protocol connection is attached.
             let mut receiver = self.method_rx.lock().unwrap();
@@ -61492,10 +61524,13 @@ mod tests {
         for delayed in delayed_protocol_test_methods() {
             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();
+            let mut old = connect_embedded_worker_protocol(
+                source.clone(),
+                embedded_protocol_test_scope(),
+                HostExecutors::default().host().unwrap(),
+            )
+            .await
+            .unwrap();
             assert_embedded_protocol_snapshot(
                 next_embedded_protocol_event(&mut old).await,
                 &worker_ref,
@@ -61559,10 +61594,13 @@ mod tests {
                 &[worker_ref.clone()]
             );
 
-            let mut current =
-                connect_embedded_worker_protocol(source, embedded_protocol_test_scope())
-                    .await
-                    .unwrap();
+            let mut current = connect_embedded_worker_protocol(
+                source,
+                embedded_protocol_test_scope(),
+                HostExecutors::default().host().unwrap(),
+            )
+            .await
+            .unwrap();
             assert_embedded_protocol_snapshot(
                 next_embedded_protocol_event(&mut current).await,
                 &worker_ref,
@@ -61710,10 +61748,13 @@ mod tests {
         for lagged in [false, true] {
             let (runtime, worker_ref, execution) = runtime_with_worker_backend();
             let source = embedded_protocol_test_source(&runtime, &worker_ref);
-            let mut bridge =
-                connect_embedded_worker_protocol(source.clone(), embedded_protocol_test_scope())
-                    .await
-                    .unwrap();
+            let mut bridge = connect_embedded_worker_protocol(
+                source.clone(),
+                embedded_protocol_test_scope(),
+                HostExecutors::default().host().unwrap(),
+            )
+            .await
+            .unwrap();
             assert_embedded_protocol_snapshot(
                 next_embedded_protocol_event(&mut bridge).await,
                 &worker_ref,
@@ -61755,6 +61796,7 @@ mod tests {
         let mut bridge = connect_embedded_worker_protocol(
             embedded_protocol_test_source(&runtime, &worker_ref),
             embedded_protocol_test_scope(),
+            HostExecutors::default().host().unwrap(),
         )
         .await
         .unwrap();
@@ -61815,9 +61857,13 @@ mod tests {
         let wrong_scope =
             worker_runtime::RuntimeWorkspaceScope::new("other-workspace", "local-token");
         assert!(
-            connect_embedded_worker_protocol(source.clone(), wrong_scope.clone())
-                .await
-                .is_err()
+            connect_embedded_worker_protocol(
+                source.clone(),
+                wrong_scope.clone(),
+                HostExecutors::default().host().unwrap()
+            )
+            .await
+            .is_err()
         );
         let (mut socket, server) =
             spawn_embedded_protocol_test_ws(source.clone(), wrong_scope, false).await;

crates/workspace-server/src/server/tests/embedded_owned_execution_tests.rs

@@ -0,0 +1,197 @@
+use super::*;
+use crate::db_executor::{
+    DbExecutor, DbExecutorLimits, DbOperationLifetime, with_operation_lifetime,
+};
+use std::{
+    num::NonZeroUsize,
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        mpsc,
+    },
+    time::Duration,
+};
+
+fn owner() -> DbExecutor {
+    DbExecutor::new(DbExecutorLimits {
+        workers: NonZeroUsize::new(1).unwrap(),
+        queued: NonZeroUsize::new(1).unwrap(),
+    })
+    .unwrap()
+}
+async fn observed<T>(future: impl std::future::Future<Output = T>) -> T {
+    tokio::time::timeout(Duration::from_secs(2), future)
+        .await
+        .expect("owned protocol stage did not progress")
+}
+struct Release(Option<mpsc::Sender<()>>);
+impl Release {
+    fn release(&mut self) {
+        if let Some(tx) = self.0.take() {
+            let _ = tx.send(());
+        }
+    }
+}
+impl Drop for Release {
+    fn drop(&mut self) {
+        self.release();
+    }
+}
+struct Lifetime(Arc<AtomicUsize>);
+struct Retained(Arc<AtomicUsize>);
+impl Drop for Retained {
+    fn drop(&mut self) {
+        self.0.fetch_sub(1, Ordering::SeqCst);
+    }
+}
+impl DbOperationLifetime for Lifetime {
+    fn retain(&self) -> Box<dyn Send> {
+        self.0.fetch_add(1, Ordering::SeqCst);
+        Box::new(Retained(self.0.clone()))
+    }
+}
+
+#[tokio::test(flavor = "current_thread")]
+async fn cancelled_embedded_endpoint_capture_keeps_actual_host_budget_off_tokio() {
+    let (runtime, worker, execution) = runtime_with_worker_backend();
+    let (started, held) = tokio::sync::oneshot::channel();
+    let (tx, rx) = mpsc::channel();
+    let mut release = Release(Some(tx));
+    let thread = std::thread::spawn(move || {
+        let _guard = execution.contexts.lock().unwrap();
+        let _ = started.send(());
+        let _ = rx.recv_timeout(Duration::from_secs(5));
+    });
+    observed(held).await.unwrap();
+    let executor = owner();
+    let leases = Arc::new(AtomicUsize::new(0));
+    let mut capture = Box::pin(with_operation_lifetime(
+        Arc::new(Lifetime(leases.clone())),
+        connect_embedded_worker_protocol(
+            embedded_protocol_test_source(&runtime, &worker),
+            embedded_protocol_test_scope(),
+            executor.clone(),
+        ),
+    ));
+    assert!(
+        futures::poll!(&mut capture).is_pending(),
+        "endpoint Runtime lock must not run on Tokio"
+    );
+    observed(tokio::time::sleep(Duration::from_millis(1))).await;
+    assert_eq!(
+        observed(tokio::spawn(async { "control" })).await.unwrap(),
+        "control"
+    );
+    drop(capture);
+    assert_eq!(leases.load(Ordering::SeqCst), 1);
+    executor.dispatch(|| {}).unwrap(); // only the one queued slot is left
+    let error = executor.run(|| Ok::<_, Error>(())).await.unwrap_err();
+    assert!(matches!(error, Error::DbCapacityExceeded));
+    release.release();
+    thread.join().unwrap();
+    observed(executor.drain()).await.unwrap();
+    assert_eq!(leases.load(Ordering::SeqCst), 0);
+}
+
+#[tokio::test(flavor = "current_thread")]
+async fn cancelled_embedded_shutdown_commits_fs_receipt_after_the_pump_disappears() {
+    let dir = tempfile::tempdir().unwrap();
+    let execution = Arc::new(DeterministicExecutionBackend::default());
+    let options = worker_runtime::fs_store::FsRuntimeStoreOptions {
+        root: dir.path().join("runtime"),
+        runtime_id: "owned-protocol".into(),
+        display_name: None,
+    };
+    let runtime = worker_runtime::Runtime::with_fs_store_and_execution_backend(
+        options.clone(),
+        execution.clone(),
+    )
+    .unwrap();
+    runtime.store_config_bundle(runtime_test_bundle()).unwrap();
+    let worker = runtime
+        .create_worker_scoped(&embedded_protocol_test_scope(), runtime_create_request())
+        .unwrap()
+        .worker_ref;
+    let executor = owner();
+    let mut connection = connect_embedded_worker_protocol(
+        embedded_protocol_test_source(&runtime, &worker),
+        embedded_protocol_test_scope(),
+        executor.clone(),
+    )
+    .await
+    .unwrap();
+    assert!(matches!(
+        observed(connection.events.recv()).await.unwrap(),
+        protocol::Event::Snapshot { .. }
+    ));
+    let (entered, held) = tokio::sync::oneshot::channel();
+    let (tx, rx) = mpsc::channel();
+    let mut release = Release(Some(tx));
+    let endpoint = execution
+        .contexts
+        .lock()
+        .unwrap()
+        .get(&worker)
+        .unwrap()
+        .clone();
+    let gate_finished = Arc::new(std::sync::atomic::AtomicBool::new(false));
+    let finished = gate_finished.clone();
+    *endpoint.dispatch_hook.lock().unwrap() = Some(Box::new(move || {
+        let _ = entered.send(());
+        let _ = rx.recv_timeout(Duration::from_secs(5));
+        finished.store(true, Ordering::SeqCst);
+    }));
+    connection
+        .methods
+        .send(protocol::Method::Shutdown {
+            command: protocol::WorkerCommandEnvelope::new(700),
+        })
+        .await
+        .unwrap();
+    observed(held).await.unwrap(); // actual Runtime method, after its durable stop admission
+    observed(tokio::time::sleep(Duration::from_millis(1))).await;
+    assert_eq!(
+        observed(tokio::spawn(async { "control" })).await.unwrap(),
+        "control"
+    );
+    assert!(
+        !gate_finished.load(Ordering::SeqCst),
+        "watchdog bailout is not Tokio progress"
+    );
+    connection
+        .termination
+        .terminate(WorkerProtocolEnd::Disconnected);
+    drop(connection);
+    executor.dispatch(|| {}).unwrap();
+    assert!(matches!(
+        executor.run(|| Ok::<_, Error>(())).await.unwrap_err(),
+        Error::DbCapacityExceeded
+    ));
+    release.release();
+    observed(executor.drain()).await.unwrap();
+    assert_eq!(
+        runtime.worker_detail(&worker).unwrap().status,
+        worker_runtime::catalog::WorkerStatus::Stopped
+    );
+    // Read the actual fs-store aggregate written by the accepted method. The
+    // deterministic execution fixture has no restorable Profile archive, so
+    // restarting it would test archive recovery rather than this commit owner.
+    let record: serde_json::Value = serde_json::from_slice(
+        &std::fs::read(
+            options
+                .root
+                .join("workers")
+                .join(worker.worker_id.to_string())
+                .join("worker.json"),
+        )
+        .unwrap(),
+    )
+    .unwrap();
+    assert_eq!(
+        record["status"], "stopped",
+        "cancelled pump lost its fs-store commit: {record}"
+    );
+    assert_ne!(
+        record["execution_state"]["state"],
+        "reconciliation_required"
+    );
+}

crates/workspace-server/src/server/tests/protocol_backpressure_tests.rs

@@ -11,6 +11,7 @@ async fn embedded_event_overflow_terminates_without_losing_the_admitted_prefix()
     let mut connection = connect_embedded_worker_protocol(
         embedded_protocol_test_source(&runtime, &worker_ref),
         embedded_protocol_test_scope(),
+        HostExecutors::default().host().unwrap(),
     )
     .await
     .unwrap();
@@ -69,6 +70,7 @@ async fn embedded_protocol_delivers_snapshot_then_ordered_replies_and_live_event
     let mut connection = connect_embedded_worker_protocol(
         embedded_protocol_test_source(&runtime, &worker_ref),
         embedded_protocol_test_scope(),
+        HostExecutors::default().host().unwrap(),
     )
     .await
     .unwrap();
@@ -105,9 +107,13 @@ async fn embedded_protocol_delivers_snapshot_then_ordered_replies_and_live_event
 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();
+    let mut old = connect_embedded_worker_protocol(
+        source.clone(),
+        embedded_protocol_test_scope(),
+        HostExecutors::default().host().unwrap(),
+    )
+    .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 {
@@ -127,9 +133,13 @@ async fn embedded_overflow_requires_explicit_reconnect_and_stop_ends_the_new_con
 
     // 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();
+    let mut fresh = connect_embedded_worker_protocol(
+        source,
+        embedded_protocol_test_scope(),
+        HostExecutors::default().host().unwrap(),
+    )
+    .await
+    .unwrap();
     assert_embedded_protocol_snapshot(next_embedded_protocol_event(&mut fresh).await, &worker_ref);
     fresh
         .methods
@@ -797,6 +807,7 @@ async fn embedded_disconnect_and_last_owner_drop_release_the_execution_event_sub
         let mut connection = connect_embedded_worker_protocol(
             embedded_protocol_test_source(&runtime, &worker_ref),
             embedded_protocol_test_scope(),
+            HostExecutors::default().host().unwrap(),
         )
         .await
         .unwrap();

crates/workspace-server/src/server_residual_execution_tests.rs

@@ -841,3 +841,64 @@ async fn ticket_tracking_and_case_authority_runtime_lookup_leave_db_and_tokio_fr
         );
     }
 }
+
+#[tokio::test(flavor = "current_thread")]
+async fn cancelled_feature_shutdown_keeps_sql_close_owned_without_blocking_control() {
+    let fixture = fixture().await;
+    let (database, _) = fixture.feature_database();
+    let mutation = database.clone();
+    let (entered, waiting) = oneshot::channel();
+    let (release, permit) = mpsc::channel();
+    let released = Arc::new(AtomicBool::new(false));
+    let flag = released.clone();
+    let thread = std::thread::spawn(move || {
+        mutation
+            .with_connection(|conn| {
+                conn.execute_batch("BEGIN IMMEDIATE; CREATE TABLE close_owner_probe(id INTEGER)")?;
+                let _ = entered.send(());
+                let _ = permit.recv_timeout(WATCHDOG);
+                flag.store(true, Ordering::SeqCst);
+                conn.execute_batch("ROLLBACK")?;
+                Ok(())
+            })
+            .unwrap();
+    });
+    let mut lock = WriteLock {
+        release: Some(release),
+        thread: Some(thread),
+        released,
+    };
+    observed(waiting).await.unwrap(); // actual SQL transaction owns the Feature connection
+    let leases = Arc::new(std::sync::atomic::AtomicUsize::new(0));
+    let mut close = Box::pin(crate::db_executor::with_operation_lifetime(
+        Arc::new(CountJobLifetime(leases.clone())),
+        shutdown_workspace_feature_storage(&fixture.api),
+    ));
+    assert!(
+        futures::poll!(&mut close).is_pending(),
+        "Feature connection close must not wait on Tokio"
+    );
+    runtime_progress_while_locked(&lock).await;
+    assert_eq!(
+        observed(fixture.host_executor.run(|| Ok::<_, Error>("control")))
+            .await
+            .unwrap(),
+        "control"
+    );
+    fixture.executor.dispatch(|| {}).unwrap();
+    assert!(matches!(
+        shutdown_workspace_feature_storage(&fixture.api)
+            .await
+            .unwrap_err(),
+        Error::DbCapacityExceeded
+    ));
+    drop(close);
+    assert_eq!(leases.load(Ordering::SeqCst), 1);
+    lock.release();
+    observed(fixture.executor.drain()).await.unwrap();
+    assert_eq!(leases.load(Ordering::SeqCst), 0);
+    assert!(
+        database.with_connection(|_| Ok(())).is_err(),
+        "accepted shutdown must fence later use"
+    );
+}