matrixcode_core/agent/

run.rs

//! Agent run loop and public methods.

use anyhow::Result;
use std::sync::Arc;
use std::sync::atomic::{AtomicU8, Ordering};
use tokio::sync::mpsc;

use crate::approval::ApproveMode;
use crate::cancel::CancellationToken;
use crate::compress::{
    CompressionStrategy, compress_messages, estimate_total_tokens, should_compress,
};
use crate::event::{AgentEvent, EventData, EventType};
use crate::prompt;
use crate::providers::{ChatRequest, Message, MessageContent, Role};
use crate::tools::Tool;
use crate::tools::ToolDefinition;
use crate::tools::toolproxy::{ProxyToolDef, ProxyToolExecutor};

use super::types::{Agent, AgentBuilder, MAX_ITERATIONS};

impl Agent {
    pub(crate) fn new(builder: AgentBuilder) -> Self {
        let event_tx = builder.event_tx.unwrap_or_else(|| {
            let (tx, _) = mpsc::channel(100);
            tx
        });

        Self {
            provider: builder.provider,
            model_name: builder.model_name,
            tools: builder.tools,
            messages: Vec::new(),
            system_prompt: builder.system_prompt,
            max_tokens: builder.max_tokens,
            context_size_override: builder.context_size_override,
            think: builder.think,
            approve_mode: Arc::new(AtomicU8::new(builder.approve_mode.to_u8())),
            event_tx,
            skills: builder.skills,
            profile: builder.profile,
            project_overview: builder.project_overview,
            memory_summary: builder.memory_summary,
            project_path: builder.project_path,
            total_input_tokens: std::sync::atomic::AtomicU64::new(0),
            total_output_tokens: std::sync::atomic::AtomicU64::new(0),
            last_input_tokens: std::sync::atomic::AtomicU64::new(0),
            cancel_token: None,
            compression_config: crate::compress::CompressionConfig::default(),
            ask_rx: None,
            proxy_tool_defs: builder.proxy_tool_defs,
            proxy_executor: builder.proxy_executor,
            mcp_registry: builder.mcp_registry,
            pending_input_rx: builder.pending_input_rx,
            pending_inputs: Vec::new(),
            previewed_tool_inputs: std::collections::HashSet::new(),
            todo_reminder_count: std::collections::HashMap::new(),
        }
    }

    /// Effective context window size, preferring explicit configuration over model inference.
    pub(crate) fn effective_context_size(&self) -> Option<u32> {
        self.context_size_override
            .or_else(|| self.provider.context_size())
    }

    /// Get event sender for streaming
    pub fn event_sender(&self) -> mpsc::Sender<AgentEvent> {
        self.event_tx.clone()
    }

    /// Set ask response channel (for TUI mode)
    pub fn set_ask_channel(&mut self, rx: mpsc::Receiver<String>) {
        self.ask_rx = Some(rx);
    }

    /// 设置代理工具执行器
    pub fn set_proxy_executor(
        &mut self,
        executor: Arc<dyn ProxyToolExecutor>,
        tool_defs: Vec<ProxyToolDef>,
    ) {
        self.proxy_executor = Some(executor);
        self.proxy_tool_defs = tool_defs;
    }

    /// Set cancellation token
    pub fn set_cancel_token(&mut self, token: CancellationToken) {
        self.cancel_token = Some(token);
    }

    /// Set approve mode at runtime
    pub fn set_approve_mode(&mut self, mode: ApproveMode) {
        let old = ApproveMode::from_u8(self.approve_mode.load(Ordering::Relaxed));
        log::info!("Agent approve mode changed: {} -> {}", old, mode);
        self.approve_mode.store(mode.to_u8(), Ordering::Relaxed);
    }

    /// Get a shared reference to the approve mode atomic.
    pub fn approve_mode_shared(&self) -> Arc<AtomicU8> {
        self.approve_mode.clone()
    }

    /// Replace the internal approve mode with an externally-created shared atomic.
    pub fn set_approve_mode_shared(&mut self, shared: Arc<AtomicU8>) {
        self.approve_mode = shared;
    }

