codex_mobile_bridge/

server.rs

use std::sync::Arc;

use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
use axum::extract::{Query, State};
use axum::http::{HeaderMap, StatusCode};
use axum::response::{IntoResponse, Response};
use axum::routing::get;
use axum::{Json, Router};
use futures_util::{SinkExt, StreamExt};
use serde::Deserialize;
use serde_json::{Value, json};
use tracing::{info, warn};

use crate::bridge_protocol::{
    ApiError, ClientEnvelope, RuntimeStatusSnapshot, RuntimeSummary, ServerEnvelope,
    error_response, event_envelope, ok_response,
};
use crate::state::BridgeState;

#[derive(Debug, Deserialize, Default)]
struct WsQuery {
    token: Option<String>,
}

pub fn build_router(state: Arc<BridgeState>) -> Router {
    Router::new()
        .route("/health", get(health_handler))
        .route("/ws", get(ws_handler))
        .with_state(state)
}

async fn health_handler(State(state): State<Arc<BridgeState>>) -> Json<Value> {
    let runtime = state.runtime_snapshot_for_client().await;
    let runtimes = state.runtime_summaries_for_client().await;
    Json(build_health_payload(&runtime, &runtimes))
}

fn build_health_payload(runtime: &RuntimeStatusSnapshot, runtimes: &[RuntimeSummary]) -> Value {
    let primary_runtime_id = runtimes
        .iter()
        .find(|item| item.is_primary)
        .map(|item| item.runtime_id.clone());

    json!({
        "ok": true,
        "bridgeVersion": crate::BRIDGE_VERSION,
        "buildHash": crate::BRIDGE_BUILD_HASH,
        "protocolVersion": crate::BRIDGE_PROTOCOL_VERSION,
        "runtimeCount": runtimes.len(),
        "primaryRuntimeId": primary_runtime_id,
        "runtime": runtime,
    })
}

async fn ws_handler(
    State(state): State<Arc<BridgeState>>,
    Query(query): Query<WsQuery>,
    headers: HeaderMap,
    ws: WebSocketUpgrade,
) -> Response {
    match authorize(&state, &query, &headers) {
        Ok(()) => ws
            .on_upgrade(move |socket| handle_socket(state, socket))
            .into_response(),
        Err(error) => (StatusCode::UNAUTHORIZED, error).into_response(),
    }
}

fn authorize(
    state: &BridgeState,
    query: &WsQuery,
    headers: &HeaderMap,
) -> Result<(), &'static str> {
    let token = query
        .token
        .clone()
        .or_else(|| {
            headers
                .get(axum::http::header::AUTHORIZATION)
                .and_then(|value| value.to_str().ok())
                .and_then(|value| value.strip_prefix("Bearer "))
                .map(ToOwned::to_owned)
        })
        .ok_or("missing token")?;

    if token == state.config_token() {
        Ok(())
    } else {
        Err("invalid token")
    }
}

async fn handle_socket(state: Arc<BridgeState>, socket: WebSocket) {
    let (mut sender, mut receiver) = socket.split();
    let mut event_rx = state.subscribe_events();
    let mut device_id: Option<String> = None;

    loop {
        tokio::select! {
            incoming = receiver.next() => {
                let Some(incoming) = incoming else {
                    info!(
                        "bridge ws 对端已断开 device_id={}",
                        device_id.as_deref().unwrap_or("<pending>")
                    );
                    break;
                };

                let Ok(message) = incoming else {
                    warn!(
                        "bridge ws 接收失败 device_id={}: {:?}",
                        device_id.as_deref().unwrap_or("<pending>"),
                        incoming.err()
                    );
                    break;
                };

                match handle_incoming_message(&state, &mut sender, &mut device_id, message).await {
                    Ok(should_continue) if should_continue => {}
                    Ok(_) => break,
                    Err(error) => {
                        warn!(
                            "bridge ws 处理消息失败 device_id={}: {error}",
                            device_id.as_deref().unwrap_or("<pending>")
                        );
                        break;
                    }
                }
            }
            broadcast_result = event_rx.recv(), if device_id.is_some() => {
                match broadcast_result {
                    Ok(event) => {
                        if send_json(&mut sender, &event_envelope(event)).await.is_err() {
                            break;
                        }
                    }
                    Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {
                        let envelope = ServerEnvelope::Response {
                            request_id: "system".to_string(),
                            success: false,
                            data: None,
                            error: Some(ApiError::new("lagged", "事件流丢失，请重新连接")),
                        };
                        let _ = send_json(&mut sender, &envelope).await;
                        break;
                    }
                    Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
                }
            }
        }
    }
}

