poimen-memory/crates/mem-core/src/domain.rs

use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::fmt;
use std::str::FromStr;
use time::OffsetDateTime;

/// Memory node level in the hierarchy.
/// Closed. L0 evidence, L1 per-query memory, L2 project synthesis.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "UPPERCASE")]
pub enum Level {
    #[serde(rename = "L0")]
    L0,
    #[serde(rename = "L1")]
    L1,
    #[serde(rename = "L2")]
    L2,
}

impl fmt::Display for Level {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Level::L0 => write!(f, "L0"),
            Level::L1 => write!(f, "L1"),
            Level::L2 => write!(f, "L2"),
        }
    }
}

impl FromStr for Level {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "L0" => Ok(Level::L0),
            "L1" => Ok(Level::L1),
            "L2" => Ok(Level::L2),
            _ => Err(format!("Invalid level: {}", s)),
        }
    }
}

/// The role of a record in a conversation.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "PascalCase")]
pub enum Role {
    User,
    Assistant,
    ToolResult,
    System,
}

/// Source identification and offset.
#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct Provenance {
    pub source_id: String,
    pub offset: u64,
}

/// A normalized unit from any source.
/// Adapters produce these; nothing downstream learns whether it came from pi,
/// claude, or a socket.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Record {
    pub role: Role,
    pub text: String,
    #[serde(with = "time::serde::rfc3339")]
    pub timestamp: OffsetDateTime,
    pub provenance: Provenance,
}

/// Content hash as a hex string.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize, PartialOrd, Ord)]
pub struct Sha256Hash([u8; 32]);

impl Sha256Hash {
    /// Create a hash from a byte array.
    pub fn from_bytes(bytes: [u8; 32]) -> Self {
        Sha256Hash(bytes)
    }

    /// Create a hash from a hex string.
    pub fn from_hex(hex: &str) -> Result<Self, String> {
        if hex.len() != 64 {
            return Err("Hash must be 64 hex characters".to_string());
        }
        let mut bytes = [0u8; 32];
        for (i, chunk) in hex.as_bytes().chunks(2).enumerate() {
            bytes[i] = u8::from_str_radix(std::str::from_utf8(chunk).unwrap(), 16)
                .map_err(|_| "Invalid hex character".to_string())?;
        }
        Ok(Sha256Hash(bytes))
    }

    /// Convert to hex string.
    pub fn to_hex(&self) -> String {
        hex::encode(self.0)
    }

    /// Get the raw bytes.
    pub fn as_bytes(&self) -> &[u8] {
        &self.0
    }
}

impl fmt::Display for Sha256Hash {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.to_hex())
    }
}

/// Newtype wrappers with no Default.
#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct ProjectId(String);

impl ProjectId {
    pub fn new(id: String) -> Result<Self, String> {
        if id.is_empty() {
            return Err("ProjectId cannot be empty".to_string());
        }
        Ok(ProjectId(id))
    }

    pub fn as_str(&self) -> &str {
        &self.0
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct QueryId(String);

impl QueryId {
    pub fn new(id: String) -> Result<Self, String> {
        if id.is_empty() {
            return Err("QueryId cannot be empty".to_string());
        }
        Ok(QueryId(id))
    }

    pub fn as_str(&self) -> &str {
        &self.0
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct RunId(String);

impl RunId {
    pub fn new(id: String) -> Result<Self, String> {
        if id.is_empty() {
            return Err("RunId cannot be empty".to_string());
        }
        Ok(RunId(id))
    }

    pub fn as_str(&self) -> &str {
        &self.0
    }
}

/// One or more Records, under the token budget, never split mid-Record.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Chunk {
    pub t: u32,  // 1-based turn index within a run
    pub records: Vec<Record>,
    pub tokens: usize,
    #[serde(skip)]
    sha256: Option<Sha256Hash>,
}

impl Chunk {
    pub fn new(t: u32, records: Vec<Record>, tokens: usize) -> Self {
        Chunk {
            t,
            records,
            tokens,
            sha256: None,
        }
    }

    /// Compute canonical hash over concatenated record texts and their provenance.
    /// Must not include timestamp or run_id to ensure rebuild idempotence.
    pub fn content_hash(&mut self) -> Sha256Hash {
        if let Some(hash) = self.sha256 {
            return hash;
        }

        let mut hasher = Sha256::new();

        // Concatenate record texts and provenance
        for record in &self.records {
            hasher.update(record.role.to_string().as_bytes());
            hasher.update(b"\x00");
            hasher.update(record.text.as_bytes());
            hasher.update(b"\x00");
            hasher.update(record.provenance.source_id.as_bytes());
            hasher.update(b"\x00");
            hasher.update(record.provenance.offset.to_le_bytes());
            hasher.update(b"\x00");
        }

        let bytes: [u8; 32] = hasher.finalize().into();
        let hash = Sha256Hash::from_bytes(bytes);
        self.sha256 = Some(hash);
        hash
    }

    pub fn sha256(&mut self) -> Sha256Hash {
        self.content_hash()
    }
}

impl fmt::Display for Role {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Role::User => write!(f, "User"),
            Role::Assistant => write!(f, "Assistant"),
            Role::ToolResult => write!(f, "ToolResult"),
            Role::System => write!(f, "System"),
        }
    }
}

/// A memory node at any level.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct MemoryNode {
    pub level: Level,
    pub project: ProjectId,
    pub query_id: Option<QueryId>,  // None at L2
    pub run_id: RunId,
    pub t: u32,
    pub text: String,
    #[serde(skip)]
    sha256: Option<Sha256Hash>,
    pub parents: Vec<Sha256Hash>,
}