    /// Update memory summary and rebuild system prompt.
    /// Note: Uses build_system_prompt (without project_path) to preserve cache.
    pub fn update_memory_summary(&mut self, summary: Option<String>) {
        self.memory_summary = summary;
        // Preserve cache by using build_system_prompt (no dynamic CodeGraph injection)
        self.system_prompt = prompt::build_system_prompt(
            &self.profile,
            &self.skills,
            self.project_overview.as_deref(),
            self.memory_summary.as_deref(),
        );
    }

    /// Refresh CodeGraph tools after /init or codegraph init.
    /// This rebuilds both tools and system prompt with project_path.
    /// Call this only when CodeGraph state changes (not every request) to preserve cache.
    pub fn refresh_codegraph_tools(&mut self) {
        if let Some(path) = &self.project_path {
            // Check if CodeGraph should be injected now
            let should_have_codegraph =
                crate::tools::codegraph::should_inject_codegraph_tools(path);

            // Check if we currently have CodeGraph tools
            let has_codegraph = self.tools.iter().any(|t| {
                let name = t.definition().name;
                name.starts_with("code_") && name != "code_review"
            });

            // Only update if state changed
            if should_have_codegraph != has_codegraph {
                // Add or remove CodeGraph tools
                if should_have_codegraph {
                    let codegraph_tools = crate::tools::codegraph::codegraph_tools(path);
                    for tool in codegraph_tools {
                        self.tools.push(Arc::from(tool));
                    }
                    // Update system prompt to include CodeGraph rules
                    self.system_prompt = prompt::build_system_prompt_with_workflows(
                        &self.profile,
                        &self.skills,
                        self.project_overview.as_deref(),
                        self.memory_summary.as_deref(),
                        Some(path),
                        None, // LSP servers not available in agent context
                    );
                } else {
                    // Remove CodeGraph tools
                    self.tools.retain(|t| {
                        let name = t.definition().name;
                        !name.starts_with("code_") || name == "code_review"
                    });
                    // Update system prompt to remove CodeGraph rules
                    self.system_prompt = prompt::build_system_prompt_with_workflows(
                        &self.profile,
                        &self.skills,
                        self.project_overview.as_deref(),
                        self.memory_summary.as_deref(),
                        Some(path),
                        None, // LSP servers not available in agent context
                    );
                }
            }
        }
    }

