systemprompt_models/managed/
bundle.rs1use serde::{Deserialize, Serialize};
7use std::collections::{BTreeMap, BTreeSet};
8use systemprompt_identifiers::ResourceRevisionId;
9
10use super::assets::{AssetDigest, AssetFile, RevisionFiles};
11use super::error::{RevisionBundleError, invalid};
12use super::manifest::RevisionManifest;
13
14type Result<T> = std::result::Result<T, RevisionBundleError>;
15
16pub const MAX_REVISIONS: usize = 64;
17pub const MAX_FILES: usize = 256;
18pub const MAX_BYTES: usize = 8 * 1024 * 1024;
19pub const ASSEMBLER_VERSION: &str = "managed-bundle-v1";
20
21#[derive(Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
26#[serde(deny_unknown_fields)]
27pub struct RevisionBundle {
28 pub schema_version: u32,
29 pub assembler_version: String,
30 pub root: ResourceRevisionId,
31 pub revisions: BTreeMap<ResourceRevisionId, RevisionManifest>,
32 pub assets: BTreeMap<AssetDigest, Vec<u8>>,
33}
34
35impl std::fmt::Debug for RevisionBundle {
36 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
37 f.debug_struct("RevisionBundle")
38 .field("schema_version", &self.schema_version)
39 .field("root", &self.root)
40 .field("revisions", &self.revisions.len())
41 .field("assets", &self.assets.len())
42 .finish_non_exhaustive()
43 }
44}
45
46impl RevisionBundle {
47 pub fn verify(&self) -> Result<()> {
48 if self.schema_version != 1
49 || self.assembler_version != ASSEMBLER_VERSION
50 || self.revisions.is_empty()
51 || self.revisions.len() > MAX_REVISIONS
52 {
53 return Err(invalid("Unsupported or excessive revision bundle"));
54 }
55 let mut used_assets = BTreeSet::new();
56 let mut file_count = 0usize;
57 let mut expanded_bytes = 0usize;
58 for manifest in self.revisions.values() {
59 if manifest.schema_version != 1 {
60 return Err(invalid("Unsupported revision manifest"));
61 }
62 file_count += manifest.files.len();
63 if file_count > MAX_FILES {
64 return Err(invalid("Bundle exceeds 256 files"));
65 }
66 let mut files = RevisionFiles::default();
67 for (path, entry) in &manifest.files {
68 let bytes = self
69 .assets
70 .get(&entry.digest)
71 .ok_or(RevisionBundleError::Integrity)?;
72 expanded_bytes = expanded_bytes
73 .checked_add(bytes.len())
74 .ok_or(RevisionBundleError::Integrity)?;
75 if expanded_bytes > MAX_BYTES {
76 return Err(invalid("Bundle exceeds 8 MiB expanded content"));
77 }
78 if bytes.len() as u64 != entry.bytes || AssetDigest::of(bytes) != entry.digest {
79 return Err(RevisionBundleError::Integrity);
80 }
81 used_assets.insert(entry.digest.clone());
82 files.0.insert(
83 path.clone(),
84 AssetFile {
85 bytes: bytes.clone(),
86 media_type: entry.media_type.clone(),
87 executable: entry.executable,
88 },
89 );
90 }
91 let rebuilt = RevisionManifest::from_files(
92 manifest.snapshot_id.clone(),
93 manifest.parent_id.clone(),
94 &files,
95 manifest.dependencies.clone(),
96 )?;
97 if &rebuilt != manifest {
98 return Err(RevisionBundleError::Integrity);
99 }
100 }
101 if used_assets.len() != self.assets.len() {
102 return Err(RevisionBundleError::Integrity);
103 }
104 let mut visited = BTreeSet::new();
105 self.visit(&self.root, &mut BTreeSet::new(), &mut visited)?;
106 if visited.len() != self.revisions.len() {
107 return Err(RevisionBundleError::Integrity);
108 }
109 Ok(())
110 }
111
112 fn visit(
113 &self,
114 id: &ResourceRevisionId,
115 active: &mut BTreeSet<ResourceRevisionId>,
116 visited: &mut BTreeSet<ResourceRevisionId>,
117 ) -> Result<()> {
118 if visited.contains(id) {
119 return Ok(());
120 }
121 if !active.insert(id.clone()) {
122 return Err(invalid("Cyclic revision dependencies"));
123 }
124 let manifest = self
125 .revisions
126 .get(id)
127 .ok_or(RevisionBundleError::Integrity)?;
128 for dependency in manifest.dependencies.values() {
129 let target = self
130 .revisions
131 .get(&dependency.revision_id)
132 .ok_or(RevisionBundleError::Integrity)?;
133 if target.digest()? != dependency.digest {
134 return Err(RevisionBundleError::Integrity);
135 }
136 self.visit(&dependency.revision_id, active, visited)?;
137 }
138 active.remove(id);
139 visited.insert(id.clone());
140 Ok(())
141 }
142
143 pub fn canonical_bytes(&self) -> Result<Vec<u8>> {
144 self.verify()?;
145 Ok(serde_jcs::to_vec(self)?)
146 }
147
148 pub fn digest(&self) -> Result<AssetDigest> {
149 Ok(AssetDigest::of(&self.canonical_bytes()?))
150 }
151
152 pub fn revision_files(&self, id: &ResourceRevisionId) -> Result<RevisionFiles> {
153 self.verify()?;
154 let manifest = self
155 .revisions
156 .get(id)
157 .ok_or_else(|| RevisionBundleError::MissingRevision(id.clone()))?;
158 Ok(RevisionFiles(
159 manifest
160 .files
161 .iter()
162 .map(|(path, entry)| {
163 (
164 path.clone(),
165 AssetFile {
166 bytes: self.assets[&entry.digest].clone(),
167 media_type: entry.media_type.clone(),
168 executable: entry.executable,
169 },
170 )
171 })
172 .collect(),
173 ))
174 }
175}