clawmates/crates/cm-api/src/routes/terminal.rs

//! The Terminal computer app: an interactive PTY (`zsh -l`) in the agent's
//! themed terminal container, bridged to the browser's xterm.js over a
//! WebSocket. Because a browser WS handshake can't carry the bearer header, the
//! client first POSTs (through the authed proxy) for a short-lived single-use
//! ticket, then opens the WS with `?ticket=`. The WS itself is routed straight
//! to this server by the edge (Traefik), bypassing the HTTP-only Next proxy.

use std::sync::Arc;

use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
use axum::extract::{Path, Query, State};
use axum::http::StatusCode;
use axum::response::{IntoResponse, Response};
use axum::Json;
use cm_db::repo::audit::Actor;
use cm_domain::{AgentId, NodeId, WorkspaceId};
use futures::{SinkExt, StreamExt};
use serde::{Deserialize, Serialize};
use serde_json::json;
use sqlx::{PgPool, Row};
use tokio::io::AsyncWriteExt;

use crate::fleet::NodeHub;
use crate::routes::claws::workspace_agent;
use crate::{ApiError, AppState, Authed};

/// Single-use WS tickets live in Postgres (hashed), so any server replica can
/// redeem one — not just the instance that minted it. A ticket is deleted on
/// redeem; expired rows are swept opportunistically. TTL is 30s (in the SQL).
fn hash_ticket(token: &str) -> String {
    use sha2::{Digest, Sha256};
    let mut h = Sha256::new();
    h.update(token.as_bytes());
    format!("{:x}", h.finalize())
}

/// Mint a short-lived single-use ticket; returns the raw token.
async fn issue_ticket(
    pool: &PgPool,
    agent_id: AgentId,
    workspace_id: WorkspaceId,
    label: &str,
) -> Result<String, sqlx::Error> {
    let token = uuid::Uuid::new_v4().simple().to_string();
    let _ = sqlx::query("DELETE FROM terminal_ws_tickets WHERE expires_at < now()")
        .execute(pool)
        .await;
    sqlx::query(
        "INSERT INTO terminal_ws_tickets (token_hash, agent_id, workspace_id, label, expires_at)
         VALUES ($1, $2, $3, $4, now() + interval '30 seconds')",
    )
    .bind(hash_ticket(&token))
    .bind(agent_id.as_uuid())
    .bind(workspace_id.as_uuid())
    .bind(label)
    .execute(pool)
    .await?;
    Ok(token)
}

/// Consume a ticket; returns (workspace, greeting label) iff valid for this agent.
async fn redeem_ticket(
    pool: &PgPool,
    token: &str,
    agent_id: AgentId,
) -> Option<(WorkspaceId, String)> {
    let row = sqlx::query(
        "DELETE FROM terminal_ws_tickets
         WHERE token_hash = $1 AND agent_id = $2 AND expires_at > now()
         RETURNING workspace_id, label",
    )
    .bind(hash_ticket(token))
    .bind(agent_id.as_uuid())
    .fetch_optional(pool)
    .await
    .ok()
    .flatten()?;
    Some((
        WorkspaceId::from(row.get::<uuid::Uuid, _>("workspace_id")),
        row.get::<String, _>("label"),
    ))
}

#[derive(Serialize)]
pub struct TicketResponse {
    pub ticket: String,
    /// The fleet-node UUID this agent's terminal is placed on, or null = gateway-
    /// local. The browser uses the WebRTC transport for a node, plain WS for local.
    pub node: Option<String>,
}