    /// Run chat loop with tool execution (streaming version).
    pub async fn run(&mut self, user_input: String) -> Result<Vec<AgentEvent>> {
        self.emit(AgentEvent::session_started())?;

        self.messages.push(Message {
            role: Role::User,
            content: MessageContent::Text(user_input.clone()),
        });

        let mut iterations = 0;
        let mut should_continue = true;
        const ITERATION_WARNING_THRESHOLD: usize = MAX_ITERATIONS - 10;

        while should_continue && iterations < MAX_ITERATIONS {
            iterations += 1;

            // Check for pending inputs BEFORE building request
            // This ensures appended messages are sent in this iteration's API call
            self.drain_pending_inputs();
            if self.has_pending_inputs() {
                let pending = self.take_pending_inputs();
                let count = pending.len();
                let merged = pending.join("\n\n---\n\n");
                log::info!("Adding {} pending input messages to request", count);

                // Send queue processed event to TUI with messages content
                self.emit(AgentEvent::queue_processed(count, pending.clone()))?;

                self.messages.push(Message {
                    role: Role::User,
                    content: MessageContent::Text(merged),
                });
            }

            if let Some(token) = &self.cancel_token
                && token.is_cancelled()
            {
                self.emit(AgentEvent::error(
                    prompt::MSG_OPERATION_CANCELLED.to_string(),
                    None,
                    None,
                ))?;
                break;
            }

            // Warn when approaching iteration limit (UI only, not in messages history)
            if iterations == ITERATION_WARNING_THRESHOLD {
                self.emit(AgentEvent::progress(
                    prompt::MSG_ITERATION_WARNING_UI
                        .replace("{iterations}", &iterations.to_string())
                        .replace("{max_iterations}", &MAX_ITERATIONS.to_string()),
                    None,
                ))?;
            }

            // Proactive compression: check context size BEFORE API call
            // For long conversations, compress early to avoid timeout issues
            let context_size = self.effective_context_size();
            let estimated_tokens = estimate_total_tokens(&self.messages);

            if should_compress(estimated_tokens, context_size, &self.compression_config) {
                self.emit(AgentEvent::progress("⚠️ 上下文过大，正在预压缩...", None))?;

                match compress_messages(
                    &self.messages,
                    CompressionStrategy::SlidingWindow,
                    &self.compression_config,
                ) {
                    Ok(compressed) => {
                        let compressed_tokens = estimate_total_tokens(&compressed);
                        self.messages = compressed;
                        crate::debug::debug_log().compression(
                            estimated_tokens,
                            compressed_tokens,
                            compressed_tokens as f32 / estimated_tokens as f32,
                        );
                    }
                    Err(e) => {
                        self.emit(AgentEvent::progress(format!("预压缩失败: {}", e), None))?;
                    }
                }
            }

            // Build request with current messages (including any pending inputs)
            let tool_defs: Vec<ToolDefinition> = {
                let mut defs: Vec<ToolDefinition> = self
                    .tools
                    .iter()
                    .map(|t| {
                        let def = t.definition();
                        let description = def.description_for_llm();
                        ToolDefinition {
                            name: def.name,
                            description,
                            parameters: def.parameters,
                            is_priority: def.is_priority,
                        }
                    })
                    .collect();
                // 添加代理工具定义
                defs.extend(self.proxy_tool_defs.iter().map(|t| {
                    let def = &t.definition;
                    let description = def.description_for_llm();
                    ToolDefinition {
                        name: def.name.clone(),
                        description,
                        parameters: def.parameters.clone(),
                        is_priority: def.is_priority,
                    }
                }));
                defs
            };
            let request = ChatRequest {
                system: Some(self.system_prompt.clone()),
                messages: self.messages.clone(),
                max_tokens: self.max_tokens,
                tools: tool_defs,
                think: self.think,
                enable_caching: true,
                server_tools: Vec::new(),
            };

            let response = self.call_streaming(&request).await?;

            self.track_usage(&response.usage);

            crate::debug::debug_log().api_call(
                &self.model_name,
                response.usage.input_tokens,
                response.usage.cache_read_input_tokens > 0,
            );

            should_continue = self.process_response(&response).await?;

            // If model wants to stop, check for pending inputs first (higher priority than todos)
            // This ensures appended messages are processed before session ends
            if !should_continue && iterations < MAX_ITERATIONS - 1 {
                // Final drain of pending inputs before checking todos
                self.drain_pending_inputs();

                if self.has_pending_inputs() {
                    log::info!("Agent: found pending inputs at session end, continuing loop");
                    should_continue = true;
                    continue; // Will be processed at start of next iteration
                }

                // Then check for pending todos
                // First check if we just sent a reminder (prevent immediate duplicate)
                if self.last_message_was_todo_reminder() {
                    log::info!("Skipping todo check: reminder already sent in recent messages");
                } else {
                    const MAX_TODO_REMINDERS: usize = 2;
                    
                    // Clone todo_reminder_count to avoid borrow conflict
                    let reminder_count_clone = self.todo_reminder_count.clone();
                    let (pending, all_at_limit) = self.get_pending_todos_with_limit(
                        &reminder_count_clone,
                        MAX_TODO_REMINDERS
                    );
                    
                    if !pending.is_empty() {
                        // Update reminder counts for todos we're about to remind about
                        for (_, content) in &pending {
                            *self.todo_reminder_count.entry(content.clone()).or_insert(0) += 1;
                        }
                        
                        let pending_list = pending
                            .iter()
                            .map(|(status, content)| {
                                let marker = match status.as_str() {
                                    "in_progress" => "[~]",
                                    "pending" => "[ ]",
                                    _ => "[?]",
                                };
                                format!("  {} {}", marker, content)
                            })
                            .collect::<Vec<_>>()
                            .join("\n");

                        let reminder = format!(
                            "📋 任务尚未完成。以下待办项需要处理：\n{}\n\n请继续执行，或在 todo_write 中标记为 completed。如遇阻塞请说明原因。",
                            pending_list
                        );

                        self.messages.push(Message {
                            role: Role::User,
                            content: MessageContent::Text(reminder),
                        });
                        should_continue = true;
                    } else if all_at_limit && !self.todo_reminder_count.is_empty() {
                        // All todos have reached reminder limit, allow session to end
                        // but inform user that todos remain incomplete
                        let remaining_count = self.todo_reminder_count.len();
                        self.emit(AgentEvent::progress(
                            format!(
                                "⚠️ 会话结束：{} 个待办项未完成（已提醒 {} 次，达到上限）",
                                remaining_count, MAX_TODO_REMINDERS
                            ),
                            None,
                        ))?;
                        log::warn!(
                            "Session ending with {} incomplete todos (reminder limit reached)",
                            remaining_count
                        );
                    }
                }
            }

            let context_size = self.effective_context_size();
            let api_tokens = self.last_input_tokens.load(Ordering::Relaxed) as u32;
            let estimated_tokens = estimate_total_tokens(&self.messages);

            let current_tokens = if api_tokens > 0 && api_tokens >= estimated_tokens / 2 {
                api_tokens
            } else {
                estimated_tokens
            };

            // Only log compression check when context is getting full (> 30%)
            // This avoids cluttering debug panel with meaningless checks
            if let Some(ctx_size) = context_size {
                // Send context size to TUI for accurate display
                self.emit(AgentEvent::with_data(
                    EventType::ContextSize,
                    EventData::ContextSize {
                        context_size: ctx_size as u64,
                    },
                ))?;

                let usage_ratio = current_tokens as f64 / ctx_size as f64;
                if usage_ratio >= 0.3 {
                    crate::debug::debug_log().log(
                        "checkcompress",
                        &format!(
                            "usage={:.1}%, tokens={}, context={}, threshold={}%",
                            usage_ratio * 100.0,
                            current_tokens,
                            ctx_size,
                            self.compression_config.threshold * 100.0
                        ),
                    );
                }
            }

            if should_compress(current_tokens, context_size, &self.compression_config) {
                self.emit(AgentEvent::progress(prompt::MSG_COMPRESSING_CONTEXT, None))?;

                let original_tokens = current_tokens;

                match compress_messages(
                    &self.messages,
                    CompressionStrategy::SlidingWindow,
                    &self.compression_config,
                ) {
                    Ok(compressed) => {
                        let compressed_tokens = estimate_total_tokens(&compressed);
                        self.messages = compressed;
                        self.total_input_tokens
                            .store(compressed_tokens as u64, Ordering::Relaxed);
                        self.last_input_tokens
                            .store(compressed_tokens as u64, Ordering::Relaxed);

                        let ratio = compressed_tokens as f32 / original_tokens as f32;
                        crate::debug::debug_log().compression(
                            original_tokens,
                            compressed_tokens,
                            ratio,
                        );

                        self.emit(AgentEvent::with_data(
                            EventType::CompressionCompleted,
                            EventData::Compression {
                                original_tokens: original_tokens as u64,
                                compressed_tokens: compressed_tokens as u64,
                                ratio: compressed_tokens as f32 / original_tokens as f32,
                            },
                        ))?;
                    }
                    Err(e) => {
                        self.emit(AgentEvent::progress(
                            format!("{}{}", prompt::MSG_COMPRESSION_FAILED, e),
                            None,
                        ))?;
                    }
                }
            }
        }

        // Check if we stopped due to reaching MAX_ITERATIONS
        if iterations >= MAX_ITERATIONS && should_continue {
            self.emit(AgentEvent::error(
                prompt::MSG_MAX_ITERATIONS_REACHED
                    .replace("{max_iterations}", &MAX_ITERATIONS.to_string())
                    .replace("{iterations}", &iterations.to_string()),
                Some("MAX_ITERATIONS_REACHED".to_string()),
                Some("agent/run.rs".to_string()),
            ))?;
        }

        self.emit(AgentEvent::usage_with_cache(
            self.total_input_tokens.load(Ordering::Relaxed),
            self.total_output_tokens.load(Ordering::Relaxed),
            0,
            0,
        ))?;

        self.emit(AgentEvent::session_ended())?;

        Ok(Vec::new())
    }

