alkhttp/tests/ws_upgrade_session.rs

//! WS upgrade + channels session integration tests (the ws-upgrade-session
//! acceptance gates): tokio-tungstenite client ↔ axum WS server
//! (upgrade → byte-adapter → ChannelsAdapter + channel-0 Dispatcher).
//! Grows from the ws-byte-adapter POC's validated suite.

use alkcall::core::auth::{Identity, IdentityProvider};
use alkcall::protocol::wire::{EventEnvelope, ResponseEnvelope, EVENT_RESPONDED};
use alkcall::registry::discovery::{services_list_handler, services_list_spec};
use alkcall::registry::registration::{
    make_handler, Handler, HandlerKind, HandlerRegistration, OperationProvenance, OperationRegistry,
};
use alkcall::registry::spec::{AccessControl, OperationSpec, OperationType, Visibility};
use alkhttp::websocket::{
    frame_channel0_chunk, ChunkAssembler, FrameAssembler, WsClient, INBOUND_WS_MESSAGE_CAP,
    WS_MESSAGE_CAP,
};
use std::collections::HashMap;
use std::sync::Arc;

fn identity(id: &str, scopes: &[&str]) -> Identity {
    Identity {
        id: id.to_string(),
        scopes: scopes.iter().map(|s| s.to_string()).collect(),
        resources: HashMap::new(),
    }
}

fn echo_handler() -> Handler {
    make_handler(|input, ctx| async move { ResponseEnvelope::ok(ctx.request_id, input) })
}

fn echo_registry() -> Arc<OperationRegistry> {
    let mut registry = OperationRegistry::new();
    registry
        .register(HandlerRegistration::new(
            OperationSpec::new(
                "echo/run",
                OperationType::Query,
                Visibility::External,
                serde_json::json!({}),
                serde_json::json!({}),
                vec![],
                AccessControl::default(),
                None,
            ),
            HandlerKind::Once(echo_handler()),
            OperationProvenance::Local,
            None,
            None,
            alkcall::core::types::Capabilities::new(),
        ))
        .unwrap();
    Arc::new(registry)
}

fn slow_and_echo_registry() -> Arc<OperationRegistry> {
    let mut registry = OperationRegistry::new();
    registry
        .register(HandlerRegistration::new(
            OperationSpec::new(
                "slow/op",
                OperationType::Query,
                Visibility::External,
                serde_json::json!({}),
                serde_json::json!({}),
                vec![],
                AccessControl::default(),
                None,
            ),
            HandlerKind::Once(make_handler(|_input, _ctx| async move {
                tokio::time::sleep(std::time::Duration::from_secs(30)).await;
                ResponseEnvelope::ok("never", serde_json::json!({}))
            })),
            OperationProvenance::Local,
            None,
            None,
            alkcall::core::types::Capabilities::new(),
        ))
        .unwrap();
    registry
        .register(HandlerRegistration::new(
            OperationSpec::new(
                "echo/run",
                OperationType::Query,
                Visibility::External,
                serde_json::json!({}),
                serde_json::json!({}),
                vec![],
                AccessControl::default(),
                None,
            ),
            HandlerKind::Once(echo_handler()),
            OperationProvenance::Local,
            None,
            None,
            alkcall::core::types::Capabilities::new(),
        ))
        .unwrap();
    Arc::new(registry)
}

fn internal_registry() -> Arc<OperationRegistry> {
    let mut registry = OperationRegistry::new();
    registry
        .register(HandlerRegistration::new(
            OperationSpec::new(
                "secret/op",
                OperationType::Query,
                Visibility::Internal,
                serde_json::json!({}),
                serde_json::json!({}),
                vec![],
                AccessControl::default(),
                None,
            ),
            HandlerKind::Once(echo_handler()),
            OperationProvenance::Local,
            None,
            None,
            alkcall::core::types::Capabilities::new(),
        ))
        .unwrap();
    Arc::new(registry)
}

