klieo-provenance 3.5.0

//! Local [`EpisodicMemory`] decorator that mirrors each episode write into
//! a [`ProvenanceRepository`] hash-chain.
//!
//! # Write amplification
//!
//! Every [`EpisodicMemory::record`] performed through
//! [`ProvenanceEpisodic`] also issues one [`ProvenanceRepository::append`].
//! This per-episode append is the accepted, documented cost of the durable
//! provenance trail — agents recording many fine-grained episodes pay one
//! extra chain append per episode.
//!
//! # Observability side-write never fails the primary write
//!
//! The wrapped store's `record` runs first and its result is propagated
//! verbatim. The provenance append is best-effort: an append failure is
//! logged (structured `warn`) and swallowed — it never turns a successful
//! episode write into an error, and never skips the inner write.

use std::sync::Arc;

use async_trait::async_trait;
use sha2::{Digest, Sha256};

use klieo_core::{Episode, EpisodicMemory, MemoryError, RunFilter, RunId, RunSummary};

use crate::chain::{ProvenanceEvent, ProvenanceEventKind};
use crate::repository::ProvenanceRepository;

/// Decorator that wraps an [`EpisodicMemory`] and mirrors every episode
/// write into a [`ProvenanceRepository`] append-only chain.
///
/// Construct via [`ProvenanceEpisodic::new`]. The `actor` is a **stable
/// role label** (e.g. `"klieo-agent"`) recorded on every provenance event —
/// never operator PII.
///
/// See the module docs for the write-amplification cost and the
/// side-write-never-fails contract.
#[non_exhaustive]
pub struct ProvenanceEpisodic {
    inner: Arc<dyn EpisodicMemory>,
    repo: Arc<dyn ProvenanceRepository>,
    scope: String,
    actor: String,
}

impl ProvenanceEpisodic {
    /// Wrap `inner`, mirroring each recorded episode into `repo` under
    /// `scope`. `actor` is a stable, non-PII role label stamped on every
    /// provenance event.
    pub fn new(
        inner: Arc<dyn EpisodicMemory>,
        repo: Arc<dyn ProvenanceRepository>,
        scope: impl Into<String>,
        actor: impl Into<String>,
    ) -> Self {
        Self {
            inner,
            repo,
            scope: scope.into(),
            actor: actor.into(),
        }
    }
}

/// Map an [`Episode`] to the [`ProvenanceEvent`] mirrored into the chain.
///
/// `actor` is recorded verbatim as the stable role label. The episode's
/// run id is the `resource_ref`, linking each chain entry back to its run.
/// The episode payload is hashed (not stored) into `payload_hash_hex` so the
/// chain carries integrity coverage without admitting episode content — which
/// may carry sensitive data — into the provenance store.
pub(crate) fn episode_to_event(run_id: RunId, episode: &Episode, actor: &str) -> ProvenanceEvent {
    ProvenanceEvent {
        kind: episode_kind(episode),
        actor: actor.to_string(),
        resource_ref: run_id.to_string(),
        payload_hash_hex: payload_hash_hex(episode),
        metadata: serde_json::Value::Null,
    }
}

fn episode_kind(episode: &Episode) -> ProvenanceEventKind {
    match episode {
        Episode::Started { .. } | Episode::RunAttributed { .. } | Episode::RunOrigin { .. } => {
            ProvenanceEventKind::AgentRunStarted
        }
        Episode::LlmCall { .. } | Episode::SummaryCheckpoint { .. } => ProvenanceEventKind::LlmCall,
        Episode::ToolCall { .. } => ProvenanceEventKind::ToolCall,
        Episode::BusPublish { .. } | Episode::BusReceive { .. } | Episode::BusCausalLink { .. } => {
            ProvenanceEventKind::MemoryWrite
        }
        Episode::Completed | Episode::Failed { .. } => ProvenanceEventKind::AgentRunCompleted,
        Episode::Ops(_) => ProvenanceEventKind::Custom("ops".to_string()),
        _ => ProvenanceEventKind::Custom("episode".to_string()),
    }
}

/// Hex SHA-256 of the episode's canonical JSON. On a serialization failure
/// (no `Episode` variant fails today) the Err branch logs a structured `warn`
/// and falls back to the empty-input digest, keeping the side-write infallible
/// at this seam.
fn payload_hash_hex(episode: &Episode) -> String {
    let bytes = serde_json::to_vec(episode).unwrap_or_else(|err| {
        tracing::warn!(%err, "episode serialization failed; provenance payload hash will use the empty-input digest");
        Vec::new()
    });
    let digest = Sha256::digest(&bytes);
    hex::encode(digest)
}

#[async_trait]
impl EpisodicMemory for ProvenanceEpisodic {
    async fn record(&self, run: RunId, event: Episode) -> Result<(), MemoryError> {
        let provenance_event = episode_to_event(run, &event, &self.actor);
        self.inner.record(run, event).await?;
        if let Err(err) = self.repo.append(&self.scope, provenance_event).await {
            tracing::warn!(%err, run_id = %run, "provenance append failed");
        }
        Ok(())
    }

    async fn replay(&self, run: RunId) -> Result<Vec<Episode>, MemoryError> {
        self.inner.replay(run).await
    }

    async fn replay_many(&self, runs: &[RunId]) -> Result<Vec<(RunId, Vec<Episode>)>, MemoryError> {
        self.inner.replay_many(runs).await
    }

