endpoint-libs 3.4.3

use eyre::Result;
use futures::StreamExt;
use futures::future::{FutureExt, LocalBoxFuture};
use futures::stream::FuturesUnordered;
use std::collections::HashSet;
use std::sync::Arc;
use tracing::*;

use crate::libs::signal::Shutdown;

/// What the session loop does after one inbound frame.
///
/// Handler bodies are polled in place on this task (`Task`) so a slow hook
/// cannot stall the read loop, without `spawn_local` and the LocalSet it
/// needs. `serve_with` already drives connections the same way.
enum Dispatch {
    Close,
    Keep,
    Task(LocalBoxFuture<'static, ()>),
}

use crate::libs::ws::WsMessage as Message;
use crate::libs::ws::outbound::Receiver;

use crate::libs::error_code::ErrorCode;
use crate::libs::toolbox::{RequestContext, TOOLBOX};

use super::mcp::{
    self, JsonRpcError, JsonRpcId, JsonRpcRequest, McpAction, McpCallCtx, McpState,
    encode_tool_error, jsonrpc_error, jsonrpc_result,
};
use super::{
    MessageStream, RequestOutcome, StreamError, WebsocketServer, WsConnection, WsRequestValue,
    WsResponseError, WsResponseValue,
};

pub struct WsClientSession {
    conn_info: Arc<WsConnection>,
    conn: Box<dyn MessageStream>,
    rx: Receiver<Message>,
    server: Arc<WebsocketServer>,
    /// Policy close, shared with [`crate::libs::ws::WsStreamState::end`].
    /// A session built with [`WsClientSession::new`] and never [`WsClientSession::bind_end`]
    /// holds a flag nobody else can set.
    end: Arc<Shutdown>,
    last_received: std::rc::Rc<std::cell::Cell<std::time::Instant>>,
}

impl WsClientSession {
    pub fn new(
        conn_info: Arc<WsConnection>,
        conn: Box<dyn MessageStream>,
        rx: Receiver<Message>,
        server: Arc<WebsocketServer>,
    ) -> Self {
        Self {
            conn_info,
            conn,
            rx,
            server,
            end: Arc::new(Shutdown::new()),
            last_received: std::rc::Rc::new(std::cell::Cell::new(std::time::Instant::now())),
        }
    }

    /// Use the connection's policy-close flag. [`WebsocketServer`] does this
    /// after `WebsocketStates::insert` has created the flag.
    pub fn bind_end(&mut self, end: Arc<Shutdown>) {
        self.end = end;
    }

    pub fn conn(&self) -> &dyn MessageStream {
        self.conn.as_ref()
    }

    pub async fn run(mut self) {
        let addr = self.conn_info.peer.display();
        let conn_id = self.conn_info.connection_id;
        let last_received = self.last_received.clone();
        let idle_timeout = self.server.config.idle_timeout;
        let idle = async move {
            let Some(timeout) = idle_timeout else {
                return futures::future::pending::<()>().await;
            };
            loop {
                let remaining = timeout.saturating_sub(last_received.get().elapsed());
                if remaining.is_zero() {
                    return;
                }
                nagoya::sleep(remaining).await;
            }
        };
        let run = self.run_loop();
        futures::pin_mut!(idle, run);
        let result = match futures::future::select(idle, run).await {
            futures::future::Either::Left(_) => Ok(()),
            futures::future::Either::Right((result, _)) => result,
        };
        if let Err(err) = result {
            error!(
                ws_server = true,
                ?err,
                ?addr,
                ?conn_id,
                "Failed to run websocket session"
            );
        }
    }