async fn handle_incoming_message(
    state: &BridgeState,
    sender: &mut futures_util::stream::SplitSink<WebSocket, Message>,
    device_id: &mut Option<String>,
    message: Message,
) -> anyhow::Result<bool> {
    let text = match message {
        Message::Text(text) => text,
        Message::Close(frame) => {
            let detail = frame
                .as_ref()
                .map(|close| format!("code={} reason={}", close.code, close.reason))
                .unwrap_or_else(|| "no close frame".to_string());
            info!(
                "bridge ws 收到 close 帧 device_id={}: {detail}",
                device_id.as_deref().unwrap_or("<pending>")
            );
            return Ok(false);
        }
        _ => return Ok(true),
    };

    let envelope = parse_client_envelope(&text).map_err(|error| {
        anyhow::anyhow!(
            "解析客户端消息失败: {error}; payload={}",
            truncate_text(&text, 240)
        )
    })?;
    match envelope {
        ClientEnvelope::Hello {
            device_id: next_device_id,
            last_ack_seq,
        } => {
            info!(
                "bridge ws 收到 hello device_id={} last_ack_seq={last_ack_seq:?}",
                next_device_id
            );
            let (runtime, runtimes, workspaces, pending_requests, replay_events) =
                state.hello_payload(&next_device_id, last_ack_seq).await?;
            *device_id = Some(next_device_id);
            let connected_device_id = device_id.as_deref().unwrap_or("<pending>");

            send_json(
                sender,
                &ServerEnvelope::Hello {
                    bridge_version: crate::BRIDGE_VERSION.to_string(),
                    protocol_version: crate::BRIDGE_PROTOCOL_VERSION,
                    runtime,
                    runtimes,
                    workspaces,
                    pending_requests,
                },
            )
            .await?;

            info!(
                "bridge ws hello 已完成 device_id={} replay_events={}",
                connected_device_id,
                replay_events.len()
            );
            for event in replay_events {
                send_json(sender, &event_envelope(event)).await?;
            }
        }
        ClientEnvelope::Request {
            request_id,
            action,
            payload,
        } => {
            if device_id.is_none() {
                send_json(
                    sender,
                    &error_response(
                        request_id,
                        ApiError::new("handshake_required", "请先发送 hello"),
                    ),
                )
                .await?;
                return Ok(true);
            }

            let response = match state.handle_request(&action, payload).await {
                Ok(data) => ok_response(request_id, data),
                Err(error) => error_response(
                    request_id,
                    ApiError::new("request_failed", error.to_string()),
                ),
            };
            send_json(sender, &response).await?;
        }
        ClientEnvelope::AckEvents { last_seq } => {
            if let Some(device_id) = device_id.as_deref() {
                state.ack_events(device_id, last_seq)?;
            }
        }
        ClientEnvelope::Ping => {
            send_json(
                sender,
                &ServerEnvelope::Pong {
                    server_time_ms: crate::bridge_protocol::now_millis(),
                },
            )
            .await?;
        }
    }

    Ok(true)
}

fn parse_client_envelope(text: &str) -> anyhow::Result<ClientEnvelope> {
    match serde_json::from_str::<ClientEnvelope>(text) {
        Ok(envelope) => Ok(envelope),
        Err(primary_error) => {
            let nested_payload = serde_json::from_str::<String>(text).map_err(|_| primary_error)?;
            serde_json::from_str::<ClientEnvelope>(&nested_payload).map_err(Into::into)
        }
    }
}

