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    pub activation: ActivationSettings,
78    #[serde(default, skip_serializing_if = "Option::is_none")]
79    pub digest: Option<String>,
80    #[serde(default, skip_serializing_if = "Option::is_none")]
81    pub path: Option<String>,
82    #[serde(default)]
83    pub verified: bool,
84}
85
86#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
87#[serde(rename_all = "camelCase")]
88pub struct ActivationSettings {
89    pub runtime: String,
90    pub mode: String,
91    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
92    pub defaults: BTreeMap<String, String>,
93}
94
95#[derive(Debug, Clone)]
96pub struct ResolvedGraph {
97    pub root: PackageRef,
98    pub registry: RegistryConfig,
99    pub packages: Vec<LockedPackage>,
100    pub optional_skipped: Vec<String>,
101}
102
103pub fn install(config: &Config, root: &PackageRef, dry_run: bool, lock_only: bool) -> Result<()> {
104    let graph = resolve_from_config(config, root)?;
105    print_plan(&graph);
106    if dry_run {
107        return Ok(());
108    }
109
110    let package_lock = package_lock_for_graph(&graph);
111    write_package_lock(&package_lock)?;
112
113    if is_omegon_runtime_root(&graph.root.kind) {
114        let activation_lock = activation_lock_for_graph(&graph)?;
115        write_activation_lock(&activation_lock)?;
116    }
117
118    println!("  wrote {}", package_lock_path()?.display());
119    if is_omegon_runtime_root(&graph.root.kind) {
120        println!("  wrote {}", activation_lock_path()?.display());
121    }
122    if lock_only {
123        println!("  lock-only: run `nex lock materialize` to fetch packages");
124        return Ok(());
125    }
126    let records = crate::armory_store::materialize_lock()?;
127    for record in records {
128        println!(
129            "  {} -> {} ({})",
130            record.package_ref,
131            record.path.display(),
132            if record.verified {
133                "verified"
134            } else {
135                "unverified"
136            }
137        );
138    }
139    Ok(())
140}
141
142pub fn refresh(config: &Config) -> Result<()> {
143    let existing = read_package_lock()?;
144    for root in existing.roots {
145        let package_ref = PackageRef::parse(&root.package_ref)?;
146        install(config, &package_ref, false, true)?;
147    }
148    Ok(())
149}
150
151pub fn remove_root(root: &PackageRef, dry_run: bool) -> Result<()> {
152    let mut lock = read_package_lock()?;
153    let root_string = root.to_string();
154    let original_roots = lock.roots.len();
155    lock.roots
156        .retain(|locked| locked.package_ref != root_string);
157    if lock.roots.len() == original_roots {
158        bail!("Armory root {root} is not installed");
159    }
160    let reachable = reachable_packages(&lock)?;
161    lock.packages
162        .retain(|package| reachable.contains_key(&package.package_ref));
163
164    println!("removing Armory root {root}");
165    if dry_run {
166        return Ok(());
167    }
168
169    write_package_lock(&lock)?;
170    if let Some(activation_lock) = activation_lock_for_package_lock(&lock)? {
171        write_activation_lock(&activation_lock)?;
172    } else if activation_lock_path()?.exists() {
173        std::fs::remove_file(activation_lock_path()?)?;
174    }
175    println!("  removed Armory root {root}");
176    Ok(())
177}
178
179fn reachable_packages(lock: &PackageLock) -> Result<BTreeMap<String, ()>> {
180    let mut reachable = BTreeMap::new();
181    let packages: BTreeMap<&str, &LockedPackage> = lock
182        .packages
183        .iter()
184        .map(|package| (package.package_ref.as_str(), package))
185        .collect();
186    for root in &lock.roots {
187        mark_reachable(&root.package_ref, &packages, &mut reachable)?;
188    }
189    Ok(reachable)
190}
191
192fn mark_reachable<'a>(
193    package_ref: &str,
194    packages: &BTreeMap<&'a str, &'a LockedPackage>,
195    reachable: &mut BTreeMap<String, ()>,
196) -> Result<()> {
197    if reachable.insert(package_ref.to_string(), ()).is_some() {
198        return Ok(());
199    }
200    let Some(package) = packages.get(package_ref) else {
201        bail!("package lock missing package {package_ref}");
202    };
203    for dependency in &package.dependencies {
204        mark_reachable(dependency, packages, reachable)?;
205    }
206    Ok(())
207}
208
209pub fn resolve_from_config(config: &Config, root: &PackageRef) -> Result<ResolvedGraph> {
210    for registry in &config.registries {
211        let index = armory::fetch_index(registry)?;
212        if armory::find(&index, root).is_some() {
213            return resolve_graph(registry, &index, root);
214        }
215    }
216    bail!("Armory package {root} not found in configured registries")
217}
218
219pub fn resolve_graph(
220    registry: &RegistryConfig,
221    index: &ArmoryIndex,
222    root: &PackageRef,
223) -> Result<ResolvedGraph> {
224    let mut packages = BTreeMap::new();
225    let mut optional_skipped = Vec::new();
226    let mut visiting = Vec::new();
227    resolve_one(
228        registry,
229        index,
230        root,
231        &mut packages,
232        &mut optional_skipped,
233        &mut visiting,
234    )?;
235
236    Ok(ResolvedGraph {
237        root: root.clone(),
238        registry: registry.clone(),
239        packages: packages.into_values().collect(),
240        optional_skipped,
241    })
242}
243
244fn resolve_one(
245    registry: &RegistryConfig,
246    index: &ArmoryIndex,
247    package_ref: &PackageRef,
248    packages: &mut BTreeMap<String, LockedPackage>,
249    optional_skipped: &mut Vec<String>,
250    visiting: &mut Vec<String>,
251) -> Result<()> {
252    let key = package_ref.to_string();
253    if packages.contains_key(&key) {
254        return Ok(());
255    }
256    if let Some(pos) = visiting.iter().position(|item| item == &key) {
257        let mut cycle = visiting[pos..].to_vec();
258        cycle.push(key);
259        bail!("Armory dependency cycle: {}", cycle.join(" -> "));
260    }
261
262    let package = armory::find(index, package_ref)
263        .with_context(|| format!("missing required Armory dependency {package_ref}"))?;
264
265    visiting.push(key.clone());
266    let dependencies = required_dependency_refs(package);
267    for dependency in &dependencies {
268        let dep_ref = PackageRef::parse(dependency)?;
269        resolve_one(
270            registry,
271            index,
272            &dep_ref,
273            packages,
274            optional_skipped,
275            visiting,
276        )?;
277    }
278    for optional in optional_dependency_refs(package) {
279        optional_skipped.push(optional);
280    }
281    visiting.pop();
282
283    let locked = LockedPackage {
284        package_ref: key.clone(),
285        version: package.version.clone(),
286        registry: registry.name.clone(),
287        oci_ref: package.oci_ref.clone(),
288        digest: package.digest.clone(),
289        dependencies,
290        path: None,
291        verified: false,
292        installed_at: None,
293    };
294
295    if let Some(existing) = packages.get(&key) {
296        if existing.version != locked.version || existing.digest != locked.digest {
297            bail!("conflicting Armory package records for {key}");
298        }
299    }
300    packages.insert(key, locked);
301    Ok(())
302}
303
304fn required_dependency_refs(package: &ArmoryPackage) -> Vec<String> {
305    package
306        .dependencies
307        .iter()
308        .filter(|dep| !dep.optional.unwrap_or(false))
309        .filter_map(|dep| dep.display_ref().map(ToString::to_string))
310        .collect()
311}
312
313fn optional_dependency_refs(package: &ArmoryPackage) -> Vec<String> {
314    let mut refs: Vec<String> = package
315        .dependencies
316        .iter()
317        .filter(|dep| dep.optional.unwrap_or(false))
318        .filter_map(|dep| dep.display_ref().map(ToString::to_string))
319        .collect();
320    refs.extend(
321        package
322            .optional_dependencies
323            .iter()
324            .filter_map(|dep| dep.display_ref().map(ToString::to_string)),
325    );
326    refs
327}
328
329fn package_lock_for_graph(graph: &ResolvedGraph) -> PackageLock {
330    PackageLock {
331        schema: PACKAGE_LOCK_SCHEMA.to_string(),
332        registries: vec![LockedRegistry {
333            name: graph.registry.name.clone(),
334            url: graph.registry.url.clone(),
335            trust: graph.registry.trust.clone(),
336        }],
337        roots: vec![LockedRoot {
338            package_ref: graph.root.to_string(),
339        }],
340        packages: graph.packages.clone(),
341    }
342}
343
344fn activation_lock_for_graph(graph: &ResolvedGraph) -> Result<OmegonActivationLock> {
345    activation_lock_from_packages(&graph.root.to_string(), &graph.packages)
346}
347
348pub(crate) fn activation_lock_for_package_lock(
349    lock: &PackageLock,
350) -> Result<Option<OmegonActivationLock>> {
351    let Some(root) = lock.roots.first() else {
352        return Ok(None);
353    };
354    let root_ref = PackageRef::parse(&root.package_ref)?;
355    if !is_omegon_runtime_root(&root_ref.kind) {
356        return Ok(None);
357    }
358    Ok(Some(activation_lock_from_packages(
359        &root.package_ref,
360        &lock.packages,
361    )?))
362}
363
364fn activation_lock_from_packages(
365    root_package_ref: &str,
366    packages: &[LockedPackage],
367) -> Result<OmegonActivationLock> {
368    let root_ref = PackageRef::parse(root_package_ref)?;
369    let root_package = packages
370        .iter()
371        .find(|package| package.package_ref == root_package_ref)
372        .context("package lock missing root package")?;
373    Ok(OmegonActivationLock {
374        schema: ACTIVATION_LOCK_SCHEMA.to_string(),
375        root: ActivationRoot {
376            kind: root_ref.kind,
377            id: root_ref.id,
378            version: root_package.version.clone(),
379        },
380        packages: packages
381            .iter()
382            .map(activation_package_for_locked_package)
383            .collect::<Result<Vec<_>>>()?,
384    })
385}
386
387fn activation_package_for_locked_package(package: &LockedPackage) -> Result<ActivationPackage> {
388    let package_ref = PackageRef::parse(&package.package_ref)?;
389    let status = if package.verified && package.path.is_some() {
390        "installed"
391    } else {
392        "pending"
393    };
394    Ok(ActivationPackage {
395        activation: activation_settings_for_kind(&package_ref.kind),
396        kind: package_ref.kind,
397        id: package_ref.id,
398        version: package.version.clone(),
399        status: status.to_string(),
400        digest: package.digest.clone(),
401        path: package.path.clone(),
402        verified: package.verified,
403    })
404}
405
406fn activation_settings_for_kind(kind: &str) -> ActivationSettings {
407    let (runtime, mode, defaults) = match kind {
408        "skill" => ("skill", "load", BTreeMap::new()),
409        "persona" => ("persona", "load", BTreeMap::new()),
410        "tone" => ("tone", "load", BTreeMap::new()),
411        "profile" => (
412            "omegon-agent-profile",
413            "defaults",
414            BTreeMap::from([("profileRole".to_string(), "agent-defaults".to_string())]),
415        ),
416        "agent" => (
417            "omegon-agent",
418            "instantiate",
419            BTreeMap::from([("enabled".to_string(), "true".to_string())]),
420        ),
421        "extension" => (
422            "omegon-extension",
423            "register",
424            BTreeMap::from([("enabled".to_string(), "true".to_string())]),
425        ),
426        "workstation" => (
427            "omegon-workstation",
428            "composite",
429            BTreeMap::from([("materialization".to_string(), "nex".to_string())]),
430        ),
431        _ => ("artifact", "load", BTreeMap::new()),
432    };
433    ActivationSettings {
434        runtime: runtime.to_string(),
435        mode: mode.to_string(),
436        defaults,
437    }
438}
439
440fn is_omegon_runtime_root(kind: &str) -> bool {
441    matches!(
442        kind,
443        "skill" | "persona" | "tone" | "profile" | "agent" | "extension" | "workstation"
444    )
445}
446
447fn print_plan(graph: &ResolvedGraph) {
448    println!("Armory install plan: {}", graph.root);
449    for package in &graph.packages {
450        let version = package.version.as_deref().unwrap_or("unknown");
451        println!("  {} {}", package.package_ref, version);
452    }
453    for optional in &graph.optional_skipped {
454        println!("  optional skipped: {optional}");
455    }
456}
457
458pub(crate) fn write_package_lock(lock: &PackageLock) -> Result<()> {
459    let path = package_lock_path()?;
460    if let Some(parent) = path.parent() {
461        std::fs::create_dir_all(parent)?;
462    }
463    let bytes = serde_json::to_vec_pretty(lock)?;
464    std::fs::write(path, bytes)?;
465    Ok(())
466}
467
468pub(crate) fn write_activation_lock(lock: &OmegonActivationLock) -> Result<()> {
469    let path = activation_lock_path()?;
470    if let Some(parent) = path.parent() {
471        std::fs::create_dir_all(parent)?;
472    }
473    let bytes = serde_json::to_vec_pretty(lock)?;
474    std::fs::write(path, bytes)?;
475    Ok(())
476}
477
478pub(crate) fn read_package_lock() -> Result<PackageLock> {
479    let path = package_lock_path()?;
480    let bytes = std::fs::read(&path).with_context(|| format!("reading {}", path.display()))?;
481    serde_json::from_slice(&bytes).context("parsing package lock")
482}
483
484pub fn package_lock_path() -> Result<PathBuf> {
485    Ok(state_dir()?.join("packages.lock.json"))
486}
487
488pub fn activation_lock_path() -> Result<PathBuf> {
489    Ok(state_dir()?.join("omegon-activation-lock.json"))
490}
491
492fn state_dir() -> Result<PathBuf> {
493    let home = std::env::var_os("HOME").context("HOME is not set")?;
494    Ok(PathBuf::from(home).join(".local/state/nex"))
495}
496
497#[cfg(test)]
498mod tests {
499    use super::{
500        activation_lock_for_graph, activation_lock_for_package_lock, reachable_packages,
501        resolve_graph,
502    };
503    use crate::armory::parse_index;
504    use crate::armory::PackageRef;
505    use crate::config::RegistryConfig;
506
507    fn registry() -> RegistryConfig {
508        RegistryConfig {
509            name: "test".to_string(),
510            url: "https://example.test/index.json".to_string(),
511            trust: None,
512        }
513    }
514
515    #[test]
516    fn resolves_required_dependencies() {
517        let index = parse_index(
518            br#"{"packages":[
519                {"packageRef":"profile/root","version":"1.0.0","dependencies":[{"packageRef":"skill/rust"}]},
520                {"packageRef":"skill/rust","version":"1.0.0"}
521            ]}"#,
522        )
523        .expect("index");
524        let graph = resolve_graph(
525            &registry(),
526            &index,
527            &PackageRef::parse("profile/root").expect("ref"),
528        )
529        .expect("graph");
530        assert_eq!(graph.packages.len(), 2);
531    }
532
533    #[test]
534    fn missing_dependency_fails() {
535        let index = parse_index(
536            br#"{"packages":[{"packageRef":"profile/root","dependencies":[{"packageRef":"skill/missing"}]}]}"#,
537        )
538        .expect("index");
539        let error = resolve_graph(
540            &registry(),
541            &index,
542            &PackageRef::parse("profile/root").expect("ref"),
543        )
544        .expect_err("missing dependency");
545        assert!(format!("{error:#}").contains("skill/missing"));
546    }
547
548    #[test]
549    fn dependency_cycle_fails() {
550        let index = parse_index(
551            br#"{"packages":[
552                {"packageRef":"profile/a","dependencies":[{"packageRef":"profile/b"}]},
553                {"packageRef":"profile/b","dependencies":[{"packageRef":"profile/a"}]}
554            ]}"#,
555        )
556        .expect("index");
557        let error = resolve_graph(
558            &registry(),
559            &index,
560            &PackageRef::parse("profile/a").expect("ref"),
561        )
562        .expect_err("cycle");
563        assert!(format!("{error:#}").contains("profile/a -> profile/b -> profile/a"));
564    }
565
566    #[test]
567    fn activation_lock_is_pending_until_materialized() {
568        let index =
569            parse_index(br#"{"packages":[{"packageRef":"profile/root","version":"1.0.0"}]}"#)
570                .expect("index");
571        let graph = resolve_graph(
572            &registry(),
573            &index,
574            &PackageRef::parse("profile/root").expect("ref"),
575        )
576        .expect("graph");
577        let lock = activation_lock_for_graph(&graph).expect("activation lock");
578        assert_eq!(lock.packages[0].status, "pending");
579    }
580
581    #[test]
582    fn activation_lock_from_package_lock_uses_local_paths_without_registry() {
583        let lock = super::PackageLock {
584            schema: super::PACKAGE_LOCK_SCHEMA.to_string(),
585            registries: vec![super::LockedRegistry {
586                name: "offline".to_string(),
587                url: "https://offline.invalid/index.json".to_string(),
588                trust: None,
589            }],
590            roots: vec![super::LockedRoot {
591                package_ref: "profile/root".to_string(),
592            }],
593            packages: vec![super::LockedPackage {
594                package_ref: "profile/root".to_string(),
595                version: Some("1.0.0".to_string()),
596                registry: "offline".to_string(),
597                oci_ref: Some("oci://offline/root".to_string()),
598                digest: Some("sha256:abc".to_string()),
599                dependencies: Vec::new(),
600                path: Some("/nex/store/root".to_string()),
601                verified: true,
602                installed_at: Some("1".to_string()),
603            }],
604        };
605
606        let activation = activation_lock_for_package_lock(&lock)
607            .expect("activation lock result")
608            .expect("runtime root activation lock");
609
610        assert_eq!(activation.packages[0].status, "installed");
611        assert_eq!(
612            activation.packages[0].path.as_deref(),
613            Some("/nex/store/root")
614        );
615        assert!(activation.packages[0].verified);
616    }
617
618    #[test]
619    fn activation_lock_encodes_runtime_defaults_for_root_kinds() {
620        for (kind, runtime, mode, default_key) in [
621            ("profile", "omegon-agent-profile", "defaults", "profileRole"),
622            ("agent", "omegon-agent", "instantiate", "enabled"),
623            ("extension", "omegon-extension", "register", "enabled"),
624            (
625                "workstation",
626                "omegon-workstation",
627                "composite",
628                "materialization",
629            ),
630        ] {
631            let package_ref = format!("{kind}/root");
632            let lock = super::PackageLock {
633                schema: super::PACKAGE_LOCK_SCHEMA.to_string(),
634                registries: Vec::new(),
635                roots: vec![super::LockedRoot {
636                    package_ref: package_ref.clone(),
637                }],
638                packages: vec![super::LockedPackage {
639                    package_ref,
640                    version: None,
641                    registry: "test".to_string(),
642                    oci_ref: None,
643                    digest: None,
644                    dependencies: Vec::new(),
645                    path: None,
646                    verified: false,
647                    installed_at: None,
648                }],
649            };
650
651            let activation = activation_lock_for_package_lock(&lock)
652                .expect("activation lock result")
653                .expect("runtime root activation lock");
654
655            assert_eq!(activation.packages[0].activation.runtime, runtime);
656            assert_eq!(activation.packages[0].activation.mode, mode);
657            assert!(activation.packages[0]
658                .activation
659                .defaults
660                .contains_key(default_key));
661        }
662    }
663
664    #[test]
665    fn reachable_packages_prunes_orphaned_dependencies() {
666        let lock = super::PackageLock {
667            schema: super::PACKAGE_LOCK_SCHEMA.to_string(),
668            registries: Vec::new(),
669            roots: vec![super::LockedRoot {
670                package_ref: "profile/root".to_string(),
671            }],
672            packages: vec![
673                super::LockedPackage {
674                    package_ref: "profile/root".to_string(),
675                    version: None,
676                    registry: "test".to_string(),
677                    oci_ref: None,
678                    digest: None,
679                    dependencies: vec!["skill/used".to_string()],
680                    path: None,
681                    verified: false,
682                    installed_at: None,
683                },
684                super::LockedPackage {
685                    package_ref: "skill/used".to_string(),
686                    version: None,
687                    registry: "test".to_string(),
688                    oci_ref: None,
689                    digest: None,
690                    dependencies: Vec::new(),
691                    path: None,
692                    verified: false,
693                    installed_at: None,
694                },
695                super::LockedPackage {
696                    package_ref: "skill/orphan".to_string(),
697                    version: None,
698                    registry: "test".to_string(),
699                    oci_ref: None,
700                    digest: None,
701                    dependencies: Vec::new(),
702                    path: None,
703                    verified: false,
704                    installed_at: None,
705                },
706            ],
707        };
708
709        let reachable = reachable_packages(&lock).expect("reachable");
710
711        assert!(reachable.contains_key("profile/root"));
712        assert!(reachable.contains_key("skill/used"));
713        assert!(!reachable.contains_key("skill/orphan"));
714    }
715}