/// `POST /api/terminal/{id}/ticket` — owner-gated; mints a short-lived ticket.
pub async fn ticket(
    State(state): State<AppState>,
    Authed(user): Authed,
    Path(agent_id): Path<AgentId>,
) -> Result<Json<TicketResponse>, ApiError> {
    // Tenant isolation: a foreign agent looks non-existent.
    workspace_agent(&state, &user, agent_id).await?;
    // Quota: gate only when this would spin up a NEW container (a reconnect to an
    // already-running terminal is always allowed).
    let existing = cm_db::repo::agent_containers::get(&state.pool, agent_id, "terminal")
        .await
        .ok()
        .flatten();
    if existing.is_none() {
        crate::quota::enforce_new_container(&state, user.workspace_id).await?;
    }
    // Greet by display name, falling back to the email local-part.
    let label = match cm_db::repo::users::get(&state.pool, user.user_id).await {
        Ok(u) if !u.display_name.trim().is_empty() => u.display_name.trim().to_string(),
        Ok(u) => u
            .email
            .split('@')
            .next()
            .filter(|s| !s.is_empty())
            .unwrap_or("there")
            .to_string(),
        Err(_) => "there".to_string(),
    };
    let ticket = issue_ticket(&state.pool, agent_id, user.workspace_id, &label).await?;
    // Where the terminal is (or would be) placed — null = gateway-local.
    let node = match state.runtime.terminals() {
        Some(tm) => match tm.placement_node_for(agent_id).await {
            n if n == "local" => None,
            n => Some(n),
        },
        None => None,
    };
    cm_db::repo::audit::append(
        &state.pool,
        user.workspace_id,
        Actor::User(user.user_id),
        "terminal.ticket_issued",
        "agent",
        &agent_id.to_string(),
        json!({}),
    )
    .await?;
    Ok(Json(TicketResponse { ticket, node }))
}

/// GET /api/terminal/{id}/tabs — the user's saved tab layout (or null).
pub async fn get_tabs(
    State(state): State<AppState>,
    Authed(user): Authed,
    Path(agent_id): Path<AgentId>,
) -> Result<Json<serde_json::Value>, ApiError> {
    workspace_agent(&state, &user, agent_id).await?;
    let layout = cm_db::repo::terminal_tabs::get(&state.pool, user.user_id, agent_id).await?;
    Ok(Json(layout.unwrap_or(serde_json::Value::Null)))
}

#[derive(Deserialize)]
pub struct SaveTabsRequest {
    layout: serde_json::Value,
}

/// PUT /api/terminal/{id}/tabs — save (replace) the user's tab layout, so the
/// named tabs + their tmux sessions survive logout / a new device.
pub async fn save_tabs(
    State(state): State<AppState>,
    Authed(user): Authed,
    Path(agent_id): Path<AgentId>,
    Json(body): Json<SaveTabsRequest>,
) -> Result<StatusCode, ApiError> {
    workspace_agent(&state, &user, agent_id).await?;
    cm_db::repo::terminal_tabs::upsert(
        &state.pool,
        user.workspace_id,
        user.user_id,
        agent_id,
        &body.layout,
    )
    .await?;
    cm_db::repo::audit::append(
        &state.pool,
        user.workspace_id,
        Actor::User(user.user_id),
        "terminal.tabs_saved",
        "agent",
        &agent_id.to_string(),
        json!({}),
    )
    .await?;
    Ok(StatusCode::NO_CONTENT)
}

#[derive(Deserialize)]
pub struct WsQuery {
    ticket: String,
    /// tmux session name → one per terminal tab (same container, shared drives).
    #[serde(default)]
    session: Option<String>,
}

/// Keep tmux-session names safe ([A-Za-z0-9_-], <=24) — they all share the
/// agent's one container, so different names = independent tabs.
fn sanitize_session(raw: Option<&str>) -> String {
    let s: String = raw
        .unwrap_or("main")
        .chars()
        .filter(|c| c.is_ascii_alphanumeric() || *c == '-' || *c == '_')
        .take(24)
        .collect();
    if s.is_empty() {
        "main".to_string()
    } else {
        s
    }
}

#[derive(Deserialize)]
struct ControlMsg {
    #[serde(rename = "type")]
    kind: String,
    #[serde(default)]
    cols: u16,
    #[serde(default)]
    rows: u16,
}

