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 crate::ports::ProviderService;
use super::ToolExecutor;

/// 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,
        }
    }
    
    fn push_event(queue: &Arc<Mutex<VecDeque<TurnEvent>>>, event: TurnEvent) {
        if let Ok(mut q) = queue.lock() {
            q.push_back(event);
        }
    }

    fn mark_done(flag: &Arc<AtomicBool>) {
        flag.store(false, Ordering::SeqCst);
    }
}

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

        let mut sys_prompt = "You are Zesdex, an AI coding assistant. You have access to various tools via native function calling to help the user.\n\n\
             CRITICAL DIRECTIVES & PRIORITY HIERARCHY:\n\
             1. WORKFLOW FIRST: For any multi-step, complex, or non-trivial task, you MUST prioritize using `workflow_run` (to construct and execute a multi-phase YAML workflow) or `hive_mind` (to orchestrate parallel autonomous agents). Workflows are your primary strategy.\n\
             2. PLANNING & TODOS: Use `plan_enter` to establish high-level architectural plans and `todowrite` to maintain granular task checklists.\n\
             3. REASONING: Use `seq_think` for deep step-by-step analysis.\n\
             4. TOOL EXECUTION: Execute individual tools (file edits, terminal commands) within or guided by your workflows. If an error occurs, analyze and fix it.\n\n\
             Respond conversationally, concisely, and helpfully.".to_string();

        let sys_msg = ChatMessage::system(sys_prompt);

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

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

            Self::push_event(&params.turn_events, TurnEvent::StreamStart);

            let mut req_messages = params.messages.clone();
            req_messages.insert(0, sys_msg.clone());

            let abort_clone = Arc::clone(&params.abort);
            let turn_events_clone = Arc::clone(&params.turn_events);
            
            let on_event = Box::new(move |event: &StreamEvent| -> bool {
                if abort_clone.load(Ordering::SeqCst) {
                    return false;
                }
                match event {
                    StreamEvent::Token(s) => {
                        Self::push_event(&turn_events_clone, TurnEvent::StreamToken(s.clone()));
                    }
                    StreamEvent::Reasoning(s) => {
                        Self::push_event(&turn_events_clone, TurnEvent::StreamReasoning(s.clone()));
                    }
                    _ => {}
                }
                true
            });

            let result = self.provider.chat_stream(
                &req_messages,
                Some(self.tool_defs.clone()),
                Some(4096),
                Some(0.7),
                on_event,
            ).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();

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

                    if let Some((tokens_in, tokens_out)) = usage {
                        Self::push_event(
                            &params.turn_events,
                            TurnEvent::Usage {
                                tokens_in,
                                tokens_out,
                            },
                        );
                    }

                    if tool_calls.is_empty() {
                        params.messages.push(ChatMessage::assistant(Some(content)));
                        break;
                    }

                    params.messages.push(assistant_msg);

                    for tc in &tool_calls {
                        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 self.tool_executor.execute(name, &args).await {
                            Ok(o) => o,
                            Err(e) => format!("Error: {e}"),
                        };

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

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

                        params
                            .messages
                            .push(ChatMessage::tool(tc.id.clone(), output.clone()));
                    }
                }
                Err(e) => {
                    warn!("LLM call failed: {e}");
                    Self::push_event(
                        &params.turn_events,
                        TurnEvent::Error(format!("LLM error: {e}")),
                    );
                    break;
                }
            }
        }

        Self::push_event(
            &params.turn_events,
            TurnEvent::Compacted(params.messages.clone()),
        );
        Self::push_event(&params.turn_events, TurnEvent::Done);
        Self::mark_done(&params.in_flight);
        
        Ok(())
    }
}

/// Compacts conversation history using AI summarization.
pub async fn compact_messages_with_ai<P: ProviderService>(
    messages: &mut Vec<ChatMessage>,
    provider: &P,
) -> anyhow::Result<()> {
    const KEEP_TAIL: usize = 6;
    if messages.len() <= KEEP_TAIL + 2 {
        return Ok(()); // Not enough messages to compact
    }

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

    let mut summary_prompt = vec![
        ChatMessage::system(
            "You are a helpful assistant summarizing conversation history. \
             Provide a concise summary of the key user requests, decisions, tools executed, and modified files. \
             Format as a clear bulleted list."
                .to_string(),
        ),
    ];
    summary_prompt.extend(evicted);
    summary_prompt.push(ChatMessage::user(
        "Please summarize 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 summarized.".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 summarization 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(())
        }
    }
}