mbx-cache-core 0.20.6

use super::{MAX_TASK_ACTION_PREDICTIONS, MAX_TASK_MANIFEST_BYTES, TASK_ACTION_MANIFEST_VERSION};
use crate::{ActionPrediction, CacheDigest, TaskActionManifest, canonical_json};
use eyre::{Context, Result, bail};
use std::collections::{BTreeMap, BTreeSet};
use std::fs;
use std::path::{Path, PathBuf};

/// Whether `task` is a well-formed task action identity.
///
/// Identities name files and directories in the store, so anything that reads
/// the store back has to be able to tell an identity from whatever else a user
/// left lying there.
pub fn is_task_identity(task: &str) -> bool {
    task.len() == 64
        && task
            .bytes()
            .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
}

pub(super) fn validate_task_identity(task: &str) -> Result<()> {
    if !is_task_identity(task) {
        bail!("invalid task action identity");
    }
    Ok(())
}

/// Where a store keeps its task prediction manifests.
pub(super) fn task_manifest_dir(store: &Path) -> PathBuf {
    store.join("task-manifests").join("v1")
}

/// The action digests a task's prediction manifest recorded.
///
/// Read straight off disk rather than through an agent, because a collector
/// needs the action set of tasks no session is running. A manifest that is
/// missing or no longer parseable yields no actions rather than an error: this
/// is a prediction index, so the worst a thin answer costs is a cold prefetch,
/// or an object collected earlier than it deserved.
pub fn task_manifest_actions(store: &Path, task: &str) -> Result<Vec<CacheDigest>> {
    validate_task_identity(task)?;
    let path = task_manifest_dir(store).join(format!("{task}.json"));
    let bytes = match fs::read(&path) {
        Ok(bytes) => bytes,
        Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
        Err(error) => {
            return Err(error).wrap_err_with(|| format!("failed to read {}", path.display()));
        }
    };
    let Ok(manifest) = serde_json::from_slice::<TaskActionManifest>(&bytes) else {
        return Ok(Vec::new());
    };
    if validate_task_manifest(&manifest, task).is_err() {
        return Ok(Vec::new());
    }
    Ok(manifest
        .predictions
        .into_iter()
        .map(|prediction| prediction.action)
        .collect())
}

/// Room for the manifest's version, task identity, and braces around its predictions.
const TASK_MANIFEST_OVERHEAD_BYTES: usize = 256;

pub(super) fn validate_task_manifest(manifest: &TaskActionManifest, task: &str) -> Result<()> {
    if manifest.task == task && manifest.validate() {
        Ok(())
    } else {
        bail!("invalid task action manifest")
    }
}

/// Refresh preferred values at the back of the durable LRU, then retire the
/// oldest unprotected history needed to keep the manifest bounded.
pub(super) fn update_task_predictions(
    predictions: Vec<ActionPrediction>,
    updates: &BTreeMap<CacheDigest, ActionPrediction>,
    protected: &BTreeSet<CacheDigest>,
) -> Result<Vec<ActionPrediction>> {
    update_task_predictions_in_order(predictions, updates.values().cloned().collect(), protected)
}

/// Append ordered task prediction updates and keep the manifest bounded.
///
/// `predictions` and `updates` are ordered from oldest to newest. An update
/// replaces any existing prediction for the same invocation and keeps the
/// update's position in the supplied order. If `updates` repeats an invocation,
/// its last value and position win. Invocations in `protected` are not retired;
/// protecting more entries than the manifest limit, or than fit in
/// [`MAX_TASK_MANIFEST_BYTES`], returns an error.
pub fn update_task_predictions_in_order(
    predictions: Vec<ActionPrediction>,
    mut updates: Vec<ActionPrediction>,
    protected: &BTreeSet<CacheDigest>,
) -> Result<Vec<ActionPrediction>> {
    if protected.len() > MAX_TASK_ACTION_PREDICTIONS {
        bail!("this task run contains too many action predictions");
    }
    let mut update_invocations = BTreeSet::new();
    updates = updates
        .into_iter()
        .rev()
        .filter(|prediction| update_invocations.insert(prediction.invocation.clone()))
        .collect();
    updates.reverse();
    let mut predictions: Vec<_> = predictions
        .into_iter()
        .filter(|prediction| !update_invocations.contains(&prediction.invocation))
        .collect();
    predictions.extend(updates);
    let mut excess = predictions
        .len()
        .saturating_sub(MAX_TASK_ACTION_PREDICTIONS);
    let mut kept = Vec::with_capacity(predictions.len() - excess);
    for prediction in predictions {
        if excess > 0 && !protected.contains(&prediction.invocation) {
            excess -= 1;
        } else {
            kept.push(prediction);
        }
    }
    // The count bound alone does not bound the serialized size, because a
    // payload is adapter-defined. The remote reader refuses anything over its
    // own byte limit, so retire the oldest unprotected entries until what is
    // published fits.
    let sizes = kept
        .iter()
        .map(|prediction| Ok(canonical_json(prediction)?.len() + 1))
        .collect::<Result<Vec<_>>>()?;
    let mut total: usize = TASK_MANIFEST_OVERHEAD_BYTES + sizes.iter().sum::<usize>();
    if total > MAX_TASK_MANIFEST_BYTES {
        let mut retained = Vec::with_capacity(kept.len());
        for (prediction, size) in kept.into_iter().zip(sizes) {
            if total > MAX_TASK_MANIFEST_BYTES && !protected.contains(&prediction.invocation) {
                total -= size;
            } else {
                retained.push(prediction);
            }
        }
        if total > MAX_TASK_MANIFEST_BYTES {
            bail!("this task run contains too many action predictions");
        }
        kept = retained;
    }
    Ok(kept)
}

/// Merge two manifest snapshots through the same bounded durable-LRU rule.
///
/// `required` values from the update side win over the base. `fallbacks` win
/// only when the base has no newer value for the invocation. Other overlap is
/// inherited from the base, while update-only history remains eligible for
/// pruning before either kind of preferred value.
pub(super) fn merge_task_manifests(
    task: &str,
    base: TaskActionManifest,
    update: TaskActionManifest,
    required: &BTreeSet<CacheDigest>,
    fallbacks: &BTreeSet<CacheDigest>,
) -> Result<TaskActionManifest> {
    validate_task_manifest(&base, task)?;
    validate_task_manifest(&update, task)?;
    let base_predictions: BTreeMap<_, _> = base
        .predictions
        .iter()
        .map(|prediction| (&prediction.invocation, prediction))
        .collect();
    let updates: Vec<_> = update
        .predictions
        .iter()
        .filter(|prediction| {
            required.contains(&prediction.invocation)
                || (fallbacks.contains(&prediction.invocation)
                    && base_predictions
                        .get(&prediction.invocation)
                        .is_none_or(|base| *base == *prediction))
        })
        .cloned()
        .collect();
    let protected = required
        .iter()
        .cloned()
        .chain(
            updates
                .iter()
                .map(|prediction| prediction.invocation.clone()),
        )
        .collect();
    let mut predictions: Vec<_> = update
        .predictions
        .into_iter()
        .filter(|prediction| !base_predictions.contains_key(&prediction.invocation))
        .collect();
    predictions.extend(base.predictions);
    let predictions = update_task_predictions_in_order(predictions, updates, &protected)?;
    let manifest = TaskActionManifest {
        version: TASK_ACTION_MANIFEST_VERSION,
        task: task.to_owned(),
        predictions,
    };
    validate_task_manifest(&manifest, task)?;
    Ok(manifest)
}