harnessd 0.1.0

//! The `/v1/ws` WebSocket handler: handshake, subscribe/replay/live-tail, and the
//! client command channel (design doc §5.1–5.2).
//!
//! One task per connection. It multiplexes two sources with `select!`:
//!   * inbound client [`ClientMessage`]s (subscribe, send_message, approve, …), and
//!   * the runtime's broadcast bus of [`EventEnvelope`]s, filtered to the sessions this
//!     connection subscribed to.

use crate::server::AppState;
use axum::extract::ws::{Message, WebSocket};
use axum::extract::{State, WebSocketUpgrade};
use axum::response::IntoResponse;
use futures::{SinkExt, StreamExt};
use harness_proto::{ClientMessage, Hello, ServerMessage, UserInfo, Welcome, PROTOCOL_VERSION};
use std::collections::HashSet;
use tracing::{debug, info, warn};

pub async fn handler(ws: WebSocketUpgrade, State(st): State<AppState>) -> impl IntoResponse {
    ws.on_upgrade(move |socket| connection(socket, st))
}

/// Serialize a [`ServerMessage`] and send it. Returns `Err` if the socket is closed.
async fn send(
    sink: &mut futures::stream::SplitSink<WebSocket, Message>,
    msg: &ServerMessage,
) -> anyhow::Result<()> {
    let text = serde_json::to_string(msg)?;
    sink.send(Message::Text(text)).await?;
    Ok(())
}

async fn connection(socket: WebSocket, st: AppState) {
    let (mut sink, mut stream) = socket.split();

    // ── Handshake: first frame must be `hello`. ──────────────────────────────
    let hello: Hello = match stream.next().await {
        Some(Ok(Message::Text(t))) => match serde_json::from_str::<ClientMessage>(&t) {
            Ok(ClientMessage::Hello(h)) => h,
            _ => {
                let _ = send(
                    &mut sink,
                    &ServerMessage::Error {
                        code: "bad_handshake".into(),
                        message: "first frame must be `hello`".into(),
                    },
                )
                .await;
                return;
            }
        },
        _ => return,
    };

    // Version negotiation: accept current and current-1.
    if hello.protocol > PROTOCOL_VERSION || hello.protocol + 1 < PROTOCOL_VERSION {
        let _ = send(
            &mut sink,
            &ServerMessage::Error {
                code: "unsupported_protocol".into(),
                message: format!("daemon speaks protocol {PROTOCOL_VERSION}"),
            },
        )
        .await;
        return;
    }

    if !st.accepts_token(hello.device_token.as_deref()) {
        let _ = send(
            &mut sink,
            &ServerMessage::Error {
                code: "unauthorized".into(),
                message: "device token rejected".into(),
            },
        )
        .await;
        return;
    }

    info!(client = %hello.client, "client connected");
    let welcome = ServerMessage::Welcome(Welcome {
        protocol: PROTOCOL_VERSION,
        daemon: format!("harnessd/{}", st.version),
        capabilities: vec!["multi_provider".into(), "event_log".into()],
        user: UserInfo {
            name: st.runtime.config().user_name.clone(),
        },
    });
    if send(&mut sink, &welcome).await.is_err() {
        return;
    }

    // ── Live loop. ───────────────────────────────────────────────────────────
    // Device label recorded on approval events (`by`). Uses the client's self-reported
    // name until pairing (M4) assigns a stable per-device identity.
    let device = hello.client.clone();
    let mut bus = st.runtime.subscribe();
    let mut subscribed: HashSet<String> = HashSet::new();

    loop {
        tokio::select! {
            // A broadcast event: forward if the client subscribed to its session.
            ev = bus.recv() => {
                match ev {
                    Ok(env) => {
                        if subscribed.contains(&env.session)
                            && send(&mut sink, &ServerMessage::Event(env)).await.is_err()
                        {
                            break;
                        }
                    }
                    // Lagged: the client fell behind the broadcast buffer. It can
                    // re-subscribe with a from_seq to catch up from the durable log.
                    Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
                        warn!(skipped = n, "client lagged the event bus");
                    }
                    Err(_) => break,
                }
            }

            // An inbound client frame.
            frame = stream.next() => {
                match frame {
                    Some(Ok(Message::Text(t))) => {
                        match serde_json::from_str::<ClientMessage>(&t) {
                            Ok(msg) => {
                                if handle_client_msg(&st, &mut sink, &mut subscribed, &device, msg).await.is_err() {
                                    break;
                                }
                            }
                            Err(e) => {
                                debug!(error = %e, "ignoring unparseable client frame");
                            }
                        }
                    }
                    Some(Ok(Message::Close(_))) | None => break,
                    Some(Ok(_)) => {} // ignore binary/ping/pong
                    Some(Err(e)) => { debug!(error = %e, "ws error"); break; }
                }
            }
        }
    }

    info!(client = %hello.client, "client disconnected");
}

