klieo_core/

memory.rs

//! Memory traits — short-term, long-term, episodic.

use crate::error::MemoryError;
use crate::ids::{FactId, RunId, ThreadId};
use crate::llm::Message;
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};

/// Outcome of a tool invocation as recorded in the episodic event stream.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "outcome", rename_all = "snake_case")]
pub enum ToolResult {
    /// Tool returned a successful JSON result.
    Ok {
        /// Result payload.
        value: serde_json::Value,
    },
    /// Tool returned an error message.
    Err {
        /// Error string.
        message: String,
    },
}

impl ToolResult {
    /// Build a successful result. Convenience for the
    /// `ToolResult::Ok { value }` struct-variant ceremony.
    pub fn ok(value: serde_json::Value) -> Self {
        Self::Ok { value }
    }

    /// Build an error result. Convenience for the
    /// `ToolResult::Err { message }` struct-variant ceremony.
    pub fn err(message: impl Into<String>) -> Self {
        Self::Err {
            message: message.into(),
        }
    }
}

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

    #[test]
    fn ok_factory_builds_ok_variant() {
        let v = serde_json::json!({"hit": true});
        let r = ToolResult::ok(v.clone());
        match r {
            ToolResult::Ok { value } => assert_eq!(value, v),
            _ => panic!("expected ToolResult::Ok variant"),
        }
    }

    #[test]
    fn err_factory_builds_err_variant() {
        let r = ToolResult::err("boom");
        match r {
            ToolResult::Err { message } => assert_eq!(message, "boom"),
            _ => panic!("expected ToolResult::Err variant"),
        }
    }
}

/// Conversation buffer scoped to a single thread.
///
/// ```
/// # tokio_test::block_on(async {
/// use klieo_core::test_utils::InMemoryShortTerm;
/// use klieo_core::{ShortTermMemory, Message, Role, ThreadId};
///
/// let m = InMemoryShortTerm::default();
/// let thread = ThreadId::new("t1");
/// m.append(thread.clone(), Message {
///     role: Role::User, content: "hi".into(),
///     tool_calls: vec![], tool_call_id: None,
/// }).await.unwrap();
/// let loaded = m.load(thread, 1024).await.unwrap();
/// assert_eq!(loaded.len(), 1);
/// # });
/// ```
#[async_trait]
pub trait ShortTermMemory: Send + Sync {
    /// Append a message to the thread's history.
    async fn append(&self, thread: ThreadId, msg: Message) -> Result<(), MemoryError>;

    /// Append a batch of messages, oldest first. The default appends each in
    /// turn; backends with bulk-insert support should override to collapse the
    /// N writes into one round-trip (the resume path replays a full history).
    async fn append_batch(
        &self,
        thread: ThreadId,
        messages: Vec<Message>,
    ) -> Result<(), MemoryError> {
        for msg in messages {
            self.append(thread.clone(), msg).await?;
        }
        Ok(())
    }

    /// Load up to `max_tokens` of the most-recent messages, oldest first.
    /// Implementations approximate token counts (provider-specific).
    async fn load(&self, thread: ThreadId, max_tokens: usize) -> Result<Vec<Message>, MemoryError>;

    /// Drop all messages for `thread`.
    async fn clear(&self, thread: ThreadId) -> Result<(), MemoryError>;
}

/// Namespacing for long-term memory facts.
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum Scope {
    /// Workspace-scoped (multi-agent shared).
    Workspace(String),
    /// Per-agent scoped.
    Agent(String),
    /// Process-global (use sparingly).
    Global,
}

/// One stored fact in long-term memory.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[non_exhaustive]
pub struct Fact {
    /// Plain-text body, embedded for retrieval.
    pub text: String,
    /// Caller-supplied metadata, opaque to the store.
    #[serde(default)]
    pub metadata: serde_json::Value,
}

impl Fact {
    /// Prefer this over struct literals outside `klieo-core` (`#[non_exhaustive]`);
    /// metadata defaults to JSON null — attach it with [`Fact::with_metadata`].
    pub fn new(text: impl Into<String>) -> Self {
        Self {
            text: text.into(),
            metadata: serde_json::Value::Null,
        }
    }

    /// Override the default JSON-null metadata with an opaque caller value.
    pub fn with_metadata(mut self, metadata: serde_json::Value) -> Self {
        self.metadata = metadata;
        self
    }
}

