zesdex/apps/application/src/agent/turn_service.rs

use std::collections::VecDeque;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use tracing::{debug, info, warn};

use zesdex_domain::agent::{AgentTurnParams, TurnEvent};
use zesdex_domain::core::{ChatMessage, StreamEvent, ToolDef};
use zesdex_domain::main_agent_prompt;

use crate::ports::ProviderService;
use super::ToolExecutor;

/// Maximum tool-call iterations per agent turn before forcing termination.
const MAX_TURN_ITERATIONS: u32 = 50;

// ---------------------------------------------------------------------------
// Helper: push a TurnEvent onto the shared queue.
// ---------------------------------------------------------------------------

fn push_event(queue: &Arc<Mutex<VecDeque<TurnEvent>>>, event: TurnEvent) {
    if let Ok(mut q) = queue.lock() {
        q.push_back(event);
    }
}

// ---------------------------------------------------------------------------
// Helper: stream-event callback that forwards tokens to the turn-event queue
// and checks the abort flag on each emission.
// ---------------------------------------------------------------------------

fn make_stream_callback(
    abort: &Arc<AtomicBool>,
    turn_events: &Arc<Mutex<VecDeque<TurnEvent>>>,
) -> Box<dyn FnMut(&StreamEvent) -> bool + Send> {
    let abort_clone = Arc::clone(abort);
    let events_clone = Arc::clone(turn_events);
    Box::new(move |event: &StreamEvent| -> bool {
        if abort_clone.load(Ordering::SeqCst) {
            return false;
        }
        match event {
            StreamEvent::Token(s) => {
                push_event(&events_clone, TurnEvent::StreamToken(s.clone()));
            }
            StreamEvent::Reasoning(s) => {
                push_event(&events_clone, TurnEvent::StreamReasoning(s.clone()));
            }
            _ => {}
        }
        true
    })
}

// ---------------------------------------------------------------------------
// Helper: execute a single tool call, push events, return the result string.
// ---------------------------------------------------------------------------

async fn execute_tool_call<T: ToolExecutor>(
    tool_executor: &T,
    turn_events: &Arc<Mutex<VecDeque<TurnEvent>>>,
    tc: &zesdex_domain::core::ToolCall,
) -> String {
    let name = &tc.function.name;
    let args = zesdex_domain::core::tool_call::sanitize_tool_arguments(&tc.function.arguments);

    debug!("executing tool: {name}");

    let output = match tool_executor.execute(name, &args).await {
        Ok(o) => o,
        Err(e) => format!("Error: {e}"),
    };

    let is_error = output.starts_with("Error:");

    push_event(
        turn_events,
        TurnEvent::ToolResult {
            tool_call_id: tc.id.clone(),
            tool_name: name.clone(),
            output: output.clone(),
            is_error,
            path: None,
        },
    );

    output
}

// ---------------------------------------------------------------------------
// Helper: emit usage event from optional LLM response metadata.
// ---------------------------------------------------------------------------

fn emit_usage(turn_events: &Arc<Mutex<VecDeque<TurnEvent>>>, usage: Option<(u64, u64)>) {
    if let Some((tokens_in, tokens_out)) = usage {
        push_event(
            turn_events,
            TurnEvent::Usage {
                tokens_in,
                tokens_out,
            },
        );
    }
}

// ---------------------------------------------------------------------------
// Service implementation
// ---------------------------------------------------------------------------

/// Service implementation for executing an agent turn asynchronously.
pub struct AgentTurnServiceImpl<P: ProviderService, T: ToolExecutor> {
    provider: Arc<P>,
    tool_executor: Arc<T>,
    tool_defs: Vec<ToolDef>,
}

impl<P: ProviderService, T: ToolExecutor> AgentTurnServiceImpl<P, T> {
    pub fn new(provider: Arc<P>, tool_executor: Arc<T>, tool_defs: Vec<ToolDef>) -> Self {
        Self {
            provider,
            tool_executor,
            tool_defs,
        }
    }

    /// Execute a single LLM call with the current message list, handling
    /// streaming events and error reporting.
    async fn call_llm(
        &self,
        messages: &[ChatMessage],
        abort: &Arc<AtomicBool>,
        turn_events: &Arc<Mutex<VecDeque<TurnEvent>>>,
    ) -> Result<(ChatMessage, Option<(u64, u64)>), String> {
        let on_event = make_stream_callback(abort, turn_events);

        self.provider
            .chat_stream(
                messages,
                Some(self.tool_defs.clone()),
                Some(4096),
                Some(0.7),
                on_event,
            )
            .await
            .map_err(|e| format!("LLM error: {e}"))
    }
}