    /// Restore message history (for session continue/resume)
    pub fn set_messages(&mut self, messages: Vec<Message>) {
        self.messages = messages;
    }

    /// Get current messages (for session saving)
    pub fn get_messages(&self) -> &[Message] {
        &self.messages
    }

    /// Get available tools
    pub fn get_tools(&self) -> &[Arc<dyn Tool>] {
        &self.tools
    }

    /// Get system prompt
    pub fn get_system_prompt(&self) -> &str {
        &self.system_prompt
    }

    /// Get current token counts
    pub fn get_token_counts(&self) -> (u64, u64) {
        (
            self.total_input_tokens.load(Ordering::Relaxed),
            self.total_output_tokens.load(Ordering::Relaxed),
        )
    }

    /// Clear message history
    pub fn clear_history(&mut self) {
        self.messages.clear();
        self.total_input_tokens.store(0, Ordering::Relaxed);
        self.total_output_tokens.store(0, Ordering::Relaxed);
        self.last_input_tokens.store(0, Ordering::Relaxed);
    }

    /// Get message count
    pub fn message_count(&self) -> usize {
        self.messages.len()
    }

    // ========================================================================
    // MCP Runtime Management
    // ========================================================================

    /// 动态添加 MCP 服务器
    ///
    /// # Example
    /// ```ignore
    /// use matrixcode_core::mcp::McpServerConfig;
    ///
    /// let config = McpServerConfig::stdio("npx", vec!["-y", "@playwright/mcp@latest"]);
    /// agent.add_mcp_server("playwright", config).await?;
    /// ```
    pub async fn add_mcp_server(
        &mut self,
        name: &str,
        config: crate::mcp::McpServerConfig,
    ) -> Result<()> {
        if let Some(registry) = &self.mcp_registry {
            let mut reg = registry.write().await;
            reg.add_server(name.to_string(), config);
            log::info!("MCP server '{}' added to registry", name);
        } else {
            log::warn!("MCP registry not initialized, cannot add server '{}'", name);
        }
        Ok(())
    }