/// `GET /api/terminal/{id}/ws?ticket=…` — upgrades to a WebSocket bridged to an
/// interactive PTY. Ticket-authed (no bearer header on a browser WS handshake).
pub async fn ws(
    State(state): State<AppState>,
    Path(agent_id): Path<AgentId>,
    Query(q): Query<WsQuery>,
    upgrade: WebSocketUpgrade,
) -> Response {
    let (workspace_id, label) = match redeem_ticket(&state.pool, &q.ticket, agent_id).await {
        Some(v) => v,
        None => return (StatusCode::UNAUTHORIZED, "invalid or expired ticket").into_response(),
    };
    let tm = match state.runtime.terminals() {
        Some(tm) => tm,
        None => {
            return (
                StatusCode::SERVICE_UNAVAILABLE,
                "terminal engine unavailable",
            )
                .into_response()
        }
    };
    // MOTD greeting injected into the PTY session (read by the container .zshrc).
    let env = vec![format!("CLAWMATES_USER={label}")];
    let session = sanitize_session(q.session.as_deref());
    // Node-placed agent → bridge to the node's daemon (WebRTC direct + WS relay),
    // execing into the agent's terminal container there. Local stays unchanged.
    if tm.placement_node_for(agent_id).await != "local" {
        if let Ok((node_id, container)) = tm.placement_for(workspace_id, agent_id).await {
            if let Ok(uuid) = node_id.parse::<uuid::Uuid>() {
                let hub = state.node_hub.clone();
                let nid = NodeId::from(uuid);
                return upgrade
                    .on_upgrade(move |socket| bridge_node(socket, hub, nid, container, session));
            }
        }
        // Couldn't place remotely (node dropped, etc.) — fall through to local.
    }
    upgrade.on_upgrade(move |socket| bridge(socket, tm, workspace_id, agent_id, env, session))
}

/// Pumps bytes both ways between the WebSocket and the container PTY. Binary
/// frames are stdin; text frames are control JSON ({"type":"resize",cols,rows}).
async fn bridge(
    socket: WebSocket,
    tm: Arc<cm_runtime::TerminalManager>,
    workspace_id: WorkspaceId,
    agent_id: AgentId,
    env: Vec<String>,
    tmux_session: String,
) {
    let (mut ws_tx, mut ws_rx) = socket.split();
    let mut session = match tm
        .attach(workspace_id, agent_id, 80, 24, &env, &tmux_session)
        .await
    {
        Ok(s) => s,
        Err(e) => {
            let _ = ws_tx
                .send(Message::Text(
                    format!("\r\n\x1b[31mterminal error: {e}\x1b[0m\r\n").into(),
                ))
                .await;
            tm.detach(agent_id).await;
            return;
        }
    };
    let exec_id = session.exec_id.clone();

    // PTY → WS: stream the shell's combined TTY output as binary frames.
    let pty_to_ws = async {
        while let Some(chunk) = session.output.next().await {
            match chunk {
                Ok(bytes) => {
                    if ws_tx.send(Message::Binary(bytes.into())).await.is_err() {
                        break;
                    }
                }
                Err(_) => break,
            }
        }
    };

    // WS → PTY: binary = keystrokes; text = control (resize).
    let ws_to_pty = async {
        while let Some(msg) = ws_rx.next().await {
            let msg = match msg {
                Ok(m) => m,
                Err(_) => break,
            };
            match msg {
                Message::Binary(b) => {
                    if session.input.write_all(&b).await.is_err() {
                        break;
                    }
                    let _ = session.input.flush().await;
                }
                Message::Text(t) => {
                    if let Ok(ctrl) = serde_json::from_str::<ControlMsg>(t.as_str()) {
                        if ctrl.kind == "resize" && ctrl.cols > 0 && ctrl.rows > 0 {
                            let _ = tm.resize(&exec_id, ctrl.cols, ctrl.rows).await;
                            continue;
                        }
                    }
                    if session.input.write_all(t.as_bytes()).await.is_err() {
                        break;
                    }
                    let _ = session.input.flush().await;
                }
                Message::Close(_) => break,
                _ => {}
            }
        }
    };

    tokio::select! {
        _ = pty_to_ws => {}
        _ = ws_to_pty => {}
    }
    tm.detach(agent_id).await;
}

