swink-agent-memory 0.13.1

//! Canonical persistence codecs used by all storage backends.
//!
//! This module centralises the shared encode / decode rules for:
//! - plain [`AgentMessage`] persistence
//! - JSONL message lines, including legacy `_custom` and `_state` records
//! - rich [`SessionEntry`] persistence used by filtered/history-aware loading
//!
//! Backends should only deal with their transport shape (line-based JSONL,
//! `(kind, data)` rows in `SQLite`, etc.), not reimplement message codecs.

use std::io;

use swink_agent::{
    AgentMessage, CustomMessageRegistry, LlmMessage, restore_single_custom,
    serialize_custom_message,
};

use crate::entry::SessionEntry;

/// Discriminator tag for an encoded agent message.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum MessageKind {
    /// An [`LlmMessage`] (user, assistant, tool result).
    Llm,
    /// A custom message serialised via `swink_agent::CustomMessage::to_json`.
    Custom,
}

impl MessageKind {
    /// The string tag stored in the `SQLite` `kind` column.
    #[allow(dead_code)]
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::Llm => "llm",
            Self::Custom => "custom",
        }
    }

    /// Parse a stored tag back to a [`MessageKind`]. Returns `None` for
    /// unknown tags so callers can skip unrecognised rows gracefully.
    pub fn parse(s: &str) -> Option<Self> {
        match s {
            "llm" => Some(Self::Llm),
            "custom" => Some(Self::Custom),
            _ => None,
        }
    }
}

/// Encode an [`AgentMessage`] to a `(kind, json)` pair.
///
/// Returns `None` and emits a `tracing::warn` for custom messages that cannot
/// be serialised (i.e. their `swink_agent::CustomMessage::to_json` /
/// `swink_agent::CustomMessage::type_name`
/// implementations return `None`). LLM messages always succeed unless
/// `serde_json` itself errors, which should not happen in practice.
#[allow(dead_code)]
fn encode(msg: &AgentMessage, session_id: &str) -> Option<(MessageKind, String)> {
    match msg {
        AgentMessage::Llm(llm) => serde_json::to_string(llm)
            .ok()
            .map(|json| (MessageKind::Llm, json)),
        AgentMessage::Custom(custom) => {
            let Some(envelope) = serialize_custom_message(custom.as_ref()) else {
                tracing::warn!(
                    session_id,
                    type_name = custom.type_name().unwrap_or("<unknown>"),
                    "skipping non-serializable CustomMessage"
                );
                return None;
            };
            serde_json::to_string(&envelope)
                .ok()
                .map(|json| (MessageKind::Custom, json))
        }
        // `AgentMessage` is `#[non_exhaustive]`: a variant this build
        // doesn't know how to encode can't be mapped to an existing
        // `MessageKind`, so skip it the same way an unserializable custom
        // message is skipped above.
        _ => {
            tracing::warn!(
                session_id,
                "skipping AgentMessage variant not recognised by this build"
            );
            None
        }
    }
}

/// Decode a `(kind, json)` pair back to an [`AgentMessage`].
///
/// - `Llm`: deserialises the JSON directly; returns `Err` on malformed data.
/// - `Custom`: returns `Ok(None)` when no `registry` is provided (the caller
///   chose not to restore custom messages). Returns `Ok(None)` on a
///   deserialisation error so callers can skip without propagating.
fn decode(
    kind: MessageKind,
    data: &str,
    registry: Option<&CustomMessageRegistry>,
) -> io::Result<Option<AgentMessage>> {
    match kind {
        MessageKind::Llm => serde_json::from_str::<LlmMessage>(data)
            .map(|m| Some(AgentMessage::Llm(m)))
            .map_err(io::Error::other),
        MessageKind::Custom => {
            let envelope: serde_json::Value =
                serde_json::from_str(data).map_err(io::Error::other)?;
            restore_single_custom(registry, &envelope)
                .map(|opt| opt.map(AgentMessage::Custom))
                .map_err(io::Error::other)
        }
    }
}

/// Encode an [`AgentMessage`] as a JSONL line.
///
/// LLM messages are stored as raw `LlmMessage` JSON for backward
/// compatibility. Custom messages are stored as their serialized envelope with
/// a `_custom: true` marker so `JSONL` can distinguish them from plain JSON
/// objects.
pub fn encode_jsonl_message_line(msg: &AgentMessage, session_id: &str) -> Option<String> {
    match msg {
        AgentMessage::Llm(llm) => serde_json::to_string(llm).ok(),
        AgentMessage::Custom(custom) => {
            let Some(envelope) = serialize_custom_message(custom.as_ref()) else {
                tracing::warn!(
                    session_id,
                    type_name = custom.type_name().unwrap_or("<unknown>"),
                    "skipping non-serializable CustomMessage"
                );
                return None;
            };

            let mut envelope = envelope;
            envelope
                .as_object_mut()
                .expect("custom message envelope must be an object")
                .insert("_custom".to_string(), serde_json::Value::Bool(true));

            serde_json::to_string(&envelope).ok()
        }
        // Same rationale as `encode` above: an unrecognised `AgentMessage`
        // variant has no JSONL encoding this build knows how to produce.
        _ => {
            tracing::warn!(
                session_id,
                "skipping AgentMessage variant not recognised by this build"
            );
            None
        }
    }
}