    fn handle_message(&mut self, msg: Message) -> Result<Dispatch> {
        let bytes = match &msg {
            Message::Text(t) => t.len(),
            Message::Binary(b) => b.len(),
            _ => 0,
        };
        self.server
            .limits
            .observe_message(&self.server.config, bytes);
        let addr = self.conn_info.peer.display();
        let mut context = RequestContext::from_conn(&self.conn_info);

        // MCP: route JSON-RPC 2.0 frames to the MCP adapter when enabled.
        // Detection is unambiguous: legacy frames ({method: u32, seq, params})
        // can never carry a top-level "jsonrpc": "2.0" member.
        if let Some(mcp) = &self.server.mcp {
            let payload = match &msg {
                Message::Text(t) => Some(t.as_str()),
                Message::Binary(b) => std::str::from_utf8(b).ok(),
                _ => None,
            };
            if let Some(frame) = payload.and_then(mcp::try_parse_jsonrpc) {
                let mcp = Arc::clone(mcp);
                if let Some(task) = self.handle_mcp_frame(mcp, frame, context) {
                    return Ok(Dispatch::Task(task));
                }
                return Ok(Dispatch::Keep);
            }
        }

        if self.server.config.mcp_only {
            self.server.toolbox.send_raw(
                context.connection_id,
                jsonrpc_error(
                    &None,
                    JsonRpcError::new(
                        mcp::INVALID_REQUEST,
                        "This WebSocket accepts MCP JSON-RPC 2.0 frames only",
                    ),
                )
                .to_string(),
            );
            return Ok(Dispatch::Keep);
        }

        #[allow(unreachable_patterns)]
        let obj: Result<WsRequestValue, _> = match msg {
            Message::Text(t) => {
                debug!(
                    ws_server = true,
                    ?addr,
                    bytes = t.len(),
                    "Handling text request"
                );
                serde_json::from_str(&t)
            }
            Message::Binary(b) => {
                debug!(
                    ws_server = true,
                    ?addr,
                    bytes = b.len(),
                    "Handling binary request"
                );
                serde_json::from_slice(&b)
            }
            Message::Ping(_) => {
                return Ok(Dispatch::Keep);
            }
            Message::Pong(_) => {
                return Ok(Dispatch::Keep);
            }
            Message::Close(_) => {
                debug!(ws_server = true, ?addr, "Receive side terminated");
                return Ok(Dispatch::Close);
            }
            _ => {
                warn!(
                    ws_server = true,
                    ?addr,
                    "Ignoring unsupported WebSocket frame"
                );
                return Ok(Dispatch::Keep);
            }
        };
        let req = match obj {
            Ok(req) => req,
            Err(err) => {
                self.server.toolbox.send(
                    context.connection_id,
                    WsResponseValue::Error(WsResponseError {
                        method: context.method,
                        code: ErrorCode::BAD_REQUEST.to_u32(),
                        seq: context.seq,
                        log_id: context.log_id.to_string(),
                        params: serde_json::json!({
                            "kind": ErrorCode::BAD_REQUEST.kind(),
                            "message": err.to_string(),
                        }),
                    }),
                );
                return Ok(Dispatch::Keep);
            }
        };
        debug!(
            ws_server = true,
            ?addr,
            method = req.method,
            seq = req.seq,
            "Parsed WebSocket request"
        );
        context.seq = req.seq;
        context.method = req.method;
        context.user_id = self.conn_info.get_user_id();
        context.roles = self.conn_info.get_roles();

        let Some(endpoint) = self.server.handlers.get(&req.method) else {
            self.server.toolbox.send(
                context.connection_id,
                WsResponseValue::Error(WsResponseError {
                    method: context.method,
                    code: ErrorCode::NOT_IMPLEMENTED.to_u32(),
                    seq: context.seq,
                    log_id: context.log_id.to_string(),
                    params: serde_json::json!({
                        "kind": ErrorCode::NOT_IMPLEMENTED.kind(),
                        "message": "Method not implemented",
                    }),
                }),
            );
            return Ok(Dispatch::Keep);
        };

        if !check_roles(&context.roles, &endpoint.allowed_roles) {
            self.server.toolbox.send(
                context.connection_id,
                WsResponseValue::Error(WsResponseError {
                    method: context.method,
                    code: ErrorCode::FORBIDDEN.to_u32(),
                    seq: context.seq,
                    log_id: context.log_id.to_string(),
                    params: serde_json::json!({
                        "kind": ErrorCode::FORBIDDEN.kind(),
                        "message": "Forbidden",
                    }),
                }),
            );
            return Ok(Dispatch::Keep);
        }

        let Some(slot) = self.server.limits.acquire_request(&self.server.config) else {
            self.server.refuse_request(&context, "server busy");
            return Ok(Dispatch::Keep);
        };

        let handler = endpoint.handler.clone();
        let toolbox = self.server.toolbox.clone();
        let hooks = self.server.hooks.clone();
        let schema = endpoint.schema.clone();
        let request_timeout = self.server.config.request_timeout;
        let server = Arc::clone(&self.server);
        Ok(Dispatch::Task(
            async move {
                let _slot = slot;
                let started = std::time::Instant::now();
                let timeout_context = context.clone();
                let responded = std::cell::Cell::new(false);
                let work = async {
                    let mut context = context;
                    // Polled alongside the session read loop so a slow hook cannot
                    // stall it, and after check_roles so they only see calls that
                    // were already allowed to reach this endpoint.
                    if let Err(custom) = hooks.run_before(&mut context, &schema, &req.params).await {
                        responded.set(true);
                        let code = custom.code.to_u32();
                        toolbox.send(
                            context.connection_id,
                            WsResponseValue::Error(WsResponseError {
                                method: context.method,
                                code,
                                seq: context.seq,
                                log_id: context.log_id.to_string(),
                                params: custom.params.clone(),
                            }),
                        );
                        hooks
                            .run_after(&context, &schema, &RequestOutcome::PublicErr { code })
                            .await;
                        return;
                    }

                    let handler = TOOLBOX.scope(
                        toolbox.clone(),
                        handler.handle(&toolbox, context.clone(), req.params),
                    );
                    handler.await;
                    responded.set(true);
                    // The erased handler reports its own outcome through the toolbox, so
                    // AfterRequest observes completion rather than the specific result here.
                    hooks
                        .run_after(&context, &schema, &RequestOutcome::Ok)
                        .await;

                };
                let expired = match request_timeout {
                    Some(timeout) => nagoya::timeout(timeout, work).await.is_err(),
                    None => { work.await; false }
                };
                server.limits.observe_request_time(&server.config, started.elapsed());
                if expired && !responded.get() {
                    let context = timeout_context;
                    toolbox.send(context.connection_id, WsResponseValue::Error(WsResponseError {
                        method: context.method,
                        code: ErrorCode::REQUEST_TIMEOUT.to_u32(),
                        seq: context.seq,
                        log_id: context.log_id.to_string(),
                        params: serde_json::json!({"kind": "timeout", "message": "request timed out"}),
                    }));
                }
            }
            .boxed_local(),
        ))
    }

