hypercall_api/websocket/

rfq.rs

use crate::rfq::indicative_quote_cache::{AggregatedIndicativeQuote, IndicativeQuoteEntry};
use crate::rfq::qp_ws_state::{QpWsState, RfqWebsocketPublisher};
use axum::{
    extract::{
        ws::{Message, WebSocket, WebSocketUpgrade},
        State,
    },
    http::HeaderMap,
    response::Response,
};
use futures::{SinkExt, StreamExt};
use hypercall_auth::SignatureRecovery;
use hypercall_types::ws_protocol::{
    WsIndicativeMarketData, WsMessage, WsRfqQuoteEntry, WsRfqQuotes, WsRfqStatusUpdate,
};
use hypercall_types::WalletAddress;
use hypercall_ws_protocol::{
    QpClientMessage as QpInboundMessage, QpServerMessage as QpOutboundMessage,
};
use metrics::counter;
use rust_decimal::Decimal;
use std::sync::Arc;
use tokio::sync::mpsc;
use tracing::{info, warn};

/// Fallback `valid_for_ms` applied when a quote provider omits the field on
/// an inbound RFQ response. Chosen to match the rfq-mm default and stay under
/// `RfqConfig::max_valid_for_ms` (30s).
const DEFAULT_QUOTE_VALID_FOR_MS: u64 = 25_000;
const GATEWAY_TOKEN_HEADER: &str = "x-hypercall-gateway-token";

/// Axum handler for QP WebSocket upgrade at `/ws/quotes`.
pub async fn qp_websocket_handler(
    ws: WebSocketUpgrade,
    headers: HeaderMap,
    State(state): State<QpWsState>,
) -> Response {
    let gateway_resume_authorized = gateway_resume_authorized(&headers, &state);
    ws.on_upgrade(move |socket| handle_qp_socket(socket, state, gateway_resume_authorized))
}

fn gateway_resume_authorized(headers: &HeaderMap, state: &QpWsState) -> bool {
    let Some(expected) = state.gateway_resume_token.as_deref() else {
        return false;
    };
    headers
        .get(GATEWAY_TOKEN_HEADER)
        .and_then(|value| value.to_str().ok())
        .is_some_and(|actual| actual == expected)
}

async fn handle_qp_socket(socket: WebSocket, state: QpWsState, gateway_resume_authorized: bool) {
    let (mut ws_sender, mut ws_receiver) = socket.split();

    // Step 1: Wait for first frame (ConnectQuoteProvider auth)
    let auth_result =
        match tokio::time::timeout(std::time::Duration::from_secs(10), ws_receiver.next()).await {
            Ok(Some(Ok(Message::Text(text)))) => {
                authenticate_qp(&text, &state, gateway_resume_authorized).await
            }
            Ok(Some(Ok(Message::Binary(bytes)))) => {
                if let Ok(text) = String::from_utf8(bytes.to_vec()) {
                    authenticate_qp(&text, &state, gateway_resume_authorized).await
                } else {
                    Err("Invalid binary frame".to_string())
                }
            }
            _ => Err("Timeout or connection closed before auth".to_string()),
        };

    let qp_wallet = match auth_result {
        Ok(wallet) => {
            let msg = QpOutboundMessage::Authenticated {
                wallet: wallet.as_hex(),
            };
            let json = serde_json::to_string(&msg).unwrap();
            if ws_sender.send(Message::Text(json)).await.is_err() {
                return;
            }
            wallet
        }
        Err(reason) => {
            let msg = QpOutboundMessage::AuthFailed { reason };
            let json = serde_json::to_string(&msg).unwrap();
            let _ = ws_sender.send(Message::Text(json)).await;
            return;
        }
    };

    info!("QP authenticated: {}", qp_wallet.as_hex());

    // Step 2: Register session
    let (outbound_tx, mut outbound_rx) = mpsc::unbounded_channel::<QpOutboundMessage>();
    let session_id = state
        .session_manager
        .register_session(qp_wallet, outbound_tx);

    // Step 3: Message loop
    loop {
        tokio::select! {
            // Outbound: engine -> QP
            Some(msg) = outbound_rx.recv() => {
                let json = serde_json::to_string(&msg).unwrap();
                if ws_sender.send(Message::Text(json)).await.is_err() {
                    break;
                }
            }
            // Inbound: QP -> engine
            Some(Ok(frame)) = ws_receiver.next() => {
                match frame {
                    Message::Text(text) => {
                        handle_qp_message(&text, &qp_wallet, &state).await;
                    }
                    Message::Binary(bytes) => {
                        if let Ok(text) = String::from_utf8(bytes.to_vec()) {
                            handle_qp_message(&text, &qp_wallet, &state).await;
                        }
                    }
                    Message::Ping(_) => {
                        let _ = ws_sender.send(Message::Pong(vec![])).await;
                    }
                    Message::Close(_) => break,
                    _ => {}
                }
            }
            else => break,
        }
    }

    // Step 4: Cleanup on disconnect
    // Only evict this QP's indicative quotes if the disconnecting socket
    // was the LAST session for the wallet. `QpSessionManager` supports
    // multiple concurrent sessions per wallet (e.g. two tabs, rolling
    // deploy with overlapping sessions), and unconditional eviction on
    // any single disconnect would clobber quotes that the surviving
    // sessions are still actively updating.
    let fully_disconnected = state.session_manager.remove_session(&qp_wallet, session_id);
    if fully_disconnected {
        let updates = state.indicative_cache.evict_qp(&qp_wallet);
        publish_indicative_updates(&state.pubsub, updates);
    }
    info!("QP disconnected: {}", qp_wallet.as_hex());
}

