tftio-clanker 0.2.0

//! Fail-closed prompter composition and harness-specific prompt injection.
//!
//! Any composition or injection failure aborts the launch; `--no-prompt` is
//! the only way to launch without a composed prompt.

use std::collections::{BTreeMap, BTreeSet};
use std::ffi::OsString;
use std::fs::{self, OpenOptions};
use std::io::Write;
use std::os::unix::fs::OpenOptionsExt;
use std::path::{Path, PathBuf};
use std::time::{Duration, SystemTime, UNIX_EPOCH};

use crate::config::{Config, HarnessConfig, InjectionConfig, ModelFamily};
use crate::launch::expand_path;

const FILE_PLACEHOLDER: &str = "{file}";
const TEXT_PLACEHOLDER: &str = "{text}";

/// Prompt material and process mutations produced for one launch.
#[derive(Debug, Default)]
pub struct PromptInjection {
    /// Rendered prompt text, when composition and injection succeeded.
    pub prompt: Option<String>,
    /// Arguments prepended before model and harness arguments.
    pub arguments: Vec<OsString>,
    /// Environment applied for env-file injection.
    pub environment: BTreeMap<OsString, OsString>,
    /// Prompt cache path used by file-based injection.
    pub cache_path: Option<PathBuf>,
}

/// Failures composing or injecting the launch prompt.
#[derive(Debug, thiserror::Error)]
pub enum PromptError {
    /// A configured prompter path is invalid.
    #[error("prompter path: {0}")]
    Path(String),
    /// The prompter bundle could not be located or loaded.
    #[error("prompter unavailable: {0}")]
    Unavailable(String),
    /// Requested profiles are absent from the prompter bundle.
    #[error("unknown prompter profile(s): {}", profiles.join(", "))]
    UnknownProfiles {
        /// Every requested profile missing from the bundle.
        profiles: Vec<String>,
    },
    /// Prompt rendering failed.
    #[error("prompt rendering failed: {0}")]
    Render(String),
    /// Rendered prompt bytes are not valid UTF-8.
    #[error("rendered prompt is not valid UTF-8: {0}")]
    NonUtf8(#[from] std::string::FromUtf8Error),
    /// Prompt cache directory could not be created.
    #[error("failed to create prompt cache {path}: {source}")]
    CacheDirectory {
        /// Cache directory path.
        path: PathBuf,
        /// Underlying I/O error.
        source: std::io::Error,
    },
    /// Prompt file could not be created or written.
    #[error("failed to write prompt file {path}: {source}")]
    CacheFile {
        /// Prompt file path.
        path: PathBuf,
        /// Underlying I/O error.
        source: std::io::Error,
    },
}

/// Compose configured profiles and build the harness injection mutation.
///
/// # Errors
/// Returns [`PromptError`] when the bundle is unavailable, any requested
/// profile is unknown, rendering fails, or the prompt file cannot be written.
pub fn compose_and_inject(
    profiles: &[String],
    family: &ModelFamily,
    harness: &HarnessConfig,
    config: &Config,
    home: &Path,
    persist_file: bool,
) -> Result<PromptInjection, PromptError> {
    let bundle = expand_path(&config.defaults.prompter_bundle, home)
        .map_err(|error| PromptError::Path(error.to_string()))?;
    let bundle_config = bundle.join("config.toml");
    let available = prompter::available_profiles(Some(&bundle_config))
        .map_err(|error| PromptError::Unavailable(error.to_string()))?
        .into_iter()
        .collect::<BTreeSet<_>>();

    let unknown: Vec<String> = profiles
        .iter()
        .filter(|profile| !available.contains(*profile))
        .cloned()
        .collect();
    if !unknown.is_empty() {
        return Err(PromptError::UnknownProfiles { profiles: unknown });
    }

    let mut injection = PromptInjection::default();
    if profiles.is_empty() {
        return Ok(injection);
    }

    let prompt = prompter::render_to_vec(
        profiles,
        Some(family.as_prompter_family()),
        Some(&bundle_config),
    )
    .map_err(|error| PromptError::Render(error.to_string()))?;
    let prompt = String::from_utf8(prompt)?;

    match &harness.injection {
        InjectionConfig::ArgText { args } => {
            injection.arguments = args
                .iter()
                .map(|argument| OsString::from(argument.replace(TEXT_PLACEHOLDER, &prompt)))
                .collect();
        }
        InjectionConfig::ArgFile { args } => {
            let path = prepare_prompt_file(config, home, &prompt, persist_file)?;
            let display = path.display().to_string();
            injection.arguments = args
                .iter()
                .map(|argument| OsString::from(argument.replace(FILE_PLACEHOLDER, &display)))
                .collect();
            injection.cache_path = Some(path);
        }
        InjectionConfig::EnvFile { environment } => {
            let path = prepare_prompt_file(config, home, &prompt, persist_file)?;
            injection.environment.insert(
                OsString::from(environment),
                OsString::from(path.as_os_str()),
            );
            injection.cache_path = Some(path);
        }
    }
    injection.prompt = Some(prompt);
    Ok(injection)
}

fn prepare_prompt_file(
    config: &Config,
    home: &Path,
    prompt: &str,
    persist: bool,
) -> Result<PathBuf, PromptError> {
    let cache = expand_path(&config.defaults.prompt_cache, home)
        .map_err(|error| PromptError::Path(error.to_string()))?;
    if !persist {
        return Ok(cache.join("dry-run-prompt.md"));
    }

    fs::create_dir_all(&cache).map_err(|source| PromptError::CacheDirectory {
        path: cache.clone(),
        source,
    })?;
    let now = SystemTime::now();
    gc_prompt_cache(
        &cache,
        Duration::from_secs(config.defaults.prompt_cache_ttl_seconds),
        now,
    );
    let timestamp = now
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_nanos();
    let path = cache.join(format!("{}-{timestamp}.md", std::process::id()));
    let mut file = OpenOptions::new()
        .write(true)
        .create_new(true)
        .mode(0o600)
        .open(&path)
        .map_err(|source| PromptError::CacheFile {
            path: path.clone(),
            source,
        })?;
    file.write_all(prompt.as_bytes())
        .map_err(|source| PromptError::CacheFile {
            path: path.clone(),
            source,
        })?;
    Ok(path)
}

fn gc_prompt_cache(cache: &Path, ttl: Duration, now: SystemTime) {
    let Ok(entries) = fs::read_dir(cache) else {
        return;
    };
    for entry in entries.flatten() {
        let path = entry.path();
        let Ok(metadata) = entry.metadata() else {
            continue;
        };
        let Ok(modified) = metadata.modified() else {
            continue;
        };
        let Ok(age) = now.duration_since(modified) else {
            continue;
        };
        if age > ttl {
            let _ignored = fs::remove_file(path);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::TempDir;

    #[test]
    fn prompt_cache_gc_removes_only_expired_files() {
        let temp = TempDir::new().unwrap();
        let expired = temp.path().join("expired.md");
        let current = temp.path().join("current.md");
        fs::write(&expired, "old").unwrap();
        std::thread::sleep(Duration::from_millis(5));
        fs::write(&current, "new").unwrap();
        let now = SystemTime::now();
        let expired_modified = fs::metadata(&expired).unwrap().modified().unwrap();
        let current_modified = fs::metadata(&current).unwrap().modified().unwrap();
        let threshold = now
            .duration_since(expired_modified)
            .unwrap()
            .checked_sub(Duration::from_millis(1))
            .unwrap();
        assert!(now.duration_since(current_modified).unwrap() < threshold);

        gc_prompt_cache(temp.path(), threshold, now);

        assert!(!expired.exists());
        assert!(current.exists());
    }
}