    /// Handles one parsed JSON-RPC frame: lifecycle methods are answered
    /// inline from [`McpState`]; `tools/call` dispatches to the endpoint's
    /// request handler via [`RequestHandlerErased::handle_mcp`].
    ///
    /// [`RequestHandlerErased::handle_mcp`]: crate::libs::handler::RequestHandlerErased::handle_mcp
    fn handle_mcp_frame(
        &mut self,
        mcp: Arc<McpState>,
        frame: Result<JsonRpcRequest, serde_json::Value>,
        mut context: RequestContext,
    ) -> Option<LocalBoxFuture<'static, ()>> {
        let conn_id = context.connection_id;
        let req = match frame {
            Ok(req) => req,
            Err(error_frame) => {
                self.server
                    .toolbox
                    .send_raw(conn_id, error_frame.to_string());
                return None;
            }
        };

        if req.is_notification() {
            return None;
        }
        context.user_id = self.conn_info.get_user_id();
        context.roles = self.conn_info.get_roles();

        match mcp.route(req, &context.roles) {
            McpAction::Respond(frame) => {
                self.server.toolbox.send_raw(conn_id, frame.to_string());
                None
            }
            McpAction::Ignore => None,
            McpAction::ToolCall {
                id,
                method_code,
                arguments,
            } => {
                context.method = method_code;
                // Numeric ids that fit u32 double as the legacy seq for logging.
                if let Some(JsonRpcId::Num(n)) = &id
                    && let Ok(seq) = u32::try_from(*n)
                {
                    context.seq = seq;
                }

                let Some(endpoint) = self.server.handlers.get(&method_code) else {
                    // Unreachable in practice: McpState is built from handlers.
                    self.server.toolbox.send_raw(
                        conn_id,
                        jsonrpc_error(
                            &id,
                            JsonRpcError::new(mcp::METHOD_NOT_FOUND, "Method not found"),
                        )
                        .to_string(),
                    );
                    return None;
                };

                let Some(slot) = self.server.limits.acquire_request(&self.server.config) else {
                    if id.is_some() {
                        self.server.toolbox.send_raw(
                            conn_id,
                            jsonrpc_error(&id, JsonRpcError::new(-32000, "server busy"))
                                .to_string(),
                        );
                    }
                    return None;
                };
                let request_timeout = self.server.config.request_timeout;
                let server = Arc::clone(&self.server);
                let handler = endpoint.handler.clone();
                let toolbox = self.server.toolbox.clone();
                let hooks = self.server.hooks.clone();
                let schema = endpoint.schema.clone();
                Some(
                    async move {
                        let _slot = slot;
                        let started = std::time::Instant::now();
                        let responded = std::cell::Cell::new(false);
                        let work = async {
                            let mut context = context;
                            // Same placement as the legacy path, but the rejection has to go
                            // back in the MCP envelope as a tool error.
                            if let Err(custom) =
                                hooks.run_before(&mut context, &schema, &arguments).await
                            {
                                responded.set(true);
                                let code = custom.code.to_u32();
                                toolbox.send_raw(
                                    conn_id,
                                    jsonrpc_result(
                                        &id,
                                        encode_tool_error(custom.code, &custom.params),
                                    )
                                    .to_string(),
                                );
                                hooks
                                    .run_after(
                                        &context,
                                        &schema,
                                        &RequestOutcome::PublicErr { code },
                                    )
                                    .await;
                                return;
                            }

                            TOOLBOX
                                .scope(
                                    toolbox.clone(),
                                    handler.handle_mcp(
                                        &toolbox,
                                        context.clone(),
                                        McpCallCtx { id: id.clone() },
                                        arguments,
                                    ),
                                )
                                .await;

                            responded.set(true);
                            hooks
                                .run_after(&context, &schema, &RequestOutcome::Ok)
                                .await;
                        };
                        let expired = match request_timeout {
                            Some(timeout) => nagoya::timeout(timeout, work).await.is_err(),
                            None => {
                                work.await;
                                false
                            }
                        };
                        server
                            .limits
                            .observe_request_time(&server.config, started.elapsed());
                        if expired && !responded.get() && id.is_some() {
                            toolbox.send_raw(
                                conn_id,
                                jsonrpc_error(&id, JsonRpcError::new(-32001, "request timed out"))
                                    .to_string(),
                            );
                        }
                    }
                    .boxed_local(),
                )
            }
        }
    }

