rightkit_ort/

lib.rs

//! Product-neutral ONNX Runtime dynamic-library resolution, environment and
//! session setup for Right Suite apps.
//!
//! Merged from ScrapeRight `ort_common` (installed-bundle layout, system-DLL
//! hazard) and HeardRight `heardright-onnx-asr` (environment, execution
//! providers, session builder). Product environment variables and data roots
//! stay in app adapters.

/// The pinned `ort` crate, re-exported so apps that run their own inference reach
/// `rightkit_ort::ort::{value::Tensor, session::Session, ...}` through this crate.
/// They never declare a direct `ort` dependency, which keeps the release ownership
/// scanner from seeing one and keeps the whole suite on a single `ort` pin.
///
/// ```no_run
/// use rightkit_ort::ort::{session::Session, value::Tensor};
///
/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
/// // Configure the runtime and environment first (see `configure_runtime`).
/// let mut session = Session::builder()?.commit_from_file("model.onnx")?;
/// let x = Tensor::from_array(([1usize, 4], vec![0.0f32; 4]))?;
/// let _outputs = session.run(rightkit_ort::ort::inputs!["x" => x])?;
/// # Ok(())
/// # }
/// ```
#[cfg(feature = "session")]
pub use ort;

/// `half` (f16/bf16) as used by ort's `half` feature: `rightkit_ort::half::f16`.
#[cfg(feature = "half")]
pub use half;

#[cfg(feature = "session")]
pub mod environment;
#[cfg(feature = "session")]
pub mod probe;
#[cfg(feature = "session")]
pub mod session;

#[cfg(feature = "session")]
pub use environment::{
    cpu_thread_budget, init_environment, init_environment_with_options, shared_pool_active,
    EnvironmentInitOptions, EnvironmentInitReport, EnvironmentOptions, EnvironmentReport,
    EnvironmentStatus, GlobalPool,
};
#[cfg(feature = "session")]
pub use probe::{probe_providers, ProviderDiagnostic, ProviderKind};
#[cfg(feature = "session")]
pub use session::{BuiltSession, ExecutionProvider, SessionError, SessionOptions};

use std::fmt;
use std::path::{Path, PathBuf};
use std::sync::Mutex;

static CONFIGURED_RUNTIME: Mutex<Option<PathBuf>> = Mutex::new(None);

/// Clear inherited `ORT_DYLIB_PATH` in release builds; no-op in debug builds.
/// This keeps shell-exported development paths from steering shipped apps,
/// consistent with the policy that inherited `ORT_DYLIB_PATH` fails closed.
///
/// Call at process startup, before any thread spawns, runtime binding or ORT
/// use. Removing a process environment variable requires exclusive access to
/// the environment; this function does not unload an already-loaded runtime.
pub fn clear_inherited_runtime_for_release() {
    clear_inherited_runtime_with(|| std::env::remove_var("ORT_DYLIB_PATH"));
}

// Injected removal keeps release/debug tests free of process env mutation.
fn clear_inherited_runtime_with(remove_path: impl FnOnce()) {
    if cfg!(not(debug_assertions)) {
        remove_path();
    }
}

/// Read a caller-named development override as an explicit runtime candidate.
/// Unset or empty values return `None`; release builds always return `None`
/// without reading the variable. Path validation remains in runtime resolution.
pub fn developer_override_candidate(env_name: &str) -> Option<RuntimeCandidate> {
    if cfg!(debug_assertions) {
        nonempty_env_path(env_name)
            .map(|path| RuntimeCandidate::new(path, CandidateSource::Explicit))
    } else {
        None
    }
}