async fn authenticate_qp(
    text: &str,
    state: &QpWsState,
    gateway_resume_authorized: bool,
) -> Result<WalletAddress, String> {
    let msg: QpInboundMessage =
        serde_json::from_str(text).map_err(|e| format!("Invalid JSON: {}", e))?;

    match msg {
        QpInboundMessage::ConnectQuoteProvider {
            wallet,
            timestamp,
            nonce,
            signature,
        } => {
            let wallet_addr = WalletAddress::from_str(&wallet)
                .map_err(|_| format!("Invalid wallet address: {}", wallet))?;

            let ts = timestamp
                .parse::<u64>()
                .map_err(|_| "Invalid timestamp".to_string())?;

            let recovered = SignatureRecovery::recover_connect_quote_provider_signer(
                wallet_addr,
                alloy::primitives::U256::from(ts),
                nonce,
                &signature,
                state.signing_chain_id,
            )
            .map_err(|e| format!("Signature recovery failed: {}", e))?;

            if recovered != wallet_addr.inner() {
                return Err("Signature does not match wallet".to_string());
            }

            // Replay protection: advance the engine nonce watermark. The
            // engine rejects if nonce <= max_seen[wallet] (strict monotonic).
            let (tx, rx) = tokio::sync::oneshot::channel();
            let nonce_req = hypercall_runtime_api::NonceCheckRequest {
                wallet: wallet_addr,
                nonce,
                applied_tx: tx,
            };
            state
                .nonce_check_sender
                .send(nonce_req)
                .await
                .map_err(|_| "Engine nonce channel closed".to_string())?;
            match tokio::time::timeout(std::time::Duration::from_secs(5), rx).await {
                Ok(Ok(Ok(()))) => {}
                Ok(Ok(Err(e))) => return Err(format!("Handshake nonce rejected: {}", e)),
                Ok(Err(_)) => return Err("Engine dropped nonce response channel".to_string()),
                Err(_) => return Err("Timed out waiting for nonce check".to_string()),
            }

            if !state.qp_cache.is_active(&wallet_addr).await {
                return Err("Wallet is not a registered active quote provider".to_string());
            }

            Ok(wallet_addr)
        }
        QpInboundMessage::GatewayResumeQuoteProvider {
            wallet,
            timestamp,
            nonce,
            signature,
        } => {
            if !gateway_resume_authorized {
                return Err(
                    "gateway_resume_quote_provider is only accepted from the internal gateway"
                        .to_string(),
                );
            }
            let wallet_addr =
                verify_qp_connect_signature(&wallet, &timestamp, nonce, &signature, state)?;

            if !state.qp_cache.is_active(&wallet_addr).await {
                return Err("Wallet is not a registered active quote provider".to_string());
            }

            Ok(wallet_addr)
        }
        _ => Err(
            "First message must be ConnectQuoteProvider or GatewayResumeQuoteProvider".to_string(),
        ),
    }
}

fn verify_qp_connect_signature(
    wallet: &str,
    timestamp: &str,
    nonce: u64,
    signature: &str,
    state: &QpWsState,
) -> Result<WalletAddress, String> {
    let wallet_addr = WalletAddress::from_str(wallet)
        .map_err(|_| format!("Invalid wallet address: {}", wallet))?;

    let ts = timestamp
        .parse::<u64>()
        .map_err(|_| "Invalid timestamp".to_string())?;

    let recovered = SignatureRecovery::recover_connect_quote_provider_signer(
        wallet_addr,
        alloy::primitives::U256::from(ts),
        nonce,
        signature,
        state.signing_chain_id,
    )
    .map_err(|e| format!("Signature recovery failed: {}", e))?;

    if recovered != wallet_addr.inner() {
        return Err("Signature does not match wallet".to_string());
    }

    Ok(wallet_addr)
}