fn restricted_registry() -> Arc<OperationRegistry> {
    let mut registry = OperationRegistry::new();
    registry
        .register(HandlerRegistration::new(
            OperationSpec::new(
                "admin/run",
                OperationType::Query,
                Visibility::External,
                serde_json::json!({}),
                serde_json::json!({}),
                vec![],
                AccessControl {
                    required_scopes: vec!["admin".to_string()],
                    ..Default::default()
                },
                None,
            ),
            HandlerKind::Once(echo_handler()),
            OperationProvenance::Local,
            None,
            None,
            alkcall::core::types::Capabilities::new(),
        ))
        .unwrap();
    Arc::new(registry)
}

fn registry_with_services_list(inner_ops: Vec<HandlerRegistration>) -> Arc<OperationRegistry> {
    let mut inner = OperationRegistry::new();
    for op in inner_ops {
        inner.register(op).unwrap();
    }
    let inner = Arc::new(inner);
    let mut registry = OperationRegistry::new();
    registry
        .register(HandlerRegistration::new(
            services_list_spec(),
            HandlerKind::Once(services_list_handler(Arc::clone(&inner))),
            OperationProvenance::Local,
            None,
            None,
            alkcall::core::types::Capabilities::new(),
        ))
        .unwrap();
    for spec in inner.list_operations() {
        let name = spec.name.clone();
        let reg = inner.registration(&name).unwrap();
        registry
            .register(HandlerRegistration::new(
                reg.spec.clone(),
                reg.handler.clone(),
                reg.provenance,
                reg.composition_authority.clone(),
                reg.scoped_env.clone(),
                reg.capabilities.clone(),
            ))
            .unwrap();
    }
    Arc::new(registry)
}

struct StaticTokens {
    tokens: std::sync::Mutex<HashMap<String, Identity>>,
}

impl IdentityProvider for StaticTokens {
    fn resolve_from_fingerprint(&self, _: &str) -> Option<Identity> {
        None
    }
    fn resolve_from_token(&self, token: &alkcall::core::auth::AuthToken) -> Option<Identity> {
        let s = String::from_utf8_lossy(&token.raw).to_string();
        self.tokens.lock().unwrap().get(&s).cloned()
    }
}

fn provider_with(tokens: Vec<(&str, Identity)>) -> Arc<dyn IdentityProvider> {
    let map: HashMap<String, Identity> = tokens
        .into_iter()
        .map(|(t, i)| (t.to_string(), i))
        .collect();
    Arc::new(StaticTokens {
        tokens: std::sync::Mutex::new(map),
    })
}

async fn spawn_ws_server(
    registry: Arc<OperationRegistry>,
    provider: Arc<dyn IdentityProvider>,
) -> String {
    let app = axum::Router::new()
        .route(
            "/alk/channels",
            axum::routing::get(alkhttp::websocket::ws_upgrade_handler),
        )
        .layer(axum::middleware::from_fn_with_state(
            provider,
            alkhttp::websocket::ws_bearer_auth,
        ))
        .with_state(registry);
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = format!("ws://{}", listener.local_addr().unwrap());
    tokio::spawn(async move {
        axum::serve(listener, app).await.unwrap();
    });
    addr
}

async fn call_and_await(
    ws: &mut WsClient,
    request_id: &str,
    operation: &str,
    input: serde_json::Value,
) -> EventEnvelope {
    let frame = EventEnvelope::requested(
        request_id,
        serde_json::json!({ "operationId": operation, "input": input }),
    );
    ws.send_binary(frame_channel0_chunk(&frame)).await;

    let mut chunks = ChunkAssembler::new();
    let mut frames = FrameAssembler::new();
    let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(5);
    loop {
        assert!(
            tokio::time::Instant::now() < deadline,
            "timed out waiting for response to {request_id}"
        );
        if let Some(env) = frames.next_frame() {
            return env;
        }
        let bin = ws.next_binary(std::time::Duration::from_millis(500)).await;
        match bin {
            Some(bytes) => {
                chunks.push(&bytes);
                while let Some((channel_id, payload)) = chunks.next_chunk() {
                    assert_eq!(channel_id, 0, "response must ride channel 0");
                    frames.push(&payload);
                }
            }
            None => panic!("ws closed unexpectedly while awaiting {request_id}"),
        }
    }
}