/// Long-term semantic memory.
///
/// ```
/// # tokio_test::block_on(async {
/// use klieo_core::test_utils::InMemoryLongTerm;
/// use klieo_core::{Fact, LongTermMemory, Scope};
///
/// let m = InMemoryLongTerm::default();
/// let scope = Scope::Workspace("ws".into());
/// m.remember(scope.clone(), Fact::new("the sky is blue")).await.unwrap();
/// let hits = m.recall(scope, "sky", 1).await.unwrap();
/// assert_eq!(hits.len(), 1);
/// # });
/// ```
#[async_trait]
pub trait LongTermMemory: Send + Sync {
    /// Store a fact under `scope`. Returns a stable id.
    async fn remember(&self, scope: Scope, fact: Fact) -> Result<FactId, MemoryError>;

    /// Top-`k` semantic recall under `scope` for the supplied query.
    async fn recall(&self, scope: Scope, query: &str, k: usize) -> Result<Vec<Fact>, MemoryError>;

    /// Remove a stored fact.
    async fn forget(&self, id: FactId) -> Result<(), MemoryError>;
}

/// One event in the episodic event stream of a single agent run.
///
/// Marked `#[non_exhaustive]` so additive variants (e.g.
/// [`Self::SummaryCheckpoint`]) can be introduced without forcing a
/// SemVer-major bump. Match arms in downstream crates must include a
/// fallback `_ => …`.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[non_exhaustive]
pub enum Episode {
    /// Run started.
    Started {
        /// Agent name.
        agent: String,
    },
    /// LLM call completed.
    ///
    /// `provider`/`model`/`prompt_tokens`/`completion_tokens` are
    /// `Option` so legacy emit sites can leave them `None`; downstream
    /// projectors fall back to the `LlmIo` sidecar in `klieo-runlog`
    /// when the structured fields are absent.
    LlmCall {
        /// Total tokens reported by the provider (prompt + completion).
        tokens: u32,
        /// Wall-clock latency in milliseconds.
        latency_ms: u32,
        /// Provider identifier — e.g. `"ollama"`, `"openai"`,
        /// `"anthropic"`, `"gemini"`. `None` for older records or
        /// providers that don't expose a stable name.
        ///
        /// `#[serde(default)]` so legacy episodes serialised under
        /// klieo 0.6.x deserialise as `None`.
        #[serde(default)]
        provider: Option<String>,
        /// Model identifier — e.g. `"qwen2.5:14b"`, `"gpt-4o-mini"`.
        #[serde(default)]
        model: Option<String>,
        /// Prompt-side token count when the provider splits the
        /// breakdown; `None` falls back to `tokens` for total-only
        /// reports.
        #[serde(default)]
        prompt_tokens: Option<u32>,
        /// Completion-side token count when the provider splits the
        /// breakdown.
        #[serde(default)]
        completion_tokens: Option<u32>,
    },
    /// Tool call completed.
    ToolCall {
        /// Tool name.
        name: String,
        /// JSON arguments.
        args: serde_json::Value,
        /// Tool outcome.
        result: ToolResult,
    },
    /// Agent published a bus message.
    BusPublish {
        /// Subject.
        subject: String,
    },
    /// Agent received a bus message.
    BusReceive {
        /// Subject.
        subject: String,
    },
    /// Run completed successfully.
    Completed,
    /// Run failed.
    Failed {
        /// Error message.
        error: String,
    },
    /// Summarizer checkpoint completed.
    ///
    /// Emitted by [`crate::summarize::summarize_history`] in lieu of
    /// [`Self::LlmCall`] so the audit trail can distinguish summarizer
    /// overhead from substantive agent reasoning. Downstream
    /// observability (e.g. `klieo-runlog`) typically projects this as
    /// a separate step kind so cost / latency attribution stays
    /// faithful.
    SummaryCheckpoint {
        /// Number of older messages folded into the summary call.
        input_message_count: u32,
        /// Length of the resulting summary, in Unicode scalar values.
        summary_chars: u32,
        /// Wall-clock latency of the summarizer call.
        latency_ms: u32,
        /// Total tokens reported by the summarizer LLM (prompt +
        /// completion).
        tokens: u32,
    },
    /// Operational-layer event (klieo-ops). Body is an opaque
    /// `serde_json::Value` to keep klieo-core free of an ops dependency.
    /// klieo-ops provides typed serde conversion helpers via `OpsEvent`.
    Ops(serde_json::Value),
    /// Non-PII tenant attribution stamped at run entry when an external
    /// caller drives the run.
    ///
    /// `tenant_label` is a derived identifier (e.g. truncated SHA-256
    /// of the caller's `sub`) — never the raw principal, which lives
    /// only in server-side tracing/authorization. Emitted at most once
    /// per run, adjacent to [`Self::Started`], so the audit trail can
    /// attribute each run to its driving tenant without admitting PII
    /// into agent memory or LLM-visible context.
    RunAttributed {
        /// Derived non-PII attribution label for the driving caller.
        tenant_label: String,
    },
    /// Cross-hop provenance origin stamped at run entry when an
    /// authenticated external caller supplies a parent-chain anchor.
    ///
    /// `parent_anchor` is the caller's own provenance chain-entry id (or
    /// its run's episodic-root hash) — recorded **verbatim** so the value
    /// equals the caller's identifier and downstream tooling can stitch
    /// klieo→klieo lineage across deployments. It is a **caller-asserted,
    /// unverified** claim (klieo does not own or validate the caller's
    /// chain); it is co-recorded with [`Self::RunAttributed`] so the
    /// claim is attributable to the authenticated principal that made it.
    /// Emitted at most once per run, adjacent to [`Self::Started`]; never
    /// admitted into agent memory or LLM-visible context.
    RunOrigin {
        /// Verbatim caller-supplied cross-hop provenance anchor.
        parent_anchor: String,
    },
}

