artifact_persistence.rs

//! Bench artifact persistence.
//!
//! Boundary: the bench command owns run lifecycle and result presentation; this
//! module owns persisting bench scenario/run artifacts into the observation
//! store — resolving artifact file pointers (including shared-state and
//! preserved-invocation paths), recording file/directory/URL artifacts,
//! rewriting artifact links back onto the results, and recording budget
//! findings. It operates on the core `ActiveObservation`, not command types.

use std::collections::BTreeMap;
use std::fs;
use std::path::{Component, Path, PathBuf};

use crate::bench::{
    BenchArtifact, BenchDiagnostic, BenchDiagnosticSource, BenchResults, BenchRunWorkflowResult,
};
use homeboy_core::artifacts as artifact_links;
use homeboy_core::engine::run_dir::{self, RunDir};
use homeboy_core::observation::{finding_records_from_budget, ActiveObservation};
use homeboy_lifecycle_contract::ArtifactRecord;

/// Persist every bench artifact referenced by `workflow`'s results into the
/// observation store, rewrite the results file with the recorded links, record
/// the standard bench run-dir artifacts, and record budget findings.
pub fn record_bench_observation_artifacts(
    observation: &ActiveObservation,
    workflow: &mut BenchRunWorkflowResult,
    run_dir: &RunDir,
) -> bool {
    if let Some(results) = workflow.results.as_mut() {
        workflow
            .diagnostics
            .extend(persist_bench_result_artifact_paths(
                observation,
                results,
                run_dir,
            ));
        rewrite_bench_results_file(results, run_dir);
    }

    record_if_exists(
        observation,
        "bench_results",
        run_dir.step_file(run_dir::files::BENCH_RESULTS),
    );
    record_if_exists(
        observation,
        "resource_summary",
        run_dir.step_file(run_dir::files::RESOURCE_SUMMARY),
    );
    record_memory_timeline_artifacts(observation, run_dir);
    record_failure_artifacts(observation, workflow, run_dir);

    let Some(results) = workflow.results.as_ref() else {
        return true;
    };
    observation.record_findings(&finding_records_from_budget(
        observation.run_id(),
        &results.budget_findings,
    ));
    all_bench_artifacts_are_promoted(results)
}

/// Persist diagnostics emitted by failed nested workloads before their temporary
/// invocation directories are removed. The artifact roots are generic runner
/// output, not workload-specific paths.
fn record_failure_artifacts(
    observation: &ActiveObservation,
    workflow: &BenchRunWorkflowResult,
    run_dir: &RunDir,
) {
    let Some(failure) = workflow.failure.as_ref() else {
        return;
    };

    let stderr_path = run_dir.step_file("bench-failure-stderr.txt");
    if fs::write(&stderr_path, &failure.stderr_tail).is_ok() {
        let _ = observation.store().record_artifact_with_metadata(
            observation.run_id(),
            "bench_failure_stderr",
            &stderr_path,
            serde_json::json!({
                "failure_diagnostic": true,
                "failure_diagnostic_rank": 0,
                "source": "bench",
            }),
        );
    }

    for root in [
        run_dir.path().join("artifacts"),
        run_dir.path().join("invocations"),
    ] {
        record_failure_artifact_tree(observation, &root);
    }
}

fn record_failure_artifact_tree(observation: &ActiveObservation, path: &Path) {
    let Ok(entries) = fs::read_dir(path) else {
        return;
    };
    for entry in entries.flatten() {
        let path = entry.path();
        if path.is_dir() {
            record_failure_artifact_tree(observation, &path);
            continue;
        }
        if !path.is_file() {
            continue;
        }
        let name = path
            .file_name()
            .and_then(|name| name.to_str())
            .unwrap_or_default();
        let diagnostic = name.contains("diagnostic") || name.contains("failure");
        let _ = observation.store().record_artifact_with_metadata(
            observation.run_id(),
            if diagnostic {
                "bench_failure_diagnostic"
            } else {
                "bench_failure_artifact"
            },
            &path,
            serde_json::json!({
                "failure_diagnostic": diagnostic,
                "failure_diagnostic_rank": diagnostic.then(|| path.components().count()).unwrap_or_default(),
                "source": "bench_nested_workload",
                "original_path": path,
            }),
        );
    }
}