#[tokio::test]
async fn end_to_end_call_round_trip_over_ws() {
    let addr = spawn_ws_server(
        echo_registry(),
        provider_with(vec![("tok-1", identity("alice", &[]))]),
    )
    .await;

    let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1")
        .await
        .unwrap();

    let env = call_and_await(
        &mut ws,
        "req-1",
        "echo/run",
        serde_json::json!({ "hello": "world" }),
    )
    .await;
    assert_eq!(env.r#type, EVENT_RESPONDED, "got {}", env.r#type);
    assert_eq!(env.id, "req-1");
    assert_eq!(env.payload["output"]["hello"], "world");
    ws.close().await;
}

#[tokio::test]
async fn upgrade_without_token_rejected_401() {
    let addr = spawn_ws_server(
        echo_registry(),
        provider_with(vec![("tok-1", identity("alice", &[]))]),
    )
    .await;
    let status = WsClient::connect_status(&format!("{addr}/alk/channels"), None).await;
    assert_eq!(status, Some(401));
}

#[tokio::test]
async fn upgrade_with_unresolvable_token_rejected_401() {
    let addr = spawn_ws_server(
        echo_registry(),
        provider_with(vec![("tok-1", identity("alice", &[]))]),
    )
    .await;
    let status =
        WsClient::connect_status(&format!("{addr}/alk/channels"), Some("wrong-token")).await;
    assert_eq!(status, Some(401));
}

#[tokio::test]
async fn oversized_payload_splits_across_ws_messages_and_round_trips() {
    let addr = spawn_ws_server(
        echo_registry(),
        provider_with(vec![("tok-1", identity("alice", &[]))]),
    )
    .await;
    let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1")
        .await
        .unwrap();

    let payload_len = 3 * 1024 * 1024;
    let blob = "x".repeat(payload_len);
    let frame = EventEnvelope::requested(
        "req-big",
        serde_json::json!({ "operationId": "echo/run", "input": { "blob": blob } }),
    );
    let chunk = frame_channel0_chunk(&frame);
    assert!(
        chunk.len() > WS_MESSAGE_CAP,
        "test precondition: request exceeds message cap"
    );

    for piece in chunk.chunks(WS_MESSAGE_CAP) {
        ws.send_binary_piece(piece).await;
    }

    let mut chunks = ChunkAssembler::new();
    let mut frames = FrameAssembler::new();
    let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(30);
    loop {
        assert!(
            tokio::time::Instant::now() < deadline,
            "timed out on big payload"
        );
        if let Some(env) = frames.next_frame() {
            assert_eq!(env.r#type, EVENT_RESPONDED, "got {}", env.r#type);
            assert_eq!(env.id, "req-big");
            let out = env.payload["output"]["blob"].as_str().unwrap();
            assert_eq!(out.len(), payload_len, "blob round-tripped intact");
            assert!(out.bytes().all(|b| b == b'x'), "byte integrity");
            ws.close().await;
            return;
        }
        let bytes = ws
            .next_binary(std::time::Duration::from_secs(2))
            .await
            .expect("ws closed unexpectedly on big payload");
        chunks.push(&bytes);
        while let Some((channel_id, payload)) = chunks.next_chunk() {
            assert_eq!(channel_id, 0);
            frames.push(&payload);
        }
    }
}

#[tokio::test]
async fn interleaved_calls_reassemble_without_corruption() {
    let addr = spawn_ws_server(
        echo_registry(),
        provider_with(vec![("tok-1", identity("alice", &[]))]),
    )
    .await;
    let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1")
        .await
        .unwrap();

    const N: usize = 20;
    for i in 0..N {
        let frame = EventEnvelope::requested(
            format!("req-{i}"),
            serde_json::json!({ "operationId": "echo/run", "input": { "i": i } }),
        );
        ws.send_binary(frame_channel0_chunk(&frame)).await;
    }

    let mut chunks = ChunkAssembler::new();
    let mut frames = FrameAssembler::new();
    let mut seen: Vec<String> = Vec::new();
    let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(10);
    while seen.len() < N {
        assert!(
            tokio::time::Instant::now() < deadline,
            "only got {} of {}",
            seen.len(),
            N
        );
        let bytes = ws
            .next_binary(std::time::Duration::from_millis(500))
            .await
            .expect("ws closed unexpectedly");
        chunks.push(&bytes);
        while let Some((channel_id, payload)) = chunks.next_chunk() {
            assert_eq!(channel_id, 0);
            frames.push(&payload);
            while let Some(env) = frames.next_frame() {
                assert_eq!(env.r#type, EVENT_RESPONDED, "got {}", env.r#type);
                let i: usize = env.payload["output"]["i"].as_u64().unwrap() as usize;
                assert_eq!(env.id, format!("req-{i}"), "correlation intact");
                seen.push(env.id);
            }
        }
    }
    seen.sort();
    let mut expected: Vec<String> = (0..N).map(|i| format!("req-{i}")).collect();
    expected.sort();
    assert_eq!(seen, expected);
    ws.close().await;
}

#[tokio::test]
async fn disconnect_mid_call_does_not_hang_the_server() {
    let addr = spawn_ws_server(
        slow_and_echo_registry(),
        provider_with(vec![("tok-1", identity("alice", &[]))]),
    )
    .await;
    let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1")
        .await
        .unwrap();

    let frame = EventEnvelope::requested(
        "req-abandon",
        serde_json::json!({ "operationId": "slow/op", "input": {} }),
    );
    ws.send_binary(frame_channel0_chunk(&frame)).await;

    // Abrupt client disconnect (drop, no close handshake) while the
    // slow handler is still running.
    drop(ws);

    // The server must not hang: a fresh session completes an echo call
    // promptly — the abandoned session's teardown (pending failure,
    // overlay drop) did not wedge the runtime or the listener.
    let mut ws2 = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1")
        .await
        .unwrap();
    let env = call_and_await(
        &mut ws2,
        "req-after",
        "echo/run",
        serde_json::json!({"ok": true}),
    )
    .await;
    assert_eq!(env.r#type, EVENT_RESPONDED);
    assert_eq!(env.payload["output"]["ok"], true);
    ws2.close().await;
}

#[tokio::test]
async fn acl_filtered_call_forbidden_over_ws() {
    let addr = spawn_ws_server(
        restricted_registry(),
        provider_with(vec![("tok-1", identity("alice", &["user"]))]),
    )
    .await;
    let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1")
        .await
        .unwrap();

    let env = call_and_await(&mut ws, "req-403", "admin/run", serde_json::json!({})).await;
    assert_eq!(env.r#type, "call.error", "got {}", env.r#type);
    assert_eq!(env.id, "req-403");
    assert_eq!(env.payload["code"], "FORBIDDEN");
    ws.close().await;
}

#[tokio::test]
async fn text_message_rejected_with_protocol_close() {
    let addr = spawn_ws_server(
        echo_registry(),
        provider_with(vec![("tok-1", identity("alice", &[]))]),
    )
    .await;
    let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1")
        .await
        .unwrap();

    ws.send_text("hello as text").await;
    let close = ws
        .next_close(std::time::Duration::from_secs(5))
        .await
        .expect("expected a close frame");
    assert_eq!(close, Some(1002));
}

#[tokio::test]
async fn internal_op_not_callable_from_wire() {
    let addr = spawn_ws_server(
        internal_registry(),
        provider_with(vec![("tok-1", identity("alice", &[]))]),
    )
    .await;
    let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1")
        .await
        .unwrap();

    let env = call_and_await(&mut ws, "req-secret", "secret/op", serde_json::json!({})).await;
    assert_eq!(env.r#type, "call.error", "got {}", env.r#type);
    assert_eq!(env.id, "req-secret");
    assert_eq!(env.payload["code"], "NOT_FOUND");
    ws.close().await;
}

#[tokio::test]
async fn services_list_over_channel0_is_access_control_filtered() {
    let registry = registry_with_services_list(vec![
        HandlerRegistration::new(
            OperationSpec::new(
                "public/echo",
                OperationType::Query,
                Visibility::External,
                serde_json::json!({}),
                serde_json::json!({}),
                vec![],
                AccessControl::default(),
                None,
            ),
            HandlerKind::Once(echo_handler()),
            OperationProvenance::Local,
            None,
            None,
            alkcall::core::types::Capabilities::new(),
        ),
        HandlerRegistration::new(
            OperationSpec::new(
                "admin/secret",
                OperationType::Query,
                Visibility::External,
                serde_json::json!({}),
                serde_json::json!({}),
                vec![],
                AccessControl {
                    required_scopes: vec!["admin".to_string()],
                    ..Default::default()
                },
                None,
            ),
            HandlerKind::Once(echo_handler()),
            OperationProvenance::Local,
            None,
            None,
            alkcall::core::types::Capabilities::new(),
        ),
    ]);
    let addr = spawn_ws_server(
        registry,
        provider_with(vec![("tok-1", identity("regular", &["user"]))]),
    )
    .await;
    let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1")
        .await
        .unwrap();

    let env = call_and_await(&mut ws, "req-list", "services/list", serde_json::json!({})).await;
    assert_eq!(env.r#type, EVENT_RESPONDED, "got {}", env.r#type);
    assert_eq!(env.id, "req-list");
    let ops = env.payload["output"]["operations"]
        .as_array()
        .expect("operations array");
    let names: Vec<&str> = ops.iter().filter_map(|o| o["name"].as_str()).collect();
    assert!(names.contains(&"public/echo"), "got: {names:?}");
    assert!(!names.contains(&"admin/secret"), "got: {names:?}");
    ws.close().await;
}

#[tokio::test]
async fn inbound_message_over_cap_fails_the_connection() {
    let addr = spawn_ws_server(
        echo_registry(),
        provider_with(vec![("tok-1", identity("alice", &[]))]),
    )
    .await;
    let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1")
        .await
        .unwrap();

    // One WS binary message larger than INBOUND_WS_MESSAGE_CAP: the
    // server's explicit max_message_size (WS-06) must fail the
    // connection instead of accepting it (the byte-adapter boundary is
    // the chunk header, so the adapter never dissects messages; the WS
    // library cap is the enforcement point).
    let oversized = vec![b'x'; INBOUND_WS_MESSAGE_CAP + 1];
    ws.send_binary_piece(&oversized).await;

    let close = ws
        .next_close(std::time::Duration::from_secs(5))
        .await
        .expect("server must fail the oversized-inbound connection (WS-06)");
    assert!(
        close.is_some(),
        "expected a close or stream end, got {close:?}"
    );
}

/// WS-08 acceptance: the upgrade retains each session's pump handle in
/// the shared `WsSessions` registry; `abort()` evicts a stuck session
/// (the peer's socket closes) and the entry is removed when the
/// session task ends, so the registry only tracks live sessions.
#[tokio::test]
async fn ws_sessions_registry_tracks_and_aborts_live_sessions() {
    use alkhttp::websocket::WsSessions;

    let sessions = WsSessions::new();
    let registry = echo_registry();
    let provider = provider_with(vec![("tok-1", identity("alice", &[]))]);

    let app = axum::Router::new()
        .route(
            "/alk/channels",
            axum::routing::get(alkhttp::websocket::ws_upgrade_handler),
        )
        .layer(axum::middleware::from_fn_with_state(
            Arc::clone(&provider),
            alkhttp::websocket::ws_bearer_auth,
        ))
        .with_state(registry)
        .layer(axum::Extension(sessions.clone()));
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = format!("ws://{}", listener.local_addr().unwrap());
    tokio::spawn(async move {
        axum::serve(listener, app).await.unwrap();
    });

    let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1")
        .await
        .unwrap();
    let env = call_and_await(&mut ws, "req-1", "echo/run", serde_json::json!({})).await;
    assert_eq!(env.r#type, EVENT_RESPONDED, "session established");

    // The upgrade registered the session.
    let saw_session = tokio::time::timeout(std::time::Duration::from_secs(5), async {
        loop {
            if !sessions.is_empty() {
                return true;
            }
            tokio::time::sleep(std::time::Duration::from_millis(20)).await;
        }
    })
    .await
    .expect("session registered in the shared registry");
    assert!(saw_session);
    assert_eq!(sessions.len(), 1);

    // Forced teardown: abort() closes the peer's socket.
    sessions.abort();
    let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(10);
    let close = loop {
        assert!(
            tokio::time::Instant::now() < deadline,
            "abort never reached the peer"
        );
        match ws.next_close(std::time::Duration::from_millis(250)).await {
            Some(x) => break Some(x),
            None => continue,
        }
    };
    assert!(close.is_some(), "peer observes the teardown: {close:?}");

    // The self-removing guard drops the entry once the session task ends.
    let drained = tokio::time::timeout(std::time::Duration::from_secs(5), async {
        loop {
            if sessions.is_empty() {
                return true;
            }
            tokio::time::sleep(std::time::Duration::from_millis(20)).await;
        }
    })
    .await
    .expect("session entry removed after the session task ends");
    assert!(drained);
    ws.close().await;
}

/// WS-09 acceptance: the session cap is enforced on the upgrade — a
/// caller over the configured cap is rejected with 503, and an ended
/// session frees its slot for the next caller.
#[tokio::test]
async fn session_cap_rejects_over_limit_with_503_and_frees_slots_on_end() {
    use alkhttp::server::HttpAdapter;

    let mut registry_val = OperationRegistry::new();
    registry_val
        .register(HandlerRegistration::new(
            OperationSpec::new(
                "echo/run",
                OperationType::Query,
                Visibility::External,
                serde_json::json!({}),
                serde_json::json!({}),
                vec![],
                AccessControl::default(),
                None,
            ),
            HandlerKind::Once(echo_handler()),
            OperationProvenance::Local,
            None,
            None,
            alkcall::core::types::Capabilities::new(),
        ))
        .unwrap();
    let registry = Arc::new(registry_val);
    struct StaticTok;
    impl IdentityProvider for StaticTok {
        fn resolve_from_fingerprint(&self, _: &str) -> Option<Identity> {
            None
        }
        fn resolve_from_token(&self, token: &alkcall::core::auth::AuthToken) -> Option<Identity> {
            let s = String::from_utf8_lossy(&token.raw).to_string();
            (s == "tok-1").then(|| identity("alice", &[]))
        }
    }

    // The adapter's built-in surface: WS route + bearer middleware, the
    // whole deployment shape the cap knob configures.
    let adapter =
        HttpAdapter::new(Arc::new(StaticTok), Arc::clone(&registry)).with_ws_max_sessions(1);
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = format!("ws://{}", listener.local_addr().unwrap());
    let app: axum::Router = adapter.router().clone();
    tokio::spawn(async move {
        axum::serve(listener, app).await.unwrap();
    });

    let url = format!("{addr}/alk/channels");

    // First session: admitted.
    let mut first = WsClient::connect_authorized(&url, "tok-1").await.unwrap();
    let env = call_and_await(&mut first, "req-1", "echo/run", serde_json::json!({})).await;
    assert_eq!(env.r#type, EVENT_RESPONDED, "first session admitted");

    // Second session over the cap: rejected with 503 at the upgrade.
    let status = WsClient::connect_status(&url, Some("tok-1")).await;
    assert_eq!(status, Some(503), "second session rejected over the cap");

    // Ending the first session frees the slot: the next caller is admitted.
    first.close().await;
    tokio::time::sleep(std::time::Duration::from_millis(200)).await;
    let mut third = WsClient::connect_authorized(&url, "tok-1").await.unwrap();
    let env = call_and_await(&mut third, "req-2", "echo/run", serde_json::json!({})).await;
    assert_eq!(env.r#type, EVENT_RESPONDED, "slot freed after session end");
    third.close().await;
}

/// WS-13 acceptance on the axum path (the integration mirror; also the
/// COV-11b dark-knob gate for `with_ws_idle_timeout`): a client
/// dribbling a declared chunk's payload one byte per message, each gap
/// inside the idle window, is evicted with 1001 — message arrival does
/// not reset the deadline, only a completed chunk would.
#[tokio::test]
async fn idle_progress_knob_evicts_forever_dribble_over_axum_upgrade() {
    use alkhttp::server::HttpAdapter;

    let registry_val = OperationRegistry::new();
    let registry = Arc::new(registry_val);
    struct StaticTok;
    impl IdentityProvider for StaticTok {
        fn resolve_from_fingerprint(&self, _: &str) -> Option<Identity> {
            None
        }
        fn resolve_from_token(&self, token: &alkcall::core::auth::AuthToken) -> Option<Identity> {
            let s = String::from_utf8_lossy(&token.raw).to_string();
            (s == "tok-1").then(|| identity("alice", &[]))
        }
    }
    let adapter = HttpAdapter::new(Arc::new(StaticTok), Arc::clone(&registry))
        .with_ws_idle_timeout(Some(std::time::Duration::from_millis(150)));
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = format!("ws://{}", listener.local_addr().unwrap());
    let app: axum::Router = adapter.router().clone();
    tokio::spawn(async move {
        axum::serve(listener, app).await.unwrap();
    });

    let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1")
        .await
        .unwrap();

    let dribbler = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1")
        .await
        .unwrap();
    let dribble = tokio::spawn(async move {
        let mut dribbler = dribbler;
        let mut header = vec![0u8; 8];
        header[4..8].copy_from_slice(&64u32.to_be_bytes());
        dribbler.send_binary(header).await;
        loop {
            tokio::time::sleep(std::time::Duration::from_millis(40)).await;
            dribbler.send_binary(vec![0u8]).await;
        }
    });

    let close = tokio::time::timeout(std::time::Duration::from_secs(10), async {
        loop {
            match ws.next_close(std::time::Duration::from_millis(250)).await {
                Some(Some(code)) => break code,
                Some(None) => break 1006,
                None => continue,
            }
        }
    })
    .await
    .expect("dribbling client must be evicted within 10 s");
    dribble.abort();
    assert_eq!(
        close,
        alkhttp::websocket::WS_GOING_AWAY,
        "the forever-dribble hits the progress deadline despite arriving messages"
    );
}

/// WS-13 survivor side on the axum path: a session whose traffic keeps
/// completing chunks (each chunk inside the window — the scaled-down
/// "messages flowing that make progress" shape) is NOT evicted, and
/// calls keep round-tripping across many windows.
#[tokio::test]
async fn idle_progress_knob_survives_productive_sessions_over_axum_upgrade() {
    use alkhttp::server::HttpAdapter;

    let registry = echo_registry();
    struct StaticTok;
    impl IdentityProvider for StaticTok {
        fn resolve_from_fingerprint(&self, _: &str) -> Option<Identity> {
            None
        }
        fn resolve_from_token(&self, token: &alkcall::core::auth::AuthToken) -> Option<Identity> {
            let s = String::from_utf8_lossy(&token.raw).to_string();
            (s == "tok-1").then(|| identity("alice", &[]))
        }
    }
    let adapter = HttpAdapter::new(Arc::new(StaticTok), Arc::clone(&registry))
        .with_ws_idle_timeout(Some(std::time::Duration::from_millis(150)));
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = format!("ws://{}", listener.local_addr().unwrap());
    let app: axum::Router = adapter.router().clone();
    tokio::spawn(async move {
        axum::serve(listener, app).await.unwrap();
    });

    let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1")
        .await
        .unwrap();

    for i in 0..6 {
        tokio::time::sleep(std::time::Duration::from_millis(75)).await;
        let n = i * 7 + 1;
        let env = call_and_await(
            &mut ws,
            &format!("req-{n}"),
            "echo/run",
            serde_json::json!({ "n": n }),
        )
        .await;
        assert_eq!(env.r#type, EVENT_RESPONDED);
        assert_eq!(env.payload["output"]["n"], n);
    }
    ws.close().await;
}

/// WS-13 `None` arm on the axum path: with the knob disabled a
/// dribbling client that completes no chunk is never evicted.
#[tokio::test]
async fn idle_progress_knob_none_disables_eviction_over_axum_upgrade() {
    use alkhttp::server::HttpAdapter;

    let registry = echo_registry();
    struct StaticTok;
    impl IdentityProvider for StaticTok {
        fn resolve_from_fingerprint(&self, _: &str) -> Option<Identity> {
            None
        }
        fn resolve_from_token(&self, token: &alkcall::core::auth::AuthToken) -> Option<Identity> {
            let s = String::from_utf8_lossy(&token.raw).to_string();
            (s == "tok-1").then(|| identity("alice", &[]))
        }
    }
    let adapter =
        HttpAdapter::new(Arc::new(StaticTok), Arc::clone(&registry)).with_ws_idle_timeout(None);
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = format!("ws://{}", listener.local_addr().unwrap());
    let app: axum::Router = adapter.router().clone();
    tokio::spawn(async move {
        axum::serve(listener, app).await.unwrap();
    });

    let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1")
        .await
        .unwrap();

    let mut header = vec![0u8; 8];
    header[4..8].copy_from_slice(&8u32.to_be_bytes());
    ws.send_binary(header).await;
    for _ in 0..4 {
        tokio::time::sleep(std::time::Duration::from_millis(150)).await;
        ws.send_binary_piece(&[0u8]).await;
    }
    let no_close = ws.next_close(std::time::Duration::from_millis(400)).await;
    assert!(
        !matches!(no_close, Some(Some(_))),
        "knob disabled: no eviction arrives, got {no_close:?}"
    );
    ws.close().await;
}

/// WS-17 acceptance: a bare-registry upgrade route takes the pump
/// knobs per request via the `WsTimeouts` extension (mirroring
/// `ChannelsPolicy`) — a client that completes no chunk is evicted
/// with 1001 inside the extension's short idle window, not the 60 s
/// default.
#[tokio::test]
async fn ws_timeouts_extension_sets_the_idle_window_on_a_bare_registry_route() {
    let registry = std::sync::Arc::new(OperationRegistry::new());
    struct StaticTok;
    impl IdentityProvider for StaticTok {
        fn resolve_from_fingerprint(&self, _: &str) -> Option<Identity> {
            None
        }
        fn resolve_from_token(&self, token: &alkcall::core::auth::AuthToken) -> Option<Identity> {
            let s = String::from_utf8_lossy(&token.raw).to_string();
            (s == "tok-1").then(|| identity("alice", &[]))
        }
    }
    let app = axum::Router::new()
        .route(
            "/alk/channels",
            axum::routing::get(alkhttp::websocket::ws_upgrade_handler),
        )
        .layer(axum::middleware::from_fn_with_state(
            std::sync::Arc::new(StaticTok) as std::sync::Arc<dyn IdentityProvider>,
            alkhttp::websocket::ws_bearer_auth,
        ))
        .layer(axum::Extension(alkhttp::websocket::WsTimeouts {
            idle: Some(std::time::Duration::from_millis(150)),
            write: None,
        }))
        .with_state(registry);
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = format!("ws://{}", listener.local_addr().unwrap());
    tokio::spawn(async move {
        axum::serve(listener, app).await.unwrap();
    });

    let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1")
        .await
        .unwrap();

    let close = ws.next_close(std::time::Duration::from_secs(5)).await;
    assert_eq!(
        close,
        Some(Some(alkhttp::websocket::WS_GOING_AWAY)),
        "evicted with 1001 inside the extension's idle window"
    );
    ws.close().await;
}