Initial boilerplate scaffold for continuum-ai-worker

src/config.rs

@@ -0,0 +1,42 @@
+use std::env;
+
+#[derive(Debug, Clone)]
+pub struct Config {
+    pub redis_url: String,
+    pub frame_stream: String,
+    pub consumer_group: String,
+    pub consumer_name: String,
+    pub stream_block_ms: usize,
+    pub stream_batch_size: usize,
+    pub database_url: String,
+    pub model_path: String,
+    pub model_input_size: u32,
+    pub confidence_threshold: f32,
+    pub intra_op_threads: usize,
+}
+
+impl Config {
+    pub fn from_env() -> anyhow::Result<Self> {
+        Ok(Self {
+            redis_url: env_or("CONTINUUM_REDIS_URL", "redis://localhost:6379"),
+            frame_stream: env_or("CONTINUUM_FRAME_STREAM", "stream:camera_frames"),
+            consumer_group: env_or("CONTINUUM_CONSUMER_GROUP", "ai-worker"),
+            consumer_name: env_or("CONTINUUM_CONSUMER_NAME", "ai-worker-1"),
+            stream_block_ms: env_or("CONTINUUM_STREAM_BLOCK_MS", "5000").parse()?,
+            stream_batch_size: env_or("CONTINUUM_STREAM_BATCH_SIZE", "16").parse()?,
+            database_url: require("DATABASE_URL")?,
+            model_path: env_or("CONTINUUM_MODEL_PATH", "./models/failure-detector.onnx"),
+            model_input_size: env_or("CONTINUUM_MODEL_INPUT_SIZE", "640").parse()?,
+            confidence_threshold: env_or("CONTINUUM_CONFIDENCE_THRESHOLD", "0.75").parse()?,
+            intra_op_threads: env_or("CONTINUUM_INTRA_OP_THREADS", "4").parse()?,
+        })
+    }
+}
+
+fn require(key: &str) -> anyhow::Result<String> {
+    env::var(key).map_err(|_| anyhow::anyhow!("missing required env var {key}"))
+}
+
+fn env_or(key: &str, default: &str) -> String {
+    env::var(key).unwrap_or_else(|_| default.to_string())
+}

src/inference/detector.rs

@@ -0,0 +1,80 @@
+use std::sync::Arc;
+
+use image::DynamicImage;
+use ort::session::builder::GraphOptimizationLevel;
+use ort::session::Session;
+use ort::value::Value;
+use serde::Serialize;
+use tokio::sync::Mutex;
+
+use super::preprocess::to_input_tensor;
+
+#[derive(Debug, Clone, Serialize)]
+pub struct Detection {
+    pub class_id: u32,
+    pub label: String,
+    pub confidence: f32,
+    /// Normalized [0, 1] bounding box in (x_min, y_min, x_max, y_max).
+    pub bbox: [f32; 4],
+}
+
+const CLASS_LABELS: &[&str] = &["spaghetti", "warping", "layer_shift", "bed_adhesion_failure"];
+
+/// Thread-safe wrapper around an ONNX Runtime session. `ort::Session` is not
+/// `Sync` on its own for concurrent `run()` calls, so callers share this
+/// behind a `Mutex` and an `Arc` when fanning out across worker tasks.
+pub struct FailureDetector {
+    session: Mutex<Session>,
+    input_size: u32,
+    confidence_threshold: f32,
+}
+
+impl FailureDetector {
+    pub fn load(model_path: &str, input_size: u32, confidence_threshold: f32, intra_threads: usize) -> anyhow::Result<Arc<Self>> {
+        let session = Session::builder()?
+            .with_optimization_level(GraphOptimizationLevel::Level3)?
+            .with_intra_threads(intra_threads)?
+            .commit_from_file(model_path)?;
+
+        Ok(Arc::new(Self {
+            session: Mutex::new(session),
+            input_size,
+            confidence_threshold,
+        }))
+    }
+
+    /// Runs inference on a single decoded frame and returns detections whose
+    /// confidence clears `confidence_threshold`.
+    pub async fn detect(&self, frame: &DynamicImage) -> anyhow::Result<Vec<Detection>> {
+        let tensor = to_input_tensor(frame, self.input_size);
+        let input = Value::from_array(tensor)?;
+
+        let mut session = self.session.lock().await;
+        let outputs = session.run(ort::inputs!["images" => input]?)?;
+
+        // Exported with NMS baked into the graph: output0 is [1, N, 6]
+        // rows of (x1, y1, x2, y2, confidence, class_id) in input-pixel space.
+        let (shape, data) = outputs["output0"].try_extract_tensor::<f32>()?;
+        let num_boxes = shape[1] as usize;
+        let size = self.input_size as f32;
+
+        let mut detections = Vec::new();
+        for i in 0..num_boxes {
+            let row = &data[i * 6..i * 6 + 6];
+            let confidence = row[4];
+            if confidence < self.confidence_threshold {
+                continue;
+            }
+
+            let class_id = row[5] as u32;
+            detections.push(Detection {
+                class_id,
+                label: CLASS_LABELS.get(class_id as usize).copied().unwrap_or("unknown").to_string(),
+                confidence,
+                bbox: [row[0] / size, row[1] / size, row[2] / size, row[3] / size],
+            });
+        }
+
+        Ok(detections)
+    }
+}

