aura-gui 0.2.2

//! Historial de una conversación y su máquina de estados: acumula lo que llega del modelo, pide las herramientas
//! y arma la siguiente petición. No hace red ni toca el editor, así que se prueba entero con eventos fabricados.

use crate::ai::types::{Block, Event, Message, Request, Role, StopReason, ToolSpec, Usage};
use serde_json::Value;

/// Llamada a herramienta pendiente de ejecutar.
#[derive(Clone, Debug, PartialEq)]
pub struct ToolCall {
    pub id: String,
    pub name: String,
    pub input: Value,
}

/// Qué sigue cuando el modelo termina de hablar.
#[derive(Clone, Debug, PartialEq)]
pub enum Turn {
    /// Respondió sin pedir herramientas.
    Finished,
    /// Pidió herramientas: cuando estén todas resueltas hay que volver a llamar al modelo.
    NeedsTools(Vec<ToolCall>),
}

#[derive(Default)]
pub struct Thread {
    pub messages: Vec<Message>,
    pub usage: Usage,
    /// Uso de la última respuesta: es lo que ocupa hoy la conversación en la ventana de contexto
    /// (`usage` suma todos los turnos y sirve para el gasto, no para el contexto).
    pub context: Usage,
    /// Bloques del mensaje del asistente que se está recibiendo.
    current: Vec<Block>,
    /// Herramientas pedidas aún sin resultado, en orden de llamada.
    pending: Vec<ToolCall>,
    /// Resultados ya recibidos, en el mismo orden que `pending` cuando estén completos.
    results: Vec<Block>,
    /// Imágenes que acompañan a los resultados (capturas de pantalla): van después de ellos en el mismo mensaje.
    extras: Vec<Block>,
}

impl Thread {
    pub fn push_user(&mut self, text: impl Into<String>) {
        self.messages.push(Message::user(text));
    }

    /// Mensaje del usuario con bloques ya armados (texto, imágenes adjuntas).
    pub fn push_user_blocks(&mut self, blocks: Vec<Block>) {
        self.messages.push(Message { role: Role::User, blocks });
    }

    /// Texto del mensaje del asistente que va llegando (para pintarlo mientras se recibe).
    pub fn streaming_text(&self) -> String {
        self.current
            .iter()
            .filter_map(|block| if let Block::Text(text) = block { Some(text.as_str()) } else { None })
            .collect()
    }

    pub fn pending_tools(&self) -> &[ToolCall] {
        &self.pending
    }

    pub fn on_event(&mut self, event: Event) -> Option<Turn> {
        match event {
            Event::Text(text) => match self.current.last_mut() {
                Some(Block::Text(last)) => last.push_str(&text),
                _ => self.current.push(Block::Text(text)),
            },
            Event::ToolUse { id, name, input } => {
                self.current.push(Block::ToolUse { id: id.clone(), name: name.clone(), input: input.clone() });
                self.pending.push(ToolCall { id, name, input });
            }
            Event::Usage(usage) => {
                self.context = usage;
                self.usage.input_tokens += usage.input_tokens;
                self.usage.output_tokens += usage.output_tokens;
            }
            Event::Done(reason) => return Some(self.finish_turn(reason)),
        }
        None
    }

    fn finish_turn(&mut self, reason: StopReason) -> Turn {
        if !self.current.is_empty() {
            self.messages.push(Message::assistant(std::mem::take(&mut self.current)));
        }
        if self.pending.is_empty() || reason != StopReason::ToolUse && reason != StopReason::EndTurn {
            // Sin herramientas, o cortado (límite de tokens): no hay nada que ejecutar.
            self.drop_unanswered_tool_uses();
            return Turn::Finished;
        }
        Turn::NeedsTools(self.pending.clone())
    }

    /// Entrega el resultado de una herramienta. `true` cuando ya están todas: el resultado queda en el historial
    /// y hay que pedir la siguiente respuesta al modelo.
    pub fn submit_tool_result(&mut self, id: &str, content: String, is_error: bool) -> bool {
        self.submit_tool_result_with(id, content, is_error, Vec::new())
    }

    /// Como `submit_tool_result`, con bloques extra (una captura de pantalla) que el modelo ve junto al resultado.
    pub fn submit_tool_result_with(&mut self, id: &str, content: String, is_error: bool, extra: Vec<Block>) -> bool {
        if !self.pending.iter().any(|call| call.id == id)
            || self.results.iter().any(|result| matches!(result, Block::ToolResult { id: done, .. } if done == id))
        {
            return false;
        }
        self.results.push(Block::ToolResult { id: id.to_owned(), content, is_error });
        self.extras.extend(extra);
        if self.results.len() < self.pending.len() {
            return false;
        }
        // La API exige los resultados en el orden en que se pidieron las herramientas.
        let mut ordered = Vec::with_capacity(self.pending.len());
        for call in self.pending.drain(..) {
            if let Some(position) =
                self.results.iter().position(|result| matches!(result, Block::ToolResult { id, .. } if *id == call.id))
            {
                ordered.push(self.results.swap_remove(position));
            }
        }
        self.results.clear();
        ordered.append(&mut self.extras);
        self.messages.push(Message { role: Role::User, blocks: ordered });
        true
    }