    async fn list_runs(&self, filter: RunFilter) -> Result<Vec<RunSummary>, MemoryError> {
        self.inner.list_runs(filter).await
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use klieo_core::ToolResult;
    use std::collections::HashMap;
    use std::sync::Mutex;

    use chrono::Utc;

    use crate::chain::ChainEntry;
    use crate::error::ProvenanceError;

    const TEST_SCOPE: &str = "test-scope";
    const TEST_ACTOR: &str = "klieo-agent";

    #[derive(Default)]
    struct FakeEpisodic {
        runs: Mutex<HashMap<RunId, Vec<Episode>>>,
    }

    #[async_trait]
    impl EpisodicMemory for FakeEpisodic {
        async fn record(&self, run: RunId, event: Episode) -> Result<(), MemoryError> {
            self.runs
                .lock()
                .unwrap()
                .entry(run)
                .or_default()
                .push(event);
            Ok(())
        }

        async fn replay(&self, run: RunId) -> Result<Vec<Episode>, MemoryError> {
            Ok(self
                .runs
                .lock()
                .unwrap()
                .get(&run)
                .cloned()
                .unwrap_or_default())
        }

        async fn list_runs(&self, _filter: RunFilter) -> Result<Vec<RunSummary>, MemoryError> {
            Ok(Vec::new())
        }
    }

    /// A [`ProvenanceRepository`] whose `append` always fails — proves the
    /// decorator's side-write never breaks the primary episode write.
    struct FailingRepo;

    #[async_trait]
    impl ProvenanceRepository for FailingRepo {
        async fn append(
            &self,
            _scope: &str,
            _event: ProvenanceEvent,
        ) -> Result<ChainEntry, ProvenanceError> {
            Err(ProvenanceError::Repository("boom".to_string()))
        }

        async fn head(&self, _scope: &str) -> Result<Option<ChainEntry>, ProvenanceError> {
            Ok(None)
        }

        async fn list_range(
            &self,
            _scope: &str,
            _from: u64,
            _to: u64,
        ) -> Result<Vec<ChainEntry>, ProvenanceError> {
            Ok(Vec::new())
        }

        async fn list_by_time(
            &self,
            _scope: &str,
            _from: chrono::DateTime<Utc>,
            _to: chrono::DateTime<Utc>,
        ) -> Result<Vec<ChainEntry>, ProvenanceError> {
            Ok(Vec::new())
        }
    }

    #[cfg(feature = "sqlite")]
    #[tokio::test]
    async fn record_delegates_and_appends_to_chain() {
        let inner = Arc::new(FakeEpisodic::default());
        let repo = Arc::new(crate::SqliteProvenanceRepository::open_in_memory().unwrap());
        let decorator =
            ProvenanceEpisodic::new(inner.clone(), repo.clone(), TEST_SCOPE, TEST_ACTOR);
        let run = RunId::new();

        decorator
            .record(run, Episode::Started { agent: "a".into() })
            .await
            .unwrap();

        let replayed = inner.replay(run).await.unwrap();
        assert_eq!(replayed.len(), 1, "inner store must hold the episode");

        let head = repo.head(TEST_SCOPE).await.unwrap();
        assert!(head.is_some(), "provenance chain must have a head entry");
        let head = head.unwrap();
        assert_eq!(head.event.actor, TEST_ACTOR);
        assert_eq!(head.event.resource_ref, run.to_string());
    }

    #[tokio::test]
    async fn append_failure_never_fails_record() {
        let inner = Arc::new(FakeEpisodic::default());
        let repo = Arc::new(FailingRepo);
        let decorator = ProvenanceEpisodic::new(inner.clone(), repo, TEST_SCOPE, TEST_ACTOR);
        let run = RunId::new();

        let result = decorator
            .record(run, Episode::Started { agent: "a".into() })
            .await;

        assert!(result.is_ok(), "append failure must not fail the record");
        let replayed = inner.replay(run).await.unwrap();
        assert_eq!(replayed.len(), 1, "inner write must still have happened");
    }

    #[test]
    fn episode_to_event_maps_variants_with_actor_and_run_linkage() {
        let run = RunId::new();

        let started = episode_to_event(run, &Episode::Started { agent: "a".into() }, TEST_ACTOR);
        assert_eq!(started.kind, ProvenanceEventKind::AgentRunStarted);
        assert_eq!(started.actor, TEST_ACTOR);
        assert_eq!(started.resource_ref, run.to_string());

        let llm = episode_to_event(run, &Episode::llm_call(42, 17), TEST_ACTOR);
        assert_eq!(llm.kind, ProvenanceEventKind::LlmCall);
        assert_eq!(llm.resource_ref, run.to_string());

        let completed = episode_to_event(run, &Episode::Completed, TEST_ACTOR);
        assert_eq!(completed.kind, ProvenanceEventKind::AgentRunCompleted);

        let tool = episode_to_event(
            run,
            &Episode::ToolCall {
                name: "search".into(),
                args: serde_json::Value::Null,
                result: ToolResult::ok(serde_json::Value::Null),
            },
            TEST_ACTOR,
        );
        assert_eq!(tool.kind, ProvenanceEventKind::ToolCall);

        let bus = episode_to_event(
            run,
            &Episode::BusPublish {
                subject: "topic".into(),
            },
            TEST_ACTOR,
        );
        assert_eq!(bus.kind, ProvenanceEventKind::MemoryWrite);

        let ops = episode_to_event(run, &Episode::Ops(serde_json::Value::Null), TEST_ACTOR);
        assert_eq!(ops.kind, ProvenanceEventKind::Custom("ops".to_string()));

        // Payload hash is content-derived and differs across variants.
        assert_ne!(started.payload_hash_hex, completed.payload_hash_hex);
    }
}