/// Return the last path successfully bound by this crate for this process,
/// falling back to a non-empty `ORT_DYLIB_PATH` when no binding is recorded.
/// Read-only: no resolution, file check, environment mutation or ORT loading.
/// Missing-mode bindings are included; this is not proof of a loaded library.
pub fn configured_runtime() -> Option<PathBuf> {
    let recorded = CONFIGURED_RUNTIME
        .lock()
        .unwrap_or_else(|poisoned| poisoned.into_inner())
        .clone();
    recorded.or_else(|| nonempty_env_path("ORT_DYLIB_PATH"))
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CandidateSource {
    Explicit,
    Bundled,
    AppData,
    /// Caller-opted-in `<workspace>/tools/bin/<runtime filename>`.
    LegacyDefault,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RuntimeCandidate {
    pub path: PathBuf,
    pub source: CandidateSource,
}

impl RuntimeCandidate {
    pub fn new(path: impl Into<PathBuf>, source: CandidateSource) -> Self {
        Self {
            path: path.into(),
            source,
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CandidateDiagnostic {
    pub path: PathBuf,
    pub source: CandidateSource,
    pub absolute: bool,
    pub expected_filename: bool,
    pub exists: bool,
    pub is_file: bool,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RuntimeSelection {
    pub path: PathBuf,
    pub source: CandidateSource,
    pub diagnostics: Vec<CandidateDiagnostic>,
}

/// Binding policy; defaults retain strict file validation & conflict refusal.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct RuntimeBindingOptions {
    /// Overwrite ORT_DYLIB_PATH, including an invalid existing value. This does
    /// not replace a library already loaded by ort. Use before any ORT use.
    pub force_replace: bool,
    /// If no regular file resolves, bind the first missing candidate. Existing
    /// directories, invalid filenames & unsafe paths remain errors.
    pub allow_missing: bool,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RuntimeBindingStatus {
    Ready,
    /// Path recorded & bound, but absent. Callers own degraded-mode messaging.
    Missing,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RuntimeBinding {
    pub selection: RuntimeSelection,
    pub status: RuntimeBindingStatus,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum RuntimeError {
    NoCandidates,
    UnsafePath {
        path: PathBuf,
    },
    UnexpectedFilename {
        path: PathBuf,
        expected: &'static str,
    },
    NotFound {
        diagnostics: Vec<CandidateDiagnostic>,
    },
    Canonicalize {
        path: PathBuf,
        message: String,
    },
    AlreadyConfigured {
        configured: PathBuf,
        selected: PathBuf,
    },
    /// The Windows-ML copy in `System32` hangs at session init; apps must
    /// supply their bundled runtime instead.
    SystemRuntimeRejected {
        path: PathBuf,
    },
}

impl fmt::Display for RuntimeError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::NoCandidates => write!(f, "no ONNX Runtime candidates were supplied"),
            Self::UnsafePath { path } => {
                write!(f, "runtime path must be absolute: {}", path.display())
            }
            Self::UnexpectedFilename { path, expected } => {
                write!(f, "runtime path {} must end in {expected}", path.display())
            }
            Self::NotFound { .. } => {
                write!(f, "no supplied ONNX Runtime candidate is a regular file")
            }
            Self::Canonicalize { path, message } => {
                write!(f, "cannot canonicalize {}: {message}", path.display())
            }
            Self::SystemRuntimeRejected { path } => write!(
                f,
                "runtime {} is the OS-provided copy; bundle and supply an app runtime",
                path.display()
            ),
            Self::AlreadyConfigured {
                configured,
                selected,
            } => write!(
                f,
                "ORT_DYLIB_PATH is already {}, refusing replacement with {}",
                configured.display(),
                selected.display()
            ),
        }
    }
}

impl std::error::Error for RuntimeError {}

pub fn runtime_filename() -> &'static str {
    #[cfg(target_os = "windows")]
    {
        "onnxruntime.dll"
    }
    #[cfg(target_os = "macos")]
    {
        "libonnxruntime.dylib"
    }
    #[cfg(all(not(target_os = "windows"), not(target_os = "macos")))]
    {
        "libonnxruntime.so"
    }
}

pub fn inspect_candidate(candidate: &RuntimeCandidate) -> CandidateDiagnostic {
    CandidateDiagnostic {
        path: candidate.path.clone(),
        source: candidate.source,
        absolute: candidate.path.is_absolute(),
        expected_filename: candidate.path.file_name().and_then(|name| name.to_str())
            == Some(runtime_filename()),
        exists: candidate.path.exists(),
        is_file: candidate.path.is_file(),
    }
}

pub fn resolve_runtime(
    candidates: impl IntoIterator<Item = RuntimeCandidate>,
) -> Result<RuntimeSelection, RuntimeError> {
    let candidates: Vec<_> = candidates.into_iter().collect();
    if candidates.is_empty() {
        return Err(RuntimeError::NoCandidates);
    }
    let mut diagnostics = Vec::with_capacity(candidates.len());
    for candidate in candidates {
        let diagnostic = inspect_candidate(&candidate);
        if !diagnostic.absolute {
            return Err(RuntimeError::UnsafePath {
                path: candidate.path,
            });
        }
        if !diagnostic.expected_filename {
            return Err(RuntimeError::UnexpectedFilename {
                path: candidate.path,
                expected: runtime_filename(),
            });
        }
        if is_system_runtime(&candidate.path) {
            return Err(RuntimeError::SystemRuntimeRejected {
                path: candidate.path,
            });
        }
        diagnostics.push(diagnostic.clone());
        if diagnostic.is_file {
            let path =
                candidate
                    .path
                    .canonicalize()
                    .map_err(|error| RuntimeError::Canonicalize {
                        path: candidate.path,
                        message: error.to_string(),
                    })?;
            return Ok(RuntimeSelection {
                path,
                source: candidate.source,
                diagnostics,
            });
        }
    }
    Err(RuntimeError::NotFound { diagnostics })
}

/// True when the path sits under a Windows `System32`/`SysWOW64` directory.
/// The OS copy of `onnxruntime.dll` (Windows ML) hangs at session init, so it
/// is never an acceptable candidate (ScrapeRight `ort_common` finding).
pub fn is_system_runtime(path: &Path) -> bool {
    path.components().any(|component| {
        component
            .as_os_str()
            .to_str()
            .map(|name| {
                name.eq_ignore_ascii_case("system32") || name.eq_ignore_ascii_case("syswow64")
            })
            .unwrap_or(false)
    })
}

/// Resources directory of an installed application bundle, derived from its
/// executable only (no checkout or working-directory fallback):
/// macOS `Foo.app/Contents/MacOS/foo` -> `Contents/Resources`; Windows the
/// executable directory; elsewhere `<exe dir>/resources`.
pub fn installed_resource_dir(executable: &Path) -> Option<PathBuf> {
    if !executable.is_absolute() {
        return None;
    }
    let dir = executable.parent()?.to_path_buf();
    let is = |p: &Path, name: &str| {
        p.file_name()
            .and_then(|n| n.to_str())
            .map(|n| n.eq_ignore_ascii_case(name))
            .unwrap_or(false)
    };
    #[cfg(target_os = "macos")]
    {
        if is(&dir, "MacOS") {
            if let Some(contents) = dir.parent() {
                if is(contents, "Contents") {
                    return Some(contents.join("Resources"));
                }
            }
        }
        if is(&dir, "Resources") {
            return Some(dir);
        }
        Some(dir.join("Resources"))
    }
    #[cfg(target_os = "windows")]
    {
        let _ = is;
        Some(dir)
    }
    #[cfg(not(any(target_os = "macos", target_os = "windows")))]
    {
        if is(&dir, "resources") {
            Some(dir)
        } else {
            Some(dir.join("resources"))
        }
    }
}

/// Bundled-runtime candidate for an installed app: `<resources>/<subdir>/<runtime file>`.
/// Apps pass their own `subdir` (ScrapeRight/HeardRight/CodeRight all use `runtime`).
pub fn installed_runtime_candidate(executable: &Path, subdir: &str) -> Option<RuntimeCandidate> {
    let dir = installed_resource_dir(executable)?;
    Some(RuntimeCandidate::new(
        dir.join(subdir).join(runtime_filename()),
        CandidateSource::Bundled,
    ))
}

/// Explicit legacy candidate; never searched unless the caller supplies it.
/// Relative roots are rejected rather than resolved against process cwd.
pub fn legacy_runtime_candidate(workspace_root: &Path) -> Result<RuntimeCandidate, RuntimeError> {
    if !workspace_root.is_absolute() {
        return Err(RuntimeError::UnsafePath {
            path: workspace_root.to_path_buf(),
        });
    }
    Ok(RuntimeCandidate::new(
        workspace_root
            .join("tools")
            .join("bin")
            .join(runtime_filename()),
        CandidateSource::LegacyDefault,
    ))
}

/// Configure ORT only after resolving a caller-supplied, absolute regular file.
/// Existing configuration is preserved and must canonicalize to the same file.
pub fn configure_runtime(
    candidates: impl IntoIterator<Item = RuntimeCandidate>,
) -> Result<RuntimeSelection, RuntimeError> {
    configure_runtime_with_options(candidates, &RuntimeBindingOptions::default())
        .map(|binding| binding.selection)
}

/// Bind explicit candidates under caller-selected policy. Available regular
/// files win in caller order; missing mode falls back to the first absent path
/// only when no file resolves. Missing paths retain their supplied absolute
/// spelling because they cannot be canonicalized.
///
/// Call before starting worker threads or using any ORT API: this mutates a
/// process environment variable, not ort's already-loaded library handle.
pub fn configure_runtime_with_options(
    candidates: impl IntoIterator<Item = RuntimeCandidate>,
    options: &RuntimeBindingOptions,
) -> Result<RuntimeBinding, RuntimeError> {
    let binding = configure_runtime_with_env(
        candidates,
        options,
        nonempty_env_path("ORT_DYLIB_PATH"),
        |path| std::env::set_var("ORT_DYLIB_PATH", path),
    )?;
    *CONFIGURED_RUNTIME
        .lock()
        .unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(binding.selection.path.clone());
    Ok(binding)
}

// Injected environment writer avoids process-global mutation in unit tests.
fn configure_runtime_with_env(
    candidates: impl IntoIterator<Item = RuntimeCandidate>,
    options: &RuntimeBindingOptions,
    configured: Option<PathBuf>,
    set_path: impl FnOnce(&Path),
) -> Result<RuntimeBinding, RuntimeError> {
    let binding = match resolve_runtime(candidates) {
        Ok(selection) => RuntimeBinding {
            selection,
            status: RuntimeBindingStatus::Ready,
        },
        Err(RuntimeError::NotFound { diagnostics }) if options.allow_missing => {
            let Some(missing) = diagnostics.iter().find(|candidate| !candidate.exists) else {
                return Err(RuntimeError::NotFound { diagnostics });
            };
            RuntimeBinding {
                selection: RuntimeSelection {
                    path: missing.path.clone(),
                    source: missing.source,
                    diagnostics,
                },
                status: RuntimeBindingStatus::Missing,
            }
        }
        Err(error) => return Err(error),
    };
    if let Some(configured) = configured.filter(|_| !options.force_replace) {
        if options.allow_missing && configured == binding.selection.path {
            return Ok(binding);
        }
        let configured = canonical_if_file(&configured)?;
        if configured != binding.selection.path {
            return Err(RuntimeError::AlreadyConfigured {
                configured,
                selected: binding.selection.path,
            });
        }
        return Ok(binding);
    }
    set_path(&binding.selection.path);
    Ok(binding)
}

fn nonempty_env_path(name: &str) -> Option<PathBuf> {
    std::env::var_os(name)
        .filter(|value| !value.is_empty())
        .map(PathBuf::from)
}

fn canonical_if_file(path: &Path) -> Result<PathBuf, RuntimeError> {
    if !path.is_absolute() {
        return Err(RuntimeError::UnsafePath {
            path: path.to_path_buf(),
        });
    }
    path.canonicalize()
        .map_err(|error| RuntimeError::Canonicalize {
            path: path.to_path_buf(),
            message: error.to_string(),
        })
}

#[cfg(test)]
mod binding_tests {
    use super::*;
    use std::sync::atomic::{AtomicUsize, Ordering};

    #[cfg(debug_assertions)]
    #[test]
    fn clearing_inherited_runtime_is_noop_in_debug() {
        clear_inherited_runtime_with(|| panic!("debug must preserve inherited runtime"));
    }

    #[cfg(not(debug_assertions))]
    #[test]
    fn clearing_inherited_runtime_removes_it_in_release() {
        let mut inherited = Some(PathBuf::from("shell-exported-dev-runtime"));
        clear_inherited_runtime_with(|| inherited = None);
        assert_eq!(inherited, None);
    }

    struct Fixture(PathBuf);

    impl Fixture {
        fn new() -> Self {
            static NEXT: AtomicUsize = AtomicUsize::new(0);
            let root = std::env::temp_dir().join(format!(
                "rightkit-ort-binding-{}-{}",
                std::process::id(),
                NEXT.fetch_add(1, Ordering::Relaxed)
            ));
            std::fs::create_dir_all(&root).unwrap();
            Self(root.canonicalize().unwrap())
        }

        fn candidate(&self, directory: &str, exists: bool) -> RuntimeCandidate {
            let path = self.0.join(directory).join(runtime_filename());
            if exists {
                std::fs::create_dir_all(path.parent().unwrap()).unwrap();
                std::fs::write(&path, b"not a real ORT binary").unwrap();
            }
            RuntimeCandidate::new(path, CandidateSource::Explicit)
        }
    }

    impl Drop for Fixture {
        fn drop(&mut self) {
            let _ = std::fs::remove_dir_all(&self.0);
        }
    }

    #[test]
    fn binding_defaults_preserve_conflicts_and_require_files() {
        let options = RuntimeBindingOptions::default();
        assert!(!options.force_replace);
        assert!(!options.allow_missing);
        let fixture = Fixture::new();
        let first = fixture.candidate("first", true);
        let second = fixture.candidate("second", true);
        let error = configure_runtime_with_env([second], &options, Some(first.path), |_| {
            panic!("must not overwrite")
        })
        .unwrap_err();
        assert!(matches!(error, RuntimeError::AlreadyConfigured { .. }));
        let error = configure_runtime_with_env(
            [fixture.candidate("missing", false)],
            &options,
            None,
            |_| panic!("must not bind missing path by default"),
        )
        .unwrap_err();
        assert!(matches!(error, RuntimeError::NotFound { .. }));
    }

    #[test]
    fn forced_replacement_overwrites_existing_or_invalid_configuration() {
        let fixture = Fixture::new();
        let first = fixture.candidate("first", true);
        let selected = fixture.candidate("second", true);
        for configured in [first.path, PathBuf::from("invalid-relative-path")] {
            let mut written = None;
            let report = configure_runtime_with_env(
                [selected.clone()],
                &RuntimeBindingOptions {
                    force_replace: true,
                    ..Default::default()
                },
                Some(configured),
                |path| written = Some(path.to_path_buf()),
            )
            .unwrap();
            let expected = selected.path.canonicalize().unwrap();
            assert_eq!(written, Some(expected.clone()));
            assert_eq!(report.selection.path, expected);
            assert_eq!(report.status, RuntimeBindingStatus::Ready);
        }
    }

    #[test]
    fn missing_mode_binds_records_and_reuses_first_absent_candidate() {
        let fixture = Fixture::new();
        let first = fixture.candidate("missing-first", false);
        let second = fixture.candidate("missing-second", false);
        let options = RuntimeBindingOptions {
            allow_missing: true,
            ..Default::default()
        };
        let mut written = None;
        let report = configure_runtime_with_env([first.clone(), second], &options, None, |path| {
            written = Some(path.to_path_buf())
        })
        .unwrap();
        assert_eq!(report.status, RuntimeBindingStatus::Missing);
        assert_eq!(report.selection.path, first.path);
        assert_eq!(written, Some(first.path.clone()));
        assert_eq!(report.selection.diagnostics.len(), 2);
        assert!(report
            .selection
            .diagnostics
            .iter()
            .all(|d| !d.exists && !d.is_file));
        let repeated =
            configure_runtime_with_env([first.clone()], &options, Some(first.path), |_| {
                panic!("matching binding must be preserved")
            })
            .unwrap();
        assert_eq!(repeated.status, RuntimeBindingStatus::Missing);
    }

    #[test]
    fn missing_mode_prefers_available_files_and_never_binds_directories() {
        let fixture = Fixture::new();
        let available = fixture.candidate("available", true);
        let options = RuntimeBindingOptions {
            allow_missing: true,
            ..Default::default()
        };
        let report = configure_runtime_with_env(
            [fixture.candidate("absent", false), available.clone()],
            &options,
            None,
            |_| {},
        )
        .unwrap();
        assert_eq!(report.status, RuntimeBindingStatus::Ready);
        assert_eq!(
            report.selection.path,
            available.path.canonicalize().unwrap()
        );
        let directory = fixture.candidate("directory", false);
        std::fs::create_dir_all(&directory.path).unwrap();
        let error = configure_runtime_with_env([directory], &options, None, |_| {
            panic!("directory must not be bound")
        })
        .unwrap_err();
        assert!(matches!(error, RuntimeError::NotFound { .. }));
    }

    #[test]
    fn forced_missing_binding_overwrites_existing_configuration() {
        let fixture = Fixture::new();
        let previous = fixture.candidate("previous", true);
        let missing = fixture.candidate("missing", false);
        let mut written = None;
        let report = configure_runtime_with_env(
            [missing.clone()],
            &RuntimeBindingOptions {
                force_replace: true,
                allow_missing: true,
            },
            Some(previous.path),
            |path| written = Some(path.to_path_buf()),
        )
        .unwrap();
        assert_eq!(report.status, RuntimeBindingStatus::Missing);
        assert_eq!(written, Some(missing.path));
    }

    #[test]
    fn permissive_options_still_reject_unsafe_paths_and_system_runtime() {
        let fixture = Fixture::new();
        let options = RuntimeBindingOptions {
            force_replace: true,
            allow_missing: true,
        };
        for (candidate, expected) in [
            (
                RuntimeCandidate::new(runtime_filename(), CandidateSource::Explicit),
                "relative",
            ),
            (
                RuntimeCandidate::new(fixture.0.join("wrong.bin"), CandidateSource::Explicit),
                "filename",
            ),
            (fixture.candidate("System32", false), "system"),
        ] {
            let error = configure_runtime_with_env([candidate], &options, None, |_| {
                panic!("unsafe candidate must not be bound")
            })
            .unwrap_err();
            assert!(match expected {
                "relative" => matches!(error, RuntimeError::UnsafePath { .. }),
                "filename" => matches!(error, RuntimeError::UnexpectedFilename { .. }),
                _ => matches!(error, RuntimeError::SystemRuntimeRejected { .. }),
            });
        }
    }
}