fix(middleware): drain websocket session end streams

architecture/sandbox.md

@@ -169,6 +169,10 @@ the remote adapter materializes an owned HTTP evaluation only when a request
 crosses that transport boundary. Both paths support bounded bidirectional
 WebSocket sessions, so a manifest advertises capabilities independently of
 transport.
+When a stage ends, the remote adapter sends its terminal event, half-closes the
+request stream, and briefly drains the response stream before releasing the
+transport. This keeps a queued terminal event from being canceled with the
+bidirectional RPC.
 The runtime keeps three states distinct: host selection attaches policy configs,
 manifest operation and phase bindings select the active chain, and the parsed
 message type determines whether that chain can inspect an individual payload.

crates/openshell-supervisor-middleware/src/lib.rs

@@ -5735,7 +5735,9 @@ mod tests {
             "all-skip preflight must not retain session capacity"
         );
         assert_eq!(
-            session_ends_rx.recv().await,
+            tokio::time::timeout(Duration::from_secs(1), session_ends_rx.recv())
+                .await
+                .expect("skipped stage must receive session_end"),
             Some(openshell_core::proto::WebSocketSessionEndReason::StageSkipped)
         );
         assert!(

crates/openshell-supervisor-middleware/src/websocket.rs

@@ -30,6 +30,7 @@ use super::{
 };
 
 const STREAM_CHANNEL_CAPACITY: usize = 4;
+const SESSION_END_TIMEOUT: Duration = Duration::from_millis(10);
 const MAX_REQUESTED_SUBPROTOCOLS: usize = 32;
 const MAX_SUBPROTOCOL_BYTES: usize = 4 * 1024;
 const MAX_SELECTED_SUBPROTOCOL_BYTES: usize = 256;
@@ -145,6 +146,42 @@ struct WebSocketStageTransport {
     responses: super::WebSocketResponseStream,
 }
 
+impl WebSocketStageTransport {
+    async fn end(self, reason: WebSocketSessionEndReason) {
+        let _ = tokio::time::timeout(SESSION_END_TIMEOUT, self.end_inner(reason)).await;
+    }
+
+    async fn end_inner(self, reason: WebSocketSessionEndReason) {
+        if self.sender.send(session_end_request(reason)).await.is_err() {
+            return;
+        }
+        self.drain().await;
+    }
+
+    async fn drain(self) {
+        let Self {
+            sender,
+            mut responses,
+        } = self;
+        // Keep the response handle alive while half-closing the request stream.
+        // Dropping both handles together schedules an HTTP/2 CANCEL, which may
+        // discard the buffered session_end before the middleware receives it.
+        drop(sender);
+        while responses.next().await.is_some() {}
+    }
+
+    fn end_now(self, reason: WebSocketSessionEndReason) {
+        if self.sender.try_send(session_end_request(reason)).is_err() {
+            return;
+        }
+        if let Ok(runtime) = tokio::runtime::Handle::try_current() {
+            drop(runtime.spawn(async move {
+                let _ = tokio::time::timeout(SESSION_END_TIMEOUT, self.drain()).await;
+            }));
+        }
+    }
+}
+
 struct WebSocketStage {
     entry: DescribedChainEntry,
     transport: Option<WebSocketStageTransport>,
@@ -161,11 +198,7 @@ impl WebSocketStage {
 
     async fn end(&mut self, reason: WebSocketSessionEndReason) {
         if let Some(transport) = self.transport.take() {
-            let _ = tokio::time::timeout(
-                Duration::from_millis(10),
-                transport.sender.send(session_end_request(reason)),
-            )
-            .await;
+            transport.end(reason).await;
         }
     }
 }
@@ -965,25 +998,25 @@ async fn open_stage(entry: DescribedChainEntry, input: WebSocketPreflightInput)
         Err(_) => return OpenStage::Failed(entry, "middleware_timeout".into()),
     };
     let Some(response) = response else {
-        let _ = sender.try_send(session_end_request(
-            WebSocketSessionEndReason::MiddlewareFailure,
-        ));
+        WebSocketStageTransport { sender, responses }
+            .end(WebSocketSessionEndReason::MiddlewareFailure)
+            .await;
         return OpenStage::Failed(entry, "missing_preflight_decision".into());
     };
     let Some(web_socket_session_event_result::Result::PreflightDecision(decision)) =
         response.result
     else {
-        let _ = sender.try_send(session_end_request(
-            WebSocketSessionEndReason::MiddlewareFailure,
-        ));
+        WebSocketStageTransport { sender, responses }
+            .end(WebSocketSessionEndReason::MiddlewareFailure)
+            .await;
         return OpenStage::Failed(entry, "invalid_preflight_decision".into());
     };
     let decision = match validate_preflight_decision(decision) {
         Ok(decision) => decision,
         Err(reason) => {
-            let _ = sender.try_send(session_end_request(
-                WebSocketSessionEndReason::MiddlewareFailure,
-            ));
+            WebSocketStageTransport { sender, responses }
+                .end(WebSocketSessionEndReason::MiddlewareFailure)
+                .await;
             return OpenStage::Failed(entry, reason.into());
         }
     };
@@ -1015,7 +1048,9 @@ async fn open_stage(entry: DescribedChainEntry, input: WebSocketPreflightInput)
         WebSocketPreflightAction::Skip => {
             let outcome =
                 preflight_stage_outcome(&entry, WebSocketInvocationOutcome::Skip, decision);
-            let _ = sender.try_send(session_end_request(WebSocketSessionEndReason::StageSkipped));
+            WebSocketStageTransport { sender, responses }
+                .end(WebSocketSessionEndReason::StageSkipped)
+                .await;
             OpenStage::Skip(outcome)
         }
         WebSocketPreflightAction::Unspecified => {
@@ -1308,7 +1343,7 @@ async fn end_stages(stages: &mut [WebSocketStage], reason: WebSocketSessionEndRe
 fn end_stages_now(stages: &mut [WebSocketStage], reason: WebSocketSessionEndReason) {
     for stage in stages {
         if let Some(transport) = stage.transport.take() {
-            let _ = transport.sender.try_send(session_end_request(reason));
+            transport.end_now(reason);
         }
     }
 }

docs/extensibility/supervisor-middleware.mdx

@@ -27,7 +27,7 @@ For an RFC 6455 upgrade over `ws://` or `wss://`, the supervisor first finds eve
 1. A preflight before the upgrade is sent upstream. The stage chooses `INSPECT`, voluntary `SKIP`, or authoritative `DENY` and may return a bounded diagnostic reason, stable reason code, findings, and metadata. OpenShell runs selected preflights concurrently; any `DENY` rejects the upgrade regardless of `on_error`.
 2. A session-start event after the upstream accepts the upgrade, including the negotiated subprotocol.
 3. Complete client-to-upstream text messages in sequence order. OpenShell reassembles fragmented messages and decompresses negotiated `permessage-deflate` messages before evaluation.
-4. A best-effort session-end event when the stage stream remains writable. OpenShell attempts at most one terminal event for each opened stream, including streams opened during a preflight that rejects the upgrade before session start.
+4. A best-effort session-end event when the stage stream remains writable. OpenShell attempts at most one terminal event for each opened stream, including streams opened during a preflight that rejects the upgrade before session start. It then half-closes the request stream and briefly drains the response stream so the terminal event can leave the local transport before the RPC closes. Middleware services should finish their response stream after the request stream reaches EOF.
 
 The protobuf represents each logical message with a `text` or `binary` payload variant. Text uses the protobuf `string` type, so invalid UTF-8 cannot enter the middleware contract. Results use an optional matching replacement variant: absence preserves the input, while presence represents a replacement even when its content is empty. OpenShell rejects attempts to change the message type. Allowed replacements are re-framed, re-compressed when required, and forwarded. Binary messages, control frames, and upstream-to-client traffic remain uninspected. Binary messages pass through under both `on_error` modes. For each active selected stage, OpenShell emits an informational `unsupported_message_type` coverage event and advances the session-global sequence; the next text message can therefore reach the stage with a valid sequence gap.