limit_cli/

agent_bridge.rs

use crate::error::CliError;
use crate::system_prompt::SYSTEM_PROMPT;
use crate::tools::{
    AstGrepTool, BashTool, FileEditTool, FileReadTool, FileWriteTool, GitAddTool, GitCloneTool,
    GitCommitTool, GitDiffTool, GitLogTool, GitPullTool, GitPushTool, GitStatusTool, GrepTool,
    LspTool,
};
use futures::StreamExt;
use limit_agent::executor::{ToolCall, ToolExecutor};
use limit_agent::registry::ToolRegistry;
use limit_llm::providers::LlmProvider;
use limit_llm::types::{Message, Role, Tool as LlmTool, ToolCall as LlmToolCall};
use limit_llm::ProviderFactory;
use limit_llm::ProviderResponseChunk;
use limit_llm::TrackingDb;
use serde_json::json;
use tokio::sync::mpsc;
use tracing::{debug, instrument};

/// Event types for streaming from agent to REPL
#[derive(Debug, Clone)]
#[allow(dead_code)]
pub enum AgentEvent {
    Thinking,
    ToolStart {
        name: String,
        args: serde_json::Value,
    },
    ToolComplete {
        name: String,
        result: String,
    },
    ContentChunk(String),
    Done,
    Error(String),
    TokenUsage {
        input_tokens: u64,
        output_tokens: u64,
    },
}

/// Bridge connecting limit-cli REPL to limit-agent executor and limit-llm client
pub struct AgentBridge {
    /// LLM client for communicating with LLM providers
    llm_client: Box<dyn LlmProvider>,
    /// Tool executor for running tool calls
    executor: ToolExecutor,
    /// List of registered tool names
    tool_names: Vec<&'static str>,
    /// Configuration loaded from ~/.limit/config.toml
    config: limit_llm::Config,
    /// Event sender for streaming events to REPL
    event_tx: Option<mpsc::UnboundedSender<AgentEvent>>,
    /// Token usage tracking database
    tracking_db: TrackingDb,
}

impl AgentBridge {
    /// Create a new AgentBridge with the given configuration
    ///
    /// # Arguments
    /// * `config` - LLM configuration (API key, model, etc.)
    ///
    /// # Returns
    /// A new AgentBridge instance or an error if initialization fails
    pub fn new(config: limit_llm::Config) -> Result<Self, CliError> {
        let llm_client = ProviderFactory::create_provider(&config)
            .map_err(|e| CliError::ConfigError(e.to_string()))?;

        let mut tool_registry = ToolRegistry::new();
        Self::register_tools(&mut tool_registry);

        // Create executor (which takes ownership of registry as Arc)
        let executor = ToolExecutor::new(tool_registry);

        // Generate tool definitions before giving ownership to executor
        let tool_names = vec![
            "file_read",
            "file_write",
            "file_edit",
            "bash",
            "git_status",
            "git_diff",
            "git_log",
            "git_add",
            "git_commit",
            "git_push",
            "git_pull",
            "git_clone",
            "grep",
            "ast_grep",
            "lsp",
        ];

        Ok(Self {
            llm_client,
            executor,
            tool_names,
            config,
            event_tx: None,
            tracking_db: TrackingDb::new().map_err(|e| CliError::ConfigError(e.to_string()))?,
        })
    }

    /// Set the event channel sender for streaming events
    pub fn set_event_tx(&mut self, tx: mpsc::UnboundedSender<AgentEvent>) {
        self.event_tx = Some(tx);
    }

