soma_provider_adapters/python/
materializer.rs

1//! Atomic materialization and frozen startup for planned Python environments.
2
3use std::{
4    ffi::{OsStr, OsString},
5    fs, io,
6    path::{Path, PathBuf},
7    process::Command,
8    sync::atomic::{AtomicU64, Ordering},
9};
10
11use serde::{Deserialize, Serialize};
12use sha2::{Digest, Sha256};
13use thiserror::Error;
14
15use super::environment::{
16    Pep723Metadata, PythonEnvironmentPlan, PythonRuntimeFingerprint, PythonWheelTag,
17};
18
19pub(super) const READY_FILE: &str = "soma-environment.json";
20pub(super) const READY_SCHEMA_VERSION: u32 = 3;
21static STAGING_SEQUENCE: AtomicU64 = AtomicU64::new(0);
22
23#[path = "materializer_repair.rs"]
24mod repair;
25pub use repair::{
26    PythonEnvironmentRepairError, PythonEnvironmentRepairOutcome, PythonEnvironmentRepairReport,
27};
28#[path = "materializer_update.rs"]
29mod update;
30pub use update::{
31    PythonEnvironmentUpdateError, PythonEnvironmentUpdateOutcome, PythonEnvironmentUpdateReport,
32    PythonEnvironmentUpdateRequest,
33};
34
35#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
36pub struct PreparedPythonEnvironment {
37    pub key: String,
38    pub directory: PathBuf,
39    pub python: PathBuf,
40    pub lockfile: PathBuf,
41    pub plan_version: u32,
42    pub dependency_count: usize,
43    pub runtime: PythonRuntimeFingerprint,
44    pub sdk_wheel_tag: PythonWheelTag,
45    pub sdk_wheel_sha256: String,
46    pub uv_version: String,
47    pub lock_sha256: String,
48    pub provider_source_sha256: Option<String>,
49    pub input_plan_key: Option<String>,
50}
51
52impl PreparedPythonEnvironment {
53    pub fn environment_plan(&self) -> PythonEnvironmentPlan {
54        PythonEnvironmentPlan {
55            key: self.key.clone(),
56            directory: self.directory.clone(),
57            plan_version: self.plan_version,
58            dependency_count: self.dependency_count,
59            runtime: self.runtime.clone(),
60            sdk_wheel_tag: self.sdk_wheel_tag.clone(),
61            sdk_wheel_sha256: self.sdk_wheel_sha256.clone(),
62            uv_version: self.uv_version.clone(),
63        }
64    }
65}
66
67#[derive(Debug, Clone, Copy)]
68pub struct PythonMaterializationRequest<'a> {
69    pub metadata: Option<&'a Pep723Metadata>,
70    pub python_executable: &'a Path,
71    pub sdk_wheel: &'a Path,
72    pub offline: bool,
73}
74
75#[derive(Debug, Error)]
76pub enum PythonMaterializationError {
77    #[error("Python environment is not cached for offline startup: {0}")]
78    OfflineCacheMiss(String),
79    #[error("SDK wheel digest does not match the environment plan")]
80    SdkDigestMismatch,
81    #[error("Python environment cache entry is incomplete: {0}")]
82    IncompleteCache(String),
83    #[error("uv command failed during {operation}: {message}")]
84    Uv {
85        operation: &'static str,
86        message: String,
87    },
88    #[error("Python environment I/O failed: {0}")]
89    Io(#[from] io::Error),
90    #[error("Python environment marker is invalid: {0}")]
91    InvalidMarker(String),
92}
93
94pub trait UvRunner: Send + Sync {
95    fn run(&self, program: &Path, args: &[OsString], current_dir: &Path) -> Result<(), String>;
96
97    fn verify_identity(&self, _program: &Path, _expected_version: &str) -> Result<(), String> {
98        Ok(())
99    }
100
101    fn verify_python(
102        &self,
103        _program: &Path,
104        _expected: &PythonRuntimeFingerprint,
105    ) -> Result<(), String> {
106        Ok(())
107    }
108}
109
110#[derive(Debug, Default, Clone, Copy)]
111pub struct SystemUvRunner;
112
113impl UvRunner for SystemUvRunner {
114    fn verify_identity(&self, program: &Path, expected_version: &str) -> Result<(), String> {
115        let output = Command::new(program)
116            .arg("--version")
117            .output()
118            .map_err(|error| error.to_string())?;
119        if !output.status.success() {
120            return Err(String::from_utf8_lossy(&output.stderr).trim().to_owned());
121        }
122        let output = String::from_utf8(output.stdout).map_err(|error| error.to_string())?;
123        let actual = output
124            .trim()
125            .strip_prefix("uv ")
126            .unwrap_or(output.trim())
127            .split_whitespace()
128            .next()
129            .unwrap_or_default();
130        if actual != expected_version {
131            return Err(format!(
132                "uv identity mismatch: expected {expected_version:?}, got {actual:?}"
133            ));
134        }
135        Ok(())
136    }
137
138    fn verify_python(
139        &self,
140        program: &Path,
141        expected: &PythonRuntimeFingerprint,
142    ) -> Result<(), String> {
143        let probe = concat!(
144            "import platform,sys\n",
145            "system={'Darwin':'macos','Windows':'windows'}.get(platform.system(),",
146            "platform.system().lower())\n",
147            "machine={'AMD64':'x86_64'}.get(platform.machine(),",
148            "platform.machine().lower())\n",
149            "print(sys.implementation.name+'\\t'+platform.python_version()+'\\t'+system+'-'+machine)\n"
150        );
151        let output = Command::new(program)
152            .args(["-I", "-c", probe])
153            .output()
154            .map_err(|error| error.to_string())?;
155        if !output.status.success() {
156            return Err(String::from_utf8_lossy(&output.stderr).trim().to_owned());
157        }
158        let actual = String::from_utf8(output.stdout).map_err(|error| error.to_string())?;
159        let expected = format!(
160            "{}\t{}\t{}",
161            expected.implementation, expected.version, expected.platform
162        );
163        if actual.trim() != expected {
164            return Err(format!(
165                "Python runtime identity mismatch: expected {expected:?}, got {:?}",
166                actual.trim()
167            ));
168        }
169        Ok(())
170    }
171
172    fn run(&self, program: &Path, args: &[OsString], current_dir: &Path) -> Result<(), String> {
173        let output = Command::new(program)
174            .args(args)
175            .current_dir(current_dir)
176            .env("UV_NO_PROGRESS", "1")
177            .output()
178            .map_err(|error| error.to_string())?;
179        if output.status.success() {
180            return Ok(());
181        }
182        let stderr = String::from_utf8_lossy(&output.stderr);
183        Err(stderr.trim().to_owned())
184    }
185}
186
187#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
188pub(super) struct ReadyMarker {
189    pub(super) schema_version: u32,
190    pub(super) environment_key: String,
191    pub(super) plan_version: u32,
192    pub(super) dependency_count: usize,
193    pub(super) runtime: PythonRuntimeFingerprint,
194    pub(super) sdk_wheel_tag: PythonWheelTag,
195    pub(super) sdk_wheel_sha256: String,
196    pub(super) uv_version: String,
197    pub(super) lock_sha256: String,
198    #[serde(default)]
199    pub(super) provider_source_sha256: Option<String>,
200    #[serde(default)]
201    pub(super) input_plan_key: Option<String>,
202}
203
204pub struct PythonEnvironmentMaterializer<R = SystemUvRunner> {
205    uv_program: PathBuf,
206    runner: R,
207}
208
209impl PythonEnvironmentMaterializer<SystemUvRunner> {
210    pub fn new(uv_program: impl Into<PathBuf>) -> Self {
211        Self {
212            uv_program: uv_program.into(),
213            runner: SystemUvRunner,
214        }
215    }
216}
217
218impl<R: UvRunner> PythonEnvironmentMaterializer<R> {
219    pub fn with_runner(uv_program: impl Into<PathBuf>, runner: R) -> Self {
220        Self {
221            uv_program: uv_program.into(),
222            runner,
223        }
224    }
225
226    pub fn open_frozen(
227        &self,
228        plan: &PythonEnvironmentPlan,
229    ) -> Result<PreparedPythonEnvironment, PythonMaterializationError> {
230        self.open_verified(plan)?
231            .ok_or_else(|| PythonMaterializationError::OfflineCacheMiss(plan.key.clone()))
232    }
233
234    pub(super) fn open_verified(
235        &self,
236        plan: &PythonEnvironmentPlan,
237    ) -> Result<Option<PreparedPythonEnvironment>, PythonMaterializationError> {
238        let environment = open_ready(plan)?;
239        if let Some(environment) = &environment {
240            self.runner
241                .verify_python(&environment.python, &plan.runtime)
242                .map_err(|message| PythonMaterializationError::Uv {
243                    operation: "Python runtime identity verification",
244                    message,
245                })?;
246        }
247        Ok(environment)
248    }
249
250    pub fn validate_prepared(
251        &self,
252        expected: &PreparedPythonEnvironment,
253    ) -> Result<PreparedPythonEnvironment, PythonMaterializationError> {
254        let reopened = self.open_frozen(&expected.environment_plan())?;
255        if reopened != *expected {
256            return Err(PythonMaterializationError::InvalidMarker(
257                "prepared environment identity changed since it was selected".to_owned(),
258            ));
259        }
260        Ok(reopened)
261    }
262
263    pub fn prepare(
264        &self,
265        plan: &PythonEnvironmentPlan,
266        request: PythonMaterializationRequest<'_>,
267    ) -> Result<PreparedPythonEnvironment, PythonMaterializationError> {
268        if let Some(environment) = self.open_verified(plan)? {
269            return Ok(environment);
270        }
271        if request.offline {
272            return Err(PythonMaterializationError::OfflineCacheMiss(
273                plan.key.clone(),
274            ));
275        }
276        self.runner
277            .verify_identity(&self.uv_program, &plan.uv_version)
278            .map_err(|message| PythonMaterializationError::Uv {
279                operation: "identity verification",
280                message,
281            })?;
282        let sdk_wheel_bytes = read_verified_sdk(request.sdk_wheel, &plan.sdk_wheel_sha256)?;
283
284        let parent = plan.directory.parent().ok_or_else(|| {
285            PythonMaterializationError::IncompleteCache("cache plan has no parent".to_owned())
286        })?;
287        fs::create_dir_all(parent)?;
288        let staging = staging_path(&plan.directory);
289        fs::create_dir(&staging)?;
290        let sdk_wheel_name = request.sdk_wheel.file_name().ok_or_else(|| {
291            PythonMaterializationError::IncompleteCache(
292                "configured SDK wheel has no file name".to_owned(),
293            )
294        })?;
295        let staged_sdk_wheel = staging.join(sdk_wheel_name);
296        fs::write(&staged_sdk_wheel, sdk_wheel_bytes)?;
297        let staged_request = PythonMaterializationRequest {
298            metadata: request.metadata,
299            python_executable: request.python_executable,
300            sdk_wheel: &staged_sdk_wheel,
301            offline: request.offline,
302        };
303
304        let result = self.materialize_staging(&staging, plan, staged_request);
305        if let Err(error) = result {
306            let _ = fs::remove_dir_all(&staging);
307            return Err(error);
308        }
309
310        match fs::rename(&staging, &plan.directory) {
311            Ok(()) => self.open_frozen(plan),
312            Err(_) if plan.directory.exists() => {
313                let _ = fs::remove_dir_all(&staging);
314                self.open_frozen(plan)
315            }
316            Err(error) => {
317                let _ = fs::remove_dir_all(&staging);
318                Err(error.into())
319            }
320        }
321    }
322
323    fn materialize_staging(
324        &self,
325        staging: &Path,
326        plan: &PythonEnvironmentPlan,
327        request: PythonMaterializationRequest<'_>,
328    ) -> Result<(), PythonMaterializationError> {
329        fs::write(
330            staging.join("pyproject.toml"),
331            render_project(request.metadata),
332        )?;
333        self.uv(
334            "lock",
335            staging,
336            [
337                OsString::from("lock"),
338                OsString::from("--project"),
339                OsString::from("."),
340                OsString::from("--python"),
341                request.python_executable.as_os_str().to_owned(),
342            ],
343        )?;
344        self.uv(
345            "sync",
346            staging,
347            [
348                OsString::from("sync"),
349                OsString::from("--project"),
350                OsString::from("."),
351                OsString::from("--locked"),
352                OsString::from("--no-install-project"),
353                OsString::from("--python"),
354                request.python_executable.as_os_str().to_owned(),
355            ],
356        )?;
357        let python = environment_python(staging);
358        self.runner
359            .verify_python(&python, &plan.runtime)
360            .map_err(|message| PythonMaterializationError::Uv {
361                operation: "Python runtime identity verification",
362                message,
363            })?;
364        self.uv(
365            "SDK install",
366            staging,
367            [
368                OsString::from("pip"),
369                OsString::from("install"),
370                OsString::from("--python"),
371                python.as_os_str().to_owned(),
372                OsString::from("--offline"),
373                OsString::from("--no-deps"),
374                request.sdk_wheel.as_os_str().to_owned(),
375            ],
376        )?;
377        let lockfile = staging.join("uv.lock");
378        if !python.is_file() || !lockfile.is_file() {
379            return Err(PythonMaterializationError::IncompleteCache(
380                "uv did not create both .venv Python and uv.lock".to_owned(),
381            ));
382        }
383        let lock_sha256 = sha256_hex(&fs::read(&lockfile)?);
384        let marker = ReadyMarker {
385            schema_version: READY_SCHEMA_VERSION,
386            environment_key: plan.key.clone(),
387            plan_version: plan.plan_version,
388            dependency_count: plan.dependency_count,
389            runtime: plan.runtime.clone(),
390            sdk_wheel_tag: plan.sdk_wheel_tag.clone(),
391            sdk_wheel_sha256: plan.sdk_wheel_sha256.clone(),
392            uv_version: plan.uv_version.clone(),
393            lock_sha256,
394            provider_source_sha256: None,
395            input_plan_key: None,
396        };
397        fs::write(
398            staging.join(READY_FILE),
399            serde_json::to_vec_pretty(&marker).unwrap(),
400        )?;
401        Ok(())
402    }
403
404    fn uv<const N: usize>(
405        &self,
406        operation: &'static str,
407        current_dir: &Path,
408        args: [OsString; N],
409    ) -> Result<(), PythonMaterializationError> {
410        self.runner
411            .run(&self.uv_program, &args, current_dir)
412            .map_err(|message| PythonMaterializationError::Uv { operation, message })
413    }
414}
415
416fn open_ready(
417    plan: &PythonEnvironmentPlan,
418) -> Result<Option<PreparedPythonEnvironment>, PythonMaterializationError> {
419    let directory_metadata = match fs::symlink_metadata(&plan.directory) {
420        Ok(metadata) => metadata,
421        Err(error) if error.kind() == io::ErrorKind::NotFound => return Ok(None),
422        Err(error) => return Err(error.into()),
423    };
424    if directory_metadata.file_type().is_symlink() || !directory_metadata.is_dir() {
425        return Err(PythonMaterializationError::IncompleteCache(
426            "cache path is not a real directory".to_owned(),
427        ));
428    }
429    let marker_path = plan.directory.join(READY_FILE);
430    let marker_metadata = match fs::symlink_metadata(&marker_path) {
431        Ok(metadata) => metadata,
432        Err(error) if error.kind() == io::ErrorKind::NotFound => {
433            return Err(PythonMaterializationError::IncompleteCache(
434                "cache directory exists without readiness marker".to_owned(),
435            ));
436        }
437        Err(error) => return Err(error.into()),
438    };
439    if marker_metadata.file_type().is_symlink() || !marker_metadata.is_file() {
440        return Err(PythonMaterializationError::IncompleteCache(
441            "readiness marker is not a regular file".to_owned(),
442        ));
443    }
444    let marker_bytes = fs::read(&marker_path)?;
445    let marker: ReadyMarker = serde_json::from_slice(&marker_bytes)
446        .map_err(|error| PythonMaterializationError::InvalidMarker(error.to_string()))?;
447    if marker.schema_version != READY_SCHEMA_VERSION {
448        return Err(PythonMaterializationError::InvalidMarker(format!(
449            "unsupported readiness schema version {}; expected {READY_SCHEMA_VERSION}",
450            marker.schema_version
451        )));
452    }
453    if !marker_matches_plan(&marker, plan) {
454        return Err(PythonMaterializationError::IncompleteCache(
455            "readiness marker does not match the plan".to_owned(),
456        ));
457    }
458    let python = environment_python(&plan.directory);
459    if !python.is_file() {
460        return Err(PythonMaterializationError::IncompleteCache(
461            "readiness marker exists without .venv Python".to_owned(),
462        ));
463    }
464    let lockfile = plan.directory.join("uv.lock");
465    let lock_metadata = match fs::symlink_metadata(&lockfile) {
466        Ok(metadata) => metadata,
467        Err(error) if error.kind() == io::ErrorKind::NotFound => {
468            return Err(PythonMaterializationError::IncompleteCache(
469                "readiness marker exists without uv.lock".to_owned(),
470            ));
471        }
472        Err(error) => return Err(error.into()),
473    };
474    if lock_metadata.file_type().is_symlink() || !lock_metadata.is_file() {
475        return Err(PythonMaterializationError::IncompleteCache(
476            "uv.lock is not a regular file".to_owned(),
477        ));
478    }
479    let lock_sha256 = sha256_hex(&fs::read(&lockfile)?);
480    if lock_sha256 != marker.lock_sha256 {
481        return Err(PythonMaterializationError::IncompleteCache(
482            "uv.lock digest does not match the readiness marker".to_owned(),
483        ));
484    }
485    Ok(Some(PreparedPythonEnvironment {
486        key: plan.key.clone(),
487        directory: plan.directory.clone(),
488        python,
489        lockfile,
490        plan_version: plan.plan_version,
491        dependency_count: plan.dependency_count,
492        runtime: plan.runtime.clone(),
493        sdk_wheel_tag: plan.sdk_wheel_tag.clone(),
494        sdk_wheel_sha256: plan.sdk_wheel_sha256.clone(),
495        uv_version: plan.uv_version.clone(),
496        lock_sha256,
497        provider_source_sha256: marker.provider_source_sha256,
498        input_plan_key: marker.input_plan_key,
499    }))
500}
501
502fn marker_matches_plan(marker: &ReadyMarker, plan: &PythonEnvironmentPlan) -> bool {
503    marker.environment_key == plan.key
504        && marker.plan_version == plan.plan_version
505        && marker.dependency_count == plan.dependency_count
506        && marker.runtime == plan.runtime
507        && marker.sdk_wheel_tag == plan.sdk_wheel_tag
508        && marker.sdk_wheel_sha256 == plan.sdk_wheel_sha256
509        && marker.uv_version == plan.uv_version
510}
511
512fn read_verified_sdk(path: &Path, expected: &str) -> Result<Vec<u8>, PythonMaterializationError> {
513    let bytes = fs::read(path)?;
514    let actual = sha256_hex(&bytes);
515    if actual != expected.trim().to_ascii_lowercase() {
516        return Err(PythonMaterializationError::SdkDigestMismatch);
517    }
518    Ok(bytes)
519}
520
521fn sha256_hex(bytes: &[u8]) -> String {
522    const HEX: &[u8; 16] = b"0123456789abcdef";
523    let digest = Sha256::digest(bytes);
524    let mut encoded = String::with_capacity(digest.len() * 2);
525    for byte in digest {
526        encoded.push(HEX[(byte >> 4) as usize] as char);
527        encoded.push(HEX[(byte & 0x0f) as usize] as char);
528    }
529    encoded
530}
531
532fn staging_path(target: &Path) -> PathBuf {
533    let sequence = STAGING_SEQUENCE.fetch_add(1, Ordering::Relaxed);
534    let name = target
535        .file_name()
536        .unwrap_or_else(|| OsStr::new("environment"));
537    target.with_file_name(format!(
538        ".{}.tmp-{}-{sequence}",
539        name.to_string_lossy(),
540        std::process::id()
541    ))
542}
543
544fn environment_python(directory: &Path) -> PathBuf {
545    let unix = directory.join(".venv/bin/python");
546    if unix.is_file() {
547        unix
548    } else {
549        directory.join(".venv/Scripts/python.exe")
550    }
551}
552
553fn render_project(metadata: Option<&Pep723Metadata>) -> String {
554    let metadata = metadata.cloned().unwrap_or_default();
555    let mut project =
556        String::from("[project]\nname = \"soma-provider-environment\"\nversion = \"0\"\n");
557    if let Some(requires_python) = metadata.requires_python {
558        project.push_str(&format!(
559            "requires-python = {}\n",
560            toml::Value::String(requires_python)
561        ));
562    }
563    project.push_str("dependencies = [\n");
564    for dependency in metadata.dependencies {
565        project.push_str(&format!("  {},\n", toml::Value::String(dependency)));
566    }
567    project.push_str("]\n");
568    if let Some(uv) = metadata.uv {
569        project.push_str("\n[tool.uv]\n");
570        if let Some(table) = uv.as_table() {
571            for (key, value) in table {
572                project.push_str(&format!("{key} = {value}\n"));
573            }
574        }
575    }
576    project
577}
578
579#[cfg(test)]
580#[path = "materializer_tests.rs"]
581mod tests;
soma_provider_adapters/python/materializer.rs

soma_provider_adapters/python/
materializer.rs