/// Decode one JSONL line back into an [`AgentMessage`].
///
/// Returns `Ok(None)` for state records, for entry-tagged records that are not
/// messages, and when a custom message cannot be restored because no registry
/// was supplied.
pub fn decode_jsonl_message_line(
    line: &str,
    registry: Option<&CustomMessageRegistry>,
) -> io::Result<Option<AgentMessage>> {
    let value: serde_json::Value = serde_json::from_str(line).map_err(io::Error::other)?;
    if value.get("_state").and_then(serde_json::Value::as_bool) == Some(true) {
        return Ok(None);
    }

    if value.get("_custom").and_then(serde_json::Value::as_bool) == Some(true) {
        return decode(MessageKind::Custom, line, registry);
    }

    if let Ok(message) = serde_json::from_value::<LlmMessage>(value) {
        return Ok(Some(AgentMessage::Llm(message)));
    }

    match SessionEntry::parse(line).map_err(io::Error::other)? {
        SessionEntry::Message(message) => Ok(Some(AgentMessage::Llm(message))),
        _ => Ok(None),
    }
}

/// Encode a rich [`SessionEntry`] for row-based backends such as `SQLite`.
///
/// The returned `kind` is a lightweight discriminator column, while `data`
/// contains the canonical JSON payload. Message entries keep their
/// `entry_type = "message"` tagged representation for compatibility with
/// existing rich-entry storage.
#[allow(dead_code)]
fn encode_session_entry(entry: &SessionEntry) -> io::Result<(String, String)> {
    let kind = entry.entry_type_name().to_string();
    let data = serde_json::to_string(entry).map_err(io::Error::other)?;
    Ok((kind, data))
}

/// Decode a rich [`SessionEntry`] from a row-based backend.
///
/// Returns `Ok(None)` for stored custom-agent-message rows (`kind = "custom"`),
/// which are not representable as `SessionEntry`.
#[allow(dead_code)]
fn decode_session_entry(kind: &str, data: &str) -> io::Result<Option<SessionEntry>> {
    match MessageKind::parse(kind) {
        Some(MessageKind::Llm) => serde_json::from_str::<LlmMessage>(data)
            .map(SessionEntry::Message)
            .map(Some)
            .map_err(io::Error::other),
        Some(MessageKind::Custom) => serde_json::from_str::<SessionEntry>(data)
            .map_or_else(|_| Ok(None), |entry| Ok(Some(entry))),
        None => serde_json::from_str::<SessionEntry>(data)
            .map(Some)
            .map_err(io::Error::other),
    }
}

#[cfg(test)]
mod tests {
    use std::io::{self, Write};
    use std::sync::{Arc, Mutex};

    use super::*;
    use swink_agent::{ContentBlock, CustomMessage, LlmMessage, UserMessage};
    use tracing_subscriber::fmt::MakeWriter;

    const LEAK_SENTINEL: &str = "LEAK_SENTINEL_717";

    #[derive(Debug)]
    struct NamedNonSerializableCustom {
        _secret: String,
    }

    impl CustomMessage for NamedNonSerializableCustom {
        fn as_any(&self) -> &dyn std::any::Any {
            self
        }

        fn type_name(&self) -> Option<&str> {
            Some("named_non_serializable")
        }
    }

    #[derive(Clone, Default)]
    struct SharedLogBuffer(Arc<Mutex<Vec<u8>>>);

    impl SharedLogBuffer {
        fn contents(&self) -> String {
            let bytes = self.0.lock().unwrap().clone();
            String::from_utf8(bytes).unwrap()
        }
    }

    struct SharedLogWriter(Arc<Mutex<Vec<u8>>>);

    impl Write for SharedLogWriter {
        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
            self.0.lock().unwrap().extend_from_slice(buf);
            Ok(buf.len())
        }

