aura-gui 0.2.2

//! Bucle de un subagente: conversa con su modelo, ejecuta las herramientas que pida con `execute` y devuelve el
//! informe final. No sabe nada del editor ni del proveedor, así sirve con cualquier modelo (también uno de pruebas).
//! Es bloqueante: se llama desde un hilo de trabajo.

use crate::ai::provider::LanguageModel;
use crate::ai::thread::{Thread, ToolCall, Turn};
use crate::ai::types::{AiError, Block, Role, ToolSpec, Usage};
use crate::core::CancelToken;

/// Lo que el subagente necesita para empezar.
pub struct Brief {
    pub model: String,
    pub system: String,
    pub prompt: String,
    pub tools: Vec<ToolSpec>,
    pub max_tokens: u32,
    pub max_turns: u32,
}

#[derive(Debug, PartialEq)]
pub struct Report {
    pub text: String,
    pub usage: Usage,
    pub turns: u32,
    /// Agotó `max_turns` sin terminar: el texto es lo último que dijo.
    pub exhausted: bool,
}

fn last_text(thread: &Thread) -> String {
    let last = thread.messages.iter().rev().find(|message| message.role == Role::Assistant);
    last.map(|message| {
        let texts: Vec<&str> = message
            .blocks
            .iter()
            .filter_map(|block| if let Block::Text(text) = block { Some(text.as_str()) } else { None })
            .collect();
        texts.join("\n")
    })
    .unwrap_or_default()
}

pub fn run(
    model: &dyn LanguageModel,
    brief: Brief,
    cancel: &CancelToken,
    execute: &mut dyn FnMut(&ToolCall) -> (String, bool),
) -> Result<Report, AiError> {
    let mut thread = Thread::default();
    thread.push_user(brief.prompt);
    for turn in 1..=brief.max_turns.max(1) {
        // En la última vuelta no se ofrecen herramientas: el modelo tiene que contestar con lo que ya sabe.
        let tools = if turn == brief.max_turns { Vec::new() } else { brief.tools.clone() };
        let request = thread.request(&brief.model, Some(brief.system.clone()), tools, brief.max_tokens);
        let mut outcome = None;
        model.stream(&request, cancel, &mut |event| {
            if let Some(next) = thread.on_event(event) {
                outcome = Some(next);
            }
        })?;
        let Some(Turn::NeedsTools(calls)) = outcome else {
            return Ok(Report { text: last_text(&thread), usage: thread.usage, turns: turn, exhausted: false });
        };
        for call in calls {
            if cancel.is_cancelled() {
                return Err(AiError::Cancelled);
            }
            let (content, is_error) = execute(&call);
            thread.submit_tool_result(&call.id, content, is_error);
        }
    }
    let turns = brief.max_turns.max(1);
    Ok(Report { text: last_text(&thread), usage: thread.usage, turns, exhausted: true })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ai::types::{Event, Request, StopReason};
    use serde_json::json;
    use std::sync::Mutex;

    /// Modelo guionado: cada llamada entrega la siguiente lista de eventos y guarda la petición.
    struct Scripted(Mutex<Vec<Vec<Event>>>, Mutex<Vec<Request>>);

    impl LanguageModel for Scripted {
        fn stream(&self, request: &Request, _: &CancelToken, emit: &mut dyn FnMut(Event)) -> Result<(), AiError> {
            self.1.lock().unwrap().push(request.clone());
            let events = self.0.lock().unwrap().remove(0);
            events.into_iter().for_each(emit);
            Ok(())
        }
    }

    fn brief(max_turns: u32) -> Brief {
        let tools = vec![ToolSpec { name: "read_file".into(), description: String::new(), schema: json!({}) }];
        Brief { model: "m".into(), system: "s".into(), prompt: "mira".into(), tools, max_tokens: 10, max_turns }
    }

    #[test]
    fn ejecuta_herramientas_y_devuelve_el_informe() {
        let model = Scripted(
            Mutex::new(vec![
                vec![
                    Event::ToolUse { id: "1".into(), name: "read_file".into(), input: json!({"path": "a"}) },
                    Event::Done(StopReason::ToolUse),
                ],
                vec![Event::Text("todo bien".into()), Event::Done(StopReason::EndTurn)],
            ]),
            Mutex::new(Vec::new()),
        );
        let mut seen = Vec::new();
        let report = run(&model, brief(5), &CancelToken::default(), &mut |call| {
            seen.push(call.name.clone());
            ("contenido".into(), false)
        })
        .unwrap();
        assert_eq!((report.text.as_str(), report.turns, report.exhausted), ("todo bien", 2, false));
        assert_eq!(seen, ["read_file"]);
        let requests = model.1.lock().unwrap();
        assert!(
            matches!(&requests[1].messages[2].blocks[0], Block::ToolResult { content, .. } if content == "contenido")
        );
    }

    #[test]
    fn la_ultima_vuelta_va_sin_herramientas() {
        let call = || {
            vec![
                Event::ToolUse { id: "x".into(), name: "read_file".into(), input: json!({}) },
                Event::Done(StopReason::ToolUse),
            ]
        };
        let model = Scripted(
            Mutex::new(vec![call(), vec![Event::Text("fin".into()), Event::Done(StopReason::EndTurn)]]),
            Mutex::new(Vec::new()),
        );
        let report = run(&model, brief(2), &CancelToken::default(), &mut |_| (String::new(), false)).unwrap();
        assert_eq!(report.text, "fin");
        assert!(model.1.lock().unwrap()[1].tools.is_empty());
    }
}