async fn handle_qp_message(text: &str, qp_wallet: &WalletAddress, state: &QpWsState) {
    let msg: QpInboundMessage = match serde_json::from_str(text) {
        Ok(m) => m,
        Err(e) => {
            warn!("Invalid QP message from {}: {}", qp_wallet.as_hex(), e);
            return;
        }
    };

    match msg {
        QpInboundMessage::ConnectQuoteProvider { .. }
        | QpInboundMessage::GatewayResumeQuoteProvider { .. } => {
            warn!(
                "Ignoring QP auth frame after authentication from {}",
                qp_wallet.as_hex()
            );
        }
        QpInboundMessage::IndicativeQuoteUpdate { quotes } => {
            let now = std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_millis() as u64;
            let quote_count = quotes.len();

            let entries: Vec<IndicativeQuoteEntry> = quotes
                .into_iter()
                .filter_map(|q| {
                    if !state.instruments_cache.allows_rfq(&q.instrument) {
                        return None;
                    }
                    Some(IndicativeQuoteEntry {
                        instrument: q.instrument,
                        bid_price: q.bid_price.parse::<Decimal>().ok()?,
                        ask_price: q.ask_price.parse::<Decimal>().ok()?,
                        max_bid_size: q.max_bid_size.parse::<Decimal>().ok()?,
                        max_ask_size: q.max_ask_size.parse::<Decimal>().ok()?,
                        updated_at: now,
                    })
                })
                .collect();

            if quote_count == 0 || !entries.is_empty() {
                let updates = state.indicative_cache.update(*qp_wallet, entries);
                publish_indicative_updates(&state.pubsub, updates);
            }
        }
        QpInboundMessage::RfqResponse {
            rfq_id,
            action,
            legs,
            net_premium,
            valid_for_ms,
            signature,
            nonce,
        } => {
            if action == "decline" {
                counter!("ht_qp_quote_rejection_total", "reason" => "decline").increment(1);
                info!("QP {} declined RFQ {}", qp_wallet.as_hex(), rfq_id);
                if let (Some(ref rfq_manager), Ok(rfq_uuid)) =
                    (&state.rfq_manager, uuid::Uuid::parse_str(&rfq_id))
                {
                    rfq_manager.record_qp_decline(rfq_uuid, *qp_wallet);
                }
                return;
            }

            if let Some(ref rfq_manager) = state.rfq_manager {
                let rfq_uuid = match uuid::Uuid::parse_str(&rfq_id) {
                    Ok(u) => u,
                    Err(e) => {
                        warn!("Invalid rfq_id from QP {}: {}", qp_wallet.as_hex(), e);
                        return;
                    }
                };

                // Re-check QP is still active before accepting quote
                if !state.qp_cache.is_active(qp_wallet).await {
                    warn!(
                        "QP {} is no longer active, rejecting RFQ response for {}",
                        qp_wallet.as_hex(),
                        rfq_id
                    );
                    rfq_manager.record_qp_invalid_response(rfq_uuid, *qp_wallet);
                    return;
                }

                // Strictly validate every QP leg. Silently filtering out
                // invalid legs (the previous behavior) would let a QP
                // submit a partial or zero/negative quote that still
                // flows into margin + settlement; a negative price in
                // particular flips `taker_premium_delta` and produces
                // economically invalid fills that still pass signature
                // checks (the QP signature binds `net_premium`, not the
                // per-leg price/size). Reject the whole quote on any
                // parse error, unknown side, or non-positive value.
                let mut quote_legs: Vec<crate::rfq::rfq_manager::QuoteLeg> =
                    Vec::with_capacity(legs.as_ref().map(|l| l.len()).unwrap_or(0));
                let mut quote_legs_valid = true;
                for l in legs.unwrap_or_default() {
                    let side = match l.side.to_lowercase().as_str() {
                        "buy" => hypercall_types::Side::Buy,
                        "sell" => hypercall_types::Side::Sell,
                        other => {
                            warn!(
                                "QP {} sent invalid leg side {:?} for RFQ {}, rejecting",
                                qp_wallet.as_hex(),
                                other,
                                rfq_id
                            );
                            quote_legs_valid = false;
                            break;
                        }
                    };
                    let price = match l.price.parse::<Decimal>() {
                        Ok(p) if p > Decimal::ZERO => p,
                        _ => {
                            warn!(
                                "QP {} sent non-positive/invalid leg price {:?} for RFQ {}, rejecting",
                                qp_wallet.as_hex(),
                                l.price,
                                rfq_id
                            );
                            quote_legs_valid = false;
                            break;
                        }
                    };
                    let size = match l.size.parse::<Decimal>() {
                        Ok(s) if s > Decimal::ZERO => s,
                        _ => {
                            warn!(
                                "QP {} sent non-positive/invalid leg size {:?} for RFQ {}, rejecting",
                                qp_wallet.as_hex(),
                                l.size,
                                rfq_id
                            );
                            quote_legs_valid = false;
                            break;
                        }
                    };
                    quote_legs.push(crate::rfq::rfq_manager::QuoteLeg {
                        instrument: l.instrument,
                        side,
                        price,
                        size,
                    });
                }
                if !quote_legs_valid {
                    rfq_manager.record_qp_invalid_response(rfq_uuid, *qp_wallet);
                    return;
                }

                // Cross-check the QP's returned legs against the ORIGINAL
                // RFQ legs that `legs_hash` was computed from. The EIP-712
                // signature check below only binds `legs_hash + net_premium`,
                // and `legs_hash` is over the taker's originally-requested
                // (instrument, side, size) tuples — NOT the per-leg metadata
                // the QP echoes back. Without this cross-check, a malicious
                // or buggy QP can emit a validly signed response while
                // swapping instrument / side / size in the websocket
                // payload, and `execute_rfq_inner` will happily run the
                // altered legs once the taker accepts the (still-valid)
                // signature. The QP may only contribute `price`; everything
                // else must match the original request exactly.
                let rfq_record = match rfq_manager.get_rfq(&rfq_uuid) {
                    Some(r) => r,
                    None => {
                        warn!(
                            "QP {} sent response for unknown RFQ {}, rejecting",
                            qp_wallet.as_hex(),
                            rfq_id
                        );
                        rfq_manager.record_qp_invalid_response(rfq_uuid, *qp_wallet);
                        return;
                    }
                };
                if quote_legs.len() != rfq_record.legs.len() {
                    warn!(
                        "QP {} sent {} legs for RFQ {} (expected {}), rejecting",
                        qp_wallet.as_hex(),
                        quote_legs.len(),
                        rfq_id,
                        rfq_record.legs.len()
                    );
                    rfq_manager.record_qp_invalid_response(rfq_uuid, *qp_wallet);
                    return;
                }
                let mut legs_match = true;
                for (qp_leg, original) in quote_legs.iter().zip(rfq_record.legs.iter()) {
                    if qp_leg.instrument != original.instrument
                        || qp_leg.side != original.side
                        || qp_leg.size != original.size
                    {
                        warn!(
                            "QP {} sent leg mismatched with original RFQ {} (got instrument={}, side={:?}, size={}; expected instrument={}, side={:?}, size={}), rejecting",
                            qp_wallet.as_hex(),
                            rfq_id,
                            qp_leg.instrument, qp_leg.side, qp_leg.size,
                            original.instrument, original.side, original.size,
                        );
                        legs_match = false;
                        break;
                    }
                }
                if !legs_match {
                    rfq_manager.record_qp_invalid_response(rfq_uuid, *qp_wallet);
                    return;
                }

                let quote = crate::rfq::rfq_manager::QpQuote {
                    quote_id: uuid::Uuid::now_v7(),
                    qp_wallet: *qp_wallet,
                    legs: quote_legs,
                    net_premium: match net_premium.and_then(|s| s.parse().ok()) {
                        Some(p) => p,
                        None => {
                            warn!(
                                "QP {} sent RFQ response with missing/invalid net_premium for RFQ {}, rejecting",
                                qp_wallet.as_hex(), rfq_id
                            );
                            rfq_manager.record_qp_invalid_response(rfq_uuid, *qp_wallet);
                            return;
                        }
                    },
                    valid_for_ms: valid_for_ms.unwrap_or(DEFAULT_QUOTE_VALID_FOR_MS),
                    qp_signature: match signature {
                        Some(s) if !s.is_empty() => s,
                        _ => {
                            warn!(
                                "QP {} sent RFQ response without signature for RFQ {}, rejecting",
                                qp_wallet.as_hex(),
                                rfq_id
                            );
                            rfq_manager.record_qp_invalid_response(rfq_uuid, *qp_wallet);
                            return;
                        }
                    },
                    qp_nonce: nonce.unwrap_or(0),
                    received_at: std::time::SystemTime::now()
                        .duration_since(std::time::UNIX_EPOCH)
                        .unwrap()
                        .as_millis() as u64,
                };

                // Verify QP EIP-712 signature (SubmitRFQResponse)
                if let Some(rfq_record) = rfq_manager.get_rfq(&rfq_uuid) {
                    let mut rfq_id_bytes = [0u8; 32];
                    rfq_id_bytes[16..].copy_from_slice(rfq_uuid.as_bytes());

                    // Fail closed on premium conversion errors. Silent fallbacks
                    // to zero would let a QP sign over a zero-premium payload
                    // and still pass verification while transmitting an
                    // un-representable `net_premium`, decoupling the accepted
                    // quote from the cryptographic binding.
                    use rust_decimal::prelude::ToPrimitive;
                    let premium_scaled =
                        quote.net_premium * rust_decimal::Decimal::new(1_000_000, 0);
                    let premium_micro = match premium_scaled.to_i64() {
                        Some(v) => v,
                        None => {
                            warn!(
                                "QP {} net_premium {} not representable as i64 for RFQ {}, rejecting",
                                qp_wallet.as_hex(),
                                quote.net_premium,
                                rfq_id
                            );
                            rfq_manager.record_qp_invalid_response(rfq_uuid, *qp_wallet);
                            return;
                        }
                    };
                    let net_premium_i256 = match alloy::primitives::I256::try_from(premium_micro) {
                        Ok(v) => v,
                        Err(e) => {
                            warn!(
                                "QP {} net_premium {} failed I256 conversion for RFQ {}: {}, rejecting",
                                qp_wallet.as_hex(),
                                quote.net_premium,
                                rfq_id,
                                e
                            );
                            rfq_manager.record_qp_invalid_response(rfq_uuid, *qp_wallet);
                            return;
                        }
                    };
                    let valid_for_u256 = alloy::primitives::U256::from(quote.valid_for_ms);

                    match hypercall_auth::SignatureRecovery::recover_submit_rfq_response_signer(
                        rfq_id_bytes,
                        rfq_record.legs_hash,
                        net_premium_i256,
                        valid_for_u256,
                        *qp_wallet,
                        quote.qp_nonce,
                        &quote.qp_signature,
                        state.signing_chain_id,
                    ) {
                        Ok(recovered) => {
                            if recovered != qp_wallet.inner() {
                                warn!(
                                    "QP {} signature mismatch for RFQ {}, rejecting",
                                    qp_wallet.as_hex(),
                                    rfq_id
                                );
                                rfq_manager.record_qp_invalid_response(rfq_uuid, *qp_wallet);
                                return;
                            }
                        }
                        Err(e) => {
                            warn!(
                                "QP {} signature recovery failed for RFQ {}: {}, rejecting",
                                qp_wallet.as_hex(),
                                rfq_id,
                                e
                            );
                            rfq_manager.record_qp_invalid_response(rfq_uuid, *qp_wallet);
                            return;
                        }
                    }
                }

                match rfq_manager.handle_qp_response(rfq_uuid, quote) {
                    Ok(result) => {
                        use crate::rfq::rfq_manager::HandleQpResponseResult;
                        match result {
                            HandleQpResponseResult::QuoteStored(record) => {
                                info!("QP {} quote stored for RFQ {}", qp_wallet.as_hex(), rfq_id);
                                if let Some(ref pubsub) = state.pubsub {
                                    let ws_quotes = record
                                        .quotes
                                        .iter()
                                        .map(|q| WsRfqQuoteEntry {
                                            quote_id: q.quote_id.to_string(),
                                            net_premium: q.net_premium,
                                            expires_at: q.received_at + q.valid_for_ms,
                                        })
                                        .collect();
                                    pubsub.publish_rfq_quotes(WsRfqQuotes {
                                        rfq_id: record.rfq_id.to_string(),
                                        quotes: ws_quotes,
                                        status: record.status.as_str().to_string(),
                                        taker_wallet: record.taker_wallet,
                                    });
                                }
                            }
                            HandleQpResponseResult::AutoExecuted {
                                record,
                                accepted_quote_id,
                            } => {
                                info!(
                                    "QP {} quote auto-executing for RFQ {} (quote={})",
                                    qp_wallet.as_hex(),
                                    rfq_id,
                                    accepted_quote_id,
                                );
                                // Push quotes update to taker first
                                if let Some(ref pubsub) = state.pubsub {
                                    let ws_quotes = record
                                        .quotes
                                        .iter()
                                        .map(|q| WsRfqQuoteEntry {
                                            quote_id: q.quote_id.to_string(),
                                            net_premium: q.net_premium,
                                            expires_at: q.received_at + q.valid_for_ms,
                                        })
                                        .collect();
                                    pubsub.publish_rfq_quotes(WsRfqQuotes {
                                        rfq_id: record.rfq_id.to_string(),
                                        quotes: ws_quotes,
                                        status: record.status.as_str().to_string(),
                                        taker_wallet: record.taker_wallet,
                                    });
                                }

                                // Spawn async execution task. The rfq_manager
                                // already transitioned the status to Accepted
                                // atomically; this task dispatches to the engine
                                // and updates to Executed/Failed.
                                let rfq_manager = Arc::clone(rfq_manager);
                                let pubsub = state.pubsub.clone();
                                let taker_wallet = record.taker_wallet;
                                let taker_sig = record.taker_signature.clone();
                                let rfq_id_str = record.rfq_id.to_string();
                                tokio::spawn(async move {
                                    match rfq_manager
                                        .auto_accept_quote(rfq_uuid, accepted_quote_id, taker_sig)
                                        .await
                                    {
                                        Ok(exec_result) => {
                                            let (status, fill_id) = match &exec_result {
                                                hypercall_runtime_api::RfqExecuteResult::Success { fill_id } => {
                                                    (hypercall_types::RfqStatus::Executed, Some(fill_id.clone()))
                                                }
                                                hypercall_runtime_api::RfqExecuteResult::Failed { .. } => {
                                                    (hypercall_types::RfqStatus::Failed, None)
                                                }
                                            };
                                            if let Some(ref pubsub) = pubsub {
                                                pubsub.publish_to_channel(
                                                    "rfq",
                                                    WsMessage::RfqAcceptResult {
                                                        rfq_id: rfq_id_str.clone(),
                                                        quote_id: accepted_quote_id.to_string(),
                                                        status: status.as_str().to_string(),
                                                        fill_id,
                                                        taker_wallet: Some(taker_wallet),
                                                    },
                                                );
                                                pubsub.publish_to_channel(
                                                    "rfq",
                                                    WsMessage::RfqStatusUpdate(WsRfqStatusUpdate {
                                                        rfq_id: rfq_id_str,
                                                        status: status.as_str().to_string(),
                                                        taker_wallet,
                                                    }),
                                                );
                                            }
                                            info!(
                                                "Auto-execute completed for RFQ: status={:?}",
                                                status
                                            );
                                        }
                                        Err(e) => {
                                            warn!(
                                                "Auto-execute failed for RFQ {}: {}",
                                                rfq_id_str, e
                                            );
                                            counter!(
                                                "ht_rfq_auto_execute_total",
                                                "result" => "engine_error"
                                            )
                                            .increment(1);
                                            // Notify the taker that auto-execution failed so
                                            // they can retry or pick another quote.
                                            if let Some(ref pubsub) = pubsub {
                                                pubsub.publish_to_channel(
                                                    "rfq",
                                                    WsMessage::RfqStatusUpdate(WsRfqStatusUpdate {
                                                        rfq_id: rfq_id_str.clone(),
                                                        status: "failed".to_string(),
                                                        taker_wallet,
                                                    }),
                                                );
                                            }
                                        }
                                    }
                                });
                            }
                            HandleQpResponseResult::AutoExecuteSkipped {
                                record,
                                quote_premium,
                                limit,
                            } => {
                                info!(
                                    "QP {} quote outside auto-execute limit for RFQ {} (premium={}, limit={})",
                                    qp_wallet.as_hex(),
                                    rfq_id,
                                    quote_premium,
                                    limit,
                                );
                                if let Some(ref pubsub) = state.pubsub {
                                    let ws_quotes = record
                                        .quotes
                                        .iter()
                                        .map(|q| WsRfqQuoteEntry {
                                            quote_id: q.quote_id.to_string(),
                                            net_premium: q.net_premium,
                                            expires_at: q.received_at + q.valid_for_ms,
                                        })
                                        .collect();
                                    pubsub.publish_rfq_quotes(WsRfqQuotes {
                                        rfq_id: record.rfq_id.to_string(),
                                        quotes: ws_quotes,
                                        status: "outside_limit".to_string(),
                                        taker_wallet: record.taker_wallet,
                                    });
                                }
                            }
                            HandleQpResponseResult::RpiQuoteStored(record) => {
                                info!(
                                    "QP {} RPI quote stored for RFQ {}",
                                    qp_wallet.as_hex(),
                                    rfq_id
                                );
                                if let Some(ref pubsub) = state.pubsub {
                                    let ws_quotes = record
                                        .quotes
                                        .iter()
                                        .map(|q| WsRfqQuoteEntry {
                                            quote_id: q.quote_id.to_string(),
                                            net_premium: q.net_premium,
                                            expires_at: q.received_at + q.valid_for_ms,
                                        })
                                        .collect();
                                    pubsub.publish_rfq_quotes(WsRfqQuotes {
                                        rfq_id: record.rfq_id.to_string(),
                                        quotes: ws_quotes,
                                        status: record.status.as_str().to_string(),
                                        taker_wallet: record.taker_wallet,
                                    });
                                }
                            }
                            HandleQpResponseResult::RpiQuoteRejected { record, reason } => {
                                info!(
                                    "QP {} RPI quote rejected for RFQ {}: {}",
                                    qp_wallet.as_hex(),
                                    rfq_id,
                                    reason
                                );
                                if let Some(ref pubsub) = state.pubsub {
                                    pubsub.publish_to_channel(
                                        "rfq",
                                        WsMessage::RfqStatusUpdate(WsRfqStatusUpdate {
                                            rfq_id: record.rfq_id.to_string(),
                                            status: record.status.as_str().to_string(),
                                            taker_wallet: record.taker_wallet,
                                        }),
                                    );
                                }
                            }
                        }
                    }
                    Err(e) => {
                        // Check if the error is because the RFQ is already
                        // accepted/executed (auto-accept race). Notify the QP
                        // immediately so it knows it lost, even if execution
                        // hasn't completed yet. The QP shouldn't wait for
                        // terminal execution before learning it lost the race.
                        if e.contains("invalid state for quotes") {
                            if let Some(record) = rfq_manager.get_rfq(&rfq_uuid) {
                                if record.status == hypercall_types::RfqStatus::Accepted
                                    || record.status == hypercall_types::RfqStatus::Executed
                                {
                                    let winning_quote_id = record
                                        .accepted_quote_id
                                        .map(|id| id.to_string())
                                        .unwrap_or_default();
                                    state.session_manager.send_to_qp(
                                        qp_wallet,
                                        QpOutboundMessage::RfqAlreadyFilled {
                                            rfq_id: rfq_id.clone(),
                                            filled_by_quote_id: winning_quote_id,
                                        },
                                    );
                                }
                            }
                        }
                        warn!(
                            "QP {} quote rejected for RFQ {}: {}",
                            qp_wallet.as_hex(),
                            rfq_id,
                            e
                        );
                    }
                }
            } else {
                warn!(
                    "Received RFQ response from QP {} but RfqManager not wired",
                    qp_wallet.as_hex()
                );
            }
        }
    }
}

