1use crate::readiness::{
9 LocalFileIntegrity, LocalFileStatus, ModelProvenance, ModelReadinessReport,
10};
11use crate::trust::{DEFAULT_TRUSTED_MODEL, TrustedModelManifest, trust_root_binding};
12use std::path::PathBuf;
13
14pub const DEFAULT_MODEL_DOWNLOAD_CONCURRENCY: usize = 2;
18pub const MAX_MODEL_DOWNLOAD_CONCURRENCY: usize = 2;
19
20#[derive(Debug, Clone, PartialEq, Eq)]
24pub enum DownloadAction {
25 Skip,
27 Download,
29 Replace,
31 Retry,
33}
34
35impl DownloadAction {
36 pub fn requires_download(&self) -> bool {
38 !matches!(self, DownloadAction::Skip)
39 }
40}
41
42#[derive(Debug, Clone)]
46pub struct ModelFilePlan {
47 pub logical_name: &'static str,
48 pub relative_path: &'static str,
49 pub remote_url: &'static str,
50 pub expected_sha256: &'static str,
51 pub exact_size_bytes: u64,
52 pub max_transfer_bytes: u64,
53 pub local_status: LocalFileStatus,
54 pub action: DownloadAction,
55 pub temp_path_suffix: &'static str,
57}
58
59impl ModelFilePlan {
60 pub fn temp_path(&self, model_dir: &std::path::Path) -> PathBuf {
62 model_dir.join(format!("{}.part", self.relative_path))
63 }
64
65 pub fn final_path(&self, model_dir: &std::path::Path) -> PathBuf {
67 model_dir.join(self.relative_path)
68 }
69}
70
71#[derive(Debug, Clone)]
75pub struct DownloadPlan {
76 pub manifest_id: &'static str,
78 pub max_concurrent: usize,
80 pub files: Vec<ModelFilePlan>,
81}
82
83impl DownloadPlan {
84 pub fn files_to_download(&self) -> Vec<&ModelFilePlan> {
86 self.files
87 .iter()
88 .filter(|f| f.action.requires_download())
89 .collect()
90 }
91
92 pub fn files_to_skip(&self) -> Vec<&ModelFilePlan> {
94 self.files
95 .iter()
96 .filter(|f| f.action == DownloadAction::Skip)
97 .collect()
98 }
99
100 pub fn has_work(&self) -> bool {
102 self.files.iter().any(|f| f.action.requires_download())
103 }
104}
105
106#[derive(Debug, Clone, Copy, PartialEq, Eq)]
109pub enum DownloadPlanError {
110 UserSuppliedFolder,
111 TrustRootMismatch,
112 ReadinessManifestMismatch,
113 IncoherentReadiness,
114 CannotCheckFiles,
115 InvalidTrustedManifest,
116}
117
118pub fn build_download_plan(
124 report: &ModelReadinessReport,
125) -> Result<DownloadPlan, DownloadPlanError> {
126 build_download_plan_for_manifest(report, &DEFAULT_TRUSTED_MODEL)
127}
128
129#[cfg(test)]
131pub(crate) fn build_download_plan_against(
132 report: &ModelReadinessReport,
133 manifest: &'static TrustedModelManifest,
134) -> Result<DownloadPlan, DownloadPlanError> {
135 build_download_plan_for_manifest(report, manifest)
136}
137
138fn build_download_plan_for_manifest(
139 report: &ModelReadinessReport,
140 manifest: &'static TrustedModelManifest,
141) -> Result<DownloadPlan, DownloadPlanError> {
142 manifest
143 .validate()
144 .map_err(|_| DownloadPlanError::InvalidTrustedManifest)?;
145 if report.provenance() != ModelProvenance::AppManaged {
146 return Err(DownloadPlanError::UserSuppliedFolder);
147 }
148 if !report.matches_trust_root(trust_root_binding(manifest)) {
149 return Err(DownloadPlanError::TrustRootMismatch);
150 }
151 if report.files().len() != manifest.files.len()
152 || report.files().iter().any(|readiness| {
153 manifest
154 .file_by_path(readiness.relative_path())
155 .is_none_or(|trusted| trusted.logical_name != readiness.logical_name())
156 })
157 || manifest.files.iter().any(|trusted| {
158 report
159 .files()
160 .iter()
161 .filter(|readiness| {
162 readiness.relative_path() == trusted.relative_path
163 && readiness.logical_name() == trusted.logical_name
164 })
165 .count()
166 != 1
167 })
168 {
169 return Err(DownloadPlanError::ReadinessManifestMismatch);
170 }
171
172 let files = report
173 .files()
174 .iter()
175 .map(|readiness| -> Result<ModelFilePlan, DownloadPlanError> {
176 let trusted = manifest
177 .file_by_path(readiness.relative_path())
178 .expect("readiness file set was validated against the manifest");
179 let action = plan_action(readiness.status(), readiness.integrity())?;
180 Ok(ModelFilePlan {
181 logical_name: trusted.logical_name,
182 relative_path: trusted.relative_path,
183 remote_url: trusted.url,
184 expected_sha256: trusted.sha256,
185 exact_size_bytes: trusted.exact_size_bytes,
186 max_transfer_bytes: trusted.max_transfer_bytes,
187 local_status: readiness.status(),
188 action,
189 temp_path_suffix: ".part",
190 })
191 })
192 .collect::<Result<Vec<_>, _>>()?;
193
194 Ok(DownloadPlan {
195 manifest_id: manifest.manifest_id,
196 max_concurrent: DEFAULT_MODEL_DOWNLOAD_CONCURRENCY,
197 files,
198 })
199}
200
201fn plan_action(
202 status: LocalFileStatus,
203 integrity: LocalFileIntegrity,
204) -> Result<DownloadAction, DownloadPlanError> {
205 match (status, integrity) {
206 (LocalFileStatus::Ready, LocalFileIntegrity::TrustedDigest) => Ok(DownloadAction::Skip),
207 (LocalFileStatus::Missing, LocalFileIntegrity::NotAvailable) => {
208 Ok(DownloadAction::Download)
209 }
210 (LocalFileStatus::Partial, LocalFileIntegrity::NotAvailable) => Ok(DownloadAction::Retry),
211 (LocalFileStatus::Invalid, LocalFileIntegrity::Mismatch) => Ok(DownloadAction::Replace),
212 (LocalFileStatus::CannotCheck, _) => Err(DownloadPlanError::CannotCheckFiles),
213 _ => Err(DownloadPlanError::IncoherentReadiness),
214 }
215}
216
217#[derive(Debug, Clone, PartialEq, Eq)]
221pub enum FileDownloadStatus {
222 Pending,
223 Checking,
224 Downloading,
225 Validating,
226 Complete,
227 Failed,
228 Skipped,
229}
230
231#[derive(Debug, Clone)]
233pub struct FileDownloadProgress {
234 pub relative_path: &'static str,
235 pub downloaded_bytes: u64,
236 pub total_bytes: Option<u64>,
237 pub status: FileDownloadStatus,
238}
239
240#[derive(Debug, Clone)]
242pub enum OverallDownloadProgress {
243 Known {
244 downloaded_bytes: u64,
245 total_bytes: u64,
246 },
247 Step {
248 completed_files: usize,
249 total_files: usize,
250 },
251 Indeterminate,
252}
253
254#[derive(Debug, Clone)]
256pub enum FriendlyDownloadProblem {
257 NetworkUnavailable,
258 ServerBusy,
259 NotEnoughSpace,
260 CannotWriteFiles,
261 CannotCheckFiles,
262 ValidationFailed,
263 Unexpected,
264}
265
266impl FriendlyDownloadProblem {
267 pub fn user_message(&self) -> &'static str {
269 match self {
270 FriendlyDownloadProblem::NetworkUnavailable => {
271 "Download did not finish. Please check your connection and try again."
272 }
273 FriendlyDownloadProblem::ServerBusy => {
274 "The download is taking longer than expected. Please try again later."
275 }
276 FriendlyDownloadProblem::NotEnoughSpace => {
277 "More space is needed to finish the download."
278 }
279 FriendlyDownloadProblem::CannotWriteFiles => {
280 "orbok could not save the search helper files here. Please choose another location or check folder permissions."
281 }
282 FriendlyDownloadProblem::CannotCheckFiles => {
283 "orbok could not check the search helper files. Please choose the folder again."
284 }
285 FriendlyDownloadProblem::ValidationFailed => {
286 "Some downloaded files could not be used. orbok can download them again."
287 }
288 FriendlyDownloadProblem::Unexpected => "Download did not finish. Please try again.",
289 }
290 }
291}