alkhttp 0.5.0

//! The WS upgrade route (`/alk/channels`, ADR-067) and the channels
//! session it establishes — the server producer half.
//!
//! Per `docs/architecture/websocket.md`: bearer auth on the upgrade
//! request (`401` without a resolvable token — the same
//! `resolve_from_token` path as any HTTP request), then upgrade →
//! WS↔byte-stream adapter → `Connection::from_bidi(ws_stream,
//! b"alk/channels")` → alkcall `ChannelsAdapter` (the channels accept
//! path). The `install_channel_zero` hook forks the deployment's base
//! registry, registers the per-session ops on the fork (the generic
//! channel lifecycle ops, the deployment's openable ALPNs, the
//! bootstrap discovery set, and `op/register` — alkcall ADR-047 §4
//! amendment #2 + ADR-022 amendment), constructs channel 0's
//! `CallConnection` (the identity rides the channels-layer
//! `Connection`), retains it in the session registry (WS-26), and
//! runs the shared `Dispatcher::run_loop_single_stream` over the fork.
//!
//! ## Detached task lifetime semantics (WS-10)
//!
//! Two task families spawned here are **detached by design** and
//! outlive the session task that spawned them:
//!
//! - the WS pump tasks (the `WsPumps` pair behind the byte adapter,
//!   re-exported from `crate::websocket`), and
//! - the channel-0 dispatcher task spawned by `install_channel_zero`
//!   per accepted channels connection.
//!
//! If the channels session task (`run_channels_session`) dies —
//! upgrade-time early return, an adapter fault, or its own task being
//! cancelled — these tasks keep running **self-healing**: each ends on
//! its own when its stream half closes (peer disconnect, peer close
//! frame, read error, or the local teardown arms —
//! `WsSessions::abort`, the idle-read timeout, `AsyncWrite::shutdown`).
//! They are never leaked unconditionally: the leak window is bounded
//! by peer behavior (a peer that holds the socket open keeps the pump
//! and dispatcher tasks alive with it) and additionally by the WS-01
//! idle-read timeout when configured, and every session's pumps stay
//! force-evictable through the WS-08 registry
//! ([`WsSessions::abort`]) for the session's whole lifetime. The
//! dispatcher task specifically ends when channel 0's `BiStream` read
//! side hits EOF (its only exit condition), i.e. when the underlying
//! WS connection ends by any of the paths above.

use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;

use alkcall::channels::adapter::ChannelsAdapter;
use alkcall::channels::operations::{ChannelCore, OpenEstablisher, OpenHandler};
use alkcall::channels::policy::{ChannelLifecyclePolicy, NoCap};
use alkcall::core::auth::{AuthContext, Identity};
use alkcall::core::types::{Connection, ProtocolHandler};
use alkcall::registry::registration::OperationRegistry;
use alkcall::registry::spec::OperationSpec;
use axum::extract::ws::WebSocketUpgrade;
use axum::http::StatusCode;
use axum::response::{IntoResponse, Response};
use parking_lot::Mutex;

use super::byte_adapter::{
    split_ws_to_bytes_idle_with_write, WsPumps, INBOUND_WS_FRAME_CAP, INBOUND_WS_MESSAGE_CAP,
};

/// Registry of live WS session pump handles (WS-08). The upgrade
/// handler retains each session's [`WsPumps`] handle here for the
/// session's lifetime so a stuck session is evictable in-crate:
/// [`WsSessions::abort`] forces the pumps' tasks to end (the socket
/// halves close and the channels session unwinds). The handle for a
/// session is removed when the session task finishes (self-removing
/// guard), so the registry only holds live sessions. Assembly layers
/// share one instance via the adapter's `RouterState` (the upgrade
/// handler registers against it) or per-route request extensions (an
/// extension clone takes precedence).
///
/// Un-registered (default) — the upgrade runs fine and simply keeps no
/// eviction lever, matching the pre-WS-08 behavior.
#[derive(Clone, Default)]
pub struct WsSessions {
    counter: Arc<AtomicU64>,
    sessions: Arc<Mutex<HashMap<u64, Arc<WsPumps>>>>,
    connections: Arc<Mutex<HashMap<u64, Arc<alkcall::protocol::connection::CallConnection>>>>,
}

/// Default bound on concurrent WS sessions (WS-09): the semaphore
/// initial capacity in [`SessionState`]; a deployment overrides it with
/// `HttpAdapter::with_ws_max_sessions`.
pub const DEFAULT_WS_MAX_SESSIONS: usize = 64;

/// The upgrade handler's state slice: what it needs beyond the request
/// itself. Axum lifts it via `FromRef` from either full router state —
/// the adapter's `RouterState` (carrying the shared [`WsSessions`]
/// instance and the configured session cap) or a bare
/// `Arc<OperationRegistry>` (custom upgrade routes / integration tests
/// get a handler-private registry, a **per-request** session-cap
/// semaphore and the default knobs; eviction still works in-crate,
/// just not shared — and the per-request semaphore bounds nothing
/// across requests, review 007 WS-30: a route that needs an effective
/// shared cap inserts a [`SessionSlots`] extension; `FromRef`
/// extraction is per request with no caching).
#[derive(Clone)]
pub struct SessionState {
    registry: Arc<OperationRegistry>,
    sessions: Arc<WsSessions>,
    /// Session cap (WS-09): acquired post-auth, pre-upgrade; a caller
    /// over the cap is rejected with 503.
    session_slots: Arc<tokio::sync::Semaphore>,
    /// Idle-read timeout (WS-01): `None` disables the knob.
    idle_timeout: Option<std::time::Duration>,
    /// The openable-ALPN set (WS-22): `None` declares no openables.
    openable_alpns: Option<Arc<[OpenableAlpn]>>,
    /// The `op/register` surface's `AccessControl` (review 007 WS-29).
    op_register_acl: alkcall::registry::spec::AccessControl,
}