fn publish_indicative_updates(
    pubsub: &Option<Arc<dyn RfqWebsocketPublisher>>,
    updates: Vec<(String, Option<AggregatedIndicativeQuote>)>,
) {
    let Some(pubsub) = pubsub else {
        return;
    };

    for (instrument, aggregate) in updates {
        let message = match aggregate {
            Some(aggregate) => WsIndicativeMarketData {
                instrument: aggregate.instrument,
                best_bid: Some(aggregate.best_bid),
                best_ask: Some(aggregate.best_ask),
                bid_iv: None,
                ask_iv: None,
                indicative_bid_size: Some(aggregate.indicative_bid_size),
                indicative_ask_size: Some(aggregate.indicative_ask_size),
                num_providers: aggregate.num_providers,
                timestamp: aggregate.updated_at,
            },
            None => WsIndicativeMarketData {
                instrument,
                best_bid: None,
                best_ask: None,
                bid_iv: None,
                ask_iv: None,
                indicative_bid_size: None,
                indicative_ask_size: None,
                num_providers: 0,
                timestamp: hypercall_types::utils::get_timestamp_millis(),
            },
        };

        pubsub.publish_indicative_market_data(message);
    }
}

use std::str::FromStr;

#[cfg(test)]
mod tests {
    use super::*;
    use crate::rfq::qp_sessions::QpSessionManager;
    use crate::websocket::{uuid, PubSubManager, WsMessage};
    use hypercall_ws_protocol::QpRfqLeg;
    use rust_decimal_macros::dec;
    use std::sync::Arc;