    /// Register all CLI tools into the tool registry
    fn register_tools(registry: &mut ToolRegistry) {
        // File tools
        registry
            .register(FileReadTool::new())
            .expect("Failed to register file_read");
        registry
            .register(FileWriteTool::new())
            .expect("Failed to register file_write");
        registry
            .register(FileEditTool::new())
            .expect("Failed to register file_edit");

        // Bash tool
        registry
            .register(BashTool::new())
            .expect("Failed to register bash");

        // Git tools
        registry
            .register(GitStatusTool::new())
            .expect("Failed to register git_status");
        registry
            .register(GitDiffTool::new())
            .expect("Failed to register git_diff");
        registry
            .register(GitLogTool::new())
            .expect("Failed to register git_log");
        registry
            .register(GitAddTool::new())
            .expect("Failed to register git_add");
        registry
            .register(GitCommitTool::new())
            .expect("Failed to register git_commit");
        registry
            .register(GitPushTool::new())
            .expect("Failed to register git_push");
        registry
            .register(GitPullTool::new())
            .expect("Failed to register git_pull");
        registry
            .register(GitCloneTool::new())
            .expect("Failed to register git_clone");

        // Analysis tools
        registry
            .register(GrepTool::new())
            .expect("Failed to register grep");
        registry
            .register(AstGrepTool::new())
            .expect("Failed to register ast_grep");
        registry
            .register(LspTool::new())
            .expect("Failed to register lsp");
    }