src/inference/mod.rs

@@ -0,0 +1,4 @@
+mod detector;
+mod preprocess;
+
+pub use detector::{Detection, FailureDetector};

src/inference/preprocess.rs

@@ -0,0 +1,24 @@
+use image::{DynamicImage, GenericImageView};
+use ndarray::Array4;
+
+/// Resizes an RGB frame to a square `size x size` input tensor in CHW layout,
+/// normalized to [0, 1] — the standard preprocessing for YOLO-family ONNX
+/// export graphs.
+pub fn to_input_tensor(image: &DynamicImage, size: u32) -> Array4<f32> {
+    let resized = image.resize_exact(size, size, image::imageops::FilterType::Triangle);
+    let rgb = resized.to_rgb8();
+
+    let mut tensor = Array4::<f32>::zeros((1, 3, size as usize, size as usize));
+    for (x, y, pixel) in rgb.enumerate_pixels() {
+        let [r, g, b] = pixel.0;
+        tensor[[0, 0, y as usize, x as usize]] = r as f32 / 255.0;
+        tensor[[0, 1, y as usize, x as usize]] = g as f32 / 255.0;
+        tensor[[0, 2, y as usize, x as usize]] = b as f32 / 255.0;
+    }
+
+    tensor
+}
+
+pub fn frame_dimensions(image: &DynamicImage) -> (u32, u32) {
+    image.dimensions()
+}

src/main.rs

@@ -0,0 +1,149 @@
+mod config;
+mod inference;
+mod reporter;
+
+use std::collections::HashMap;
+
+use redis::streams::{StreamKey, StreamReadOptions, StreamReadReply};
+use redis::AsyncCommands;
+use tracing::{error, info, warn};
+use tracing_subscriber::EnvFilter;
+
+use crate::config::Config;
+use crate::inference::FailureDetector;
+use crate::reporter::Reporter;
+
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
+    tracing_subscriber::fmt()
+        .with_env_filter(EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")))
+        .json()
+        .init();
+
+    dotenvy_load();
+    let config = Config::from_env()?;
+
+    info!(model = %config.model_path, "loading ONNX Runtime failure-detection model");
+    let detector = FailureDetector::load(
+        &config.model_path,
+        config.model_input_size,
+        config.confidence_threshold,
+        config.intra_op_threads,
+    )?;
+
+    let mut reporter = Reporter::connect(&config.redis_url, &config.database_url).await?;
+
+    let client = redis::Client::open(config.redis_url.clone())?;
+    let mut conn = client.get_multiplexed_tokio_connection().await?;
+
+    ensure_consumer_group(&mut conn, &config).await?;
+
+    info!(
+        stream = %config.frame_stream,
+        group = %config.consumer_group,
+        consumer = %config.consumer_name,
+        "starting frame consumer loop"
+    );
+
+    let read_opts = StreamReadOptions::default()
+        .group(&config.consumer_group, &config.consumer_name)
+        .block(config.stream_block_ms)
+        .count(config.stream_batch_size);
+
+    loop {
+        let reply: StreamReadReply = match conn
+            .xread_options(&[&config.frame_stream], &[">"], &read_opts)
+            .await
+        {
+            Ok(reply) => reply,
+            Err(err) => {
+                error!(?err, "XREADGROUP failed, backing off");
+                tokio::time::sleep(std::time::Duration::from_secs(2)).await;
+                continue;
+            }
+        };
+
+        for StreamKey { key, ids } in reply.keys {
+            for entry in ids {
+                let fields = entry_fields(&entry.map);
+                match process_entry(&detector, &mut reporter, &fields).await {
+                    Ok(()) => {}
+                    Err(err) => warn!(?err, id = %entry.id, "failed to process frame entry"),
+                }
+
+                let _: redis::RedisResult<i64> = conn
+                    .xack(&key, &config.consumer_group, &[entry.id.as_str()])
+                    .await;
+            }
+        }
+    }
+}
+
+async fn ensure_consumer_group(
+    conn: &mut redis::aio::MultiplexedConnection,
+    config: &Config,
+) -> anyhow::Result<()> {
+    let result: redis::RedisResult<String> = conn
+        .xgroup_create_mkstream(&config.frame_stream, &config.consumer_group, "$")
+        .await;
+
+    if let Err(err) = result {
+        // BUSYGROUP means the group already exists — fine, everything else is real.
+        if !err.to_string().contains("BUSYGROUP") {
+            return Err(err.into());
+        }
+    }
+
+    Ok(())
+}
+
+fn entry_fields(map: &HashMap<String, redis::Value>) -> HashMap<String, String> {
+    map.iter()
+        .filter_map(|(k, v)| match v {
+            redis::Value::BulkString(bytes) => {
+                Some((k.clone(), String::from_utf8_lossy(bytes).to_string()))
+            }
+            _ => None,
+        })
+        .collect()
+}
+
+async fn process_entry(
+    detector: &FailureDetector,
+    reporter: &mut Reporter,
+    fields: &HashMap<String, String>,
+) -> anyhow::Result<()> {
+    let printer_id = fields
+        .get("printer_id")
+        .ok_or_else(|| anyhow::anyhow!("frame entry missing printer_id field"))?;
+    let job_id = fields.get("job_id").map(String::as_str);
+    let frame_b64 = fields
+        .get("frame")
+        .ok_or_else(|| anyhow::anyhow!("frame entry missing frame field"))?;
+
+    let bytes = base64::Engine::decode(&base64::engine::general_purpose::STANDARD, frame_b64)?;
+    let image = image::load_from_memory(&bytes)?;
+
+    let detections = detector.detect(&image).await?;
+    if !detections.is_empty() {
+        reporter.report_failure(printer_id, job_id, &detections).await?;
+    }
+
+    Ok(())
+}
+
+fn dotenvy_load() {
+    if let Ok(contents) = std::fs::read_to_string(".env") {
+        for line in contents.lines() {
+            let line = line.trim();
+            if line.is_empty() || line.starts_with('#') {
+                continue;
+            }
+            if let Some((key, value)) = line.split_once('=') {
+                if std::env::var(key).is_err() {
+                    std::env::set_var(key, value);
+                }
+            }
+        }
+    }
+}

