cortex_llm/

replay.rs

//! Deterministic LLM adapter that replays canned responses from on-disk
//! fixtures.
//!
//! `ReplayAdapter` is the default adapter for CI and any other context where
//! we refuse to spend tokens or leak prompts. It loads a directory of JSON
//! fixtures plus a sibling `INDEX.toml` and:
//!
//! 1. Verifies every fixture path is inside the fixtures directory.
//! 2. Verifies every fixture's on-disk BLAKE3 hash matches the value pinned
//!    in `INDEX.toml` (Lane 1.D / threat row T-RM-1: a fixture-keyed
//!    reflection is otherwise trivially forgeable by any actor with write
//!    access to the fixtures directory).
//! 3. On each `complete()` call, looks up `(model, prompt_hash)` and returns
//!    the pinned response.
//!
//! Fixture format is documented in `crates/cortex-llm/fixtures/schema.json`.

use std::collections::HashMap;
use std::fs;
use std::path::{Path, PathBuf};

use async_trait::async_trait;
use serde::{Deserialize, Serialize};

use crate::adapter::{blake3_hex, LlmAdapter, LlmError, LlmRequest, LlmResponse, TokenUsage};

/// One row of the on-disk `INDEX.toml`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IndexEntry {
    /// Path to the fixture, relative to the fixtures directory and without
    /// any `..` components. The path is rooted inside the fixtures directory
    /// at load time.
    pub path: String,
    /// Lowercase hex BLAKE3 of the fixture's exact on-disk bytes.
    pub blake3: String,
}

/// On-disk representation of the signed fixture manifest.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct FixtureIndex {
    /// Every fixture under the fixtures directory that the adapter trusts.
    #[serde(default, rename = "fixture")]
    pub fixtures: Vec<IndexEntry>,
}

/// Match criteria a fixture declares against incoming [`LlmRequest`]s.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FixtureMatch {
    /// `LlmRequest::model` to match exactly.
    pub model: String,
    /// `LlmRequest::prompt_hash()` to match exactly.
    pub prompt_hash: String,
}

/// Response payload pinned by a fixture.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FixtureResponse {
    /// The reply text returned to the caller.
    pub text: String,
    /// Optional structured form when the original call asked for JSON.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub parsed_json: Option<serde_json::Value>,
    /// Model the fixture claims emitted the reply (defaults to the request's
    /// model on absence — most fixtures echo the requested model).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub model: Option<String>,
    /// Token-usage echo, optional.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub usage: Option<TokenUsage>,
}

/// Parsed shape of a single fixture file.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FixtureFile {
    /// Match criteria.
    pub request_match: FixtureMatch,
    /// Pinned response.
    pub response: FixtureResponse,
}

/// Deterministic adapter that returns fixture-pinned responses.
///
/// Thread-safe: holds an in-memory map keyed by `(model, prompt_hash)`.
#[derive(Debug)]
pub struct ReplayAdapter {
    fixtures_dir: PathBuf,
    by_key: HashMap<(String, String), FixtureFile>,
}