    fn test_wallet(id: u8) -> WalletAddress {
        let mut bytes = [0u8; 20];
        bytes[19] = id;
        WalletAddress::from(alloy::primitives::Address::from(bytes))
    }

    #[test]
    fn test_session_manager_register_remove() {
        let mgr = QpSessionManager::new();
        let wallet = test_wallet(1);
        let (tx, _rx) = mpsc::unbounded_channel();

        let session_id = mgr.register_session(wallet, tx);
        assert_eq!(mgr.get_connected_wallets().len(), 1);

        mgr.remove_session(&wallet, session_id);
        assert_eq!(mgr.get_connected_wallets().len(), 0);
    }

    #[test]
    fn test_session_manager_multiple_sessions_per_wallet() {
        let mgr = QpSessionManager::new();
        let wallet = test_wallet(1);
        let (tx1, _rx1) = mpsc::unbounded_channel();
        let (tx2, _rx2) = mpsc::unbounded_channel();

        let s1 = mgr.register_session(wallet, tx1);
        let s2 = mgr.register_session(wallet, tx2);

        // One wallet, two sessions
        assert_eq!(mgr.get_connected_wallets().len(), 1);

        // Removing one session while another remains must NOT report
        // fully-disconnected — downstream cleanup (e.g. evicting
        // indicative quotes from the cache) would otherwise clobber
        // state the surviving session is still actively updating.
        assert!(
            !mgr.remove_session(&wallet, s1),
            "remove_session must return false while other sessions remain"
        );
        assert_eq!(mgr.get_connected_wallets().len(), 1);

        assert!(
            mgr.remove_session(&wallet, s2),
            "remove_session must return true when the last session drops"
        );
        assert_eq!(mgr.get_connected_wallets().len(), 0);
    }

