agent-orchestrator/agor-pro/src/memory.rs

// SPDX-License-Identifier: LicenseRef-Commercial
// Persistent Agent Memory — project-scoped structured fragments that survive sessions.

use rusqlite::params;
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct MemoryFragment {
    pub id: i64,
    pub project_id: String,
    pub content: String,
    pub source: String,
    pub trust: String,
    pub confidence: f64,
    pub created_at: i64,
    pub ttl_days: i64,
    pub tags: String,
}

fn ensure_tables(conn: &rusqlite::Connection) -> Result<(), String> {
    conn.execute_batch(
        "CREATE TABLE IF NOT EXISTS pro_memories (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            project_id TEXT NOT NULL,
            content TEXT NOT NULL,
            source TEXT NOT NULL DEFAULT '',
            trust TEXT NOT NULL DEFAULT 'agent',
            confidence REAL NOT NULL DEFAULT 1.0,
            created_at INTEGER NOT NULL,
            ttl_days INTEGER NOT NULL DEFAULT 90,
            tags TEXT NOT NULL DEFAULT ''
        );
        CREATE VIRTUAL TABLE IF NOT EXISTS pro_memories_fts USING fts5(
            content, tags, content=pro_memories, content_rowid=id
        );
        CREATE TRIGGER IF NOT EXISTS pro_memories_ai AFTER INSERT ON pro_memories BEGIN
            INSERT INTO pro_memories_fts(rowid, content, tags) VALUES (new.id, new.content, new.tags);
        END;
        CREATE TRIGGER IF NOT EXISTS pro_memories_ad AFTER DELETE ON pro_memories BEGIN
            INSERT INTO pro_memories_fts(pro_memories_fts, rowid, content, tags)
                VALUES ('delete', old.id, old.content, old.tags);
        END;
        CREATE TRIGGER IF NOT EXISTS pro_memories_au AFTER UPDATE ON pro_memories BEGIN
            INSERT INTO pro_memories_fts(pro_memories_fts, rowid, content, tags)
                VALUES ('delete', old.id, old.content, old.tags);
            INSERT INTO pro_memories_fts(rowid, content, tags) VALUES (new.id, new.content, new.tags);
        END;"
    ).map_err(|e| format!("Failed to create memory tables: {e}"))
}

fn now_epoch() -> i64 {
    super::analytics::now_epoch()
}

fn prune_expired(conn: &rusqlite::Connection) -> Result<(), String> {
    let now = now_epoch();
    conn.execute(
        "DELETE FROM pro_memories WHERE created_at + (ttl_days * 86400) < ?1",
        params![now],
    ).map_err(|e| format!("Prune failed: {e}"))?;
    Ok(())
}

fn row_to_fragment(row: &rusqlite::Row) -> rusqlite::Result<MemoryFragment> {
    Ok(MemoryFragment {
        id: row.get("id")?,
        project_id: row.get("project_id")?,
        content: row.get("content")?,
        source: row.get("source")?,
        trust: row.get("trust")?,
        confidence: row.get("confidence")?,
        created_at: row.get("created_at")?,
        ttl_days: row.get("ttl_days")?,
        tags: row.get("tags")?,
    })
}

#[tauri::command]
pub fn pro_memory_add(
    project_id: String,
    content: String,
    source: Option<String>,
    tags: Option<String>,
) -> Result<i64, String> {
    if content.len() > 10000 {
        return Err("Memory content too long (max 10000 chars)".into());
    }
    let conn = super::open_sessions_db()?;
    ensure_tables(&conn)?;

    // Per-project memory cap
    let count: i64 = conn.query_row(
        "SELECT COUNT(*) FROM pro_memories WHERE project_id = ?1",
        params![project_id], |row| row.get(0)
    ).unwrap_or(0);
    if count >= 1000 {
        return Err("Memory limit reached for this project (max 1000 fragments)".into());
    }

    let ts = now_epoch();
    let src = source.unwrap_or_default();
    let tgs = tags.unwrap_or_default();
    conn.execute(
        "INSERT INTO pro_memories (project_id, content, source, created_at, tags) VALUES (?1, ?2, ?3, ?4, ?5)",
        params![project_id, content, src, ts, tgs],
    ).map_err(|e| format!("Failed to add memory: {e}"))?;
    Ok(conn.last_insert_rowid())
}

