solstice-ci/crates/runner-integration/src/streamer.rs

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

use tracing::{debug, warn};
use uuid::Uuid;

use crate::connect::ConnectClient;
use crate::proto::runner::v1::{LogRow, UpdateLogRequest};
use crate::state::{RunnerState, StepInfo};

const POLL_INTERVAL: Duration = Duration::from_secs(3);

/// Categories that belong to the "Set up job" phase — must match reporter.rs.
const SETUP_CATEGORIES: &[&str] = &["boot", "default", "env", "env_setup", "tool_check"];

/// Log category summary from logs-service.
#[derive(serde::Deserialize)]
struct LogCategorySummary {
    category: String,
    count: i64,
}

/// Streams logs from logs-service to Forgejo while a job is in-flight.
///
/// Categories are ordered: setup categories first, then step categories in
/// KDL workflow order (matching the YAML step order), then any remaining.
pub async fn stream_logs(
    client: Arc<ConnectClient>,
    state: Arc<RunnerState>,
    request_id: Uuid,
    task_id: i64,
    logs_base: String,
    steps: Vec<StepInfo>,
    mut stop: tokio::sync::watch::Receiver<bool>,
) -> i64 {
    let http = reqwest::Client::new();
    let mut log_index: i64 = 0;

    loop {
        if *stop.borrow() {
            log_index = poll_and_send(
                &client, &state, &http, &logs_base, request_id, task_id, log_index, &steps,
            )
            .await;
            break;
        }

        log_index = poll_and_send(
            &client, &state, &http, &logs_base, request_id, task_id, log_index, &steps,
        )
        .await;

        tokio::select! {
            _ = tokio::time::sleep(POLL_INTERVAL) => {}
            _ = stop.changed() => {
                log_index = poll_and_send(
                    &client, &state, &http, &logs_base,
                    request_id, task_id, log_index, &steps,
                ).await;
                break;
            }
        }
    }

    debug!(task_id, log_index, "log streamer stopped");
    log_index
}

/// Sort categories: setup first, then step categories in KDL order, then any remaining.
fn sort_categories(categories: &mut [LogCategorySummary], steps: &[StepInfo]) {
    categories.sort_by_key(|c| {
        if SETUP_CATEGORIES.contains(&c.category.as_str()) {
            // Setup categories come first, sub-sorted alphabetically
            (0, 0, c.category.clone())
        } else if let Some(pos) = steps.iter().position(|s| s.log_category == c.category) {
            // Step categories in KDL workflow order
            (1, pos, c.category.clone())
        } else {
            // Any remaining categories at the end
            (2, 0, c.category.clone())
        }
    });
}

async fn poll_and_send(
    client: &ConnectClient,
    state: &RunnerState,
    http: &reqwest::Client,
    logs_base: &str,
    request_id: Uuid,
    task_id: i64,
    current_index: i64,
    steps: &[StepInfo],
) -> i64 {
    let categories_url = format!(
        "{}/jobs/{}/logs",
        logs_base.trim_end_matches('/'),
        request_id
    );

    let mut categories = match http.get(&categories_url).send().await {
        Ok(resp) if resp.status().is_success() => resp
            .json::<Vec<LogCategorySummary>>()
            .await
            .unwrap_or_default(),
        _ => return current_index,
    };

    if categories.is_empty() {
        return current_index;
    }

    sort_categories(&mut categories, steps);

    // Build the full ordered log by fetching each category
    let mut all_lines: Vec<String> = Vec::new();
    for cat in &categories {
        let url = format!(
            "{}/jobs/{}/logs/{}",
            logs_base.trim_end_matches('/'),
            request_id,
            cat.category
        );

        let text = match http.get(&url).send().await {
            Ok(resp) if resp.status().is_success() => match resp.text().await {
                Ok(t) => t,
                Err(_) => continue,
            },
            _ => continue,
        };

        for line in text.lines() {
            all_lines.push(line.to_string());
        }
    }

    let total = all_lines.len() as i64;
    if total <= current_index {
        return current_index;
    }

    let new_lines = &all_lines[current_index as usize..];

    let now = prost_types::Timestamp {
        seconds: time::OffsetDateTime::now_utc().unix_timestamp(),
        nanos: 0,
    };

    let rows: Vec<LogRow> = new_lines
        .iter()
        .map(|line| LogRow {
            time: Some(now.clone()),
            content: line.clone(),
        })
        .collect();

    let count = rows.len();

    let req = UpdateLogRequest {
        task_id,
        index: current_index,
        rows,
        no_more: false,
    };

    match client
        .update_log(&req, &state.identity.uuid, &state.identity.token)
        .await
    {
        Ok(resp) => {
            debug!(
                task_id,
                new_lines = count,
                ack_index = resp.ack_index,
                "streamed logs"
            );
            resp.ack_index
        }
        Err(e) => {
            warn!(error = %e, task_id, "failed to stream logs");
            current_index
        }
    }
}