src/reporter/mod.rs

@@ -0,0 +1,43 @@
+mod postgres;
+mod redis_pub;
+
+use crate::inference::Detection;
+use sqlx::PgPool;
+use tracing::info;
+
+pub struct Reporter {
+    redis: redis::aio::MultiplexedConnection,
+    pg: PgPool,
+}
+
+impl Reporter {
+    pub async fn connect(redis_url: &str, database_url: &str) -> anyhow::Result<Self> {
+        let redis_client = redis::Client::open(redis_url)?;
+        let redis = redis_client.get_multiplexed_tokio_connection().await?;
+        let pg = PgPool::connect(database_url).await?;
+
+        Ok(Self { redis, pg })
+    }
+
+    /// Fast-paths a failure alert: publishes to Redis immediately for any
+    /// live UI subscribers, then persists to Postgres for the audit trail
+    /// and downstream farm-manager notifications.
+    pub async fn report_failure(
+        &mut self,
+        printer_id: &str,
+        job_id: Option<&str>,
+        detections: &[Detection],
+    ) -> anyhow::Result<()> {
+        let best = detections
+            .iter()
+            .max_by(|a, b| a.confidence.total_cmp(&b.confidence))
+            .expect("report_failure called with no detections");
+
+        info!(printer_id, label = %best.label, confidence = best.confidence, "reporting print failure");
+
+        redis_pub::publish_alert(&mut self.redis, printer_id, best).await?;
+        postgres::insert_alert(&self.pg, printer_id, job_id, detections).await?;
+
+        Ok(())
+    }
+}

src/reporter/postgres.rs

@@ -0,0 +1,31 @@
+use sqlx::PgPool;
+use uuid::Uuid;
+
+use crate::inference::Detection;
+
+/// Persists the full detection set for a failure event to the shared
+/// Postgres database (the same instance `continuum-backend` writes to via
+/// Drizzle), for audit history and the farm-manager review queue.
+pub async fn insert_alert(
+    pool: &PgPool,
+    printer_id: &str,
+    job_id: Option<&str>,
+    detections: &[Detection],
+) -> anyhow::Result<()> {
+    let detections_json = serde_json::to_value(detections)?;
+
+    sqlx::query(
+        r#"
+        INSERT INTO print_failure_alerts (id, printer_id, print_job_id, detections, created_at)
+        VALUES ($1, $2, $3, $4, now())
+        "#,
+    )
+    .bind(Uuid::new_v4())
+    .bind(printer_id)
+    .bind(job_id)
+    .bind(detections_json)
+    .execute(pool)
+    .await?;
+
+    Ok(())
+}

src/reporter/redis_pub.rs

@@ -0,0 +1,26 @@
+use redis::AsyncCommands;
+use serde_json::json;
+
+use crate::inference::Detection;
+
+const ALERT_CHANNEL: &str = "channel:print_alerts";
+
+/// Publishes a low-latency alert for anything subscribed to
+/// `channel:print_alerts` — primarily `continuum-backend`'s WebSocket
+/// fan-out to connected operator consoles.
+pub async fn publish_alert(
+    conn: &mut redis::aio::MultiplexedConnection,
+    printer_id: &str,
+    detection: &Detection,
+) -> anyhow::Result<()> {
+    let payload = json!({
+        "type": "printer.failure_detected",
+        "printerId": printer_id,
+        "label": detection.label,
+        "confidence": detection.confidence,
+        "bbox": detection.bbox,
+    });
+
+    conn.publish::<_, _, ()>(ALERT_CHANNEL, payload.to_string()).await?;
+    Ok(())
+}