    async fn run_loop(&mut self) -> Result<()> {
        let conn_id = self.conn_info.connection_id;
        let mut handlers = FuturesUnordered::new();
        // Owned so the recv futures can drop before the arms borrow `self` again.
        // `pin_mut` holds those borrows to the end of its scope.
        enum Ready {
            Outbound(Outbound),
            Inbound(Option<std::result::Result<Message, StreamError>>),
            Handler,
            End,
        }
        loop {
            if self.idle_expired() {
                return Ok(());
            }
            while let Ok(msg) = self.rx.try_recv() {
                if self.idle_expired() {
                    return Ok(());
                }
                if !self.send_message(msg).await {
                    return Ok(());
                }
                if self.server.config.header_only {
                    return Ok(());
                }
            }
            // After the queue is drained. A payload already queued is sent
            // above; the flag only wins once there is nothing left to send.
            if self.end.is_cancelled() {
                debug!(ws_server = true, ?conn_id, "Session end flagged");
                return Ok(());
            }

            // A finished handler, then outbound, then inbound, then the policy
            // flag. `futures::future::select` returns on its left arm the
            // moment that arm is ready and never polls the right one, so this
            // order is a strict priority. `tokio::select!` picked at random,
            // so nothing here inherits its fairness and the order has to earn
            // itself.
            //
            // Handlers go first because they are the only arm that cannot
            // feed itself. The set is finite, each entry completes once, and
            // new entries arrive only from inbound, so draining it cannot
            // starve anything: it empties, and then the arm is disabled
            // again. Behind inbound it starves instead -- a peer that keeps a
            // frame available makes inbound ready on every poll, and the set
            // grows without ever being polled, so a queued request never
            // completes and never answers.
            //
            // The handler arm stays disabled while the set is empty, because
            // `FuturesUnordered::next` on an empty set is `Ready(None)` and
            // would spin.
            // The flag is last so a frame that is already ready is not dropped
            // on the floor when the flag trips in the same poll.
            let ready = {
                let outbound = self.rx.recv();
                let inbound = self.conn.recv();
                let end = self.end.cancelled();
                let handler = if handlers.is_empty() {
                    None
                } else {
                    Some(handlers.next())
                };
                match priority4(handler, outbound, inbound, end).await {
                    Arm::Handler(_) => Ready::Handler,
                    Arm::Outbound(msg) => Ready::Outbound(classify_outbound(msg)),
                    Arm::Inbound(msg) => Ready::Inbound(msg),
                    Arm::End(()) => Ready::End,
                }
            };
            match ready {
                Ready::Outbound(Outbound::Frame(msg)) => {
                    if !self.send_message(msg).await {
                        break;
                    }
                    if self.server.config.header_only {
                        break;
                    }
                }
                Ready::Outbound(Outbound::Closed) => {
                    debug!(ws_server = true, ?conn_id, "Outbound channel closed");
                    break;
                }
                Ready::End => {
                    debug!(ws_server = true, ?conn_id, "Session end flagged");
                    break;
                }
                Ready::Inbound(Some(msg_result)) => {
                    let msg = match msg_result {
                        Ok(m) => m,
                        Err(StreamError::Closed) => {
                            debug!(ws_server = true, ?conn_id, "WS receive: connection closed");
                            break;
                        }
                        Err(StreamError::Protocol(e)) => {
                            warn!(ws_server = true, ?conn_id, err=%e, "WS protocol error on receive");
                            break;
                        }
                        Err(StreamError::WriteBufferFull) => {
                            warn!(
                                ws_server = true,
                                ?conn_id,
                                "WS write buffer full on receive"
                            );
                            break;
                        }
                        Err(StreamError::Other(e)) => {
                            error!(ws_server = true, ?conn_id, err=%e, "WS receive error");
                            break;
                        }
                    };
                    self.last_received.set(std::time::Instant::now());
                    match self.handle_message(msg)? {
                        Dispatch::Close => break,
                        Dispatch::Keep => {}
                        Dispatch::Task(task) => handlers.push(task),
                    }
                }
                Ready::Inbound(None) => {
                    debug!(ws_server = true, ?conn_id, "Inbound stream ended");
                    break;
                }
                Ready::Handler => {}
            }
        }

        Ok(())
    }

