fix(server): preserve main worker responses after request cancellation

crates/base/src/server.rs

@@ -235,9 +235,10 @@ impl Service<Request<Body>> for WorkerService {
         }
       };
 
-      // If the token has already been canceled, return 503 instead of
-      // dropping the socket connection without a response.
-      if cancel.is_cancelled() {
+      // User-worker teardown also cancels this shared token. The main worker
+      // can recover from that failure and return a valid response, so only use
+      // the cancellation fallback when it could not produce a response.
+      if res.is_err() && cancel.is_cancelled() {
         error!("connection aborted (uri: {:?})", req_uri.to_string());
         return Ok(
           Response::builder()
@@ -1005,3 +1006,91 @@ fn accept_stream<I>(
     }
   });
 }
+
+#[cfg(test)]
+mod tests {
+  use super::*;
+
+  #[tokio::test]
+  async fn cancelled_request_preserves_main_worker_response() {
+    let (tx, mut rx) = mpsc::unbounded_channel::<WorkerRequestMsg>();
+    let (mut service, connection_token) =
+      WorkerService::new(SharedMetricSource::default(), tx);
+
+    // Force teardown to win the race, without relying on scheduling or CPU
+    // timing. The main worker can still recover and provide a complete response.
+    let main_worker = tokio::spawn(async move {
+      let msg = rx.recv().await.unwrap();
+      msg.conn_token.unwrap().cancel();
+      let (mut body_tx, body) = Body::channel();
+      msg
+        .res_tx
+        .send(Ok(
+          Response::builder()
+            .status(http_v02::StatusCode::INTERNAL_SERVER_ERROR)
+            .header("content-type", "application/json")
+            .header("x-main-worker", "recovered")
+            .body(body)
+            .unwrap(),
+        ))
+        .unwrap();
+      body_tx.send_data("{\"msg\":".into()).await.unwrap();
+      body_tx
+        .send_data("\"worker failed\"}".into())
+        .await
+        .unwrap();
+    });
+
+    let response = service.call(Request::new(Body::empty())).await.unwrap();
+    assert_eq!(
+      response.status(),
+      http_v02::StatusCode::INTERNAL_SERVER_ERROR
+    );
+    assert_eq!(response.headers()["x-main-worker"], "recovered");
+    assert_eq!(response.headers()["content-type"], "application/json");
+    assert!(!response.headers().contains_key("x-served-by"));
+    let body = hyper_v014::body::to_bytes(response.into_body())
+      .await
+      .unwrap();
+    assert_eq!(body.as_ref(), br#"{"msg":"worker failed"}"#);
+    main_worker.await.unwrap();
+    assert!(!connection_token.is_cancelled());
+  }
+
+  #[tokio::test]
+  async fn response_drop_cancels_request_but_not_keep_alive_connection() {
+    let (tx, mut rx) = mpsc::unbounded_channel::<WorkerRequestMsg>();
+    let (mut service, connection_token) =
+      WorkerService::new(SharedMetricSource::default(), tx);
+    let main_worker = tokio::spawn(async move {
+      let msg = rx.recv().await.unwrap();
+      let token = msg.conn_token.unwrap();
+      msg
+        .res_tx
+        .send(Ok(Response::new(Body::from("ok"))))
+        .unwrap();
+      token
+    });
+    let response = service.call(Request::new(Body::empty())).await.unwrap();
+    let request_token = main_worker.await.unwrap();
+    assert!(!request_token.is_cancelled());
+    drop(response);
+    assert!(request_token.is_cancelled());
+    assert!(!connection_token.is_cancelled());
+  }
+
+  #[tokio::test]
+  async fn connection_close_still_cancels_pending_request() {
+    let (tx, mut rx) = mpsc::unbounded_channel::<WorkerRequestMsg>();
+    let (mut service, connection_token) =
+      WorkerService::new(SharedMetricSource::default(), tx);
+    let response = tokio::spawn(service.call(Request::new(Body::empty())));
+    let msg = rx.recv().await.unwrap();
+    let request_token = msg.conn_token.unwrap();
+    assert!(!request_token.is_cancelled());
+    connection_token.cancel();
+    assert!(request_token.is_cancelled());
+    drop(msg.res_tx);
+    response.await.unwrap().unwrap();
+  }
+}

crates/base/test_cases/main_worker_recovery/index.ts

@@ -0,0 +1,30 @@
+// Keep a failed handle to exercise dispatch after teardown as well as a fresh
+// worker. HTTP keep-alive is controlled independently by the integration test.
+let failedWorker;
+
+Deno.serve(async (req: Request) => {
+  const scenario = new URL(req.url).pathname.slice(1);
+  try {
+    const worker = scenario === 'retired' ? failedWorker : await EdgeRuntime.userWorkers.create({
+      servicePath: scenario === 'cpu'
+        ? '../../examples/cpu'
+        : scenario === 'rejection'
+        ? './test_cases/unhandled-rejection'
+        : './test_cases/module-scope-throw',
+      memoryLimitMb: 150,
+      workerTimeoutMs: 10_000,
+      cpuTimeSoftLimitMs: 100,
+      cpuTimeHardLimitMs: 200,
+      forceCreate: true,
+      noModuleCache: false,
+      envVars: [],
+    });
+    failedWorker = worker;
+    return await worker.fetch(req);
+  } catch (e) {
+    return Response.json({ msg: e.toString() }, {
+      status: 500,
+      headers: { 'x-main-worker': 'recovered' },
+    });
+  }
+});

crates/base/test_cases/module-scope-throw/index.ts

@@ -0,0 +1 @@
+throw new Error('module-scope regression fixture');

crates/base/test_cases/unhandled-rejection/index.ts

@@ -0,0 +1,4 @@
+Deno.serve(() => {
+  Promise.reject(new Error('unhandled rejection regression fixture'));
+  return new Promise<Response>(() => {});
+});

crates/base/tests/integration_tests.rs

@@ -1263,6 +1263,75 @@ async fn req_failure_case_intentional_peer_reset_secure() {
   req_failure_case_intentional_peer_reset(new_localhost_tls(true)).await;
 }
 
+#[tokio::test]
+#[serial]
+async fn req_failure_case_preserves_main_worker_response() {
+  integration_test!(
+    "./test_cases/main_worker_recovery",
+    NON_SECURE_PORT,
+    "cpu",
+    None,
+    None,
+    None,
+    (
+      |(port, _, _, _, ..)| async move {
+        for keep_alive in [false, true] {
+          let client = Client::builder()
+            .pool_max_idle_per_host(if keep_alive { 1 } else { 0 })
+            .timeout(Duration::from_secs(10))
+            .build()
+            .unwrap();
+          for scenario in
+            ["cpu", "retired", "throw", "retired", "rejection", "retired"]
+          {
+            let res = client
+              .post(format!("http://localhost:{port}/{scenario}"))
+              .header("content-type", "application/json")
+              .body(r#"{"base":100}"#)
+              .send()
+              .await
+              .unwrap();
+            assert_eq!(res.status().as_u16(), 500, "{scenario}");
+            assert_eq!(res.headers()["x-main-worker"], "recovered");
+            assert!(!res.headers().contains_key("x-served-by"));
+            let payload = res.json::<ErrorResponsePayload>().await.unwrap();
+            let expected = match scenario {
+              "cpu" => payload.msg.starts_with("WorkerRequestCancelled:"),
+              // The pool may have removed the worker, marked it retired, or
+              // still retain its original failure. Preserve the main response
+              // in each case, independently of pool cleanup timing.
+              "retired" => {
+                payload.msg.starts_with("WorkerAlreadyRetired:")
+                  || payload.msg.starts_with("InvalidWorkerResponse:")
+              }
+              _ => payload.msg.starts_with("InvalidWorkerResponse:"),
+            };
+            assert!(expected, "{scenario}: {}", payload.msg);
+          }
+        }
+        Some(
+          Client::new()
+            .get(format!("http://localhost:{port}/throw"))
+            .send()
+            .await,
+        )
+      },
+      |resp| async {
+        let res = resp.unwrap();
+        assert_eq!(res.status().as_u16(), 500);
+        assert_eq!(res.headers()["x-main-worker"], "recovered");
+        assert!(res
+          .json::<ErrorResponsePayload>()
+          .await
+          .unwrap()
+          .msg
+          .starts_with("InvalidWorkerResponse:"));
+      }
+    ),
+    TerminationToken::new()
+  );
+}
+
 #[tokio::test]
 #[serial]
 async fn req_failure_case_op_cancel_from_server_due_to_cpu_resource_limit() {
@@ -1291,58 +1360,21 @@ async fn req_failure_case_op_cancel_from_server_due_to_cpu_resource_limit_2() {
   .await;
 }
 
-/// When the supervisor tears down a user worker that exceeded its CPU limit,
-/// two paths race each other and both outcomes are correct:
-///
-/// 1. The connection token is canceled before the main worker's response is
-///    handed back to the server, so the server serves 503 on its own.
-/// 2. The main worker's `WorkerRequestCancelled` handler wins the race, and its
-///    own 500 payload is relayed to the client untouched.
-///
-/// Which one wins depends on scheduling alone, so accept either, but keep
-/// asserting the shape of the response so that unrelated failures (most
-/// notably the request body being detached from its receiver) are still caught.
+/// Worker teardown can cancel the shared token before the main worker replies.
+/// A complete main-worker response must still reach the client in that order.
 async fn assert_op_cancel_from_server_response(
   resp: Result<Response, reqwest::Error>,
 ) {
   let res = resp.unwrap();
-  let status = res.status().as_u16();
-  let served_by = res
-    .headers()
-    .get("x-served-by")
-    .map(|v| v.to_str().unwrap().to_owned());
-
-  match status {
-    503 => {
-      assert_eq!(
-        served_by.as_deref(),
-        Some(concat!(env!("CARGO_PKG_NAME"), "/server"))
-      );
-    }
-
-    500 => {
-      assert_eq!(served_by, None);
-
-      let payload = res.json::<ErrorResponsePayload>().await;
+  assert_eq!(res.status().as_u16(), 500);
+  assert!(!res.headers().contains_key("x-served-by"));
 
-      assert!(payload.is_ok());
-
-      let msg = payload.unwrap().msg;
-
-      assert!(
-        !msg.starts_with("TypeError: request body receiver not connected"),
-        "unexpected error message: {msg}"
-      );
-      assert!(
-        msg
-          == "WorkerRequestCancelled: request has been cancelled by supervisor"
-          || msg == "broken pipe",
-        "unexpected error message: {msg}"
-      );
-    }
-
-    _ => panic!("unexpected status code: {status}"),
-  }
+  let msg = res.json::<ErrorResponsePayload>().await.unwrap().msg;
+  assert!(
+    msg == "WorkerRequestCancelled: request has been cancelled by supervisor"
+      || msg == "broken pipe",
+    "unexpected error message: {msg}"
+  );
 }
 
 async fn test_oak_file_upload<F, R>(