DevFlow/crates/df-im-gateway/src/commander.rs

//! WS Commander 客户端(连 u-relay /ws/commander)
//!
//! 与桌面端 df-tunnel(Agent 角色)对称:网关作为 Commander 发 TunnelCommand、收 TunnelEvent。
//! 协议类型直接复用 df_tunnel::events(编译期一致)。
//!
//! 连接生命周期:hello {kind:"commander", channel_id, token} → (relay 校验) → 双向消息。
//! 断线重连:指数退避(1/2/4/8/16/30s 封顶,与桌面端 df-tunnel 同款策略)。
//!
//! 注:relay 对同 channel 多 Commander 的投递行为(P0 必验项 ①)待实测——
//! 过渡期约定同 channel 单活跃控制端。

use std::time::Duration;

use futures_util::{SinkExt, StreamExt};
use tokio::sync::mpsc;
use tokio_tungstenite::tungstenite::Message;
use tracing::{info, warn};

use df_tunnel::events::{Hello, HelloKind, RelayControlError};
use df_tunnel::{TunnelCommand, TunnelEvent, TunnelMessage};

/// 指数退避序列(秒,封顶 30)
fn backoff_secs(attempt: u32) -> u64 {
    [1, 2, 4, 8, 16, 30]
        .get(attempt.min(5) as usize)
        .copied()
        .unwrap_or(30)
}

/// 启动 Commander 连接循环(阻塞直到 shutdown)。重连在内部循环,指令/事件双向透传。
///
/// * `relay_url` — ws://host:9180/ws/commander
/// * `channel_id` — 配对的桌面端 channel
/// * `token` — U_RELAY_TOKEN 同源
/// * `ev_tx` — 事件出口(桌面端 → 网关核心)
/// * `cmd_rx` — 指令入口(网关核心 → 桌面端)
/// * `shutdown` — 关闭信号
pub async fn run_commander(
    relay_url: String,
    channel_id: String,
    token: String,
    ev_tx: mpsc::Sender<TunnelEvent>,
    mut cmd_rx: mpsc::Receiver<TunnelCommand>,
    shutdown: tokio::sync::watch::Receiver<bool>,
) -> anyhow::Result<()> {
    let mut attempt: u32 = 0;
    loop {
        if *shutdown.borrow() {
            info!("[commander] 收到关闭信号,退出");
            return Ok(());
        }
        match connect_once(&relay_url, &channel_id, &token, &ev_tx, &mut cmd_rx, shutdown.clone()).await {
            Ok(()) => {
                info!("[commander] 连接正常关闭");
                return Ok(());
            }
            Err(e) => {
                let wait = backoff_secs(attempt);
                attempt += 1;
                warn!("[commander] 连接断开({e}),{}s 后重连(第 {attempt} 次)", wait);
                let mut shutdown_w = shutdown.clone();
                tokio::select! {
                    _ = tokio::time::sleep(Duration::from_secs(wait)) => {}
                    _ = shutdown_w.changed() => return Ok(()),
                }
            }
        }
    }
}

/// 单次连接会话(hello 握手 + 消息循环),断开/出错返回 Err 触发外层重连
async fn connect_once(
    relay_url: &str,
    channel_id: &str,
    token: &str,
    ev_tx: &mpsc::Sender<TunnelEvent>,
    cmd_rx: &mut mpsc::Receiver<TunnelCommand>,
    mut shutdown: tokio::sync::watch::Receiver<bool>,
) -> anyhow::Result<()> {
    let (ws, _resp) = tokio_tungstenite::connect_async(relay_url).await?;
    let (mut sink, mut stream) = ws.split();

    // hello 握手(kind=commander;字段与 u-relay::relay::Hello 严格一致)
    let hello = Hello {
        kind: HelloKind::Commander,
        channel_id: channel_id.to_string(),
        token: token.to_string(),
    };
    sink.send(Message::Text(serde_json::to_string(&hello)?)).await?;

    // 等待 hello_ack(首条文本帧;control error 帧则报错)
    let ack = tokio::time::timeout(Duration::from_secs(10), stream.next()).await
        .map_err(|_| anyhow::anyhow!("hello_ack 10s 超时"))?
        .ok_or_else(|| anyhow::anyhow!("连接在握手期关闭"))?;
    match ack? {
        Message::Text(t) => {
            // relay 控制帧:{"kind":"control","error":"..."} → 握手被拒
            if let Ok(ctrl_err) = serde_json::from_str::<RelayControlError>(&t) {
                if ctrl_err.kind == "control" {
                    anyhow::bail!("握手被拒: {}", ctrl_err.error);
                }
            }
            info!("[commander] hello_ack: channel={channel_id} 握手成功");
        }
        other => anyhow::bail!("hello_ack 非文本帧: {other:?}"),
    }

    // 心跳:30s ping(relay 90s 读死线)
    let mut ping = tokio::time::interval(Duration::from_secs(30));
    ping.tick().await; // 首个立即返回,跳过

    // 消息循环
    loop {
        tokio::select! {
            _ = shutdown.changed() => {
                let _ = sink.close().await;
                return Ok(());
            }
            _ = ping.tick() => {
                sink.send(Message::Ping(vec![])).await?;
            }
            cmd = cmd_rx.recv() => {
                let Some(cmd) = cmd else {
                    // 核心侧通道关闭:正常收尾
                    let _ = sink.close().await;
                    return Ok(());
                };
                let frame = TunnelMessage::Command(cmd);
                sink.send(Message::Text(serde_json::to_string(&frame)?)).await?;
            }
            msg = stream.next() => {
                match msg {
                    Some(Ok(Message::Text(t))) => {
                        match serde_json::from_str::<TunnelMessage>(&t) {
                            Ok(TunnelMessage::Event(e)) => {
                                if ev_tx.send(e).await.is_err() {
                                    return Err(anyhow::anyhow!("核心事件通道关闭"));
                                }
                            }
                            Ok(TunnelMessage::Command(_)) => {
                                warn!("[commander] commander 收到 Command 帧(协议错向,忽略)");
                            }
                            Err(_) => warn!("[commander] 无法解析帧(非 TunnelMessage,忽略): {} 字节", t.len()),
                        }
                    }
                    Some(Ok(Message::Ping(p))) => { let _ = sink.send(Message::Pong(p)).await; }
                    Some(Ok(_)) => {}
                    Some(Err(e)) => return Err(e.into()),
                    None => return Err(anyhow::anyhow!("stream 结束")),
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn backoff_sequence() {
        assert_eq!(backoff_secs(0), 1);
        assert_eq!(backoff_secs(1), 2);
        assert_eq!(backoff_secs(2), 4);
        assert_eq!(backoff_secs(3), 8);
        assert_eq!(backoff_secs(4), 16);
        assert_eq!(backoff_secs(5), 30);
        assert_eq!(backoff_secs(99), 30, "封顶 30s");
    }

    #[test]
    fn tunnel_message_roundtrip() {
        let cmd = TunnelMessage::Command(TunnelCommand::Send {
            conv_id: "c1".into(),
            content: "你好".into(),
        });
        let json = serde_json::to_string(&cmd).unwrap();
        assert!(json.contains("\"command\""));
        let back: TunnelMessage = serde_json::from_str(&json).unwrap();
        match back {
            TunnelMessage::Command(TunnelCommand::Send { conv_id, content }) => {
                assert_eq!(conv_id, "c1");
                assert_eq!(content, "你好");
            }
            _ => panic!("roundtrip 失败"),
        }
    }
}