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::{ExploreService, 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.
///
/// # Explore phase
///
/// Before the main LLM loop begins, [`AgentTurnServiceImpl`] runs a mandatory
/// explore phase that spawns ≥3 parallel subagents (code structure, symbol
/// index, semantic context) and injects their consolidated findings as a
/// system message. See [`ExploreService`] for the trait contract.
pub struct AgentTurnServiceImpl<P: ProviderService, T: ToolExecutor> {
    provider: Arc<P>,
    tool_executor: Arc<T>,
    tool_defs: Vec<ToolDef>,
    /// Optional explore-phase service. When `Some`, the explore phase runs
    /// before every turn; when `None` it is skipped (tests, daemon mode).
    explore_service: Option<Arc<dyn ExploreService>>,
}

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,
            explore_service: None,
        }
    }

    /// Attach an optional explore-phase service.
    ///
    /// When set, every call to `run_turn` will first run the explore phase
    /// and inject the consolidated context as a system message.
    pub fn with_explore(mut self, service: Arc<dyn ExploreService>) -> Self {
        self.explore_service = Some(service);
        self
    }

    /// 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
        );

        // ── Phase 0: Mandatory explore ──────────────────────────────────
        // Spawn ≥3 parallel subagents to discover code structure, symbols,
        // and semantic context. The consolidated summary is injected as a
        // system message before the main agent prompt.
        if let Some(ref explorer) = self.explore_service {
            // Determine workspace root from the first message's context or
            // the first workspace root in params.
            let user_query = params
                .messages
                .last()
                .map(|m| m.content.clone().unwrap_or_default())
                .unwrap_or_default();
            let workspace_root = params
                .workspace_roots
                .first()
                .map(|p| p.to_string_lossy().to_string())
                .unwrap_or_else(|| ".".to_string());

            push_event(
                &params.turn_events,
                TurnEvent::SystemNote {
                    kind: "info".into(),
                    message: "🔍 Exploring codebase structure...".into(),
                },
            );

            match explorer.explore(&user_query, &workspace_root, &params.turn_events).await {
                Ok(output) => {
                    // Insert each context message as a system message.
                    // They go at index 0 and are removed after the turn
                    // like the main agent prompt.
                    for ctx_msg in &output.context_messages {
                        params
                            .messages
                            .insert(0, ChatMessage::system(ctx_msg.clone()));
                    }
                    info!(
                        "Explore phase complete: {} context messages, {}",
                        output.context_messages.len(),
                        output.summary
                    );
                }
                Err(e) => {
                    warn!("Explore phase failed (non-fatal): {e}");
                    push_event(
                        &params.turn_events,
                        TurnEvent::SystemNote {
                            kind: "warn".into(),
                            message: format!("Explore phase failed: {e}"),
                        },
                    );
                }
            }
        }

        // 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(())
        }
    }
}