impl SessionState {
    /// From the plain registry state (custom upgrade routes /
    /// integration tests): sessions default to a fresh [`WsSessions`]
    /// private to the handler (eviction still functional in-crate but
    /// not shared with the assembly layer) and a **per-request**
    /// default session-cap semaphore — built by `FromRef` per request
    /// (no caching), so it bounds nothing across requests (review 007
    /// WS-30); a route needing an effective shared cap inserts a
    /// [`SessionSlots`] extension.
    pub(crate) fn from_registry(registry: &Arc<OperationRegistry>) -> Self {
        Self {
            registry: Arc::clone(registry),
            sessions: Arc::new(WsSessions::new()),
            session_slots: Arc::new(tokio::sync::Semaphore::new(DEFAULT_WS_MAX_SESSIONS)),
            idle_timeout: Some(crate::websocket::DEFAULT_WS_IDLE_TIMEOUT),
            openable_alpns: None,
            op_register_acl: alkcall::registry::spec::AccessControl::default(),
        }
    }

    /// From the adapter's router state: the shared [`WsSessions`]
    /// instance the assembly layer can hold for eviction, the
    /// pre-built session-cap semaphore (one per `HttpAdapter`, shared
    /// across requests), the deployment's openable-ALPN set, and the
    /// `op/register` ACL.
    pub(crate) fn new(
        registry: Arc<OperationRegistry>,
        sessions: Arc<WsSessions>,
        session_slots: Arc<tokio::sync::Semaphore>,
        idle_timeout: Option<std::time::Duration>,
        openable_alpns: Option<Arc<[OpenableAlpn]>>,
        op_register_acl: alkcall::registry::spec::AccessControl,
    ) -> Self {
        Self {
            registry,
            sessions,
            session_slots,
            idle_timeout,
            openable_alpns,
            op_register_acl,
        }
    }

    pub(crate) fn registry(&self) -> &Arc<OperationRegistry> {
        &self.registry
    }

    pub(crate) fn sessions(&self) -> &Arc<WsSessions> {
        &self.sessions
    }

    pub(crate) fn session_slots(&self) -> &Arc<tokio::sync::Semaphore> {
        &self.session_slots
    }

    pub(crate) fn idle_timeout(&self) -> Option<std::time::Duration> {
        self.idle_timeout
    }

    pub(crate) fn openable_alpns(&self) -> Option<Arc<[OpenableAlpn]>> {
        self.openable_alpns.clone()
    }

    pub(crate) fn op_register_acl(&self) -> alkcall::registry::spec::AccessControl {
        self.op_register_acl.clone()
    }
}

/// `FromRef` chain: a bare `Arc<OperationRegistry>` router state lifts
/// into the handler's [`SessionState`]; a full `RouterState` carries
/// the shared [`WsSessions`] instance and lifts through its own impl.
impl axum::extract::FromRef<Arc<OperationRegistry>> for SessionState {
    fn from_ref(registry: &Arc<OperationRegistry>) -> Self {
        SessionState::from_registry(registry)
    }
}

impl WsSessions {
    /// A fresh, empty session registry.
    pub fn new() -> Self {
        Self::default()
    }

    /// Abort every live session's pump tasks (forced teardown).
    pub fn abort(&self) {
        for (_, pumps) in self.sessions.lock().drain() {
            pumps.abort();
        }
    }

    /// Number of live sessions currently tracked.
    pub fn len(&self) -> usize {
        self.sessions.lock().len()
    }

    /// Whether no sessions are tracked.
    pub fn is_empty(&self) -> bool {
        self.sessions.lock().is_empty()
    }

    /// The live session's channel-0
    /// [`CallConnection`](alkcall::protocol::connection::CallConnection)
    /// handles (WS-26):
    /// the deployment-visible surface the assembly layer (or the hub)
    /// reaches browser/session-side ops through — connection-local
    /// overlay composition, `announce_op`-style pushes, session
    /// enumeration. The handle is registered when the session's
    /// channel-0 connection is built and removed when the channel-0
    /// task ends (session teardown of any path), so the registry only
    /// holds live sessions.
    ///
    /// Un-registered (default) — no handles are retained; the upgrade
    /// runs exactly as the pre-WS-26 shape.
    pub fn live_connections(&self) -> Vec<Arc<alkcall::protocol::connection::CallConnection>> {
        self.connections.lock().values().cloned().collect()
    }

    /// Number of live channel-0 connection handles currently tracked.
    pub fn live_connection_count(&self) -> usize {
        self.connections.lock().len()
    }

    fn insert(&self, pumps: Arc<WsPumps>) -> u64 {
        let id = self.counter.fetch_add(1, Ordering::Relaxed);
        self.sessions.lock().insert(id, pumps);
        id
    }

    fn remove(&self, id: u64) {
        self.sessions.lock().remove(&id);
    }

    fn insert_connection(&self, conn: Arc<alkcall::protocol::connection::CallConnection>) -> u64 {
        let id = self.counter.fetch_add(1, Ordering::Relaxed);
        self.connections.lock().insert(id, conn);
        id
    }

    fn remove_connection(&self, id: u64) {
        self.connections.lock().remove(&id);
    }
}

