klieo_provenance/

chain.rs

//! Append-only Merkle chain primitives.

use chrono::{DateTime, SecondsFormat, Timelike, Utc};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use uuid::Uuid;

use crate::error::ProvenanceError;

/// Identifies an entry in the provenance chain.
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct ChainEntryId(
    /// Underlying UUID-v4.
    pub Uuid,
);

impl ChainEntryId {
    /// Mint a fresh random id.
    pub fn new() -> Self {
        Self(Uuid::new_v4())
    }
}

impl Default for ChainEntryId {
    fn default() -> Self {
        Self::new()
    }
}

impl std::fmt::Display for ChainEntryId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.0)
    }
}

/// What kind of operation this provenance entry records.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
#[serde(rename_all = "snake_case")]
pub enum ProvenanceEventKind {
    /// An agent run started.
    AgentRunStarted,
    /// An LLM was invoked.
    LlmCall,
    /// A tool was invoked.
    ToolCall,
    /// Memory was written.
    MemoryWrite,
    /// An agent run reached a terminal state.
    AgentRunCompleted,
    /// Scheduled liveness marker.
    Heartbeat,
    /// Catch-all for callers that need to emit a kind not yet promoted.
    Custom(String),
}

/// A single appendable record in the chain.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct ProvenanceEvent {
    /// What kind of operation this entry records.
    pub kind: ProvenanceEventKind,
    /// Who initiated the operation (e.g. `agent:hello`).
    pub actor: String,
    /// Strong reference to the resource (run id, spec id, etc.).
    pub resource_ref: String,
    /// Hash-only payload to avoid storing potentially sensitive content directly.
    pub payload_hash_hex: String,
    /// Free-form structured metadata.
    pub metadata: serde_json::Value,
}

impl ProvenanceEvent {
    /// Canonical bytes used as input to the entry hash.
    ///
    /// Map keys are sorted recursively so logically-equal events with
    /// different insertion orders produce identical bytes.
    pub fn canonical_bytes(&self) -> Vec<u8> {
        let mut buf = Vec::new();
        let mut ser =
            serde_json::Serializer::with_formatter(&mut buf, serde_json::ser::CompactFormatter);
        let value: serde_json::Value = serde_json::to_value(self).expect("event serialisable");
        let canonical = canonicalise(&value);
        canonical.serialize(&mut ser).expect("canonical write");
        buf
    }
}

fn canonicalise(value: &serde_json::Value) -> serde_json::Value {
    use serde_json::Value;
    match value {
        Value::Object(m) => {
            let mut sorted = std::collections::BTreeMap::new();
            for (k, v) in m {
                sorted.insert(k.clone(), canonicalise(v));
            }
            let mut out = serde_json::Map::with_capacity(sorted.len());
            for (k, v) in sorted {
                out.insert(k, v);
            }
            Value::Object(out)
        }
        Value::Array(a) => Value::Array(a.iter().map(canonicalise).collect()),
        other => other.clone(),
    }
}

/// Current chain canonicalisation version emitted by
/// [`ProvenanceChain::append`]. See ADR-037 (`docs/adr/adr-037-chain-canonicalisation-v2.md`).
pub const CURRENT_CHAIN_VERSION: u8 = 2;

fn default_chain_version_v1() -> u8 {
    1
}

/// One entry in the append-only chain.
///
/// `#[non_exhaustive]` from klieo-provenance 2.0.0 onwards — new optional
/// fields can land in a minor without forcing struct-literal updates. New
/// callers construct via [`ProvenanceChain::append`] (or
/// [`ProvenanceChain::append_with_parents`] for v2 multi-parent entries);
/// reading the public fields stays unchanged.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
#[non_exhaustive]
pub struct ChainEntry {
    /// Random id for this entry.
    pub id: ChainEntryId,
    /// Free-form scope this chain belongs to.
    pub scope: String,
    /// 0-based monotonic position in the chain.
    pub sequence: u64,
    /// Hex SHA-256 of the prior entry, or [`GENESIS_HASH_HEX`] for `sequence == 0`.
    pub previous_hash_hex: String,
    /// The audit event this entry records.
    pub event: ProvenanceEvent,
    /// Wall-clock at which the entry was minted, rounded to ms precision.
    pub recorded_at: DateTime<Utc>,
    /// Hex SHA-256 of this entry's canonical inputs.
    pub entry_hash_hex: String,
    /// Additional semantic-derivation parents (v2+). Empty for linear chains.
    /// Non-empty when an entry derives from multiple prior facts (e.g. M5
    /// PathRAG verbalisation across several path hops). Sorted lexicographically
    /// before hashing for determinism (see [`Self::compute_hash_v2`]).
    /// Legacy v1 JSON deserialises with this field empty via `#[serde(default)]`.
    #[serde(default)]
    pub parent_hashes_hex: Vec<String>,
    /// Canonicalisation version that produced [`Self::entry_hash_hex`].
    /// `1` for entries minted before ADR-037 (legacy JSON has no field;
    /// the `#[serde(default)]` hook supplies 1); `2` for entries minted
    /// by [`ProvenanceChain::append`] under klieo-provenance 2.0+.
    /// Verification dispatches per-entry on this value, so chains may
    /// contain a v1 prefix and a v2 suffix and still verify end-to-end.
    #[serde(default = "default_chain_version_v1")]
    pub chain_version: u8,
}