    /// El usuario detuvo la respuesta: se conserva el texto recibido y se descartan las herramientas sin resolver,
    /// para que el historial siga siendo válido para la API.
    pub fn abort_turn(&mut self) {
        if !self.current.is_empty() {
            self.messages.push(Message::assistant(std::mem::take(&mut self.current)));
        }
        self.drop_unanswered_tool_uses();
    }

    fn drop_unanswered_tool_uses(&mut self) {
        self.pending.clear();
        self.results.clear();
        self.extras.clear();
        if let Some(last) = self.messages.last_mut().filter(|message| message.role == Role::Assistant) {
            last.blocks.retain(|block| !matches!(block, Block::ToolUse { .. }));
        }
        if self.messages.last().is_some_and(|message| message.role == Role::Assistant && message.blocks.is_empty()) {
            self.messages.pop();
        }
    }

    pub fn request(&self, model: &str, system: Option<String>, tools: Vec<ToolSpec>, max_tokens: u32) -> Request {
        Request { model: model.to_owned(), system, messages: self.messages.clone(), tools, max_tokens }
    }
}

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

    fn tool(id: &str) -> Event {
        Event::ToolUse { id: id.into(), name: "read_file".into(), input: json!({"path": id}) }
    }

    #[test]
    fn el_texto_en_trozos_forma_un_solo_bloque() {
        let mut thread = Thread::default();
        thread.push_user("hola");
        thread.on_event(Event::Text("Ho".into()));
        thread.on_event(Event::Text("la".into()));
        assert_eq!(thread.streaming_text(), "Hola");
        assert_eq!(thread.on_event(Event::Done(StopReason::EndTurn)), Some(Turn::Finished));
        assert_eq!(thread.messages[1], Message::assistant(vec![Block::Text("Hola".into())]));
    }

    #[test]
    fn ciclo_de_herramientas_conserva_el_orden_de_llamada() {
        let mut thread = Thread::default();
        thread.push_user("revisa");
        thread.on_event(Event::Text("miro".into()));
        thread.on_event(tool("a"));
        thread.on_event(tool("b"));
        let turn = thread.on_event(Event::Done(StopReason::ToolUse));
        assert!(matches!(turn, Some(Turn::NeedsTools(calls)) if calls.len() == 2));
        assert!(!thread.submit_tool_result("b", "B".into(), false), "falta una");
        assert!(!thread.submit_tool_result("b", "otra vez".into(), false), "duplicados ignorados");
        assert!(thread.submit_tool_result("a", "A".into(), false));
        let Some(last) = thread.messages.last() else { panic!("sin mensajes") };
        assert_eq!(last.role, Role::User);
        let ids: Vec<&str> = last
            .blocks
            .iter()
            .filter_map(|b| if let Block::ToolResult { id, .. } = b { Some(id.as_str()) } else { None })
            .collect();
        assert_eq!(ids, ["a", "b"], "primero el resultado de la primera llamada");
        assert!(thread.pending_tools().is_empty());
    }

    #[test]
    fn detener_descarta_herramientas_sin_resultado_y_mantiene_el_texto() {
        let mut thread = Thread::default();
        thread.push_user("hazlo");
        thread.on_event(Event::Text("voy".into()));
        thread.on_event(tool("a"));
        thread.on_event(Event::Done(StopReason::ToolUse));
        thread.abort_turn();
        assert_eq!(thread.messages.last(), Some(&Message::assistant(vec![Block::Text("voy".into())])));
        assert!(thread.pending_tools().is_empty());
    }

    #[test]
    fn mensaje_solo_de_herramientas_cancelado_no_deja_hueco() {
        let mut thread = Thread::default();
        thread.push_user("x");
        thread.on_event(tool("a"));
        thread.abort_turn();
        assert_eq!(thread.messages.len(), 1, "el mensaje vacío del asistente se descarta");
    }

    #[test]
    fn el_uso_se_acumula_entre_turnos() {
        let mut thread = Thread::default();
        thread.on_event(Event::Usage(Usage { input_tokens: 10, output_tokens: 2 }));
        thread.on_event(Event::Usage(Usage { input_tokens: 5, output_tokens: 3 }));
        assert_eq!(thread.usage, Usage { input_tokens: 15, output_tokens: 5 });
        assert_eq!(thread.context, Usage { input_tokens: 5, output_tokens: 3 }, "el contexto es el último turno");
    }

    #[test]
    fn la_captura_va_despues_de_los_resultados_en_el_mismo_mensaje() {
        let mut thread = Thread::default();
        thread.on_event(Event::ToolUse { id: "a".into(), name: "computer".into(), input: Value::Null });
        thread.on_event(Event::Done(StopReason::ToolUse));
        let image = Block::Image { media_type: "image/jpeg".into(), data_base64: "eA==".into(), name: None };
        assert!(thread.submit_tool_result_with("a", "captura".into(), false, vec![image.clone()]));
        let last = thread.messages.last().unwrap();
        assert!(matches!(&last.blocks[0], Block::ToolResult { id, .. } if id == "a"));
        assert_eq!(last.blocks[1], image);
    }
}