/// One deployment-declared openable ALPN for the WS path (WS-22):
/// the per-ALPN open-op `OperationSpec` (with the `channel_open`
/// marker set via `OperationSpec::with_channel_open`), the
/// ALPN-specific [`OpenHandler`] the data-plane protocol runs on the
/// allocated channel's `Connection`, and the optional establishment
/// phase (ADR-049) the open-op wrapper awaits — bounded — before the
/// reply. The ALPN-specific handler and establisher stay in the ALPN
/// crates (alktty et al.); this crate only ferries the registration
/// onto each session's fork.
#[derive(Clone)]
pub struct OpenableAlpn {
    /// The open-op spec (Query/Mutation/Sub with the `channel_open`
    /// marker; a `Pub` open op resolves `channel:pub_open_not_implemented`
    /// per the upstream C-08 stub until channel adoption lands).
    pub spec: OperationSpec,
    /// The data-plane protocol handler spawned on the allocated
    /// channel's `Connection` — `Fn(Value, Option<ChannelPlan>,
    /// Connection, AuthContext) -> JoinHandle<()>`. With an
    /// establisher attached, the plan (alkcall 0.6.0 /
    /// ADR-049 amendment 2) is the establisher's `Establishment.plan`:
    /// the establisher and the handler agree on the concrete type
    /// (downcast in the ALPN crate); `None` when no establisher is
    /// registered or it returned `Establishment::default()`
    /// (`alkcall::channels::operations::Establishment::default`).
    pub open_handler: OpenHandler,
    /// The awaited establishment phase (ADR-049 §1): validate params
    /// semantically, consult ownership, prepare/dial the backend —
    /// before the open reply. `None` (the default) = an always-OK
    /// establisher (the pre-ADR-049 shape; existing registrations
    /// behave unchanged). A successful establisher returns
    /// `Establishment::new(plan)` (to deliver a channel plan to the
    /// pump handler) or `Establishment::default()` when it only
    /// validates — both in
    /// `alkcall::channels::operations`.
    pub establisher: Option<OpenEstablisher>,
    /// The per-registration bound on the establisher await
    /// (ADR-049 §2). `None` = [`ESTABLISHMENT_TIMEOUT`] (10s) when the
    /// dispatch carries no deadline; the effective bound is the
    /// earlier of the dispatch deadline and this override.
    ///
    /// [`ESTABLISHMENT_TIMEOUT`]: alkcall::channels::operations::ESTABLISHMENT_TIMEOUT
    pub establisher_timeout: Option<std::time::Duration>,
}

impl OpenableAlpn {
    /// Declare one openable ALPN (no establisher — the pre-ADR-049
    /// shape; the open op replies as soon as the pump handler is
    /// spawned).
    pub fn new(spec: OperationSpec, open_handler: OpenHandler) -> Self {
        Self {
            spec,
            open_handler,
            establisher: None,
            establisher_timeout: None,
        }
    }

    /// Attach an establishment phase (ADR-049 §1): awaited by the
    /// open-op wrapper — bounded by the dispatch deadline, this
    /// crate's [`ESTABLISHMENT_TIMEOUT`] default, or
    /// `timeout` when set — before the reply; on failure the open op
    /// resolves `channel:open_failed` with `details.reason` and the
    /// channel never exists consumer-side.
    ///
    /// [`ESTABLISHMENT_TIMEOUT`]: alkcall::channels::operations::ESTABLISHMENT_TIMEOUT
    pub fn with_establisher(
        mut self,
        establisher: OpenEstablisher,
        timeout: Option<std::time::Duration>,
    ) -> Self {
        self.establisher = Some(establisher);
        self.establisher_timeout = timeout;
        self
    }
}

/// The channels session for an upgraded socket: adapt → `Connection`
/// (identity attached) → `ChannelsAdapter::handle`. `policy` gates
/// data-channel opens (ADR-041). When `sessions` is `Some`, the
/// session's pump handle is registered for the session's lifetime —
/// the WS-08 eviction lever ([`WsSessions::abort`]). `idle_timeout`
/// bounds the read stall (WS-01): `None` disables the knob, `Some(d)`
/// closes the read with 1001 after `d` without inbound chunk
/// progress. `write_timeout` bounds one outbound WS send (WS-18):
/// `None` = the crate default window, `Some(d)` a deployment-set
/// window. `openable_alpns` (WS-22) is the deployment's openable-ALPN
/// set registered per session (see the `install_channel_zero` hook).
/// `op_register_acl` (review 007 WS-29) is the `op/register`
/// surface's `AccessControl` — `AccessControl::default()` (the
/// permissive crate default) unless a deployment restricts which
/// authenticated peers may announce ops.
#[allow(clippy::too_many_arguments)]
pub async fn run_channels_session(
    socket: axum::extract::ws::WebSocket,
    registry: Arc<OperationRegistry>,
    identity: Identity,
    policy: Arc<dyn ChannelLifecyclePolicy>,
    sessions: Option<WsSessions>,
    idle_timeout: Option<std::time::Duration>,
    write_timeout: Option<std::time::Duration>,
    openable_alpns: Option<Arc<[OpenableAlpn]>>,
    op_register_acl: alkcall::registry::spec::AccessControl,
) {
    let (byte_stream, pumps) =
        split_ws_to_bytes_idle_with_write(socket, idle_timeout, write_timeout);
    let pumps = Arc::new(pumps);

    let _guard = sessions.as_ref().map(|s| {
        let id = s.insert(Arc::clone(&pumps));
        SessionGuard {
            sessions: s.clone(),
            id,
        }
    });

    let conn = Connection::from_bidi(byte_stream, b"alk/channels".to_vec(), None);
    let _ = conn.set_identity(identity.clone());

    let adapter = ChannelsAdapter::new(
        install_channel_zero(
            registry,
            sessions,
            Arc::clone(&policy),
            openable_alpns,
            op_register_acl,
        ),
        policy,
    );
    let auth = AuthContext {
        identity: Some(identity),
        alpn: b"alk/channels".to_vec(),
        remote_addr: None,
        tls_client_fingerprint: None,
    };
    if let Err(e) = ProtocolHandler::handle(&adapter, conn, &auth).await {
        tracing::warn!(error = %e, "channels session ended");
    }
}