impl<P: ProviderService, T: ToolExecutor> super::AgentTurnService for AgentTurnServiceImpl<P, T> {
    async fn run_turn(&self, mut params: AgentTurnParams) -> anyhow::Result<()> {
        info!(
            "Starting async agent turn with {} messages (model: {})",
            params.messages.len(),
            params.model
        );

        // Insert system prompt at position 0 once and keep it there for the
        // entire turn, avoiding per-iteration clones of the full message list.
        // It is removed before emitting the Compacted event so persistence
        // does not store the prompt redundantly.
        params.messages.insert(0, ChatMessage::system(main_agent_prompt()));
        let original_count = params.messages.len();

        for iteration in 0..MAX_TURN_ITERATIONS {
            // ── Check abort flag ────────────────────────────────────────
            if params.abort.load(Ordering::SeqCst) {
                params.abort.store(false, Ordering::SeqCst);
                push_event(
                    &params.turn_events,
                    TurnEvent::SystemNote {
                        kind: "info".into(),
                        message: "Turn aborted by user".into(),
                    },
                );
                break;
            }

            debug!("agent turn iteration {iteration}");

            // ── Stream start + call LLM ─────────────────────────────────
            push_event(&params.turn_events, TurnEvent::StreamStart);

            // Uses params.messages directly (sys_msg[0] already in place
            // from the insert above) — no per-iteration clone needed.
            let result = self
                .call_llm(&params.messages, &params.abort, &params.turn_events)
                .await;

            match result {
                Ok((assistant_msg, usage)) => {
                    let content = assistant_msg.content.clone().unwrap_or_default();
                    let tool_calls = assistant_msg.tool_calls.clone().unwrap_or_default();

                    push_event(
                        &params.turn_events,
                        TurnEvent::StreamDone(assistant_msg.clone()),
                    );

                    emit_usage(&params.turn_events, usage);

                    // ── No tool calls → assistant is done ──────────────
                    if tool_calls.is_empty() {
                        params.messages.push(ChatMessage::assistant(Some(content)));
                        break;
                    }

                    params.messages.push(assistant_msg);

                    // ── Execute each tool call ──────────────────────────
                    for tc in &tool_calls {
                        let output =
                            execute_tool_call(self.tool_executor.as_ref(), &params.turn_events, tc).await;
                        params
                            .messages
                            .push(ChatMessage::tool(tc.id.clone(), output));
                    }
                }
                Err(e) => {
                    warn!("{e}");
                    push_event(
                        &params.turn_events,
                        TurnEvent::Error(e),
                    );
                    break;
                }
            }
        }

        // Remove the synthetic sys_msg before shipping events to the TUI
        // so the transcript shows only the actual user/assistant/tool exchange.
        let compacted: Vec<ChatMessage> = params.messages.drain(original_count - 1..).collect();
        push_event(
            &params.turn_events,
            TurnEvent::Compacted(compacted),
        );
        push_event(&params.turn_events, TurnEvent::Done);
        params.in_flight.store(false, Ordering::SeqCst);

        Ok(())
    }
}

// ---------------------------------------------------------------------------
// Conversation compaction
// ---------------------------------------------------------------------------

/// Maximum number of recent messages to preserve during compaction.
const COMPACT_KEEP_TAIL: usize = 6;

/// Compacts conversation history using AI summarisation.
///
/// Flow: if the message count exceeds `KEEP_TAIL + 2`, the oldest messages
/// are drained and summarised by the LLM. The summary is inserted as a
/// system message at the head of the remaining history.
pub async fn compact_messages_with_ai<P: ProviderService>(
    messages: &mut Vec<ChatMessage>,
    provider: &P,
) -> anyhow::Result<()> {
    if messages.len() <= COMPACT_KEEP_TAIL + 2 {
        return Ok(()); // Not enough messages to compact
    }

    let split_idx = messages.len() - COMPACT_KEEP_TAIL;
    let evicted: Vec<_> = messages.drain(..split_idx).collect();

    let mut summary_prompt = vec![
        ChatMessage::system(zesdex_domain::compaction_prompt()),
    ];
    summary_prompt.extend(evicted);
    summary_prompt.push(ChatMessage::user(
        "Please summarise our previous conversation above for context continuity.".to_string(),
    ));

    match provider.chat(&summary_prompt, None, Some(1024), Some(0.3)).await {
        Ok((summary_msg, _)) => {
            let summary_text = summary_msg
                .content
                .unwrap_or_else(|| "Previous context summarised.".to_string());
            let summary_node = ChatMessage::system(format!(
                "[AI Summary of Previous Conversation]\n{}",
                summary_text.trim()
            ));
            messages.insert(0, summary_node);
            Ok(())
        }
        Err(e) => {
            warn!("AI summarisation failed during compact, falling back to simple notice: {e}");
            messages.insert(
                0,
                ChatMessage::system(
                    "[Earlier conversation messages compacted to save context window]".to_string(),
                ),
            );
            Ok(())
        }
    }
}