    #[test]
    fn test_send_to_qp_fans_out() {
        let mgr = QpSessionManager::new();
        let wallet = test_wallet(1);
        let (tx1, mut rx1) = mpsc::unbounded_channel();
        let (tx2, mut rx2) = mpsc::unbounded_channel();

        mgr.register_session(wallet, tx1);
        mgr.register_session(wallet, tx2);

        mgr.send_to_qp(
            &wallet,
            QpOutboundMessage::Authenticated {
                wallet: "0x01".to_string(),
            },
        );

        assert!(rx1.try_recv().is_ok());
        assert!(rx2.try_recv().is_ok());
    }

    #[test]
    fn test_send_rfq_to_eligible() {
        let mgr = QpSessionManager::new();
        let w1 = test_wallet(1);
        let w2 = test_wallet(2);
        let w3 = test_wallet(3);
        let (tx1, mut rx1) = mpsc::unbounded_channel();
        let (tx2, mut rx2) = mpsc::unbounded_channel();
        let (tx3, mut rx3) = mpsc::unbounded_channel();

        mgr.register_session(w1, tx1);
        mgr.register_session(w2, tx2);
        mgr.register_session(w3, tx3);

        // Only send to w1 and w3
        mgr.send_rfq_to_eligible(
            &[w1, w3],
            QpOutboundMessage::RfqRequest {
                rfq_id: "test".to_string(),
                legs: vec![],
                taker_wallet: "0x00".to_string(),
                request_timestamp: 0,
                response_deadline_ms: 5000,
                auto_accept_limit: None,
                auto_execute: false,
                taker_limit_price: None,
                reference_price: None,
                min_improvement_tick: None,
                auction_deadline_ms: None,
                requires_price_improvement: false,
            },
        );

        assert!(rx1.try_recv().is_ok());
        assert!(rx2.try_recv().is_err()); // w2 not eligible
        assert!(rx3.try_recv().is_ok());
    }

