continuum-proxy/src/uplink/mod.rs

//! The cloud connection: a WebSocket to `continuum-backend`'s `/ws/edge/v1`.
//!
//! This is where the proxy earns its name. Until now it connected, said
//! hello, and sent heartbeats into the void while the printer fleet lived in
//! a completely separate part of the program with nothing routing between
//! them. Both directions are wired up now:
//!
//! - **Up:** every `TELEMETRY_INTERVAL`, one snapshot per printer, batched
//!   into a single `Telemetry` message.
//! - **Down:** a `Command` from the cloud is looked up by printer id,
//!   executed, and answered with a `CommandResult`.
//!
//! Still deliberately the simplest version of a reconnecting client: a fixed
//! retry delay rather than a backoff curve, and no local buffering, so
//! telemetry produced while the uplink is down is simply lost (re-adding the
//! SQLite buffer cut in the simplification pass is the fix, and it is a
//! separate piece of work - see PLAN.md).

mod download;
mod protocol;

pub use protocol::{ClientMessage, PrinterCommand, PrinterTelemetry, ServerMessage};

use std::sync::Arc;
use std::time::Duration;

use futures_util::{SinkExt, StreamExt};
use tokio::sync::Mutex;
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::tungstenite::handshake::client::{Request, generate_key};
use tracing::{info, warn};

use crate::config::Config;
use crate::printer::{GenericPrinter, PrinterHandle};

const RETRY_DELAY: Duration = Duration::from_secs(5);
const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(15);
/// Faster than the heartbeat, slower than the printers push. A Bambu on v2
/// emits a full status roughly every second; forwarding at that rate would be
/// a lot of traffic to describe a machine whose state changes on the scale of
/// a layer, so this samples the cached report instead of chasing every push.
const TELEMETRY_INTERVAL: Duration = Duration::from_secs(5);

/// The printer fleet, shared between `main` (which owns it and keeps every
/// connection alive) and this module (which reads status and dispatches
/// commands).
///
/// `Arc` is shared ownership - several places hold the same value and the
/// last one to let go drops it, roughly how every TS object reference
/// already behaves. `Mutex` is what makes mutation across tasks safe: Rust
/// will not let two tasks hold `&mut` to the same data, so a lock is how you
/// take turns. `tokio::sync::Mutex` rather than `std::sync::Mutex`
/// specifically because the code below holds the lock across an `.await`,
/// which the std one is not designed for.
///
/// One lock for the whole `Vec` is coarse: a slow command to one printer
/// stalls telemetry for all of them. Fine at the handful-of-printers scale
/// this runs at now, and the shape to revisit (a lock per printer, or an
/// actor task per printer owning its own handle) before a farm gets large.
pub type SharedFleet = Arc<Mutex<Vec<PrinterHandle>>>;

/// The `Host` header for a websocket URL: authority without the scheme.
///
/// Needed because building the upgrade request by hand (to attach the edge
/// token) means tokio-tungstenite no longer derives the handshake headers
/// from a URL for us. An error rather than a guess if the URL has no
/// authority: a `Host` of the wrong thing fails the handshake with something
/// far less obvious than this does.
fn host_header(url: &str) -> anyhow::Result<String> {
    let authority = url
        .split_once("://")
        .map(|(_, rest)| rest)
        .unwrap_or(url)
        .split(['/', '?', '#'])
        .next()
        .unwrap_or("");

    if authority.is_empty() {
        anyhow::bail!("the uplink URL has no host to connect to: {url}");
    }
    Ok(authority.to_string())
}

/// Runs forever: connect, serve the connection until it drops, wait, retry.
pub async fn run(config: Config, fleet: SharedFleet) {
    loop {
        info!(url = %config.uplink_url, "connecting to cloud uplink");

        if let Err(err) = connect_and_serve(&config, &fleet).await {
            warn!(?err, "uplink connection failed, retrying in {RETRY_DELAY:?}");
        }

        tokio::time::sleep(RETRY_DELAY).await;
    }
}

