cli/sunbeam/src/code/agent.rs

//! Agent service — async message bus between TUI and Sol gRPC session.
//!
//! The TUI sends `AgentRequest`s and receives `AgentEvent`s through
//! crossbeam channels. The gRPC session runs on a background tokio task,
//! so the UI thread never blocks on network I/O.
//!
//! This module is designed to be usable as a library — nothing here
//! depends on ratatui or terminal state.

use crossbeam_channel::{Receiver, Sender};

use super::client::{self, CodeSession};

/// Turn raw internal errors into something a human can read.
fn friendly_error(e: &str) -> String {
    let lower = e.to_lowercase();
    if lower.contains("broken pipe") || lower.contains("stream closed") || lower.contains("h2 protocol") {
        "sol disconnected — try again or restart with /exit".into()
    } else if lower.contains("channel closed") || lower.contains("send on closed") {
        "connection to sol lost".into()
    } else if lower.contains("timed out") || lower.contains("timeout") {
        "request timed out — sol may be overloaded".into()
    } else if lower.contains("connection refused") {
        "can't reach sol — is it running?".into()
    } else if lower.contains("not found") && lower.contains("agent") {
        "sol's agent was reset — reconnect with /exit".into()
    } else if lower.contains("invalid_request_error") {
        // Extract the actual message from Mistral API errors
        if let Some(start) = e.find("\"msg\":\"") {
            let rest = &e[start + 7..];
            if let Some(end) = rest.find('"') {
                return rest[..end].to_string();
            }
        }
        "request error from sol".into()
    } else {
        // Truncate long errors and strip Rust debug formatting
        let clean = e.replace("\\n", " ").replace("\\\"", "'");
        if clean.len() > 120 {
            format!("{}…", &clean[..117])
        } else {
            clean
        }
    }
}

// ── Requests (TUI → Agent) ──────────────────────────────────────────────

/// A request from the UI to the agent backend.
pub enum AgentRequest {
    /// Send a chat message to Sol.
    Chat { text: String },
    /// End the session gracefully.
    End,
}

// ── Events (Agent → TUI) ───────────────────────────────────────────────

/// An event from the agent backend to the UI.
#[derive(Clone, Debug)]
pub enum AgentEvent {
    /// Sol started generating a response.
    Generating,
    /// A tool started executing.
    ToolStart { name: String, detail: String },
    /// A tool finished executing.
    ToolDone { name: String, success: bool },
    /// Sol's full response text.
    Response { text: String },
    /// A non-fatal error from Sol.
    Error { message: String },
    /// Status update (shown in title bar).
    Status { message: String },
    /// Connection health: true = reachable, false = unreachable.
    Health { connected: bool },
    /// Session ended.
    SessionEnded,
}

// ── Agent handle (owned by TUI) ────────────────────────────────────────

/// Handle for the TUI to communicate with the background agent task.
pub struct AgentHandle {
    tx: Sender<AgentRequest>,
    pub rx: Receiver<AgentEvent>,
}

impl AgentHandle {
    /// Send a chat message. Non-blocking.
    pub fn chat(&self, text: &str) {
        let _ = self.tx.try_send(AgentRequest::Chat {
            text: text.to_string(),
        });
    }

    /// Request session end. Non-blocking.
    pub fn end(&self) {
        let _ = self.tx.try_send(AgentRequest::End);
    }

    /// Drain all pending events. Non-blocking.
    pub fn poll_events(&self) -> Vec<AgentEvent> {
        let mut events = Vec::new();
        while let Ok(event) = self.rx.try_recv() {
            events.push(event);
        }
        events
    }
}

// ── Spawn ──────────────────────────────────────────────────────────────

/// Spawn the agent background task. Returns a handle for the TUI.
pub fn spawn(session: CodeSession, endpoint: String) -> AgentHandle {
    let (req_tx, req_rx) = crossbeam_channel::bounded::<AgentRequest>(32);
    let (evt_tx, evt_rx) = crossbeam_channel::bounded::<AgentEvent>(256);

    tokio::spawn(agent_loop(session, req_rx, evt_tx.clone()));
    tokio::spawn(heartbeat_loop(endpoint, evt_tx));

    AgentHandle {
        tx: req_tx,
        rx: evt_rx,
    }
}

/// Ping the gRPC endpoint every second to check if Sol is reachable.
async fn heartbeat_loop(endpoint: String, evt_tx: Sender<AgentEvent>) {
    use sunbeam_proto::sunbeam_code_v1::code_agent_client::CodeAgentClient;

    let mut last_state = true; // assume connected initially (we just connected)
    let _ = evt_tx.try_send(AgentEvent::Health { connected: true });

    loop {
        tokio::time::sleep(std::time::Duration::from_secs(1)).await;

        let connected = CodeAgentClient::connect(endpoint.clone()).await.is_ok();

        if connected != last_state {
            let _ = evt_tx.try_send(AgentEvent::Health { connected });
            last_state = connected;
        }
    }
}

/// The background agent loop. Reads requests, calls gRPC, emits events.
async fn agent_loop(
    mut session: CodeSession,
    req_rx: Receiver<AgentRequest>,
    evt_tx: Sender<AgentEvent>,
) {
    loop {
        // Block on the crossbeam channel from a tokio context
        let req = match tokio::task::block_in_place(|| req_rx.recv()) {
            Ok(req) => req,
            Err(_) => break, // TUI dropped the handle
        };

        match req {
            AgentRequest::Chat { text } => {
                let _ = evt_tx.try_send(AgentEvent::Generating);

                match session.chat(&text).await {
                    Ok(resp) => {
                        // Emit tool events
                        for event in &resp.events {
                            let agent_event = match event {
                                client::ChatEvent::ToolStart { name, detail } => {
                                    AgentEvent::ToolStart {
                                        name: name.clone(),
                                        detail: detail.clone(),
                                    }
                                }
                                client::ChatEvent::ToolDone { name, success } => {
                                    AgentEvent::ToolDone {
                                        name: name.clone(),
                                        success: *success,
                                    }
                                }
                                client::ChatEvent::Status(msg) => AgentEvent::Status {
                                    message: msg.clone(),
                                },
                                client::ChatEvent::Error(msg) => AgentEvent::Error {
                                    message: friendly_error(msg),
                                },
                            };
                            let _ = evt_tx.try_send(agent_event);
                        }

                        let _ = evt_tx.try_send(AgentEvent::Response { text: resp.text });
                    }
                    Err(e) => {
                        let _ = evt_tx.try_send(AgentEvent::Error {
                            message: friendly_error(&e.to_string()),
                        });
                    }
                }
            }
            AgentRequest::End => {
                let _ = session.end().await;
                let _ = evt_tx.try_send(AgentEvent::SessionEnded);
                break;
            }
        }
    }
}