        fn flush(&mut self) -> io::Result<()> {
            Ok(())
        }
    }

    impl<'a> MakeWriter<'a> for SharedLogBuffer {
        type Writer = SharedLogWriter;

        fn make_writer(&'a self) -> Self::Writer {
            SharedLogWriter(Arc::clone(&self.0))
        }
    }

    fn user_msg() -> AgentMessage {
        AgentMessage::Llm(LlmMessage::User(
            UserMessage::new(vec![ContentBlock::Text {
                text: "hello".to_string(),
            }])
            .with_timestamp(1),
        ))
    }

    #[test]
    fn encode_decode_llm_roundtrip() {
        let msg = user_msg();
        let (kind, json) = encode(&msg, "test-session").expect("llm encodes");
        assert_eq!(kind, MessageKind::Llm);

        let decoded = decode(MessageKind::Llm, &json, None)
            .expect("no io error")
            .expect("decoded to Some");
        assert!(matches!(decoded, AgentMessage::Llm(LlmMessage::User(_))));
    }

    #[test]
    fn message_kind_parse_and_as_str() {
        assert_eq!(MessageKind::parse("llm"), Some(MessageKind::Llm));
        assert_eq!(MessageKind::parse("custom"), Some(MessageKind::Custom));
        assert_eq!(MessageKind::parse("unknown"), None);
        assert_eq!(MessageKind::Llm.as_str(), "llm");
        assert_eq!(MessageKind::Custom.as_str(), "custom");
    }

    #[test]
    fn decode_llm_errors_on_malformed_json() {
        let result = decode(MessageKind::Llm, "not-valid-json", None);
        assert!(result.is_err());
    }

    #[test]
    fn decode_custom_returns_none_without_registry() {
        // A minimal valid custom-message envelope.
        let envelope = serde_json::json!({
            "type_name": "MyType",
            "data": {}
        });
        let json = serde_json::to_string(&envelope).unwrap();
        let result = decode(MessageKind::Custom, &json, None).expect("no io error");
        assert!(result.is_none(), "no registry → None, not an error");
    }

    #[test]
    fn decode_jsonl_message_line_reads_tagged_message_entry() {
        let entry = SessionEntry::Message(LlmMessage::User(
            UserMessage::new(vec![ContentBlock::Text {
                text: "from-entry".to_string(),
            }])
            .with_timestamp(7),
        ));
        let line = serde_json::to_string(&entry).unwrap();

        let decoded = decode_jsonl_message_line(&line, None)
            .expect("parse succeeds")
            .expect("message preserved");
        assert!(matches!(decoded, AgentMessage::Llm(LlmMessage::User(_))));
    }

    #[test]
    fn decode_jsonl_message_line_skips_state_records() {
        let line = serde_json::json!({
            "_state": true,
            "data": { "cursor": 42 }
        })
        .to_string();

        let decoded = decode_jsonl_message_line(&line, None).expect("parse succeeds");
        assert!(decoded.is_none());
    }

    #[test]
    fn session_entry_codec_roundtrips_message_entries() {
        let entry = SessionEntry::Message(LlmMessage::User(
            UserMessage::new(vec![ContentBlock::Text {
                text: "hello".to_string(),
            }])
            .with_timestamp(1),
        ));

        let (kind, data) = encode_session_entry(&entry).unwrap();
        let decoded = decode_session_entry(&kind, &data)
            .expect("decode succeeds")
            .expect("entry preserved");

        assert!(matches!(
            decoded,
            SessionEntry::Message(LlmMessage::User(_))
        ));
    }

    #[test]
    fn session_entry_codec_preserves_custom_entries_despite_custom_kind_overlap() {
        let entry = SessionEntry::Custom {
            type_name: "audit".to_string(),
            data: serde_json::json!({"ok": true}),
            timestamp: 11,
        };

        let (kind, data) = encode_session_entry(&entry).unwrap();
        let decoded = decode_session_entry(&kind, &data)
            .expect("decode succeeds")
            .expect("entry preserved");

        assert!(matches!(decoded, SessionEntry::Custom { .. }));
    }

    #[test]
    fn encode_warning_redacts_non_serializable_custom_debug_payload() {
        let msg = AgentMessage::Custom(Box::new(NamedNonSerializableCustom {
            _secret: LEAK_SENTINEL.to_string(),
        }));

        let logs = SharedLogBuffer::default();
        let subscriber = tracing_subscriber::fmt()
            .with_max_level(tracing::Level::WARN)
            .without_time()
            .with_writer(logs.clone())
            .finish();
        let _guard = tracing::subscriber::set_default(subscriber);

        let encoded = encode(&msg, "session-123");
        let log_output = logs.contents();

        assert!(encoded.is_none());
        assert!(
            !log_output.contains(LEAK_SENTINEL),
            "warning log leaked debug payload: {log_output}"
        );
        assert!(
            log_output.contains("named_non_serializable"),
            "warning log should include custom type name: {log_output}"
        );
        assert!(
            log_output.contains("session-123"),
            "warning log should include session context: {log_output}"
        );
    }
}