struct SessionGuard {
    sessions: WsSessions,
    id: u64,
}

impl Drop for SessionGuard {
    fn drop(&mut self) {
        self.sessions.remove(self.id);
    }
}

/// The channel-0 connection handle's self-removing guard (WS-26): the
/// drop removes the retained [`CallConnection`] from the session
/// registry, so the registry only holds live sessions' handles.
struct ConnectionGuard {
    sessions: WsSessions,
    id: u64,
}

impl Drop for ConnectionGuard {
    fn drop(&mut self) {
        self.sessions.remove_connection(self.id);
    }
}

/// The `install_channel_zero` hook (WS-20/21/22/25/26, review 006
/// Unit 2): fork the deployment's base registry, register the
/// per-session ops on the fork — the generic channel lifecycle ops
/// (`ChannelOperations::register_on`: `channel/close`,
/// `channel/control`, `channel/resources/subscribe`), the
/// deployment's openable ALPNs (`ChannelCore::register_openable_with_establisher`),
/// the bootstrap discovery set closed over the fork
/// (`install_bootstrap_discovery`, so `services/list` sees the
/// session's own openables — the F-06 shape), and `op/register`
/// (alkcall ADR-022 amendment; the collision set is the fork) — then
/// run the dispatcher over the fork (alkcall ADR-047 §4 amendment #2,
/// the per-session-fork mechanism; the alkcall e2e reference shape).
///
/// `policy` is the session's channel cap policy (the `ChannelsPolicy`
/// extension resolution, `NoCap` default): the open-op wrappers the
/// hook registers consult it (`check_open` per identity, ledger
/// decrement on teardown), and the adapter's demux loop decrements it
/// on connection-drop teardown — one policy instance across both
/// halves.
///
/// The hook also retains the channel-0 `CallConnection` in the
/// session registry (WS-26) when a shared [`WsSessions`] instance is
/// in play: the deployment reaches the session's connection-local
/// overlay and can hub→session-call announced or imported ops. The
/// handle is removed when the channel-0 task ends (any teardown path).
///
/// `op_register_acl` (review 007 WS-29) is the `AccessControl` the
/// per-session `op/register` op is registered with — the announce
/// surface's gate. The default is `AccessControl::default()` (the
/// SRV-10 permissive-crate-default precedent, review-006 UP-02's
/// recorded posture); a deployment restricting which authenticated
/// peers may announce ops threads a stricter value —
/// `HttpAdapter::with_ws_op_register_acl` on the built-in surface or
/// the [`OpRegisterAcl`] request extension on bare-registry/custom
/// upgrade routes (both mirroring the openables/policy threading).
///
/// The dispatcher's token resolver is a no-op: the WS identity is
/// attached to the connection at upgrade time and
/// `Dispatcher::resolve_identity` falls back to it when the payload
/// carries no `auth_token` — the correct accept-side behavior (the
/// identity was established at upgrade time, not per-call).
fn install_channel_zero(
    registry: Arc<OperationRegistry>,
    sessions: Option<WsSessions>,
    policy: Arc<dyn ChannelLifecyclePolicy>,
    openable_alpns: Option<Arc<[OpenableAlpn]>>,
    op_register_acl: alkcall::registry::spec::AccessControl,
) -> alkcall::channels::adapter::InstallChannelZero {
    Arc::new(move |manager, channel0_conn, auth| {
        let registry = Arc::clone(&registry);
        let sessions = sessions.clone();
        let policy = Arc::clone(&policy);
        let openable_alpns = openable_alpns.clone();
        let op_register_acl = op_register_acl.clone();
        tokio::spawn(async move {
            // The WS identity rides the upgrade request; propagate it to
            // channel 0's `CallConnection` so the dispatcher's
            // `resolve_identity` (and thus `AccessControl::check` on
            // `services/list` and every operation) sees it. Without this
            // the freshly constructed channel-0 `Connection` carries no
            // identity and all ACL-restricted ops look unauthenticated.
            if let Some(identity) = auth.identity.clone() {
                let _ = channel0_conn.set_identity(identity);
            }
            let channel0_bidi = match channel0_conn.accept_bi().await {
                Ok(s) => s,
                Err(_) => return,
            };
            let (writer, reader) =
                alkcall::protocol::connection::split_single_stream(channel0_bidi);
            let call_connection = Arc::new(
                alkcall::protocol::connection::CallConnection::new_single_stream(
                    channel0_conn,
                    Arc::clone(&writer),
                ),
            );

            // The per-session dispatch registry (WS-24's decided
            // mechanism): fork the base, register the session's ops on
            // the fork, dispatch over the fork. Registration failures
            // (a spec whose publish_schema fails to compile — the only
            // error class the registry can produce here, since the
            // openable specs were accepted at declaration) kill the
            // session before the loop starts: a session that would
            // mis-discover must not silently lose the failing op.
            let fork = Arc::new(registry.fork());

            let register_result: Result<(), String> = (|| {
                alkcall::channels::operations::ChannelOperations::new(
                    manager.clone(),
                    Arc::clone(&policy),
                )
                .register_on(&fork)?;
                if let Some(openables) = openable_alpns.as_ref() {
                    let core = ChannelCore::new(manager, Arc::clone(&policy));
                    for openable in openables.iter() {
                        core.register_openable_with_establisher(
                            openable.spec.clone(),
                            openable.establisher.clone(),
                            Arc::clone(&openable.open_handler),
                            &fork,
                            auth.clone(),
                            openable.establisher_timeout,
                        )?;
                    }
                }
                alkcall::registry::discovery::install_bootstrap_discovery(&fork)?;
                fork.register(alkcall::registry::registration::HandlerRegistration::new(
                    alkcall::registry::op_register::op_register_spec(op_register_acl),
                    alkcall::registry::registration::HandlerKind::Once(
                        alkcall::registry::op_register::op_register_handler(
                            Arc::clone(&call_connection),
                            Arc::clone(&fork),
                        ),
                    ),
                    alkcall::registry::registration::OperationProvenance::Local,
                    None,
                    None,
                    alkcall::core::types::Capabilities::new(),
                ))?;
                Ok(())
            })();
            if let Err(e) = register_result {
                tracing::error!(error = %e, "channel-0 session registry setup failed");
                return;
            }

            let _conn_guard = sessions.as_ref().map(|sessions| {
                let id = sessions.insert_connection(Arc::clone(&call_connection));
                ConnectionGuard {
                    sessions: sessions.clone(),
                    id,
                }
            });
            // `conn_guard` lives in the dispatcher task's frame: its drop
            // removes the handle when the loop below returns — any
            // teardown path (peer close, idle eviction, abort, EOF).

            let dispatcher = alkcall::protocol::dispatch::Dispatcher::new(
                fork,
                std::sync::Arc::new(NoopProvider),
            );
            dispatcher
                .run_loop_single_stream(call_connection, reader, writer)
                .await;
        })
    })
}