impl ChainEntry {
    /// Canonical timestamp form — pinned to ms precision + Z suffix.
    pub fn canonical_timestamp(recorded_at: DateTime<Utc>) -> String {
        recorded_at.to_rfc3339_opts(SecondsFormat::Millis, true)
    }

    /// Canonical SHA-256 hash for the legacy v1 entry shape.
    ///
    /// Canonical input: `prev_hex_bytes || seq_be_bytes || ts_bytes || event_canonical_bytes`.
    /// Bit-identical to the pre-ADR-037 `compute_hash`; preserved verbatim
    /// so v1 entries serialised before the migration still verify.
    pub fn compute_hash_v1(
        previous_hash_hex: &str,
        sequence: u64,
        event: &ProvenanceEvent,
        recorded_at: DateTime<Utc>,
    ) -> String {
        let mut hasher = Sha256::new();
        hasher.update(previous_hash_hex.as_bytes());
        hasher.update(sequence.to_be_bytes());
        hasher.update(Self::canonical_timestamp(recorded_at).as_bytes());
        hasher.update(event.canonical_bytes());
        hex::encode(hasher.finalize())
    }

    /// Canonical SHA-256 hash for the v2 entry shape (ADR-037).
    ///
    /// Canonical input: `0x02 || prev_hex_bytes || seq_be_bytes || ts_bytes ||
    /// event_canonical_bytes || sorted_parents_separated`.
    /// `parent_hashes_hex` is sorted lexicographically before hashing so
    /// the call-site order is irrelevant. A null byte separates adjacent
    /// parent strings to make parent boundaries unambiguous. The leading
    /// `0x02` distinguishes v2 entries from v1 entries with identical
    /// payloads and empty parent sets.
    pub fn compute_hash_v2(
        previous_hash_hex: &str,
        sequence: u64,
        event: &ProvenanceEvent,
        recorded_at: DateTime<Utc>,
        parent_hashes_hex: &[String],
    ) -> String {
        let mut sorted_parents: Vec<&str> = parent_hashes_hex.iter().map(String::as_str).collect();
        sorted_parents.sort_unstable();

        let mut hasher = Sha256::new();
        hasher.update([CURRENT_CHAIN_VERSION]);
        hasher.update(previous_hash_hex.as_bytes());
        hasher.update(sequence.to_be_bytes());
        hasher.update(Self::canonical_timestamp(recorded_at).as_bytes());
        hasher.update(event.canonical_bytes());
        for parent in sorted_parents {
            hasher.update(parent.as_bytes());
            hasher.update([0x00_u8]);
        }
        hex::encode(hasher.finalize())
    }

    /// Dispatch to the version-specific hash algorithm.
    ///
    /// Returns `Err(ProvenanceError::UnsupportedChainVersion)` for any
    /// `chain_version` outside `{1, 2}`. Verifiers MUST refuse unknown
    /// versions rather than fall back to a default — a chain with an
    /// unknown version cannot be cryptographically validated.
    pub fn compute_hash_for_version(
        chain_version: u8,
        previous_hash_hex: &str,
        sequence: u64,
        event: &ProvenanceEvent,
        recorded_at: DateTime<Utc>,
        parent_hashes_hex: &[String],
    ) -> Result<String, ProvenanceError> {
        match chain_version {
            1 => Ok(Self::compute_hash_v1(
                previous_hash_hex,
                sequence,
                event,
                recorded_at,
            )),
            2 => Ok(Self::compute_hash_v2(
                previous_hash_hex,
                sequence,
                event,
                recorded_at,
                parent_hashes_hex,
            )),
            other => Err(ProvenanceError::UnsupportedChainVersion(other)),
        }
    }
}

/// Genesis hash used as the predecessor of the first entry in any chain.
pub const GENESIS_HASH_HEX: &str =
    "0000000000000000000000000000000000000000000000000000000000000000";

