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::{ProxyToolExecutor, ProxyToolDef};

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

    /// 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
            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.provider.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
                let pending = self.get_pending_todos();
                if !pending.is_empty() {
                    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;
                }
            }

            let context_size = self.provider.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"))
        }
    }
}