klieo 3.3.1

//! Recall-recording wrapper for [`LongTermMemory`].
//!
//! GraphRAG recall queries can carry PII (a claimant name, an IBAN typed
//! into a search box) and are otherwise unbounded in length. Wrapping the
//! long-term store in [`RecallRecorder`] ensures every recorded
//! [`Episode::MemoryRecall`] is redacted and length-bounded *before* it
//! reaches durable episodic storage, rather than trusting each call site
//! to remember to sanitize.

use std::sync::Arc;

use async_trait::async_trait;

use klieo_core::{
    AuditRedactor, Episode, EpisodicMemory, Fact, FactId, LongTermMemory, MemoryError, RunId, Scope,
};

/// Max stored chars of a recall query in a [`Episode::MemoryRecall`].
///
/// Public so callers wiring up a [`RecallRecorder`] have a named default
/// instead of hardcoding the bound at each call site.
pub const MAX_RECALL_QUERY_CHARS: usize = 512;

/// Wraps a [`LongTermMemory`] so every [`recall`](LongTermMemory::recall)
/// also records a redacted, length-bounded [`Episode::MemoryRecall`] under
/// the wrapped run id. `remember`/`forget` pass straight through.
///
/// # Best-effort recording
///
/// The inner recall runs first and its facts are returned to the caller
/// unconditionally. The episode write is best-effort: a recording failure
/// is logged (structured `warn`) and swallowed rather than turning a
/// successful recall into an error, mirroring the side-write-never-fails
/// contract `klieo_provenance::ProvenanceEpisodic` uses for the same
/// reason — an observability write should never take down the read path
/// it's observing.
#[non_exhaustive]
pub struct RecallRecorder {
    inner: Arc<dyn LongTermMemory>,
    episodic: Arc<dyn EpisodicMemory>,
    run_id: RunId,
    redactor: Arc<dyn AuditRedactor>,
    max_query_chars: usize,
}

impl RecallRecorder {
    /// Wrap `inner`, recording each `recall` as an [`Episode::MemoryRecall`]
    /// under `run_id` into `episodic`. The recorded query is redacted via
    /// `redactor` and bounded to `max_query_chars` Unicode scalar values.
    pub fn new(
        inner: Arc<dyn LongTermMemory>,
        episodic: Arc<dyn EpisodicMemory>,
        run_id: RunId,
        redactor: Arc<dyn AuditRedactor>,
        max_query_chars: usize,
    ) -> Self {
        Self {
            inner,
            episodic,
            run_id,
            redactor,
            max_query_chars,
        }
    }
}

/// Sentinel stored when a redactor returns a non-string projection.
///
/// A redactor that turns the wrapped query into anything other than a
/// string (e.g. `Value::Null`, or an opaque-digest object) is signalling
/// that it wanted to suppress the content, not reshape it. Storing a fixed
/// marker represents that honestly, rather than the JSON rendering of the
/// value — which could collide with a genuine query (`"null"`) or leak the
/// digest structure into the episode's `query` field.
const NON_STRING_REDACTION_MARKER: &str = "[redacted]";

/// Redact `query` through `redactor` and bound the result to
/// `max_chars` Unicode scalar values (never bytes, so multibyte queries
/// truncate on a character boundary instead of panicking).
///
/// [`AuditRedactor::redact`] only operates on `serde_json::Value`, so the
/// query is wrapped as [`serde_json::Value::String`] for the call. Most
/// redactors preserve the string shape; a redactor that returns a
/// non-string projection is treated as suppression and stored as
/// [`NON_STRING_REDACTION_MARKER`].
fn redact_and_bound_query(redactor: &dyn AuditRedactor, query: &str, max_chars: usize) -> String {
    let redacted = match redactor.redact(&serde_json::Value::String(query.to_string())) {
        serde_json::Value::String(s) => s,
        _ => NON_STRING_REDACTION_MARKER.to_string(),
    };
    if redacted.chars().count() <= max_chars {
        redacted
    } else {
        redacted.chars().take(max_chars).collect()
    }
}

#[async_trait]
impl LongTermMemory for RecallRecorder {
    async fn recall(&self, scope: Scope, query: &str, k: usize) -> Result<Vec<Fact>, MemoryError> {
        let facts = self.inner.recall(scope, query, k).await?;
        let redacted_query =
            redact_and_bound_query(self.redactor.as_ref(), query, self.max_query_chars);
        let episode = Episode::MemoryRecall {
            query: redacted_query,
            k: k as u32,
            // `Fact` carries no id at recall time — `LongTermMemory::recall`
            // returns bare facts, not the ids they were stored under — so
            // there is nothing to populate here. The run-view click-through
            // re-runs the recall live rather than looking up historical ids.
            returned_fact_ids: Vec::new(),
        };
        if let Err(err) = self.episodic.record(self.run_id, episode).await {
            tracing::warn!(
                %err,
                run_id = %self.run_id,
                "failed to record MemoryRecall episode; recall still returned",
            );
        }
        Ok(facts)
    }