async fn connect_and_serve(config: &Config, fleet: &SharedFleet) -> anyhow::Result<()> {
    // The edge token goes on the upgrade request, not in the first message.
    // The backend verifies it before the socket exists, so a bad token is an
    // HTTP 401 here rather than a connection that opens and is then closed -
    // and the credential stays out of the URL, where it would land in every
    // access and proxy log along the way. `Authorization` is the header
    // those tools already know to redact.
    let request = Request::builder()
        .uri(&config.uplink_url)
        .header("Authorization", format!("Bearer {}", config.gateway_token))
        // tokio-tungstenite fills these in for a plain URL, but not once we
        // hand it a request we built ourselves - and a handshake missing any
        // of them is rejected by the server rather than ignored.
        .header("Host", host_header(&config.uplink_url)?)
        .header("Connection", "Upgrade")
        .header("Upgrade", "websocket")
        .header("Sec-WebSocket-Version", "13")
        .header("Sec-WebSocket-Key", generate_key())
        .body(())?;

    let (ws_stream, _response) = tokio_tungstenite::connect_async(request).await?;
    let (mut write, mut read) = ws_stream.split();

    let hello = ClientMessage::Hello {
        gateway_id: config.gateway_id.clone(),
        farm_id: config.farm_id.clone(),
        stream_base_url: config.stream_base_url.clone(),
    };
    write.send(Message::Text(serde_json::to_string(&hello)?)).await?;

    let mut heartbeat = tokio::time::interval(HEARTBEAT_INTERVAL);
    let mut telemetry = tokio::time::interval(TELEMETRY_INTERVAL);
    let mut seq: u64 = 0;

    loop {
        tokio::select! {
            _ = heartbeat.tick() => {
                seq += 1;
                let msg = ClientMessage::Heartbeat { seq };
                write.send(Message::Text(serde_json::to_string(&msg)?)).await?;
            }

            _ = telemetry.tick() => {
                // Read first, then ask for the next one. A v1 Bambu answers
                // `refresh_status` over MQTT some time after this returns, so
                // collecting first means every tick reports the previous
                // tick's refresh - one interval of lag, rather than a
                // guaranteed miss from reading before the answer could
                // possibly have arrived.
                let printers = collect_telemetry(fleet).await;
                refresh_fleet(fleet).await;
                // Nothing to say is not worth a frame. Happens whenever no
                // printer has reported yet - every vendor but Bambu is still
                // a stub, so a Prusa-only fleet is silent here by design.
                if !printers.is_empty() {
                    let msg = ClientMessage::Telemetry { printers };
                    write.send(Message::Text(serde_json::to_string(&msg)?)).await?;
                }
            }

            frame = read.next() => {
                match frame {
                    Some(Ok(Message::Text(text))) => {
                        match serde_json::from_str::<ServerMessage>(&text) {
                            Ok(ServerMessage::Command { command_id, printer_id, command }) => {
                                // `start_print` is the one command that cannot
                                // be answered inline. See `spawn_long_command`.
                                let msg = if command.runs_for_minutes() {
                                    spawn_long_command(fleet.clone(), printer_id.clone(), command);
                                    ClientMessage::CommandResult {
                                        command_id, ok: true, error: None,
                                    }
                                } else {
                                    match dispatch(fleet, &printer_id, command).await {
                                        Ok(()) => ClientMessage::CommandResult {
                                            command_id, ok: true, error: None,
                                        },
                                        Err(err) => {
                                            warn!(%printer_id, %err, "command failed");
                                            ClientMessage::CommandResult {
                                                command_id, ok: false, error: Some(err),
                                            }
                                        }
                                    }
                                };
                                write.send(Message::Text(serde_json::to_string(&msg)?)).await?;
                            }
                            // Retrying will not help: the token is wrong, not
                            // the network. Logged at error level and returned
                            // so it is obvious which of the two it is.
                            Ok(ServerMessage::HelloError { message }) => {
                                tracing::error!(%message, "cloud rejected this gateway's handshake - check CONTINUUM_GATEWAY_TOKEN");
                                return Ok(());
                            }
                            Ok(other) => info!(?other, "received from cloud"),
                            // Loudly. This used to be an `if let Ok(..)` with
                            // no else, which is the same silent-drop bug that
                            // hid a real `print.push_status` failure for
                            // weeks - and here it would mean commands
                            // vanishing with the cloud never being told.
                            Err(err) => warn!(%err, payload = %text, "unparseable message from cloud, ignoring"),
                        }
                    }
                    Some(Ok(Message::Close(_))) | None => return Ok(()),
                    Some(Err(err)) => return Err(err.into()),
                    _ => {}
                }
            }
        }
    }
}

/// Asks every printer to report in full, for the next tick to pick up. Failure
/// is logged and otherwise ignored: a printer that cannot be asked right now
/// is a printer whose telemetry goes stale, not a reason to tear down the
/// cloud connection for the whole farm.
async fn refresh_fleet(fleet: &SharedFleet) {
    let fleet = fleet.lock().await;
    for printer in fleet.iter() {
        if let Err(err) = printer.refresh_status().await {
            warn!(printer = %printer.base().id, %err, "failed to request a status refresh");
        }
    }
}

/// One snapshot per printer that has something to report. Translates Bambu's
/// `print.push_status` into the vendor-neutral shape the cloud consumes - see
/// `PrinterTelemetry` for why this isn't just forwarded raw.
async fn collect_telemetry(fleet: &SharedFleet) -> Vec<PrinterTelemetry> {
    let fleet = fleet.lock().await;
    let mut out = Vec::new();

    for printer in fleet.iter() {
        let base = printer.base();
        let Some(report) = printer.print_status().await else {
            continue;
        };

        out.push(PrinterTelemetry {
            printer_id: base.id.clone(),
            name: base.name.clone(),
            state: report.gcode_state.clone(),
            percent: report.mc_percent,
            stage: report.stg_cur,
            layer: report.layer_num,
            total_layers: report.total_layer_num,
            remaining_minutes: report.mc_remaining_time,
            // An idle printer reports `""` rather than omitting the field;
            // empty-string-as-absent is worth normalising here so the cloud
            // doesn't have to know that.
            job_name: report.subtask_name.clone().filter(|s| !s.is_empty()),
            nozzle_temp: report.nozzle_temper,
            nozzle_target: report.nozzle_target_temper,
            bed_temp: report.bed_temper,
            bed_target: report.bed_target_temper,
            print_error: report.print_error,
            active_alerts: report.hms.as_ref().map(|h| h.len()),
            // Derived from the same reported URL `camera::sources` builds
            // the go2rtc config from, so a printer is advertised as
            // streamable here only if it is also actually being served
            // there. Asking `camera` rather than keeping a second copy of
            // the rule means the two cannot drift apart.
            camera_stream: crate::camera::source_for(
                printer,
                report.ipcam.as_ref().and_then(|c| c.rtsp_url.as_deref()),
            )
            .map(|source| source.stream_name),
        });
    }

    out
}