fn all_bench_artifacts_are_promoted(results: &BenchResults) -> bool {
    results.scenarios.iter().all(|scenario| {
        scenario
            .artifacts
            .values()
            .chain(
                scenario
                    .runs
                    .iter()
                    .flatten()
                    .flat_map(|run| run.artifacts.values()),
            )
            .filter(|artifact| artifact.required_durable)
            .all(|artifact| artifact.observation_artifact_id.is_some())
    })
}

/// Record an artifact when the file at `path` exists.
pub fn record_if_exists(observation: &ActiveObservation, kind: &str, path: PathBuf) {
    observation.record_artifact_if_file(kind, &path);
}

/// Record any `bench-memory-timeline` JSON/CSV artifacts found in the run dir.
pub fn record_memory_timeline_artifacts(observation: &ActiveObservation, run_dir: &RunDir) {
    let Ok(entries) = fs::read_dir(run_dir.path()) else {
        return;
    };
    for entry in entries.flatten() {
        let path = entry.path();
        let Some(name) = path.file_name().and_then(|name| name.to_str()) else {
            continue;
        };
        if name.starts_with("bench-memory-timeline")
            && (name.ends_with(".json") || name.ends_with(".csv"))
        {
            observation.record_artifact_if_file("bench_memory_timeline", &path);
        }
    }
}

fn persist_bench_result_artifact_paths(
    observation: &ActiveObservation,
    results: &mut BenchResults,
    run_dir: &RunDir,
) -> Vec<BenchDiagnostic> {
    let mut diagnostics = Vec::new();
    let shared_state = results
        .run_metadata
        .as_ref()
        .and_then(|metadata| metadata.shared_state.as_deref());
    for scenario in &mut results.scenarios {
        for (name, artifact) in &mut scenario.artifacts {
            if let Some(diagnostic) = persist_bench_artifact(
                observation,
                &scenario.id,
                None,
                name,
                artifact,
                run_dir,
                shared_state,
            ) {
                diagnostics.push(diagnostic);
            }
        }
        if let Some(runs) = &mut scenario.runs {
            for (run_index, run) in runs.iter_mut().enumerate() {
                for (name, artifact) in &mut run.artifacts {
                    if let Some(diagnostic) = persist_bench_artifact(
                        observation,
                        &scenario.id,
                        Some(run_index),
                        name,
                        artifact,
                        run_dir,
                        shared_state,
                    ) {
                        diagnostics.push(diagnostic);
                    }
                }
            }
        }
    }
    diagnostics
}

