cortex-agent 0.5.2

Self-learning AI agent with persistent memory, skills, project awareness, and a beautiful terminal UI
use serde::{Deserialize, Serialize};
use std::path::Path;

/// Redact sensitive patterns (API keys, tokens, passwords) from a string.
/// Replaces matches with `***` while preserving surrounding context.
pub fn redact_sensitive(text: &str) -> String {
    // Common patterns: sk-..., api_key=..., token=..., password=...
    let patterns: [(&str, &str); 6] = [
        // OpenAI-style keys: sk-... (up to 51 chars)
        (r"(?i)(sk-[a-zA-Z0-9]{20,51})", "sk-***"),
        // Bearer tokens
        (r"(?i)(Bearer\s+)[a-zA-Z0-9._-]{20,}", "${1}***"),
        // api_key=xxx, api_key: xxx
        (r"(?i)(api_key[\s:=]+)[a-zA-Z0-9._-]{8,}", "${1}***"),
        // token=xxx, token: xxx
        (r"(?i)(token[\s:=]+)[a-zA-Z0-9._-]{8,}", "${1}***"),
        // password=xxx
        (r"(?i)(password[\s:=]+)[a-zA-Z0-9!@#$%^&*()_+=-]{4,}", "${1}***"),
        // Authorization: Bearer xxx
        (r"(?i)(Authorization:\s*Bearer\s+)[a-zA-Z0-9._-]+", "${1}***"),
    ];
    let mut result = text.to_string();
    for (pattern, replacement) in &patterns {
        if let Ok(re) = regex::Regex::new(pattern) {
            result = re.replace_all(&result, *replacement).to_string();
        }
    }
    result
}

/// Runtime session state โ€” last used provider/model, saved between runs.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionState {
    pub last_provider: String,
    pub last_model: String,
    pub last_session_turns: u64,
    pub last_session_tokens: u64,
    pub last_session_duration_secs: u64,
}

impl Default for SessionState {
    fn default() -> Self {
        Self {
            last_provider: String::new(),
            last_model: String::new(),
            last_session_turns: 0,
            last_session_tokens: 0,
            last_session_duration_secs: 0,
        }
    }
}

fn session_path() -> String {
    let home = std::env::var("HOME").unwrap_or_else(|_| ".".into());
    format!("{}/.cortex/session.json", home)
}

/// Save session state to `~/.cortex/session.json`.
pub fn save_session_state(state: &SessionState) {
    let path = session_path();
    if let Some(parent) = Path::new(&path).parent() {
        let _ = std::fs::create_dir_all(parent);
    }
    if let Ok(json) = serde_json::to_string_pretty(state) {
        let _ = std::fs::write(&path, &json);
    }
}

/// Load session state from `~/.cortex/session.json`.
pub fn load_session_state() -> SessionState {
    let path = session_path();
    std::fs::read_to_string(&path)
        .ok()
        .and_then(|s| serde_json::from_str(&s).ok())
        .unwrap_or_default()
}

/// Persist a named session (conversation messages) to `~/.cortex/sessions/<name>.json`.
pub fn save_named_session(name: &str, messages: &[crate::messages::Message]) -> Result<(), String> {
    let home = std::env::var("HOME").map_err(|_| "HOME not set".to_string())?;
    let dir = Path::new(&home).join(".cortex").join("sessions");
    std::fs::create_dir_all(&dir).map_err(|e| format!("Cannot create sessions dir: {}", e))?;
    let path = dir.join(format!("{}.json", name));
    let data = serde_json::to_string_pretty(messages).map_err(|e| format!("Serialize: {}", e))?;
    std::fs::write(&path, &data).map_err(|e| format!("Write: {}", e))?;
    Ok(())
}

/// Load a named session from `~/.cortex/sessions/<name>.json`.
pub fn load_named_session(name: &str) -> Result<Vec<crate::messages::Message>, String> {
    let home = std::env::var("HOME").map_err(|_| "HOME not set".to_string())?;
    let path = Path::new(&home).join(".cortex").join("sessions").join(format!("{}.json", name));
    let data = std::fs::read_to_string(&path).map_err(|e| format!("Cannot read session '{}': {}", name, e))?;
    serde_json::from_str(&data).map_err(|e| format!("Parse: {}", e))
}