    fn idle_expired(&self) -> bool {
        self.server
            .config
            .idle_timeout
            .is_some_and(|timeout| self.last_received.get().elapsed() >= timeout)
    }

    async fn send_message(&mut self, msg: Message) -> bool {
        let conn_id = self.conn_info.connection_id;
        match self.conn.send(msg).await {
            Ok(()) => true,
            Err(StreamError::Closed) => {
                debug!(ws_server = true, ?conn_id, "WS send: connection closed");
                false
            }
            Err(StreamError::WriteBufferFull) => {
                warn!(ws_server = true, ?conn_id, "WS send: write buffer full");
                false
            }
            Err(StreamError::Protocol(e)) => {
                warn!(ws_server = true, ?conn_id, err=%e, "WS send: protocol error");
                false
            }
            Err(StreamError::Other(e)) => {
                error!(ws_server = true, ?conn_id, err=%e, "WS send error");
                false
            }
        }
    }
}

/// What an outbound `recv` means. `Closed` is the teardown edge: every sender
/// is gone. It is not a frame and it ends the session.
pub(crate) enum Outbound {
    Frame(Message),
    Closed,
}

pub(crate) fn classify_outbound(msg: Option<Message>) -> Outbound {
    match msg {
        Some(frame) => Outbound::Frame(frame),
        None => Outbound::Closed,
    }
}