    /// Process a user message through the LLM and execute any tool calls
    ///
    /// # Arguments
    /// * `user_input` - The user's message to process
    /// * `messages` - The conversation history (will be updated in place)
    ///
    /// # Returns
    /// The final response from the LLM or an error
    #[instrument(skip(self, _messages))]
    pub async fn process_message(
        &mut self,
        user_input: &str,
        _messages: &mut Vec<Message>,
    ) -> Result<String, CliError> {
        // Add system message if this is the first message in the conversation
        // Note: Some providers (z.ai) don't support system role, but OpenAI-compatible APIs generally do
        if _messages.is_empty() {
            let system_message = Message {
                role: Role::System,
                content: Some(SYSTEM_PROMPT.to_string()),
                tool_calls: None,
                tool_call_id: None,
            };
            _messages.push(system_message);
        }

        // Add user message to history
        let user_message = Message {
            role: Role::User,
            content: Some(user_input.to_string()),
            tool_calls: None,
            tool_call_id: None,
        };
        _messages.push(user_message);

        // Get tool definitions
        let tool_definitions = self.get_tool_definitions();

        // Main processing loop
        let mut full_response = String::new();
        let mut tool_calls: Vec<LlmToolCall> = Vec::new();
        let max_iterations = 30; // Allow enough iterations for complex tasks
        let mut iteration = 0;

        while iteration < max_iterations {
            iteration += 1;
            debug!("Agent loop iteration {}", iteration);

            // Send thinking event
            self.send_event(AgentEvent::Thinking);

            // Track timing for token usage
            let request_start = std::time::Instant::now();

            // Call LLM
            let mut stream = self
                .llm_client
                .send(_messages.clone(), tool_definitions.clone())
                .await
                .map_err(|e| CliError::ConfigError(e.to_string()))?;

            tool_calls.clear();
            let mut current_content = String::new();
            // Track tool calls: (id) -> (name, args)
            let mut accumulated_calls: std::collections::HashMap<
                String,
                (String, serde_json::Value),
            > = std::collections::HashMap::new();

            // Process stream chunks
            while let Some(chunk_result) = stream.next().await {
                match chunk_result {
                    Ok(ProviderResponseChunk::ContentDelta(text)) => {
                        current_content.push_str(&text);
                        self.send_event(AgentEvent::ContentChunk(text));
                    }
                    Ok(ProviderResponseChunk::ReasoningDelta(_)) => {
                        // Ignore reasoning chunks for now
                    }
                    Ok(ProviderResponseChunk::ToolCallDelta {
                        id,
                        name,
                        arguments,
                    }) => {
                        debug!("ToolCallDelta: id={}, name={}", id, name);
                        // Store/merge tool call arguments
                        accumulated_calls.insert(id.clone(), (name.clone(), arguments.clone()));
                    }
                    Ok(ProviderResponseChunk::Done(usage)) => {
                        // Track token usage
                        let duration_ms = request_start.elapsed().as_millis() as u64;
                        let cost =
                            calculate_cost(self.model(), usage.input_tokens, usage.output_tokens);
                        let _ = self.tracking_db.track_request(
                            self.model(),
                            usage.input_tokens,
                            usage.output_tokens,
                            cost,
                            duration_ms,
                        );
                        // Emit token usage event for TUI display
                        self.send_event(AgentEvent::TokenUsage {
                            input_tokens: usage.input_tokens,
                            output_tokens: usage.output_tokens,
                        });
                        break;
                    }
                    Err(e) => {
                        let error_msg = format!("LLM error: {}", e);
                        self.send_event(AgentEvent::Error(error_msg.clone()));
                        return Err(CliError::ConfigError(error_msg));
                    }
                }
            }

            // Convert accumulated calls to Vec<ToolCall>
            tool_calls = accumulated_calls
                .into_iter()
                .map(|(id, (name, args))| LlmToolCall {
                    id,
                    tool_type: "function".to_string(),
                    function: limit_llm::types::FunctionCall {
                        name,
                        arguments: args.to_string(),
                    },
                })
                .collect();
            full_response.push_str(&current_content);

            debug!(
                "After iter {}: content.len()={}, tool_calls={}, response.len()={}",
                iteration,
                current_content.len(),
                tool_calls.len(),
                full_response.len()
            );

            // If no tool calls, we're done
            if tool_calls.is_empty() {
                break;
            }

            // Execute tool calls - add assistant message with tool_calls
            // Note: Per OpenAI API spec, when tool_calls are present, content should be null
            let assistant_message = Message {
                role: Role::Assistant,
                content: None, // Don't include content when tool_calls are present
                tool_calls: Some(tool_calls.clone()),
                tool_call_id: None,
            };
            _messages.push(assistant_message);

            // Convert LLM tool calls to executor tool calls
            let executor_calls: Vec<ToolCall> = tool_calls
                .iter()
                .map(|tc| {
                    let args: serde_json::Value =
                        serde_json::from_str(&tc.function.arguments).unwrap_or_default();
                    ToolCall::new(&tc.id, &tc.function.name, args)
                })
                .collect();

            // Send ToolStart event for each tool BEFORE execution
            for tc in &tool_calls {
                let args: serde_json::Value =
                    serde_json::from_str(&tc.function.arguments).unwrap_or_default();
                self.send_event(AgentEvent::ToolStart {
                    name: tc.function.name.clone(),
                    args,
                });
            }
            // Execute tools
            let results = self.executor.execute_tools(executor_calls).await;

            // Add tool results to messages (OpenAI format: role=tool, tool_call_id, content)
            for result in results {
                let tool_call = tool_calls.iter().find(|tc| tc.id == result.call_id);
                if let Some(tool_call) = tool_call {
                    let output_json = match &result.output {
                        Ok(value) => {
                            serde_json::to_string(value).unwrap_or_else(|_| "{}".to_string())
                        }
                        Err(e) => json!({ "error": e.to_string() }).to_string(),
                    };

                    self.send_event(AgentEvent::ToolComplete {
                        name: tool_call.function.name.clone(),
                        result: output_json.clone(),
                    });

                    // OpenAI tool result format
                    let tool_result_message = Message {
                        role: Role::Tool,
                        content: Some(output_json),
                        tool_calls: None,
                        tool_call_id: Some(result.call_id),
                    };
                    _messages.push(tool_result_message);
                }
            }
        }

        // If we hit max iterations, make one final request to get a response (no tools = forced text)
        if iteration >= max_iterations && !_messages.is_empty() {
            debug!("Making final LLM call after hitting max iterations (forcing text response)");

            // Add constraint message to force text response
            let constraint_message = Message {
                role: Role::User,
                content: Some(
                    "We've reached the iteration limit. Please provide a summary of:\n\
                    1. What you've completed so far\n\
                    2. What remains to be done\n\
                    3. Recommended next steps for the user to continue"
                        .to_string(),
                ),
                tool_calls: None,
                tool_call_id: None,
            };
            _messages.push(constraint_message);

            // Send with NO tools to force text response
            let no_tools: Vec<LlmTool> = vec![];
            let mut stream = self
                .llm_client
                .send(_messages.clone(), no_tools)
                .await
                .map_err(|e| CliError::ConfigError(e.to_string()))?;

            while let Some(chunk_result) = stream.next().await {
                match chunk_result {
                    Ok(ProviderResponseChunk::ContentDelta(text)) => {
                        full_response.push_str(&text);
                        self.send_event(AgentEvent::ContentChunk(text));
                    }
                    Ok(ProviderResponseChunk::Done(_)) => {
                        break;
                    }
                    Err(e) => {
                        debug!("Error in final LLM call: {}", e);
                        break;
                    }
                    _ => {}
                }
            }
        }

        self.send_event(AgentEvent::Done);
        Ok(full_response)
    }