/// List all named sessions.
pub fn list_named_sessions() -> Vec<String> {
    let home = match std::env::var("HOME") {
        Ok(h) => h,
        Err(_) => return vec![],
    };
    let dir = Path::new(&home).join(".cortex").join("sessions");
    let _ = std::fs::create_dir_all(&dir);
    let mut sessions: Vec<String> = match std::fs::read_dir(&dir) {
        Ok(entries) => entries
            .filter_map(|e| e.ok())
            .filter(|e| e.path().extension().map_or(false, |ext| ext == "json"))
            .filter_map(|e| e.path().file_stem().map(|s| s.to_string_lossy().to_string()))
            .collect(),
        Err(_) => return vec![],
    };
    sessions.sort();
    sessions
}

/// Path to the auto-recovery file for crash recovery.
fn recovery_path() -> String {
    let home = std::env::var("HOME").unwrap_or_else(|_| ".".into());
    format!("{}/.cortex/recovery.json", home)
}

/// Auto-save current messages to a recovery file (crash recovery).
pub fn auto_save_messages(messages: &[crate::messages::Message]) {
    let path = recovery_path();
    if let Some(parent) = Path::new(&path).parent() {
        let _ = std::fs::create_dir_all(parent);
    }
    if let Ok(json) = serde_json::to_string(messages) {
        let _ = std::fs::write(&path, &json);
    }
}

/// Load messages from the recovery file (after a crash).
pub fn load_recovery_messages() -> Option<Vec<crate::messages::Message>> {
    let path = recovery_path();
    let data = std::fs::read_to_string(&path).ok()?;
    serde_json::from_str(&data).ok()
}

/// Delete the recovery file after a clean exit.
pub fn clear_recovery_file() {
    let path = recovery_path();
    let _ = std::fs::remove_file(&path);
}

/// Check if a recovery file exists from a previous session.
pub fn has_recovery_file() -> bool {
    Path::new(&recovery_path()).exists()
}

/// Export the current session as a markdown transcript.
pub fn export_session_as_md(messages: &[crate::messages::Message]) -> String {
    let mut md = String::new();
    md.push_str("# Cortex Session Transcript\n\n");
    for msg in messages {
        match msg.role.as_str() {
            "system" => {
                let preview = msg.content.as_deref().unwrap_or("").chars().take(100).collect::<String>();
                md.push_str(&format!("> **System:** {}โ€ฆ\n\n", redact_sensitive(&preview)));
            }
            "user" => {
                md.push_str(&format!("**User:**\n{}\n\n", redact_sensitive(msg.content.as_deref().unwrap_or(""))));
            }
            "assistant" => {
                if let Some(ref content) = msg.content {
                    md.push_str(&format!("**Assistant:**\n{}\n\n", redact_sensitive(content)));
                }
                if let Some(ref calls) = msg.tool_calls {
                    for tc in calls {
                        md.push_str(&format!("_Tool: {}_\n\n", tc.name));
                    }
                }
            }
            "tool" => {
                let preview = msg.content.as_deref().unwrap_or("").chars().take(80).collect::<String>();
                md.push_str(&format!("> **Tool ({}):** {}โ€ฆ\n\n", msg.name.as_deref().unwrap_or("?"), redact_sensitive(&preview)));
            }
            _ => {}
        }
    }
    md
}

