mbx_cache_core/agent/
manifest.rs

1use super::TASK_ACTION_MANIFEST_VERSION;
2use crate::{CacheDigest, TaskActionManifest};
3use eyre::{Context, Result, bail};
4use log::warn;
5use std::collections::BTreeMap;
6use std::fs;
7use std::path::{Path, PathBuf};
8
9/// Whether `task` is a well-formed task action identity.
10///
11/// Identities name files and directories in the store, so anything that reads
12/// the store back has to be able to tell an identity from whatever else a user
13/// left lying there.
14pub fn is_task_identity(task: &str) -> bool {
15    task.len() == 64
16        && task
17            .bytes()
18            .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
19}
20
21pub(super) fn validate_task_identity(task: &str) -> Result<()> {
22    if !is_task_identity(task) {
23        bail!("invalid task action identity");
24    }
25    Ok(())
26}
27
28/// Where a store keeps its task prediction manifests.
29pub(super) fn task_manifest_dir(store: &Path) -> PathBuf {
30    store.join("task-manifests").join("v1")
31}
32
33/// The action digests a task's prediction manifest recorded.
34///
35/// Read straight off disk rather than through an agent, because a collector
36/// needs the action set of tasks no session is running. A manifest that is
37/// missing or no longer parseable yields no actions rather than an error: this
38/// is a prediction index, so the worst a thin answer costs is a cold prefetch,
39/// or an object collected earlier than it deserved.
40pub fn task_manifest_actions(store: &Path, task: &str) -> Result<Vec<CacheDigest>> {
41    validate_task_identity(task)?;
42    let path = task_manifest_dir(store).join(format!("{task}.json"));
43    let bytes = match fs::read(&path) {
44        Ok(bytes) => bytes,
45        Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
46        Err(error) => {
47            return Err(error).wrap_err_with(|| format!("failed to read {}", path.display()));
48        }
49    };
50    let Ok(manifest) = serde_json::from_slice::<TaskActionManifest>(&bytes) else {
51        return Ok(Vec::new());
52    };
53    if validate_task_manifest(&manifest, task).is_err() {
54        return Ok(Vec::new());
55    }
56    Ok(manifest
57        .predictions
58        .into_iter()
59        .map(|prediction| prediction.action)
60        .collect())
61}
62
63pub(super) fn validate_task_manifest(manifest: &TaskActionManifest, task: &str) -> Result<()> {
64    if manifest.task == task && manifest.validate() {
65        Ok(())
66    } else {
67        bail!("invalid task action manifest")
68    }
69}
70
71pub(super) fn merge_task_manifests(
72    task: &str,
73    base: Option<TaskActionManifest>,
74    update: TaskActionManifest,
75) -> Result<TaskActionManifest> {
76    validate_task_manifest(&update, task)?;
77    let mut predictions = BTreeMap::new();
78    if let Some(base) = base {
79        validate_task_manifest(&base, task)?;
80        predictions.extend(
81            base.predictions
82                .into_iter()
83                .map(|prediction| (prediction.invocation.clone(), prediction)),
84        );
85    }
86    predictions.extend(
87        update
88            .predictions
89            .into_iter()
90            .map(|prediction| (prediction.invocation.clone(), prediction)),
91    );
92    let manifest = TaskActionManifest {
93        version: TASK_ACTION_MANIFEST_VERSION,
94        task: task.to_owned(),
95        predictions: predictions.into_values().collect(),
96    };
97    validate_task_manifest(&manifest, task)?;
98    Ok(manifest)
99}
100
101pub(super) fn merge_remote_task_manifest(
102    task: &str,
103    remote: TaskActionManifest,
104    local: TaskActionManifest,
105) -> (TaskActionManifest, bool) {
106    match merge_task_manifests(task, Some(remote), local.clone()) {
107        Ok(manifest) => (manifest, true),
108        Err(error) => {
109            warn!("remote task action manifest merge failed for {task}: {error}");
110            (local, false)
111        }
112    }
113}
mbx_cache_core/agent/manifest.rs

mbx_cache_core/agent/
manifest.rs