/// Browser→server control frames on the node-placed terminal WS: resize, the
/// `fallback` to open the WS-relay PTY, or WebRTC signaling (offer/ICE/close).
#[derive(Deserialize)]
struct NodeCtrl {
    #[serde(rename = "type")]
    kind: String,
    cols: Option<u16>,
    rows: Option<u16>,
    sdp: Option<String>,
    candidate: Option<String>,
    sdp_mid: Option<String>,
    sdp_mline_index: Option<u16>,
}

/// Bridge the browser terminal to a NODE-placed agent terminal container: relays
/// PTY bytes + WebRTC signaling to/from the node's daemon, which `docker exec`s a
/// tmux into `container`. Mirrors the Infra node terminal bridge, but targets the
/// agent's container (so the shell shares the agent's node-local ~/drives) and
/// carries the tmux `session` (one per tab). The same WS is the WebRTC signaling
/// channel and the WS-relay fallback path.
async fn bridge_node(
    socket: WebSocket,
    hub: Arc<NodeHub>,
    node_id: NodeId,
    container: String,
    session: String,
) {
    let Some((sid, mut pty_rx, mut sig_rx)) = hub.open_session(node_id).await else {
        return;
    };
    let (mut ws_tx, mut ws_rx) = socket.split();
    let to_browser = async {
        loop {
            tokio::select! {
                bytes = pty_rx.recv() => match bytes {
                    Some(b) => { if ws_tx.send(Message::Binary(b.into())).await.is_err() { break; } }
                    None => break,
                },
                text = sig_rx.recv() => match text {
                    Some(t) => { if ws_tx.send(Message::Text(t.into())).await.is_err() { break; } }
                    None => break,
                },
            }
        }
    };
    let to_node = async {
        while let Some(Ok(msg)) = ws_rx.next().await {
            match msg {
                // Binary = keystrokes over the WS fallback (the DataChannel carries
                // its own input when direct).
                Message::Binary(b) => hub.terminal_input(node_id, sid, b.as_ref()).await,
                Message::Text(t) => {
                    let Ok(c) = serde_json::from_str::<NodeCtrl>(t.as_str()) else {
                        continue;
                    };
                    let cols = c.cols.unwrap_or(80);
                    let rows = c.rows.unwrap_or(24);
                    match c.kind.as_str() {
                        "resize" => hub.terminal_resize(node_id, sid, cols, rows).await,
                        "fallback" => {
                            hub.open_pty(node_id, sid, cols, rows, Some(&container), Some(&session), None)
                                .await
                        }
                        "webrtc_offer" => {
                            hub.webrtc_offer(
                                node_id,
                                sid,
                                c.sdp.as_deref().unwrap_or(""),
                                Some(&container),
                                Some(&session),
                            )
                            .await
                        }
                        "webrtc_ice" => {
                            hub.webrtc_ice(
                                node_id,
                                sid,
                                c.candidate.as_deref().unwrap_or(""),
                                c.sdp_mid.as_deref(),
                                c.sdp_mline_index,
                            )
                            .await
                        }
                        "webrtc_close" => hub.webrtc_close(node_id, sid).await,
                        _ => {}
                    }
                }
                Message::Close(_) => break,
                _ => {}
            }
        }
    };
    tokio::select! {
        _ = to_browser => {},
        _ = to_node => {},
    }
    hub.terminal_close(node_id, sid).await;
}