sovereign_browser_rust/src/agent_loop.rs

//! Agent main loop — orchestrates LLM calls and tool execution.
//!
//! Port of `agent_loop.c` from the C project:
//!   1. Append the user message to the chat store.
//!   2. Build `[system] + history` and call the LLM with the tool catalog.
//!   3. Store the assistant reply (with any tool_calls).
//!   4. Execute each tool call, storing the results as `tool` messages.
//!   5. Repeat until there are no tool calls, the iteration cap is hit,
//!      or the user cancels.
//!
//! The loop runs on its own OS thread. The UI polls
//! `sovereign://agents/status` + `sovereign://agents/messages`.

use once_cell::sync::Lazy;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Mutex;

use crate::agent_chat_store;
use crate::agent_llm;
use crate::agent_tools;

/// Loop state, serialized as the strings the chat UI expects.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AgentLoopState {
    Idle,
    Thinking,
    ToolCall,
    Complete,
    Error,
    Cancelled,
}

impl AgentLoopState {
    pub fn as_str(&self) -> &'static str {
        match self {
            AgentLoopState::Idle => "idle",
            AgentLoopState::Thinking => "thinking",
            AgentLoopState::ToolCall => "tool_call",
            AgentLoopState::Complete => "complete",
            AgentLoopState::Error => "error",
            AgentLoopState::Cancelled => "cancelled",
        }
    }

    fn is_running(&self) -> bool {
        matches!(self, AgentLoopState::Thinking | AgentLoopState::ToolCall)
    }
}

/// Status snapshot returned by `sovereign://agents/status`.
#[derive(Debug, Clone)]
pub struct AgentLoopStatus {
    pub state: AgentLoopState,
    pub iteration: i32,
    pub current_tool: Option<String>,
    pub last_message: Option<String>,
    pub error: Option<String>,
}

impl AgentLoopStatus {
    pub fn to_json(&self) -> serde_json::Value {
        serde_json::json!({
            "state": self.state.as_str(),
            "iteration": self.iteration,
            "current_tool": self.current_tool,
            "last_message": self.last_message,
            "error": self.error,
        })
    }
}

static G_STATUS: Lazy<Mutex<AgentLoopStatus>> = Lazy::new(|| {
    Mutex::new(AgentLoopStatus {
        state: AgentLoopState::Idle,
        iteration: 0,
        current_tool: None,
        last_message: None,
        error: None,
    })
});
static G_CANCEL: AtomicBool = AtomicBool::new(false);

/// Max characters of a single tool result fed back to the LLM (keeps a
/// huge page dump from blowing the context window).
const MAX_TOOL_RESULT_CHARS: usize = 60_000;

fn set_status(state: AgentLoopState, iteration: i32, tool: Option<String>) {
    let mut s = G_STATUS.lock().unwrap();
    s.state = state;
    s.iteration = iteration;
    s.current_tool = if state == AgentLoopState::ToolCall { tool } else { None };
    if state != AgentLoopState::Error {
        s.error = None;
    }
}

fn set_error(msg: &str) {
    eprintln!("[agent-loop] error: {}", msg);
    let mut s = G_STATUS.lock().unwrap();
    s.state = AgentLoopState::Error;
    s.current_tool = None;
    s.error = Some(msg.to_string());
}

/// Current status snapshot.
pub fn agent_loop_get_status() -> AgentLoopStatus {
    G_STATUS.lock().unwrap().clone()
}

/// Whether the loop is currently running.
pub fn agent_loop_is_running() -> bool {
    G_STATUS.lock().unwrap().state.is_running()
}

/// Request cancellation (takes effect between LLM calls / tool calls).
pub fn agent_loop_cancel() {
    if agent_loop_is_running() {
        G_CANCEL.store(true, Ordering::SeqCst);
        // Reflect the cancel immediately so the UI's stop button resets;
        // the worker exits at its next checkpoint.
        set_status(AgentLoopState::Cancelled, agent_loop_get_status().iteration, None);
    }
}

fn truncate(s: &str, max: usize) -> String {
    if s.chars().count() <= max {
        s.to_string()
    } else {
        let head: String = s.chars().take(max).collect();
        format!("{}\n…[truncated {} chars]", head, s.chars().count() - max)
    }
}

/// OpenAI function-tool definitions for the whole catalog.
fn tool_definitions() -> Vec<serde_json::Value> {
    let mut tools = agent_tools::agent_tools_get_all();
    tools.sort_by(|a, b| a.name.cmp(&b.name));
    tools
        .iter()
        .map(|t| {
            serde_json::json!({
                "type": "function",
                "function": {
                    "name": t.name,
                    "description": t.description,
                    "parameters": t.input_schema,
                }
            })
        })
        .collect()
}

