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, 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;
}