fn round_to_millis(t: DateTime<Utc>) -> DateTime<Utc> {
    let nanos = t.timestamp_subsec_nanos();
    let trimmed = (nanos / 1_000_000) * 1_000_000;
    t.with_nanosecond(trimmed).unwrap_or(t)
}

/// Provenance chain aggregate. Append-only.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct ProvenanceChain {
    /// The scope this chain belongs to.
    pub scope: String,
    /// Entries in monotonically-increasing sequence.
    pub entries: Vec<ChainEntry>,
}

impl ProvenanceChain {
    /// Construct an empty chain for the given scope.
    pub fn new(scope: String) -> Self {
        Self {
            scope,
            entries: Vec::new(),
        }
    }

    /// Hex SHA-256 of the most recent entry, or [`GENESIS_HASH_HEX`] if empty.
    pub fn head_hash_hex(&self) -> &str {
        self.entries
            .last()
            .map(|e| e.entry_hash_hex.as_str())
            .unwrap_or(GENESIS_HASH_HEX)
    }

    /// Sequence the next appended entry will use.
    pub fn next_sequence(&self) -> u64 {
        self.entries.last().map(|e| e.sequence + 1).unwrap_or(0)
    }

    /// Append a linear-chain event (empty semantic parents). Mints a v2
    /// entry. Use [`Self::append_with_parents`] when the entry derives
    /// from multiple prior facts.
    pub fn append(&mut self, event: ProvenanceEvent) -> ChainEntry {
        self.append_with_parents(event, Vec::new())
    }

    /// Append an event whose canonical hash includes the supplied
    /// semantic-derivation parents. Parents are sorted lexicographically
    /// in the canonical input (see [`ChainEntry::compute_hash_v2`]).
    pub fn append_with_parents(
        &mut self,
        event: ProvenanceEvent,
        parent_hashes_hex: Vec<String>,
    ) -> ChainEntry {
        let recorded_at = round_to_millis(Utc::now());
        let sequence = self.next_sequence();
        let previous_hash_hex = self.head_hash_hex().to_string();
        let entry_hash_hex = ChainEntry::compute_hash_v2(
            &previous_hash_hex,
            sequence,
            &event,
            recorded_at,
            &parent_hashes_hex,
        );
        let entry = ChainEntry {
            id: ChainEntryId::new(),
            scope: self.scope.clone(),
            sequence,
            previous_hash_hex,
            event,
            recorded_at,
            entry_hash_hex,
            parent_hashes_hex,
            chain_version: CURRENT_CHAIN_VERSION,
        };
        self.entries.push(entry.clone());
        entry
    }