    /// 移除 MCP 服务器
    pub async fn remove_mcp_server(&mut self, name: &str) -> Result<()> {
        if let Some(registry) = &self.mcp_registry {
            let mut reg = registry.write().await;
            reg.remove_server(name).await?;
            log::info!("MCP server '{}' removed from registry", name);
        }
        Ok(())
    }

    /// 获取 MCP 服务器状态列表
    pub async fn mcp_server_status(&self) -> Vec<crate::mcp::ServerStatus> {
        if let Some(registry) = &self.mcp_registry {
            let reg = registry.read().await;
            reg.server_status().await.values().cloned().collect()
        } else {
            Vec::new()
        }
    }

    /// 启动指定的 MCP 服务器
    pub async fn start_mcp_server(
        &self,
        name: &str,
    ) -> Result<Vec<Arc<crate::mcp::McpToolWrapper>>> {
        if let Some(registry) = &self.mcp_registry {
            let reg = registry.read().await;
            if let Some(placeholder) = reg.get_server(name) {
                let tools = placeholder.start().await?;
                log::info!("MCP server '{}' started with {} tools", name, tools.len());
                Ok(tools)
            } else {
                Err(anyhow::anyhow!(
                    "MCP server '{}' not found in registry",
                    name
                ))
            }
        } else {
            Err(anyhow::anyhow!("MCP registry not initialized"))
        }
    }
}