/// The bare-registry variant of the hook for the from_wss test
/// server: no session registry (no WS-26 handle retention — the test
/// server's pump slot is its own teardown lever) and no openables.
#[cfg(all(test, feature = "server", feature = "wss"))]
pub(crate) fn adapter_install_channel_zero(
    registry: Arc<OperationRegistry>,
) -> alkcall::channels::adapter::InstallChannelZero {
    install_channel_zero(
        registry,
        None,
        Arc::new(NoCap),
        None,
        alkcall::registry::spec::AccessControl::default(),
    )
}

struct NoopProvider;

impl alkcall::core::auth::IdentityProvider for NoopProvider {
    fn resolve_from_fingerprint(&self, _: &str) -> Option<Identity> {
        None
    }
    fn resolve_from_token(&self, _: &alkcall::core::auth::AuthToken) -> Option<Identity> {
        None
    }
}

/// Extension wrapper for the channels-policy injection point (SRV-10):
/// a deployment inserts `ChannelsPolicy(Arc<dyn ChannelLifecyclePolicy>)`
/// into the request extensions (a route layer on the WS route) to gate
/// data-channel opens per identity (ADR-041). Without the extension the
/// upgrade defaults to [`NoCap`] — the crate default for the built-in
/// surface (POC/trusted-peer semantics); assembly layers that build
/// their own upgrade route pass a stricter policy directly to
/// [`run_channels_session`].
#[derive(Clone)]
pub struct ChannelsPolicy(pub Arc<dyn ChannelLifecyclePolicy>);

/// Per-request WS pump timeouts (WS-17): a deployment inserts
/// `WsTimeouts` into the request extensions (a route layer on the WS
/// route, mirroring [`ChannelsPolicy`]) to set the pump knobs per
/// route instead of the router-state defaults. `idle` (WS-01) bounds
/// the read staleness (no completed inbound chunk for the window →
/// 1001 eviction); `write` (WS-18) bounds one outbound WS send (a peer
/// that stops reading is evicted once a single send outlasts it);
/// `None` disables a knob. Without the extension the
/// [`SessionState`] values apply — the built-in surface carries
/// `HttpAdapter::with_ws_idle_timeout` /
/// `DEFAULT_WS_IDLE_TIMEOUT` for the read side and
/// [`crate::websocket::DEFAULT_WS_WRITE_TIMEOUT`] for the write side,
/// which is also what bare-registry routes (custom upgrade routes on
/// a plain `Arc<OperationRegistry>` state) get: 60 s idle + 60 s
/// write windows, a handler-private [`WsSessions`], and a per-request
/// (no effective) session cap — see [`SessionSlots`] for the
/// shared-cap surface.
#[derive(Clone, Copy, Debug)]
pub struct WsTimeouts {
    /// Idle-read window (WS-01); `None` disables the eviction.
    pub idle: Option<std::time::Duration>,
    /// Write-progress window (WS-18); `None` selects the crate
    /// default (`DEFAULT_WS_WRITE_TIMEOUT`).
    pub write: Option<std::time::Duration>,
}

/// Per-request openable-ALPN override (WS-22): a deployment inserts
/// `OpenableAlpns` into the request extensions (a route layer on the WS
/// route, mirroring [`ChannelsPolicy`] / [`WsTimeouts`]) to set the
/// session's openable set per route instead of the router-state
/// default. Without the extension the [`SessionState`] value applies —
/// [`HttpAdapter::with_ws_openable_alpns`](crate::server::HttpAdapter::with_ws_openable_alpns)
/// for the built-in surface; bare-registry routes carry none.
#[derive(Clone)]
pub struct OpenableAlpns(pub Arc<[OpenableAlpn]>);

