kernel/install/
ollama_pull.rs

1//! Folding an Ollama `/api/pull` ndjson stream into install progress. Each line is
2//! a status object; digest lines carry per-layer byte totals that aggregate into
3//! overall progress, and a `success` line ends the pull.
4
5use std::collections::BTreeMap;
6
7use serde::Deserialize;
8
9use crate::install::bytes::saturating_sum;
10use crate::install::error::InstallError;
11use crate::install::event::InstallProgress;
12
13/// What folding one pull line produced.
14#[derive(Debug, Clone, PartialEq, Eq, Hash)]
15pub enum Outcome {
16    /// The line carried nothing new (unparseable, blank, or a repeated status).
17    Ignored,
18    /// A new status line.
19    Status(String),
20    /// Updated aggregate download progress.
21    Progress(InstallProgress),
22    /// The pull completed.
23    Success,
24}
25
26/// Accumulates per-layer byte counts across pull lines into overall progress.
27#[derive(Debug, Default)]
28pub struct Aggregator {
29    totals: BTreeMap<String, i64>,
30    completed: BTreeMap<String, i64>,
31    last_status: Option<String>,
32}
33
34impl Aggregator {
35    /// A fresh aggregator.
36    pub fn new() -> Self {
37        Self::default()
38    }
39
40    /// Fold one ndjson line, updating internal totals. Returns what to emit, or a
41    /// [`InstallError::TransferFailed`] if the line carried an error.
42    pub fn fold(&mut self, line: &str) -> Result<Outcome, InstallError> {
43        let Ok(decoded) = serde_json::from_str::<Line>(line) else {
44            return Ok(Outcome::Ignored);
45        };
46        if let Some(error) = decoded.error.as_deref().filter(|error| !error.is_empty()) {
47            return Err(InstallError::TransferFailed(format!("ollama: {error}")));
48        }
49        let Some(status) = decoded
50            .status
51            .as_deref()
52            .filter(|status| !status.is_empty())
53        else {
54            return Ok(Outcome::Ignored);
55        };
56        if status == "success" {
57            return Ok(Outcome::Success);
58        }
59        if let (Some(digest), Some(total)) = (decoded.digest.as_ref(), decoded.total) {
60            self.totals.insert(digest.clone(), total);
61            // Keep the last known completed count for this layer when a line omits it.
62            let completed = decoded
63                .completed
64                .or_else(|| self.completed.get(digest).copied())
65                .unwrap_or(0);
66            self.completed.insert(digest.clone(), completed);
67            return Ok(Outcome::Progress(self.aggregate()));
68        }
69        if self.last_status.as_deref() == Some(status) {
70            return Ok(Outcome::Ignored);
71        }
72        self.last_status = Some(status.to_owned());
73        Ok(Outcome::Status(status.to_owned()))
74    }
75
76    fn aggregate(&self) -> InstallProgress {
77        InstallProgress {
78            bytes_downloaded: saturating_sum(self.completed.values().copied()),
79            total_bytes: Some(saturating_sum(self.totals.values().copied())),
80            total_is_partial: true,
81            current_file: None,
82        }
83    }
84}
85
86#[derive(Deserialize)]
87struct Line {
88    status: Option<String>,
89    digest: Option<String>,
90    total: Option<i64>,
91    completed: Option<i64>,
92    error: Option<String>,
93}
kernel/install/ollama_pull.rs

kernel/install/
ollama_pull.rs