fn persist_bench_artifact(
    observation: &ActiveObservation,
    scenario_id: &str,
    run_index: Option<usize>,
    name: &str,
    artifact: &mut BenchArtifact,
    run_dir: &RunDir,
    shared_state: Option<&str>,
) -> Option<BenchDiagnostic> {
    // Results can be retried or hydrated from a prior run. Only this attempt's
    // persisted record may satisfy the current durability declaration.
    artifact.observation_artifact_id = None;
    let metadata = bench_artifact_metadata(scenario_id, run_index, name, artifact);

    let original_path = artifact.path.clone();
    if let Some(original_path) = original_path.as_deref() {
        if bench_artifact_path_is_blocked(original_path) {
            return Some(bench_artifact_diagnostic(
                scenario_id,
                run_index,
                name,
                "bench_artifact_path_blocked",
                format!("blocked bench artifact path `{original_path}` for `{name}`"),
                serde_json::json!({ "path": original_path }),
            ));
        }
    }

    let Some(path) = resolve_bench_artifact_file_pointer(
        original_path.as_deref(),
        artifact.url.as_deref(),
        run_dir,
        shared_state,
    ) else {
        if let Some(url) = artifact.url.clone() {
            if !artifact.required_durable {
                return match observation.store().record_url_artifact_with_metadata(
                    observation.run_id(),
                    artifact.kind.as_deref().unwrap_or(name),
                    &url,
                    metadata,
                ) {
                    Ok(record) => apply_recorded_bench_artifact_links(
                        scenario_id,
                        run_index,
                        name,
                        artifact,
                        &record,
                    ),
                    Err(error) => Some(bench_artifact_diagnostic(
                        scenario_id,
                        run_index,
                        name,
                        "bench_artifact_url_record_failed",
                        format!("failed to record bench URL artifact `{name}`: {error}"),
                        serde_json::json!({ "url": url }),
                    )),
                };
            }
            return Some(bench_artifact_diagnostic(
                scenario_id,
                run_index,
                name,
                "bench_artifact_durable_source_missing",
                format!(
                    "bench artifact `{name}` has only URL source `{url}`; reviewer-facing artifacts require a local file or directory for durable promotion"
                ),
                serde_json::json!({ "url": url }),
            ));
        }

        let original_path = original_path?;
        return Some(bench_artifact_diagnostic(
            scenario_id,
            run_index,
            name,
            "bench_artifact_path_missing",
            format!("bench artifact path for `{name}` was not found: {original_path}"),
            serde_json::json!({
                "path": original_path,
                "resolved_path": resolve_bench_artifact_path(&original_path, run_dir, shared_state).to_string_lossy().to_string(),
            }),
        ));
    };
    let record = if path.is_file() {
        observation.store().record_artifact_with_metadata(
            observation.run_id(),
            "bench_artifact",
            &path,
            metadata,
        )
    } else if path.is_dir() {
        observation.store().record_directory_artifact_with_metadata(
            observation.run_id(),
            "bench_artifact",
            &path,
            metadata,
        )
    } else {
        unreachable!("resolved bench artifact file pointers must be files or directories");
    };

    match record {
        Ok(record) => {
            // The source URL can be a tunnel or runner-local address. Keep it
            // in the persisted provenance metadata, but never present it as a
            // reviewer-facing reference after durable promotion.
            artifact.url = None;
            artifact.path = Some(record.path.clone());
            apply_recorded_bench_artifact_links(scenario_id, run_index, name, artifact, &record)
        }
        Err(error) => Some(bench_artifact_diagnostic(
            scenario_id,
            run_index,
            name,
            "bench_artifact_record_failed",
            format!("failed to persist bench artifact `{name}`: {error}"),
            serde_json::json!({
                "path": original_path,
                "resolved_path": path.to_string_lossy().to_string(),
            }),
        )),
    }
}

/// Attach the recorded artifact's observation id, public URL, and validated
/// viewer links back onto the bench artifact, returning a diagnostic when the
/// public URL is unreachable.
pub fn apply_recorded_bench_artifact_links(
    scenario_id: &str,
    run_index: Option<usize>,
    name: &str,
    artifact: &mut BenchArtifact,
    record: &ArtifactRecord,
) -> Option<BenchDiagnostic> {
    artifact.observation_artifact_id = Some(record.id.clone());
    let public_url = artifact_links::public_artifact_url(record)?;
    if artifact_links::public_artifact_url_is_reachable_or_legacy(record) {
        artifact.public_url = Some(public_url.clone());
        artifact.viewer_refs.viewer_links =
            artifact_links::cached_validated_viewer_links(record, &public_url);
        artifact.viewer_refs.viewer_url = artifact
            .viewer_refs
            .viewer_links
            .first()
            .map(|link| link.url.clone());
        None
    } else {
        let validation = record
            .metadata_json
            .get("public_url_validation")
            .expect("unreachable new artifact has validation metadata");
        let error = validation
            .get("error")
            .and_then(serde_json::Value::as_str)
            .unwrap_or("public artifact URL was not reachable");
        Some(bench_artifact_diagnostic(
                scenario_id,
                run_index,
                name,
                "bench_public_artifact_url_unreachable",
                format!(
                    "public artifact URL for bench artifact `{name}` is not reachable; viewer links were not published: {error}"
                ),
                serde_json::json!({
                    "url": validation.get("url").cloned().unwrap_or(serde_json::Value::Null),
                    "status_code": validation.get("status_code").cloned().unwrap_or(serde_json::Value::Null),
                    "error": validation.get("error").cloned().unwrap_or(serde_json::Value::Null),
                }),
            ))
    }
}

fn bench_artifact_metadata(
    scenario_id: &str,
    run_index: Option<usize>,
    name: &str,
    artifact: &BenchArtifact,
) -> serde_json::Value {
    serde_json::json!({
        "source": "bench",
        "scenario_id": scenario_id,
        "run_index": run_index,
        "name": name,
        "kind": artifact.kind,
        "label": artifact.label,
        "original_path": artifact.path,
        "url": artifact.url,
        "viewer": artifact.viewer,
    })
}