impl Episode {
    /// Construct an [`Episode::LlmCall`] with the legacy two-field shape
    /// (`tokens` + `latency_ms`), leaving the 0.7-added
    /// `provider`/`model`/`prompt_tokens`/`completion_tokens` fields as
    /// `None`.
    ///
    /// Use this when the emit site does not have the enriched fields to
    /// hand — e.g. test fixtures, providers that only report total
    /// tokens, or call sites being migrated incrementally. For full
    /// 0.7 emit semantics, construct the struct variant directly.
    ///
    /// ```
    /// use klieo_core::Episode;
    /// let ep = Episode::llm_call(42, 17);
    /// match ep {
    ///     Episode::LlmCall { tokens, latency_ms, provider, .. } => {
    ///         assert_eq!(tokens, 42);
    ///         assert_eq!(latency_ms, 17);
    ///         assert!(provider.is_none());
    ///     }
    ///     _ => unreachable!(),
    /// }
    /// ```
    pub fn llm_call(tokens: u32, latency_ms: u32) -> Self {
        Episode::LlmCall {
            tokens,
            latency_ms,
            provider: None,
            model: None,
            prompt_tokens: None,
            completion_tokens: None,
        }
    }
}

/// Filter passed to `EpisodicMemory::list_runs`.
#[derive(Debug, Clone, Default)]
pub struct RunFilter {
    /// Filter by agent name (substring match).
    pub agent: Option<String>,
    /// Inclusive lower bound on `started_at`.
    pub since: Option<DateTime<Utc>>,
    /// Inclusive upper bound on `started_at`.
    pub until: Option<DateTime<Utc>>,
    /// Maximum rows returned.
    pub limit: Option<usize>,
}

/// Index-row summary returned by `list_runs`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RunSummary {
    /// Run id.
    pub run_id: RunId,
    /// Agent name.
    pub agent: String,
    /// First-event timestamp.
    pub started_at: DateTime<Utc>,
    /// Last-event timestamp, if completed/failed.
    pub finished_at: Option<DateTime<Utc>>,
    /// Number of episodes.
    pub episode_count: u32,
}

/// Append-only event log of agent runs.
///
/// ```
/// # tokio_test::block_on(async {
/// use klieo_core::test_utils::InMemoryEpisodic;
/// use klieo_core::{Episode, EpisodicMemory, RunId};
///
/// let m = InMemoryEpisodic::default();
/// let run = RunId::new();
/// m.record(run, Episode::Started { agent: "a".into() }).await.unwrap();
/// let events = m.replay(run).await.unwrap();
/// assert_eq!(events.len(), 1);
/// # });
/// ```
#[async_trait]
pub trait EpisodicMemory: Send + Sync {
    /// Record an episode for `run`.
    async fn record(&self, run: RunId, event: Episode) -> Result<(), MemoryError>;