    /// Verify chain integrity. Returns `ChainBroken` on first
    /// inconsistency. Per-entry hash computation dispatches on
    /// `entry.chain_version` so chains containing a v1 prefix + v2
    /// suffix verify end-to-end.
    pub fn verify(&self) -> Result<(), ProvenanceError> {
        let mut expected_prev = GENESIS_HASH_HEX.to_string();
        for (i, entry) in self.entries.iter().enumerate() {
            if entry.sequence != i as u64 {
                return Err(ProvenanceError::ChainBroken {
                    sequence: entry.sequence,
                    reason: format!("sequence gap at index {i}"),
                });
            }
            if entry.previous_hash_hex != expected_prev {
                return Err(ProvenanceError::ChainBroken {
                    sequence: entry.sequence,
                    reason: "previous_hash mismatch".into(),
                });
            }
            let recomputed = ChainEntry::compute_hash_for_version(
                entry.chain_version,
                &entry.previous_hash_hex,
                entry.sequence,
                &entry.event,
                entry.recorded_at,
                &entry.parent_hashes_hex,
            )?;
            if recomputed != entry.entry_hash_hex {
                return Err(ProvenanceError::ChainBroken {
                    sequence: entry.sequence,
                    reason: format!(
                        "hash mismatch (recomputed {recomputed} vs stored {})",
                        entry.entry_hash_hex
                    ),
                });
            }
            expected_prev = entry.entry_hash_hex.clone();
        }
        Ok(())
    }
}

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

    fn run_started(run_id: &str) -> ProvenanceEvent {
        ProvenanceEvent {
            kind: ProvenanceEventKind::AgentRunStarted,
            actor: "agent:hello".into(),
            resource_ref: run_id.into(),
            payload_hash_hex: "ab".repeat(32),
            metadata: serde_json::json!({"agent_name": "hello"}),
        }
    }

    #[test]
    fn empty_chain_has_genesis_head() {
        let chain = ProvenanceChain::new("run-1".into());
        assert_eq!(chain.head_hash_hex(), GENESIS_HASH_HEX);
        assert_eq!(chain.next_sequence(), 0);
    }

    #[test]
    fn append_first_entry_links_to_genesis() {
        let mut chain = ProvenanceChain::new("run-1".into());
        let entry = chain.append(run_started("run-1"));
        assert_eq!(entry.sequence, 0);
        assert_eq!(entry.previous_hash_hex, GENESIS_HASH_HEX);
        assert_ne!(entry.entry_hash_hex, GENESIS_HASH_HEX);
    }

    #[test]
    fn append_extends_chain_and_links_correctly() {
        let mut chain = ProvenanceChain::new("run-1".into());
        let a = chain.append(run_started("run-1"));
        let b = chain.append(run_started("run-1"));
        assert_eq!(b.sequence, 1);
        assert_eq!(b.previous_hash_hex, a.entry_hash_hex);
    }

    #[test]
    fn verify_passes_on_clean_chain() {
        let mut chain = ProvenanceChain::new("run-1".into());
        chain.append(run_started("run-1"));
        chain.append(run_started("run-1"));
        chain.append(run_started("run-1"));
        chain.verify().unwrap();
    }

    #[test]
    fn verify_detects_broken_link() {
        let mut chain = ProvenanceChain::new("run-1".into());
        chain.append(run_started("run-1"));
        chain.append(run_started("run-1"));
        chain.entries[1].previous_hash_hex = "ff".repeat(32);
        let err = chain.verify().unwrap_err();
        match err {
            ProvenanceError::ChainBroken { sequence, reason } => {
                assert_eq!(sequence, 1);
                assert!(reason.contains("previous_hash mismatch"));
            }
            other => panic!("unexpected: {other:?}"),
        }
    }

    fn fixture_event() -> ProvenanceEvent {
        ProvenanceEvent {
            kind: ProvenanceEventKind::AgentRunStarted,
            actor: "agent:fixture".into(),
            resource_ref: "run-fixture".into(),
            payload_hash_hex: "ab".repeat(32),
            metadata: serde_json::json!({"agent_name": "fixture"}),
        }
    }

    fn fixture_recorded_at() -> DateTime<Utc> {
        "2026-06-04T12:00:00.000Z".parse().unwrap()
    }

    #[test]
    fn append_mints_v2_entries() {
        let mut chain = ProvenanceChain::new("run-1".into());
        let entry = chain.append(run_started("run-1"));
        assert_eq!(entry.chain_version, CURRENT_CHAIN_VERSION);
        assert_eq!(entry.chain_version, 2);
        assert!(entry.parent_hashes_hex.is_empty());
    }

    #[test]
    fn append_with_parents_records_supplied_parents() {
        let mut chain = ProvenanceChain::new("run-1".into());
        let parents = vec!["aa".repeat(32), "bb".repeat(32)];
        let entry = chain.append_with_parents(run_started("run-1"), parents.clone());
        assert_eq!(entry.chain_version, 2);
        assert_eq!(entry.parent_hashes_hex, parents);
        chain.verify().expect("v2 chain with parents verifies");
    }

    #[test]
    fn compute_hash_v1_matches_pre_extension_fixture() {
        // Frozen pre-ADR-037 wire format: no `chain_version`, no
        // `parent_hashes_hex` in the JSON. `#[serde(default)]` fills both.
        let recorded_at = fixture_recorded_at();
        let event = fixture_event();
        let expected_hex = ChainEntry::compute_hash_v1(GENESIS_HASH_HEX, 0, &event, recorded_at);
        let payload_hash = "ab".repeat(32);

        let json = format!(
            r#"{{
                "id": "00000000-0000-4000-8000-000000000001",
                "scope": "run-fixture",
                "sequence": 0,
                "previous_hash_hex": "{GENESIS_HASH_HEX}",
                "event": {{
                    "kind": "agent_run_started",
                    "actor": "agent:fixture",
                    "resource_ref": "run-fixture",
                    "payload_hash_hex": "{payload_hash}",
                    "metadata": {{"agent_name": "fixture"}}
                }},
                "recorded_at": "2026-06-04T12:00:00.000Z",
                "entry_hash_hex": "{expected_hex}"
            }}"#
        );

        let entry: ChainEntry = serde_json::from_str(&json).expect("legacy JSON deserialises");
        assert_eq!(
            entry.chain_version, 1,
            "missing chain_version defaults to v1"
        );
        assert!(
            entry.parent_hashes_hex.is_empty(),
            "missing parents default to empty"
        );

        let recomputed = ChainEntry::compute_hash_for_version(
            entry.chain_version,
            &entry.previous_hash_hex,
            entry.sequence,
            &entry.event,
            entry.recorded_at,
            &entry.parent_hashes_hex,
        )
        .expect("v1 dispatcher accepts version 1");

        assert_eq!(
            recomputed, entry.entry_hash_hex,
            "v1 hash via dispatcher must match the embedded entry_hash_hex"
        );
    }

    #[test]
    fn compute_hash_v2_includes_parent_hashes() {
        let recorded_at = fixture_recorded_at();
        let event = fixture_event();
        let h_a = ChainEntry::compute_hash_v2(
            GENESIS_HASH_HEX,
            0,
            &event,
            recorded_at,
            &["aa".repeat(32)],
        );
        let h_b = ChainEntry::compute_hash_v2(
            GENESIS_HASH_HEX,
            0,
            &event,
            recorded_at,
            &["bb".repeat(32)],
        );
        assert_ne!(h_a, h_b, "different parents must produce different hashes");
    }

    #[test]
    fn compute_hash_v2_parent_hashes_order_independent() {
        let recorded_at = fixture_recorded_at();
        let event = fixture_event();
        let asc = ChainEntry::compute_hash_v2(
            GENESIS_HASH_HEX,
            0,
            &event,
            recorded_at,
            &["aa".repeat(32), "bb".repeat(32)],
        );
        let desc = ChainEntry::compute_hash_v2(
            GENESIS_HASH_HEX,
            0,
            &event,
            recorded_at,
            &["bb".repeat(32), "aa".repeat(32)],
        );
        assert_eq!(
            asc, desc,
            "parent order at call site must not change the hash"
        );
    }

    #[test]
    fn compute_hash_v1_distinct_from_v2_with_empty_parents() {
        let recorded_at = fixture_recorded_at();
        let event = fixture_event();
        let v1 = ChainEntry::compute_hash_v1(GENESIS_HASH_HEX, 0, &event, recorded_at);
        let v2 = ChainEntry::compute_hash_v2(GENESIS_HASH_HEX, 0, &event, recorded_at, &[]);
        assert_ne!(
            v1, v2,
            "the v2 version-byte prefix must distinguish v1 from v2 even with empty parents"
        );
    }

    #[test]
    fn compute_hash_for_version_rejects_unknown_version() {
        let recorded_at = fixture_recorded_at();
        let event = fixture_event();
        let err =
            ChainEntry::compute_hash_for_version(99, GENESIS_HASH_HEX, 0, &event, recorded_at, &[])
                .expect_err("unknown version must be rejected");
        match err {
            ProvenanceError::UnsupportedChainVersion(v) => assert_eq!(v, 99),
            other => panic!("expected UnsupportedChainVersion, got {other:?}"),
        }
    }

    #[test]
    fn chain_verifies_mixed_versions() {
        // Build a chain whose first entry is a hand-minted v1 entry
        // (simulating pre-ADR-037 wire) followed by v2 entries via append().
        let recorded_at = fixture_recorded_at();
        let event = fixture_event();
        let v1_hash = ChainEntry::compute_hash_v1(GENESIS_HASH_HEX, 0, &event, recorded_at);
        let v1_entry = ChainEntry {
            id: ChainEntryId::new(),
            scope: "mixed".into(),
            sequence: 0,
            previous_hash_hex: GENESIS_HASH_HEX.to_string(),
            event,
            recorded_at,
            entry_hash_hex: v1_hash,
            parent_hashes_hex: Vec::new(),
            chain_version: 1,
        };
        let mut chain = ProvenanceChain::new("mixed".into());
        chain.entries.push(v1_entry);
        chain.append(run_started("mixed"));
        chain.append(run_started("mixed"));
        chain
            .verify()
            .expect("mixed v1-prefix + v2-suffix chain verifies end-to-end");
        // Sanity: head is v2.
        assert_eq!(chain.entries.last().unwrap().chain_version, 2);
        assert_eq!(chain.entries[0].chain_version, 1);
    }

    #[test]
    fn verify_rejects_unknown_chain_version_on_an_entry() {
        let mut chain = ProvenanceChain::new("run-1".into());
        chain.append(run_started("run-1"));
        chain.entries[0].chain_version = 99;
        let err = chain
            .verify()
            .expect_err("verify must surface unsupported version");
        match err {
            ProvenanceError::UnsupportedChainVersion(v) => assert_eq!(v, 99),
            other => panic!("expected UnsupportedChainVersion, got {other:?}"),
        }
    }
}