/// Per-request `op/register` ACL (review 007 WS-29): a deployment
/// inserts `OpRegisterAcl` into the request extensions (a route layer
/// on the WS route, mirroring [`ChannelsPolicy`] / [`WsTimeouts`] /
/// [`OpenableAlpns`]) to gate the announce surface per route — the
/// per-session `op/register` op is registered with this
/// `AccessControl`, so a peer whose identity does not satisfy it gets
/// `FORBIDDEN` on the announce. Without the extension the
/// [`SessionState`] value applies —
/// [`HttpAdapter::with_ws_op_register_acl`](crate::server::HttpAdapter::with_ws_op_register_acl)
/// for the built-in surface; the default everywhere is
/// `AccessControl::default()` (any authenticated peer may announce).
#[derive(Clone, Debug)]
pub struct OpRegisterAcl(pub alkcall::registry::spec::AccessControl);

/// Per-request session-cap slots (review 007 WS-30): a deployment
/// inserts `SessionSlots` into the request extensions (a route layer on
/// the WS route, mirroring [`ChannelsPolicy`] / [`WsTimeouts`] /
/// [`OpenableAlpns`] / [`OpRegisterAcl`]) to bound the route's
/// concurrent WS sessions with a shared `Arc<Semaphore>` — one permit
/// per upgrade, held for the session's lifetime; over-cap callers are
/// rejected with 503 (WS-09). Without the extension the
/// [`SessionState`] value applies — the built-in surface's
/// adapter-wide semaphore
/// (`HttpAdapter::with_ws_max_sessions`, one per `HttpAdapter`,
/// shared across requests). A bare-registry route's state is built per
/// request (`FromRef`, no caching), so its semaphore bounds nothing
/// across requests; such a route inserts this extension to get an
/// effective shared cap. Wrap the semaphore in
/// `SessionSlots(Arc::new(tokio::sync::Semaphore::new(n)))` at the
/// route layer.
#[derive(Clone)]
pub struct SessionSlots(pub Arc<tokio::sync::Semaphore>);

/// The upgrade handler. Requires the resolved identity in request
/// extensions (stashed by [`ws_bearer_auth`]) — a WS session without
/// an identity cannot run `AccessControl::check`.
///
/// The channel lifecycle policy comes from the
/// [`ChannelsPolicy`] request extension when present, else `NoCap`.
/// The session pump handles are retained in the [`WsSessions`]
/// registry — the shared instance from the router state
/// (`RouterState::ws_sessions`) unless a request extension carries
/// one, so a stuck session stays evictable (WS-08).
///
/// The pump timeout knobs (WS-01 read / WS-18 write) come from the
/// [`WsTimeouts`] request extension when present, else from the
/// router state (`HttpAdapter::with_ws_idle_timeout` for the read;
/// the write side is fixed at
/// [`crate::websocket::DEFAULT_WS_WRITE_TIMEOUT`] on the built-in
/// surface). Bare-registry routes (`Arc<OperationRegistry>` state)
/// carry no configurable state — they run the documented defaults
/// above; a deployment overriding them inserts the `WsTimeouts`
/// extension on its upgrade route.
///
/// The openable-ALPN set (WS-22) comes from the [`OpenableAlpns`]
/// request extension when present, else from the router state
/// (`HttpAdapter::with_ws_openable_alpns`); the default is no
/// openables (channel 0 only).
///
/// The `op/register` ACL (review 007 WS-29) comes from the
/// [`OpRegisterAcl`] request extension when present, else from the
/// router state (`HttpAdapter::with_ws_op_register_acl`); the default
/// is `AccessControl::default()` (any authenticated peer may announce).
///
/// Session cap (WS-09): one semaphore permit is acquired per upgrade,
/// post-auth and pre-upgrade; when the configured cap
/// (`HttpAdapter::with_ws_max_sessions`, default
/// [`DEFAULT_WS_MAX_SESSIONS`]) is exhausted the upgrade is rejected
/// with **503 Service Unavailable** — holding the permit for the
/// session's lifetime, so ended sessions free their slot. The permit
/// source is the [`SessionSlots`] request extension when present, else
/// the router state's semaphore; a bare-registry route's state is
/// built per request, so its semaphore bounds nothing across requests
/// (review 007 WS-30) — a route that needs an effective shared cap
/// inserts the `SessionSlots` extension.
#[allow(clippy::too_many_arguments)]
pub async fn ws_upgrade_handler(
    sessions: Option<axum::Extension<WsSessions>>,
    axum::extract::State(state): axum::extract::State<SessionState>,
    axum::Extension(identity): axum::Extension<Identity>,
    policy: Option<axum::Extension<ChannelsPolicy>>,
    timeouts: Option<axum::Extension<WsTimeouts>>,
    openables: Option<axum::Extension<OpenableAlpns>>,
    op_register_acl: Option<axum::Extension<OpRegisterAcl>>,
    session_slots: Option<axum::Extension<SessionSlots>>,
    ws_upgrade: WebSocketUpgrade,
) -> Response {
    let session_slots = session_slots
        .map(|axum::Extension(s)| s.0)
        .unwrap_or_else(|| Arc::clone(state.session_slots()));
    let Ok(permit) = session_slots.try_acquire_owned() else {
        return (
            StatusCode::SERVICE_UNAVAILABLE,
            "503 Service Unavailable: WS session cap reached",
        )
            .into_response();
    };
    let sessions = Some(
        sessions
            .map(|axum::Extension(s)| s)
            .unwrap_or_else(|| WsSessions::clone(state.sessions())),
    );
    let policy = policy
        .map(|axum::Extension(p)| p.0)
        .unwrap_or_else(|| Arc::new(NoCap));
    let registry = Arc::clone(state.registry());
    let idle_timeout = match timeouts {
        Some(axum::Extension(t)) => t.idle,
        None => state.idle_timeout(),
    };
    let write_timeout = timeouts.and_then(|t| t.write);
    let openable_alpns = openables
        .map(|axum::Extension(o)| o.0)
        .or_else(|| state.openable_alpns());
    let op_register_acl = op_register_acl
        .map(|axum::Extension(a)| a.0)
        .unwrap_or_else(|| state.op_register_acl());
    ws_upgrade
        .max_frame_size(INBOUND_WS_FRAME_CAP)
        .max_message_size(INBOUND_WS_MESSAGE_CAP)
        .on_upgrade(move |socket| async move {
            let _permit = permit;
            run_channels_session(
                socket,
                registry,
                identity,
                policy,
                sessions,
                idle_timeout,
                Some(write_timeout.unwrap_or(crate::websocket::DEFAULT_WS_WRITE_TIMEOUT)),
                openable_alpns,
                op_register_acl,
            )
            .await
        })
}