    /// Get tool definitions formatted for the LLM
    pub fn get_tool_definitions(&self) -> Vec<LlmTool> {
        self.tool_names
            .iter()
            .map(|name| {
                let (description, parameters) = Self::get_tool_schema(name);
                LlmTool {
                    tool_type: "function".to_string(),
                    function: limit_llm::types::ToolFunction {
                        name: name.to_string(),
                        description,
                        parameters,
                    },
                }
            })
            .collect()
    }

    /// Get the schema (description and parameters) for a tool
    fn get_tool_schema(name: &str) -> (String, serde_json::Value) {
        match name {
            "file_read" => (
                "Read the contents of a file".to_string(),
                json!({
                    "type": "object",
                    "properties": {
                        "path": {
                            "type": "string",
                            "description": "Path to the file to read"
                        }
                    },
                    "required": ["path"]
                }),
            ),
            "file_write" => (
                "Write content to a file, creating parent directories if needed".to_string(),
                json!({
                    "type": "object",
                    "properties": {
                        "path": {
                            "type": "string",
                            "description": "Path to the file to write"
                        },
                        "content": {
                            "type": "string",
                            "description": "Content to write to the file"
                        }
                    },
                    "required": ["path", "content"]
                }),
            ),
            "file_edit" => (
                "Replace text in a file with new text".to_string(),
                json!({
                    "type": "object",
                    "properties": {
                        "path": {
                            "type": "string",
                            "description": "Path to the file to edit"
                        },
                        "old_text": {
                            "type": "string",
                            "description": "Text to find and replace"
                        },
                        "new_text": {
                            "type": "string",
                            "description": "New text to replace with"
                        }
                    },
                    "required": ["path", "old_text", "new_text"]
                }),
            ),
            "bash" => (
                "Execute a bash command in a shell".to_string(),
                json!({
                    "type": "object",
                    "properties": {
                        "command": {
                            "type": "string",
                            "description": "Bash command to execute"
                        },
                        "workdir": {
                            "type": "string",
                            "description": "Working directory (default: current directory)"
                        },
                        "timeout": {
                            "type": "integer",
                            "description": "Timeout in seconds (default: 60)"
                        }
                    },
                    "required": ["command"]
                }),
            ),
            "git_status" => (
                "Get git repository status".to_string(),
                json!({
                    "type": "object",
                    "properties": {},
                    "required": []
                }),
            ),
            "git_diff" => (
                "Get git diff".to_string(),
                json!({
                    "type": "object",
                    "properties": {},
                    "required": []
                }),
            ),
            "git_log" => (
                "Get git commit log".to_string(),
                json!({
                    "type": "object",
                    "properties": {
                        "count": {
                            "type": "integer",
                            "description": "Number of commits to show (default: 10)"
                        }
                    },
                    "required": []
                }),
            ),
            "git_add" => (
                "Add files to git staging area".to_string(),
                json!({
                    "type": "object",
                    "properties": {
                        "files": {
                            "type": "array",
                            "items": {"type": "string"},
                            "description": "List of file paths to add"
                        }
                    },
                    "required": ["files"]
                }),
            ),
            "git_commit" => (
                "Create a git commit".to_string(),
                json!({
                    "type": "object",
                    "properties": {
                        "message": {
                            "type": "string",
                            "description": "Commit message"
                        }
                    },
                    "required": ["message"]
                }),
            ),
            "git_push" => (
                "Push commits to remote repository".to_string(),
                json!({
                    "type": "object",
                    "properties": {
                        "remote": {
                            "type": "string",
                            "description": "Remote name (default: origin)"
                        },
                        "branch": {
                            "type": "string",
                            "description": "Branch name (default: current branch)"
                        }
                    },
                    "required": []
                }),
            ),
            "git_pull" => (
                "Pull changes from remote repository".to_string(),
                json!({
                    "type": "object",
                    "properties": {
                        "remote": {
                            "type": "string",
                            "description": "Remote name (default: origin)"
                        },
                        "branch": {
                            "type": "string",
                            "description": "Branch name (default: current branch)"
                        }
                    },
                    "required": []
                }),
            ),
            "git_clone" => (
                "Clone a git repository".to_string(),
                json!({
                    "type": "object",
                    "properties": {
                        "url": {
                            "type": "string",
                            "description": "Repository URL to clone"
                        },
                        "directory": {
                            "type": "string",
                            "description": "Directory to clone into (optional)"
                        }
                    },
                    "required": ["url"]
                }),
            ),
            "grep" => (
                "Search for text patterns in files using regex".to_string(),
                json!({
                    "type": "object",
                    "properties": {
                        "pattern": {
                            "type": "string",
                            "description": "Regex pattern to search for"
                        },
                        "path": {
                            "type": "string",
                            "description": "Path to search in (default: current directory)"
                        }
                    },
                    "required": ["pattern"]
                }),
            ),
            "ast_grep" => (
                "Search code using AST patterns (structural code matching)".to_string(),
                json!({
                    "type": "object",
                    "properties": {
                        "pattern": {
                            "type": "string",
                            "description": "AST pattern to match"
                        },
                        "language": {
                            "type": "string",
                            "description": "Programming language (rust, typescript, python)"
                        },
                        "path": {
                            "type": "string",
                            "description": "Path to search in (default: current directory)"
                        }
                    },
                    "required": ["pattern", "language"]
                }),
            ),
            "lsp" => (
                "Perform Language Server Protocol operations (goto_definition, find_references)"
                    .to_string(),
                json!({
                    "type": "object",
                    "properties": {
                        "command": {
                            "type": "string",
                            "description": "LSP command: goto_definition or find_references"
                        },
                        "file_path": {
                            "type": "string",
                            "description": "Path to the file"
                        },
                        "position": {
                            "type": "object",
                            "description": "Position in the file (line, character)",
                            "properties": {
                                "line": {"type": "integer"},
                                "character": {"type": "integer"}
                            },
                            "required": ["line", "character"]
                        }
                    },
                    "required": ["command", "file_path", "position"]
                }),
            ),
            _ => (
                format!("Tool: {}", name),
                json!({
                    "type": "object",
                    "properties": {},
                    "required": []
                }),
            ),
        }
    }