fn truncate_text(text: &str, max_chars: usize) -> String {
    let mut truncated = String::new();
    for (index, character) in text.chars().enumerate() {
        if index >= max_chars {
            truncated.push('…');
            return truncated;
        }
        truncated.push(character);
    }
    truncated
}

async fn send_json(
    sender: &mut futures_util::stream::SplitSink<WebSocket, Message>,
    envelope: &ServerEnvelope,
) -> anyhow::Result<()> {
    let text = serde_json::to_string(envelope)?;
    sender.send(Message::Text(text.into())).await?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use std::env;
    use std::fs;
    use std::path::PathBuf;
    use std::sync::Arc;

    use axum::extract::State;
    use serde_json::{Value, json};
    use tokio::time::{Duration, timeout};
    use uuid::Uuid;

    use super::build_health_payload;
    use super::health_handler;
    use super::parse_client_envelope;
    use crate::bridge_protocol::{
        ClientEnvelope, RuntimeRecord, RuntimeStatusSnapshot, RuntimeSummary,
    };
    use crate::config::Config;
    use crate::state::BridgeState;

    #[test]
    fn build_health_payload_contains_bridge_metadata_and_primary_runtime() {
        let runtime = RuntimeStatusSnapshot {
            runtime_id: "primary".to_string(),
            status: "running".to_string(),
            codex_home: Some("/srv/codex-home".to_string()),
            user_agent: Some("codex-mobile".to_string()),
            platform_family: Some("linux".to_string()),
            platform_os: Some("ubuntu".to_string()),
            last_error: None,
            pid: Some(4242),
            updated_at_ms: 1234,
        };
        let runtime_record = RuntimeRecord {
            runtime_id: "primary".to_string(),
            display_name: "Primary".to_string(),
            codex_home: Some("/srv/codex-home".to_string()),
            codex_binary: "codex".to_string(),
            is_primary: true,
            auto_start: true,
            created_at_ms: 1000,
            updated_at_ms: 1000,
        };
        let runtimes = vec![RuntimeSummary::from_parts(&runtime_record, runtime.clone())];

        let payload = build_health_payload(&runtime, &runtimes);

        assert_eq!(payload["ok"], Value::Bool(true));
        assert_eq!(
            payload["bridgeVersion"],
            Value::String(crate::BRIDGE_VERSION.to_string())
        );
        assert_eq!(
            payload["buildHash"],
            Value::String(crate::BRIDGE_BUILD_HASH.to_string())
        );
        assert_eq!(
            payload["protocolVersion"],
            Value::Number(crate::BRIDGE_PROTOCOL_VERSION.into())
        );
        assert_eq!(payload["runtimeCount"], Value::Number(1.into()));
        assert_eq!(
            payload["primaryRuntimeId"],
            Value::String("primary".to_string())
        );
        assert_eq!(
            payload["runtime"]["runtimeId"],
            Value::String("primary".to_string())
        );
        assert_eq!(
            payload["runtime"]["status"],
            Value::String("running".to_string())
        );
    }

    #[test]
    fn parse_client_envelope_accepts_plain_hello_payload() {
        let envelope = parse_client_envelope(
            r#"{"kind":"hello","device_id":"device-alpha","last_ack_seq":7}"#,
        )
        .expect("hello payload 应可解析");

        match envelope {
            ClientEnvelope::Hello {
                device_id,
                last_ack_seq,
            } => {
                assert_eq!(device_id, "device-alpha");
                assert_eq!(last_ack_seq, Some(7));
            }
            _ => panic!("应解析为 hello"),
        }
    }

    #[test]
    fn parse_client_envelope_accepts_double_encoded_hello_payload() {
        let envelope = parse_client_envelope(
            r#""{\"kind\":\"hello\",\"device_id\":\"device-beta\",\"last_ack_seq\":9}""#,
        )
        .expect("双重编码 hello payload 应可解析");

        match envelope {
            ClientEnvelope::Hello {
                device_id,
                last_ack_seq,
            } => {
                assert_eq!(device_id, "device-beta");
                assert_eq!(last_ack_seq, Some(9));
            }
            _ => panic!("应解析为 hello"),
        }
    }

    #[tokio::test]
    async fn runtime_snapshot_returns_without_hanging() {
        let state = bootstrap_test_state().await;

        let snapshot = timeout(Duration::from_secs(2), state.runtime_snapshot())
            .await
            .expect("runtime_snapshot 超时");

        assert_eq!(snapshot.runtime_id, "primary");
    }

    #[tokio::test]
    async fn runtime_summaries_return_without_hanging() {
        let state = bootstrap_test_state().await;

        let summaries = timeout(Duration::from_secs(2), state.runtime_summaries())
            .await
            .expect("runtime_summaries 超时");

        assert!(!summaries.is_empty());
        assert_eq!(summaries[0].runtime_id, "primary");
    }

    #[tokio::test]
    async fn health_handler_returns_without_hanging() {
        let state = bootstrap_test_state().await;

        let _ = timeout(
            Duration::from_secs(2),
            health_handler(State(Arc::clone(&state))),
        )
        .await
        .expect("/health handler 超时");
    }

    #[tokio::test]
    async fn hello_payload_returns_without_hanging() {
        let state = bootstrap_test_state().await;

        let (runtime, runtimes, ..) = timeout(
            Duration::from_secs(2),
            state.hello_payload("device-test", None),
        )
        .await
        .expect("hello_payload 超时")
        .expect("hello_payload 返回错误");

        assert_eq!(runtime.runtime_id, "primary");
        assert!(!runtimes.is_empty());
        assert_eq!(runtimes[0].runtime_id, "primary");
    }

    #[tokio::test]
    async fn list_runtimes_request_returns_without_hanging() {
        let state = bootstrap_test_state().await;

        let response = timeout(
            Duration::from_secs(2),
            state.handle_request("list_runtimes", json!({})),
        )
        .await
        .expect("list_runtimes 超时")
        .expect("list_runtimes 返回错误");

        let runtimes = response["runtimes"].as_array().expect("runtimes 应为数组");
        assert!(!runtimes.is_empty());
        assert_eq!(
            runtimes[0]["runtimeId"],
            Value::String("primary".to_string())
        );
    }

    #[tokio::test]
    async fn get_runtime_status_request_returns_without_hanging() {
        let state = bootstrap_test_state().await;

        let response = timeout(
            Duration::from_secs(2),
            state.handle_request("get_runtime_status", json!({ "runtimeId": "primary" })),
        )
        .await
        .expect("get_runtime_status 超时")
        .expect("get_runtime_status 返回错误");

        assert_eq!(
            response["runtime"]["runtimeId"],
            Value::String("primary".to_string())
        );
    }

    async fn bootstrap_test_state() -> Arc<BridgeState> {
        let base_dir = env::temp_dir().join(format!("codex-mobile-bridge-test-{}", Uuid::new_v4()));
        fs::create_dir_all(&base_dir).expect("创建测试目录失败");
        let db_path = base_dir.join("bridge.db");

        let config = Config {
            listen_addr: "127.0.0.1:0".to_string(),
            token: "test-token".to_string(),
            runtime_limit: 4,
            db_path,
            codex_home: None,
            codex_binary: resolve_true_binary(),
            workspace_roots: Vec::new(),
        };

        BridgeState::bootstrap(config)
            .await
            .expect("bootstrap 测试 BridgeState 失败")
    }

    fn resolve_true_binary() -> String {
        for candidate in ["/usr/bin/true", "/bin/true"] {
            if PathBuf::from(candidate).exists() {
                return candidate.to_string();
            }
        }
        "true".to_string()
    }
}