klieo-core 0.8.1

//! 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>;

    /// 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)]
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,
}

/// 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 {
///     text: "the sky is blue".into(),
///     metadata: serde_json::Value::Null,
/// }).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),
}

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>;
}

#[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) {}

    /// 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:?}"),
        }
    }
}