    async fn remember(&self, scope: Scope, fact: Fact) -> Result<FactId, MemoryError> {
        self.inner.remember(scope, fact).await
    }

    async fn forget(&self, id: FactId) -> Result<(), MemoryError> {
        self.inner.forget(id).await
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use klieo_core::test_utils::{InMemoryEpisodic, InMemoryLongTerm};

    /// Token a real secret would contain. Distinct from any redaction
    /// marker, so its absence from the stored query proves redaction ran.
    const SECRET_TOKEN: &str = "AKIA1234567890";

    /// Minimal [`AuditRedactor`] that scrubs [`SECRET_TOKEN`] from string
    /// values, leaving everything else untouched. Enough to prove
    /// [`RecallRecorder`] routes the query through the redactor before
    /// storage — a full masking redactor's shape-preservation is already
    /// covered by `klieo-tools-mcp`'s PII redaction tests.
    struct TokenScrubbingRedactor;

    impl AuditRedactor for TokenScrubbingRedactor {
        fn redact(&self, value: &serde_json::Value) -> serde_json::Value {
            match value {
                serde_json::Value::String(s) => {
                    serde_json::Value::String(s.replace(SECRET_TOKEN, "[REDACTED]"))
                }
                other => other.clone(),
            }
        }
    }

    fn test_redactor() -> Arc<dyn AuditRedactor> {
        Arc::new(TokenScrubbingRedactor)
    }

    /// [`AuditRedactor`] that collapses every input to [`serde_json::Value::Null`],
    /// modelling a redactor that suppresses content by returning a non-string
    /// projection. Proves [`RecallRecorder`] stores the suppression marker
    /// rather than the literal `"null"` rendering.
    struct NullRedactor;

    impl AuditRedactor for NullRedactor {
        fn redact(&self, _value: &serde_json::Value) -> serde_json::Value {
            serde_json::Value::Null
        }
    }

    /// [`EpisodicMemory`] whose `record` always fails, proving the
    /// best-effort contract: a recording failure is swallowed and the
    /// recall still returns its facts.
    struct FailingEpisodic;

    #[async_trait]
    impl EpisodicMemory for FailingEpisodic {
        async fn record(&self, _run: RunId, _event: Episode) -> Result<(), MemoryError> {
            Err(MemoryError::Store("boom".to_string()))
        }

        async fn replay(&self, _run: RunId) -> Result<Vec<Episode>, MemoryError> {
            Ok(Vec::new())
        }

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

    fn memory_recalls(episodes: &[Episode]) -> Vec<&Episode> {
        episodes
            .iter()
            .filter(|e| matches!(e, Episode::MemoryRecall { .. }))
            .collect()
    }

    #[tokio::test]
    async fn records_one_redacted_memory_recall_per_recall() {
        let inner = Arc::new(InMemoryLongTerm::default());
        let episodic = Arc::new(InMemoryEpisodic::default());
        let run = RunId::new();
        let rec = RecallRecorder::new(
            inner.clone(),
            episodic.clone(),
            run,
            test_redactor(),
            MAX_RECALL_QUERY_CHARS,
        );

        rec.recall(Scope::Global, "secret token AKIA1234567890", 3)
            .await
            .unwrap();

        let episodes = episodic.replay(run).await.unwrap();
        let recalls = memory_recalls(&episodes);
        assert_eq!(
            recalls.len(),
            1,
            "exactly one MemoryRecall must be recorded"
        );
        match recalls[0] {
            Episode::MemoryRecall {
                query,
                k,
                returned_fact_ids,
            } => {
                assert!(
                    !query.contains(SECRET_TOKEN),
                    "query must be redacted before storage"
                );
                assert_eq!(*k, 3);
                assert!(
                    returned_fact_ids.is_empty(),
                    "Fact carries no id at recall time"
                );
            }
            other => panic!("expected MemoryRecall, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn no_recall_no_episode() {
        let episodic = Arc::new(InMemoryEpisodic::default());
        let run = RunId::new();
        let _rec = RecallRecorder::new(
            Arc::new(InMemoryLongTerm::default()),
            episodic.clone(),
            run,
            test_redactor(),
            MAX_RECALL_QUERY_CHARS,
        );

        let episodes = episodic.replay(run).await.unwrap();
        assert!(episodes.is_empty(), "no recall must record no episode");
    }

    async fn stored_query_for_recall(query: &str) -> String {
        let episodic = Arc::new(InMemoryEpisodic::default());
        let run = RunId::new();
        let rec = RecallRecorder::new(
            Arc::new(InMemoryLongTerm::default()),
            episodic.clone(),
            run,
            test_redactor(),
            MAX_RECALL_QUERY_CHARS,
        );

        rec.recall(Scope::Global, query, 1).await.unwrap();

        let episodes = episodic.replay(run).await.unwrap();
        match memory_recalls(&episodes)[0] {
            Episode::MemoryRecall { query, .. } => query.clone(),
            other => panic!("expected MemoryRecall, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn query_is_length_bounded() {
        let long_query = "q".repeat(MAX_RECALL_QUERY_CHARS + 100);

        let stored = stored_query_for_recall(&long_query).await;

        assert!(
            stored.chars().count() <= MAX_RECALL_QUERY_CHARS,
            "stored query must be bounded to MAX_RECALL_QUERY_CHARS"
        );
    }

    /// A multibyte query longer than the bound must truncate on a Unicode
    /// scalar boundary — a byte-slice bound would split a 4-byte scalar and
    /// panic. Emoji are 4 UTF-8 bytes each, so the byte length far exceeds
    /// the char count.
    #[tokio::test]
    async fn multibyte_query_is_char_bounded_without_panic() {
        let long_query = "🦀".repeat(MAX_RECALL_QUERY_CHARS + 100);

        let stored = stored_query_for_recall(&long_query).await;

        assert!(
            stored.chars().count() <= MAX_RECALL_QUERY_CHARS,
            "stored query must be bounded by Unicode scalar count, not bytes"
        );
    }

    /// A redactor returning a non-string projection (here `Value::Null`)
    /// signals suppression. The stored query must be the fixed suppression
    /// marker, never the literal `"null"` rendering — which would be
    /// indistinguishable from a genuine query containing the word "null".
    #[tokio::test]
    async fn non_string_redaction_stores_marker_not_literal_null() {
        let inner = Arc::new(InMemoryLongTerm::default());
        let episodic = Arc::new(InMemoryEpisodic::default());
        let run = RunId::new();
        let rec = RecallRecorder::new(
            inner,
            episodic.clone(),
            run,
            Arc::new(NullRedactor),
            MAX_RECALL_QUERY_CHARS,
        );

        rec.recall(Scope::Global, "anything", 1).await.unwrap();

        let episodes = episodic.replay(run).await.unwrap();
        match memory_recalls(&episodes)[0] {
            Episode::MemoryRecall { query, .. } => {
                assert_eq!(query, NON_STRING_REDACTION_MARKER);
                assert_ne!(query, "null", "must not store the JSON rendering");
            }
            other => panic!("expected MemoryRecall, got {other:?}"),
        }
    }

    /// The best-effort contract: when the episodic `record` fails, the
    /// failure is swallowed (logged) and the recall still returns its
    /// inner facts rather than propagating the recording error.
    #[tokio::test]
    async fn recording_failure_does_not_fail_recall() {
        let inner = Arc::new(InMemoryLongTerm::default());
        inner
            .remember(Scope::Global, Fact::new("the sky is blue"))
            .await
            .unwrap();
        let rec = RecallRecorder::new(
            inner,
            Arc::new(FailingEpisodic),
            RunId::new(),
            test_redactor(),
            MAX_RECALL_QUERY_CHARS,
        );

        let facts = rec
            .recall(Scope::Global, "sky", 1)
            .await
            .expect("recall must succeed despite the recording failure");

        assert_eq!(facts.len(), 1);
        assert_eq!(facts[0].text, "the sky is blue");
    }

    #[tokio::test]
    async fn recall_returns_inner_facts_unchanged() {
        let inner = Arc::new(InMemoryLongTerm::default());
        inner
            .remember(Scope::Global, Fact::new("the sky is blue"))
            .await
            .unwrap();
        let episodic = Arc::new(InMemoryEpisodic::default());
        let run = RunId::new();
        let rec = RecallRecorder::new(
            inner,
            episodic,
            run,
            test_redactor(),
            MAX_RECALL_QUERY_CHARS,
        );

        let facts = rec.recall(Scope::Global, "sky", 1).await.unwrap();

        assert_eq!(facts.len(), 1);
        assert_eq!(facts[0].text, "the sky is blue");
    }

    #[tokio::test]
    async fn remember_and_forget_pass_through_to_inner() {
        let inner = Arc::new(InMemoryLongTerm::default());
        let episodic = Arc::new(InMemoryEpisodic::default());
        let run = RunId::new();
        let rec = RecallRecorder::new(
            inner.clone(),
            episodic.clone(),
            run,
            test_redactor(),
            MAX_RECALL_QUERY_CHARS,
        );

        let id = rec
            .remember(Scope::Global, Fact::new("alice likes tea"))
            .await
            .unwrap();
        assert_eq!(
            inner.recall(Scope::Global, "tea", 1).await.unwrap().len(),
            1
        );

        rec.forget(id).await.unwrap();
        assert!(inner
            .recall(Scope::Global, "tea", 1)
            .await
            .unwrap()
            .is_empty());

        assert!(
            episodic.replay(run).await.unwrap().is_empty(),
            "remember/forget must not record any episode"
        );
    }
}