/// Start the agent loop for `user_message` on a background thread.
/// Returns an error if it is already running.
pub fn agent_loop_run(user_message: &str) -> Result<(), String> {
    {
        let mut s = G_STATUS.lock().unwrap();
        if s.state.is_running() {
            return Err("Agent is already running".to_string());
        }
        s.state = AgentLoopState::Thinking;
        s.iteration = 0;
        s.current_tool = None;
        s.error = None;
        s.last_message = None;
    }
    G_CANCEL.store(false, Ordering::SeqCst);

    // Store the user message synchronously so it shows up immediately.
    // If it can't be stored, don't start: the LLM would never see it.
    if let Err(e) = agent_chat_store::agent_chat_store_add_user_message(user_message) {
        set_error(&e);
        return Err(e);
    }

    std::thread::Builder::new()
        .name("agent-loop".into())
        .spawn(run_loop)
        .map_err(|e| {
            set_error(&format!("Failed to spawn agent thread: {}", e));
            "Failed to start agent loop".to_string()
        })?;
    Ok(())
}

fn cancelled(iter: i32) -> bool {
    if G_CANCEL.load(Ordering::SeqCst) {
        set_status(AgentLoopState::Cancelled, iter, None);
        println!("[agent-loop] cancelled at iteration {}", iter);
        true
    } else {
        false
    }
}

fn run_loop() {
    run_loop_inner();
    // Persist the conversation (derives a title from the first user
    // message if it has none, and syncs to Nostr when a signer is
    // available). The chat page also auto-saves, but only if its poll
    // happened to observe the running state — a fast reply, or a message
    // sent via the URL bar's `; ` shortcut, would otherwise never be saved.
    let sid = agent_chat_store::agent_chat_store_get_session();
    let _ = crate::agent_conversations::agent_conversations_save(Some(&sid), None);
}

fn run_loop_inner() {
    let cfg = agent_llm::agent_config_get();
    let base_url = cfg.base_url();
    let api_key = cfg.api_key();
    let model = cfg.model.clone();
    let max_iter = cfg.effective_max_iterations();
    let system_prompt = crate::agent_skills::agent_skills_build_system_prompt();
    let tools = tool_definitions();

    println!(
        "[agent-loop] base_url={} model={} api_key={} max_iter={}",
        base_url,
        model,
        if api_key.is_empty() { "(empty)" } else { "(set)" },
        max_iter
    );

    for iter in 0..max_iter {
        if cancelled(iter) {
            return;
        }
        set_status(AgentLoopState::Thinking, iter, None);

        let mut messages = vec![serde_json::json!({"role": "system", "content": system_prompt})];
        messages.extend(agent_chat_store::agent_chat_store_get_messages());

        let reply = match agent_llm::agent_llm_chat(&base_url, &api_key, &model, &messages, &tools) {
            Ok(r) => r,
            Err(e) => {
                set_error(&e);
                return;
            }
        };
        if cancelled(iter) {
            return;
        }

        let content = reply["content"].as_str().unwrap_or("").to_string();
        let tool_calls: Vec<serde_json::Value> = reply["tool_calls"].as_array().cloned().unwrap_or_default();
        if let Err(e) = agent_chat_store::agent_chat_store_add_assistant_message(&content, reply.get("tool_calls")) {
            set_error(&e);
            return;
        }
        if !content.is_empty() {
            G_STATUS.lock().unwrap().last_message = Some(content.clone());
        }

        if tool_calls.is_empty() {
            set_status(AgentLoopState::Complete, iter, None);
            println!("[agent-loop] complete after {} iteration(s)", iter + 1);
            return;
        }

        for call in &tool_calls {
            if cancelled(iter) {
                // Keep the transcript valid for the next request: every
                // tool_call must have a matching tool message.
                let id = call["id"].as_str().unwrap_or("");
                let _ = agent_chat_store::agent_chat_store_add_tool_result(id, "{\"error\":\"cancelled by user\"}");
                continue;
            }
            let id = call["id"].as_str().unwrap_or("").to_string();
            let name = call["function"]["name"].as_str().unwrap_or("").to_string();
            let args_str = call["function"]["arguments"].as_str().unwrap_or("{}");
            let args: serde_json::Value = if args_str.trim().is_empty() {
                serde_json::json!({})
            } else {
                serde_json::from_str(args_str).unwrap_or_else(|_| serde_json::json!({}))
            };

            set_status(AgentLoopState::ToolCall, iter, Some(name.clone()));
            println!("[agent-loop] iter {}: tool {} {}", iter, name, truncate(args_str, 200));

            let result = match agent_tools::agent_tools_execute(&name, &args) {
                Ok(v) => serde_json::to_string(&v).unwrap_or_default(),
                Err(e) => serde_json::json!({"error": e}).to_string(),
            };
            if let Err(e) = agent_chat_store::agent_chat_store_add_tool_result(&id, &truncate(&result, MAX_TOOL_RESULT_CHARS)) {
                set_error(&e);
                return;
            }
        }
        if G_CANCEL.load(Ordering::SeqCst) {
            set_status(AgentLoopState::Cancelled, iter, None);
            return;
        }
    }

    set_error(&format!("Reached the iteration limit ({})", max_iter));
}