/// Runs a command that takes minutes, off this connection's task, and answers
/// the cloud immediately.
///
/// Everything else here is an MQTT publish that returns in milliseconds, so
/// awaiting it inline is fine. `start_print` is not: it downloads tens of
/// megabytes and then transfers them to a printer over FTPS at a measured
/// 0.2 MiB/s, which is about four minutes for a 50 MB project. Awaited
/// inline, that blocks this whole `select!` loop - no heartbeats and, far
/// worse, **no telemetry** for the duration.
///
/// No telemetry is the one that makes this mandatory rather than tidy. The
/// cloud derives a job's state from the telemetry that follows a start, so a
/// gateway that goes silent for four minutes during every job start is
/// withholding exactly the information the job depends on, while the backend
/// is also timing the command out after ten seconds and concluding the start
/// failed.
///
/// So the `ok` this produces means "accepted, and working on it" - which is
/// all the cloud ever trusted a start reply to mean anyway. A failure after
/// this point is reported by the job never reaching `printing`, which the
/// backend reaps; it is not reported as a `CommandResult`, because answering
/// later would need an outbound channel this loop does not have yet. That is
/// the next thing to build here, and until it exists the gateway's log is
/// where the reason lives.
fn spawn_long_command(fleet: SharedFleet, printer_id: String, command: PrinterCommand) {
    tokio::spawn(async move {
        match dispatch(&fleet, &printer_id, command).await {
            Ok(()) => info!(%printer_id, "long-running command finished"),
            // Loudly, and at error level: nothing else will say why. The
            // cloud sees only a job that never started printing.
            Err(err) => tracing::error!(
                %printer_id, %err,
                "long-running command failed after the cloud was already told it was accepted"
            ),
        }
    });
}

/// Finds the printer by our own id and runs the command. The error is a
/// `String` rather than a typed error because its only destination is a
/// `CommandResult` going back to the cloud as text.
async fn dispatch(fleet: &SharedFleet, printer_id: &str, command: PrinterCommand) -> Result<(), String> {
    // Fetched before the fleet lock is taken: a download is bounded by
    // minutes of network time, and holding the lock across it would stall
    // telemetry for every other printer. The transfer to the printer that
    // follows has to hold it (it needs `&mut` access), but that is no reason
    // to hold it through this part too. A `start_print` for a printer that
    // does not exist is therefore found out after the download rather than
    // before; the lookup below is cheap, so it is done first for that case.
    let downloaded = match &command {
        PrinterCommand::StartPrint { url, .. } => {
            if !fleet.lock().await.iter().any(|p| p.base().id == printer_id) {
                return Err(format!("no printer with id '{printer_id}' in this gateway's fleet"));
            }
            Some(
                download::fetch_job_file(url)
                    .await
                    .map_err(|err| format!("download failed (nothing was sent to the printer): {err}"))?,
            )
        }
        _ => None,
    };

    let mut fleet = fleet.lock().await;
    let printer = fleet
        .iter_mut()
        .find(|p| p.base().id == printer_id)
        .ok_or_else(|| format!("no printer with id '{printer_id}' in this gateway's fleet"))?;

    info!(%printer_id, ?command, "dispatching cloud command");

    match command {
        PrinterCommand::SetFanSpeed { percent } => {
            printer.set_fan_speed(percent).await.map_err(|err| err.to_string())
        }
        PrinterCommand::SendGcodeLine { line } => {
            printer.send_gcode_line(line).await.map_err(|err| err.to_string())
        }
        PrinterCommand::ControlPrint { control } => {
            printer.control_print(control).await.map_err(|err| err.to_string())
        }
        PrinterCommand::StartPrint { job_id, file_name, .. } => {
            let bytes = downloaded.ok_or_else(|| "internal error: start_print reached dispatch with no downloaded file".to_string())?;
            let byte_count = bytes.len();
            info!(%printer_id, %job_id, %file_name, byte_count, "job file downloaded, sending to printer");
            // One call covers deriving the filament mapping, the FTPS
            // transfer, the truncation check and the start, and reports them
            // all as one error - so the stage here is "printer", not finer.
            // The printer's own message says which part (mapping, upload of
            // the file, or start) it was.
            printer
                .upload_project_file(file_name, bytes)
                .await
                .map_err(|err| format!("printer upload/start failed (download of {byte_count} bytes succeeded): {err}"))
        }
    }
}