/// Which arm of [`priority4`] fired, carrying that arm's output.
pub(crate) enum Arm<H, O, I, E> {
    Handler(H),
    Outbound(O),
    Inbound(I),
    End(E),
}

/// Await four futures in strict priority order: `handler`, then `outbound`,
/// then `inbound`, then `end`.
///
/// `futures::future::select` returns on its left arm the moment that arm is
/// ready and never polls the right one, so nesting it is a priority and not
/// the random choice `tokio::select!` made. Every arm that can be made ready
/// on demand by a peer therefore has to sit below one that cannot, or it
/// starves.
///
/// `handler` is [`None`] when there is nothing to wait for, which parks that
/// arm forever instead of disabling it by hand.
/// `FuturesUnordered::next` on an empty set is `Ready(None)`, so passing it
/// here while empty would take this arm on every poll and spin.
pub(crate) async fn priority4<HF, OF, IF, EF>(
    handler: Option<HF>,
    outbound: OF,
    inbound: IF,
    end: EF,
) -> Arm<HF::Output, OF::Output, IF::Output, EF::Output>
where
    HF: std::future::Future,
    OF: std::future::Future,
    IF: std::future::Future,
    EF: std::future::Future,
{
    use futures::future::{Either, select};

    let handler = async move {
        match handler {
            Some(fut) => fut.await,
            None => futures::future::pending::<HF::Output>().await,
        }
    };
    futures::pin_mut!(handler);
    futures::pin_mut!(outbound);
    futures::pin_mut!(inbound);
    futures::pin_mut!(end);

    match select(handler, select(outbound, select(inbound, end))).await {
        Either::Left((fired, _)) => Arm::Handler(fired),
        Either::Right((inner, _)) => match inner {
            Either::Left((fired, _)) => Arm::Outbound(fired),
            Either::Right((inner, _)) => match inner {
                Either::Left((fired, _)) => Arm::Inbound(fired),
                Either::Right((fired, _)) => Arm::End(fired),
            },
        },
    }
}

fn check_roles(actual_roles: &[u32], allowed_roles: &HashSet<u32>) -> bool {
    if allowed_roles.is_empty() || actual_roles.is_empty() {
        return false;
    }
    actual_roles.iter().any(|role| allowed_roles.contains(role))
}

#[cfg(test)]
mod tests {
    #[test]
    fn check_roles_allowed() {
        use super::check_roles;
        use std::collections::HashSet;

        let allowed_roles: HashSet<u32> = [1, 2, 3].iter().cloned().collect();
        assert!(check_roles(&[1], &allowed_roles.clone()));
        assert!(check_roles(&[2], &allowed_roles.clone()));
        assert!(check_roles(&[1, 2], &allowed_roles.clone()));
        assert!(check_roles(&[4, 2], &allowed_roles.clone()));

        assert!(!check_roles(&[4], &allowed_roles.clone()));
    }