impl ReplayAdapter {
    /// Construct a new replay adapter rooted at `fixtures_dir`.
    ///
    /// The directory MUST contain an `INDEX.toml` listing every trusted
    /// fixture and its expected BLAKE3 hash. Any fixture in the directory
    /// that is **not** listed in the index, or whose on-disk bytes hash to a
    /// different value than the index records, causes
    /// [`LlmError::FixtureIntegrityFailed`].
    ///
    /// The check is upfront — once `new()` returns, the in-memory map is
    /// trusted for the lifetime of the adapter.
    pub fn new<P: Into<PathBuf>>(fixtures_dir: P) -> Result<Self, LlmError> {
        let fixtures_dir = fixtures_dir.into();
        let canonical_root = fs::canonicalize(&fixtures_dir)
            .map_err(|e| LlmError::Io(format!("fixtures dir {}: {e}", fixtures_dir.display())))?;

        let index_path = canonical_root.join("INDEX.toml");
        let index_text = fs::read_to_string(&index_path).map_err(|e| {
            LlmError::FixtureIntegrityFailed(format!(
                "INDEX.toml not readable at {}: {e}",
                index_path.display()
            ))
        })?;
        let index: FixtureIndex = toml::from_str(&index_text).map_err(|e| {
            LlmError::FixtureIntegrityFailed(format!("INDEX.toml parse error: {e}"))
        })?;

        // Build the trusted fixture set first (path → expected hash). This
        // also catches duplicate INDEX entries.
        let mut trusted: HashMap<PathBuf, String> = HashMap::new();
        for entry in &index.fixtures {
            let resolved = resolve_under(&canonical_root, &entry.path)?;
            if trusted
                .insert(resolved.clone(), entry.blake3.clone())
                .is_some()
            {
                return Err(LlmError::FixtureIntegrityFailed(format!(
                    "duplicate INDEX entry: {}",
                    entry.path
                )));
            }
        }

        // Reject any *.json fixture that is not in the trusted set
        // (defence against an attacker dropping a fresh fixture into the
        // directory after the index was committed). `INDEX.toml` itself is
        // skipped, as is `schema.json`.
        for dirent in fs::read_dir(&canonical_root)
            .map_err(|e| LlmError::Io(format!("scanning {}: {e}", canonical_root.display())))?
        {
            let dirent = dirent.map_err(|e| LlmError::Io(format!("dirent: {e}")))?;
            let path = dirent.path();
            if !path.is_file() {
                continue;
            }
            let name = path
                .file_name()
                .and_then(|s| s.to_str())
                .unwrap_or_default();
            if name == "INDEX.toml" || name == "schema.json" {
                continue;
            }
            if !name.ends_with(".json") {
                continue;
            }
            if !trusted.contains_key(&path) {
                return Err(LlmError::FixtureIntegrityFailed(format!(
                    "unsigned fixture present (not in INDEX.toml): {}",
                    path.display()
                )));
            }
        }

        // Verify each trusted fixture's on-disk hash and load it.
        let mut by_key: HashMap<(String, String), FixtureFile> = HashMap::new();
        for (path, expected_hash) in trusted {
            let bytes = fs::read(&path).map_err(|e| {
                LlmError::FixtureIntegrityFailed(format!("read {}: {e}", path.display()))
            })?;
            let actual = blake3_hex(&bytes);
            if !constant_time_eq(actual.as_bytes(), expected_hash.as_bytes()) {
                return Err(LlmError::FixtureIntegrityFailed(format!(
                    "hash mismatch for {} (expected {expected_hash}, got {actual})",
                    path.display()
                )));
            }
            let fixture: FixtureFile = serde_json::from_slice(&bytes).map_err(|e| {
                LlmError::FixtureIntegrityFailed(format!("fixture {} parse: {e}", path.display()))
            })?;
            let key = (
                fixture.request_match.model.clone(),
                fixture.request_match.prompt_hash.clone(),
            );
            if by_key.insert(key, fixture).is_some() {
                return Err(LlmError::FixtureIntegrityFailed(format!(
                    "duplicate (model, prompt_hash) match in fixtures dir {}",
                    canonical_root.display()
                )));
            }
        }

        Ok(Self {
            fixtures_dir: canonical_root,
            by_key,
        })
    }

    /// Path the adapter is rooted at.
    #[must_use]
    pub fn fixtures_dir(&self) -> &Path {
        &self.fixtures_dir
    }

    /// Number of trusted fixtures currently loaded.
    #[must_use]
    pub fn fixture_count(&self) -> usize {
        self.by_key.len()
    }
}

#[async_trait]
impl LlmAdapter for ReplayAdapter {
    fn adapter_id(&self) -> &'static str {
        "replay"
    }

    async fn complete(&self, req: LlmRequest) -> Result<LlmResponse, LlmError> {
        let prompt_hash = req.prompt_hash();
        let key = (req.model.clone(), prompt_hash.clone());
        let Some(fixture) = self.by_key.get(&key) else {
            return Err(LlmError::NoFixture {
                model: req.model,
                prompt_hash,
            });
        };

        let text = fixture.response.text.clone();
        Ok(LlmResponse {
            text: text.clone(),
            parsed_json: fixture.response.parsed_json.clone(),
            model: fixture
                .response
                .model
                .clone()
                .unwrap_or_else(|| req.model.clone()),
            usage: fixture.response.usage.clone(),
            raw_hash: blake3_hex(text.as_bytes()),
        })
    }
}

/// Resolve `relative` under `root`, refusing any value that escapes the root
/// after canonicalization. Returns the absolute, canonical path on success.
fn resolve_under(root: &Path, relative: &str) -> Result<PathBuf, LlmError> {
    let candidate = root.join(relative);
    if candidate
        .components()
        .any(|c| matches!(c, std::path::Component::ParentDir))
    {
        return Err(LlmError::FixtureIntegrityFailed(format!(
            "fixture path escapes fixtures dir: {relative}"
        )));
    }
    let canonical = fs::canonicalize(&candidate).map_err(|e| {
        LlmError::FixtureIntegrityFailed(format!(
            "fixture path {} not resolvable: {e}",
            candidate.display()
        ))
    })?;
    if !canonical.starts_with(root) {
        return Err(LlmError::FixtureIntegrityFailed(format!(
            "fixture path {} escapes fixtures dir {}",
            canonical.display(),
            root.display()
        )));
    }
    Ok(canonical)
}