    #[test]
    fn test_qp_inbound_message_deserialize() {
        let json = r#"{"type":"indicative_quote_update","quotes":[{"instrument":"ETH-C-3000","bid_price":"120.5","ask_price":"125.0","max_bid_size":"10","max_ask_size":"10"}]}"#;
        let msg: QpInboundMessage = serde_json::from_str(json).unwrap();
        assert!(matches!(
            msg,
            QpInboundMessage::IndicativeQuoteUpdate { .. }
        ));
    }

    #[test]
    fn test_publish_indicative_updates_sends_market_data() {
        let pubsub = PubSubManager::new();
        let client_id = uuid::Uuid::new_v4();
        let (tx, mut rx) = mpsc::unbounded_channel::<Arc<WsMessage>>();
        pubsub.add_client(client_id, tx, false, None);
        pubsub
            .subscribe(
                client_id,
                "indicative_market_data".to_string(),
                None,
                None,
                None,
            )
            .unwrap();

        publish_indicative_updates(
            &Some(Arc::new(pubsub)),
            vec![(
                "ETH-20260131-3000-C".to_string(),
                Some(AggregatedIndicativeQuote {
                    instrument: "ETH-20260131-3000-C".to_string(),
                    best_bid: dec!(120.5),
                    best_ask: dec!(125.0),
                    indicative_bid_size: dec!(10),
                    indicative_ask_size: dec!(12),
                    num_providers: 2,
                    updated_at: 1_737_331_200_000,
                }),
            )],
        );

        let message = rx.try_recv().expect("indicative update should publish");
        match message.as_ref() {
            WsMessage::IndicativeMarketData(update) => {
                assert_eq!(update.instrument, "ETH-20260131-3000-C");
                assert_eq!(update.best_bid, Some(dec!(120.5)));
                assert_eq!(update.best_ask, Some(dec!(125.0)));
                assert_eq!(update.indicative_bid_size, Some(dec!(10)));
                assert_eq!(update.indicative_ask_size, Some(dec!(12)));
                assert_eq!(update.num_providers, 2);
                assert_eq!(update.timestamp, 1_737_331_200_000);
            }
            other => panic!("expected IndicativeMarketData, got {other:?}"),
        }
    }