    #[test]
    fn check_roles_empty() {
        use super::check_roles;
        use std::collections::HashSet;

        let allowed_roles: HashSet<u32> = HashSet::new();
        assert!(!check_roles(&[1], &allowed_roles));
    }

    #[test]
    fn closed_outbound_is_not_a_frame() {
        use super::{Outbound, classify_outbound};
        use crate::libs::ws::WsMessage as Message;

        assert!(matches!(classify_outbound(None), Outbound::Closed));
        assert!(matches!(
            classify_outbound(Some(Message::from("hi"))),
            Outbound::Frame(_)
        ));
    }

    /// A ready handler beats a ready inbound frame.
    ///
    /// This is the arm that cannot feed itself: new entries arrive only from
    /// inbound, so it drains. Below inbound it starves instead, because a
    /// peer that keeps a frame available makes inbound ready on every poll
    /// and the handler set is then never polled at all.
    ///
    /// `nagoya::block_on` rather than a runtime: every arm below is `ready` or
    /// `pending`, so `priority4` resolves on its first poll and an executor is
    /// providing nothing but that poll.
    #[test]
    fn a_ready_handler_beats_a_ready_inbound() {
        use futures::future::{pending, ready};

        use super::{Arm, priority4};

        let fired = nagoya::block_on(priority4(
            Some(ready("handler")),
            pending::<()>(),
            ready("inbound"),
            ready(()),
        ));
        assert!(
            matches!(fired, Arm::Handler("handler")),
            "a finished handler lost to an inbound frame that was also ready"
        );
    }

    #[test]
    fn an_absent_handler_does_not_take_the_arm() {
        use futures::future::{Ready, pending, ready};

        use super::{Arm, priority4};

        let absent: Option<Ready<()>> = None;
        let fired = nagoya::block_on(priority4(
            absent,
            ready("outbound"),
            pending::<()>(),
            pending::<()>(),
        ));
        assert!(
            matches!(fired, Arm::Outbound("outbound")),
            "an empty handler set took the arm instead of parking it"
        );
    }

    #[test]
    fn outbound_beats_inbound_and_both_beat_the_flag() {
        use futures::future::{Ready, pending, ready};

        use super::{Arm, priority4};

        let absent: Option<Ready<()>> = None;
        let fired = nagoya::block_on(priority4(
            absent,
            ready("outbound"),
            ready("inbound"),
            ready(()),
        ));
        assert!(matches!(fired, Arm::Outbound("outbound")));

        let absent: Option<Ready<()>> = None;
        let fired = nagoya::block_on(priority4(
            absent,
            pending::<()>(),
            ready("inbound"),
            ready(()),
        ));
        assert!(
            matches!(fired, Arm::Inbound("inbound")),
            "the policy flag dropped a frame that was already ready"
        );
    }

    struct PendingStream {
        parked: std::sync::Arc<std::sync::atomic::AtomicBool>,
    }

    #[async_trait::async_trait(?Send)]
    impl super::MessageStream for PendingStream {
        async fn send(
            &mut self,
            _msg: super::Message,
        ) -> std::result::Result<(), super::StreamError> {
            Ok(())
        }

        async fn recv(
            &mut self,
        ) -> Option<std::result::Result<super::Message, super::StreamError>> {
            self.parked
                .store(true, std::sync::atomic::Ordering::Release);
            std::future::pending().await
        }
    }

    fn test_conn() -> std::sync::Arc<super::WsConnection> {
        use std::sync::Arc;
        use std::sync::atomic::AtomicU64;

        use parking_lot::RwLock;

        use crate::libs::peer::{Extensions, PeerIdentity};

        Arc::new(super::WsConnection {
            connection_id: 1,
            user_id: AtomicU64::new(0),
            roles: Arc::new(RwLock::new(Arc::new(Vec::new()))),
            peer: PeerIdentity::Unknown,
            extensions: Extensions::new(),
            log_id: 0,
        })
    }

