Skip to main content

nex_pkg/
armory_lock.rs

1use std::collections::BTreeMap;
2use std::path::PathBuf;
3
4use anyhow::{bail, Context, Result};
5use serde::{Deserialize, Serialize};
6
7use crate::armory::{self, ArmoryIndex, ArmoryPackage, PackageRef};
8use crate::config::{Config, RegistryConfig};
9
10pub const PACKAGE_LOCK_SCHEMA: &str = "io.styrene.nex.package-lock.v1";
11pub const ACTIVATION_LOCK_SCHEMA: &str = "io.styrene.nex.omegon-activation-lock.v1";
12
13#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
14#[serde(rename_all = "camelCase")]
15pub struct PackageLock {
16    pub schema: String,
17    pub registries: Vec<LockedRegistry>,
18    pub roots: Vec<LockedRoot>,
19    pub packages: Vec<LockedPackage>,
20}
21
22#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
23#[serde(rename_all = "camelCase")]
24pub struct LockedRegistry {
25    pub name: String,
26    pub url: String,
27    #[serde(default, skip_serializing_if = "Option::is_none")]
28    pub trust: Option<String>,
29}
30
31#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
32#[serde(rename_all = "camelCase")]
33pub struct LockedRoot {
34    pub package_ref: String,
35}
36
37#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
38#[serde(rename_all = "camelCase")]
39pub struct LockedPackage {
40    pub package_ref: String,
41    pub version: Option<String>,
42    pub registry: String,
43    pub oci_ref: Option<String>,
44    pub digest: Option<String>,
45    pub dependencies: Vec<String>,
46    #[serde(default, skip_serializing_if = "Option::is_none")]
47    pub path: Option<String>,
48    #[serde(default)]
49    pub verified: bool,
50    #[serde(default, skip_serializing_if = "Option::is_none")]
51    pub installed_at: Option<String>,
52}
53
54#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
55#[serde(rename_all = "camelCase")]
56pub struct OmegonActivationLock {
57    pub schema: String,
58    pub root: ActivationRoot,
59    pub packages: Vec<ActivationPackage>,
60}
61
62#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
63#[serde(rename_all = "camelCase")]
64pub struct ActivationRoot {
65    pub kind: String,
66    pub id: String,
67    pub version: Option<String>,
68}
69
70#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
71#[serde(rename_all = "camelCase")]
72pub struct ActivationPackage {
73    pub kind: String,
74    pub id: String,
75    pub version: Option<String>,
76    pub status: String,
77    #[serde(default, skip_serializing_if = "Option::is_none")]
78    pub digest: Option<String>,
79    #[serde(default, skip_serializing_if = "Option::is_none")]
80    pub path: Option<String>,
81    #[serde(default)]
82    pub verified: bool,
83}
84
85#[derive(Debug, Clone)]
86pub struct ResolvedGraph {
87    pub root: PackageRef,
88    pub registry: RegistryConfig,
89    pub packages: Vec<LockedPackage>,
90    pub optional_skipped: Vec<String>,
91}
92
93pub fn install(config: &Config, root: &PackageRef, dry_run: bool, lock_only: bool) -> Result<()> {
94    let graph = resolve_from_config(config, root)?;
95    print_plan(&graph);
96    if dry_run {
97        return Ok(());
98    }
99
100    let package_lock = package_lock_for_graph(&graph);
101    write_package_lock(&package_lock)?;
102
103    if is_omegon_runtime_root(&graph.root.kind) {
104        let activation_lock = activation_lock_for_graph(&graph)?;
105        write_activation_lock(&activation_lock)?;
106    }
107
108    println!("  wrote {}", package_lock_path()?.display());
109    if is_omegon_runtime_root(&graph.root.kind) {
110        println!("  wrote {}", activation_lock_path()?.display());
111    }
112    if lock_only {
113        println!("  lock-only: run `nex lock materialize` to fetch packages");
114        return Ok(());
115    }
116    let records = crate::armory_store::materialize_lock()?;
117    for record in records {
118        println!(
119            "  {} -> {} ({})",
120            record.package_ref,
121            record.path.display(),
122            if record.verified {
123                "verified"
124            } else {
125                "unverified"
126            }
127        );
128    }
129    Ok(())
130}
131
132pub fn refresh(config: &Config) -> Result<()> {
133    let existing = read_package_lock()?;
134    for root in existing.roots {
135        let package_ref = PackageRef::parse(&root.package_ref)?;
136        install(config, &package_ref, false, true)?;
137    }
138    Ok(())
139}
140
141pub fn remove_root(root: &PackageRef, dry_run: bool) -> Result<()> {
142    let mut lock = read_package_lock()?;
143    let root_string = root.to_string();
144    let original_roots = lock.roots.len();
145    lock.roots
146        .retain(|locked| locked.package_ref != root_string);
147    if lock.roots.len() == original_roots {
148        bail!("Armory root {root} is not installed");
149    }
150    let reachable = reachable_packages(&lock)?;
151    lock.packages
152        .retain(|package| reachable.contains_key(&package.package_ref));
153
154    println!("removing Armory root {root}");
155    if dry_run {
156        return Ok(());
157    }
158
159    write_package_lock(&lock)?;
160    if let Some(activation_lock) = activation_lock_for_package_lock(&lock)? {
161        write_activation_lock(&activation_lock)?;
162    } else if activation_lock_path()?.exists() {
163        std::fs::remove_file(activation_lock_path()?)?;
164    }
165    println!("  removed Armory root {root}");
166    Ok(())
167}
168
169fn reachable_packages(lock: &PackageLock) -> Result<BTreeMap<String, ()>> {
170    let mut reachable = BTreeMap::new();
171    let packages: BTreeMap<&str, &LockedPackage> = lock
172        .packages
173        .iter()
174        .map(|package| (package.package_ref.as_str(), package))
175        .collect();
176    for root in &lock.roots {
177        mark_reachable(&root.package_ref, &packages, &mut reachable)?;
178    }
179    Ok(reachable)
180}
181
182fn mark_reachable<'a>(
183    package_ref: &str,
184    packages: &BTreeMap<&'a str, &'a LockedPackage>,
185    reachable: &mut BTreeMap<String, ()>,
186) -> Result<()> {
187    if reachable.insert(package_ref.to_string(), ()).is_some() {
188        return Ok(());
189    }
190    let Some(package) = packages.get(package_ref) else {
191        bail!("package lock missing package {package_ref}");
192    };
193    for dependency in &package.dependencies {
194        mark_reachable(dependency, packages, reachable)?;
195    }
196    Ok(())
197}
198
199pub fn resolve_from_config(config: &Config, root: &PackageRef) -> Result<ResolvedGraph> {
200    for registry in &config.registries {
201        let index = armory::fetch_index(registry)?;
202        if armory::find(&index, root).is_some() {
203            return resolve_graph(registry, &index, root);
204        }
205    }
206    bail!("Armory package {root} not found in configured registries")
207}
208
209pub fn resolve_graph(
210    registry: &RegistryConfig,
211    index: &ArmoryIndex,
212    root: &PackageRef,
213) -> Result<ResolvedGraph> {
214    let mut packages = BTreeMap::new();
215    let mut optional_skipped = Vec::new();
216    let mut visiting = Vec::new();
217    resolve_one(
218        registry,
219        index,
220        root,
221        &mut packages,
222        &mut optional_skipped,
223        &mut visiting,
224    )?;
225
226    Ok(ResolvedGraph {
227        root: root.clone(),
228        registry: registry.clone(),
229        packages: packages.into_values().collect(),
230        optional_skipped,
231    })
232}
233
234fn resolve_one(
235    registry: &RegistryConfig,
236    index: &ArmoryIndex,
237    package_ref: &PackageRef,
238    packages: &mut BTreeMap<String, LockedPackage>,
239    optional_skipped: &mut Vec<String>,
240    visiting: &mut Vec<String>,
241) -> Result<()> {
242    let key = package_ref.to_string();
243    if packages.contains_key(&key) {
244        return Ok(());
245    }
246    if let Some(pos) = visiting.iter().position(|item| item == &key) {
247        let mut cycle = visiting[pos..].to_vec();
248        cycle.push(key);
249        bail!("Armory dependency cycle: {}", cycle.join(" -> "));
250    }
251
252    let package = armory::find(index, package_ref)
253        .with_context(|| format!("missing required Armory dependency {package_ref}"))?;
254
255    visiting.push(key.clone());
256    let dependencies = required_dependency_refs(package);
257    for dependency in &dependencies {
258        let dep_ref = PackageRef::parse(dependency)?;
259        resolve_one(
260            registry,
261            index,
262            &dep_ref,
263            packages,
264            optional_skipped,
265            visiting,
266        )?;
267    }
268    for optional in optional_dependency_refs(package) {
269        optional_skipped.push(optional);
270    }
271    visiting.pop();
272
273    let locked = LockedPackage {
274        package_ref: key.clone(),
275        version: package.version.clone(),
276        registry: registry.name.clone(),
277        oci_ref: package.oci_ref.clone(),
278        digest: package.digest.clone(),
279        dependencies,
280        path: None,
281        verified: false,
282        installed_at: None,
283    };
284
285    if let Some(existing) = packages.get(&key) {
286        if existing.version != locked.version || existing.digest != locked.digest {
287            bail!("conflicting Armory package records for {key}");
288        }
289    }
290    packages.insert(key, locked);
291    Ok(())
292}
293
294fn required_dependency_refs(package: &ArmoryPackage) -> Vec<String> {
295    package
296        .dependencies
297        .iter()
298        .filter(|dep| !dep.optional.unwrap_or(false))
299        .filter_map(|dep| dep.display_ref().map(ToString::to_string))
300        .collect()
301}
302
303fn optional_dependency_refs(package: &ArmoryPackage) -> Vec<String> {
304    let mut refs: Vec<String> = package
305        .dependencies
306        .iter()
307        .filter(|dep| dep.optional.unwrap_or(false))
308        .filter_map(|dep| dep.display_ref().map(ToString::to_string))
309        .collect();
310    refs.extend(
311        package
312            .optional_dependencies
313            .iter()
314            .filter_map(|dep| dep.display_ref().map(ToString::to_string)),
315    );
316    refs
317}
318
319fn package_lock_for_graph(graph: &ResolvedGraph) -> PackageLock {
320    PackageLock {
321        schema: PACKAGE_LOCK_SCHEMA.to_string(),
322        registries: vec![LockedRegistry {
323            name: graph.registry.name.clone(),
324            url: graph.registry.url.clone(),
325            trust: graph.registry.trust.clone(),
326        }],
327        roots: vec![LockedRoot {
328            package_ref: graph.root.to_string(),
329        }],
330        packages: graph.packages.clone(),
331    }
332}
333
334fn activation_lock_for_graph(graph: &ResolvedGraph) -> Result<OmegonActivationLock> {
335    let root_package = graph
336        .packages
337        .iter()
338        .find(|package| package.package_ref == graph.root.to_string())
339        .context("resolved graph missing root package")?;
340    Ok(OmegonActivationLock {
341        schema: ACTIVATION_LOCK_SCHEMA.to_string(),
342        root: ActivationRoot {
343            kind: graph.root.kind.clone(),
344            id: graph.root.id.clone(),
345            version: root_package.version.clone(),
346        },
347        packages: graph
348            .packages
349            .iter()
350            .map(|package| {
351                let package_ref = PackageRef::parse(&package.package_ref)?;
352                Ok(ActivationPackage {
353                    kind: package_ref.kind,
354                    id: package_ref.id,
355                    version: package.version.clone(),
356                    status: if package.verified && package.path.is_some() {
357                        "installed".to_string()
358                    } else {
359                        "pending".to_string()
360                    },
361                    digest: package.digest.clone(),
362                    path: package.path.clone(),
363                    verified: package.verified,
364                })
365            })
366            .collect::<Result<Vec<_>>>()?,
367    })
368}
369
370pub(crate) fn activation_lock_for_package_lock(
371    lock: &PackageLock,
372) -> Result<Option<OmegonActivationLock>> {
373    let Some(root) = lock.roots.first() else {
374        return Ok(None);
375    };
376    let root_ref = PackageRef::parse(&root.package_ref)?;
377    if !is_omegon_runtime_root(&root_ref.kind) {
378        return Ok(None);
379    }
380    let root_package = lock
381        .packages
382        .iter()
383        .find(|package| package.package_ref == root.package_ref)
384        .context("package lock missing root package")?;
385    Ok(Some(OmegonActivationLock {
386        schema: ACTIVATION_LOCK_SCHEMA.to_string(),
387        root: ActivationRoot {
388            kind: root_ref.kind,
389            id: root_ref.id,
390            version: root_package.version.clone(),
391        },
392        packages: lock
393            .packages
394            .iter()
395            .map(|package| {
396                let package_ref = PackageRef::parse(&package.package_ref)?;
397                Ok(ActivationPackage {
398                    kind: package_ref.kind,
399                    id: package_ref.id,
400                    version: package.version.clone(),
401                    status: if package.verified && package.path.is_some() {
402                        "installed".to_string()
403                    } else {
404                        "pending".to_string()
405                    },
406                    digest: package.digest.clone(),
407                    path: package.path.clone(),
408                    verified: package.verified,
409                })
410            })
411            .collect::<Result<Vec<_>>>()?,
412    }))
413}
414
415fn is_omegon_runtime_root(kind: &str) -> bool {
416    matches!(
417        kind,
418        "skill" | "persona" | "tone" | "profile" | "agent" | "extension" | "workstation"
419    )
420}
421
422fn print_plan(graph: &ResolvedGraph) {
423    println!("Armory install plan: {}", graph.root);
424    for package in &graph.packages {
425        let version = package.version.as_deref().unwrap_or("unknown");
426        println!("  {} {}", package.package_ref, version);
427    }
428    for optional in &graph.optional_skipped {
429        println!("  optional skipped: {optional}");
430    }
431}
432
433pub(crate) fn write_package_lock(lock: &PackageLock) -> Result<()> {
434    let path = package_lock_path()?;
435    if let Some(parent) = path.parent() {
436        std::fs::create_dir_all(parent)?;
437    }
438    let bytes = serde_json::to_vec_pretty(lock)?;
439    std::fs::write(path, bytes)?;
440    Ok(())
441}
442
443pub(crate) fn write_activation_lock(lock: &OmegonActivationLock) -> Result<()> {
444    let path = activation_lock_path()?;
445    if let Some(parent) = path.parent() {
446        std::fs::create_dir_all(parent)?;
447    }
448    let bytes = serde_json::to_vec_pretty(lock)?;
449    std::fs::write(path, bytes)?;
450    Ok(())
451}
452
453pub(crate) fn read_package_lock() -> Result<PackageLock> {
454    let path = package_lock_path()?;
455    let bytes = std::fs::read(&path).with_context(|| format!("reading {}", path.display()))?;
456    serde_json::from_slice(&bytes).context("parsing package lock")
457}
458
459pub fn package_lock_path() -> Result<PathBuf> {
460    Ok(state_dir()?.join("packages.lock.json"))
461}
462
463pub fn activation_lock_path() -> Result<PathBuf> {
464    Ok(state_dir()?.join("omegon-activation-lock.json"))
465}
466
467fn state_dir() -> Result<PathBuf> {
468    let home = std::env::var_os("HOME").context("HOME is not set")?;
469    Ok(PathBuf::from(home).join(".local/state/nex"))
470}
471
472#[cfg(test)]
473mod tests {
474    use super::{activation_lock_for_graph, reachable_packages, resolve_graph};
475    use crate::armory::parse_index;
476    use crate::armory::PackageRef;
477    use crate::config::RegistryConfig;
478
479    fn registry() -> RegistryConfig {
480        RegistryConfig {
481            name: "test".to_string(),
482            url: "https://example.test/index.json".to_string(),
483            trust: None,
484        }
485    }
486
487    #[test]
488    fn resolves_required_dependencies() {
489        let index = parse_index(
490            br#"{"packages":[
491                {"packageRef":"profile/root","version":"1.0.0","dependencies":[{"packageRef":"skill/rust"}]},
492                {"packageRef":"skill/rust","version":"1.0.0"}
493            ]}"#,
494        )
495        .expect("index");
496        let graph = resolve_graph(
497            &registry(),
498            &index,
499            &PackageRef::parse("profile/root").expect("ref"),
500        )
501        .expect("graph");
502        assert_eq!(graph.packages.len(), 2);
503    }
504
505    #[test]
506    fn missing_dependency_fails() {
507        let index = parse_index(
508            br#"{"packages":[{"packageRef":"profile/root","dependencies":[{"packageRef":"skill/missing"}]}]}"#,
509        )
510        .expect("index");
511        let error = resolve_graph(
512            &registry(),
513            &index,
514            &PackageRef::parse("profile/root").expect("ref"),
515        )
516        .expect_err("missing dependency");
517        assert!(format!("{error:#}").contains("skill/missing"));
518    }
519
520    #[test]
521    fn dependency_cycle_fails() {
522        let index = parse_index(
523            br#"{"packages":[
524                {"packageRef":"profile/a","dependencies":[{"packageRef":"profile/b"}]},
525                {"packageRef":"profile/b","dependencies":[{"packageRef":"profile/a"}]}
526            ]}"#,
527        )
528        .expect("index");
529        let error = resolve_graph(
530            &registry(),
531            &index,
532            &PackageRef::parse("profile/a").expect("ref"),
533        )
534        .expect_err("cycle");
535        assert!(format!("{error:#}").contains("profile/a -> profile/b -> profile/a"));
536    }
537
538    #[test]
539    fn activation_lock_is_pending_until_materialized() {
540        let index =
541            parse_index(br#"{"packages":[{"packageRef":"profile/root","version":"1.0.0"}]}"#)
542                .expect("index");
543        let graph = resolve_graph(
544            &registry(),
545            &index,
546            &PackageRef::parse("profile/root").expect("ref"),
547        )
548        .expect("graph");
549        let lock = activation_lock_for_graph(&graph).expect("activation lock");
550        assert_eq!(lock.packages[0].status, "pending");
551    }
552
553    #[test]
554    fn reachable_packages_prunes_orphaned_dependencies() {
555        let lock = super::PackageLock {
556            schema: super::PACKAGE_LOCK_SCHEMA.to_string(),
557            registries: Vec::new(),
558            roots: vec![super::LockedRoot {
559                package_ref: "profile/root".to_string(),
560            }],
561            packages: vec![
562                super::LockedPackage {
563                    package_ref: "profile/root".to_string(),
564                    version: None,
565                    registry: "test".to_string(),
566                    oci_ref: None,
567                    digest: None,
568                    dependencies: vec!["skill/used".to_string()],
569                    path: None,
570                    verified: false,
571                    installed_at: None,
572                },
573                super::LockedPackage {
574                    package_ref: "skill/used".to_string(),
575                    version: None,
576                    registry: "test".to_string(),
577                    oci_ref: None,
578                    digest: None,
579                    dependencies: Vec::new(),
580                    path: None,
581                    verified: false,
582                    installed_at: None,
583                },
584                super::LockedPackage {
585                    package_ref: "skill/orphan".to_string(),
586                    version: None,
587                    registry: "test".to_string(),
588                    oci_ref: None,
589                    digest: None,
590                    dependencies: Vec::new(),
591                    path: None,
592                    verified: false,
593                    installed_at: None,
594                },
595            ],
596        };
597
598        let reachable = reachable_packages(&lock).expect("reachable");
599
600        assert!(reachable.contains_key("profile/root"));
601        assert!(reachable.contains_key("skill/used"));
602        assert!(!reachable.contains_key("skill/orphan"));
603    }
604}