    #[test]
    fn test_publish_indicative_updates_sends_removal() {
        let pubsub = PubSubManager::new();
        let client_id = uuid::Uuid::new_v4();
        let (tx, mut rx) = mpsc::unbounded_channel::<Arc<WsMessage>>();
        pubsub.add_client(client_id, tx, false, None);
        pubsub
            .subscribe(
                client_id,
                "indicative_market_data".to_string(),
                None,
                None,
                None,
            )
            .unwrap();

        publish_indicative_updates(
            &Some(Arc::new(pubsub)),
            vec![("ETH-20260131-3000-C".to_string(), None)],
        );

        let message = rx.try_recv().expect("indicative removal should publish");
        match message.as_ref() {
            WsMessage::IndicativeMarketData(update) => {
                assert_eq!(update.instrument, "ETH-20260131-3000-C");
                assert_eq!(update.best_bid, None);
                assert_eq!(update.best_ask, None);
                assert_eq!(update.indicative_bid_size, None);
                assert_eq!(update.indicative_ask_size, None);
                assert_eq!(update.num_providers, 0);
                assert!(update.timestamp > 0);
            }
            other => panic!("expected IndicativeMarketData removal, got {other:?}"),
        }
    }

    #[test]
    fn test_qp_outbound_message_serialize() {
        let msg = QpOutboundMessage::RfqRequest {
            rfq_id: "abc".to_string(),
            legs: vec![QpRfqLeg {
                instrument: "ETH-C-3000".to_string(),
                side: "buy".to_string(),
                size: "5.0".to_string(),
            }],
            taker_wallet: "0x123".to_string(),
            request_timestamp: 1000,
            response_deadline_ms: 5000,
            auto_accept_limit: None,
            auto_execute: false,
            taker_limit_price: None,
            reference_price: None,
            min_improvement_tick: None,
            auction_deadline_ms: None,
            requires_price_improvement: false,
        };
        let json = serde_json::to_string(&msg).unwrap();
        assert!(json.contains("rfq_request"));
        assert!(json.contains("ETH-C-3000"));
        assert!(!json.contains("auto_accept_limit"));
    }

    #[test]
    fn test_qp_outbound_message_serializes_auto_accept_limit() {
        let msg = QpOutboundMessage::RfqRequest {
            rfq_id: "abc".to_string(),
            legs: vec![QpRfqLeg {
                instrument: "BTC-20260501-90000-C".to_string(),
                side: "buy".to_string(),
                size: "1".to_string(),
            }],
            taker_wallet: "0x123".to_string(),
            request_timestamp: 1000,
            response_deadline_ms: 5000,
            auto_accept_limit: Some("3999".to_string()),
            auto_execute: true,
            taker_limit_price: Some("3999".to_string()),
            reference_price: Some("3999".to_string()),
            min_improvement_tick: Some("0.0001".to_string()),
            auction_deadline_ms: Some(2000),
            requires_price_improvement: true,
        };
        let json = serde_json::to_string(&msg).unwrap();
        assert!(json.contains(r#""auto_accept_limit":"3999""#));
        assert!(json.contains(r#""auto_execute":true"#));
        assert!(json.contains(r#""requires_price_improvement":true"#));
    }
}