    /// Replay all episodes for `run` in order.
    async fn replay(&self, run: RunId) -> Result<Vec<Episode>, MemoryError>;

    /// List run summaries matching `filter`.
    async fn list_runs(&self, filter: RunFilter) -> Result<Vec<RunSummary>, MemoryError>;
}

/// Resolved trio of memory handles ready to drop into an
/// [`crate::agent::AgentContext`] or an `App`.
///
/// Impl crates (`klieo-memory-sqlite`, `klieo-memory-neo4j`,
/// `klieo-memory-qdrant`) provide `From` conversions where they cover
/// the full trio. Crates that only carry a subset (Neo4j covers
/// short + episodic; Qdrant covers long) compose via the `App` builder's
/// per-trait setters rather than a direct `From`.
#[derive(Clone)]
pub struct MemoryHandles {
    /// Short-term conversation memory.
    pub short_term: std::sync::Arc<dyn ShortTermMemory>,
    /// Long-term semantic memory.
    pub long_term: std::sync::Arc<dyn LongTermMemory>,
    /// Episodic event log.
    pub episodic: std::sync::Arc<dyn EpisodicMemory>,
}

impl MemoryHandles {
    /// Build directly from three already-`Arc`-wrapped handles.
    /// Most callers go through an impl crate's `From` instead.
    pub fn new(
        short_term: std::sync::Arc<dyn ShortTermMemory>,
        long_term: std::sync::Arc<dyn LongTermMemory>,
        episodic: std::sync::Arc<dyn EpisodicMemory>,
    ) -> Self {
        Self {
            short_term,
            long_term,
            episodic,
        }
    }
}

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

    #[test]
    fn fact_new_defaults_metadata_null() {
        let f = Fact::new("alice likes tea");
        assert_eq!(f.text, "alice likes tea");
        assert_eq!(f.metadata, serde_json::Value::Null);
        let f2 = Fact::new("x").with_metadata(serde_json::json!({"k":"v"}));
        assert_eq!(f2.metadata, serde_json::json!({"k":"v"}));
    }
}

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

    #[allow(dead_code)]
    fn _assert_dyn_short(_: &dyn ShortTermMemory) {}
    #[allow(dead_code)]
    fn _assert_dyn_long(_: &dyn LongTermMemory) {}
    #[allow(dead_code)]
    fn _assert_dyn_episodic(_: &dyn EpisodicMemory) {}

    /// Maps each variant to its published snake_case wire discriminant.
    ///
    /// The exhaustive match is the drift guard: a new `Episode` variant
    /// fails to compile here until its discriminant is added, which is the
    /// signal to also publish a payload schema under `docs/schemas/runlog/`
    /// and a fixture in `tests/schema_drift.rs`. `#[non_exhaustive]` does not
    /// force a wildcard inside the defining crate, so this stays exhaustive.
    fn kind_discriminant(episode: &Episode) -> &'static str {
        match episode {
            Episode::Started { .. } => "started",
            Episode::LlmCall { .. } => "llm_call",
            Episode::ToolCall { .. } => "tool_call",
            Episode::BusPublish { .. } => "bus_publish",
            Episode::BusReceive { .. } => "bus_receive",
            Episode::Completed => "completed",
            Episode::Failed { .. } => "failed",
            Episode::SummaryCheckpoint { .. } => "summary_checkpoint",
            Episode::Ops(_) => "ops",
            Episode::RunAttributed { .. } => "run_attributed",
            Episode::RunOrigin { .. } => "run_origin",
        }
    }

    fn one_sample_per_variant() -> Vec<Episode> {
        vec![
            Episode::Started {
                agent: String::new(),
            },
            Episode::LlmCall {
                tokens: 0,
                latency_ms: 0,
                provider: None,
                model: None,
                prompt_tokens: None,
                completion_tokens: None,
            },
            Episode::ToolCall {
                name: String::new(),
                args: serde_json::Value::Null,
                result: ToolResult::Ok {
                    value: serde_json::Value::Null,
                },
            },
            Episode::BusPublish {
                subject: String::new(),
            },
            Episode::BusReceive {
                subject: String::new(),
            },
            Episode::Completed,
            Episode::Failed {
                error: String::new(),
            },
            Episode::SummaryCheckpoint {
                input_message_count: 0,
                summary_chars: 0,
                latency_ms: 0,
                tokens: 0,
            },
            Episode::Ops(serde_json::Value::Null),
            Episode::RunAttributed {
                tenant_label: String::new(),
            },
            Episode::RunOrigin {
                parent_anchor: String::new(),
            },
        ]
    }

    /// The set of Rust `Episode` discriminants must equal the `kind` enum
    /// published in the envelope schema. Compile-time exhaustiveness of
    /// [`kind_discriminant`] plus this runtime equality close the drift loop
    /// from the type side; `tests/schema_drift.rs` validates payload shapes.
    #[test]
    fn episode_discriminants_match_published_envelope_schema() {
        let mut discriminants: Vec<&str> = one_sample_per_variant()
            .iter()
            .map(kind_discriminant)
            .collect();
        discriminants.sort_unstable();

        let schema_path = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR"))
            .join("../../docs/schemas/runlog/episode.schema.json");
        let text = std::fs::read_to_string(&schema_path)
            .unwrap_or_else(|err| panic!("read {}: {err}", schema_path.display()));
        let schema: serde_json::Value =
            serde_json::from_str(&text).expect("envelope schema parses");
        let mut published: Vec<&str> = schema["properties"]["kind"]["enum"]
            .as_array()
            .expect("envelope schema declares a kind enum")
            .iter()
            .map(|value| value.as_str().expect("kind enum is strings"))
            .collect();
        published.sort_unstable();

        assert_eq!(
            discriminants, published,
            "Episode variants have drifted from the published envelope kind enum",
        );
    }

    /// Episodes serialised under klieo 0.6.x carry only `tokens` and
    /// `latency_ms`. The four 0.7 fields must decode as `None` via
    /// `#[serde(default)]`. Without that attribute, every persisted
    /// 0.6 row breaks replay on upgrade.
    #[test]
    fn legacy_llm_call_json_deserialises_with_none_for_new_fields() {
        let legacy = serde_json::json!({
            "LlmCall": {
                "tokens": 42,
                "latency_ms": 17
            }
        });
        let ep: Episode = serde_json::from_value(legacy).expect("legacy LlmCall decodes");
        match ep {
            Episode::LlmCall {
                tokens,
                latency_ms,
                provider,
                model,
                prompt_tokens,
                completion_tokens,
            } => {
                assert_eq!(tokens, 42);
                assert_eq!(latency_ms, 17);
                assert!(provider.is_none());
                assert!(model.is_none());
                assert!(prompt_tokens.is_none());
                assert!(completion_tokens.is_none());
            }
            other => panic!("expected LlmCall, got {other:?}"),
        }
    }

    /// [`Episode::llm_call`] returns the struct variant with all four
    /// 0.7-added enrichment fields as `None`, preserving the legacy
    /// emit shape for callers that don't have provider/model split.
    #[test]
    fn llm_call_ctor_leaves_enrichment_fields_none() {
        match Episode::llm_call(42, 17) {
            Episode::LlmCall {
                tokens,
                latency_ms,
                provider,
                model,
                prompt_tokens,
                completion_tokens,
            } => {
                assert_eq!(tokens, 42);
                assert_eq!(latency_ms, 17);
                assert!(provider.is_none());
                assert!(model.is_none());
                assert!(prompt_tokens.is_none());
                assert!(completion_tokens.is_none());
            }
            other => panic!("expected LlmCall, got {other:?}"),
        }
    }

    /// 0.7 emit sites with the full field set round-trip through serde
    /// unchanged.
    #[test]
    fn enriched_llm_call_round_trips() {
        let original = Episode::LlmCall {
            tokens: 60,
            latency_ms: 17,
            provider: Some("ollama".into()),
            model: Some("qwen2.5:14b".into()),
            prompt_tokens: Some(40),
            completion_tokens: Some(20),
        };
        let json = serde_json::to_value(&original).expect("serialises");
        let back: Episode = serde_json::from_value(json).expect("deserialises");
        match back {
            Episode::LlmCall {
                provider, model, ..
            } => {
                assert_eq!(provider.as_deref(), Some("ollama"));
                assert_eq!(model.as_deref(), Some("qwen2.5:14b"));
            }
            other => panic!("expected LlmCall, got {other:?}"),
        }
    }
}