/// Constant-time byte comparison used for the hash check. Pure stdlib; we do
/// not pull in `subtle` for one helper.
fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
    if a.len() != b.len() {
        return false;
    }
    let mut diff: u8 = 0;
    for (x, y) in a.iter().zip(b.iter()) {
        diff |= x ^ y;
    }
    diff == 0
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::adapter::{LlmMessage, LlmRole};
    use std::fs::File;
    use std::io::Write;
    use tempfile::TempDir;

    fn sample_request(content: &str) -> LlmRequest {
        LlmRequest {
            model: "claude-3-5-sonnet-20240620".into(),
            system: "you are a test".into(),
            messages: vec![LlmMessage {
                role: LlmRole::User,
                content: content.to_string(),
            }],
            temperature: 0.0,
            max_tokens: 256,
            json_schema: None,
            timeout_ms: 30_000,
        }
    }

    /// Write a single fixture under `dir/name`, produce its expected match
    /// criteria from `req`, and return both the resolved path and the
    /// fixture-file struct that was written.
    fn write_fixture(
        dir: &Path,
        name: &str,
        req: &LlmRequest,
        reply: &str,
    ) -> (PathBuf, FixtureFile) {
        let fixture = FixtureFile {
            request_match: FixtureMatch {
                model: req.model.clone(),
                prompt_hash: req.prompt_hash(),
            },
            response: FixtureResponse {
                text: reply.into(),
                parsed_json: None,
                model: None,
                usage: None,
            },
        };
        let path = dir.join(name);
        let bytes = serde_json::to_vec_pretty(&fixture).unwrap();
        let mut f = File::create(&path).unwrap();
        f.write_all(&bytes).unwrap();
        (path, fixture)
    }

    fn write_index(dir: &Path, entries: &[(&str, &str)]) {
        let mut s = String::new();
        for (path, hash) in entries {
            s.push_str(&format!(
                "[[fixture]]\npath = \"{path}\"\nblake3 = \"{hash}\"\n\n"
            ));
        }
        fs::write(dir.join("INDEX.toml"), s).unwrap();
    }

    fn hash_file(p: &Path) -> String {
        blake3_hex(&fs::read(p).unwrap())
    }

    #[tokio::test]
    async fn replay_returns_matching_fixture() {
        let tmp = TempDir::new().unwrap();
        let dir = tmp.path();

        let req_a = sample_request("hello");
        let req_b = sample_request("world");
        let (path_a, _) = write_fixture(dir, "a.json", &req_a, "hi from A");
        let (path_b, _) = write_fixture(dir, "b.json", &req_b, "hi from B");
        write_index(
            dir,
            &[
                ("a.json", &hash_file(&path_a)),
                ("b.json", &hash_file(&path_b)),
            ],
        );

        let adapter = ReplayAdapter::new(dir).unwrap();
        assert_eq!(adapter.fixture_count(), 2);
        let resp = adapter.complete(req_b).await.unwrap();
        assert_eq!(resp.text, "hi from B");
        assert_eq!(resp.model, "claude-3-5-sonnet-20240620");
    }

    #[tokio::test]
    async fn replay_rejects_unsigned_fixture() {
        let tmp = TempDir::new().unwrap();
        let dir = tmp.path();

        let req = sample_request("hello");
        let (path, _) = write_fixture(dir, "a.json", &req, "ok");
        // Drop a SECOND fixture that is NOT listed in the index.
        let (_path_b, _) = write_fixture(dir, "b-unsigned.json", &sample_request("world"), "boom");
        // INDEX only signs the first one.
        write_index(dir, &[("a.json", &hash_file(&path))]);

        let err = ReplayAdapter::new(dir).unwrap_err();
        match err {
            LlmError::FixtureIntegrityFailed(msg) => {
                assert!(
                    msg.contains("unsigned fixture present"),
                    "unexpected message: {msg}"
                );
            }
            other => panic!("expected FixtureIntegrityFailed, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn replay_rejects_hash_mismatch() {
        let tmp = TempDir::new().unwrap();
        let dir = tmp.path();

        let req = sample_request("hello");
        let (path, _) = write_fixture(dir, "a.json", &req, "ok");
        let original_hash = hash_file(&path);
        write_index(dir, &[("a.json", &original_hash)]);

        // Mutate one byte AFTER the index is written.
        let mut bytes = fs::read(&path).unwrap();
        // flip the last byte (which is `\n` or `}` either way: still valid bytes)
        let last = bytes.len() - 1;
        bytes[last] = bytes[last].wrapping_add(1);
        fs::write(&path, bytes).unwrap();

        let err = ReplayAdapter::new(dir).unwrap_err();
        match err {
            LlmError::FixtureIntegrityFailed(msg) => {
                assert!(msg.contains("hash mismatch"), "unexpected message: {msg}");
            }
            other => panic!("expected FixtureIntegrityFailed, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn replay_returns_no_fixture_when_unmatched() {
        let tmp = TempDir::new().unwrap();
        let dir = tmp.path();

        let req = sample_request("hello");
        let (path, _) = write_fixture(dir, "a.json", &req, "ok");
        write_index(dir, &[("a.json", &hash_file(&path))]);

        let adapter = ReplayAdapter::new(dir).unwrap();
        let other = sample_request("not in any fixture");
        let err = adapter.complete(other).await.unwrap_err();
        assert!(matches!(err, LlmError::NoFixture { .. }));
    }

    #[test]
    fn missing_index_is_integrity_failure() {
        let tmp = TempDir::new().unwrap();
        let err = ReplayAdapter::new(tmp.path()).unwrap_err();
        assert!(matches!(err, LlmError::FixtureIntegrityFailed(_)));
    }
}