/// Handle one client command. Returns `Err` only on a fatal socket write failure.
async fn handle_client_msg(
    st: &AppState,
    sink: &mut futures::stream::SplitSink<WebSocket, Message>,
    subscribed: &mut HashSet<String>,
    device: &str,
    msg: ClientMessage,
) -> anyhow::Result<()> {
    match msg {
        ClientMessage::Hello(_) => { /* already handshook; ignore */ }

        ClientMessage::Subscribe { session, from_seq } => {
            // Replay the durable log first, then mark live so the select loop tails it.
            match st.runtime.store().events_since(&session, from_seq) {
                Ok(events) => {
                    for env in events {
                        send(sink, &ServerMessage::Event(env)).await?;
                    }
                    subscribed.insert(session);
                    send(sink, &ServerMessage::Ack { ok: true }).await?;
                }
                Err(e) => {
                    send(
                        sink,
                        &ServerMessage::Error {
                            code: e.code().into(),
                            message: e.to_string(),
                        },
                    )
                    .await?;
                }
            }
        }

        ClientMessage::Unsubscribe { session } => {
            subscribed.remove(&session);
        }

        ClientMessage::SendMessage { session, content } => {
            // Run the turn in the background so the socket stays responsive for
            // approvals/cancel while the model streams. Events reach this client via
            // the bus (assuming it subscribed).
            let rt = st.runtime.clone();
            tokio::spawn(async move {
                if let Err(e) = rt.handle_user_message(&session, &content).await {
                    warn!(error = %e, %session, "turn failed");
                }
            });
            send(sink, &ServerMessage::Ack { ok: true }).await?;
        }

        // Approval gate commands (M2). The resulting tool.approved / tool.denied
        // events reach subscribed clients over the bus.
        ClientMessage::Approve {
            session,
            call_id,
            always,
        } => {
            let ok = st
                .runtime
                .approve(&session, &call_id, device, always)
                .is_ok();
            send(sink, &ServerMessage::Ack { ok }).await?;
        }
        ClientMessage::Deny {
            session,
            call_id,
            note,
        } => {
            let ok = st.runtime.deny(&session, &call_id, device, note).is_ok();
            send(sink, &ServerMessage::Ack { ok }).await?;
        }
        ClientMessage::CancelTurn { session } => {
            let ok = st.runtime.cancel_turn(&session).is_ok();
            send(sink, &ServerMessage::Ack { ok }).await?;
        }
        ClientMessage::ResolvePlan {
            session,
            plan_id,
            approve,
            note,
        } => {
            let ok = st
                .runtime
                .resolve_plan(&session, &plan_id, approve, device, note)
                .is_ok();
            send(sink, &ServerMessage::Ack { ok }).await?;
        }
        ClientMessage::SetAutoApprove { session, enabled } => {
            st.runtime.set_auto_approve(&session, enabled);
            send(sink, &ServerMessage::Ack { ok: true }).await?;
        }
    }
    Ok(())
}