/// Bearer-auth middleware for the WS upgrade route: resolves the token
/// via the shared [`crate::server::auth`] path and stashes the identity
/// for the upgrade handler. No token / unresolvable token → `401`
/// before the upgrade.
pub async fn ws_bearer_auth(
    axum::extract::State(provider): axum::extract::State<
        Arc<dyn alkcall::core::auth::IdentityProvider>,
    >,
    mut req: axum::http::Request<axum::body::Body>,
    next: axum::middleware::Next,
) -> Response {
    let identity = crate::server::auth::extract_bearer_identity(&req, provider.as_ref());
    match identity {
        Some(identity) => {
            req.extensions_mut().insert(identity);
            next.run(req).await
        }
        None => (StatusCode::UNAUTHORIZED, "401 Unauthorized").into_response(),
    }
}

/// Test support: a minimal tokio-tungstenite WS client speaking raw
/// channels framing (`frame_channel0_chunk` / `ChunkAssembler` /
/// `FrameAssembler`). Gated behind the `test-support` feature so it
/// never ships in release builds; used by this crate's integration
/// tests and by downstream consumers testing their deployments.
#[cfg(any(test, feature = "test-support"))]
pub mod test_support {
    use alkcall::protocol::wire::EventEnvelope;
    use futures::StreamExt;
    use tokio_tungstenite::tungstenite::client::IntoClientRequest;
    use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode;

    /// Frame one `EventEnvelope` as a channel-0 chunk (8-byte chunk
    /// header + 4-byte length prefix + JSON body) — the client-side
    /// framing channel 0 uses over any transport.
    ///
    /// # Panics
    ///
    /// Panics only if `serde_json` cannot serialize the envelope —
    /// unreachable for the acyclic wire type (no non-string map keys,
    /// no untagged ambiguities), which is why this returns `Vec<u8>`
    /// rather than `Result`: a test helper returning `Result` for an
    /// impossible case is worse ergonomics than a documented panic
    /// (review 001 HY-04, kept-as-is decision — the item ships behind
    /// the opt-in `test-support` feature, the crate's documented
    /// exception to no-panics-in-library-code).
    pub fn frame_channel0_chunk(envelope: &EventEnvelope) -> Vec<u8> {
        let body = serde_json::to_vec(envelope).unwrap();
        let mut out = Vec::with_capacity(8 + 4 + body.len());
        out.extend_from_slice(&0u32.to_be_bytes());
        out.extend_from_slice(&((body.len() + 4) as u32).to_be_bytes());
        out.extend_from_slice(&(body.len() as u32).to_be_bytes());
        out.extend_from_slice(&body);
        out
    }

    /// Accumulate WS binary messages into bytes and extract complete
    /// chunks: (channel_id, payload).
    #[derive(Default)]
    pub struct ChunkAssembler {
        buf: Vec<u8>,
    }

    impl ChunkAssembler {
        /// A fresh, empty assembler.
        pub fn new() -> Self {
            Self::default()
        }

        /// Append raw bytes to the assembly buffer.
        pub fn push(&mut self, bytes: &[u8]) {
            self.buf.extend_from_slice(bytes);
        }

        /// Extract the next complete `(channel_id, payload)` chunk, if
        /// a full one is buffered (8-byte header + declared payload
        /// length).
        pub fn next_chunk(&mut self) -> Option<(u32, Vec<u8>)> {
            if self.buf.len() < 8 {
                return None;
            }
            let channel_id =
                u32::from_be_bytes([self.buf[0], self.buf[1], self.buf[2], self.buf[3]]);
            let len =
                u32::from_be_bytes([self.buf[4], self.buf[5], self.buf[6], self.buf[7]]) as usize;
            if self.buf.len() < 8 + len {
                return None;
            }
            let payload = self.buf.drain(..8 + len).skip(8).collect();
            Some((channel_id, payload))
        }
    }

    /// Reassemble length-prefixed frames from the concatenated byte
    /// stream of channel-0 chunk payloads. POC finding (OQ-01): one
    /// call frame may arrive as multiple chunks (`write_frame`'s
    /// prefix and body surface as separate mux payloads), so frame
    /// parsing must run over the reassembled byte stream — never over
    /// individual chunks.
    #[derive(Default)]
    pub struct FrameAssembler {
        buf: Vec<u8>,
    }

    impl FrameAssembler {
        /// A fresh, empty assembler.
        pub fn new() -> Self {
            Self::default()
        }

        /// Append raw bytes to the assembly buffer.
        pub fn push(&mut self, bytes: &[u8]) {
            self.buf.extend_from_slice(bytes);
        }