fn bench_artifact_path_is_blocked(path: &str) -> bool {
    Path::new(path)
        .components()
        .any(|component| matches!(component, Component::ParentDir))
}

fn bench_artifact_diagnostic(
    scenario_id: &str,
    run_index: Option<usize>,
    name: &str,
    class: &str,
    message: String,
    metadata: serde_json::Value,
) -> BenchDiagnostic {
    let mut diagnostic_metadata = BTreeMap::new();
    diagnostic_metadata.insert("artifact_name".to_string(), serde_json::json!(name));
    if let Some(object) = metadata.as_object() {
        for (key, value) in object {
            diagnostic_metadata.insert(key.clone(), value.clone());
        }
    }
    BenchDiagnostic {
        class: class.to_string(),
        message: Some(message),
        severity: None,
        source: Some(match run_index {
            Some(run_index) => BenchDiagnosticSource::ScenarioRun {
                scenario_id: scenario_id.to_string(),
                run_index,
            },
            None => BenchDiagnosticSource::Scenario {
                scenario_id: scenario_id.to_string(),
            },
        }),
        metadata: diagnostic_metadata,
    }
}

fn rewrite_bench_results_file(results: &BenchResults, run_dir: &RunDir) {
    let Ok(json) = serde_json::to_vec_pretty(results) else {
        return;
    };
    let _ = fs::write(run_dir.step_file(run_dir::files::BENCH_RESULTS), json);
}

fn resolve_bench_artifact_path(
    path: &str,
    run_dir: &RunDir,
    shared_state: Option<&str>,
) -> PathBuf {
    let artifact_path = PathBuf::from(path);
    if artifact_path.exists() {
        return artifact_path;
    }
    if artifact_path.is_absolute() {
        if let Some(shared_state_path) = resolve_shared_state_artifact(&artifact_path, shared_state)
        {
            return shared_state_path;
        }
        if let Some(preserved_path) = resolve_preserved_invocation_artifact(&artifact_path, run_dir)
        {
            return preserved_path;
        }
        return artifact_path;
    }
    let run_dir_path = run_dir.path().join(path);
    if run_dir_path.exists() {
        return run_dir_path;
    }
    artifact_path
}

fn resolve_bench_artifact_file_pointer(
    path: Option<&str>,
    url: Option<&str>,
    run_dir: &RunDir,
    shared_state: Option<&str>,
) -> Option<PathBuf> {
    if let Some(path) = path {
        let resolved = resolve_bench_artifact_path(path, run_dir, shared_state);
        if resolved.is_file() || resolved.is_dir() {
            return Some(resolved);
        }
    }

    let filename = url.and_then(artifact_filename_from_url)?;
    [
        run_dir.path().join("artifacts").join(&filename),
        run_dir.path().join(&filename),
    ]
    .into_iter()
    .find(|candidate| candidate.is_file() || candidate.is_dir())
}

fn artifact_filename_from_url(url: &str) -> Option<String> {
    let parsed = reqwest::Url::parse(url).ok()?;
    parsed
        .path_segments()?
        .next_back()
        .filter(|segment| !segment.is_empty() && !bench_artifact_path_is_blocked(segment))
        .map(str::to_string)
}

fn resolve_shared_state_artifact(path: &Path, shared_state: Option<&str>) -> Option<PathBuf> {
    let shared_state = shared_state?;
    let relative = path.strip_prefix("/bench-shared-state").ok()?;
    let candidate = Path::new(shared_state).join(relative);
    if candidate.exists() {
        return Some(candidate);
    }
    None
}

fn resolve_preserved_invocation_artifact(path: &Path, run_dir: &RunDir) -> Option<PathBuf> {
    let mut components = path.components().peekable();
    while let Some(component) = components.next() {
        let name = component.as_os_str().to_string_lossy();
        let Some(short_id) = name.strip_suffix(".a") else {
            continue;
        };

        let mut preserved = run_dir
            .path()
            .join("invocations")
            .join(format!("inv-{short_id}"))
            .join("artifacts");
        for rest in components {
            preserved.push(rest.as_os_str());
        }
        if preserved.exists() {
            return Some(preserved);
        }
        return None;
    }
    None
}