/// Export current session messages as a standalone HTML file.
pub fn export_session_as_html(messages: &[crate::messages::Message]) -> String {
    let provider = std::env::var("ACTIVE_PROVIDER").unwrap_or_default();
    let model = std::env::var("ACTIVE_MODEL").unwrap_or_default();
    let mut html = String::new();
    html.push_str(r#"<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>Cortex Session Transcript</title>
<style>
  :root { --bg: #1e1e2e; --surface: #181825; --text: #cdd6f4; --subtext: #a6adc8;
    --user: #89b4fa; --assistant: #a6e3a1; --system: #f9e2af; --tool: #fab387;
    --border: #313244; }
  body { font-family: 'SF Mono', 'Fira Code', 'Cascadia Code', monospace;
    background: var(--bg); color: var(--text); max-width: 860px; margin: 0 auto;
    padding: 2rem; line-height: 1.6; }
  h1 { color: var(--text); border-bottom: 2px solid var(--border); padding-bottom: 0.5rem; }
  .meta { color: var(--subtext); font-size: 0.85rem; margin-bottom: 2rem; }
  .msg { margin: 1rem 0; padding: 0.75rem 1rem; border-radius: 8px; background: var(--surface); }
  .msg .role { font-weight: 700; font-size: 0.8rem; text-transform: uppercase; letter-spacing: 0.05em; margin-bottom: 0.4rem; }
  .msg.user { border-left: 3px solid var(--user); }
  .msg.user .role { color: var(--user); }
  .msg.assistant { border-left: 3px solid var(--assistant); }
  .msg.assistant .role { color: var(--assistant); }
  .msg.system { border-left: 3px solid var(--system); }
  .msg.system .role { color: var(--system); }
  .msg.tool { border-left: 3px solid var(--tool); }
  .msg.tool .role { color: var(--tool); }
  .content { white-space: pre-wrap; word-break: break-word; }
  .tool-call { font-size: 0.85rem; color: var(--subtext); margin-top: 0.4rem; }
  code { background: var(--bg); padding: 0.1rem 0.3rem; border-radius: 4px; font-size: 0.9em; }
  pre { background: var(--bg); padding: 1rem; border-radius: 8px; overflow-x: auto; }
</style>
</head>
<body>
<h1>๐Ÿง  Cortex Session</h1>
<div class="meta">
"#);
    if !provider.is_empty() {
        html.push_str(&format!("Provider: {} | Model: {}<br>\n", escape_html(&provider), escape_html(&model)));
    }
    html.push_str(&format!("Messages: {} | Exported: {}</div>\n", messages.len(), chrono::Local::now().format("%Y-%m-%d %H:%M")));

    for msg in messages {
        let role_class = match msg.role.as_str() {
            "system" => "system",
            "user" => "user",
            "assistant" => "assistant",
            "tool" => "tool",
            _ => "system",
        };
        html.push_str(&format!("<div class=\"msg {}\">\n<div class=\"role\">{}</div>\n", role_class, escape_html(&msg.role)));
        if let Some(ref content) = msg.content {
            let escaped = escape_html(content);
            if content.contains("```") {
                // Simple code block rendering
                let parts: Vec<&str> = content.split("```").collect();
                for (i, part) in parts.iter().enumerate() {
                    if i % 2 == 0 {
                        if !part.trim().is_empty() {
                            html.push_str(&format!("<div class=\"content\">{}</div>\n", escape_html(part)));
                        }
                    } else {
                        let (lang, code) = if let Some(pos) = part.find('\n') {
                            (escape_html(&part[..pos]), escape_html(&part[pos + 1..]))
                        } else {
                            (String::new(), escape_html(part))
                        };
                        if lang.is_empty() {
                            html.push_str(&format!("<pre>{}</pre>\n", code));
                        } else {
                            html.push_str(&format!("<pre><code class=\"language-{}\">{}</code></pre>\n", lang, code));
                        }
                    }
                }
            } else {
                html.push_str(&format!("<div class=\"content\">{}</div>\n", escaped));
            }
        }
        if let Some(ref calls) = msg.tool_calls {
            for tc in calls {
                html.push_str(&format!("<div class=\"tool-call\">๐Ÿ›  {} ยท {}</div>\n",
                    escape_html(&tc.name), escape_html(&format!("{:?}", tc.arguments))));
            }
        }
        html.push_str("</div>\n");
    }
    html.push_str("</body>\n</html>\n");
    html
}

fn escape_html(s: &str) -> String {
    s.replace('&', "&amp;")
        .replace('<', "&lt;")
        .replace('>', "&gt;")
        .replace('"', "&quot;")
}

/// Ping a provider's /v1/models endpoint to verify connectivity.
/// Returns (success, model_count_or_error_message).
pub async fn health_check(base_url: &str, api_key: &str) -> (bool, String) {
    let url = format!("{}/models", base_url.trim_end_matches('/'));
    let client = reqwest::Client::builder()
        .timeout(std::time::Duration::from_secs(5))
        .build()
        .unwrap_or_default();

    if api_key.is_empty() || api_key.contains("${") {
        return (false, "API key not configured".into());
    }

    match client
        .get(&url)
        .header("Authorization", format!("Bearer {}", api_key))
        .header("Content-Type", "application/json")
        .send()
        .await
    {
        Ok(resp) if resp.status().is_success() => {
            match resp.json::<serde_json::Value>().await {
                Ok(data) => {
                    let count = data["data"].as_array().map(|a| a.len()).unwrap_or(0);
                    (true, format!("{} models available", count))
                }
                Err(_) => (false, "Failed to parse model list".into()),
            }
        }
        Ok(resp) => (false, format!("HTTP {}", resp.status())),
        Err(e) => (false, format!("Connection failed: {}", e)),
    }
}