    /// The teardown edge: no sender left, so the session has to return rather
    /// than wait on a channel nobody can write to.
    ///
    /// `nagoya::block_on` polls on this thread, which is what `run` needs --
    /// the stream is `?Send` and the future is not movable between threads.
    /// The timeout is the whole assertion: a session that fails to notice the
    /// closed channel parks forever, so without a deadline this test hangs
    /// instead of failing.
    #[test]
    fn dropping_the_outbound_sender_ends_the_session() {
        use std::time::Duration;

        use super::WsClientSession;
        use crate::libs::ws::outbound;
        use crate::libs::ws::{WebsocketServer, WsServerConfig};

        let (tx, rx) = outbound::channel(1);
        drop(tx);
        let session = WsClientSession::new(
            test_conn(),
            Box::new(PendingStream {
                parked: std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)),
            }),
            rx,
            std::sync::Arc::new(WebsocketServer::new(WsServerConfig::default())),
        );
        let finished = nagoya::block_on(nagoya::timeout(Duration::from_secs(1), session.run()));
        assert!(
            finished.is_ok(),
            "session kept running after every outbound sender was dropped"
        );
    }

    /// The policy-close edge: the session is parked on `recv` with nothing to
    /// read, and setting the flag still ends it.
    ///
    /// The session and the thing that flags it have to run concurrently, and
    /// both are `?Send`, which is what a `LocalSet` and `spawn_local` used to
    /// buy. `futures::future::join` buys the same concurrency without a
    /// spawner: two futures, one task, polled on this thread by
    /// `nagoya::block_on`. Nothing here needs a second thread -- the flag is
    /// set from inside the same task, by the arm that watches `parked` -- so
    /// single-threaded is not a weakening, it just removes the runtime.
    ///
    /// `parked` is the ordering that makes the test mean anything: flagging
    /// before the session has reached `recv` would prove only that a session
    /// that never started can stop. `flagged` then says which of the two
    /// deadlines the timeout hit, which the two separate `timeout` calls in
    /// the tokio version reported by having two messages.
    #[test]
    fn policy_flag_ends_a_session_blocked_on_recv() {
        use std::sync::Arc;
        use std::sync::atomic::{AtomicBool, Ordering};
        use std::time::Duration;

        use super::WsClientSession;
        use crate::libs::signal::Shutdown;
        use crate::libs::ws::outbound;
        use crate::libs::ws::{WebsocketServer, WsServerConfig};

        let (_tx, rx) = outbound::channel(1);
        let parked = Arc::new(AtomicBool::new(false));
        let end = Arc::new(Shutdown::new());
        let mut session = WsClientSession::new(
            test_conn(),
            Box::new(PendingStream {
                parked: Arc::clone(&parked),
            }),
            rx,
            Arc::new(WebsocketServer::new(WsServerConfig::default())),
        );
        session.bind_end(Arc::clone(&end));

        let flagged = Arc::new(AtomicBool::new(false));
        let flag_once_parked = {
            let parked = Arc::clone(&parked);
            let flagged = Arc::clone(&flagged);
            let end = Arc::clone(&end);
            async move {
                // `yield_now` wakes before it returns `Pending`, so this hands
                // the poll back to the session arm rather than stalling on a
                // wake that is nobody's job to send.
                while !parked.load(Ordering::Acquire) {
                    nagoya::yield_now().await;
                }
                end.cancel();
                flagged.store(true, Ordering::Release);
            }
        };

        let finished = nagoya::block_on(nagoya::timeout(
            Duration::from_secs(1),
            futures::future::join(session.run(), flag_once_parked),
        ));
        assert!(
            finished.is_ok(),
            "{}",
            if flagged.load(Ordering::Acquire) {
                "session kept running after the policy flag was set"
            } else {
                "session never waited on the connection"
            }
        );
    }
}