        /// Extract the next complete `EventEnvelope` frame, if a full
        /// length-prefixed frame is buffered and parses. Unparseable
        /// frames are dropped (test-only surface).
        pub fn next_frame(&mut self) -> Option<EventEnvelope> {
            if self.buf.len() < 4 {
                return None;
            }
            let len =
                u32::from_be_bytes([self.buf[0], self.buf[1], self.buf[2], self.buf[3]]) as usize;
            if self.buf.len() < 4 + len {
                return None;
            }
            let frame: Vec<u8> = self.buf.drain(..4 + len).collect();
            serde_json::from_slice(&frame[4..]).ok()
        }
    }

    /// Minimal WS client for tests: connect with/without a bearer
    /// token, send/recv binary + text, await the close frame.
    pub struct WsClient {
        sink: futures::stream::SplitSink<
            tokio_tungstenite::WebSocketStream<
                tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
            >,
            tokio_tungstenite::tungstenite::Message,
        >,
        stream: futures::stream::SplitStream<
            tokio_tungstenite::WebSocketStream<
                tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
            >,
        >,
    }

    impl WsClient {
        /// Connect with a `Bearer` token header; the full WS stream is
        /// returned (upgrade succeeded).
        pub async fn connect_authorized(url: &str, token: &str) -> Result<Self, String> {
            let mut request = url
                .into_client_request()
                .map_err(|e| format!("bad url: {e}"))?;
            request.headers_mut().insert(
                http::header::AUTHORIZATION,
                http::HeaderValue::from_str(&format!("Bearer {token}"))
                    .map_err(|e| format!("bad token: {e}"))?,
            );
            let (stream, _resp): (
                tokio_tungstenite::WebSocketStream<
                    tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
                >,
                _,
            ) = tokio_tungstenite::connect_async(request)
                .await
                .map_err(|e| format!("connect failed: {e}"))?;
            Ok(Self::from_stream(stream))
        }

        /// Connect and return just the HTTP status (for negative tests).
        pub async fn connect_status(url: &str, token: Option<&str>) -> Option<u16> {
            let mut request = url.into_client_request().ok()?;
            if let Some(t) = token {
                request.headers_mut().insert(
                    http::header::AUTHORIZATION,
                    http::HeaderValue::from_str(&format!("Bearer {t}")).ok()?,
                );
            }
            type WsStream = tokio_tungstenite::WebSocketStream<
                tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
            >;
            type WsConnectResult = Result<
                (
                    WsStream,
                    tokio_tungstenite::tungstenite::http::Response<Option<Vec<u8>>>,
                ),
                tokio_tungstenite::tungstenite::Error,
            >;
            let result: WsConnectResult = tokio_tungstenite::connect_async(request).await;
            match result {
                Ok((stream, resp)) => {
                    drop(stream);
                    Some(resp.status().as_u16())
                }
                Err(tokio_tungstenite::tungstenite::Error::Http(resp)) => {
                    Some(resp.status().as_u16())
                }
                Err(_) => None,
            }
        }

        fn from_stream(
            stream: tokio_tungstenite::WebSocketStream<
                tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
            >,
        ) -> Self {
            let (sink, stream) = stream.split();
            Self { sink, stream }
        }

        /// Send one binary WS message.
        ///
        /// # Panics
        ///
        /// Panics on a socket write failure — a test client that cannot
        /// send has a broken test, not a recoverable runtime state.
        pub async fn send_binary(&mut self, bytes: Vec<u8>) {
            self.send_binary_piece(&bytes).await;
        }

        /// Send one binary WS message, in pieces (for split-frame
        /// tests). Same panic contract as [`Self::send_binary`].
        pub async fn send_binary_piece(&mut self, bytes: &[u8]) {
            use futures::SinkExt;
            self.sink
                .send(tokio_tungstenite::tungstenite::Message::Binary(
                    bytes.to_vec().into(),
                ))
                .await
                .unwrap();
        }

        /// Send one text WS message. Same panic contract as
        /// [`Self::send_binary`].
        pub async fn send_text(&mut self, text: &str) {
            use futures::SinkExt;
            self.sink
                .send(tokio_tungstenite::tungstenite::Message::Text(
                    text.to_string().into(),
                ))
                .await
                .unwrap();
        }

        /// Next binary message, or `None` on timeout/close/error.
        pub async fn next_binary(&mut self, timeout: std::time::Duration) -> Option<Vec<u8>> {
            use futures::StreamExt;
            loop {
                match tokio::time::timeout(timeout, self.stream.next()).await {
                    Err(_) => return None,
                    Ok(None) => return None,
                    Ok(Some(Err(_))) => return None,
                    Ok(Some(Ok(m))) => match m {
                        tokio_tungstenite::tungstenite::Message::Binary(b) => {
                            return Some(b.to_vec())
                        }
                        tokio_tungstenite::tungstenite::Message::Close(_) => return None,
                        _ => continue,
                    },
                }
            }
        }

        /// Await the close frame: `Some(Some(code))` close with code,
        /// `Some(None)` stream ended without a close frame, `None`
        /// timed out.
        pub async fn next_close(&mut self, timeout: std::time::Duration) -> Option<Option<u16>> {
            use futures::StreamExt;
            match tokio::time::timeout(timeout, self.stream.next()).await {
                Err(_) => None,
                Ok(None) => Some(None),
                Ok(Some(Ok(tokio_tungstenite::tungstenite::Message::Close(cf)))) => {
                    Some(cf.map(|f| match f.code {
                        CloseCode::Error => 1011,
                        other => other.into(),
                    }))
                }
                Ok(Some(Ok(_))) => Box::pin(self.next_close(timeout)).await,
                Ok(Some(Err(_))) => Some(None),
            }
        }

        /// Close the WS with a normal-close frame.
        pub async fn close(&mut self) {
            use futures::SinkExt;
            let _ = self.sink.close().await;
        }
    }
}