    /// Send an event through the event channel
    fn send_event(&self, event: AgentEvent) {
        if let Some(ref tx) = self.event_tx {
            let _ = tx.send(event);
        }
    }

    /// Check if the bridge is ready to process messages
    #[allow(dead_code)]
    pub fn is_ready(&self) -> bool {
        self.config
            .providers
            .get(&self.config.provider)
            .map(|p| p.api_key_or_env(&self.config.provider).is_some())
            .unwrap_or(false)
    }

    /// Get the current model name
    pub fn model(&self) -> &str {
        self.config
            .providers
            .get(&self.config.provider)
            .map(|p| p.model.as_str())
            .unwrap_or("")
    }

    /// Get the max tokens setting
    pub fn max_tokens(&self) -> u32 {
        self.config
            .providers
            .get(&self.config.provider)
            .map(|p| p.max_tokens)
            .unwrap_or(4096)
    }

    /// Get the timeout setting
    pub fn timeout(&self) -> u64 {
        self.config
            .providers
            .get(&self.config.provider)
            .map(|p| p.timeout)
            .unwrap_or(60)
    }
}
/// Calculate cost based on model pricing (per 1M tokens)
fn calculate_cost(model: &str, input_tokens: u64, output_tokens: u64) -> f64 {
    let (input_price, output_price) = match model {
        // Claude 3.5 Sonnet: $3/1M input, $15/1M output
        "claude-3-5-sonnet-20241022" | "claude-3-5-sonnet" => (3.0, 15.0),
        // GPT-4: $30/1M input, $60/1M output
        "gpt-4" => (30.0, 60.0),
        // GPT-4 Turbo: $10/1M input, $30/1M output
        "gpt-4-turbo" | "gpt-4-turbo-preview" => (10.0, 30.0),
        // Default: no cost tracking
        _ => (0.0, 0.0),
    };
    (input_tokens as f64 * input_price / 1_000_000.0)
        + (output_tokens as f64 * output_price / 1_000_000.0)
}