#[tauri::command]
pub fn pro_memory_list(project_id: String, limit: Option<i64>) -> Result<Vec<MemoryFragment>, String> {
    let conn = super::open_sessions_db()?;
    ensure_tables(&conn)?;
    prune_expired(&conn)?;
    let lim = limit.unwrap_or(50);

    let mut stmt = conn.prepare(
        "SELECT id, project_id, content, source, trust, confidence, created_at, ttl_days, tags
         FROM pro_memories WHERE project_id = ?1 ORDER BY created_at DESC LIMIT ?2"
    ).map_err(|e| format!("Query failed: {e}"))?;

    let rows = stmt.query_map(params![project_id, lim], row_to_fragment)
        .map_err(|e| format!("Query failed: {e}"))?
        .collect::<Result<Vec<_>, _>>()
        .map_err(|e| format!("Row read failed: {e}"))?;

    Ok(rows)
}

#[tauri::command]
pub fn pro_memory_search(project_id: String, query: String) -> Result<Vec<MemoryFragment>, String> {
    let conn = super::open_sessions_db()?;
    ensure_tables(&conn)?;
    prune_expired(&conn)?;

    // Sanitize query to prevent FTS5 operator injection
    let safe_query = format!("\"{}\"", query.replace('"', "\"\""));

    let mut stmt = conn.prepare(
        "SELECT m.id, m.project_id, m.content, m.source, m.trust, m.confidence, m.created_at, m.ttl_days, m.tags
         FROM pro_memories m
         JOIN pro_memories_fts f ON m.id = f.rowid
         WHERE f.pro_memories_fts MATCH ?1 AND m.project_id = ?2
         ORDER BY rank LIMIT 20"
    ).map_err(|e| format!("Search query failed: {e}"))?;

    let rows = stmt.query_map(params![safe_query, project_id], row_to_fragment)
        .map_err(|e| format!("Search failed: {e}"))?
        .collect::<Result<Vec<_>, _>>()
        .map_err(|e| format!("Row read failed: {e}"))?;

    Ok(rows)
}

#[tauri::command]
pub fn pro_memory_update(
    id: i64,
    content: Option<String>,
    trust: Option<String>,
    confidence: Option<f64>,
) -> Result<(), String> {
    let conn = super::open_sessions_db()?;
    ensure_tables(&conn)?;

    let tx = conn.unchecked_transaction()
        .map_err(|e| format!("Transaction failed: {e}"))?;

    if let Some(c) = content {
        tx.execute("UPDATE pro_memories SET content = ?2 WHERE id = ?1", params![id, c])
            .map_err(|e| format!("Update content failed: {e}"))?;
    }
    if let Some(t) = trust {
        tx.execute("UPDATE pro_memories SET trust = ?2 WHERE id = ?1", params![id, t])
            .map_err(|e| format!("Update trust failed: {e}"))?;
    }
    if let Some(c) = confidence {
        tx.execute("UPDATE pro_memories SET confidence = ?2 WHERE id = ?1", params![id, c])
            .map_err(|e| format!("Update confidence failed: {e}"))?;
    }

    tx.commit().map_err(|e| format!("Commit failed: {e}"))?;
    Ok(())
}

#[tauri::command]
pub fn pro_memory_delete(id: i64) -> Result<(), String> {
    let conn = super::open_sessions_db()?;
    ensure_tables(&conn)?;
    conn.execute("DELETE FROM pro_memories WHERE id = ?1", params![id])
        .map_err(|e| format!("Delete failed: {e}"))?;
    Ok(())
}

#[tauri::command]
pub fn pro_memory_inject(project_id: String, max_tokens: Option<i64>) -> Result<String, String> {
    let conn = super::open_sessions_db()?;
    ensure_tables(&conn)?;
    prune_expired(&conn)?;

    let max_chars = (max_tokens.unwrap_or(2000) * 3) as usize; // ~3 chars per token heuristic

    let mut stmt = conn.prepare(
        "SELECT content, trust, confidence FROM pro_memories
         WHERE project_id = ?1 ORDER BY confidence DESC, created_at DESC"
    ).map_err(|e| format!("Query failed: {e}"))?;

    let entries: Vec<(String, String, f64)> = stmt
        .query_map(params![project_id], |row| Ok((row.get("content")?, row.get("trust")?, row.get("confidence")?)))
        .map_err(|e| format!("Query failed: {e}"))?
        .collect::<Result<Vec<_>, _>>()
        .map_err(|e| format!("Row read failed: {e}"))?;

    let mut md = String::from("## Project Memory\n\n");
    let mut chars = md.len();

    for (content, trust, confidence) in &entries {
        let line = format!("- [{}|{:.1}] {}\n", trust, confidence, content);
        if chars + line.len() > max_chars {
            break;
        }
        md.push_str(&line);
        chars += line.len();
    }

    Ok(md)
}

