From b40f9617f83e0660b310bb006978ca7554a4f692 Mon Sep 17 00:00:00 2001 From: Piotr Mlocek Date: Wed, 2 Sep 2026 12:03:27 -0700 Subject: [PATCH 1/3] fix(middleware): drain websocket session end streams Signed-off-by: Piotr Mlocek --- architecture/sandbox.md | 4 ++ .../src/lib.rs | 4 +- .../src/websocket.rs | 67 ++++++++++++++----- docs/extensibility/supervisor-middleware.mdx | 2 +- 4 files changed, 59 insertions(+), 18 deletions(-) diff --git a/architecture/sandbox.md b/architecture/sandbox.md index dd9621a09d..8900d81679 100644 --- a/architecture/sandbox.md +++ b/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. diff --git a/crates/openshell-supervisor-middleware/src/lib.rs b/crates/openshell-supervisor-middleware/src/lib.rs index 902b5165c2..0f5f1d2de6 100644 --- a/crates/openshell-supervisor-middleware/src/lib.rs +++ b/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!( diff --git a/crates/openshell-supervisor-middleware/src/websocket.rs b/crates/openshell-supervisor-middleware/src/websocket.rs index 1fd95021a8..e61460946a 100644 --- a/crates/openshell-supervisor-middleware/src/websocket.rs +++ b/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, @@ -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); } } } diff --git a/docs/extensibility/supervisor-middleware.mdx b/docs/extensibility/supervisor-middleware.mdx index 68cdbd29a5..6d2e1868e2 100644 --- a/docs/extensibility/supervisor-middleware.mdx +++ b/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. From 342e05e0943e8cb87ddfe043870fb3fd262a80fc Mon Sep 17 00:00:00 2001 From: Piotr Mlocek Date: Wed, 2 Sep 2026 12:16:28 -0700 Subject: [PATCH 2/3] docs(middleware): diagram websocket stream shutdown Signed-off-by: Piotr Mlocek --- docs/extensibility/supervisor-middleware.mdx | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/docs/extensibility/supervisor-middleware.mdx b/docs/extensibility/supervisor-middleware.mdx index 6d2e1868e2..5d61ed12b5 100644 --- a/docs/extensibility/supervisor-middleware.mdx +++ b/docs/extensibility/supervisor-middleware.mdx @@ -29,6 +29,25 @@ For an RFC 6455 upgrade over `ws://` or `wss://`, the supervisor first finds eve 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. 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. +```mermaid +sequenceDiagram + participant S as Supervisor + participant H as tonic / h2 + participant M as Middleware + + Note over S,M: Previous shutdown + S->>H: Queue session_end + S-xH: Drop request and response handles + H--xM: RST_STREAM(CANCEL) may discard session_end + + Note over S,M: Graceful shutdown + S->>H: Queue session_end + S->>H: Drop request sender to half-close + H->>M: Deliver session_end, then request EOF + M-->>H: Finish response stream + H-->>S: Drain completes or reaches the 10 ms bound +``` + 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. The network supervisor reserves process-wide assembly capacity before buffering every parsed WebSocket text message, even when no middleware is selected. At most 32 assemblies run while 64 additional callers wait without buffering payload bytes. When both bounds are full, OpenShell closes the WebSocket with code `1013` before reading the new message payload. A text message may contain at most 4,096 fragments, must make input progress within 30 seconds, and must finish assembly within 2 minutes. Forwarding the completed text frame must finish within another 2 minutes. The assembly budget lasts for the supervisor process lifetime, so policy reloads do not reset its capacity. From 8d19cdc99812ec61e047679f0a24a90b4ed5ee90 Mon Sep 17 00:00:00 2001 From: Piotr Mlocek Date: Wed, 2 Sep 2026 12:18:43 -0700 Subject: [PATCH 3/3] docs(middleware): keep shutdown diagram in pull request Signed-off-by: Piotr Mlocek --- docs/extensibility/supervisor-middleware.mdx | 19 ------------------- 1 file changed, 19 deletions(-) diff --git a/docs/extensibility/supervisor-middleware.mdx b/docs/extensibility/supervisor-middleware.mdx index 5d61ed12b5..6d2e1868e2 100644 --- a/docs/extensibility/supervisor-middleware.mdx +++ b/docs/extensibility/supervisor-middleware.mdx @@ -29,25 +29,6 @@ For an RFC 6455 upgrade over `ws://` or `wss://`, the supervisor first finds eve 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. 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. -```mermaid -sequenceDiagram - participant S as Supervisor - participant H as tonic / h2 - participant M as Middleware - - Note over S,M: Previous shutdown - S->>H: Queue session_end - S-xH: Drop request and response handles - H--xM: RST_STREAM(CANCEL) may discard session_end - - Note over S,M: Graceful shutdown - S->>H: Queue session_end - S->>H: Drop request sender to half-close - H->>M: Deliver session_end, then request EOF - M-->>H: Finish response stream - H-->>S: Drain completes or reaches the 10 ms bound -``` - 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. The network supervisor reserves process-wide assembly capacity before buffering every parsed WebSocket text message, even when no middleware is selected. At most 32 assemblies run while 64 additional callers wait without buffering payload bytes. When both bounds are full, OpenShell closes the WebSocket with code `1013` before reading the new message payload. A text message may contain at most 4,096 fragments, must make input progress within 30 seconds, and must finish assembly within 2 minutes. Forwarding the completed text frame must finish within another 2 minutes. The assembly budget lasts for the supervisor process lifetime, so policy reloads do not reset its capacity.