#[cfg(test)]
mod tests {
    use super::*;
    use limit_llm::{Config as LlmConfig, ProviderConfig};
    use std::collections::HashMap;

    #[tokio::test]
    async fn test_agent_bridge_new() {
        let mut providers = HashMap::new();
        providers.insert(
            "anthropic".to_string(),
            ProviderConfig {
                api_key: Some("test-key".to_string()),
                model: "claude-3-5-sonnet-20241022".to_string(),
                base_url: None,
                max_tokens: 4096,
                timeout: 60,
            },
        );
        let config = LlmConfig {
            provider: "anthropic".to_string(),
            providers,
        };

        let bridge = AgentBridge::new(config).unwrap();
        assert!(bridge.is_ready());
    }

    #[tokio::test]
    async fn test_agent_bridge_new_no_api_key() {
        let mut providers = HashMap::new();
        providers.insert(
            "anthropic".to_string(),
            ProviderConfig {
                api_key: None,
                model: "claude-3-5-sonnet-20241022".to_string(),
                base_url: None,
                max_tokens: 4096,
                timeout: 60,
            },
        );
        let config = LlmConfig {
            provider: "anthropic".to_string(),
            providers,
        };

        let result = AgentBridge::new(config);
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_get_tool_definitions() {
        let mut providers = HashMap::new();
        providers.insert(
            "anthropic".to_string(),
            ProviderConfig {
                api_key: Some("test-key".to_string()),
                model: "claude-3-5-sonnet-20241022".to_string(),
                base_url: None,
                max_tokens: 4096,
                timeout: 60,
            },
        );
        let config = LlmConfig {
            provider: "anthropic".to_string(),
            providers,
        };

        let bridge = AgentBridge::new(config).unwrap();
        let definitions = bridge.get_tool_definitions();

        assert_eq!(definitions.len(), 15);

        // Check file_read tool definition
        let file_read = definitions
            .iter()
            .find(|d| d.function.name == "file_read")
            .unwrap();
        assert_eq!(file_read.tool_type, "function");
        assert_eq!(file_read.function.name, "file_read");
        assert!(file_read.function.description.contains("Read"));

        // Check bash tool definition
        let bash = definitions
            .iter()
            .find(|d| d.function.name == "bash")
            .unwrap();
        assert_eq!(bash.function.name, "bash");
        assert!(bash.function.parameters["required"]
            .as_array()
            .unwrap()
            .contains(&"command".into()));
    }

    #[test]
    fn test_get_tool_schema() {
        let (desc, params) = AgentBridge::get_tool_schema("file_read");
        assert!(desc.contains("Read"));
        assert_eq!(params["properties"]["path"]["type"], "string");
        assert!(params["required"]
            .as_array()
            .unwrap()
            .contains(&"path".into()));

        let (desc, params) = AgentBridge::get_tool_schema("bash");
        assert!(desc.contains("bash"));
        assert_eq!(params["properties"]["command"]["type"], "string");

        let (desc, _params) = AgentBridge::get_tool_schema("unknown_tool");
        assert!(desc.contains("unknown_tool"));
    }

    #[test]
    fn test_is_ready() {
        let mut providers = HashMap::new();
        providers.insert(
            "anthropic".to_string(),
            ProviderConfig {
                api_key: Some("test-key".to_string()),
                model: "claude-3-5-sonnet-20241022".to_string(),
                base_url: None,
                max_tokens: 4096,
                timeout: 60,
            },
        );
        let config_with_key = LlmConfig {
            provider: "anthropic".to_string(),
            providers,
        };

        let bridge = AgentBridge::new(config_with_key).unwrap();
        assert!(bridge.is_ready());
    }
}