#[tauri::command]
pub fn pro_memory_extract_from_session(
    project_id: String,
    session_messages_json: String,
) -> Result<Vec<MemoryFragment>, String> {
    let conn = super::open_sessions_db()?;
    ensure_tables(&conn)?;

    let messages: Vec<serde_json::Value> = serde_json::from_str(&session_messages_json)
        .map_err(|e| format!("Invalid JSON: {e}"))?;

    let ts = now_epoch();
    let mut extracted = Vec::new();

    // Patterns to extract: decisions, file references, errors
    let decision_patterns = ["decision:", "chose ", "decided to ", "instead of "];
    let error_patterns = ["error:", "failed:", "Error:", "panic", "FAILED"];

    for msg in &messages {
        let content = msg.get("content").and_then(|c| c.as_str()).unwrap_or("");

        // Extract decisions
        for pattern in &decision_patterns {
            if content.contains(pattern) {
                let fragment_content = extract_surrounding(content, pattern, 200);
                conn.execute(
                    "INSERT INTO pro_memories (project_id, content, source, trust, confidence, created_at, tags)
                     VALUES (?1, ?2, 'auto-extract', 'auto', 0.7, ?3, 'decision')",
                    params![project_id, fragment_content, ts],
                ).map_err(|e| format!("Insert failed: {e}"))?;
                let id = conn.last_insert_rowid();
                extracted.push(MemoryFragment {
                    id,
                    project_id: project_id.clone(),
                    content: fragment_content,
                    source: "auto-extract".into(),
                    trust: "auto".into(),
                    confidence: 0.7,
                    created_at: ts,
                    ttl_days: 90,
                    tags: "decision".into(),
                });
                break; // One extraction per message
            }
        }

        // Extract errors
        for pattern in &error_patterns {
            if content.contains(pattern) {
                let fragment_content = extract_surrounding(content, pattern, 300);
                conn.execute(
                    "INSERT INTO pro_memories (project_id, content, source, trust, confidence, created_at, tags)
                     VALUES (?1, ?2, 'auto-extract', 'auto', 0.6, ?3, 'error')",
                    params![project_id, fragment_content, ts],
                ).map_err(|e| format!("Insert failed: {e}"))?;
                let id = conn.last_insert_rowid();
                extracted.push(MemoryFragment {
                    id,
                    project_id: project_id.clone(),
                    content: fragment_content,
                    source: "auto-extract".into(),
                    trust: "auto".into(),
                    confidence: 0.6,
                    created_at: ts,
                    ttl_days: 90,
                    tags: "error".into(),
                });
                break;
            }
        }
    }

    Ok(extracted)
}

/// Extract surrounding text around a pattern match, up to max_chars.
fn extract_surrounding(text: &str, pattern: &str, max_chars: usize) -> String {
    if let Some(pos) = text.find(pattern) {
        let start = pos.saturating_sub(50);
        let end = (pos + max_chars).min(text.len());
        // Ensure valid UTF-8 boundaries
        let start = text.floor_char_boundary(start);
        let end = text.ceil_char_boundary(end);
        text[start..end].to_string()
    } else {
        text.chars().take(max_chars).collect()
    }
}

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

    #[test]
    fn test_memory_fragment_serializes_camel_case() {
        let f = MemoryFragment {
            id: 1,
            project_id: "proj1".into(),
            content: "We decided to use SQLite".into(),
            source: "session-abc".into(),
            trust: "agent".into(),
            confidence: 0.9,
            created_at: 1710000000,
            ttl_days: 90,
            tags: "decision,architecture".into(),
        };
        let json = serde_json::to_string(&f).unwrap();
        assert!(json.contains("projectId"));
        assert!(json.contains("createdAt"));
        assert!(json.contains("ttlDays"));
    }

    #[test]
    fn test_memory_fragment_deserializes() {
        let json = r#"{"id":1,"projectId":"p","content":"test","source":"s","trust":"human","confidence":1.0,"createdAt":0,"ttlDays":30,"tags":"t"}"#;
        let f: MemoryFragment = serde_json::from_str(json).unwrap();
        assert_eq!(f.project_id, "p");
        assert_eq!(f.trust, "human");
    }

    #[test]
    fn test_extract_surrounding() {
        let text = "We chose SQLite instead of PostgreSQL for simplicity";
        let result = extract_surrounding(text, "chose ", 100);
        assert!(result.contains("chose SQLite"));
    }
}