impl MemoryNode {
    pub fn new(
        level: Level,
        project: ProjectId,
        query_id: Option<QueryId>,
        run_id: RunId,
        t: u32,
        text: String,
        parents: Vec<Sha256Hash>,
    ) -> Self {
        MemoryNode {
            level,
            project,
            query_id,
            run_id,
            t,
            text,
            sha256: None,
            parents,
        }
    }

    /// Compute canonical hash over (level, project, query_id, text).
    /// Must not include timestamp or run_id to ensure rebuild idempotence.
    pub fn content_hash(&mut self) -> Sha256Hash {
        if let Some(hash) = self.sha256 {
            return hash;
        }

        let mut hasher = Sha256::new();

        hasher.update(self.level.to_string().as_bytes());
        hasher.update(b"\x00");
        hasher.update(self.project.as_str().as_bytes());
        hasher.update(b"\x00");
        
        if let Some(query_id) = &self.query_id {
            hasher.update(query_id.as_str().as_bytes());
        }
        hasher.update(b"\x00");
        
        hasher.update(self.text.as_bytes());

        let bytes: [u8; 32] = hasher.finalize().into();
        let hash = Sha256Hash::from_bytes(bytes);
        self.sha256 = Some(hash);
        hash
    }

    pub fn sha256(&mut self) -> Sha256Hash {
        self.content_hash()
    }
}

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

    #[test]
    fn test_level_serialization() {
        assert_eq!(serde_json::to_string(&Level::L0).unwrap(), "\"L0\"");
        assert_eq!(serde_json::to_string(&Level::L1).unwrap(), "\"L1\"");
        assert_eq!(serde_json::to_string(&Level::L2).unwrap(), "\"L2\"");
    }

    #[test]
    fn test_level_round_trip() {
        for level in &[Level::L0, Level::L1, Level::L2] {
            let json = serde_json::to_string(level).unwrap();
            let deserialized: Level = serde_json::from_str(&json).unwrap();
            assert_eq!(level, &deserialized);
        }
    }

    #[test]
    fn test_role_round_trip() {
        for role in &[Role::User, Role::Assistant, Role::ToolResult, Role::System] {
            let json = serde_json::to_string(role).unwrap();
            let deserialized: Role = serde_json::from_str(&json).unwrap();
            assert_eq!(role, &deserialized);
        }
    }

    #[test]
    fn test_sha256_hash_round_trip() {
        let original = Sha256Hash::from_bytes([1; 32]);
        let json = serde_json::to_string(&original).unwrap();
        let deserialized: Sha256Hash = serde_json::from_str(&json).unwrap();
        assert_eq!(original, deserialized);
    }

    #[test]
    fn test_project_id_creation() {
        let id = ProjectId::new("project1".to_string()).unwrap();
        assert_eq!(id.as_str(), "project1");

        let result = ProjectId::new("".to_string());
        assert!(result.is_err());
    }

    #[test]
    fn test_record_round_trip() {
        let record = Record {
            role: Role::User,
            text: "Hello, world!".to_string(),
            timestamp: datetime!(2024-08-20 12:00:00 UTC),
            provenance: Provenance {
                source_id: "session1".to_string(),
                offset: 0,
            },
        };

        let json = serde_json::to_string(&record).unwrap();
        let deserialized: Record = serde_json::from_str(&json).unwrap();
        assert_eq!(record.role, deserialized.role);
        assert_eq!(record.text, deserialized.text);
        assert_eq!(record.provenance, deserialized.provenance);
    }

    #[test]
    fn test_chunk_round_trip() {
        let record = Record {
            role: Role::User,
            text: "Hello".to_string(),
            timestamp: datetime!(2024-08-20 12:00:00 UTC),
            provenance: Provenance {
                source_id: "session1".to_string(),
                offset: 0,
            },
        };

        let chunk = Chunk {
            t: 1,
            records: vec![record],
            tokens: 2,
            sha256: None,
        };

        let json = serde_json::to_string(&chunk).unwrap();
        let deserialized: Chunk = serde_json::from_str(&json).unwrap();
        assert_eq!(chunk.t, deserialized.t);
        assert_eq!(chunk.tokens, deserialized.tokens);
    }

    #[test]
    fn test_memory_node_round_trip() {
        let node = MemoryNode {
            level: Level::L0,
            project: ProjectId::new("p1".to_string()).unwrap(),
            query_id: Some(QueryId::new("q1".to_string()).unwrap()),
            run_id: RunId::new("r1".to_string()).unwrap(),
            t: 1,
            text: "test".to_string(),
            sha256: None,
            parents: vec![],
        };

        let json = serde_json::to_string(&node).unwrap();
        let deserialized: MemoryNode = serde_json::from_str(&json).unwrap();
        assert_eq!(node.level, deserialized.level);
        assert_eq!(node.t, deserialized.t);
        assert_eq!(node.text, deserialized.text);
    }
}