orchai/src-tauri/src/models/live_agent.rs

use rusqlite::{params, Connection, Result};
use serde::{Deserialize, Serialize};
use uuid::Uuid;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LiveSession {
    pub id: String,
    pub project_id: String,
    pub agent_id: String,
    pub title: String,
    pub status: String,
    pub created_at: String,
    pub updated_at: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LiveMessage {
    pub id: String,
    pub session_id: String,
    pub sender: String,
    pub content: String,
    pub created_at: String,
}

fn session_from_row(row: &rusqlite::Row) -> rusqlite::Result<LiveSession> {
    Ok(LiveSession {
        id: row.get(0)?,
        project_id: row.get(1)?,
        agent_id: row.get(2)?,
        title: row.get(3)?,
        status: row.get(4)?,
        created_at: row.get(5)?,
        updated_at: row.get(6)?,
    })
}

fn message_from_row(row: &rusqlite::Row) -> rusqlite::Result<LiveMessage> {
    Ok(LiveMessage {
        id: row.get(0)?,
        session_id: row.get(1)?,
        sender: row.get(2)?,
        content: row.get(3)?,
        created_at: row.get(4)?,
    })
}

impl LiveSession {
    pub fn create(
        conn: &Connection,
        project_id: &str,
        agent_id: &str,
        title: &str,
    ) -> Result<LiveSession> {
        let id = Uuid::new_v4().to_string();
        let now = chrono::Utc::now().to_rfc3339();

        conn.execute(
            "INSERT INTO project_live_sessions (id, project_id, agent_id, title, status, created_at, updated_at)
             VALUES (?1, ?2, ?3, ?4, 'active', ?5, ?6)",
            params![id, project_id, agent_id, title, now, now],
        )?;

        Ok(LiveSession {
            id,
            project_id: project_id.to_string(),
            agent_id: agent_id.to_string(),
            title: title.to_string(),
            status: "active".to_string(),
            created_at: now.clone(),
            updated_at: now,
        })
    }

    pub fn list_by_project(conn: &Connection, project_id: &str) -> Result<Vec<LiveSession>> {
        let mut stmt = conn.prepare(
            "SELECT id, project_id, agent_id, title, status, created_at, updated_at
             FROM project_live_sessions
             WHERE project_id = ?1
             ORDER BY CASE status WHEN 'active' THEN 0 ELSE 1 END, updated_at DESC",
        )?;
        let rows = stmt.query_map(params![project_id], session_from_row)?;
        rows.collect()
    }

    pub fn get_by_id(conn: &Connection, id: &str) -> Result<LiveSession> {
        conn.query_row(
            "SELECT id, project_id, agent_id, title, status, created_at, updated_at
             FROM project_live_sessions
             WHERE id = ?1",
            params![id],
            session_from_row,
        )
    }

    pub fn touch(conn: &Connection, id: &str) -> Result<()> {
        let now = chrono::Utc::now().to_rfc3339();
        conn.execute(
            "UPDATE project_live_sessions SET updated_at = ?1 WHERE id = ?2",
            params![now, id],
        )?;
        Ok(())
    }

    pub fn set_archived(conn: &Connection, id: &str, archived: bool) -> Result<LiveSession> {
        let now = chrono::Utc::now().to_rfc3339();
        let status = if archived { "archived" } else { "active" };

        let affected = conn.execute(
            "UPDATE project_live_sessions SET status = ?1, updated_at = ?2 WHERE id = ?3",
            params![status, now, id],
        )?;

        if affected == 0 {
            return Err(rusqlite::Error::QueryReturnedNoRows);
        }

        Self::get_by_id(conn, id)
    }
}

impl LiveMessage {
    pub fn insert(
        conn: &Connection,
        session_id: &str,
        sender: &str,
        content: &str,
    ) -> Result<LiveMessage> {
        let id = Uuid::new_v4().to_string();
        let now = chrono::Utc::now().to_rfc3339();

        conn.execute(
            "INSERT INTO project_live_messages (id, session_id, sender, content, created_at)
             VALUES (?1, ?2, ?3, ?4, ?5)",
            params![id, session_id, sender, content, now],
        )?;

        LiveSession::touch(conn, session_id)?;

        Ok(LiveMessage {
            id,
            session_id: session_id.to_string(),
            sender: sender.to_string(),
            content: content.to_string(),
            created_at: now,
        })
    }

    pub fn list_by_session(conn: &Connection, session_id: &str) -> Result<Vec<LiveMessage>> {
        let mut stmt = conn.prepare(
            "SELECT id, session_id, sender, content, created_at
             FROM project_live_messages
             WHERE session_id = ?1
             ORDER BY created_at ASC",
        )?;
        let rows = stmt.query_map(params![session_id], message_from_row)?;
        rows.collect()
    }

    pub fn list_recent_by_session(
        conn: &Connection,
        session_id: &str,
        limit: usize,
    ) -> Result<Vec<LiveMessage>> {
        let mut stmt = conn.prepare(
            "SELECT id, session_id, sender, content, created_at
             FROM project_live_messages
             WHERE session_id = ?1
             ORDER BY created_at DESC
             LIMIT ?2",
        )?;
        let mut messages: Vec<LiveMessage> = stmt
            .query_map(params![session_id, limit as i64], message_from_row)?
            .collect::<Result<Vec<_>>>()?;
        messages.reverse();
        Ok(messages)
    }

    pub fn update_content(conn: &Connection, id: &str, content: &str) -> Result<()> {
        conn.execute(
            "UPDATE project_live_messages SET content = ?1 WHERE id = ?2",
            params![content, id],
        )?;
        Ok(())
    }

    pub fn delete(conn: &Connection, id: &str) -> Result<()> {
        conn.execute(
            "DELETE FROM project_live_messages WHERE id = ?1",
            params![id],
        )?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::db;
    use crate::models::agent::{Agent, AgentRole, AgentTool};
    use crate::models::project::Project;

    fn setup() -> (Connection, String, String) {
        let conn = db::init_in_memory().expect("db init should succeed");
        let project = Project::insert(&conn, "Live Project", "/tmp/live-project", None, "main")
            .expect("project insert should succeed");
        let agent = Agent::insert(
            &conn,
            "Live Agent",
            AgentRole::Analyst,
            AgentTool::Codex,
            "",
        )
        .expect("agent insert should succeed");

        (conn, project.id, agent.id)
    }

    #[test]
    fn test_create_session_and_list_messages() {
        let (conn, project_id, agent_id) = setup();
        let session = LiveSession::create(&conn, &project_id, &agent_id, "Session 1")
            .expect("session create");

        let user_message =
            LiveMessage::insert(&conn, &session.id, "user", "Bonjour").expect("message insert");
        let agent_message =
            LiveMessage::insert(&conn, &session.id, "agent", "Salut").expect("message insert");

        let sessions =
            LiveSession::list_by_project(&conn, &project_id).expect("session list should work");
        let messages =
            LiveMessage::list_by_session(&conn, &session.id).expect("message list should work");

        assert_eq!(sessions.len(), 1);
        assert_eq!(messages.len(), 2);
        assert_eq!(messages[0].id, user_message.id);
        assert_eq!(messages[1].id, agent_message.id);
    }

    #[test]
    fn test_update_and_delete_message() {
        let (conn, project_id, agent_id) = setup();
        let session = LiveSession::create(&conn, &project_id, &agent_id, "Session 2")
            .expect("session create");
        let message = LiveMessage::insert(&conn, &session.id, "agent", "")
            .expect("message insert should work");

        LiveMessage::update_content(&conn, &message.id, "Streaming done")
            .expect("message update should work");

        let messages =
            LiveMessage::list_by_session(&conn, &session.id).expect("message list should work");
        assert_eq!(messages.len(), 1);
        assert_eq!(messages[0].content, "Streaming done");

        LiveMessage::delete(&conn, &message.id).expect("message delete should work");
        let messages_after_delete =
            LiveMessage::list_by_session(&conn, &session.id).expect("message list should work");
        assert!(messages_after_delete.is_empty());
    }

    #[test]
    fn test_archive_session_updates_status_and_sort_order() {
        let (conn, project_id, agent_id) = setup();
        let archived =
            LiveSession::create(&conn, &project_id, &agent_id, "Archivee").expect("session create");
        let active =
            LiveSession::create(&conn, &project_id, &agent_id, "Active").expect("session create");

        let archived_session =
            LiveSession::set_archived(&conn, &archived.id, true).expect("archive should work");
        assert_eq!(archived_session.status, "archived");

        let sessions =
            LiveSession::list_by_project(&conn, &project_id).expect("session list should work");
        assert_eq!(sessions.len(), 2);
        assert_eq!(sessions[0].id, active.id);
        assert_eq!(sessions[0].status, "active");
        assert_eq!(sessions[1].id, archived.id);
        assert_eq!(sessions[1].status, "archived");
    }
}